diff --git a/src/mcp/mcp-tool-service.ts b/src/mcp/mcp-tool-service.ts index ea265a9ed..63f35364b 100644 --- a/src/mcp/mcp-tool-service.ts +++ b/src/mcp/mcp-tool-service.ts @@ -14,6 +14,8 @@ import type { McpServer, McpServerStore } from "./mcp-server-store.ts"; const REFRESH_INTERVAL_MS = 5 * 60_000; const MAX_TOOLS_PER_SERVER = 64; const MAX_RESULT_CHARS = 60_000; +const MAX_SCHEMA_NODES = 20_000; +const MAX_SCHEMA_DEPTH = 100; export interface McpToolDescriptor { /** Namespaced tool name exposed to the model, e.g. "salesforce_query". */ @@ -37,6 +39,73 @@ export interface McpToolService { close(): void; } +function resolvePointer(root: unknown, ref: string): unknown { + if (!ref.startsWith("#/")) return undefined; + let node: unknown = root; + for (const raw of ref.slice(2).split("/")) { + let key: string; + try { + key = decodeURIComponent(raw).replace(/~1/g, "/").replace(/~0/g, "~"); + } catch { + return undefined; + } + if (Array.isArray(node)) node = node[Number(key)]; + else if (node && typeof node === "object") node = (node as Record)[key]; + else return undefined; + } + return node; +} + +function expandRef( + ref: string, + root: unknown, + stack: string[], + budget: { left: number }, + depth: number, +): Record { + if (stack.includes(ref)) return {}; + const target = resolvePointer(root, ref); + if (typeof target === "boolean") return target ? {} : { not: {} }; + if (!target || typeof target !== "object" || Array.isArray(target)) return {}; + return inlineRefs(target, root, [...stack, ref], budget, depth) as Record; +} + +function inlineRefs( + node: unknown, + root: unknown, + stack: string[], + budget: { left: number }, + depth = 0, +): unknown { + if (budget.left <= 0 || depth > MAX_SCHEMA_DEPTH) return Array.isArray(node) ? [] : {}; + budget.left -= 1; + if (Array.isArray(node)) return node.map((item) => inlineRefs(item, root, stack, budget, depth + 1)); + if (!node || typeof node !== "object") return node; + const obj = node as Record; + const siblings: Record = {}; + for (const [k, v] of Object.entries(obj)) { + if (k === "$ref" || k === "$defs" || k === "definitions" || k === "__proto__") continue; + siblings[k] = inlineRefs(v, root, stack, budget, depth + 1); + } + if (typeof obj.$ref !== "string") return siblings; + return { ...expandRef(obj.$ref, root, stack, budget, depth + 1), ...siblings }; +} + +export function sanitizeToolSchema(schema: Record): Record { + const inlined = inlineRefs(schema, schema, [], { left: MAX_SCHEMA_NODES }); + if (!inlined || typeof inlined !== "object" || Array.isArray(inlined)) return objectSchemaFallback(); + const out = inlined as Record; + if (out.type === "object") return out; + if (out.type === undefined && out.properties && typeof out.properties === "object") { + return { ...out, type: "object" }; + } + return objectSchemaFallback(); +} + +function objectSchemaFallback(): Record { + return { type: "object", properties: {}, additionalProperties: true }; +} + function authOf(server: McpServer): McpAuth { if (server.auth === "bearer") return { mode: "bearer", token: server.bearerToken ?? "" }; if (server.auth === "client-credentials") @@ -92,7 +161,7 @@ export function createMcpToolService(opts: { serverId: server.id, remoteName: tool.name, description: tool.description || `${tool.name} on ${server.name}`, - inputSchema: tool.inputSchema, + inputSchema: sanitizeToolSchema(tool.inputSchema), readOnly: server.readOnly, }); } diff --git a/test/mcp-connectors.test.ts b/test/mcp-connectors.test.ts index 2158ed32f..19458979f 100644 --- a/test/mcp-connectors.test.ts +++ b/test/mcp-connectors.test.ts @@ -2,7 +2,7 @@ import { test } from "node:test"; import assert from "node:assert/strict"; import { createMcpClient, mcpResultText, type McpFetch } from "../src/mcp/mcp-client.ts"; import { createMcpServerStore, isValidMcpServerId, type McpServer } from "../src/mcp/mcp-server-store.ts"; -import { createMcpToolService } from "../src/mcp/mcp-tool-service.ts"; +import { createMcpToolService, sanitizeToolSchema } from "../src/mcp/mcp-tool-service.ts"; import { createMemoryMap } from "../src/persistence/durable-map.ts"; function jsonResponse(body: unknown, status = 200, contentType = "application/json") { @@ -124,3 +124,150 @@ test("unknown tool call rejects", async () => { await assert.rejects(() => service.call("nope_tool", {}), /unknown MCP tool/); service.close(); }); + +test("tool schemas ingest with internal refs inlined", async () => { + const store = createMcpServerStore(createMemoryMap()); + const schema = { + type: "object", + $defs: { id: { type: "string", minLength: 1 } }, + properties: { + board: { $ref: "#/$defs/id" }, + parent: { $ref: "#/properties/board" }, + missing: { $ref: "#/$defs/nope" }, + external: { $ref: "https://example.com/schema.json#/x" }, + }, + }; + const fetch: McpFetch = async (_url, init) => { + const req = JSON.parse(init.body) as { id: number; method: string }; + if (req.method === "tools/list") { + return jsonResponse({ + jsonrpc: "2.0", + id: req.id, + result: { tools: [{ name: "create", description: "Create", inputSchema: schema }] }, + }); + } + return jsonResponse({ jsonrpc: "2.0", id: req.id, result: {} }); + }; + const service = createMcpToolService({ servers: store, fetchImpl: fetch, refreshIntervalMs: 3600_000 }); + await store.put(server()); + await service.refresh(); + const def = service.toolDefs().find((d) => d.name === "crm_create"); + assert.ok(def); + const props = (def.inputSchema as { properties: Record }).properties; + assert.equal(JSON.stringify(def.inputSchema).includes("$ref"), false); + assert.deepEqual(props.board, { type: "string", minLength: 1 }); + assert.deepEqual(props.parent, { type: "string", minLength: 1 }); + assert.deepEqual(props.missing, {}); + assert.deepEqual(props.external, {}); + assert.equal((def.inputSchema as Record).$defs, undefined); + service.close(); +}); + +test("sanitizeToolSchema keeps keywords that sit beside a $ref", () => { + const out = sanitizeToolSchema({ + type: "object", + $defs: { int: { type: "integer" } }, + properties: { + count: { $ref: "#/$defs/int", minimum: 1, maximum: 10, description: "1..10 only" }, + }, + }) as { properties: Record> }; + assert.deepEqual(out.properties.count, { + type: "integer", + minimum: 1, + maximum: 10, + description: "1..10 only", + }); +}); + +test("sanitizeToolSchema always yields an object schema", () => { + const fallback = { type: "object", properties: {}, additionalProperties: true }; + assert.deepEqual(sanitizeToolSchema({ $ref: "#/$defs/Input" }), fallback); + assert.deepEqual(sanitizeToolSchema({ $ref: "#" }), fallback); + assert.deepEqual(sanitizeToolSchema({ type: "string" }), fallback); + assert.deepEqual(sanitizeToolSchema({ properties: { a: { type: "string" } } }), { + type: "object", + properties: { a: { type: "string" } }, + }); +}); + +test("sanitizeToolSchema bounds expansion instead of exploding", () => { + const width = 8; + const depth = 8; + const defs: Record = { leaf: { type: "string" } }; + for (let level = 1; level <= depth; level += 1) { + const properties: Record = {}; + const child = level === 1 ? "#/$defs/leaf" : `#/$defs/level${level - 1}`; + for (let field = 0; field < width; field += 1) properties[`f${field}`] = { $ref: child }; + defs[`level${level}`] = { type: "object", properties }; + } + const schema = { type: "object", $defs: defs, properties: { root: { $ref: `#/$defs/level${depth}` } } }; + const started = Date.now(); + const out = sanitizeToolSchema(schema); + assert.ok(Date.now() - started < 2000); + assert.ok(JSON.stringify(out).length < 2_000_000); +}); + +test("sanitizeToolSchema survives deep nesting without refs", () => { + let node: Record = { type: "string" }; + for (let i = 0; i < 5000; i += 1) node = { type: "object", properties: { next: node } }; + assert.equal(typeof sanitizeToolSchema(node), "object"); +}); + +test("sanitizeToolSchema resolves escaped and indexed pointers", () => { + const out = sanitizeToolSchema({ + type: "object", + $defs: { "a/b": { type: "number" }, "c~d": { type: "boolean" }, "e f": { type: "null" } }, + prefixItems: [{ type: "integer" }], + properties: { + slash: { $ref: "#/$defs/a~1b" }, + tilde: { $ref: "#/$defs/c~0d" }, + spaced: { $ref: "#/$defs/e%20f" }, + indexed: { $ref: "#/prefixItems/0" }, + }, + }) as { properties: Record }; + assert.deepEqual(out.properties.slash, { type: "number" }); + assert.deepEqual(out.properties.tilde, { type: "boolean" }); + assert.deepEqual(out.properties.spaced, { type: "null" }); + assert.deepEqual(out.properties.indexed, { type: "integer" }); +}); + +test("sanitizeToolSchema stops cyclic refs, direct and mutual", () => { + const direct = sanitizeToolSchema({ + type: "object", + properties: { self: { $ref: "#/properties/self" } }, + }) as { properties: Record }; + assert.deepEqual(direct.properties.self, {}); + const mutual = sanitizeToolSchema({ + type: "object", + $defs: { a: { $ref: "#/$defs/b" }, b: { $ref: "#/$defs/a" } }, + properties: { start: { $ref: "#/$defs/a" } }, + }) as { properties: Record }; + assert.deepEqual(mutual.properties.start, {}); +}); + +test("sanitizeToolSchema drops nested $defs and a __proto__ key", () => { + const out = sanitizeToolSchema({ + type: "object", + properties: { + nested: { type: "object", $defs: { x: { type: "string" } }, definitions: { y: { type: "string" } } }, + }, + }) as { properties: Record | undefined> }; + assert.equal(out.properties.nested?.$defs, undefined); + assert.equal(out.properties.nested?.definitions, undefined); + assert.deepEqual(out.properties.nested, { type: "object" }); + + const hostile = sanitizeToolSchema( + JSON.parse('{"type":"object","properties":{"__proto__":{"type":"string"},"ok":{"type":"string"}}}'), + ) as { properties: Record }; + assert.deepEqual(Object.keys(hostile.properties), ["ok"]); + assert.equal(Object.getPrototypeOf(hostile.properties), Object.prototype); +}); + +test("sanitizeToolSchema keeps a false subschema restrictive", () => { + const out = sanitizeToolSchema({ + type: "object", + $defs: { never: false }, + properties: { blocked: { $ref: "#/$defs/never" } }, + }) as { properties: Record }; + assert.deepEqual(out.properties.blocked, { not: {} }); +});