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

mcp/streamable.go

@@ -217,6 +217,22 @@ 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 the response
+	// stream of a subscriptions/listen request whenever it has carried no
+	// bytes for this duration, so that idle-timeout intermediaries do not
+	// sever the long-lived stream. The 2026-07-28 Streamable HTTP
+	// specification encourages this keep-alive; SSE clients ignore comment
+	// lines, so it has no protocol-level effect. Other SSE responses are not
+	// kept alive.
+	//
+	// The keep-alive starts only after the listen acknowledgment has
+	// committed the response headers, since until then 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
@@ -427,6 +443,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),
 	}
@@ -653,11 +670,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:
@@ -822,6 +840,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
@@ -846,6 +871,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),
@@ -881,6 +907,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
@@ -983,6 +1013,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 parks on it instead of polling.
+	committed chan struct{}
+
 	// protocolVersion is the protocol version for this stream.
 	protocolVersion string
 
@@ -1032,6 +1071,86 @@ 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.
+func (s *stream) startKeepAliveLocked(ctx context.Context, interval time.Duration) {
+	// Nothing has been written yet, so a SEP-2575 status override may still
+	// be needed (see deliverLocked): 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
@@ -1100,6 +1219,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
@@ -1135,6 +1255,7 @@ 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
 }
@@ -1687,6 +1808,8 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 	if useSSE {
 		w.Header().Set("Content-Type", "text/event-stream")
 		w.Header().Set("Connection", "keep-alive")
+		// The spec recommends this for SSE: ask reverse proxies not to buffer.
+		w.Header().Set("X-Accel-Buffering", "no")
 	} else {
 		w.Header().Set("Content-Type", "application/json")
 	}
@@ -1749,6 +1872,11 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 				c.logger.Warn(fmt.Sprintf("Writing priming event: %v", err))
 			}
 		}
+		if c.streamKeepAlive > 0 && stream.isListen {
+			stream.mu.Lock()
+			stream.startKeepAliveLocked(req.Context(), c.streamKeepAlive)
+			stream.mu.Unlock()
+		}
 	}
 
 	// Publish incoming messages.