From 39cb17b3f253c5881045c33f400e3cfc74afee09 Mon Sep 17 00:00:00 2001 From: DavidBabinec Date: Sat, 11 Jul 2026 18:54:56 +0200 Subject: [PATCH] fix(mcp): clean up disconnected editor bridges --- server/ai/mcp/editorBridge.test.ts | 37 +++++++++++ server/ai/mcp/editorBridge.ts | 68 +++++++++++++-------- server/ai/mcp/server.ts | 11 ++-- server/ai/runtime/transport.ts | 6 ++ src/__tests__/ai/bridgeTimeoutAbort.test.ts | 15 +++++ 5 files changed, 109 insertions(+), 28 deletions(-) diff --git a/server/ai/mcp/editorBridge.test.ts b/server/ai/mcp/editorBridge.test.ts index e381d1e38..d7e5e5a54 100644 --- a/server/ai/mcp/editorBridge.test.ts +++ b/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) + }) }) diff --git a/server/ai/mcp/editorBridge.ts b/server/ai/mcp/editorBridge.ts index ee2df5707..e1bc46edf 100644 --- a/server/ai/mcp/editorBridge.ts +++ b/server/ai/mcp/editorBridge.ts @@ -53,62 +53,82 @@ export function createEditorBridgeStream( scope: EditorBridgeScope, signal: AbortSignal, ): ReadableStream { + let closeStream: (() => void) | null = null + return new ReadableStream({ start(controller) { let closed = false const encoder = new TextEncoder() + let bridgeId = '' + let destroyBridge = (): void => {} + let heartbeat: ReturnType | 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() 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 }, }) } diff --git a/server/ai/mcp/server.ts b/server/ai/mcp/server.ts index 5da41bbaf..b2a29a36f 100644 --- a/server/ai/mcp/server.ts +++ b/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 diff --git a/server/ai/runtime/transport.ts b/server/ai/runtime/transport.ts index 7b1f755d5..c6cebf784 100644 --- a/server/ai/runtime/transport.ts +++ b/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((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) { diff --git a/src/__tests__/ai/bridgeTimeoutAbort.test.ts b/src/__tests__/ai/bridgeTimeoutAbort.test.ts index 4bc993448..338daa63a 100644 --- a/src/__tests__/ai/bridgeTimeoutAbort.test.ts +++ b/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)