test(sse): cover delivery across rotation

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_rotating_sse_connections?: Map<string, unknown> }
-    ).__mcp_rotating_sse_connections?.clear()
+      globalThis as unknown as { __mcp_sse_connections?: Map<string, unknown> }
+    ).__mcp_sse_connections?.clear()
     ;(globalThis as unknown as { __mcp_sse_subscribed?: Set<string> }).__mcp_sse_subscribed?.clear()
   })
 

apps/sim/hooks/queries/mcp.ts

@@ -543,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_rotating_sse_connections' as const
+const SSE_KEY = '__mcp_sse_connections' as const
 
-type SseEntry = { connection: RotatingEventSourceConnection; refs: number }
+type SseEntry = { source: RotatingEventSourceConnection; refs: number }
 
 const sseConnections: Map<string, SseEntry> =
   ((globalThis as Record<string, unknown>)[SSE_KEY] as Map<string, SseEntry>) ??
@@ -582,7 +582,7 @@ export function useMcpToolsEvents(workspaceId: string) {
     if (!entry) {
       const isResubscribe = sseEverSubscribed.has(workspaceId)
       sseEverSubscribed.add(workspaceId)
-      const connection = createRotatingEventSource({
+      const source = createRotatingEventSource({
         url: `/api/mcp/events?workspaceId=${encodeURIComponent(workspaceId)}`,
         events: {
           tools_changed: (event) => {
@@ -602,7 +602,7 @@ export function useMcpToolsEvents(workspaceId: string) {
         },
       })
 
-      entry = { connection, refs: 0 }
+      entry = { source, refs: 0 }
       sseConnections.set(workspaceId, entry)
     }
 
@@ -614,7 +614,7 @@ export function useMcpToolsEvents(workspaceId: string) {
 
       current.refs--
       if (current.refs <= 0) {
-        current.connection.close()
+        current.source.close()
         sseConnections.delete(workspaceId)
       }
     }

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

@@ -68,13 +68,15 @@ describe('createRotatingEventSource', () => {
 
   it('keeps the current source open until its replacement connects', () => {
     const reasons: EventSourceOpenReason[] = []
+    const onMessage = vi.fn()
     const connection = createRotatingEventSource({
       url: '/api/events',
-      events: {},
+      events: { message: onMessage },
       onOpen: (reason) => reasons.push(reason),
     })
     const first = MockEventSource.instances[0]
     first.open()
+    first.emit('message')
 
     first.emit('rotate')
 
@@ -83,28 +85,34 @@ describe('createRotatingEventSource', () => {
 
     const second = MockEventSource.instances[1]
     second.open()
+    second.emit('message')
 
     expect(first.readyState).toBe(MockEventSource.CLOSED)
     expect(second.readyState).toBe(MockEventSource.OPEN)
     expect(reasons).toEqual(['initial', 'rotation'])
+    expect(onMessage).toHaveBeenCalledTimes(2)
     connection.close()
   })
 
   it('classifies a replacement as reconnecting when the old source dropped first', () => {
     const reasons: EventSourceOpenReason[] = []
+    const onMessage = vi.fn()
     const connection = createRotatingEventSource({
       url: '/api/events',
-      events: {},
+      events: { message: onMessage },
       onOpen: (reason) => reasons.push(reason),
     })
     const first = MockEventSource.instances[0]
     first.open()
     first.emit('rotate')
     first.error()
 
-    MockEventSource.instances[1].open()
+    const second = MockEventSource.instances[1]
+    second.open()
+    second.emit('message')
 
     expect(reasons).toEqual(['initial', 'reconnect'])
+    expect(onMessage).toHaveBeenCalledTimes(1)
     connection.close()
   })