diff --git a/mcpcompat/server/keepalive.go b/mcpcompat/server/keepalive.go new file mode 100644 index 0000000..97d1967 --- /dev/null +++ b/mcpcompat/server/keepalive.go @@ -0,0 +1,193 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "log/slog" + "mime" + "net/http" + "sync" + "time" +) + +// keepAliveComment is the SSE comment written on each heartbeat tick. SSE +// comments (lines beginning with ':') are ignored by every spec-conforming SSE +// parser: go-sdk's client-side scanEvents cuts each line on the first ':', and +// an empty key yields no event. The payload MUST contain the colon — a line +// without one is a malformed SSE event. This deliberately diverges from +// mcp-go, which sent a full JSON-RPC ping REQUEST as the heartbeat: a +// conforming client would answer that ping with a response carrying an ID the +// go-sdk server never issued, generating unknown-request-ID handling on every +// tick. A comment is invisible at the JSON-RPC layer and costs the client +// nothing. +var keepAliveComment = []byte(": keep-alive\n\n") + +// keepAliveWriter wraps the http.ResponseWriter handed to the go-sdk GET handler +// and injects keepAliveComment at each heartbeat interval once — and only once — +// the response is established as an SSE stream (status 200 and Content-Type +// text/event-stream). All writes and flushes, the handler's and the keep-alive +// tick's alike, are serialized by mu. go-sdk writes one whole SSE frame per +// Write call (writeEvent buffers the full frame before issuing a single Write), +// so serializing at the call boundary preserves frame atomicity: a comment can +// never land in the middle of an event. This depends on go-sdk v1.6.1's +// frame-atomic writeEvent; the module is version-pinned and the end-to-end test +// would surface any regression (the client-side SSE scanner errors loudly on a +// malformed line). +// +// The wrapper forwards http.Flusher (routing go-sdk's ResponseController.Flush +// through mu) and Unwrap (preserving ResponseController deadline capabilities on +// the underlying writer). It deliberately does not forward http.Hijacker or +// http.Pusher: go-sdk's streamable GET path uses neither. +type keepAliveWriter struct { + interval time.Duration + logger *slog.Logger // may be nil; guarded at every use + + mu sync.Mutex + w http.ResponseWriter + rc *http.ResponseController // controller over the underlying w; used only under mu + wroteHeader bool // an explicit or implicit WriteHeader happened + streaming bool // SSE stream established; ticker running + stopped bool // stopKeepAlive called; no new tick write may start + + stop chan struct{} // closed by stopKeepAlive (via stopOnce) + stopOnce sync.Once + done chan struct{} // closed when the ticker goroutine exits (lifecycle state, assertable in tests) +} + +// newKeepAliveWriter wraps w for the passive SSE keep-alive at the given +// interval. The ticker is not started until the response is established as an +// SSE stream (see maybeStartLocked). +func newKeepAliveWriter(w http.ResponseWriter, interval time.Duration, logger *slog.Logger) *keepAliveWriter { + return &keepAliveWriter{ + interval: interval, + logger: logger, + w: w, + rc: http.NewResponseController(w), + stop: make(chan struct{}), + done: make(chan struct{}), + } +} + +// Header delegates to the underlying writer. No lock is needed: the handler +// mutates the header map only on its own goroutine, before WriteHeader (which +// takes mu), and the ticker goroutine never touches the header. +func (k *keepAliveWriter) Header() http.Header { + return k.w.Header() +} + +// WriteHeader delegates the status to the underlying writer and, on an SSE 200, +// starts the keep-alive ticker. +func (k *keepAliveWriter) WriteHeader(code int) { + k.mu.Lock() + defer k.mu.Unlock() + k.w.WriteHeader(code) + k.wroteHeader = true + k.maybeStartLocked(code) +} + +// Write delegates to the underlying writer. A Write before any WriteHeader is +// treated as an implicit WriteHeader(200) (net/http semantics), which also +// covers go-sdk's replay-stream path that writes without an explicit +// WriteHeader. +func (k *keepAliveWriter) Write(b []byte) (int, error) { + k.mu.Lock() + defer k.mu.Unlock() + if !k.wroteHeader { + k.wroteHeader = true + k.maybeStartLocked(http.StatusOK) + } + return k.w.Write(b) +} + +// Flush flushes the underlying writer under mu. Implementing http.Flusher +// directly is what routes go-sdk's http.ResponseController.Flush through our +// mutex (ResponseController resolves http.Flusher before Unwrap), guaranteeing +// a handler flush cannot interleave with a keep-alive write. Errors are ignored, +// matching go-sdk's own best-effort flush handling. +func (k *keepAliveWriter) Flush() { + k.mu.Lock() + defer k.mu.Unlock() + _ = k.rc.Flush() +} + +// Unwrap returns the underlying writer so http.ResponseController can reach its +// deadline/flush capabilities. This does not create a lock bypass for writes: +// ResponseController resolves http.Flusher (our locked Flush) before Unwrap, and +// the deadline setters reached through Unwrap do not write response bytes. +func (k *keepAliveWriter) Unwrap() http.ResponseWriter { + return k.w +} + +// stopKeepAlive releases the keep-alive. It is idempotent and is the single +// authoritative teardown path: the caller defers it, so it runs when the +// handler's ServeHTTP returns (context cancel, stream completion, session close, +// or panic unwind). After it returns, any in-flight tick write completed before +// it acquired mu (strictly before ServeHTTP returned to net/http, which is +// legal), and no subsequent tick write can start because stopped is set under +// the same mu every write checks. The ticker goroutine exits promptly via the +// closed stop channel. +func (k *keepAliveWriter) stopKeepAlive() { + k.stopOnce.Do(func() { close(k.stop) }) + k.mu.Lock() + k.stopped = true + k.mu.Unlock() +} + +// maybeStartLocked starts the ticker goroutine exactly once, and only for an +// established SSE stream. Precondition: mu is held. It is a no-op if the ticker +// already runs, the keep-alive was stopped, the interval is non-positive, the +// status is not 200, or the Content-Type is not text/event-stream — so error +// responses (denial 403, session 404/409/503, any http.Error) and non-SSE GETs +// never spawn a goroutine. +func (k *keepAliveWriter) maybeStartLocked(status int) { + if k.streaming || k.stopped || k.interval <= 0 || status != http.StatusOK { + return + } + mediaType, _, err := mime.ParseMediaType(k.w.Header().Get("Content-Type")) + if err != nil || mediaType != "text/event-stream" { + return + } + k.streaming = true + go k.run() +} + +// run is the per-connection ticker goroutine. It emits the first comment after +// one full interval (mcp-go parity: a positive interval opts in, the first +// heartbeat is not immediate) and exits when the keep-alive is stopped or a +// write fails (client gone, ServeHTTP unwinding). +func (k *keepAliveWriter) run() { + defer close(k.done) + t := time.NewTicker(k.interval) + defer t.Stop() + for { + select { + case <-k.stop: + return + case <-t.C: + if !k.writeComment() { + return + } + } + } +} + +// writeComment writes and flushes one keep-alive comment under mu. It returns +// false — signaling run to exit — if the keep-alive was stopped or the +// underlying write failed (the stream is dead; the deferred stopKeepAlive is +// about to run, so there is no point spinning). +func (k *keepAliveWriter) writeComment() bool { + k.mu.Lock() + defer k.mu.Unlock() + if k.stopped { + return false + } + if _, err := k.w.Write(keepAliveComment); err != nil { + if k.logger != nil { + k.logger.Debug("keep-alive write failed; stopping heartbeat", "error", err) + } + return false + } + _ = k.rc.Flush() + return true +} diff --git a/mcpcompat/server/keepalive_internal_test.go b/mcpcompat/server/keepalive_internal_test.go new file mode 100644 index 0000000..baef883 --- /dev/null +++ b/mcpcompat/server/keepalive_internal_test.go @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "bytes" + "errors" + "net/http" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// recordingWriter is a fake http.ResponseWriter that records writes, flushes and +// the status code. It is safe for concurrent use because keepAliveWriter +// serializes every Write/WriteHeader/Flush through its own mutex; the internal +// lock here guards the recording state against test-goroutine reads. +type recordingWriter struct { + mu sync.Mutex + hdr http.Header + buf bytes.Buffer + status int + flushes int + writeN int + writeErr error // when non-nil, Write returns it after recording nothing +} + +func newRecordingWriter() *recordingWriter { + return &recordingWriter{hdr: make(http.Header)} +} + +func (w *recordingWriter) Header() http.Header { return w.hdr } + +func (w *recordingWriter) WriteHeader(code int) { + w.mu.Lock() + defer w.mu.Unlock() + w.status = code +} + +func (w *recordingWriter) Write(b []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.writeErr != nil { + return 0, w.writeErr + } + w.writeN++ + return w.buf.Write(b) +} + +func (w *recordingWriter) Flush() { + w.mu.Lock() + defer w.mu.Unlock() + w.flushes++ +} + +func (w *recordingWriter) bytes() []byte { + w.mu.Lock() + defer w.mu.Unlock() + return append([]byte(nil), w.buf.Bytes()...) +} + +func (w *recordingWriter) commentCount() int { + return bytes.Count(w.bytes(), keepAliveComment) +} + +func (w *recordingWriter) flushCount() int { + w.mu.Lock() + defer w.mu.Unlock() + return w.flushes +} + +func (w *recordingWriter) setSSEHeaders() { + w.hdr.Set("Content-Type", "text/event-stream") +} + +// waitFor polls cond until it is true or a generous fixed deadline elapses; it +// fails the test on timeout so a broken expectation cannot hang the suite. +func waitFor(t *testing.T, cond func() bool, msg string) { + t.Helper() + const timeout = 2 * time.Second + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(time.Millisecond) + } + require.Failf(t, "timeout waiting for condition", "%s (waited %s)", msg, timeout) +} + +// requireClosed asserts ch is closed within a generous fixed window, so a stuck +// goroutine fails the test rather than hanging the suite. +func requireClosed(t *testing.T, ch <-chan struct{}, msg string) { + t.Helper() + select { + case <-ch: + case <-time.After(2 * time.Second): + require.FailNow(t, "timeout waiting for channel close", msg) + } +} + +// requireNotClosed asserts ch is NOT closed within the window. +func requireNotClosed(t *testing.T, ch <-chan struct{}, window time.Duration, msg string) { + t.Helper() + select { + case <-ch: + require.FailNow(t, "channel closed unexpectedly", msg) + case <-time.After(window): + } +} + +func TestKeepAlive_SSE200_EmitsCommentsWithFlush(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + k := newKeepAliveWriter(rw, 10*time.Millisecond, nil) + + k.WriteHeader(http.StatusOK) + + // Wait for several comments, then quiesce the ticker BEFORE asserting on the + // buffer. The test-goroutine reads (commentCount/flushCount/bytes) each take + // the recordingWriter lock independently, so reading them while the ticker is + // live is a TOCTOU race: a tick landing between the Write and Flush of one + // comment could skew flushCount below commentCount. stopKeepAlive + + // requireClosed(done) guarantee no further writes, so the state is stable. + waitFor(t, func() bool { return rw.commentCount() >= 3 }, + "expected at least 3 keep-alive comments on an SSE 200 stream") + k.stopKeepAlive() + requireClosed(t, k.done, "goroutine must exit after stop") + + // Stable state: every comment was followed by a flush, and the only bytes on + // the wire are keep-alive comments. + assert.Equal(t, rw.commentCount(), rw.flushCount(), + "each keep-alive comment must be followed by exactly one flush") + assert.Equal(t, rw.commentCount()*len(keepAliveComment), len(rw.bytes()), + "stream must carry only keep-alive comments") +} + +func TestKeepAlive_ImplicitHeader_ReplayPath(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() // Content-Type pre-set, as go-sdk's replay path does + k := newKeepAliveWriter(rw, 10*time.Millisecond, nil) + defer k.stopKeepAlive() + + // A Write with no prior WriteHeader is an implicit 200 (replay-stream path). + _, err := k.Write([]byte("event: message\ndata: {}\n\n")) + require.NoError(t, err) + + waitFor(t, func() bool { return rw.commentCount() >= 2 }, + "implicit-header SSE write must start the ticker") + requireClosed(t, doneAfterStop(k), "goroutine must exit after stop") +} + +// doneAfterStop stops the keep-alive and returns its done channel, for the +// common "assert the goroutine exits" pattern. +func doneAfterStop(k *keepAliveWriter) <-chan struct{} { + k.stopKeepAlive() + return k.done +} + +func TestKeepAlive_NeverStarts(t *testing.T) { + t.Parallel() + cases := []struct { + name string + setup func(rw *recordingWriter, k *keepAliveWriter) + }{ + { + name: "non-SSE 200 (application/json)", + setup: func(rw *recordingWriter, k *keepAliveWriter) { + rw.hdr.Set("Content-Type", "application/json") + k.WriteHeader(http.StatusOK) + }, + }, + { + name: "SSE content-type but non-200 status", + setup: func(rw *recordingWriter, k *keepAliveWriter) { + rw.setSSEHeaders() + k.WriteHeader(http.StatusForbidden) + }, + }, + { + name: "SSE content-type set but handler returns without writing", + setup: func(rw *recordingWriter, _ *keepAliveWriter) { + // SSE Content-Type is present, but the handler calls neither + // WriteHeader nor Write (e.g. it returns immediately): with no + // established response the ticker must never start. + rw.setSSEHeaders() + }, + }, + { + name: "empty content-type 200", + setup: func(_ *recordingWriter, k *keepAliveWriter) { + k.WriteHeader(http.StatusOK) + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + k := newKeepAliveWriter(rw, 5*time.Millisecond, nil) + defer k.stopKeepAlive() + + tc.setup(rw, k) + + // Wait well past several intervals: the ticker must never have run. + requireNotClosed(t, k.done, 60*time.Millisecond, "ticker goroutine must never start") + assert.Zero(t, rw.commentCount(), "no keep-alive comment must be written") + assert.False(t, k.streaming, "streaming must stay false when no SSE stream is established") + }) + } +} + +// TestKeepAlive_FirstCommentNotImmediate locks in mcp-go parity: run() waits one +// full interval before the first comment, so no comment may appear immediately +// on SSE establishment. A regression to an immediate first tick would fail here. +func TestKeepAlive_FirstCommentNotImmediate(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + // A long interval makes the "nothing yet" window unambiguous. + k := newKeepAliveWriter(rw, time.Second, nil) + defer k.stopKeepAlive() + + k.WriteHeader(http.StatusOK) + + // Shortly after the SSE 200 header — well before the first interval elapses — + // there must be no comment yet. + time.Sleep(50 * time.Millisecond) + assert.True(t, k.streaming, "ticker must be running after an SSE 200") + assert.Zero(t, rw.commentCount(), "first keep-alive comment must not be immediate (fires after one full interval)") +} + +// TestKeepAlive_ContentTypeWithParams verifies maybeStartLocked parses the +// media type (mime.ParseMediaType) rather than comparing the raw header, so an +// SSE Content-Type carrying parameters still starts the ticker. +func TestKeepAlive_ContentTypeWithParams(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.hdr.Set("Content-Type", "text/event-stream; charset=utf-8") + k := newKeepAliveWriter(rw, 5*time.Millisecond, nil) + defer k.stopKeepAlive() + + k.WriteHeader(http.StatusOK) + + waitFor(t, func() bool { return rw.commentCount() >= 1 }, + "a parameterized text/event-stream Content-Type must start the ticker") +} + +// TestKeepAlive_StopBeforeStart verifies the stopped guard in maybeStartLocked: +// if the keep-alive is stopped before the SSE stream is established, a later +// WriteHeader(200) on an SSE stream must NOT start the ticker. +func TestKeepAlive_StopBeforeStart(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + k := newKeepAliveWriter(rw, 5*time.Millisecond, nil) + + k.stopKeepAlive() // stop first + k.WriteHeader(http.StatusOK) // then establish the SSE stream + + requireNotClosed(t, k.done, 60*time.Millisecond, "ticker must not start after stopKeepAlive") + assert.Zero(t, rw.commentCount(), "no comment may be written when stopped before start") + assert.False(t, k.streaming, "streaming must stay false when stopped before start") +} + +func TestKeepAlive_StopPreventsLaterWrites(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + k := newKeepAliveWriter(rw, 10*time.Millisecond, nil) + k.WriteHeader(http.StatusOK) + + waitFor(t, func() bool { return rw.commentCount() >= 1 }, + "expected at least one comment before stop") + + k.stopKeepAlive() + requireClosed(t, k.done, "goroutine must exit promptly after stop") + + // No bytes may be written after stopKeepAlive returns. + after := len(rw.bytes()) + time.Sleep(50 * time.Millisecond) // > 2 intervals + assert.Equal(t, after, len(rw.bytes()), "no bytes may be written after stopKeepAlive returns") + + // stopKeepAlive is idempotent. + assert.NotPanics(t, k.stopKeepAlive, "stopKeepAlive must be idempotent") +} + +func TestKeepAlive_WriteErrorStopsGoroutine(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + rw.writeErr = errors.New("client gone") + k := newKeepAliveWriter(rw, 5*time.Millisecond, nil) + defer k.stopKeepAlive() + + // Establish the stream via an explicit WriteHeader (does not write bytes), + // so the ticker starts and the first tick's Write fails. + k.WriteHeader(http.StatusOK) + + requireClosed(t, k.done, "goroutine must exit when the underlying write errors") + assert.Zero(t, rw.commentCount(), "a failed write records no bytes") +} + +// TestKeepAlive_ConcurrentFramesNoInterleaving hammers Write with whole SSE +// frames while the ticker runs at a tiny interval and asserts that no frame or +// comment is ever split or interleaved. This is the core concurrency guarantee +// and is meaningful under the -race detector. +func TestKeepAlive_ConcurrentFramesNoInterleaving(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + k := newKeepAliveWriter(rw, time.Millisecond, nil) + + // Establish the SSE stream and start the ticker. + k.WriteHeader(http.StatusOK) + + const ( + writers = 8 + framesPerWriter = 200 + ) + frame := []byte("event: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/x\"}\n\n") + + var wg sync.WaitGroup + wg.Add(writers) + for i := 0; i < writers; i++ { + go func() { + defer wg.Done() + for j := 0; j < framesPerWriter; j++ { + _, err := k.Write(frame) + assert.NoError(t, err) + } + }() + } + wg.Wait() + k.stopKeepAlive() + requireClosed(t, k.done, "goroutine must exit after stop") + + // Every emitted token must be one of exactly two whole units: the frame or + // the keep-alive comment. Walk the buffer and match one whole unit at a time; + // any split or interleaving leaves an unmatched prefix and fails. + out := rw.bytes() + var frames, comments int + for len(out) > 0 { + switch { + case bytes.HasPrefix(out, frame): + out = out[len(frame):] + frames++ + case bytes.HasPrefix(out, keepAliveComment): + out = out[len(keepAliveComment):] + comments++ + default: + require.Failf(t, "interleaved output", + "output is not a clean concatenation of whole frames/comments; %d bytes unmatched", len(out)) + } + } + assert.Equal(t, writers*framesPerWriter, frames, "every whole frame must appear intact") + assert.Positive(t, comments, "the ticker should have emitted at least one comment") +} + +func TestKeepAlive_DisabledInterval(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + // A non-positive interval must never start the ticker even on a valid SSE 200. + k := newKeepAliveWriter(rw, 0, nil) + defer k.stopKeepAlive() + + k.WriteHeader(http.StatusOK) + requireNotClosed(t, k.done, 40*time.Millisecond, "ticker must not start when interval <= 0") + assert.Zero(t, rw.commentCount(), "disabled keep-alive must write no comments") +} + +// TestKeepAlive_UnwrapReturnsUnderlying verifies Unwrap exposes the underlying +// writer so http.ResponseController can reach it, and Flush routes through the +// wrapper under lock. +func TestKeepAlive_UnwrapAndFlush(t *testing.T) { + t.Parallel() + rw := newRecordingWriter() + rw.setSSEHeaders() + k := newKeepAliveWriter(rw, time.Hour, nil) // interval large: no tick during test + defer k.stopKeepAlive() + + assert.Same(t, rw, k.Unwrap(), "Unwrap must return the underlying writer") + + // A ResponseController built over the wrapper must resolve Flush through our + // locked Flush (Flusher is checked before Unwrap). + rc := http.NewResponseController(k) + require.NoError(t, rc.Flush()) + assert.Equal(t, 1, rw.flushCount(), "Flush via ResponseController must reach the underlying writer once") +} diff --git a/mcpcompat/server/keepalive_test.go b/mcpcompat/server/keepalive_test.go new file mode 100644 index 0000000..02ac228 --- /dev/null +++ b/mcpcompat/server/keepalive_test.go @@ -0,0 +1,168 @@ +// SPDX-FileCopyrightText: Copyright 2026 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server_test + +import ( + "bufio" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/server" +) + +// keepAliveMarker is the SSE comment the shim writes as its passive keep-alive. +// It is asserted on the wire; it must match keepAliveComment in keepalive.go. +const keepAliveMarker = ": keep-alive" + +// openGETStream opens the standalone GET SSE stream for a session and returns +// the response. The caller must close the body. +func openGETStream(ctx context.Context, t *testing.T, url, sid string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("MCP-Protocol-Version", "2025-06-18") + if sid != "" { + req.Header.Set("Mcp-Session-Id", sid) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// awaitKeepAliveComment scans an SSE stream until it reads a keep-alive comment +// line or the deadline (carried by the response body's request context) elapses. +// It returns true if a comment was seen. Running the scan in a goroutine with a +// timeout guard means a silent stream cannot hang the test. +func awaitKeepAliveComment(t *testing.T, resp *http.Response, timeout time.Duration) bool { + t.Helper() + found := make(chan bool, 1) + go func() { + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for sc.Scan() { + if strings.HasPrefix(sc.Text(), keepAliveMarker) { + found <- true + return + } + } + found <- false + }() + select { + case ok := <-found: + return ok + case <-time.After(timeout): + return false + } +} + +// TestKeepAlive_EndToEnd_StreamGetsComments verifies that with +// WithHeartbeatInterval set, an idle standalone GET SSE stream receives +// keep-alive comments, while the initialize POST response body carries none. +func TestKeepAlive_EndToEnd_StreamGetsComments(t *testing.T) { + t.Parallel() + stream := server.NewMCPServer("ka", "1.0.0") + addGreetTool(stream) + s := server.NewStreamableHTTPServer(stream, server.WithHeartbeatInterval(50*time.Millisecond)) + ts := httptest.NewServer(s) + defer ts.Close() + + // Establish the session via the shared handshake helper. + sid := initSession(t, ts.URL) + + // POST purity: a POST (JSON) response body must never carry keep-alive bytes. + // A unique request id keeps this literal distinct from the other tests'. + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + postResp := postRPC(ctx, t, ts.URL, sid, `{"jsonrpc":"2.0","id":8801,"method":"tools/list","params":{}}`) + postBody, err := io.ReadAll(postResp.Body) + require.NoError(t, err) + _ = postResp.Body.Close() + assert.NotContains(t, string(postBody), keepAliveMarker, + "a POST (JSON) response must not carry keep-alive bytes") + + // Open the standalone GET stream: it must receive keep-alive comments. + streamCtx, streamCancel := context.WithTimeout(t.Context(), 5*time.Second) + defer streamCancel() + getResp := openGETStream(streamCtx, t, ts.URL, sid) + defer func() { _ = getResp.Body.Close() }() + require.Equal(t, http.StatusOK, getResp.StatusCode) + require.True(t, strings.HasPrefix(getResp.Header.Get("Content-Type"), "text/event-stream"), + "GET stream must be an SSE stream") + + assert.True(t, awaitKeepAliveComment(t, getResp, 3*time.Second), + "idle GET stream must receive a keep-alive comment when a heartbeat is configured") +} + +// TestKeepAlive_Disabled_IdleStreamSilent verifies mcp-go parity: with no +// heartbeat configured, an idle GET stream receives no keep-alive comments. +func TestKeepAlive_Disabled_IdleStreamSilent(t *testing.T) { + t.Parallel() + stream := server.NewMCPServer("noka", "1.0.0") + addGreetTool(stream) + // No WithHeartbeatInterval: keep-alive disabled. + s := server.NewStreamableHTTPServer(stream) + ts := httptest.NewServer(s) + defer ts.Close() + + sid := initSession(t, ts.URL) + + streamCtx, streamCancel := context.WithTimeout(t.Context(), 2*time.Second) + defer streamCancel() + getResp := openGETStream(streamCtx, t, ts.URL, sid) + defer func() { _ = getResp.Body.Close() }() + require.Equal(t, http.StatusOK, getResp.StatusCode) + + // Within a bounded window the idle stream must stay silent (no comment). + assert.False(t, awaitKeepAliveComment(t, getResp, 400*time.Millisecond), + "idle GET stream must be silent when no heartbeat is configured") +} + +// TestKeepAlive_RehydratedStreamGetsComments verifies the keep-alive covers the +// cross-replica serveRehydrated path: a session initialized on replica A gets a +// GET stream on replica B (rehydrated) that receives keep-alive comments. +func TestKeepAlive_RehydratedStreamGetsComments(t *testing.T) { + t.Parallel() + mgr := newSharedSessionManager() + + streamA := server.NewMCPServer("A", "1.0.0") + addGreetTool(streamA) + sA := server.NewStreamableHTTPServer(streamA, server.WithSessionIdManager(mgr)) + tsA := httptest.NewServer(sA) + defer tsA.Close() + + // Replica B carries the heartbeat. + streamB := server.NewMCPServer("B", "1.0.0") + addGreetTool(streamB) + sB := server.NewStreamableHTTPServer(streamB, + server.WithSessionIdManager(mgr), + server.WithHeartbeatInterval(50*time.Millisecond)) + tsB := httptest.NewServer(sB) + defer tsB.Close() + + // Initialize on A, then open the GET stream on B (which must rehydrate). + sid := initSession(t, tsA.URL) + // Prime the rehydrated session on B with a request so the stream attaches to + // a known session. + require.Contains(t, listToolNames(t, tsB.URL, sid), "greet") + + streamCtx, streamCancel := context.WithTimeout(t.Context(), 5*time.Second) + defer streamCancel() + getResp := openGETStream(streamCtx, t, tsB.URL, sid) + defer func() { _ = getResp.Body.Close() }() + require.Equal(t, http.StatusOK, getResp.StatusCode) + require.True(t, strings.HasPrefix(getResp.Header.Get("Content-Type"), "text/event-stream"), + "rehydrated GET stream must be an SSE stream") + + assert.True(t, awaitKeepAliveComment(t, getResp, 3*time.Second), + "rehydrated GET stream must receive a keep-alive comment") +} diff --git a/mcpcompat/server/server.go b/mcpcompat/server/server.go index e1997db..111833a 100644 --- a/mcpcompat/server/server.go +++ b/mcpcompat/server/server.go @@ -38,7 +38,6 @@ import ( "net/http" "strings" "sync" - "time" "github.com/modelcontextprotocol/go-sdk/jsonrpc" gosdk "github.com/modelcontextprotocol/go-sdk/mcp" @@ -410,22 +409,7 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) { // middleware installed by this function syncs that session's overlay tools and // resources onto its own server once the OnRegisterSession hooks have run. This // mirrors mcp-go, whose per-session tools were dispatched per connection. -func (s *MCPServer) buildServer(genSessionID func() string, _ time.Duration) (*gosdk.Server, error) { - // NOTE: the keepalive parameter is intentionally unused. Wave 3 wired - // WithHeartbeatInterval to go-sdk's ServerOptions.KeepAlive, but go-sdk's - // KeepAlive sends an active ping REQUEST (server→client) each interval and - // session.Close()s on failure. Under JSONResponse mode that ping routes to - // the standalone SSE stream; a JSON-only client (no GET stream, the shim - // client's own default) has that stream unconnected → ping rejected → - // session closed on the first tick. mcp-go's heartbeat was passive (pings - // written only to an existing GET stream, never awaited, never - // terminating). Wiring KeepAlive therefore evicts every healthy JSON-only - // client session at the heartbeat interval. The interval is stored by - // WithHeartbeatInterval (see transports.go) but NOT wired to KeepAlive. - // - // TODO(issue #156): implement a passive keep-alive (SSE comment injection - // via a ResponseWriter wrapper around the go-sdk handler) matching - // mcp-go's design, rather than the active ping go-sdk's KeepAlive performs. +func (s *MCPServer) buildServer(genSessionID func() string) (*gosdk.Server, error) { s.mu.RLock() tools := make(map[string]ServerTool, len(s.tools)) for k, v := range s.tools { @@ -452,10 +436,10 @@ func (s *MCPServer) buildServer(genSessionID func() string, _ time.Duration) (*g // DefaultPageSize (1000) in place, preserving the pre-existing behavior // for callers that do not set WithPageSize. PageSize: s.pageSize, - // KeepAlive is intentionally NOT set: see the note on buildServer's - // keepalive parameter above. go-sdk's KeepAlive sends an active ping - // request that closes sessions without a connected standalone SSE - // stream, incompatible with JSONResponse mode + JSON-only clients. + // KeepAlive is deliberately NOT set. go-sdk's KeepAlive is an active ping + // request that evicts JSON-mode sessions without a connected standalone + // SSE stream (incompatible with JSONResponse + JSON-only clients); the + // shim's passive keep-alive lives in keepalive.go instead. InitializedHandler: func(ctx context.Context, req *gosdk.InitializedRequest) { if req == nil || req.Session == nil { return @@ -542,13 +526,11 @@ func addGlobalTool(srv *gosdk.Server, gt *gosdk.Tool, h gosdk.ToolHandler, name // getServerFunc returns a getServer callback for the go-sdk HTTP/SSE handlers, // which invoke it once per new client session. genSessionID (may be nil) is the -// session-ID generator to install on each per-session server. The keepalive -// parameter is retained for signature stability but is NOT wired to go-sdk's -// KeepAlive (see buildServer); it is effectively ignored. On a build error it -// logs and returns nil, which the go-sdk handler surfaces as an HTTP 400. -func (s *MCPServer) getServerFunc(genSessionID func() string, keepalive time.Duration) func(*http.Request) *gosdk.Server { +// session-ID generator to install on each per-session server. On a build error +// it logs and returns nil, which the go-sdk handler surfaces as an HTTP 400. +func (s *MCPServer) getServerFunc(genSessionID func() string) func(*http.Request) *gosdk.Server { return func(*http.Request) *gosdk.Server { - srv, err := s.buildServer(genSessionID, keepalive) + srv, err := s.buildServer(genSessionID) if err != nil { if s.logger != nil { s.logger.Error("building per-session MCP server", "error", err) diff --git a/mcpcompat/server/server_internal_test.go b/mcpcompat/server/server_internal_test.go index 54ee940..390325f 100644 --- a/mcpcompat/server/server_internal_test.go +++ b/mcpcompat/server/server_internal_test.go @@ -98,7 +98,7 @@ func TestForgetSession_ClosesNotifChannel(t *testing.T) { func TestSetSessionTools_NonObjectSchemaDoesNotPanic(t *testing.T) { t.Parallel() s := NewMCPServer("s", "1") - srv, err := s.buildServer(nil, 0) + srv, err := s.buildServer(nil) require.NoError(t, err) cs := s.sessionFor("sid-bad-schema") @@ -172,7 +172,7 @@ func TestBuildServer_GlobalAndSessionTools(t *testing.T) { // Building the global server (with the globally-registered tool) must succeed. // Per-session overlays are no longer baked in here; they are synced onto the // per-session server by syncSessionTools once the session registers. - srv, err := s.buildServer(nil, 0) + srv, err := s.buildServer(nil) require.NoError(t, err) require.NotNil(t, srv) @@ -195,7 +195,7 @@ func TestBuildServer_WithSessionIDGenerator(t *testing.T) { called := false gen := func() string { called = true; return "generated-id" } - srv, err := s.buildServer(gen, 0) + srv, err := s.buildServer(gen) require.NoError(t, err) require.NotNil(t, srv) // The generator is installed on the server (invoked by the SDK per new diff --git a/mcpcompat/server/transports.go b/mcpcompat/server/transports.go index cd276c8..2054898 100644 --- a/mcpcompat/server/transports.go +++ b/mcpcompat/server/transports.go @@ -60,7 +60,7 @@ type stdioConfig struct{} // ServeStdio runs the MCP server over stdio until the context is done. It // mirrors mcp-go's server.ServeStdio. func ServeStdio(server *MCPServer, _ ...StdioOption) error { - srv, err := server.buildServer(nil, 0) + srv, err := server.buildServer(nil) if err != nil { return err } @@ -84,7 +84,9 @@ type StreamableHTTPServer struct { contextFunc HTTPContextFunc sessionIDMgr SessionIdManager callGate CallGate - heartbeat time.Duration + // heartbeat is the passive SSE keep-alive comment interval for GET streams; + // ≤ 0 disables it. See WithHeartbeatInterval. + heartbeat time.Duration // disableLocalhostProtection turns off go-sdk's DNS-rebinding/localhost // protection (which 403s requests on a loopback listener with a non-localhost // Host header). mcp-go had no such protection; local proxies with custom Host @@ -160,31 +162,36 @@ func WithSessionIdManager(manager SessionIdManager) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.sessionIDMgr = manager } } -// WithHeartbeatInterval sets the keep-alive ping interval. This option is -// currently a NO-OP: it stores the interval but does NOT wire it to go-sdk's -// ServerOptions.KeepAlive. go-sdk's KeepAlive sends an active ping REQUEST -// (server→client) each interval and session.Close()s on failure; under -// JSONResponse mode that ping routes to the standalone SSE stream, and a -// JSON-only client (no GET stream, the shim client's own default) has that -// stream unconnected → ping rejected → session closed on the first tick. That -// would evict every healthy JSON-only client session at the interval. mcp-go's -// heartbeat was passive (pings written only to an existing GET stream, never -// awaited, never terminating). +// WithHeartbeatInterval sets the passive SSE keep-alive interval for GET +// streams. A positive interval enables the keep-alive; an interval ≤ 0 disables +// it (mcp-go parity: mcp-go's default was no heartbeat, a positive interval +// opts in). The keep-alive is implemented at the HTTP layer by keepAliveWriter, +// which wraps the GET stream's ResponseWriter and injects an SSE comment once +// per interval, on BOTH the local go-sdk handler path and the cross-replica +// serveRehydrated path. The first comment is emitted after one full interval. // -// NOTE: the passive keep-alive is more load-bearing than "nice to have" now -// that elicitation and all server→client notifications (list_changed, -// progress, logging, resources/updated, elicitation/complete) structurally -// depend on the client's idle standalone SSE stream staying open. Any -// intermediary (LB, reverse proxy, ToolHive's own proxy) will reap that idle -// stream on timeout — after which elicitation fails (ErrRejected) and -// server-initiated notifications are silently dropped, with no reconnect. It -// also feeds unbounded session growth (abandoned sessions are never reaped). +// The payload is an SSE comment (": keep-alive"), NOT a JSON-RPC ping. This +// diverges deliberately from mcp-go, which sent a full ping REQUEST event: a +// conforming client answers a ping with a response POST carrying an ID the +// go-sdk server never issued, so every tick would generate unknown-request-ID +// handling. An SSE comment is ignored by every conforming SSE parser and is +// invisible at the JSON-RPC layer (see keepAliveComment). // -// TODO(issue #156): implement a passive keep-alive (SSE comment injection via -// a ResponseWriter wrapper around the go-sdk handler) matching mcp-go's design -// and wire this option to it. This should be prioritized before internet-facing -// vMCP use; a max-session cap / idle sweep would also help. Until then the -// value is stored but unused. +// This does NOT wire go-sdk's ServerOptions.KeepAlive, and must not: go-sdk's +// KeepAlive sends an ACTIVE ping request (server→client) each interval and +// session.Close()s on failure. Under JSONResponse mode that ping routes to the +// standalone SSE stream, so a JSON-only client (no GET stream — the shim +// client's own default) would have its session evicted on the first tick. The +// passive comment-based keep-alive avoids that by writing only to an +// already-open GET stream and never awaiting a response. +// +// The keep-alive is load-bearing, not merely nice to have: elicitation and all +// server→client notifications (list_changed, progress, logging, +// resources/updated, elicitation/complete) structurally depend on the client's +// idle standalone SSE stream staying open. Any intermediary (LB, reverse proxy, +// ToolHive's own proxy) reaps an idle stream on timeout — after which +// elicitation fails (ErrRejected) and server-initiated notifications are +// silently dropped, with no reconnect. func WithHeartbeatInterval(interval time.Duration) StreamableHTTPOption { return func(s *StreamableHTTPServer) { s.heartbeat = interval } } @@ -241,9 +248,8 @@ func (s *StreamableHTTPServer) build() { gen = s.sessionIDMgr.Generate } // Validate the server configuration once up-front so a bad registration - // surfaces as a clean 500 rather than a per-request nil. Pass 0 for - // keepalive: WithHeartbeatInterval is a documented no-op (see its doc). - if _, err := s.mcp.buildServer(gen, 0); err != nil { + // surfaces as a clean 500 rather than a per-request nil. + if _, err := s.mcp.buildServer(gen); err != nil { s.buildErr = err return } @@ -261,9 +267,10 @@ func (s *StreamableHTTPServer) build() { } // A fresh go-sdk server per client session lets each session carry its own // tool/resource overlay (mcp-go's per-session projection), synced by the - // registration middleware buildServer installs. Keepalive is 0 - // (WithHeartbeatInterval is a no-op; see its doc). - s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen, 0), opts) + // registration middleware buildServer installs. The heartbeat keep-alive + // is applied at the HTTP layer (keepAliveWriter, see WithHeartbeatInterval), + // not through go-sdk's active KeepAlive. + s.handler = gosdk.NewStreamableHTTPHandler(s.mcp.getServerFunc(gen), opts) }) } @@ -276,6 +283,12 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) http.Error(w, fmt.Sprintf("building server: %v", s.buildErr), http.StatusInternalServerError) return } + // Wrap the GET stream with the passive keep-alive before either dispatch + // branch (local handler or serveRehydrated) so both are covered. The wrap is + // pure observation; stopKA is the single authoritative teardown and must run + // on every return path, hence the defer here at the top. + w, stopKA := s.wrapKeepAlive(w, r) + defer stopKA() ensureAcceptMediaTypes(r) if s.contextFunc != nil { r = r.WithContext(s.contextFunc(r.Context(), r)) @@ -408,6 +421,21 @@ func (s *StreamableHTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) s.handler.ServeHTTP(w, r) } +// wrapKeepAlive wraps w with the passive SSE keep-alive for GET requests when a +// heartbeat interval is configured. All other requests (POST/DELETE, or the +// heartbeat disabled) pass through unchanged with a no-op stop. Keeping the +// method gate here — rather than a branch in ServeHTTP — holds the gocyclo +// budget flat and keeps the keep-alive logic in one place. POST is doubly +// safe: it is excluded here by method, and even if wrapped the Content-Type +// gate (application/json under JSONResponse) would keep the ticker off. +func (s *StreamableHTTPServer) wrapKeepAlive(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, func()) { + if r.Method != http.MethodGet || s.heartbeat <= 0 { + return w, func() {} + } + k := newKeepAliveWriter(w, s.heartbeat, s.mcp.logger) + return k, k.stopKeepAlive +} + // serveRehydrated routes a request for a session created on another replica. // It validates the session ID against the shared SessionIdManager and serves it // through a locally-reconstructed session, matching mcp-go's behavior where any @@ -505,7 +533,7 @@ func (s *StreamableHTTPServer) rehydrate(r *http.Request, sid string) (*rehydrat return rt, nil } - srv, err := s.mcp.buildServer(nil, 0) + srv, err := s.mcp.buildServer(nil) if err != nil { return nil, err } @@ -612,11 +640,11 @@ func WithMessageEndpoint(endpoint string) SSEOption { func (s *SSEServer) build() { s.once.Do(func() { - if _, err := s.mcp.buildServer(nil, 0); err != nil { + if _, err := s.mcp.buildServer(nil); err != nil { s.buildErr = err return } - s.handler = gosdk.NewSSEHandler(s.mcp.getServerFunc(nil, 0), nil) + s.handler = gosdk.NewSSEHandler(s.mcp.getServerFunc(nil), nil) }) }