diff --git a/.changeset/memoize-standard-schema-conversion.md b/.changeset/memoize-standard-schema-conversion.md new file mode 100644 index 0000000000..e2ed922ce9 --- /dev/null +++ b/.changeset/memoize-standard-schema-conversion.md @@ -0,0 +1,6 @@ +--- +'@modelcontextprotocol/client': patch +'@modelcontextprotocol/server': patch +--- + +Memoize Standard Schema → JSON Schema conversion per schema instance, across `McpServer` instances. In the stateless `createMcpHandler(() => buildServer())` pattern the app builds a fresh `McpServer` per request and re-registers its tools, so the per-instance `_toolInputSchemaJson` memo never hit and every request re-converted every registered tool's schema (once eagerly in `registerTool`, again per `tools/list`). `standardSchemaToJsonSchema` now caches successful conversions process-wide, keyed by schema identity (a `WeakMap`, so entries stay collectible with their schema) and `io` direction: an app that hoists its schemas to module scope converts each schema once per process instead of once per request (53 hoisted tools on a fresh server: ~14–19 ms → ~1 ms). Conversion failures are not cached, and repeat calls return the same object, which callers must treat as read-only. Fixes #2838. diff --git a/packages/core-internal/src/util/standardSchema.ts b/packages/core-internal/src/util/standardSchema.ts index d904c7f5fa..b9a3bfa199 100644 --- a/packages/core-internal/src/util/standardSchema.ts +++ b/packages/core-internal/src/util/standardSchema.ts @@ -169,6 +169,19 @@ let warnedZodFallback = false; /** JSON Schema draft targeted by every conversion; shared so pattern references above stay in lockstep. */ export const JSON_SCHEMA_CONVERSION_TARGET = 'draft-2020-12'; +/** + * Process-wide memo for {@linkcode standardSchemaToJsonSchema}. In the + * per-request-factory `createMcpHandler` model the app builds a fresh + * `McpServer` per request, so the per-instance `_toolInputSchemaJson` memo + * never hits and every request re-converts every registered tool's schema. + * Apps in that model hoist their schema definitions to module scope, so keying + * by schema identity converts each schema once per process instead of once per + * `McpServer` instance. The WeakMap keeps entries collectible with their + * schema, so per-request schema objects (no reuse) cost nothing beyond one + * un-hittable entry that is collected with the schema. + */ +const jsonSchemaConversionMemo = new WeakMap>>>(); + /** * Converts a StandardSchema to JSON Schema for use as an MCP tool/prompt schema. * @@ -179,8 +192,27 @@ export const JSON_SCHEMA_CONVERSION_TARGET = 'draft-2020-12'; * and throws on an explicit non-object `type` (e.g. `z.string()`). For * `io: 'output'` a non-object root is returned as-is; the `"object"` default is * applied only when the root is provably object-shaped. + * + * Successful conversions are memoized process-wide, keyed by schema identity + * and `io` direction (see {@linkcode jsonSchemaConversionMemo}). Repeat calls + * with the same schema instance return the SAME object — callers must treat + * the result as read-only. A conversion that throws is not memoized, so a + * throwing schema keeps throwing from the same call sites it always has. */ export function standardSchemaToJsonSchema(schema: StandardJSONSchemaV1, io: 'input' | 'output' = 'input'): Record { + const memoized = jsonSchemaConversionMemo.get(schema); + const hit = memoized?.[io]; + if (hit !== undefined) return hit; + const result = convertStandardSchemaToJsonSchema(schema, io); + if (memoized === undefined) { + jsonSchemaConversionMemo.set(schema, { [io]: result }); + } else { + memoized[io] = result; + } + return result; +} + +function convertStandardSchemaToJsonSchema(schema: StandardJSONSchemaV1, io: 'input' | 'output'): Record { const std = schema['~standard']; let result: Record; if (std.jsonSchema) { diff --git a/packages/core-internal/test/util/standardSchema.test.ts b/packages/core-internal/test/util/standardSchema.test.ts index 8856592ff0..4e30c4a866 100644 --- a/packages/core-internal/test/util/standardSchema.test.ts +++ b/packages/core-internal/test/util/standardSchema.test.ts @@ -1,7 +1,31 @@ import * as z from 'zod/v4'; +import type { StandardSchemaWithJSON } from '../../src/util/standardSchema'; import { standardSchemaToJsonSchema } from '../../src/util/standardSchema'; +/** Minimal vendor-neutral Standard Schema whose JSON Schema conversions count invocations. */ +function makeCountingSchema(root: Record): { schema: StandardSchemaWithJSON; calls: { input: number; output: number } } { + const calls = { input: 0, output: 0 }; + const schema: StandardSchemaWithJSON = { + '~standard': { + version: 1, + vendor: 'counting-fake', + validate: value => ({ value }), + jsonSchema: { + input: () => { + calls.input++; + return { ...root }; + }, + output: () => { + calls.output++; + return { ...root }; + } + } + } + }; + return { schema, calls }; +} + describe('standardSchemaToJsonSchema', () => { test('emits type:object for plain z.object schemas', () => { const schema = z.object({ name: z.string(), age: z.number() }); @@ -39,4 +63,61 @@ describe('standardSchemaToJsonSchema', () => { expect(keys.filter(k => k === 'type')).toHaveLength(1); expect(result.type).toBe('object'); }); + + describe('memoization', () => { + test('converts a given schema instance at most once per io direction', () => { + const { schema, calls } = makeCountingSchema({ type: 'object', properties: { a: { type: 'string' } } }); + + const first = standardSchemaToJsonSchema(schema, 'input'); + expect(standardSchemaToJsonSchema(schema, 'input')).toBe(first); + expect(standardSchemaToJsonSchema(schema, 'input')).toBe(first); + expect(calls.input).toBe(1); + + // The other direction is a separate conversion, itself memoized. + const output = standardSchemaToJsonSchema(schema, 'output'); + expect(standardSchemaToJsonSchema(schema, 'output')).toBe(output); + expect(calls.output).toBe(1); + expect(calls.input).toBe(1); + }); + + test('memoizes the post-stamping result for typeless roots', () => { + const { schema, calls } = makeCountingSchema({ properties: { a: { type: 'string' } } }); + + const first = standardSchemaToJsonSchema(schema, 'input'); + expect(first.type).toBe('object'); + expect(standardSchemaToJsonSchema(schema, 'input')).toBe(first); + expect(calls.input).toBe(1); + }); + + test('returns the identical object on repeat conversion of a hoisted zod schema', () => { + // The per-request-factory `createMcpHandler` pattern: one module-scope + // schema, many McpServer instances — conversion must run once per process. + const schema = z.object({ name: z.string() }); + expect(standardSchemaToJsonSchema(schema, 'input')).toBe(standardSchemaToJsonSchema(schema, 'input')); + expect(standardSchemaToJsonSchema(schema, 'output')).toBe(standardSchemaToJsonSchema(schema, 'output')); + }); + + test('does not memoize conversion failures', () => { + let calls = 0; + const schema: StandardSchemaWithJSON = { + '~standard': { + version: 1, + vendor: 'flaky-fake', + validate: value => ({ value }), + jsonSchema: { + input: () => { + calls++; + if (calls === 1) throw new Error('transient conversion failure'); + return { type: 'object' as const }; + }, + output: () => ({ type: 'object' as const }) + } + } + }; + + expect(() => standardSchemaToJsonSchema(schema, 'input')).toThrow('transient conversion failure'); + expect(standardSchemaToJsonSchema(schema, 'input')).toEqual({ type: 'object' }); + expect(calls).toBe(2); + }); + }); }); diff --git a/packages/server/test/server/toolSchemaMemoization.test.ts b/packages/server/test/server/toolSchemaMemoization.test.ts new file mode 100644 index 0000000000..87689b8b0b --- /dev/null +++ b/packages/server/test/server/toolSchemaMemoization.test.ts @@ -0,0 +1,59 @@ +/** + * Regression test for https://github.com/modelcontextprotocol/typescript-sdk/issues/2838 + * + * In the stateless `createMcpHandler(() => buildServer())` pattern the app builds + * a fresh `McpServer` per request and re-registers its tools on it. The per-server + * `_toolInputSchemaJson` memo never hits across instances, so every request paid a + * full zod→JSON-Schema conversion for every tool. With the conversion memoized + * process-wide (keyed by schema identity), an app that hoists its schemas to + * module scope converts each one once, no matter how many `McpServer` instances + * register it — and `tools/list` on a fresh instance reuses the same conversion. + */ +import type { StandardSchemaWithJSON } from '@modelcontextprotocol/core-internal'; +import { describe, expect, it } from 'vitest'; + +import { invoke } from '../../src/server/invoke'; +import { McpServer } from '../../src/server/mcp'; + +const LEGACY = { classification: { era: 'legacy' as const } }; + +describe('registerTool schema conversion memoization (#2838)', () => { + it('converts a hoisted schema once across per-request McpServer instances and tools/list calls', async () => { + let inputConversions = 0; + // Module-scope ("hoisted") schema, as in the recommended stateless pattern. + // A structural Standard Schema double keeps the count observable; zod's own + // `~standard.jsonSchema` converter cannot be spied on. + const hoisted: StandardSchemaWithJSON = { + '~standard': { + version: 1, + vendor: 'stateless-repro', + validate: value => ({ value }), + jsonSchema: { + input: () => { + inputConversions++; + return { type: 'object', properties: { value: { type: 'string' } } }; + }, + output: () => { + throw new Error('output conversion must not run for an input-only tool'); + } + } + } + }; + + const REQUESTS = 25; + const servers: McpServer[] = []; + for (let i = 0; i < REQUESTS; i++) { + const server = new McpServer({ name: 'stateless', version: '0' }); + server.registerTool('echo', { inputSchema: hoisted }, async () => ({ content: [] })); + servers.push(server); + } + expect(inputConversions).toBe(1); + + // `tools/list` on any of those instances reuses the memoized conversion. + for (const server of servers) { + const response = await invoke(server, { jsonrpc: '2.0', id: 1, method: 'tools/list', params: {} }, LEGACY); + expect(response.status).toBe(200); + } + expect(inputConversions).toBe(1); + }); +});