diff --git a/apps/web/src/components/ai/shared/AiUsageMonitor.tsx b/apps/web/src/components/ai/shared/AiUsageMonitor.tsx index 7fd8ec32a..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 { @@ -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; 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(() => {