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