diff --git a/mcp/streamable.go b/mcp/streamable.go index d1462d9f..036b0af4 100644 --- a/mcp/streamable.go +++ b/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() } } diff --git a/mcp/streamable_keepalive_test.go b/mcp/streamable_keepalive_test.go new file mode 100644 index 00000000..e925beff --- /dev/null +++ b/mcp/streamable_keepalive_test.go @@ -0,0 +1,540 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "sync/atomic" + "testing" + "time" +) + +// listenRequest returns a raw 2026-07-28 subscriptions/listen POST for uri. +func listenRequest(t *testing.T, ctx context.Context, url, uri string) *http.Request { + t.Helper() + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": methodSubscriptionsListen, + "params": map[string]any{ + "_meta": map[string]any{ + MetaKeyProtocolVersion: protocolVersion20260728, + MetaKeyClientInfo: map[string]any{"name": "new-proto-client", "version": "9.9"}, + MetaKeyClientCapabilities: map[string]any{}, + }, + "notifications": map[string]any{"resourceSubscriptions": []string{uri}}, + }, + }) + if err != nil { + t.Fatal(err) + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, methodSubscriptionsListen) + return req +} + +// TestStreamKeepAlive_ListenStream checks that a quiet subscriptions/listen +// stream carries periodic SSE comments once headers are committed by the +// acknowledgment, and that the stream still tears down normally when the +// client goes away. +func TestStreamKeepAlive_ListenStream(t *testing.T) { + const interval = 25 * time.Millisecond + const window = 16 * interval + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + server := resourceSubServer(t, subCh, unsubCh) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: interval}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + resp, err := http.DefaultClient.Do(listenRequest(t, ctx, httpServer.URL, "file:///r1")) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, body) + } + if got := resp.Header.Get("X-Accel-Buffering"); got != "no" { + t.Errorf("X-Accel-Buffering = %q, want %q", got, "no") + } + + // Read the stream for a while: expect the acknowledgment event, then + // keep-alive comments and nothing else. + deadline := time.After(window) + lines := make(chan string) + go func() { + defer close(lines) + sc := bufio.NewScanner(resp.Body) + for sc.Scan() { + select { + case lines <- sc.Text(): + case <-ctx.Done(): + return + } + } + }() + var comments, events int +loop: + for { + select { + case line, ok := <-lines: + if !ok { + t.Fatal("stream ended early") + } + switch { + case line == ": keepalive": + comments++ + case strings.HasPrefix(line, "event: "): + events++ + case line == "" || strings.HasPrefix(line, "data: ") || strings.HasPrefix(line, "id: "): + default: + t.Errorf("unexpected line %q", line) + } + case <-deadline: + break loop + } + } + if events != 1 { + t.Errorf("got %d events, want 1 (the acknowledgment)", events) + } + if comments < 3 { + t.Errorf("got %d keep-alive comments in %v, want at least 3", comments, window) + } + + // Closing the request still unwinds the listen handler. + cancel() + select { + case <-unsubCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for UnsubscribeHandler after client disconnect") + } +} + +// TestStreamKeepAlive_WaitsForFirstEvent checks that on a >= 2026-07-28 +// stream no comment is written while the response headers are uncommitted, so +// a SEP-2575 error override can still set the HTTP status. A slow tool call +// produces an SSE stream whose only content is the final response. +func TestStreamKeepAlive_WaitsForFirstEvent(t *testing.T) { + const interval = 10 * time.Millisecond + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "slow"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + time.Sleep(10 * interval) + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: interval}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + req, err := http.NewRequest(http.MethodPost, httpServer.URL, bytes.NewReader(newProtocolBody(t, "slow", struct{}{}))) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(protocolVersionHeader, protocolVersion20260728) + req.Header.Set(methodHeader, "tools/call") + req.Header.Set(nameHeader, "slow") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200; body = %s", resp.StatusCode, body) + } + if bytes.Contains(body, []byte(": keepalive")) { + t.Errorf("keep-alive written before the first event:\n%s", body) + } +} + +// recordingWriter is an http.ResponseWriter that records writes and can be +// made to fail. +type recordingWriter struct { + header http.Header + buf bytes.Buffer + err error +} + +func (w *recordingWriter) Header() http.Header { return w.header } +func (w *recordingWriter) WriteHeader(int) {} +func (w *recordingWriter) Write(p []byte) (int, error) { + if w.err != nil { + return 0, w.err + } + return w.buf.Write(p) +} + +// TestStreamKeepAlive_IdleReset checks the timer semantics directly: on a +// >= 2026-07-28 stream nothing is written before the first event, an event +// written between ticks defers the next comment by a full interval, and a +// failed write closes the stream. +func TestStreamKeepAlive_IdleReset(t *testing.T) { + const interval = 60 * time.Millisecond + + w := &recordingWriter{header: http.Header{}} + done := make(chan struct{}) + s := &stream{ + id: "s", + logger: ensureLogger(nil), + w: w, + done: done, + protocolVersion: protocolVersion20260728, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + s.mu.Lock() + s.startKeepAliveLocked(ctx, interval) + s.mu.Unlock() + + // Headers uncommitted: nothing is written however long we wait. + time.Sleep(3 * interval) + s.mu.Lock() + if got := w.buf.Len(); got != 0 { + t.Errorf("wrote %d bytes before the first event", got) + } + // Simulate a first event and let a comment land. + s.markWrittenLocked() + s.mu.Unlock() + time.Sleep(2 * interval) + s.mu.Lock() + if got := w.buf.String(); !strings.Contains(got, ": keepalive\n\n") { + t.Errorf("after the idle interval, wrote %q, want a comment", got) + } + // A fresh event resets the idle timer: the next comment must not arrive + // within the following interval. + w.buf.Reset() + s.markWrittenLocked() + s.mu.Unlock() + time.Sleep(interval / 2) + s.mu.Lock() + if got := w.buf.Len(); got != 0 { + t.Errorf("comment written %v after an event, before the idle interval elapsed", interval/2) + } + // Fail the next write: the stream is closed so the hanging request ends. + w.err = errors.New("peer gone") + s.mu.Unlock() + select { + case <-done: + case <-time.After(5 * interval): + t.Fatal("stream not closed after a failed keep-alive write") + } + s.mu.Lock() + if s.done != nil { + t.Error("done not cleared after close") + } + s.mu.Unlock() +} + +// idleTimeoutProxy forwards requests to upstream and, like nginx's +// proxy_read_timeout, drops a response whose body has been silent for idle. +func idleTimeoutProxy(t *testing.T, upstream string, idle time.Duration) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + out, err := http.NewRequestWithContext(req.Context(), req.Method, upstream+req.URL.RequestURI(), req.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + out.Header = req.Header.Clone() + resp, err := http.DefaultTransport.RoundTrip(out) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + for k, vs := range resp.Header { + w.Header()[k] = vs + } + w.WriteHeader(resp.StatusCode) + rc := http.NewResponseController(w) + gone := make(chan struct{}) + defer close(gone) + chunks := make(chan []byte) + go func() { + defer close(chunks) + buf := make([]byte, 4096) + for { + n, err := resp.Body.Read(buf) + if n > 0 { + select { + case chunks <- append([]byte(nil), buf[:n]...): + case <-gone: + return + } + } + if err != nil { + return + } + } + }() + for { + select { + case chunk, ok := <-chunks: + if !ok { + return + } + if _, err := w.Write(chunk); err != nil { + return + } + _ = rc.Flush() + case <-time.After(idle): + return // idle timeout: closes both the downstream response and, via defer, the upstream body + } + } + })) +} + +// TestStreamKeepAlive_SurvivesIdleTimeoutProxy is the scenario from the +// issue: behind an intermediary that drops silent responses, a quiet listen +// stream dies without the keep-alive and outlives the timeout with it, still +// delivering the next real notification. The SDK client is used end to end, +// which also checks that it ignores the comment lines. +func TestStreamKeepAlive_SurvivesIdleTimeoutProxy(t *testing.T) { + const idle = 600 * time.Millisecond + + for _, tc := range []struct { + name string + keepAlive time.Duration + survives bool + }{ + {"without keep-alive", 0, false}, + {"with keep-alive", idle / 6, true}, + } { + t.Run(tc.name, func(t *testing.T) { + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + server := resourceSubServer(t, subCh, unsubCh) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: tc.keepAlive}, + ) + upstream := httptest.NewServer(mustNotPanic(t, handler)) + defer upstream.Close() + proxy := idleTimeoutProxy(t, upstream.URL, idle) + defer proxy.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + events := make(chan string, 8) + client := NewClient(testImpl, &ClientOptions{ + ResourceUpdatedHandler: func(_ context.Context, req *ResourceUpdatedNotificationRequest) { + events <- req.Params.URI + }, + }) + cs, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: proxy.URL, MaxRetries: -1}, + &ClientSessionOptions{ProtocolVersion: protocolVersion20260728}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + if err := cs.Subscribe(ctx, &SubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatal(err) + } + <-subCh + + // Stay quiet for several idle periods. + time.Sleep(3 * idle) + + server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + select { + case <-events: + if !tc.survives { + t.Fatal("notification delivered although the proxy should have dropped the idle stream") + } + case <-time.After(idle): + if tc.survives { + t.Fatal("notification not delivered: the keep-alive did not keep the stream open") + } + } + if !tc.survives { + // The drop reached the server: the listen handler unwound. + select { + case <-unsubCh: + case <-time.After(5 * time.Second): + t.Fatal("UnsubscribeHandler not called after the proxy dropped the stream") + } + } + }) + } +} + +// commentCounter is an http.RoundTripper that counts SSE comment lines on +// every text/event-stream response body it sees. +type commentCounter struct { + next http.RoundTripper + comments atomic.Int64 +} + +func (c *commentCounter) RoundTrip(req *http.Request) (*http.Response, error) { + resp, err := c.next.RoundTrip(req) + if err != nil || !strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") { + return resp, err + } + pr, pw := io.Pipe() + body := resp.Body + go func() { + sc := bufio.NewScanner(io.TeeReader(body, pw)) + for sc.Scan() { + if strings.HasPrefix(sc.Text(), ":") { + c.comments.Add(1) + } + } + pw.CloseWithError(sc.Err()) + }() + resp.Body = struct { + io.Reader + io.Closer + }{pr, body} + return resp, nil +} + +// TestStreamKeepAlive_StatefulGETStream covers the acquireStream path: on a +// stateful server the standalone GET stream is kept alive too, and the SDK +// client keeps working through the comments. +func TestStreamKeepAlive_StatefulGETStream(t *testing.T) { + const interval = 25 * time.Millisecond + + server := NewServer(testImpl, nil) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{StreamKeepAlive: interval}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + counter := &commentCounter{next: http.DefaultTransport} + cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + HTTPClient: &http.Client{Transport: counter}, + }, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + time.Sleep(16 * interval) + if err := cs.Ping(ctx, nil); err != nil { + t.Fatalf("ping after keep-alives: %v", err) + } + // One ": ok" on connect, then keep-alives. + if got := counter.comments.Load(); got < 4 { + t.Errorf("saw %d comment lines on the GET stream in %v, want at least 4", got, 16*interval) + } +} + +// TestStreamKeepAlive_LegacyStreamFromStart checks that a stream on a protocol +// version before 2026-07-28 — which has no HTTP status to protect — is kept +// alive from the start, so a long tool call with no events is covered. +func TestStreamKeepAlive_LegacyStreamFromStart(t *testing.T) { + const interval = 25 * time.Millisecond + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "slow"}, + func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) { + time.Sleep(12 * interval) + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil + }) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: interval}, + ) + httpServer := httptest.NewServer(handler) + defer httpServer.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + counter := &commentCounter{next: http.DefaultTransport} + cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + HTTPClient: &http.Client{Transport: counter}, + }, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125}) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + if _, err := cs.CallTool(ctx, &CallToolParams{Name: "slow"}); err != nil { + t.Fatal(err) + } + if got := counter.comments.Load(); got < 3 { + t.Errorf("saw %d comment lines during a %v tool call, want at least 3", got, 12*interval) + } +} + +// TestStreamKeepAlive_NoGoroutineLeak checks that keep-alive goroutines end +// with their streams. +func TestStreamKeepAlive_NoGoroutineLeak(t *testing.T) { + const interval = 10 * time.Millisecond + + subCh := make(chan string, 8) + unsubCh := make(chan string, 8) + server := resourceSubServer(t, subCh, unsubCh) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: interval}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + for range 5 { + ctx, cancel := context.WithCancel(context.Background()) + resp, err := http.DefaultClient.Do(listenRequest(t, ctx, httpServer.URL, "file:///r1")) + if err != nil { + t.Fatal(err) + } + <-subCh + time.Sleep(3 * interval) + cancel() + resp.Body.Close() + <-unsubCh + } + + deadline := time.Now().Add(5 * time.Second) + for { + buf := make([]byte, 1<<20) + n := runtime.Stack(buf, true) + if !bytes.Contains(buf[:n], []byte("(*stream).keepAlive")) { + return + } + if time.Now().After(deadline) { + t.Fatalf("keep-alive goroutines still running after their streams ended:\n%s", buf[:n]) + } + time.Sleep(10 * time.Millisecond) + } +}