From 367e08c73ba7dcc609c8594675d39537328fc58b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 16 Feb 2026 23:39:51 +0000 Subject: [PATCH 1/2] fix: resolve socket race condition in AI usage monitor and billing counter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Subscribe to socket state reactively from the Zustand store instead of calling getSocket() synchronously after the async connect(). Previously, connect() initiated an async token fetch while getSocket() was called immediately — returning null before the socket existed. This meant the usage:updated listener was never registered, leaving the UI stale. Split the single effect into two: one for initiating the connection, and one reactive to the socket reference for registering event listeners. https://claude.ai/code/session_01B5GJTEZGDpfPTcJ3DqpsTE --- .../components/ai/shared/AiUsageMonitor.tsx | 44 +++++++------- .../src/components/billing/UsageCounter.tsx | 58 +++++++++---------- 2 files changed, 52 insertions(+), 50 deletions(-) diff --git a/apps/web/src/components/ai/shared/AiUsageMonitor.tsx b/apps/web/src/components/ai/shared/AiUsageMonitor.tsx index 7fd8ec32a..ef766d766 100644 --- a/apps/web/src/components/ai/shared/AiUsageMonitor.tsx +++ b/apps/web/src/components/ai/shared/AiUsageMonitor.tsx @@ -28,7 +28,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 +40,30 @@ 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 = () => { + 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; diff --git a/apps/web/src/components/billing/UsageCounter.tsx b/apps/web/src/components/billing/UsageCounter.tsx index 3cac056d0..881ba1472 100644 --- a/apps/web/src/components/billing/UsageCounter.tsx +++ b/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('/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(() => { From 5ccb33c9ce5e720c72047ff44a5cf64bffb580d9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 17 Feb 2026 00:14:50 +0000 Subject: [PATCH 2/2] fix: filter usage:updated events by conversation/page ID in AiUsageMonitor Accept the UsageEventPayload in handleUsageUpdated (matching the pattern in UsageCounter.tsx) and skip events whose conversationId or pageId don't match the local instance. This avoids unnecessary SWR revalidations when multiple AiUsageMonitor instances are mounted. The filter is forward-compatible: the current payload lacks these fields, so all events pass through today; once the backend includes them, filtering activates automatically. https://claude.ai/code/session_01B5GJTEZGDpfPTcJ3DqpsTE --- apps/web/src/components/ai/shared/AiUsageMonitor.tsx | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/apps/web/src/components/ai/shared/AiUsageMonitor.tsx b/apps/web/src/components/ai/shared/AiUsageMonitor.tsx index ef766d766..425a207ff 100644 --- a/apps/web/src/components/ai/shared/AiUsageMonitor.tsx +++ b/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 { @@ -50,7 +51,11 @@ export function AiUsageMonitor({ conversationId, pageId, className, compact = fa useEffect(() => { if (!socket) return; - const handleUsageUpdated = () => { + 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) {