mcp: don't persist failed streamable sessions (#605)

If a stateful streamable session fails to initialize, it is unusable.
Avoid allocating resources by closing the session immediately.

Also:
- Fix a bug where InitializeParams is not guarded with its mutex.
- Fix a resource leak where sessions are persisted with session id "",
even though they are unaddressable.
- Add a relevant benchmark, and update tests.

Fixes #578

mcp/streamable_bench_test.go

@@ -6,9 +6,15 @@ package mcp_test
 
 import (
 	"context"
+	"flag"
+	"log"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"reflect"
+	"runtime"
+	"runtime/pprof"
+	"strings"
 	"testing"
 
 	"github.com/google/jsonschema-go/jsonschema"
@@ -65,3 +71,62 @@ func BenchmarkStreamableServing(b *testing.B) {
 		}
 	}
 }
+
+var streamableHeap = flag.String("streamable_heap", "", "if set, write streamable heap profiles with this prefix")
+
+func BenchmarkStreamableServing_BadSessions(b *testing.B) {
+	server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
+
+	handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
+		return server
+	}, &mcp.StreamableHTTPOptions{JSONResponse: true})
+
+	httpServer := httptest.NewServer(handler)
+	defer httpServer.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	if *streamableHeap != "" {
+		writeHeap := func(file string) {
+			// GC a couple times to ensure accurate heap.
+			runtime.GC()
+			runtime.GC()
+			f, err := os.Create(file)
+			if err != nil {
+				log.Fatal("could not create memory profile: ", err)
+			}
+			defer func() {
+				if err := f.Close(); err != nil {
+					b.Errorf("writing heap file %q: %v", file, err)
+				}
+			}()
+			if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil {
+				b.Errorf("could not write heap profile: %v", err)
+			}
+		}
+		writeHeap(*streamableHeap + ".before")
+		defer writeHeap(*streamableHeap + ".after")
+	}
+
+	b.ResetTimer()
+	for range b.N {
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, strings.NewReader("{}"))
+		if err != nil {
+			b.Fatal(err)
+		}
+		req.Header.Add("Accept", "application/json")
+		req.Header.Add("Accept", "text/event-stream")
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			b.Fatal(err)
+		}
+		if got, want := resp.StatusCode, http.StatusBadRequest; got != want {
+			b.Fatalf("POST got status %d, want %d", got, want)
+		}
+		if got := resp.Header.Get("Mcp-Session-Id"); got != "" {
+			b.Fatalf("POST got unexpected session ID")
+		}
+		resp.Body.Close()
+	}
+}