diff --git a/mcp/streamable.go b/mcp/streamable.go index 813a3824..a206bfb1 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1072,6 +1072,11 @@ type stream struct { // goroutine parks on it instead of polling. committed chan struct{} + // keepAliveInterval is the interval the keep-alive writes at, and zero when + // none is running for the current w. Writes push the write deadline out by + // twice it; see markWrittenLocked. Reset by release. + keepAliveInterval time.Duration + // protocolVersion is the protocol version for this stream. protocolVersion string @@ -1123,13 +1128,20 @@ func (s *stream) release() { s.done = nil // may already be nil, if the stream is done or closed s.lastWrite = time.Time{} s.committed = nil + s.keepAliveInterval = 0 } -// markWrittenLocked records a write to s.w, for the keep-alive. +// markWrittenLocked records a write to s.w, for the keep-alive, and moves its +// write deadline past the next keep-alive. // // s.mu must be held. func (s *stream) markWrittenLocked() { s.lastWrite = time.Now() + if s.keepAliveInterval > 0 && s.w != nil { + // Best-effort: a writer managing its own deadlines reports + // ErrNotSupported, and a server with no WriteTimeout has none to move. + _ = http.NewResponseController(s.w).SetWriteDeadline(time.Now().Add(2 * s.keepAliveInterval)) + } if s.committed != nil { close(s.committed) s.committed = nil @@ -1145,6 +1157,8 @@ func (s *stream) startKeepAliveLocked(ctx context.Context, interval time.Duratio // be needed (see deliverLocked): wait for the first event. committed := make(chan struct{}) s.committed = committed + // Before the first write: markWrittenLocked reads it to move the deadline. + s.keepAliveInterval = interval go s.keepAlive(ctx, interval, committed) } diff --git a/mcp/streamable_writedeadline_test.go b/mcp/streamable_writedeadline_test.go new file mode 100644 index 00000000..232189fb --- /dev/null +++ b/mcp/streamable_writedeadline_test.go @@ -0,0 +1,97 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package mcp + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// TestStreamWriteDeadline_SurvivesServerWriteTimeout checks that a listen +// stream outlives an http.Server.WriteTimeout shorter than the stream: a write +// must land past the first deadline, however many do. +func TestStreamWriteDeadline_SurvivesServerWriteTimeout(t *testing.T) { + const interval = 25 * time.Millisecond + // Longer than one interval, so early keep-alives land inside the original + // deadline and only later ones need it moved. + const writeTimeout = 3 * interval + const window = 5 * writeTimeout + + // Buffered: the handlers send unconditionally, so a nil or full channel + // would block the subscribe this stream waits on. + server := resourceSubServer(t, make(chan string, 8), make(chan string, 8)) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true, StreamKeepAlive: interval}, + ) + httpServer := httptest.NewUnstartedServer(mustNotPanic(t, handler)) + httpServer.Config.WriteTimeout = writeTimeout + httpServer.Start() + 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) + } + + type line struct { + text string + at time.Duration + } + start := time.Now() + lines := make(chan line) + go func() { + defer close(lines) + sc := bufio.NewScanner(resp.Body) + for sc.Scan() { + select { + case lines <- line{text: sc.Text(), at: time.Since(start)}: + case <-ctx.Done(): + return + } + } + }() + + var last time.Duration + var afterDeadline int + deadline := time.After(window) +readLoop: + for { + select { + case l, ok := <-lines: + if !ok { + t.Fatalf("the stream ended inside the %v window, with the last keep-alive at %v: a write past the %v deadline failed", + window, last, writeTimeout) + } + if !strings.HasPrefix(l.text, ":") { + continue + } + last = l.at + if l.at > writeTimeout { + afterDeadline++ + } + case <-deadline: + break readLoop + } + } + + if afterDeadline == 0 { + t.Errorf("no write landed after the %v deadline in a %v window; the stream stopped being written to at %v", + writeTimeout, window, last) + } +}