mcp: add StreamableHTTPOptions.StreamKeepAlive for periodic SSE comments on idle streams

mcp/streamable.go

@@ -218,6 +218,21 @@ type StreamableHTTPOptions struct {
 	// Requests using older protocol versions (including those routed through
 	// the allowsessionsinstateless compatibility path) are unaffected.
 	PropagateRequestCancellation bool
+
+	// StreamKeepAlive, if non-zero, writes an SSE comment to any SSE response
+	// stream that has carried no bytes for this duration, so that idle-timeout
+	// intermediaries do not sever long-lived streams such as the response to a
+	// subscriptions/listen request. The 2026-07-28 Streamable HTTP
+	// specification encourages this keep-alive; SSE clients ignore comment
+	// lines, so it has no protocol-level effect.
+	//
+	// On streams using protocol version 2026-07-28 or later, the keep-alive
+	// starts only after a first event has committed the response headers,
+	// since the HTTP status may still have to change (see #1229). A write
+	// failure ends the stream as a disconnect.
+	//
+	// If StreamKeepAlive is the zero value, no keep-alive is written.
+	StreamKeepAlive time.Duration
 }
 
 // DefaultMaxRequestBodyBytes is the default value used for
@@ -452,6 +467,7 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.
 		Stateless:                   true,
 		EventStore:                  h.opts.EventStore,
 		jsonResponse:                h.opts.JSONResponse,
+		streamKeepAlive:             h.opts.StreamKeepAlive,
 		logger:                      h.opts.Logger,
 		shouldPropagateCancellation: info.usesNewProtocol && (info.isSubscriptionsListen || h.opts.PropagateRequestCancellation),
 	}
@@ -678,11 +694,12 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht
 	sessionID = server.opts.GetSessionID()
 
 	transport := &StreamableServerTransport{
-		SessionID:    sessionID,
-		Stateless:    false,
-		EventStore:   h.opts.EventStore,
-		jsonResponse: h.opts.JSONResponse,
-		logger:       h.opts.Logger,
+		SessionID:       sessionID,
+		Stateless:       false,
+		EventStore:      h.opts.EventStore,
+		jsonResponse:    h.opts.JSONResponse,
+		streamKeepAlive: h.opts.StreamKeepAlive,
+		logger:          h.opts.Logger,
 	}
 
 	// Sessions without a session ID (GetSessionID returned "") are ephemeral:
@@ -847,6 +864,13 @@ type StreamableServerTransport struct {
 	// to write their own streamable HTTP handler.
 	jsonResponse bool
 
+	// streamKeepAlive is the idle interval after which an SSE comment is
+	// written to a stream; see [StreamableHTTPOptions.StreamKeepAlive].
+	//
+	// TODO: streamKeepAlive should be exported, like jsonResponse and logger,
+	// once users can write their own streamable HTTP handler.
+	streamKeepAlive time.Duration
+
 	// optional logger provided through the [StreamableHTTPOptions.Logger].
 	//
 	// TODO(rfindley): logger should be exported, since we want to allow users
@@ -871,6 +895,7 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er
 		stateless:                   t.Stateless,
 		eventStore:                  t.EventStore,
 		jsonResponse:                t.jsonResponse,
+		streamKeepAlive:             t.streamKeepAlive,
 		logger:                      ensureLogger(t.logger), // see #556: must be non-nil
 		shouldPropagateCancellation: t.shouldPropagateCancellation,
 		incoming:                    make(chan jsonrpc.Message, 10),
@@ -906,6 +931,10 @@ type streamableServerConn struct {
 	jsonResponse bool
 	eventStore   EventStore
 
+	// streamKeepAlive is the idle interval for SSE keep-alive comments; zero
+	// disables them. See [StreamableHTTPOptions.StreamKeepAlive].
+	streamKeepAlive time.Duration
+
 	// shouldPropagateCancellation is true when the underlying HTTP request's
 	// lifetime IS the connection's cancellation signal (e.g., a stateless
 	// POST that owns a long-lived subscriptions/listen stream). It is read
@@ -1008,6 +1037,15 @@ type stream struct {
 	// It starts at -1 since indices start at 0.
 	lastIdx int
 
+	// lastWrite is when bytes were last written to w. The zero value means
+	// nothing has been written to the current w, so its headers are still
+	// uncommitted and the HTTP status can still be changed. Reset by release.
+	lastWrite time.Time
+
+	// committed, if non-nil, is closed by the first write to w. The keep-alive
+	// goroutine of a >= 2026-07-28 stream parks on it instead of polling.
+	committed chan struct{}
+
 	// protocolVersion is the protocol version for this stream.
 	protocolVersion string
 
@@ -1057,6 +1095,89 @@ func (s *stream) release() {
 	defer s.mu.Unlock()
 	s.w = nil
 	s.done = nil // may already be nil, if the stream is done or closed
+	s.lastWrite = time.Time{}
+	s.committed = nil
+}
+
+// markWrittenLocked records a write to s.w, for the keep-alive.
+//
+// s.mu must be held.
+func (s *stream) markWrittenLocked() {
+	s.lastWrite = time.Now()
+	if s.committed != nil {
+		close(s.committed)
+		s.committed = nil
+	}
+}
+
+// startKeepAliveLocked starts the keep-alive goroutine for the HTTP request
+// currently claiming the stream. ctx is that request's context.
+//
+// s.mu must be held, and s.protocolVersion must be set.
+func (s *stream) startKeepAliveLocked(ctx context.Context, interval time.Duration) {
+	var committed chan struct{}
+	if s.lastWrite.IsZero() && s.protocolVersion >= protocolVersion20260728 {
+		// Headers uncommitted: a SEP-2575 status override may still be needed
+		// (see deliverLocked), so wait for the first event.
+		committed = make(chan struct{})
+		s.committed = committed
+	}
+	go s.keepAlive(ctx, interval, committed)
+}
+
+// keepAlive writes an SSE comment to the stream whenever it has been idle for
+// interval, until ctx is done or the stream is released or closed. A failed
+// write closes the stream, releasing the hanging request so that a dead peer
+// is noticed within one interval.
+func (s *stream) keepAlive(ctx context.Context, interval time.Duration, committed chan struct{}) {
+	if committed != nil {
+		select {
+		case <-ctx.Done():
+			return
+		case <-committed:
+		}
+	}
+	timer := time.NewTimer(interval)
+	defer timer.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-timer.C:
+		}
+		if ctx.Err() != nil {
+			// The request ended; don't touch a stream that a later request may
+			// have re-acquired.
+			return
+		}
+		s.mu.Lock()
+		if s.done == nil {
+			s.mu.Unlock()
+			return
+		}
+		if wait := interval - time.Since(s.lastWrite); !s.lastWrite.IsZero() && wait > 0 {
+			s.mu.Unlock()
+			timer.Reset(wait)
+			continue
+		}
+		_, err := fmt.Fprint(s.w, ": keepalive\n\n")
+		if err == nil {
+			// Ignore returned error as flushing is best-effort.
+			_ = http.NewResponseController(s.w).Flush()
+			s.markWrittenLocked()
+		} else {
+			close(s.done)
+			s.done = nil
+		}
+		s.mu.Unlock()
+		if err != nil {
+			// A client that closes its connection cancels ctx before any write
+			// fails, so reaching this means the peer vanished without closing.
+			s.logger.Warn(fmt.Sprintf("Writing keep-alive: %v", err))
+			return
+		}
+		timer.Reset(interval)
+	}
 }
 
 // extractErrorStatus reports the HTTP status to send when the given
@@ -1125,6 +1246,7 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I
 	// SSE framing.
 	if overrideStatus != 0 {
 		s.w.Header().Set("Content-Type", "application/json")
+		s.w.Header().Del("X-Accel-Buffering")
 		s.w.WriteHeader(overrideStatus)
 		if _, err := s.w.Write(data); err != nil {
 			return done, err
@@ -1160,10 +1282,20 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I
 		if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil {
 			return done, err
 		}
+		s.markWrittenLocked()
 	}
 	return done, nil
 }
 
+// setSSEHeaders sets the response headers for an SSE stream. Accept was
+// checked in [StreamableHTTPHandler]. X-Accel-Buffering asks reverse proxies
+// not to buffer the response, as the spec recommends for SSE.
+func setSSEHeaders(h http.Header) {
+	h.Set("Content-Type", "text/event-stream")
+	h.Set("Connection", "keep-alive")
+	h.Set("X-Accel-Buffering", "no")
+}
+
 // doneLocked reports whether the stream is logically complete.
 //
 // s.requests was populated when reading the POST body, requests are deleted as
@@ -1379,9 +1511,10 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
 	}
 
 	w.Header().Set("Cache-Control", "no-cache, no-transform")
-	w.Header().Set("Content-Type", "text/event-stream") // Accept checked in [StreamableHTTPHandler]
-	w.Header().Set("Connection", "keep-alive")
+	setSSEHeaders(w.Header())
 
+	// written records that headers are committed, for the keep-alive.
+	written := false
 	if s.id == "" {
 		// Issue #410: the standalone SSE stream is likely not to receive messages
 		// for a long time. Ensure that headers are flushed.
@@ -1405,6 +1538,7 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
 		rc := http.NewResponseController(w)
 		// Ignore returned error as flushing is best-effort.
 		_ = rc.Flush()
+		written = true
 	}
 
 	for _, data := range toReplay {
@@ -1416,6 +1550,7 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
 		if _, err := writeEvent(w, e); err != nil {
 			return nil, nil
 		}
+		written = true
 	}
 
 	if tempStream || s.doneLocked() {
@@ -1429,6 +1564,12 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
 	s.done = make(chan struct{})
 	s.lastIdx = lastIdx
 	s.protocolVersion = protocolVersion
+	if written {
+		s.markWrittenLocked()
+	}
+	if c.streamKeepAlive > 0 {
+		s.startKeepAliveLocked(ctx, c.streamKeepAlive)
+	}
 	return s, s.done
 }
 
@@ -1710,8 +1851,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 	// Set response headers. Accept was checked in [StreamableHTTPHandler].
 	w.Header().Set("Cache-Control", "no-cache, no-transform")
 	if useSSE {
-		w.Header().Set("Content-Type", "text/event-stream")
-		w.Header().Set("Connection", "keep-alive")
+		setSSEHeaders(w.Header())
 	} else {
 		w.Header().Set("Content-Type", "application/json")
 	}
@@ -1773,6 +1913,14 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 			if _, err := writeEvent(w, e); err != nil {
 				c.logger.Warn(fmt.Sprintf("Writing priming event: %v", err))
 			}
+			stream.mu.Lock()
+			stream.markWrittenLocked()
+			stream.mu.Unlock()
+		}
+		if c.streamKeepAlive > 0 {
+			stream.mu.Lock()
+			stream.startKeepAliveLocked(req.Context(), c.streamKeepAlive)
+			stream.mu.Unlock()
 		}
 	}