import { useEffect } from 'react' import { createLogger } from '@sim/logger' import type { QueryClient } from '@tanstack/react-query' import { useQueryClient } from '@tanstack/react-query' import { getLiveAssistantMessageId } from '@/lib/copilot/chat/effective-transcript' import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/mothership-chats' const logger = createLogger('MothershipChatEvents') const CHAT_STATUS_TYPES = ['started', 'completed', 'created', 'deleted', 'renamed'] as const type ChatStatusEventType = (typeof CHAT_STATUS_TYPES)[number] const CHAT_STATUS_TYPE_SET = new Set(CHAT_STATUS_TYPES) interface ChatStatusEventPayload { chatId?: string type?: ChatStatusEventType streamId?: string } const DETAIL_INVALIDATING_CHAT_STATUS_TYPES = new Set([ 'started', 'completed', 'renamed', ]) function isChatStatusEventType(value: unknown): value is ChatStatusEventType { return typeof value === 'string' && CHAT_STATUS_TYPE_SET.has(value) } function isLocalOptimisticActiveStream(current: MothershipChatHistory | undefined) { if (!current?.activeStreamId) return false const liveAssistantId = getLiveAssistantMessageId(current.activeStreamId) return current.messages.some((message) => message.id === liveAssistantId) } /** * Returns true when the cached active stream is known to be later in the * chronological transcript than the stream that emitted this status event. * If either stream is absent from the transcript, callers should refetch * instead of inferring order from incomplete cache state. */ function hasNewerKnownActiveStream(current: MothershipChatHistory | undefined, streamId: string) { if (!current?.activeStreamId || current.activeStreamId === streamId) return false const activeIndex = current.messages.findIndex((message) => message.id === current.activeStreamId) const eventStreamIndex = current.messages.findIndex((message) => message.id === streamId) if (activeIndex === -1) return false if (eventStreamIndex === -1) return false return activeIndex > eventStreamIndex } function shouldSkipDetailInvalidationForStreamEvent( current: MothershipChatHistory | undefined, payload: ChatStatusEventPayload ) { if (payload.type !== 'started' && payload.type !== 'completed') return false if (!current?.activeStreamId) return false if (!payload.streamId) return isLocalOptimisticActiveStream(current) if (payload.type === 'started' && current.activeStreamId === payload.streamId) return true if (current.activeStreamId === payload.streamId) return false if (hasNewerKnownActiveStream(current, payload.streamId)) return true return ( payload.type === 'completed' && isLocalOptimisticActiveStream(current) && !current.messages.some((message) => message.id === payload.streamId) ) } function parseChatStatusEventPayload(data: unknown): ChatStatusEventPayload | null { let parsed = data if (typeof parsed === 'string') { try { parsed = JSON.parse(parsed) } catch { return null } } if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { return null } const record = parsed as Record return { ...(typeof record.chatId === 'string' ? { chatId: record.chatId } : {}), ...(isChatStatusEventType(record.type) ? { type: record.type } : {}), ...(typeof record.streamId === 'string' ? { streamId: record.streamId } : {}), } } export function handleMothershipChatStatusEvent( queryClient: Pick, workspaceId: string, data: unknown ): void { const payload = parseChatStatusEventPayload(data) if (!payload) { logger.warn('Received invalid task_status payload') return } queryClient.invalidateQueries({ queryKey: mothershipChatKeys.list(workspaceId) }) if (!payload.chatId) return if (payload.type === 'deleted') { queryClient.removeQueries({ queryKey: mothershipChatKeys.detail(payload.chatId) }) return } if (payload.type === 'started' || payload.type === 'completed') { const current = queryClient.getQueryData( mothershipChatKeys.detail(payload.chatId) ) if (shouldSkipDetailInvalidationForStreamEvent(current, payload)) { return } } if (payload.type && DETAIL_INVALIDATING_CHAT_STATUS_TYPES.has(payload.type)) { queryClient.invalidateQueries({ queryKey: mothershipChatKeys.detail(payload.chatId) }) } } /** * Subscribes to chat status SSE events and invalidates chat caches on changes. * The SSE event name remains `task_status` for wire compatibility. */ export function useMothershipChatEvents(workspaceId: string | undefined) { const queryClient = useQueryClient() useEffect(() => { if (!workspaceId) return const eventSource = new EventSource( `/api/mothership/events?workspaceId=${encodeURIComponent(workspaceId)}` ) eventSource.addEventListener('task_status', (event) => { handleMothershipChatStatusEvent( queryClient, workspaceId, event instanceof MessageEvent ? event.data : undefined ) }) eventSource.onerror = () => { logger.warn(`SSE connection error for workspace ${workspaceId}`) } return () => { eventSource.close() } }, [workspaceId, queryClient]) }