Refactor WebSocketService authentication and cleanup

Replaces getUserId with headerBuilder for authentication, improves connection logic by validating user ID, and removes redundant comments and unused code for clarity. Enhances error handling and streamlines reconnection and subscription management.

Author: Fr4nc3

src/frontend/src/services/WebSocketService.tsx

@@ -1,4 +1,4 @@
-import { getUserId } from '../api/config';
+import { headerBuilder } from '../api/config';
 
 export interface StreamMessage {
     type: 'plan_update' | 'step_update' | 'agent_message' | 'error' | 'connection_status' | 'plan_approval_request' | 'final_result';
@@ -19,7 +19,6 @@ export interface StreamingPlanUpdate {
     timestamp?: number;
 }
 
-// Add these new interfaces after StreamingPlanUpdate
 export interface PlanApprovalRequestData {
     plan_id: string;
     session_id: string;
@@ -45,49 +44,25 @@ export interface PlanApprovalResponseData {
 
 class WebSocketService {
     private ws: WebSocket | null = null;
-    private processId: string | null = null; // Add this to store the process ID
     private reconnectAttempts = 0;
     private maxReconnectAttempts = 5;
-    private reconnectDelay = 12000; // Changed from 1000ms to 12000ms (12 seconds)
+    private reconnectDelay = 12000;
     private listeners: Map<string, Set<(message: StreamMessage) => void>> = new Map();
     private planSubscriptions: Set<string> = new Set();
-    private reconnectTimer: NodeJS.Timeout | null = null; // Add timer tracking
-    private isConnecting = false; // Add connection state tracking
+    private reconnectTimer: NodeJS.Timeout | null = null;
+    private isConnecting = false;
 
     /**
      * Connect to WebSocket server
      */
     connect(): Promise<void> {
-        // If already connected, return resolved promise
-        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
-            console.log('WebSocket already connected');
-            return Promise.resolve();
-        }
-
-        // If already connecting, wait for that connection
-        if (this.ws && this.ws.readyState === WebSocket.CONNECTING) {
-            console.log('WebSocket connection already in progress');
-            return new Promise((resolve, reject) => {
-                const checkConnection = () => {
-                    if (this.ws?.readyState === WebSocket.OPEN) {
-                        resolve();
-                    } else if (this.ws?.readyState === WebSocket.CLOSED) {
-                        reject(new Error('Connection failed'));
-                    } else {
-                        setTimeout(checkConnection, 100);
-                    }
-                };
-                checkConnection();
-            });
-        }
         return new Promise((resolve, reject) => {
-            // Prevent multiple simultaneous connection attempts
             if (this.isConnecting) {
                 console.log('Connection attempt already in progress');
+                resolve();
                 return;
             }
 
-            // Clear any existing reconnection timer
             if (this.reconnectTimer) {
                 clearTimeout(this.reconnectTimer);
                 this.reconnectTimer = null;
@@ -96,19 +71,27 @@ class WebSocketService {
             try {
                 this.isConnecting = true;
 
-                // Get WebSocket URL from environment or default to localhost
                 const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
                 const wsHost = process.env.REACT_APP_WS_HOST || '127.0.0.1:8000';
-                const processId = crypto.randomUUID(); // Generate unique process ID for this session
+                const processId = crypto.randomUUID();
 
-                // Build WebSocket URL with authentication headers as query parameters
-                const userId = getUserId(); // Import this from config
+                const authHeaders = headerBuilder();
+                const userId = authHeaders['x-ms-client-principal-id'];
+                
+                if (!userId) {
+                    console.error('No user ID available for WebSocket connection');
+                    this.isConnecting = false;
+                    reject(new Error('Authentication required for WebSocket connection'));
+                    return;
+                }
+
+                // Use query parameter for WebSocket authentication (as backend expects)
                 const wsUrl = `${wsProtocol}//${wsHost}/api/v3/socket/${processId}?user_id=${encodeURIComponent(userId)}`;
 
                 console.log('Connecting to WebSocket:', wsUrl);
 
                 this.ws = new WebSocket(wsUrl);
-
+                
                 this.ws.onopen = () => {
                     console.log('WebSocket connected successfully');
                     this.reconnectAttempts = 0;
@@ -132,8 +115,7 @@ class WebSocketService {
                     this.isConnecting = false;
                     this.emit('connection_status', { connected: false });
 
-                    // Only attempt reconnect if it wasn't a manual disconnect
-                    if (event.code !== 1000) { // 1000 = normal closure
+                    if (event.code !== 1000) {
                         this.attemptReconnect();
                     }
                 };
@@ -158,16 +140,15 @@ class WebSocketService {
     disconnect(): void {
         console.log('Manually disconnecting WebSocket');
 
-        // Clear any pending reconnection attempts
         if (this.reconnectTimer) {
             clearTimeout(this.reconnectTimer);
             this.reconnectTimer = null;
         }
 
-        this.reconnectAttempts = this.maxReconnectAttempts; // Prevent reconnection
+        this.reconnectAttempts = this.maxReconnectAttempts;
 
         if (this.ws) {
-            this.ws.close(1000, 'Manual disconnect'); // Use normal closure code
+            this.ws.close(1000, 'Manual disconnect');
             this.ws = null;
         }
         this.planSubscriptions.clear();
@@ -216,7 +197,6 @@ class WebSocketService {
 
         this.listeners.get(eventType)!.add(callback);
 
-        // Return unsubscribe function
         return () => {
             const eventListeners = this.listeners.get(eventType);
             if (eventListeners) {
@@ -269,17 +249,14 @@ class WebSocketService {
     private handleMessage(message: StreamMessage): void {
         console.log('WebSocket message received:', message);
 
-        // Emit to specific event listeners
         if (message.type) {
             this.emit(message.type, message.data);
         }
 
-        // Handle plan approval requests specifically
         if (message.type === 'plan_approval_request') {
             console.log('Plan approval request received via WebSocket:', message.data);
         }
 
-        // Emit to general message listeners
         this.emit('message', message);
     }
 
@@ -293,14 +270,12 @@ class WebSocketService {
             return;
         }
 
-        // Prevent multiple simultaneous reconnection attempts
         if (this.isConnecting || this.reconnectTimer) {
             console.log('Reconnection attempt already in progress');
             return;
         }
 
         this.reconnectAttempts++;
-        // Use exponential backoff: 12s, 24s, 48s, 96s, 192s
         const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1);
 
         console.log(`Scheduling reconnection attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts} in ${delay/1000}s`);
@@ -312,15 +287,12 @@ class WebSocketService {
             this.connect()
                 .then(() => {
                     console.log('Reconnection successful - re-subscribing to plans');
-                    // Re-subscribe to all plans
                     this.planSubscriptions.forEach(planId => {
                         this.subscribeToPlan(planId);
                     });
                 })
                 .catch((error) => {
                     console.error('Reconnection failed:', error);
-                    // The connect() method will trigger another reconnection attempt
-                    // through the onclose handler if needed
                 });
         }, delay);
     }
@@ -367,6 +339,5 @@ class WebSocketService {
     }
 }
 
-// Export singleton instance
 export const webSocketService = new WebSocketService();
 export default webSocketService;