fix: harden SSE keep-alive lifecycle (v1.x backport of the #2541 review fixes)

.changeset/keepalive-lifecycle-hardening.md

@@ -0,0 +1,5 @@
+---
+'@modelcontextprotocol/sdk': patch
+---
+
+Hardens the Streamable HTTP server transport's SSE lifecycle: deferred work cannot register streams after transport close, error cleanup preserves successor request mappings, invalid timer delays safely disable keep-alive, and SSE responses disable proxy buffering.

src/server/webStandardStreamableHttp.ts

@@ -156,7 +156,8 @@
      *
      * Comment frames are ignored by SSE parsers and never surface as messages.
      * Defaults to 15000 (per the WHATWG SSE spec recommendation of roughly every
-     * 15 seconds). Set to 0 to disable keep-alive frames.
+     * 15 seconds). Set to 0 to disable keep-alive frames; values below 1, above
+     * 2147483647, or non-finite values also disable the timer.
      */
     keepAliveMs?: number;
 }
@@ -227,6 +228,7 @@
     // when sessionId is not set (undefined), it means the transport is in stateless mode
     private sessionIdGenerator: (() => string) | undefined;
     private _started: boolean = false;
+    private _closed: boolean = false;
     private _hasHandledRequest: boolean = false;
     private _streamMapping: Map<string, StreamMapping> = new Map();
     private _requestToStreamMapping: Map<RequestId, string> = new Map();
@@ -271,7 +273,11 @@
      * clears itself if a write fails (stream already closed/cancelled).
      */
     private startKeepAlive(streamId: string, controller: ReadableStreamDefaultController<Uint8Array>, encoder: TextEncoder): void {
-        if (this._keepAliveMs <= 0) {
+        // A deferred arm (e.g. after an event-store await that straddled
+        // close()) must not outlive the transport: close()'s timer sweep has
+        // already run. Invalid timer delays disable keep-alive rather than
+        // letting setInterval clamp them to ~1ms and flood every stream.
+        if (!Number.isFinite(this._keepAliveMs) || this._keepAliveMs < 1 || this._keepAliveMs > 2_147_483_647 || this._closed) {
             return;
         }
         this.stopKeepAlive(streamId);
@@ -376,6 +382,10 @@
      * Returns a Response object (Web Standard)
      */
     async handleRequest(req: Request, options?: HandleRequestOptions): Promise<Response> {
+        if (this._closed) {
+            return this.createJsonErrorResponse(404, -32001, 'Session not found');
+        }
+
         // In stateless mode (no sessionIdGenerator), each request must use a fresh transport.
         // Reusing a stateless transport causes message ID collisions between clients.
         if (!this.sessionIdGenerator && this._hasHandledRequest) {
@@ -488,7 +498,8 @@
         const headers: Record<string, string> = {
             'Content-Type': 'text/event-stream',
             'Cache-Control': 'no-cache, no-transform',
-            Connection: 'keep-alive'
+            Connection: 'keep-alive',
+            'X-Accel-Buffering': 'no'
         };
 
         // After initialization, always include the session ID if we have one
@@ -547,7 +558,8 @@
             const headers: Record<string, string> = {
                 'Content-Type': 'text/event-stream',
                 'Cache-Control': 'no-cache, no-transform',
-                Connection: 'keep-alive'
+                Connection: 'keep-alive',
+                'X-Accel-Buffering': 'no'
             };
 
             if (this.sessionId !== undefined) {
@@ -583,6 +595,21 @@
                 }
             });
 
+            // The transport may have closed while the replay await was parked:
+            // its cleanup sweep ran before this stream was registered, so
+            // registering now would strand a mapping (and an open controller)
+            // on a dead transport, hang the client on a stream that never
+            // ends, and 409-block a later resume of this stream id. End the
+            // stream instead so the client observes termination.
+            if (this._closed) {
+                try {
+                    streamController!.close();
+                } catch {
+                    // Controller might already be closed
+                }
+                return new Response(readable, { headers });
+            }
+
             this._streamMapping.set(replayedStreamId, {
                 controller: streamController!,
                 encoder,
@@ -658,6 +685,12 @@
      * Handles POST requests containing JSON-RPC messages
      */
     private async handlePostRequest(req: Request, options?: HandleRequestOptions): Promise<Response> {
+        // Set once the SSE stream bookkeeping has been registered, so the
+        // catch below can reclaim it: an error after registration (a failed
+        // priming event write, a throwing message handler) returns an error
+        // response, leaving nothing that could ever cancel the discarded
+        // stream or retire the request mappings.
+        let reclaimSseBookkeeping: (() => void) | undefined;
         try {
             // Validate the Accept header
             const acceptHeader = req.headers.get('accept');
@@ -750,6 +783,12 @@
                 }
             }
 
+            // Request parsing and session initialization may await user/runtime
+            // work. Do not register or dispatch after close() has swept state.
+            if (this._closed) {
+                return this.createJsonErrorResponse(404, -32001, 'Session not found');
+            }
+
             // check if it contains requests
             const hasRequests = messages.some(isJSONRPCRequest);
 
@@ -812,8 +851,9 @@
 
             const headers: Record<string, string> = {
                 'Content-Type': 'text/event-stream',
-                'Cache-Control': 'no-cache',
-                Connection: 'keep-alive'
+                'Cache-Control': 'no-cache, no-transform',
+                Connection: 'keep-alive',
+                'X-Accel-Buffering': 'no'
             };
 
             // After initialization, always include the session ID if we have one
@@ -842,7 +882,14 @@
                 }
             }
 
-            this.startKeepAlive(streamId, streamController!, encoder);
+            reclaimSseBookkeeping = () => {
+                this._streamMapping.get(streamId)?.cleanup();
+                for (const message of messages) {
+                    if (isJSONRPCRequest(message) && this._requestToStreamMapping.get(message.id) === streamId) {
+                        this._requestToStreamMapping.delete(message.id);
+                    }
+                }
+            };
 
             // Write priming event if event store is configured (after mapping is set up)
             await this.writePrimingEvent(streamController!, encoder, streamId, clientProtocolVersion);
@@ -869,10 +916,19 @@
             // The server SHOULD NOT close the SSE stream before sending all JSON-RPC responses
             // This will be handled by the send() method when responses are ready
 
+            // Arm keep-alive only after the fallible awaits above — an error
+            // path returning 400 discards the Response, so nothing could ever
+            // cancel the stream and clear an already-armed timer. Skip if the
+            // responses already completed and cleaned the stream up.
+            if (this._streamMapping.get(streamId)?.controller === streamController!) {
+                this.startKeepAlive(streamId, streamController!, encoder);
+            }
+
             return new Response(readable, { status: 200, headers });
         } catch (error) {
             // return JSON-RPC formatted error
             this.onerror?.(error as Error);
+            reclaimSseBookkeeping?.();
             return this.createJsonErrorResponse(400, -32700, 'Parse error', { data: String(error) });
         }
     }
@@ -961,6 +1017,8 @@
     }
 
     async close(): Promise<void> {
+        this._closed = true;
+
         // Close all SSE connections
         this._streamMapping.forEach(({ cleanup }) => {
             cleanup();