Fix: add tool calls appropriately

Author: jahooma

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

@@ -137,16 +137,16 @@ export async function runProgrammaticStep(
   if (!generator) {
     const createLogMethod =
       (level: 'debug' | 'info' | 'warn' | 'error') =>
-      (data: any, msg?: string) => {
-        logger[level](data, msg) // Log to backend
-        handleStepsLogChunk({
-          userInputId,
-          runId: agentState.runId ?? 'undefined',
-          level,
-          data,
-          message: msg,
-        })
-      }
+        (data: any, msg?: string) => {
+          logger[level](data, msg) // Log to backend
+          handleStepsLogChunk({
+            userInputId,
+            runId: agentState.runId ?? 'undefined',
+            level,
+            data,
+            message: msg,
+          })
+        }
 
     const streamingLogger = {
       debug: createLogMethod('debug'),
@@ -243,7 +243,7 @@ export async function runProgrammaticStep(
       if (!parseResult.success) {
         throw new Error(
           `Invalid yield value from handleSteps in agent ${template.id}: ${parseResult.error.message}. ` +
-            `Received: ${JSON.stringify(result.value)}`,
+          `Received: ${JSON.stringify(result.value)}`,
         )
       }
 
@@ -334,9 +334,8 @@ export async function runProgrammaticStep(
   } catch (error) {
     endTurn = true
 
-    const errorMessage = `Error executing handleSteps for agent ${template.id}: ${
-      error instanceof Error ? error.message : 'Unknown error'
-    }`
+    const errorMessage = `Error executing handleSteps for agent ${template.id}: ${error instanceof Error ? error.message : 'Unknown error'
+      }`
     logger.error(
       { error: getErrorObject(error), template: template.id },
       errorMessage,
@@ -485,6 +484,7 @@ async function executeSingleToolCall(
     // })
   }
 
+  const toolResultsToAddToMessageHistory: ToolMessage[] = []
   // Execute the tool call
   await executeToolCall({
     ...params,
@@ -494,7 +494,9 @@ async function executeSingleToolCall(
     excludeToolFromMessageHistory,
     fromHandleSteps: true,
     toolCallId,
-    toolResultsToAddToMessageHistory: [],
+    toolCalls: [],
+    toolCallsToAddToMessageHistory: [],
+    toolResultsToAddToMessageHistory,
 
     onResponseChunk: (chunk: string | PrintModeEvent) => {
       if (typeof chunk === 'string') {
@@ -539,6 +541,8 @@ async function executeSingleToolCall(
     },
   })
 
+  agentState.messageHistory.push(...toolResultsToAddToMessageHistory)
+
   // Get the latest tool result
   return toolResults[toolResults.length - 1]?.content
 }

packages/agent-runtime/src/tools/stream-parser.ts

@@ -87,6 +87,7 @@ export async function processStream(
   const toolResults: ToolMessage[] = []
   const toolResultsToAddToMessageHistory: ToolMessage[] = []
   const toolCalls: (CodebuffToolCall | CustomToolCall)[] = []
+  const toolCallsToAddToMessageHistory: (CodebuffToolCall | CustomToolCall)[] = []
   const assistantMessages: Message[] = []
   let hadToolCallError = false
   const errorMessages: Message[] = []
@@ -108,13 +109,7 @@ export async function processStream(
   function createResponseHandler(isXmlMode: boolean) {
     return (chunk: string | PrintModeEvent) => {
       if (typeof chunk !== 'string') {
-        if (chunk.type === 'tool_call') {
-          if (chunk.includeToolCall !== false) {
-            assistantMessages.push(
-              assistantMessage({ ...chunk, type: 'tool-call' }),
-            )
-          }
-        } else if (isXmlMode && chunk.type === 'tool_result') {
+        if (isXmlMode && chunk.type === 'tool_result') {
           const toolResultMessage: ToolMessage = {
             role: 'tool',
             toolName: chunk.toolName,
@@ -188,6 +183,7 @@ export async function processStream(
             previousToolCallFinished: previousPromise,
             toolCallId,
             toolCalls,
+            toolCallsToAddToMessageHistory,
             toolResults,
             toolResultsToAddToMessageHistory: resultsArray,
             excludeToolFromMessageHistory: false,
@@ -206,6 +202,7 @@ export async function processStream(
             previousToolCallFinished: previousPromise,
             toolCallId,
             toolCalls,
+            toolCallsToAddToMessageHistory,
             toolResults,
             toolResultsToAddToMessageHistory: resultsArray,
             excludeToolFromMessageHistory: false,
@@ -330,6 +327,7 @@ export async function processStream(
   agentState.messageHistory = buildArray<Message>([
     ...messageHistoryBeforeStream,
     ...assistantMessages,
+    ...toolCallsToAddToMessageHistory.map((toolCall) => assistantMessage({ ...toolCall, type: 'tool-call' })),
     ...toolResultsToAddToMessageHistory,
     ...errorMessages,
   ])

packages/agent-runtime/src/tools/tool-executor.ts

@@ -33,7 +33,7 @@ import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ToolMessage } from '@codebuff/common/types/messages/codebuff-message'
 import type { ToolResultOutput } from '@codebuff/common/types/messages/content-part'
 import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
-import type { AgentTemplateType , AgentState, Subgoal } from '@codebuff/common/types/session-state'
+import type { AgentTemplateType, AgentState, Subgoal } from '@codebuff/common/types/session-state'
 import type {
   CustomToolDefinitions,
   ProjectFileContext,
@@ -119,6 +119,7 @@ export type ExecuteToolCallParams<T extends string = ToolName> = {
   tools: ToolSet
   toolCallId: string | undefined
   toolCalls: (CodebuffToolCall | CustomToolCall)[]
+  toolCallsToAddToMessageHistory: (CodebuffToolCall | CustomToolCall)[]
   toolResults: ToolMessage[]
   toolResultsToAddToMessageHistory: ToolMessage[]
   userId: string | undefined
@@ -144,6 +145,7 @@ export async function executeToolCall<T extends ToolName>(
     logger,
     previousToolCallFinished,
     toolCalls,
+    toolCallsToAddToMessageHistory,
     toolResults,
     toolResultsToAddToMessageHistory,
     userInputId,
@@ -298,8 +300,6 @@ export async function executeToolCall<T extends ToolName>(
     includeToolCall: !excludeToolFromMessageHistory,
   })
 
-  toolCalls.push(toolCall)
-
   // Cast to any to avoid type errors
   const handler = codebuffToolHandlers[
     toolName
@@ -311,6 +311,12 @@ export async function executeToolCall<T extends ToolName>(
       ? { ...toolCall, input: effectiveInput }
       : toolCall
 
+  toolCalls.push(finalToolCall)
+  if (!excludeToolFromMessageHistory) {
+    toolCallsToAddToMessageHistory.push(finalToolCall)
+  }
+
+
   const toolResultPromise = handler({
     ...params,
     toolCall: finalToolCall,
@@ -448,6 +454,7 @@ export async function executeCustomToolCall(
     requestToolCall,
     toolCallId,
     toolCalls,
+    toolCallsToAddToMessageHistory,
     toolResults,
     toolResultsToAddToMessageHistory,
     userInputId,
@@ -512,6 +519,9 @@ export async function executeCustomToolCall(
   })
 
   toolCalls.push(toolCall)
+  if (!excludeToolFromMessageHistory) {
+    toolCallsToAddToMessageHistory.push(toolCall)
+  }
 
   return previousToolCallFinished
     .then(async () => {