mcp: extend the write deadline on every write to a keep-alive stream

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)
 }
 

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)
+	}
+}