Skip to content

Commit 3a36eb8

Browse files
findleyrfindleyr
committed
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). For #190
1 parent 22f86c4 commit 3a36eb8

File tree

4 files changed

+261
-1
lines changed

4 files changed

+261
-1
lines changed

‎examples/server/weather/main.go‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package main
6+
7+
import (
8+
"context"
9+
"log"
10+
"time"
11+
12+
"github.com/google/jsonschema-go/jsonschema"
13+
"github.com/modelcontextprotocol/go-sdk/mcp"
14+
)
15+
16+
type WeatherInput struct {
17+
Location Location `json:"location" jsonschema:"user location"`
18+
Days int `json:"days" jsonschema:"number of days to forecast"`
19+
}
20+
21+
type Location struct {
22+
Name string `json:"name"`
23+
Latitude *float64 `json:"latitude,omitempty"`
24+
Longitude *float64 `json:"longitude,omitempty"`
25+
}
26+
27+
type Forecast struct {
28+
Forecast string `json:"forecast" jsonschema:"description of the day's weather"`
29+
Type WeatherType `json:"type" jsonschema:"type of weather"`
30+
Rain float64 `json:"rain" jsonschema:"probability of rain, between 0 and 1"`
31+
High float64 `json:"high" jsonschema:"high temperature"`
32+
Low float64 `json:"low" jsonschema:"low temperature"`
33+
}
34+
35+
type WeatherType string
36+
37+
const (
38+
Sunny WeatherType = "sun"
39+
PartlyCloudy WeatherType = "partly_cloudy"
40+
Cloudy WeatherType = "clouds"
41+
Rainy WeatherType = "rain"
42+
Snowy WeatherType = "snow"
43+
)
44+
45+
type Probability float64
46+
47+
type WeatherOutput struct {
48+
Summary string `json:"summary" jsonschema:"a summary of the weather forecast"`
49+
Confidence Probability `json:"confidence" jsonschema:"confidence, between 0 and 1"`
50+
AsOf time.Time `json:"asOf" jsonschema:"the time the weather was computed"`
51+
DailyForecast []Forecast `json:"dailyForecast" jsonschema:"the daily forecast"`
52+
Source string `json:"source,omitempty" jsonschema:"the organization providing the weather forecast"`
53+
}
54+
55+
func WeatherTool(ctx context.Context, req *mcp.CallToolRequest, in WeatherInput) (*mcp.CallToolResult, WeatherOutput, error) {
56+
perfectWeather := WeatherOutput{
57+
Summary: "perfect",
58+
Confidence: 1.0,
59+
AsOf: time.Now(),
60+
}
61+
for range in.Days {
62+
perfectWeather.DailyForecast = append(perfectWeather.DailyForecast, Forecast{
63+
Forecast: "another perfect day",
64+
Type: Sunny,
65+
Rain: 0.0,
66+
High: 72.0,
67+
Low: 72.0,
68+
})
69+
}
70+
return nil, perfectWeather, nil
71+
}
72+
73+
func main() {
74+
customSchemas := map[any]*jsonschema.Schema{
75+
Probability(0): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
76+
WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
77+
}
78+
opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
79+
in, err := jsonschema.For[WeatherInput](opts)
80+
if err != nil {
81+
log.Fatal(err)
82+
}
83+
out, err := jsonschema.For[WeatherOutput](opts)
84+
if err != nil {
85+
log.Fatal(err)
86+
}
87+
88+
server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
89+
mcp.AddTool(server, &mcp.Tool{
90+
Name: "weather",
91+
InputSchema: in,
92+
OutputSchema: out,
93+
}, WeatherTool)
94+
95+
server.Run(context.Background(), &mcp.StdioTransport{})
96+
}

‎mcp/event_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ func BenchmarkMemoryEventStore(b *testing.B) {
287287
payload := make([]byte, test.datasize)
288288
start := time.Now()
289289
b.ResetTimer()
290-
for i := 0; i < b.N; i++ {
290+
for i := range b.N {
291291
sessionID := sessionIDs[i%len(sessionIDs)]
292292
streamID := streamIDs[i%len(sessionIDs)][i%3]
293293
store.Append(ctx, sessionID, streamID, payload)

‎mcp/streamable_bench_test.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp_test
6+
7+
import (
8+
"context"
9+
"net/http"
10+
"net/http/httptest"
11+
"testing"
12+
13+
"github.com/google/jsonschema-go/jsonschema"
14+
"github.com/modelcontextprotocol/go-sdk/mcp"
15+
)
16+
17+
func BenchmarkStreamableServing(b *testing.B) {
18+
customSchemas := map[any]*jsonschema.Schema{
19+
Probability(0): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
20+
WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
21+
}
22+
opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
23+
in, err := jsonschema.For[WeatherInput](opts)
24+
if err != nil {
25+
b.Fatal(err)
26+
}
27+
out, err := jsonschema.For[WeatherOutput](opts)
28+
if err != nil {
29+
b.Fatal(err)
30+
}
31+
32+
server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
33+
mcp.AddTool(server, &mcp.Tool{
34+
Name: "weather",
35+
InputSchema: in,
36+
OutputSchema: out,
37+
}, WeatherTool)
38+
39+
handler := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server {
40+
return server
41+
}, &mcp.StreamableHTTPOptions{JSONResponse: true})
42+
httpServer := httptest.NewServer(handler)
43+
defer httpServer.Close()
44+
45+
ctx, cancel := context.WithCancel(context.Background())
46+
defer cancel()
47+
session, err := mcp.NewClient(testImpl, nil).Connect(ctx, &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil)
48+
if err != nil {
49+
b.Fatal(err)
50+
}
51+
defer session.Close()
52+
b.ResetTimer()
53+
for range b.N {
54+
if _, err := session.CallTool(ctx, &mcp.CallToolParams{
55+
Name: "weather",
56+
Arguments: WeatherInput{
57+
Location: Location{Name: "somewhere"},
58+
Days: 7,
59+
},
60+
}); err != nil {
61+
b.Errorf("CallTool failed: %v", err)
62+
}
63+
}
64+
}

‎mcp/tool_example_test.go‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,106 @@ func ExampleAddTool_customTypeSchema() {
8383
// }
8484
}
8585

86+
type WeatherInput struct {
87+
Location Location `json:"location" jsonschema:"user location"`
88+
Days int `json:"days" jsonschema:"number of days to forecast"`
89+
}
90+
91+
type Location struct {
92+
Name string `json:"name"`
93+
Latitude *float64 `json:"latitude,omitempty"`
94+
Longitude *float64 `json:"longitude,omitempty"`
95+
}
96+
97+
type Forecast struct {
98+
Forecast string `json:"forecast" jsonschema:"description of the day's weather"`
99+
Type WeatherType `json:"type" jsonschema:"type of weather"`
100+
Rain float64 `json:"rain" jsonschema:"probability of rain, between 0 and 1"`
101+
High float64 `json:"high" jsonschema:"high temperature"`
102+
Low float64 `json:"low" jsonschema:"low temperature"`
103+
}
104+
105+
type WeatherType string
106+
107+
const (
108+
Sunny WeatherType = "sun"
109+
PartlyCloudy WeatherType = "partly_cloudy"
110+
Cloudy WeatherType = "clouds"
111+
Rainy WeatherType = "rain"
112+
Snowy WeatherType = "snow"
113+
)
114+
115+
type Probability float64
116+
117+
type WeatherOutput struct {
118+
Summary string `json:"summary" jsonschema:"a summary of the weather forecast"`
119+
Confidence Probability `json:"confidence" jsonschema:"confidence, between 0 and 1"`
120+
AsOf time.Time `json:"asOf" jsonschema:"the time the weather was computed"`
121+
DailyForecast []Forecast `json:"dailyForecast" jsonschema:"the daily forecast"`
122+
Source string `json:"source,omitempty" jsonschema:"the organization providing the weather forecast"`
123+
}
124+
125+
func WeatherTool(ctx context.Context, req *mcp.CallToolRequest, in WeatherInput) (*mcp.CallToolResult, WeatherOutput, error) {
126+
perfectWeather := WeatherOutput{
127+
Summary: "perfect",
128+
Confidence: 1.0,
129+
AsOf: time.Now(),
130+
}
131+
for range in.Days {
132+
perfectWeather.DailyForecast = append(perfectWeather.DailyForecast, Forecast{
133+
Forecast: "another perfect day",
134+
Type: Sunny,
135+
Rain: 0.0,
136+
High: 72.0,
137+
Low: 72.0,
138+
})
139+
}
140+
return nil, perfectWeather, nil
141+
}
142+
143+
func ExampleAddTool_complexSchema() {
144+
customSchemas := map[any]*jsonschema.Schema{
145+
Probability(0): {Type: "number", Minimum: jsonschema.Ptr(0.0), Maximum: jsonschema.Ptr(1.0)},
146+
WeatherType(""): {Type: "string", Enum: []any{Sunny, PartlyCloudy, Cloudy, Rainy, Snowy}},
147+
}
148+
opts := &jsonschema.ForOptions{TypeSchemas: customSchemas}
149+
in, err := jsonschema.For[WeatherInput](opts)
150+
if err != nil {
151+
log.Fatal(err)
152+
}
153+
out, err := jsonschema.For[WeatherOutput](opts)
154+
if err != nil {
155+
log.Fatal(err)
156+
}
157+
158+
server := mcp.NewServer(&mcp.Implementation{Name: "server", Version: "v0.0.1"}, nil)
159+
mcp.AddTool(server, &mcp.Tool{
160+
Name: "weather",
161+
InputSchema: in,
162+
OutputSchema: out,
163+
}, WeatherTool)
164+
165+
ctx := context.Background()
166+
session, err := connect(ctx, server) // create an in-memory connection
167+
if err != nil {
168+
log.Fatal(err)
169+
}
170+
defer session.Close()
171+
172+
for t, err := range session.Tools(ctx, nil) {
173+
if err != nil {
174+
log.Fatal(err)
175+
}
176+
schemaJSON, err := json.MarshalIndent(t.OutputSchema, "", "\t")
177+
if err != nil {
178+
log.Fatal(err)
179+
}
180+
fmt.Println(t.Name, string(schemaJSON))
181+
}
182+
// Output:
183+
// nothing
184+
}
185+
86186
func connect(ctx context.Context, server *mcp.Server) (*mcp.ClientSession, error) {
87187
t1, t2 := mcp.NewInMemoryTransports()
88188
if _, err := server.Connect(ctx, t1, nil); err != nil {

0 commit comments

Comments
 (0)