From 0e68e73f657be1505375d79399b29d7502964503 Mon Sep 17 00:00:00 2001 From: 0xamlab <79715912+0xamlab@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:26:11 -0400 Subject: [PATCH] =?UTF-8?q?fix(core):=20memoize=20Standard=20Schema=20?= =?UTF-8?q?=E2=86=92=20JSON=20Schema=20conversion=20per=20schema=20instanc?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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: every request re-converted every registered tool's schema (once eagerly in registerTool, again per tools/list). standardSchemaToJsonSchema now memoizes successful conversions process-wide in a WeakMap keyed by schema identity and io direction, so an app that hoists its schemas to module scope (the natural pattern in that model) converts each schema once per process instead of once per McpServer instance — 53 hoisted tools on a fresh server: ~14-19 ms -> ~1 ms. Conversion failures are not memoized, entries stay collectible with their schema, and repeat calls return the same (read-only) object. Fixes #2838 --- .../memoize-standard-schema-conversion.md | 6 ++ .../core-internal/src/util/standardSchema.ts | 32 ++++++++ .../test/util/standardSchema.test.ts | 81 +++++++++++++++++++ .../test/server/toolSchemaMemoization.test.ts | 59 ++++++++++++++ 4 files changed, 178 insertions(+) create mode 100644 .changeset/memoize-standard-schema-conversion.md create mode 100644 packages/server/test/server/toolSchemaMemoization.test.ts 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); + }); +});