fix(sse): rotate workspace streams without gaps

Author: waleedlatif1

apps/sim/hooks/queries/mcp.test.tsx

@@ -127,8 +127,8 @@ describe('useMcpToolsQuery', () => {
     // mcp.ts captured these Map/Set instances in module consts at import, so reassigning the
     // globalThis property wouldn't reset what the module uses — clear the shared instances.
     ;(
-      globalThis as unknown as { __mcp_sse_connections?: Map<string, unknown> }
-    ).__mcp_sse_connections?.clear()
+      globalThis as unknown as { __mcp_rotating_sse_connections?: Map<string, unknown> }
+    ).__mcp_rotating_sse_connections?.clear()
     ;(globalThis as unknown as { __mcp_sse_subscribed?: Set<string> }).__mcp_sse_subscribed?.clear()
   })
 

apps/sim/hooks/queries/mcp.ts

@@ -27,6 +27,10 @@ import {
   testMcpServerConnectionContract,
   updateMcpServerContract,
 } from '@/lib/api/contracts/mcp'
+import {
+  createRotatingEventSource,
+  type RotatingEventSourceConnection,
+} from '@/lib/events/rotating-event-source'
 import { sanitizeForHttp, sanitizeHeaders } from '@/lib/mcp/shared'
 import type {
   McpAuthType,
@@ -539,9 +543,9 @@ export function useStoredMcpTools(workspaceId: string, options?: { enabled?: boo
  * Reference-counted so the connection is closed when the last consumer unmounts.
  * Attached to `globalThis` so connections survive HMR in development.
  */
-const SSE_KEY = '__mcp_sse_connections' as const
+const SSE_KEY = '__mcp_rotating_sse_connections' as const
 
-type SseEntry = { source: EventSource; refs: number }
+type SseEntry = { connection: RotatingEventSourceConnection; refs: number }
 
 const sseConnections: Map<string, SseEntry> =
   ((globalThis as Record<string, unknown>)[SSE_KEY] as Map<string, SseEntry>) ??
@@ -576,40 +580,29 @@ export function useMcpToolsEvents(workspaceId: string) {
     let entry = sseConnections.get(workspaceId)
 
     if (!entry) {
-      const source = new EventSource(`/api/mcp/events?workspaceId=${workspaceId}`)
-
-      source.addEventListener('tools_changed', (e) => {
-        let serverId: string | undefined
-        try {
-          const parsed = JSON.parse((e as MessageEvent).data) as { serverId?: string }
-          serverId = parsed.serverId
-        } catch {
-          // Non-JSON payload → workspace-wide fallback.
-        }
-        invalidate(serverId)
-      })
-
-      // EventSource fires `onopen` on the initial connect and on every auto-reconnect. Re-sync
-      // the workspace whenever we could have missed a `tools_changed` event: on any reconnect,
-      // on the first open of a RE-subscription (leaving the tab tears the connection down), and
-      // on a first open that only succeeded after an earlier connection error (the initial tools
-      // query may have failed during that gap and won't retry itself). Skip only a clean first
-      // subscription — the queries fetch fresh on their own initial mount.
       const isResubscribe = sseEverSubscribed.has(workspaceId)
       sseEverSubscribed.add(workspaceId)
-      let opened = false
-      let erroredBeforeOpen = false
-      source.onopen = () => {
-        if (opened || isResubscribe || erroredBeforeOpen) invalidate()
-        opened = true
-      }
-
-      source.onerror = () => {
-        if (!opened) erroredBeforeOpen = true
-        logger.warn(`SSE connection error for workspace ${workspaceId}`)
-      }
+      const connection = createRotatingEventSource({
+        url: `/api/mcp/events?workspaceId=${encodeURIComponent(workspaceId)}`,
+        events: {
+          tools_changed: (event) => {
+            let serverId: string | undefined
+            try {
+              const parsed = JSON.parse((event as MessageEvent).data) as { serverId?: string }
+              serverId = parsed.serverId
+            } catch {}
+            invalidate(serverId)
+          },
+        },
+        onOpen: (reason) => {
+          if (reason === 'reconnect' || (reason === 'initial' && isResubscribe)) invalidate()
+        },
+        onError: () => {
+          logger.warn(`SSE connection error for workspace ${workspaceId}`)
+        },
+      })
 
-      entry = { source, refs: 0 }
+      entry = { connection, refs: 0 }
       sseConnections.set(workspaceId, entry)
     }
 
@@ -621,7 +614,7 @@ export function useMcpToolsEvents(workspaceId: string) {
 
       current.refs--
       if (current.refs <= 0) {
-        current.source.close()
+        current.connection.close()
         sseConnections.delete(workspaceId)
       }
     }

apps/sim/hooks/use-mothership-chat-events.ts

@@ -5,6 +5,7 @@ import { useQueryClient } from '@tanstack/react-query'
 import { getLiveAssistantMessageId } from '@/lib/copilot/chat/effective-transcript'
 import { isChatEnabled } from '@/lib/core/config/env-flags'
 import { suspendDesktopChatScopes } from '@/lib/desktop/chat-scope'
+import { createRotatingEventSource } from '@/lib/events/rotating-event-source'
 import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/mothership-chats'
 
 const logger = createLogger('MothershipChatEvents')
@@ -167,45 +168,31 @@ export function useMothershipChatEvents(workspaceId: string | undefined) {
   useEffect(() => {
     if (!workspaceId || !isChatEnabled) 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
-      )
-    })
-
-    // `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}`)
-    }
+    const connection = createRotatingEventSource({
+      url: `/api/mothership/events?workspaceId=${encodeURIComponent(workspaceId)}`,
+      events: {
+        task_status: (event) => {
+          handleMothershipChatStatusEvent(
+            queryClient,
+            workspaceId,
+            event instanceof MessageEvent ? event.data : undefined
+          )
+        },
+      },
+      onOpen: (reason) => {
+        if (reason === 'reconnect' || (reason === 'initial' && isResubscribe)) {
+          resyncMothershipChatCaches(queryClient, workspaceId)
+        }
+      },
+      onError: () => {
+        logger.warn(`SSE connection error for workspace ${workspaceId}`)
+      },
+    })
 
     return () => {
-      eventSource.close()
+      connection.close()
     }
   }, [workspaceId, queryClient])
 }

apps/sim/lib/events/rotating-event-source.test.ts

@@ -0,0 +1,154 @@
+/**
+ * @vitest-environment node
+ */
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import {
+  createRotatingEventSource,
+  type EventSourceOpenReason,
+} from '@/lib/events/rotating-event-source'
+
+class MockEventSource {
+  static readonly CONNECTING = 0
+  static readonly OPEN = 1
+  static readonly CLOSED = 2
+  static instances: MockEventSource[] = []
+  static failNextConstruction = false
+
+  readonly listeners = new Map<string, Array<(event: Event) => void>>()
+  readyState = MockEventSource.CONNECTING
+  onopen: (() => void) | null = null
+  onerror: (() => void) | null = null
+
+  constructor(readonly url: string) {
+    if (MockEventSource.failNextConstruction) {
+      MockEventSource.failNextConstruction = false
+      throw new Error('construction failed')
+    }
+    MockEventSource.instances.push(this)
+  }
+
+  addEventListener(eventName: string, listener: (event: Event) => void): void {
+    const listeners = this.listeners.get(eventName) ?? []
+    listeners.push(listener)
+    this.listeners.set(eventName, listeners)
+  }
+
+  close(): void {
+    this.readyState = MockEventSource.CLOSED
+  }
+
+  open(): void {
+    this.readyState = MockEventSource.OPEN
+    this.onopen?.()
+  }
+
+  error(): void {
+    this.readyState = MockEventSource.CONNECTING
+    this.onerror?.()
+  }
+
+  emit(eventName: string): void {
+    for (const listener of this.listeners.get(eventName) ?? []) {
+      listener(new Event(eventName))
+    }
+  }
+}
+
+describe('createRotatingEventSource', () => {
+  beforeEach(() => {
+    MockEventSource.instances = []
+    MockEventSource.failNextConstruction = false
+    vi.stubGlobal('EventSource', MockEventSource)
+  })
+
+  afterEach(() => {
+    vi.unstubAllGlobals()
+  })
+
+  it('keeps the current source open until its replacement connects', () => {
+    const reasons: EventSourceOpenReason[] = []
+    const connection = createRotatingEventSource({
+      url: '/api/events',
+      events: {},
+      onOpen: (reason) => reasons.push(reason),
+    })
+    const first = MockEventSource.instances[0]
+    first.open()
+
+    first.emit('rotate')
+
+    expect(MockEventSource.instances).toHaveLength(2)
+    expect(first.readyState).toBe(MockEventSource.OPEN)
+
+    const second = MockEventSource.instances[1]
+    second.open()
+
+    expect(first.readyState).toBe(MockEventSource.CLOSED)
+    expect(second.readyState).toBe(MockEventSource.OPEN)
+    expect(reasons).toEqual(['initial', 'rotation'])
+    connection.close()
+  })
+
+  it('classifies a replacement as reconnecting when the old source dropped first', () => {
+    const reasons: EventSourceOpenReason[] = []
+    const connection = createRotatingEventSource({
+      url: '/api/events',
+      events: {},
+      onOpen: (reason) => reasons.push(reason),
+    })
+    const first = MockEventSource.instances[0]
+    first.open()
+    first.emit('rotate')
+    first.error()
+
+    MockEventSource.instances[1].open()
+
+    expect(reasons).toEqual(['initial', 'reconnect'])
+    connection.close()
+  })
+
+  it('classifies an automatic EventSource recovery as a reconnect', () => {
+    const reasons: EventSourceOpenReason[] = []
+    const connection = createRotatingEventSource({
+      url: '/api/events',
+      events: {},
+      onOpen: (reason) => reasons.push(reason),
+    })
+    const source = MockEventSource.instances[0]
+    source.open()
+    source.error()
+    source.open()
+
+    expect(reasons).toEqual(['initial', 'reconnect'])
+    connection.close()
+  })
+
+  it('closes both sources when disposed during rotation', () => {
+    const connection = createRotatingEventSource({ url: '/api/events', events: {} })
+    const first = MockEventSource.instances[0]
+    first.open()
+    first.emit('rotate')
+    const second = MockEventSource.instances[1]
+
+    connection.close()
+
+    expect(first.readyState).toBe(MockEventSource.CLOSED)
+    expect(second.readyState).toBe(MockEventSource.CLOSED)
+  })
+
+  it('keeps the current source when opening its replacement fails', () => {
+    const onError = vi.fn()
+    const connection = createRotatingEventSource({ url: '/api/events', events: {}, onError })
+    const first = MockEventSource.instances[0]
+    first.open()
+    MockEventSource.failNextConstruction = true
+
+    first.emit('rotate')
+
+    expect(first.readyState).toBe(MockEventSource.OPEN)
+    expect(MockEventSource.instances).toHaveLength(1)
+    expect(onError).toHaveBeenCalledTimes(1)
+    connection.close()
+  })
+})