From c6e5a95b80335c5efa3927d4ad67db208a672622 Mon Sep 17 00:00:00 2001 From: Bhavesh Patel Date: Thu, 17 Sep 2026 00:53:46 +0530 Subject: [PATCH] fix: don't close a shared MCP transport under in-flight sibling calls --- .changeset/age-2253-mcp-reset-inflight.md | 5 + .../trueforge-core/src/core/mcp/RemoteMCP.ts | 96 ++++++++++++++----- .../tests/core/mcp/remoteMcpServer.test.ts | 79 ++++++++++++++- 3 files changed, 154 insertions(+), 26 deletions(-) create mode 100644 .changeset/age-2253-mcp-reset-inflight.md diff --git a/.changeset/age-2253-mcp-reset-inflight.md b/.changeset/age-2253-mcp-reset-inflight.md new file mode 100644 index 000000000..2b04d7f68 --- /dev/null +++ b/.changeset/age-2253-mcp-reset-inflight.md @@ -0,0 +1,5 @@ +--- +"@truefoundry/trueforge-core": patch +--- + +Keep a shared remote MCP transport open while sibling tool calls are in flight so a session-expired reset cannot fail them with a non-retried close error. diff --git a/packages/trueforge-core/src/core/mcp/RemoteMCP.ts b/packages/trueforge-core/src/core/mcp/RemoteMCP.ts index ef4a7c312..13a1037f8 100644 --- a/packages/trueforge-core/src/core/mcp/RemoteMCP.ts +++ b/packages/trueforge-core/src/core/mcp/RemoteMCP.ts @@ -67,6 +67,8 @@ export class RemoteMCP implements ToolSource { private sessionId: string | null | undefined; private resolvedTransportType?: RemoteMcpTransportType | undefined; private cachedTools?: AgentToolSchema[] | undefined; + private inflight = 0; + private pendingClose: RemoteMcpConnection | undefined; constructor(params: { name: string; @@ -101,14 +103,29 @@ export class RemoteMCP implements ToolSource { return this.sessionId ?? undefined; } - private get connection(): RemoteMcpConnection { - if (!this._connection) { - throw new Error(`Remote MCP '${this.name}' not connected - connectIfNeeded() must run first`); + // Concurrent callTool/listTools share one transport. Session-expired retry must not close it + // while a sibling is still using it — close() aborts the sibling with "Connection closed", + // which is not treated as session-expired, so that call is never retried. + // + // Example: A (long) and B share socket S. B gets session-expired. + // Detach S (pendingClose), B retries on a new socket, A finishes on S, + // last caller closes pendingClose when inflight hits 0. + // + // connectAndRun: + // conn = this._connection // capture; op does not re-read this._connection + // inflight++ + // try: return op(conn) + // catch sessionExpired: + // detach conn // _connection = undefined; do not close if inflight > 0 + // pendingClose = conn + // reconnect and retry once + // finally: + // inflight-- + // if inflight == 0: close(pendingClose) + private async resetConnection(expired?: RemoteMcpConnection): Promise { + if (expired !== undefined && this._connection !== expired) { + return; } - return this._connection; - } - - private async resetConnection(): Promise { this.isConnected = false; this.sessionId = undefined; this.cachedTools = undefined; @@ -117,10 +134,19 @@ export class RemoteMCP implements ToolSource { } private async closeAndClearConnection(): Promise { - await this._connection?.close().catch(() => { + const connection = this._connection; + this._connection = undefined; + if (!connection) { + return; + } + if (this.inflight > 0) { + // One leftover socket; a second session-expiry while the first is still pending can leak it. + this.pendingClose = connection; + return; + } + await connection.close().catch(() => { /* no-op */ }); - this._connection = undefined; } private async resolveHeaders(): Promise { @@ -130,7 +156,9 @@ export class RemoteMCP implements ToolSource { return await this.headers(); } - private async executeWithSessionRetry(operation: () => Promise): Promise> { + private async executeWithSessionRetry( + operation: (connection: RemoteMcpConnection) => Promise, + ): Promise> { // Auth is re-checked on every operation, not only on the first connect: a registered server's OAuth // can be revoked or expire mid-request, and callers must get authRequired rather than a generic // upstream failure. When already connected the resolved headers are unused (connect is skipped). @@ -143,29 +171,47 @@ export class RemoteMCP implements ToolSource { private async connectAndRun( headers: Record, - operation: () => Promise, + operation: (connection: RemoteMcpConnection) => Promise, canRetry: boolean, ): Promise> { + let used: RemoteMcpConnection | undefined; try { const initInfo = await this.connectIfNeeded(headers); - return { result: await operation(), wasInitialized: initInfo }; + const connection = this._connection; + if (!connection) { + throw new Error(`Remote MCP '${this.name}' not connected - connectIfNeeded() must run first`); + } + used = connection; + this.inflight += 1; + return { result: await operation(connection), wasInitialized: initInfo }; } catch (error) { - if (canRetry && isSessionExpiredError(error)) { - this.logger.info(`Session expired for remote MCP ${this.name}, reinitializing...`); - await this.resetConnection(); - return this.connectAndRun(headers, operation, false); + if (!(canRetry && isSessionExpiredError(error))) { + throw error; + } + this.logger.info(`Session expired for remote MCP ${this.name}, reinitializing...`); + await this.resetConnection(used); + } finally { + if (used) { + this.inflight -= 1; + if (this.inflight === 0 && this.pendingClose) { + const stale = this.pendingClose; + this.pendingClose = undefined; + await stale.close().catch(() => { + /* no-op */ + }); + } } - throw error; } + return this.connectAndRun(headers, operation, false); } - private async loadTools(): Promise<{ tools: AgentToolSchema[] }> { + private async loadTools(connection: RemoteMcpConnection): Promise<{ tools: AgentToolSchema[] }> { return this.tracing.withRemoteMcpToolSpan( { method: 'tools/list', serverName: this.name, serverId: this.id, serverUrl: this.traceUrl, enabled: true }, async span => { const tools = await paginateWithCursorGuard( async cursor => { - const page = await this.connection.listTools(cursor); + const page = await connection.listTools(cursor); return { items: page.tools, nextCursor: page.nextCursor }; }, this.name, @@ -183,7 +229,7 @@ export class RemoteMCP implements ToolSource { if (this.cachedTools) { return { result: { tools: this.cachedTools }, wasInitialized: undefined }; } - const response = await this.executeWithSessionRetry(() => this.loadTools()); + const response = await this.executeWithSessionRetry(connection => this.loadTools(connection)); if ('authRequired' in response) { return response; } @@ -191,7 +237,7 @@ export class RemoteMCP implements ToolSource { } async callTool(params: CallToolRequest['params']): Promise { - const response = await this.executeWithSessionRetry(() => + const response = await this.executeWithSessionRetry(connection => this.tracing.withRemoteMcpToolSpan( { method: 'tools/call', @@ -203,7 +249,7 @@ export class RemoteMCP implements ToolSource { enabled: true, }, async span => { - const result = await this.connection.callTool(params); + const result = await connection.callTool(params); span.setOutput(JSON.stringify(result)); return result; }, @@ -244,6 +290,7 @@ export class RemoteMCP implements ToolSource { connection = await this.tracing.withRemoteMcpToolSpan( { method: 'initialize', serverName: this.name, serverId: this.id, serverUrl: this.traceUrl, enabled: true }, async span => { + const attached: { current: RemoteMcpConnection | undefined } = { current: undefined }; const conn = await connectRemoteMcp({ url: this.url, headers, @@ -254,7 +301,9 @@ export class RemoteMCP implements ToolSource { connectTimeoutMs: this.connectTimeoutMs, signal: this.signal, onClose: () => { - this.isConnected = false; + if (attached.current !== undefined && this._connection === attached.current) { + this.isConnected = false; + } }, onError: error => { const fields = extractErrorLogFields(error); @@ -267,6 +316,7 @@ export class RemoteMCP implements ToolSource { }, }); span.setOutput(JSON.stringify({ transport: conn.transportType, stateful: conn.sessionId !== null })); + attached.current = conn; return conn; }, ); diff --git a/packages/trueforge-core/tests/core/mcp/remoteMcpServer.test.ts b/packages/trueforge-core/tests/core/mcp/remoteMcpServer.test.ts index 50881302f..1ecb1dd56 100644 --- a/packages/trueforge-core/tests/core/mcp/remoteMcpServer.test.ts +++ b/packages/trueforge-core/tests/core/mcp/remoteMcpServer.test.ts @@ -48,6 +48,7 @@ interface FakeConnectionState { callToolCalls: number; closes: number; queueCallToolError(error: unknown): void; + holdNextSuccessfulCall(): { started: Promise; release: () => void; fail: (error: unknown) => void }; } function toError(error: unknown): Error { @@ -58,11 +59,18 @@ function installFakeConnection( opts: { tools?: ToolSchema[]; sessionId?: string | null; connectError?: unknown } = {}, ): FakeConnectionState { const callToolErrors: unknown[] = []; + const holds: Array<{ started: () => void; wait: Promise }> = []; const state: FakeConnectionState = { connectCalls: 0, callToolCalls: 0, closes: 0, queueCallToolError: (error: unknown) => callToolErrors.push(error), + holdNextSuccessfulCall: () => { + const started = Promise.withResolvers(); + const wait = Promise.withResolvers(); + holds.push({ started: started.resolve, wait: wait.promise }); + return { started: started.promise, release: wait.resolve, fail: wait.reject }; + }, }; mockConnect.mockImplementation(() => { @@ -70,20 +78,41 @@ function installFakeConnection( if (opts.connectError) { return Promise.reject(toError(opts.connectError)); } + let closed = false; + const abortHeld: Array<(error: Error) => void> = []; return Promise.resolve({ transportType: 'streamable-http' as const, sessionId: opts.sessionId ?? null, listTools: () => Promise.resolve({ tools: opts.tools ?? [READ_TOOL, WRITE_TOOL] }), - callTool: callParams => { + callTool: async callParams => { state.callToolCalls += 1; const queued = callToolErrors.shift(); if (queued) { - return Promise.reject(toError(queued)); + throw toError(queued); } - return Promise.resolve({ content: [{ type: 'text', text: `called ${callParams.name}` }] }); + const hold = holds.shift(); + if (hold) { + hold.started(); + await Promise.race([ + hold.wait, + new Promise((_, reject) => { + abortHeld.push(reject); + }), + ]); + } + if (closed) { + throw new Error('Connection closed'); + } + return { content: [{ type: 'text', text: `called ${callParams.name}` }] }; }, close: () => { + closed = true; state.closes += 1; + const err = new Error('Connection closed'); + for (const reject of abortHeld) { + reject(err); + } + abortHeld.length = 0; return Promise.resolve(); }, }); @@ -282,6 +311,50 @@ describe('RemoteMCP + ToolSet', () => { expect(state.callToolCalls).toBe(2); }); + it('does not close the shared transport while a sibling call is in flight', async () => { + const state = installFakeConnection({ sessionId: 'sess-1' }); + const server = makeServer({}); + await server.listTools(); + + const hold = state.holdNextSuccessfulCall(); + const hanging = server.callTool({ name: 'read_thing', arguments: {} }); + await hold.started; + + state.queueCallToolError(new Error('session-expired')); + const retried = await server.callTool({ name: 'write_thing', arguments: {} }); + if (!isCallToolResponseResult(retried)) throw new Error('expected result response'); + expect(retried.result.isError).toBeFalsy(); + expect(state.connectCalls).toBe(2); + expect(state.closes).toBe(0); + + hold.release(); + const hung = await hanging; + if (!isCallToolResponseResult(hung)) throw new Error('expected result response'); + expect(hung.result.isError).toBeFalsy(); + expect(state.closes).toBe(1); + }); + + it('does not reset a newer transport when a stale sibling later expires', async () => { + const state = installFakeConnection({ sessionId: 'sess-1' }); + const server = makeServer({}); + await server.listTools(); + + const hold = state.holdNextSuccessfulCall(); + const hanging = server.callTool({ name: 'read_thing', arguments: {} }); + await hold.started; + + state.queueCallToolError(new Error('session-expired')); + const retried = await server.callTool({ name: 'write_thing', arguments: {} }); + if (!isCallToolResponseResult(retried)) throw new Error('expected result response'); + expect(state.connectCalls).toBe(2); + + hold.fail(new Error('session-expired')); + const hung = await hanging; + if (!isCallToolResponseResult(hung)) throw new Error('expected result response'); + expect(hung.result.isError).toBeFalsy(); + expect(state.connectCalls).toBe(2); + }); + it('wraps a non-auth connect failure in a server-named McpConnectionError preserving the status hint', async () => { installFakeConnection({ connectError: new McpConnectionError('upstream down', 502) }); const server = makeServer({});