Auto-repair tool calls with agent ids to spawn_agents tool calls!

Author: jahooma

common/src/types/contracts/llm.ts

@@ -6,6 +6,7 @@ import type { ParamsExcluding } from '../function-params'
 import type { Logger } from './logger'
 import type { Model } from '../../old-constants'
 import type { Message } from '../messages/codebuff-message'
+import type { AgentTemplate } from '../agent-template'
 import type { generateText, streamText, ToolCallPart } from 'ai'
 import type z from 'zod/v4'
 
@@ -42,6 +43,10 @@ export type PromptAiSdkStreamFn = (
     onCostCalculated?: (credits: number) => Promise<void>
     includeCacheControl?: boolean
     agentProviderOptions?: OpenRouterProviderRoutingOptions
+    /** List of agents that can be spawned - used to transform agent tool calls */
+    spawnableAgents?: string[]
+    /** Map of locally available agent templates - used to transform agent tool calls */
+    localAgentTemplates?: Record<string, AgentTemplate>
     sendAction: SendActionFn
     logger: Logger
     trackEvent: TrackEventFn

packages/agent-runtime/src/prompt-agent-stream.ts

@@ -21,6 +21,7 @@ export const getAgentStreamFromTemplate = (params: {
   fingerprintId: string
   includeCacheControl?: boolean
   liveUserInputRecord: UserInputRecord
+  localAgentTemplates: Record<string, AgentTemplate>
   logger: Logger
   messages: Message[]
   runId: string
@@ -42,6 +43,7 @@ export const getAgentStreamFromTemplate = (params: {
     fingerprintId,
     includeCacheControl,
     liveUserInputRecord,
+    localAgentTemplates,
     logger,
     messages,
     runId,
@@ -71,12 +73,14 @@ export const getAgentStreamFromTemplate = (params: {
     includeCacheControl,
     logger,
     liveUserInputRecord,
+    localAgentTemplates,
     maxOutputTokens: 32_000,
     maxRetries: 3,
     messages,
     model,
     runId,
     sessionConnections,
+    spawnableAgents: template.spawnableAgents,
     stopSequences: [globalStopSequence],
     tools,
     userId,

sdk/src/impl/llm.ts

@@ -18,7 +18,13 @@ import {
   OpenAICompatibleChatLanguageModel,
   VERSION,
 } from '@codebuff/internal/openai-compatible/index'
-import { streamText, APICallError, generateText, generateObject } from 'ai'
+import {
+  streamText,
+  APICallError,
+  generateText,
+  generateObject,
+  NoSuchToolError,
+} from 'ai'
 
 import { WEBSITE_URL } from '../constants'
 import { NetworkError, PaymentRequiredError, ErrorCodes } from '../errors'
@@ -217,6 +223,115 @@ export async function* promptAiSdkStream(
       ...params,
       agentProviderOptions: params.agentProviderOptions,
     }),
+    // Transform agent tool calls (e.g. 'file-picker') into spawn_agents calls
+    experimental_repairToolCall: async ({ toolCall, tools, error }) => {
+      // Only handle NoSuchToolError - when model tries to call a non-existent tool
+      if (!NoSuchToolError.isInstance(error)) {
+        return null
+      }
+
+      const { spawnableAgents = [], localAgentTemplates = {} } = params
+
+      // Check if spawn_agents tool is available
+      if (!('spawn_agents' in tools)) {
+        return null
+      }
+
+      // Check if the tool name matches a spawnable agent or local agent template
+      const toolName = toolCall.toolName
+      const isSpawnableAgent = spawnableAgents.some((agentId) => {
+        // Match by exact name or by the agent part of publisher/agent@version format
+        const withoutVersion = agentId.split('@')[0]
+        const parts = withoutVersion.split('/')
+        const agentName = parts[parts.length - 1]
+        return agentName === toolName || agentId === toolName
+      })
+      const isLocalAgent = toolName in localAgentTemplates
+
+      if (!isSpawnableAgent && !isLocalAgent) {
+        return null
+      }
+
+      // Parse the input - it comes as a JSON string from the AI SDK
+      // We also need to recursively parse any nested JSON strings
+      const deepParseJson = (value: unknown): unknown => {
+        if (typeof value === 'string') {
+          try {
+            const parsed = JSON.parse(value)
+            // Recursively parse the result in case it contains more JSON strings
+            return deepParseJson(parsed)
+          } catch {
+            // Not valid JSON, return as-is
+            return value
+          }
+        }
+        if (Array.isArray(value)) {
+          return value.map(deepParseJson)
+        }
+        if (value !== null && typeof value === 'object') {
+          const result: Record<string, unknown> = {}
+          for (const [k, v] of Object.entries(value)) {
+            result[k] = deepParseJson(v)
+          }
+          return result
+        }
+        return value
+      }
+
+      let input: Record<string, unknown> = {}
+      try {
+        const rawInput =
+          typeof toolCall.input === 'string'
+            ? JSON.parse(toolCall.input)
+            : (toolCall.input as Record<string, unknown>)
+        // Deep parse to handle any nested JSON strings
+        input = deepParseJson(rawInput) as Record<string, unknown>
+      } catch {
+        // If parsing fails, use empty object
+      }
+
+      // Extract prompt from input if it exists and is a string
+      const prompt =
+        typeof input.prompt === 'string' ? input.prompt : undefined
+
+      // All other input parameters become the params object
+      // These should NOT be stringified - they should remain as objects
+      const agentParams: Record<string, unknown> = {}
+      for (const [key, value] of Object.entries(input)) {
+        if (key === 'prompt' && typeof value === 'string') {
+          continue
+        }
+        agentParams[key] = value
+      }
+
+      // Transform into spawn_agents call
+      // The input must be a JSON string for the AI SDK, but the nested objects
+      // should remain as proper objects (not stringified)
+      const spawnAgentsInput = {
+        agents: [
+          {
+            agent_type: toolName,
+            ...(prompt !== undefined && { prompt }),
+            ...(Object.keys(agentParams).length > 0 && { params: agentParams }),
+          },
+        ],
+      }
+
+      logger.info(
+        {
+          originalToolName: toolName,
+          transformedInput: spawnAgentsInput,
+        },
+        'Transformed agent tool call to spawn_agents',
+      )
+
+      // Return the repaired tool call - input must be a JSON string
+      return {
+        ...toolCall,
+        toolName: 'spawn_agents',
+        input: JSON.stringify(spawnAgentsInput),
+      }
+    },
   })
 
   let content = ''