diff --git a/.changeset/inbound-processing-completion.md b/.changeset/inbound-processing-completion.md new file mode 100644 index 0000000000..adebceee2d --- /dev/null +++ b/.changeset/inbound-processing-completion.md @@ -0,0 +1,6 @@ +--- +'@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. diff --git a/docs/inbound-processing-completion.md b/docs/inbound-processing-completion.md new file mode 100644 index 0000000000..63afe6d4a0 --- /dev/null +++ b/docs/inbound-processing-completion.md @@ -0,0 +1,7 @@ +# 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 not completion. The stdio entry waits for cancellation or teardown, including a pending acknowledgement and any graceful terminal write it owns. `supportsInboundMessageCompletion` is exported from `@modelcontextprotocol/server/stdio` for adapters that require this contract before accepting input. + +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. diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 8118b60cea..e102e7708e 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -2158,7 +2158,7 @@ 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 { // 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 @@ -2218,7 +2218,7 @@ export class Client extends Protocol { return; } } - super._onnotification(raw, extra); + return super._onnotification(raw, extra); } /** @@ -2231,7 +2231,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) { @@ -2260,7 +2260,7 @@ export class Client extends Protocol { } return; } - super._onresponse(response); + return super._onresponse(response); } /** diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 9d60be8efe..2ad4e6e96f 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -570,7 +570,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(); @@ -810,15 +810,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(); + } + }); } }; @@ -874,7 +907,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 @@ -930,12 +963,12 @@ export abstract class Protocol { } // Starting with Promise.resolve() puts any synchronous errors into the monad as well. - Promise.resolve() + return Promise.resolve() .then(() => (handler === undefined ? fallback!(notification) : handler(notification, codec))) .catch(error => this._onerror(new Error(`Uncaught error in notification handler: ${error}`))); } - 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. @@ -967,7 +1000,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 @@ -988,7 +1023,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 @@ -999,7 +1034,6 @@ export abstract class Protocol { supported: this._supportedProtocolVersions, requested }); - return; } } @@ -1015,15 +1049,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` @@ -1037,8 +1069,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 @@ -1108,11 +1139,10 @@ export abstract class Protocol { }, http: extra?.authInfo ? { authInfo: extra.authInfo } : undefined }; - const ctx = this.buildContext(baseCtx, extra); // Starting with Promise.resolve() puts any synchronous errors into the monad as well. - Promise.resolve() - .then(() => handler(request, ctx)) + return Promise.resolve() + .then(() => handler(request, this.buildContext(baseCtx, extra))) .then( async result => { if (abortController.signal.aborted) { @@ -1133,8 +1163,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'); } const response: JSONRPCResponse = { @@ -1210,7 +1239,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); @@ -1224,10 +1253,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); } } @@ -1563,11 +1592,14 @@ export abstract class Protocol { legOptions ) }; - return resolve(this._resolveNonCompleteResult(decoded, flow) as Promise>); + return (this._resolveNonCompleteResult(decoded, flow) as Promise>).then( + resolve, + reject + ); } const result = decoded.result; - validateStandardSchema(resultSchema, result).then(parseResult => { + return validateStandardSchema(resultSchema, result).then(parseResult => { if (parseResult.success) { resolve(parseResult.data); } else { 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..2b94d24c0e --- /dev/null +++ b/packages/core-internal/test/shared/inboundProcessing.test.ts @@ -0,0 +1,138 @@ +import { vi } from 'vitest'; +import * as z from 'zod/v4'; +import type { BaseContext } from '../../src/shared/protocol'; +import { Protocol } 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() { + const protocol = new TestProtocol(); + const transport: Transport = { + start: async () => {}, + close: async () => {}, + send: vi.fn(async () => {}) + }; + await protocol.connect(transport); + const receive = (message: JSONRPCMessage, extra?: MessageExtraInfo) => transport.onmessage!(message, extra); + return { protocol, transport, receive }; +} +const request = { jsonrpc: '2.0' as const, id: 1, method: 'test/work', params: {} }; + +test('completion waits for handler, response send and controller cleanup', async () => { + const { protocol, transport, receive } = await fixture(); + const handler = deferred(); + const write = deferred(); + const complete = vi.fn(); + protocol.setRequestHandler('test/work', { params: z.object({}), result: z.object({}) }, async () => { + await handler.promise; + 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()); + 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('notification completion waits for its asynchronous handler', async () => { + const { protocol, receive } = await fixture(); + 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(); + handler.resolve(); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); +}); + +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('inbound response completion joins asynchronous SDK result validation', async () => { + const f = await fixture(); + 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(); + held.resolve(); + expect(await result).toEqual({ value: 'accepted' }); + await vi.waitFor(() => expect(complete).toHaveBeenCalledOnce()); +}); diff --git a/packages/server/src/server/serveStdio.ts b/packages/server/src/server/serveStdio.ts index 1cabf8ed9b..4f6d5fc792 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,7 @@ class StdioConnectionChannel implements Transport { // sends are dropped. return; } - if (this._outboundIntercept?.(message) === 'handled') { + if ((await this._outboundIntercept?.(message)) === 'handled') { return; } return this._wire.send(message, options); @@ -199,6 +199,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 +415,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 +424,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 +449,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 +489,22 @@ 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 }) + const acknowledgement = Promise.resolve() + .then(() => wire.send('error' in reply ? reply : { jsonrpc: '2.0', method: reply.method, params: reply.params })) .catch(error => reportError(toError(error))); + if (!('error' in reply)) { + listenCompletions.set(message.id, { complete: extra?.onMessageProcessed, acknowledgement, deliveries: new Set() }); + } + await acknowledgement; return true; } if (isJSONRPCNotification(message) && message.method === 'notifications/cancelled') { @@ -476,6 +514,7 @@ 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)) { + await completeListen(cancelledId); return true; } } @@ -569,175 +608,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 }); + + 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; } - // 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.era === 'modern' && (await tryServeListen(message, extra))) { 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 }); - 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 +809,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 +818,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 +835,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 +857,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 +865,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 +892,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..14539690b8 --- /dev/null +++ b/packages/server/test/server/inboundProcessing.test.ts @@ -0,0 +1,329 @@ +import { Readable, Writable } from 'node:stream'; + +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('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(); +}); 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(); +});