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..038a2f6897 100644 --- a/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts +++ b/packages/agent-runtime/src/__tests__/prompts-schema-handling.test.ts @@ -12,6 +12,8 @@ import { handleLookupAgentInfo } from '../tools/handlers/tool/lookup-agent-info' import { ensureZodSchema, buildToolDescription, + cloneCustomToolDefinitions, + cloneToolDefinition, getToolSet, } from '../tools/prompts' @@ -510,3 +512,83 @@ describe('getToolSet: commit-attribution suppression', () => { ) }) }) + +describe('cloneToolDefinition: Zod schema preservation (issue #1306)', () => { + const mcpJsonSchema = { + type: 'object', + properties: { + query: { type: 'string' }, + max_results: { type: 'number' }, + }, + required: ['query'], + additionalProperties: false, + } + + test('keeps genuine Zod schemas by reference with _zod intact', () => { + const mcpSchema = convertJsonSchemaToZod(mcpJsonSchema as any) as any + const def = { + description: 'Search the web', + inputSchema: mcpSchema, + endsAgentStep: true, + } + + const cloned = cloneToolDefinition(def) + + // New container, same live schema — lodash cloneDeep would strip the + // non-enumerable `_zod` engine and break all later conversions. + expect(cloned).not.toBe(def) + expect(cloned.inputSchema).toBe(mcpSchema) + expect((cloned.inputSchema as any)._zod).toBeDefined() + expect(cloned.inputSchema instanceof z.ZodType).toBe(true) + const jsonSchema = z.toJSONSchema(cloned.inputSchema, { + io: 'input', + }) as any + expect(Object.keys(jsonSchema.properties ?? {})).toContain('query') + }) + + test('deep-clones plain JSON Schema objects so edits do not alias', () => { + const defs = { + json_tool: { + description: 'json tool', + inputSchema: { + type: 'object', + properties: { q: { type: 'string' } }, + } as Record, + endsAgentStep: false, + }, + } + + const cloned = cloneCustomToolDefinitions(defs) + + expect(cloned.json_tool).not.toBe(defs.json_tool) + expect(cloned.json_tool.inputSchema).not.toBe(defs.json_tool.inputSchema) + expect(cloned.json_tool.inputSchema).toEqual(defs.json_tool.inputSchema) + }) + + test('getToolSet serves MCP params after the run-agent-step clone path', async () => { + const mcpSchema = convertJsonSchemaToZod(mcpJsonSchema as any) as any + // Simulates run-agent-step.ts additionalToolDefinitions: the defs map is + // cloned before MCP data is merged in. + const defs = cloneCustomToolDefinitions({ + web__search: { + description: 'Search the web', + inputSchema: mcpSchema, + endsAgentStep: true, + }, + }) + + const toolSet = await getToolSet({ + toolNames: [], + windowedFileReads: false, + additionalToolDefinitions: async () => defs, + agentTools: {}, + skills: {}, + }) + + const served = (toolSet as any).web__search.inputSchema + const jsonSchema = z.toJSONSchema(served, { io: 'input' }) as any + const props = jsonSchema.properties ?? {} + expect(Object.keys(props)).toContain('query') + expect(Object.keys(props)).toContain('max_results') + }) +}) diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index e956a468da..9268d2769d 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -22,7 +22,7 @@ import { userMessage, } from '@codebuff/common/util/messages' import { type ToolSet } from 'ai' -import { cloneDeep, mapValues } from 'lodash' +import { mapValues } from 'lodash' import z from 'zod/v4' import { maybeCompactHistory } from './compact-history' @@ -41,7 +41,7 @@ import { import { getAgentTemplate } from './templates/agent-registry' import { buildAgentToolSet } from './templates/prompts' import { getAgentPrompt } from './templates/strings' -import { getToolSet } from './tools/prompts' +import { cloneCustomToolDefinitions, getToolSet } from './tools/prompts' import { processStream } from './tools/stream-parser' import { getAgentOutput } from './util/agent-output' import { @@ -151,7 +151,9 @@ async function additionalToolDefinitions( ): Promise { const { agentTemplate, fileContext } = params - const defs = cloneDeep( + // Zod-aware clone: plain cloneDeep would strip Zod v4's non-enumerable + // `_zod` engine from MCP tool schemas and break them downstream. + const defs = cloneCustomToolDefinitions( Object.fromEntries( Object.entries(fileContext.customToolDefinitions).filter(([toolName]) => agentTemplate!.toolNames.includes(toolName), diff --git a/packages/agent-runtime/src/tools/prompts.ts b/packages/agent-runtime/src/tools/prompts.ts index d3d9110665..f0d688771c 100644 --- a/packages/agent-runtime/src/tools/prompts.ts +++ b/packages/agent-runtime/src/tools/prompts.ts @@ -38,6 +38,49 @@ export function ensureZodSchema( return convertJsonSchemaToZod(schema as Record) } +/** + * Clone a custom tool definition without destroying Zod schemas. + * + * lodash `cloneDeep` only copies enumerable own properties, but Zod v4 keeps + * its engine in the non-enumerable `_zod` field. Deep-cloning a Zod schema + * therefore produces a broken lookalike: `safeParse` is copied, so + * `ensureZodSchema` still accepts it, but any `z.toJSONSchema()` call + * crashes on the missing `_zod` and the tool's params collapse to `{}`. + * + * Shallow-copy the definition, keep genuine Zod schemas by reference (they + * are treated as immutable downstream), and deep-clone only plain JSON + * Schema objects plus the remaining fields. + */ +export function cloneToolDefinition< + T extends { inputSchema?: z.ZodType | Record }, +>(toolDefinition: T): T { + const { inputSchema, ...rest } = toolDefinition + const clonedRest = cloneDeep(rest) + const clonedSchema = + inputSchema instanceof z.ZodType || inputSchema == null + ? inputSchema + : cloneDeep(inputSchema) + return { ...clonedRest, inputSchema: clonedSchema } as T +} + +/** + * Clone a map of custom tool definitions, preserving Zod schema identity + * (see `cloneToolDefinition`). + */ +export function cloneCustomToolDefinitions< + T extends Record< + string, + { inputSchema?: z.ZodType | Record } + >, +>(toolDefinitions: T): T { + return Object.fromEntries( + Object.entries(toolDefinitions).map(([toolName, toolDefinition]) => [ + toolName, + cloneToolDefinition(toolDefinition), + ]), + ) as T +} + function ensureJsonSchemaCompatible(schema: z.ZodType): z.ZodType { try { z.toJSONSchema(schema, { io: 'input' }) @@ -430,7 +473,9 @@ export async function getToolSet(params: { const toolDefinitions = await additionalToolDefinitions() for (const [toolName, toolDefinition] of Object.entries(toolDefinitions)) { - const clonedDef = cloneDeep(toolDefinition) + // Zod-aware clone: plain cloneDeep would strip Zod v4's non-enumerable + // `_zod` engine and break the schema (see cloneToolDefinition). + const clonedDef = cloneToolDefinition(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) diff --git a/packages/agent-runtime/src/tools/tool-executor.ts b/packages/agent-runtime/src/tools/tool-executor.ts index 36c4708752..a72c22be65 100644 --- a/packages/agent-runtime/src/tools/tool-executor.ts +++ b/packages/agent-runtime/src/tools/tool-executor.ts @@ -1,7 +1,6 @@ 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 { getMCPToolData } from '../mcp' import { MCP_TOOL_SEPARATOR } from '../mcp-constants' @@ -11,7 +10,7 @@ import { codebuffToolHandlers } from './handlers/list' import { getMatchingSpawn } from './handlers/tool/spawn-agent-utils' import { getAgentTemplate } from '../templates/agent-registry' import { resolveGravityIndexLink } from './gravity-index-cta' -import { ensureZodSchema } from './prompts' +import { cloneCustomToolDefinitions, ensureZodSchema } from './prompts' import type { AgentTemplate } from '../templates/types' import type { CodebuffToolHandlerFunction } from './handlers/handler-function-type' @@ -675,7 +674,9 @@ export async function executeCustomToolCall( ...params, toolNames: agentTemplate.toolNames, mcpServers: agentTemplate.mcpServers, - writeTo: cloneDeep(fileContext.customToolDefinitions), + // Zod-aware clone: plain cloneDeep would strip Zod v4's + // non-enumerable `_zod` engine and break validation downstream. + writeTo: cloneCustomToolDefinitions(fileContext.customToolDefinitions), }), rawToolCall: { toolName,