Refactor socket store access and split useEffect hooks

apps/web/src/components/ai/shared/AiUsageMonitor.tsx

@@ -7,6 +7,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/comp
 import { Activity, DollarSign, Database, AlertCircle } from 'lucide-react';
 import { cn } from '@/lib/utils';
 import { useSocketStore } from '@/stores/useSocketStore';
+import type { UsageEventPayload } from '@/lib/websocket';
 import { getUserFacingModelName } from '@/lib/ai/core/ai-providers-config';
 
 interface AiUsageMonitorProps {
@@ -28,7 +29,7 @@ interface AiUsageMonitorProps {
  */
 export function AiUsageMonitor({ conversationId, pageId, className, compact = false }: AiUsageMonitorProps) {
   const connect = useSocketStore((state) => state.connect);
-  const getSocket = useSocketStore((state) => state.getSocket);
+  const socket = useSocketStore((state) => state.socket);
 
   // Use conversation-based tracking for Global Assistant
   const { usage: conversationUsage, isLoading: conversationLoading, isError: conversationError, mutate: mutateConversation } = useAiUsage(
@@ -40,28 +41,34 @@ export function AiUsageMonitor({ conversationId, pageId, className, compact = fa
     !conversationId ? pageId : null // Only query if no conversationId
   );
 
-  // Socket.IO listener for real-time usage updates
+  // Ensure socket connection is initiated
   useEffect(() => {
     connect();
-    const socket = getSocket();
-
-    if (socket) {
-      const handleUsageUpdated = () => {
-        // Trigger refetch when usage updates
-        if (conversationId) {
-          mutateConversation();
-        } else if (pageId) {
-          mutatePage();
-        }
-      };
-
-      socket.on('usage:updated', handleUsageUpdated);
-
-      return () => {
-        socket.off('usage:updated', handleUsageUpdated);
-      };
-    }
-  }, [connect, getSocket, conversationId, pageId, mutateConversation, mutatePage]);
+  }, [connect]);
+
+  // Socket.IO listener for real-time usage updates
+  // Reactive to `socket` so the listener is registered once the async connection completes
+  useEffect(() => {
+    if (!socket) return;
+
+    const handleUsageUpdated = (payload: UsageEventPayload & { conversationId?: string; pageId?: string }) => {
+      // Skip events targeting a different conversation or page
+      if (payload.conversationId && payload.conversationId !== conversationId) return;
+      if (payload.pageId && payload.pageId !== pageId) return;
+
+      if (conversationId) {
+        mutateConversation();
+      } else if (pageId) {
+        mutatePage();
+      }
+    };
+
+    socket.on('usage:updated', handleUsageUpdated);
+
+    return () => {
+      socket.off('usage:updated', handleUsageUpdated);
+    };
+  }, [socket, conversationId, pageId, mutateConversation, mutatePage]);
 
   // Determine which data to use
   const usage = conversationId ? conversationUsage : pageUsage;

apps/web/src/components/billing/UsageCounter.tsx

@@ -41,7 +41,7 @@ const fetcher = async (url: string) => {
 export function UsageCounter() {
   const router = useRouter();
   const connect = useSocketStore((state) => state.connect);
-  const getSocket = useSocketStore((state) => state.getSocket);
+  const socket = useSocketStore((state) => state.socket);
   const { showBilling } = useBillingVisibility();
 
   const { data: usage, error, mutate } = useSWR<UsageData>('/api/subscriptions/usage', fetcher, {
@@ -61,36 +61,36 @@ export function UsageCounter() {
     router.push('/settings/billing');
   };
 
-  // Connect to Socket.IO and listen for usage events
+  // Ensure socket connection is initiated
   useEffect(() => {
     connect();
-    const socket = getSocket();
-
-    if (socket) {
-      const handleUsageUpdated = (payload: UsageEventPayload) => {
-        usageLogger.debug('Received usage update payload', {
-          userId: maskIdentifier(payload.userId),
-          subscriptionTier: payload.subscriptionTier,
-          standard: payload.standard,
-          pro: payload.pro,
-        });
-
-        // Update SWR cache with new usage data
-        mutate({
-          subscriptionTier: payload.subscriptionTier,
-          standard: payload.standard,
-          pro: payload.pro
-        }, false); // Don't revalidate, trust the real-time data
-      };
-
-      socket.on('usage:updated', handleUsageUpdated);
-
-      // Cleanup listener on unmount
-      return () => {
-        socket.off('usage:updated', handleUsageUpdated);
-      };
-    }
-  }, [connect, getSocket, mutate]);
+  }, [connect]);
+
+  // Listen for usage events once socket is available
+  useEffect(() => {
+    if (!socket) return;
+
+    const handleUsageUpdated = (payload: UsageEventPayload) => {
+      usageLogger.debug('Received usage update payload', {
+        userId: maskIdentifier(payload.userId),
+        subscriptionTier: payload.subscriptionTier,
+        standard: payload.standard,
+        pro: payload.pro,
+      });
+
+      mutate({
+        subscriptionTier: payload.subscriptionTier,
+        standard: payload.standard,
+        pro: payload.pro
+      }, false);
+    };
+
+    socket.on('usage:updated', handleUsageUpdated);
+
+    return () => {
+      socket.off('usage:updated', handleUsageUpdated);
+    };
+  }, [socket, mutate]);
 
   // Refresh usage when AI conversations complete (fallback)
   useEffect(() => {