fix(sse): bound workspace SSE connection lifetime

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<QueryClient, 'invalidateQueries'>
+
+  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() })
+    )
+  })
+})

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<string>()
+
 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<string>(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<QueryClient, 'invalidateQueries'>,
+  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}`)
     }
 

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<Uint8Array>, unsubscribe }
+}
+
+/** Resolves once the stream closes. */
+async function drain(body: ReadableStream<Uint8Array>): Promise<void> {
+  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)
+  })
+})

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<string, unknown>) => {
           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')
       },
     })