fix(react): websocket message event handler (#286)

RBND-studio · Jan 23, 2025 · b6ad4f6 · b6ad4f6
1 parent 01c8e34
commit b6ad4f6
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 26 deletions.
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/workspaces/e2e/package.json b/workspaces/e2e/package.json
@@ -15,7 +15,7 @@
     "@flows/shared": "workspace:*"
   },
   "devDependencies": {
-    "@playwright/test": "^1.49.1",
+    "@playwright/test": "^1.50.0",
     "@types/node": "^20",
     "@types/react": "19.0.7",
     "@types/react-dom": "19.0.3",

diff --git a/workspaces/e2e/tests/websocket.spec.ts b/workspaces/e2e/tests/websocket.spec.ts
@@ -30,6 +30,7 @@ const run = (packageName: string) => {
       route.fulfill({ json: { blocks: [] } });
     });
     await page.goto(`/${packageName}.html`);
+    await expect(page.locator("h1")).toBeVisible();
     await expect(page.getByText("Hello world")).toBeHidden();
     const payload: BlockUpdatesPayload = {
       exitedBlockIds: [],

diff --git a/workspaces/react/src/hooks/use-websocket.ts b/workspaces/react/src/hooks/use-websocket.ts
@@ -16,6 +16,14 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
     onOpenRef.current = onOpen;
   }, [onOpen]);
 
+  const onMessageRef = useRef(onMessage);
+  useEffect(() => {
+    onMessageRef.current = onMessage;
+  }, [onMessage]);
+  const handleMessage = useCallback((event: MessageEvent<unknown>) => {
+    onMessageRef.current(event);
+  }, []);
+
   const connect = useCallback(() => {
     if (cleanupRef.current) {
       cleanupRef.current();
@@ -35,10 +43,12 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
     };
     socket.addEventListener("open", handleOpen);
     socket.addEventListener("close", handleClose);
+    socket.addEventListener("message", handleMessage);
 
     const cleanup = (): void => {
       socket.removeEventListener("open", handleOpen);
       socket.removeEventListener("close", handleClose);
+      socket.removeEventListener("message", handleMessage);
 
       if (socket.readyState === WebSocket.CONNECTING) {
         socket.addEventListener("open", () => {
@@ -51,7 +61,7 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
 
     cleanupRef.current = cleanup;
     return cleanup;
-  }, [url]);
+  }, [handleMessage, url]);
 
   useEffect(() => {
     const cleanup = connect();
@@ -75,15 +85,6 @@ export const useWebsocket = ({ url, onMessage, onOpen }: Props): void => {
     };
   }, [reconnectAttempts, ws, connect]);
 
-  useEffect(() => {
-    if (!ws) return;
-
-    ws.addEventListener("message", onMessage);
-    return () => {
-      ws.removeEventListener("message", onMessage);
-    };
-  }, [onMessage, ws]);
-
   useEffect(() => {
     return () => {
       if (cleanupRef.current) {