mcp: prevent leaks of ServerSession on StatefulServer (#1137)

Fixes #1136

mcp/server.go

@@ -897,9 +897,17 @@ func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams]
 		Capabilities:    req.ClientCapabilities(),
 		ClientInfo:      req.ClientInfo(),
 	}
-	req.Session.updateState(func(state *ServerSessionState) {
-		state.InitializeParams = init
-	})
+	// Only persist InitializeParams when the transport can actually serve
+	// the new protocol. On transports that cannot (notably stateful
+	// StreamableHTTPHandler), a discover request creates a session that
+	// is never surfaced to the client via Mcp-Session-Id; leaving
+	// InitializeParams nil lets serveStatefulPOST's safety-net cleanup
+	// close it instead of leaking.
+	if supportedVersion := negotiateMutuallySupportedVersion(versions); supportedVersion >= protocolVersion20260728 {
+		req.Session.updateState(func(state *ServerSessionState) {
+			state.InitializeParams = init
+		})
+	}
 	res := &DiscoverResult{
 		SupportedVersions: versions,
 		Capabilities:      s.capabilities(),

mcp/streamable_test.go

@@ -3682,6 +3682,70 @@ func TestStreamableStateful_AcceptsDiscover(t *testing.T) {
 	}
 }
 
+// On a stateful StreamableHTTPHandler, a v1.7.0-style client always
+// probes server/discover before falling back to initialize. That probe used
+// to create a ServerSession that was registered in the handler's session
+// map but never surfaced via Mcp-Session-Id, so the client could not DELETE
+// it and it was leaked for the lifetime of the process.
+//
+// The fix in Server.discover skips populating InitializeParams when the
+// transport cannot serve the new protocol, which lets the safety-net
+// cleanup in serveStatefulPOST close the otherwise-unaddressable session.
+func TestStreamableStateful_DiscoverDoesNotLeakSession(t *testing.T) {
+	ctx := context.Background()
+	server := NewServer(testImpl, nil)
+	AddTool(server, &Tool{Name: "noop", Description: "noop"},
+		func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) {
+			return &CallToolResult{}, nil, nil
+		})
+
+	// Stateful (default). SessionTimeout unset so nothing masks the leak
+	// via idle eviction.
+	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
+	httpServer := httptest.NewServer(handler)
+	defer httpServer.Close()
+
+	countSessions := func() int {
+		n := 0
+		for range server.Sessions() {
+			n++
+		}
+		return n
+	}
+
+	const clients = 3
+	for i := 1; i <= clients; i++ {
+		c := NewClient(&Implementation{Name: "c", Version: "1.0.0"}, nil)
+		sess, err := c.Connect(ctx, &StreamableClientTransport{
+			Endpoint:             httpServer.URL,
+			DisableStandaloneSSE: true,
+			MaxRetries:           -1,
+		}, nil)
+		if err != nil {
+			t.Fatalf("client %d Connect: %v", i, err)
+		}
+		if _, err := sess.ListTools(ctx, nil); err != nil {
+			t.Fatalf("client %d ListTools: %v", i, err)
+		}
+		// Exactly one addressable session should exist while this client
+		// is connected. Anything more means a discover-only session
+		// slipped through.
+		if got := countSessions(); got != 1 {
+			t.Errorf("client %d connected: sessions = %d, want 1", i, got)
+		}
+		if err := sess.Close(); err != nil {
+			t.Fatalf("client %d Close: %v", i, err)
+		}
+		if got := countSessions(); got != 0 {
+			t.Errorf("client %d closed: sessions = %d, want 0", i, got)
+		}
+	}
+
+	if got := countSessions(); got != 0 {
+		t.Errorf("final leaked sessions = %d, want 0", got)
+	}
+}
+
 // TestStreamableHTTP_E2E_DiscoverSuccess is a full end-to-end smoke test for
 // SEP-2575 over the streamable HTTP transport.
 func TestStreamableHTTP_E2E_DiscoverSuccess(t *testing.T) {