Skip to content

Commit 5209f26

Browse files
brandonkachenbrandonkachen
committed
fix: subagent/tool call structure
1 parent c2ec6c0 commit 5209f26

File tree

10 files changed

+498
-105
lines changed

10 files changed

+498
-105
lines changed

‎backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,17 @@ import type {
3131
import type { LanguageModel } from 'ai'
3232
import type { z } from 'zod/v4'
3333

34+
export type StreamChunk =
35+
| {
36+
type: 'text'
37+
text: string
38+
agentId?: string
39+
}
40+
| {
41+
type: 'reasoning'
42+
text: string
43+
}
44+
| { type: 'error'; message: string }
3445
// TODO: We'll want to add all our models here!
3546
const modelToAiSDKModel = (model: Model): LanguageModel => {
3647
if (
@@ -57,6 +68,9 @@ export async function* promptAiSdkStream(
5768
params: ParamsOf<PromptAiSdkStreamFn>,
5869
): ReturnType<PromptAiSdkStreamFn> {
5970
const { logger } = params
71+
const agentChunkMetadata =
72+
params.agentId != null ? { agentId: params.agentId } : undefined
73+
6074
if (
6175
!checkLiveUserInput({ ...params, clientSessionId: params.clientSessionId })
6276
) {
@@ -92,6 +106,7 @@ export async function* promptAiSdkStream(
92106
yield {
93107
type: 'text',
94108
text: flushed,
109+
...(agentChunkMetadata ?? {}),
95110
}
96111
}
97112
}
@@ -144,6 +159,7 @@ export async function* promptAiSdkStream(
144159
yield {
145160
type: 'text',
146161
text: chunk.text,
162+
...(agentChunkMetadata ?? {}),
147163
}
148164
}
149165
continue
@@ -155,6 +171,7 @@ export async function* promptAiSdkStream(
155171
yield {
156172
type: 'text',
157173
text: stopSequenceResult.text,
174+
...(agentChunkMetadata ?? {}),
158175
}
159176
}
160177
}
@@ -165,6 +182,7 @@ export async function* promptAiSdkStream(
165182
yield {
166183
type: 'text',
167184
text: flushed,
185+
...(agentChunkMetadata ?? {}),
168186
}
169187
}
170188

‎backend/src/run-agent-step.ts‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ export const runAgentStep = async (
137137
spawnParams,
138138
system,
139139
logger,
140+
promptAiSdkStream,
140141
} = params
141142
let agentState = params.agentState
142143

@@ -259,8 +260,11 @@ export const runAgentStep = async (
259260
const { model } = agentTemplate
260261

261262
const { getStream } = getAgentStreamFromTemplate({
262-
...params,
263-
agentId: agentState.agentId,
263+
clientSessionId,
264+
fingerprintId,
265+
userInputId,
266+
userId,
267+
agentId: agentState.parentId ? agentState.agentId : undefined,
264268
template: agentTemplate,
265269
onCostCalculated: async (credits: number) => {
266270
try {
@@ -279,6 +283,8 @@ export const runAgentStep = async (
279283
)
280284
}
281285
},
286+
promptAiSdkStream,
287+
logger,
282288
includeCacheControl: supportsCacheControl(agentTemplate.model),
283289
})
284290

‎backend/src/run-programmatic-step.ts‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -268,12 +268,6 @@ export async function runProgrammaticStep(
268268
role: 'assistant' as const,
269269
content: toolCallString,
270270
})
271-
state.sendSubagentChunk({
272-
userInputId,
273-
agentId: state.agentState.agentId,
274-
agentType: state.agentState.agentType!,
275-
chunk: toolCallString,
276-
})
277271
}
278272

279273
// Execute the tool synchronously and get the result immediately

‎backend/src/tools/handlers/tool/spawn-agent-utils.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,8 @@ export async function executeSubagent(
334334

335335
onResponseChunk({
336336
type: 'subagent_start',
337-
agentId: agentTemplate.id,
337+
agentId: agentState.agentId,
338+
agentType: agentTemplate.id,
338339
displayName: agentTemplate.displayName,
339340
onlyChild: isOnlyChild,
340341
})
@@ -349,7 +350,8 @@ export async function executeSubagent(
349350

350351
onResponseChunk({
351352
type: 'subagent_finish',
352-
agentId: agentTemplate.id,
353+
agentId: agentState.agentId,
354+
agentType: agentTemplate.id,
353355
displayName: agentTemplate.displayName,
354356
onlyChild: isOnlyChild,
355357
})

‎backend/src/tools/stream-parser.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@ export async function processStreamWithTools(
274274
reasoning = false
275275
onResponseChunk(`"\n}${endToolTag}\n\n`)
276276
}
277-
onResponseChunk(chunk.text)
277+
onResponseChunk(chunk)
278278
fullResponseChunks.push(chunk.text)
279279
} else if (chunk.type === 'error') {
280280
onResponseChunk(chunk)

‎cli/src/components/message-block.tsx‎

Lines changed: 79 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -181,8 +181,8 @@ export const MessageBlock = ({
181181
: ''
182182

183183
const finishedPreview =
184-
!isStreaming && isCollapsed
185-
? lastLine.replace(/[#*_`~\[\]()]/g, '').trim()
184+
!isStreaming && isCollapsed && block.initialPrompt
185+
? block.initialPrompt.replace(/[#*_`~\[\]()]/g, '').trim()
186186
: ''
187187

188188
const agentCodeBlockWidth = Math.max(10, availableWidth - 12)
@@ -195,16 +195,88 @@ export const MessageBlock = ({
195195
codeBlockWidth: agentCodeBlockWidth,
196196
palette: agentPalette,
197197
}
198-
const displayContent = hasMarkdown(block.content)
199-
? renderMarkdown(block.content, agentMarkdownOptions)
200-
: block.content
198+
199+
const displayContent = block.content
200+
? (hasMarkdown(block.content)
201+
? renderMarkdown(block.content, agentMarkdownOptions)
202+
: block.content)
203+
: ''
204+
205+
const nestedToolBlocks = block.blocks && block.blocks.length > 0 && !isCollapsed
206+
? block.blocks.map((nestedBlock, nestedIdx) => {
207+
if (nestedBlock.type === 'tool') {
208+
const displayInfo = getToolDisplayInfo(nestedBlock.toolName)
209+
const isNestedCollapsed = collapsedAgents.has(nestedBlock.toolCallId)
210+
const isNestedStreaming = streamingAgents.has(nestedBlock.toolCallId)
211+
212+
const inputContent = `\`\`\`json\n${JSON.stringify(nestedBlock.input, null, 2)}\n\`\`\``
213+
const codeBlockLang =
214+
nestedBlock.toolName === 'run_terminal_command' ? '' : 'yaml'
215+
const resultContent = nestedBlock.output
216+
? `\n\n**Result:**\n\`\`\`${codeBlockLang}\n${nestedBlock.output}\n\`\`\``
217+
: ''
218+
const fullNestedContent = inputContent + resultContent
219+
220+
const nestedLines = fullNestedContent.split('\n').filter((line) => line.trim())
221+
const firstNestedLine = nestedLines[0] || ''
222+
const lastNestedLine = nestedLines[nestedLines.length - 1] || firstNestedLine
223+
224+
const nestedStreamingPreview = isNestedStreaming
225+
? firstNestedLine.replace(/[#*_`~\[\]()]/g, '').trim() + '...'
226+
: ''
227+
228+
let nestedFinishedPreview = ''
229+
if (!isNestedStreaming && isNestedCollapsed) {
230+
if (nestedBlock.toolName === 'run_terminal_command' && nestedBlock.output) {
231+
const outputLines = nestedBlock.output
232+
.split('\n')
233+
.filter((line) => line.trim())
234+
const lastThreeLines = outputLines.slice(-3)
235+
const hasMoreLines = outputLines.length > 3
236+
nestedFinishedPreview = hasMoreLines
237+
? '...\n' + lastThreeLines.join('\n')
238+
: lastThreeLines.join('\n')
239+
} else {
240+
nestedFinishedPreview = lastNestedLine
241+
.replace(/[#*_`~\[\]()]/g, '')
242+
.trim()
243+
}
244+
}
245+
246+
const nestedDisplayContent = hasMarkdown(fullNestedContent)
247+
? renderMarkdown(fullNestedContent, agentMarkdownOptions)
248+
: fullNestedContent
249+
250+
const nextNestedBlock = block.blocks![nestedIdx + 1]
251+
const isLastNestedBranch = !nextNestedBlock
252+
const nestedBranchChar = isLastNestedBranch ? ' └─ ' : ' ├─ '
253+
254+
return (
255+
<box key={`${messageId}-agent-${block.agentId}-tool-${nestedBlock.toolCallId}`}>
256+
<BranchItem
257+
name={displayInfo.name}
258+
content={nestedDisplayContent}
259+
isCollapsed={isNestedCollapsed}
260+
isStreaming={isNestedStreaming}
261+
branchChar={nestedBranchChar}
262+
streamingPreview={nestedStreamingPreview}
263+
finishedPreview={nestedFinishedPreview}
264+
theme={theme}
265+
onToggle={() => onToggleCollapsed(nestedBlock.toolCallId)}
266+
/>
267+
</box>
268+
)
269+
}
270+
return null
271+
})
272+
: null
201273

202274
const nextBlock = blocks[idx + 1]
203275
const isLastBranch = !nextBlock || nextBlock.type === 'text'
204276
const branchChar = isLastBranch ? '└─ ' : '├─ '
205277

206278
return (
207-
<box key={`${messageId}-agent-${block.agentId}`}>
279+
<box key={`${messageId}-agent-${block.agentId}`} style={{ flexDirection: 'column', gap: 0 }}>
208280
<BranchItem
209281
name={block.agentName}
210282
content={displayContent}
@@ -216,6 +288,7 @@ export const MessageBlock = ({
216288
theme={theme}
217289
onToggle={() => onToggleCollapsed(block.agentId)}
218290
/>
291+
{nestedToolBlocks}
219292
</box>
220293
)
221294
}

0 commit comments

Comments
 (0)