mcp: make schema caching opt-in via ServerOptions.SchemaCache

Add ServerOptions.SchemaCache to enable optional caching of JSON schemas
for tools. When set, the cache avoids repeated reflection-based schema
generation and resolution, which significantly improves performance for
stateless server deployments where tools are re-registered on every request.

This change addresses reviewer feedback to make the caching optimization
opt-in, avoiding subtle behavior changes for users who may re-use and
modify schemas between requests.

Usage:
    cache := mcp.NewSchemaCache()
    server := mcp.NewServer(impl, &mcp.ServerOptions{
        SchemaCache: cache,
    })
    mcp.AddTool(server, tool, handler)

Author: SamMorrowDrums

mcp/schema_cache.go

@@ -18,6 +18,8 @@ import (
 // This cache significantly improves performance for stateless server deployments
 // where tools are re-registered on every request. Without caching, each AddTool
 // call would trigger expensive reflection-based schema generation and resolution.
+//
+// Create a cache using [NewSchemaCache] and pass it to [ServerOptions.SchemaCache].
 type schemaCache struct {
 	// byType caches schemas generated from Go types via jsonschema.ForType.
 	// Key: reflect.Type, Value: *cachedSchema
@@ -36,9 +38,11 @@ type cachedSchema struct {
 	resolved *jsonschema.Resolved
 }
 
-// globalSchemaCache is the package-level cache used by setSchema.
-// It is unbounded since typical MCP servers have <100 tools.
-var globalSchemaCache = &schemaCache{}
+// NewSchemaCache creates a new schema cache for use with [ServerOptions.SchemaCache].
+// Safe for concurrent use, unbounded.
+func NewSchemaCache() *schemaCache {
+	return &schemaCache{}
+}
 
 // getByType retrieves a cached schema by Go type.
 // Returns the schema, resolved schema, and whether the cache hit.
@@ -68,9 +72,3 @@ func (c *schemaCache) getBySchema(schema *jsonschema.Schema) (*jsonschema.Resolv
 func (c *schemaCache) setBySchema(schema *jsonschema.Schema, resolved *jsonschema.Resolved) {
 	c.bySchema.Store(schema, resolved)
 }
-
-// resetForTesting clears the cache. Only for use in tests.
-func (c *schemaCache) resetForTesting() {
-	c.byType.Clear()
-	c.bySchema.Clear()
-}

mcp/schema_cache_benchmark_test.go

@@ -34,14 +34,16 @@ func BenchmarkAddToolTypedHandler(b *testing.B) {
 		Description: "Search for items",
 	}
 
-	// Reset cache to simulate cold start for first iteration
-	globalSchemaCache.resetForTesting()
+	// Create a shared cache for caching benefit
+	cache := NewSchemaCache()
 
 	b.ResetTimer()
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
-		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, nil)
+		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, &ServerOptions{
+			SchemaCache: cache,
+		})
 		AddTool(s, tool, handler)
 	}
 }
@@ -69,9 +71,6 @@ func BenchmarkAddToolPreDefinedSchema(b *testing.B) {
 		InputSchema: schema, // Pre-defined schema like github-mcp-server
 	}
 
-	// Reset cache to simulate cold start for first iteration
-	globalSchemaCache.resetForTesting()
-
 	b.ResetTimer()
 	b.ReportAllocs()
 
@@ -108,9 +107,7 @@ func BenchmarkAddToolTypedHandlerNoCache(b *testing.B) {
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
-		// Reset cache every iteration to simulate no caching
-		globalSchemaCache.resetForTesting()
-
+		// No cache - each iteration generates new schemas
 		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, nil)
 		AddTool(s, tool, handler)
 	}
@@ -146,14 +143,16 @@ func BenchmarkAddToolMultipleTools(b *testing.B) {
 	tool2 := &Tool{Name: "tool2", Description: "Tool 2"}
 	tool3 := &Tool{Name: "tool3", Description: "Tool 3"}
 
-	// Reset cache before benchmark
-	globalSchemaCache.resetForTesting()
+	// Create a shared cache for caching benefit
+	cache := NewSchemaCache()
 
 	b.ResetTimer()
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i++ {
-		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, nil)
+		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, &ServerOptions{
+			SchemaCache: cache,
+		})
 		AddTool(s, tool1, handler1)
 		AddTool(s, tool2, handler2)
 		AddTool(s, tool3, handler3)

mcp/schema_cache_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 func TestSchemaCache_ByType(t *testing.T) {
-	cache := &schemaCache{}
+	cache := NewSchemaCache()
 
 	type TestInput struct {
 		Name string `json:"name"`
@@ -49,7 +49,7 @@ func TestSchemaCache_ByType(t *testing.T) {
 }
 
 func TestSchemaCache_BySchema(t *testing.T) {
-	cache := &schemaCache{}
+	cache := NewSchemaCache()
 
 	schema := &jsonschema.Schema{
 		Type: "object",
@@ -89,7 +89,7 @@ func TestSchemaCache_BySchema(t *testing.T) {
 }
 
 func TestSetSchema_CachesGeneratedSchemas(t *testing.T) {
-	globalSchemaCache.resetForTesting()
+	cache := NewSchemaCache()
 
 	type TestInput struct {
 		Query string `json:"query"`
@@ -100,21 +100,21 @@ func TestSetSchema_CachesGeneratedSchemas(t *testing.T) {
 	// First call should generate and cache
 	var sfield1 any
 	var rfield1 *jsonschema.Resolved
-	_, err := setSchema[TestInput](&sfield1, &rfield1)
+	_, err := setSchema[TestInput](&sfield1, &rfield1, cache)
 	if err != nil {
 		t.Fatalf("setSchema failed: %v", err)
 	}
 
 	// Verify it's in cache
-	cachedSchema, cachedResolved, ok := globalSchemaCache.getByType(rt)
+	cachedSchema, cachedResolved, ok := cache.getByType(rt)
 	if !ok {
 		t.Fatal("schema not cached after first setSchema call")
 	}
 
 	// Second call should hit cache
 	var sfield2 any
 	var rfield2 *jsonschema.Resolved
-	_, err = setSchema[TestInput](&sfield2, &rfield2)
+	_, err = setSchema[TestInput](&sfield2, &rfield2, cache)
 	if err != nil {
 		t.Fatalf("setSchema failed on second call: %v", err)
 	}
@@ -129,7 +129,7 @@ func TestSetSchema_CachesGeneratedSchemas(t *testing.T) {
 }
 
 func TestSetSchema_CachesProvidedSchemas(t *testing.T) {
-	globalSchemaCache.resetForTesting()
+	cache := NewSchemaCache()
 
 	// This simulates the github-mcp-server pattern:
 	// schema is created once and reused across requests
@@ -143,13 +143,13 @@ func TestSetSchema_CachesProvidedSchemas(t *testing.T) {
 	// First call should resolve and cache
 	var sfield1 any = schema
 	var rfield1 *jsonschema.Resolved
-	_, err := setSchema[map[string]any](&sfield1, &rfield1)
+	_, err := setSchema[map[string]any](&sfield1, &rfield1, cache)
 	if err != nil {
 		t.Fatalf("setSchema failed: %v", err)
 	}
 
 	// Verify it's in cache
-	cachedResolved, ok := globalSchemaCache.getBySchema(schema)
+	cachedResolved, ok := cache.getBySchema(schema)
 	if !ok {
 		t.Fatal("resolved schema not cached after first setSchema call")
 	}
@@ -160,7 +160,7 @@ func TestSetSchema_CachesProvidedSchemas(t *testing.T) {
 	// Second call with same schema pointer should hit cache
 	var sfield2 any = schema
 	var rfield2 *jsonschema.Resolved
-	_, err = setSchema[map[string]any](&sfield2, &rfield2)
+	_, err = setSchema[map[string]any](&sfield2, &rfield2, cache)
 	if err != nil {
 		t.Fatalf("setSchema failed on second call: %v", err)
 	}
@@ -170,8 +170,39 @@ func TestSetSchema_CachesProvidedSchemas(t *testing.T) {
 	}
 }
 
+func TestSetSchema_NoCacheWhenNil(t *testing.T) {
+	type TestInput struct {
+		Query string `json:"query"`
+	}
+
+	// First call without cache
+	var sfield1 any
+	var rfield1 *jsonschema.Resolved
+	_, err := setSchema[TestInput](&sfield1, &rfield1, nil)
+	if err != nil {
+		t.Fatalf("setSchema failed: %v", err)
+	}
+
+	// Second call without cache - should still generate a new schema
+	var sfield2 any
+	var rfield2 *jsonschema.Resolved
+	_, err = setSchema[TestInput](&sfield2, &rfield2, nil)
+	if err != nil {
+		t.Fatalf("setSchema failed on second call: %v", err)
+	}
+
+	// Both calls should succeed, schemas should be equivalent but not same pointer
+	// (since no caching is happening)
+	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 TestAddTool_CachesBetweenCalls(t *testing.T) {
-	globalSchemaCache.resetForTesting()
+	cache := NewSchemaCache()
 
 	type GreetInput struct {
 		Name string `json:"name" jsonschema:"the name to greet"`
@@ -190,15 +221,17 @@ func TestAddTool_CachesBetweenCalls(t *testing.T) {
 		Description: "Greet someone",
 	}
 
-	// Simulate stateless server pattern: create new server each time
+	// Simulate stateless server pattern: create new server each time, but share cache
 	for i := 0; i < 3; i++ {
-		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, nil)
+		s := NewServer(&Implementation{Name: "test", Version: "1.0"}, &ServerOptions{
+			SchemaCache: cache,
+		})
 		AddTool(s, tool, handler)
 	}
 
 	// Verify schema was cached by type
 	rt := reflect.TypeFor[GreetInput]()
-	_, _, ok := globalSchemaCache.getByType(rt)
+	_, _, ok := cache.getByType(rt)
 	if !ok {
 		t.Error("expected schema to be cached by type after multiple AddTool calls")
 	}

mcp/server.go

@@ -89,6 +89,10 @@ type ServerOptions struct {
 	// If true, advertises the tools capability during initialization,
 	// even if no tools have been registered.
 	HasTools bool
+	// SchemaCache, if non-nil, enables caching of JSON schemas for tools.
+	// This can significantly improve performance for stateless server
+	// deployments where tools are re-registered on every request.
+	SchemaCache *schemaCache
 
 	// GetSessionID provides the next session ID to use for an incoming request.
 	// If nil, a default randomly generated ID will be used.
@@ -239,7 +243,7 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
 		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.
@@ -248,7 +252,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)
 	}
 
@@ -263,7 +267,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)
 		}
@@ -364,9 +368,11 @@ 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) {
+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()
@@ -377,37 +383,43 @@ func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err
 
 	if *sfield == nil {
 		// Case 1: No schema provided - check type cache first
-		if schema, resolved, ok := globalSchemaCache.getByType(rt); ok {
-			*sfield = schema
-			*rfield = resolved
-			return zero, nil
+		if cache != nil {
+			if schema, resolved, ok := cache.getByType(rt); ok {
+				*sfield = schema
+				*rfield = resolved
+				return zero, nil
+			}
 		}
 
-		// Generate schema via reflection (expensive, but cached for next time)
+		// Generate schema via reflection (expensive, but cached for next time if cache is set)
 		internalSchema, err = jsonschema.ForType(rt, &jsonschema.ForOptions{})
 		if err != nil {
 			return zero, err
 		}
 		*sfield = internalSchema
 
-		// Resolve and cache
+		// Resolve and optionally cache
 		resolved, err := internalSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
 		if err != nil {
 			return zero, err
 		}
 		*rfield = resolved
-		globalSchemaCache.setByType(rt, internalSchema, resolved)
+		if cache != nil {
+			cache.setByType(rt, internalSchema, resolved)
+		}
 		return zero, nil
 	}
 
 	// Case 2: Schema was provided
 	// Check if it's a *jsonschema.Schema we can cache by pointer
 	if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok {
-		if resolved, ok := globalSchemaCache.getBySchema(providedSchema); ok {
-			*rfield = resolved
-			return zero, nil
+		if cache != nil {
+			if resolved, ok := cache.getBySchema(providedSchema); ok {
+				*rfield = resolved
+				return zero, nil
+			}
 		}
-		// Need to resolve and cache
+		// Need to resolve and optionally cache
 		internalSchema = providedSchema
 	} else {
 		// Schema provided as different type (e.g., map) - need to remarshal
@@ -424,8 +436,10 @@ func setSchema[T any](sfield *any, rfield **jsonschema.Resolved) (zero any, err
 	*rfield = resolved
 
 	// Cache by schema pointer if we got a direct *jsonschema.Schema
-	if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok {
-		globalSchemaCache.setBySchema(providedSchema, resolved)
+	if cache != nil {
+		if providedSchema, ok := (*sfield).(*jsonschema.Schema); ok {
+			cache.setBySchema(providedSchema, resolved)
+		}
 	}
 
 	return zero, nil
@@ -451,7 +465,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))
 	}

mcp/server_test.go

@@ -562,7 +562,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)
 	}