From 7002f0190a33a5279c7cde523367c812f89600f0 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 10:14:04 -0800 Subject: [PATCH 01/30] Convert text responses to tool calls --- .../agent-runtime/src/tool-stream-parser.ts | 5 +-- .../agent-runtime/src/tools/stream-parser.ts | 39 ++++++++++++++++++- 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/packages/agent-runtime/src/tool-stream-parser.ts b/packages/agent-runtime/src/tool-stream-parser.ts index 0191596c4e..e7e07ca433 100644 --- a/packages/agent-runtime/src/tool-stream-parser.ts +++ b/packages/agent-runtime/src/tool-stream-parser.ts @@ -13,7 +13,6 @@ import type { Logger } from '@codebuff/common/types/contracts/logger' import type { PrintModeError, PrintModeText, - PrintModeToolCall, } from '@codebuff/common/types/print-mode' const toolExtractionPattern = new RegExp( @@ -37,9 +36,7 @@ export async function* processStreamWithTags(params: { onTagEnd: (tagName: string, params: Record) => void } onError: (tagName: string, errorMessage: string) => void - onResponseChunk: ( - chunk: PrintModeText | PrintModeToolCall | PrintModeError, - ) => void + onResponseChunk: (chunk: PrintModeText | PrintModeError) => void logger: Logger loggerOptions?: { userId?: string diff --git a/packages/agent-runtime/src/tools/stream-parser.ts b/packages/agent-runtime/src/tools/stream-parser.ts index f47d41a67a..ab2b2d3338 100644 --- a/packages/agent-runtime/src/tools/stream-parser.ts +++ b/packages/agent-runtime/src/tools/stream-parser.ts @@ -80,12 +80,14 @@ export async function processStreamWithTools( runId, signal, userId, + logger, } = params const fullResponseChunks: string[] = [fullResponse] const toolResults: ToolMessage[] = [] const toolResultsToAddAfterStream: ToolMessage[] = [] const toolCalls: (CodebuffToolCall | CustomToolCall)[] = [] + const assistantMessages: Message[] = [] const { promise: streamDonePromise, resolve: resolveStreamDonePromise } = Promise.withResolvers() let previousToolCallFinished = streamDonePromise @@ -122,6 +124,14 @@ export async function processStreamWithTools( toolResultsToAddAfterStream, onCostCalculated, + onResponseChunk: (chunk) => { + if (typeof chunk !== 'string' && chunk.type === 'tool_call') { + assistantMessages.push( + assistantMessage({ ...chunk, type: 'tool-call' }), + ) + } + return onResponseChunk(chunk) + }, }) }, } @@ -147,6 +157,15 @@ export async function processStreamWithTools( toolCalls, toolResults, toolResultsToAddAfterStream, + + onResponseChunk: (chunk) => { + if (typeof chunk !== 'string' && chunk.type === 'tool_call') { + assistantMessages.push( + assistantMessage({ ...chunk, type: 'tool-call' }), + ) + } + return onResponseChunk(chunk) + }, }) }, } @@ -179,6 +198,21 @@ export async function processStreamWithTools( model: agentTemplate.model, agentName: agentTemplate.id, }, + onResponseChunk: (chunk) => { + if (chunk.type === 'text') { + if (chunk.text) { + assistantMessages.push(assistantMessage(chunk.text)) + } + } else if (chunk.type === 'error') { + // do nothing + } else { + chunk satisfies never + throw new Error( + `Internal error: unhandled chunk type: ${(chunk as any).type}`, + ) + } + return onResponseChunk(chunk) + }, }) let messageId: string | null = null @@ -211,11 +245,12 @@ export async function processStreamWithTools( agentState.messageHistory = buildArray([ ...expireMessages(agentState.messageHistory, 'agentStep'), - fullResponseChunks.length > 0 && - assistantMessage(fullResponseChunks.join('')), + ...assistantMessages, ...toolResultsToAddAfterStream, ]) + logger.info({ messages: agentState.messageHistory }, 'asdf messages') + if (!signal.aborted) { resolveStreamDonePromise() await previousToolCallFinished From c939465e0a3d9cfbf87a765bad386838fc138d18 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:06:06 -0800 Subject: [PATCH 02/30] input tools into ai sdk --- .../agent-runtime/src/prompt-agent-stream.ts | 27 +++++++++------- packages/agent-runtime/src/run-agent-step.ts | 32 +++++++++++++++---- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/packages/agent-runtime/src/prompt-agent-stream.ts b/packages/agent-runtime/src/prompt-agent-stream.ts index 3447b59481..04e46378e8 100644 --- a/packages/agent-runtime/src/prompt-agent-stream.ts +++ b/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 = { + 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, } diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index 2d54df692b..e41565da87 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/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) { From 1ecbe4ed4b811b3ed09d94e2cbf2d8d6d57fb0b4 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:16:07 -0800 Subject: [PATCH 03/30] do not convert cb tool messages --- common/src/util/messages.ts | 52 ++++++++++++++----------------------- 1 file changed, 20 insertions(+), 32 deletions(-) diff --git a/common/src/util/messages.ts b/common/src/util/messages.ts index b8387a96f2..478a2d4d22 100644 --- a/common/src/util/messages.ts +++ b/common/src/util/messages.ts @@ -1,7 +1,5 @@ import { cloneDeep, has, isEqual } from 'lodash' -import { getToolCallString } from '../tools/utils' - import type { JSONValue } from '../types/json' import type { AssistantMessage, @@ -100,21 +98,21 @@ function assistantToCodebuffMessage( content: Exclude[number] }, ): AssistantMessage { - if (message.content.type === 'tool-call') { - return cloneDeep({ - ...message, - content: [ - { - type: 'text', - text: getToolCallString( - message.content.toolName, - message.content.input, - false, - ), - }, - ], - }) - } + // if (message.content.type === 'tool-call') { + // return cloneDeep({ + // ...message, + // content: [ + // { + // type: 'text', + // text: getToolCallString( + // message.content.toolName, + // message.content.input, + // false, + // ), + // }, + // ], + // }) + // } return cloneDeep({ ...message, content: [message.content] }) } @@ -123,20 +121,10 @@ function convertToolResultMessage( ): ModelMessageWithAuxiliaryData[] { return message.content.map((c) => { if (c.type === 'json') { - const toolResult = { - toolName: message.toolName, - toolCallId: message.toolCallId, - output: c.value, - } - return cloneDeep({ + return cloneDeep({ ...message, - role: 'user', - content: [ - { - type: 'text', - text: `\n${JSON.stringify(toolResult, null, 2)}\n`, - }, - ], + role: 'tool', + content: [{ ...message, output: c, type: 'tool-result' }], }) } if (c.type === 'media') { @@ -147,8 +135,8 @@ function convertToolResultMessage( }) } c satisfies never - const oAny = c as any - throw new Error(`Invalid tool output type: ${oAny.type}`) + const cAny = c as any + throw new Error(`Invalid tool output type: ${cAny.type}`) }) } From de15639532ad4bff241b4b9fa262f2b76a77a70c Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:25:27 -0800 Subject: [PATCH 04/30] pass tool-call chunk through promptAiSdkStream --- sdk/src/impl/llm.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sdk/src/impl/llm.ts b/sdk/src/impl/llm.ts index ab7c4708fc..27d0e88299 100644 --- a/sdk/src/impl/llm.ts +++ b/sdk/src/impl/llm.ts @@ -339,6 +339,9 @@ export async function* promptAiSdkStream( } } } + if (chunk.type === 'tool-call') { + yield chunk + } } const flushed = stopSequenceHandler.flush() if (flushed) { From eae4f02587389d2eb53f78c1ec187a30af35cb7f Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:26:22 -0800 Subject: [PATCH 05/30] fork tool-stream-parser and do not parse text --- common/src/types/contracts/llm.ts | 8 +- .../src/__tests__/malformed-tool-call.test.ts | 16 +- .../src/__tests__/tool-stream-parser.test.ts | 26 +-- packages/agent-runtime/src/run-agent-step.ts | 7 +- .../src/tool-stream-parser.old.ts | 217 ++++++++++++++++++ .../agent-runtime/src/tool-stream-parser.ts | 141 ++---------- .../agent-runtime/src/tools/stream-parser.ts | 18 +- 7 files changed, 282 insertions(+), 151 deletions(-) create mode 100644 packages/agent-runtime/src/tool-stream-parser.old.ts diff --git a/common/src/types/contracts/llm.ts b/common/src/types/contracts/llm.ts index 23fc5ede72..ac32359954 100644 --- a/common/src/types/contracts/llm.ts +++ b/common/src/types/contracts/llm.ts @@ -1,12 +1,12 @@ import type { TrackEventFn } from './analytics' import type { SendActionFn } from './client' import type { CheckLiveUserInputFn } from './live-user-input' +import type { OpenRouterProviderRoutingOptions } from '../agent-template' 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 { OpenRouterProviderRoutingOptions } from '../agent-template' -import type { generateText, streamText } from 'ai' +import type { generateText, streamText, ToolCallPart } from 'ai' import type z from 'zod/v4' export type StreamChunk = @@ -19,6 +19,10 @@ export type StreamChunk = type: 'reasoning' text: string } + | Pick< + ToolCallPart, + 'type' | 'toolCallId' | 'toolName' | 'input' | 'providerOptions' + > | { type: 'error'; message: string } export type PromptAiSdkStreamFn = ( diff --git a/packages/agent-runtime/src/__tests__/malformed-tool-call.test.ts b/packages/agent-runtime/src/__tests__/malformed-tool-call.test.ts index 8b32ea54a2..a7b9472814 100644 --- a/packages/agent-runtime/src/__tests__/malformed-tool-call.test.ts +++ b/packages/agent-runtime/src/__tests__/malformed-tool-call.test.ts @@ -16,7 +16,7 @@ import { } from 'bun:test' import { mockFileContext } from './test-utils' -import { processStreamWithTools } from '../tools/stream-parser' +import { processStream } from '../tools/stream-parser' import type { AgentTemplate } from '../templates/types' import type { @@ -34,7 +34,7 @@ let agentRuntimeImpl: AgentRuntimeDeps = { ...TEST_AGENT_RUNTIME_IMPL } describe('malformed tool call error handling', () => { let testAgent: AgentTemplate let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps - let defaultParams: ParamsOf + let defaultParams: ParamsOf beforeEach(() => { agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL } @@ -139,7 +139,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - await processStreamWithTools({ + await processStream({ ...defaultParams, stream, }) @@ -177,7 +177,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - await processStreamWithTools({ + await processStream({ ...defaultParams, stream, }) @@ -204,7 +204,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - const result = await processStreamWithTools({ + const result = await processStream({ ...defaultParams, stream, }) @@ -235,7 +235,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - await processStreamWithTools({ + await processStream({ ...defaultParams, stream, }) @@ -268,7 +268,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - await processStreamWithTools({ + await processStream({ ...defaultParams, requestFiles: async ({ filePaths }) => { return Object.fromEntries( @@ -307,7 +307,7 @@ describe('malformed tool call error handling', () => { const stream = createMockStream(chunks) - await processStreamWithTools({ + await processStream({ ...defaultParams, stream, }) diff --git a/packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts b/packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts index b5c5dfb23a..0cff771dab 100644 --- a/packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts +++ b/packages/agent-runtime/src/__tests__/tool-stream-parser.test.ts @@ -4,7 +4,7 @@ import { getToolCallString } from '@codebuff/common/tools/utils' import { beforeEach, describe, expect, it } from 'bun:test' import { globalStopSequence } from '../constants' -import { processStreamWithTags } from '../tool-stream-parser' +import { processStreamWithTools } from '../tool-stream-parser' import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime' @@ -61,7 +61,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -129,7 +129,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -206,7 +206,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -282,7 +282,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -349,7 +349,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -415,7 +415,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -488,7 +488,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -555,7 +555,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -612,7 +612,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -655,7 +655,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -716,7 +716,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, @@ -787,7 +787,7 @@ describe('processStreamWithTags', () => { } } - for await (const chunk of processStreamWithTags({ + for await (const chunk of processStreamWithTools({ ...agentRuntimeImpl, stream, processors, diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index e41565da87..e6a111f694 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -15,7 +15,7 @@ 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 { processStream } from './tools/stream-parser' import { getAgentOutput } from './util/agent-output' import { withSystemInstructionTags, @@ -107,7 +107,7 @@ export const runAgentStep = async ( trackEvent: TrackEventFn promptAiSdk: PromptAiSdkFn } & ParamsExcluding< - typeof processStreamWithTools, + typeof processStream, | 'agentContext' | 'agentState' | 'agentStepId' @@ -338,6 +338,7 @@ export const runAgentStep = async ( let fullResponse = '' const toolResults: ToolMessage[] = [] + // Raw stream from AI SDK const stream = getAgentStreamFromTemplate({ ...params, agentId: agentState.parentId ? agentState.agentId : undefined, @@ -353,7 +354,7 @@ export const runAgentStep = async ( messageId, toolCalls, toolResults: newToolResults, - } = await processStreamWithTools({ + } = await processStream({ ...params, agentContext, agentState, diff --git a/packages/agent-runtime/src/tool-stream-parser.old.ts b/packages/agent-runtime/src/tool-stream-parser.old.ts new file mode 100644 index 0000000000..e7e07ca433 --- /dev/null +++ b/packages/agent-runtime/src/tool-stream-parser.old.ts @@ -0,0 +1,217 @@ +import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' +import { + endsAgentStepParam, + endToolTag, + startToolTag, + toolNameParam, +} from '@codebuff/common/tools/constants' + +import type { Model } from '@codebuff/common/old-constants' +import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics' +import type { StreamChunk } from '@codebuff/common/types/contracts/llm' +import type { Logger } from '@codebuff/common/types/contracts/logger' +import type { + PrintModeError, + PrintModeText, +} from '@codebuff/common/types/print-mode' + +const toolExtractionPattern = new RegExp( + `${startToolTag}(.*?)${endToolTag}`, + 'gs', +) + +const completionSuffix = `${JSON.stringify(endsAgentStepParam)}: true\n}${endToolTag}` + +export async function* processStreamWithTags(params: { + stream: AsyncGenerator + processors: Record< + string, + { + onTagStart: (tagName: string, attributes: Record) => void + onTagEnd: (tagName: string, params: Record) => void + } + > + defaultProcessor: (toolName: string) => { + onTagStart: (tagName: string, attributes: Record) => void + onTagEnd: (tagName: string, params: Record) => void + } + onError: (tagName: string, errorMessage: string) => void + onResponseChunk: (chunk: PrintModeText | PrintModeError) => void + logger: Logger + loggerOptions?: { + userId?: string + model?: Model + agentName?: string + } + trackEvent: TrackEventFn +}): AsyncGenerator { + const { + stream, + processors, + defaultProcessor, + onError, + onResponseChunk, + logger, + loggerOptions, + trackEvent, + } = params + + let streamCompleted = false + let buffer = '' + let autocompleted = false + + function extractToolCalls(): string[] { + const matches: string[] = [] + let lastIndex = 0 + for (const match of buffer.matchAll(toolExtractionPattern)) { + if (match.index > lastIndex) { + onResponseChunk({ + type: 'text', + text: buffer.slice(lastIndex, match.index), + }) + } + lastIndex = match.index + match[0].length + matches.push(match[1]) + } + + buffer = buffer.slice(lastIndex) + return matches + } + + function processToolCallContents(contents: string): void { + let parsedParams: any + try { + parsedParams = JSON.parse(contents) + } catch (error: any) { + trackEvent({ + event: AnalyticsEvent.MALFORMED_TOOL_CALL_JSON, + userId: loggerOptions?.userId ?? '', + properties: { + contents: JSON.stringify(contents), + model: loggerOptions?.model, + agent: loggerOptions?.agentName, + error: { + name: error.name, + message: error.message, + stack: error.stack, + }, + autocompleted, + }, + logger, + }) + const shortenedContents = + contents.length < 200 + ? contents + : contents.slice(0, 100) + '...' + contents.slice(-100) + const errorMessage = `Invalid JSON: ${JSON.stringify(shortenedContents)}\nError: ${error.message}` + onResponseChunk({ + type: 'error', + message: errorMessage, + }) + onError('parse_error', errorMessage) + return + } + + const toolName = parsedParams[toolNameParam] as keyof typeof processors + const processor = + typeof toolName === 'string' + ? processors[toolName] ?? defaultProcessor(toolName) + : undefined + if (!processor) { + trackEvent({ + event: AnalyticsEvent.UNKNOWN_TOOL_CALL, + userId: loggerOptions?.userId ?? '', + properties: { + contents, + toolName, + model: loggerOptions?.model, + agent: loggerOptions?.agentName, + autocompleted, + }, + logger, + }) + onError( + 'parse_error', + `Unknown tool ${JSON.stringify(toolName)} for tool call: ${contents}`, + ) + return + } + + trackEvent({ + event: AnalyticsEvent.TOOL_USE, + userId: loggerOptions?.userId ?? '', + properties: { + toolName, + contents, + parsedParams, + autocompleted, + model: loggerOptions?.model, + agent: loggerOptions?.agentName, + }, + logger, + }) + delete parsedParams[toolNameParam] + + processor.onTagStart(toolName, {}) + processor.onTagEnd(toolName, parsedParams) + } + + function extractToolsFromBufferAndProcess(forceFlush = false) { + const matches = extractToolCalls() + matches.forEach(processToolCallContents) + if (forceFlush) { + onResponseChunk({ + type: 'text', + text: buffer, + }) + buffer = '' + } + } + + function* processChunk( + chunk: StreamChunk | undefined, + ): Generator { + if (chunk !== undefined && chunk.type === 'text') { + buffer += chunk.text + } + extractToolsFromBufferAndProcess() + + if (chunk === undefined) { + streamCompleted = true + if (buffer.includes(startToolTag)) { + buffer += completionSuffix + chunk = { + type: 'text', + text: completionSuffix, + } + autocompleted = true + } + extractToolsFromBufferAndProcess(true) + } + + if (chunk) { + yield chunk + } + } + + let messageId: string | null = null + while (true) { + const { value, done } = await stream.next() + if (done) { + messageId = value + break + } + if (streamCompleted) { + break + } + + yield* processChunk(value) + } + + if (!streamCompleted) { + // After the stream ends, try parsing one last time in case there's leftover text + yield* processChunk(undefined) + } + + return messageId +} diff --git a/packages/agent-runtime/src/tool-stream-parser.ts b/packages/agent-runtime/src/tool-stream-parser.ts index e7e07ca433..e5fd319332 100644 --- a/packages/agent-runtime/src/tool-stream-parser.ts +++ b/packages/agent-runtime/src/tool-stream-parser.ts @@ -1,10 +1,4 @@ import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events' -import { - endsAgentStepParam, - endToolTag, - startToolTag, - toolNameParam, -} from '@codebuff/common/tools/constants' import type { Model } from '@codebuff/common/old-constants' import type { TrackEventFn } from '@codebuff/common/types/contracts/analytics' @@ -15,14 +9,7 @@ import type { PrintModeText, } from '@codebuff/common/types/print-mode' -const toolExtractionPattern = new RegExp( - `${startToolTag}(.*?)${endToolTag}`, - 'gs', -) - -const completionSuffix = `${JSON.stringify(endsAgentStepParam)}: true\n}${endToolTag}` - -export async function* processStreamWithTags(params: { +export async function* processStreamWithTools(params: { stream: AsyncGenerator processors: Record< string, @@ -55,87 +42,18 @@ export async function* processStreamWithTags(params: { loggerOptions, trackEvent, } = params - let streamCompleted = false let buffer = '' let autocompleted = false - function extractToolCalls(): string[] { - const matches: string[] = [] - let lastIndex = 0 - for (const match of buffer.matchAll(toolExtractionPattern)) { - if (match.index > lastIndex) { - onResponseChunk({ - type: 'text', - text: buffer.slice(lastIndex, match.index), - }) - } - lastIndex = match.index + match[0].length - matches.push(match[1]) - } + function processToolCallObject(params: { + toolName: string + input: any + contents?: string + }): void { + const { toolName, input, contents } = params - buffer = buffer.slice(lastIndex) - return matches - } - - function processToolCallContents(contents: string): void { - let parsedParams: any - try { - parsedParams = JSON.parse(contents) - } catch (error: any) { - trackEvent({ - event: AnalyticsEvent.MALFORMED_TOOL_CALL_JSON, - userId: loggerOptions?.userId ?? '', - properties: { - contents: JSON.stringify(contents), - model: loggerOptions?.model, - agent: loggerOptions?.agentName, - error: { - name: error.name, - message: error.message, - stack: error.stack, - }, - autocompleted, - }, - logger, - }) - const shortenedContents = - contents.length < 200 - ? contents - : contents.slice(0, 100) + '...' + contents.slice(-100) - const errorMessage = `Invalid JSON: ${JSON.stringify(shortenedContents)}\nError: ${error.message}` - onResponseChunk({ - type: 'error', - message: errorMessage, - }) - onError('parse_error', errorMessage) - return - } - - const toolName = parsedParams[toolNameParam] as keyof typeof processors - const processor = - typeof toolName === 'string' - ? processors[toolName] ?? defaultProcessor(toolName) - : undefined - if (!processor) { - trackEvent({ - event: AnalyticsEvent.UNKNOWN_TOOL_CALL, - userId: loggerOptions?.userId ?? '', - properties: { - contents, - toolName, - model: loggerOptions?.model, - agent: loggerOptions?.agentName, - autocompleted, - }, - logger, - }) - onError( - 'parse_error', - `Unknown tool ${JSON.stringify(toolName)} for tool call: ${contents}`, - ) - return - } + const processor = processors[toolName] ?? defaultProcessor(toolName) trackEvent({ event: AnalyticsEvent.TOOL_USE, @@ -143,55 +61,47 @@ export async function* processStreamWithTags(params: { properties: { toolName, contents, - parsedParams, + parsedParams: input, autocompleted, model: loggerOptions?.model, agent: loggerOptions?.agentName, }, logger, }) - delete parsedParams[toolNameParam] processor.onTagStart(toolName, {}) - processor.onTagEnd(toolName, parsedParams) + processor.onTagEnd(toolName, input) } - function extractToolsFromBufferAndProcess(forceFlush = false) { - const matches = extractToolCalls() - matches.forEach(processToolCallContents) - if (forceFlush) { + function flush() { + if (buffer) { onResponseChunk({ type: 'text', text: buffer, }) - buffer = '' } + buffer = '' } function* processChunk( chunk: StreamChunk | undefined, ): Generator { - if (chunk !== undefined && chunk.type === 'text') { - buffer += chunk.text - } - extractToolsFromBufferAndProcess() - if (chunk === undefined) { streamCompleted = true - if (buffer.includes(startToolTag)) { - buffer += completionSuffix - chunk = { - type: 'text', - text: completionSuffix, - } - autocompleted = true - } - extractToolsFromBufferAndProcess(true) + return } - if (chunk) { - yield chunk + if (chunk.type === 'text') { + buffer += chunk.text + } else { + flush() } + + if (chunk.type === 'tool-call') { + processToolCallObject(chunk) + } + + yield chunk } let messageId: string | null = null @@ -204,14 +114,11 @@ export async function* processStreamWithTags(params: { if (streamCompleted) { break } - yield* processChunk(value) } - if (!streamCompleted) { // After the stream ends, try parsing one last time in case there's leftover text yield* processChunk(undefined) } - return messageId } diff --git a/packages/agent-runtime/src/tools/stream-parser.ts b/packages/agent-runtime/src/tools/stream-parser.ts index ab2b2d3338..224df9533e 100644 --- a/packages/agent-runtime/src/tools/stream-parser.ts +++ b/packages/agent-runtime/src/tools/stream-parser.ts @@ -7,7 +7,7 @@ import { import { generateCompactId } from '@codebuff/common/util/string' import { cloneDeep } from 'lodash' -import { processStreamWithTags } from '../tool-stream-parser' +import { processStreamWithTools } from '../tool-stream-parser' import { executeCustomToolCall, executeToolCall } from './tool-executor' import { expireMessages } from '../util/messages' @@ -33,7 +33,7 @@ export type ToolCallError = { error: string } & Omit -export async function processStreamWithTools( +export async function processStream( params: { agentContext: Record agentTemplate: AgentTemplate @@ -65,7 +65,7 @@ export async function processStreamWithTools( | 'toolResultsToAddAfterStream' > & ParamsExcluding< - typeof processStreamWithTags, + typeof processStreamWithTools, 'processors' | 'defaultProcessor' | 'onError' | 'loggerOptions' >, ) { @@ -171,14 +171,13 @@ export async function processStreamWithTools( } } - const streamWithTags = processStreamWithTags({ + const streamWithTags = processStreamWithTools({ ...params, processors: Object.fromEntries([ ...toolNames.map((toolName) => [toolName, toolCallback(toolName)]), - ...Object.keys(fileContext.customToolDefinitions ?? {}).map((toolName) => [ - toolName, - customToolCallback(toolName), - ]), + ...Object.keys(fileContext.customToolDefinitions ?? {}).map( + (toolName) => [toolName, customToolCallback(toolName)], + ), ]), defaultProcessor: customToolCallback, onError: (toolName, error) => { @@ -238,8 +237,11 @@ export async function processStreamWithTools( fullResponseChunks.push(chunk.text) } else if (chunk.type === 'error') { onResponseChunk(chunk) + } else if (chunk.type === 'tool-call') { + // Do nothing, the onResponseChunk for tool is handled in the processor } else { chunk satisfies never + throw new Error(`Unhandled chunk type: ${(chunk as any).type}`) } } From 619fc93002bb5c4e8998d4b2189a9185941ea670 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:34:20 -0800 Subject: [PATCH 06/30] change example tool call format --- common/src/tools/params/tool/add-message.ts | 7 ++- common/src/tools/params/tool/add-subgoal.ts | 4 +- common/src/tools/params/tool/ask-user.ts | 60 ++++++++++++++----- common/src/tools/params/tool/browser-logs.ts | 4 +- common/src/tools/params/tool/code-search.ts | 14 ++--- common/src/tools/params/tool/create-plan.ts | 4 +- common/src/tools/params/tool/end-turn.ts | 11 ++-- common/src/tools/params/tool/find-files.ts | 4 +- common/src/tools/params/tool/glob.ts | 4 +- .../src/tools/params/tool/list-directory.ts | 6 +- .../tools/params/tool/lookup-agent-info.ts | 4 +- common/src/tools/params/tool/read-docs.ts | 6 +- common/src/tools/params/tool/read-files.ts | 4 +- common/src/tools/params/tool/read-subtree.ts | 4 +- .../params/tool/run-file-change-hooks.ts | 4 +- .../tools/params/tool/run-terminal-command.ts | 6 +- common/src/tools/params/tool/set-messages.ts | 7 ++- common/src/tools/params/tool/set-output.ts | 4 +- .../tools/params/tool/spawn-agent-inline.ts | 7 ++- common/src/tools/params/tool/spawn-agents.ts | 4 +- common/src/tools/params/tool/str-replace.ts | 4 +- .../src/tools/params/tool/task-completed.ts | 11 ++-- common/src/tools/params/tool/think-deeply.ts | 7 ++- .../src/tools/params/tool/update-subgoal.ts | 10 ++-- common/src/tools/params/tool/web-search.ts | 6 +- common/src/tools/params/tool/write-file.ts | 6 +- common/src/tools/params/tool/write-todos.ts | 4 +- common/src/tools/params/utils.ts | 14 +++++ 28 files changed, 147 insertions(+), 83 deletions(-) diff --git a/common/src/tools/params/tool/add-message.ts b/common/src/tools/params/tool/add-message.ts index 2866cc2d3d..7312fdbd53 100644 --- a/common/src/tools/params/tool/add-message.ts +++ b/common/src/tools/params/tool/add-message.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -16,7 +19,7 @@ const inputSchema = z ) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/add-subgoal.ts b/common/src/tools/params/tool/add-subgoal.ts index ed592797b9..0630e76de7 100644 --- a/common/src/tools/params/tool/add-subgoal.ts +++ b/common/src/tools/params/tool/add-subgoal.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -32,7 +32,7 @@ const inputSchema = z ) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/ask-user.ts b/common/src/tools/params/tool/ask-user.ts index 8a228de460..dc83e1618f 100644 --- a/common/src/tools/params/tool/ask-user.ts +++ b/common/src/tools/params/tool/ask-user.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -10,11 +10,16 @@ export const questionSchema = z.object({ .string() .max(12) .optional() - .describe('Short label (max 12 chars) displayed as a chip/tag. Example: "Auth method"'), + .describe( + 'Short label (max 12 chars) displayed as a chip/tag. Example: "Auth method"', + ), options: z .object({ label: z.string().describe('The display text for this option'), - description: z.string().optional().describe('Explanation shown when option is focused'), + description: z + .string() + .optional() + .describe('Explanation shown when option is focused'), }) .array() .refine((opts) => opts.length >= 2, { @@ -30,10 +35,22 @@ export const questionSchema = z.object({ ), validation: z .object({ - maxLength: z.number().optional().describe('Maximum length for "Other" text input'), - minLength: z.number().optional().describe('Minimum length for "Other" text input'), - pattern: z.string().optional().describe('Regex pattern for "Other" text input'), - patternError: z.string().optional().describe('Custom error message when pattern fails'), + maxLength: z + .number() + .optional() + .describe('Maximum length for "Other" text input'), + minLength: z + .number() + .optional() + .describe('Minimum length for "Other" text input'), + pattern: z + .string() + .optional() + .describe('Regex pattern for "Other" text input'), + patternError: z + .string() + .optional() + .describe('Custom error message when pattern fails'), }) .optional() .describe('Validation rules for "Other" text input'), @@ -67,14 +84,20 @@ const outputSchema = z.object({ .array(z.string()) .optional() .describe('Array of selected option texts (multi-select mode)'), - otherText: z.string().optional().describe('Custom text input (if user typed their own answer)'), + otherText: z + .string() + .optional() + .describe('Custom text input (if user typed their own answer)'), }), ) .optional() .describe( 'Array of user answers, one per question. Each answer has either selectedOption (single), selectedOptions (multi), or otherText.', ), - skipped: z.boolean().optional().describe('True if user skipped the questions'), + skipped: z + .boolean() + .optional() + .describe('True if user skipped the questions'), }) const description = ` @@ -87,7 +110,7 @@ The user can either: - Skip the questions to provide different instructions instead Single-select example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -96,9 +119,18 @@ ${$getToolCallString({ question: 'Which authentication method should we use?', header: 'Auth method', options: [ - { label: 'JWT tokens', description: 'Stateless tokens stored in localStorage' }, - { label: 'Session cookies', description: 'Server-side sessions with httpOnly cookies' }, - { label: 'OAuth2', description: 'Third-party authentication (Google, GitHub, etc.)' }, + { + label: 'JWT tokens', + description: 'Stateless tokens stored in localStorage', + }, + { + label: 'Session cookies', + description: 'Server-side sessions with httpOnly cookies', + }, + { + label: 'OAuth2', + description: 'Third-party authentication (Google, GitHub, etc.)', + }, ], }, ], @@ -107,7 +139,7 @@ ${$getToolCallString({ })} Multi-select example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/browser-logs.ts b/common/src/tools/params/tool/browser-logs.ts index acb4d51d94..742c2168ca 100644 --- a/common/src/tools/params/tool/browser-logs.ts +++ b/common/src/tools/params/tool/browser-logs.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { BrowserResponseSchema } from '../../../browser-actions' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -64,7 +64,7 @@ Navigate: - \`waitUntil\`: (required) One of 'load', 'domcontentloaded', 'networkidle0' Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/code-search.ts b/common/src/tools/params/tool/code-search.ts index 876ea29349..2f5d827910 100644 --- a/common/src/tools/params/tool/code-search.ts +++ b/common/src/tools/params/tool/code-search.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -85,37 +85,37 @@ RESULT LIMITING: - If the global limit is reached, remaining files will be skipped Examples: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'foo' }, endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'foo\\.bar = 1\\.0' }, endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'import.*foo', cwd: 'src' }, endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'function.*authenticate', flags: '-i -t ts -t js' }, endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'TODO', flags: '-n --type-not py' }, endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { pattern: 'getUserData', maxResults: 10 }, diff --git a/common/src/tools/params/tool/create-plan.ts b/common/src/tools/params/tool/create-plan.ts index 1aca1d6cee..56c027da28 100644 --- a/common/src/tools/params/tool/create-plan.ts +++ b/common/src/tools/params/tool/create-plan.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { updateFileResultSchema } from './str-replace' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -52,7 +52,7 @@ After creating the plan, you should end turn to let the user review the plan. Important: Use this tool sparingly. Do not use this tool more than once in a conversation, unless in ask mode. Examples: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/end-turn.ts b/common/src/tools/params/tool/end-turn.ts index 16d21a6720..fff9669116 100644 --- a/common/src/tools/params/tool/end-turn.ts +++ b/common/src/tools/params/tool/end-turn.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -20,14 +23,14 @@ Only use this tool to hand control back to the user. - Effect: Signals the UI to wait for the user's reply; any pending tool results will be ignored. *INCORRECT USAGE*: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName: 'some_tool_that_produces_results', inputSchema: null, input: { query: 'some example search term' }, endsAgentStep: false, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: {}, @@ -37,7 +40,7 @@ ${$getToolCallString({ *CORRECT USAGE*: All done! Would you like some more help with xyz? -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: {}, diff --git a/common/src/tools/params/tool/find-files.ts b/common/src/tools/params/tool/find-files.ts index 4b46e15ec2..3a931b3423 100644 --- a/common/src/tools/params/tool/find-files.ts +++ b/common/src/tools/params/tool/find-files.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { fileContentsSchema } from './read-files' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -21,7 +21,7 @@ const inputSchema = z ) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/glob.ts b/common/src/tools/params/tool/glob.ts index e98dc67986..b944dd73ec 100644 --- a/common/src/tools/params/tool/glob.ts +++ b/common/src/tools/params/tool/glob.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -26,7 +26,7 @@ const inputSchema = z ) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/list-directory.ts b/common/src/tools/params/tool/list-directory.ts index 4031799818..d70590f375 100644 --- a/common/src/tools/params/tool/list-directory.ts +++ b/common/src/tools/params/tool/list-directory.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -19,7 +19,7 @@ const description = ` Lists all files and directories in the specified path. Useful for exploring directory structure and finding files. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -28,7 +28,7 @@ ${$getToolCallString({ endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/lookup-agent-info.ts b/common/src/tools/params/tool/lookup-agent-info.ts index 4f1ee5cc5a..029668ec4f 100644 --- a/common/src/tools/params/tool/lookup-agent-info.ts +++ b/common/src/tools/params/tool/lookup-agent-info.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { jsonValueSchema } from '../../../types/json' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -18,7 +18,7 @@ const description = ` Retrieve information about an agent by ID for proper spawning. Use this when you see a request with a full agent ID like "@publisher/agent-id@version" to validate the agent exists and get its metadata. Only agents that are published under a publisher and version are supported for this tool. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/read-docs.ts b/common/src/tools/params/tool/read-docs.ts index 235c3faee7..25e5ee06ba 100644 --- a/common/src/tools/params/tool/read-docs.ts +++ b/common/src/tools/params/tool/read-docs.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -50,7 +50,7 @@ Use cases: The tool will search for the library and return the most relevant documentation content. If a topic is specified, it will focus the results on that specific area. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -61,7 +61,7 @@ ${$getToolCallString({ endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/read-files.ts b/common/src/tools/params/tool/read-files.ts index 948f44af2f..7c286d5bd8 100644 --- a/common/src/tools/params/tool/read-files.ts +++ b/common/src/tools/params/tool/read-files.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -38,7 +38,7 @@ const description = ` Note: DO NOT call this tool for files you've already read! There's no need to read them again — any changes to the files will be surfaced to you as a file update tool result. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/read-subtree.ts b/common/src/tools/params/tool/read-subtree.ts index 3156d8ca74..09f0c1f58d 100644 --- a/common/src/tools/params/tool/read-subtree.ts +++ b/common/src/tools/params/tool/read-subtree.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -28,7 +28,7 @@ const inputSchema = z ) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/run-file-change-hooks.ts b/common/src/tools/params/tool/run-file-change-hooks.ts index 1b13799821..e69c211d66 100644 --- a/common/src/tools/params/tool/run-file-change-hooks.ts +++ b/common/src/tools/params/tool/run-file-change-hooks.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { terminalCommandOutputSchema } from './run-terminal-command' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -25,7 +25,7 @@ Use cases: The client will run only the hooks whose filePattern matches the provided files. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/run-terminal-command.ts b/common/src/tools/params/tool/run-terminal-command.ts index c89e16e570..4bd53f0c23 100644 --- a/common/src/tools/params/tool/run-terminal-command.ts +++ b/common/src/tools/params/tool/run-terminal-command.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -156,7 +156,7 @@ Notes: ${gitCommitGuidePrompt} Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -165,7 +165,7 @@ ${$getToolCallString({ endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/set-messages.ts b/common/src/tools/params/tool/set-messages.ts index bb062cadff..a381f4ca7f 100644 --- a/common/src/tools/params/tool/set-messages.ts +++ b/common/src/tools/params/tool/set-messages.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -13,7 +16,7 @@ const inputSchema = z .describe(`Set the conversation history to the provided messages.`) const description = ` Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/set-output.ts b/common/src/tools/params/tool/set-output.ts index f86c94f800..6b976ce0da 100644 --- a/common/src/tools/params/tool/set-output.ts +++ b/common/src/tools/params/tool/set-output.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString } from '../utils' +import { $getNativeToolCallExampleString } from '../utils' import type { $ToolParams } from '../../constants' @@ -16,7 +16,7 @@ You must use this tool as it is the only way to report any findings to the user. Please set the output with all the information and analysis you want to pass on to the user. If you just want to send a simple message, use an object with the key "message" and value of the message you want to send. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/spawn-agent-inline.ts b/common/src/tools/params/tool/spawn-agent-inline.ts index 6ee9a9d442..8b3b682ad9 100644 --- a/common/src/tools/params/tool/spawn-agent-inline.ts +++ b/common/src/tools/params/tool/spawn-agent-inline.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -31,7 +34,7 @@ This is useful for: - Managing message history (e.g., summarization) The agent will run until it calls end_turn, then control returns to you. There is no tool result for this tool. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/spawn-agents.ts b/common/src/tools/params/tool/spawn-agents.ts index f7da5e5d7d..2c83c8b5ba 100644 --- a/common/src/tools/params/tool/spawn-agents.ts +++ b/common/src/tools/params/tool/spawn-agents.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { jsonObjectSchema } from '../../../types/json' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -36,7 +36,7 @@ Use this tool to spawn agents to help you complete the user request. Each agent The prompt field is a simple string, while params is a JSON object that gets validated against the agent's schema. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/str-replace.ts b/common/src/tools/params/tool/str-replace.ts index 5aee745fec..b02ce1e81c 100644 --- a/common/src/tools/params/tool/str-replace.ts +++ b/common/src/tools/params/tool/str-replace.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -61,7 +61,7 @@ Important: If you are making multiple edits in a row to a file, use only one str_replace call with multiple replacements instead of multiple str_replace tool calls. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/task-completed.ts b/common/src/tools/params/tool/task-completed.ts index a8c35d1c68..7ea2a4f856 100644 --- a/common/src/tools/params/tool/task-completed.ts +++ b/common/src/tools/params/tool/task-completed.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -34,19 +37,19 @@ Use this tool to signal that the task is complete. All changes have been implemented and tested successfully! -${$getToolCallString({ toolName, inputSchema, input: {}, endsAgentStep })} +${$getNativeToolCallExampleString({ toolName, inputSchema, input: {}, endsAgentStep })} OR I need more information to proceed. Which database schema should I use for this migration? -${$getToolCallString({ toolName, inputSchema, input: {}, endsAgentStep })} +${$getNativeToolCallExampleString({ toolName, inputSchema, input: {}, endsAgentStep })} OR I can't get the tests to pass after several different attempts. I need help from the user to proceed. -${$getToolCallString({ toolName, inputSchema, input: {}, endsAgentStep })} +${$getNativeToolCallExampleString({ toolName, inputSchema, input: {}, endsAgentStep })} `.trim() export const taskCompletedParams = { diff --git a/common/src/tools/params/tool/think-deeply.ts b/common/src/tools/params/tool/think-deeply.ts index 4292332fa5..e84a076019 100644 --- a/common/src/tools/params/tool/think-deeply.ts +++ b/common/src/tools/params/tool/think-deeply.ts @@ -1,6 +1,9 @@ import z from 'zod/v4' -import { $getToolCallString, emptyToolResultSchema } from '../utils' +import { + $getNativeToolCallExampleString, + emptyToolResultSchema, +} from '../utils' import type { $ToolParams } from '../../constants' @@ -29,7 +32,7 @@ Avoid for simple changes (e.g., single functions, minor edits). This tool does not generate a tool result. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/update-subgoal.ts b/common/src/tools/params/tool/update-subgoal.ts index 299ca9eeaf..75e778c63a 100644 --- a/common/src/tools/params/tool/update-subgoal.ts +++ b/common/src/tools/params/tool/update-subgoal.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -31,7 +31,7 @@ const description = ` Examples: Usage 1 (update status): -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -42,7 +42,7 @@ ${$getToolCallString({ })} Usage 2 (update plan): -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -53,7 +53,7 @@ ${$getToolCallString({ })} Usage 3 (add log): -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -64,7 +64,7 @@ ${$getToolCallString({ })} Usage 4 (update status and add log): -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/web-search.ts b/common/src/tools/params/tool/web-search.ts index 7a458cc01a..e87c8f2715 100644 --- a/common/src/tools/params/tool/web-search.ts +++ b/common/src/tools/params/tool/web-search.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -34,7 +34,7 @@ Use cases: The tool will return search results with titles, URLs, and content snippets. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -44,7 +44,7 @@ ${$getToolCallString({ endsAgentStep, })} -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/write-file.ts b/common/src/tools/params/tool/write-file.ts index 00ec71c6d2..cf50fee058 100644 --- a/common/src/tools/params/tool/write-file.ts +++ b/common/src/tools/params/tool/write-file.ts @@ -1,7 +1,7 @@ import z from 'zod/v4' import { updateFileResultSchema } from './str-replace' -import { $getToolCallString, jsonToolResultSchema } from '../utils' +import { $getNativeToolCallExampleString, jsonToolResultSchema } from '../utils' import type { $ToolParams } from '../../constants' @@ -39,7 +39,7 @@ Do not use this tool to delete or rename a file. Instead run a terminal command Examples: Example 1 - Simple file creation: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { @@ -51,7 +51,7 @@ ${$getToolCallString({ })} Example 2 - Editing with placeholder comments: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/tool/write-todos.ts b/common/src/tools/params/tool/write-todos.ts index 7b7489a6f2..ae73e72a1c 100644 --- a/common/src/tools/params/tool/write-todos.ts +++ b/common/src/tools/params/tool/write-todos.ts @@ -1,6 +1,6 @@ import z from 'zod/v4' -import { $getToolCallString } from '../utils' +import { $getNativeToolCallExampleString } from '../utils' import type { $ToolParams } from '../../constants' @@ -30,7 +30,7 @@ After completing each todo step, call this tool again to update the list and mar Use this tool frequently as you work through tasks to update the list of todos with their current status. Doing this is extremely useful because it helps you stay on track and complete all the requirements of the user's request. It also helps inform the user of your plans and the current progress, which they want to know at all times. Example: -${$getToolCallString({ +${$getNativeToolCallExampleString({ toolName, inputSchema, input: { diff --git a/common/src/tools/params/utils.ts b/common/src/tools/params/utils.ts index cbf79d327f..951ee3a61a 100644 --- a/common/src/tools/params/utils.ts +++ b/common/src/tools/params/utils.ts @@ -34,6 +34,20 @@ export function $getToolCallString(params: { return [startToolTag, JSON.stringify(obj, null, 2), endToolTag].join('') } +export function $getNativeToolCallExampleString(params: { + toolName: string + inputSchema: z.ZodType | null + input: Input + endsAgentStep?: boolean // unused +}): string { + const { toolName, input } = params + return [ + `<${toolName}_params_example>\n`, + JSON.stringify(input, null, 2), + `\n`, + ].join('') +} + /** Generates the zod schema for a single JSON tool result. */ export function jsonToolResultSchema( valueSchema: z.ZodType, From 532a3b42f6ff1b4f80f4cd94e8855754b23da100 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:39:19 -0800 Subject: [PATCH 07/30] remove tool instructions from prompts --- packages/agent-runtime/src/templates/strings.ts | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/packages/agent-runtime/src/templates/strings.ts b/packages/agent-runtime/src/templates/strings.ts index 2f7c4e75f2..0fd3bc20ce 100644 --- a/packages/agent-runtime/src/templates/strings.ts +++ b/packages/agent-runtime/src/templates/strings.ts @@ -11,11 +11,6 @@ import { getProjectFileTreePrompt, getSystemInfoPrompt, } from '../system-prompt/prompts' -import { - fullToolList, - getShortToolInstructions, - getToolsInstructions, -} from '../tools/prompts' import { parseUserMessage } from '../util/messages' import type { AgentTemplate, PlaceholderValue } from './types' @@ -113,8 +108,7 @@ export async function formatPrompt( [PLACEHOLDER.REMAINING_STEPS]: () => `${agentState.stepsRemaining!}`, [PLACEHOLDER.PROJECT_ROOT]: () => fileContext.projectRoot, [PLACEHOLDER.SYSTEM_INFO_PROMPT]: () => getSystemInfoPrompt(fileContext), - [PLACEHOLDER.TOOLS_PROMPT]: async () => - getToolsInstructions(tools, (await additionalToolDefinitions()) ?? {}), + [PLACEHOLDER.TOOLS_PROMPT]: async () => '', [PLACEHOLDER.AGENTS_PROMPT]: () => buildSpawnableAgentsDescription(params), [PLACEHOLDER.USER_CWD]: () => fileContext.cwd, [PLACEHOLDER.USER_INPUT_PROMPT]: () => escapeString(lastUserInput ?? ''), @@ -204,15 +198,7 @@ export async function getAgentPrompt( // Add tool instructions, spawnable agents, and output schema prompts to instructionsPrompt if (promptType.type === 'instructionsPrompt' && agentState.agentType) { - const toolsInstructions = agentTemplate.inheritParentSystemPrompt - ? fullToolList(agentTemplate.toolNames, await additionalToolDefinitions()) - : getShortToolInstructions( - agentTemplate.toolNames, - await additionalToolDefinitions(), - ) addendum += - '\n\n' + - toolsInstructions + '\n\n' + (await buildSpawnableAgentsDescription({ ...params, From 908664e452884f29754b4b97316acdbc297e327f Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:56:24 -0800 Subject: [PATCH 08/30] remove debug message --- packages/agent-runtime/src/tools/stream-parser.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/agent-runtime/src/tools/stream-parser.ts b/packages/agent-runtime/src/tools/stream-parser.ts index 224df9533e..435afebd25 100644 --- a/packages/agent-runtime/src/tools/stream-parser.ts +++ b/packages/agent-runtime/src/tools/stream-parser.ts @@ -251,8 +251,6 @@ export async function processStream( ...toolResultsToAddAfterStream, ]) - logger.info({ messages: agentState.messageHistory }, 'asdf messages') - if (!signal.aborted) { resolveStreamDonePromise() await previousToolCallFinished From e0876894f9965d9c0ddc9ea993dbadafb70219a9 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 13:56:57 -0800 Subject: [PATCH 09/30] have runProgrammaticStep use native tool calls --- packages/agent-runtime/src/run-agent-step.ts | 67 +++++++++---------- .../src/run-programmatic-step.ts | 67 +++++++++++-------- .../agent-runtime/src/tools/tool-executor.ts | 42 ++++++------ 3 files changed, 93 insertions(+), 83 deletions(-) diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index e6a111f694..9dd3b3d6e8 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -470,37 +470,35 @@ export const runAgentStep = async ( export async function loopAgentSteps( params: { - userInputId: string - agentType: AgentTemplateType + addAgentStep: AddAgentStepFn agentState: AgentState - prompt: string | undefined + agentType: AgentTemplateType + clearUserPromptMessagesAfterResponse?: boolean + clientSessionId: string content?: Array - spawnParams: Record | undefined fileContext: ProjectFileContext + finishAgentRun: FinishAgentRunFn localAgentTemplates: Record - clearUserPromptMessagesAfterResponse?: boolean + logger: Logger parentSystemPrompt?: string + prompt: string | undefined signal: AbortSignal - - userId: string | undefined - clientSessionId: string - + spawnParams: Record | undefined startAgentRun: StartAgentRunFn - finishAgentRun: FinishAgentRunFn - addAgentStep: AddAgentStepFn - logger: Logger + userId: string | undefined + userInputId: string } & ParamsExcluding & ParamsExcluding< typeof runProgrammaticStep, - | 'runId' | 'agentState' - | 'template' + | 'onCostCalculated' | 'prompt' - | 'toolCallParams' - | 'stepsComplete' + | 'runId' | 'stepNumber' + | 'stepsComplete' | 'system' - | 'onCostCalculated' + | 'template' + | 'toolCallParams' > & ParamsExcluding & ParamsExcluding< @@ -546,23 +544,23 @@ export async function loopAgentSteps( output: AgentOutput }> { const { - userInputId, - agentType, + addAgentStep, agentState, - prompt, + agentType, + clearUserPromptMessagesAfterResponse = true, + clientSessionId, content, - spawnParams, fileContext, + finishAgentRun, localAgentTemplates, - userId, - clientSessionId, - clearUserPromptMessagesAfterResponse = true, + logger, parentSystemPrompt, + prompt, signal, + spawnParams, startAgentRun, - finishAgentRun, - addAgentStep, - logger, + userId, + userInputId, } = params const agentTemplate = await getAgentTemplate({ @@ -724,20 +722,21 @@ export async function loopAgentSteps( if (agentTemplate.handleSteps) { const programmaticResult = await runProgrammaticStep({ ...params, - runId, + agentState: currentAgentState, - template: agentTemplate, localAgentTemplates, - prompt: currentPrompt, - toolCallParams: currentParams, - system, - stepsComplete: shouldEndTurn, - stepNumber: totalSteps, nResponses, onCostCalculated: async (credits: number) => { agentState.creditsUsed += credits agentState.directCreditsUsed += credits }, + prompt: currentPrompt, + runId, + stepNumber: totalSteps, + stepsComplete: shouldEndTurn, + system, + template: agentTemplate, + toolCallParams: currentParams, }) const { agentState: programmaticAgentState, diff --git a/packages/agent-runtime/src/run-programmatic-step.ts b/packages/agent-runtime/src/run-programmatic-step.ts index f8cda7edf2..8af7234755 100644 --- a/packages/agent-runtime/src/run-programmatic-step.ts +++ b/packages/agent-runtime/src/run-programmatic-step.ts @@ -1,4 +1,3 @@ -import { getToolCallString } from '@codebuff/common/tools/utils' import { getErrorObject } from '@codebuff/common/util/error' import { assistantMessage } from '@codebuff/common/util/messages' import { cloneDeep } from 'lodash' @@ -21,7 +20,10 @@ import type { AddAgentStepFn } from '@codebuff/common/types/contracts/database' import type { Logger } from '@codebuff/common/types/contracts/logger' import type { ParamsExcluding } from '@codebuff/common/types/function-params' import type { ToolMessage } from '@codebuff/common/types/messages/codebuff-message' -import type { ToolResultOutput } from '@codebuff/common/types/messages/content-part' +import type { + ToolCallPart, + ToolResultOutput, +} from '@codebuff/common/types/messages/content-part' import type { PrintModeEvent } from '@codebuff/common/types/print-mode' import type { AgentState } from '@codebuff/common/types/session-state' @@ -40,26 +42,26 @@ export function clearAgentGeneratorCache(params: { logger: Logger }) { // Function to handle programmatic agents export async function runProgrammaticStep( params: { + addAgentStep: AddAgentStepFn agentState: AgentState - template: AgentTemplate + clientSessionId: string + fingerprintId: string + handleStepsLogChunk: HandleStepsLogChunkFn + localAgentTemplates: Record + logger: Logger + nResponses?: string[] + onResponseChunk: (chunk: string | PrintModeEvent) => void prompt: string | undefined - toolCallParams: Record | undefined - system: string | undefined - userId: string | undefined repoId: string | undefined repoUrl: string | undefined - userInputId: string - fingerprintId: string - clientSessionId: string - onResponseChunk: (chunk: string | PrintModeEvent) => void - localAgentTemplates: Record - stepsComplete: boolean stepNumber: number - handleStepsLogChunk: HandleStepsLogChunkFn + stepsComplete: boolean + template: AgentTemplate + toolCallParams: Record | undefined sendAction: SendActionFn - addAgentStep: AddAgentStepFn - logger: Logger - nResponses?: string[] + system: string | undefined + userId: string | undefined + userInputId: string } & Omit< ExecuteToolCallParams, | 'toolName' @@ -274,20 +276,27 @@ export async function runProgrammaticStep( const excludeToolFromMessageHistory = toolCall?.includeToolCall === false // Add assistant message with the tool call before executing it if (!excludeToolFromMessageHistory) { - const toolCallString = getToolCallString( - toolCall.toolName, - toolCall.input, - ) - onResponseChunk(toolCallString) - agentState.messageHistory.push(assistantMessage(toolCallString)) + const toolCallPart: ToolCallPart = { + type: 'tool-call', + toolCallId, + toolName: toolCall.toolName, + input: toolCall.input, + } + // onResponseChunk({ + // ...toolCallPart, + // type: 'tool_call', + // agentId: agentState.agentId, + // parentAgentId: agentState.parentId, + // }) + agentState.messageHistory.push(assistantMessage(toolCallPart)) // Optional call handles both top-level and nested agents - sendSubagentChunk({ - userInputId, - agentId: agentState.agentId, - agentType: agentState.agentType!, - chunk: toolCallString, - forwardToPrompt: !agentState.parentId, - }) + // sendSubagentChunk({ + // userInputId, + // agentId: agentState.agentId, + // agentType: agentState.agentType!, + // chunk: toolCallString, + // forwardToPrompt: !agentState.parentId, + // }) } // Execute the tool synchronously and get the result immediately diff --git a/packages/agent-runtime/src/tools/tool-executor.ts b/packages/agent-runtime/src/tools/tool-executor.ts index 1baa2b774d..4304daed2a 100644 --- a/packages/agent-runtime/src/tools/tool-executor.ts +++ b/packages/agent-runtime/src/tools/tool-executor.ts @@ -3,7 +3,6 @@ import { toolParams } from '@codebuff/common/tools/list' import { jsonToolResult } from '@codebuff/common/util/messages' import { generateCompactId } from '@codebuff/common/util/string' import { cloneDeep } from 'lodash' -import z from 'zod/v4' import { checkLiveUserInput } from '../live-user-inputs' import { getMCPToolData } from '../mcp' @@ -66,24 +65,28 @@ export function parseRawToolCall(params: { } const validName = toolName as T - const processedParameters: Record = {} - for (const [param, val] of Object.entries(rawToolCall.input ?? {})) { - processedParameters[param] = val - } + // const processedParameters: Record = {} + // for (const [param, val] of Object.entries(rawToolCall.input ?? {})) { + // processedParameters[param] = val + // } // Add the required codebuff_end_step parameter with the correct value for this tool if requested - if (autoInsertEndStepParam) { - processedParameters[endsAgentStepParam] = - toolParams[validName].endsAgentStep - } + // if (autoInsertEndStepParam) { + // processedParameters[endsAgentStepParam] = + // toolParams[validName].endsAgentStep + // } + + // const paramsSchema = toolParams[validName].endsAgentStep + // ? ( + // toolParams[validName].inputSchema satisfies z.ZodObject as z.ZodObject + // ).extend({ + // [endsAgentStepParam]: z.literal(toolParams[validName].endsAgentStep), + // }) + // : toolParams[validName].inputSchema + + const processedParameters = rawToolCall.input + const paramsSchema = toolParams[validName].inputSchema - const paramsSchema = toolParams[validName].endsAgentStep - ? ( - toolParams[validName].inputSchema satisfies z.ZodObject as z.ZodObject - ).extend({ - [endsAgentStepParam]: z.literal(toolParams[validName].endsAgentStep), - }) - : toolParams[validName].inputSchema const result = paramsSchema.safeParse(processedParameters) if (!result.success) { @@ -178,10 +181,9 @@ export function executeToolCall( toolCallId, toolName, input, - // Only include agentId for subagents (agents with a parent) - ...(agentState.parentId && { agentId: agentState.agentId }), - // Include includeToolCall flag if explicitly set to false - ...(excludeToolFromMessageHistory && { includeToolCall: false }), + agentId: agentState.agentId, + parentAgentId: agentState.parentId, + includeToolCall: !excludeToolFromMessageHistory, }) const toolCall: CodebuffToolCall | ToolCallError = parseRawToolCall({ From cf76b1268e0629082e5d9b1f0ded677c0d6363dd Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Tue, 25 Nov 2025 14:58:19 -0800 Subject: [PATCH 10/30] fix typecheck --- evals/scaffolding.ts | 26 ++++++++++--------- .../src/__tests__/n-parameter.test.ts | 1 + .../src/__tests__/read-docs-tool.test.ts | 1 + .../__tests__/run-agent-step-tools.test.ts | 6 ++--- .../src/__tests__/web-search-tool.test.ts | 6 ++--- 5 files changed, 22 insertions(+), 18 deletions(-) diff --git a/evals/scaffolding.ts b/evals/scaffolding.ts index 6250a2f0ba..fcc8cc1a68 100644 --- a/evals/scaffolding.ts +++ b/evals/scaffolding.ts @@ -206,13 +206,15 @@ export async function runAgentStepScaffolding( const result = await runAgentStep({ ...EVALS_AGENT_RUNTIME_IMPL, ...agentRuntimeScopedImpl, + additionalToolDefinitions: () => Promise.resolve({}), - textOverride: null, - runId: 'test-run-id', - userId: TEST_USER_ID, - userInputId: generateCompactId(), + agentState, + agentType, + ancestorRunIds: [], clientSessionId: sessionId, + fileContext, fingerprintId: 'test-fingerprint-id', + localAgentTemplates, onResponseChunk: (chunk: string | PrintModeEvent) => { if (typeof chunk !== 'string') { return @@ -222,17 +224,17 @@ export async function runAgentStepScaffolding( } fullResponse += chunk }, - agentType, - fileContext, - localAgentTemplates, - agentState, prompt, - ancestorRunIds: [], - spawnParams: undefined, - repoUrl: undefined, repoId: undefined, - system: 'Test system prompt', + repoUrl: undefined, + runId: 'test-run-id', signal: new AbortController().signal, + spawnParams: undefined, + system: 'Test system prompt', + textOverride: null, + tools: {}, + userId: TEST_USER_ID, + userInputId: generateCompactId(), }) return { diff --git a/packages/agent-runtime/src/__tests__/n-parameter.test.ts b/packages/agent-runtime/src/__tests__/n-parameter.test.ts index c30ef339fb..897a520417 100644 --- a/packages/agent-runtime/src/__tests__/n-parameter.test.ts +++ b/packages/agent-runtime/src/__tests__/n-parameter.test.ts @@ -122,6 +122,7 @@ describe('n parameter and GENERATE_N functionality', () => { spawnParams: undefined, system: 'Test system', signal: new AbortController().signal, + tools: {} } }) diff --git a/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts b/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts index 4b62cb588b..ea41ff9f6a 100644 --- a/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts +++ b/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts @@ -89,6 +89,7 @@ describe('read_docs tool with researcher agent (via web API facade)', () => { agentType: 'researcher', spawnParams: undefined, signal: new AbortController().signal, + tools: {}, } }) diff --git a/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts b/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts index 62e026c0f7..0d15eae292 100644 --- a/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts +++ b/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts @@ -116,10 +116,12 @@ describe('runAgentStep - set_output tool', () => { runAgentStepBaseParams = { ...agentRuntimeImpl, + additionalToolDefinitions: () => Promise.resolve({}), ancestorRunIds: [], clientSessionId: 'test-session', fileContext: mockFileContext, fingerprintId: 'test-fingerprint', + onResponseChunk: () => {}, repoId: undefined, repoUrl: undefined, runId: 'test-run-id', @@ -127,11 +129,9 @@ describe('runAgentStep - set_output tool', () => { spawnParams: undefined, system: 'Test system prompt', textOverride: null, + tools: {}, userId: TEST_USER_ID, userInputId: 'test-input', - - additionalToolDefinitions: () => Promise.resolve({}), - onResponseChunk: () => {}, } }) diff --git a/packages/agent-runtime/src/__tests__/web-search-tool.test.ts b/packages/agent-runtime/src/__tests__/web-search-tool.test.ts index 417f57819f..7c3b6801b8 100644 --- a/packages/agent-runtime/src/__tests__/web-search-tool.test.ts +++ b/packages/agent-runtime/src/__tests__/web-search-tool.test.ts @@ -61,11 +61,13 @@ describe('web_search tool with researcher agent (via web API facade)', () => { runAgentStepBaseParams = { ...agentRuntimeImpl, + additionalToolDefinitions: () => Promise.resolve({}), agentType: 'researcher', ancestorRunIds: [], clientSessionId: 'test-session', fileContext: mockFileContext, fingerprintId: 'test-fingerprint', + onResponseChunk: () => {}, repoId: undefined, repoUrl: undefined, runId: 'test-run-id', @@ -73,11 +75,9 @@ describe('web_search tool with researcher agent (via web API facade)', () => { spawnParams: undefined, system: 'Test system prompt', textOverride: null, + tools: {}, userId: TEST_USER_ID, userInputId: 'test-input', - - additionalToolDefinitions: () => Promise.resolve({}), - onResponseChunk: () => {}, } // Mock analytics and tracing From 6eff7d45177f673858e1f83b396140569f332c38 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Tue, 25 Nov 2025 17:40:39 -0800 Subject: [PATCH 11/30] fix: cli uses parentAgentId to distinguish tools from subagents --- cli/src/hooks/use-send-message.ts | 36 ++++++++++++++++++++----------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/cli/src/hooks/use-send-message.ts b/cli/src/hooks/use-send-message.ts index 613abe396e..80b98c9ead 100644 --- a/cli/src/hooks/use-send-message.ts +++ b/cli/src/hooks/use-send-message.ts @@ -28,7 +28,12 @@ import { import type { ElapsedTimeTracker } from './use-elapsed-time' import type { StreamStatus } from './use-message-queue' -import type { ChatMessage, ContentBlock, ToolContentBlock, AskUserContentBlock } from '../types/chat' +import type { + ChatMessage, + ContentBlock, + ToolContentBlock, + AskUserContentBlock, +} from '../types/chat' import type { SendMessageFn } from '../types/contracts/send-message' import type { ParamsOf } from '../types/function-params' import type { SetElement } from '../types/utils' @@ -1123,7 +1128,7 @@ export const useSendMessage = ({ ] of spawnAgentsMapRef.current.entries()) { const eventType = event.agentType || '' const storedType = info.agentType || '' - + // Extract base names without version or scope // e.g., 'codebuff/file-picker@0.0.2' -> 'file-picker' // 'file-picker' -> 'file-picker' @@ -1135,10 +1140,10 @@ export const useSendMessage = ({ // Handle simple names, possibly with version return type.split('@')[0] } - + const eventBaseName = getBaseName(eventType) const storedBaseName = getBaseName(storedType) - + // Match if base names are the same const isMatch = eventBaseName === storedBaseName if (isMatch) { @@ -1416,6 +1421,7 @@ export const useSendMessage = ({ input, agentId, includeToolCall, + parentAgentId, } = event if (toolName === 'spawn_agents' && input?.agents) { @@ -1487,7 +1493,7 @@ export const useSendMessage = ({ } // If this tool call belongs to a subagent, add it to that agent's blocks - if (agentId) { + if (parentAgentId && agentId) { applyMessageUpdate((prev) => prev.map((msg) => { if (msg.id !== aiMessageId || !msg.blocks) { @@ -1557,18 +1563,24 @@ export const useSendMessage = ({ } setStreamingAgents((prev) => new Set(prev).add(toolCallId)) - } else if (event.type === 'tool_result' && event.toolCallId) { + } else if (event.type === 'tool_result' && event.toolCallId) { const { toolCallId } = event // Handle ask_user result transformation - applyMessageUpdate((prev) => + applyMessageUpdate((prev) => prev.map((msg) => { if (msg.id !== aiMessageId || !msg.blocks) return msg // Recursively check for tool blocks to transform - const transformAskUser = (blocks: ContentBlock[]): ContentBlock[] => { + const transformAskUser = ( + blocks: ContentBlock[], + ): ContentBlock[] => { return blocks.map((block) => { - if (block.type === 'tool' && block.toolCallId === toolCallId && block.toolName === 'ask_user') { + if ( + block.type === 'tool' && + block.toolCallId === toolCallId && + block.toolName === 'ask_user' + ) { const resultValue = (event.output?.[0] as any)?.value const skipped = resultValue?.skipped const answers = resultValue?.answers @@ -1587,7 +1599,7 @@ export const useSendMessage = ({ skipped, } as AskUserContentBlock } - + if (block.type === 'agent' && block.blocks) { const updatedBlocks = transformAskUser(block.blocks) if (updatedBlocks !== block.blocks) { @@ -1600,10 +1612,10 @@ export const useSendMessage = ({ const newBlocks = transformAskUser(msg.blocks) if (newBlocks !== msg.blocks) { - return { ...msg, blocks: newBlocks } + return { ...msg, blocks: newBlocks } } return msg - }) + }), ) // Check if this is a spawn_agents result From 17a4e5768a0e3d6d5f666bd5ad960521f90341db Mon Sep 17 00:00:00 2001 From: James Grugett Date: Tue, 25 Nov 2025 18:34:59 -0800 Subject: [PATCH 12/30] For "last_message" output mode, return the entire assistant response, including tools and results --- common/src/types/session-state.ts | 2 +- .../src/__tests__/cost-aggregation.test.ts | 10 ++-- .../spawn-agents-message-history.test.ts | 2 +- .../spawn-agents-permissions.test.ts | 2 +- .../src/__tests__/subagent-streaming.test.ts | 2 +- .../agent-runtime/src/util/agent-output.ts | 50 ++++++++++++++++--- 6 files changed, 53 insertions(+), 15 deletions(-) diff --git a/common/src/types/session-state.ts b/common/src/types/session-state.ts index ca04c73121..267e9b34fb 100644 --- a/common/src/types/session-state.ts +++ b/common/src/types/session-state.ts @@ -48,7 +48,7 @@ export const AgentOutputSchema = z.discriminatedUnion('type', [ }), z.object({ type: z.literal('lastMessage'), - value: z.any(), + value: z.array(z.any()), // Array of assistant and tool messages from the last turn, including tool results }), z.object({ type: z.literal('allMessages'), diff --git a/packages/agent-runtime/src/__tests__/cost-aggregation.test.ts b/packages/agent-runtime/src/__tests__/cost-aggregation.test.ts index 5c73a01b3e..b46fee77ed 100644 --- a/packages/agent-runtime/src/__tests__/cost-aggregation.test.ts +++ b/packages/agent-runtime/src/__tests__/cost-aggregation.test.ts @@ -159,7 +159,7 @@ describe('Cost Aggregation System', () => { stepsRemaining: 10, creditsUsed: 75, // First subagent uses 75 credits }, - output: { type: 'lastMessage', value: 'Sub-agent 1 response' }, + output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] }, }) .mockResolvedValueOnce({ agentState: { @@ -169,7 +169,7 @@ describe('Cost Aggregation System', () => { stepsRemaining: 10, creditsUsed: 100, // Second subagent uses 100 credits }, - output: { type: 'lastMessage', value: 'Sub-agent 2 response' }, + output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] }, }) const mockToolCall = { @@ -223,7 +223,7 @@ describe('Cost Aggregation System', () => { stepsRemaining: 10, creditsUsed: 50, // Successful agent }, - output: { type: 'lastMessage', value: 'Successful response' }, + output: { type: 'lastMessage', value: [assistantMessage('Successful response')] }, }) .mockRejectedValueOnce( (() => { @@ -370,7 +370,7 @@ describe('Cost Aggregation System', () => { stepsRemaining: 10, creditsUsed: subAgent1Cost, } as AgentState, - output: { type: 'lastMessage', value: 'Sub-agent 1 response' }, + output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 1 response')] }, }) .mockResolvedValueOnce({ agentState: { @@ -381,7 +381,7 @@ describe('Cost Aggregation System', () => { stepsRemaining: 10, creditsUsed: subAgent2Cost, } as AgentState, - output: { type: 'lastMessage', value: 'Sub-agent 2 response' }, + output: { type: 'lastMessage', value: [assistantMessage('Sub-agent 2 response')] }, }) const mockToolCall = { diff --git a/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts b/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts index b472231299..90715389b9 100644 --- a/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts +++ b/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts @@ -52,7 +52,7 @@ describe('Spawn Agents Message History', () => { assistantMessage('Mock agent response'), ], }, - output: { type: 'lastMessage', value: 'Mock agent response' }, + output: { type: 'lastMessage', value: [assistantMessage('Mock agent response')] }, } }) diff --git a/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts b/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts index 3f827e2a65..dc8d322522 100644 --- a/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts +++ b/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts @@ -85,7 +85,7 @@ describe('Spawn Agents Permissions', () => { ...options.agentState, messageHistory: [assistantMessage('Mock agent response')], }, - output: { type: 'lastMessage', value: 'Mock agent response' }, + output: { type: 'lastMessage', value: [assistantMessage('Mock agent response')] }, } }) }) diff --git a/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts b/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts index 1bd1a69705..134cadd8ba 100644 --- a/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts +++ b/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts @@ -96,7 +96,7 @@ describe('Subagent Streaming', () => { ...options.agentState, messageHistory: [assistantMessage('Test response from subagent')], }, - output: { type: 'lastMessage', value: 'Test response from subagent' }, + output: { type: 'lastMessage', value: [assistantMessage('Test response from subagent')] }, } }) diff --git a/packages/agent-runtime/src/util/agent-output.ts b/packages/agent-runtime/src/util/agent-output.ts index 624e3ca632..fe3a8da0a6 100644 --- a/packages/agent-runtime/src/util/agent-output.ts +++ b/packages/agent-runtime/src/util/agent-output.ts @@ -1,10 +1,49 @@ import type { AgentTemplate } from '@codebuff/common/types/agent-template' -import type { AssistantMessage } from '@codebuff/common/types/messages/codebuff-message' +import type { Message } from '@codebuff/common/types/messages/codebuff-message' import type { AgentState, AgentOutput, } from '@codebuff/common/types/session-state' +/** + * Get the last assistant turn messages, which includes the last assistant message + * and any subsequent tool messages that are responses to its tool calls. + */ +function getLastAssistantTurnMessages(messageHistory: Message[]): Message[] { + // Find the index of the last assistant message + let lastAssistantIndex = -1 + for (let i = messageHistory.length - 1; i >= 0; i--) { + if (messageHistory[i].role === 'assistant') { + lastAssistantIndex = i + break + } + } + + for (let i = lastAssistantIndex; i >= 0; i--) { + if (messageHistory[i].role === 'assistant') { + lastAssistantIndex = i + } else break + } + + if (lastAssistantIndex === -1) { + return [] + } + + // Collect the assistant message and all subsequent tool messages + const result: Message[] = [] + for (let i = lastAssistantIndex; i < messageHistory.length; i++) { + const message = messageHistory[i] + if (message.role === 'assistant' || message.role === 'tool') { + result.push(message) + } else { + // Stop if we hit a user or system message + break + } + } + + return result +} + export function getAgentOutput( agentState: AgentState, agentTemplate: AgentTemplate, @@ -16,11 +55,10 @@ export function getAgentOutput( } } if (agentTemplate.outputMode === 'last_message') { - const assistantMessages = agentState.messageHistory.filter( - (message): message is AssistantMessage => message.role === 'assistant', + const lastTurnMessages = getLastAssistantTurnMessages( + agentState.messageHistory, ) - const lastAssistantMessage = assistantMessages[assistantMessages.length - 1] - if (!lastAssistantMessage) { + if (lastTurnMessages.length === 0) { return { type: 'error', message: 'No response from agent', @@ -28,7 +66,7 @@ export function getAgentOutput( } return { type: 'lastMessage', - value: lastAssistantMessage.content, + value: lastTurnMessages, } } if (agentTemplate.outputMode === 'all_messages') { From 32197cb75b5e9a1775554a8fc4aac7860eb7b34a Mon Sep 17 00:00:00 2001 From: James Grugett Date: Tue, 25 Nov 2025 19:35:53 -0800 Subject: [PATCH 13/30] fix stream parser bug by adding flush() --- packages/agent-runtime/src/tool-stream-parser.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/agent-runtime/src/tool-stream-parser.ts b/packages/agent-runtime/src/tool-stream-parser.ts index e5fd319332..2f096695dc 100644 --- a/packages/agent-runtime/src/tool-stream-parser.ts +++ b/packages/agent-runtime/src/tool-stream-parser.ts @@ -87,6 +87,7 @@ export async function* processStreamWithTools(params: { chunk: StreamChunk | undefined, ): Generator { if (chunk === undefined) { + flush() streamCompleted = true return } From 4487c4e9d037ad963c935a9910ca6b0083a3818e Mon Sep 17 00:00:00 2001 From: James Grugett Date: Tue, 25 Nov 2025 21:16:46 -0800 Subject: [PATCH 14/30] Fix file picker --- .../file-explorer.integration.test.ts | 348 ++++++++++++++++++ .agents/file-explorer/file-picker.ts | 98 +++-- 2 files changed, 423 insertions(+), 23 deletions(-) create mode 100644 .agents/__tests__/file-explorer.integration.test.ts diff --git a/.agents/__tests__/file-explorer.integration.test.ts b/.agents/__tests__/file-explorer.integration.test.ts new file mode 100644 index 0000000000..0aa3cc3f6a --- /dev/null +++ b/.agents/__tests__/file-explorer.integration.test.ts @@ -0,0 +1,348 @@ +import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants' +import { describe, expect, it } from 'bun:test' + +import { CodebuffClient } from '@codebuff/sdk' +import filePickerDefinition from '../file-explorer/file-picker' +import fileListerDefinition from '../file-explorer/file-lister' + +import type { PrintModeEvent } from '@codebuff/common/types/print-mode' + +/** + * Integration tests for agents that use the read_subtree tool. + * These tests verify that the SDK properly initializes the session state + * with project files and that agents can access the file tree through + * the read_subtree tool. + * + * The file-lister agent is used directly instead of file-picker because: + * - file-lister directly uses the read_subtree tool + * - file-picker spawns file-lister as a subagent, adding complexity + * - Testing file-lister directly verifies the core functionality + */ +describe('File Lister Agent Integration - read_subtree tool', () => { + it( + 'should find relevant files using read_subtree tool', + async () => { + const apiKey = process.env[API_KEY_ENV_VAR] + if (!apiKey) { + throw new Error('API key not found') + } + + // Create mock project files that the file-lister should be able to find + const projectFiles: Record = { + 'src/index.ts': ` +import { UserService } from './services/user-service' +import { AuthService } from './services/auth-service' + +export function main() { + const userService = new UserService() + const authService = new AuthService() + console.log('Application started') +} +`, + 'src/services/user-service.ts': ` +export class UserService { + async getUser(id: string) { + return { id, name: 'John Doe' } + } + + async createUser(name: string) { + return { id: 'new-user-id', name } + } + + async deleteUser(id: string) { + console.log('User deleted:', id) + } +} +`, + 'src/services/auth-service.ts': ` +export class AuthService { + async login(email: string, password: string) { + return { token: 'mock-token' } + } + + async logout() { + console.log('Logged out') + } + + async validateToken(token: string) { + return token === 'mock-token' + } +} +`, + 'src/utils/logger.ts': ` +export function log(message: string) { + console.log('[LOG]', message) +} + +export function error(message: string) { + console.error('[ERROR]', message) +} +`, + 'src/types/user.ts': ` +export interface User { + id: string + name: string + email?: string +} +`, + 'package.json': JSON.stringify({ + name: 'test-project', + version: '1.0.0', + dependencies: {}, + }), + 'README.md': + '# Test Project\n\nA simple test project for integration testing.', + } + + const client = new CodebuffClient({ + apiKey, + cwd: '/tmp/test-project', + projectFiles, + }) + + const events: PrintModeEvent[] = [] + + // Run the file-lister agent to find files related to user service + // The file-lister agent uses the read_subtree tool directly + const run = await client.run({ + agent: 'file-lister', + prompt: 'Find files related to user authentication and user management', + handleEvent: (event) => { + events.push(event) + }, + }) + + // The output should not be an error + expect(run.output.type).not.toEqual('error') + + // Verify we got some output + expect(run.output).toBeDefined() + + // The file-lister should have found relevant files + const outputStr = + typeof run.output === 'string' ? run.output : JSON.stringify(run.output) + + // Verify that the file-lister found some relevant files + const relevantFiles = [ + 'user-service', + 'auth-service', + 'user', + 'auth', + 'services', + ] + const foundRelevantFile = relevantFiles.some((file) => + outputStr.toLowerCase().includes(file.toLowerCase()), + ) + + expect(foundRelevantFile).toBe(true) + }, + { timeout: 60_000 }, + ) + + it( + 'should use the file tree from session state', + async () => { + const apiKey = process.env[API_KEY_ENV_VAR] + if (!apiKey) { + throw new Error('API key not found') + } + + // Create a different set of project files with a specific structure + const projectFiles: Record = { + 'packages/core/src/index.ts': 'export const VERSION = "1.0.0"', + 'packages/core/src/api/server.ts': + 'export function startServer() { console.log("started") }', + 'packages/core/src/api/routes.ts': + 'export const routes = { health: "/health" }', + 'packages/utils/src/helpers.ts': + 'export function formatDate(d: Date) { return d.toISOString() }', + 'docs/api.md': '# API Documentation\n\nAPI docs here.', + 'package.json': JSON.stringify({ name: 'mono-repo', version: '2.0.0' }), + } + + const client = new CodebuffClient({ + apiKey, + cwd: '/tmp/test-project', + projectFiles, + }) + + const events: PrintModeEvent[] = [] + + // Run file-lister to find API-related files + const run = await client.run({ + agent: 'file-lister', + prompt: 'Find files related to the API server implementation', + handleEvent: (event) => { + events.push(event) + }, + }) + + expect(run.output.type).not.toEqual('error') + + const outputStr = + typeof run.output === 'string' ? run.output : JSON.stringify(run.output) + + // Should find API-related files + const apiRelatedTerms = ['server', 'routes', 'api', 'core'] + const foundApiFile = apiRelatedTerms.some((term) => + outputStr.toLowerCase().includes(term.toLowerCase()), + ) + + expect(foundApiFile).toBe(true) + }, + { timeout: 60_000 }, + ) + + it( + 'should respect directories parameter', + async () => { + const apiKey = process.env[API_KEY_ENV_VAR] + if (!apiKey) { + throw new Error('API key not found') + } + + // Create project with multiple top-level directories + const projectFiles: Record = { + 'frontend/src/App.tsx': + 'export function App() { return
App
}', + 'frontend/src/components/Button.tsx': + 'export function Button() { return }', + 'backend/src/server.ts': + 'export function start() { console.log("started") }', + 'backend/src/routes/users.ts': + 'export function getUsers() { return [] }', + 'shared/types/common.ts': 'export type ID = string', + 'package.json': JSON.stringify({ name: 'full-stack-app' }), + } + + const client = new CodebuffClient({ + apiKey, + cwd: '/tmp/test-project', + projectFiles, + }) + + // Run file-lister with directories parameter to limit to frontend only + const run = await client.run({ + agent: 'file-lister', + prompt: 'Find React component files', + params: { + directories: ['frontend'], + }, + handleEvent: () => {}, + }) + + expect(run.output.type).not.toEqual('error') + + const outputStr = + typeof run.output === 'string' ? run.output : JSON.stringify(run.output) + + // Should find frontend files + const frontendTerms = ['app', 'button', 'component', 'frontend'] + const foundFrontendFile = frontendTerms.some((term) => + outputStr.toLowerCase().includes(term.toLowerCase()), + ) + + expect(foundFrontendFile).toBe(true) + }, + { timeout: 60_000 }, + ) +}) + +/** + * Integration tests for the file-picker agent that spawns subagents. + * The file-picker spawns file-lister as a subagent to find files. + * This tests the spawn_agents tool functionality through the SDK. + */ +describe('File Picker Agent Integration - spawn_agents tool', () => { + // Note: This test requires the local agent definitions to be used for both + // file-picker AND its spawned file-lister subagent. Currently, the spawned + // agent may resolve to the server version which has the old parsing bug. + // Skip until we have a way to ensure spawned agents use local definitions. + it.skip( + 'should spawn file-lister subagent and find relevant files', + async () => { + const apiKey = process.env[API_KEY_ENV_VAR] + if (!apiKey) { + throw new Error('API key not found') + } + + // Create mock project files + const projectFiles: Record = { + 'src/index.ts': ` +import { UserService } from './services/user-service' +export function main() { + const userService = new UserService() + console.log('Application started') +} +`, + 'src/services/user-service.ts': ` +export class UserService { + async getUser(id: string) { + return { id, name: 'John Doe' } + } +} +`, + 'src/services/auth-service.ts': ` +export class AuthService { + async login(email: string, password: string) { + return { token: 'mock-token' } + } +} +`, + 'package.json': JSON.stringify({ + name: 'test-project', + version: '1.0.0', + }), + } + + // Use local agent definitions to test the updated handleSteps + const localFilePickerDef = filePickerDefinition as unknown as any + const localFileListerDef = fileListerDefinition as unknown as any + + const client = new CodebuffClient({ + apiKey, + cwd: '/tmp/test-project-picker', + projectFiles, + agentDefinitions: [localFilePickerDef, localFileListerDef], + }) + + const events: PrintModeEvent[] = [] + + // Run the file-picker agent which spawns file-lister as a subagent + const run = await client.run({ + agent: localFilePickerDef.id, + prompt: 'Find files related to user authentication', + handleEvent: (event) => { + events.push(event) + }, + }) + + // Check for errors in the output + if (run.output.type === 'error') { + console.error('File picker error:', run.output) + } + + console.log('File picker output type:', run.output.type) + console.log('File picker output:', JSON.stringify(run.output, null, 2)) + + // The output should not be an error + expect(run.output.type).not.toEqual('error') + + // Verify we got some output + expect(run.output).toBeDefined() + + // The file-picker should have found relevant files via its spawned file-lister + const outputStr = + typeof run.output === 'string' ? run.output : JSON.stringify(run.output) + + // Verify that the file-picker found some relevant files + const relevantFiles = ['user', 'auth', 'service'] + const foundRelevantFile = relevantFiles.some((file) => + outputStr.toLowerCase().includes(file.toLowerCase()), + ) + + expect(foundRelevantFile).toBe(true) + }, + { timeout: 90_000 }, + ) +}) diff --git a/.agents/file-explorer/file-picker.ts b/.agents/file-explorer/file-picker.ts index 25f7b60080..048d904d30 100644 --- a/.agents/file-explorer/file-picker.ts +++ b/.agents/file-explorer/file-picker.ts @@ -64,17 +64,22 @@ Do not use any further tools or spawn any further agents. }, } satisfies ToolCall - const filesResult = - extractSpawnResults<{ text: string }[]>(fileListerResults)[0] - if (!Array.isArray(filesResult)) { + const spawnResults = extractSpawnResults(fileListerResults) + const firstResult = spawnResults[0] + const fileListText = extractLastMessageText(firstResult) + + if (!fileListText) { + const errorMessage = extractErrorMessage(firstResult) yield { type: 'STEP_TEXT', - text: filesResult.errorMessage, + text: errorMessage + ? `Error from file-lister: ${errorMessage}` + : 'Error: Could not extract file list from spawned agent', } satisfies StepText return } - const paths = filesResult[0].text.split('\n').filter(Boolean) + const paths = fileListText.split('\n').filter(Boolean) yield { toolName: 'read_files', @@ -85,24 +90,71 @@ Do not use any further tools or spawn any further agents. yield 'STEP' - function extractSpawnResults( - results: any[] | undefined, - ): (T | { errorMessage: string })[] { - if (!results) return [] - const spawnedResults = results - .filter((result) => result.type === 'json') - .map((result) => result.value) - .flat() as { - agentType: string - value: { value?: T; errorMessage?: string } - }[] - return spawnedResults.map( - (result) => - result.value.value ?? { - errorMessage: - result.value.errorMessage ?? 'Error extracting spawn results', - }, - ) + /** + * Extracts the array of subagent results from spawn_agents tool output. + * + * The spawn_agents tool result structure is: + * [{ type: 'json', value: [{ agentName, agentType, value: AgentOutput }] }] + * + * Returns an array of agent outputs, one per spawned agent. + */ + function extractSpawnResults(results: any[] | undefined): any[] { + if (!results || results.length === 0) return [] + + // Find the json result containing spawn results + const jsonResult = results.find((r) => r.type === 'json') + if (!jsonResult?.value) return [] + + // Get the spawned agent results array + const spawnedResults = Array.isArray(jsonResult.value) ? jsonResult.value : [jsonResult.value] + + // Extract the value (AgentOutput) from each result + return spawnedResults.map((result: any) => result?.value).filter(Boolean) + } + + /** + * Extracts the text content from a 'lastMessage' AgentOutput. + * + * For agents with outputMode: 'last_message', the output structure is: + * { type: 'lastMessage', value: [{ role: 'assistant', content: [{ type: 'text', text: '...' }] }] } + * + * Returns the text from the last assistant message, or null if not found. + */ + function extractLastMessageText(agentOutput: any): string | null { + if (!agentOutput) return null + + // Handle 'lastMessage' output mode - the value contains an array of messages + if (agentOutput.type === 'lastMessage' && Array.isArray(agentOutput.value)) { + // Find the last assistant message with text content + for (let i = agentOutput.value.length - 1; i >= 0; i--) { + const message = agentOutput.value[i] + if (message.role === 'assistant' && Array.isArray(message.content)) { + // Find text content in the message + for (const part of message.content) { + if (part.type === 'text' && typeof part.text === 'string') { + return part.text + } + } + } + } + } + + return null + } + + /** + * Extracts the error message from an AgentOutput if it's an error type. + * + * Returns the error message string, or null if not an error output. + */ + function extractErrorMessage(agentOutput: any): string | null { + if (!agentOutput) return null + + if (agentOutput.type === 'error') { + return agentOutput.message ?? agentOutput.value ?? null + } + + return null } }, } From 45c91d0f1e354dbe5b61921eff1a04d58da575ea Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 10:13:11 -0800 Subject: [PATCH 15/30] fix typecheck --- sdk/src/__tests__/run-with-retry.test.ts | 108 ++++++++++++++--------- 1 file changed, 68 insertions(+), 40 deletions(-) diff --git a/sdk/src/__tests__/run-with-retry.test.ts b/sdk/src/__tests__/run-with-retry.test.ts index cf0351cf5a..e240b8cffd 100644 --- a/sdk/src/__tests__/run-with-retry.test.ts +++ b/sdk/src/__tests__/run-with-retry.test.ts @@ -1,10 +1,12 @@ +import { assistantMessage } from '@codebuff/common/util/messages' import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test' -import { ErrorCodes, NetworkError } from '../errors' +import { ErrorCodes } from '../errors' import { run } from '../run' import * as runModule from '../run' import type { RunState } from '../run-state' +import type { SessionState } from '@codebuff/common/types/session-state' const baseOptions = { apiKey: 'test-key', @@ -19,8 +21,13 @@ describe('run retry wrapper', () => { }) it('returns immediately on success without retrying', async () => { - const expectedState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'hi' } } as RunState - const runSpy = spyOn(runModule, 'runOnce').mockResolvedValueOnce(expectedState) + const expectedState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('hi')] }, + } + const runSpy = spyOn(runModule, 'runOnce').mockResolvedValueOnce( + expectedState, + ) const result = await run(baseOptions) @@ -29,11 +36,14 @@ describe('run retry wrapper', () => { }) it('retries once on retryable error output and then succeeds', async () => { - const errorState = { - sessionState: {} as any, - output: { type: 'error', message: 'NetworkError: Service unavailable' } - } as RunState - const successState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'hi' } } as RunState + const errorState: RunState = { + sessionState: {} as SessionState, + output: { type: 'error', message: 'NetworkError: Service unavailable' }, + } + const successState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('hi')] }, + } const runSpy = spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) @@ -51,7 +61,7 @@ describe('run retry wrapper', () => { it('stops after max retries are exhausted and returns error output', async () => { const errorState = { sessionState: {} as any, - output: { type: 'error', message: 'NetworkError: Connection timeout' } + output: { type: 'error', message: 'NetworkError: Connection timeout' }, } as RunState const runSpy = spyOn(runModule, 'runOnce').mockResolvedValue(errorState) @@ -73,7 +83,7 @@ describe('run retry wrapper', () => { it('does not retry non-retryable error outputs', async () => { const errorState = { sessionState: {} as any, - output: { type: 'error', message: 'Invalid input' } + output: { type: 'error', message: 'Invalid input' }, } as RunState const runSpy = spyOn(runModule, 'runOnce').mockResolvedValue(errorState) @@ -91,7 +101,7 @@ describe('run retry wrapper', () => { it('skips retry when retry is false even for retryable error outputs', async () => { const errorState = { sessionState: {} as any, - output: { type: 'error', message: 'NetworkError: Connection failed' } + output: { type: 'error', message: 'NetworkError: Connection failed' }, } as RunState const runSpy = spyOn(runModule, 'runOnce').mockResolvedValue(errorState) @@ -106,11 +116,14 @@ describe('run retry wrapper', () => { }) it('retries when provided custom retryableErrorCodes set', async () => { - const errorState = { + const errorState: RunState = { sessionState: {} as any, - output: { type: 'error', message: 'Server error (500)' } - } as RunState - const successState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'hi' } } as RunState + output: { type: 'error', message: 'Server error (500)' }, + } + const successState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('hi')] }, + } const runSpy = spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) @@ -149,11 +162,14 @@ describe('run retry wrapper', () => { }) it('calls onRetry callback with correct parameters on error output', async () => { - const errorState = { - sessionState: {} as any, - output: { type: 'error', message: 'Service unavailable (503)' } - } as RunState - const successState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'done' } } as RunState + const errorState: RunState = { + sessionState: {} as SessionState, + output: { type: 'error', message: 'Service unavailable (503)' }, + } + const successState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('done')] }, + } const runSpy = spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) @@ -178,7 +194,7 @@ describe('run retry wrapper', () => { it('calls onRetryExhausted after all retries fail', async () => { const errorState = { sessionState: {} as any, - output: { type: 'error', message: 'NetworkError: timeout' } + output: { type: 'error', message: 'NetworkError: timeout' }, } as RunState spyOn(runModule, 'runOnce').mockResolvedValue(errorState) @@ -200,7 +216,7 @@ describe('run retry wrapper', () => { it('returns error output without sessionState on first attempt failure', async () => { const errorState = { - output: { type: 'error', message: 'Not retryable' } + output: { type: 'error', message: 'Not retryable' }, } as RunState spyOn(runModule, 'runOnce').mockResolvedValue(errorState) @@ -216,14 +232,14 @@ describe('run retry wrapper', () => { it('preserves sessionState from previousRun on retry', async () => { const previousSession = { fileContext: { cwd: '/test' } } as any - const errorState = { - sessionState: { fileContext: { cwd: '/new' } } as any, - output: { type: 'error', message: 'Service unavailable' } - } as RunState - const successState = { - sessionState: { fileContext: { cwd: '/final' } } as any, - output: { type: 'lastMessage', value: 'ok' } - } as RunState + const errorState: RunState = { + sessionState: { fileContext: { cwd: '/new' } } as SessionState, + output: { type: 'error', message: 'Service unavailable' }, + } + const successState: RunState = { + sessionState: { fileContext: { cwd: '/final' } } as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('ok')] }, + } const runSpy = spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) @@ -231,7 +247,10 @@ describe('run retry wrapper', () => { const result = await run({ ...baseOptions, - previousRun: { sessionState: previousSession, output: { type: 'lastMessage', value: 'prev' } }, + previousRun: { + sessionState: previousSession, + output: { type: 'lastMessage', value: [assistantMessage('prev')] }, + }, retry: { backoffBaseMs: 1, backoffMaxMs: 2 }, }) @@ -240,11 +259,17 @@ describe('run retry wrapper', () => { }) it('handles 503 Service Unavailable errors as retryable', async () => { - const errorState = { - sessionState: {} as any, - output: { type: 'error', message: 'Error from AI SDK: 503 Service Unavailable' } - } as RunState - const successState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'ok' } } as RunState + const errorState: RunState = { + sessionState: {} as SessionState, + output: { + type: 'error', + message: 'Error from AI SDK: 503 Service Unavailable', + }, + } + const successState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('ok')] }, + } const runSpy = spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) @@ -260,11 +285,14 @@ describe('run retry wrapper', () => { }) it('applies exponential backoff correctly', async () => { - const errorState = { - sessionState: {} as any, - output: { type: 'error', message: 'NetworkError: Connection refused' } + const errorState: RunState = { + sessionState: {} as SessionState, + output: { type: 'error', message: 'NetworkError: Connection refused' }, } as RunState - const successState = { sessionState: {} as any, output: { type: 'lastMessage', value: 'ok' } } as RunState + const successState: RunState = { + sessionState: {} as SessionState, + output: { type: 'lastMessage', value: [assistantMessage('ok')] }, + } spyOn(runModule, 'runOnce') .mockResolvedValueOnce(errorState) From 576c44f570e4d24aadb8ce48651e72c673f684ce Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 11:08:31 -0800 Subject: [PATCH 16/30] Remove spawn agents example tool call in old format --- packages/agent-runtime/src/templates/prompts.ts | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/packages/agent-runtime/src/templates/prompts.ts b/packages/agent-runtime/src/templates/prompts.ts index e1cb77d0a2..ab86aaad01 100644 --- a/packages/agent-runtime/src/templates/prompts.ts +++ b/packages/agent-runtime/src/templates/prompts.ts @@ -73,17 +73,6 @@ Notes: - There are two types of input arguments for agents: prompt and params. The prompt is a string, and the params is a json object. Some agents require only one or the other, some require both, and some require none. - Below are the *only* available agents by their agent_type. Other agents may be referenced earlier in the conversation, but they are not available to you. -Example: - -${getToolCallString('spawn_agents', { - agents: [ - { - agent_type: 'example-agent', - prompt: 'Do an example task for me', - }, - ], -})} - Spawn only the below agents: ${agentsDescription}` From a23a6980de2f06d909dc7710b1a32ff30618cf1f Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 12:58:21 -0800 Subject: [PATCH 17/30] Include stringified error for more detail --- backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts | 1 + packages/agent-runtime/src/run-agent-step.ts | 6 +++++- sdk/src/impl/llm.ts | 1 + 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts index 9a197a8658..4384286ecd 100644 --- a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts +++ b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts @@ -118,6 +118,7 @@ export async function* promptAiSdkStream( chunk: { ...chunk, error: undefined }, error: getErrorObject(chunk.error), model: params.model, + errorStr: JSON.stringify(chunk.error), }, 'Error from AI SDK', ) diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index d19c5eb673..b4d891b1e5 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -870,6 +870,7 @@ export async function loopAgentSteps( logger.error( { error: getErrorObject(error), + errorStr: JSON.stringify(error), agentType, agentId: currentAgentState.agentId, runId, @@ -881,7 +882,10 @@ export async function loopAgentSteps( ) // Re-throw NetworkError and PaymentRequiredError to allow SDK retry wrapper to handle it - if (error instanceof Error && (error.name === 'NetworkError' || error.name === 'PaymentRequiredError')) { + if ( + error instanceof Error && + (error.name === 'NetworkError' || error.name === 'PaymentRequiredError') + ) { throw error } diff --git a/sdk/src/impl/llm.ts b/sdk/src/impl/llm.ts index f743a70b7c..e149ef477e 100644 --- a/sdk/src/impl/llm.ts +++ b/sdk/src/impl/llm.ts @@ -240,6 +240,7 @@ export async function* promptAiSdkStream( chunk: { ...chunk, error: undefined }, error: getErrorObject(chunk.error), model: params.model, + errorStr: JSON.stringify(chunk.error), }, 'Error from AI SDK', ) From 8a1ba97cdf60319cf64d6049dad604b449215bb3 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 13:02:16 -0800 Subject: [PATCH 18/30] update editor best of n max + add unit tests (not fully working yet tho) --- .../editor-best-of-n.integration.test.ts | 91 +++++++++++++++ .agents/editor/best-of-n/editor-best-of-n.ts | 106 +++++++++++------- .../editor/best-of-n/editor-implementor.ts | 2 +- 3 files changed, 157 insertions(+), 42 deletions(-) create mode 100644 .agents/__tests__/editor-best-of-n.integration.test.ts diff --git a/.agents/__tests__/editor-best-of-n.integration.test.ts b/.agents/__tests__/editor-best-of-n.integration.test.ts new file mode 100644 index 0000000000..b4f8fcce0c --- /dev/null +++ b/.agents/__tests__/editor-best-of-n.integration.test.ts @@ -0,0 +1,91 @@ +import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants' +import { describe, expect, it } from 'bun:test' + +import { CodebuffClient } from '@codebuff/sdk' + +import type { PrintModeEvent } from '@codebuff/common/types/print-mode' + +/** + * Integration tests for the editor-best-of-n-max agent. + * These tests verify that the best-of-n editor workflow works correctly: + * 1. Spawns multiple implementor agents in parallel + * 2. Collects their implementation proposals + * 3. Uses a selector agent to choose the best implementation + * 4. Applies the chosen implementation + */ +describe('Editor Best-of-N Max Agent Integration', () => { + it( + 'should generate and select the best implementation for a simple edit', + async () => { + const apiKey = process.env[API_KEY_ENV_VAR] + if (!apiKey) { + throw new Error('API key not found') + } + + // Create mock project files with a simple TypeScript file to edit + const projectFiles: Record = { + 'src/utils/math.ts': ` +export function add(a: number, b: number): number { + return a + b +} + +export function subtract(a: number, b: number): number { + return a - b +} +`, + 'src/index.ts': ` +import { add, subtract } from './utils/math' + +console.log(add(1, 2)) +console.log(subtract(5, 3)) +`, + 'package.json': JSON.stringify({ + name: 'test-project', + version: '1.0.0', + dependencies: {}, + }), + } + + const client = new CodebuffClient({ + apiKey, + cwd: '/tmp/test-best-of-n-project', + projectFiles, + }) + + const events: PrintModeEvent[] = [] + + // Run the editor-best-of-n-max agent with a simple task + // Using n=2 to keep the test fast while still testing the best-of-n workflow + const run = await client.run({ + agent: 'editor-best-of-n-max', + prompt: + 'Add a multiply function to src/utils/math.ts that takes two numbers and returns their product', + params: { n: 2 }, + handleEvent: (event) => { + console.log(event) + events.push(event) + }, + }) + + // The output should not be an error + expect(run.output.type).not.toEqual('error') + + // Verify we got some output + expect(run.output).toBeDefined() + + // The output should contain the implementation response + const outputStr = + typeof run.output === 'string' ? run.output : JSON.stringify(run.output) + console.log('Output:', outputStr) + + // Should contain evidence of the multiply function being added + const relevantTerms = ['multiply', 'product', 'str_replace', 'write_file'] + const foundRelevantTerm = relevantTerms.some((term) => + outputStr.toLowerCase().includes(term.toLowerCase()), + ) + + expect(foundRelevantTerm).toBe(true) + }, + { timeout: 120_000 }, // 2 minute timeout for best-of-n workflow + ) +}) diff --git a/.agents/editor/best-of-n/editor-best-of-n.ts b/.agents/editor/best-of-n/editor-best-of-n.ts index d9dd526344..6ba452f0c8 100644 --- a/.agents/editor/best-of-n/editor-best-of-n.ts +++ b/.agents/editor/best-of-n/editor-best-of-n.ts @@ -39,11 +39,11 @@ export function createBestOfNEditor( spawnableAgents: buildArray( 'best-of-n-selector', 'best-of-n-selector-opus', - isDefault && 'best-of-n-selector-gemini', + 'best-of-n-selector-gemini', 'editor-implementor', 'editor-implementor-opus', - isDefault && 'editor-implementor-gemini', - isMax && 'editor-implementor-gpt-5', + 'editor-implementor-gemini', + 'editor-implementor-gpt-5', ), inputSchema: { @@ -230,6 +230,7 @@ function* handleStepsDefault({ } function* handleStepsMax({ params, + logger, }: AgentStepContext): ReturnType< NonNullable > { @@ -269,8 +270,9 @@ function* handleStepsMax({ } satisfies ToolCall<'spawn_agents'> // Extract spawn results - const spawnedImplementations = - extractSpawnResults<{ text: string }[]>(implementorResults) + const spawnedImplementations = extractSpawnResults( + implementorResults, + ) as any[] // Extract all the plans from the structured outputs const letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' @@ -280,9 +282,14 @@ function* handleStepsMax({ content: 'errorMessage' in result ? `Error: ${result.errorMessage}` - : result[0].text, + : extractLastMessageText(result), })) + logger.info( + { spawnedImplementations, implementations }, + 'spawnedImplementations', + ) + // Spawn selector with implementations as params const { toolResult: selectorResult } = yield { toolName: 'spawn_agents', @@ -321,15 +328,9 @@ function* handleStepsMax({ return } - // Apply the chosen implementation using STEP_TEXT (only tool calls, no commentary) - const toolCallsOnly = extractToolCallsOnly( - typeof chosenImplementation.content === 'string' - ? chosenImplementation.content - : '', - ) const { agentState: postEditsAgentState } = yield { type: 'STEP_TEXT', - text: toolCallsOnly, + text: chosenImplementation.content, } as StepText const { messageHistory } = postEditsAgentState const lastAssistantMessageIndex = messageHistory.findLastIndex( @@ -352,37 +353,60 @@ function* handleStepsMax({ includeToolCall: false, } satisfies ToolCall<'set_output'> - function extractSpawnResults( - results: any[] | undefined, - ): (T | { errorMessage: string })[] { - if (!results) return [] - const spawnedResults = results - .filter((result) => result.type === 'json') - .map((result) => result.value) - .flat() as { - agentType: string - value: { value?: T; errorMessage?: string } - }[] - return spawnedResults.map( - (result) => - result.value.value ?? { - errorMessage: - result.value.errorMessage ?? 'Error extracting spawn results', - }, - ) + /** + * Extracts the array of subagent results from spawn_agents tool output. + * + * The spawn_agents tool result structure is: + * [{ type: 'json', value: [{ agentName, agentType, value: AgentOutput }] }] + * + * Returns an array of agent outputs, one per spawned agent. + */ + function extractSpawnResults(results: any[] | undefined): T[] { + if (!results || results.length === 0) return [] + + // Find the json result containing spawn results + const jsonResult = results.find((r) => r.type === 'json') + if (!jsonResult?.value) return [] + + // Get the spawned agent results array + const spawnedResults = Array.isArray(jsonResult.value) + ? jsonResult.value + : [jsonResult.value] + + // Extract the value (AgentOutput) from each result + return spawnedResults.map((result: any) => result?.value).filter(Boolean) } - // Extract only tool calls from text, removing any commentary - function extractToolCallsOnly(text: string): string { - const toolExtractionPattern = - /\n(.*?)\n<\/codebuff_tool_call>/gs - const matches: string[] = [] - - for (const match of text.matchAll(toolExtractionPattern)) { - matches.push(match[0]) // Include the full tool call with tags + /** + * Extracts the text content from a 'lastMessage' AgentOutput. + * + * For agents with outputMode: 'last_message', the output structure is: + * { type: 'lastMessage', value: [{ role: 'assistant', content: [{ type: 'text', text: '...' }] }] } + * + * Returns the text from the last assistant message, or null if not found. + */ + function extractLastMessageText(agentOutput: any): string | null { + if (!agentOutput) return null + + // Handle 'lastMessage' output mode - the value contains an array of messages + if ( + agentOutput.type === 'lastMessage' && + Array.isArray(agentOutput.value) + ) { + // Find the last assistant message with text content + for (let i = agentOutput.value.length - 1; i >= 0; i--) { + const message = agentOutput.value[i] + if (message.role === 'assistant' && Array.isArray(message.content)) { + // Find text content in the message + for (const part of message.content) { + if (part.type === 'text' && typeof part.text === 'string') { + return part.text + } + } + } + } } - - return matches.join('\n') + return null } } diff --git a/.agents/editor/best-of-n/editor-implementor.ts b/.agents/editor/best-of-n/editor-implementor.ts index f159df2ce6..630803a47d 100644 --- a/.agents/editor/best-of-n/editor-implementor.ts +++ b/.agents/editor/best-of-n/editor-implementor.ts @@ -37,7 +37,7 @@ export const createBestOfNImplementor = (options: { Your task is to write out ALL the code changes needed to complete the user's request in a single comprehensive response. -Important: You can not make any other tool calls besides editing files. You cannot read more files, write todos, or spawn agents. +Important: You can not make any other tool calls besides editing files. You cannot read more files, write todos, or spawn agents. Do not call any of these tools! Write out what changes you would make using the tool call format below. Use this exact format for each file change: From 9f9f464ea51bfc4a3881c747968941466f98579e Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 13:36:33 -0800 Subject: [PATCH 19/30] Revert "Include stringified error for more detail" This reverts commit a23a6980de2f06d909dc7710b1a32ff30618cf1f. --- backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts | 1 - packages/agent-runtime/src/run-agent-step.ts | 6 +----- sdk/src/impl/llm.ts | 1 - 3 files changed, 1 insertion(+), 7 deletions(-) diff --git a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts index 4384286ecd..9a197a8658 100644 --- a/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts +++ b/backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts @@ -118,7 +118,6 @@ export async function* promptAiSdkStream( chunk: { ...chunk, error: undefined }, error: getErrorObject(chunk.error), model: params.model, - errorStr: JSON.stringify(chunk.error), }, 'Error from AI SDK', ) diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index b4d891b1e5..d19c5eb673 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -870,7 +870,6 @@ export async function loopAgentSteps( logger.error( { error: getErrorObject(error), - errorStr: JSON.stringify(error), agentType, agentId: currentAgentState.agentId, runId, @@ -882,10 +881,7 @@ export async function loopAgentSteps( ) // Re-throw NetworkError and PaymentRequiredError to allow SDK retry wrapper to handle it - if ( - error instanceof Error && - (error.name === 'NetworkError' || error.name === 'PaymentRequiredError') - ) { + if (error instanceof Error && (error.name === 'NetworkError' || error.name === 'PaymentRequiredError')) { throw error } diff --git a/sdk/src/impl/llm.ts b/sdk/src/impl/llm.ts index e149ef477e..f743a70b7c 100644 --- a/sdk/src/impl/llm.ts +++ b/sdk/src/impl/llm.ts @@ -240,7 +240,6 @@ export async function* promptAiSdkStream( chunk: { ...chunk, error: undefined }, error: getErrorObject(chunk.error), model: params.model, - errorStr: JSON.stringify(chunk.error), }, 'Error from AI SDK', ) From 71c285bebb1de7b90466134478469f7d87ef4b1a Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 13:37:42 -0800 Subject: [PATCH 20/30] web: Pass open router errors through --- web/src/app/api/v1/chat/completions/_post.ts | 7 ++ web/src/llm-api/openrouter.ts | 101 +++++++++++++++++-- 2 files changed, 97 insertions(+), 11 deletions(-) diff --git a/web/src/app/api/v1/chat/completions/_post.ts b/web/src/app/api/v1/chat/completions/_post.ts index 0762f3f1b7..45f99d675a 100644 --- a/web/src/app/api/v1/chat/completions/_post.ts +++ b/web/src/app/api/v1/chat/completions/_post.ts @@ -12,6 +12,7 @@ import { import { handleOpenRouterNonStream, handleOpenRouterStream, + OpenRouterError, } from '@/llm-api/openrouter' import { extractApiKeyFromHeader } from '@/util/auth' @@ -339,6 +340,12 @@ export async function postChatCompletions(params: { }, logger, }) + + // Pass through OpenRouter provider-specific errors + if (error instanceof OpenRouterError) { + return NextResponse.json(error.toJSON(), { status: error.statusCode }) + } + return NextResponse.json( { error: 'Failed to process request' }, { status: 500 }, diff --git a/web/src/llm-api/openrouter.ts b/web/src/llm-api/openrouter.ts index d9a85ed640..173eb9bfc6 100644 --- a/web/src/llm-api/openrouter.ts +++ b/web/src/llm-api/openrouter.ts @@ -6,7 +6,10 @@ import { extractRequestMetadata, insertMessageToBigQuery, } from './helpers' -import { OpenRouterStreamChatCompletionChunkSchema } from './type/openrouter' +import { + OpenRouterErrorResponseSchema, + OpenRouterStreamChatCompletionChunkSchema, +} from './type/openrouter' import type { UsageData } from './helpers' import type { OpenRouterStreamChatCompletionChunk } from './type/openrouter' @@ -14,7 +17,6 @@ import type { InsertMessageBigqueryFn } from '@codebuff/common/types/contracts/b import type { Logger } from '@codebuff/common/types/contracts/logger' type StreamState = { responseText: string; reasoningText: string } - function createOpenRouterRequest(params: { body: any openrouterApiKey: string | null @@ -93,9 +95,9 @@ export async function handleOpenRouterNonStream({ const responses = await Promise.all(requests) if (responses.every((r) => !r.ok)) { - throw new Error( - `Failed to make all ${n} requests: ${responses.map((r) => r.statusText).join(', ')}`, - ) + // Return provider-specific error from the first failed response + const firstFailedResponse = responses[0] + throw await parseOpenRouterError(firstFailedResponse) } const allData = await Promise.all(responses.map((r) => r.json())) @@ -183,9 +185,7 @@ export async function handleOpenRouterNonStream({ }) if (!response.ok) { - throw new Error( - `OpenRouter API error (${response.statusText}): ${await response.text()}`, - ) + throw await parseOpenRouterError(response) } const data = await response.json() @@ -261,9 +261,7 @@ export async function handleOpenRouterStream({ }) if (!response.ok) { - throw new Error( - `OpenRouter API error (${response.statusText}): ${await response.text()}`, - ) + throw await parseOpenRouterError(response) } const reader = response.body?.getReader() @@ -532,3 +530,84 @@ async function handleStreamChunk({ state.reasoningText += choice.delta?.reasoning ?? '' return state } + +/** + * Custom error class for OpenRouter API errors that preserves provider-specific details. + */ +export class OpenRouterError extends Error { + constructor( + public readonly statusCode: number, + public readonly statusText: string, + public readonly errorBody: { + error: { + message: string + code: string | number | null + type?: string | null + param?: unknown + metadata?: { + raw?: string + provider_name?: string + } + } + }, + ) { + super(errorBody.error.message) + this.name = 'OpenRouterError' + } + + /** + * Returns the error in a format suitable for API responses. + */ + toJSON() { + return { + error: { + message: this.errorBody.error.message, + code: this.errorBody.error.code, + type: this.errorBody.error.type, + param: this.errorBody.error.param, + metadata: this.errorBody.error.metadata, + }, + } + } +} + +/** + * Parses an error response from OpenRouter and returns an OpenRouterError. + */ +async function parseOpenRouterError( + response: Response, +): Promise { + const errorText = await response.text() + let errorBody: OpenRouterError['errorBody'] + try { + const parsed = JSON.parse(errorText) + const validated = OpenRouterErrorResponseSchema.safeParse(parsed) + if (validated.success) { + errorBody = { + error: { + message: validated.data.error.message, + code: validated.data.error.code ?? null, + type: validated.data.error.type, + param: validated.data.error.param, + // metadata is not in the schema but OpenRouter includes it for provider errors + metadata: (parsed as any).error?.metadata, + }, + } + } else { + errorBody = { + error: { + message: errorText || response.statusText, + code: response.status, + }, + } + } + } catch { + errorBody = { + error: { + message: errorText || response.statusText, + code: response.status, + }, + } + } + return new OpenRouterError(response.status, response.statusText, errorBody) +} From d5a5381472eb352bbfa3143680d3284c56302b1e Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 13:55:09 -0800 Subject: [PATCH 21/30] fix cost-aggregation integration tests --- .../cost-aggregation.integration.test.ts | 84 ++++++++++++++----- 1 file changed, 61 insertions(+), 23 deletions(-) diff --git a/backend/src/__tests__/cost-aggregation.integration.test.ts b/backend/src/__tests__/cost-aggregation.integration.test.ts index 3206f3e5d0..5dd6a5cd83 100644 --- a/backend/src/__tests__/cost-aggregation.integration.test.ts +++ b/backend/src/__tests__/cost-aggregation.integration.test.ts @@ -4,6 +4,7 @@ import * as agentRegistry from '@codebuff/agent-runtime/templates/agent-registry import { TEST_USER_ID } from '@codebuff/common/old-constants' import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime' import { getInitialSessionState } from '@codebuff/common/types/session-state' +import { generateCompactId } from '@codebuff/common/util/string' import { spyOn, beforeEach, @@ -22,6 +23,7 @@ import type { AgentRuntimeScopedDeps, } from '@codebuff/common/types/contracts/agent-runtime' import type { SendActionFn } from '@codebuff/common/types/contracts/client' +import type { StreamChunk } from '@codebuff/common/types/contracts/llm' import type { ParamsExcluding } from '@codebuff/common/types/function-params' import type { ProjectFileContext } from '@codebuff/common/util/file' import type { Mock } from 'bun:test' @@ -149,15 +151,30 @@ describe('Cost Aggregation Integration Tests', () => { if (callCount === 1) { // Main agent spawns a subagent yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write a simple hello world file"}]}\n', - } + type: 'tool-call', + toolName: 'spawn_agents', + toolCallId: generateCompactId('test-id-'), + input: { + agents: [ + { + agent_type: 'editor', + prompt: 'Write a simple hello world file', + }, + ], + }, + } satisfies StreamChunk } else { // Subagent writes a file yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "write_file", "path": "hello.txt", "instructions": "Create hello world file", "content": "Hello, World!"}\n', - } + type: 'tool-call', + toolName: 'write_file', + toolCallId: generateCompactId('test-id-'), + input: { + path: 'hello.txt', + instructions: 'Create hello world file', + content: 'Hello, World!', + }, + } satisfies StreamChunk } return 'mock-message-id' }, @@ -252,8 +269,8 @@ describe('Cost Aggregation Integration Tests', () => { // Verify the total cost includes both main agent and subagent costs const finalCreditsUsed = result.sessionState.mainAgentState.creditsUsed - // The actual cost is higher than expected due to multiple steps in agent execution - expect(finalCreditsUsed).toEqual(73) + // 10 for the first call, 7 for the subagent, 7*9 for the next 9 calls + expect(finalCreditsUsed).toEqual(80) // Verify the cost breakdown makes sense expect(finalCreditsUsed).toBeGreaterThan(0) @@ -307,21 +324,35 @@ describe('Cost Aggregation Integration Tests', () => { if (callCount === 1) { // Main agent spawns first-level subagent yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Create files"}]}\n', - } + type: 'tool-call', + toolName: 'spawn_agents', + toolCallId: generateCompactId('test-id-'), + input: { + agents: [{ agent_type: 'editor', prompt: 'Create files' }], + }, + } satisfies StreamChunk } else if (callCount === 2) { // First-level subagent spawns second-level subagent yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write specific file"}]}\n', - } + type: 'tool-call', + toolName: 'spawn_agents', + toolCallId: generateCompactId('test-id-'), + input: { + agents: [{ agent_type: 'editor', prompt: 'Write specific file' }], + }, + } satisfies StreamChunk } else { // Second-level subagent does actual work yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "write_file", "path": "nested.txt", "instructions": "Create nested file", "content": "Nested content"}\n', - } + type: 'tool-call', + toolName: 'write_file', + toolCallId: generateCompactId('test-id-'), + input: { + path: 'nested.txt', + instructions: 'Create nested file', + content: 'Nested content', + }, + } satisfies StreamChunk } return 'mock-message-id' @@ -348,8 +379,8 @@ describe('Cost Aggregation Integration Tests', () => { // Should aggregate costs from all levels: main + sub1 + sub2 const finalCreditsUsed = result.sessionState.mainAgentState.creditsUsed - // Multi-level agents should have higher costs than simple ones - expect(finalCreditsUsed).toEqual(50) + // 10 calls from base agent, 1 from first subagent, 1 from second subagent: 12 calls total + expect(finalCreditsUsed).toEqual(60) }) it('should maintain cost integrity when subagents fail', async () => { @@ -365,12 +396,19 @@ describe('Cost Aggregation Integration Tests', () => { if (callCount === 1) { // Main agent spawns subagent yield { - type: 'text' as const, - text: '\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "This will fail"}]}\n', - } + type: 'tool-call', + toolName: 'spawn_agents', + toolCallId: generateCompactId('test-id-'), + input: { + agents: [{ agent_type: 'editor', prompt: 'This will fail' }], + }, + } satisfies StreamChunk } else { // Subagent fails after incurring cost - yield { type: 'text' as const, text: 'Some response' } + yield { + type: 'text', + text: 'Some response', + } satisfies StreamChunk throw new Error('Subagent execution failed') } From ca0399f833cc11d71b9224b8d2842b3ddf764216 Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 13:56:53 -0800 Subject: [PATCH 22/30] fix typecheck? --- .agents/package.json | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.agents/package.json b/.agents/package.json index 436971bdec..b995f9b5d3 100644 --- a/.agents/package.json +++ b/.agents/package.json @@ -6,5 +6,8 @@ "scripts": { "typecheck": "bun x tsc --noEmit -p tsconfig.json", "test": "bun test" + }, + "dependencies": { + "@codebuff/sdk": "workspace:*" } } From 839ef3e50d78759d060012c0707540ca0e83a860 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 14:12:43 -0800 Subject: [PATCH 23/30] editor: don't include spawn agents tool call so anthropic api doesn't error --- .agents/editor/best-of-n/editor-best-of-n.ts | 23 ++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/.agents/editor/best-of-n/editor-best-of-n.ts b/.agents/editor/best-of-n/editor-best-of-n.ts index 6ba452f0c8..6c58a327d2 100644 --- a/.agents/editor/best-of-n/editor-best-of-n.ts +++ b/.agents/editor/best-of-n/editor-best-of-n.ts @@ -229,6 +229,7 @@ function* handleStepsDefault({ } } function* handleStepsMax({ + agentState, params, logger, }: AgentStepContext): ReturnType< @@ -255,6 +256,28 @@ function* handleStepsMax({ 'editor-implementor-opus', ] as const + // Only keep messages up to just before the last spawn agent tool call. + const { messageHistory: initialMessageHistory } = agentState + const lastSpawnAgentMessageIndex = initialMessageHistory.findLastIndex( + (message) => + message.role === 'assistant' && + Array.isArray(message.content) && + message.content.length > 0 && + message.content[0].type === 'tool-call' && + message.content[0].toolName === 'spawn_agents', + ) + const updatedMessageHistory = initialMessageHistory.slice( + 0, + lastSpawnAgentMessageIndex, + ) + yield { + toolName: 'set_messages', + input: { + messages: updatedMessageHistory, + }, + includeToolCall: false, + } satisfies ToolCall<'set_messages'> + // Spawn implementor agents using the model pattern const implementorAgents = MAX_MODEL_PATTERN.slice(0, n).map((agent_type) => ({ agent_type, From e10ff729401861caa9912e7fb8286608272f259d Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 14:14:46 -0800 Subject: [PATCH 24/30] fix typecheck for .agents --- .agents/package.json | 3 --- .agents/tsconfig.json | 1 + 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/.agents/package.json b/.agents/package.json index b995f9b5d3..436971bdec 100644 --- a/.agents/package.json +++ b/.agents/package.json @@ -6,8 +6,5 @@ "scripts": { "typecheck": "bun x tsc --noEmit -p tsconfig.json", "test": "bun test" - }, - "dependencies": { - "@codebuff/sdk": "workspace:*" } } diff --git a/.agents/tsconfig.json b/.agents/tsconfig.json index 4387f3d664..dbb372c162 100644 --- a/.agents/tsconfig.json +++ b/.agents/tsconfig.json @@ -5,6 +5,7 @@ "skipLibCheck": true, "types": ["bun", "node"], "paths": { + "@codebuff/sdk": ["../sdk/src/index.ts"], "@codebuff/common/*": ["../common/src/*"] } }, From 8994cb74d04632ed0ea4e0bac3523c452d5675b8 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 14:21:32 -0800 Subject: [PATCH 25/30] Validate handleSteps yield values with zod! --- common/src/types/agent-template.ts | 30 ++- .../__tests__/run-programmatic-step.test.ts | 235 ++++++++++++++++++ .../src/run-programmatic-step.ts | 12 + 3 files changed, 276 insertions(+), 1 deletion(-) diff --git a/common/src/types/agent-template.ts b/common/src/types/agent-template.ts index 77989fc6da..9cd57c24d9 100644 --- a/common/src/types/agent-template.ts +++ b/common/src/types/agent-template.ts @@ -5,6 +5,8 @@ * It imports base types from the user-facing template to eliminate duplication. */ +import { z } from 'zod/v4' + import type { MCPConfig } from './mcp' import type { Model } from '../old-constants' import type { ToolResultOutput } from './messages/content-part' @@ -15,7 +17,6 @@ import type { } from '../templates/initial-agents-dir/types/agent-definition' import type { Logger } from '../templates/initial-agents-dir/types/util-types' import type { ToolName } from '../tools/constants' -import type { z } from 'zod/v4' export type AgentId = `${string}/${string}@${number}.${number}.${number}` @@ -141,6 +142,33 @@ export type AgentTemplate< export type StepText = { type: 'STEP_TEXT'; text: string } export type GenerateN = { type: 'GENERATE_N'; n: number } +// Zod schemas for handleSteps yield values +export const StepTextSchema = z.object({ + type: z.literal('STEP_TEXT'), + text: z.string(), +}) + +export const GenerateNSchema = z.object({ + type: z.literal('GENERATE_N'), + n: z.number().int().positive(), +}) + +export const HandleStepsToolCallSchema = z.object({ + toolName: z.string().min(1), + input: z.record(z.string(), z.any()), + includeToolCall: z.boolean().optional(), +}) + +export const HandleStepsYieldValueSchema = z.union([ + z.literal('STEP'), + z.literal('STEP_ALL'), + StepTextSchema, + GenerateNSchema, + HandleStepsToolCallSchema, +]) + +export type HandleStepsYieldValue = z.infer + export type StepGenerator = Generator< Omit | 'STEP' | 'STEP_ALL' | StepText | GenerateN, // Generic tool call type void, diff --git a/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts b/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts index df7ded81d2..777ce76b42 100644 --- a/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts +++ b/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts @@ -1433,6 +1433,241 @@ describe('runProgrammaticStep', () => { }) }) + describe('yield value validation', () => { + it('should reject invalid yield values', async () => { + const mockGenerator = (function* () { + yield { invalid: 'value' } as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const responseChunks: any[] = [] + mockParams.onResponseChunk = (chunk) => responseChunks.push(chunk) + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject yield values with wrong types', async () => { + const mockGenerator = (function* () { + yield { type: 'STEP_TEXT', text: 123 } as any // text should be string + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const responseChunks: any[] = [] + mockParams.onResponseChunk = (chunk) => responseChunks.push(chunk) + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject GENERATE_N with non-positive n', async () => { + const mockGenerator = (function* () { + yield { type: 'GENERATE_N', n: 0 } as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const responseChunks: any[] = [] + mockParams.onResponseChunk = (chunk) => responseChunks.push(chunk) + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject GENERATE_N with negative n', async () => { + const mockGenerator = (function* () { + yield { type: 'GENERATE_N', n: -5 } as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const responseChunks: any[] = [] + mockParams.onResponseChunk = (chunk) => responseChunks.push(chunk) + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should accept valid STEP literal', async () => { + const mockGenerator = (function* () { + yield 'STEP' + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(false) + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should accept valid STEP_ALL literal', async () => { + const mockGenerator = (function* () { + yield 'STEP_ALL' + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(false) + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should accept valid STEP_TEXT object', async () => { + const mockGenerator = (function* () { + yield { type: 'STEP_TEXT', text: 'Custom response text' } + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(false) + expect(result.textOverride).toBe('Custom response text') + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should accept valid GENERATE_N object', async () => { + const mockGenerator = (function* () { + yield { type: 'GENERATE_N', n: 3 } + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(false) + expect(result.generateN).toBe(3) + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should accept valid tool call object', async () => { + const mockGenerator = (function* () { + yield { toolName: 'read_files', input: { paths: ['test.txt'] } } + yield { toolName: 'end_turn', input: {} } + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should accept tool call with includeToolCall option', async () => { + const mockGenerator = (function* () { + yield { + toolName: 'read_files', + input: { paths: ['test.txt'] }, + includeToolCall: false, + } + yield { toolName: 'end_turn', input: {} } + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toBeUndefined() + }) + + it('should reject random string values', async () => { + const mockGenerator = (function* () { + yield 'INVALID_STEP' as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject null yield values', async () => { + const mockGenerator = (function* () { + yield null as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject undefined yield values', async () => { + const mockGenerator = (function* () { + yield undefined as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject tool call without toolName', async () => { + const mockGenerator = (function* () { + yield { input: { paths: ['test.txt'] } } as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + + it('should reject tool call without input', async () => { + const mockGenerator = (function* () { + yield { toolName: 'read_files' } as any + })() as StepGenerator + + mockTemplate.handleSteps = () => mockGenerator + + const result = await runProgrammaticStep(mockParams) + + expect(result.endTurn).toBe(true) + expect(result.agentState.output?.error).toContain( + 'Invalid yield value from handleSteps', + ) + }) + }) + describe('logging and context', () => { it('should log agent execution start', async () => { const mockGenerator = (function* () { diff --git a/packages/agent-runtime/src/run-programmatic-step.ts b/packages/agent-runtime/src/run-programmatic-step.ts index 8af7234755..d3426e8314 100644 --- a/packages/agent-runtime/src/run-programmatic-step.ts +++ b/packages/agent-runtime/src/run-programmatic-step.ts @@ -7,6 +7,8 @@ import { executeToolCall } from './tools/tool-executor' import type { FileProcessingState } from './tools/handlers/tool/write-file' import type { ExecuteToolCallParams } from './tools/tool-executor' import type { CodebuffToolCall } from '@codebuff/common/tools/list' +import { HandleStepsYieldValueSchema } from '@codebuff/common/types/agent-template' + import type { AgentTemplate, StepGenerator, @@ -234,6 +236,16 @@ export async function runProgrammaticStep( endTurn = true break } + + // Validate the yield value from handleSteps + const parseResult = HandleStepsYieldValueSchema.safeParse(result.value) + if (!parseResult.success) { + throw new Error( + `Invalid yield value from handleSteps in agent ${template.id}: ${parseResult.error.message}. ` + + `Received: ${JSON.stringify(result.value)}`, + ) + } + if (result.value === 'STEP') { break } From 3d2a1db521e164fc94bd7d3eae3349f15198dcff Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 14:31:44 -0800 Subject: [PATCH 26/30] fix common unit tests --- common/src/util/__tests__/messages.test.ts | 45 ++++++++++++---------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/common/src/util/__tests__/messages.test.ts b/common/src/util/__tests__/messages.test.ts index 53e1cb7225..72658d1a0a 100644 --- a/common/src/util/__tests__/messages.test.ts +++ b/common/src/util/__tests__/messages.test.ts @@ -13,6 +13,7 @@ import { } from '../messages' import type { Message } from '../../types/messages/codebuff-message' +import type { AssistantModelMessage, ToolResultPart } from 'ai' describe('withCacheControl', () => { it('should add cache control to object without providerOptions', () => { @@ -189,12 +190,6 @@ describe('convertCbToModelMessages', () => { describe('tool message conversion', () => { it('should convert tool messages with JSON output', () => { - const toolResult = [ - { - type: 'json', - value: { result: 'success' }, - }, - ] const messages: Message[] = [ { role: 'tool', @@ -211,15 +206,17 @@ describe('convertCbToModelMessages', () => { expect(result).toEqual([ expect.objectContaining({ - role: 'user', + role: 'tool', content: [ expect.objectContaining({ - type: 'text', - }), + type: 'tool-result', + toolCallId: 'call_123', + toolName: 'test_tool', + output: { type: 'json', value: { result: 'success' } }, + } satisfies ToolResultPart), ], }), ]) - expect((result as any)[0].content[0].text).toContain('') }) it('should convert tool messages with media output', () => { @@ -270,14 +267,15 @@ describe('convertCbToModelMessages', () => { includeCacheControl: false, }) - console.dir({ result }, { depth: null }) // Multiple tool outputs are aggregated into one user message expect(result).toEqual([ expect.objectContaining({ - role: 'user', + role: 'tool', + }), + expect.objectContaining({ + role: 'tool', }), ]) - expect(result[0].content).toHaveLength(2) }) }) @@ -806,14 +804,19 @@ describe('convertCbToModelMessages', () => { includeCacheControl: false, }) - expect(result).toHaveLength(1) - expect(result[0].role).toBe('assistant') - if (typeof result[0].content !== 'string') { - expect(result[0].content[0].type).toBe('text') - if (result[0].content[0].type === 'text') { - expect(result[0].content[0].text).toContain('test_tool') - } - } + expect(result).toEqual([ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call_123', + toolName: 'test_tool', + input: { param: 'value' }, + }, + ], + } satisfies AssistantModelMessage, + ]) }) it('should preserve message metadata during conversion', () => { From 567d6615a009445034b66529bc4a127194bee976 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 14:52:01 -0800 Subject: [PATCH 27/30] work around readonly messageHistory --- packages/agent-runtime/src/run-programmatic-step.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/agent-runtime/src/run-programmatic-step.ts b/packages/agent-runtime/src/run-programmatic-step.ts index d3426e8314..98ef557831 100644 --- a/packages/agent-runtime/src/run-programmatic-step.ts +++ b/packages/agent-runtime/src/run-programmatic-step.ts @@ -300,6 +300,8 @@ export async function runProgrammaticStep( // agentId: agentState.agentId, // parentAgentId: agentState.parentId, // }) + // NOTE(James): agentState.messageHistory is readonly for some reason (?!). Recreating the array is a workaround. We should figure out why it's frozen. + agentState.messageHistory = [...agentState.messageHistory] agentState.messageHistory.push(assistantMessage(toolCallPart)) // Optional call handles both top-level and nested agents // sendSubagentChunk({ From b08c8f039535a6efbeb1f58ff1b7170a50a0c003 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 15:09:14 -0800 Subject: [PATCH 28/30] tweak best of n selector to have the impl --- .agents/editor/best-of-n/editor-best-of-n.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/.agents/editor/best-of-n/editor-best-of-n.ts b/.agents/editor/best-of-n/editor-best-of-n.ts index 6c58a327d2..7fed3a8694 100644 --- a/.agents/editor/best-of-n/editor-best-of-n.ts +++ b/.agents/editor/best-of-n/editor-best-of-n.ts @@ -328,8 +328,10 @@ function* handleStepsMax({ } satisfies ToolCall<'spawn_agents'> const selectorOutput = extractSpawnResults<{ - implementationId: string - reasoning: string + value: { + implementationId: string + reasoning: string + } }>(selectorResult)[0] if ('errorMessage' in selectorOutput) { @@ -339,7 +341,7 @@ function* handleStepsMax({ } satisfies ToolCall<'set_output'> return } - const { implementationId } = selectorOutput + const { implementationId } = selectorOutput.value const chosenImplementation = implementations.find( (implementation) => implementation.id === implementationId, ) From 72c617a441ad2d0aa817e1f1dbb9b0ea15015d08 Mon Sep 17 00:00:00 2001 From: James Grugett Date: Wed, 26 Nov 2025 16:38:16 -0800 Subject: [PATCH 29/30] Implement STEP_TEXT within custom parsing within run-programmatic-step. Get best of n editor working! --- .agents/editor/best-of-n/editor-best-of-n.ts | 18 +- .../editor/best-of-n/editor-implementor.ts | 2 +- evals/scaffolding.ts | 1 - .../src/__tests__/n-parameter.test.ts | 1 - .../src/__tests__/read-docs-tool.test.ts | 2 - .../__tests__/run-agent-step-tools.test.ts | 1 - .../__tests__/run-programmatic-step.test.ts | 1 - .../src/__tests__/web-search-tool.test.ts | 1 - .../agent-runtime/src/prompt-agent-stream.ts | 10 - packages/agent-runtime/src/run-agent-step.ts | 4 - .../src/run-programmatic-step.ts | 313 +++++++++------ .../parse-tool-calls-from-text.test.ts | 363 ++++++++++++++++++ .../src/util/parse-tool-calls-from-text.ts | 117 ++++++ 13 files changed, 680 insertions(+), 154 deletions(-) create mode 100644 packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts create mode 100644 packages/agent-runtime/src/util/parse-tool-calls-from-text.ts diff --git a/.agents/editor/best-of-n/editor-best-of-n.ts b/.agents/editor/best-of-n/editor-best-of-n.ts index 7fed3a8694..7592a17856 100644 --- a/.agents/editor/best-of-n/editor-best-of-n.ts +++ b/.agents/editor/best-of-n/editor-best-of-n.ts @@ -314,7 +314,7 @@ function* handleStepsMax({ ) // Spawn selector with implementations as params - const { toolResult: selectorResult } = yield { + const { toolResult: selectorResult, agentState: selectorAgentState } = yield { toolName: 'spawn_agents', input: { agents: [ @@ -353,27 +353,19 @@ function* handleStepsMax({ return } + const numMessagesBeforeStepText = selectorAgentState.messageHistory.length + const { agentState: postEditsAgentState } = yield { type: 'STEP_TEXT', text: chosenImplementation.content, } as StepText const { messageHistory } = postEditsAgentState - const lastAssistantMessageIndex = messageHistory.findLastIndex( - (message) => message.role === 'assistant', - ) - const editToolResults = messageHistory - .slice(lastAssistantMessageIndex) - .filter((message) => message.role === 'tool') - .flatMap((message) => message.content) - .filter((output) => output.type === 'json') - .map((output) => output.value) - // Set output with the chosen implementation and reasoning + // Set output with the messages from running the step text of the chosen implementation yield { toolName: 'set_output', input: { - response: chosenImplementation.content, - toolResults: editToolResults, + messages: messageHistory.slice(numMessagesBeforeStepText), }, includeToolCall: false, } satisfies ToolCall<'set_output'> diff --git a/.agents/editor/best-of-n/editor-implementor.ts b/.agents/editor/best-of-n/editor-implementor.ts index 630803a47d..c27af72a23 100644 --- a/.agents/editor/best-of-n/editor-implementor.ts +++ b/.agents/editor/best-of-n/editor-implementor.ts @@ -37,7 +37,7 @@ export const createBestOfNImplementor = (options: { Your task is to write out ALL the code changes needed to complete the user's request in a single comprehensive response. -Important: You can not make any other tool calls besides editing files. You cannot read more files, write todos, or spawn agents. Do not call any of these tools! +Important: You can not make any other tool calls besides editing files. You cannot read more files, write todos, spawn agents, or set output. Do not call any of these tools! Write out what changes you would make using the tool call format below. Use this exact format for each file change: diff --git a/evals/scaffolding.ts b/evals/scaffolding.ts index fcc8cc1a68..a86b7b4e3c 100644 --- a/evals/scaffolding.ts +++ b/evals/scaffolding.ts @@ -231,7 +231,6 @@ export async function runAgentStepScaffolding( signal: new AbortController().signal, spawnParams: undefined, system: 'Test system prompt', - textOverride: null, tools: {}, userId: TEST_USER_ID, userInputId: generateCompactId(), diff --git a/packages/agent-runtime/src/__tests__/n-parameter.test.ts b/packages/agent-runtime/src/__tests__/n-parameter.test.ts index 897a520417..6cecb22f57 100644 --- a/packages/agent-runtime/src/__tests__/n-parameter.test.ts +++ b/packages/agent-runtime/src/__tests__/n-parameter.test.ts @@ -104,7 +104,6 @@ describe('n parameter and GENERATE_N functionality', () => { runAgentStepBaseParams = { ...agentRuntimeImpl, additionalToolDefinitions: () => Promise.resolve({}), - textOverride: null, runId: 'test-run-id', ancestorRunIds: [], repoId: undefined, diff --git a/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts b/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts index ea41ff9f6a..65660004cf 100644 --- a/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts +++ b/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts @@ -75,7 +75,6 @@ describe('read_docs tool with researcher agent (via web API facade)', () => { runAgentStepBaseParams = { ...agentRuntimeImpl, additionalToolDefinitions: () => Promise.resolve({}), - textOverride: null, runId: 'test-run-id', ancestorRunIds: [], repoId: undefined, @@ -215,7 +214,6 @@ describe('read_docs tool with researcher agent (via web API facade)', () => { const { agentState: newAgentState } = await runAgentStep({ ...runAgentStepBaseParams, - textOverride: null, fileContext: mockFileContextWithAgents, localAgentTemplates: agentTemplates, agentState, diff --git a/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts b/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts index 0d15eae292..f8da2e23b1 100644 --- a/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts +++ b/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts @@ -128,7 +128,6 @@ describe('runAgentStep - set_output tool', () => { signal: new AbortController().signal, spawnParams: undefined, system: 'Test system prompt', - textOverride: null, tools: {}, userId: TEST_USER_ID, userInputId: 'test-input', diff --git a/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts b/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts index 777ce76b42..c69886945e 100644 --- a/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts +++ b/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts @@ -1542,7 +1542,6 @@ describe('runProgrammaticStep', () => { const result = await runProgrammaticStep(mockParams) expect(result.endTurn).toBe(false) - expect(result.textOverride).toBe('Custom response text') expect(result.agentState.output?.error).toBeUndefined() }) diff --git a/packages/agent-runtime/src/__tests__/web-search-tool.test.ts b/packages/agent-runtime/src/__tests__/web-search-tool.test.ts index 7c3b6801b8..badc3e126e 100644 --- a/packages/agent-runtime/src/__tests__/web-search-tool.test.ts +++ b/packages/agent-runtime/src/__tests__/web-search-tool.test.ts @@ -74,7 +74,6 @@ describe('web_search tool with researcher agent (via web API facade)', () => { signal: new AbortController().signal, spawnParams: undefined, system: 'Test system prompt', - textOverride: null, tools: {}, userId: TEST_USER_ID, userInputId: 'test-input', diff --git a/packages/agent-runtime/src/prompt-agent-stream.ts b/packages/agent-runtime/src/prompt-agent-stream.ts index 04e46378e8..ecf4a691cf 100644 --- a/packages/agent-runtime/src/prompt-agent-stream.ts +++ b/packages/agent-runtime/src/prompt-agent-stream.ts @@ -26,7 +26,6 @@ export const getAgentStreamFromTemplate = (params: { runId: string sessionConnections: SessionRecord template: AgentTemplate - textOverride: string | null tools: ToolSet userId: string | undefined userInputId: string @@ -48,7 +47,6 @@ export const getAgentStreamFromTemplate = (params: { runId, sessionConnections, template, - textOverride, tools, userId, userInputId, @@ -59,14 +57,6 @@ export const getAgentStreamFromTemplate = (params: { trackEvent, } = params - if (textOverride !== null) { - async function* stream(): ReturnType { - yield { type: 'text', text: textOverride!, agentId } - return crypto.randomUUID() - } - return stream() - } - if (!template) { throw new Error('Agent template is null/undefined') } diff --git a/packages/agent-runtime/src/run-agent-step.ts b/packages/agent-runtime/src/run-agent-step.ts index d19c5eb673..fa0040ef83 100644 --- a/packages/agent-runtime/src/run-agent-step.ts +++ b/packages/agent-runtime/src/run-agent-step.ts @@ -526,7 +526,6 @@ export async function loopAgentSteps( | 'runId' | 'spawnParams' | 'system' - | 'textOverride' | 'tools' > & ParamsExcluding< @@ -716,7 +715,6 @@ export async function loopAgentSteps( const startTime = new Date() // 1. Run programmatic step first if it exists - let textOverride = null let n: number | undefined = undefined if (agentTemplate.handleSteps) { @@ -744,7 +742,6 @@ export async function loopAgentSteps( stepNumber, generateN, } = programmaticResult - textOverride = programmaticResult.textOverride n = generateN currentAgentState = programmaticAgentState @@ -808,7 +805,6 @@ export async function loopAgentSteps( runId, spawnParams: currentParams, system, - textOverride: textOverride, tools, additionalToolDefinitions: async () => { diff --git a/packages/agent-runtime/src/run-programmatic-step.ts b/packages/agent-runtime/src/run-programmatic-step.ts index 98ef557831..d7b3fcd56c 100644 --- a/packages/agent-runtime/src/run-programmatic-step.ts +++ b/packages/agent-runtime/src/run-programmatic-step.ts @@ -3,6 +3,9 @@ import { assistantMessage } from '@codebuff/common/util/messages' import { cloneDeep } from 'lodash' import { executeToolCall } from './tools/tool-executor' +import { parseTextWithToolCalls } from './util/parse-tool-calls-from-text' + +import type { ParsedSegment } from './util/parse-tool-calls-from-text' import type { FileProcessingState } from './tools/handlers/tool/write-file' import type { ExecuteToolCallParams } from './tools/tool-executor' @@ -28,7 +31,6 @@ import type { } from '@codebuff/common/types/messages/content-part' import type { PrintModeEvent } from '@codebuff/common/types/print-mode' import type { AgentState } from '@codebuff/common/types/session-state' - // Maintains generator state for all agents. Generator state can't be serialized, so we store it in memory. const runIdToGenerator: Record = {} export const runIdToStepAll: Set = new Set() @@ -93,7 +95,6 @@ export async function runProgrammaticStep( >, ): Promise<{ agentState: AgentState - textOverride: string | null endTurn: boolean stepNumber: number generateN?: number @@ -174,7 +175,7 @@ export async function runProgrammaticStep( // Clear the STEP_ALL mode. Stepping can continue if handleSteps doesn't return. runIdToStepAll.delete(agentState.runId) } else { - return { agentState, textOverride: null, endTurn: false, stepNumber } + return { agentState, endTurn: false, stepNumber } } } @@ -209,7 +210,6 @@ export async function runProgrammaticStep( let toolResult: ToolResultOutput[] | undefined = undefined let endTurn = false - let textOverride: string | null = null let generateN: number | undefined = undefined let startTime = new Date() @@ -255,7 +255,25 @@ export async function runProgrammaticStep( } if ('type' in result.value && result.value.type === 'STEP_TEXT') { - textOverride = result.value.text + // Parse text and tool calls, preserving interleaved order + const segments = parseTextWithToolCalls(result.value.text) + + if (segments.length > 0) { + // Execute segments (text and tool calls) in order + toolResult = await executeSegmentsArray(segments, { + ...params, + agentContext, + agentStepId, + agentTemplate: template, + agentState, + fileProcessingState, + fullResponse: '', + previousToolCallFinished: Promise.resolve(), + toolCalls, + toolResults, + onResponseChunk, + }) + } break } @@ -268,130 +286,22 @@ export async function runProgrammaticStep( } // Process tool calls yielded by the generator - const toolCallWithoutId = result.value - const toolCallId = crypto.randomUUID() - const toolCall = { - ...toolCallWithoutId, - toolCallId, - } as CodebuffToolCall & { - includeToolCall?: boolean - } + const toolCall = result.value as ToolCallToExecute - // Note: We don't check if the tool is available for the agent template anymore. - // You can run any tool from handleSteps now! - // if (!template.toolNames.includes(toolCall.toolName)) { - // throw new Error( - // `Tool ${toolCall.toolName} is not available for agent ${template.id}. Available tools: ${template.toolNames.join(', ')}`, - // ) - // } - - const excludeToolFromMessageHistory = toolCall?.includeToolCall === false - // Add assistant message with the tool call before executing it - if (!excludeToolFromMessageHistory) { - const toolCallPart: ToolCallPart = { - type: 'tool-call', - toolCallId, - toolName: toolCall.toolName, - input: toolCall.input, - } - // onResponseChunk({ - // ...toolCallPart, - // type: 'tool_call', - // agentId: agentState.agentId, - // parentAgentId: agentState.parentId, - // }) - // NOTE(James): agentState.messageHistory is readonly for some reason (?!). Recreating the array is a workaround. We should figure out why it's frozen. - agentState.messageHistory = [...agentState.messageHistory] - agentState.messageHistory.push(assistantMessage(toolCallPart)) - // Optional call handles both top-level and nested agents - // sendSubagentChunk({ - // userInputId, - // agentId: agentState.agentId, - // agentType: agentState.agentType!, - // chunk: toolCallString, - // forwardToPrompt: !agentState.parentId, - // }) - } - - // Execute the tool synchronously and get the result immediately - // Wrap onResponseChunk to add parentAgentId to nested agent events - await executeToolCall({ + toolResult = await executeSingleToolCall(toolCall, { ...params, - toolName: toolCall.toolName, - input: toolCall.input, - autoInsertEndStepParam: true, - excludeToolFromMessageHistory, - fromHandleSteps: true, - agentContext, agentStepId, agentTemplate: template, + agentState, fileProcessingState, fullResponse: '', previousToolCallFinished: Promise.resolve(), - toolCallId, toolCalls, toolResults, - toolResultsToAddAfterStream: [], - - onResponseChunk: (chunk: string | PrintModeEvent) => { - if (typeof chunk === 'string') { - onResponseChunk(chunk) - return - } - - // Only add parentAgentId if this programmatic agent has a parent (i.e., it's nested) - // This ensures we don't add parentAgentId to top-level spawns - if (agentState.parentId) { - const parentAgentId = agentState.agentId - - switch (chunk.type) { - case 'subagent_start': - case 'subagent_finish': - if (!chunk.parentAgentId) { - onResponseChunk({ - ...chunk, - parentAgentId, - }) - return - } - break - case 'tool_call': - case 'tool_result': { - if (!chunk.parentAgentId) { - const debugPayload = - chunk.type === 'tool_call' - ? { - eventType: chunk.type, - agentId: chunk.agentId, - parentId: parentAgentId, - } - : { - eventType: chunk.type, - parentId: parentAgentId, - } - onResponseChunk({ - ...chunk, - parentAgentId, - }) - return - } - break - } - default: - break - } - } - - // For other events or top-level spawns, send as-is - onResponseChunk(chunk) - }, + onResponseChunk, }) - // Get the latest tool result - const latestToolResult = toolResults[toolResults.length - 1] - toolResult = latestToolResult?.content - if (agentState.runId) { await addAgentStep({ ...params, @@ -416,7 +326,6 @@ export async function runProgrammaticStep( return { agentState, - textOverride, endTurn, stepNumber, generateN, @@ -460,7 +369,6 @@ export async function runProgrammaticStep( return { agentState, - textOverride: null, endTurn, stepNumber, generateN: undefined, @@ -485,3 +393,170 @@ export const getPublicAgentState = ( output, } } + +/** + * Represents a tool call to be executed. + * Can optionally include `includeToolCall: false` to exclude from message history. + */ +type ToolCallToExecute = { + toolName: string + input: Record + includeToolCall?: boolean +} + +/** + * Parameters for executing an array of tool calls. + */ +type ExecuteToolCallsArrayParams = Omit< + ExecuteToolCallParams, + | 'toolName' + | 'input' + | 'autoInsertEndStepParam' + | 'excludeToolFromMessageHistory' + | 'toolCallId' + | 'toolResultsToAddAfterStream' +> & { + agentState: AgentState + onResponseChunk: (chunk: string | PrintModeEvent) => void +} + +/** + * Executes a single tool call. + * Adds the tool call as an assistant message and then executes it. + * + * @returns The tool result from the executed tool call. + */ +async function executeSingleToolCall( + toolCallToExecute: ToolCallToExecute, + params: ExecuteToolCallsArrayParams, +): Promise { + const { agentState, 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! + // if (!template.toolNames.includes(toolCall.toolName)) { + // throw new Error( + // `Tool ${toolCall.toolName} is not available for agent ${template.id}. Available tools: ${template.toolNames.join(', ')}`, + // ) + // } + + const toolCallId = crypto.randomUUID() + const excludeToolFromMessageHistory = + toolCallToExecute.includeToolCall === false + + // Add assistant message with the tool call before executing it + if (!excludeToolFromMessageHistory) { + const toolCallPart: ToolCallPart = { + type: 'tool-call', + toolCallId, + toolName: toolCallToExecute.toolName, + input: toolCallToExecute.input, + } + // onResponseChunk({ + // ...toolCallPart, + // type: 'tool_call', + // agentId: agentState.agentId, + // parentAgentId: agentState.parentId, + // }) + // NOTE(James): agentState.messageHistory is readonly for some reason (?!). Recreating the array is a workaround. + agentState.messageHistory = [...agentState.messageHistory] + agentState.messageHistory.push(assistantMessage(toolCallPart)) + // Optional call handles both top-level and nested agents + // sendSubagentChunk({ + // userInputId, + // agentId: agentState.agentId, + // agentType: agentState.agentType!, + // chunk: toolCallString, + // forwardToPrompt: !agentState.parentId, + // }) + } + + // Execute the tool call + await executeToolCall({ + ...params, + toolName: toolCallToExecute.toolName as any, + input: toolCallToExecute.input, + autoInsertEndStepParam: true, + excludeToolFromMessageHistory, + fromHandleSteps: true, + toolCallId, + toolResultsToAddAfterStream: [], + + onResponseChunk: (chunk: string | PrintModeEvent) => { + if (typeof chunk === 'string') { + onResponseChunk(chunk) + return + } + + // Only add parentAgentId if this programmatic agent has a parent (i.e., it's nested) + // This ensures we don't add parentAgentId to top-level spawns + if (agentState.parentId) { + const parentAgentId = agentState.agentId + + switch (chunk.type) { + case 'subagent_start': + case 'subagent_finish': + if (!chunk.parentAgentId) { + onResponseChunk({ + ...chunk, + parentAgentId, + }) + return + } + break + case 'tool_call': + case 'tool_result': { + if (!chunk.parentAgentId) { + onResponseChunk({ + ...chunk, + parentAgentId, + }) + return + } + break + } + default: + break + } + } + + // For other events or top-level spawns, send as-is + onResponseChunk(chunk) + }, + }) + + // Get the latest tool result + return toolResults[toolResults.length - 1]?.content +} + +/** + * Executes an array of segments (text and tool calls) sequentially. + * Text segments are added as assistant messages. + * Tool calls are added as assistant messages and then executed. + * + * @returns The tool result from the last executed tool call. + */ +async function executeSegmentsArray( + segments: ParsedSegment[], + params: ExecuteToolCallsArrayParams, +): Promise { + const { agentState } = params + + let toolResults: ToolResultOutput[] = [] + + for (const segment of segments) { + if (segment.type === 'text') { + // Add text as an assistant message + agentState.messageHistory = [...agentState.messageHistory] + agentState.messageHistory.push(assistantMessage(segment.text)) + } else { + // Handle tool call segment + const toolResult = await executeSingleToolCall(segment, params) + if (toolResult) { + toolResults.push(...toolResult) + } + } + } + + return toolResults +} diff --git a/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts b/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts new file mode 100644 index 0000000000..a61e82703f --- /dev/null +++ b/packages/agent-runtime/src/util/__tests__/parse-tool-calls-from-text.test.ts @@ -0,0 +1,363 @@ +import { describe, expect, it } from 'bun:test' + +import { + parseToolCallsFromText, + parseTextWithToolCalls, +} from '../parse-tool-calls-from-text' + +describe('parseToolCallsFromText', () => { + it('should parse a single tool call', () => { + const text = ` +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"] +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0]).toEqual({ + toolName: 'read_files', + input: { paths: ['test.ts'] }, + }) + }) + + it('should parse multiple tool calls', () => { + const text = `Some commentary before + + +{ + "cb_tool_name": "read_files", + "paths": ["file1.ts"] +} + + +Some text between + + +{ + "cb_tool_name": "str_replace", + "path": "file1.ts", + "replacements": [{"old": "foo", "new": "bar"}] +} + + +Some commentary after` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(2) + expect(result[0]).toEqual({ + toolName: 'read_files', + input: { paths: ['file1.ts'] }, + }) + expect(result[1]).toEqual({ + toolName: 'str_replace', + input: { + path: 'file1.ts', + replacements: [{ old: 'foo', new: 'bar' }], + }, + }) + }) + + it('should remove cb_tool_name from input', () => { + const text = ` +{ + "cb_tool_name": "write_file", + "path": "test.ts", + "content": "console.log('hello')" +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0].input).not.toHaveProperty('cb_tool_name') + expect(result[0].input).toEqual({ + path: 'test.ts', + content: "console.log('hello')", + }) + }) + + it('should remove cb_easp from input', () => { + const text = ` +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"], + "cb_easp": true +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0].input).not.toHaveProperty('cb_easp') + expect(result[0].input).toEqual({ paths: ['test.ts'] }) + }) + + it('should skip malformed JSON', () => { + const text = ` +{ + "cb_tool_name": "read_files", + "paths": ["test.ts" +} + + + +{ + "cb_tool_name": "write_file", + "path": "good.ts", + "content": "valid" +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0].toolName).toBe('write_file') + }) + + it('should skip tool calls without cb_tool_name', () => { + const text = ` +{ + "paths": ["test.ts"] +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(0) + }) + + it('should return empty array for text without tool calls', () => { + const text = 'Just some regular text without any tool calls' + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(0) + }) + + it('should return empty array for empty string', () => { + const result = parseToolCallsFromText('') + + expect(result).toHaveLength(0) + }) + + it('should handle complex nested objects in input', () => { + const text = ` +{ + "cb_tool_name": "spawn_agents", + "agents": [ + { + "agent_type": "file-picker", + "prompt": "Find relevant files" + }, + { + "agent_type": "code-searcher", + "params": { + "searchQueries": [ + {"pattern": "function test"} + ] + } + } + ] +} +` + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0].toolName).toBe('spawn_agents') + expect(result[0].input.agents).toHaveLength(2) + }) + + it('should handle tool calls with escaped characters in strings', () => { + const text = + '\n' + + '{\n' + + ' "cb_tool_name": "str_replace",\n' + + ' "path": "test.ts",\n' + + ' "replacements": [{"old": "console.log(\\"hello\\")", "new": "console.log(\'world\')"}]\n' + + '}\n' + + '' + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + const replacements = result[0].input.replacements as Array<{ + old: string + new: string + }> + expect(replacements[0].old).toBe('console.log("hello")') + }) + + it('should handle tool calls with newlines in content', () => { + const text = + '\n' + + '{\n' + + ' "cb_tool_name": "write_file",\n' + + ' "path": "test.ts",\n' + + ' "content": "line1\\nline2\\nline3"\n' + + '}\n' + + '' + + const result = parseToolCallsFromText(text) + + expect(result).toHaveLength(1) + expect(result[0].input.content).toBe('line1\nline2\nline3') + }) +}) + +describe('parseTextWithToolCalls', () => { + it('should parse interleaved text and tool calls', () => { + const text = `Some commentary before + + +{ + "cb_tool_name": "read_files", + "paths": ["file1.ts"] +} + + +Some text between + + +{ + "cb_tool_name": "write_file", + "path": "file2.ts", + "content": "test" +} + + +Some commentary after` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(5) + expect(result[0]).toEqual({ type: 'text', text: 'Some commentary before' }) + expect(result[1]).toEqual({ + type: 'tool_call', + toolName: 'read_files', + input: { paths: ['file1.ts'] }, + }) + expect(result[2]).toEqual({ type: 'text', text: 'Some text between' }) + expect(result[3]).toEqual({ + type: 'tool_call', + toolName: 'write_file', + input: { path: 'file2.ts', content: 'test' }, + }) + expect(result[4]).toEqual({ type: 'text', text: 'Some commentary after' }) + }) + + it('should return only tool call when no surrounding text', () => { + const text = ` +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"] +} +` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(1) + expect(result[0]).toEqual({ + type: 'tool_call', + toolName: 'read_files', + input: { paths: ['test.ts'] }, + }) + }) + + it('should return only text when no tool calls', () => { + const text = 'Just some regular text without any tool calls' + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(1) + expect(result[0]).toEqual({ + type: 'text', + text: 'Just some regular text without any tool calls', + }) + }) + + it('should return empty array for empty string', () => { + const result = parseTextWithToolCalls('') + + expect(result).toHaveLength(0) + }) + + it('should handle text only before tool call', () => { + const text = `Introduction text + + +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"] +} +` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(2) + expect(result[0]).toEqual({ type: 'text', text: 'Introduction text' }) + expect(result[1].type).toBe('tool_call') + }) + + it('should handle text only after tool call', () => { + const text = ` +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"] +} + + +Conclusion text` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(2) + expect(result[0].type).toBe('tool_call') + expect(result[1]).toEqual({ type: 'text', text: 'Conclusion text' }) + }) + + it('should skip malformed tool calls but keep surrounding text', () => { + const text = `Before text + + +{ + "cb_tool_name": "read_files", + "paths": ["test.ts" +} + + +After text` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(2) + expect(result[0]).toEqual({ type: 'text', text: 'Before text' }) + expect(result[1]).toEqual({ type: 'text', text: 'After text' }) + }) + + it('should trim whitespace from text segments', () => { + const text = ` + Text with whitespace + + +{ + "cb_tool_name": "read_files", + "paths": ["test.ts"] +} + + + More text + ` + + const result = parseTextWithToolCalls(text) + + expect(result).toHaveLength(3) + expect(result[0]).toEqual({ type: 'text', text: 'Text with whitespace' }) + expect(result[1].type).toBe('tool_call') + expect(result[2]).toEqual({ type: 'text', text: 'More text' }) + }) +}) diff --git a/packages/agent-runtime/src/util/parse-tool-calls-from-text.ts b/packages/agent-runtime/src/util/parse-tool-calls-from-text.ts new file mode 100644 index 0000000000..4f9900a9ee --- /dev/null +++ b/packages/agent-runtime/src/util/parse-tool-calls-from-text.ts @@ -0,0 +1,117 @@ +import { + startToolTag, + endToolTag, + toolNameParam, +} from '@codebuff/common/tools/constants' + +export type ParsedToolCallFromText = { + type: 'tool_call' + toolName: string + input: Record +} + +export type ParsedTextSegment = { + type: 'text' + text: string +} + +export type ParsedSegment = ParsedToolCallFromText | ParsedTextSegment + +/** + * Parses text containing tool calls in the XML format, + * returning interleaved text and tool call segments in order. + * + * Example input: + * ``` + * Some text before + * + * { + * "cb_tool_name": "read_files", + * "paths": ["file.ts"] + * } + * + * Some text after + * ``` + * + * @param text - The text containing tool calls in XML format + * @returns Array of segments (text and tool calls) in order of appearance + */ +export function parseTextWithToolCalls(text: string): ParsedSegment[] { + const segments: ParsedSegment[] = [] + + // Match ... blocks + const toolExtractionPattern = new RegExp( + `${escapeRegex(startToolTag)}([\\s\\S]*?)${escapeRegex(endToolTag)}`, + 'gs', + ) + + let lastIndex = 0 + + for (const match of text.matchAll(toolExtractionPattern)) { + // Add any text before this tool call + if (match.index !== undefined && match.index > lastIndex) { + const textBefore = text.slice(lastIndex, match.index).trim() + if (textBefore) { + segments.push({ type: 'text', text: textBefore }) + } + } + + const jsonContent = match[1].trim() + + try { + const parsed = JSON.parse(jsonContent) + const toolName = parsed[toolNameParam] + + if (typeof toolName === 'string') { + // Remove the tool name param from the input + const input = { ...parsed } + delete input[toolNameParam] + + // Also remove cb_easp if present + delete input['cb_easp'] + + segments.push({ + type: 'tool_call', + toolName, + input, + }) + } + } catch { + // Skip malformed JSON - don't add segment + } + + // Update lastIndex to after this match + if (match.index !== undefined) { + lastIndex = match.index + match[0].length + } + } + + // Add any remaining text after the last tool call + if (lastIndex < text.length) { + const textAfter = text.slice(lastIndex).trim() + if (textAfter) { + segments.push({ type: 'text', text: textAfter }) + } + } + + return segments +} + +/** + * Parses tool calls from text in the XML format. + * This is a convenience function that returns only tool calls (no text segments). + * + * @param text - The text containing tool calls in XML format + * @returns Array of parsed tool calls with toolName and input + */ +export function parseToolCallsFromText( + text: string, +): Omit[] { + return parseTextWithToolCalls(text) + .filter((segment): segment is ParsedToolCallFromText => segment.type === 'tool_call') + .map(({ toolName, input }) => ({ toolName, input })) +} + +function escapeRegex(string: string): string { + return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') +} From 3b59c8797394ccf0f5e0963d65c05a294e7ae7ce Mon Sep 17 00:00:00 2001 From: Charles Lien Date: Wed, 26 Nov 2025 17:32:04 -0800 Subject: [PATCH 30/30] require json output for tool results --- .agents/types/util-types.ts | 19 +++-- .../initial-agents-dir/types/util-types.ts | 19 +++-- common/src/types/messages/codebuff-message.ts | 2 +- common/src/types/messages/content-part.ts | 23 ++++--- common/src/util/messages.ts | 69 ++++++------------- 5 files changed, 53 insertions(+), 79 deletions(-) diff --git a/.agents/types/util-types.ts b/.agents/types/util-types.ts index c6bc95d73b..79b4e81f52 100644 --- a/.agents/types/util-types.ts +++ b/.agents/types/util-types.ts @@ -74,16 +74,15 @@ export type ToolCallPart = { providerExecuted?: boolean } -export type ToolResultOutput = - | { - type: 'json' - value: JSONValue - } - | { - type: 'media' - data: string - mediaType: string - } +export type MediaToolResultOutputSchema = { + data: string + mediaType: string +} + +export type ToolResultOutput = { + value: JSONValue + media?: MediaToolResultOutputSchema[] +} // ===== Message Types ===== type AuxiliaryData = { diff --git a/common/src/templates/initial-agents-dir/types/util-types.ts b/common/src/templates/initial-agents-dir/types/util-types.ts index c6bc95d73b..79b4e81f52 100644 --- a/common/src/templates/initial-agents-dir/types/util-types.ts +++ b/common/src/templates/initial-agents-dir/types/util-types.ts @@ -74,16 +74,15 @@ export type ToolCallPart = { providerExecuted?: boolean } -export type ToolResultOutput = - | { - type: 'json' - value: JSONValue - } - | { - type: 'media' - data: string - mediaType: string - } +export type MediaToolResultOutputSchema = { + data: string + mediaType: string +} + +export type ToolResultOutput = { + value: JSONValue + media?: MediaToolResultOutputSchema[] +} // ===== Message Types ===== type AuxiliaryData = { diff --git a/common/src/types/messages/codebuff-message.ts b/common/src/types/messages/codebuff-message.ts index 0ce1708b31..d222e29469 100644 --- a/common/src/types/messages/codebuff-message.ts +++ b/common/src/types/messages/codebuff-message.ts @@ -41,7 +41,7 @@ export type ToolMessage = { role: 'tool' toolCallId: string toolName: string - content: ToolResultOutput[] + content: ToolResultOutput } & AuxiliaryMessageData export type Message = diff --git a/common/src/types/messages/content-part.ts b/common/src/types/messages/content-part.ts index c4692e42aa..e78a2cada1 100644 --- a/common/src/types/messages/content-part.ts +++ b/common/src/types/messages/content-part.ts @@ -45,15 +45,16 @@ export const toolCallPartSchema = z.object({ }) export type ToolCallPart = z.infer -export const toolResultOutputSchema = z.discriminatedUnion('type', [ - z.object({ - type: z.literal('json'), - value: jsonValueSchema, - }), - z.object({ - type: z.literal('media'), - data: z.string(), - mediaType: z.string(), - }), -]) +export const mediaToolResultOutputSchema = z.object({ + data: z.string(), + mediaType: z.string(), +}) +export type MediaToolResultOutputSchema = z.infer< + typeof mediaToolResultOutputSchema +> + +export const toolResultOutputSchema = z.object({ + value: jsonValueSchema, + media: mediaToolResultOutputSchema.array().optional(), +}) export type ToolResultOutput = z.infer diff --git a/common/src/util/messages.ts b/common/src/util/messages.ts index 478a2d4d22..54f89fb92a 100644 --- a/common/src/util/messages.ts +++ b/common/src/util/messages.ts @@ -1,6 +1,5 @@ import { cloneDeep, has, isEqual } from 'lodash' -import type { JSONValue } from '../types/json' import type { AssistantMessage, AuxiliaryMessageData, @@ -9,7 +8,6 @@ import type { ToolMessage, UserMessage, } from '../types/messages/codebuff-message' -import type { ToolResultOutput } from '../types/messages/content-part' import type { ProviderMetadata } from '../types/messages/provider-metadata' import type { AssistantModelMessage, @@ -119,25 +117,31 @@ function assistantToCodebuffMessage( function convertToolResultMessage( message: ToolMessage, ): ModelMessageWithAuxiliaryData[] { - return message.content.map((c) => { - if (c.type === 'json') { - return cloneDeep({ - ...message, - role: 'tool', - content: [{ ...message, output: c, type: 'tool-result' }], - }) - } - if (c.type === 'media') { - return cloneDeep({ + const messages: ModelMessageWithAuxiliaryData[] = [ + cloneDeep({ + ...message, + role: 'tool', + content: [ + { + ...message, + output: { type: 'json', value: message.content.value }, + type: 'tool-result', + }, + ], + }), + ] + + for (const c of message.content.media ?? []) { + messages.push( + cloneDeep({ ...message, role: 'user', content: [{ type: 'file', data: c.data, mediaType: c.mediaType }], - }) - } - c satisfies never - const cAny = c as any - throw new Error(`Invalid tool output type: ${cAny.type}`) - }) + }), + ) + } + + return messages } function convertToolMessage(message: Message): ModelMessageWithAuxiliaryData[] { @@ -411,32 +415,3 @@ export function assistantMessage( content: assistantContent(params), } } - -export function jsonToolResult( - value: T, -): [ - Extract & { - value: T - }, -] { - return [ - { - type: 'json', - value, - }, - ] -} - -export function mediaToolResult(params: { - data: string - mediaType: string -}): [Extract] { - const { data, mediaType } = params - return [ - { - type: 'media', - data, - mediaType, - }, - ] -}