fix(mcp): clean up disconnected editor bridges

server/ai/mcp/editorBridge.test.ts

@@ -94,4 +94,41 @@ describe('editor bridge', () => {
     await contentReader.read().catch(() => {})
     expect(hasEditorBridge(userId, 'content')).toBe(false)
   })
+
+  it('clears a bridge when the stream consumer cancels without aborting the request', async () => {
+    const userId = `u_${Math.floor(performance.now())}_cancelled`
+    const ctrl = new AbortController()
+    const reader = createEditorBridgeStream(userId, 'site', ctrl.signal).getReader()
+
+    await readUntil(reader, (e) => e.type === 'bridgeReady')
+    expect(hasEditorBridge(userId, 'site')).toBe(true)
+
+    await reader.cancel()
+    expect(ctrl.signal.aborted).toBe(false)
+    expect(hasEditorBridge(userId, 'site')).toBe(false)
+  })
+
+  it('does not evict the newest bridge when a superseded consumer disconnects', async () => {
+    const userId = `u_${Math.floor(performance.now())}_superseded`
+    const firstReader = createEditorBridgeStream(
+      userId,
+      'site',
+      new AbortController().signal,
+    ).getReader()
+    await readUntil(firstReader, (e) => e.type === 'bridgeReady')
+
+    const secondReader = createEditorBridgeStream(
+      userId,
+      'site',
+      new AbortController().signal,
+    ).getReader()
+    await readUntil(secondReader, (e) => e.type === 'bridgeReady')
+    expect(hasEditorBridge(userId, 'site')).toBe(true)
+
+    await firstReader.cancel()
+    expect(hasEditorBridge(userId, 'site')).toBe(true)
+
+    await secondReader.cancel()
+    expect(hasEditorBridge(userId, 'site')).toBe(false)
+  })
 })

server/ai/mcp/editorBridge.ts

@@ -53,62 +53,82 @@ export function createEditorBridgeStream(
   scope: EditorBridgeScope,
   signal: AbortSignal,
 ): ReadableStream<Uint8Array> {
+  let closeStream: (() => void) | null = null
+
   return new ReadableStream<Uint8Array>({
     start(controller) {
       let closed = false
       const encoder = new TextEncoder()
 
+      let bridgeId = ''
+      let destroyBridge = (): void => {}
+      let heartbeat: ReturnType<typeof setInterval> | null = null
+
+      const cleanup = () => {
+        if (closed) return
+        closed = true
+        if (heartbeat) clearInterval(heartbeat)
+        signal.removeEventListener('abort', cleanup)
+        destroyBridge()
+
+        // Only evict if we're still the current bridge for this scope. Keep
+        // the user's other workspace registered until its own stream closes.
+        const liveUserBridges = byUser.get(userId)
+        if (liveUserBridges?.get(scope)?.bridgeId === bridgeId) {
+          liveUserBridges.delete(scope)
+          if (liveUserBridges.size === 0) byUser.delete(userId)
+        }
+        try {
+          controller.close()
+        } catch {
+          /* already closed or cancelled */
+        }
+      }
+      closeStream = cleanup
+
       const emit = (event: AiStreamEvent): void => {
         if (closed) return
         try {
           controller.enqueue(encodeStreamEvent(event))
         } catch {
-          closed = true
+          cleanup()
         }
       }
 
-      const { bridgeId, bridge, destroy } = createBridge(emit, signal)
+      const created = createBridge(emit, signal)
+      bridgeId = created.bridgeId
+      destroyBridge = created.destroy
 
       // Newest instance of this workspace wins. The user's other workspace
       // remains connected, so Site and Content may serve MCP simultaneously.
       const userBridges = byUser.get(userId) ?? new Map<EditorBridgeScope, EditorBridgeEntry>()
       const previous = userBridges.get(scope)
       if (previous) previous.destroy()
-      userBridges.set(scope, { bridgeId, bridge, destroy })
+      userBridges.set(scope, { bridgeId, bridge: created.bridge, destroy: destroyBridge })
       byUser.set(userId, userBridges)
 
       emit({ type: 'bridgeReady', bridgeId })
 
       // Heartbeat blank line keeps proxies from idling the connection;
       // `readNdjsonStream` skips empty lines.
-      const heartbeat = setInterval(() => {
+      heartbeat = setInterval(() => {
         if (closed) return
         try {
           controller.enqueue(encoder.encode('\n'))
         } catch {
-          closed = true
+          cleanup()
         }
       }, 25_000)
 
-      const cleanup = () => {
-        if (closed) return
-        closed = true
-        clearInterval(heartbeat)
-        destroy()
-        // Only evict if we're still the current bridge for this scope. Keep
-        // the user's other workspace registered until its own stream closes.
-        const liveUserBridges = byUser.get(userId)
-        if (liveUserBridges?.get(scope)?.bridgeId === bridgeId) {
-          liveUserBridges.delete(scope)
-          if (liveUserBridges.size === 0) byUser.delete(userId)
-        }
-        try {
-          controller.close()
-        } catch {
-          /* already closed */
-        }
-      }
-      signal.addEventListener('abort', cleanup, { once: true })
+      if (signal.aborted) cleanup()
+      else signal.addEventListener('abort', cleanup, { once: true })
+    },
+    cancel() {
+      // Bun cancels the response body when the browser tab/context closes, but
+      // that transport cancellation does not abort the server Request signal.
+      // Tear down the heartbeat + registry entry from either lifecycle signal.
+      closeStream?.()
+      closeStream = null
     },
   })
 }

server/ai/mcp/server.ts

@@ -91,11 +91,12 @@ export function buildMcpServer(ctx: McpServerContext): Server {
           content: [{ type: 'text', text: `Browser tool "${tool.name}" has unsupported scope "${tool.scope}".` }],
         }
       }
-      const live = getEditorBridgeForUser(ctx.userId, tool.scope)
+      const browserScope: EditorBridgeScope = tool.scope
+      const live = getEditorBridgeForUser(ctx.userId, browserScope)
       if (!live) {
-        return { isError: true, content: [{ type: 'text', text: NO_WORKSPACE_MESSAGE[tool.scope] }] }
+        return { isError: true, content: [{ type: 'text', text: NO_WORKSPACE_MESSAGE[browserScope] }] }
       }
-      bridge = tool.scope === 'content'
+      bridge = browserScope === 'content'
         ? {
             callBrowser: async (toolName, input) => {
               await authorizeMcpContentTool(
@@ -105,7 +106,9 @@ export function buildMcpServer(ctx: McpServerContext): Server {
                 toolName,
                 input,
               )
-              return live.callBrowser(toolName, input)
+              const current = getEditorBridgeForUser(ctx.userId, browserScope)
+              if (!current) throw new Error(NO_WORKSPACE_MESSAGE[browserScope])
+              return current.callBrowser(toolName, input)
             },
           }
         : live

server/ai/runtime/transport.ts

@@ -87,9 +87,13 @@ export function createBridge(
   const bridgeId = nanoid()
   const entry: BridgeEntry = { pending: new Map(), emit, onSnapshot }
   activeBridges.set(bridgeId, entry)
+  let destroyed = false
 
   const bridge: AiBrowserBridge = {
     callBrowser(toolName, input) {
+      if (destroyed) {
+        return Promise.reject(new Error('AI chat stream ended before tool result arrived.'))
+      }
       const requestId = nanoid()
       return new Promise<AiToolOutput>((resolve, reject) => {
         // Settle (and remove) the pending wait on timeout or client disconnect
@@ -125,6 +129,8 @@ export function createBridge(
   }
 
   const destroy = () => {
+    if (destroyed) return
+    destroyed = true
     const live = activeBridges.get(bridgeId)
     if (!live) return
     if (live.pending.size > 0) {

src/__tests__/ai/bridgeTimeoutAbort.test.ts

@@ -16,6 +16,21 @@ afterEach(() => __destroyAllBridgesForTesting())
  * must settle on abort and on timeout.
  */
 describe('bridge tool-call settlement', () => {
+  test('rejects calls made after destroy without creating an orphan waiter', async () => {
+    let emitted = false
+    const { bridgeId, bridge, destroy } = createBridge(() => {
+      emitted = true
+    })
+
+    destroy()
+
+    await expect(bridge.callBrowser('cms.write', { x: 1 })).rejects.toThrow(
+      'AI chat stream ended before tool result arrived.',
+    )
+    expect(emitted).toBe(false)
+    expect(__listActiveBridgesForTesting()).not.toContain(bridgeId)
+  })
+
   test('rejects a pending tool call when the abort signal fires', async () => {
     const controller = new AbortController()
     const { bridge, destroy } = createBridge(() => {}, controller.signal)