diff --git a/host/src/services/mcp.ts b/host/src/services/mcp.ts index e9af8f2ce..82f5fcc53 100644 --- a/host/src/services/mcp.ts +++ b/host/src/services/mcp.ts @@ -26,6 +26,85 @@ export function isMcpRequestAuthorized(context: AuthPluginContext): boolean { return context.isAuthenticated === true; } +const COMPONENT_REF_PREFIX = "#/components/schemas/"; + +function collectComponentRefs(value: unknown, names: Set): void { + if (Array.isArray(value)) { + for (const item of value) collectComponentRefs(item, names); + return; + } + if (!value || typeof value !== "object") return; + for (const [key, child] of Object.entries(value)) { + if (key === "$ref" && typeof child === "string" && child.startsWith(COMPONENT_REF_PREFIX)) { + names.add(child.slice(COMPONENT_REF_PREFIX.length)); + } else { + collectComponentRefs(child, names); + } + } +} + +function rewriteComponentRefs(value: unknown): unknown { + if (Array.isArray(value)) return value.map(rewriteComponentRefs); + if (!value || typeof value !== "object") return value; + return Object.fromEntries( + Object.entries(value).map(([key, child]) => + key === "$ref" && typeof child === "string" && child.startsWith(COMPONENT_REF_PREFIX) + ? [key, `#/$defs/${child.slice(COMPONENT_REF_PREFIX.length)}`] + : [key, rewriteComponentRefs(child)], + ), + ); +} + +export function buildToolInputSchema( + operation: any, + componentSchemas: Record = {}, +): Record { + const properties: Record = {}; + const required: string[] = []; + + for (const param of operation.parameters ?? []) { + properties[param.name] = param.schema ?? { type: "string" }; + if (param.required) required.push(param.name); + } + + const bodySchema = operation.requestBody?.content?.["application/json"]?.schema; + if (bodySchema) { + if (bodySchema.type === "object" && bodySchema.properties) { + for (const [key, value] of Object.entries(bodySchema.properties)) { + properties[key] = value; + } + if (Array.isArray(bodySchema.required)) { + required.push(...bodySchema.required); + } + } else { + properties.body = bodySchema; + required.push("body"); + } + } + + const referenced = new Set(); + collectComponentRefs(properties, referenced); + const defs: Record = {}; + for (const name of referenced) { + if (!(name in componentSchemas)) continue; + defs[name] = componentSchemas[name]; + collectComponentRefs(componentSchemas[name], referenced); + } + const missing = [...referenced].filter((name) => !(name in componentSchemas)); + if (missing.length > 0) { + throw new Error( + `references component schema(s) not present in the OpenAPI spec: ${missing.join(", ")}`, + ); + } + + return { + type: "object", + properties: rewriteComponentRefs(properties), + ...(required.length > 0 ? { required } : {}), + ...(Object.keys(defs).length > 0 ? { $defs: rewriteComponentRefs(defs) } : {}), + }; +} + export function closeMcpServer(): Promise { if (!mcpTransport) return Promise.resolve(); const t = mcpTransport; @@ -64,6 +143,7 @@ export async function mountMcpRoute( const registeredTools = new Set(); const paths = (spec as any).paths ?? {}; + const componentSchemas = (spec as any).components?.schemas ?? {}; for (const [path, methods] of Object.entries(paths)) { for (const [method, operation] of Object.entries(methods as Record)) { @@ -80,6 +160,14 @@ export async function mountMcpRoute( logger.warn(`[MCP] Skipping duplicate tool "${operationId}"`); continue; } + let inputSchema: any; + try { + inputSchema = fromJsonSchema(buildToolInputSchema(operation, componentSchemas) as any); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn(`[MCP] Skipping tool "${operationId}": ${message}`); + continue; + } registeredTools.add(operationId); const methodUpper = method.toUpperCase(); @@ -92,42 +180,13 @@ export async function mountMcpRoute( .filter((p: any) => p.in === "query") .map((p: any) => p.name as string); - const properties: Record = {}; - const required: string[] = []; - - for (const param of parameters) { - properties[param.name] = param.schema ?? { type: "string" }; - if (param.required) required.push(param.name); - } - - const bodySchema = operation.requestBody?.content?.["application/json"]?.schema; - if (bodySchema) { - if (bodySchema.type === "object" && bodySchema.properties) { - for (const [key, value] of Object.entries(bodySchema.properties)) { - properties[key] = value; - } - if (Array.isArray(bodySchema.required)) { - required.push(...bodySchema.required); - } - } else { - properties.body = bodySchema; - required.push("body"); - } - } - - const inputSchema = { - type: "object" as const, - properties, - ...(required.length > 0 ? { required } : {}), - }; - const description = operation.description ?? operation.summary ?? `${methodUpper} ${path}`; server.registerTool( operationId, { description, - inputSchema: fromJsonSchema(inputSchema as any) as any, + inputSchema, }, async (args: Record) => { const store = mcpRequestStorage.getStore(); diff --git a/host/tests/unit/mcp-tool-schema.test.ts b/host/tests/unit/mcp-tool-schema.test.ts new file mode 100644 index 000000000..c9d7cbc58 --- /dev/null +++ b/host/tests/unit/mcp-tool-schema.test.ts @@ -0,0 +1,68 @@ +import "@orpc/openapi/extensions/route"; +import { fromJsonSchema } from "@modelcontextprotocol/server"; +import { oc } from "@orpc/contract"; +import { OpenAPIGenerator } from "@orpc/openapi"; +import { ZodToJsonSchemaConverter } from "@orpc/zod"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { buildToolInputSchema } from "../../src/services/mcp"; + +async function operationAndComponents(contract: Record, path: string) { + const spec: any = await new OpenAPIGenerator({ + converters: [new ZodToJsonSchemaConverter()], + }).generate(contract as any, { version: "3.1.1", base: { info: { title: "t", version: "1" } } }); + return { operation: spec.paths[path].post, components: spec.components?.schemas ?? {} }; +} + +function validate(schema: Record, value: unknown) { + return (fromJsonSchema(schema as any) as any)["~standard"].validate(value); +} + +describe("buildToolInputSchema", () => { + it("makes a recursive body schema self-contained so the tool compiles", async () => { + const { operation, components } = await operationAndComponents( + { + submit: oc + .route({ method: "POST", path: "/events" }) + .input(z.object({ name: z.string(), payload: z.json() })), + }, + "/events", + ); + + const schema = buildToolInputSchema(operation, components); + + expect(JSON.stringify(schema)).not.toContain("#/components/"); + expect(validate(schema, { name: "a", payload: { nested: [1, { deep: null }] } })).toEqual({ + value: { name: "a", payload: { nested: [1, { deep: null }] } }, + }); + expect(validate(schema, { payload: {} }).issues).toBeDefined(); + }); + + it("leaves schemas without component refs free of $defs", async () => { + const { operation, components } = await operationAndComponents( + { + submit: oc.route({ method: "POST", path: "/items" }).input(z.object({ name: z.string() })), + }, + "/items", + ); + + expect(buildToolInputSchema(operation, components)).toEqual({ + type: "object", + properties: { name: { type: "string" } }, + required: ["name"], + }); + }); + + it("names the missing component when a ref cannot be resolved", () => { + const operation = { + parameters: [], + requestBody: { + content: { "application/json": { schema: { $ref: "#/components/schemas/Gone" } } }, + }, + }; + + expect(() => buildToolInputSchema(operation as any, {})).toThrow( + /not present in the OpenAPI spec.*Gone/, + ); + }); +});