perf: add schema caching to avoid repeated reflection (#685)

## Summary This PR adds schema caching that dramatically reduces the cost of tool registration in stateless server patterns. **Contributors:** Thanks to @omgitsads and @almaleksia for collaborating on and researching this solution! ## Problem In stateless deployments like `github/github-mcp-server`, a new `*mcp.Server` is created for each incoming request. This means `AddTool` is called repeatedly for the same tools, causing: 1. **For typed handlers**: `jsonschema.ForType()` reflection is called every time 2. **For pre-defined schemas**: `schema.Resolve()` is called every time ## Solution Add a `schemaCache` that stores: - Schemas by `reflect.Type` (for auto-generated schemas from typed handlers) - Resolved schemas by schema pointer (for pre-defined schemas) The cache is: - Concurrent-safe using `sync.Map` - Unbounded (typical MCP servers have <100 tools) - Global (lives across server instances) ## Benchmark Results ``` BenchmarkAddToolTypedHandler-22 977080 1223 ns/op 1208 B/op 21 allocs/op BenchmarkAddToolTypedHandlerNoCache-22 6764 161463 ns/op 39262 B/op 1072 allocs/op ``` | Metric | With Cache | Without Cache | Improvement | |--------|------------|---------------|-------------| | Time | 1,223 ns/op | 161,463 ns/op | **132x faster** | | Allocations | 21 allocs | 1,072 allocs | **51x fewer** | | Memory | 1,208 B/op | 39,262 B/op | **32x less** | ## Files Changed - `mcp/schema_cache.go` - New cache implementation - `mcp/server.go` - Modified `setSchema` to use cache - `mcp/schema_cache_test.go` - Unit tests for caching behavior - `mcp/schema_cache_benchmark_test.go` - Benchmarks ## Impact for Integrators **Automatic** - no code changes required. Integrators using: - Typed handlers (`AddTool[In, Out]`) → cache by type - Pre-defined schemas (`Tool{InputSchema: schema}`) → cache by pointer Both patterns benefit from caching after the first call. --- ## Real-World Performance Validation The following benchmarks were conducted using `github/github-mcp-server` (a production MCP server with ~130 tools) to validate the performance impact in a real-world scenario. ### Test Environment - **Server:** github-mcp-server (stateless HTTP deployment) - **Tools registered per request:** ~130 tools - **Test methodology:** Python benchmark client, 30 iterations (latency), 100 iterations (stress) - **Date:** December 4, 2025 ### Configurations Tested | Configuration | Description | |--------------|-------------| | **main (mcp-go)** | Original implementation using mcp-go library | | **go-sdk (no cache)** | go-sdk WITHOUT schema caching (broken) | | **go-sdk (with cache)** | go-sdk WITH this PR's schema caching fix | ### Latency Test Results (n=30) #### Operation: `initialize` | Configuration | P50 | P99 | Status | |--------------|-----|-----|--------| | main (mcp-go) | 11.48ms | 14.00ms | ✅ Baseline | | go-sdk (no cache) | 20.47ms | 25.30ms | 🔴 **+78% REGRESSION** | | go-sdk (with cache) | 10.36ms | 14.33ms | ✅ **FIXED (-10%)** | #### Operation: `tools/list` | Configuration | P50 | P99 | Status | |--------------|-----|-----|--------| | main (mcp-go) | 13.46ms | 22.47ms | ✅ Baseline | | go-sdk (no cache) | 22.98ms | 29.11ms | 🔴 **+71% REGRESSION** | | go-sdk (with cache) | 14.15ms | 15.91ms | ✅ **FIXED (+5%)** | #### Operation: `prompts/list` | Configuration | P50 | P99 | Status | |--------------|-----|-----|--------| | main (mcp-go) | 11.56ms | 13.81ms | ✅ Baseline | | go-sdk (no cache) | 21.00ms | 26.42ms | 🔴 **+82% REGRESSION** | | go-sdk (with cache) | 10.77ms | 15.25ms | ✅ **FIXED (-7%)** | ### Stress Test Results (n=100) | Configuration | P50 (initialize) | P50 (tools/list) | P50 (prompts/list) | |--------------|------------------|------------------|-------------------| | main (mcp-go) | 12.06ms | 14.30ms | 11.34ms | | go-sdk (no cache) | 20.06ms | 23.44ms | 19.59ms | | go-sdk (with cache) | **11.83ms** | **14.58ms** | **11.05ms** | ### Memory/Allocation Comparison (from pprof) | Configuration | Total Allocations | Comparison | |--------------|-------------------|------------| | main (mcp-go) | 355.91 MB | Baseline | | go-sdk (no cache) | **1208.70 MB** | 🔴 **3.4x MORE allocations** | #### Top Allocation Sources - go-sdk WITHOUT cache (broken) | Function | Size | % | Issue | |----------|------|---|-------| | `jsonschema.UnmarshalJSON` | 324.60 MB | 27% | 🚨 Schema re-parsing | | `encoding/json.Unmarshal` | 341.10 MB | 28% | JSON deserialization | | `jsonschema.resolve` | 219.51 MB | 18% | 🚨 Schema re-resolution | | `jsonschema.MarshalJSON` | 92.52 MB | 8% | Schema JSON encoding | ### Root Cause Analysis The `google/jsonschema-go` library regenerates JSON schemas on every request instead of caching them. In a server with ~130 tools, this causes: - **70-80% latency regression** on all MCP operations - **3.4x more memory allocations** per request - **~70% of all allocations** are schema-related operations ### Key Findings 1. ✅ **REGRESSION CONFIRMED:** go-sdk without schema caching is 70-80% slower than mcp-go 2. ✅ **FIX VERIFIED:** Schema caching restores performance to baseline (or better) 3. ✅ **MEMORY IMPACT:** 3.4x reduction in allocations with caching 4. ✅ **PRODUCTION READY:** Fixed version performs equivalently to mcp-go baseline --------- Co-authored-by: Adam Holt <4619+omgitsads@users.noreply.github.com> Co-authored-by: Ksenia Bobrova <1885174+almaleksia@users.noreply.github.com>
2026-01-21 02:16:46 +01:00 · 2026-01-21 02:16:46 +01:00 · c2c7edc3f0
commit c2c7edc3f0
parent 3381035a50
6 changed files with 379 additions and 17 deletions
--- a/docs/server.md
+++ b/docs/server.md
@ -403,6 +403,23 @@ _See [mcp/tool_example_test.go](../mcp/tool_example_test.go) for the full
 example, or [examples/server/toolschemas](examples/server/toolschemas/main.go)
 for more examples of customizing tool schemas._

+**Stateless server deployments:** Some deployments create a new
+[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server)
+for each incoming request, re-registering tools every time. To avoid repeated
+schema generation, create a
+[`SchemaCache`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SchemaCache)
+and share it across server instances:
+
+```go
+var schemaCache = mcp.NewSchemaCache() // create once at startup
+
+func handleRequest(w http.ResponseWriter, r *http.Request) {
+    s := mcp.NewServer(impl, &mcp.ServerOptions{SchemaCache: schemaCache})
+    mcp.AddTool(s, myTool, myHandler)
+    // ...
+}
+```
+
 ## Utilities

 ### Completion
--- a/internal/docs/server.src.md
+++ b/internal/docs/server.src.md
@ -193,6 +193,23 @@ _See [mcp/tool_example_test.go](../mcp/tool_example_test.go) for the full
 example, or [examples/server/toolschemas](examples/server/toolschemas/main.go)
 for more examples of customizing tool schemas._

+**Stateless server deployments:** Some deployments create a new
+[`Server`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#Server)
+for each incoming request, re-registering tools every time. To avoid repeated
+schema generation, create a
+[`SchemaCache`](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk/mcp#SchemaCache)
+and share it across server instances:
+
+```go
+var schemaCache = mcp.NewSchemaCache() // create once at startup
+
+func handleRequest(w http.ResponseWriter, r *http.Request) {
+    s := mcp.NewServer(impl, &mcp.ServerOptions{SchemaCache: schemaCache})
+    mcp.AddTool(s, myTool, myHandler)
+    // ...
+}
+```
+
 ## Utilities

 ### Completion
--- a/mcp/schema_cache.go
+++ b/mcp/schema_cache.go
@ -0,0 +1,69 @@
+// 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
+
+import (
+	"reflect"
+	"sync"
+
+	"github.com/google/jsonschema-go/jsonschema"
+)
+
+// A SchemaCache caches JSON schemas to avoid repeated reflection and resolution.
+//
+// This is useful for stateless server deployments (one [Server] per request)
+// where tools are re-registered on every request. Without caching, each
+// [AddTool] call triggers expensive reflection-based schema generation.
+//
+// A SchemaCache is safe for concurrent use by multiple goroutines.
+//
+// # Trade-offs
+//
+// The cache is unbounded: it stores one entry per unique Go type or schema
+// pointer. For typical MCP servers with a fixed set of tools, memory usage
+// is negligible. However, if tool input types are generated dynamically,
+// the cache will grow without bound.
+//
+// The cache uses pointer identity for pre-defined schemas. If a schema's
+// contents change but the pointer remains the same, stale resolved schemas
+// may be returned. In practice, this is not an issue because tool schemas
+// are typically defined once at startup.
+type SchemaCache struct {
+	byType   sync.Map // reflect.Type -> *cachedSchema
+	bySchema sync.Map // *jsonschema.Schema -> *jsonschema.Resolved
+}
+
+type cachedSchema struct {
+	schema   *jsonschema.Schema
+	resolved *jsonschema.Resolved
+}
+
+// NewSchemaCache creates a new [SchemaCache].
+func NewSchemaCache() *SchemaCache {
+	return &SchemaCache{}
+}
+
+func (c *SchemaCache) getByType(t reflect.Type) (*jsonschema.Schema, *jsonschema.Resolved, bool) {
+	if v, ok := c.byType.Load(t); ok {
+		cs := v.(*cachedSchema)
+		return cs.schema, cs.resolved, true
+	}
+	return nil, nil, false
+}
+
+func (c *SchemaCache) setByType(t reflect.Type, schema *jsonschema.Schema, resolved *jsonschema.Resolved) {
+	c.byType.Store(t, &cachedSchema{schema: schema, resolved: resolved})
+}
+
+func (c *SchemaCache) getBySchema(schema *jsonschema.Schema) (*jsonschema.Resolved, bool) {
+	if v, ok := c.bySchema.Load(schema); ok {
+		return v.(*jsonschema.Resolved), true
+	}
+	return nil, false
+}
+
+func (c *SchemaCache) setBySchema(schema *jsonschema.Schema, resolved *jsonschema.Resolved) {
+	c.bySchema.Store(schema, resolved)
+}
--- a/mcp/schema_cache_test.go
+++ b/mcp/schema_cache_test.go
@ -0,0 +1,208 @@
+// 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
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	"github.com/google/jsonschema-go/jsonschema"
+)
+
+func TestSchemaCacheByType(t *testing.T) {
+	cache := NewSchemaCache()
+
+	type TestInput struct {
+		Name string `json:"name"`
+	}
+
+	rt := reflect.TypeFor[TestInput]()
+
+	if _, _, ok := cache.getByType(rt); ok {
+		t.Error("expected cache miss for new type")
+	}
+
+	schema := &jsonschema.Schema{Type: "object"}
+	resolved, err := schema.Resolve(nil)
+	if err != nil {
+		t.Fatalf("failed to resolve schema: %v", err)
+	}
+	cache.setByType(rt, schema, resolved)
+
+	gotSchema, gotResolved, ok := cache.getByType(rt)
+	if !ok {
+		t.Error("expected cache hit after set")
+	}
+	if gotSchema != schema {
+		t.Error("schema mismatch")
+	}
+	if gotResolved != resolved {
+		t.Error("resolved schema mismatch")
+	}
+}
+
+func TestSchemaCacheBySchema(t *testing.T) {
+	cache := NewSchemaCache()
+
+	schema := &jsonschema.Schema{
+		Type: "object",
+		Properties: map[string]*jsonschema.Schema{
+			"query": {Type: "string"},
+		},
+	}
+
+	if _, ok := cache.getBySchema(schema); ok {
+		t.Error("expected cache miss for new schema")
+	}
+
+	resolved, err := schema.Resolve(nil)
+	if err != nil {
+		t.Fatalf("failed to resolve schema: %v", err)
+	}
+	cache.setBySchema(schema, resolved)
+
+	gotResolved, ok := cache.getBySchema(schema)
+	if !ok {
+		t.Error("expected cache hit after set")
+	}
+	if gotResolved != resolved {
+		t.Error("resolved schema mismatch")
+	}
+
+	// Different pointer should miss (cache uses pointer identity).
+	schema2 := &jsonschema.Schema{Type: "object"}
+	if _, ok = cache.getBySchema(schema2); ok {
+		t.Error("expected cache miss for different schema pointer")
+	}
+}
+
+func TestSetSchemaCachesGeneratedSchemas(t *testing.T) {
+	cache := NewSchemaCache()
+
+	type TestInput struct {
+		Query string `json:"query"`
+	}
+
+	rt := reflect.TypeFor[TestInput]()
+
+	var sfield1 any
+	var rfield1 *jsonschema.Resolved
+	if _, err := setSchema[TestInput](&sfield1, &rfield1, cache); err != nil {
+		t.Fatalf("setSchema failed: %v", err)
+	}
+
+	cachedSchema, cachedResolved, ok := cache.getByType(rt)
+	if !ok {
+		t.Fatal("schema not cached after first setSchema call")
+	}
+
+	var sfield2 any
+	var rfield2 *jsonschema.Resolved
+	if _, err := setSchema[TestInput](&sfield2, &rfield2, cache); err != nil {
+		t.Fatalf("setSchema failed on second call: %v", err)
+	}
+
+	if sfield2.(*jsonschema.Schema) != cachedSchema {
+		t.Error("expected cached schema to be returned")
+	}
+	if rfield2 != cachedResolved {
+		t.Error("expected cached resolved schema to be returned")
+	}
+}
+
+func TestSetSchemaCachesProvidedSchemas(t *testing.T) {
+	cache := NewSchemaCache()
+
+	schema := &jsonschema.Schema{
+		Type: "object",
+		Properties: map[string]*jsonschema.Schema{
+			"query": {Type: "string"},
+		},
+	}
+
+	var sfield1 any = schema
+	var rfield1 *jsonschema.Resolved
+	if _, err := setSchema[map[string]any](&sfield1, &rfield1, cache); err != nil {
+		t.Fatalf("setSchema failed: %v", err)
+	}
+
+	cachedResolved, ok := cache.getBySchema(schema)
+	if !ok {
+		t.Fatal("resolved schema not cached after first setSchema call")
+	}
+	if rfield1 != cachedResolved {
+		t.Error("expected same resolved schema")
+	}
+
+	var sfield2 any = schema
+	var rfield2 *jsonschema.Resolved
+	if _, err := setSchema[map[string]any](&sfield2, &rfield2, cache); err != nil {
+		t.Fatalf("setSchema failed on second call: %v", err)
+	}
+
+	if rfield2 != cachedResolved {
+		t.Error("expected cached resolved schema to be returned")
+	}
+}
+
+func TestSetSchemaNilCache(t *testing.T) {
+	type TestInput struct {
+		Query string `json:"query"`
+	}
+
+	var sfield1 any
+	var rfield1 *jsonschema.Resolved
+	if _, err := setSchema[TestInput](&sfield1, &rfield1, nil); err != nil {
+		t.Fatalf("setSchema failed: %v", err)
+	}
+
+	var sfield2 any
+	var rfield2 *jsonschema.Resolved
+	if _, err := setSchema[TestInput](&sfield2, &rfield2, nil); err != nil {
+		t.Fatalf("setSchema failed on second call: %v", err)
+	}
+
+	if sfield1 == nil || sfield2 == nil {
+		t.Error("expected schemas to be generated")
+	}
+	if rfield1 == nil || rfield2 == nil {
+		t.Error("expected resolved schemas to be generated")
+	}
+}
+
+func TestAddToolWithSharedCache(t *testing.T) {
+	cache := NewSchemaCache()
+
+	type GreetInput struct {
+		Name string `json:"name" jsonschema:"the name to greet"`
+	}
+
+	type GreetOutput struct {
+		Message string `json:"message"`
+	}
+
+	handler := func(ctx context.Context, req *CallToolRequest, in GreetInput) (*CallToolResult, GreetOutput, error) {
+		return &CallToolResult{}, GreetOutput{Message: "Hello, " + in.Name}, nil
+	}
+
+	tool := &Tool{
+		Name:        "greet",
+		Description: "Greet someone",
+	}
+
+	// Simulate stateless server pattern: new server per request, shared cache.
+	for i := 0; i < 3; i++ {
+		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, &ServerOptions{
+			SchemaCache: cache,
+		})
+		AddTool(s, tool, handler)
+	}
+
+	rt := reflect.TypeFor[GreetInput]()
+	if _, _, ok := cache.getByType(rt); !ok {
+		t.Error("expected schema to be cached by type after multiple AddTool calls")
+	}
+}
--- a/mcp/server.go
+++ b/mcp/server.go
@ -127,6 +127,11 @@ type ServerOptions struct {
 	//
 	// Deprecated: Use Capabilities instead.
 	HasTools bool
+	// SchemaCache, if non-nil, caches JSON schemas to avoid repeated
+	// reflection. This is useful for stateless server deployments where
+	// a new [Server] is created for each request. See [SchemaCache] for
+	// trade-offs and usage guidance.
+	SchemaCache *SchemaCache

 	// GetSessionID provides the next session ID to use for an incoming request.
 	// If nil, a default randomly generated ID will be used.
@ -275,7 +280,7 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 	s.changeAndNotify(notificationToolListChanged, func() bool { s.tools.add(st); return true })
 }

-func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHandler, error) {
+func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out], cache *SchemaCache) (*Tool, ToolHandler, error) {
 	tt := *t

 	// Special handling for an "any" input: treat as an empty object.
@ -284,7 +289,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
 	}

 	var inputResolved *jsonschema.Resolved
-	if _, err := setSchema[In](&tt.InputSchema, &inputResolved); err != nil {
+	if _, err := setSchema[In](&tt.InputSchema, &inputResolved, cache); err != nil {
 		return nil, nil, fmt.Errorf("input schema: %w", err)
 	}

@ -299,7 +304,7 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
 	)
 	if t.OutputSchema != nil || reflect.TypeFor[Out]() != reflect.TypeFor[any]() {
 		var err error
-		elemZero, err = setSchema[Out](&tt.OutputSchema, &outputResolved)
+		elemZero, err = setSchema[Out](&tt.OutputSchema, &outputResolved, cache)
 		if err != nil {
 			return nil, nil, fmt.Errorf("output schema: %v", err)
 		}
@ -400,29 +405,75 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
 // pointer: if the user provided the schema, they may have intentionally
 // derived it from the pointer type, and handling of zero values is up to them.
 //
+// If cache is non-nil, schemas are cached to avoid repeated reflection.
+//
 // TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we
 // should have a jsonschema.Zero(schema) helper?
-func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err error) {
-	var internalSchema *jsonschema.Schema
-	if *sfield == nil {
-		rt := reflect.TypeFor[T]()
-		if rt.Kind() == reflect.Pointer {
-			rt = rt.Elem()
-			zero = reflect.Zero(rt).Interface()
-		}
-		// TODO: we should be able to pass nil opts here.
-		internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
-		if err == nil {
-			*sfield = internalSchema
-		}
-	} else if err := remarshal(*sfield, &internalSchema); err != nil {
-		return zero, err
+func setSchema[T any](sfield *any, rfield **jsonschema.Resolved, cache *SchemaCache) (zero any, err error) {
+	rt := reflect.TypeFor[T]()
+	if rt.Kind() == reflect.Pointer {
+		rt = rt.Elem()
+		zero = reflect.Zero(rt).Interface()
 	}
+
+	var internalSchema *jsonschema.Schema
+
+	if *sfield == nil {
+		// No schema provided: check cache, or generate via reflection.
+		if cache != nil {
+			if schema, resolved, ok := cache.getByType(rt); ok {
+				*sfield = schema
+				*rfield = resolved
+				return zero, nil
+			}
+		}
+
+		internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
+		if err != nil {
+			return zero, err
+		}
+		*sfield = internalSchema
+
+		resolved, err := internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
+		if err != nil {
+			return zero, err
+		}
+		*rfield = resolved
+		if cache != nil {
+			cache.setByType(rt, internalSchema, resolved)
+		}
+		return zero, nil
+	}
+
+	// Schema was provided: check cache by pointer, or resolve it.
+	if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok {
+		if cache != nil {
+			if resolved, ok := cache.getBySchema(providedSchema); ok {
+				*rfield = resolved
+				return zero, nil
+			}
+		}
+		internalSchema = providedSchema
+	} else {
+		// Schema provided as different type (e.g., map): remarshal to *Schema.
+		if err := remarshal(*sfield, &internalSchema); err != nil {
+			return zero, err
+		}
+	}
+
+	resolved, err := internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
 	if err != nil {
 		return zero, err
 	}
-	*rfield, err = internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
-	return zero, err
+	*rfield = resolved
+
+	if cache != nil {
+		if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok {
+			cache.setBySchema(providedSchema, resolved)
+		}
+	}
+
+	return zero, nil
 }

 // AddTool adds a tool and typed tool handler to the server.
@ -445,7 +496,7 @@ func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err
 // tools to conform to the MCP spec. See [ToolHandlerFor] for a detailed
 // description of this automatic behavior.
 func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
-	tt, hh, err := toolForErr(t, h)
+	tt, hh, err := toolForErr(t, h, s.opts.SchemaCache)
 	if err != nil {
 		panic(fmt.Sprintf("AddTool: tool %q: %v", t.Name, err))
 	}
--- a/mcp/server_test.go
+++ b/mcp/server_test.go
@ -640,7 +640,7 @@ func testToolForSchema[In, Out any](t *testing.T, tool *Tool, in string, out Out
 	th := func(context.Context, *CallToolRequest, In) (*CallToolResult, Out, error) {
 		return nil, out, nil
 	}
-	gott, goth, err := toolForErr(tool, th)
+	gott, goth, err := toolForErr(tool, th, nil)
 	if err != nil {
 		t.Fatal(err)
 	}