diff --git a/.changeset/close-retires-in-flight-request-state.md b/.changeset/close-retires-in-flight-request-state.md new file mode 100644 index 0000000000..23f17ca8f4 --- /dev/null +++ b/.changeset/close-retires-in-flight-request-state.md @@ -0,0 +1,19 @@ +--- +'@modelcontextprotocol/server': patch +--- + +Fix `WebStandardStreamableHTTPServerTransport.close()` leaving in-flight request state behind. + +`close()` cleared `_streamMapping` and `_requestResponseMap` but never cleared +`_requestToStreamMapping`, so request ids stayed resolvable after the transport was +closed. A `send()` suspended in `eventStore.storeEvent()` across a `close()` resumed +against the retired correlation and recorded a response that could never be delivered, +leaking both the response and its correlation for the lifetime of the process. `close()` +now retires the correlations, and `send()` re-checks the correlation after the +`storeEvent()` await — the same reason the stream is already re-read there. + +`close()` also now settles JSON-response-mode POSTs that are parked waiting for their +responses, answering `404 Session not found` instead of leaving the HTTP request hanging +until the platform's own timeout fires. + +Both paths also apply to `NodeStreamableHTTPServerTransport`, which wraps this transport. diff --git a/packages/server/src/server/streamableHttp.ts b/packages/server/src/server/streamableHttp.ts index c0f48560a2..896a1aec15 100644 --- a/packages/server/src/server/streamableHttp.ts +++ b/packages/server/src/server/streamableHttp.ts @@ -1054,14 +1054,22 @@ export class WebStandardStreamableHTTPServerTransport implements Transport { } this._closed = true; - // Close all SSE connections - for (const { cleanup } of this._streamMapping.values()) { - cleanup(); + // Close all SSE connections. JSON-response mode parks the POST's + // `Response` promise in `resolveJson` until every response is ready — + // settle it here, otherwise closing the transport mid-request leaves + // the HTTP request hanging until the platform's own timeout fires. + for (const stream of this._streamMapping.values()) { + stream.resolveJson?.(this.createJsonErrorResponse(404, -32_001, 'Session not found')); + stream.cleanup(); } this._streamMapping.clear(); - // Clear any pending responses + // Clear any pending responses, and the in-flight request correlations + // they belong to. Leaving `_requestToStreamMapping` populated keeps + // retired request ids resolvable in `send()`, so a closed transport + // still accepts responses it can never deliver. this._requestResponseMap.clear(); + this._requestToStreamMapping.clear(); this.onclose?.(); } @@ -1159,6 +1167,15 @@ export class WebStandardStreamableHTTPServerTransport implements Transport { // storeEvent() may have registered a resumed stream under this // streamId (mirrors the standalone path's post-await read). stream = this._streamMapping.get(streamId); + // The correlation itself can also disappear across the await — + // close() retires every in-flight request id. Re-check it for + // the same reason the stream is re-read above: without this, + // the response below is recorded into `_requestResponseMap` on + // a transport that can never deliver it, and both the entry and + // its correlation leak for the lifetime of the process. + if (!this._requestToStreamMapping.has(requestId)) { + throw new Error(`No connection established for request ID: ${String(requestId)}`); + } } // Write the event to the response stream — unless the resumed // stream's replay already delivered this exact eventId (the store diff --git a/packages/server/test/server/streamableHttp.test.ts b/packages/server/test/server/streamableHttp.test.ts index 9ec6baf46c..9996ebdeb3 100644 --- a/packages/server/test/server/streamableHttp.test.ts +++ b/packages/server/test/server/streamableHttp.test.ts @@ -1406,6 +1406,96 @@ describe('Zod v4', () => { expect(cleanupCalls).toEqual(['stream-1']); }); }); + + describe('close() retires in-flight request state', () => { + const tick = (): Promise => new Promise(resolve => setTimeout(resolve, 10)); + + /** + * Event store whose `storeEvent` can be parked mid-write, so a close() + * can be interleaved with an in-flight send(). + */ + function createGatedEventStore(): EventStore & { park: () => void; release: () => void } { + let gate: Promise | undefined; + let open: () => void = () => {}; + let counter = 0; + + return { + park(): void { + gate = new Promise(resolve => (open = resolve)); + }, + release(): void { + gate = undefined; + open(); + }, + async storeEvent(streamId: StreamId): Promise { + if (gate) { + await gate; + } + return `${streamId}_${counter++}`; + }, + async replayEventsAfter(): Promise { + return ''; + } + }; + } + + const request = (id: number): JSONRPCMessage => ({ jsonrpc: '2.0', id, method: 'ping' }) as JSONRPCMessage; + const response = (id: number): JSONRPCMessage => ({ jsonrpc: '2.0', id, result: {} }) as JSONRPCMessage; + + it('should reject a send parked in storeEvent across close() rather than recording an undeliverable response', async () => { + const eventStore = createGatedEventStore(); + const transport = new WebStandardStreamableHTTPServerTransport({ sessionIdGenerator: undefined, eventStore }); + transport.onmessage = () => {}; + + const res = await transport.handleRequest(createRequest('POST', [request(1), request(2)])); + expect(res.status).toBe(200); + + // First response lands; the batch is not complete, so nothing is retired yet. + await transport.send(response(1)); + + // Second response parks inside storeEvent, and close() sweeps while it is suspended. + eventStore.park(); + const parked = transport.send(response(2)); + await tick(); + await transport.close(); + eventStore.release(); + + await expect(parked).rejects.toThrow(/No connection established for request ID: 2/); + + // @ts-expect-error accessing private map for test purposes + expect(transport._requestToStreamMapping.size).toBe(0); + // @ts-expect-error accessing private map for test purposes + expect(transport._requestResponseMap.size).toBe(0); + }); + + it('should settle a JSON-mode POST parked across close() instead of leaving the HTTP request hanging', async () => { + const transport = new WebStandardStreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: true + }); + transport.onmessage = () => {}; + + // JSON mode parks the Response promise until every response is ready. + const pending = transport.handleRequest(createRequest('POST', request(1))); + await tick(); + await transport.close(); + + const res = await Promise.race([pending, tick().then(() => 'still pending' as const)]); + expect(res).not.toBe('still pending'); + expect((res as Response).status).toBe(404); + expectErrorResponse(await (res as Response).json(), -32_001, /Session not found/); + }); + + it('should not accept a response for a request id retired by close()', async () => { + const transport = new WebStandardStreamableHTTPServerTransport({ sessionIdGenerator: undefined }); + transport.onmessage = () => {}; + + await transport.handleRequest(createRequest('POST', request(1))); + await transport.close(); + + await expect(transport.send(response(1))).rejects.toThrow(/No connection established for request ID: 1/); + }); + }); }); describe('WebStandardStreamableHTTPServerTransport SSE keep-alive', () => {