diff --git a/frontend/src/api/websocket.ts b/frontend/src/api/websocket.ts index 3db7cdfa6..a8459b50a 100644 --- a/frontend/src/api/websocket.ts +++ b/frontend/src/api/websocket.ts @@ -190,3 +190,12 @@ export class WebSocketClient { } } } + +let sharedClient: WebSocketClient | null = null; + +export function getSharedWebSocketClient(): WebSocketClient { + if (sharedClient) return sharedClient; + const basePath = (typeof window !== 'undefined' && window.X_UI_BASE_PATH) || ''; + sharedClient = new WebSocketClient(basePath); + return sharedClient; +} diff --git a/frontend/src/api/websocketBridge.ts b/frontend/src/api/websocketBridge.ts index b6d942249..21d8f8f7f 100644 --- a/frontend/src/api/websocketBridge.ts +++ b/frontend/src/api/websocketBridge.ts @@ -1,34 +1,19 @@ import { useEffect } from 'react'; import { useQueryClient } from '@tanstack/react-query'; -import { WebSocketClient } from '@/api/websocket'; +import { getSharedWebSocketClient } from '@/api/websocket'; import { keys } from '@/api/queryKeys'; import { isRecentLocalInvalidate } from '@/api/invalidationTracker'; type Handler = (payload: unknown) => void; -interface SharedClient { - connect(): void; - on(event: string, fn: Handler): void; - off(event: string, fn: Handler): void; -} - -let sharedClient: SharedClient | null = null; - -function getSharedClient(): SharedClient { - if (sharedClient) return sharedClient; - const basePath = (typeof window !== 'undefined' && window.X_UI_BASE_PATH) || ''; - sharedClient = new WebSocketClient(basePath) as SharedClient; - return sharedClient; -} - let invalidateTimer: number | null = null; export function useWebSocketBridge() { const queryClient = useQueryClient(); useEffect(() => { - const client = getSharedClient(); + const client = getSharedWebSocketClient(); const onInvalidate: Handler = (payload) => { const p = payload as { type?: string } | undefined; diff --git a/frontend/src/hooks/useWebSocket.ts b/frontend/src/hooks/useWebSocket.ts index 4a5e036e9..1eccb41d4 100644 --- a/frontend/src/hooks/useWebSocket.ts +++ b/frontend/src/hooks/useWebSocket.ts @@ -1,26 +1,11 @@ import { useEffect } from 'react'; -import { WebSocketClient } from '@/api/websocket'; +import { getSharedWebSocketClient } from '@/api/websocket'; type Handler = (payload: unknown) => void; -interface SharedClient { - connect(): void; - on(event: string, fn: Handler): void; - off(event: string, fn: Handler): void; -} - -let sharedClient: SharedClient | null = null; - -function getSharedClient(): SharedClient { - if (sharedClient) return sharedClient; - const basePath = (typeof window !== 'undefined' && window.X_UI_BASE_PATH) || ''; - sharedClient = new WebSocketClient(basePath) as SharedClient; - return sharedClient; -} - export function useWebSocket(handlers: Record) { useEffect(() => { - const client = getSharedClient(); + const client = getSharedWebSocketClient(); const entries = Object.entries(handlers); for (const [event, fn] of entries) client.on(event, fn); client.connect(); diff --git a/frontend/src/test/websocket-shared-client.test.tsx b/frontend/src/test/websocket-shared-client.test.tsx new file mode 100644 index 000000000..81f9a5497 --- /dev/null +++ b/frontend/src/test/websocket-shared-client.test.tsx @@ -0,0 +1,49 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { renderHook } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { ReactNode } from 'react'; + +import { useWebSocket } from '@/hooks/useWebSocket'; +import { useWebSocketBridge } from '@/api/websocketBridge'; + +describe('shared WebSocket connection', () => { + let socketCount = 0; + let originalWS: typeof WebSocket; + + beforeEach(() => { + socketCount = 0; + originalWS = globalThis.WebSocket; + class FakeWebSocket { + static readonly CONNECTING = 0; + static readonly OPEN = 1; + static readonly CLOSING = 2; + static readonly CLOSED = 3; + readyState = FakeWebSocket.CONNECTING; + constructor() { + socketCount += 1; + } + addEventListener() {} + removeEventListener() {} + close() {} + send() {} + } + globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket; + }); + + afterEach(() => { + globalThis.WebSocket = originalWS; + vi.restoreAllMocks(); + }); + + it('opens a single socket when the bridge and a page hook are both mounted', () => { + const queryClient = new QueryClient(); + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + + renderHook(() => useWebSocketBridge(), { wrapper }); + renderHook(() => useWebSocket({ traffic: () => {} })); + + expect(socketCount).toBe(1); + }); +});