From a7d8f3e90179b674ab0298135512a681b40691fb Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 19 May 2026 12:34:56 +0000 Subject: [PATCH] Update WebSocketProvider to use detectionAlert and BehaviorSubject - getAiAlertStream -> getDetectionAlertStream - websocketUrl.mlAlert -> websocketUrl.detectionAlert - Subject -> BehaviorSubject (late subscribers get last cached value) - Add getLatestFlowData / getLatestSystemHealth for sync cache reads - share() operator for multicast on each stream https://claude.ai/code/session_01Wen51749mAnwBEfh7XmvUw --- .../frontend-refactor/WebSocketProvider.tsx | 109 ++++++++++++------ 1 file changed, 71 insertions(+), 38 deletions(-) diff --git a/.research/frontend-refactor/WebSocketProvider.tsx b/.research/frontend-refactor/WebSocketProvider.tsx index 268ba6d..9d37509 100644 --- a/.research/frontend-refactor/WebSocketProvider.tsx +++ b/.research/frontend-refactor/WebSocketProvider.tsx @@ -1,7 +1,8 @@ -import React, { createContext, useState, useEffect, useRef } from 'react' +import React, { createContext, useState, useEffect, useRef, useCallback } from 'react' import { websocketUrl, urls } from '../config' import { createWebSocket, fetchData } from '../utils/connectionUtils' -import { Subject } from 'rxjs' +import { BehaviorSubject } from 'rxjs' +import { share } from 'rxjs/operators' interface WebSocketContextType { bootTime: number | null; @@ -9,6 +10,8 @@ interface WebSocketContextType { getIPv4FlowStream: (direction: string, flow_direction: string, timeType: string) => any; getIPv6FlowStream: (direction: string, flow_direction: string, timeType: string) => any; getSystemHealthStream: () => any; + getLatestFlowData: (protocol: 'ipv4' | 'ipv6', direction: string, flow_direction: string, timeType: string) => any; + getLatestSystemHealth: () => any; } export const WebsocketContext = createContext({ @@ -17,17 +20,16 @@ export const WebsocketContext = createContext({ getIPv4FlowStream: () => null, getIPv6FlowStream: () => null, getSystemHealthStream: () => null, + getLatestFlowData: () => null, + getLatestSystemHealth: () => null, }) -interface WebSocketProviderProps { - children: React.ReactNode -} - -export const WebSocketProvider: React.FC = ({ children }) => { +export const WebSocketProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => { const wsRefs = useRef<{ [key: string]: WebSocket }>({}) const [bootTime, setBootTime] = useState(null) const retryDelays = useRef<{ [key: string]: number }>({}) - const dataSubjects = useRef<{ [key: string]: Subject }>({}) + const dataSubjects = useRef<{ [key: string]: BehaviorSubject }>({}) + const latestDataCache = useRef<{ [key: string]: any }>({}) const fetchBootTime = async () => { try { @@ -43,14 +45,24 @@ export const WebSocketProvider: React.FC = ({ children } } } - const getWebSocketStream = (url: string) => { + const getOrCreateSubject = (url: string) => { if (!dataSubjects.current[url]) { - dataSubjects.current[url] = new Subject() + dataSubjects.current[url] = new BehaviorSubject(null) } - return dataSubjects.current[url].asObservable() + return dataSubjects.current[url] + } + + const getWebSocketStream = (url: string) => { + const subject = getOrCreateSubject(url) + return subject.asObservable().pipe(share()) } const setupWebSocket = (url: string, retryCount: number = 0) => { + if (wsRefs.current[url] && wsRefs.current[url].readyState === WebSocket.OPEN) { + console.log(`WebSocket already connected: ${url}`) + return wsRefs.current[url] + } + if (wsRefs.current[url]) { wsRefs.current[url].close() } @@ -58,15 +70,15 @@ export const WebSocketProvider: React.FC = ({ children } const ws = createWebSocket( url, (data) => { - if (dataSubjects.current[url]) { - dataSubjects.current[url].next(data) - } + latestDataCache.current[url] = data + getOrCreateSubject(url).next(data) }, (error) => console.error(`WebSocket error for ${url}:`, error) ) ws.onclose = () => { console.warn(`WebSocket for ${url} closed. Retrying...`) + delete wsRefs.current[url] let delay = retryDelays.current[url] || 1000 retryDelays.current[url] = Math.min(delay * 2, 30000) setTimeout(() => setupWebSocket(url, retryCount + 1), delay) @@ -81,25 +93,42 @@ export const WebSocketProvider: React.FC = ({ children } return ws } - const getDetectionAlertStream = () => { - return getWebSocketStream(websocketUrl.detectionAlert) - } + const getLatestFlowData = useCallback(( + protocol: 'ipv4' | 'ipv6', + direction: string, + flow_direction: string, + timeType: string + ) => { + const urlFunc = protocol === 'ipv4' ? websocketUrl.ipv4FlowStats : websocketUrl.ipv6FlowStats + const url = urlFunc(direction, flow_direction, timeType) + return latestDataCache.current[url] || null + }, []) - const getIPv4FlowStream = (direction: string, flow_direction: string, timeType: string) => { + const getLatestSystemHealth = useCallback(() => { + return latestDataCache.current[websocketUrl.systemHealth] || null + }, []) + + const getDetectionAlertStream = useCallback(() => { + return getWebSocketStream(websocketUrl.detectionAlert) + }, []) + + const getIPv4FlowStream = useCallback((direction: string, flow_direction: string, timeType: string) => { const url = websocketUrl.ipv4FlowStats(direction, flow_direction, timeType) return getWebSocketStream(url) - } + }, []) - const getIPv6FlowStream = (direction: string, flow_direction: string, timeType: string) => { + const getIPv6FlowStream = useCallback((direction: string, flow_direction: string, timeType: string) => { const url = websocketUrl.ipv6FlowStats(direction, flow_direction, timeType) return getWebSocketStream(url) - } + }, []) - const getSystemHealthStream = () => { + const getSystemHealthStream = useCallback(() => { return getWebSocketStream(websocketUrl.systemHealth) - } + }, []) + + const initializeWebSockets = useCallback(() => { + console.log("Initializing all WebSocket connections...") - const initializeWebSockets = () => { setupWebSocket(websocketUrl.detectionAlert) setupWebSocket(websocketUrl.systemHealth) @@ -115,41 +144,45 @@ export const WebSocketProvider: React.FC = ({ children } }) }) }) - } - const closeWebSockets = () => { + console.log(`Total WebSocket connections: ${Object.keys(wsRefs.current).length}`) + }, []) + + const closeWebSockets = useCallback(() => { + console.log("Closing all WebSocket connections...") + Object.entries(wsRefs.current).forEach(([url, ws]) => { try { ws.close() - delete wsRefs.current[url] } catch (error) { console.error(`Failed to close WebSocket ${url}:`, error) } }) + wsRefs.current = {} Object.values(dataSubjects.current).forEach(subject => { subject.complete() }) dataSubjects.current = {} - } + latestDataCache.current = {} + }, []) useEffect(() => { fetchBootTime() initializeWebSockets() - return () => closeWebSockets() }, []) - const contextValue: WebSocketContextType = { - bootTime, - getDetectionAlertStream, - getIPv4FlowStream, - getIPv6FlowStream, - getSystemHealthStream, - } - return ( - + {children} )