diff --git a/src/pool/connection-pool.ts b/src/pool/connection-pool.ts index 1033411..50e8b8a 100644 --- a/src/pool/connection-pool.ts +++ b/src/pool/connection-pool.ts @@ -287,6 +287,13 @@ export class ConnectionPool extends EventEmitter { return false; } + // Transport must be in CONNECTED state. A transport that reached ERROR + // (SSE reconnect exhausted, HTTP request failure, process exit) must not + // be reused — otherwise the pool keeps handing out a dead connection. + if (!connection.transport.isConnected()) { + return false; + } + // For stdio transport, check if the process is still running if (connection.transport.getType() === 'stdio') { const stdioTransport = connection.transport as { @@ -396,17 +403,68 @@ export class ConnectionPool extends EventEmitter { /** * Find an idle connection in the pool * - * @returns Idle connection or undefined if none available + * Only healthy idle connections are returned. If an idle connection is no + * longer healthy (e.g. its transport reached ERROR), it is removed from the + * pool so the next acquire creates a fresh connection instead of reusing a + * dead one or letting it occupy a slot up to maxConnections. + * + * @returns Healthy idle connection or undefined if none available */ private findIdleConnection(): Connection | undefined { + let unhealthy: Connection | undefined; + for (const connection of this.connections.values()) { - if (isIdle(connection)) { + if (!isIdle(connection)) { + continue; + } + + if (this.isConnectionHealthy(connection)) { + // If we passed an unhealthy idle connection on the way, remove it first + if (unhealthy) { + void this.invalidateConnection(unhealthy, new Error('Connection is unhealthy')); + } return connection; } + + if (!unhealthy) { + unhealthy = connection; + } } + + if (unhealthy) { + void this.invalidateConnection(unhealthy, new Error('Connection is unhealthy')); + } + return undefined; } + /** + * Remove a connection from the pool. + * + * Idempotent: drops the connection from the pool immediately (so a dead + * connection does not occupy a maxConnections slot), closes its transport, + * and awakes any queued requests so they get a fresh connection. + */ + private async invalidateConnection(connection: Connection, error: Error): Promise { + if (!this.connections.has(connection.id)) { + return; + } + + // Drop from the pool synchronously before awaiting the transport close so + // concurrent health checks / request failures cannot double-remove it. + this.connections.delete(connection.id); + this.emit('connectionFailed', connection.id, error); + + try { + await connection.transport.close(); + this.emit('connectionClosed', connection.id); + } catch (closeError) { + this.emit('error', closeError); + } + + this.processQueue(); + } + /** * Create a new connection * @@ -468,11 +526,15 @@ export class ConnectionPool extends EventEmitter { const connection = createConnection(id, transport); // Attach the listener before initialization so an early process exit does not - // surface as an unhandled EventEmitter error. + // surface as an unhandled EventEmitter error. A transport error (SSE reconnect + // exhausted, HTTP request failure, process exit) makes the connection unusable: + // invalidate it in the pool so it is not reused or counted against maxConnections. transport.on('error', (error: unknown) => { const errorMessage = error instanceof Error ? error.message : String(error); log.warn(`[${this.service.name}] Transport error: ${errorMessage}`); this.emit('error', error); + const cause = error instanceof Error ? error : new Error(String(error)); + void this.invalidateConnection(connection, cause).catch(() => {}); }); await this.withTimeout( diff --git a/src/transport/stdio.ts b/src/transport/stdio.ts index 18f7d7b..2488529 100644 --- a/src/transport/stdio.ts +++ b/src/transport/stdio.ts @@ -3,7 +3,7 @@ */ import { ChildProcess, spawn } from 'child_process'; -import { BaseTransport, TransportError } from './base.js'; +import { BaseTransport, TransportError, TransportState } from './base.js'; import type { JsonRpcMessage } from '../types/jsonrpc.js'; import type { TransportType } from '../types/service.js'; import * as log from '../utils/logger.js'; @@ -246,16 +246,22 @@ export class StdioTransport extends BaseTransport { /** * Handle process exit + * + * Any exit not caused by our own close() means the backend process is gone and + * the connection is dead — even a clean code-0 exit or an external signal like + * SIGTERM/SIGKILL. Mark the transport ERROR so the connection pool stops reusing + * it. During an intentional close() the state is CLOSING/CLOSED, so skip then. */ private handleProcessExit(code: number | null, signal: NodeJS.Signals | null): void { const exitInfo = signal ? `signal ${signal}` : `code ${code}`; - // Only treat non-zero exit codes as errors - // Signals like SIGTERM are normal termination - if (code !== null && code !== 0) { + if (this.state !== TransportState.CLOSING && this.state !== TransportState.CLOSED) { this.handleError(new TransportError(`Process exited with ${exitInfo}`, 'PROCESS_EXITED')); } + // Close the receive side and settle all waiting receivers + this.handleStreamEnd(); + // Reject all waiting receivers const error = new TransportError(`Process exited with ${exitInfo}`, 'PROCESS_EXITED'); while (this.rejectQueue.length > 0) { diff --git a/tests/integration/backend-recovery.test.ts b/tests/integration/backend-recovery.test.ts new file mode 100644 index 0000000..85c268a --- /dev/null +++ b/tests/integration/backend-recovery.test.ts @@ -0,0 +1,198 @@ +/** + * Integration tests: backend death does not permanently break tool discovery. + * + * These tests drive the REAL ConnectionPool with a REAL StdioTransport that + * spawns a real backend process (the mock-stdio-mcp fixture). They cover the + * end-to-end recovery chain the unit tests mock away: + * + * backend process dies → transport reaches ERROR → isConnected() flips false + * → pool invalidates the connection → next acquire creates a fresh connection + * → tools/list succeeds again + * + * This is the actual user-facing scenario from MERC-3: "long-running server + * mode loses tools after the backend connection drops". + */ + +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import * as path from 'path'; +import { ConnectionPool } from '../../src/pool/connection-pool.js'; +import type { Connection } from '../../src/pool/connection.js'; +import type { JsonRpcMessage } from '../../src/types/jsonrpc.js'; +import type { ServiceDefinition, ConnectionPoolConfig } from '../../src/types/service.js'; + +const FIXTURE_PATH = path.resolve(__dirname, 'fixtures/mock-stdio-mcp.cjs'); + +const EXPECTED_TOOL_COUNT = 2; + +/** + * Send a tools/list request over a connection and return the tools array. + * Throws if the transport returns an error or the stream ends first — which is + * exactly the TRANSPORT_ERROR-equivalent failure the fix is meant to prevent. + */ +async function listTools(connection: Connection): Promise { + const id = `tools-list-${Date.now()}-${Math.random().toString(36).slice(2)}`; + await connection.transport.send({ + jsonrpc: '2.0', + id, + method: 'tools/list', + params: {}, + } as JsonRpcMessage); + + const iterator = connection.transport.receive(); + try { + for (;;) { + const next = await iterator.next(); + if (next.done || !next.value) { + throw new Error('transport stream ended before tools/list response'); + } + const message = next.value as unknown as Record; + if ('id' in message && String(message['id']) === String(id)) { + const error = message['error'] as { code?: number; message?: string } | undefined; + if (error) { + throw new Error(`tools/list returned error code ${error.code}: ${error.message}`); + } + const result = message['result'] as { tools?: unknown[] } | undefined; + return result?.tools ?? []; + } + } + } finally { + await iterator.return?.(undefined as unknown as JsonRpcMessage); + } +} + +function waitForExit(pool: ConnectionPool, timeoutMs = 5000): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + pool.off('connectionFailed', onFailed); + reject(new Error('Timed out waiting for backend process exit / pool invalidation')); + }, timeoutMs); + const onFailed = () => { + clearTimeout(timer); + pool.off('connectionFailed', onFailed); + resolve(); + }; + pool.on('connectionFailed', onFailed); + }); +} + +describe('Backend death recovery (integration)', () => { + let pool: ConnectionPool; + let service: ServiceDefinition; + let poolConfig: ConnectionPoolConfig; + + beforeEach(() => { + // Swallow pool-level error re-emits (transport errors surface here too). + vi.spyOn(console, 'error').mockImplementation(() => {}); + + service = { + name: 'mock-backend', + enabled: true, + tags: [], + transport: 'stdio', + command: process.execPath, + args: [FIXTURE_PATH], + env: { ONEMCP_FIXTURE_EXIT_AFTER_LIST: '1' }, + connectionPool: { + maxConnections: 2, + idleTimeout: 60000, + connectionTimeout: 10000, + }, + }; + + poolConfig = { + maxConnections: 2, + idleTimeout: 60000, + connectionTimeout: 10000, + }; + + pool = new ConnectionPool(service, poolConfig); + pool.on('error', () => {}); + }); + + afterEach(async () => { + if (pool) { + await pool.closeAll().catch(() => {}); + } + vi.restoreAllMocks(); + }); + + it('recovers tools/list after the stdio backend exits with code 0', async () => { + // First acquire: pool spawns backend A, runs the MCP handshake, returns conn1. + const conn1 = await pool.acquire(); + const tools1 = await listTools(conn1); + expect(tools1).toHaveLength(EXPECTED_TOOL_COUNT); + + // The fixture exits with code 0 right after responding — wait for the pool + // to invalidate the now-dead connection. + await waitForExit(pool); + + // The dead connection must no longer occupy a pool slot. + expect(pool.getStats().total).toBe(0); + + // Second acquire must return a FRESH connection (a new backend process), + // not the dead one — and tools/list must succeed again, not TRANSPORT_ERROR. + const conn2 = await pool.acquire(); + expect(conn2.id).not.toBe(conn1.id); + expect(pool.isConnectionHealthy(conn2)).toBe(true); + + const tools2 = await listTools(conn2); + expect(tools2).toHaveLength(EXPECTED_TOOL_COUNT); + }, 30000); + + it('recovers tools/list after the stdio backend is killed by a signal', async () => { + // Use a backend that stays alive so we control the death moment precisely. + const liveService: ServiceDefinition = { + ...service, + env: {}, // no EXIT_AFTER_LIST -> stays alive + }; + const livePool = new ConnectionPool(liveService, poolConfig); + livePool.on('error', () => {}); + try { + const conn1 = await livePool.acquire(); + const tools1 = await listTools(conn1); + expect(tools1).toHaveLength(EXPECTED_TOOL_COUNT); + livePool.release(conn1); + + // Kill the backend child process externally (simulates OOM / signal kill). + const proc = (conn1.transport as unknown as { process: { kill: (s: string) => void } }) + .process; + const exitPromise = waitForExit(livePool); + proc.kill('SIGKILL'); + + await exitPromise; + expect(livePool.getStats().total).toBe(0); + + // Next acquire creates a fresh connection; tools/list succeeds again. + const conn2 = await livePool.acquire(); + expect(conn2.id).not.toBe(conn1.id); + const tools2 = await listTools(conn2); + expect(tools2).toHaveLength(EXPECTED_TOOL_COUNT); + } finally { + await livePool.closeAll().catch(() => {}); + } + }, 30000); + + it('does not reuse a dead idle connection when maxConnections would be exhausted', async () => { + // With maxConnections=1, a dead idle connection must be evicted (not held) + // so the next acquire can create a replacement rather than hanging. + const singleConfig: ConnectionPoolConfig = { ...poolConfig, maxConnections: 1 }; + const singlePool = new ConnectionPool(service, singleConfig); + singlePool.on('error', () => {}); + try { + const conn1 = await singlePool.acquire(); + await listTools(conn1); + singlePool.release(conn1); + + await waitForExit(singlePool); + expect(singlePool.getStats().total).toBe(0); + + // Would hang / time out if the dead connection kept its slot. + const conn2 = await singlePool.acquire(); + const tools2 = await listTools(conn2); + expect(tools2).toHaveLength(EXPECTED_TOOL_COUNT); + expect(conn2.id).not.toBe(conn1.id); + } finally { + await singlePool.closeAll().catch(() => {}); + } + }, 30000); +}); diff --git a/tests/integration/fixtures/mock-stdio-mcp.cjs b/tests/integration/fixtures/mock-stdio-mcp.cjs new file mode 100644 index 0000000..a1e722b --- /dev/null +++ b/tests/integration/fixtures/mock-stdio-mcp.cjs @@ -0,0 +1,100 @@ +/** + * Minimal stdio MCP backend used by integration tests. + * + * Speaks NDJSON over stdin/stdout (OneMCP's stdio transport framing): + * - initialize -> success result + * - notifications/initialized -> no response (notification) + * - tools/list -> N tools + * - tools/call -> success result + * + * When ONEMCP_FIXTURE_EXIT_AFTER_LIST=1, the process exits with code 0 a short + * time after responding to tools/list, simulating a backend that dies mid-run. + * The small delay ensures the response is flushed to the reader before exit. + */ +'use strict'; + +const readline = require('readline'); + +const TOOLS = [ + { + name: 'echo', + description: 'Echo back the input text', + inputSchema: { + type: 'object', + properties: { text: { type: 'string' } }, + required: ['text'], + }, + }, + { + name: 'add', + description: 'Add two numbers', + inputSchema: { + type: 'object', + properties: { a: { type: 'number' }, b: { type: 'number' } }, + required: ['a', 'b'], + }, + }, +]; + +const EXIT_AFTER_LIST = process.env.ONEMCP_FIXTURE_EXIT_AFTER_LIST === '1'; + +function send(message) { + process.stdout.write(JSON.stringify(message) + '\n'); +} + +const rl = readline.createInterface({ input: process.stdin, crlfDelay: Infinity }); + +rl.on('line', (line) => { + let request; + try { + request = JSON.parse(line); + } catch { + return; // ignore malformed lines + } + if (!request || !request.method) { + return; + } + + switch (request.method) { + case 'initialize': + send({ + jsonrpc: '2.0', + id: request.id, + result: { + protocolVersion: '2024-11-05', + capabilities: {}, + serverInfo: { name: 'mock-stdio-mcp', version: '1.0.0' }, + }, + }); + break; + case 'notifications/initialized': + // Notification — no response. + break; + case 'tools/list': + send({ jsonrpc: '2.0', id: request.id, result: { tools: TOOLS } }); + if (EXIT_AFTER_LIST) { + rl.close(); + setTimeout(() => process.exit(0), 50); + } + break; + case 'tools/call': + send({ + jsonrpc: '2.0', + id: request.id, + result: { content: [{ type: 'text', text: 'ok' }] }, + }); + break; + default: + if (request.id !== undefined && request.id !== null) { + send({ + jsonrpc: '2.0', + id: request.id, + error: { code: -32601, message: 'Method not found' }, + }); + } + break; + } +}); + +// Keep stderr quiet so it doesn't interfere with the transport. +process.stderr.on('error', () => {}); diff --git a/tests/unit/pool/connection-pool.test.ts b/tests/unit/pool/connection-pool.test.ts index f19fba2..0f2005c 100644 --- a/tests/unit/pool/connection-pool.test.ts +++ b/tests/unit/pool/connection-pool.test.ts @@ -3,6 +3,7 @@ */ import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { EventEmitter } from 'events'; import { ConnectionPool, ConnectionPoolError } from '../../../src/pool/connection-pool.js'; import type { ServiceDefinition, ConnectionPoolConfig } from '../../../src/types/service.js'; import { StdioTransport } from '../../../src/transport/stdio.js'; @@ -780,5 +781,166 @@ describe('ConnectionPool', () => { await httpPool.closeAll(); }, 15000); + + it('should treat a transport in ERROR state as unhealthy', async () => { + const connection = await pool.acquire(); + + // Simulate transport reaching ERROR (e.g. SSE reconnect exhausted, + // HTTP request failure, process exit) — isConnected() flips to false. + (connection.transport as { isConnected: () => boolean }).isConnected = () => false; + + expect(pool.isConnectionHealthy(connection)).toBe(false); + }, 15000); + + it('should not reuse an unhealthy idle connection and should create a new one', async () => { + const conn1 = await pool.acquire(); + // Simulate the backend connection dying while idle in the pool. + (conn1.transport as { isConnected: () => boolean }).isConnected = () => false; + pool.release(conn1); + + const conn2 = await pool.acquire(); + + // The dead connection must not be handed out again — a fresh one is created. + expect(conn2.id).not.toBe(conn1.id); + expect(pool.isConnectionHealthy(conn2)).toBe(true); + + // The unhealthy connection should no longer occupy a pool slot. + const stats = pool.getStats(); + expect(stats.total).toBe(1); + + await pool.closeAll(); + }, 15000); + + it('should remove an unhealthy idle connection from the pool', async () => { + const conn1 = await pool.acquire(); + (conn1.transport as { isConnected: () => boolean }).isConnected = () => false; + pool.release(conn1); + + const statsBefore = pool.getStats(); + expect(statsBefore.total).toBe(1); + expect(statsBefore.idle).toBe(1); + + // Trigger a path that scans idle connections (acquire). + await pool.acquire(); + + const statsAfter = pool.getStats(); + expect(statsAfter.total).toBe(1); // one fresh connection, the dead one is gone + expect(statsAfter.idle).toBe(0); // it's busy (acquired) and healthy + }, 15000); + }); + + describe('transport error invalidates pool connection', () => { + // The previous tests mock isConnected() directly. These tests exercise the + // real "transport emits 'error' → pool invalidates → next acquire is fresh" + // path that the fix added in createConnectionWithTimeout, using a transport + // whose `on/emit` behave like a real EventEmitter so listeners actually fire. + beforeEach(() => { + vi.mocked(StdioTransport).mockImplementation(function (this: any) { + const ee = new EventEmitter(); + this.on = (event: string, listener: (...args: unknown[]) => void) => { + ee.on(event, listener); + return this; + }; + this.emit = (event: string, ...args: unknown[]) => ee.emit(event, ...args); + this.send = vi.fn().mockResolvedValue(undefined); + this.receive = vi.fn().mockReturnValue({ + async next() { + return { value: { jsonrpc: '2.0', id: 1, result: {} }, done: false }; + }, + async return(value?: any) { + return { value, done: true }; + }, + async throw(error?: any) { + throw error; + }, + [Symbol.asyncIterator]() { + return this; + }, + }); + this.close = vi.fn().mockResolvedValue(undefined); + this.getType = vi.fn().mockReturnValue('stdio'); + this.isConnected = vi.fn().mockReturnValue(true); + this.process = { killed: false, exitCode: null }; + return this; + }); + }); + + it('should emit connectionFailed and remove the connection when a pooled transport emits error', async () => { + const conn1 = await pool.acquire(); + pool.release(conn1); // back to idle in the pool + + const failedSpy = vi.fn(); + const closedSpy = vi.fn(); + pool.on('connectionFailed', failedSpy); + pool.on('connectionClosed', closedSpy); + // The pool re-emits transport errors on itself; swallow them so Node's + // EventEmitter does not throw on an unhandled 'error' event. + pool.on('error', () => {}); + + const error = new Error('backend process died'); + // Fire the transport-level 'error' event — the "instant invalidation" path. + (conn1.transport as unknown as { emit: (e: string, ...a: unknown[]) => void }).emit( + 'error', + error + ); + + // invalidateConnection awaits transport.close() before emitting connectionClosed. + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + expect(failedSpy).toHaveBeenCalledWith(conn1.id, error); + expect(closedSpy).toHaveBeenCalledWith(conn1.id); + + // The dead connection no longer occupies a maxConnections slot. + const stats = pool.getStats(); + expect(stats.total).toBe(0); + }, 15000); + + it('should hand out a fresh connection on the next acquire after a transport error', async () => { + const conn1 = await pool.acquire(); + pool.release(conn1); + + pool.on('error', () => {}); // swallow pool-level error re-emit + + (conn1.transport as unknown as { emit: (e: string, ...a: unknown[]) => void }).emit( + 'error', + new Error('backend died') + ); + // Let invalidateConnection finish. + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + const conn2 = await pool.acquire(); + + // A fresh connection is created, not the dead one reused. + expect(conn2.id).not.toBe(conn1.id); + expect(pool.isConnectionHealthy(conn2)).toBe(true); + + await pool.closeAll(); + }, 15000); + + it('should not double-invalidate if both a transport error and a failed request occur', async () => { + const conn1 = await pool.acquire(); + pool.release(conn1); + + const failedSpy = vi.fn(); + pool.on('connectionFailed', failedSpy); + pool.on('error', () => {}); // swallow pool-level error re-emit + + (conn1.transport as unknown as { emit: (e: string, ...a: unknown[]) => void }).emit( + 'error', + new Error('transport error') + ); + await new Promise((resolve) => setImmediate(resolve)); + await new Promise((resolve) => setImmediate(resolve)); + + // A later markConnectionFailed on the same (already removed) connection is a no-op. + await pool.markConnectionFailed(conn1, new Error('request failed')); + + expect(failedSpy).toHaveBeenCalledTimes(1); + expect(pool.getStats().total).toBe(0); + + await pool.closeAll(); + }, 15000); }); }); diff --git a/tests/unit/transport/stdio.test.ts b/tests/unit/transport/stdio.test.ts index e7038e2..c21ed35 100644 --- a/tests/unit/transport/stdio.test.ts +++ b/tests/unit/transport/stdio.test.ts @@ -386,17 +386,20 @@ describe('StdioTransport', () => { }); }); - it('should handle process exit with code 0', async () => { - const exitPromise = new Promise((resolve) => { - transport.on('error', () => { - // Should not emit error for clean exit - throw new Error('Should not emit error for exit code 0'); + it('should handle process exit with code 0 as an error', async () => { + // A code-0 exit the pool did not initiate still means the backend process + // is gone — the transport must go ERROR so the connection is not reused. + const errorPromise = new Promise((resolve) => { + transport.on('error', (error) => { + resolve(error); }); - setTimeout(resolve, 100); }); mockProcess.emit('exit', 0, null); - await exitPromise; + + const error = await errorPromise; + expect(error).toBeInstanceOf(TransportError); + expect(error.message).toContain('code 0'); }); it('should handle process exit with non-zero code', async () => { @@ -413,22 +416,20 @@ describe('StdioTransport', () => { expect(error.message).toContain('code 1'); }); - it('should handle process exit with signal', async () => { - // SIGTERM is a normal termination signal, should not emit error - const errorPromise = new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - resolve(); // No error is expected - }, 100); - - transport.on('error', (_error) => { - clearTimeout(timeout); - reject(new Error('Should not emit error for SIGTERM')); + it('should handle process exit with signal as an error', async () => { + // A signal (SIGTERM/SIGKILL) that was not caused by our own close() means + // the process died externally — mark ERROR so the pool rebuilds it. + const errorPromise = new Promise((resolve) => { + transport.on('error', (error) => { + resolve(error); }); }); mockProcess.emit('exit', null, 'SIGTERM'); - await errorPromise; + const error = await errorPromise; + expect(error).toBeInstanceOf(TransportError); + expect(error.message).toContain('signal SIGTERM'); }); it('should handle process errors', async () => {