diff --git a/mcp/streamable.go b/mcp/streamable.go index f24badbb..3855a754 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -217,8 +217,29 @@ type StreamableHTTPOptions struct { // Requests using older protocol versions (including those routed through // the allowsessionsinstateless compatibility path) are unaffected. PropagateRequestCancellation bool + + // StreamKeepAlive 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 zero, [DefaultStreamKeepAlive] is used. A negative value disables + // keep-alives. + StreamKeepAlive time.Duration } +// DefaultStreamKeepAlive is the default value used for +// [StreamableHTTPOptions.StreamKeepAlive] when it is left at zero. +const DefaultStreamKeepAlive = 30 * time.Second + // DefaultMaxRequestBodyBytes is the default value used for // [StreamableHTTPOptions.MaxRequestBodyBytes] when it is left at zero. const DefaultMaxRequestBodyBytes = 4 << 20 // 4 MiB @@ -242,6 +263,9 @@ func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *Strea if h.opts.MaxRequestBodyBytes == 0 { h.opts.MaxRequestBodyBytes = DefaultMaxRequestBodyBytes } + if h.opts.StreamKeepAlive == 0 { + h.opts.StreamKeepAlive = DefaultStreamKeepAlive + } return h } @@ -427,6 +451,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 +678,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 +848,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 +879,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 +915,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 +1021,20 @@ 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{} + + // writeDeadline is the interval used to extend the HTTP write deadline + // before writing an SSE event or comment. It is zero when keep-alives are + // disabled. + writeDeadline time.Duration + // protocolVersion is the protocol version for this stream. protocolVersion string @@ -1032,6 +1084,98 @@ 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 + } +} + +// extendWriteDeadlineLocked keeps the server's slow-write guard active while +// allowing an idle SSE stream to remain open until its next keep-alive. +// +// s.mu must be held. +func (s *stream) extendWriteDeadlineLocked() { + if s.writeDeadline <= 0 { + return + } + _ = http.NewResponseController(s.w).SetWriteDeadline(time.Now().Add(2 * s.writeDeadline)) +} + +// 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 + } + s.extendWriteDeadlineLocked() + _, err := fmt.Fprint(s.w, ":\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 +1244,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 @@ -1132,9 +1277,11 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I } else { // SSE mode: write event to response writer. s.lastIdx++ + s.extendWriteDeadlineLocked() if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil { return done, err } + s.markWrittenLocked() } return done, nil } @@ -1674,6 +1821,9 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques } stream.isListen = isSubscriptionsListen stream.isBatch = isBatch + if stream.isListen { + stream.writeDeadline = c.streamKeepAlive + } // subscriptions/listen is inherently a long-lived SSE endpoint (SEP-2575): // it has no synchronous result, the response stream stays open until the @@ -1687,6 +1837,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 +1901,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. diff --git a/mcp/streamable_keepalive_test.go b/mcp/streamable_keepalive_test.go new file mode 100644 index 00000000..583c1f02 --- /dev/null +++ b/mcp/streamable_keepalive_test.go @@ -0,0 +1,572 @@ +// 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" + "net/http" + "net/http/httptest" + "runtime" + "strings" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" +) + +// 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 == ":": + 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_OnlyListenStreams checks that the keep-alive is +// confined to subscriptions/listen: a slow tools/call produces an SSE stream +// whose only content is the final response, however long it stays silent. +func TestStreamKeepAlive_OnlyListenStreams(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(":\n\n")) { + t.Errorf("keep-alive written on a tools/call stream:\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 + deadlines []time.Time + deadlineAtWrite []time.Time +} + +func (w *recordingWriter) Header() http.Header { return w.header } +func (w *recordingWriter) WriteHeader(int) {} +func (w *recordingWriter) SetWriteDeadline(d time.Time) error { + w.deadlines = append(w.deadlines, d) + return nil +} +func (w *recordingWriter) Write(p []byte) (int, error) { + if len(w.deadlines) > 0 { + w.deadlineAtWrite = append(w.deadlineAtWrite, w.deadlines[len(w.deadlines)-1]) + } + if w.err != nil { + return 0, w.err + } + return w.buf.Write(p) +} + +func TestStreamKeepAliveExtendsWriteDeadline(t *testing.T) { + const interval = 20 * time.Millisecond + + w := &recordingWriter{header: http.Header{}} + s := &stream{ + logger: ensureLogger(nil), + w: w, + done: make(chan struct{}), + writeDeadline: interval, + } + before := time.Now() + s.mu.Lock() + _, err := s.deliverLocked([]byte(`{"jsonrpc":"2.0","method":"test"}`), "", jsonrpc.ID{}, 0) + s.mu.Unlock() + if err != nil { + t.Fatal(err) + } + if len(w.deadlines) != 1 { + t.Fatalf("SetWriteDeadline calls = %d, want 1", len(w.deadlines)) + } + if got := w.deadlines[0]; got.Before(before.Add(interval)) || got.After(time.Now().Add(3*interval)) { + t.Errorf("write deadline = %v, want approximately %v after now", got, 2*interval) + } + if len(w.deadlineAtWrite) != 1 || !w.deadlineAtWrite[0].Equal(w.deadlines[0]) { + t.Errorf("write used deadline %v, want %v", w.deadlineAtWrite, w.deadlines) + } +} + +func TestStreamKeepAliveDisabledDoesNotExtendWriteDeadline(t *testing.T) { + w := &recordingWriter{header: http.Header{}} + s := &stream{ + logger: ensureLogger(nil), + w: w, + done: make(chan struct{}), + writeDeadline: -1, + } + s.mu.Lock() + _, err := s.deliverLocked([]byte(`{"jsonrpc":"2.0","method":"test"}`), "", jsonrpc.ID{}, 0) + s.mu.Unlock() + if err != nil { + t.Fatal(err) + } + if len(w.deadlines) != 0 { + t.Errorf("SetWriteDeadline calls = %d, want 0", len(w.deadlines)) + } +} + +func TestStreamKeepAliveDefault(t *testing.T) { + h := NewStreamableHTTPHandler(func(*http.Request) *Server { return nil }, nil) + if got := h.opts.StreamKeepAlive; got != DefaultStreamKeepAlive { + t.Errorf("default StreamKeepAlive = %v, want %v", got, DefaultStreamKeepAlive) + } + h = NewStreamableHTTPHandler(func(*http.Request) *Server { return nil }, &StreamableHTTPOptions{StreamKeepAlive: -1}) + if got := h.opts.StreamKeepAlive; got != -1 { + t.Errorf("disabled StreamKeepAlive = %v, want -1", got) + } +} + +func TestStreamKeepAliveSurvivesWriteTimeout(t *testing.T) { + const writeTimeout = 50 * time.Millisecond + + subCh := make(chan string, 1) + unsubCh := make(chan string, 1) + server := resourceSubServer(t, subCh, unsubCh) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + httpServer := &http.Server{Handler: handler, WriteTimeout: writeTimeout} + go func() { _ = httpServer.Serve(listener) }() + + ctx, cancel := context.WithCancel(t.Context()) + resp, err := http.DefaultClient.Do(listenRequest(t, ctx, "http://"+listener.Addr().String(), "file:///r1")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + cancel() + resp.Body.Close() + httpServer.Close() + }) + + lines := make(chan string, 16) + go func() { + defer close(lines) + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + select { + case lines <- scanner.Text(): + case <-ctx.Done(): + return + } + } + }() + + select { + case <-subCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for subscription") + } + time.Sleep(3 * writeTimeout) + if err := server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}); err != nil { + t.Fatal(err) + } + + deadline := time.After(5 * time.Second) + for { + select { + case line, ok := <-lines: + if !ok { + t.Fatal("stream ended before resource update") + } + if strings.Contains(line, `"method":"notifications/resources/updated"`) { + return + } + case <-deadline: + t.Fatal("timed out waiting for resource update") + } + } +} + +// TestStreamKeepAlive_IdleReset checks the timer semantics directly: 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, + writeDeadline: interval, + } + 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, ":\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") + } + if len(w.deadlines) < 2 { + t.Errorf("SetWriteDeadline calls = %d, want at least 2", len(w.deadlines)) + } + 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", -1, 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") + } + } + }) + } +} + +// 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) + } +}