diff --git a/src/common/utils/tools/optionalNullSchema.test.ts b/src/common/utils/tools/optionalNullSchema.test.ts new file mode 100644 index 0000000000..3634c7a63a --- /dev/null +++ b/src/common/utils/tools/optionalNullSchema.test.ts @@ -0,0 +1,182 @@ +import { describe, expect, test } from "bun:test"; + +import { + createOptionalNullSchemaContract, + stripSyntheticNulls, + widenOptionalPropertiesToNullable, +} from "./optionalNullSchema"; + +describe("optional null JSON Schema contract", () => { + test("round trips a Linear-shaped optional argument schema", () => { + const source = { + type: "object", + required: ["issueId"], + properties: { + issueId: { type: "string" }, + cursor: { type: "string" }, + statusUpdateType: { type: "string", enum: ["project", "initiative"] }, + nullableNote: { type: ["string", "null"] }, + }, + additionalProperties: false, + }; + + const modelSchema = widenOptionalPropertiesToNullable(source); + + expect(modelSchema).toEqual({ + ...source, + properties: { + issueId: { type: "string" }, + cursor: { anyOf: [{ type: "string" }, { type: "null" }] }, + statusUpdateType: { + anyOf: [{ type: "string", enum: ["project", "initiative"] }, { type: "null" }], + }, + nullableNote: { type: ["string", "null"] }, + }, + }); + expect(source.properties.cursor).toEqual({ type: "string" }); + expect( + stripSyntheticNulls(source, { + issueId: "CODAGT-709", + cursor: "", + statusUpdateType: null, + nullableNote: null, + }) + ).toEqual({ issueId: "CODAGT-709", cursor: "", nullableNote: null }); + }); + + test("preserves optional-property annotations on the widened schema", () => { + const source = { + type: "object", + properties: { + cursor: { type: "string", title: "Cursor", description: "Continue from this cursor" }, + }, + }; + + expect(widenOptionalPropertiesToNullable(source)).toMatchObject({ + properties: { + cursor: { + title: "Cursor", + description: "Continue from this cursor", + anyOf: [source.properties.cursor, { type: "null" }], + }, + }, + }); + }); + + test("restores nested optional values in arrays and unions", () => { + const source = { + type: "object", + required: ["values"], + properties: { + values: { + anyOf: [ + { + type: "array", + items: { + type: "object", + properties: { label: { type: "string" } }, + additionalProperties: false, + }, + }, + ], + }, + }, + additionalProperties: false, + }; + + expect(widenOptionalPropertiesToNullable(source)).toMatchObject({ + properties: { + values: { + anyOf: [ + { + items: { + properties: { + label: { anyOf: [{ type: "string" }, { type: "null" }] }, + }, + }, + }, + ], + }, + }, + }); + expect(stripSyntheticNulls(source, { values: [{ label: null }] })).toEqual({ values: [{}] }); + }); + + test("preserves a raw value accepted by another union branch", () => { + const source = { + anyOf: [ + { + type: "object", + properties: { value: { type: "string" } }, + additionalProperties: false, + }, + { + type: "object", + required: ["value"], + properties: { value: { type: ["string", "null"] } }, + additionalProperties: false, + }, + ], + }; + + expect(stripSyntheticNulls(source, { value: null })).toEqual({ value: null }); + }); + + test.each(["$ref", "$dynamicRef", "$recursiveRef"])( + "falls back to non-strict decoding for schemas with %s", + (keyword) => { + const source = { + type: "object", + properties: { value: { [keyword]: "#/$defs/value" } }, + $defs: { value: { type: "string" } }, + }; + const contract = createOptionalNullSchemaContract(source); + + expect(contract.strict).toBe(false); + expect(contract.modelSchema).toEqual(source); + expect(contract.restore({ value: null })).toEqual({ value: null }); + } + ); + + test("handles boolean schemas without changing explicit valid nulls", () => { + const source = { + type: "object", + properties: { anything: true, impossible: false }, + }; + + expect(widenOptionalPropertiesToNullable(source)).toEqual({ + ...source, + properties: { + anything: true, + impossible: { anyOf: [false, { type: "null" }] }, + }, + }); + expect(stripSyntheticNulls(source, { anything: null, impossible: null })).toEqual({ + anything: null, + }); + }); + + test("applies root property constraints before matching a union branch", () => { + const source = { + type: "object", + properties: { value: { type: "string" } }, + anyOf: [{ type: "object" }], + }; + + expect(stripSyntheticNulls(source, { value: null })).toEqual({}); + }); + + test.each(["allOf", "anyOf"] as const)( + "preserves a parent-required property declared inside %s", + (keyword) => { + const source = { + type: "object", + required: ["value"], + [keyword]: [{ properties: { value: { type: "string" } } }], + }; + + expect(widenOptionalPropertiesToNullable(source)).toEqual(source); + expect(stripSyntheticNulls(source, { value: null })).toEqual({ value: null }); + } + ); +}); diff --git a/src/common/utils/tools/optionalNullSchema.ts b/src/common/utils/tools/optionalNullSchema.ts new file mode 100644 index 0000000000..a72505a762 --- /dev/null +++ b/src/common/utils/tools/optionalNullSchema.ts @@ -0,0 +1,275 @@ +import { validateJsonSchemaSubset } from "@/common/utils/jsonSchemaSubset"; + +function isRecord(value: unknown): value is Record { + return value != null && typeof value === "object" && !Array.isArray(value); +} + +function containsReferenceKeyword(value: unknown): boolean { + if (Array.isArray(value)) { + return value.some(containsReferenceKeyword); + } + if (!isRecord(value)) { + return false; + } + if (["$ref", "$dynamicRef", "$recursiveRef"].some((key) => Object.hasOwn(value, key))) { + return true; + } + return Object.values(value).some(containsReferenceKeyword); +} + +function getRequiredProperties(schema: Record): Set { + const required = new Set( + Array.isArray(schema.required) + ? schema.required.filter((key): key is string => typeof key === "string") + : [] + ); + if (Array.isArray(schema.allOf)) { + for (const subSchema of schema.allOf) { + if (!isRecord(subSchema)) { + continue; + } + for (const key of getRequiredProperties(subSchema)) { + required.add(key); + } + } + } + return required; +} + +type Nullability = "allows" | "rejects" | "unknown"; + +function combineAll(states: Nullability[]): Nullability { + if (states.includes("rejects")) { + return "rejects"; + } + return states.includes("unknown") ? "unknown" : "allows"; +} + +function getNullability(schema: unknown): Nullability { + if (schema === true) { + return "allows"; + } + if (schema === false) { + return "rejects"; + } + if (!isRecord(schema)) { + return "unknown"; + } + + const constraints: Nullability[] = []; + if (Object.hasOwn(schema, "$ref")) { + constraints.push("unknown"); + } + if (typeof schema.type === "string") { + constraints.push(schema.type === "null" ? "allows" : "rejects"); + } else if (Array.isArray(schema.type)) { + constraints.push(schema.type.includes("null") ? "allows" : "rejects"); + } + if (Array.isArray(schema.enum)) { + constraints.push(schema.enum.includes(null) ? "allows" : "rejects"); + } + if (Object.hasOwn(schema, "const")) { + constraints.push(schema.const === null ? "allows" : "rejects"); + } + if (Array.isArray(schema.anyOf)) { + const states = schema.anyOf.map(getNullability); + constraints.push( + states.includes("allows") ? "allows" : states.includes("unknown") ? "unknown" : "rejects" + ); + } + if (Array.isArray(schema.oneOf)) { + const states = schema.oneOf.map(getNullability); + const allowedCount = states.filter((state) => state === "allows").length; + constraints.push( + states.includes("unknown") ? "unknown" : allowedCount === 1 ? "allows" : "rejects" + ); + } + if (Array.isArray(schema.allOf)) { + constraints.push(combineAll(schema.allOf.map(getNullability))); + } + if (Object.hasOwn(schema, "not")) { + const state = getNullability(schema.not); + constraints.push(state === "allows" ? "rejects" : state === "rejects" ? "allows" : "unknown"); + } + if (Object.hasOwn(schema, "if")) { + constraints.push("unknown"); + } + + return constraints.length === 0 ? "allows" : combineAll(constraints); +} + +function makeNullableSchema(schema: unknown): Record { + const annotations = isRecord(schema) + ? { + ...(typeof schema.title === "string" ? { title: schema.title } : {}), + ...(typeof schema.description === "string" ? { description: schema.description } : {}), + } + : {}; + return { ...annotations, anyOf: [schema, { type: "null" }] }; +} + +function widenSchemaNode(schema: unknown, inheritedRequired = new Set()): void { + if (!isRecord(schema)) { + return; + } + + const required = new Set([...inheritedRequired, ...getRequiredProperties(schema)]); + if (isRecord(schema.properties)) { + for (const [propertyName, propertySchema] of Object.entries(schema.properties)) { + const modelSchema = + !required.has(propertyName) && getNullability(propertySchema) === "rejects" + ? makeNullableSchema(propertySchema) + : propertySchema; + schema.properties[propertyName] = modelSchema; + widenSchemaNode(modelSchema); + } + } + + const items = schema.items; + if (Array.isArray(items)) { + for (const itemSchema of items) { + widenSchemaNode(itemSchema); + } + } else { + widenSchemaNode(items); + } + + for (const keyword of ["anyOf", "oneOf", "allOf"] as const) { + const branches = schema[keyword]; + if (Array.isArray(branches)) { + for (const branch of branches) { + widenSchemaNode(branch, required); + } + } + } +} + +export interface OptionalNullSchemaContract { + modelSchema: unknown; + strict: false | undefined; + restore: (value: unknown) => unknown; +} + +export function createOptionalNullSchemaContract(schema: unknown): OptionalNullSchemaContract { + if (containsReferenceKeyword(schema)) { + return { + modelSchema: structuredClone(schema), + strict: false, + restore: (value) => value, + }; + } + return { + modelSchema: widenOptionalPropertiesToNullable(schema), + strict: undefined, + restore: (value) => stripSyntheticNulls(schema, value), + }; +} + +/** + * Apply Mux's nullish optional-property convention to a third-party JSON Schema. + * The model contract only widens the source contract, so every provider can use it. + */ +export function widenOptionalPropertiesToNullable(schema: unknown): unknown { + const modelSchema = structuredClone(schema); + widenSchemaNode(modelSchema); + return modelSchema; +} + +function stripProperties( + value: Record, + properties: Record, + required: Set +): Record { + const stripped = { ...value }; + for (const [propertyName, propertySchema] of Object.entries(properties)) { + if (!(propertyName in stripped)) { + continue; + } + if ( + stripped[propertyName] === null && + !required.has(propertyName) && + getNullability(propertySchema) === "rejects" + ) { + delete stripped[propertyName]; + continue; + } + stripped[propertyName] = stripSyntheticNulls(propertySchema, stripped[propertyName]); + } + return stripped; +} + +function schemaAcceptsValue(schema: unknown, value: unknown): boolean { + if (schema === true) { + return true; + } + if (schema === false) { + return false; + } + return validateJsonSchemaSubset(schema, value).success; +} + +function stripMatchingUnionBranch( + schema: Record, + value: unknown, + inheritedRequired: ReadonlySet +): unknown { + for (const keyword of ["anyOf", "oneOf"] as const) { + const branches = schema[keyword]; + if (!Array.isArray(branches)) { + continue; + } + for (const branch of branches) { + if (schemaAcceptsValue(branch, value)) { + return value; + } + } + for (const branch of branches) { + const stripped = stripSyntheticNullsNode(branch, value, inheritedRequired); + if (schemaAcceptsValue(branch, stripped)) { + return stripped; + } + } + } + return null; +} + +function stripSyntheticNullsNode( + schema: unknown, + value: unknown, + inheritedRequired: ReadonlySet +): unknown { + if (!isRecord(schema) || schemaAcceptsValue(schema, value)) { + return value; + } + + const required = new Set([...inheritedRequired, ...getRequiredProperties(schema)]); + if (Array.isArray(value)) { + const itemSchema = schema.items; + const stripped = Array.isArray(itemSchema) + ? value.map((item, index) => stripSyntheticNullsNode(itemSchema[index], item, new Set())) + : value.map((item) => stripSyntheticNullsNode(itemSchema, item, new Set())); + return stripMatchingUnionBranch(schema, stripped, required) ?? stripped; + } + if (!isRecord(value)) { + return value; + } + + let stripped = { ...value }; + if (isRecord(schema.properties)) { + stripped = stripProperties(stripped, schema.properties, required); + } + if (Array.isArray(schema.allOf)) { + for (const subSchema of schema.allOf) { + stripped = stripSyntheticNullsNode(subSchema, stripped, required) as Record; + } + } + return stripMatchingUnionBranch(schema, stripped, required) ?? stripped; +} + +/** + * Restore a third-party executor contract after a model uses `null` to represent + * an omitted optional property. Explicit source-nullable values remain unchanged. + */ +export function stripSyntheticNulls(schema: unknown, value: unknown): unknown { + return stripSyntheticNullsNode(schema, value, new Set()); +} diff --git a/src/common/utils/tools/schemaSanitizer.test.ts b/src/common/utils/tools/schemaSanitizer.test.ts index 4d57c622cd..60ea7ec191 100644 --- a/src/common/utils/tools/schemaSanitizer.test.ts +++ b/src/common/utils/tools/schemaSanitizer.test.ts @@ -188,10 +188,14 @@ describe("schemaSanitizer", () => { required: ["content"], }; + const validate = (value: unknown) => ({ success: true as const, value }); + const mcpTool = { type: "dynamic", description: "MCP test tool", + strict: false, inputSchema: { + validate, // Simulate the jsonSchema getter that the MCP tool adapter creates get jsonSchema() { return jsonSchema; @@ -203,6 +207,10 @@ describe("schemaSanitizer", () => { const sanitized = sanitizeToolSchemaForOpenAI(mcpTool); const schema = getInputSchema(sanitized); + expect(sanitized.strict).toBe(false); + const sanitizedInputSchema = sanitized.inputSchema as { validate?: unknown }; + expect(sanitizedInputSchema.validate).toBe(validate); + // Unsupported properties should be stripped expect(schema.properties.content).toEqual({ type: "string" }); expect(schema.properties.count).toEqual({ type: "number" }); diff --git a/src/common/utils/tools/workflowReportPayload.ts b/src/common/utils/tools/workflowReportPayload.ts deleted file mode 100644 index 9069653bf2..0000000000 --- a/src/common/utils/tools/workflowReportPayload.ts +++ /dev/null @@ -1,137 +0,0 @@ -import { validateJsonSchemaSubset } from "@/common/utils/jsonSchemaSubset"; - -function isRecord(value: unknown): value is Record { - return value != null && typeof value === "object" && !Array.isArray(value); -} - -function getRequiredProperties(schema: Record): Set { - const required = new Set( - Array.isArray(schema.required) - ? schema.required.filter((key): key is string => typeof key === "string") - : [] - ); - if (Array.isArray(schema.allOf)) { - for (const subSchema of schema.allOf) { - if (!isRecord(subSchema)) { - continue; - } - for (const key of getRequiredProperties(subSchema)) { - required.add(key); - } - } - } - return required; -} - -function schemaAllowsNull(schema: unknown): boolean { - if (!isRecord(schema)) { - return true; - } - if (schema.type === "null") { - return true; - } - if (Array.isArray(schema.type)) { - return schema.type.includes("null"); - } - if (Array.isArray(schema.enum)) { - return schema.enum.includes(null); - } - for (const keyword of ["anyOf", "oneOf"] as const) { - const options = schema[keyword]; - if (Array.isArray(options) && options.some((option) => schemaAllowsNull(option))) { - return true; - } - } - return false; -} - -function normalizeProperties( - value: Record, - properties: Record, - required: Set -): Record { - const normalized = { ...value }; - for (const [propertyName, propertySchema] of Object.entries(properties)) { - if (!(propertyName in normalized)) { - continue; - } - if ( - normalized[propertyName] === null && - !required.has(propertyName) && - !schemaAllowsNull(propertySchema) - ) { - delete normalized[propertyName]; - continue; - } - normalized[propertyName] = normalizeWorkflowAgentReportPayloadForHostSchema( - propertySchema, - normalized[propertyName] - ); - } - return normalized; -} - -function normalizeMatchingUnionBranch( - schema: Record, - value: Record -): Record | null { - for (const keyword of ["anyOf", "oneOf"] as const) { - const options = schema[keyword]; - if (!Array.isArray(options)) { - continue; - } - for (const option of options) { - const normalized = normalizeWorkflowAgentReportPayloadForHostSchema(option, value); - if (!isRecord(normalized)) { - continue; - } - if (validateJsonSchemaSubset(option, normalized).success) { - return normalized; - } - } - } - return null; -} - -/** - * OpenAI strict tool schemas require every object property and represent originally-optional - * properties as nullable. For host validation/persistence, a `null` value for an optional - * non-nullable field means "the model omitted it", not an explicit workflow value. - */ -export function normalizeWorkflowAgentReportPayloadForHostSchema( - schema: unknown, - value: unknown -): unknown { - if (!isRecord(schema)) { - return value; - } - if (Array.isArray(value)) { - const itemSchema = schema.items; - if (Array.isArray(itemSchema)) { - return value.map((item, index) => - normalizeWorkflowAgentReportPayloadForHostSchema(itemSchema[index], item) - ); - } - return value.map((item) => normalizeWorkflowAgentReportPayloadForHostSchema(itemSchema, item)); - } - if (!isRecord(value)) { - return value; - } - - let normalized: Record = { ...value }; - const properties = isRecord(schema.properties) ? schema.properties : null; - if (properties != null) { - normalized = normalizeProperties(normalized, properties, getRequiredProperties(schema)); - } - - if (Array.isArray(schema.allOf)) { - for (const subSchema of schema.allOf) { - normalized = normalizeWorkflowAgentReportPayloadForHostSchema( - subSchema, - normalized - ) as Record; - } - } - - return normalizeMatchingUnionBranch(schema, normalized) ?? normalized; -} diff --git a/src/node/services/mcpClient.test.ts b/src/node/services/mcpClient.test.ts new file mode 100644 index 0000000000..f672c63457 --- /dev/null +++ b/src/node/services/mcpClient.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, test } from "bun:test"; + +import { createMCPToolContract, createMCPToolInputSchema } from "./mcpClient"; + +describe("createMCPToolInputSchema", () => { + test("exposes a nullable model contract and restores the server contract", async () => { + const inputSchema = createMCPToolInputSchema({ + type: "object", + required: ["issueId"], + properties: { + issueId: { type: "string" }, + cursor: { type: "string" }, + statusUpdateType: { type: "string", enum: ["project", "initiative"] }, + }, + additionalProperties: false, + }); + + expect(inputSchema.jsonSchema).toMatchObject({ + required: ["issueId"], + additionalProperties: false, + properties: { + issueId: { type: "string" }, + cursor: { anyOf: [{ type: "string" }, { type: "null" }] }, + statusUpdateType: { + anyOf: [{ type: "string", enum: ["project", "initiative"] }, { type: "null" }], + }, + }, + }); + + expect( + await inputSchema.validate?.({ + issueId: "CODAGT-709", + cursor: "", + statusUpdateType: null, + }) + ).toEqual({ + success: true, + value: { issueId: "CODAGT-709", cursor: "" }, + }); + }); + + test.each(["$ref", "$dynamicRef", "$recursiveRef"])( + "uses the non-strict fallback for schemas with %s", + async (keyword) => { + const source = { + type: "object", + properties: { value: { [keyword]: "#/$defs/value" } }, + $defs: { value: { type: "string" } }, + }; + const contract = createMCPToolContract(source); + + expect(contract.strict).toBe(false); + expect(contract.inputSchema.jsonSchema as Record).toEqual({ + ...source, + additionalProperties: false, + }); + expect(await contract.inputSchema.validate?.({ value: null })).toEqual({ + success: true, + value: { value: null }, + }); + } + ); + + test("does not close a composed root schema with synthetic empty properties", () => { + const inputSchema = createMCPToolInputSchema({ + type: "object", + allOf: [ + { + type: "object", + properties: { value: { type: "string" } }, + }, + ], + }); + + expect(inputSchema.jsonSchema).not.toHaveProperty("additionalProperties"); + expect(inputSchema.jsonSchema).not.toHaveProperty("properties"); + }); + + test("preserves dictionary schemas", () => { + const inputSchema = createMCPToolInputSchema({ + type: "object", + additionalProperties: { type: "string" }, + }); + + expect(inputSchema.jsonSchema).toEqual({ + type: "object", + properties: {}, + additionalProperties: { type: "string" }, + }); + }); +}); diff --git a/src/node/services/mcpClient.ts b/src/node/services/mcpClient.ts index d95bf7bff6..590449ba41 100644 --- a/src/node/services/mcpClient.ts +++ b/src/node/services/mcpClient.ts @@ -7,6 +7,7 @@ import { type Transport, } from "@modelcontextprotocol/client"; import { dynamicTool, jsonSchema, type JSONSchema7, type Tool } from "ai"; +import { createOptionalNullSchemaContract } from "@/common/utils/tools/optionalNullSchema"; import assert from "@/common/utils/assert"; /** @@ -203,6 +204,43 @@ function mcpToModelOutput({ return { type: "content", value: convertedContent }; } +function isJsonSchemaProperties(value: unknown): value is NonNullable { + return value != null && typeof value === "object" && !Array.isArray(value); +} + +function createMCPSourceSchema(inputSchema: Record | undefined): JSONSchema7 { + const sourceSchema: JSONSchema7 = { ...(inputSchema ?? { type: "object" }) }; + const hasComposition = ["$ref", "$dynamicRef", "$recursiveRef", "allOf", "anyOf", "oneOf"].some( + (key) => Object.hasOwn(sourceSchema, key) + ); + if (isJsonSchemaProperties(inputSchema?.properties)) { + sourceSchema.properties = inputSchema.properties; + sourceSchema.additionalProperties = inputSchema.additionalProperties ?? false; + } else if (!hasComposition || inputSchema?.additionalProperties !== undefined) { + sourceSchema.properties = {}; + sourceSchema.additionalProperties = inputSchema?.additionalProperties ?? false; + } + return sourceSchema; +} + +export function createMCPToolContract(inputSchema: Record | undefined) { + const contract = createOptionalNullSchemaContract(createMCPSourceSchema(inputSchema)); + return { + strict: contract.strict, + inputSchema: jsonSchema(contract.modelSchema as JSONSchema7, { + validate: (value) => ({ + success: true as const, + value: contract.restore(value), + }), + }), + }; +} + +/** Build the provider contract and restore the MCP server contract during input parsing. */ +export function createMCPToolInputSchema(inputSchema: Record | undefined) { + return createMCPToolContract(inputSchema).inputSchema; +} + /** * Connect to an MCP server and return a handle exposing AI SDK tools. * @@ -238,21 +276,12 @@ export async function createMCPClient(config: MCPClientConfig): Promise = {}; for (const definition of listResult.tools) { - const inputSchema = definition.inputSchema ?? { type: "object" }; + const contract = createMCPToolContract(definition.inputSchema); tools[definition.name] = dynamicTool({ description: definition.description, title: definition.title ?? definition.annotations?.title, - // Server-provided JSON schemas are runtime data; the SDK types them - // as loose JSON values, so narrow to the AI SDK's JSONSchema7 shape. - // Preserve a server-declared additionalProperties (map-style params - // like env vars or labels rely on it); default to false only when - // absent, matching the previous @ai-sdk/mcp behavior for strict - // providers. Provider-specific strictness stays in schemaSanitizer. - inputSchema: jsonSchema({ - ...inputSchema, - properties: inputSchema.properties ?? {}, - additionalProperties: inputSchema.additionalProperties ?? false, - } as JSONSchema7), + strict: contract.strict, + inputSchema: contract.inputSchema, execute: async (args: unknown, options: { abortSignal?: AbortSignal }) => { options.abortSignal?.throwIfAborted(); return await client.callTool( diff --git a/src/node/services/mcpResultTransform.test.ts b/src/node/services/mcpResultTransform.test.ts index ffdd2c03a5..96a40bfa4e 100644 --- a/src/node/services/mcpResultTransform.test.ts +++ b/src/node/services/mcpResultTransform.test.ts @@ -1,5 +1,28 @@ import { describe, it, expect } from "bun:test"; -import { transformMCPResult, MAX_IMAGE_DATA_BYTES } from "./mcpResultTransform"; +import { + describeMCPErrorResult, + transformMCPResult, + MAX_IMAGE_DATA_BYTES, +} from "./mcpResultTransform"; + +describe("describeMCPErrorResult", () => { + it("omits binary payloads from tool error messages", () => { + const description = describeMCPErrorResult({ + isError: true, + content: [ + { + type: "image", + data: "x".repeat(MAX_IMAGE_DATA_BYTES + 1), + mimeType: "image/png", + }, + ], + }); + + expect(description).toContain("Image omitted"); + expect(description).toContain("per-image guard"); + expect(description.length).toBeLessThan(200); + }); +}); describe("transformMCPResult", () => { describe("image data overflow handling", () => { @@ -109,38 +132,6 @@ describe("transformMCPResult", () => { expect(transformMCPResult("serena")).toBe("serena"); }); - it("should pass through text-only error results unchanged", () => { - const errorResult = { - isError: true, - content: [{ type: "text" as const, text: "Error!" }], - }; - expect(transformMCPResult(errorResult)).toBe(errorResult); - }); - - it("should convert binary content in error results and mark the error", () => { - // Error results carrying binary payloads must not bypass the media - // conversion (or the size guard); the error flag is surfaced as text. - const bigData = "x".repeat(9 * 1024 * 1024); - const errorResult = { - isError: true, - content: [ - { type: "text" as const, text: "capture failed" }, - { type: "image" as const, data: "abc123", mimeType: "image/png" }, - { type: "image" as const, data: bigData, mimeType: "image/png" }, - ], - }; - const result = transformMCPResult(errorResult) as { - type: string; - value: Array<{ type: string; text?: string; data?: string; mediaType?: string }>; - }; - expect(result.type).toBe("content"); - expect(result.value[0]).toEqual({ type: "text", text: "[Tool reported an error]" }); - expect(result.value[1]).toEqual({ type: "text", text: "capture failed" }); - expect(result.value[2]).toEqual({ type: "media", data: "abc123", mediaType: "image/png" }); - expect(result.value[3].type).toBe("text"); - expect(result.value[3].text).toContain("Image omitted"); - }); - it("should pass through toolResult unchanged", () => { const toolResult = { toolResult: { foo: "bar" } }; expect(transformMCPResult(toolResult)).toBe(toolResult); diff --git a/src/node/services/mcpResultTransform.ts b/src/node/services/mcpResultTransform.ts index ecaf70eac1..3d75d32d08 100644 --- a/src/node/services/mcpResultTransform.ts +++ b/src/node/services/mcpResultTransform.ts @@ -9,6 +9,7 @@ import { log } from "@/node/services/log"; * pass normal screenshots while preventing pathological payloads. */ export const MAX_IMAGE_DATA_BYTES = 8 * 1024 * 1024; // 8MB guard per image +const MAX_ERROR_DESCRIPTION_CHARACTERS = 64 * 1024; /** * MCP CallToolResult content types (MCP spec wire shapes) @@ -43,6 +44,65 @@ export interface MCPCallToolResult { toolResult?: unknown; } +export function isMCPErrorResult(value: unknown): value is MCPCallToolResult & { isError: true } { + return ( + value != null && typeof value === "object" && (value as { isError?: unknown }).isError === true + ); +} + +export function describeMCPErrorResult(result: MCPCallToolResult): string { + const readableParts = (result.content ?? []).flatMap((item) => { + if (item.type === "text") { + return item.text; + } + if (item.type === "resource") { + return item.resource.text ?? item.resource.uri; + } + return []; + }); + if (readableParts.length > 0) { + return readableParts.join("\n"); + } + + const binaryParts = (result.content ?? []).flatMap((item) => { + if (item.type === "image") { + return describeBinaryErrorPart("Image", item.data, item.mimeType); + } + if (item.type === "audio") { + return describeBinaryErrorPart("Audio", item.data, item.mimeType); + } + return []; + }); + if (binaryParts.length > 0) { + return binaryParts.join("\n"); + } + + return stringifyMCPErrorValue(result.toolResult ?? result.content ?? result); +} + +function stringifyMCPErrorValue(value: unknown): string { + try { + const serialized = JSON.stringify(value); + if (serialized == null) { + return "MCP tool call failed"; + } + if (serialized.length <= MAX_ERROR_DESCRIPTION_CHARACTERS) { + return serialized; + } + return `${serialized.slice(0, MAX_ERROR_DESCRIPTION_CHARACTERS)}\n[MCP error details truncated]`; + } catch { + return "MCP tool call failed"; + } +} + +function describeBinaryErrorPart(kind: string, data: string, mediaType: string): string { + const dataLength = data.length; + if (dataLength > MAX_IMAGE_DATA_BYTES) { + return `[${kind} omitted: ${formatBytesSI(dataLength)} exceeds per-${kind.toLowerCase()} guard of ${formatBytesSI(MAX_IMAGE_DATA_BYTES)}.]`; + } + return `[${kind} omitted from MCP error text: ${formatBytesSI(dataLength)}, ${mediaType}.]`; +} + /** * AI SDK LanguageModelV2ToolResultOutput content types */ @@ -107,9 +167,8 @@ export function transformMCPResult(result: unknown): unknown { return result; } - // Only rewrite results carrying binary payloads; text-only results - // (including text-only errors) pass through in MCP shape (converted by the - // tool's toModelOutput, which keeps the isError flag visible to the model). + // Only rewrite results carrying binary payloads. Text-only results pass + // through in MCP shape for the tool's toModelOutput conversion. const hasBinaryContent = typed.content.some( (c) => c.type === "image" || @@ -152,12 +211,5 @@ export function transformMCPResult(result: unknown): unknown { return { type: "text" as const, text: JSON.stringify(item) }; }); - // The model-output "content" shape has no error flag, so error results - // carrying binary payloads get an explicit text marker instead of bypassing - // the media conversion (and its size guard). - if (typed.isError) { - transformedContent.unshift({ type: "text", text: "[Tool reported an error]" }); - } - return { type: "content", value: transformedContent }; } diff --git a/src/node/services/mcpServerManager.test.ts b/src/node/services/mcpServerManager.test.ts index bf09028fa9..e06132441d 100644 --- a/src/node/services/mcpServerManager.test.ts +++ b/src/node/services/mcpServerManager.test.ts @@ -6,6 +6,7 @@ import * as path from "node:path"; import * as mcpSdk from "@/node/services/mcpClient"; import { MCPServerManager, + MCPToolCallError, isClosedClientError, prepareStdioLaunch, runMCPToolWithDeadline, @@ -1592,6 +1593,66 @@ describe("isClosedClientError", () => { }); describe("wrapMCPTools", () => { + test("converts MCP application failures into tool errors without recycling the client", async () => { + const onActivity = mock(() => undefined); + const onClosed = mock(() => undefined); + const tool = { + execute: mock(() => + Promise.resolve({ + isError: true, + content: [ + { type: "text", text: "statusUpdateType requires statusUpdateId" }, + { + type: "resource", + resource: { uri: "linear://issue/CODAGT-709", text: "Linear rejected the call" }, + }, + ], + }) + ), + parameters: {}, + } as unknown as Tool; + + const wrapped = wrapMCPTools({ myTool: tool }, { onActivity, onClosed }); + + let caught: unknown; + try { + await wrapped.myTool.execute!({}, {} as never); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(MCPToolCallError); + expect((caught as Error).message).toContain("statusUpdateType requires statusUpdateId"); + expect((caught as Error).message).toContain("Linear rejected the call"); + expect(onActivity).toHaveBeenCalledTimes(1); + expect(onClosed).not.toHaveBeenCalled(); + }); + + test("does not recycle the client when application error text resembles a closed client", async () => { + const onClosed = mock(() => undefined); + const tool = { + execute: mock(() => + Promise.resolve({ + isError: true, + content: [{ type: "text", text: "Not connected to the external account" }], + }) + ), + parameters: {}, + } as unknown as Tool; + + const wrapped = wrapMCPTools({ myTool: tool }, { onClosed }); + + let caught: unknown; + try { + await Promise.resolve(wrapped.myTool.execute!({}, {} as never)); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(MCPToolCallError); + expect(onClosed).not.toHaveBeenCalled(); + }); + for (const [message, expectedOnClosedCalls] of [ ["Attempted to send a request from a closed client", 1], ["some other failure", 0], diff --git a/src/node/services/mcpServerManager.ts b/src/node/services/mcpServerManager.ts index e5396bcab8..caf17cfb4d 100644 --- a/src/node/services/mcpServerManager.ts +++ b/src/node/services/mcpServerManager.ts @@ -34,7 +34,12 @@ import { type McpOauthService, } from "@/node/services/mcpOauthService"; import { createRuntime } from "@/node/runtime/runtimeFactory"; -import { transformMCPResult, type MCPCallToolResult } from "@/node/services/mcpResultTransform"; +import { + describeMCPErrorResult, + isMCPErrorResult, + transformMCPResult, + type MCPCallToolResult, +} from "@/node/services/mcpResultTransform"; import { buildMcpToolName } from "@/common/utils/tools/mcpToolName"; import { getErrorMessage } from "@/common/utils/errors"; @@ -80,6 +85,13 @@ class MCPDeadlineError extends Error { } } +export class MCPToolCallError extends Error { + constructor(message: string) { + super(message); + this.name = "MCPToolCallError"; + } +} + /** * Wraps errors raised while connecting a freshly-spawned stdio MCP client. * Typed so the negotiation retry loop can distinguish "the connect (possibly @@ -163,6 +175,9 @@ export async function runMCPToolWithDeadline( } function shouldRecycleClientAfterToolError(error: unknown): boolean { + if (error instanceof MCPToolCallError) { + return false; + } return isClosedClientError(error) || error instanceof MCPDeadlineError; } @@ -201,6 +216,9 @@ export function wrapMCPTools( () => Promise.resolve(originalExecute(args, context)) as Promise, { toolName, timeoutMs: MCP_TOOL_CALL_TIMEOUT_MS, signal: abortSignal } ); + if (isMCPErrorResult(result)) { + throw new MCPToolCallError(describeMCPErrorResult(result)); + } return transformMCPResult(result as MCPCallToolResult); } catch (error) { if (shouldRecycleClientAfterToolError(error)) { diff --git a/src/node/services/taskService.ts b/src/node/services/taskService.ts index 0016d6f244..010fdda873 100644 --- a/src/node/services/taskService.ts +++ b/src/node/services/taskService.ts @@ -158,7 +158,7 @@ import { } from "@/node/services/terminalAttentionStore"; import { readAgentWorkflowRunReferences } from "@/node/services/agentWorkflowRunReferences"; import { isWorkflowRunTaskId } from "@/node/services/tools/taskId"; -import { normalizeWorkflowAgentReportPayloadForHostSchema } from "@/common/utils/tools/workflowReportPayload"; +import { stripSyntheticNulls } from "@/common/utils/tools/optionalNullSchema"; import { formatJsonSchemaValidationErrors, validateJsonSchemaSubset, @@ -425,10 +425,7 @@ function normalizeWorkflowAgentReportArgsForWorkflowTask( } return { ...reportArgs, - structuredOutput: normalizeWorkflowAgentReportPayloadForHostSchema( - workflowTask.outputSchema, - reportArgs.structuredOutput - ), + structuredOutput: stripSyntheticNulls(workflowTask.outputSchema, reportArgs.structuredOutput), }; } @@ -456,7 +453,7 @@ function validateWorkflowAgentReportStructuredOutput(params: { }); } - const structuredOutput = normalizeWorkflowAgentReportPayloadForHostSchema( + const structuredOutput = stripSyntheticNulls( workflowTask.outputSchema, params.reportArgs.structuredOutput ); diff --git a/src/node/services/tools/agent_report.ts b/src/node/services/tools/agent_report.ts index 391ebba553..a6ca93a12e 100644 --- a/src/node/services/tools/agent_report.ts +++ b/src/node/services/tools/agent_report.ts @@ -6,7 +6,7 @@ import { validateJsonSchemaSubsetSchema, type JsonSchemaValidationError, } from "@/common/utils/jsonSchemaSubset"; -import { normalizeWorkflowAgentReportPayloadForHostSchema } from "@/common/utils/tools/workflowReportPayload"; +import { stripSyntheticNulls } from "@/common/utils/tools/optionalNullSchema"; import { sanitizeWorkflowAgentReportSchemaForOpenAI } from "@/common/utils/tools/schemaSanitizer"; import type { ToolConfiguration, ToolFactory } from "@/common/utils/tools/tools"; import { @@ -80,10 +80,7 @@ function validateStructuredOutput(config: ToolConfiguration, structuredOutput: u return null; } - const normalizedOutput = normalizeWorkflowAgentReportPayloadForHostSchema( - outputSchema, - structuredOutput - ); + const normalizedOutput = stripSyntheticNulls(outputSchema, structuredOutput); const validation = validateJsonSchemaSubset(outputSchema, normalizedOutput); return validation.success ? null @@ -103,7 +100,7 @@ function buildInlineInputSchema(config: ToolConfiguration) { ) as JSONSchema7; return jsonSchema(providerFacingSchema, { validate: (value) => { - const normalizedValue = normalizeWorkflowAgentReportPayloadForHostSchema(outputSchema, value); + const normalizedValue = stripSyntheticNulls(outputSchema, value); const validation = validateStructuredOutput(config, normalizedValue); if (validation) { return { success: false, error: new Error(validation.message) }; @@ -119,10 +116,7 @@ function parseProgressReport( ): { report: AgentProgressReport } | { failure: AgentReportFailureResult } { const workflowOutputSchema = getWorkflowAgentOutputSchema(config); if (workflowOutputSchema != null) { - const normalizedArgs = normalizeWorkflowAgentReportPayloadForHostSchema( - workflowOutputSchema, - rawArgs - ); + const normalizedArgs = stripSyntheticNulls(workflowOutputSchema, rawArgs); const structuredValidation = validateStructuredOutput(config, normalizedArgs); if (structuredValidation) { return { failure: structuredValidation };