Skip to content

Commit c3aea29

Browse files
charlesliencharleslien
committed
Return messageId instead of resolving in stream
1 parent 173ec67 commit c3aea29

15 files changed

+103
-107
lines changed

‎backend/src/__tests__/agent-run.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,7 @@ describe('Agent Run Database Functions', () => {
327327
agentRunId: 'run-123',
328328
stepNumber: 2,
329329
startTime,
330+
messageId: null,
330331
})
331332

332333
expect(mockValues).toHaveBeenCalledWith({
@@ -357,6 +358,7 @@ describe('Agent Run Database Functions', () => {
357358
status: 'skipped',
358359
errorMessage: 'Step failed validation',
359360
startTime,
361+
messageId: null,
360362
})
361363

362364
expect(mockValues).toHaveBeenCalledWith(
@@ -380,6 +382,7 @@ describe('Agent Run Database Functions', () => {
380382
stepNumber: 4,
381383
status: 'running',
382384
startTime,
385+
messageId: null,
383386
})
384387

385388
expect(mockValues).toHaveBeenCalledWith(
@@ -402,6 +405,7 @@ describe('Agent Run Database Functions', () => {
402405
stepNumber: 5,
403406
credits: 0, // Zero credits
404407
startTime,
408+
messageId: null,
405409
})
406410

407411
expect(mockValues).toHaveBeenCalledWith(
@@ -425,6 +429,7 @@ describe('Agent Run Database Functions', () => {
425429
agentRunId: 'run-123',
426430
stepNumber: 6,
427431
startTime,
432+
messageId: null,
428433
}),
429434
).rejects.toThrow('Insert failed')
430435

@@ -451,6 +456,7 @@ describe('Agent Run Database Functions', () => {
451456
stepNumber: 1,
452457
credits: 123.456789, // High precision number
453458
startTime: new Date(),
459+
messageId: null,
454460
})
455461

456462
expect(mockValues).toHaveBeenCalledWith(
@@ -472,6 +478,7 @@ describe('Agent Run Database Functions', () => {
472478
agentRunId: 'run-123',
473479
stepNumber: 1,
474480
startTime: specificStartTime,
481+
messageId: null,
475482
})
476483

477484
expect(mockValues).toHaveBeenCalledWith(

‎backend/src/__tests__/cost-aggregation.integration.test.ts‎

Lines changed: 27 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -185,14 +185,18 @@ describe('Cost Aggregation Integration Tests', () => {
185185
// Simulate different responses based on call
186186
if (callCount === 1) {
187187
// Main agent spawns a subagent
188-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write a simple hello world file"}]}\n</codebuff_tool_call>' }
188+
yield {
189+
type: 'text' as const,
190+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write a simple hello world file"}]}\n</codebuff_tool_call>',
191+
}
189192
} else {
190193
// Subagent writes a file
191-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "hello.txt", "instructions": "Create hello world file", "content": "Hello, World!"}\n</codebuff_tool_call>' }
192-
}
193-
if (options.resolveMessageId) {
194-
options.resolveMessageId('mock-message-id')
194+
yield {
195+
type: 'text' as const,
196+
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "hello.txt", "instructions": "Create hello world file", "content": "Hello, World!"}\n</codebuff_tool_call>',
197+
}
195198
}
199+
return 'mock-message-id'
196200
},
197201
)
198202

@@ -342,18 +346,25 @@ describe('Cost Aggregation Integration Tests', () => {
342346

343347
if (callCount === 1) {
344348
// Main agent spawns first-level subagent
345-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Create files"}]}\n</codebuff_tool_call>' }
349+
yield {
350+
type: 'text' as const,
351+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Create files"}]}\n</codebuff_tool_call>',
352+
}
346353
} else if (callCount === 2) {
347354
// First-level subagent spawns second-level subagent
348-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write specific file"}]}\n</codebuff_tool_call>' }
355+
yield {
356+
type: 'text' as const,
357+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "Write specific file"}]}\n</codebuff_tool_call>',
358+
}
349359
} else {
350360
// Second-level subagent does actual work
351-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "nested.txt", "instructions": "Create nested file", "content": "Nested content"}\n</codebuff_tool_call>' }
361+
yield {
362+
type: 'text' as const,
363+
text: '<codebuff_tool_call>\n{"cb_tool_name": "write_file", "path": "nested.txt", "instructions": "Create nested file", "content": "Nested content"}\n</codebuff_tool_call>',
364+
}
352365
}
353366

354-
if (options.resolveMessageId) {
355-
options.resolveMessageId('mock-message-id')
356-
}
367+
return 'mock-message-id'
357368
},
358369
)
359370

@@ -401,16 +412,17 @@ describe('Cost Aggregation Integration Tests', () => {
401412

402413
if (callCount === 1) {
403414
// Main agent spawns subagent
404-
yield { type: 'text' as const, text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "This will fail"}]}\n</codebuff_tool_call>' }
415+
yield {
416+
type: 'text' as const,
417+
text: '<codebuff_tool_call>\n{"cb_tool_name": "spawn_agents", "agents": [{"agent_type": "editor", "prompt": "This will fail"}]}\n</codebuff_tool_call>',
418+
}
405419
} else {
406420
// Subagent fails after incurring cost
407421
yield { type: 'text' as const, text: 'Some response' }
408422
throw new Error('Subagent execution failed')
409423
}
410424

411-
if (options.resolveMessageId) {
412-
options.resolveMessageId('mock-message-id')
413-
}
425+
return 'mock-message-id'
414426
},
415427
)
416428

‎backend/src/__tests__/loop-agent-steps.test.ts‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -103,14 +103,13 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
103103
})),
104104
} as any)
105105

106-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
107-
resolveMessageId,
108-
}) {
106+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
109107
llmCallCount++
110-
yield { type: 'text' as const, text: `LLM response\n\n${getToolCallString('end_turn', {})}` }
111-
if (resolveMessageId) {
112-
resolveMessageId('mock-message-id')
108+
yield {
109+
type: 'text' as const,
110+
text: `LLM response\n\n${getToolCallString('end_turn', {})}`,
113111
}
112+
return 'mock-message-id'
114113
})
115114

116115
// Mock analytics

‎backend/src/__tests__/main-prompt.test.ts‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,9 @@ import type { ProjectFileContext } from '@codebuff/common/util/file'
3636
import type { WebSocket } from 'ws'
3737

3838
const mockAgentStream = (streamOutput: string) => {
39-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
40-
resolveMessageId,
41-
}) {
39+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
4240
yield { type: 'text' as const, text: streamOutput }
43-
if (resolveMessageId) {
44-
resolveMessageId('mock-message-id')
45-
}
41+
return 'mock-message-id'
4642
})
4743
}
4844

‎backend/src/__tests__/read-docs-tool.test.ts‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,18 +33,14 @@ import * as websocketAction from '../websockets/websocket-action'
3333
import type { WebSocket } from 'ws'
3434

3535
function mockAgentStream(content: string | string[]) {
36-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
37-
resolveMessageId,
38-
}) {
36+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
3937
if (typeof content === 'string') {
4038
content = [content]
4139
}
4240
for (const chunk of content) {
4341
yield { type: 'text' as const, text: chunk }
4442
}
45-
if (resolveMessageId) {
46-
resolveMessageId('mock-message-id')
47-
}
43+
return 'mock-message-id'
4844
})
4945
}
5046

‎backend/src/__tests__/run-agent-step-tools.test.ts‎

Lines changed: 18 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -176,13 +176,9 @@ describe('runAgentStep - set_output tool', () => {
176176
'\n\n' +
177177
getToolCallString('end_turn', {})
178178

179-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
180-
resolveMessageId,
181-
}) {
179+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
182180
yield { type: 'text' as const, text: mockResponse }
183-
if (resolveMessageId) {
184-
resolveMessageId('mock-message-id')
185-
}
181+
return 'mock-message-id'
186182
})
187183

188184
const sessionState = getInitialSessionState(mockFileContext)
@@ -222,13 +218,9 @@ describe('runAgentStep - set_output tool', () => {
222218
findings: ['Bug in auth.ts', 'Missing validation'],
223219
}) + getToolCallString('end_turn', {})
224220

225-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
226-
resolveMessageId,
227-
}) {
221+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
228222
yield { type: 'text' as const, text: mockResponse }
229-
if (resolveMessageId) {
230-
resolveMessageId('mock-message-id')
231-
}
223+
return 'mock-message-id'
232224
})
233225

234226
const sessionState = getInitialSessionState(mockFileContext)
@@ -269,13 +261,9 @@ describe('runAgentStep - set_output tool', () => {
269261
existingField: 'updated value',
270262
}) + getToolCallString('end_turn', {})
271263

272-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
273-
resolveMessageId,
274-
}) {
264+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
275265
yield { type: 'text' as const, text: mockResponse }
276-
if (resolveMessageId) {
277-
resolveMessageId('mock-message-id')
278-
}
266+
return 'mock-message-id'
279267
})
280268

281269
const sessionState = getInitialSessionState(mockFileContext)
@@ -316,13 +304,9 @@ describe('runAgentStep - set_output tool', () => {
316304
const mockResponse =
317305
getToolCallString('set_output', {}) + getToolCallString('end_turn', {})
318306

319-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
320-
resolveMessageId,
321-
}) {
307+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
322308
yield { type: 'text' as const, text: mockResponse }
323-
if (resolveMessageId) {
324-
resolveMessageId('mock-message-id')
325-
}
309+
return 'mock-message-id'
326310
})
327311

328312
const sessionState = getInitialSessionState(mockFileContext)
@@ -400,13 +384,9 @@ describe('runAgentStep - set_output tool', () => {
400384
)
401385

402386
// Mock the LLM stream to return a response that doesn't end the turn
403-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
404-
resolveMessageId,
405-
}) {
387+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
406388
yield { type: 'text' as const, text: 'Continuing with the analysis...' } // Non-empty response, no tool calls
407-
if (resolveMessageId) {
408-
resolveMessageId('mock-message-id')
409-
}
389+
return 'mock-message-id'
410390
})
411391

412392
const sessionState = getInitialSessionState(mockFileContext)
@@ -550,16 +530,15 @@ describe('runAgentStep - set_output tool', () => {
550530
}
551531

552532
// Mock the LLM stream to spawn the inline agent
553-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
554-
resolveMessageId,
555-
}) {
556-
yield { type: 'text' as const, text: getToolCallString('spawn_agent_inline', {
557-
agent_type: 'message-deleter-agent',
558-
prompt: 'Delete the last two assistant messages',
559-
}) }
560-
if (resolveMessageId) {
561-
resolveMessageId('mock-message-id')
533+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
534+
yield {
535+
type: 'text' as const,
536+
text: getToolCallString('spawn_agent_inline', {
537+
agent_type: 'message-deleter-agent',
538+
prompt: 'Delete the last two assistant messages',
539+
}),
562540
}
541+
return 'mock-message-id'
563542
})
564543

565544
const sessionState = getInitialSessionState(mockFileContext)

‎backend/src/__tests__/web-search-tool.test.ts‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,18 +36,14 @@ import * as websocketAction from '../websockets/websocket-action'
3636
import type { WebSocket } from 'ws'
3737

3838
function mockAgentStream(content: string | string[]) {
39-
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({
40-
resolveMessageId,
41-
}) {
39+
spyOn(aisdk, 'promptAiSdkStream').mockImplementation(async function* ({}) {
4240
if (typeof content === 'string') {
4341
content = [content]
4442
}
4543
for (const chunk of content) {
4644
yield { type: 'text' as const, text: chunk }
4745
}
48-
if (resolveMessageId) {
49-
resolveMessageId('mock-message-id')
50-
}
46+
return 'mock-message-id'
5147
})
5248
}
5349

‎backend/src/__tests__/xml-stream-parser.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ describe('processStreamWithTags', () => {
1010
for (const chunk of chunks) {
1111
yield { type: 'text' as const, text: chunk }
1212
}
13+
14+
return 'mock-message-id'
1315
}
1416

1517
it('should handle basic tool call parsing', async () => {

‎backend/src/agent-run.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ export async function addAgentStep({
106106
stepNumber: number
107107
credits?: number
108108
childRunIds?: string[]
109-
messageId?: string
109+
messageId: string | null
110110
status?: 'running' | 'completed' | 'skipped'
111111
errorMessage?: string
112112
startTime: Date

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

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,8 @@ export const promptAiSdkStream = async function* (
8282
maxRetries?: number
8383
onCostCalculated?: (credits: number) => Promise<void>
8484
includeCacheControl?: boolean
85-
resolveMessageId: (messageId: string | undefined) => unknown
8685
} & Omit<Parameters<typeof streamText>[0], 'model' | 'messages'>,
87-
): AsyncGenerator<StreamChunk> {
86+
): AsyncGenerator<StreamChunk, string | null> {
8887
if (
8988
!checkLiveUserInput(
9089
options.userId,
@@ -100,7 +99,7 @@ export const promptAiSdkStream = async function* (
10099
},
101100
'Skipping stream due to canceled user input',
102101
)
103-
return
102+
return null
104103
}
105104
const startTime = Date.now()
106105

@@ -151,10 +150,7 @@ export const promptAiSdkStream = async function* (
151150
message: errorMessage,
152151
}
153152

154-
// Important: we need to resolve the message id before returning.
155-
options.resolveMessageId(undefined)
156-
157-
return
153+
return null
158154
}
159155
if (chunk.type === 'reasoning-delta') {
160156
if (
@@ -200,11 +196,6 @@ export const promptAiSdkStream = async function* (
200196
}
201197
}
202198

203-
const messageId = (await response.response).id
204-
if (options.resolveMessageId) {
205-
options.resolveMessageId(messageId)
206-
}
207-
208199
const providerMetadata = (await response.providerMetadata) ?? {}
209200
const usage = await response.usage
210201
let inputTokens = usage.inputTokens || 0
@@ -236,6 +227,7 @@ export const promptAiSdkStream = async function* (
236227
}
237228
}
238229

230+
const messageId = (await response.response).id
239231
const creditsUsedPromise = saveMessage({
240232
messageId,
241233
userId: options.userId,
@@ -261,6 +253,8 @@ export const promptAiSdkStream = async function* (
261253
const creditsUsed = await creditsUsedPromise
262254
await options.onCostCalculated(creditsUsed)
263255
}
256+
257+
return messageId
264258
}
265259

266260
// TODO: figure out a nice way to unify stream & non-stream versions maybe?

0 commit comments

Comments
 (0)