From 13c5b9d5a8c2258b9085e86f47ed9177c678eed3 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 9 Jul 2026 11:38:29 +0000 Subject: [PATCH 1/4] test: pin connection lifecycle behavior and document teardown gaps --- .../client/test/client/sessionResume.test.ts | 109 ++++++ .../test/shared/connectionLifecycle.test.ts | 342 ++++++++++++++++++ 2 files changed, 451 insertions(+) create mode 100644 packages/client/test/client/sessionResume.test.ts create mode 100644 packages/core-internal/test/shared/connectionLifecycle.test.ts diff --git a/packages/client/test/client/sessionResume.test.ts b/packages/client/test/client/sessionResume.test.ts new file mode 100644 index 0000000000..9ff1b9bde3 --- /dev/null +++ b/packages/client/test/client/sessionResume.test.ts @@ -0,0 +1,109 @@ +/** + * Session resumption across a client close(): reconnecting with a transport + * that carries a sessionId (and the protocol version persisted alongside it) + * must push that version onto the transport so HTTP requests carry the + * required mcp-protocol-version header. + */ +import type { JSONRPCMessage } from '@modelcontextprotocol/core-internal'; +import { isJSONRPCRequest } from '@modelcontextprotocol/core-internal'; +import { describe, expect, test } from 'vitest'; + +import { Client } from '../../src/client/client'; + +const LEGACY = '2025-11-25'; + +/** + * A sessionful legacy server transport: answers initialize like a 2025 + * server and assigns a session id with the initialize response, mirroring + * Streamable HTTP's mcp-session-id assignment. + */ +class SessionfulTransport { + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: JSONRPCMessage) => void; + sessionId?: string; + protocolVersion?: string; + + sent: JSONRPCMessage[] = []; + setProtocolVersionCalls: string[] = []; + + constructor(carried?: { sessionId: string; protocolVersion: string }) { + this.sessionId = carried?.sessionId; + this.protocolVersion = carried?.protocolVersion; + } + + async start(): Promise {} + + async send(message: JSONRPCMessage): Promise { + this.sent.push(message); + if (isJSONRPCRequest(message) && message.method === 'initialize') { + this.sessionId = 'session-123'; + const reply: JSONRPCMessage = { + jsonrpc: '2.0', + id: message.id, + result: { + protocolVersion: LEGACY, + capabilities: {}, + serverInfo: { name: 'sessionful-legacy-server', version: '1.0.0' } + } + }; + queueMicrotask(() => this.onmessage?.(reply)); + } + } + + async close(): Promise { + this.onclose?.(); + } + + setProtocolVersion(version: string): void { + this.setProtocolVersionCalls.push(version); + this.protocolVersion = version; + } +} + +describe('r4: session resume after close()', () => { + // close() wipes the instance's negotiated state, so the resume branches + // fall back to the version the resuming transport itself carries + // (Transport.protocolVersion) and re-push it. Pinned here so the + // connection-state consolidation must preserve it. + test('reconnecting after close() with a carried sessionId + protocolVersion pushes the version onto the new transport', async () => { + const client = new Client({ name: 'resume-client', version: '1.0.0' }); + const first = new SessionfulTransport(); + + await client.connect(first); + expect(first.setProtocolVersionCalls).toEqual([LEGACY]); + expect(first.sessionId).toBe('session-123'); + + await client.close(); + + // The consumer persisted the session id and negotiated version (the + // documented Streamable HTTP resumption flow) and reconnects with + // both carried on the new transport. + const resumed = new SessionfulTransport({ sessionId: 'session-123', protocolVersion: LEGACY }); + await client.connect(resumed); + + // No re-initialize on a resume... + expect(resumed.sent.filter(m => isJSONRPCRequest(m) && m.method === 'initialize')).toHaveLength(0); + // ...and the carried version is pushed to the transport. + expect(resumed.setProtocolVersionCalls).toEqual([LEGACY]); + + await client.close(); + }); + + test('resume on the same instance without close() re-pushes the negotiated version (pin)', async () => { + const client = new Client({ name: 'resume-client', version: '1.0.0' }); + const first = new SessionfulTransport(); + + await client.connect(first); + // The transport drops (server side went away) — no client.close(). + await first.close(); + + const resumed = new SessionfulTransport({ sessionId: 'session-123', protocolVersion: LEGACY }); + await client.connect(resumed); + + expect(resumed.sent.filter(m => isJSONRPCRequest(m) && m.method === 'initialize')).toHaveLength(0); + expect(resumed.setProtocolVersionCalls).toEqual([LEGACY]); + + await client.close(); + }); +}); diff --git a/packages/core-internal/test/shared/connectionLifecycle.test.ts b/packages/core-internal/test/shared/connectionLifecycle.test.ts new file mode 100644 index 0000000000..fcf60738e6 --- /dev/null +++ b/packages/core-internal/test/shared/connectionLifecycle.test.ts @@ -0,0 +1,342 @@ +/** + * Connection lifecycle acceptance suite. + * + * Pins the behaviors any refactor of Protocol's connection-scoped state must + * preserve (g1-g4, r1, r2), and documents today's remaining teardown gap as + * an expected-failing test (r3 — flipped to a plain test by the fix). r2: + * the failed-start unwind restores the transport's own callbacks, so a late + * event from a failed transport cannot torpedo a later connection. + */ +import { describe, expect, test } from 'vitest'; +import * as z from 'zod/v4'; + +import { SdkError, SdkErrorCode } from '../../src/errors/sdkErrors'; +import type { BaseContext } from '../../src/shared/protocol'; +import { Protocol } from '../../src/shared/protocol'; +import type { Transport } from '../../src/shared/transport'; +import type { JSONRPCMessage } from '../../src/types/index'; + +class MockTransport implements Transport { + id: string; + onclose?: () => void; + onerror?: (error: Error) => void; + onmessage?: (message: unknown) => void; + sentMessages: JSONRPCMessage[] = []; + /** When true, close() resolves without firing onclose (see r1). */ + closeWithoutOnclose = false; + + constructor(id: string) { + this.id = id; + } + + async start(): Promise {} + + async close(): Promise { + if (!this.closeWithoutOnclose) { + this.onclose?.(); + } + } + + async send(message: JSONRPCMessage): Promise { + this.sentMessages.push(message); + } +} + +function createProtocol(options?: ConstructorParameters[0]): Protocol { + return new (class extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + protected buildContext(ctx: BaseContext): BaseContext { + return ctx; + } + })(options); +} + +const flushMicrotasks = () => new Promise(resolve => setImmediate(resolve)); + +/** Holds a request handler open and captures its context. */ +function installBlockingHandler(protocol: Protocol): { + entered: Promise; + release: () => void; + ctx: () => BaseContext; +} { + let ctxRef: BaseContext | undefined; + let enteredResolve!: () => void; + const entered = new Promise(resolve => { + enteredResolve = resolve; + }); + let release!: () => void; + const gate = new Promise(resolve => { + release = resolve; + }); + protocol.setRequestHandler('ping', async (_request, ctx) => { + ctxRef = ctx; + enteredResolve(); + await gate; + return {}; + }); + return { entered, release, ctx: () => ctxRef! }; +} + +describe('g1: aborted-handler send gates fire before era resolution', () => { + test('an aborted ctx.mcpReq.send with a spec method rejects ConnectionClosed — never a synchronous era throw', async () => { + const protocol = createProtocol(); + const transportA = new MockTransport('A'); + const { entered, ctx } = installBlockingHandler(protocol); + + await protocol.connect(transportA); + transportA.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + + await protocol.close(); + expect(ctx().mcpReq.signal.aborted).toBe(true); + + // 'subscriptions/listen' is a spec method the (default) legacy era + // does not define: were era resolution consulted before the abort + // gate, this call would THROW MethodNotSupportedByProtocolVersion + // synchronously. The gate must win: an async ConnectionClosed + // rejection, evaluated against nothing. + let syncThrow: unknown; + let pending: Promise | undefined; + try { + pending = ctx().mcpReq.send({ method: 'subscriptions/listen' }); + } catch (error) { + syncThrow = error; + } + expect(syncThrow).toBeUndefined(); + await expect(pending).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + + // Same ordering on the in-era spec-method path (no result schema): + // the gate short-circuits before the codec's result-validator lookup. + await expect(ctx().mcpReq.send({ method: 'ping' })).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + + // Explicit-schema path rejects identically, and notify() resolves as + // a no-op. + await expect(ctx().mcpReq.send({ method: 'custom/probe' }, z.object({}))).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + await expect( + ctx().mcpReq.notify({ method: 'notifications/progress', params: { progressToken: 1, progress: 1 } }) + ).resolves.toBeUndefined(); + }); +}); + +describe('g2: send-failure cleanup', () => { + test('a transport.send rejection cleans up the progress handler and the request timeout', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + transport.send = async () => { + throw new Error('wire is down'); + }; + await protocol.connect(transport); + + const errors: Error[] = []; + protocol.onerror = error => errors.push(error); + + let progressCalls = 0; + await expect( + protocol.request({ method: 'custom/slow' }, z.object({}), { + timeout: 20, + onprogress: () => { + progressCalls++; + } + }) + ).rejects.toThrow('wire is down'); + + // The progress handler was removed: a progress notification for the + // dead request's token surfaces as an unknown-token onerror instead + // of invoking the stale callback. + transport.onmessage?.({ + jsonrpc: '2.0', + method: 'notifications/progress', + params: { progressToken: 0, progress: 1 } + }); + await flushMicrotasks(); + expect(progressCalls).toBe(0); + expect(errors.some(e => e.message.includes('unknown token'))).toBe(true); + + // The timeout was cleaned up: past the deadline nothing fires (no + // cancellation attempt, no late onerror beyond the one above). + const errorCountAfterProgress = errors.length; + await new Promise(resolve => setTimeout(resolve, 40)); + expect(errors.length).toBe(errorCountAfterProgress); + }); +}); + +describe('g3: replies route to the transport that delivered the request', () => { + test('a handler whose connection was replaced mid-flight answers on the captured transport, never the new one', async () => { + const protocol = createProtocol(); + const transportA = new MockTransport('A'); + const transportB = new MockTransport('B'); + + // Two in-flight requests sharing one JSON-RPC id: the second SET on + // `_requestHandlerAbortControllers` replaces the first entry, so + // close() aborts only the second handler. The first handler survives + // un-aborted across close()+connect(B) — its response must still go + // to the transport that delivered it (captured at dispatch), not to + // the live transport read at completion time. + let firstEnteredResolve!: () => void; + const firstEntered = new Promise(resolve => { + firstEnteredResolve = resolve; + }); + let releaseFirst!: () => void; + const firstGate = new Promise(resolve => { + releaseFirst = resolve; + }); + let handlerCalls = 0; + protocol.setRequestHandler('ping', async () => { + handlerCalls++; + if (handlerCalls === 1) { + firstEnteredResolve(); + await firstGate; + return {}; + } + await new Promise(() => {}); // second stays parked + return {}; + }); + + await protocol.connect(transportA); + transportA.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 7 }); + await firstEntered; + transportA.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 7 }); + await flushMicrotasks(); + + await protocol.close(); + await protocol.connect(transportB); + + releaseFirst(); + await flushMicrotasks(); + + // The surviving handler's response went to A (which delivered the + // request) — transport B never sees a message it has no request for. + expect(transportB.sentMessages).toHaveLength(0); + expect(transportA.sentMessages).toHaveLength(1); + expect(transportA.sentMessages[0]).toMatchObject({ jsonrpc: '2.0', id: 7, result: {} }); + }); +}); + +describe('g4: teardown ordering', () => { + test('user onclose fires before in-flight handler aborts; pending responses settle ConnectionClosed', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + const { entered, ctx } = installBlockingHandler(protocol); + + await protocol.connect(transport); + + const events: string[] = []; + protocol.onclose = () => events.push('onclose'); + + // A pending outbound request that teardown must settle. + const pending = protocol.request({ method: 'custom/pending' }, z.object({}), { timeout: 60_000 }).catch((error: unknown) => error); + await flushMicrotasks(); + + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + ctx().mcpReq.signal.addEventListener('abort', () => events.push('abort')); + + await protocol.close(); + const settled = await pending; + + // The pending response settled with the typed ConnectionClosed... + expect(settled).toSatisfy((error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed); + // ...the in-flight handler's abort reason carries the same typed + // error (the response-handlers-then-abort-controllers ordering inside + // teardown shares one error instance; both run after user onclose)... + expect(ctx().mcpReq.signal.reason).toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + // ...and the user onclose callback ran strictly before the abort + // signal fired (both are synchronous within teardown). + expect(events).toEqual(['onclose', 'abort']); + }); +}); + +describe('r1: close() on a transport that never fires onclose (pin)', () => { + // close() runs the teardown itself when the transport's onclose callback + // has not (it used to rely on the onclose round-trip alone, wedging the + // instance on the AlreadyConnected guard). Pinned here so the teardown + // consolidation must preserve it. + test('after await close(), the instance accepts a new connect() even when the transport never fired onclose', async () => { + const protocol = createProtocol(); + const transportA = new MockTransport('A'); + transportA.closeWithoutOnclose = true; + const transportB = new MockTransport('B'); + + await protocol.connect(transportA); + await protocol.close(); + + await protocol.connect(transportB); + expect(protocol.transport).toBe(transportB); + }); +}); + +describe('r2: late events from a failed transport cannot torpedo a later connection (pin)', () => { + test('transport A start() rejects; a late A.onclose leaves connection B fully live', async () => { + const protocol = createProtocol(); + const failing = new MockTransport('failing'); + failing.start = async () => { + throw new Error('spawn failed'); + }; + await expect(protocol.connect(failing)).rejects.toThrow('spawn failed'); + + const transportB = new MockTransport('B'); + await protocol.connect(transportB); + let sawClose = false; + protocol.onclose = () => { + sawClose = true; + }; + + const pending = protocol.request({ method: 'custom/probe' }, z.object({ ok: z.boolean() })); + await flushMicrotasks(); + expect(transportB.sentMessages).toHaveLength(1); + const requestId = (transportB.sentMessages[0] as { id: number }).id; + + // The late event from the failed transport (e.g. a child process + // 'close' arriving after a failed spawn already rejected start()). + failing.onclose?.(); + + // B's connection is unaffected: no teardown ran, and the in-flight + // request on B still completes normally. + expect(sawClose).toBe(false); + expect(protocol.transport).toBe(transportB); + transportB.onmessage?.({ jsonrpc: '2.0', id: requestId, result: { ok: true } }); + await expect(pending).resolves.toEqual({ ok: true }); + }); +}); + +describe('r3: debounced notifications are connection-scoped', () => { + // DEFECT (today): the debounce microtask checks transport PRESENCE + // (`if (!this._transport) return`), not identity. A debounced + // notification scheduled on connection A, with close()+connect(B) both + // landing inside the microtask window, is delivered on B — a connection + // it was never sent on. + test.fails( + 'a notification debounced on connection A must not be sent on a connection B established inside the microtask window', + async () => { + const protocol = createProtocol({ debouncedNotificationMethods: ['test/debounced'] }); + const transportA = new MockTransport('A'); + const transportB = new MockTransport('B'); + + await protocol.connect(transportA); + + // Schedules the send as a microtask on connection A... + void protocol.notification({ method: 'test/debounced' }); + // ...and replaces the connection before that microtask runs. + // MockTransport.close() fires onclose synchronously, so close() has + // already torn A down when connect(B) installs the new transport; + // both happen before the debounce microtask fires. + const closing = protocol.close(); + const connecting = protocol.connect(transportB); + await Promise.all([closing, connecting]); + await flushMicrotasks(); + + expect(transportA.sentMessages).toHaveLength(0); + expect(transportB.sentMessages).toHaveLength(0); + } + ); +}); From 4dc3507089d0016d37a572b28f9338a41d8b7cd9 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 9 Jul 2026 11:44:12 +0000 Subject: [PATCH 2/4] refactor: extract connection-scoped state into a private Connection class --- packages/client/test/client/listen.test.ts | 12 +- packages/core-internal/src/shared/protocol.ts | 358 +++++++++++------- 2 files changed, 227 insertions(+), 143 deletions(-) diff --git a/packages/client/test/client/listen.test.ts b/packages/client/test/client/listen.test.ts index 88d5a68f68..68e42280ff 100644 --- a/packages/client/test/client/listen.test.ts +++ b/packages/client/test/client/listen.test.ts @@ -20,6 +20,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { Client } from '../../src/client/client'; const MODERN = '2026-07-28'; + +/** White-box probe: pending response handlers on the client's live connection (0 when disconnected). */ +function responseHandlerCount(client: unknown): number { + const connection = (client as { _conn?: { responseHandlers: Map } })._conn; + return connection?.responseHandlers.size ?? 0; +} const flush = () => new Promise(r => setTimeout(r, 10)); async function scriptedModern(onListen?: (id: number | string, filter: unknown, send: (m: JSONRPCMessage) => void) => void) { @@ -374,7 +380,7 @@ describe('Client.listen()', () => { // registered before send threw; listen() never registers in // Protocol's `_responseHandlers` so there is nothing to leak there. expect((client as unknown as { _listenState: Map })._listenState.size).toBe(0); - expect((client as unknown as { _responseHandlers: Map })._responseHandlers.size).toBe(0); + expect(responseHandlerCount(client)).toBe(0); clientTx.send = realSend; await client.close(); }); @@ -739,7 +745,7 @@ describe('Client.listen()', () => { expect(error).toBeInstanceOf(Error); expect(Date.now() - t0).toBeLessThan(1000); expect((client as unknown as { _listenState: Map })._listenState.size).toBe(0); - expect((client as unknown as { _responseHandlers: Map })._responseHandlers.size).toBe(0); + expect(responseHandlerCount(client)).toBe(0); }); it("transport closes WHILE the subscription is open: closed resolves 'remote'; close() is a no-op", async () => { @@ -991,7 +997,7 @@ describe('Client.listen() — ack timeout (fake timers)', () => { expect(cancelled).toMatchObject({ params: { requestId: listenId } }); // No leaked state. expect((client as unknown as { _listenState: Map })._listenState.size).toBe(0); - expect((client as unknown as { _responseHandlers: Map })._responseHandlers.size).toBe(0); + expect(responseHandlerCount(client)).toBe(0); // Restore real timers before close to avoid hanging on transport timers. vi.useRealTimers(); await client.close(); diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 66900a8313..11f4da8e3b 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -11,6 +11,7 @@ import type { ElicitResult, HandlerResultTypeMap, JSONRPCErrorResponse, + JSONRPCMessage, JSONRPCNotification, JSONRPCRequest, JSONRPCResponse, @@ -521,6 +522,120 @@ type TimeoutInfo = { onTimeout: () => void; }; +/** + * Owns the state scoped to a single transport connection: the transport + * itself with its wrapped callbacks, the response/progress-handler and + * request-timeout registries, the in-flight request-handler abort + * controllers, and the pending debounced notifications. + */ +class Connection { + /** `undefined` only on the idle placeholder that absorbs post-teardown operations. */ + readonly transport?: Transport; + + responseHandlers: Map void> = new Map(); + progressHandlers: Map = new Map(); + timeoutInfo: Map = new Map(); + requestHandlerAbortControllers: Map = new Map(); + pendingDebouncedNotifications = new Set(); + + private _originalOnClose?: Transport['onclose']; + private _originalOnError?: Transport['onerror']; + private _originalOnMessage?: Transport['onmessage']; + + constructor(transport: Transport | undefined) { + this.transport = transport; + } + + /** + * Wraps the transport's callbacks so its events also feed the owning + * Protocol, preserving any callbacks the transport already had (they run + * first). + */ + installCallbacks(hooks: { + onclose: () => void; + onerror: (error: Error) => void; + onmessage: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; + }): void { + const transport = this.transport; + if (!transport) { + return; + } + const _onclose = (this._originalOnClose = transport.onclose); + transport.onclose = () => { + try { + _onclose?.(); + } finally { + hooks.onclose(); + } + }; + + const _onerror = (this._originalOnError = transport.onerror); + transport.onerror = (error: Error) => { + _onerror?.(error); + hooks.onerror(error); + }; + + const _onmessage = (this._originalOnMessage = transport.onmessage); + transport.onmessage = (message, extra) => { + _onmessage?.(message, extra); + hooks.onmessage(message, extra); + }; + } + + /** Undoes {@linkcode installCallbacks}: puts the transport's own callbacks back. */ + restoreCallbacks(): void { + if (!this.transport) { + return; + } + this.transport.onclose = this._originalOnClose; + this.transport.onerror = this._originalOnError; + this.transport.onmessage = this._originalOnMessage; + } + + setupTimeout( + messageId: number, + timeout: number, + maxTotalTimeout: number | undefined, + onTimeout: () => void, + resetTimeoutOnProgress: boolean = false + ): void { + this.timeoutInfo.set(messageId, { + timeoutId: setTimeout(onTimeout, timeout), + startTime: Date.now(), + timeout, + maxTotalTimeout, + resetTimeoutOnProgress, + onTimeout + }); + } + + resetTimeout(messageId: number): boolean { + const info = this.timeoutInfo.get(messageId); + if (!info) return false; + + const totalElapsed = Date.now() - info.startTime; + if (info.maxTotalTimeout && totalElapsed >= info.maxTotalTimeout) { + this.timeoutInfo.delete(messageId); + throw new SdkError(SdkErrorCode.RequestTimeout, 'Maximum total timeout exceeded', { + maxTotalTimeout: info.maxTotalTimeout, + totalElapsed + }); + } + + clearTimeout(info.timeoutId); + info.timeoutId = setTimeout(info.onTimeout, info.timeout); + return true; + } + + cleanupTimeout(messageId: number): void { + const info = this.timeoutInfo.get(messageId); + if (info) { + clearTimeout(info.timeoutId); + this.timeoutInfo.delete(messageId); + } + } +} + /* * Package-internal write access to Protocol's negotiated-protocol-version state. * @@ -555,15 +670,20 @@ export function setNegotiatedProtocolVersion( * implementations most code should use. */ export abstract class Protocol { - private _transport?: Transport; + private _connection?: Connection; + /** + * Absorbs connection-scoped operations that land after teardown (a + * request's finally-cleanup, late cancellations): reads and deletes + * against it are no-ops, matching the fresh maps `_onclose` used to + * install. Never carries live traffic. + */ + private readonly _idleConnectionState = new Connection(undefined); + private get _conn(): Connection { + return this._connection ?? this._idleConnectionState; + } private _requestMessageId = 0; private _requestHandlers: Map Promise> = new Map(); - private _requestHandlerAbortControllers: Map = new Map(); private _notificationHandlers: Map Promise> = new Map(); - private _responseHandlers: Map void> = new Map(); - private _progressHandlers: Map = new Map(); - private _timeoutInfo: Map = new Map(); - private _pendingDebouncedNotifications = new Set(); /** * The protocol version negotiated for the current connection (`undefined` @@ -727,59 +847,16 @@ export abstract class Protocol { return; } // Handle request cancellation - const controller = this._requestHandlerAbortControllers.get(notification.params.requestId); + const controller = this._conn.requestHandlerAbortControllers.get(notification.params.requestId); controller?.abort(notification.params.reason); } - private _setupTimeout( - messageId: number, - timeout: number, - maxTotalTimeout: number | undefined, - onTimeout: () => void, - resetTimeoutOnProgress: boolean = false - ) { - this._timeoutInfo.set(messageId, { - timeoutId: setTimeout(onTimeout, timeout), - startTime: Date.now(), - timeout, - maxTotalTimeout, - resetTimeoutOnProgress, - onTimeout - }); - } - - private _resetTimeout(messageId: number): boolean { - const info = this._timeoutInfo.get(messageId); - if (!info) return false; - - const totalElapsed = Date.now() - info.startTime; - if (info.maxTotalTimeout && totalElapsed >= info.maxTotalTimeout) { - this._timeoutInfo.delete(messageId); - throw new SdkError(SdkErrorCode.RequestTimeout, 'Maximum total timeout exceeded', { - maxTotalTimeout: info.maxTotalTimeout, - totalElapsed - }); - } - - clearTimeout(info.timeoutId); - info.timeoutId = setTimeout(info.onTimeout, info.timeout); - return true; - } - - private _cleanupTimeout(messageId: number) { - const info = this._timeoutInfo.get(messageId); - if (info) { - clearTimeout(info.timeoutId); - this._timeoutInfo.delete(messageId); - } - } - /** * Throws if already bound to a transport. Subclass connect overrides that * mutate per-connection state before `super.connect()` must call this first. */ protected assertNotConnected(): void { - if (this._transport) { + if (this._connection) { throw new SdkError( SdkErrorCode.AlreadyConnected, 'Already connected to a transport. Call close() before connecting to a new transport, or use a separate Protocol instance per connection.' @@ -795,54 +872,47 @@ export abstract class Protocol { async connect(transport: Transport): Promise { this.assertNotConnected(); - this._transport = transport; - const _onclose = this.transport?.onclose; - this._transport.onclose = () => { - try { - _onclose?.(); - } finally { + const connection = new Connection(transport); + this._connection = connection; + connection.installCallbacks({ + onclose: () => { // Identity guard: close() may have already torn this // connection down (and a replacement may be live) by the time // a transport fires a deferred onclose. - if (this._transport === transport) { + if (this._connection === connection) { this._onclose(); } + }, + onerror: error => this._onerror(error), + onmessage: (message, extra) => { + if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { + this._onresponse(message); + } else if (isJSONRPCRequest(message)) { + this._onrequest(message, extra); + } else if (isJSONRPCNotification(message)) { + this._onnotification(message, extra); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); + } } - }; - - const _onerror = this.transport?.onerror; - this._transport.onerror = (error: Error) => { - _onerror?.(error); - this._onerror(error); - }; - - const _onmessage = this._transport?.onmessage; - this._transport.onmessage = (message, extra) => { - _onmessage?.(message, extra); - if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { - this._onresponse(message); - } else if (isJSONRPCRequest(message)) { - this._onrequest(message, extra); - } else if (isJSONRPCNotification(message)) { - this._onnotification(message, extra); - } else { - this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); - } - }; + }); // Pass supported protocol versions to transport for header validation transport.setSupportedProtocolVersions?.(this._supportedProtocolVersions); try { - await this._transport.start(); + await transport.start(); } catch (error) { - // Unwind a failed start so connect() can be retried: restore the - // transport's own callbacks — a late event from the failed - // transport must not tear down a successful retry's connection. - transport.onclose = _onclose; - transport.onerror = _onerror; - transport.onmessage = _onmessage; - this._transport = undefined; + // A transport that failed to start was never connected: unwind so + // the instance can retry connect() — the reuse guard above would + // otherwise lock the instance onto a transport that never carried + // traffic. Restore the transport's own callbacks too: leaving the + // wrapped closures installed would let a late event from the + // failed transport (e.g. a child process 'close' after a failed + // spawn) tear down whatever connection a successful retry + // established. + connection.restoreCallbacks(); + this._connection = undefined; throw error; } } @@ -853,22 +923,23 @@ export abstract class Protocol { * timeout clearing, in-flight request abort) does not run otherwise. */ protected _onclose(): void { - const responseHandlers = this._responseHandlers; - this._responseHandlers = new Map(); - this._progressHandlers.clear(); - this._pendingDebouncedNotifications.clear(); + const connection = this._conn; + const responseHandlers = connection.responseHandlers; + connection.responseHandlers = new Map(); + connection.progressHandlers.clear(); + connection.pendingDebouncedNotifications.clear(); - for (const info of this._timeoutInfo.values()) { + for (const info of connection.timeoutInfo.values()) { clearTimeout(info.timeoutId); } - this._timeoutInfo.clear(); + connection.timeoutInfo.clear(); - const requestHandlerAbortControllers = this._requestHandlerAbortControllers; - this._requestHandlerAbortControllers = new Map(); + const requestHandlerAbortControllers = connection.requestHandlerAbortControllers; + connection.requestHandlerAbortControllers = new Map(); const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed'); - this._transport = undefined; + this._connection = undefined; try { this.onclose?.(); @@ -978,7 +1049,7 @@ export abstract class Protocol { } // Capture the current transport at request time to ensure responses go to the correct client - const capturedTransport = this._transport; + const capturedTransport = this.transport; const sendErrorResponse = (code: number, message: string, data?: unknown) => { const errorResponse: JSONRPCErrorResponse = { @@ -1056,15 +1127,18 @@ export abstract class Protocol { } const abortController = new AbortController(); - this._requestHandlerAbortControllers.set(request.id, abortController); + this._conn.requestHandlerAbortControllers.set(request.id, abortController); // Related sends resolve through the SAME instance era as every other // sender (the per-request/instance asymmetry is deliberately gone): // the codec is resolved at send time from the connection state. // - // After the handler aborts, related sends must not reach the transport - // (read live at send time, it may belong to a later connection): - // notifications no-op, requests reject with ConnectionClosed. + // Once this request's handler has been aborted (connection closed or + // request cancelled), its related sends must not reach the transport: + // the transport is read live at send time, so a handler still + // running after close() could otherwise write onto a transport + // connected later. Notifications no-op; requests reject with + // ConnectionClosed. const sendNotification = (notification: Notification, options?: NotificationOptions): Promise => { if (abortController.signal.aborted) { return Promise.resolve(); @@ -1205,8 +1279,8 @@ export abstract class Protocol { ) .catch(error => this._onerror(new Error(`Failed to send response: ${error}`))) .finally(() => { - if (this._requestHandlerAbortControllers.get(request.id) === abortController) { - this._requestHandlerAbortControllers.delete(request.id); + if (this._conn.requestHandlerAbortControllers.get(request.id) === abortController) { + this._conn.requestHandlerAbortControllers.delete(request.id); } }); } @@ -1215,23 +1289,23 @@ export abstract class Protocol { const { progressToken, ...params } = notification.params; const messageId = Number(progressToken); - const handler = this._progressHandlers.get(messageId); + const handler = this._conn.progressHandlers.get(messageId); if (!handler) { this._onerror(new Error(`Received a progress notification for an unknown token: ${JSON.stringify(notification)}`)); return; } - const responseHandler = this._responseHandlers.get(messageId); - const timeoutInfo = this._timeoutInfo.get(messageId); + const responseHandler = this._conn.responseHandlers.get(messageId); + const timeoutInfo = this._conn.timeoutInfo.get(messageId); if (timeoutInfo && responseHandler && timeoutInfo.resetTimeoutOnProgress) { try { - this._resetTimeout(messageId); + this._conn.resetTimeout(messageId); } catch (error) { // Clean up if maxTotalTimeout was exceeded - this._responseHandlers.delete(messageId); - this._progressHandlers.delete(messageId); - this._cleanupTimeout(messageId); + this._conn.responseHandlers.delete(messageId); + this._conn.progressHandlers.delete(messageId); + this._conn.cleanupTimeout(messageId); responseHandler(error as Error); return; } @@ -1249,15 +1323,15 @@ export abstract class Protocol { protected _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void { const messageId = Number(response.id); - const handler = this._responseHandlers.get(messageId); + const handler = this._conn.responseHandlers.get(messageId); if (handler === undefined) { this._onerror(new Error(`Received a response for an unknown message ID: ${JSON.stringify(response)}`)); return; } - this._responseHandlers.delete(messageId); - this._cleanupTimeout(messageId); - this._progressHandlers.delete(messageId); + this._conn.responseHandlers.delete(messageId); + this._conn.cleanupTimeout(messageId); + this._conn.progressHandlers.delete(messageId); if (isJSONRPCResultResponse(response)) { handler(response); @@ -1268,19 +1342,19 @@ export abstract class Protocol { } get transport(): Transport | undefined { - return this._transport; + return this._connection?.transport; } /** * Closes the connection. */ async close(): Promise { - const transport = this._transport; - await transport?.close(); + const connection = this._connection; + await connection?.transport?.close(); // A transport whose close() resolves without firing onclose would // otherwise leave the instance wedged (still "connected" per the // connect() guard). Run the teardown ourselves if it hasn't happened. - if (transport !== undefined && this._transport === transport) { + if (connection !== undefined && this._connection === connection) { this._onclose(); } } @@ -1438,7 +1512,7 @@ export abstract class Protocol { reject(error); }; - if (!this._transport) { + if (!this.transport) { earlyReject(new Error('Not connected')); return; } @@ -1473,7 +1547,7 @@ export abstract class Protocol { // combination — legacy era on any transport, modern era on stdio / // in-memory — keeps today's `notifications/cancelled` POST path // unchanged. - const streamCloseCancels = codec.era === MODERN_WIRE_REVISION && this._transport.hasPerRequestStream === true; + const streamCloseCancels = codec.era === MODERN_WIRE_REVISION && this.transport?.hasPerRequestStream === true; const requestAbort = streamCloseCancels ? new AbortController() : undefined; const messageId = this._requestMessageId++; @@ -1485,7 +1559,7 @@ export abstract class Protocol { }; if (options?.onprogress) { - this._progressHandlers.set(messageId, options.onprogress); + this._conn.progressHandlers.set(messageId, options.onprogress); jsonrpcRequest.params = { ...request.params, _meta: { @@ -1507,10 +1581,10 @@ export abstract class Protocol { if (responseReceived) { return; } - this._progressHandlers.delete(messageId); + this._conn.progressHandlers.delete(messageId); if (requestAbort === undefined) { - this._transport + this.transport ?.send( this._envelopeOutbound({ jsonrpc: '2.0', @@ -1537,7 +1611,7 @@ export abstract class Protocol { reject(error); }; - this._responseHandlers.set(messageId, response => { + this._conn.responseHandlers.set(messageId, response => { if (options?.signal?.aborted) { return; } @@ -1613,14 +1687,18 @@ export abstract class Protocol { const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC; const timeoutHandler = () => cancel(new SdkError(SdkErrorCode.RequestTimeout, 'Request timed out', { timeout })); - this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); + this._conn.setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); - this._transport - .send(outbound, { relatedRequestId, resumptionToken, onresumptiontoken, headers, requestSignal: requestAbort?.signal }) - .catch(error => { - this._progressHandlers.delete(messageId); - reject(error); - }); + this.transport!.send(outbound, { + relatedRequestId, + resumptionToken, + onresumptiontoken, + headers, + requestSignal: requestAbort?.signal + }).catch(error => { + this._conn.progressHandlers.delete(messageId); + reject(error); + }); }).finally(() => { // Per-request cleanup that must run on every exit path. Consolidated // here so new exit paths added to the promise body can't forget it. @@ -1630,8 +1708,8 @@ export abstract class Protocol { options?.signal?.removeEventListener('abort', onAbort); } if (cleanupMessageId !== undefined) { - this._responseHandlers.delete(cleanupMessageId); - this._cleanupTimeout(cleanupMessageId); + this._conn.responseHandlers.delete(cleanupMessageId); + this._conn.cleanupTimeout(cleanupMessageId); } }); } @@ -1649,7 +1727,7 @@ export abstract class Protocol { * resolve through the instance's negotiated era at send time. */ private async _notificationViaCodec(codec: WireCodec, notification: Notification, options?: NotificationOptions): Promise { - if (!this._transport) { + if (!this.transport) { throw new SdkError(SdkErrorCode.NotConnected, 'Not connected'); } @@ -1674,34 +1752,34 @@ export abstract class Protocol { if (canDebounce) { // If a notification of this type is already scheduled, do nothing. - if (this._pendingDebouncedNotifications.has(notification.method)) { + if (this._conn.pendingDebouncedNotifications.has(notification.method)) { return; } // Mark this notification type as pending. - this._pendingDebouncedNotifications.add(notification.method); + this._conn.pendingDebouncedNotifications.add(notification.method); // Schedule the actual send to happen in the next microtask. // This allows all synchronous calls in the current event loop tick to be coalesced. Promise.resolve().then(() => { // Un-mark the notification so the next one can be scheduled. - this._pendingDebouncedNotifications.delete(notification.method); + this._conn.pendingDebouncedNotifications.delete(notification.method); // SAFETY CHECK: If the connection was closed while this was pending, abort. - if (!this._transport) { + if (!this.transport) { return; } // Send the notification, but don't await it here to avoid blocking. // Handle potential errors with a .catch(). - this._transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error)); + this.transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error)); }); // Return immediately. return; } - await this._transport.send(jsonrpcNotification, options); + await this.transport!.send(jsonrpcNotification, options); } /** From 50fd34ba2922249fc6f8f75b8f568df440c56c43 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 9 Jul 2026 11:59:11 +0000 Subject: [PATCH 3/4] refactor: consolidate connection teardown; connection-scoped sends fail structurally after dispose --- packages/client/src/client/client.ts | 46 +- .../client/test/client/sessionResume.test.ts | 10 +- packages/core-internal/src/shared/protocol.ts | 452 ++++++++++++------ .../test/shared/connectionLifecycle.test.ts | 62 ++- packages/server/src/server/server.ts | 15 +- test/e2e/requirements.ts | 16 +- 6 files changed, 391 insertions(+), 210 deletions(-) diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 6746f63820..2538c5bb0c 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -518,6 +518,18 @@ export class Client extends Protocol { private _autoOpenedSubscription?: McpSubscription; /** Backing store for {@linkcode getDiscoverResult}. Per-connection. */ private _discoverResult?: DiscoverResult; + /** + * The protocol version to restore when resuming a server session on a new + * transport. Written when a connection establishes (the negotiated + * version must outlive the connection — a transport drop does not end the + * server session); cleared by {@linkcode _resetConnectionState} + * (close / fresh connect). The resume branches also fall back to the + * version the resuming transport itself carries + * ({@linkcode Transport.protocolVersion}), so resuming AFTER close() — + * where the consumer persisted sessionId + protocolVersion — still pushes + * the header. + */ + private _sessionResumption?: { version: string }; /** * Clears every per-connection field in one place. Called at the start of @@ -527,6 +539,7 @@ export class Client extends Protocol { */ private _resetConnectionState(): void { this._negotiatedProtocolVersion = undefined; + this._sessionResumption = undefined; this._serverCapabilities = undefined; this._serverVersion = undefined; this._instructions = undefined; @@ -945,10 +958,15 @@ export class Client extends Protocol { // the originally negotiated version — from this instance, or from the // transport itself after close() wiped the instance's copy. if (transport.sessionId !== undefined) { - const version = this._negotiatedProtocolVersion ?? transport.protocolVersion; - if (version !== undefined) { - this._negotiatedProtocolVersion = version; - transport.setProtocolVersion?.(version); + // Resuming keeps the original negotiation: the version recorded + // when the session's connection established (or, after close(), + // carried by the resuming transport itself) seeds the new + // connection's era and is pushed onto the transport again. + const negotiatedProtocolVersion = this._sessionResumption?.version ?? transport.protocolVersion; + if (negotiatedProtocolVersion !== undefined) { + this._negotiatedProtocolVersion = negotiatedProtocolVersion; + this._sessionResumption = { version: negotiatedProtocolVersion }; + transport.setProtocolVersion?.(negotiatedProtocolVersion); } return; } @@ -1020,6 +1038,10 @@ export class Client extends Protocol { // EXCHANGE is the legacy handshake by definition and completes on // that era. this._negotiatedProtocolVersion = result.protocolVersion; + // Recorded for later session-resuming reconnects: the negotiated + // version must outlive this connection (a transport drop does not + // end the server session), but not a close()/fresh handshake. + this._sessionResumption = { version: result.protocolVersion }; // Set up list changed handlers now that we know server capabilities if (this._listChangedConfig) { @@ -1046,10 +1068,16 @@ export class Client extends Protocol { // this instance or, after close() wiped it, from the transport itself. if (transport.sessionId !== undefined) { await super.connect(transport); - const version = this._negotiatedProtocolVersion ?? transport.protocolVersion; - if (version !== undefined) { - this._negotiatedProtocolVersion = version; - transport.setProtocolVersion?.(version); + // Same restore as the legacy resume branch: recorded at session + // establishment, or carried by the resuming transport after a + // close(). + const negotiatedProtocolVersion = this._sessionResumption?.version ?? transport.protocolVersion; + if (negotiatedProtocolVersion !== undefined) { + this._negotiatedProtocolVersion = negotiatedProtocolVersion; + this._sessionResumption = { version: negotiatedProtocolVersion }; + if (transport.setProtocolVersion) { + transport.setProtocolVersion(negotiatedProtocolVersion); + } } return; } @@ -1090,6 +1118,7 @@ export class Client extends Protocol { this._discoverResult = result.discover; // Modern selection: the same connection state the legacy handshake completion sets. this._negotiatedProtocolVersion = result.version; + this._sessionResumption = { version: result.version }; // The single setProtocolVersion call site on this path, mirroring the legacy path after initialize. if (transport.setProtocolVersion) { transport.setProtocolVersion(result.version); @@ -1210,6 +1239,7 @@ export class Client extends Protocol { this._cache.setServerIdentity(this._deriveServerIdentity(transport)); this._instructions = prior.instructions; this._negotiatedProtocolVersion = version; + this._sessionResumption = { version }; transport.setProtocolVersion?.(version); // No auto-opened listen stream on this path (request-only workers). diff --git a/packages/client/test/client/sessionResume.test.ts b/packages/client/test/client/sessionResume.test.ts index 9ff1b9bde3..5665926df4 100644 --- a/packages/client/test/client/sessionResume.test.ts +++ b/packages/client/test/client/sessionResume.test.ts @@ -62,10 +62,12 @@ class SessionfulTransport { } describe('r4: session resume after close()', () => { - // close() wipes the instance's negotiated state, so the resume branches - // fall back to the version the resuming transport itself carries - // (Transport.protocolVersion) and re-push it. Pinned here so the - // connection-state consolidation must preserve it. + // The resume branch seeds the new connection from the session-resumption + // record or, after close() cleared it, from the version the resuming + // transport itself carries — so setProtocolVersion is pushed either way. + // (It used to read the negotiated version, which close() had reset, and + // silently skip the push: the resumed session's HTTP requests went out + // without the required mcp-protocol-version header.) test('reconnecting after close() with a carried sessionId + protocolVersion pushes the version onto the new transport', async () => { const client = new Client({ name: 'resume-client', version: '1.0.0' }); const first = new SessionfulTransport(); diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 11f4da8e3b..74a7bf7686 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -538,6 +538,14 @@ class Connection { requestHandlerAbortControllers: Map = new Map(); pendingDebouncedNotifications = new Set(); + /** + * The protocol version negotiated on THIS connection (`undefined` before + * negotiation completes) — connection state, so a stale era can never + * leak onto a successor connection. + */ + negotiatedProtocolVersion?: string; + + private _disposed = false; private _originalOnClose?: Transport['onclose']; private _originalOnError?: Transport['onerror']; private _originalOnMessage?: Transport['onmessage']; @@ -546,43 +554,137 @@ class Connection { this.transport = transport; } + /** `true` once {@linkcode dispose} ran: sends reject, transport events no-op. */ + get disposed(): boolean { + return this._disposed; + } + /** - * Wraps the transport's callbacks so its events also feed the owning - * Protocol, preserving any callbacks the transport already had (they run - * first). + * Sends on this connection's transport; rejects with the typed + * ConnectionClosed error once the connection is disposed, so senders that + * captured this connection fail structurally instead of writing onto + * whatever connection replaced it. */ - installCallbacks(hooks: { - onclose: () => void; - onerror: (error: Error) => void; - onmessage: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; - }): void { + send(message: JSONRPCMessage, options?: TransportSendOptions): Promise { + if (this._disposed || this.transport === undefined) { + return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); + } + return this.transport.send(message, options); + } + + /** + * Atomically brings the connection up: installs the wrapped callbacks and + * starts the transport, or — when start() rejects — fully unwinds + * (restores the transport's own callbacks and disposes this connection so + * late events from the failed transport are inert). + */ + async open( + hooks: { + onclose: () => void; + onerror: (error: Error) => void; + onmessage: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; + }, + supportedProtocolVersions: string[] + ): Promise { const transport = this.transport; if (!transport) { return; } + this._installCallbacks(transport, hooks); + + // Pass supported protocol versions to transport for header validation + transport.setSupportedProtocolVersions?.(supportedProtocolVersions); + + try { + await transport.start(); + } catch (error) { + // A transport that failed to start was never connected: unwind so + // the instance can retry connect(). Restore the transport's own + // callbacks and dispose — a late event from the failed transport + // (e.g. a child process 'close' after a failed spawn) must not + // tear down whatever connection a successful retry established. + this.restoreCallbacks(); + this.dispose(); + throw error; + } + } + + /** + * Tears down this connection's state exactly once. Returns the response + * handlers and in-flight handler abort controllers for the caller to + * settle (in that order), or `undefined` when already disposed. + */ + dispose(): + | { + responseHandlers: Map void>; + requestHandlerAbortControllers: Map; + } + | undefined { + if (this._disposed) { + return undefined; + } + this._disposed = true; + + const responseHandlers = this.responseHandlers; + this.responseHandlers = new Map(); + this.progressHandlers.clear(); + this.pendingDebouncedNotifications.clear(); + + for (const info of this.timeoutInfo.values()) { + clearTimeout(info.timeoutId); + } + this.timeoutInfo.clear(); + + const requestHandlerAbortControllers = this.requestHandlerAbortControllers; + this.requestHandlerAbortControllers = new Map(); + + return { responseHandlers, requestHandlerAbortControllers }; + } + + /** + * Wraps the transport's callbacks so its events also feed the owning + * Protocol, preserving any callbacks the transport already had (they run + * first). A disposed connection's events never reach the Protocol hooks — + * a stale transport cannot tear down or write into a successor + * connection. + */ + private _installCallbacks( + transport: Transport, + hooks: { + onclose: () => void; + onerror: (error: Error) => void; + onmessage: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; + } + ): void { const _onclose = (this._originalOnClose = transport.onclose); transport.onclose = () => { try { _onclose?.(); } finally { - hooks.onclose(); + if (!this._disposed) { + hooks.onclose(); + } } }; const _onerror = (this._originalOnError = transport.onerror); transport.onerror = (error: Error) => { _onerror?.(error); - hooks.onerror(error); + if (!this._disposed) { + hooks.onerror(error); + } }; const _onmessage = (this._originalOnMessage = transport.onmessage); transport.onmessage = (message, extra) => { _onmessage?.(message, extra); - hooks.onmessage(message, extra); + if (!this._disposed) { + hooks.onmessage(message, extra); + } }; } - /** Undoes {@linkcode installCallbacks}: puts the transport's own callbacks back. */ + /** Puts the transport's own callbacks back. */ restoreCallbacks(): void { if (!this.transport) { return; @@ -685,13 +787,33 @@ export abstract class Protocol { private _requestHandlers: Map Promise> = new Map(); private _notificationHandlers: Map Promise> = new Map(); + /** + * Seed for the negotiated protocol version written BEFORE a connection + * exists (the serving entries bind a factory instance to an era before + * connecting it); {@linkcode Protocol.connect | connect()} copies it onto + * each new connection. + */ + private _pendingNegotiatedProtocolVersion?: string; + /** * The protocol version negotiated for the current connection (`undefined` * before negotiation completes), which determines the wire era this * instance speaks. Set by the SDK's negotiation and initialize paths - * (`Client.connect`, `Server._oninitialize`). + * (`Client.connect`, `Server._oninitialize`). Lives on the connection: + * writes before connect() land on the pending seed the next connection + * starts from; reads without a live connection fall back to that seed. */ - protected _negotiatedProtocolVersion?: string; + protected get _negotiatedProtocolVersion(): string | undefined { + return this._connection === undefined ? this._pendingNegotiatedProtocolVersion : this._connection.negotiatedProtocolVersion; + } + + protected set _negotiatedProtocolVersion(version: string | undefined) { + if (this._connection === undefined) { + this._pendingNegotiatedProtocolVersion = version; + } else { + this._connection.negotiatedProtocolVersion = version; + } + } static { writeNegotiatedProtocolVersion = (instance, version) => { @@ -873,45 +995,33 @@ export abstract class Protocol { this.assertNotConnected(); const connection = new Connection(transport); + // An era bound before connect() (serving entries, the package-internal + // write channel) seeds the new connection. + connection.negotiatedProtocolVersion = this._pendingNegotiatedProtocolVersion; this._connection = connection; - connection.installCallbacks({ - onclose: () => { - // Identity guard: close() may have already torn this - // connection down (and a replacement may be live) by the time - // a transport fires a deferred onclose. - if (this._connection === connection) { - this._onclose(); - } - }, - onerror: error => this._onerror(error), - onmessage: (message, extra) => { - if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { - this._onresponse(message); - } else if (isJSONRPCRequest(message)) { - this._onrequest(message, extra); - } else if (isJSONRPCNotification(message)) { - this._onnotification(message, extra); - } else { - this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); - } - } - }); - - // Pass supported protocol versions to transport for header validation - transport.setSupportedProtocolVersions?.(this._supportedProtocolVersions); - try { - await transport.start(); + await connection.open( + { + onclose: () => this._onclose(), + onerror: error => this._onerror(error), + onmessage: (message, extra) => { + if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { + this._onresponse(message); + } else if (isJSONRPCRequest(message)) { + this._onrequest(message, extra); + } else if (isJSONRPCNotification(message)) { + this._onnotification(message, extra); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); + } + } + }, + this._supportedProtocolVersions + ); } catch (error) { - // A transport that failed to start was never connected: unwind so - // the instance can retry connect() — the reuse guard above would - // otherwise lock the instance onto a transport that never carried - // traffic. Restore the transport's own callbacks too: leaving the - // wrapped closures installed would let a late event from the - // failed transport (e.g. a child process 'close' after a failed - // spawn) tear down whatever connection a successful retry - // established. - connection.restoreCallbacks(); + // A transport that failed to start was never connected — open() + // fully unwound the connection; free the slot so the instance can + // retry connect(). this._connection = undefined; throw error; } @@ -923,33 +1033,23 @@ export abstract class Protocol { * timeout clearing, in-flight request abort) does not run otherwise. */ protected _onclose(): void { - const connection = this._conn; - const responseHandlers = connection.responseHandlers; - connection.responseHandlers = new Map(); - connection.progressHandlers.clear(); - connection.pendingDebouncedNotifications.clear(); - - for (const info of connection.timeoutInfo.values()) { - clearTimeout(info.timeoutId); - } - connection.timeoutInfo.clear(); - - const requestHandlerAbortControllers = connection.requestHandlerAbortControllers; - connection.requestHandlerAbortControllers = new Map(); + const connection = this._connection; + this._connection = undefined; + const captured = connection?.dispose(); const error = new SdkError(SdkErrorCode.ConnectionClosed, 'Connection closed'); - this._connection = undefined; - try { this.onclose?.(); } finally { - for (const handler of responseHandlers.values()) { - handler(error); - } + if (captured !== undefined) { + for (const handler of captured.responseHandlers.values()) { + handler(error); + } - for (const controller of requestHandlerAbortControllers.values()) { - controller.abort(error); + for (const controller of captured.requestHandlerAbortControllers.values()) { + controller.abort(error); + } } } } @@ -1126,40 +1226,45 @@ export abstract class Protocol { return; } + // The connection that delivered this request. Per-request senders + // capture it and send through it, so a handler that outlives its + // connection fails structurally (requests reject ConnectionClosed, + // notifications no-op) instead of writing onto whatever connection + // replaced it. + const capturedConnection = this._conn; + const abortController = new AbortController(); - this._conn.requestHandlerAbortControllers.set(request.id, abortController); + capturedConnection.requestHandlerAbortControllers.set(request.id, abortController); // Related sends resolve through the SAME instance era as every other // sender (the per-request/instance asymmetry is deliberately gone): // the codec is resolved at send time from the connection state. - // - // Once this request's handler has been aborted (connection closed or - // request cancelled), its related sends must not reach the transport: - // the transport is read live at send time, so a handler still - // running after close() could otherwise write onto a transport - // connected later. Notifications no-op; requests reject with - // ConnectionClosed. const sendNotification = (notification: Notification, options?: NotificationOptions): Promise => { - if (abortController.signal.aborted) { - return Promise.resolve(); - } - return this._notificationViaCodec(this._resolveOutboundCodec(notification.method), notification, { - ...options, - relatedRequestId: request.id - }); + return this._notificationViaCodec( + this._resolveOutboundCodec(notification.method), + notification, + { + ...options, + relatedRequestId: request.id + }, + capturedConnection + ); }; const sendRequest = ( r: Request, resultSchema: U, options?: RequestOptions ): Promise> => { - if (abortController.signal.aborted) { - return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); - } - return this._requestWithSchemaViaCodec(this._resolveOutboundCodec(r.method), r, resultSchema, { - ...options, - relatedRequestId: request.id - }); + return this._requestWithSchemaViaCodec( + this._resolveOutboundCodec(r.method), + r, + resultSchema, + { + ...options, + relatedRequestId: request.id + }, + capturedConnection + ); }; // Multi-round-trip retry material: only BARE response objects are @@ -1192,10 +1297,12 @@ export abstract class Protocol { // that overloaded property type. The cast is sound: this impl dispatches both overload paths via the // isStandardSchema guard, and sendRequest validates the result against the resolved schema either way. send: ((r: Request, schemaOrOptions?: StandardSchemaV1 | RequestOptions, maybeOptions?: RequestOptions) => { - // Abort guard FIRST, before era/codec resolution: an - // aborted handler gets ConnectionClosed, never an era - // throw evaluated against a replacement connection. - if (abortController.signal.aborted) { + // The disposed-connection check comes FIRST — before + // era/codec resolution — so a handler that outlived its + // connection always gets the documented ConnectionClosed + // rejection, never a synchronous era throw evaluated + // against whatever connection replaced this request's. + if (capturedConnection.disposed) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } // Related requests resolve through the instance era at @@ -1279,8 +1386,8 @@ export abstract class Protocol { ) .catch(error => this._onerror(new Error(`Failed to send response: ${error}`))) .finally(() => { - if (this._conn.requestHandlerAbortControllers.get(request.id) === abortController) { - this._conn.requestHandlerAbortControllers.delete(request.id); + if (capturedConnection.requestHandlerAbortControllers.get(request.id) === abortController) { + capturedConnection.requestHandlerAbortControllers.delete(request.id); } }); } @@ -1345,17 +1452,34 @@ export abstract class Protocol { return this._connection?.transport; } + /** + * Opaque liveness token for the connection serving the current dispatch. + * Role classes capture it when building per-request context senders + * (e.g. the Server's `ctx.mcpReq.elicitInput`), so a sender used after + * the delivering connection is gone rejects structurally instead of + * writing onto a successor connection. + */ + protected _liveConnectionState(): { readonly disposed: boolean } { + return this._conn; + } + /** * Closes the connection. */ async close(): Promise { const connection = this._connection; - await connection?.transport?.close(); - // A transport whose close() resolves without firing onclose would - // otherwise leave the instance wedged (still "connected" per the - // connect() guard). Run the teardown ourselves if it hasn't happened. - if (connection !== undefined && this._connection === connection) { - this._onclose(); + if (connection === undefined) { + return; + } + try { + await connection.transport?.close(); + } finally { + // A transport whose close() resolves without ever firing onclose + // must not wedge the instance: run teardown directly (a no-op + // when the transport's onclose callback already ran it). + if (!connection.disposed) { + this._onclose(); + } } } @@ -1495,7 +1619,8 @@ export abstract class Protocol { codec: WireCodec, request: Request, resultSchema: T, - options?: RequestOptions + options?: RequestOptions, + connection: Connection = this._conn ): Promise> { const { relatedRequestId, resumptionToken, onresumptiontoken, headers } = options ?? {}; // Flow start for non-complete result resolution: `maxTotalTimeout` @@ -1512,7 +1637,16 @@ export abstract class Protocol { reject(error); }; - if (!this.transport) { + // A caller holding a disposed connection (a handler that + // outlived the request's connection) fails structurally with the + // typed ConnectionClosed — never a write onto a successor + // connection's transport. + if (connection.disposed) { + earlyReject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); + return; + } + + if (!connection.transport) { earlyReject(new Error('Not connected')); return; } @@ -1547,7 +1681,7 @@ export abstract class Protocol { // combination — legacy era on any transport, modern era on stdio / // in-memory — keeps today's `notifications/cancelled` POST path // unchanged. - const streamCloseCancels = codec.era === MODERN_WIRE_REVISION && this.transport?.hasPerRequestStream === true; + const streamCloseCancels = codec.era === MODERN_WIRE_REVISION && connection.transport?.hasPerRequestStream === true; const requestAbort = streamCloseCancels ? new AbortController() : undefined; const messageId = this._requestMessageId++; @@ -1559,7 +1693,7 @@ export abstract class Protocol { }; if (options?.onprogress) { - this._conn.progressHandlers.set(messageId, options.onprogress); + connection.progressHandlers.set(messageId, options.onprogress); jsonrpcRequest.params = { ...request.params, _meta: { @@ -1581,22 +1715,28 @@ export abstract class Protocol { if (responseReceived) { return; } - this._conn.progressHandlers.delete(messageId); + connection.progressHandlers.delete(messageId); if (requestAbort === undefined) { - this.transport - ?.send( - this._envelopeOutbound({ - jsonrpc: '2.0', - method: 'notifications/cancelled', - params: { - requestId: messageId, - reason: String(reason) - } - }), - { relatedRequestId, resumptionToken, onresumptiontoken } - ) - .catch(error => this._onerror(new Error(`Failed to send cancellation: ${error}`))); + // A disposed connection swallows the cancellation + // notification: there is nothing left to cancel on a dead + // connection, and it must never reach a successor's + // transport. + if (!connection.disposed) { + connection + .send( + this._envelopeOutbound({ + jsonrpc: '2.0', + method: 'notifications/cancelled', + params: { + requestId: messageId, + reason: String(reason) + } + }), + { relatedRequestId, resumptionToken, onresumptiontoken } + ) + .catch(error => this._onerror(new Error(`Failed to send cancellation: ${error}`))); + } } else { // Modern-era per-request-stream transport: aborting the // request's underlying stream IS the spec cancel signal. @@ -1611,7 +1751,7 @@ export abstract class Protocol { reject(error); }; - this._conn.responseHandlers.set(messageId, response => { + connection.responseHandlers.set(messageId, response => { if (options?.signal?.aborted) { return; } @@ -1665,7 +1805,8 @@ export abstract class Protocol { codec, params === undefined ? { method: request.method } : { method: request.method, params }, resultSchema, - legOptions + legOptions, + connection ) }; return resolve(this._resolveNonCompleteResult(decoded, flow) as Promise>); @@ -1687,18 +1828,20 @@ export abstract class Protocol { const timeout = options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC; const timeoutHandler = () => cancel(new SdkError(SdkErrorCode.RequestTimeout, 'Request timed out', { timeout })); - this._conn.setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); - - this.transport!.send(outbound, { - relatedRequestId, - resumptionToken, - onresumptiontoken, - headers, - requestSignal: requestAbort?.signal - }).catch(error => { - this._conn.progressHandlers.delete(messageId); - reject(error); - }); + connection.setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); + + connection + .send(outbound, { + relatedRequestId, + resumptionToken, + onresumptiontoken, + headers, + requestSignal: requestAbort?.signal + }) + .catch(error => { + connection.progressHandlers.delete(messageId); + reject(error); + }); }).finally(() => { // Per-request cleanup that must run on every exit path. Consolidated // here so new exit paths added to the promise body can't forget it. @@ -1708,8 +1851,8 @@ export abstract class Protocol { options?.signal?.removeEventListener('abort', onAbort); } if (cleanupMessageId !== undefined) { - this._conn.responseHandlers.delete(cleanupMessageId); - this._conn.cleanupTimeout(cleanupMessageId); + connection.responseHandlers.delete(cleanupMessageId); + connection.cleanupTimeout(cleanupMessageId); } }); } @@ -1726,8 +1869,20 @@ export abstract class Protocol { * direct sends and related notifications (`ctx.mcpReq.notify`) alike * resolve through the instance's negotiated era at send time. */ - private async _notificationViaCodec(codec: WireCodec, notification: Notification, options?: NotificationOptions): Promise { - if (!this.transport) { + private async _notificationViaCodec( + codec: WireCodec, + notification: Notification, + options?: NotificationOptions, + connection: Connection = this._conn + ): Promise { + // A caller holding a disposed connection (a handler that outlived the + // request's connection) no-ops: the notification dies with its + // connection instead of reaching a successor's transport. + if (connection.disposed) { + return; + } + + if (!connection.transport) { throw new SdkError(SdkErrorCode.NotConnected, 'Not connected'); } @@ -1752,34 +1907,37 @@ export abstract class Protocol { if (canDebounce) { // If a notification of this type is already scheduled, do nothing. - if (this._conn.pendingDebouncedNotifications.has(notification.method)) { + if (connection.pendingDebouncedNotifications.has(notification.method)) { return; } // Mark this notification type as pending. - this._conn.pendingDebouncedNotifications.add(notification.method); + connection.pendingDebouncedNotifications.add(notification.method); // Schedule the actual send to happen in the next microtask. // This allows all synchronous calls in the current event loop tick to be coalesced. Promise.resolve().then(() => { // Un-mark the notification so the next one can be scheduled. - this._conn.pendingDebouncedNotifications.delete(notification.method); + connection.pendingDebouncedNotifications.delete(notification.method); - // SAFETY CHECK: If the connection was closed while this was pending, abort. - if (!this.transport) { + // SAFETY CHECK: the send is bound to the connection it was + // scheduled on — if that connection was closed (or replaced) + // while this was pending, abort rather than send on whatever + // connection exists now. + if (connection.disposed) { return; } // Send the notification, but don't await it here to avoid blocking. // Handle potential errors with a .catch(). - this.transport?.send(jsonrpcNotification, options).catch(error => this._onerror(error)); + connection.send(jsonrpcNotification, options).catch(error => this._onerror(error)); }); // Return immediately. return; } - await this.transport!.send(jsonrpcNotification, options); + await connection.send(jsonrpcNotification, options); } /** diff --git a/packages/core-internal/test/shared/connectionLifecycle.test.ts b/packages/core-internal/test/shared/connectionLifecycle.test.ts index fcf60738e6..ff120e8156 100644 --- a/packages/core-internal/test/shared/connectionLifecycle.test.ts +++ b/packages/core-internal/test/shared/connectionLifecycle.test.ts @@ -257,10 +257,10 @@ describe('g4: teardown ordering', () => { }); describe('r1: close() on a transport that never fires onclose (pin)', () => { - // close() runs the teardown itself when the transport's onclose callback - // has not (it used to rely on the onclose round-trip alone, wedging the - // instance on the AlreadyConnected guard). Pinned here so the teardown - // consolidation must preserve it. + // close() disposes the connection directly after transport.close() + // settles, so a transport that resolves close() without firing onclose + // cannot wedge the instance (it used to: `_transport` stayed set forever + // and every subsequent connect() threw AlreadyConnected). test('after await close(), the instance accepts a new connect() even when the transport never fired onclose', async () => { const protocol = createProtocol(); const transportA = new MockTransport('A'); @@ -310,33 +310,29 @@ describe('r2: late events from a failed transport cannot torpedo a later connect }); describe('r3: debounced notifications are connection-scoped', () => { - // DEFECT (today): the debounce microtask checks transport PRESENCE - // (`if (!this._transport) return`), not identity. A debounced - // notification scheduled on connection A, with close()+connect(B) both - // landing inside the microtask window, is delivered on B — a connection - // it was never sent on. - test.fails( - 'a notification debounced on connection A must not be sent on a connection B established inside the microtask window', - async () => { - const protocol = createProtocol({ debouncedNotificationMethods: ['test/debounced'] }); - const transportA = new MockTransport('A'); - const transportB = new MockTransport('B'); - - await protocol.connect(transportA); - - // Schedules the send as a microtask on connection A... - void protocol.notification({ method: 'test/debounced' }); - // ...and replaces the connection before that microtask runs. - // MockTransport.close() fires onclose synchronously, so close() has - // already torn A down when connect(B) installs the new transport; - // both happen before the debounce microtask fires. - const closing = protocol.close(); - const connecting = protocol.connect(transportB); - await Promise.all([closing, connecting]); - await flushMicrotasks(); - - expect(transportA.sentMessages).toHaveLength(0); - expect(transportB.sentMessages).toHaveLength(0); - } - ); + // The debounce microtask captures the connection it was scheduled on and + // aborts when that connection is disposed (it used to check transport + // PRESENCE, not identity, so a close()+connect(B) inside the microtask + // window delivered A's notification on B). + test('a notification debounced on connection A must not be sent on a connection B established inside the microtask window', async () => { + const protocol = createProtocol({ debouncedNotificationMethods: ['test/debounced'] }); + const transportA = new MockTransport('A'); + const transportB = new MockTransport('B'); + + await protocol.connect(transportA); + + // Schedules the send as a microtask on connection A... + void protocol.notification({ method: 'test/debounced' }); + // ...and replaces the connection before that microtask runs. + // MockTransport.close() fires onclose synchronously, so close() has + // already torn A down when connect(B) installs the new transport; + // both happen before the debounce microtask fires. + const closing = protocol.close(); + const connecting = protocol.connect(transportB); + await Promise.all([closing, connecting]); + await flushMicrotasks(); + + expect(transportA.sentMessages).toHaveLength(0); + expect(transportB.sentMessages).toHaveLength(0); + }); }); diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts index 661963f4ea..992e6bda22 100644 --- a/packages/server/src/server/server.ts +++ b/packages/server/src/server/server.ts @@ -360,6 +360,10 @@ export class Server extends Protocol { protected override buildContext(ctx: BaseContext, transportInfo?: MessageExtraInfo): ServerContext { // Only create http when there's actual HTTP transport info or auth info const hasHttpInfo = ctx.http || transportInfo?.request || transportInfo?.closeSSEStream || transportInfo?.closeStandaloneSSEStream; + // The connection serving this dispatch: the per-request senders below + // capture it so they fail structurally once it is gone (see the gate + // replacement on elicitInput/requestSampling). + const connection = this._liveConnectionState(); return { ...ctx, mcpReq: { @@ -400,16 +404,19 @@ export class Server extends Protocol { // session-wide stream to deliver it on. return ctx.mcpReq.notify({ method: 'notifications/message', params: { level, data, logger } }); }, - // Same abort gate as `ctx.mcpReq.send`: an aborted handler's - // senders must not reach a possibly-replaced transport. + // Same structural posture as `ctx.mcpReq.send`: once the + // connection this request arrived on is gone, its per-request + // senders must not reach the transport — `this.elicitInput` / + // `this.createMessage` read the live transport at send time, + // which may belong to a later connection. elicitInput: (params, options) => { - if (ctx.mcpReq.signal.aborted) { + if (connection.disposed) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } return this.elicitInput(params, options); }, requestSampling: (params, options) => { - if (ctx.mcpReq.signal.aborted) { + if (connection.disposed) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } return this.createMessage(params, options); diff --git a/test/e2e/requirements.ts b/test/e2e/requirements.ts index 349124727e..92d9c0c822 100644 --- a/test/e2e/requirements.ts +++ b/test/e2e/requirements.ts @@ -60,13 +60,7 @@ export const REQUIREMENTS: Record = { 'lifecycle:version:reject-unsupported': { entryExclusions: [{ arm: 'entryModern', reason: 'asserts-legacy-handshake' }], source: 'https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle#version-negotiation', - behavior: 'When server returns a protocolVersion the client does not support, connect rejects and the transport is closed.', - knownFailures: [ - { - transport: 'stdio', - note: 'connect rejects but client.transport is not cleared on stdio (other transports clear it)' - } - ] + behavior: 'When server returns a protocolVersion the client does not support, connect rejects and the transport is closed.' }, 'lifecycle:capability:experimental-passthrough': { transports: STATEFUL_TRANSPORTS, @@ -189,13 +183,7 @@ export const REQUIREMENTS: Record = { }, 'typescript:protocol:error:connection-closed': { source: 'sdk', - behavior: 'Closing the transport invokes onclose and rejects all in-flight requests with ErrorCode.ConnectionClosed.', - knownFailures: [ - { - transport: 'stdio', - note: 'in-process stdio does not fire client.onclose after close()' - } - ] + behavior: 'Closing the transport invokes onclose and rejects all in-flight requests with ErrorCode.ConnectionClosed.' }, 'protocol:error:internal-error': { source: 'https://modelcontextprotocol.io/specification/2025-11-25/basic#responses', From fc7fe83f7d718b17cd17e2de28ad4fb876c07fd6 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 9 Jul 2026 12:43:47 +0000 Subject: [PATCH 4/4] fix: gate per-request senders on request cancellation, not just connection disposal notifications/cancelled aborts a handler while its connection stays open; the structural disposed-connection check alone let a cancelled handler's per-request senders (notify/send, elicitInput/requestSampling) reach the wire. Restore the abort-signal check alongside the connection check at the sender gates, preserving the previous error shapes: notifications resolve as no-ops, requests reject with the typed ConnectionClosed error. Also clear the pending negotiated-version seed once a connection completes its own negotiation, so a post-close read cannot resurrect a stale pre-connect seed over the connection-negotiated value. Tests: pin cancellation-on-a-live-connection behavior (senders inert, nothing reaches the wire, the connection still serves fresh requests), unit-test the Connection owner's structural mechanisms (disposed send rejection, idempotent dispose, request-funnel early reject), and cover the seed lifecycle and the outbound _meta envelope attach seam. --- packages/core-internal/src/shared/protocol.ts | 30 ++- .../test/shared/connectionLifecycle.test.ts | 240 +++++++++++++++++- .../test/shared/protocol.test.ts | 38 +++ packages/server/src/server/server.ts | 16 +- 4 files changed, 303 insertions(+), 21 deletions(-) diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 74a7bf7686..b73dfc63a5 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -811,6 +811,10 @@ export abstract class Protocol { if (this._connection === undefined) { this._pendingNegotiatedProtocolVersion = version; } else { + // A connection that negotiates its own version supersedes the seed: + // clear it so a post-close read cannot resurrect a stale pre-connect + // value over what the connection actually negotiated. + this._pendingNegotiatedProtocolVersion = undefined; this._connection.negotiatedProtocolVersion = version; } } @@ -1239,7 +1243,16 @@ export abstract class Protocol { // Related sends resolve through the SAME instance era as every other // sender (the per-request/instance asymmetry is deliberately gone): // the codec is resolved at send time from the connection state. + // + // Two independent exits gate them: the delivering connection being + // gone (disposed — the structural check), and THIS request having + // been cancelled (`notifications/cancelled` aborts the handler while + // the connection stays live). Either way the sends must not reach + // the wire: notifications no-op, requests reject ConnectionClosed. const sendNotification = (notification: Notification, options?: NotificationOptions): Promise => { + if (capturedConnection.disposed || abortController.signal.aborted) { + return Promise.resolve(); + } return this._notificationViaCodec( this._resolveOutboundCodec(notification.method), notification, @@ -1255,6 +1268,9 @@ export abstract class Protocol { resultSchema: U, options?: RequestOptions ): Promise> => { + if (capturedConnection.disposed || abortController.signal.aborted) { + return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); + } return this._requestWithSchemaViaCodec( this._resolveOutboundCodec(r.method), r, @@ -1297,12 +1313,14 @@ export abstract class Protocol { // that overloaded property type. The cast is sound: this impl dispatches both overload paths via the // isStandardSchema guard, and sendRequest validates the result against the resolved schema either way. send: ((r: Request, schemaOrOptions?: StandardSchemaV1 | RequestOptions, maybeOptions?: RequestOptions) => { - // The disposed-connection check comes FIRST — before - // era/codec resolution — so a handler that outlived its - // connection always gets the documented ConnectionClosed - // rejection, never a synchronous era throw evaluated - // against whatever connection replaced this request's. - if (capturedConnection.disposed) { + // The disposed-connection / cancelled-request check comes + // FIRST — before era/codec resolution — so a handler that + // outlived its connection, or whose request was cancelled + // on a still-live connection, always gets the documented + // ConnectionClosed rejection, never a synchronous era + // throw evaluated against whatever connection replaced + // this request's. + if (capturedConnection.disposed || abortController.signal.aborted) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } // Related requests resolve through the instance era at diff --git a/packages/core-internal/test/shared/connectionLifecycle.test.ts b/packages/core-internal/test/shared/connectionLifecycle.test.ts index ff120e8156..a3e56053ef 100644 --- a/packages/core-internal/test/shared/connectionLifecycle.test.ts +++ b/packages/core-internal/test/shared/connectionLifecycle.test.ts @@ -2,19 +2,22 @@ * Connection lifecycle acceptance suite. * * Pins the behaviors any refactor of Protocol's connection-scoped state must - * preserve (g1-g4, r1, r2), and documents today's remaining teardown gap as - * an expected-failing test (r3 — flipped to a plain test by the fix). r2: - * the failed-start unwind restores the transport's own callbacks, so a late - * event from a failed transport cannot torpedo a later connection. + * preserve (g1-g5, r1, r2; r3 documented a teardown gap as an expected-failing + * test until the consolidation flipped it). r2: the failed-start unwind + * restores the transport's own callbacks, so a late event from a failed + * transport cannot torpedo a later connection. u1 unit-tests the structural + * mechanisms directly (disposed-connection send rejection, idempotent + * dispose, the request funnel's early disposed reject). */ import { describe, expect, test } from 'vitest'; import * as z from 'zod/v4'; import { SdkError, SdkErrorCode } from '../../src/errors/sdkErrors'; import type { BaseContext } from '../../src/shared/protocol'; -import { Protocol } from '../../src/shared/protocol'; +import { Protocol, setNegotiatedProtocolVersion } from '../../src/shared/protocol'; import type { Transport } from '../../src/shared/transport'; import type { JSONRPCMessage } from '../../src/types/index'; +import { codecForVersion } from '../../src/wire/codec'; class MockTransport implements Transport { id: string; @@ -55,8 +58,15 @@ function createProtocol(options?: ConstructorParameters[0]): Pr const flushMicrotasks = () => new Promise(resolve => setImmediate(resolve)); -/** Holds a request handler open and captures its context. */ -function installBlockingHandler(protocol: Protocol): { +/** + * Holds a request handler open and captures its context. `body` runs after + * release and produces the handler's outcome (default: resolve with `{}`; + * throw inside it to make the handler reject). + */ +function installBlockingHandler( + protocol: Protocol, + body: () => Record = () => ({}) +): { entered: Promise; release: () => void; ctx: () => BaseContext; @@ -74,7 +84,7 @@ function installBlockingHandler(protocol: Protocol): { ctxRef = ctx; enteredResolve(); await gate; - return {}; + return body(); }); return { entered, release, ctx: () => ctxRef! }; } @@ -126,6 +136,128 @@ describe('g1: aborted-handler send gates fire before era resolution', () => { }); }); +describe('g5: request cancellation on a live connection (pin)', () => { + // `notifications/cancelled` aborts the handler while the connection stays + // OPEN: `disposed` is false, but the cancelled request's per-request + // senders must behave exactly as they do after teardown — notify no-ops, + // send rejects ConnectionClosed, and nothing reaches the wire — while the + // connection itself keeps serving fresh requests. + test('after notifications/cancelled, the ctx senders are inert and the same connection serves a fresh request', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + const { entered, release, ctx } = installBlockingHandler(protocol); + + await protocol.connect(transport); + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + + transport.onmessage?.({ + jsonrpc: '2.0', + method: 'notifications/cancelled', + params: { requestId: 1, reason: 'client cancelled' } + }); + await flushMicrotasks(); + + // The request is cancelled but the connection is fully live. + expect(ctx().mcpReq.signal.aborted).toBe(true); + expect(protocol.transport).toBe(transport); + + // notify resolves as a no-op; send rejects ConnectionClosed on both + // the spec-method and the explicit-schema path. + await expect( + ctx().mcpReq.notify({ method: 'notifications/progress', params: { progressToken: 1, progress: 1 } }) + ).resolves.toBeUndefined(); + await expect(ctx().mcpReq.send({ method: 'ping' })).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + await expect(ctx().mcpReq.send({ method: 'custom/probe' }, z.object({}))).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + + // NOTHING reached the wire: no notification, no request, and — once + // the cancelled handler completes — no response for the cancelled + // request either. + release(); + await flushMicrotasks(); + expect(transport.sentMessages).toHaveLength(0); + + // The SAME live connection serves a fresh request normally. + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 2 }); + await flushMicrotasks(); + expect(transport.sentMessages).toHaveLength(1); + expect(transport.sentMessages[0]).toMatchObject({ jsonrpc: '2.0', id: 2, result: {} }); + }); + + test('a notifications/cancelled without a requestId is ignored', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + const { entered, release, ctx } = installBlockingHandler(protocol); + + await protocol.connect(transport); + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + + transport.onmessage?.({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { reason: 'no id' } }); + await flushMicrotasks(); + expect(ctx().mcpReq.signal.aborted).toBe(false); + + release(); + await flushMicrotasks(); + expect(transport.sentMessages).toHaveLength(1); + expect(transport.sentMessages[0]).toMatchObject({ jsonrpc: '2.0', id: 1, result: {} }); + }); + + test('a cancelled handler that rejects produces no error response on the wire', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + const { entered, release } = installBlockingHandler(protocol, () => { + throw new Error('handler failed after cancellation'); + }); + + await protocol.connect(transport); + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + transport.onmessage?.({ + jsonrpc: '2.0', + method: 'notifications/cancelled', + params: { requestId: 1, reason: 'client cancelled' } + }); + await flushMicrotasks(); + + release(); + await flushMicrotasks(); + expect(transport.sentMessages).toHaveLength(0); + }); +}); + +describe('ctx sender contract on a live request', () => { + test('ctx.mcpReq.send without a schema on a non-spec method throws TypeError (the gate lets live requests through)', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + const { entered, release, ctx } = installBlockingHandler(protocol); + + await protocol.connect(transport); + transport.onmessage?.({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await entered; + + // The overloads make this call unrepresentable at the type level (a + // non-spec method requires an explicit schema); the cast reaches the + // runtime guard behind them. + const sendUnchecked = ctx().mcpReq.send as (r: { method: string }) => Promise; + expect(() => sendUnchecked({ method: 'custom/no-schema' })).toThrow(TypeError); + + release(); + await flushMicrotasks(); + }); +}); + +describe('close() without a connection', () => { + test('close() before any connect() resolves as a no-op', async () => { + const protocol = createProtocol(); + await expect(protocol.close()).resolves.toBeUndefined(); + }); +}); + describe('g2: send-failure cleanup', () => { test('a transport.send rejection cleans up the progress handler and the request timeout', async () => { const protocol = createProtocol(); @@ -336,3 +468,95 @@ describe('r3: debounced notifications are connection-scoped', () => { expect(transportB.sentMessages).toHaveLength(0); }); }); + +describe('u1: Connection structural mechanisms (unit)', () => { + // The private Connection owner is deliberately not exported; these tests + // reach it through the instance, following the established internal-access + // pattern (see protocol.test.ts's testRequest helper). + type ConnectionInternals = { + disposed: boolean; + send(message: JSONRPCMessage): Promise; + dispose(): { responseHandlers: Map; requestHandlerAbortControllers: Map } | undefined; + }; + const internalConnection = (protocol: Protocol): ConnectionInternals => + (protocol as unknown as { _connection: ConnectionInternals })._connection; + + test('send on a disposed connection rejects ConnectionClosed and never reaches the transport', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + await protocol.connect(transport); + const connection = internalConnection(protocol); + + await protocol.close(); + expect(connection.disposed).toBe(true); + + await expect( + connection.send({ jsonrpc: '2.0', method: 'notifications/progress', params: { progressToken: 1, progress: 1 } }) + ).rejects.toSatisfy((error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed); + expect(transport.sentMessages).toHaveLength(0); + }); + + test('dispose() is idempotent: the first call returns the settlement material, the second returns undefined', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + await protocol.connect(transport); + const connection = internalConnection(protocol); + + const first = connection.dispose(); + expect(first).toBeDefined(); + expect(first!.responseHandlers).toBeInstanceOf(Map); + expect(first!.requestHandlerAbortControllers).toBeInstanceOf(Map); + + expect(connection.dispose()).toBeUndefined(); + expect(connection.disposed).toBe(true); + }); + + test('the request funnel early-rejects on a disposed connection before any transport interaction', async () => { + const protocol = createProtocol(); + const transport = new MockTransport('A'); + await protocol.connect(transport); + const connection = internalConnection(protocol); + await protocol.close(); + + const pending = ( + protocol as unknown as { + _requestWithSchemaViaCodec( + codec: unknown, + request: { method: string }, + resultSchema: unknown, + options: undefined, + connection: ConnectionInternals + ): Promise; + } + )._requestWithSchemaViaCodec(codecForVersion(undefined), { method: 'custom/probe' }, z.object({}), undefined, connection); + + await expect(pending).rejects.toSatisfy( + (error: unknown) => error instanceof SdkError && error.code === SdkErrorCode.ConnectionClosed + ); + expect(transport.sentMessages).toHaveLength(0); + }); +}); + +describe('negotiated-version seed lifecycle', () => { + test('a post-close read does not resurrect a stale seed over a connection-negotiated value', async () => { + const protocol = createProtocol(); + const readNegotiated = () => (protocol as unknown as { _negotiatedProtocolVersion: string | undefined })._negotiatedProtocolVersion; + + // An era bound before connect() lands on the pending seed... + setNegotiatedProtocolVersion(protocol, '2025-03-26'); + expect(readNegotiated()).toBe('2025-03-26'); + + // ...and seeds the new connection, which then completes its OWN + // negotiation at a different version. + const transport = new MockTransport('A'); + await protocol.connect(transport); + expect(readNegotiated()).toBe('2025-03-26'); + setNegotiatedProtocolVersion(protocol, '2025-06-18'); + expect(readNegotiated()).toBe('2025-06-18'); + + // After close there is no connection: the read must NOT fall back to + // the pre-connect seed the connection's own negotiation superseded. + await protocol.close(); + expect(readNegotiated()).toBeUndefined(); + }); +}); diff --git a/packages/core-internal/test/shared/protocol.test.ts b/packages/core-internal/test/shared/protocol.test.ts index 2ecdc40adc..a0cfd33903 100644 --- a/packages/core-internal/test/shared/protocol.test.ts +++ b/packages/core-internal/test/shared/protocol.test.ts @@ -1170,3 +1170,41 @@ describe('inbound protocol-version mismatch (−32022): the error data lists eve await protocol.close(); }); }); + +describe('outbound _meta envelope attach', () => { + // Exercises the base-class attach seam directly (the Client's modern-era + // override is covered in the client package): the envelope merges into + // existing params with user `_meta` keys winning, and materializes + // `params._meta` from nothing when the message carried no params. + class EnvelopeProtocol extends TestProtocolImpl { + protected override _outboundMetaEnvelope(): Readonly> | undefined { + return { 'io.modelcontextprotocol/test-key': 'envelope' }; + } + } + + test('envelope keys merge under user _meta and materialize params when absent', async () => { + const protocol = new EnvelopeProtocol(); + const transport = new MockTransport(); + const sent: JSONRPCMessage[] = []; + transport.send = async message => { + sent.push(message); + }; + await protocol.connect(transport); + + await protocol.notification({ + method: 'test/withParams', + params: { value: 1, _meta: { 'io.modelcontextprotocol/test-key': 'user' } } + }); + await protocol.notification({ method: 'test/bare' }); + + expect((sent[0] as JSONRPCNotification).params).toEqual({ + value: 1, + _meta: { 'io.modelcontextprotocol/test-key': 'user' } + }); + expect((sent[1] as JSONRPCNotification).params).toEqual({ + _meta: { 'io.modelcontextprotocol/test-key': 'envelope' } + }); + + await protocol.close(); + }); +}); diff --git a/packages/server/src/server/server.ts b/packages/server/src/server/server.ts index 992e6bda22..a3cda61b7b 100644 --- a/packages/server/src/server/server.ts +++ b/packages/server/src/server/server.ts @@ -404,19 +404,21 @@ export class Server extends Protocol { // session-wide stream to deliver it on. return ctx.mcpReq.notify({ method: 'notifications/message', params: { level, data, logger } }); }, - // Same structural posture as `ctx.mcpReq.send`: once the - // connection this request arrived on is gone, its per-request - // senders must not reach the transport — `this.elicitInput` / - // `this.createMessage` read the live transport at send time, - // which may belong to a later connection. + // Same double gate as `ctx.mcpReq.send`: once the connection + // this request arrived on is gone (disposed), OR this request + // was cancelled while the connection stayed live, its + // per-request senders must not reach the transport — + // `this.elicitInput` / `this.createMessage` read the live + // transport at send time, which may belong to a later + // connection. elicitInput: (params, options) => { - if (connection.disposed) { + if (connection.disposed || ctx.mcpReq.signal.aborted) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } return this.elicitInput(params, options); }, requestSampling: (params, options) => { - if (connection.disposed) { + if (connection.disposed || ctx.mcpReq.signal.aborted) { return Promise.reject(new SdkError(SdkErrorCode.ConnectionClosed, 'Request was cancelled')); } return this.createMessage(params, options);