mcp: add a streamable serving benchmark
Add a benchmark using the streamable HTTP handler, and a more realistic input/output signature. Also expose this handler as an example (both a go example and a server in the examples/ directory).
This commit is contained in:
parent
a3e753b8d8
commit
b615fa49b1
3 changed files with 186 additions and 2 deletions
|
|
@ -287,7 +287,7 @@ func BenchmarkMemoryEventStore(b *testing.B) {
|
||||||
payload := make([]byte, test.datasize)
|
payload := make([]byte, test.datasize)
|
||||||
start := time.Now()
|
start := time.Now()
|
||||||
b.ResetTimer()
|
b.ResetTimer()
|
||||||
for i := 0; i < b.N; i++ {
|
for i := range b.N {
|
||||||
sessionID := sessionIDs[i%len(sessionIDs)]
|
sessionID := sessionIDs[i%len(sessionIDs)]
|
||||||
streamID := streamIDs[i%len(sessionIDs)][i%3]
|
streamID := streamIDs[i%len(sessionIDs)][i%3]
|
||||||
store.Append(ctx, sessionID, streamID, payload)
|
store.Append(ctx, sessionID, streamID, payload)
|
||||||
|
|
|
||||||
66
mcp/streamable_bench_test.go
Normal file
66
mcp/streamable_bench_test.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -15,7 +15,7 @@ import (
|
||||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
"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
|
// 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
|
// tool, you need to customize the schema of a single helper type that's used
|
||||||
// in several places.
|
// 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) {
|
func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) {
|
||||||
t1, t2 := mcp.NewInMemoryTransports()
|
t1, t2 := mcp.NewInMemoryTransports()
|
||||||
if _, err := server.Connect(ctx, t1, nil); err != nil {
|
if _, err := server.Connect(ctx, t1, nil); err != nil {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue