Implement socket state management with Zustand and enhance connection handling

- Introduced a new Zustand store for managing socket connection status and socket ID, improving state management for real-time communication.
- Updated the ConnectionIndicator component to utilize the socket store for dynamic status updates, allowing for better integration with the socket connection state.
- Enhanced the SocketService class to update the socket connection status during various connection events, ensuring accurate status reporting.
- Refactored the TelegramLoginButton component to clarify JWT token handling in comments, improving code readability.

These changes enhance the application's real-time capabilities and provide a more robust user experience by accurately reflecting connection states.

src/components/ConnectionIndicator.tsx

@@ -1,3 +1,5 @@
+import { useSocketStore } from '../store/socketStore';
+
 interface ConnectionIndicatorProps {
   status?: 'connected' | 'disconnected' | 'connecting';
   description?: string;
@@ -5,10 +7,13 @@ interface ConnectionIndicatorProps {
 }
 
 const ConnectionIndicator = ({
-  status = 'connected',
+  status: overrideStatus,
   description = 'Your browser is now connected to the AlphaHuman AI. Keep this tab open to keep the connection alive. You can message your assistant with the button below.',
   className = '',
 }: ConnectionIndicatorProps) => {
+  // Use socket store status, but allow override via props
+  const storeStatus = useSocketStore((state) => state.status);
+  const status = overrideStatus || storeStatus;
   const statusConfig = {
     connected: {
       color: 'bg-sage-500',

src/components/TelegramLoginButton.tsx

@@ -76,9 +76,10 @@ const TelegramLoginButton = ({
           popup.close();
 
           try {
+            // This token is a JWT token generated by the backend server
             const jwtToken = event.data.token;
 
-            // Store session token in store
+            // Store JWT token in store (received from backend server)
             setToken(jwtToken);
 
             // Call onSuccess callback if provided, otherwise navigate to onboarding
@@ -110,7 +111,7 @@ const TelegramLoginButton = ({
           // Try to read popup location (will throw if cross-origin)
           const popupUrl = popup.location.href;
 
-          // If popup is on our callback URL, extract data and complete auth
+          // If popup is on our callback URL, extract data and get JWT from backend
           if (popupUrl.startsWith(callbackUrl.split('?')[0])) {
             const url = new URL(popupUrl);
             const telegramData: TelegramAuthData = {
@@ -131,7 +132,8 @@ const TelegramLoginButton = ({
               window.removeEventListener('message', messageHandler);
 
               try {
-                // Call web-complete endpoint to get handoff token
+                // Send Telegram auth data to backend to get JWT token
+                // Use web-complete to verify and then exchange for JWT
                 const webCompleteResponse = await fetch(`${BACKEND_URL}/auth/web-complete`, {
                   method: 'POST',
                   headers: {
@@ -155,7 +157,7 @@ const TelegramLoginButton = ({
                   throw new Error('No login token received from server');
                 }
 
-                // Exchange handoff token for session token
+                // Exchange handoff token for JWT session token
                 const exchangeResponse = await fetch(`${BACKEND_URL}/auth/desktop-exchange`, {
                   method: 'POST',
                   headers: {
@@ -175,10 +177,10 @@ const TelegramLoginButton = ({
                 const { sessionToken, user } = exchangeData.data;
 
                 if (!sessionToken) {
-                  throw new Error('No session token received from server');
+                  throw new Error('No JWT token received from server');
                 }
 
-                // Store session token in store
+                // Store JWT token in store (this is the JWT from the backend)
                 setToken(sessionToken);
                 
                 // Store user data in localStorage for backward compatibility

src/services/socketService.ts

@@ -1,6 +1,7 @@
 import { io, Socket } from 'socket.io-client';
 import { BACKEND_URL } from '../utils/config';
 import { useAuthStore } from '../store/authStore';
+import { useSocketStore } from '../store/socketStore';
 
 class SocketService {
   private socket: Socket | null = null;
@@ -16,6 +17,7 @@ class SocketService {
 
     if (!token) {
       console.warn('[SocketService] No token available, cannot connect');
+      useSocketStore.getState().setStatus('disconnected');
       return;
     }
 
@@ -25,6 +27,7 @@ class SocketService {
     }
 
     console.log('[SocketService] Connecting to socket server...');
+    useSocketStore.getState().setStatus('connecting');
 
     this.socket = io(BACKEND_URL, {
       auth: {
@@ -50,6 +53,8 @@ class SocketService {
       this.socket.disconnect();
       this.socket = null;
       this.reconnectAttempts = 0;
+      useSocketStore.getState().setStatus('disconnected');
+      useSocketStore.getState().setSocketId(null);
     }
   }
 
@@ -125,19 +130,35 @@ class SocketService {
     this.socket.on('connect', () => {
       console.log('[SocketService] Connected to socket server:', this.socket?.id);
       this.reconnectAttempts = 0;
+      useSocketStore.getState().setStatus('connected');
+      useSocketStore.getState().setSocketId(this.socket?.id || null);
     });
 
     this.socket.on('disconnect', (reason) => {
       console.log('[SocketService] Disconnected from socket server:', reason);
+      useSocketStore.getState().setStatus('disconnected');
+      useSocketStore.getState().setSocketId(null);
+      
+      // If disconnected due to error, set status to connecting if we're trying to reconnect
+      if (reason === 'io server disconnect' || reason === 'io client disconnect') {
+        // Manual disconnect, keep as disconnected
+      } else {
+        // Network error or other, might be reconnecting
+        if (this.socket && !this.socket.disconnected) {
+          useSocketStore.getState().setStatus('connecting');
+        }
+      }
     });
 
     this.socket.on('connect_error', (error) => {
       console.error('[SocketService] Connection error:', error.message);
       this.reconnectAttempts++;
+      useSocketStore.getState().setStatus('connecting');
 
       // If max attempts reached, try reconnecting with fresh token
       if (this.reconnectAttempts >= this.maxReconnectAttempts) {
         console.log('[SocketService] Max reconnection attempts reached, reconnecting with fresh token...');
+        useSocketStore.getState().setStatus('disconnected');
         setTimeout(() => {
           this.connect();
         }, this.reconnectDelay * 2);
@@ -146,10 +167,12 @@ class SocketService {
 
     this.socket.on('ready', () => {
       console.log('[SocketService] Server ready');
+      useSocketStore.getState().setStatus('connected');
     });
 
     this.socket.on('error', (error: { message?: string; status?: number; requestId?: string }) => {
       console.error('[SocketService] Server error:', error);
+      // Don't change status on server errors, connection might still be active
     });
   }
 }

src/store/socketStore.ts

@@ -0,0 +1,28 @@
+import { create } from 'zustand';
+
+export type SocketConnectionStatus = 'connected' | 'disconnected' | 'connecting';
+
+interface SocketState {
+  status: SocketConnectionStatus;
+  socketId: string | null;
+  setStatus: (status: SocketConnectionStatus) => void;
+  setSocketId: (socketId: string | null) => void;
+  reset: () => void;
+}
+
+export const useSocketStore = create<SocketState>((set) => ({
+  status: 'disconnected',
+  socketId: null,
+
+  setStatus: (status: SocketConnectionStatus) => {
+    set({ status });
+  },
+
+  setSocketId: (socketId: string | null) => {
+    set({ socketId });
+  },
+
+  reset: () => {
+    set({ status: 'disconnected', socketId: null });
+  },
+}));