diff --git a/src/server/toolRegistrationOptions.ts b/src/server/toolRegistrationOptions.ts index 5bf6a2788..49e13e43b 100644 --- a/src/server/toolRegistrationOptions.ts +++ b/src/server/toolRegistrationOptions.ts @@ -2,6 +2,15 @@ import { z } from "zod"; import type { ToolContract } from "../contracts/toolContracts.js"; +const ADVERTISED_INPUT_SCHEMAS = new WeakMap< + z.ZodObject, + Readonly> +>(); +const ADVERTISED_OUTPUT_SCHEMAS = new WeakMap< + z.ZodObject, + Readonly> +>(); + /** Project the public contract fields accepted by MCP tool registration. */ export const toolRegistrationOptions = (contract: ToolContract) => ({ title: contract.title, @@ -22,10 +31,8 @@ const advertisedInputSchema = (schema: z.ZodObject) => ({ vendor: "rea", validate: (input: unknown) => ({ value: input }), jsonSchema: { - input: () => - describeProperties(closeObjectSchemas(inputJsonSchema(schema))), - output: () => - describeProperties(closeObjectSchemas(inputJsonSchema(schema))), + input: () => advertisedInputJsonSchema(schema), + output: () => advertisedInputJsonSchema(schema), }, }, }); @@ -41,19 +48,42 @@ const advertisedOutputSchema = (schema: z.ZodObject) => { "~standard": { ...standard, jsonSchema: { - input: () => - describeRootProperties( - closeObjectSchemas(converter({ target: "draft-2020-12" })), - ), - output: () => - describeRootProperties( - closeObjectSchemas(converter({ target: "draft-2020-12" })), - ), + input: () => advertisedOutputJsonSchema(schema, converter), + output: () => advertisedOutputJsonSchema(schema, converter), }, }, }; }; +const advertisedInputJsonSchema = ( + schema: z.ZodObject, +): Readonly> => + cachedJsonSchema(ADVERTISED_INPUT_SCHEMAS, schema, () => + describeProperties(closeObjectSchemas(inputJsonSchema(schema))), + ); + +const advertisedOutputJsonSchema = ( + schema: z.ZodObject, + converter: NonNullable["output"], +): Readonly> => + cachedJsonSchema(ADVERTISED_OUTPUT_SCHEMAS, schema, () => + describeRootProperties( + closeObjectSchemas(converter({ target: "draft-2020-12" })), + ), + ); + +const cachedJsonSchema = ( + cache: WeakMap>>, + schema: z.ZodObject, + create: () => Record, +): Readonly> => { + const cached = cache.get(schema); + if (cached !== undefined) return cached; + const created = create(); + cache.set(schema, created); + return created; +}; + const inputJsonSchema = (schema: z.ZodObject): Record => { const converter = schema["~standard"].jsonSchema?.input; if (converter === undefined) diff --git a/tests/contractSnapshot.test.ts b/tests/contractSnapshot.test.ts index 833c0943b..964012325 100644 --- a/tests/contractSnapshot.test.ts +++ b/tests/contractSnapshot.test.ts @@ -62,6 +62,21 @@ const emptySchemaPaths = (value: unknown, path = "$"): string[] => { }; describe("tool contract surface", () => { + it("reuses finalized advertised schemas", () => { + const contract = TOOL_CONTRACTS[0]; + expect(contract).toBeDefined(); + if (contract === undefined) return; + + const first = toolRegistrationOptions(contract); + const second = toolRegistrationOptions(contract); + expect(first.inputSchema["~standard"].jsonSchema.input()).toBe( + second.inputSchema["~standard"].jsonSchema.output(), + ); + expect(first.outputSchema["~standard"].jsonSchema.input()).toBe( + second.outputSchema["~standard"].jsonSchema.output(), + ); + }); + it("gives fallback parameter guidance operational meaning", () => { const contract = TOOL_CONTRACTS.find( ({ name }) => name === "analyze_function",