From 862d78a265cb1f1b9692ff49cfb8bfd109b6ebc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benjamin=20Egelund-M=C3=BCller?= Date: Fri, 3 Apr 2026 20:17:49 +0200 Subject: [PATCH] mcp: fix race condition in `ServerSession.startKeepalive` (#856) This PR fixes a race condition where `ServerSession.keepaliveCancel` was accessed in `Close` concurrently with its initialization in `ServerSession.startKeepalive`. It also removes a now redundant test. The fix moves keepalive initialization into `ServerSession.Connect` similar to the implementation in `ClientSession.Connect`. This is safe because the MCP spec allows keepalive pings before the initialization messages. From the MCP [spec](https://modelcontextprotocol.io/specification/2025-11-25/basic/lifecycle) (emphasis added): > The client SHOULD NOT send requests **other than pings** before the server has responded to the initialize request. The PR includes a reproducer test case, which fails with the following error without this fix: ``` > go test -count=10 -race -run ^TestStreamableStatelessKeepaliveRace$ github.com/modelcontextprotocol/go-sdk/mcp 8s ================== WARNING: DATA RACE Write at 0x00c00007fbf0 by goroutine 1095: github.com/modelcontextprotocol/go-sdk/mcp.startKeepalive() /Users/benjamin/Documents/go-sdk/mcp/shared.go:590 +0x58 github.com/modelcontextprotocol/go-sdk/mcp.(*ServerSession).startKeepalive() /Users/benjamin/Documents/go-sdk/mcp/server.go:1526 +0x130 github.com/modelcontextprotocol/go-sdk/mcp.(*ServerSession).initialized() /Users/benjamin/Documents/go-sdk/mcp/server.go:1059 +0xf8 github.com/modelcontextprotocol/go-sdk/mcp.init.serverSessionMethod[go.shape.*uint8,go.shape.interface { GetMeta() map[string]interface {}; SetMeta(map[string]interface {}); github.com/modelcontextprotocol/go-sdk/mcp.isResult() }].func27() /Users/benjamin/Documents/go-sdk/mcp/shared.go:332 +0x98 github.com/modelcontextprotocol/go-sdk/mcp.newServerMethodInfo[go.shape.*github.com/modelcontextprotocol/go-sdk/mcp.InitializedParams,go.shape.interface { GetMeta() map[string]interface {}; SetMeta(map[string]interface {}); github.com/modelcontextprotocol/go-sdk/mcp.isResult() },go.shape.struct { github.com/modelcontextprotocol/go-sdk/mcp.Meta "json:\"_meta,omitempty\"" }].func2() /Users/benjamin/Documents/go-sdk/mcp/shared.go:272 +0x84 github.com/modelcontextprotocol/go-sdk/mcp.defaultReceivingMethodHandler[go.shape.*uint8]() /Users/benjamin/Documents/go-sdk/mcp/shared.go:152 +0xe8 github.com/modelcontextprotocol/go-sdk/mcp.defaultReceivingMethodHandler[*github.com/modelcontextprotocol/go-sdk/mcp.ServerSession]() /Users/benjamin/Documents/go-sdk/mcp/shared.go:146 +0x60 github.com/modelcontextprotocol/go-sdk/mcp.handleReceive[go.shape.*uint8]() /Users/benjamin/Documents/go-sdk/mcp/shared.go:169 +0x228 github.com/modelcontextprotocol/go-sdk/mcp.(*ServerSession).handle() /Users/benjamin/Documents/go-sdk/mcp/server.go:1445 +0x370 github.com/modelcontextprotocol/go-sdk/mcp.connect[go.shape.*uint8,go.shape.*uint8].func1.1() /Users/benjamin/Documents/go-sdk/mcp/transport.go:170 +0x5c github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2.HandlerFunc.Handle() /Users/benjamin/Documents/go-sdk/internal/jsonrpc2/jsonrpc2.go:84 +0x48 github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2.(*Connection).handleAsync.func3() /Users/benjamin/Documents/go-sdk/internal/jsonrpc2/conn.go:717 +0xd8 Previous read at 0x00c00007fbf0 by goroutine 1089: github.com/modelcontextprotocol/go-sdk/mcp.(*ServerSession).Close() /Users/benjamin/Documents/go-sdk/mcp/server.go:1502 +0x34 github.com/modelcontextprotocol/go-sdk/mcp.(*StreamableHTTPHandler).ServeHTTP.deferwrap1() /Users/benjamin/Documents/go-sdk/mcp/streamable.go:519 +0x34 runtime.deferreturn() /opt/homebrew/Cellar/go/1.25.4/libexec/src/runtime/panic.go:589 +0x5c github.com/modelcontextprotocol/go-sdk/mcp.TestStreamableStatelessKeepaliveRace.mustNotPanic.func2() /Users/benjamin/Documents/go-sdk/mcp/streamable_test.go:2156 +0x94 net/http.HandlerFunc.ServeHTTP() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/server.go:2322 +0x48 net/http.serverHandler.ServeHTTP() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/server.go:3340 +0x270 net/http.(*conn).serve() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/server.go:2109 +0x9b0 net/http.(*Server).Serve.gowrap3() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/server.go:3493 +0x48 Goroutine 1095 (running) created at: github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2.(*Connection).handleAsync() /Users/benjamin/Documents/go-sdk/internal/jsonrpc2/conn.go:715 +0x334 github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2.(*Connection).acceptRequest.func2.gowrap1() /Users/benjamin/Documents/go-sdk/internal/jsonrpc2/conn.go:676 +0x34 Goroutine 1089 (running) created at: net/http.(*Server).Serve() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/server.go:3493 +0x5dc net/http/httptest.(*Server).goServe.func1() /opt/homebrew/Cellar/go/1.25.4/libexec/src/net/http/httptest/server.go:311 +0x98 ================== --- FAIL: TestStreamableStatelessKeepaliveRace (0.14s) testing.go:1617: race detected during execution of test FAIL FAIL github.com/modelcontextprotocol/go-sdk/mcp 1.727s FAIL ``` --------- Co-authored-by: Maciej Kisiel --- mcp/client.go | 3 ++- mcp/server.go | 15 ++++++++----- mcp/server_test.go | 48 ------------------------------------------ mcp/streamable_test.go | 28 ++++++++++++++++++++++++ 4 files changed, 40 insertions(+), 54 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index dc7bef1..4860b62 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -340,7 +340,8 @@ func (cs *ClientSession) ID() string { // Close is idempotent and concurrency safe. func (cs *ClientSession) Close() error { // Note: keepaliveCancel access is safe without a mutex because: - // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 1. keepaliveCancel is only written once during Client.Connect (through startKeepalive), + // which happens before any code that may call Close from another goroutine // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines // 3. The keepalive goroutine calls Close on ping failure, but this is safe since // Close is idempotent and conn.Close() handles concurrent calls correctly diff --git a/mcp/server.go b/mcp/server.go index 97b5d44..2357959 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1031,6 +1031,13 @@ func (s *Server) Connect(ctx context.Context, t Transport, opts *ServerSessionOp s.opts.Logger.Error("server connect error", "error", err) return nil, err } + + // Start keepalive before returning the session to avoid race conditions with Close. + // This is safe because the spec allows sending pings before initialization (see ServerSession.handle for details). + if s.opts.KeepAlive > 0 { + ss.startKeepalive(ss.server.opts.KeepAlive) + } + return ss, nil } @@ -1058,9 +1065,6 @@ func (ss *ServerSession) initialized(ctx context.Context, params *InitializedPar ss.server.opts.Logger.Error("duplicate initialized notification") return nil, fmt.Errorf("duplicate %q received", notificationInitialized) } - if ss.server.opts.KeepAlive > 0 { - ss.startKeepalive(ss.server.opts.KeepAlive) - } if h := ss.server.opts.InitializedHandler; h != nil { h(ctx, serverRequestFor(ss, params)) } @@ -1110,7 +1114,7 @@ type ServerSession struct { server *Server conn *jsonrpc2.Connection mcpConn Connection - keepaliveCancel context.CancelFunc // TODO: theory around why keepaliveCancel need not be guarded + keepaliveCancel context.CancelFunc mu sync.Mutex state ServerSessionState @@ -1504,7 +1508,8 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLoggingLevelPara func (ss *ServerSession) Close() error { if ss.keepaliveCancel != nil { // Note: keepaliveCancel access is safe without a mutex because: - // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 1. keepaliveCancel is only written once during Server.Connect (through startKeepalive), + // which happens before any code that may call Close from another goroutine // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines // 3. The keepalive goroutine calls Close on ping failure, but this is safe since // Close is idempotent and conn.Close() handles concurrent calls correctly diff --git a/mcp/server_test.go b/mcp/server_test.go index 85c3990..1312e1d 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -508,54 +508,6 @@ func TestServerAddResourceTemplate(t *testing.T) { } } -// TestServerSessionkeepaliveCancelOverwritten is to verify that `ServerSession.keepaliveCancel` is assigned exactly once, -// ensuring that only a single goroutine is responsible for the session's keepalive ping mechanism. -func TestServerSessionkeepaliveCancelOverwritten(t *testing.T) { - // Set KeepAlive to a long duration to ensure the keepalive - // goroutine stays alive for the duration of the test without actually sending - // ping requests, since we don't have a real client connection established. - server := NewServer(testImpl, &ServerOptions{KeepAlive: 5 * time.Second}) - ss := &ServerSession{server: server} - - // 1. Initialize the session. - _, err := ss.initialize(context.Background(), &InitializeParams{}) - if err != nil { - t.Fatalf("ServerSession initialize failed: %v", err) - } - - // 2. Call 'initialized' for the first time. This should start the keepalive mechanism. - _, err = ss.initialized(context.Background(), &InitializedParams{}) - if err != nil { - t.Fatalf("First initialized call failed: %v", err) - } - if ss.keepaliveCancel == nil { - t.Fatalf("expected ServerSession.keepaliveCancel to be set after the first call of initialized") - } - - // Save the cancel function and use defer to ensure resources are cleaned up. - firstCancel := ss.keepaliveCancel - defer firstCancel() - - // 3. Manually set the field to nil. - // Do this to facilitate the test's core assertion. The goal is to verify that - // 'ss.keepaliveCancel' is not assigned a second time. By setting it to nil, - // we can easily check after the next call if a new keepalive goroutine was started. - ss.keepaliveCancel = nil - - // 4. Call 'initialized' for the second time. This should return an error. - _, err = ss.initialized(context.Background(), &InitializedParams{}) - if err == nil { - t.Fatalf("Expected 'duplicate initialized received' error on second call, got nil") - } - - // 5. Re-check the field to ensure it remains nil. - // Since 'initialized' correctly returned an error and did not call - // 'startKeepalive', the field should remain unchanged. - if ss.keepaliveCancel != nil { - t.Fatal("expected ServerSession.keepaliveCancel to be nil after we manually niled it and re-initialized") - } -} - // panicks reports whether f() panics. func panics(f func()) (b bool) { defer func() { diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index a275d3c..3600277 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -363,6 +363,34 @@ func TestStreamableServerShutdown(t *testing.T) { } } +// TestStreamableStatelessKeepaliveRace verifies that there is no data race between +// ServerSession.startKeepalive and ServerSession.Close in stateless servers. +func TestStreamableStatelessKeepaliveRace(t *testing.T) { + ctx := context.Background() + server := NewServer(testImpl, &ServerOptions{KeepAlive: time.Hour}) + AddTool(server, &Tool{Name: "greet"}, sayHi) + handler := NewStreamableHTTPHandler( + func(*http.Request) *Server { return server }, + &StreamableHTTPOptions{Stateless: true}, + ) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + for range 50 { + cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{ + Endpoint: httpServer.URL, + }, nil) + if err != nil { + t.Fatalf("NewClient() failed: %v", err) + } + _, _ = cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "world"}, + }) + _ = cs.Close() + } +} + // TestClientReplay verifies that the client can recover from a mid-stream // network failure and receive replayed messages (if replay is configured). It // uses a proxy that is killed and restarted to simulate a recoverable network