diff --git a/src/core/client/clientManager.test.ts b/src/core/client/clientManager.test.ts index a43c784f..af4bc755 100644 --- a/src/core/client/clientManager.test.ts +++ b/src/core/client/clientManager.test.ts @@ -1,4 +1,5 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; @@ -332,6 +333,140 @@ describe('ClientManager (Integration)', () => { }); }); + describe('session loss recovery', () => { + it('reconnects with a fresh (session-less) transport when the backend reports a lost session', async () => { + vi.useRealTimers(); + + const originalTransport = { + _url: new URL('https://example.com/mcp'), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(originalTransport, StreamableHTTPClientTransport.prototype); + + const freshTransport = { + _url: new URL('https://example.com/mcp'), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(freshTransport, StreamableHTTPClientTransport.prototype); + + (mockClient.connect as unknown as MockInstance).mockResolvedValue(undefined); + (mockClient.getServerVersion as unknown as MockInstance).mockResolvedValue({ + name: 'test-server', + version: '1.0.0', + }); + + const recreateForSessionLoss = vi + .spyOn((clientManager as any).transportRecreator, 'recreateForSessionLoss') + .mockReturnValue(freshTransport); + + await clientManager.createSingleClient('session-loss-client', originalTransport); + expect(clientManager.getTransport('session-loss-client')).toBe(originalTransport); + + const registeredClient = clientManager.getClient('session-loss-client').client; + registeredClient.onerror?.( + new Error( + 'Streamable HTTP error: Error POSTing to endpoint: ' + + '{"jsonrpc":"2.0","id":"server-error","error":{"code":-32600,"message":"Session not found"}}', + ), + ); + + // onerror kicks off recovery fire-and-forget — wait for it to settle + await vi.waitFor(() => { + expect(clientManager.getTransport('session-loss-client')).toBe(freshTransport); + }); + + expect(recreateForSessionLoss).toHaveBeenCalledWith(originalTransport, 'session-loss-client'); + expect(clientManager.getClient('session-loss-client').status).toBe(ClientStatus.Connected); + }); + + it('also recovers the "Could not find session ID" wording used by SSE-transport backends', async () => { + vi.useRealTimers(); + + const originalTransport = { + _url: new URL('https://example.com/sse'), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(originalTransport, SSEClientTransport.prototype); + + (mockClient.connect as unknown as MockInstance).mockResolvedValue(undefined); + (mockClient.getServerVersion as unknown as MockInstance).mockResolvedValue({ + name: 'test-server', + version: '1.0.0', + }); + + // Not mocked — exercises the real SSEClientTransport recreation path. + const recreateForSessionLoss = vi.spyOn((clientManager as any).transportRecreator, 'recreateForSessionLoss'); + + await clientManager.createSingleClient('sse-session-loss-client', originalTransport); + const registeredClient = clientManager.getClient('sse-session-loss-client').client; + registeredClient.onerror?.(new Error("Error POSTing to endpoint (HTTP 404): Could not find session ID 'abc'")); + + await vi.waitFor(() => { + expect(clientManager.getClient('sse-session-loss-client').status).toBe(ClientStatus.Connected); + }); + + expect(recreateForSessionLoss).toHaveBeenCalledWith(originalTransport, 'sse-session-loss-client'); + const recreatedTransport = clientManager.getTransport('sse-session-loss-client'); + expect(recreatedTransport).not.toBe(originalTransport); + expect(recreatedTransport).toBeInstanceOf(SSEClientTransport); + }); + + it('ignores unrelated client errors', async () => { + vi.useRealTimers(); + + const originalTransport = { + _url: new URL('https://example.com/mcp'), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(originalTransport, StreamableHTTPClientTransport.prototype); + + (mockClient.connect as unknown as MockInstance).mockResolvedValue(undefined); + (mockClient.getServerVersion as unknown as MockInstance).mockResolvedValue({ + name: 'test-server', + version: '1.0.0', + }); + + const recreateForSessionLoss = vi.spyOn((clientManager as any).transportRecreator, 'recreateForSessionLoss'); + + await clientManager.createSingleClient('unrelated-error-client', originalTransport); + const registeredClient = clientManager.getClient('unrelated-error-client').client; + registeredClient.onerror?.(new Error('ECONNRESET')); + + await Promise.resolve(); + expect(recreateForSessionLoss).not.toHaveBeenCalled(); + expect(clientManager.getTransport('unrelated-error-client')).toBe(originalTransport); + }); + + it('does not let a recreation failure escape onerror for a non-HTTP/SSE transport', async () => { + vi.useRealTimers(); + + // A stdio-style transport: TransportRecreator only supports HTTP/SSE, so + // recreateForSessionLoss throws for this one — that throw must be caught, + // not propagated out of the onerror callback. + const stdioTransport = { + name: 'stdio', + start: vi.fn(), + send: vi.fn(), + close: vi.fn(), + } as unknown as AuthProviderTransport; + + (mockClient.connect as unknown as MockInstance).mockResolvedValue(undefined); + (mockClient.getServerVersion as unknown as MockInstance).mockResolvedValue({ + name: 'test-server', + version: '1.0.0', + }); + + await clientManager.createSingleClient('stdio-client', stdioTransport); + const registeredClient = clientManager.getClient('stdio-client').client; + + expect(() => registeredClient.onerror?.(new Error('Session not found'))).not.toThrow(); + + await Promise.resolve(); + expect(logger.error).toHaveBeenCalledWith(expect.stringContaining('Cannot recover stdio-client')); + expect(clientManager.getTransport('stdio-client')).toBe(stdioTransport); + }); + }); + describe('getClient and getClients', () => { beforeEach(async () => { (mockClient.connect as unknown as MockInstance).mockResolvedValue(undefined); diff --git a/src/core/client/clientManager.ts b/src/core/client/clientManager.ts index 90e75735..9ed3388d 100644 --- a/src/core/client/clientManager.ts +++ b/src/core/client/clientManager.ts @@ -33,6 +33,20 @@ export const enum ClientManagerEvent { BackendSupervisionStateChanged = 'backend-supervision-state-changed', } +// Matches the two message shapes downstream servers are observed to send back +// when a Streamable HTTP / SSE session ID they issued no longer exists on +// their side (e.g. the backend process restarted and its in-memory session +// store was wiped): the MCP spec's own "Session not found" wording, and the +// free-form "Could not find session ID '...'" text some server SDKs use. +const SESSION_LOST_PATTERN = /session (?:id )?not found|could not find session/i; + +function isSessionLostError(error: unknown): boolean { + if (!(error instanceof Error)) { + return false; + } + return SESSION_LOST_PATTERN.test(error.message); +} + type StdioSupervisionMetadata = NonNullable; export class ClientManager extends EventEmitter { @@ -159,9 +173,55 @@ export class ClientManager extends EventEmitter { client.onerror = (error) => { logger.error(`Client ${name} error: ${error}`); + + if (isSessionLostError(error)) { + this.recoverFromSessionLoss(name, client); + } }; } + /** + * Reconnects a backend whose Streamable HTTP / SSE session was invalidated + * server-side (typically because the backend process restarted and lost its + * in-memory session store). Without this, the client keeps reusing the dead + * session ID on every subsequent request and the backend never recovers + * until the whole 1MCP process is restarted. + */ + private recoverFromSessionLoss(name: string, erroredClient: Client): void { + if (this.isShuttingDown) { + return; + } + + // The client that errored may already have been superseded by a newer + // connection attempt (e.g. a concurrent recovery, or the server coming + // back mid-retry) — only recover if it's still the one on record. + const clientInfo = this.outboundConns.get(name); + if (!clientInfo || clientInfo.client !== erroredClient) { + return; + } + + logger.warn(`Session for ${name} was lost (backend likely restarted) — reconnecting with a fresh session`); + + const staleTransport = this.transports[name] ?? clientInfo.transport; + let freshTransport: AuthProviderTransport; + try { + freshTransport = this.transportRecreator.recreateForSessionLoss(staleTransport, name); + } catch (error) { + // recreateForSessionLoss only supports HTTP/SSE transports; a session-loss-shaped + // error message from some other transport kind must not escape client.onerror. + logger.error( + `Cannot recover ${name} from session loss: ${error instanceof Error ? error.message : String(error)}`, + ); + return; + } + + void this.createSingleClient(name, freshTransport).catch((error) => { + logger.error( + `Failed to recover ${name} after session loss: ${error instanceof Error ? error.message : String(error)}`, + ); + }); + } + /** * Create multiple MCP clients in parallel with controlled concurrency * diff --git a/src/core/client/transportRecreator.test.ts b/src/core/client/transportRecreator.test.ts index c6354fff..a5c301ad 100644 --- a/src/core/client/transportRecreator.test.ts +++ b/src/core/client/transportRecreator.test.ts @@ -60,6 +60,35 @@ describe('TransportRecreator', () => { expect((newTransport as any)._sessionId).toBe('existing-session'); }); + it('should drop the session ID when preserveSessionId is false', () => { + const originalTransport = { + _url: new URL('https://example.com/mcp'), + _sessionId: 'stale-session', + oauthProvider: { token: 'test-token' }, + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(originalTransport, StreamableHTTPClientTransport.prototype); + + const newTransport = transportRecreator.recreateHttpTransport(originalTransport, 'test-server', { + preserveSessionId: false, + }); + + expect((newTransport as any)._sessionId).toBeUndefined(); + }); + + it('recreateForSessionLoss should drop the session ID', () => { + const originalTransport = { + _url: new URL('https://example.com/mcp'), + _sessionId: 'stale-session', + oauthProvider: { token: 'test-token' }, + } as unknown as AuthProviderTransport; + Object.setPrototypeOf(originalTransport, StreamableHTTPClientTransport.prototype); + + const newTransport = transportRecreator.recreateForSessionLoss(originalTransport, 'test-server'); + + expect(newTransport).not.toBe(originalTransport); + expect((newTransport as any)._sessionId).toBeUndefined(); + }); + it('should preserve oauthProvider configuration', () => { const oauthProvider = { token: 'oauth-token-123', diff --git a/src/core/client/transportRecreator.ts b/src/core/client/transportRecreator.ts index a199e1e5..73113a9c 100644 --- a/src/core/client/transportRecreator.ts +++ b/src/core/client/transportRecreator.ts @@ -5,6 +5,20 @@ import { AuthProviderTransport } from '@src/core/types/index.js'; import type { TransportRecreationState } from './transportRecreationState.js'; +export interface RecreateHttpTransportOptions { + /** + * Whether to carry the existing `sessionId` over to the new transport. + * + * Defaults to `true`, which is correct for OAuth retries (the session itself + * is still valid; only auth needs refreshing). Callers recovering from a + * server-invalidated session (e.g. the backend restarted and lost its + * in-memory session store) must pass `false` so the new transport performs + * a full `initialize` handshake and is issued a fresh session ID, instead of + * immediately failing again with the same stale one. + */ + preserveSessionId?: boolean; +} + export class TransportRecreator { public recreateForRetry(transport: AuthProviderTransport, serverName?: string): AuthProviderTransport { if (this.isHttpTransport(transport)) { @@ -14,12 +28,26 @@ export class TransportRecreator { return transport; } - public recreateHttpTransport(transport: AuthProviderTransport, serverName?: string): AuthProviderTransport { + /** + * Recreates a transport whose backend session was lost (server restarted or + * otherwise invalidated the session ID). Unlike {@link recreateForRetry}, + * this never carries the old session ID forward. + */ + public recreateForSessionLoss(transport: AuthProviderTransport, serverName?: string): AuthProviderTransport { + return this.recreateHttpTransport(transport, serverName, { preserveSessionId: false }); + } + + public recreateHttpTransport( + transport: AuthProviderTransport, + serverName?: string, + options?: RecreateHttpTransportOptions, + ): AuthProviderTransport { if (!this.isHttpTransport(transport)) { const name = serverName ? `Transport for ${serverName}` : 'Transport'; throw new Error(`${name} does not support OAuth (requires HTTP or SSE transport)`); } + const preserveSessionId = options?.preserveSessionId ?? true; const state = transport as unknown as TransportRecreationState; const authTransport = transport as AuthProviderTransport; const oauthProvider = authTransport.oauthProvider; @@ -31,7 +59,7 @@ export class TransportRecreator { requestInit: state._requestInit, fetch: state._fetch, reconnectionOptions: state._reconnectionOptions, - sessionId: state._sessionId, + sessionId: preserveSessionId ? state._sessionId : undefined, }) as AuthProviderTransport) : (new SSEClientTransport(state._url, { authProvider: oauthProvider,