socket reruns

apps/stock/web-app/src/hooks/useWebSocket.ts

@@ -0,0 +1,195 @@
+import { useEffect, useRef, useState, useCallback } from 'react';
+
+export interface WebSocketMessage {
+  type: 'connected' | 'run_update' | 'progress' | 'error' | 'completed' | 'pong';
+  runId?: string;
+  data?: any;
+  timestamp?: string;
+}
+
+interface UseWebSocketOptions {
+  runId: string | null;
+  onMessage?: (message: WebSocketMessage) => void;
+  onProgress?: (progress: number, currentDate?: string) => void;
+  onError?: (error: string) => void;
+  onCompleted?: (results?: any) => void;
+}
+
+export function useWebSocket({
+  runId,
+  onMessage,
+  onProgress,
+  onError,
+  onCompleted
+}: UseWebSocketOptions) {
+  const [isConnected, setIsConnected] = useState(false);
+  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
+  const wsRef = useRef<WebSocket | null>(null);
+  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+  const pingIntervalRef = useRef<NodeJS.Timeout | null>(null);
+  const isIntentionalDisconnect = useRef(false);
+
+  const connect = useCallback(() => {
+    if (!runId) {
+      console.log('useWebSocket: No runId provided, skipping connection');
+      return;
+    }
+    
+    // Check if already connected or connecting
+    if (wsRef.current && (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING)) {
+      console.log('useWebSocket: Already connected or connecting to runId:', runId);
+      return;
+    }
+    
+    // Reset intentional disconnect flag
+    isIntentionalDisconnect.current = false;
+    
+    // Clear any pending reconnect timeout
+    if (reconnectTimeoutRef.current) {
+      clearTimeout(reconnectTimeoutRef.current);
+      reconnectTimeoutRef.current = null;
+    }
+
+    const wsUrl = `ws://localhost:2003/ws?runId=${runId}`;
+    console.log('Connecting to WebSocket:', wsUrl);
+
+    try {
+      const ws = new WebSocket(wsUrl);
+      wsRef.current = ws;
+
+      ws.onopen = () => {
+        console.log('WebSocket connected');
+        setIsConnected(true);
+        
+        // Start ping interval
+        pingIntervalRef.current = setInterval(() => {
+          if (ws.readyState === WebSocket.OPEN) {
+            ws.send(JSON.stringify({ type: 'ping' }));
+          }
+        }, 30000); // Ping every 30 seconds
+      };
+
+      ws.onmessage = (event) => {
+        try {
+          const message: WebSocketMessage = JSON.parse(event.data);
+          console.log('WebSocket message:', message);
+          setLastMessage(message);
+
+          // Call appropriate callbacks
+          if (onMessage) {
+            onMessage(message);
+          }
+
+          switch (message.type) {
+            case 'progress':
+              if (onProgress && message.data) {
+                onProgress(message.data.progress, message.data.currentDate);
+              }
+              break;
+            case 'error':
+              if (onError && message.data?.error) {
+                onError(message.data.error);
+                console.error('Run Error:', message.data.error);
+              }
+              break;
+            case 'completed':
+              if (onCompleted) {
+                onCompleted(message.data?.results);
+              }
+              console.log('Run Completed: The backtest run has completed successfully.');
+              break;
+          }
+        } catch (error) {
+          console.error('Error parsing WebSocket message:', error);
+        }
+      };
+
+      ws.onerror = (error) => {
+        console.error('WebSocket error:', error);
+        setIsConnected(false);
+      };
+
+      ws.onclose = () => {
+        console.log('WebSocket disconnected');
+        setIsConnected(false);
+        wsRef.current = null;
+
+        // Clear ping interval
+        if (pingIntervalRef.current) {
+          clearInterval(pingIntervalRef.current);
+          pingIntervalRef.current = null;
+        }
+
+        // Attempt to reconnect after 3 seconds
+        // Only reconnect if not intentionally disconnected and still the same WebSocket
+        if (runId && wsRef.current === ws && !isIntentionalDisconnect.current) {
+          reconnectTimeoutRef.current = setTimeout(() => {
+            console.log('Attempting to reconnect...');
+            connect();
+          }, 3000);
+        }
+      };
+    } catch (error) {
+      console.error('Error creating WebSocket:', error);
+      setIsConnected(false);
+    }
+  }, [runId, onMessage, onProgress, onError, onCompleted]);
+
+  const disconnect = useCallback(() => {
+    console.log('Disconnecting WebSocket...');
+    
+    // Set flag to prevent automatic reconnection
+    isIntentionalDisconnect.current = true;
+    
+    if (reconnectTimeoutRef.current) {
+      clearTimeout(reconnectTimeoutRef.current);
+      reconnectTimeoutRef.current = null;
+    }
+    
+    if (pingIntervalRef.current) {
+      clearInterval(pingIntervalRef.current);
+      pingIntervalRef.current = null;
+    }
+
+    if (wsRef.current) {
+      // Set a flag to prevent reconnection
+      const ws = wsRef.current;
+      wsRef.current = null;
+      ws.close();
+    }
+    
+    setIsConnected(false);
+  }, []);
+
+  const sendMessage = useCallback((message: any) => {
+    if (wsRef.current?.readyState === WebSocket.OPEN) {
+      wsRef.current.send(JSON.stringify(message));
+    } else {
+      console.warn('WebSocket is not connected');
+    }
+  }, []);
+
+  useEffect(() => {
+    // Small delay to prevent rapid reconnections during React's render cycles
+    const timeoutId = setTimeout(() => {
+      if (runId) {
+        connect();
+      } else {
+        disconnect();
+      }
+    }, 100);
+
+    return () => {
+      clearTimeout(timeoutId);
+      disconnect();
+    };
+  }, [runId]); // Only depend on runId, not the functions
+
+  return {
+    isConnected,
+    lastMessage,
+    sendMessage,
+    disconnect,
+    connect
+  };
+}