fix: don't close a shared MCP transport under in-flight sibling calls

.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.

packages/trueforge-core/src/core/mcp/RemoteMCP.ts

@@ -69,6 +69,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;
@@ -105,14 +107,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<void> {
+    if (expired !== undefined && this._connection !== expired) {
+      return;
     }
-    return this._connection;
-  }
-
-  private async resetConnection(): Promise<void> {
     this.isConnected = false;
     this.sessionId = undefined;
     this.cachedTools = undefined;
@@ -121,10 +138,19 @@ export class RemoteMCP implements ToolSource {
   }
 
   private async closeAndClearConnection(): Promise<void> {
-    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<ResolveHeadersResult> {
@@ -134,7 +160,9 @@ export class RemoteMCP implements ToolSource {
     return await this.headers();
   }
 
-  private async executeWithSessionRetry<T>(operation: () => Promise<T>): Promise<ExecuteResult<T>> {
+  private async executeWithSessionRetry<T>(
+    operation: (connection: RemoteMcpConnection) => Promise<T>,
+  ): Promise<ExecuteResult<T>> {
     // 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).
@@ -147,29 +175,47 @@ export class RemoteMCP implements ToolSource {
 
   private async connectAndRun<T>(
     headers: Record<string, string>,
-    operation: () => Promise<T>,
+    operation: (connection: RemoteMcpConnection) => Promise<T>,
     canRetry: boolean,
   ): Promise<ExecuteResult<T>> {
+    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,
@@ -187,15 +233,15 @@ 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;
     }
     return { result: { tools: response.result.tools }, wasInitialized: response.wasInitialized };
   }
 
   async callTool(params: CallToolRequest['params']): Promise<CallToolResolvedResponse | AuthRequiredResponse> {
-    const response = await this.executeWithSessionRetry(() =>
+    const response = await this.executeWithSessionRetry(connection =>
       this.tracing.withRemoteMcpToolSpan(
         {
           method: 'tools/call',
@@ -207,7 +253,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;
         },
@@ -248,6 +294,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,
@@ -259,7 +306,9 @@ export class RemoteMCP implements ToolSource {
               maxResponseBytes: this.maxResponseBytes,
               signal: this.signal,
               onClose: () => {
-                this.isConnected = false;
+                if (attached.current !== undefined && this._connection === attached.current) {
+                  this.isConnected = false;
+                }
               },
               onError: error => {
                 const fields = extractErrorLogFields(error);
@@ -272,6 +321,7 @@ export class RemoteMCP implements ToolSource {
               },
             });
             span.setOutput(JSON.stringify({ transport: conn.transportType, stateful: conn.sessionId !== null }));
+            attached.current = conn;
             return conn;
           },
         );

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<void>; release: () => void; fail: (error: unknown) => void };
 }
 
 function toError(error: unknown): Error {
@@ -58,32 +59,60 @@ function installFakeConnection(
   opts: { tools?: ToolSchema[]; sessionId?: string | null; connectError?: unknown } = {},
 ): FakeConnectionState {
   const callToolErrors: unknown[] = [];
+  const holds: Array<{ started: () => void; wait: Promise<void> }> = [];
   const state: FakeConnectionState = {
     connectCalls: 0,
     callToolCalls: 0,
     closes: 0,
     queueCallToolError: (error: unknown) => callToolErrors.push(error),
+    holdNextSuccessfulCall: () => {
+      const started = Promise.withResolvers<void>();
+      const wait = Promise.withResolvers<void>();
+      holds.push({ started: started.resolve, wait: wait.promise });
+      return { started: started.promise, release: wait.resolve, fail: wait.reject };
+    },
   };
 
   mockConnect.mockImplementation(() => {
     state.connectCalls += 1;
     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<never>((_, 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({});