mcp: fix race condition in ServerSession.startKeepalive

Author: begelundmuller

mcp/server.go

@@ -1028,6 +1028,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
 }
 
@@ -1055,9 +1062,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))
 	}
@@ -1107,7 +1111,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

mcp/streamable_test.go

@@ -367,6 +367,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