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>

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))
 	}