From dcf9a2217dc7510627ced5f1d5e0fa3c35a034a1 Mon Sep 17 00:00:00 2001 From: Jason Woltje Date: Thu, 5 Feb 2026 18:58:35 -0600 Subject: [PATCH] fix(#338): Fix useWebSocket stale closure by using refs for callbacks - Use useRef to store callbacks, preventing stale closures - Remove callback functions from useEffect dependencies - Only workspaceId and token trigger reconnects now - Callback changes update the ref without causing reconnects - Add 5 new tests verifying no reconnect on callback changes Refs #338 Co-Authored-By: Claude Opus 4.5 --- apps/web/src/hooks/useWebSocket.test.tsx | 134 +++++++++++++++++++++++ apps/web/src/hooks/useWebSocket.ts | 106 +++++++++--------- 2 files changed, 190 insertions(+), 50 deletions(-) diff --git a/apps/web/src/hooks/useWebSocket.test.tsx b/apps/web/src/hooks/useWebSocket.test.tsx index 4e0f46a..242e0f9 100644 --- a/apps/web/src/hooks/useWebSocket.test.tsx +++ b/apps/web/src/hooks/useWebSocket.test.tsx @@ -215,6 +215,140 @@ describe("useWebSocket", (): void => { expect(io).toHaveBeenCalledTimes(2); }); + describe("stale closure prevention", (): void => { + it("should NOT disconnect when callback functions change", (): void => { + const onTaskCreated1 = vi.fn(); + const onTaskCreated2 = vi.fn(); + + const { rerender } = renderHook( + ({ onTaskCreated }: { onTaskCreated: (task: { id: string }) => void }) => + useWebSocket("workspace-123", "token", { onTaskCreated }), + { initialProps: { onTaskCreated: onTaskCreated1 } } + ); + + expect(io).toHaveBeenCalledTimes(1); + expect(mockSocket.disconnect).not.toHaveBeenCalled(); + + // Change the callback - this should NOT cause a reconnect + rerender({ onTaskCreated: onTaskCreated2 }); + + // Socket should NOT have been disconnected or reconnected + expect(mockSocket.disconnect).not.toHaveBeenCalled(); + expect(io).toHaveBeenCalledTimes(1); + }); + + it("should use the latest callback after callback change", async (): Promise => { + const onTaskCreated1 = vi.fn(); + const onTaskCreated2 = vi.fn(); + + const { rerender } = renderHook( + ({ onTaskCreated }: { onTaskCreated: (task: { id: string }) => void }) => + useWebSocket("workspace-123", "token", { onTaskCreated }), + { initialProps: { onTaskCreated: onTaskCreated1 } } + ); + + // Emit event with first callback + const task1 = { id: "task-1" }; + act(() => { + eventHandlers["task:created"]?.(task1); + }); + + await waitFor(() => { + expect(onTaskCreated1).toHaveBeenCalledWith(task1); + expect(onTaskCreated2).not.toHaveBeenCalled(); + }); + + // Update to new callback + rerender({ onTaskCreated: onTaskCreated2 }); + + // Emit another event - should use the new callback + const task2 = { id: "task-2" }; + act(() => { + eventHandlers["task:created"]?.(task2); + }); + + await waitFor(() => { + expect(onTaskCreated2).toHaveBeenCalledWith(task2); + // First callback should only have been called once (with task1) + expect(onTaskCreated1).toHaveBeenCalledTimes(1); + }); + }); + + it("should NOT disconnect when multiple callbacks change simultaneously", (): void => { + const callbacks1 = { + onTaskCreated: vi.fn(), + onTaskUpdated: vi.fn(), + onEventCreated: vi.fn(), + }; + const callbacks2 = { + onTaskCreated: vi.fn(), + onTaskUpdated: vi.fn(), + onEventCreated: vi.fn(), + }; + + interface CallbackProps { + onTaskCreated: (task: { id: string }) => void; + onTaskUpdated: (task: { id: string }) => void; + onEventCreated: (event: { id: string }) => void; + } + + const { rerender } = renderHook( + (props: CallbackProps) => useWebSocket("workspace-123", "token", props), + { initialProps: callbacks1 } + ); + + expect(io).toHaveBeenCalledTimes(1); + + // Change all callbacks at once - should NOT cause reconnect + rerender(callbacks2); + + expect(mockSocket.disconnect).not.toHaveBeenCalled(); + expect(io).toHaveBeenCalledTimes(1); + }); + + it("should handle callback being removed without reconnect", (): void => { + const onTaskCreated = vi.fn(); + + interface CallbackProps { + onTaskCreated?: (task: { id: string }) => void; + } + + const { rerender } = renderHook( + (props: CallbackProps) => useWebSocket("workspace-123", "token", props), + { initialProps: { onTaskCreated } as CallbackProps } + ); + + expect(io).toHaveBeenCalledTimes(1); + + // Remove callback - should NOT cause reconnect + rerender({}); + + expect(mockSocket.disconnect).not.toHaveBeenCalled(); + expect(io).toHaveBeenCalledTimes(1); + }); + + it("should handle callback being added without reconnect", (): void => { + const onTaskCreated = vi.fn(); + + interface CallbackProps { + onTaskCreated?: (task: { id: string }) => void; + } + + const { rerender } = renderHook( + (props: CallbackProps) => useWebSocket("workspace-123", "token", props), + { initialProps: {} as CallbackProps } + ); + + expect(io).toHaveBeenCalledTimes(1); + + // Add callback - should NOT cause reconnect + rerender({ onTaskCreated }); + + expect(mockSocket.disconnect).not.toHaveBeenCalled(); + expect(io).toHaveBeenCalledTimes(1); + }); + }); + it("should clean up all event listeners on unmount", (): void => { const { unmount } = renderHook(() => useWebSocket("workspace-123", "token", { diff --git a/apps/web/src/hooks/useWebSocket.ts b/apps/web/src/hooks/useWebSocket.ts index e5cea5f..eff6d39 100644 --- a/apps/web/src/hooks/useWebSocket.ts +++ b/apps/web/src/hooks/useWebSocket.ts @@ -1,4 +1,4 @@ -import { useEffect, useState } from "react"; +import { useEffect, useRef, useState } from "react"; import type { Socket } from "socket.io-client"; import { io } from "socket.io-client"; import { API_BASE_URL } from "@/lib/config"; @@ -77,15 +77,14 @@ export function useWebSocket( const [isConnected, setIsConnected] = useState(false); const [connectionError, setConnectionError] = useState(null); - const { - onTaskCreated, - onTaskUpdated, - onTaskDeleted, - onEventCreated, - onEventUpdated, - onEventDeleted, - onProjectUpdated, - } = callbacks; + // Use refs for callbacks to prevent stale closures and unnecessary reconnects + // This ensures that callback changes don't trigger useEffect re-runs + const callbacksRef = useRef(callbacks); + + // Keep the ref up-to-date with the latest callbacks + useEffect(() => { + callbacksRef.current = callbacks; + }, [callbacks]); useEffect(() => { // Use WebSocket URL from central config @@ -123,32 +122,48 @@ export function useWebSocket( setIsConnected(false); }; + // Wrapper functions that always use the latest callbacks via ref + // This prevents stale closure issues while avoiding reconnects on callback changes + const handleTaskCreated = (task: Task): void => { + callbacksRef.current.onTaskCreated?.(task); + }; + + const handleTaskUpdated = (task: Task): void => { + callbacksRef.current.onTaskUpdated?.(task); + }; + + const handleTaskDeleted = (payload: DeletePayload): void => { + callbacksRef.current.onTaskDeleted?.(payload); + }; + + const handleEventCreated = (event: Event): void => { + callbacksRef.current.onEventCreated?.(event); + }; + + const handleEventUpdated = (event: Event): void => { + callbacksRef.current.onEventUpdated?.(event); + }; + + const handleEventDeleted = (payload: DeletePayload): void => { + callbacksRef.current.onEventDeleted?.(payload); + }; + + const handleProjectUpdated = (project: Project): void => { + callbacksRef.current.onProjectUpdated?.(project); + }; + newSocket.on("connect", handleConnect); newSocket.on("disconnect", handleDisconnect); newSocket.on("connect_error", handleConnectError); - // Real-time event handlers - if (onTaskCreated) { - newSocket.on("task:created", onTaskCreated); - } - if (onTaskUpdated) { - newSocket.on("task:updated", onTaskUpdated); - } - if (onTaskDeleted) { - newSocket.on("task:deleted", onTaskDeleted); - } - if (onEventCreated) { - newSocket.on("event:created", onEventCreated); - } - if (onEventUpdated) { - newSocket.on("event:updated", onEventUpdated); - } - if (onEventDeleted) { - newSocket.on("event:deleted", onEventDeleted); - } - if (onProjectUpdated) { - newSocket.on("project:updated", onProjectUpdated); - } + // Register all event handlers - they'll check the ref for actual callbacks + newSocket.on("task:created", handleTaskCreated); + newSocket.on("task:updated", handleTaskUpdated); + newSocket.on("task:deleted", handleTaskDeleted); + newSocket.on("event:created", handleEventCreated); + newSocket.on("event:updated", handleEventUpdated); + newSocket.on("event:deleted", handleEventDeleted); + newSocket.on("project:updated", handleProjectUpdated); // Cleanup on unmount or dependency change return (): void => { @@ -156,27 +171,18 @@ export function useWebSocket( newSocket.off("disconnect", handleDisconnect); newSocket.off("connect_error", handleConnectError); - if (onTaskCreated) newSocket.off("task:created", onTaskCreated); - if (onTaskUpdated) newSocket.off("task:updated", onTaskUpdated); - if (onTaskDeleted) newSocket.off("task:deleted", onTaskDeleted); - if (onEventCreated) newSocket.off("event:created", onEventCreated); - if (onEventUpdated) newSocket.off("event:updated", onEventUpdated); - if (onEventDeleted) newSocket.off("event:deleted", onEventDeleted); - if (onProjectUpdated) newSocket.off("project:updated", onProjectUpdated); + newSocket.off("task:created", handleTaskCreated); + newSocket.off("task:updated", handleTaskUpdated); + newSocket.off("task:deleted", handleTaskDeleted); + newSocket.off("event:created", handleEventCreated); + newSocket.off("event:updated", handleEventUpdated); + newSocket.off("event:deleted", handleEventDeleted); + newSocket.off("project:updated", handleProjectUpdated); newSocket.disconnect(); }; - }, [ - workspaceId, - token, - onTaskCreated, - onTaskUpdated, - onTaskDeleted, - onEventCreated, - onEventUpdated, - onEventDeleted, - onProjectUpdated, - ]); + // Only stable values in deps - callbacks are accessed via ref + }, [workspaceId, token]); return { isConnected,