Preserve Zod v4 schemas when cloning custom tool definitions

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<string, unknown>,
+        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')
+  })
+})

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<CustomToolDefinitions> {
   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),

packages/agent-runtime/src/tools/prompts.ts

@@ -38,6 +38,49 @@ export function ensureZodSchema(
   return convertJsonSchemaToZod(schema as Record<string, unknown>)
 }
 
+/**
+ * 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<string, unknown> },
+>(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<string, unknown> }
+  >,
+>(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)

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,