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 <mkisiel@google.com>
This commit is contained in:
parent
34a19dff46
commit
862d78a265
1 changed files with 40 additions and 54 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue