Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,28 @@ describe('runProgrammaticStep', () => {
})

describe('tool execution', () => {
it('assigns deterministic per-tool ids to handleSteps tool calls', async () => {
const mockGenerator = (function* () {
yield { toolName: 'read_files', input: { paths: ['first.txt'] } }
yield { toolName: 'read_files', input: { paths: ['second.txt'] } }
yield { toolName: 'end_turn', input: {} }
})() as StepGenerator

mockTemplate.handleSteps = () => mockGenerator

await runProgrammaticStep(mockParams)

expect(executeToolCallSpy.mock.calls[0][0].toolCallId).toBe(
'functions.read_files:0',
)
expect(executeToolCallSpy.mock.calls[1][0].toolCallId).toBe(
'functions.read_files:1',
)
expect(executeToolCallSpy.mock.calls[2][0].toolCallId).toBe(
'functions.end_turn:0',
)
})

it('should not add tool call message for add_message tool', async () => {
const mockGenerator = (function* () {
yield {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -401,13 +401,16 @@ describe('tool validation error handling', () => {
)
expect(toolCallEvents.length).toBe(1)
expect(toolCallEvents[0].toolName).toBe('read_files')
expect(toolCallEvents[0].toolCallId).toBe('functions.read_files:0')

// Verify tool_result event was emitted
const toolResultEvents = responseChunks.filter(
(chunk): chunk is Extract<PrintModeEvent, { type: 'tool_result' }> =>
typeof chunk !== 'string' && chunk.type === 'tool_result',
)
expect(toolResultEvents.length).toBe(1)
expect(toolResultEvents[0].toolName).toBe('read_files')
expect(toolResultEvents[0].toolCallId).toBe('functions.read_files:0')

// Verify NO error events
const errorEvents = responseChunks.filter(
Expand Down
10 changes: 7 additions & 3 deletions packages/agent-runtime/src/run-programmatic-step.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import { cloneDeep } from 'lodash'
import { clearProposedContentForRun } from './tools/handlers/tool/proposed-content-store'
import { executeToolCall } from './tools/tool-executor'
import { parseTextWithToolCalls } from './util/parse-tool-calls-from-text'

import { createToolCallIdGenerator } from './util/tool-call-id'

import type { FileProcessingState } from './tools/handlers/tool/write-file'
import type { ExecuteToolCallParams } from './tools/tool-executor'
Expand Down Expand Up @@ -213,6 +213,7 @@ export async function runProgrammaticStep(
let toolResult: ToolResultOutput[] | undefined = undefined
let endTurn = false
let generateN: number | undefined = undefined
const getToolCallId = createToolCallIdGenerator(agentState.messageHistory)

let startTime = new Date()
let creditsBefore = agentState.directCreditsUsed
Expand Down Expand Up @@ -273,6 +274,7 @@ export async function runProgrammaticStep(
previousToolCallFinished: Promise.resolve(),
toolCalls,
toolResults,
getToolCallId,
onResponseChunk,
})
}
Expand Down Expand Up @@ -301,6 +303,7 @@ export async function runProgrammaticStep(
previousToolCallFinished: Promise.resolve(),
toolCalls,
toolResults,
getToolCallId,
onResponseChunk,
})

Expand Down Expand Up @@ -432,6 +435,7 @@ type ExecuteToolCallsArrayParams = Omit<
| 'toolResultsToAddToMessageHistory'
> & {
agentState: AgentState
getToolCallId: (toolName: string) => string
onResponseChunk: (chunk: string | PrintModeEvent) => void
}

Expand All @@ -445,7 +449,7 @@ async function executeSingleToolCall(
toolCallToExecute: ToolCallToExecute,
params: ExecuteToolCallsArrayParams,
): Promise<ToolResultOutput[] | undefined> {
const { agentState, onResponseChunk, toolResults } = params
const { agentState, getToolCallId, onResponseChunk, toolResults } = params

// Note: We don't check if the tool is available for the agent template anymore.
// You can run any tool from handleSteps now!
Expand All @@ -455,7 +459,7 @@ async function executeSingleToolCall(
// )
// }

const toolCallId = crypto.randomUUID()
const toolCallId = getToolCallId(toolCallToExecute.toolName)
const excludeToolFromMessageHistory =
toolCallToExecute.includeToolCall === false

Expand Down
4 changes: 0 additions & 4 deletions packages/agent-runtime/src/tool-stream-parser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,6 @@ export async function* processStreamWithTools(params: {
}
trackEvent: TrackEventFn
executeXmlToolCall: (params: {
toolCallId: string
toolName: string
input: Record<string, unknown>
}) => Promise<void>
Expand Down Expand Up @@ -150,12 +149,9 @@ export async function* processStreamWithTools(params: {

// Then process and yield any XML tool calls found
for (const toolCall of toolCalls) {
const toolCallId = `xml-${crypto.randomUUID().slice(0, 8)}`

// Execute the tool immediately if callback provided, pausing the stream
// The callback handles emitting tool_call and tool_result events
await executeXmlToolCall({
toolCallId,
toolName: toolCall.toolName,
input: toolCall.input,
})
Expand Down
15 changes: 8 additions & 7 deletions packages/agent-runtime/src/tools/stream-parser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import {
assistantMessage,
userMessage,
} from '@codebuff/common/util/messages'
import { generateCompactId } from '@codebuff/common/util/string'

import { processStreamWithTools } from '../tool-stream-parser'
import { INCLUDE_REASONING_IN_MESSAGE_HISTORY } from '../constants'
Expand All @@ -14,6 +13,7 @@ import {
executeToolCall,
tryTransformAgentToolCall,
} from './tool-executor'
import { createToolCallIdGenerator } from '../util/tool-call-id'
import { withSystemTags } from '../util/messages'

import type { CustomToolCall, ExecuteToolCallParams } from './tool-executor'
Expand Down Expand Up @@ -91,6 +91,7 @@ export async function processStream(
const toolCalls: (CodebuffToolCall | CustomToolCall)[] = []
const toolCallsToAddToMessageHistory: (CodebuffToolCall | CustomToolCall)[] = []
const assistantMessages: Message[] = []
const getToolCallId = createToolCallIdGenerator(params.messages)
let hadToolCallError = false
const errorMessages: Message[] = []
const { promise: streamDonePromise, resolve: resolveStreamDonePromise } =
Expand Down Expand Up @@ -137,7 +138,6 @@ export async function processStream(
if (signal.aborted) {
return
}
const toolCallId = generateCompactId()
const isNativeTool = toolNames.includes(toolName as ToolName)

// Check if this is an agent tool call that should be transformed to spawn_agents
Expand All @@ -160,19 +160,20 @@ export async function processStream(
// Determine which executor to use and with what parameters
let toolPromise: Promise<void>
if (isNativeTool || transformed) {
const effectiveToolName = transformed
? transformed.toolName
: (toolName as ToolName)
// Use executeToolCall for native tools or transformed agent calls
toolPromise = executeToolCall({
...params,
toolName: transformed
? transformed.toolName
: (toolName as ToolName),
toolName: effectiveToolName,
input: transformed ? transformed.input : input,
fromHandleSteps: false,

fileProcessingState,
fullResponse: fullResponseChunks.join(''),
previousToolCallFinished: previousPromise,
toolCallId,
toolCallId: getToolCallId(effectiveToolName),
toolCalls,
toolCallsToAddToMessageHistory,
toolResults,
Expand All @@ -191,7 +192,7 @@ export async function processStream(
fileProcessingState,
fullResponse: fullResponseChunks.join(''),
previousToolCallFinished: previousPromise,
toolCallId,
toolCallId: getToolCallId(toolName),
toolCalls,
toolCallsToAddToMessageHistory,
toolResults,
Expand Down
12 changes: 9 additions & 3 deletions packages/agent-runtime/src/tools/tool-executor.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
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'
import { getAgentShortName, getAgentToolName } from '../templates/prompts'
import { formatValueForError } from '../util/format-value'
import { createToolCallIdGenerator } from '../util/tool-call-id'
import { codebuffToolHandlers } from './handlers/list'
import { getMatchingSpawn } from './handlers/tool/spawn-agent-utils'
import { getAgentTemplate } from '../templates/agent-registry'
Expand Down Expand Up @@ -308,7 +308,9 @@ export async function executeToolCall<T extends ToolName>(
onResponseChunk,
requestToolCall,
} = params
const toolCallId = params.toolCallId ?? generateCompactId()
const toolCallId =
params.toolCallId ??
createToolCallIdGenerator(agentState.messageHistory, toolCalls)(toolName)

const toolCall: CodebuffToolCall<T> | ToolCallError = parseRawToolCall<T>({
rawToolCall: {
Expand Down Expand Up @@ -640,7 +642,11 @@ export async function executeCustomToolCall(
}),
rawToolCall: {
toolName,
toolCallId: toolCallId ?? generateCompactId(),
toolCallId:
toolCallId ??
createToolCallIdGenerator(agentState.messageHistory, toolCalls)(
toolName,
),
input,
},
autoInsertEndStepParam,
Expand Down
63 changes: 63 additions & 0 deletions packages/agent-runtime/src/util/__tests__/tool-call-id.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import { assistantMessage } from '@codebuff/common/util/messages'
import { describe, expect, it } from 'bun:test'

import {
countToolCallsByName,
createToolCallIdGenerator,
formatToolCallId,
} from '../tool-call-id'

describe('tool call ids', () => {
it('formats ids with the tool name and per-tool invocation index', () => {
expect(formatToolCallId('glob', 0)).toBe('functions.glob:0')
})

it('seeds per-tool counters from existing message history', () => {
const messages = [
assistantMessage({
type: 'tool-call',
toolName: 'glob',
toolCallId: 'functions.glob:0',
input: { pattern: '**/*.ts' },
}),
assistantMessage({
type: 'tool-call',
toolName: 'read_files',
toolCallId: 'functions.read_files:0',
input: { paths: ['src/index.ts'] },
}),
assistantMessage({
type: 'tool-call',
toolName: 'glob',
toolCallId: 'functions.glob:1',
input: { pattern: '**/*.tsx' },
}),
]

expect(countToolCallsByName(messages)).toEqual(
new Map([
['glob', 2],
['read_files', 1],
]),
)

const getToolCallId = createToolCallIdGenerator(messages)

expect(getToolCallId('glob')).toBe('functions.glob:2')
expect(getToolCallId('glob')).toBe('functions.glob:3')
expect(getToolCallId('read_files')).toBe('functions.read_files:1')
})

it('can seed counters from pending tool calls', () => {
const getToolCallId = createToolCallIdGenerator([], [
{
toolName: 'glob',
},
{
toolName: 'glob',
},
])

expect(getToolCallId('glob')).toBe('functions.glob:2')
})
})
48 changes: 48 additions & 0 deletions packages/agent-runtime/src/util/tool-call-id.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import type { Message } from '@codebuff/common/types/messages/codebuff-message'

const TOOL_CALL_ID_PREFIX = 'functions'
type ToolCallLike = { toolName: string }

export function formatToolCallId(toolName: string, index: number): string {
return `${TOOL_CALL_ID_PREFIX}.${toolName}:${index}`
}

export function countToolCallsByName(
messages: Message[],
pendingToolCalls: ToolCallLike[] = [],
): Map<string, number> {
const counts = new Map<string, number>()

for (const message of messages) {
if (message.role !== 'assistant') {
continue
}

for (const part of message.content) {
if (part.type !== 'tool-call') {
continue
}

counts.set(part.toolName, (counts.get(part.toolName) ?? 0) + 1)
}
}

for (const toolCall of pendingToolCalls) {
counts.set(toolCall.toolName, (counts.get(toolCall.toolName) ?? 0) + 1)
}

return counts
}

export function createToolCallIdGenerator(
messages: Message[],
pendingToolCalls: ToolCallLike[] = [],
) {
const counts = countToolCallsByName(messages, pendingToolCalls)

return (toolName: string): string => {
const index = counts.get(toolName) ?? 0
counts.set(toolName, index + 1)
return formatToolCallId(toolName, index)
}
}
Loading