modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp: fix race condition in ServerSession.startKeepalive (#856)

Browse source 
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 <mkisiel@google.com>
This commit is contained in:
Benjamin Egelund-Müller 2026-04-03 20:17:49 +02:00 • committed by GitHub
commit 862d78a265
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 40 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

3
mcp/client.go
View file

@ -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

15
mcp/server.go
View file

@ -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

48
mcp/server_test.go
View file

@ -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() {

28
mcp/streamable_test.go
View file

@ -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