From c6abb6d2b5f445729b243ee522347115336f91c6 Mon Sep 17 00:00:00 2001 From: K4bain <296577378+K4bain@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:26:05 +0500 Subject: [PATCH 1/2] fix(client): opt-in graceful close drains in-flight requests before transport teardown close({ drainPendingRequests: true }) waits for in-flight requests to settle before the transport closes. Without it, transport teardown aborts in-flight HTTP requests the server had already answered, which OpenTelemetry's undici instrumentation reports as UND_ERR_ABORTED on 200 OK responses (modelcontextprotocol/typescript-sdk#1231). - Protocol tracks pending request ids alongside the response-handler lifecycle and drains them before transport close when opted in - Client.close({ drainPendingRequests }) and a ClientOptions.gracefulClose constructor default expose the behavior; an explicit argument wins - Requests outstanding after the drain timeout (default 2s) settle via the normal close path; default close() behavior is unchanged --- .changeset/graceful-close-drain.md | 6 + packages/client/src/client/client.ts | 45 +++- packages/client/src/index.ts | 9 +- .../client/test/client/gracefulClose.test.ts | 204 ++++++++++++++++++ .../core-internal/src/exports/public/index.ts | 1 + packages/core-internal/src/shared/protocol.ts | 106 ++++++++- 6 files changed, 367 insertions(+), 4 deletions(-) create mode 100644 .changeset/graceful-close-drain.md create mode 100644 packages/client/test/client/gracefulClose.test.ts diff --git a/.changeset/graceful-close-drain.md b/.changeset/graceful-close-drain.md new file mode 100644 index 0000000000..65aee21f44 --- /dev/null +++ b/.changeset/graceful-close-drain.md @@ -0,0 +1,6 @@ +--- +'@modelcontextprotocol/client': patch +'@modelcontextprotocol/core-internal': patch +--- + +Add opt-in graceful close: `client.close({ drainPendingRequests: true })` (or a per-call `{ timeoutMs }`) waits for in-flight requests to settle before the transport closes, so a completed HTTP response is no longer aborted mid-read by teardown — which OpenTelemetry's undici instrumentation previously reported as `UND_ERR_ABORTED` on 200 OK responses. A `ClientOptions.gracefulClose` default covers SIGINT-style shutdowns where the caller does not know what is in flight; requests still outstanding after the drain timeout are abandoned to the normal close path, and the default `close()` behavior is unchanged. diff --git a/packages/client/src/client/client.ts b/packages/client/src/client/client.ts index 0b386a63e8..efdbf2eae7 100644 --- a/packages/client/src/client/client.ts +++ b/packages/client/src/client/client.ts @@ -345,6 +345,40 @@ export type ClientOptions = ProtocolOptions & { * regardless. The spec defines absent-or-≤0 as "immediately stale". */ defaultCacheTtlMs?: number; + + /** + * Default close posture for {@linkcode Client.close | close()}. + * + * `true` (or an object with a `timeoutMs`) makes every parameterless + * `close()` call wait for in-flight requests to settle before the + * transport closes — useful for SIGINT-style shutdowns where the caller + * does not know what is in flight. `false` or absent keeps today's + * behavior: the transport closes immediately and in-flight requests + * settle with a connection-closed error. + * + * An explicit argument to `close()` always wins over this default. + */ + gracefulClose?: boolean | { timeoutMs?: number }; +}; + +/** + * Options for {@linkcode Client.close | Client.close()}. + */ +export type ClientCloseOptions = { + /** + * Wait for in-flight requests to settle before closing the transport. + * `true` uses the default drain timeout (2s); an object sets + * `timeoutMs` explicitly. `false` closes immediately (the default + * behavior when absent, unless {@linkcode ClientOptions.gracefulClose} + * was set at construction). + * + * Without draining, the transport's teardown aborts in-flight HTTP + * requests that the server may have already answered, which + * instrumentation such as OpenTelemetry's undici instrumentation reports + * as aborted requests (`UND_ERR_ABORTED` on 200 OK responses). Draining + * lets those responses land first so telemetry reflects the real outcome. + */ + drainPendingRequests?: boolean | { timeoutMs?: number }; }; /** @@ -555,6 +589,8 @@ export class Client extends Protocol { */ private readonly _listChangedConfig?: ListChangedHandlers; private _enforceStrictCapabilities: boolean; + /** The constructor `gracefulClose` posture applied when `close()` gets no explicit argument. */ + private readonly _gracefulCloseDefault?: boolean | { timeoutMs?: number }; private _versionNegotiation?: VersionNegotiationOptions; private _supportedProtocolVersionsOption?: string[]; private _inputRequiredDriverConfig: ResolvedInputRequiredDriverConfig; @@ -615,9 +651,13 @@ export class Client extends Protocol { this._cache.resetForReconnect(); } - override async close(): Promise { + override async close(options?: ClientCloseOptions): Promise { + // An explicit argument wins over the constructor default. `false` + // (or absent with no default) drains nothing — the historical + // immediate-close behavior. + const resolved = options?.drainPendingRequests ?? this._gracefulCloseDefault; try { - await super.close(); + await super.close(resolved ? { drainPendingRequests: resolved } : undefined); } finally { // Per-connection state is cleared even when the transport's close // rejects, so a stale negotiated era / live listen state cannot @@ -637,6 +677,7 @@ export class Client extends Protocol { this._capabilities = options?.capabilities ? { ...options.capabilities } : {}; this._jsonSchemaValidator = options?.jsonSchemaValidator ?? new DefaultJsonSchemaValidator(); this._enforceStrictCapabilities = options?.enforceStrictCapabilities ?? false; + this._gracefulCloseDefault = options?.gracefulClose; this._versionNegotiation = options?.versionNegotiation; this._supportedProtocolVersionsOption = options?.supportedProtocolVersions; // Multi-round-trip auto-fulfilment driver (2026-07-28): on by default, diff --git a/packages/client/src/index.ts b/packages/client/src/index.ts index 0b5b6e86ea..14fa0981bd 100644 --- a/packages/client/src/index.ts +++ b/packages/client/src/index.ts @@ -67,7 +67,14 @@ export { PrivateKeyJwtProvider, StaticPrivateKeyJwtProvider } from './client/authExtensions'; -export type { CacheableRequestOptions, CallToolRequestOptions, ClientOptions, ConnectOptions, McpSubscription } from './client/client'; +export type { + CacheableRequestOptions, + CallToolRequestOptions, + ClientCloseOptions, + ClientOptions, + ConnectOptions, + McpSubscription +} from './client/client'; export { Client } from './client/client'; export { getSupportedElicitationModes } from './client/client'; export type { DiscoverAndRequestJwtAuthGrantOptions, JwtAuthGrantResult, RequestJwtAuthGrantOptions } from './client/crossAppAccess'; diff --git a/packages/client/test/client/gracefulClose.test.ts b/packages/client/test/client/gracefulClose.test.ts new file mode 100644 index 0000000000..0ddf503b39 --- /dev/null +++ b/packages/client/test/client/gracefulClose.test.ts @@ -0,0 +1,204 @@ +/** + * Graceful close — opt-in drain of in-flight requests (issue #1231). + * + * `close({ drainPendingRequests })` waits for in-flight requests to settle + * before the transport closes, so a completed-but-still-reading HTTP response + * is not torn down by the transport's abort (which OpenTelemetry's undici + * instrumentation reports as UND_ERR_ABORTED on a 200 OK). Default close + * behavior is unchanged. + */ +import type { JSONRPCMessage } from '@modelcontextprotocol/core-internal'; +import { InMemoryTransport } from '@modelcontextprotocol/core-internal'; +import { describe, expect, it } from 'vitest'; + +import { Client } from '../../src/client/client'; + +const flush = () => new Promise(r => setTimeout(r, 10)); + +type ScriptedServer = { + clientTx: InMemoryTransport; + serverTx: InMemoryTransport; + written: JSONRPCMessage[]; + /** Replies to the oldest outstanding non-initialize request. */ + reply: (result: Record) => void; + /** Replies to the oldest outstanding non-initialize request on a delay. */ + replyAfter: (ms: number, result: Record) => Promise; +}; + +/** + * A linked in-memory pair where the server auto-answers the legacy + * `initialize` handshake (so `connect()` resolves) but holds every other + * request until the test calls `reply()` / `replyAfter()`. + */ +async function scriptedLegacyServer(): Promise { + const [clientTx, serverTx] = InMemoryTransport.createLinkedPair(); + const written: JSONRPCMessage[] = []; + const pendingIds: (number | string)[] = []; + serverTx.onmessage = message => { + written.push(message); + const req = message as { id?: number | string; method?: string; params?: { protocolVersion?: string } }; + if (req.method === 'initialize' && req.id !== undefined) { + void serverTx.send({ + jsonrpc: '2.0', + id: req.id, + result: { + protocolVersion: req.params?.protocolVersion ?? '2025-06-18', + capabilities: {}, + serverInfo: { name: 'scripted', version: '1' } + } + }); + return; + } + if (req.method === 'notifications/initialized') { + return; + } + if (req.id !== undefined) { + pendingIds.push(req.id); + } + }; + await serverTx.start(); + const reply = (result: Record) => { + const id = pendingIds.shift(); + if (id === undefined) { + throw new Error('no pending request to reply to'); + } + void serverTx.send({ jsonrpc: '2.0', id, result }); + }; + return { + clientTx, + serverTx, + written, + reply, + replyAfter: async (ms: number, result: Record) => { + await new Promise(r => setTimeout(r, ms)); + reply(result); + } + }; +} + +async function connectClient(options?: ConstructorParameters[1]): Promise<{ client: Client; server: ScriptedServer }> { + const server = await scriptedLegacyServer(); + const client = new Client({ name: 'test-client', version: '1.0.0' }, options); + await client.connect(server.clientTx); + return { client, server }; +} + +/** Spies on the client transport's close() without changing behavior. */ +function spyTransportClose(client: Client): { closed: () => boolean } { + let closed = false; + const transport = client.transport!; + const originalClose = transport.close.bind(transport); + transport.close = async () => { + closed = true; + await originalClose(); + }; + return { closed: () => closed }; +} + +describe('Client.close graceful drain', () => { + it('default close() is unchanged: transport closes with a request in flight', async () => { + const { client } = await connectClient(); + const inFlight = client.request({ method: 'ping' }).catch(e => e); + await flush(); + await client.close(); + const settled = (await inFlight) as Error; + // The request is settled by the close itself, not by a response. + expect(settled).toBeInstanceOf(Error); + expect((settled as Error).message).toMatch(/closed/i); + }); + + it('close({ drainPendingRequests: true }) waits for the in-flight response before closing', async () => { + const { client, server } = await connectClient(); + let settled: unknown; + const inFlight = client + .request({ method: 'ping' }) + .then(r => (settled = r)) + .catch(e => (settled = e)); + await flush(); + + const spy = spyTransportClose(client); + const closing = client.close({ drainPendingRequests: true }); + + // The transport must still be open while the request is outstanding. + await flush(); + expect(spy.closed()).toBe(false); + + // The response lands on the still-open connection; the drain then + // completes and the transport closes. + await server.replyAfter(20, {}); + await inFlight; + await closing; + expect(spy.closed()).toBe(true); + expect(settled).toBeDefined(); + }); + + it('multiple in-flight requests all drain before the transport closes', async () => { + const { client, server } = await connectClient(); + const first = client.request({ method: 'ping' }).catch(e => e); + const second = client.request({ method: 'ping' }).catch(e => e); + await flush(); + + const spy = spyTransportClose(client); + const closing = client.close({ drainPendingRequests: true }); + await flush(); + expect(spy.closed()).toBe(false); + + server.reply({}); + await first; + await flush(); + // One of two requests still outstanding: no close yet. + expect(spy.closed()).toBe(false); + + server.reply({}); + await second; + await closing; + expect(spy.closed()).toBe(true); + }); + + it('falls back to a hard close after the drain timeout and requests settle with the close', async () => { + const { client } = await connectClient(); + const inFlight = client.request({ method: 'ping' }).catch(e => e); + await flush(); + + await client.close({ drainPendingRequests: { timeoutMs: 30 } }); + const settled = (await inFlight) as Error; + expect(settled).toBeInstanceOf(Error); + expect((settled as Error).message).toMatch(/closed/i); + }); + + it('drain resolves immediately when nothing is in flight', async () => { + const { client } = await connectClient(); + const start = Date.now(); + await client.close({ drainPendingRequests: true }); + expect(Date.now() - start).toBeLessThan(500); + }); + + it('ClientOptions.gracefulClose applies to parameterless close()', async () => { + const { client, server } = await connectClient({ gracefulClose: true }); + const inFlight = client.request({ method: 'ping' }).catch(e => e); + await flush(); + + const spy = spyTransportClose(client); + const closing = client.close(); + await flush(); + expect(spy.closed()).toBe(false); + + await server.replyAfter(20, {}); + await inFlight; + await closing; + expect(spy.closed()).toBe(true); + }); + + it('an explicit close({ drainPendingRequests: false }) overrides the constructor default', async () => { + const { client } = await connectClient({ gracefulClose: true }); + const inFlight = client.request({ method: 'ping' }).catch(e => e); + await flush(); + + const spy = spyTransportClose(client); + await client.close({ drainPendingRequests: false }); + expect(spy.closed()).toBe(true); + const settled = (await inFlight) as Error; + expect(settled).toBeInstanceOf(Error); + expect((settled as Error).message).toMatch(/closed/i); + }); +}); diff --git a/packages/core-internal/src/exports/public/index.ts b/packages/core-internal/src/exports/public/index.ts index bf985d4160..199d5d66ba 100644 --- a/packages/core-internal/src/exports/public/index.ts +++ b/packages/core-internal/src/exports/public/index.ts @@ -49,6 +49,7 @@ export { getDisplayName } from '../../shared/metadataUtils'; export type { BaseContext, ClientContext, + CloseOptions, NotificationOptions, ProgressCallback, ProtocolOptions, diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 637be389aa..64821b303e 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -522,6 +522,25 @@ type TimeoutInfo = { onTimeout: () => void; }; +/** + * Default wait before a graceful close gives up on in-flight requests and + * proceeds with the hard close (which settles them locally). + */ +const DEFAULT_DRAIN_TIMEOUT_MS = 2000; + +/** + * Options for {@linkcode Protocol.close | close()}. + */ +export type CloseOptions = { + /** + * Wait for in-flight requests to settle before closing the transport. + * `true` uses the default drain timeout (2s); an object sets `timeoutMs` + * explicitly. Absent or `false` closes immediately (the historical + * behavior). + */ + drainPendingRequests?: boolean | { timeoutMs?: number }; +}; + /* * Package-internal write access to Protocol's negotiated-protocol-version state. * @@ -562,6 +581,22 @@ export abstract class Protocol { private _requestHandlerAbortControllers: Map = new Map(); private _notificationHandlers: Map Promise> = new Map(); private _responseHandlers: Map void> = new Map(); + /** + * Message ids with a registered response handler whose request has not + * settled yet. Written by the request funnel before the message goes on + * the wire and deleted on every exit path (response, error, close). This + * is the drain set: {@linkcode Protocol._drainPendingRequests | + * _drainPendingRequests} waits on it so a graceful close can let + * in-flight requests finish before the transport goes down. + */ + private _pendingRequestIds: Set = new Set(); + /** + * Wake hook for an active graceful-close drain. Set by + * {@linkcode Protocol._drainPendingRequests | _drainPendingRequests} and + * fired by {@linkcode Protocol._releasePendingRequest | + * _releasePendingRequest} when the pending set reaches empty. + */ + private _drainNotify?: () => void; private _progressHandlers: Map = new Map(); private _timeoutInfo: Map = new Map(); private _pendingDebouncedNotifications = new Set(); @@ -829,6 +864,11 @@ export abstract class Protocol { this._responseHandlers = new Map(); this._progressHandlers.clear(); this._pendingDebouncedNotifications.clear(); + // The transport is down: every in-flight request is about to be + // settled with the connection-closed error below, so the drain set + // empties here and any graceful-close drain resolves immediately. + this._pendingRequestIds.clear(); + this._drainNotify?.(); for (const info of this._timeoutInfo.values()) { clearTimeout(info.timeoutId); @@ -1221,10 +1261,69 @@ export abstract class Protocol { return this._transport; } + /** + * Drops a message id from the pending set and wakes a graceful-close + * drain if that was the last outstanding request. All response-handler + * removal paths funnel through this. + */ + private _releasePendingRequest(messageId: number): void { + if (this._pendingRequestIds.delete(messageId) && this._pendingRequestIds.size === 0) { + this._drainNotify?.(); + } + } + + /** + * Waits for in-flight requests to settle before the transport closes. + * + * Resolves `true` once every pending request has left the pending set — + * it received its response or exited through its own cleanup path — and + * `false` if the drain timeout elapsed with requests still outstanding, + * in which case the caller proceeds with a hard close and the stragglers + * settle via {@linkcode Protocol._onclose | _onclose}. Never rejects. + * + * Used by the opt-in graceful close ({@linkcode Protocol.close | + * close({ drainPendingRequests })}): without draining, the transport's + * teardown aborts in-flight HTTP requests that had already been answered, + * which instrumentation such as OpenTelemetry's undici instrumentation + * reports as aborted requests (modelcontextprotocol/typescript-sdk#1231). + */ + protected _drainPendingRequests(timeoutMs: number): Promise { + if (this._pendingRequestIds.size === 0) { + return Promise.resolve(true); + } + return new Promise(resolve => { + let settled = false; + const finish = (drained: boolean) => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + this._drainNotify = undefined; + resolve(drained); + }; + const timer = setTimeout(() => finish(false), timeoutMs); + this._drainNotify = () => finish(true); + }); + } + /** * Closes the connection. + * + * With `drainPendingRequests`, in-flight requests are given the chance to + * finish before the transport closes; requests still outstanding after + * the drain timeout are abandoned to the normal close path, which settles + * them with a connection-closed error. Default behavior (no options) is + * unchanged: the transport closes immediately. */ - async close(): Promise { + async close(options?: CloseOptions): Promise { + if (options?.drainPendingRequests) { + const timeoutMs = + typeof options.drainPendingRequests === 'object' + ? (options.drainPendingRequests.timeoutMs ?? DEFAULT_DRAIN_TIMEOUT_MS) + : DEFAULT_DRAIN_TIMEOUT_MS; + await this._drainPendingRequests(timeoutMs); + } await this._transport?.close(); } @@ -1421,6 +1520,9 @@ export abstract class Protocol { const messageId = this._requestMessageId++; cleanupMessageId = messageId; + // The request is now in flight for drain purposes: registered + // before the send so a concurrent graceful close sees it. + this._pendingRequestIds.add(messageId); const jsonrpcRequest: JSONRPCRequest = { ...request, jsonrpc: '2.0', @@ -1585,6 +1687,8 @@ export abstract class Protocol { if (cleanupMessageId !== undefined) { this._responseHandlers.delete(cleanupMessageId); this._cleanupTimeout(cleanupMessageId); + // Last: wakes a graceful-close drain waiting on this id. + this._releasePendingRequest(cleanupMessageId); } }); } From 7e9700ac667e5ab90f520503c74e86ec6b4a9013 Mon Sep 17 00:00:00 2001 From: K4bain <296577378+K4bain@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:07:33 +0500 Subject: [PATCH 2/2] docs(protocol): fix typedoc inline-tag brace warning in drain JSDoc The open brace in {@linkcode Protocol.close | close({ drainPendingRequests })} is flagged by typedoc ('Encountered an open brace within an inline tag'), which fails the docs:check CI gate. Use plain backticks for the call form. --- packages/core-internal/src/shared/protocol.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/core-internal/src/shared/protocol.ts b/packages/core-internal/src/shared/protocol.ts index 64821b303e..eebc075988 100644 --- a/packages/core-internal/src/shared/protocol.ts +++ b/packages/core-internal/src/shared/protocol.ts @@ -1281,11 +1281,11 @@ export abstract class Protocol { * in which case the caller proceeds with a hard close and the stragglers * settle via {@linkcode Protocol._onclose | _onclose}. Never rejects. * - * Used by the opt-in graceful close ({@linkcode Protocol.close | - * close({ drainPendingRequests })}): without draining, the transport's - * teardown aborts in-flight HTTP requests that had already been answered, - * which instrumentation such as OpenTelemetry's undici instrumentation - * reports as aborted requests (modelcontextprotocol/typescript-sdk#1231). + * Used by the opt-in graceful close (`close({ drainPendingRequests })`): + * without draining, the transport's teardown aborts in-flight HTTP + * requests that had already been answered, which instrumentation such as + * OpenTelemetry's undici instrumentation reports as aborted requests + * (modelcontextprotocol/typescript-sdk#1231). */ protected _drainPendingRequests(timeoutMs: number): Promise { if (this._pendingRequestIds.size === 0) {