From 4f6b0b4838be1066240be2a8b2f7deadf6e66914 Mon Sep 17 00:00:00 2001
From: Matvey Kukuy <Matvey-Kuk@users.noreply.github.com>
Date: Thu, 14 Nov 2024 18:03:37 +0200
Subject: [PATCH] fix: Local pusher & alert feed (#2484)

---
 keep-ui/utils/hooks/usePusher.ts     | 35 ++++++++++++++++++----------
 keep/api/tasks/process_event_task.py | 27 ++++++++++++++-------
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/keep-ui/utils/hooks/usePusher.ts b/keep-ui/utils/hooks/usePusher.ts
index 20be77a03..f1b0c23ed 100644
--- a/keep-ui/utils/hooks/usePusher.ts
+++ b/keep-ui/utils/hooks/usePusher.ts
@@ -1,11 +1,11 @@
-import Pusher from "pusher-js";
+import Pusher, { Options as PusherOptions } from "pusher-js";
 import { useConfig } from "./useConfig";
 import { useSession } from "next-auth/react";
 import { useApiUrl } from "./useConfig";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 
 let PUSHER: Pusher | null = null;
-const POLLING_INTERVAL = 3000;
+const POLLING_INTERVAL = 1000 * 10;  // Once per 10 seconds.
 
 export const useWebsocket = () => {
   const apiUrl = useApiUrl();
@@ -26,15 +26,19 @@ export const useWebsocket = () => {
     channelName = `private-${session?.tenantId}`;
     console.log("useWebsocket: Creating new Pusher instance");
     try {
-      const isRelativeHost =
-        configData.PUSHER_HOST && !configData.PUSHER_HOST.includes("://");
-      console.log("useWebsocket: isRelativeHost:", isRelativeHost);
-      PUSHER = new Pusher(configData.PUSHER_APP_KEY, {
-        wsHost: isRelativeHost
+      const isRelativeHostAndNotLocal = 
+        configData.PUSHER_HOST && 
+        !configData.PUSHER_HOST.includes("://") && 
+        !["localhost", "127.0.0.1"].includes(configData.PUSHER_HOST);
+
+      console.log("useWebsocket: isRelativeHostAndNotLocal:", isRelativeHostAndNotLocal);
+
+      var pusherOptions: PusherOptions = {
+        wsHost: isRelativeHostAndNotLocal
           ? window.location.hostname
           : configData.PUSHER_HOST,
-        wsPath: isRelativeHost ? configData.PUSHER_HOST : "",
-        wsPort: isRelativeHost
+        wsPath: isRelativeHostAndNotLocal ? configData.PUSHER_HOST : "",
+        wsPort: isRelativeHostAndNotLocal
           ? window.location.protocol === "https:"
             ? 443
             : 80
@@ -50,8 +54,10 @@ export const useWebsocket = () => {
             Authorization: `Bearer ${session?.accessToken!}`,
           },
         },
-      });
-      console.log("useWebsocket: Pusher instance created successfully");
+      }
+      PUSHER = new Pusher(configData.PUSHER_APP_KEY, pusherOptions);
+
+      console.log("useWebsocket: Pusher instance created successfully. Options:", pusherOptions);
 
       PUSHER.connection.bind("connected", () => {
         console.log("useWebsocket: Pusher connected successfully");
@@ -61,6 +67,10 @@ export const useWebsocket = () => {
         console.error("useWebsocket: Pusher connection error:", err);
       });
 
+      PUSHER.connection.bind('state_change', function(states:any) {
+        console.log("useWebsocket: Connection state changed from", states.previous, "to", states.current);
+      });
+
       PUSHER.subscribe(channelName)
         .bind("pusher:subscription_succeeded", () => {
           console.log(
@@ -148,13 +158,14 @@ export const useAlertPolling = () => {
       `useAlertPolling: Time since last poll: ${timeSinceLastPoll}ms`
     );
 
+    const newPollValue = Math.floor(Math.random() * 10000);
+
     if (timeSinceLastPoll < POLLING_INTERVAL) {
       console.log("useAlertPolling: Ignoring poll due to short interval");
       setPollAlerts(0);
     } else {
       console.log("useAlertPolling: Updating poll alerts");
       lastPollTimeRef.current = currentTime;
-      const newPollValue = Math.floor(Math.random() * 10000);
       console.log(`useAlertPolling: New poll value: ${newPollValue}`);
       setPollAlerts(newPollValue);
     }
diff --git a/keep/api/tasks/process_event_task.py b/keep/api/tasks/process_event_task.py
index 53a064dcc..a4bec5e8c 100644
--- a/keep/api/tasks/process_event_task.py
+++ b/keep/api/tasks/process_event_task.py
@@ -411,11 +411,23 @@ def __handle_formatted_events(
             },
         )
 
+    
+    pusher_client = get_pusher_client() if notify_client else None
+
     # Tell the client to poll alerts
-    if notify_client and incidents:
-        pusher_client = get_pusher_client()
-        if not pusher_client:
+    if pusher_client:
+        try:
+            pusher_client.trigger(
+                f"private-{tenant_id}",
+                "poll-alerts",
+                "{}",
+            )
+            logger.info("Told client to poll alerts")
+        except Exception:
+            logger.exception("Failed to tell client to poll alerts")
             pass
+    
+    if incidents and pusher_client:
         try:
             pusher_client.trigger(
                 f"private-{tenant_id}",
@@ -423,14 +435,13 @@ def __handle_formatted_events(
                 {},
             )
         except Exception:
-            logger.exception("Failed to push alert to the client")
+            logger.exception("Failed to tell the client to pull incidents")
 
     # Now we need to update the presets
     # send with pusher
-    if notify_client:
-        pusher_client = get_pusher_client()
-        if not pusher_client:
-            return
+    if not pusher_client:
+        return
+    
     try:
         presets = get_all_presets(tenant_id)
         rules_engine = RulesEngine(tenant_id=tenant_id)