input tools into ai sdk

Author: charleslien

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

@@ -12,6 +12,7 @@ import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsOf } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'
 import type { OpenRouterProviderOptions } from '@codebuff/internal/openrouter-ai-sdk'
+import type { ToolSet } from 'ai'
 
 export const getAgentStreamFromTemplate = (params: {
   agentId?: string
@@ -26,6 +27,7 @@ export const getAgentStreamFromTemplate = (params: {
   sessionConnections: SessionRecord
   template: AgentTemplate
   textOverride: string | null
+  tools: ToolSet
   userId: string | undefined
   userInputId: string
 
@@ -47,6 +49,7 @@ export const getAgentStreamFromTemplate = (params: {
     sessionConnections,
     template,
     textOverride,
+    tools,
     userId,
     userInputId,
 
@@ -71,24 +74,26 @@ export const getAgentStreamFromTemplate = (params: {
   const { model } = template
 
   const aiSdkStreamParams: ParamsOf<PromptAiSdkStreamFn> = {
+    agentId,
     apiKey,
-    runId,
+    clientSessionId,
+    fingerprintId,
+    includeCacheControl,
+    logger,
+    liveUserInputRecord,
+    maxOutputTokens: 32_000,
+    maxRetries: 3,
     messages,
     model,
+    runId,
+    sessionConnections,
     stopSequences: [globalStopSequence],
-    clientSessionId,
-    fingerprintId,
-    userInputId,
+    tools,
     userId,
-    maxOutputTokens: 32_000,
+    userInputId,
+
     onCostCalculated,
-    includeCacheControl,
-    agentId,
-    maxRetries: 3,
     sendAction,
-    liveUserInputRecord,
-    sessionConnections,
-    logger,
     trackEvent,
   }
 

packages/agent-runtime/src/run-agent-step.ts

@@ -14,6 +14,7 @@ import { runProgrammaticStep } from './run-programmatic-step'
 import { additionalSystemPrompts } from './system-prompt/prompts'
 import { getAgentTemplate } from './templates/agent-registry'
 import { getAgentPrompt } from './templates/strings'
+import { getToolSet } from './tools/prompts'
 import { processStreamWithTools } from './tools/stream-parser'
 import { getAgentOutput } from './util/agent-output'
 import {
@@ -527,6 +528,7 @@ export async function loopAgentSteps(
       | 'spawnParams'
       | 'system'
       | 'textOverride'
+      | 'tools'
     > &
     ParamsExcluding<
       AddAgentStepFn,
@@ -631,6 +633,19 @@ export async function loopAgentSteps(
           },
         })) ?? ''
 
+  const tools = await getToolSet({
+    toolNames: agentTemplate.toolNames,
+    additionalToolDefinitions: async () => {
+      if (!cachedAdditionalToolDefinitions) {
+        cachedAdditionalToolDefinitions = await additionalToolDefinitions({
+          ...params,
+          agentTemplate,
+        })
+      }
+      return cachedAdditionalToolDefinitions
+    },
+  })
+
   const hasUserMessage = Boolean(
     prompt || (spawnParams && Object.keys(spawnParams).length > 0),
   )
@@ -786,6 +801,16 @@ export async function loopAgentSteps(
         nResponses: generatedResponses,
       } = await runAgentStep({
         ...params,
+
+        agentState: currentAgentState,
+        n,
+        prompt: currentPrompt,
+        runId,
+        spawnParams: currentParams,
+        system,
+        textOverride: textOverride,
+        tools,
+
         additionalToolDefinitions: async () => {
           if (!cachedAdditionalToolDefinitions) {
             cachedAdditionalToolDefinitions = await additionalToolDefinitions({
@@ -795,13 +820,6 @@ export async function loopAgentSteps(
           }
           return cachedAdditionalToolDefinitions
         },
-        textOverride: textOverride,
-        runId,
-        agentState: currentAgentState,
-        prompt: currentPrompt,
-        spawnParams: currentParams,
-        system,
-        n,
       })
 
       if (newAgentState.runId) {