diff --git a/.changeset/inbound-processing-completion.md b/.changeset/inbound-processing-completion.md new file mode 100644 index 0000000000..c085a9650c --- /dev/null +++ b/.changeset/inbound-processing-completion.md @@ -0,0 +1,8 @@ +--- +'@modelcontextprotocol/server': minor +'@modelcontextprotocol/client': minor +--- + +Add local inbound-processing completion metadata for transport adapters. Completion follows request handlers, response transmission and internal request cleanup; cancellation alone does not complete a request. The stdio entry preserves this metadata through opening and error paths and holds listen requests until subscription teardown. Export a stdio capability marker so reusable bounded admission can reject SDK versions without this contract. Completion remains independent of OpenTelemetry span outcomes and permits request-ID reuse from the callback after controller cleanup. + +Queue subscription acknowledgements before routed changes and retain pending writes without blocking unrelated cancellation. Client completion also waits for asynchronous cache invalidation while notification handlers dispatch immediately. diff --git a/docs/inbound-processing-completion.md b/docs/inbound-processing-completion.md new file mode 100644 index 0000000000..715b0fb941 --- /dev/null +++ b/docs/inbound-processing-completion.md @@ -0,0 +1,11 @@ +# Local inbound-processing completion + +Transport adapters may attach `MessageExtraInfo.onMessageProcessed` when forwarding an inbound message to a connected protocol instance or the `serveStdio` entry. The callback runs once after SDK-owned processing, including asynchronous handlers, response `send()` settlement and request-controller cleanup. Failed handlers, failed writes and entry-level rejection paths also complete after their owned work. The official stdio transport joins the callback for each Node Writable chunk, including failed callbacks. Backpressure or stream error events alone do not complete that write. A cancellation notification or transport close does not settle an unfinished handler or write. + +For `subscriptions/listen`, the opening acknowledgement is queued before routed change notifications and is not completion. The stdio entry waits for cancellation or teardown, including a pending acknowledgement and any graceful terminal write it owns. Pending subscription writes retain admission without blocking unrelated inbound cancellation. `supportsInboundMessageCompletion` is exported from `@modelcontextprotocol/server/stdio` for adapters that require this contract before accepting input. + +Client change notifications include asynchronous cache eviction in completion. Notification handlers start immediately while the cache store is pending; completion waits for both the handler and eviction from all affected cache partitions. + +The metadata is local and is never serialized. It is suitable for returning a local admission credit or reusing an inbound request identifier. It does not prove delivery to the remote peer, remote effects, resource cleanup started outside the handler's returned promise, or application-owned physical settlement. Applications must join their own work before their handlers finish. Observers should not throw; observer errors are reported through the existing error callback. + +Completion is independent of OpenTelemetry span outcomes: a span may end before a response write settles or when cancellation starts, while the message remains admitted until its owned work finishes. Completion is not FIFO across messages. An adapter may synchronously reuse a completed request ID from its callback; the previous request controller has already been removed. diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index f071b13e50..73751e4f80 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -2206,11 +2206,12 @@ export class Client extends Protocol { * `setNotificationHandler` / `fallbackNotificationHandler` instead of * being silently swallowed. */ - protected override _onnotification(raw: JSONRPCNotification, extra?: MessageExtraInfo): void { + protected override _onnotification(raw: JSONRPCNotification, extra?: MessageExtraInfo): void | Promise { + const evictions: Promise[] = []; // Response-cache invalidation: a `list_changed` notification means the - // matching cached list result is stale. Evict (do NOT refetch) before - // dispatch so a handler that reaches the cache observes the cleared - // entry. Runs regardless of whether the user + // matching cached list result is stale. Start eviction (do NOT refetch) + // before dispatch so synchronous stores are cleared when a handler runs. + // Runs regardless of whether the user // configured `listChanged` — derived views (the tool index, output // validators) must drop the stale entry either way. `raw.method` is // server-controlled; guard with `Object.hasOwn` so an inherited @@ -2222,10 +2223,9 @@ export class Client extends Protocol { // `resources/read` body for this URI is stale on receipt; drop it // from BOTH partitions so the next `readResource` for the same URI // refetches even within TTL (the documented subscribe → updated → - // re-read flow). Fire-and-forget like the `list_changed` path — - // dispatch must not block on an async store. + // re-read flow). Start eviction before dispatch and join both below. const uri = (raw.params as { uri?: unknown } | undefined)?.uri; - if (typeof uri === 'string') void this._cache.evictKey('resources/read', uri); + if (typeof uri === 'string') evictions.push(this._cache.evictKey('resources/read', uri)); } else if (evicted !== undefined) { for (const method of evicted) { // `evict()` bumps the generation FIRST and unconditionally @@ -2235,9 +2235,9 @@ export class Client extends Protocol { // store keep their entries. Store failures are reported via // `onerror` inside `evict()` and the call resolves, so // dispatch (and the user's `listChanged` handler) runs - // regardless. Fire-and-forget — dispatch must not block on - // an async store. - void this._cache.evict(method); + // regardless. Dispatch starts immediately; completion joins + // the asynchronous store work below. + evictions.push(this._cache.evict(method)); } } if (raw.method === 'notifications/subscriptions/acknowledged') { @@ -2266,7 +2266,8 @@ export class Client extends Protocol { return; } } - super._onnotification(raw, extra); + const handled = super._onnotification(raw, extra); + return evictions.length === 0 ? handled : Promise.all([...evictions, handled]).then(() => {}); } /** @@ -2279,7 +2280,7 @@ export class Client extends Protocol { * NUMBER and never holds a listen id, so passing a string-id response * through would surface as "unknown message ID" via `onerror`). */ - protected override _onresponse(response: JSONRPCResponse): void { + protected override _onresponse(response: JSONRPCResponse): void | Promise { const id = response.id; const entry = typeof id === 'string' ? this._listenState.get(id) : undefined; if (entry !== undefined) { @@ -2308,7 +2309,7 @@ export class Client extends Protocol { } return; } - super._onresponse(response); + return super._onresponse(response); } /** diff --git a/packages/client/test/client/responseCache.test.ts b/packages/client/test/client/responseCache.test.ts index e2e2326705..404b934b7f 100644 --- a/packages/client/test/client/responseCache.test.ts +++ b/packages/client/test/client/responseCache.test.ts @@ -7,9 +7,11 @@ * `toolDefinition` hit/miss and re-derivation only on a stamp change; the * generation guard skipping a stale write. */ -import type { JSONRPCMessage, JSONRPCRequest, Tool } from '@modelcontextprotocol/core-internal'; +import { setImmediate as nextTurn } from 'node:timers/promises'; + +import type { JSONRPCMessage, JSONRPCRequest, Tool, Transport } from '@modelcontextprotocol/core-internal'; import { InMemoryTransport, SdkError, SdkErrorCode } from '@modelcontextprotocol/core-internal'; -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; import { Client } from '../../src/client/client'; import type { CacheEntry, ResponseCacheStore } from '../../src/client/responseCache'; @@ -422,6 +424,50 @@ const cacheOf = (client: Client): ClientResponseCache => (client as unknown as { const toolDef = (client: Client, name: string): Promise => cacheOf(client).toolDefinition(name); describe('Client response-cache substrate', () => { + it.each([ + { method: 'notifications/tools/list_changed', params: {}, deletesPerPartition: 1 }, + { method: 'notifications/resources/list_changed', params: {}, deletesPerPartition: 2 }, + { method: 'notifications/resources/updated', params: { uri: 'file:///a' }, deletesPerPartition: 1 } + ] as const)( + 'completion joins async cache eviction for $method without delaying dispatch', + async ({ method, params, deletesPerPartition }) => { + const store = new InMemoryResponseCacheStore(); + const pending: (() => void)[] = []; + const deleted: string[] = []; + (store as ResponseCacheStore).delete = async key => { + await new Promise(resolve => pending.push(resolve)); + deleted.push(key.partition!); + }; + const transport: Transport = (await scriptedModernServer([])).clientTx; + const client = modernClient(store, { cachePartition: 'alice' }); + const handled = vi.fn(); + const complete = vi.fn(); + client.setNotificationHandler(method, handled); + await client.connect(transport); + try { + transport.onmessage!({ jsonrpc: '2.0', method, params }, { onMessageProcessed: complete }); + await nextTurn(); + expect(handled).toHaveBeenCalledOnce(); + expect(pending).toHaveLength(deletesPerPartition); + expect(complete).not.toHaveBeenCalled(); + pending.splice(0).forEach(resolve => resolve()); + await nextTurn(); + expect(deleted).toEqual(Array(deletesPerPartition).fill(part('alice'))); + expect(pending).toHaveLength(deletesPerPartition); + expect(complete).not.toHaveBeenCalled(); + pending.splice(0).forEach(resolve => resolve()); + await nextTurn(); + expect(deleted).toEqual([...Array(deletesPerPartition).fill(part('alice')), ...Array(deletesPerPartition).fill(part())]); + expect(complete).toHaveBeenCalledOnce(); + } finally { + pending.splice(0).forEach(resolve => resolve()); + await nextTurn(); + pending.splice(0).forEach(resolve => resolve()); + await client.close(); + } + } + ); + it('listTools() with no cursor reads every page, writes one cache entry; listTools({cursor}) stays per-page and does not write', async () => { const store = new InMemoryResponseCacheStore(); const { clientTx, listCount } = await scriptedModernServer([[TOOL_A], [TOOL_B]]); diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index ea86863a4f..b09d5e4569 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -589,7 +589,7 @@ export abstract class Protocol { private _customSchemaRequestMethods = new Set(); private _requestHandlerAbortControllers: Map = new Map(); private _notificationHandlers: Map Promise> = new Map(); - private _responseHandlers: Map void> = new Map(); + private _responseHandlers: Map void | Promise> = new Map(); private _progressHandlers: Map = new Map(); private _timeoutInfo: Map = new Map(); private _pendingDebouncedNotifications = new Set(); @@ -834,15 +834,48 @@ export abstract class Protocol { 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)}`)); + let processed = false; + const complete = () => { + if (processed) return; + processed = true; + try { + extra?.onMessageProcessed?.(); + } catch (error) { + this._onerror(error instanceof Error ? error : new Error('Inbound completion callback failed')); + } + }; + let work: void | Promise | undefined; + try { + _onmessage?.(message, extra); + if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) { + work = this._onresponse(message); + } else if (isJSONRPCRequest(message)) { + work = this._onrequest(message, extra); + } else if (isJSONRPCNotification(message)) { + work = this._onnotification(message, extra); + } else { + this._onerror(new Error(`Unknown message type: ${JSON.stringify(message)}`)); + } + } catch (error) { + // Without a completion observer retain the historical dispatch behavior. + if (!extra?.onMessageProcessed) throw error; + this._onerror(error instanceof Error ? error : new Error('Inbound dispatch failed')); + work = isJSONRPCRequest(message) + ? this._transport?.send({ + jsonrpc: '2.0', + id: message.id, + error: { code: ProtocolErrorCode.InternalError, message: 'Internal server error' } + }) + : undefined; + } + if (extra?.onMessageProcessed) { + void Promise.resolve(work).then(complete, error => { + try { + this._onerror(error instanceof Error ? error : new Error('Inbound processing failed')); + } finally { + complete(); + } + }); } }; @@ -898,7 +931,7 @@ export abstract class Protocol { * an override that consumes only what it owns and falls through to base * dispatch for everything else. */ - protected _onnotification(rawNotification: JSONRPCNotification, extra?: MessageExtraInfo): void { + protected _onnotification(rawNotification: JSONRPCNotification, extra?: MessageExtraInfo): void | Promise { // Hide wire-only material from notification handlers too — but ONLY // the reserved envelope `_meta` keys (the retry params names are // reserved on requests, not notifications). There is no @@ -960,7 +993,7 @@ export abstract class Protocol { ); // Starting with Promise.resolve() puts any synchronous errors into the monad as well. - runInScope(scope, () => + return runInScope(scope, () => Promise.resolve().then(() => (handler === undefined ? fallback!(notification) : handler(notification, codec))) ).then( () => scope?.end({ type: 'delivered' }), @@ -971,7 +1004,7 @@ export abstract class Protocol { ); } - private _onrequest(rawRequest: JSONRPCRequest, extra?: MessageExtraInfo): void { + private _onrequest(rawRequest: JSONRPCRequest, extra?: MessageExtraInfo): void | Promise { // Lift wire-only material before dispatch: handlers (including the // fallback handler and the per-method schema parse) see exactly the // 2025-era shape; the envelope and retry fields surface via ctx. @@ -1010,7 +1043,9 @@ export abstract class Protocol { id: request.id, error: { code, message, ...(data !== undefined && { data }) } }; - capturedTransport?.send(errorResponse).catch(error => this._onerror(new Error(`Failed to send an error response: ${error}`))); + return capturedTransport + ?.send(errorResponse) + .catch(error => this._onerror(new Error(`Failed to send an error response: ${error}`))); }; // Edge→instance handoff check: a classification that disagrees with @@ -1031,7 +1066,7 @@ export abstract class Protocol { // only the fallback for classifications without an exact // revision. const requested = extra.classification.revision ?? classified; - sendErrorResponse(ProtocolErrorCode.UnsupportedProtocolVersion, `Unsupported protocol version: ${requested}`, { + return sendErrorResponse(ProtocolErrorCode.UnsupportedProtocolVersion, `Unsupported protocol version: ${requested}`, { // Per spec, `supported` is the full list of protocol // versions the receiver supports — not just the version // this connection is on — so the peer can pick a mutually @@ -1042,7 +1077,6 @@ export abstract class Protocol { supported: this._supportedProtocolVersions, requested }); - return; } } @@ -1058,15 +1092,13 @@ export abstract class Protocol { !codec.hasRequestMethod(request.method) && !(isExtensionReusedRequestMethod(request.method) && this._customSchemaRequestMethods.has(request.method)) ) { - sendErrorResponse(ProtocolErrorCode.MethodNotFound, 'Method not found'); - return; + return sendErrorResponse(ProtocolErrorCode.MethodNotFound, 'Method not found'); } const handler = this._requestHandlers.get(request.method) ?? this.fallbackRequestHandler; if (handler === undefined) { - sendErrorResponse(ProtocolErrorCode.MethodNotFound, 'Method not found'); - return; + return sendErrorResponse(ProtocolErrorCode.MethodNotFound, 'Method not found'); } // Envelope enforcement: the 2026 era requires the per-request `_meta` @@ -1080,8 +1112,7 @@ export abstract class Protocol { // only the −32601-over-−32602 rule.) const envelopeError = codec.checkInboundEnvelope(lifted); if (envelopeError !== undefined) { - sendErrorResponse(ProtocolErrorCode.InvalidParams, envelopeError); - return; + return sendErrorResponse(ProtocolErrorCode.InvalidParams, envelopeError); } // Related sends resolve through the SAME instance era as every other @@ -1151,7 +1182,6 @@ export abstract class Protocol { }, http: extra?.authInfo ? { authInfo: extra.authInfo } : undefined }; - const ctx = this.buildContext(baseCtx, extra); if (scope !== undefined) { abortController.signal.addEventListener('abort', () => scope.end(outcomeForAbort(abortController.signal.reason)), { @@ -1161,7 +1191,7 @@ export abstract class Protocol { // Starting with Promise.resolve() puts any synchronous errors into the monad as well. // Only the handler runs in the instrumentation's context; encoding and sending do not need it. - runInScope(scope, () => Promise.resolve().then(() => handler(request, ctx))) + return runInScope(scope, () => Promise.resolve().then(() => handler(request, this.buildContext(baseCtx, extra)))) .then( async result => { if (abortController.signal.aborted) { @@ -1182,8 +1212,7 @@ export abstract class Protocol { encoded = codec.encodeResult(request.method, result, this._outboundServerInfo()); } catch (error) { this._onerror(new Error(`Failed to encode result for ${request.method}: ${error}`)); - sendErrorResponse(ProtocolErrorCode.InternalError, 'Internal error'); - return; + return sendErrorResponse(ProtocolErrorCode.InternalError, 'Internal error'); } scope?.end({ type: 'result', result: encoded }); @@ -1265,7 +1294,7 @@ export abstract class Protocol { * only what it owns and falls through to base dispatch for everything * else. */ - protected _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void { + protected _onresponse(response: JSONRPCResponse | JSONRPCErrorResponse): void | Promise { const messageId = Number(response.id); const handler = this._responseHandlers.get(messageId); @@ -1279,10 +1308,10 @@ export abstract class Protocol { this._progressHandlers.delete(messageId); if (isJSONRPCResultResponse(response)) { - handler(response); + return handler(response); } else { const error = ProtocolError.fromError(response.error.code, response.error.message, response.error.data); - handler(error); + return handler(error); } } @@ -1624,15 +1653,15 @@ export abstract class Protocol { }; // Embedded input requests are fulfilled in this leg's context, so their spans nest under it. const inputRequired = decoded; - return resolve( + return ( runInScope(scope, () => this._resolveNonCompleteResult(inputRequired, flow)) as Promise< StandardSchemaV1.InferOutput > - ); + ).then(resolve, reject); } const result = decoded.result; - validateStandardSchema(resultSchema, result).then( + return validateStandardSchema(resultSchema, result).then( parseResult => { if (parseResult.success) { scope?.end({ type: 'result', result: response.result }); diff --git a/packages/core-internal/src/types/types.ts b/packages/core-internal/src/types/types.ts index 61cc7f9b36..7996f04cfd 100644 --- a/packages/core-internal/src/types/types.ts +++ b/packages/core-internal/src/types/types.ts @@ -901,6 +901,17 @@ export interface MessageClassification { * Extra information about a message. */ export interface MessageExtraInfo { + /** + * Local transport-side completion of this inbound message's SDK processing. + * Called once after dispatch, awaited response sends and request-controller + * cleanup, including cancellation with no response. For a stdio listen, + * completion waits for subscription teardown rather than its acknowledgement. + * This does not prove downstream remote effects or consumer-detached work + * have finished. Never serialized onto the wire. Callback errors are reported + * through the receiving protocol/entry's error handler. + */ + onMessageProcessed?: () => void; + /** * The original HTTP request. */ diff --git a/packages/core-internal/test/shared/inboundProcessing.test.ts b/packages/core-internal/test/shared/inboundProcessing.test.ts new file mode 100644 index 0000000000..04a3ca2847 --- /dev/null +++ b/packages/core-internal/test/shared/inboundProcessing.test.ts @@ -0,0 +1,207 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; + +import { vi } from 'vitest'; +import * as z from 'zod/v4'; +import type { BaseContext } from '../../src/shared/protocol'; +import type { OperationScope } from '../../src/shared/instrumentation'; +import { Protocol, setProtocolInstrumentation } from '../../src/shared/protocol'; +import type { Transport } from '../../src/shared/transport'; +import type { JSONRPCMessage, MessageExtraInfo } from '../../src/types/index'; + +class TestProtocol extends Protocol { + protected assertCapabilityForMethod(): void {} + protected assertNotificationCapability(): void {} + protected assertRequestHandlerCapability(): void {} + protected buildContext(ctx: BaseContext): BaseContext { + return ctx; + } +} +function deferred() { + let resolve!: () => void; + const promise = new Promise(r => { + resolve = r; + }); + return { promise, resolve }; +} +async function fixture(instrumented = false) { + const protocol = new TestProtocol(); + const trace = new AsyncLocalStorage(); + const spanEnd = vi.fn(); + if (instrumented) { + const scope: OperationScope = { run: fn => trace.run(true, fn), end: spanEnd }; + setProtocolInstrumentation(protocol, { receive: () => scope, send: () => scope }); + } + const transport: Transport = { + start: async () => {}, + close: async () => transport.onclose?.(), + send: vi.fn(async () => {}) + }; + await protocol.connect(transport); + const receive = (message: JSONRPCMessage, extra?: MessageExtraInfo) => transport.onmessage!(message, extra); + return { protocol, transport, receive, trace, spanEnd }; +} +const request = { jsonrpc: '2.0' as const, id: 1, method: 'test/work', params: {} }; + +test.each([false, true])('completion waits for handler, response send and controller cleanup (instrumented: %s)', async instrumented => { + const { protocol, transport, receive, trace, spanEnd } = await fixture(instrumented); + const handler = deferred(); + const write = deferred(); + const complete = vi.fn(); + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, async () => { + await handler.promise; + expect(trace.getStore()).toBe(instrumented ? true : undefined); + return {}; + }); + transport.send = vi.fn(() => write.promise); + receive(request, { onMessageProcessed: complete }); + await Promise.resolve(); + expect(complete).not.toHaveBeenCalled(); + handler.resolve(); + await vi.waitFor(() => expect(transport.send).toHaveBeenCalledOnce()); + if (instrumented) expect(spanEnd).toHaveBeenCalledExactlyOnceWith({ type: 'result', result: {} }); + expect(complete).not.toHaveBeenCalled(); + write.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); +}); + +test('cancellation and close do not settle an unfinished handler', async () => { + const { protocol, transport, receive } = await fixture(); + const handler = deferred(); + const complete = vi.fn(); + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, async () => { + await handler.promise; + return {}; + }); + receive(request, { onMessageProcessed: complete }); + await Promise.resolve(); + await Promise.resolve(); + receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 1 } }); + await protocol.close(); + expect(complete).not.toHaveBeenCalled(); + handler.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + expect(transport.send).not.toHaveBeenCalled(); +}); + +test.each(['unknown', 'handler-error', 'write-error'])('completion covers %s failure', async kind => { + const { protocol, transport, receive } = await fixture(); + const complete = vi.fn(); + if (kind !== 'unknown') + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, () => { + if (kind === 'handler-error') throw new Error('failure'); + return {}; + }); + if (kind === 'write-error') { + protocol.onerror = () => {}; + transport.send = vi.fn(async () => { + throw new Error('closed'); + }); + } + receive(request, { onMessageProcessed: complete }); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); +}); + +test.each([false, true])('notification completion waits for its asynchronous handler (instrumented: %s)', async instrumented => { + const { protocol, receive, spanEnd } = await fixture(instrumented); + const handler = deferred(); + const complete = vi.fn(); + protocol.setNotificationHandler('test/notification', { params: z.object({}) }, () => handler.promise); + receive({ jsonrpc: '2.0', method: 'test/notification', params: {} }, { onMessageProcessed: complete }); + await Promise.resolve(); + expect(complete).not.toHaveBeenCalled(); + expect(spanEnd).not.toHaveBeenCalled(); + handler.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + if (instrumented) expect(spanEnd).toHaveBeenCalledExactlyOnceWith({ type: 'delivered' }); +}); + +test('callback failures are reported once after processing', async () => { + const { protocol, receive } = await fixture(); + const error = new Error('observer'); + const report = vi.fn(); + protocol.onerror = report; + receive(request, { + onMessageProcessed: () => { + throw error; + } + }); + await vi.waitFor(() => expect(report).toHaveBeenCalledWith(error)); + expect(report).toHaveBeenCalledOnce(); +}); + +test.each([false, true])('inbound response completion joins asynchronous SDK result validation (instrumented: %s)', async instrumented => { + const f = await fixture(instrumented); + const held = deferred(); + const entered = deferred(); + const complete = vi.fn(); + const schema = z.object({ value: z.string() }).superRefine(async () => { + entered.resolve(); + await held.promise; + }); + const result = f.protocol.request({ method: 'test/work' }, schema); + await vi.waitFor(() => expect(f.transport.send).toHaveBeenCalledOnce()); + const outbound = vi.mocked(f.transport.send).mock.calls[0]![0]; + if (!('id' in outbound) || outbound.id === undefined) throw new Error('expected outbound request'); + f.receive({ jsonrpc: '2.0', id: outbound.id, result: { value: 'accepted' } }, { onMessageProcessed: complete }); + await entered.promise; + expect(complete).not.toHaveBeenCalled(); + expect(f.spanEnd).not.toHaveBeenCalled(); + held.resolve(); + expect(await result).toEqual({ value: 'accepted' }); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + if (instrumented) expect(f.spanEnd).toHaveBeenCalledExactlyOnceWith({ type: 'result', result: { value: 'accepted' } }); +}); + +test.each(['cancelled', 'closed'] as const)('an instrumented %s request stays admitted until its handler settles', async outcome => { + const { protocol, transport, receive, spanEnd } = await fixture(true); + const entered = deferred(); + const handler = deferred(); + const complete = vi.fn(); + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, async () => { + entered.resolve(); + await handler.promise; + return {}; + }); + receive(request, { onMessageProcessed: complete }); + await entered.promise; + if (outcome === 'cancelled') { + receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: request.id } }); + } else { + await protocol.close(); + } + await vi.waitFor(() => expect(spanEnd).toHaveBeenCalledWith({ type: outcome })); + expect(complete).not.toHaveBeenCalled(); + handler.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + expect(transport.send).not.toHaveBeenCalled(); +}); + +test('completion can immediately reuse a request ID without retaining the previous controller', async () => { + const { protocol, transport, receive } = await fixture(true); + const entered = deferred(); + const handler = deferred(); + const signals: AbortSignal[] = []; + const secondComplete = vi.fn(); + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, async (_params, ctx) => { + signals.push(ctx.mcpReq.signal); + if (signals.length === 2) { + entered.resolve(); + await handler.promise; + } + return {}; + }); + const firstComplete = vi.fn(() => receive(request, { onMessageProcessed: secondComplete })); + receive(request, { onMessageProcessed: firstComplete }); + await entered.promise; + expect(firstComplete).toHaveBeenCalledOnce(); + expect(secondComplete).not.toHaveBeenCalled(); + expect(signals).toHaveLength(2); + await protocol.close(); + expect(signals[0]!.aborted).toBe(false); + expect(signals[1]!.aborted).toBe(true); + expect(secondComplete).not.toHaveBeenCalled(); + handler.resolve(); + await vi.waitFor(() => expect(secondComplete).toHaveBeenCalledOnce()); + expect(firstComplete).toHaveBeenCalledOnce(); + expect(transport.send).toHaveBeenCalledOnce(); +}); diff --git a/packages/server/src/server/serveStdio.ts b/packages/server/src/server/serveStdio.ts index 1cabf8ed9b..3010422d95 100644 --- a/packages/server/src/server/serveStdio.ts +++ b/packages/server/src/server/serveStdio.ts @@ -163,7 +163,7 @@ class StdioConnectionChannel implements Transport { * modern-era listen router to fan a change notification out onto the * active subscriptions instead of broadcasting it unsolicited. */ - private readonly _outboundIntercept?: (message: JSONRPCMessage) => 'handled' | undefined + private readonly _outboundIntercept?: (message: JSONRPCMessage) => 'handled' | undefined | Promise<'handled' | undefined> ) {} async start(): Promise { @@ -186,7 +186,10 @@ class StdioConnectionChannel implements Transport { // sends are dropped. return; } - if (this._outboundIntercept?.(message) === 'handled') { + if ((await this._outboundIntercept?.(message)) === 'handled') { + return; + } + if (this._closed) { return; } return this._wire.send(message, options); @@ -199,6 +202,7 @@ class StdioConnectionChannel implements Transport { /** Forwards one inbound message to the connected instance. */ deliver(message: JSONRPCMessage, extra?: MessageExtraInfo): void { if (this._closed) { + extra?.onMessageProcessed?.(); return; } if (isJSONRPCRequest(message)) { @@ -414,7 +418,7 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions const listenRouter = new StdioListenRouter(options.maxSubscriptions ?? DEFAULT_MAX_SUBSCRIPTIONS); /** Outbound intercept installed on a modern instance's channel. */ - const modernOutboundIntercept = (message: JSONRPCMessage): 'handled' | undefined => { + const modernOutboundIntercept = async (message: JSONRPCMessage): Promise<'handled' | undefined> => { if (!isJSONRPCNotification(message)) return undefined; const routed = listenRouter.routeOutbound(message); if (routed === 'passthrough') return undefined; @@ -423,9 +427,23 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions // it is dropped — the modern era never delivers an un-requested // change type unsolicited). Nothing else from the instance is // affected. - for (const stamped of routed) { - void wire.send({ jsonrpc: '2.0', ...stamped }).catch(error => reportError(toError(error))); - } + const deliveries = routed.map(stamped => { + const id = stamped.params?._meta?.['io.modelcontextprotocol/subscriptionId']; + const record = typeof id === 'number' || typeof id === 'string' ? listenCompletions.get(id) : undefined; + let delivery: Promise; + try { + delivery = wire.send({ jsonrpc: '2.0', ...stamped }).catch(error => reportError(toError(error))); + } catch (error) { + reportError(toError(error)); + delivery = Promise.resolve(); + } + record?.deliveries.add(delivery); + void delivery.then(() => record?.deliveries.delete(delivery)); + return delivery; + }); + // Start the routed snapshot without an intervening await: cancellation cannot + // retire an id and then observe a deferred new write from this same snapshot. + await Promise.all(deliveries); return 'handled'; }; @@ -434,7 +452,20 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions * Returns `true` when the message was served at the entry and must NOT * be delivered to the pinned instance. */ - const tryServeListen = async (message: JSONRPCMessage): Promise => { + const listenCompletions = new Map< + RequestId, + { complete?: () => void; acknowledgement: Promise; deliveries: Set> } + >(); + const completeListen = async (id: RequestId): Promise => { + const record = listenCompletions.get(id); + if (!record) return; + await record.acknowledgement; + await Promise.all(record.deliveries); + if (listenCompletions.get(id) !== record) return; + listenCompletions.delete(id); + record.complete?.(); + }; + const tryServeListen = async (message: JSONRPCMessage, extra?: MessageExtraInfo): Promise => { if (isJSONRPCRequest(message) && message.method === 'subscriptions/listen') { // Entry-handled listen is its own request-handling subsystem; it // applies the same per-request envelope rung the instance's @@ -461,12 +492,31 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions requested: claimedVersion ?? 'unknown' }); reply = { jsonrpc: '2.0' as const, id: message.id, error: { code: error.code, message: error.message, data: error.data } }; + } else if (listenCompletions.has(message.id)) { + reply = { + jsonrpc: '2.0' as const, + id: message.id, + error: { code: -32_602, message: 'Subscription request ID is already active' } + }; } else { reply = listenRouter.serve(message); } - await wire - .send('error' in reply ? reply : { jsonrpc: '2.0', method: reply.method, params: reply.params }) - .catch(error => reportError(toError(error))); + let acknowledge!: () => void; + const acknowledgement = new Promise(resolve => { + acknowledge = resolve; + }); + // Register completion before send, which may synchronously close the wire. + if (!('error' in reply)) { + listenCompletions.set(message.id, { complete: extra?.onMessageProcessed, acknowledgement, deliveries: new Set() }); + } + // Queue the acknowledgement before a running handler can emit a routed change. + try { + await wire.send('error' in reply ? reply : { jsonrpc: '2.0', method: reply.method, params: reply.params }); + } catch (error) { + reportError(toError(error)); + } finally { + acknowledge(); + } return true; } if (isJSONRPCNotification(message) && message.method === 'notifications/cancelled') { @@ -476,6 +526,8 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions // request). After this point nothing further is delivered for // that subscription id (post-cancel hardening). if (cancelledId !== undefined && listenRouter.cancel(cancelledId)) { + // Retain the listen's admission without blocking unrelated inbound cancellation. + void completeListen(cancelledId); return true; } } @@ -569,175 +621,200 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions } }; - const processMessage = async (message: JSONRPCMessage): Promise => { - if (state.phase === 'closed') { - return; - } - - if (state.phase === 'pinned') { - if ( - state.era === 'modern' && - isJSONRPCRequest(message) && - message.method === 'initialize' && - !carriesValidModernEnvelopeClaim(message.params) - ) { - // The modern era is confirmed for this connection; a late - // legacy handshake is answered with the version error naming - // the supported revisions (the specification recommends - // naming them in any error returned to `initialize`, and - // rules out falling back once the modern era is confirmed). - const requestedVersion = - message.params !== null && - typeof message.params === 'object' && - typeof (message.params as { protocolVersion?: unknown }).protocolVersion === 'string' - ? ((message.params as { protocolVersion: string }).protocolVersion as string) - : undefined; - await answerLegacyRejection(message, 'initialize', requestedVersion); - return; - } - if (state.era === 'modern' && (await tryServeListen(message))) { - return; - } - if (isTornDown()) { - // Closed while entry listen routing was being checked (a - // buffered final message can race stdin EOF): stay closed. - return; - } - state.instance.channel.deliver(message); - return; - } - - // Negotiation window ('opening' | 'probe'). - if (!isJSONRPCRequest(message) && !isJSONRPCNotification(message)) { - // A JSON-RPC response before any era is pinned: nothing has been - // asked of the client yet, so there is nothing it can answer. - reportError(new Error('Discarded a JSON-RPC response received before the connection negotiated an era')); - return; - } - - const opening = classifyOpeningMessage(message); - switch (opening.kind) { - case 'invalid-envelope': { - const detail = `Invalid _meta envelope for protocol revision 2026-07-28: ${opening.issue.key}: ${opening.issue.problem}`; - if (isJSONRPCRequest(message)) { - await writeErrorResponse(message.id, ProtocolErrorCode.InvalidParams, detail, { envelope: opening.issue }); - } else { - reportError(new Error(`Discarded a notification with a malformed envelope: ${detail}`)); - } - return; - } - case 'unsupported-revision': { - if (isJSONRPCRequest(message)) { - const error = new UnsupportedProtocolVersionError({ - supported: [...SUPPORTED_MODERN_PROTOCOL_VERSIONS], - requested: opening.requested - }); - reportError(error); - await writeErrorResponse(message.id, error.code, error.message, error.data); - } else { - reportError(new Error(`Discarded a notification claiming unsupported protocol revision ${opening.requested}`)); + const processMessage = async (message: JSONRPCMessage, extra?: MessageExtraInfo): Promise => { + let delegated = false; + const deliver = (channel: StdioConnectionChannel, classification?: MessageClassification) => { + channel.deliver(message, { ...extra, ...(classification ? { classification } : {}) }); + delegated = true; + }; + try { + await (async () => { + if (state.phase === 'closed') { + return; } - return; - } - case 'modern': { - if (isJSONRPCRequest(message) && message.method === 'server/discover') { - if (state.phase === 'probe') { - // A repeated probe is answered by the same optimistic - // instance and the negotiation window stays open: only - // a non-discover enveloped request commits the - // connection to the modern era, so a later fallback - // `initialize` is still served by a fresh legacy - // instance. - state.instance.channel.deliver(message, { classification: opening.classification }); - return; - } - // Probe: answer from an optimistically built modern - // instance so the advertisement reflects the real server - // definition, but do not pin the connection yet — the - // client may still fall back to `initialize` when it - // shares no modern revision with the advertisement. - const instance = await connectInstance('modern', opening.revision); - if (isTornDown()) { - // The connection was torn down while the factory was - // building the probe instance: dispose of it and stay - // closed instead of resurrecting the negotiation - // window; nothing is delivered or answered. - await disposeLateInstance(instance); + + if (state.phase === 'pinned') { + if ( + state.era === 'modern' && + isJSONRPCRequest(message) && + message.method === 'initialize' && + !carriesValidModernEnvelopeClaim(message.params) + ) { + // The modern era is confirmed for this connection; a late + // legacy handshake is answered with the version error naming + // the supported revisions (the specification recommends + // naming them in any error returned to `initialize`, and + // rules out falling back once the modern era is confirmed). + const requestedVersion = + message.params !== null && + typeof message.params === 'object' && + typeof (message.params as { protocolVersion?: unknown }).protocolVersion === 'string' + ? ((message.params as { protocolVersion: string }).protocolVersion as string) + : undefined; + await answerLegacyRejection(message, 'initialize', requestedVersion); return; } - state = { phase: 'probe', instance }; - instance.channel.deliver(message, { classification: opening.classification }); - return; - } - if (state.phase === 'probe') { - if (isJSONRPCNotification(message)) { - // An enveloped notification during the negotiation - // window (for example a notifications/cancelled for - // the probe itself) is delivered to the probe instance - // without committing the era: only a non-discover - // enveloped request pins the connection, so a later - // fallback `initialize` is still served by a fresh - // legacy instance. - state.instance.channel.deliver(message, { classification: opening.classification }); + if (state.era === 'modern' && (await tryServeListen(message, extra))) { return; } - // The probe was followed by a modern request: the client - // committed to the modern era — pin the probe instance. - state = { phase: 'pinned', era: 'modern', instance: state.instance }; - } else { - const instance = await connectInstance('modern', opening.revision); if (isTornDown()) { - // Closed while the factory was building the modern - // instance: dispose of it and stay closed. - await disposeLateInstance(instance); + // Closed while entry listen routing was being checked (a + // buffered final message can race stdin EOF): stay closed. return; } - state = { phase: 'pinned', era: 'modern', instance }; - } - if (await tryServeListen(message)) { + deliver(state.instance.channel); return; } - if (isTornDown()) { - // Closed while entry listen routing was being checked (a - // buffered final message can race stdin EOF): stay closed. + + // Negotiation window ('opening' | 'probe'). + if (!isJSONRPCRequest(message) && !isJSONRPCNotification(message)) { + // A JSON-RPC response before any era is pinned: nothing has been + // asked of the client yet, so there is nothing it can answer. + reportError(new Error('Discarded a JSON-RPC response received before the connection negotiated an era')); return; } - state.instance.channel.deliver(message, { classification: opening.classification }); - return; - } - case 'legacy': { - if (legacyMode === 'reject') { - if (isJSONRPCRequest(message)) { - await answerLegacyRejection(message, opening.reason, opening.requestedVersion); + + const opening = classifyOpeningMessage(message); + switch (opening.kind) { + case 'invalid-envelope': { + const detail = `Invalid _meta envelope for protocol revision 2026-07-28: ${opening.issue.key}: ${opening.issue.problem}`; + if (isJSONRPCRequest(message)) { + await writeErrorResponse(message.id, ProtocolErrorCode.InvalidParams, detail, { envelope: opening.issue }); + } else { + reportError(new Error(`Discarded a notification with a malformed envelope: ${detail}`)); + } + return; } - // Claim-less notifications are accepted and dropped (the - // stdio analog of the HTTP entry's 202-and-drop); the - // connection stays open for a modern opening. - return; - } - if (state.phase === 'probe') { - // Probe-then-fallback: the client probed, found no - // mutually supported modern revision, and fell back to - // the 2025 handshake on the same connection. The probe - // instance is discarded; a fresh legacy instance serves - // the handshake. - await discardProbeInstance(state.instance); - if (isTornDown()) { - // Closed while the probe was being discarded: stay closed. + case 'unsupported-revision': { + if (isJSONRPCRequest(message)) { + const error = new UnsupportedProtocolVersionError({ + supported: [...SUPPORTED_MODERN_PROTOCOL_VERSIONS], + requested: opening.requested + }); + reportError(error); + await writeErrorResponse(message.id, error.code, error.message, error.data); + } else { + reportError(new Error(`Discarded a notification claiming unsupported protocol revision ${opening.requested}`)); + } + return; + } + case 'modern': { + if (isJSONRPCRequest(message) && message.method === 'server/discover') { + if (state.phase === 'probe') { + // A repeated probe is answered by the same optimistic + // instance and the negotiation window stays open: only + // a non-discover enveloped request commits the + // connection to the modern era, so a later fallback + // `initialize` is still served by a fresh legacy + // instance. + deliver(state.instance.channel, opening.classification); + return; + } + // Probe: answer from an optimistically built modern + // instance so the advertisement reflects the real server + // definition, but do not pin the connection yet — the + // client may still fall back to `initialize` when it + // shares no modern revision with the advertisement. + const instance = await connectInstance('modern', opening.revision); + if (isTornDown()) { + // The connection was torn down while the factory was + // building the probe instance: dispose of it and stay + // closed instead of resurrecting the negotiation + // window; nothing is delivered or answered. + await disposeLateInstance(instance); + return; + } + state = { phase: 'probe', instance }; + deliver(instance.channel, opening.classification); + return; + } + if (state.phase === 'probe') { + if (isJSONRPCNotification(message)) { + // An enveloped notification during the negotiation + // window (for example a notifications/cancelled for + // the probe itself) is delivered to the probe instance + // without committing the era: only a non-discover + // enveloped request pins the connection, so a later + // fallback `initialize` is still served by a fresh + // legacy instance. + deliver(state.instance.channel, opening.classification); + return; + } + // The probe was followed by a modern request: the client + // committed to the modern era — pin the probe instance. + state = { phase: 'pinned', era: 'modern', instance: state.instance }; + } else { + const instance = await connectInstance('modern', opening.revision); + if (isTornDown()) { + // Closed while the factory was building the modern + // instance: dispose of it and stay closed. + await disposeLateInstance(instance); + return; + } + state = { phase: 'pinned', era: 'modern', instance }; + } + if (await tryServeListen(message, extra)) { + return; + } + if (isTornDown()) { + // Closed while entry listen routing was being checked (a + // buffered final message can race stdin EOF): stay closed. + return; + } + deliver(state.instance.channel, opening.classification); + return; + } + case 'legacy': { + if (legacyMode === 'reject') { + if (isJSONRPCRequest(message)) { + await answerLegacyRejection(message, opening.reason, opening.requestedVersion); + } + // Claim-less notifications are accepted and dropped (the + // stdio analog of the HTTP entry's 202-and-drop); the + // connection stays open for a modern opening. + return; + } + if (state.phase === 'probe') { + // Probe-then-fallback: the client probed, found no + // mutually supported modern revision, and fell back to + // the 2025 handshake on the same connection. The probe + // instance is discarded; a fresh legacy instance serves + // the handshake. + await discardProbeInstance(state.instance); + if (isTornDown()) { + // Closed while the probe was being discarded: stay closed. + return; + } + state = { phase: 'opening' }; + } + const instance = await connectInstance('legacy'); + if (isTornDown()) { + // Closed while the factory was building the legacy + // instance: dispose of it and stay closed. + await disposeLateInstance(instance); + return; + } + state = { phase: 'pinned', era: 'legacy', instance }; + deliver(state.instance.channel); return; } - state = { phase: 'opening' }; - } - const instance = await connectInstance('legacy'); - if (isTornDown()) { - // Closed while the factory was building the legacy - // instance: dispose of it and stay closed. - await disposeLateInstance(instance); - return; } - state = { phase: 'pinned', era: 'legacy', instance }; - state.instance.channel.deliver(message); - return; + })(); + } catch (error) { + if (isJSONRPCRequest(message)) { + await writeErrorResponse(message.id, ProtocolErrorCode.InternalError, 'Internal server error'); + } + reportError(toError(error)); + } finally { + if ( + !delegated && + !( + isJSONRPCRequest(message) && + listenCompletions.get(message.id)?.complete === extra?.onMessageProcessed && + extra?.onMessageProcessed !== undefined + ) + ) { + extra?.onMessageProcessed?.(); } } }; @@ -745,7 +822,7 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions // Inbound messages are processed strictly in arrival order: the queue // absorbs anything that arrives while the opening exchange is still being // decided (factory construction and instance connection are async). - const queue: JSONRPCMessage[] = []; + const queue: { message: JSONRPCMessage; extra?: MessageExtraInfo }[] = []; let pumping = false; const pump = async (): Promise => { if (pumping) { @@ -754,9 +831,9 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions pumping = true; try { while (queue.length > 0) { - const message = queue.shift()!; + const { message, extra } = queue.shift()!; try { - await processMessage(message); + await processMessage(message, extra); } catch (error) { // Every arm of processMessage that answers a request does // so through writeErrorResponse (which never throws — wire @@ -771,6 +848,7 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions await writeErrorResponse(message.id, ProtocolErrorCode.InternalError, 'Internal server error'); } reportError(toError(error)); + extra?.onMessageProcessed?.(); } } } finally { @@ -792,6 +870,7 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions // from a transport drop. for (const result of listenRouter.teardownAll()) { await wire.send(result).catch(error => reportError(toError(error))); + await completeListen(result.id); } if (current.phase === 'probe' || current.phase === 'pinned') { await current.instance.product.close().catch(error => reportError(toError(error))); @@ -799,8 +878,18 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions await wire.close().catch(error => reportError(toError(error))); }; - wire.onmessage = (message: JSONRPCMessage) => { - queue.push(message); + wire.onmessage = (message: JSONRPCMessage, extra?: MessageExtraInfo) => { + let processed = false; + const onMessageProcessed = () => { + if (processed) return; + processed = true; + try { + extra?.onMessageProcessed?.(); + } catch (error) { + reportError(toError(error)); + } + }; + queue.push({ message, extra: extra?.onMessageProcessed ? { ...extra, onMessageProcessed } : extra }); void pump(); }; wire.onerror = error => { @@ -816,6 +905,8 @@ export function serveStdio(factory: McpServerFactory, options: ServeStdioOptions closing = true; const current = state; state = { phase: 'closed' }; + listenRouter.teardownAll(); + for (const id of listenCompletions.keys()) void completeListen(id); if (current.phase === 'probe' || current.phase === 'pinned') { void current.instance.product.close().catch(error => reportError(toError(error))); } diff --git a/packages/server/src/server/stdio.ts b/packages/server/src/server/stdio.ts index 4895359ab3..ddb6ccfc4b 100644 --- a/packages/server/src/server/stdio.ts +++ b/packages/server/src/server/stdio.ts @@ -189,32 +189,38 @@ export class StdioServerTransport implements Transport { } return new Promise((resolve, reject) => { const json = serializeMessage(message); - + let returned = false; + let written = false; + let failure: Error | undefined; let settled = false; - const onError = (error: Error) => { - if (settled) return; - settled = true; + const cleanup = () => { this._stdout.off('error', onError); - this._stdout.off('drain', onDrain); - reject(error); }; - const onDrain = () => { - if (settled) return; + const finish = () => { + if (settled || !returned || !written) return; settled = true; - this._stdout.off('error', onError); - this._stdout.off('drain', onDrain); - resolve(); + cleanup(); + if (failure) reject(failure); + else resolve(); }; - - this._stdout.once('error', onError); - - if (this._stdout.write(json)) { - if (settled) return; - settled = true; - this._stdout.off('error', onError); - resolve(); - } else if (!settled) { - this._stdout.once('drain', onDrain); + const onError = (error: Error) => { + failure = error; + finish(); + }; + this._stdout.prependOnceListener('error', onError); + try { + // The callback owns this chunk. A drain event is global flow control + // and may never occur when the stream ends or is destroyed. + this._stdout.write(json, error => { + written = true; + if (error) failure = error; + finish(); + }); + returned = true; + finish(); + } catch (error) { + cleanup(); + reject(error); } }); } diff --git a/packages/server/src/stdio.ts b/packages/server/src/stdio.ts index 1d38a61a0d..27b3e763c7 100644 --- a/packages/server/src/stdio.ts +++ b/packages/server/src/stdio.ts @@ -9,3 +9,6 @@ export type { ServeStdioOptions, StdioServerHandle } from './server/serveStdio'; export { serveStdio } from './server/serveStdio'; export { StdioServerTransport } from './server/stdio'; + +/** This SDK propagates MessageExtraInfo.onMessageProcessed through stdio dispatch. */ +export const supportsInboundMessageCompletion = true; diff --git a/packages/server/test/server/inboundProcessing.test.ts b/packages/server/test/server/inboundProcessing.test.ts new file mode 100644 index 0000000000..9804878d27 --- /dev/null +++ b/packages/server/test/server/inboundProcessing.test.ts @@ -0,0 +1,480 @@ +import { Readable, Writable } from 'node:stream'; +import { setImmediate as nextTurn } from 'node:timers/promises'; + +import type { JSONRPCMessage, MessageExtraInfo, Transport } from '@modelcontextprotocol/core-internal'; +import { CLIENT_CAPABILITIES_META_KEY, CLIENT_INFO_META_KEY, PROTOCOL_VERSION_META_KEY } from '@modelcontextprotocol/core-internal'; +import { vi } from 'vitest'; +import * as z from 'zod/v4'; + +import { McpServer } from '../../src/server/mcp'; +import { serveStdio } from '../../src/server/serveStdio'; +import { StdioServerTransport } from '../../src/server/stdio'; + +function deferred() { + let resolve!: () => void; + const promise = new Promise(r => { + resolve = r; + }); + return { promise, resolve }; +} +function fixture(factory: () => McpServer | Promise) { + const sent: JSONRPCMessage[] = []; + const wire: Transport = { + start: async () => {}, + close: async () => {}, + send: vi.fn(async message => { + sent.push(message); + }) + }; + const handle = serveStdio(factory, { transport: wire, onerror: () => {} }); + return { wire, handle, sent, receive: (message: JSONRPCMessage, extra?: MessageExtraInfo) => wire.onmessage!(message, extra) }; +} +const init = { + jsonrpc: '2.0' as const, + id: 1, + method: 'initialize', + params: { protocolVersion: '2025-11-25', capabilities: {}, clientInfo: { name: 'test', version: '1' } } +}; +const envelope = { + [PROTOCOL_VERSION_META_KEY]: '2026-07-28', + [CLIENT_INFO_META_KEY]: { name: 'test', version: '1' }, + [CLIENT_CAPABILITIES_META_KEY]: {} +}; + +test.each(['legacy', 'modern'] as const)('closed %s channels drop sends suspended at outbound interception', async era => { + const server = new McpServer({ name: 'test', version: '1' }); + const f = fixture(() => server); + const opened = vi.fn(); + f.receive(era === 'legacy' ? init : { jsonrpc: '2.0', id: 1, method: 'ping', params: { _meta: envelope } }, { + onMessageProcessed: opened + }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + const channel = server.server.transport!; + const writesBefore = f.sent.length; + try { + const sending = channel.send({ jsonrpc: '2.0', id: 7, result: {} }); + // Close synchronously while send is suspended at the interceptor's await boundary. + const closing = channel.close(); + await Promise.all([sending, closing]); + expect(f.sent).toHaveLength(writesBefore); + } finally { + await f.handle.close(); + } +}); + +test('queued opening completion waits for the real asynchronous factory', async () => { + const gate = deferred(); + const complete = vi.fn(); + const f = fixture(async () => { + await gate.promise; + return new McpServer({ name: 'test', version: '1' }); + }); + f.receive(init, { onMessageProcessed: complete }); + await Promise.resolve(); + expect(complete).not.toHaveBeenCalled(); + gate.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + await f.handle.close(); +}); + +test('entry rejection waits for its error response send', async () => { + const complete = vi.fn(); + const write = deferred(); + const f = fixture(() => new McpServer({ name: 'test', version: '1' })); + f.wire.send = vi.fn(() => write.promise); + f.receive( + { jsonrpc: '2.0', id: 1, method: 'server/discover', params: { _meta: { [PROTOCOL_VERSION_META_KEY]: '2099-01-01' } } }, + { onMessageProcessed: complete } + ); + await vi.waitFor(() => expect(f.wire.send).toHaveBeenCalledOnce()); + expect(complete).not.toHaveBeenCalled(); + write.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + await f.handle.close(); +}); + +test('factory failure completion follows the error response', async () => { + const complete = vi.fn(); + const write = deferred(); + const f = fixture(() => { + throw new Error('factory'); + }); + f.wire.send = vi.fn(() => write.promise); + f.receive(init, { onMessageProcessed: complete }); + await vi.waitFor(() => expect(f.wire.send).toHaveBeenCalledOnce()); + expect(complete).not.toHaveBeenCalled(); + write.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + await f.handle.close(); +}); + +test('cancelled tool remains charged through actual handler completion', async () => { + const entered = deferred(); + const held = deferred(); + const complete = vi.fn(); + const f = fixture(() => { + const s = new McpServer({ name: 'test', version: '1' }); + s.registerTool('work', { inputSchema: z.object({}) }, async () => { + entered.resolve(); + await held.promise; + return { content: [] }; + }); + return s; + }); + const opened = vi.fn(); + f.receive(init, { onMessageProcessed: opened }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + f.receive({ jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'work', arguments: {} } }, { onMessageProcessed: complete }); + await entered.promise; + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 2 } }); + await f.handle.close(); + expect(complete).not.toHaveBeenCalled(); + held.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + expect(f.sent.filter(m => 'id' in m && m.id === 2)).toEqual([]); +}); + +test('listen acknowledgement does not release the subscription request', async () => { + const complete = vi.fn(); + const f = fixture(() => new McpServer({ name: 'test', version: '1' }, { capabilities: { tools: { listChanged: true } } })); + const probe = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 1, method: 'server/discover', params: { _meta: envelope } }, { onMessageProcessed: probe }); + await vi.waitFor(() => expect(probe).toHaveBeenCalledOnce()); + f.receive( + { + jsonrpc: '2.0', + id: 2, + method: 'subscriptions/listen', + params: { _meta: envelope, notifications: { tools: { listChanged: true } } } + }, + { onMessageProcessed: complete } + ); + await vi.waitFor(() => expect(f.sent.some(m => 'method' in m && m.method === 'notifications/subscriptions/acknowledged')).toBe(true)); + expect(complete).not.toHaveBeenCalled(); + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 2 } }); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + await f.handle.close(); +}); + +test.skipIf(!process.env.MCP_LIFECYCLE_LONG_RUN)( + 'both protocol generations reuse eight actual completed request IDs for more than ten minutes', + async () => { + const duration = 601_000; + await Promise.all( + ['legacy', 'modern'].map(async era => { + let active = 0; + let peak = 0; + let completed = 0; + let replies = 0; + const ids = new Set(); + const input = new Readable({ read() {} }); + const output = new Writable({ + write(chunk, _encoding, callback) { + const message: JSONRPCMessage = JSON.parse(chunk.toString()); + expect('error' in message).toBe(false); + if ('result' in message && message.id !== 0) expect(JSON.stringify(message.result)).toContain('completed'); + replies++; + setImmediate(callback); + } + }); + const official = new StdioServerTransport(input, output, { maxBufferSize: 1_048_576 }); + const completions = new Map void>(); + const wire: Transport = { + start: () => official.start(), + close: () => official.close(), + send: message => official.send(message) + }; + official.onmessage = message => { + if (!('id' in message) || typeof message.id !== 'number') throw new Error('unexpected inbound frame'); + const complete = completions.get(message.id); + if (!complete) throw new Error('unowned inbound frame'); + wire.onmessage!(message, { onMessageProcessed: complete }); + }; + let handled = 0; + const handle = serveStdio( + () => { + const server = new McpServer({ name: 'long-test', version: '1' }); + server.registerTool('work', { inputSchema: z.object({}) }, async () => { + await new Promise(resolve => setImmediate(resolve)); + handled++; + return { content: [{ type: 'text', text: 'completed' }] }; + }); + return server; + }, + { transport: wire } + ); + const request = (id: number, method: string, params: Record) => + new Promise(resolve => { + expect(ids.has(id)).toBe(false); + ids.add(id); + active++; + peak = Math.max(peak, active); + completions.set(id, () => { + expect(ids.delete(id)).toBe(true); + completions.delete(id); + active--; + completed++; + resolve(); + }); + input.push(Buffer.from(JSON.stringify({ jsonrpc: '2.0', id, method, params }) + '\n')); + }); + await (era === 'legacy' ? request(0, 'initialize', init.params) : request(0, 'server/discover', { _meta: envelope })); + const started = performance.now(); + while (performance.now() - started < duration) { + await Promise.all( + Array.from({ length: 8 }, (_, i) => + request(i + 1, 'tools/call', { name: 'work', arguments: {}, ...(era === 'modern' ? { _meta: envelope } : {}) }) + ) + ); + expect(active).toBe(0); + expect(ids.size).toBe(0); + await new Promise(resolve => setTimeout(resolve, 1000)); + } + await request(1, 'tools/call', { name: 'work', arguments: {}, ...(era === 'modern' ? { _meta: envelope } : {}) }); + expect(performance.now() - started).toBeGreaterThan(600_000); + expect(peak).toBe(8); + expect(completed).toBeGreaterThan(4000); + expect(replies).toBe(completed); + expect(handled).toBe(completed - 1); + expect(completions.size).toBe(0); + await handle.close(); + input.destroy(); + output.end(); + }) + ); + }, + 630_000 +); + +test('rejected duplicate listen messages complete without replacing the original subscription', async () => { + const f = fixture(() => new McpServer({ name: 'test', version: '1' }, { capabilities: { tools: { listChanged: true } } })); + const opened = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 1, method: 'server/discover', params: { _meta: envelope } }, { onMessageProcessed: opened }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + const original = vi.fn(); + const params = { _meta: envelope, notifications: { tools: { listChanged: true } } }; + f.receive({ jsonrpc: '2.0', id: 2, method: 'subscriptions/listen', params }, { onMessageProcessed: original }); + await vi.waitFor(() => expect(f.sent.some(m => 'method' in m && m.method === 'notifications/subscriptions/acknowledged')).toBe(true)); + const rejected = vi.fn(); + f.receive( + { jsonrpc: '2.0', id: 2, method: 'subscriptions/listen', params: { _meta: { [PROTOCOL_VERSION_META_KEY]: '2099-01-01' } } }, + { onMessageProcessed: rejected } + ); + await vi.waitFor(() => expect(rejected).toHaveBeenCalledOnce()); + expect(original).not.toHaveBeenCalled(); + const duplicate = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 2, method: 'subscriptions/listen', params }, { onMessageProcessed: duplicate }); + await vi.waitFor(() => expect(duplicate).toHaveBeenCalledOnce()); + expect(original).not.toHaveBeenCalled(); + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 2 } }); + await vi.waitFor(() => expect(original).toHaveBeenCalledOnce()); + await f.handle.close(); +}); + +test('listen teardown joins held acknowledgement and terminal transmission', async () => { + const f = fixture(() => new McpServer({ name: 'test', version: '1' }, { capabilities: { tools: { listChanged: true } } })); + const opened = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 1, method: 'server/discover', params: { _meta: envelope } }, { onMessageProcessed: opened }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + const ack = deferred(); + const terminal = deferred(); + const complete = vi.fn(); + const written: JSONRPCMessage[] = []; + f.wire.send = vi.fn(message => { + written.push(message); + return 'method' in message && message.method === 'notifications/subscriptions/acknowledged' ? ack.promise : terminal.promise; + }); + f.receive( + { + jsonrpc: '2.0', + id: 2, + method: 'subscriptions/listen', + params: { _meta: envelope, notifications: { tools: { listChanged: true } } } + }, + { onMessageProcessed: complete } + ); + await vi.waitFor(() => expect(written.length).toBe(1)); + const close = f.handle.close(); + await vi.waitFor(() => expect(written.length).toBe(2)); + expect(complete).not.toHaveBeenCalled(); + ack.resolve(); + await Promise.resolve(); + expect(complete).not.toHaveBeenCalled(); + terminal.resolve(); + await close; + expect(complete).toHaveBeenCalledOnce(); +}); + +test('cancelled listeners retain already-started routed notification writes', async () => { + let server!: McpServer; + const f = fixture(() => { + server = new McpServer({ name: 'test', version: '1' }, { capabilities: { tools: { listChanged: true } } }); + return server; + }); + const opened = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 1, method: 'server/discover', params: { _meta: envelope } }, { onMessageProcessed: opened }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + const a = vi.fn(), + b = vi.fn(); + const params = { _meta: envelope, notifications: { toolsListChanged: true } }; + f.receive({ jsonrpc: '2.0', id: 2, method: 'subscriptions/listen', params }, { onMessageProcessed: a }); + f.receive({ jsonrpc: '2.0', id: 3, method: 'subscriptions/listen', params }, { onMessageProcessed: b }); + await vi.waitFor(() => + expect(f.sent.filter(m => 'method' in m && m.method === 'notifications/subscriptions/acknowledged')).toHaveLength(2) + ); + const heldA = deferred(), + heldB = deferred(); + const written: number[] = []; + f.wire.send = vi.fn(async message => { + if ('method' in message && message.method === 'notifications/tools/list_changed') { + const id = message.params!._meta!['io.modelcontextprotocol/subscriptionId'] as number; + written.push(id); + await (id === 2 ? heldA : heldB).promise; + } + }); + server.sendToolListChanged(); + await vi.waitFor(() => expect(written).toEqual([2, 3])); + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 2 } }); + await Promise.resolve(); + expect(a).not.toHaveBeenCalled(); + heldA.resolve(); + await vi.waitFor(() => expect(a).toHaveBeenCalledOnce()); + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 3 } }); + await Promise.resolve(); + expect(b).not.toHaveBeenCalled(); + heldB.resolve(); + await vi.waitFor(() => expect(b).toHaveBeenCalledOnce()); + server.sendToolListChanged(); + await new Promise(resolve => setTimeout(resolve, 10)); + expect(written).toEqual([2, 3]); + await f.handle.close(); +}); + +test('listen acknowledgement precedes changes from an already-running handler', async () => { + const entered = deferred(); + const held = deferred(); + const complete = vi.fn(); + const f = fixture(() => { + const server = new McpServer({ name: 'test', version: '1' }); + server.registerTool('notify', { inputSchema: z.object({}) }, async () => { + entered.resolve(); + await held.promise; + await server.server.sendToolListChanged(); + return { content: [] }; + }); + return server; + }); + try { + const opened = vi.fn(); + f.receive({ jsonrpc: '2.0', id: 1, method: 'tools/list', params: { _meta: envelope } }, { onMessageProcessed: opened }); + await vi.waitFor(() => expect(opened).toHaveBeenCalledOnce()); + f.receive( + { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'notify', arguments: {}, _meta: envelope } }, + { onMessageProcessed: complete } + ); + await entered.promise; + await nextTurn(); + held.resolve(); + f.receive({ + jsonrpc: '2.0', + id: 3, + method: 'subscriptions/listen', + params: { _meta: envelope, notifications: { toolsListChanged: true } } + }); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + expect(f.sent.filter(m => 'method' in m).map(m => m.method)).toEqual([ + 'notifications/subscriptions/acknowledged', + 'notifications/tools/list_changed' + ]); + } finally { + held.resolve(); + await f.handle.close(); + } +}); + +test('listen completion waits for an acknowledgement whose send closes the wire synchronously', async () => { + const write = deferred(); + const complete = vi.fn(); + const f = fixture(() => new McpServer({ name: 'test', version: '1' })); + try { + f.wire.send = vi.fn(() => { + f.wire.onclose!(); + return write.promise; + }); + f.receive( + { + jsonrpc: '2.0', + id: 1, + method: 'subscriptions/listen', + params: { _meta: envelope, notifications: {} } + }, + { onMessageProcessed: complete } + ); + await nextTurn(); + expect(f.wire.send).toHaveBeenCalledOnce(); + expect(complete).not.toHaveBeenCalled(); + write.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + } finally { + write.resolve(); + await f.handle.close(); + } +}); + +test.each([false, true])('listen teardown does not block unrelated cancellation (completion observer: %s)', async observe => { + const entered = deferred(); + const held = deferred(); + const write = deferred(); + const writing = deferred(); + const complete = vi.fn(); + let signal!: AbortSignal; + let server!: McpServer; + const f = fixture(() => { + server = new McpServer({ name: 'test', version: '1' }); + server.registerTool('work', { inputSchema: z.object({}) }, async (_args, ctx) => { + signal = ctx.mcpReq.signal; + entered.resolve(); + await held.promise; + return { content: [] }; + }); + return server; + }); + try { + f.receive( + { + jsonrpc: '2.0', + id: 1, + method: 'subscriptions/listen', + params: { _meta: envelope, notifications: { toolsListChanged: true } } + }, + observe ? { onMessageProcessed: complete } : undefined + ); + await vi.waitFor(() => + expect(f.sent.some(m => 'method' in m && m.method === 'notifications/subscriptions/acknowledged')).toBe(true) + ); + f.receive({ jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'work', arguments: {}, _meta: envelope } }); + await entered.promise; + await nextTurn(); + const send = f.wire.send; + f.wire.send = vi.fn(message => { + if ('method' in message && message.method === 'notifications/tools/list_changed') { + writing.resolve(); + return write.promise; + } + return send(message); + }); + const notification = server.server.sendToolListChanged(); + await writing.promise; + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 1 } }); + f.receive({ jsonrpc: '2.0', method: 'notifications/cancelled', params: { requestId: 2 } }); + await vi.waitFor(() => expect(signal.aborted).toBe(true)); + expect(complete).not.toHaveBeenCalled(); + write.resolve(); + await notification; + if (observe) await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); + } finally { + write.resolve(); + held.resolve(); + await f.handle.close(); + } +}); diff --git a/packages/server/test/server/stdio.test.ts b/packages/server/test/server/stdio.test.ts index 4390360ceb..d5f06db43f 100644 --- a/packages/server/test/server/stdio.test.ts +++ b/packages/server/test/server/stdio.test.ts @@ -357,8 +357,8 @@ test('should swallow stdout errors that surface after close (late EPIPE from a p await server.start(); - // Fast-path send: write() returns true, so the per-send 'error' listener is - // detached immediately — but the OS-level write may still be pending. + // The owned chunk callback has completed, so its per-send error listener + // is detached. The closed transport still swallows subsequent stream errors. await server.send({ jsonrpc: '2.0', id: 1, method: 'ping' }); // The client hangs up: stdin EOF closes the transport. @@ -466,3 +466,89 @@ test('should fire onerror and close when ReadBuffer overflows', async () => { expect(receivedError?.message).toMatch(/ReadBuffer exceeded maximum size/); expect(closeCount).toBe(1); }); + +test('send remains pending after write accepts a chunk until its actual callback', async () => { + let complete!: () => void; + const held = new Writable({ + write(_chunk, _encoding, callback) { + complete = callback; + } + }); + const server = new StdioServerTransport(input, held); + await server.start(); + let done = false; + const sending = server.send({ jsonrpc: '2.0', id: 1, result: {} }); + void sending.then(() => { + done = true; + }); + await Promise.resolve(); + expect(done).toBe(false); + await server.close(); + await Promise.resolve(); + expect(done).toBe(false); + complete(); + await sending; + expect(done).toBe(true); +}); + +test('send error also waits for the owned write callback', async () => { + let complete!: (error?: Error) => void; + const held = new Writable({ + write(_chunk, _encoding, callback) { + complete = callback; + } + }); + const server = new StdioServerTransport(input, held); + server.onerror = () => {}; + await server.start(); + let done = false; + const sending = server.send({ jsonrpc: '2.0', id: 1, result: {} }); + const rejected = expect(sending).rejects.toThrow('held failure'); + void sending.catch(() => { + done = true; + }); + held.emit('error', new Error('held failure')); + await Promise.resolve(); + expect(done).toBe(false); + complete(new Error('held failure')); + await rejected; + expect(done).toBe(true); +}); + +test('send rejects a destroyed writable after its callback without requiring an error event', async () => { + const held = new Writable({ + write(_chunk, _encoding, callback) { + callback(); + } + }); + const server = new StdioServerTransport(input, held); + await server.start(); + held.destroy(); + await new Promise(resolve => held.once('close', resolve)); + await expect(server.send({ jsonrpc: '2.0', id: 1, result: {} })).rejects.toThrow(); + await server.close(); +}); + +test('send joins a final backpressured write without requiring drain after end', async () => { + let complete!: () => void; + const held = new Writable({ + highWaterMark: 1, + write(_chunk, _encoding, callback) { + complete = callback; + } + }); + const server = new StdioServerTransport(input, held); + await server.start(); + let done = false; + const sending = server.send({ jsonrpc: '2.0', id: 1, result: {} }); + void sending.then(() => { + done = true; + }); + held.end(); + await Promise.resolve(); + expect(done).toBe(false); + complete(); + await sending; + expect(done).toBe(true); + await server.close(); +});