From be2b4f455bdc3b110667bd26293873f72394019f Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Sun, 25 Oct 2026 14:53:52 -0700 Subject: [PATCH] feat(server): advertised MCP schemas merge tagged unions into one object Co-Authored-By: Claude Opus 5.5 --- .../src/mcp/advertisedJsonSchema.test.ts | 21 +++----- apps/server/src/mcp/advertisedJsonSchema.ts | 51 ++++++++++++++++++- 2 files changed, 57 insertions(+), 15 deletions(-) diff --git a/apps/server/src/mcp/advertisedJsonSchema.test.ts b/apps/server/src/mcp/advertisedJsonSchema.test.ts index de72d63930298..09a90556a9dd0 100644 --- a/apps/server/src/mcp/advertisedJsonSchema.test.ts +++ b/apps/server/src/mcp/advertisedJsonSchema.test.ts @@ -4,7 +4,7 @@ import { Tool } from "effect/ai"; import { advertisedJsonSchema } from "./advertisedJsonSchema.ts"; -it("drops JSON noise but keeps null where it is required", () => { +it("drops JSON noise, merges tagged unions, and keeps null where it is required", () => { const Every = Schema.Struct({ ms: Schema.Int }); const tool = Tool.make("example", { parameters: Schema.Struct({ @@ -31,18 +31,13 @@ it("drops JSON noise but keeps null where it is required", () => { limit: { type: "integer" }, every: { type: "object", properties: { ms: { type: "integer" } }, required: ["ms"] }, target: { - anyOf: [ - { - type: "object", - properties: { kind: { type: "string", enum: ["url"] }, url: { type: "string" } }, - required: ["kind", "url"], - }, - { - type: "object", - properties: { kind: { type: "string", enum: ["number"] }, number: { type: "integer" } }, - required: ["kind", "number"], - }, - ], + type: "object", + properties: { + kind: { type: "string", enum: ["url", "number"] }, + url: { type: "string" }, + number: { type: "integer" }, + }, + required: ["kind"], }, }, required: ["beforeRunId", "every"], diff --git a/apps/server/src/mcp/advertisedJsonSchema.ts b/apps/server/src/mcp/advertisedJsonSchema.ts index f1193bd7739c0..10001307f2d6c 100644 --- a/apps/server/src/mcp/advertisedJsonSchema.ts +++ b/apps/server/src/mcp/advertisedJsonSchema.ts @@ -21,6 +21,52 @@ const SAFE_INTEGER_BOUNDS: Record = { minimum: Number.MIN_SAFE_INTEGER, }; +const tagOf = (branch: JsonObject, key: string): Json => { + const property = Predicate.isObject(branch.properties) ? branch.properties[key] : undefined; + return Predicate.isObject(property) && Array.isArray(property.enum) && property.enum.length === 1 + ? property.enum[0] + : undefined; +}; + +/** One object for a tagged union of objects: each property once, the tag listing every value. */ +const mergeTaggedObjects = (branches: ReadonlyArray): JsonObject | undefined => { + if (branches.length < 2) return undefined; + const objects = branches.filter( + (branch): branch is JsonObject & { properties: JsonObject } => + Predicate.isObject(branch) && + branch.type === "object" && + Predicate.isObject(branch.properties), + ); + if (objects.length !== branches.length) return undefined; + const tag = Object.keys(objects[0]!.properties).find((key) => { + const values = objects.map((object) => tagOf(object, key)); + return values.every((value) => value !== undefined) && new Set(values).size === values.length; + }); + if (tag === undefined) return undefined; + const properties: JsonObject = { + [tag]: { type: "string", enum: objects.map((object) => tagOf(object, tag)) }, + }; + for (const object of objects) { + for (const [name, property] of Object.entries(object.properties)) { + if (name === tag) continue; + const seen = properties[name]; + if (seen === undefined) properties[name] = property; + else if (JSON.stringify(seen) !== JSON.stringify(property)) { + const options = Predicate.isObject(seen) && isBareAnyOf(seen) ? seen.anyOf : [seen]; + if (!options.some((option) => JSON.stringify(option) === JSON.stringify(property))) + properties[name] = { anyOf: [...options, property] }; + } + } + } + const [first = [], ...others] = objects.map((object) => + Array.isArray(object.required) ? object.required : [], + ); + const requiredByAll = first.filter( + (name) => name !== tag && others.every((required) => required.includes(name)), + ); + return { type: "object", properties, required: [tag, ...requiredByAll] }; +}; + const compact = (schema: Json, optional: boolean): Json => { if (Array.isArray(schema)) return schema.map((item) => compact(item, false)); if (!Predicate.isObject(schema)) return schema; @@ -44,7 +90,7 @@ const compact = (schema: Json, optional: boolean): Json => { .flatMap((branch) => (isBareAnyOf(branch) ? branch.anyOf : [branch])) .filter((branch) => !(optional && isNullSchema(branch)) && !isJsonStringSchema(branch)); const { anyOf: _anyOf, ...rest } = result; - const only = branches.length === 1 ? branches[0] : undefined; + const only = branches.length === 1 ? branches[0] : mergeTaggedObjects(branches); result = Predicate.isObject(only) ? { ...only, ...rest } : { ...rest, anyOf: branches }; } return result; @@ -53,7 +99,8 @@ const compact = (schema: Json, optional: boolean): Json => { /** * The input schema T3 advertises for a tool. Drops JSON noise that every thread pays for: * `additionalProperties: true`, safe-integer bounds, nested `anyOf`, JSON-string fallbacks, and - * the `null` branch of optional properties. Decoding still uses the full schema. + * the `null` branch of optional properties. A tagged union of objects becomes one object. + * Decoding still uses the full schema, so each tag's own required fields still apply. */ export const advertisedJsonSchema = (tool: Tool.Any) => compact(Tool.getJsonSchema(tool), false) as ReturnType;