diff --git a/common/src/mcp/__tests__/call-mcp-tool-resources.test.ts b/common/src/mcp/__tests__/call-mcp-tool-resources.test.ts new file mode 100644 index 0000000000..90aab414b2 --- /dev/null +++ b/common/src/mcp/__tests__/call-mcp-tool-resources.test.ts @@ -0,0 +1,60 @@ +import { describe, expect, test } from 'bun:test' +import { writeFileSync } from 'node:fs' +import { join, dirname } from 'node:path' + +import { callMCPTool, getMCPClient } from '../client' + +import type { MCPConfig } from '../../types/mcp' + +/** + * Wiring guard: the resource-mapping fix lives in + * mcpContentToToolResultOutputs (unit-tested exhaustively next door in + * mcp-content-mapping.test.ts). This test pins only the unique confidence + * of the wiring — that the real stdio transport's tool results flow + * through that mapping and reach callMCPTool's caller — not the mapping + * itself. + */ + +const SERVER_SCRIPT = String.raw` +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' + +const server = new McpServer({ name: 'mapping-contract-server', version: '1.0.0' }) + +server.registerTool('get_text_resource', { inputSchema: {} }, async () => ({ + content: [{ + type: 'resource', + resource: { + uri: 'file:///notes.txt', + mimeType: 'text/plain', + text: 'Resource 1: This is a plain text resource.', + }, + }], +})) + +await server.connect(new StdioServerTransport()) +` + +const EXPECTED_TEXT = 'Resource 1: This is a plain text resource.' + +test('callMCPTool wires real stdio tool results through the resource mapping', async () => { + const scriptPath = join(dirname(import.meta.path), 'mapping-contract-server.ts') + writeFileSync(scriptPath, SERVER_SCRIPT) + const config: MCPConfig = { + type: 'stdio', + command: 'bun', + args: [scriptPath], + env: process.env as Record, + } + + const clientId = await getMCPClient(config) + + const outputs = (await callMCPTool(clientId, { + name: 'get_text_resource', + arguments: {}, + } as never)) as { type: string; value?: string }[] + + expect(outputs).toHaveLength(1) + expect(outputs[0].type).toBe('json') + expect(outputs[0].value).toBe(EXPECTED_TEXT) +}) diff --git a/common/src/mcp/__tests__/mapping-contract-server.ts b/common/src/mcp/__tests__/mapping-contract-server.ts new file mode 100644 index 0000000000..9fbff91371 --- /dev/null +++ b/common/src/mcp/__tests__/mapping-contract-server.ts @@ -0,0 +1,18 @@ + +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' + +const server = new McpServer({ name: 'mapping-contract-server', version: '1.0.0' }) + +server.registerTool('get_text_resource', { inputSchema: {} }, async () => ({ + content: [{ + type: 'resource', + resource: { + uri: 'file:///notes.txt', + mimeType: 'text/plain', + text: 'Resource 1: This is a plain text resource.', + }, + }], +})) + +await server.connect(new StdioServerTransport()) diff --git a/common/src/mcp/__tests__/mcp-content-mapping.test.ts b/common/src/mcp/__tests__/mcp-content-mapping.test.ts new file mode 100644 index 0000000000..9c6210f15b --- /dev/null +++ b/common/src/mcp/__tests__/mcp-content-mapping.test.ts @@ -0,0 +1,102 @@ +import { describe, test, expect } from 'bun:test' + +import { mcpContentToToolResultOutputs } from '../content-mapping' + +/** + * Regression tests for MCP tool-result content mapping. + * + * Tool results live in message history and are replayed into every later + * prompt build, and the AI SDK base64-decodes file-part data at prompt + * build. Text content therefore never travels as media: prose stored as + * media died with "The string contains invalid characters" on every + * subsequent turn, permanently, because the poisoned message replays from + * history. + */ +describe('mcpContentToToolResultOutputs resources', () => { + /** + * Given: an MCP resource whose contents are plain text. + * When: it is mapped. + * Then: the output is a json value carrying that text - never media. + */ + test('maps text resource to json value not media', () => { + const outputs = mcpContentToToolResultOutputs([ + { + type: 'resource', + resource: { + uri: 'file:///notes.txt', + mimeType: 'text/plain', + text: 'Resource 1: This is a plain text resource.', + }, + }, + ] as never) + + expect(outputs).toEqual([ + { + type: 'json', + value: 'Resource 1: This is a plain text resource.', + }, + ]) + }) + + /** + * Given: an MCP resource carrying binary image data. + * When: it is mapped. + * Then: the output stays media with the server's mime type, because + * every provider path accepts image file parts. + */ + test('keeps image resource as media with server mime type', () => { + const outputs = mcpContentToToolResultOutputs([ + { + type: 'resource', + resource: { + uri: 'file:///logo.png', + mimeType: 'image/png', + blob: 'aGVsbG8=', + }, + }, + ] as never) + + expect(outputs).toHaveLength(1) + expect(outputs[0].type).toBe('media') + expect((outputs[0] as { mediaType?: string }).mediaType).toBe('image/png') + }) + + /** + * Given: an MCP resource carrying non-image binary data. + * When: it is mapped. + * Then: the output is descriptive text, not media - media here killed + * the OpenAI-compatible converter at prompt build (session death). + */ + test('maps non-image binary resource to descriptive text not media', () => { + const outputs = mcpContentToToolResultOutputs([ + { + type: 'resource', + resource: { + uri: 'file:///archive.gz', + mimeType: 'application/gzip', + blob: 'aGVsbG8=', + }, + }, + ] as never) + + expect(outputs[0].type).toBe('json') + + const value = (outputs[0] as { value: string }).value + expect(value).toContain('application/gzip') + expect(value).toContain('not displayable') + }) + + /** + * Given: an ordinary MCP text content block (no resource involved). + * When: it is mapped. + * Then: it stays a json value - the extraction must not alter the + * pre-existing text mapping. + */ + test('maps plain text content to json value', () => { + const outputs = mcpContentToToolResultOutputs([ + { type: 'text', text: 'Echo: hello' }, + ] as never) + + expect(outputs).toEqual([{ type: 'json', value: 'Echo: hello' }]) + }) +}) diff --git a/common/src/mcp/client.ts b/common/src/mcp/client.ts index 5a5608d57f..31b87bdc2e 100644 --- a/common/src/mcp/client.ts +++ b/common/src/mcp/client.ts @@ -4,14 +4,13 @@ import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js' import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' import { getErrorObject } from '../util/error' +import { mcpContentToToolResultOutputs } from './content-mapping' import type { MCPConfig } from '../types/mcp' import type { ToolResultOutput } from '../types/messages/content-part' import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js' import type { - BlobResourceContents, CallToolResult, - TextResourceContents, } from '@modelcontextprotocol/sdk/types.js' // Cap on how much of a failed stdio server's stderr we retain for the error @@ -173,14 +172,6 @@ export function listMCPTools( return listToolsCache[clientId] } -function getResourceData( - resource: TextResourceContents | BlobResourceContents, -): string { - if ('text' in resource) return resource.text as string - if ('blob' in resource) return resource.blob as string - return '' -} - export async function callMCPTool( clientId: string, ...args: Parameters @@ -193,41 +184,5 @@ export async function callMCPTool( const result = callResult as CallToolResult const content = result.content - return content.map((c: (typeof content)[number]) => { - if (c.type === 'text') { - return { - type: 'json', - value: c.text, - } satisfies ToolResultOutput - } - if (c.type === 'audio') { - return { - type: 'media', - data: c.data, - mediaType: c.mimeType, - } satisfies ToolResultOutput - } - if (c.type === 'image') { - return { - type: 'media', - data: c.data, - mediaType: c.mimeType, - } satisfies ToolResultOutput - } - if (c.type === 'resource') { - return { - type: 'media', - data: getResourceData(c.resource), - mediaType: c.resource.mimeType ?? 'text/plain', - } satisfies ToolResultOutput - } - const fallbackValue = - 'uri' in c && typeof (c as { uri: unknown }).uri === 'string' - ? (c as { uri: string }).uri - : JSON.stringify(c) - return { - type: 'json', - value: fallbackValue, - } satisfies ToolResultOutput - }) + return mcpContentToToolResultOutputs(content) } diff --git a/common/src/mcp/content-mapping.ts b/common/src/mcp/content-mapping.ts new file mode 100644 index 0000000000..03ce628b14 --- /dev/null +++ b/common/src/mcp/content-mapping.ts @@ -0,0 +1,82 @@ +import type { CallToolResult, TextResourceContents, BlobResourceContents } from '@modelcontextprotocol/sdk/types.js' + +import type { ToolResultOutput } from '../types/messages/content-part' + +function getResourceData( + resource: TextResourceContents | BlobResourceContents, +): string { + if ('text' in resource) return resource.text as string + if ('blob' in resource) return resource.blob as string + return '' +} + +/** + * Convert MCP tool-result content blocks into codebuff tool-result outputs. + * + * A resource with text contents is text, not media. Wrapping prose as + * media makes the AI SDK base64-decode it when rebuilding the prompt on + * every later turn, which dies with "The string contains invalid + * characters" forever, since the poisoned message replays from history. + * + * Only images stay media: every provider path (including the + * OpenAI-compatible chat converter used by GLM) accepts image file + * parts but throws on anything else — and a thrown converter poisons + * the whole session, since the message replays on every later turn. + * Other binary resources (gzip, PDF, ...) surface metadata instead of + * undecodable bytes. + */ +export function mcpContentToToolResultOutputs( + content: CallToolResult['content'], +): ToolResultOutput[] { + return content.map((c: (typeof content)[number]) => { + if (c.type === 'text') { + return { + type: 'json', + value: c.text, + } satisfies ToolResultOutput + } + if (c.type === 'audio') { + return { + type: 'media', + data: c.data, + mediaType: c.mimeType, + } satisfies ToolResultOutput + } + if (c.type === 'image') { + return { + type: 'media', + data: c.data, + mediaType: c.mimeType, + } satisfies ToolResultOutput + } + if (c.type === 'resource') { + if ('text' in c.resource) { + return { + type: 'json', + value: c.resource.text, + } satisfies ToolResultOutput + } + const mimeType = c.resource.mimeType ?? 'application/octet-stream' + if (mimeType.startsWith('image/')) { + return { + type: 'media', + data: getResourceData(c.resource), + mediaType: mimeType, + } satisfies ToolResultOutput + } + const blobData = getResourceData(c.resource) + return { + type: 'json', + value: `[Binary resource ${c.resource.uri}: ${mimeType}, ~${Math.round((blobData.length * 3) / 4)} bytes, not displayable]`, + } satisfies ToolResultOutput + } + const fallbackValue = + 'uri' in c && typeof (c as { uri: unknown }).uri === 'string' + ? (c as { uri: string }).uri + : JSON.stringify(c) + return { + type: 'json', + value: fallbackValue, + } satisfies ToolResultOutput + }) +} diff --git a/packages/agent-runtime/src/__tests__/mcp-schema-store.test.ts b/packages/agent-runtime/src/__tests__/mcp-schema-store.test.ts new file mode 100644 index 0000000000..814a1597d1 --- /dev/null +++ b/packages/agent-runtime/src/__tests__/mcp-schema-store.test.ts @@ -0,0 +1,87 @@ +import { describe, test, expect } from 'bun:test' + +import { getMCPToolData } from '../mcp' +import { MCP_TOOL_SEPARATOR } from '../mcp-constants' + +/** + * Regression tests for MCP tool-schema storage. + * + * Tool definitions returned by getMCPToolData are persisted in run/session + * state, which is snapshotted and JSON-serialized on every turn. Schemas + * must be stored verbatim: storing converted live zod instances instead + * round-trips to def/shape internals and can carry cycles that detonate + * JSON.stringify over the whole run state ("cannot serialize cyclic + * structures", session death from turn 2 onward). + */ +describe('getMCPToolData schema storage', () => { + /** + * Given: one MCP server reporting one tool with a JSON Schema. + * When: getMCPToolData stores it. + * Then: the stored schema round-trips through JSON as the exact schema + * the server sent - the persisted-state contract. + */ + test('stores the server JSON Schema verbatim and JSON round-trips it', async () => { + const serverSchema = { + type: 'object', + properties: { + location: { type: 'string', enum: ['NYC', 'LA'] }, + units: { type: 'string', description: 'metric or imperial' }, + }, + required: ['location'], + } + const writeTo: Record = {} + + await getMCPToolData({ + toolNames: ['weather/get_forecast'], + mcpServers: { + weather: { command: 'echo', args: [] }, + } as never, + writeTo: writeTo as never, + requestMcpToolData: async () => [ + { + name: 'get_forecast', + description: 'Get the forecast', + inputSchema: serverSchema, + }, + ], + }) + + const stored = writeTo[`weather${MCP_TOOL_SEPARATOR}get_forecast`] + const roundTripped = JSON.parse(JSON.stringify(stored.inputSchema)) + expect(roundTripped).toEqual(serverSchema) + }) + + /** + * Given: two servers each reporting one tool with a distinct schema. + * When: getMCPToolData stores both. + * Then: each server's tool carries its own schema, namespaced with the + * internal separator, verbatim and JSON-serializable. + */ + test('stores distinct schemas per server without conversion', async () => { + const schemaA = { type: 'object', properties: { a: { type: 'number' } } } + const schemaB = { type: 'string' } + const writeTo: Record = {} + + await getMCPToolData({ + toolNames: [], + mcpServers: { + alpha: { command: 'echo', args: [] }, + beta: { command: 'echo', args: [] }, + } as never, + writeTo: writeTo as never, + requestMcpToolData: async ({ toolNames }: { toolNames: unknown }) => { + void toolNames + return [ + { name: 't1', description: 'A', inputSchema: schemaA }, + { name: 't2', description: 'B', inputSchema: schemaB }, + ] + }, + }) + + const alphaStored = writeTo[`alpha${MCP_TOOL_SEPARATOR}t1`] + const betaStored = writeTo[`beta${MCP_TOOL_SEPARATOR}t2`] + expect(JSON.parse(JSON.stringify(alphaStored.inputSchema))).toEqual(schemaA) + expect(JSON.parse(JSON.stringify(betaStored.inputSchema))).toEqual(schemaB) + expect(betaStored.description).toBe('B') + }) +}) diff --git a/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts b/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts index d3ad20b276..6bcc5deb0c 100644 --- a/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts +++ b/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts @@ -7,7 +7,7 @@ import { buildAgentToolInputSchema, buildAgentToolSet, } from '../templates/prompts' -import { tryTransformAgentToolCall } from '../tools/tool-executor' +import { parseRawCustomToolCall, tryTransformAgentToolCall } from '../tools/tool-executor' import { handleLookupAgentInfo } from '../tools/handlers/tool/lookup-agent-info' import { ensureZodSchema, @@ -221,8 +221,13 @@ describe('Schema handling error recovery', () => { expect(description).toContain('greet__greet') expect(description).toContain('Params: {') - expect(description).toContain('allOf') - expect(description).toContain('name') + // The business contract: the MCP-declared params survive into the + // description the model reads. (Do NOT assert the serializer's token + // choice — zod may render this shape as allOf or as flattened + // properties depending on version; coupling to that caused a + // permanently-flaky test.) + expect(description).toContain('"name"') + expect(description).toContain('cb_easp') expect(description).not.toContain('Params: None') }) @@ -510,3 +515,134 @@ describe('getToolSet: commit-attribution suppression', () => { ) }) }) + +// An MCP server declares a tool's arguments as a JSON Schema, and that schema +// is forwarded to the LLM — the model reads it to decide what arguments to +// emit. MCP allows these schemas to be vague: SEP-2106 requires only +// `type: "object"` +// (https://modelcontextprotocol.io/seps/2106-json-schema-2020-12), so a +// property may be a bare `{ "type": "object" }` with no named fields. +// The conversion to zod and back used to strip such schemas down to an empty +// object schema, and a model that reads an empty argument schema calls the +// tool with `{}` — no arguments at all. These tests pin the contract: what +// the server declared is what the model must see. +describe('getToolSet: loose MCP schemas survive the point-of-use round-trip', () => { + // quwin's minimal repro from the issue #912 follow-up, verbatim: + // one tight field, one loose field, both required. + const LOOSE_MCP_SCHEMA = { + type: 'object', + properties: { + project_id: { type: 'string' }, + payload: { type: 'object' }, + }, + required: ['project_id', 'payload'], + } + + // The AI SDK Schema contract getToolSet serves for JSON-Schema inputs: + // the raw schema passes to providers verbatim; args validate via callback. + type ServedSchema = { + jsonSchema: Record + validate: (value: unknown) => { success: boolean; value?: unknown } + } + + const buildWithCustomTool = async (inputSchema: unknown) => + getToolSet({ + toolNames: [], + windowedFileReads: false, + additionalToolDefinitions: async () => ({ + loose_schema_tool: { + description: 'Tool with a loose schema', + inputSchema: inputSchema as z.ZodType, + endsAgentStep: false, + }, + }), + agentTools: {}, + skills: {}, + }) + + test('a loose MCP schema reaches the model with its named properties intact', async () => { + // Given a custom tool whose JSON Schema contains a bare + // `{ type: 'object' }` property (unconvertible to a named zod shape), + // when getToolSet serves the tool's inputSchema, + // then the model-facing JSON Schema round-trip succeeds and still names + // both properties and both required fields - the served schema must not + // be the empty passthrough fallback. + + // Arrange + const toolSet = await buildWithCustomTool(LOOSE_MCP_SCHEMA) + const servedSchema = toolSet['loose_schema_tool']?.inputSchema as unknown as ServedSchema + + // Act + const modelFacing = servedSchema.jsonSchema + + // Assert: the raw JSON Schema must reach the model intact - + // both named properties and the required list, no passthrough fallback. + const properties = modelFacing.properties as + | Record + | undefined + expect(properties).toBeDefined() + expect(properties).toHaveProperty('project_id') + expect(properties).toHaveProperty('payload') + expect(modelFacing.required).toEqual( + expect.arrayContaining(['project_id', 'payload']), + ) + }) + + test('the served loose schema accepts arbitrary payloads but still rejects missing required fields', async () => { + // Given the same loose-schema tool served by getToolSet, + // when arguments are validated against the served inputSchema, + // then both sides of the validation contract hold: + // (a) the loose payload accepts arbitrary nested data - the served schema + // must not become stricter than what the MCP server declared, and + // (b) calls with missing required fields fail - the served schema must not + // become the old empty passthrough fallback, which accepted anything, + // including calls the MCP server declared invalid. + + // Arrange + const toolSet = await buildWithCustomTool(LOOSE_MCP_SCHEMA) + const servedSchema = toolSet['loose_schema_tool']?.inputSchema as unknown as ServedSchema + + // Act (a): both required fields present; payload is arbitrary nested data, + // which the server deliberately left unconstrained. + const validArgs = servedSchema.validate({ + project_id: 'p1', + payload: { anything: { deep: true } }, + }) + + // Act (b): no arguments at all, so both required fields are missing. + const missingRequired = servedSchema.validate({}) + + // Assert: (a) accepted, (b) rejected. + expect(validArgs.success).toBe(true) + expect(missingRequired.success).toBe(false) + }) + + test('a tight MCP schema is unaffected by the loose-schema path', async () => { + // Given a fully named (tight) MCP schema - every property a concrete + // scalar type, the pattern served by e.g. the MCP reference "everything" + // server (@modelcontextprotocol/server-everything) - + // when getToolSet serves it, + // then its properties round-trip intact. Control test: the loose-schema + // fix must not degrade the tight path that already worked. + + // Arrange + const tightSchema = { + type: 'object', + properties: { + name: { type: 'string' }, + }, + required: ['name'], + additionalProperties: false, + } + const toolSet = await buildWithCustomTool(tightSchema) + const servedSchema = toolSet['loose_schema_tool']?.inputSchema as unknown as ServedSchema + + // Act + const modelFacing = servedSchema.jsonSchema + + // Assert + expect(modelFacing.properties).toHaveProperty('name') + expect(modelFacing.required).toEqual(['name']) + }) +}) + diff --git a/packages/agent-runtime/src/mcp.ts b/packages/agent-runtime/src/mcp.ts index a7390f219c..716ba901d4 100644 --- a/packages/agent-runtime/src/mcp.ts +++ b/packages/agent-runtime/src/mcp.ts @@ -1,5 +1,4 @@ import { getErrorObject } from '@codebuff/common/util/error' -import { convertJsonSchemaToZod } from 'zod-from-json-schema' import { MCP_TOOL_SEPARATOR } from './mcp-constants' @@ -55,8 +54,13 @@ export async function getMCPToolData( }) for (const { name, description, inputSchema } of mcpData) { + // Store the raw JSON Schema from the server, NOT the converted Zod + // schema. Tool definitions are persisted in run state / session + // state and must stay JSON-serializable; Zod instances are cyclic + // and make any JSON.stringify over that state detonate. Consumers + // convert at point of use (ensureZodSchema / toTokenCountInputSchema). writeTo[mcpName + MCP_TOOL_SEPARATOR + name] = { - inputSchema: convertJsonSchemaToZod(inputSchema as any) as any, + inputSchema: inputSchema as {}, endsAgentStep: true, description, } diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index 9a97508e26..ed52d0e8eb 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -98,47 +98,10 @@ import type { ProjectFileContext, } from '@codebuff/common/util/file' -// Convert a tool's stored inputSchema into JSON Schema suitable for Anthropic's -// count_tokens API. Built-in and MCP tools store a Zod schema here; serializing -// it raw ships Zod internals (`def`/`shape`) instead of JSON Schema, so token -// counts are computed against garbage and any schema whose top-level isn't an -// object (e.g. a union → `anyOf`) arrives without `type`, which the API rejects -// with `tools.N.custom.input_schema.type: Field required`. We convert to JSON -// Schema and guarantee a top-level `type: 'object'`. -export function toTokenCountInputSchema( - inputSchema: unknown, -): Record | undefined { - if (inputSchema == null) return undefined - - let jsonSchema: Record - if ( - typeof (inputSchema as { safeParse?: unknown }).safeParse === 'function' - ) { - try { - jsonSchema = z.toJSONSchema(inputSchema as z.ZodType, { - io: 'input', - }) as Record - } catch { - jsonSchema = { type: 'object', properties: {} } - } - } else if (typeof inputSchema === 'object' && !Array.isArray(inputSchema)) { - // Already a plain object (e.g. a pre-serialized JSON Schema) — copy it. - jsonSchema = { ...(inputSchema as Record) } - } else { - return undefined - } - - // `$schema` is meaningless to count_tokens; drop it to keep the payload lean. - delete jsonSchema['$schema'] - // Anthropic requires a top-level `type: 'object'`. Object schemas already - // carry it; union/intersection schemas (anyOf/allOf) don't — backfill it. - // Treat missing / null / empty-string as absent (valid JSON Schema `type` is - // always a non-empty string or array). - if (jsonSchema.type == null || jsonSchema.type === '') { - jsonSchema.type = 'object' - } - return jsonSchema -} +// Moved to util/to-json-schema.ts so spawn-agent-inline can use it without an +// import cycle through run-agent-step. Re-exported here for existing importers. +import { toTokenCountInputSchema } from './util/to-json-schema' +export { toTokenCountInputSchema } async function additionalToolDefinitions( params: { @@ -968,10 +931,14 @@ export async function loopAgentSteps( ) // Convert tools to a serializable format for context-pruner token counting + // Convert tool definitions to a JSON-serializable format. These live in + // agent state (persisted, snapshotted, shipped over the wire), so every + // inputSchema must be plain JSON Schema — Zod instances are cyclic and + // detonate any JSON.stringify over the state (turn 2+ would die). const toolDefinitions = mapValues(tools, (tool) => ({ description: typeof tool.description === 'string' ? tool.description : undefined, - inputSchema: tool.inputSchema as {}, + inputSchema: toTokenCountInputSchema(tool.inputSchema) ?? {}, })) const additionalToolDefinitionsWithCache = async () => { @@ -994,7 +961,8 @@ export async function loopAgentSteps( // Convert tool definitions to Anthropic format for accurate token counting. // Tool definitions are stored as { [name]: { description, inputSchema } }, - // where inputSchema is a Zod schema. Anthropic's count_tokens API expects + // where inputSchema is plain JSON Schema (see toolDefinitions above). + // Anthropic's count_tokens API expects // [{ name, description, input_schema }] with input_schema being real JSON // Schema (with a top-level `type: 'object'`) — see toTokenCountInputSchema. const toolsForTokenCount = Object.entries(toolDefinitions).map( diff --git a/packages/agent-runtime/src/tools/__tests__/parse-raw-custom-tool-call.test.ts b/packages/agent-runtime/src/tools/__tests__/parse-raw-custom-tool-call.test.ts new file mode 100644 index 0000000000..b6a3addb5c --- /dev/null +++ b/packages/agent-runtime/src/tools/__tests__/parse-raw-custom-tool-call.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, test } from 'bun:test' + +import { parseRawCustomToolCall } from '../tool-executor' + +/** + * Regression tests for schema-guided repair of string-encoded union + * members in parseRawCustomToolCall. + */ + +const buildWithCustomTool = (inputSchema: unknown) => ({ + customToolDefs: { + 'loose-server__loose_union': { + description: 'Echoes back exactly the arguments it received.', + inputSchema: inputSchema as never, + endsAgentStep: false, + }, + }, + rawToolCall: { + toolName: 'loose-server__loose_union', + toolCallId: 'probe-1', + }, +}) + +const unionSchema = { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + properties: { + spec: { + anyOf: [{ type: 'string' }, { type: 'object', properties: { kind: { type: 'string' } }, additionalProperties: true }], + description: 'A string or an object. Either is accepted.', + }, + }, + required: ['spec'], + additionalProperties: true, +} + +describe('parseRawCustomToolCall: schema-guided repair of string-encoded union members', () => { + test('decodes a JSON-encoded string for a union param with an object variant', () => { + const { customToolDefs, rawToolCall } = buildWithCustomTool(unionSchema) + const withInput = { ...rawToolCall, input: { spec: '{"kind": "unhinged-union-spec", "extra": 42}' } } + + const result = parseRawCustomToolCall({ customToolDefs, rawToolCall: withInput }) + + expect((result as { input: { spec: unknown } }).input.spec).toEqual({ + kind: 'unhinged-union-spec', + extra: 42, + }) + }) + + test('keeps a real object value for a union param unchanged', () => { + const { customToolDefs, rawToolCall } = buildWithCustomTool(unionSchema) + const withInput = { ...rawToolCall, input: { spec: { kind: 'plain-object' } } } + + const result = parseRawCustomToolCall({ customToolDefs, rawToolCall: withInput }) + + expect((result as { input: { spec: unknown } }).input.spec).toEqual({ kind: 'plain-object' }) + }) + + test('keeps a non-JSON string for a union param as a string', () => { + const { customToolDefs, rawToolCall } = buildWithCustomTool(unionSchema) + const withInput = { ...rawToolCall, input: { spec: 'plain-string-variant' } } + + const result = parseRawCustomToolCall({ customToolDefs, rawToolCall: withInput }) + + expect((result as { input: { spec: unknown } }).input.spec).toBe('plain-string-variant') + }) + + test('does not decode a JSON-encoded string for a plain string-typed param', () => { + const stringOnlySchema = { + type: 'object', + properties: { code: { type: 'string' } }, + required: ['code'], + additionalProperties: false, + } + const { customToolDefs, rawToolCall } = buildWithCustomTool(stringOnlySchema) + const withInput = { ...rawToolCall, input: { code: '{"looks": "like json"}' } } + + const result = parseRawCustomToolCall({ customToolDefs, rawToolCall: withInput }) + + expect((result as { input: { code: unknown } }).input.code).toBe('{"looks": "like json"}') + }) +}) diff --git a/packages/agent-runtime/src/tools/__tests__/serve-input-schema.test.ts b/packages/agent-runtime/src/tools/__tests__/serve-input-schema.test.ts new file mode 100644 index 0000000000..cac38c4108 --- /dev/null +++ b/packages/agent-runtime/src/tools/__tests__/serve-input-schema.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, test } from 'bun:test' + +import { getToolSet } from '../prompts' + +const makeToolSet = async (inputSchema: unknown) => + getToolSet({ + toolNames: [], + windowedFileReads: false, + additionalToolDefinitions: async () => + ({ + shaped_tool: { + description: 'A tool defined with a live zod schema', + inputSchema, + }, + }) as never, + agentTools: {} as never, + skills: {} as never, + }) + +describe('getToolSet serves custom tool inputSchemas', () => { + test('keeps_live_zod_schema_functional_through_clone_and_serving', async () => { + const { z } = await import('zod/v4') + const liveSchema = z.object({ path: z.string() }) + + const toolSet = await makeToolSet(liveSchema) + + const served = (toolSet['shaped_tool'] as { inputSchema: unknown }).inputSchema + const parse = (served as { safeParse?: (v: unknown) => { success: boolean } }).safeParse + expect(typeof parse).toBe('function') + expect(parse!({ path: 'a.ts' }).success).toBe(true) + expect(() => z.toJSONSchema(served as never, { io: 'input' })).not.toThrow() + }) +}) diff --git a/packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts b/packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts index 3b996cdb87..a6288f0255 100644 --- a/packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts +++ b/packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts @@ -1,5 +1,7 @@ import { mapValues } from 'lodash' +import { toTokenCountInputSchema } from '../../../util/to-json-schema' + import { validateAndGetAgentTemplate, validateAgentInput, @@ -111,10 +113,12 @@ export const handleSpawnAgentInline = (async ( }, ), systemPrompt: system, + // Subagent tool definitions also live in agent state (persisted, + // snapshotted), so inputSchemas must be plain JSON Schema here too. toolDefinitions: mapValues(parentTools, (tool) => ({ description: typeof tool.description === 'string' ? tool.description : undefined, - inputSchema: tool.inputSchema as {}, + inputSchema: toTokenCountInputSchema(tool.inputSchema) ?? {}, })), } diff --git a/packages/agent-runtime/src/tools/prompts.ts b/packages/agent-runtime/src/tools/prompts.ts index d3d9110665..d27df73658 100644 --- a/packages/agent-runtime/src/tools/prompts.ts +++ b/packages/agent-runtime/src/tools/prompts.ts @@ -8,7 +8,8 @@ import { getToolCallString } from '@codebuff/common/tools/utils' import { buildArray } from '@codebuff/common/util/array' import { formatAvailableSkillsXml } from '@codebuff/common/util/skills' import { pluralize } from '@codebuff/common/util/string' -import { cloneDeep } from 'lodash' +import { cloneDeepKeepingZod } from '../util/zod-safe-clone' +import { serveInputSchema } from './serve-input-schema' import z from 'zod/v4' import { convertJsonSchemaToZod } from 'zod-from-json-schema' @@ -430,11 +431,12 @@ export async function getToolSet(params: { const toolDefinitions = await additionalToolDefinitions() for (const [toolName, toolDefinition] of Object.entries(toolDefinitions)) { - const clonedDef = cloneDeep(toolDefinition) - // Custom tool inputSchema may be JSON Schema (from SDK) or Zod (from MCP) - // Ensure it's a Zod schema for the AI SDK - const zodSchema = ensureZodSchema(clonedDef.inputSchema) - const safeSchema = ensureJsonSchemaCompatible(zodSchema) + const clonedDef = cloneDeepKeepingZod(toolDefinition) + // Custom tool inputSchema may be JSON Schema (from SDK) or Zod (from MCP). + // JSON Schema is served verbatim (see serveInputSchema); the former + // unconditional zod round-trip stripped loose schemas to an empty + // object schema at the model. + const safeSchema = serveInputSchema(clonedDef.inputSchema) toolSet[toolName] = { ...clonedDef, inputSchema: safeSchema, diff --git a/packages/agent-runtime/src/tools/serve-input-schema.ts b/packages/agent-runtime/src/tools/serve-input-schema.ts new file mode 100644 index 0000000000..32d4f28959 --- /dev/null +++ b/packages/agent-runtime/src/tools/serve-input-schema.ts @@ -0,0 +1,88 @@ +import { jsonSchema as wrapJsonSchema } from 'ai' +import z from 'zod/v4' + +import { convertJsonSchemaToZod } from 'zod-from-json-schema' + +import type { Logger } from '@codebuff/common/types/contracts/logger' + +/** + * Ensures the inputSchema is a Zod schema. If it's a JSON Schema object + * (from SDK custom tools that were serialized), converts it to Zod. + */ +export function ensureZodSchema( + schema: z.ZodType | Record, +): z.ZodType { + // Check if it's already a Zod schema by looking for the safeParse method + if ( + schema && + typeof (schema as { safeParse?: unknown }).safeParse === 'function' + ) { + return schema as z.ZodType + } + // JSON Schema object - convert to Zod + return convertJsonSchemaToZod(schema as Record) +} + +function ensureJsonSchemaCompatible(schema: z.ZodType): z.ZodType { + try { + z.toJSONSchema(schema, { io: 'input' }) + return schema + } catch { + const fallback = z.object({}).passthrough() + return schema.description ? fallback.describe(schema.description) : fallback + } +} + +/** + * Prepares a custom tool's inputSchema for the AI SDK. The schema ends up in + * two places, with different fidelity requirements: + * + * 1. The tool definition sent to the LLM provider. The model reads this to + * decide what arguments to emit, so it must match what the MCP server + * declared. JSON Schema inputs are therefore passed through verbatim, + * wrapped in ai's jsonSchema() (a pass-through container). + * 2. Argument validation at call time (the validate callback below). + * Approximation is acceptable here — a wrong rejection is recoverable, + * the model can retry — so the zod conversion does this job. + * + * Converting the schema to zod and back would be lossy: schemas zod cannot + * represent (e.g. a property typed only `{ "type": "object" }`) come back + * as an empty object schema, and a model reading an empty argument schema + * emits `{}` — a tool call with no arguments. Zod-typed inputSchemas + * (internal tools defined in TypeScript) keep the + * ensureJsonSchemaCompatible path, which converts in one direction only. + */ +export function serveInputSchema( + inputSchema: z.ZodType | Record, + opts?: { logger?: Logger; name?: string }, +): z.ZodType | ReturnType { + if ( + inputSchema && + typeof (inputSchema as { safeParse?: unknown }).safeParse === 'function' + ) { + return ensureJsonSchemaCompatible(inputSchema as z.ZodType) + } + const rawJsonSchema = inputSchema as Record + // Validation only. The zod conversion handles checking arguments fine; + // its weakness is serializing back to JSON Schema, which we never do here. + const validationSchema = ensureZodSchema(rawJsonSchema) + const served = wrapJsonSchema( + rawJsonSchema as unknown as Parameters[0], + { + validate: (value: unknown) => { + const result = validationSchema.safeParse(value) + return result.success + ? { success: true as const, value: result.data } + : { success: false as const, error: result.error } + }, + }, + ) + if ( + typeof rawJsonSchema.description === 'string' && + rawJsonSchema.description.length > 0 + ) { + ;(served as { description?: string }).description ??= + rawJsonSchema.description + } + return served +} diff --git a/packages/agent-runtime/src/tools/tool-executor.ts b/packages/agent-runtime/src/tools/tool-executor.ts index 36c4708752..361fadd749 100644 --- a/packages/agent-runtime/src/tools/tool-executor.ts +++ b/packages/agent-runtime/src/tools/tool-executor.ts @@ -1,7 +1,7 @@ import { endsAgentStepParam, toolNames } from '@codebuff/common/tools/constants' import { toolParams } from '@codebuff/common/tools/list' import { generateCompactId } from '@codebuff/common/util/string' -import { cloneDeep } from 'lodash' +import { cloneDeepKeepingZod } from '../util/zod-safe-clone' import { getMCPToolData } from '../mcp' import { MCP_TOOL_SEPARATOR } from '../mcp-constants' @@ -10,6 +10,7 @@ import { formatValueForError } from '../util/format-value' import { codebuffToolHandlers } from './handlers/list' import { getMatchingSpawn } from './handlers/tool/spawn-agent-utils' import { getAgentTemplate } from '../templates/agent-registry' +import { repairStringEncodedUnionMembers } from '../util/repair-string-encoded-union-members' import { resolveGravityIndexLink } from './gravity-index-cta' import { ensureZodSchema } from './prompts' @@ -618,6 +619,7 @@ export function parseRawCustomToolCall(params: { const rawSchema = customToolDefs?.[toolName]?.inputSchema if (rawSchema) { + repairStringEncodedUnionMembers(processedParameters, rawSchema) const paramsSchema = ensureZodSchema(rawSchema) const result = paramsSchema.safeParse(processedParameters) @@ -635,7 +637,9 @@ export function parseRawCustomToolCall(params: { } } - const input = JSON.parse(JSON.stringify(parsedInput.input)) + // processedParameters is what the schema saw (including the union repair + // above), so it - not the untouched raw input - is what the handler gets. + const input = JSON.parse(JSON.stringify(processedParameters)) if (endsAgentStepParam in input) { delete input[endsAgentStepParam] } @@ -675,7 +679,7 @@ export async function executeCustomToolCall( ...params, toolNames: agentTemplate.toolNames, mcpServers: agentTemplate.mcpServers, - writeTo: cloneDeep(fileContext.customToolDefinitions), + writeTo: cloneDeepKeepingZod(fileContext.customToolDefinitions), }), rawToolCall: { toolName, diff --git a/packages/agent-runtime/src/util/__tests__/json-safe-state.test.ts b/packages/agent-runtime/src/util/__tests__/json-safe-state.test.ts new file mode 100644 index 0000000000..8491827127 --- /dev/null +++ b/packages/agent-runtime/src/util/__tests__/json-safe-state.test.ts @@ -0,0 +1,148 @@ +import * as analytics from '@codebuff/common/analytics' +import { TEST_USER_ID } from '@codebuff/common/old-constants' +import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime' +import { + createMockDbOperations, + setupDbSpies, +} from '@codebuff/common/testing/mocks/database' +import { getInitialSessionState } from '@codebuff/common/types/session-state' +import { promptSuccess } from '@codebuff/common/util/error' +import { afterEach, describe, expect, spyOn, test } from 'bun:test' + +import { loopAgentSteps } from '../../run-agent-step' + +import type { AgentTemplate } from '../../templates/types' +import type { DbSpies } from '@codebuff/common/testing/mocks/database' +import type { ProjectFileContext } from '@codebuff/common/util/file' + +/** + * Application-tier regression test: toolDefinitions live in agent state + * (persisted, snapshotted, shipped over the wire), so every stored + * inputSchema must be plain JSON Schema. A live zod instance in state + * serializes as {"def":{...}} internals instead of the declared schema. + */ + +const CUSTOM_TOOL_NAME = 'declared_tool' + +const baseFileContext: ProjectFileContext = { + projectRoot: '/test', + cwd: '/test', + fileTree: [], + fileTokenScores: {}, + knowledgeFiles: {}, + gitChanges: { status: '', diff: '', diffCached: '', lastCommitMessages: '' }, + changesSinceLastChat: {}, + shellConfigFiles: {}, + systemInfo: { + platform: 'test', + shell: 'test', + nodeVersion: 'test', + arch: 'test', + homedir: '/home/test', + cpus: 1, + chromeAvailable: false, + }, + agentTemplates: {}, + customToolDefinitions: {}, +} + +const makeAgent = (): AgentTemplate => ({ + id: 'json-safe-state-agent', + displayName: 'JSON Safe State Agent', + spawnerPrompt: 'Regression: state toolDefinitions stay JSON-safe', + model: 'google/gemini-2.5-flash', + inputSchema: {}, + outputMode: 'last_message' as const, + includeMessageHistory: true, + inheritParentSystemPrompt: false, + mcpServers: {}, + toolNames: [CUSTOM_TOOL_NAME], + spawnableAgents: [], + systemPrompt: 'Test system prompt', + instructionsPrompt: '', + stepPrompt: '', +}) + +const makeFileContextWithDeclaredTool = (schema: unknown): ProjectFileContext => + ({ + ...baseFileContext, + customToolDefinitions: { + [CUSTOM_TOOL_NAME]: { + description: 'A tool declared with a JSON Schema', + inputSchema: schema, + }, + }, + }) as ProjectFileContext + +const runStepToPopulation = async (fileContext: ProjectFileContext) => { + const agent = makeAgent() + const sessionState = getInitialSessionState(baseFileContext) + const agentState = sessionState.mainAgentState + agentState.messageHistory = [] + + await loopAgentSteps({ + ...(TEST_AGENT_RUNTIME_IMPL as unknown as Record), + sendAction: () => {}, + additionalToolDefinitions: () => Promise.resolve({}), + ancestorRunIds: [], + clientSessionId: 'json-safe-state-session', + fileContext, + fingerprintId: 'json-safe-state-fingerprint', + onResponseChunk: () => {}, + repoId: undefined, + repoUrl: undefined, + runId: 'json-safe-state-run', + signal: new AbortController().signal, + spawnParams: undefined, + system: 'Test system prompt', + tools: {}, + userId: TEST_USER_ID, + userInputId: 'json-safe-state-input', + promptAiSdkStream: async function* () { + yield { type: 'text' as const, text: 'response text' } + return promptSuccess('mock-message-id') + }, + agentType: agent.id, + localAgentTemplates: { [agent.id]: agent }, + agentTemplate: agent, + agentState, + prompt: 'hello', + } as never) + + return agentState +} + +describe('agent state toolDefinitions serialization', () => { + let dbSpies: DbSpies + let analyticsSpy: ReturnType + + afterEach(() => { + dbSpies.restore() + analyticsSpy.mockRestore() + }) + + test('stores_declared_json_schema_without_zod_internals', async () => { + dbSpies = setupDbSpies(createMockDbOperations()) + analyticsSpy = spyOn(analytics, 'trackEvent').mockImplementation(() => {}) + const declaredSchema = { + type: 'object', + properties: { path: { type: 'string', description: 'The file path' } }, + required: ['path'], + } + + const agentState = await runStepToPopulation( + makeFileContextWithDeclaredTool(declaredSchema), + ) + + const toolDefs = agentState.toolDefinitions as Record< + string, + { inputSchema?: unknown } + > + expect(Object.keys(toolDefs)).toContain(CUSTOM_TOOL_NAME) + + const serialized = JSON.stringify(toolDefs[CUSTOM_TOOL_NAME].inputSchema) + const roundTripped = JSON.parse(serialized) as { type?: string } + expect(roundTripped.type).toBe('object') + expect(serialized).not.toContain('"def"') + }) +}) diff --git a/packages/agent-runtime/src/util/__tests__/to-json-schema.test.ts b/packages/agent-runtime/src/util/__tests__/to-json-schema.test.ts new file mode 100644 index 0000000000..a5cc67e1dd --- /dev/null +++ b/packages/agent-runtime/src/util/__tests__/to-json-schema.test.ts @@ -0,0 +1,85 @@ +import { describe, test, expect } from 'bun:test' +import { z } from 'zod/v4' + +import { toTokenCountInputSchema } from '../to-json-schema' + +/** + * Regression tests for the persisted-state schema conversion. + * + * Tool inputSchemas are persisted into agent state, snapshotted and replayed + * on every turn, and shipped to Anthropic's count_tokens API. Every stored + * schema must therefore be plain JSON Schema with a top-level type: zod + * internals never leak into state, and foreign (already-JSON) schemas pass + * through unmangled. + */ +describe('toTokenCountInputSchema', () => { + /** + * Given: a zod object schema with an optional field. + * When: it is converted. + * Then: the result is JSON Schema with type object and the field mapped, + * not a serialized zod instance. + */ + test('converts zod object schema to JSON Schema with top level type object', () => { + const schema = z.object({ + q: z.string().describe('query'), + n: z.number().optional(), + }) + + const out = toTokenCountInputSchema(schema) as Record | undefined + + expect(out?.type).toBe('object') + expect(out?.properties.q.type).toBe('string') + }) + + /** + * Given: a union schema, which JSON Schema represents as anyOf with no + * top-level type. + * When: it is converted. + * Then: type object is backfilled, because Anthropic's count_tokens + * rejects input_schema values without a top-level type. + */ + test('backfills type object for union schemas represented as anyOf', () => { + const schema = z.union([z.object({ a: z.string() }), z.object({ b: z.number() })]) + + const out = toTokenCountInputSchema(schema) as Record | undefined + + expect(out?.type).toBe('object') + expect(out?.anyOf).toBeDefined() + }) + + /** + * Given: a schema that is already a plain JSON Schema object (the shape + * MCP servers and the SDK send). + * When: it is converted. + * Then: it is copied as-is - conversion must not mangle foreign schemas. + */ + test('copies an already plain JSON Schema object unchanged', () => { + const jsonSchema = { + type: 'object', + properties: { location: { type: 'string', enum: ['NYC', 'LA'] } }, + required: ['location'], + } + + const out = toTokenCountInputSchema(jsonSchema) + + expect(out).toEqual(jsonSchema) + }) + + /** + * Given: nullish input and a schema carrying a $schema key. + * When: they are converted. + * Then: nullish input yields undefined, and the meaningless $schema key + * is dropped to keep the token-count payload lean. + */ + test('returns undefined for nullish input and strips the schema meta key', () => { + const withMeta = { $schema: 'https://json-schema.org/x', type: 'object' } + + const nullishOut = toTokenCountInputSchema(undefined) + const metaOut = toTokenCountInputSchema(withMeta) + + expect(nullishOut).toBeUndefined() + expect(toTokenCountInputSchema(null)).toBeUndefined() + expect(metaOut?.$schema).toBeUndefined() + expect(metaOut?.type).toBe('object') + }) +}) diff --git a/packages/agent-runtime/src/util/__tests__/zod-safe-clone.test.ts b/packages/agent-runtime/src/util/__tests__/zod-safe-clone.test.ts new file mode 100644 index 0000000000..fbf378c49c --- /dev/null +++ b/packages/agent-runtime/src/util/__tests__/zod-safe-clone.test.ts @@ -0,0 +1,67 @@ +import { describe, test, expect } from 'bun:test' +import { cloneDeep } from 'lodash' +import { z } from 'zod/v4' + +import { cloneDeepKeepingZod } from '../zod-safe-clone' + +/** + * Regression tests for tool-schema cloning. + * + * Tool definitions carry live zod v4 schemas, and state boundaries + * deep-clone the surrounding data. lodash cloneDeep strips zod's + * non-enumerable _zod engine: the stripped clone still looks like a schema + * (safeParse, def, shape all present) but throws the first time zod + * internals touch it - which is how MCP and custom tool schemas silently + * became empty {} at the model. cloneDeepKeepingZod is the fix pinned here. + */ +describe('lodash cloneDeep zod amputation (the bug)', () => { + /** + * Given: a zod v4 schema. + * When: it is cloned with lodash cloneDeep. + * Then: the clone still looks like a schema (safeParse present) but its + * engine is gone: z.toJSONSchema throws on it - the production failure + * behind the empty-schema bug, and the reason the helper below exists. + */ + test('cloneDeep strips the zod engine so toJSONSchema throws on the clone', () => { + const schema = z.object({ q: z.string() }) + + const cloned = cloneDeep(schema) + + // Asserted behaviorally: the clone still parses, but conversion fails. + expect(typeof cloned.safeParse).toBe('function') + expect(() => z.toJSONSchema(cloned as never)).toThrow() + }) +}) + +describe('cloneDeepKeepingZod', () => { + /** + * Given: a plain (schema-free) nested structure. + * When: it is cloned with cloneDeepKeepingZod. + * Then: the result matches cloneDeep exactly, including fresh nested + * references - the clone helper must not change plain-data semantics. + */ + test('cloneDeepKeepingZod deep-clones plain structures exactly like cloneDeep', () => { + const input = { a: { b: [1, { c: 'd' }] }, e: null } + + const out = cloneDeepKeepingZod(input) + + expect(out).toEqual(input) + expect(out.a).not.toBe(input.a) + expect(out.a.b[1]).not.toBe(input.a.b[1]) + }) + + /** + * Given: a zod schema nested inside a collection, the shape custom tool * definitions actually arrive in. + * When: the containing structure is cloned. + * Then: the schema survives as a live instance usable by zod internals. + */ + test('cloneDeepKeepingZod preserves schemas nested inside collections', () => { + const schema = z.object({ id: z.number() }) + const input = { tools: [{ name: 'x', inputSchema: schema }] } + + const out = cloneDeepKeepingZod(input) + + expect(out.tools[0].inputSchema).toBe(schema) + expect(() => z.toJSONSchema(out.tools[0].inputSchema)).not.toThrow() + }) +}) diff --git a/packages/agent-runtime/src/util/repair-string-encoded-union-members.ts b/packages/agent-runtime/src/util/repair-string-encoded-union-members.ts new file mode 100644 index 0000000000..8026db1832 --- /dev/null +++ b/packages/agent-runtime/src/util/repair-string-encoded-union-members.ts @@ -0,0 +1,45 @@ +/** + * Repairs values the model string-encoded against its schema. When a + * parameter's declared schema is a union containing an object variant, a + * model may emit the object as a JSON-encoded string (a string is + * unambiguously valid for the union, so nothing downstream fails). The + * schema-guided decode below restores the object the model meant; plain + * strings and params without an object variant are never touched, so + * tools whose string parameters legitimately contain JSON (script + * sources, file contents) are unaffected. + */ +export function repairStringEncodedUnionMembers( + parameters: Record, + rawSchema: unknown, +): void { + if (!rawSchema || typeof rawSchema !== 'object') return + const properties = (rawSchema as { properties?: Record }) + .properties + if (!properties) return + for (const [param, value] of Object.entries(parameters)) { + if (typeof value !== 'string') continue + const propSchema = properties[param] + if (!propSchema || typeof propSchema !== 'object') continue + const union = + (propSchema as { anyOf?: unknown[] }).anyOf ?? + (propSchema as { oneOf?: unknown[] }).oneOf + if (!Array.isArray(union)) continue + const hasObjectVariant = union.some( + (variant) => + variant && + typeof variant === 'object' && + (variant as { type?: unknown }).type === 'object', + ) + if (!hasObjectVariant) continue + const trimmed = value.trim() + if (!trimmed.startsWith('{') && !trimmed.startsWith('[')) continue + try { + const decoded = JSON.parse(trimmed) + if (decoded && typeof decoded === 'object') { + parameters[param] = decoded + } + } catch { + // Not JSON after all - the string is a legitimate value. + } + } +} diff --git a/packages/agent-runtime/src/util/to-json-schema.ts b/packages/agent-runtime/src/util/to-json-schema.ts new file mode 100644 index 0000000000..1d32a34971 --- /dev/null +++ b/packages/agent-runtime/src/util/to-json-schema.ts @@ -0,0 +1,46 @@ +import z from 'zod/v4' + +// Convert a tool's stored inputSchema into JSON Schema suitable for Anthropic's +// count_tokens API. Built-in and MCP tools store a Zod schema here; serializing +// it raw ships Zod internals (`def`/`shape`) instead of JSON Schema, so token +// counts are computed against garbage and any schema whose top-level isn't an +// object (e.g. a union → `anyOf`) arrives without `type`, which the API rejects +// with `tools.N.custom.input_schema.type: Field required`. We convert to JSON +// Schema and guarantee a top-level `type: 'object'`. +// +// Lives in util/ (not run-agent-step) so spawn-agent-inline can use it without +// an import cycle through run-agent-step. +export function toTokenCountInputSchema( + inputSchema: unknown, +): Record | undefined { + if (inputSchema == null) return undefined + + let jsonSchema: Record + if ( + typeof (inputSchema as { safeParse?: unknown }).safeParse === 'function' + ) { + try { + jsonSchema = z.toJSONSchema(inputSchema as z.ZodType, { + io: 'input', + }) as Record + } catch { + jsonSchema = { type: 'object', properties: {} } + } + } else if (typeof inputSchema === 'object' && !Array.isArray(inputSchema)) { + // Already a plain object (e.g. a pre-serialized JSON Schema) — copy it. + jsonSchema = { ...(inputSchema as Record) } + } else { + return undefined + } + + // `$schema` is meaningless to count_tokens; drop it to keep the payload lean. + delete jsonSchema['$schema'] + // Anthropic requires a top-level `type: 'object'`. Object schemas already + // carry it; union/intersection schemas (anyOf/allOf) don't — backfill it. + // Treat missing / null / empty-string as absent (valid JSON Schema `type` is + // always a non-empty string or array). + if (jsonSchema.type == null || jsonSchema.type === '') { + jsonSchema.type = 'object' + } + return jsonSchema +} diff --git a/packages/agent-runtime/src/util/zod-safe-clone.ts b/packages/agent-runtime/src/util/zod-safe-clone.ts new file mode 100644 index 0000000000..6e680895c4 --- /dev/null +++ b/packages/agent-runtime/src/util/zod-safe-clone.ts @@ -0,0 +1,34 @@ +import { cloneDeepWith } from 'lodash' + +/** + * lodash cloneDeep destroys zod v4 schema instances. + * + * zod v4 stores its engine on a non-enumerable `_zod` property, and lodash + * only copies enumerable own properties. The clone therefore looks like a + * schema (has safeParse/def/type) but has no `_zod` internals, and any zod + * internal that touches `schema._zod.*` detonates with: + * "undefined is not an object (evaluating 'schema._zod.def')" + * + * This deep-clones plain data (descriptions, maps, arrays) exactly like + * cloneDeep, but passes zod schema instances through by reference so their + * internals survive. + */ +export function cloneDeepKeepingZod(value: T): T { + const cloned = cloneDeepWith(value, (node) => { + if (isZodSchemaInstance(node)) { + // Pass the live schema through untouched. + return node as T + } + // Fall through to lodash's default deep clone. + return undefined + }) + return cloned as T +} + +function isZodSchemaInstance(node: unknown): boolean { + if (typeof node !== 'object' || node === null) { + return false + } + const candidate = node as { _zod?: unknown } + return typeof candidate._zod === 'object' && candidate._zod !== null +} diff --git a/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.test.ts b/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.test.ts index 195d63b819..9827f7314d 100644 --- a/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.test.ts +++ b/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.test.ts @@ -1083,3 +1083,47 @@ describe('consecutive assistant messages', () => { ]) }) }) + +/** + * Regression tests for non-image file parts. + * + * MCP resources can put non-image file parts (e.g. gzip) into message + * history, which is replayed into every later prompt build. The + * OpenAI-compatible converter must degrade such parts to a text + * placeholder: throwing here failed the entire prompt build and, because + * the message stays in history, killed the session on every subsequent + * turn. + */ +describe('non-image file parts', () => { + // The fixture's base64 string is 20 chars; the placeholder estimates raw + // bytes as round(20 * 3 / 4) = 15. + const GZIP_FIXTURE_BASE64 = Buffer.from('Hello freebuff!').toString('base64') + const EXPECTED_BYTE_ESTIMATE = 15 + + it('degrades non-image file part to text placeholder instead of throwing', () => { + const result = convertToOpenAICompatibleChatMessages([ + { + role: 'user', + content: [ + { + type: 'file', + data: GZIP_FIXTURE_BASE64, + mediaType: 'application/gzip', + }, + ], + }, + ]) + + expect(result).toEqual([ + { + role: 'user', + content: [ + { + type: 'text', + text: `[application/gzip file part not displayable (~${EXPECTED_BYTE_ESTIMATE} bytes)]`, + }, + ], + }, + ]) + }) +}) diff --git a/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.ts b/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.ts index ead5daab11..4491f8dfaa 100644 --- a/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.ts +++ b/packages/llm-providers/src/openai-compatible/chat/convert-to-openai-compatible-chat-messages.ts @@ -14,6 +14,25 @@ function getOpenAIMetadata(message: { return message?.providerOptions?.openaiCompatible ?? {} } +/** Approximate payload size of a file part's data, for placeholder text. */ +function filePartByteLength(data: unknown): number { + let value = data + if (value && typeof value === 'object' && 'type' in value) { + if (value.type === 'data' && 'data' in value) { + value = value.data + } else if (value.type === 'url' && 'url' in value) { + value = value.url + } + } + if (typeof value === 'string') { + return Math.round((value.length * 3) / 4) + } + if (value instanceof Uint8Array) { + return value.byteLength + } + return 0 +} + function imageUrlFromData(data: unknown, mediaType: string): string { // AI SDK 7 adapts this v2 provider to v4, whose file data is tagged. The // compatibility proxy passes that v4 shape through to the v2 implementation. @@ -89,9 +108,17 @@ export function convertToOpenAICompatibleChatMessages( ...partMetadata, } } else { - throw new UnsupportedFunctionalityError({ - functionality: `file part media type ${part.mediaType}`, - }) + // Non-image file parts (e.g. application/gzip from an MCP + // resource) have no OpenAI-compatible representation. + // Degrade to a text placeholder instead of throwing: a + // throw here fails the entire prompt build and, because + // the message stays in history, kills the session on every + // subsequent turn. + return { + type: 'text', + text: `[${part.mediaType} file part not displayable (~${filePartByteLength(part.data)} bytes)]`, + ...partMetadata, + } } } }