apply zod schema to custom tool input

Author: charleslien

backend/src/tools/tool-executor.ts

@@ -335,13 +335,13 @@ export function parseRawCustomToolCall(
     }
   }
 
-  if (endsAgentStepParam in result.data) {
-    delete result.data[endsAgentStepParam]
+  const input = JSON.parse(JSON.stringify(rawToolCall.input))
+  if (endsAgentStepParam in input) {
+    delete input[endsAgentStepParam]
   }
-
   return {
     toolName: toolName,
-    input: result.data,
+    input,
     toolCallId: rawToolCall.toolCallId,
   }
 }

sdk/src/client.ts

@@ -116,7 +116,7 @@ export class CodebuffClient {
    *
    * @returns A Promise that resolves to a RunState JSON object which you can pass to a subsequent run() call to continue the run.
    */
-  public async run({
+  public async run<A extends string = string, B = any, C = any>({
     agent,
     prompt,
     params,
@@ -136,7 +136,7 @@ export class CodebuffClient {
     projectFiles?: Record<string, string>
     knowledgeFiles?: Record<string, string>
     agentDefinitions?: AgentDefinition[]
-    customToolDefinitions?: CustomToolDefinition[]
+    customToolDefinitions?: CustomToolDefinition<A, B, C>[]
     maxAgentSteps?: number
   }): Promise<RunState> {
     await this.websocketHandler.connect()
@@ -169,13 +169,15 @@ export class CodebuffClient {
             `Implementation for custom tool ${toolName} not found.`,
           )
         }
-        const handler = toolDefs[toolDefs.length - 1].handler
+        const toolDef = toolDefs[toolDefs.length - 1]
+        const handler = toolDef.handler
         try {
           return {
             success: true,
             output: {
               type: 'text',
-              value: (await handler(input)).toolResultMessage,
+              value: (await handler(toolDef.zodSchema.parse(input)))
+                .toolResultMessage,
             },
           }
         } catch (error) {

sdk/src/custom-tool.ts

@@ -13,7 +13,7 @@ export type CustomToolDefinition<
   description?: string
   endsAgentStep: boolean
   exampleInputs: Input[]
-  handler: (params: Input) => Promise<{
+  handler: (params: Output) => Promise<{
     toolResultMessage: string
   }>
 }
@@ -35,7 +35,7 @@ export function getCustomToolDefinintion<
   description?: string
   endsAgentStep?: boolean
   exampleInputs?: Input[]
-  handler: (params: Input) => Promise<{
+  handler: (params: Output) => Promise<{
     toolResultMessage: string
   }>
 }): CustomToolDefinition<ToolName, Output, Input> {