diff --git a/apps/sim/hooks/use-mothership-chat-events.test.ts b/apps/sim/hooks/use-mothership-chat-events.test.ts index 52475b0ecf6..fb01186b06c 100644 --- a/apps/sim/hooks/use-mothership-chat-events.test.ts +++ b/apps/sim/hooks/use-mothership-chat-events.test.ts @@ -14,7 +14,10 @@ vi.mock('@/lib/browser-agent/transport', () => ({ suspendBrowserScope })) vi.mock('@/lib/terminal/transport', () => ({ suspendTerminalScope })) import { mothershipChatKeys } from '@/hooks/queries/mothership-chats' -import { handleMothershipChatStatusEvent } from '@/hooks/use-mothership-chat-events' +import { + handleMothershipChatStatusEvent, + resyncMothershipChatCaches, +} from '@/hooks/use-mothership-chat-events' describe('handleMothershipChatStatusEvent', () => { const queryClient = { @@ -419,3 +422,30 @@ describe('handleMothershipChatStatusEvent', () => { expect(queryClient.removeQueries).not.toHaveBeenCalled() }) }) + +describe('resyncMothershipChatCaches', () => { + const queryClient = { + invalidateQueries: vi.fn().mockResolvedValue(undefined), + } satisfies Pick + + beforeEach(() => { + vi.clearAllMocks() + }) + + it('invalidates the workspace lists', () => { + resyncMothershipChatCaches(queryClient, 'ws-1') + + expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1) + expect(queryClient.invalidateQueries).toHaveBeenCalledWith({ + queryKey: mothershipChatKeys.workspaceLists('ws-1'), + }) + }) + + it('leaves chat details untouched so a mounted stream cannot be refetched mid-turn', () => { + resyncMothershipChatCaches(queryClient, 'ws-1') + + expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith( + expect.objectContaining({ queryKey: mothershipChatKeys.details() }) + ) + }) +}) diff --git a/apps/sim/hooks/use-mothership-chat-events.ts b/apps/sim/hooks/use-mothership-chat-events.ts index 096aaed5812..175891a4dbd 100644 --- a/apps/sim/hooks/use-mothership-chat-events.ts +++ b/apps/sim/hooks/use-mothership-chat-events.ts @@ -9,6 +9,9 @@ import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/ const logger = createLogger('MothershipChatEvents') +/** Workspaces this process has subscribed to before, so a re-subscribe can be told from a first one. */ +const everSubscribed = new Set() + 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) @@ -128,6 +131,28 @@ export function handleMothershipChatStatusEvent( } } +/** + * Re-syncs the workspace chat lists after a gap in the event stream. + * + * `task_status` events are transient — nothing replays what was published while + * no connection was open — so a reconnect may have missed a create, rename, or + * delete. The lists carry that workspace-level state, and refetching them is + * always safe. + * + * Chat details are deliberately left alone. The only one that would refetch is + * the mounted chat, which may be rendering an in-flight stream, and refetching + * there replaces the optimistic transcript with a server copy that does not yet + * hold the streaming message. Cached state cannot reliably say whether a turn is + * still running — the optimistic markers outlive it — so detail reconciliation + * stays as it is today and belongs with the streaming state that can answer it. + */ +export function resyncMothershipChatCaches( + queryClient: Pick, + workspaceId: string +): void { + queryClient.invalidateQueries({ queryKey: mothershipChatKeys.workspaceLists(workspaceId) }) +} + /** * Subscribes to chat status SSE events and invalidates chat caches on changes. * The SSE event name remains `task_status` for wire compatibility. @@ -154,7 +179,28 @@ export function useMothershipChatEvents(workspaceId: string | undefined) { ) }) + // `onopen` fires on the initial connect and on every auto-reconnect. Re-sync + // whenever a gap could have swallowed an event: on any reconnect, on the + // first open of a RE-subscription (switching workspace away and back tears + // the connection down, and the list/detail queries remount inside their + // stale times so they do not refetch on their own), and on a first open that + // only succeeded after an error (the initial queries may have failed during + // that gap and will not retry themselves). Skip only a clean first + // subscription — those queries fetch fresh on their own initial mount. + const isResubscribe = everSubscribed.has(workspaceId) + everSubscribed.add(workspaceId) + let opened = false + let erroredBeforeOpen = false + + eventSource.onopen = () => { + if (opened || isResubscribe || erroredBeforeOpen) { + resyncMothershipChatCaches(queryClient, workspaceId) + } + opened = true + } + eventSource.onerror = () => { + if (!opened) erroredBeforeOpen = true logger.warn(`SSE connection error for workspace ${workspaceId}`) } diff --git a/apps/sim/lib/events/sse-endpoint.test.ts b/apps/sim/lib/events/sse-endpoint.test.ts new file mode 100644 index 00000000000..242a9be1b7c --- /dev/null +++ b/apps/sim/lib/events/sse-endpoint.test.ts @@ -0,0 +1,111 @@ +/** + * @vitest-environment node + */ + +import { authMockFns, permissionsMock, permissionsMockFns } from '@sim/testing' +import { NextRequest } from 'next/server' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { + createWorkspaceSSE, + HEARTBEAT_INTERVAL_MS, + MAX_CONNECTION_JITTER_MS, + MAX_CONNECTION_MS, + MAX_UNDRAINED_CHUNKS, +} from '@/lib/events/sse-endpoint' + +vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock) + +const PAST_CEILING_MS = MAX_CONNECTION_MS + MAX_CONNECTION_JITTER_MS + HEARTBEAT_INTERVAL_MS + +/** Enough undrained heartbeats to trip the unread check, and no more. */ +const PAST_UNREAD_MS = (MAX_UNDRAINED_CHUNKS + 2) * HEARTBEAT_INTERVAL_MS + +async function openConnection(signal: AbortSignal = new AbortController().signal) { + const unsubscribe = vi.fn() + const handler = createWorkspaceSSE({ + label: 'test', + subscriptions: [{ subscribe: () => unsubscribe }], + }) + const request = new NextRequest(new URL('https://sim.test/api/test/events?workspaceId=ws-1'), { + signal, + }) + const response = await handler(request) + + return { body: response.body as ReadableStream, unsubscribe } +} + +/** Resolves once the stream closes. */ +async function drain(body: ReadableStream): Promise { + const reader = body.getReader() + while (true) { + const { done } = await reader.read() + if (done) return + } +} + +describe('createWorkspaceSSE', () => { + beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers() + authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-1' } }) + permissionsMockFns.mockGetUserEntityPermissions.mockResolvedValue('admin') + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('releases subscriptions and closes the stream when the connection reaches its ceiling', async () => { + const { body, unsubscribe } = await openConnection() + const drained = drain(body) + + await vi.advanceTimersByTimeAsync(PAST_CEILING_MS) + + await drained + expect(unsubscribe).toHaveBeenCalledTimes(1) + }) + + it('releases subscriptions when the consumer stops draining the stream', async () => { + const { unsubscribe } = await openConnection() + + await vi.advanceTimersByTimeAsync(PAST_UNREAD_MS) + + expect(unsubscribe).toHaveBeenCalledTimes(1) + }) + + it('keeps a drained connection alive past the unread threshold', async () => { + const { body, unsubscribe } = await openConnection() + void drain(body) + + await vi.advanceTimersByTimeAsync(PAST_UNREAD_MS) + + expect(unsubscribe).not.toHaveBeenCalled() + }) + + it('releases subscriptions when the request aborts', async () => { + const controller = new AbortController() + const { unsubscribe } = await openConnection(controller.signal) + + controller.abort() + + expect(unsubscribe).toHaveBeenCalledTimes(1) + }) + + it('releases subscriptions when the consumer cancels the stream', async () => { + const { body, unsubscribe } = await openConnection() + + await body.cancel() + + expect(unsubscribe).toHaveBeenCalledTimes(1) + }) + + it('releases subscriptions once when abort and the ceiling both elapse', async () => { + const controller = new AbortController() + const { unsubscribe } = await openConnection(controller.signal) + + controller.abort() + await vi.advanceTimersByTimeAsync(PAST_CEILING_MS) + + expect(unsubscribe).toHaveBeenCalledTimes(1) + }) +}) diff --git a/apps/sim/lib/events/sse-endpoint.ts b/apps/sim/lib/events/sse-endpoint.ts index 30cc46619e0..8e5df184e1d 100644 --- a/apps/sim/lib/events/sse-endpoint.ts +++ b/apps/sim/lib/events/sse-endpoint.ts @@ -6,6 +6,7 @@ */ import { createLogger } from '@sim/logger' +import { randomFloat } from '@sim/utils/random' import type { NextRequest } from 'next/server' import { getSession } from '@/lib/auth' import { SSE_HEADERS } from '@/lib/core/utils/sse' @@ -23,7 +24,40 @@ interface WorkspaceSSEConfig { subscriptions: SSESubscription[] } -const HEARTBEAT_INTERVAL_MS = 30_000 +const encoder = new TextEncoder() + +export const HEARTBEAT_INTERVAL_MS = 30_000 + +/** + * Defensive ceiling on one connection's lifetime; `EventSource` reconnects past + * this, so delivery continues across the boundary. + * + * `request.signal` abort and stream `cancel()` are the primary teardown paths, + * but both fire only when the runtime reports the client disconnect, and the + * unread check below only catches a consumer that has stopped draining. This + * releases whatever both miss, so retention is bounded by the ceiling instead + * of by process uptime. + * + * Matches the ceiling `lib/realtime/event-stream-route.ts` already uses for the + * same purpose. It is deliberately far longer than the unread window: a healthy + * client is drained and therefore never unread, so a short ceiling would only + * force reconnects on the connections that are working, and every reconnect is + * a window in which a transient event can be missed. + */ +export const MAX_CONNECTION_MS = 4 * 60 * 60 * 1000 + +/** Spreads reconnects so connections opened together do not expire together. */ +export const MAX_CONNECTION_JITTER_MS = 60_000 + +/** + * Undrained chunk count that marks a connection unread, which shortens the + * reclaim window for a vanished consumer that the ceiling above would + * otherwise hold for its full duration. The default queuing strategy reports + * `desiredSize` as `1 - queued`, so this trips only once the consumer has + * pulled nothing for several minutes — well beyond the transient backpressure + * of a slow but live client. + */ +export const MAX_UNDRAINED_CHUNKS = 16 export function createWorkspaceSSE(config: WorkspaceSSEConfig) { const logger = createLogger(`${config.label}-SSE`) @@ -45,21 +79,30 @@ export function createWorkspaceSSE(config: WorkspaceSSEConfig) { return new Response('Access denied to workspace', { status: 403 }) } - const encoder = new TextEncoder() - const unsubscribers: Array<() => void> = [] + const teardowns: Array<() => void> = [] let cleaned = false - const cleanup = () => { + const cleanup = (reason: string) => { if (cleaned) return cleaned = true - for (const unsub of unsubscribers) { - unsub() + for (const teardown of teardowns) { + teardown() } - logger.info(`SSE connection closed for workspace ${workspaceId}`) + teardowns.length = 0 + logger.info(`SSE connection closed for workspace ${workspaceId}`, { reason }) } const stream = new ReadableStream({ start(controller) { + const close = (reason: string) => { + cleanup(reason) + try { + controller.close() + } catch { + // Already closed + } + } + const send = (eventName: string, data: Record) => { if (cleaned) return try { @@ -72,40 +115,47 @@ export function createWorkspaceSSE(config: WorkspaceSSEConfig) { } for (const subscription of config.subscriptions) { - const unsub = subscription.subscribe(workspaceId, send) - unsubscribers.push(unsub) + teardowns.push(subscription.subscribe(workspaceId, send)) } + const deadline = Date.now() + MAX_CONNECTION_MS + randomFloat() * MAX_CONNECTION_JITTER_MS + const heartbeat = setInterval(() => { if (cleaned) { clearInterval(heartbeat) return } + if (Date.now() >= deadline) { + close('expired') + return + } + const desiredSize = controller.desiredSize + if (desiredSize !== null && desiredSize <= -MAX_UNDRAINED_CHUNKS) { + close('unread') + return + } try { controller.enqueue(encoder.encode(': heartbeat\n\n')) } catch { - clearInterval(heartbeat) + close('errored') } }, HEARTBEAT_INTERVAL_MS) - unsubscribers.push(() => clearInterval(heartbeat)) - - request.signal.addEventListener( - 'abort', - () => { - cleanup() - try { - controller.close() - } catch { - // Already closed - } - }, - { once: true } - ) + teardowns.push(() => clearInterval(heartbeat)) + + // `once` only self-removes if abort fires; the expiry and unread paths + // close the connection while the signal is still live, so the listener + // needs its own removal or it retains this whole scope. + const listenerScope = new AbortController() + request.signal.addEventListener('abort', () => close('aborted'), { + once: true, + signal: listenerScope.signal, + }) + teardowns.push(() => listenerScope.abort()) logger.info(`SSE connection opened for workspace ${workspaceId}`) }, cancel() { - cleanup() + cleanup('cancelled') }, })