mcp: add a streamable serving benchmark

mcp/event_test.go

@@ -287,7 +287,7 @@ func BenchmarkMemoryEventStore(b *testing.B) {
 			payload := make([]byte, test.datasize)
 			start := time.Now()
 			b.ResetTimer()
-			for i := 0; i < b.N; i++ {
+			for i := range b.N {
 				sessionID := sessionIDs[i%len(sessionIDs)]
 				streamID := streamIDs[i%len(sessionIDs)][i%3]
 				store.Append(ctx, sessionID, streamID, payload)

mcp/streamable_bench_test.go

@@ -0,0 +1,66 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package mcp_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+func BenchmarkStreamableServing(b *testing.B) {
+	// This benchmark measures how fast we can handle a single tool on a
+	// streamable server, including tool validation and stream management.
+	customSchemas := map[any]*jsonschema.Schema{
+		Probability(0):  {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
+		WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
+	}
+	opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
+	in, err := jsonschema.For[WeatherInput](opts)
+	if err != nil {
+		b.Fatal(err)
+	}
+	out, err := jsonschema.For[WeatherOutput](opts)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
+	mcp.AddTool(server, &mcp.Tool{
+		Name:         "weather",
+		InputSchema:  in,
+		OutputSchema: out,
+	}, WeatherTool)
+
+	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()
+	session, err := mcp.NewClient(testImpl, nil).Connect(ctx, &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer session.Close()
+	b.ResetTimer()
+	for range b.N {
+		if _, err := session.CallTool(ctx, &mcp.CallToolParams{
+			Name: "weather",
+			Arguments: WeatherInput{
+				Location: Location{Name: "somewhere"},
+				Days:     7,
+			},
+		}); err != nil {
+			b.Errorf("CallTool failed: %v", err)
+		}
+	}
+}

mcp/tool_example_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-func ExampleAddTool_customTypeSchema() {
+func ExampleAddTool_customMarshalling() {
 	// Sometimes when you want to customize the input or output schema for a
 	// tool, you need to customize the schema of a single helper type that's used
 	// in several places.
@@ -83,6 +83,124 @@ func ExampleAddTool_customTypeSchema() {
 	// }
 }
 
+type WeatherInput struct {
+	Location Location `json:"location" jsonschema:"user location"`
+	Days     int      `json:"days" jsonschema:"number of days to forecast"`
+}
+
+type Location struct {
+	Name      string   `json:"name"`
+	Latitude  *float64 `json:"latitude,omitempty"`
+	Longitude *float64 `json:"longitude,omitempty"`
+}
+
+type Forecast struct {
+	Forecast string      `json:"forecast" jsonschema:"description of the day's weather"`
+	Type     WeatherType `json:"type" jsonschema:"type of weather"`
+	Rain     float64     `json:"rain" jsonschema:"probability of rain, between 0 and 1"`
+	High     float64     `json:"high" jsonschema:"high temperature"`
+	Low      float64     `json:"low" jsonschema:"low temperature"`
+}
+
+type WeatherType string
+
+const (
+	Sunny        WeatherType = "sun"
+	PartlyCloudy WeatherType = "partly_cloudy"
+	Cloudy       WeatherType = "clouds"
+	Rainy        WeatherType = "rain"
+	Snowy        WeatherType = "snow"
+)
+
+type Probability float64
+
+type WeatherOutput struct {
+	Summary       string      `json:"summary" jsonschema:"a summary of the weather forecast"`
+	Confidence    Probability `json:"confidence" jsonschema:"confidence, between 0 and 1"`
+	AsOf          time.Time   `json:"asOf" jsonschema:"the time the weather was computed"`
+	DailyForecast []Forecast  `json:"dailyForecast" jsonschema:"the daily forecast"`
+	Source        string      `json:"source,omitempty" jsonschema:"the organization providing the weather forecast"`
+}
+
+func WeatherTool(ctx context.Context, req *mcp.CallToolRequest, in WeatherInput) (*mcp.CallToolResult, WeatherOutput, error) {
+	perfectWeather := WeatherOutput{
+		Summary:    "perfect",
+		Confidence: 1.0,
+		AsOf:       time.Now(),
+	}
+	for range in.Days {
+		perfectWeather.DailyForecast = append(perfectWeather.DailyForecast, Forecast{
+			Forecast: "another perfect day",
+			Type:     Sunny,
+			Rain:     0.0,
+			High:     72.0,
+			Low:      72.0,
+		})
+	}
+	return nil, perfectWeather, nil
+}
+
+func ExampleAddTool_complexSchema() {
+	// This example demonstrates a tool with a more 'realistic' input and output
+	// schema. We use a combination of techniques to tune our input and output
+	// schemas.
+
+	// Distinguished Go types allow custom schemas to be reused during inference.
+	customSchemas := map[any]*jsonschema.Schema{
+		Probability(0):  {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
+		WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
+	}
+	opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
+	in, err := jsonschema.For[WeatherInput](opts)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Furthermore, we can tweak the inferred schema, in this case limiting
+	// forecasts to 0-10 days.
+	daysSchema := in.Properties["days"]
+	daysSchema.Minimum = jsonschema.Ptr(0.0)
+	daysSchema.Maximum = jsonschema.Ptr(10.0)
+
+	// Output schema inference can reuse our custom schemas from input inference.
+	out, err := jsonschema.For[WeatherOutput](opts)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Now add our tool to a server. Since we've customized the schemas, we need
+	// to override the default schema inference.
+	server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
+	mcp.AddTool(server, &mcp.Tool{
+		Name:         "weather",
+		InputSchema:  in,
+		OutputSchema: out,
+	}, WeatherTool)
+
+	ctx := context.Background()
+	session, err := connect(ctx, server) // create an in-memory connection
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer session.Close()
+
+	// Check that the client observes the correct schemas.
+	for t, err := range session.Tools(ctx, nil) {
+		if err != nil {
+			log.Fatal(err)
+		}
+		// Formatting the entire schemas would be too much output.
+		// Just check that our customizations were effective.
+		fmt.Println("max days:", *t.InputSchema.Properties["days"].Maximum)
+		fmt.Println("max confidence:", *t.OutputSchema.Properties["confidence"].Maximum)
+		fmt.Println("weather types:", t.OutputSchema.Properties["dailyForecast"].Items.Properties["type"].Enum)
+	}
+	// Output:
+	// max days: 10
+	// max confidence: 1
+	// weather types: [sun partly_cloudy clouds rain snow]
+}
+
 func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) {
 	t1, t2 := mcp.NewInMemoryTransports()
 	if _, err := server.Connect(ctx, t1, nil); err != nil {