mcp: improvements for 'stateless' streamable servers; 'distributed' mode

Several improvements for the stateless streamable mode, plus support for
a 'distributed' (or rather, distributable) version of the stateless
server.

- Add a 'Stateless' option to StreamableHTTPOptions and
  StreamableServerTransport, which controls stateless behavior.
  GetSessionID may still return a non-empty session ID.
- Audit validation of stateless mode to allow requests with a session
  id. Propagate this session ID to the temporary connection.
- Peek at requests to allow 'initialize' requests to go through to the
  session, so that version negotiation can occur (FIXME: add tests).

Fixes #284
For #148

mcp/streamable_test.go

@@ -748,7 +748,7 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream
 		if !request.ignoreResponse {
 			transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() })
 			if diff := cmp.Diff(request.wantMessages, got, transform); diff != "" {
-				t.Errorf("received unexpected messages (-want +got):\n%s", diff)
+				t.Errorf("request #%d: received unexpected messages (-want +got):\n%s", i, diff)
 			}
 		}
 		sessionID.CompareAndSwap("", gotSessionID)
@@ -996,19 +996,18 @@ func TestEventID(t *testing.T) {
 }
 
 func TestStreamableStateless(t *testing.T) {
-	// This version of sayHi doesn't make a ping request (we can't respond to
+	// This version of sayHi expects
 	// that request from our client).
 	sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResult, error) {
+		if err := req.Session.Ping(ctx, nil); err == nil {
+			// ping should fail, but not break the connection
+			t.Errorf("ping succeeded unexpectedly")
+		}
 		return &CallToolResult{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
 	}
 	server := NewServer(testImpl, nil)
 	AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
 
-	// Test stateless mode.
-	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
-		GetSessionID: func() string { return "" },
-	})
-
 	requests := []streamableRequest{
 		{
 			method:         "POST",
@@ -1028,7 +1027,74 @@ func TestStreamableStateless(t *testing.T) {
 			},
 			wantSessionID: false,
 		},
+		{
+			method:         "POST",
+			wantStatusCode: http.StatusOK,
+			messages: []jsonrpc.Message{
+				req(2, "tools/call", &CallToolParams{Name: "greet", Arguments: hiParams{Name: "foo"}}),
+			},
+			wantMessages: []jsonrpc.Message{
+				resp(2, &CallToolResult{Content: []Content{&TextContent{Text: "hi foo"}}}, nil),
+			},
+			wantSessionID: false,
+		},
 	}
 
-	testStreamableHandler(t, handler, requests)
+	testClientCompatibility := func(t *testing.T, handler http.Handler) {
+		ctx := context.Background()
+		httpServer := httptest.NewServer(handler)
+		defer httpServer.Close()
+		cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		res, err := cs.CallTool(ctx, &CallToolParams{Name: "greet", Arguments: hiParams{Name: "bar"}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := textContent(t, res), "hi bar"; got != want {
+			t.Errorf("Result = %q, want %q", got, want)
+		}
+	}
+
+	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
+		GetSessionID: func() string { return "" },
+		Stateless:    true,
+	})
+
+	// Test the default stateless mode.
+	t.Run("stateless", func(t *testing.T) {
+		testStreamableHandler(t, handler, requests)
+		testClientCompatibility(t, handler)
+	})
+
+	// Test a "distributed" variant of stateless mode, where it has non-empty
+	// session IDs, but is otherwise stateless.
+	//
+	// This can be used by tools to look up application state preserved across
+	// subsequent requests.
+	for i, req := range requests {
+		// Now, we want a session for all requests.
+		req.wantSessionID = true
+		requests[i] = req
+	}
+	distributableHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
+		Stateless: true,
+	})
+	t.Run("distributed", func(t *testing.T) {
+		testStreamableHandler(t, distributableHandler, requests)
+		testClientCompatibility(t, handler)
+	})
+}
+
+func textContent(t *testing.T, res *CallToolResult) string {
+	t.Helper()
+	if len(res.Content) != 1 {
+		t.Fatalf("len(Content) = %d, want 1", len(res.Content))
+	}
+	text, ok := res.Content[0].(*TextContent)
+	if !ok {
+		t.Fatalf("Content[0] is %T, want *TextContent", res.Content[0])
+	}
+	return text.Text
 }