From 57b01f88fd555ac0a6c8cccc30fa0e2cf9b3ae1e Mon Sep 17 00:00:00 2001
From: Curry Yang <1019yanglu@gmail.com>
Date: Wed, 25 Dec 2024 16:46:03 +0800
Subject: [PATCH] feat: keep grpc&websocket connection

---
 .../insomnia/src/ui/hooks/use-close-grpc.ts   | 27 +++++++++++++++++++
 .../src/ui/hooks/use-close-websocket.ts       | 27 +++++++++++++++++++
 packages/insomnia/src/ui/routes/debug.tsx     | 12 ++++-----
 3 files changed, 59 insertions(+), 7 deletions(-)
 create mode 100644 packages/insomnia/src/ui/hooks/use-close-grpc.ts
 create mode 100644 packages/insomnia/src/ui/hooks/use-close-websocket.ts

diff --git a/packages/insomnia/src/ui/hooks/use-close-grpc.ts b/packages/insomnia/src/ui/hooks/use-close-grpc.ts
new file mode 100644
index 00000000000..8168147c827
--- /dev/null
+++ b/packages/insomnia/src/ui/hooks/use-close-grpc.ts
@@ -0,0 +1,27 @@
+import { useCallback, useEffect } from 'react';
+
+import { isGrpcRequestId } from '../../models/grpc-request';
+import uiEventBus, { UIEventType } from '../eventBus';
+
+export const useCloseGrpc = () => {
+  const closeGrpcConnection = useCallback((ids: 'all' | string[]) => {
+    if (ids === 'all') {
+      window.main.webSocket.closeAll();
+      return;
+    }
+
+    ids.forEach(id => {
+      if (isGrpcRequestId(id)) {
+        window.main.grpc.cancel(id);
+      }
+    });
+  }, []);
+
+  useEffect(() => {
+    uiEventBus.on(UIEventType.CLOSE_TAB, closeGrpcConnection);
+
+    return () => {
+      uiEventBus.off(UIEventType.CLOSE_TAB, closeGrpcConnection);
+    };
+  }, [closeGrpcConnection]);
+};
diff --git a/packages/insomnia/src/ui/hooks/use-close-websocket.ts b/packages/insomnia/src/ui/hooks/use-close-websocket.ts
new file mode 100644
index 00000000000..0d5816b9e6b
--- /dev/null
+++ b/packages/insomnia/src/ui/hooks/use-close-websocket.ts
@@ -0,0 +1,27 @@
+import { useCallback, useEffect } from 'react';
+
+import { isWebSocketRequestId } from '../../models/websocket-request';
+import uiEventBus, { UIEventType } from '../eventBus';
+
+export const useCloseWebSocket = () => {
+  const closeWebSocketConnection = useCallback((ids: 'all' | string[]) => {
+    if (ids === 'all') {
+      window.main.webSocket.closeAll();
+      return;
+    }
+
+    ids.forEach(id => {
+      if (isWebSocketRequestId(id)) {
+        window.main.webSocket.close({ requestId: id });
+      }
+    });
+  }, []);
+
+  useEffect(() => {
+    uiEventBus.on(UIEventType.CLOSE_TAB, closeWebSocketConnection);
+
+    return () => {
+      uiEventBus.off(UIEventType.CLOSE_TAB, closeWebSocketConnection);
+    };
+  }, [closeWebSocketConnection]);
+};
diff --git a/packages/insomnia/src/ui/routes/debug.tsx b/packages/insomnia/src/ui/routes/debug.tsx
index 42b8160b995..aa3e6abd276 100644
--- a/packages/insomnia/src/ui/routes/debug.tsx
+++ b/packages/insomnia/src/ui/routes/debug.tsx
@@ -105,6 +105,8 @@ import { getMethodShortHand } from '../components/tags/method-tag';
 import { RealtimeResponsePane } from '../components/websockets/realtime-response-pane';
 import { WebSocketRequestPane } from '../components/websockets/websocket-request-pane';
 import { INSOMNIA_TAB_HEIGHT } from '../constant';
+import { useCloseGrpc } from '../hooks/use-close-grpc';
+import { useCloseWebSocket } from '../hooks/use-close-websocket';
 import { useExecutionState } from '../hooks/use-execution-state';
 import { useInsomniaTab } from '../hooks/use-insomnia-tab';
 import { useReadyState } from '../hooks/use-ready-state';
@@ -452,13 +454,9 @@ export const Debug: FC = () => {
       }
     },
   });
-  // Close all websocket connections when the active environment changes
-  useEffect(() => {
-    return () => {
-      // window.main.webSocket.closeAll();
-      // window.main.grpc.closeAll();
-    };
-  }, [activeEnvironment?._id]);
+
+  useCloseWebSocket();
+  useCloseGrpc();
 
   const isRealtimeRequest =
     activeRequest &&