Skip to content

Commit 85b64ba

Browse files
brandonkachenbrandonkachen
committed
Fix text block aggregation and extend streaming coverage
1 parent e70cf19 commit 85b64ba

File tree

2 files changed

+389
-13
lines changed

2 files changed

+389
-13
lines changed

‎sdk/src/__tests__/run-text-emission.test.ts‎

Lines changed: 372 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,377 @@ describe('run() text emission', () => {
189189
})
190190
})
191191

192+
test('emits aggregated text blocks while streaming chunk deltas', async () => {
193+
const events: PrintModeEvent[] = []
194+
const streamChunks: string[] = []
195+
const runPromise = run({
196+
...baseRunOptions,
197+
handleEvent: (event) => {
198+
events.push(event)
199+
},
200+
handleStreamChunk: (chunk) => {
201+
streamChunks.push(chunk)
202+
},
203+
})
204+
205+
const handler = await waitForHandler()
206+
207+
await handler.options.onResponseChunk(
208+
responseChunk(handler, {
209+
type: 'text',
210+
text: 'Hello ',
211+
}),
212+
)
213+
await handler.options.onResponseChunk(
214+
responseChunk(handler, {
215+
type: 'text',
216+
text: 'Hello world',
217+
}),
218+
)
219+
await handler.options.onResponseChunk(
220+
responseChunk(handler, {
221+
type: 'finish',
222+
totalCost: 0,
223+
}),
224+
)
225+
226+
await resolvePrompt(handler)
227+
await runPromise
228+
229+
expect(streamChunks).toEqual(['Hello ', 'world'])
230+
231+
const textEvents = events.filter(
232+
(event): event is PrintModeEvent & { type: 'text' } =>
233+
event.type === 'text',
234+
)
235+
expect(textEvents).toEqual([
236+
expect.objectContaining({ type: 'text', text: 'Hello world' }),
237+
])
238+
})
239+
240+
test('emits combined text when raw string and structured chunks interleave', async () => {
241+
const events: PrintModeEvent[] = []
242+
const streamChunks: string[] = []
243+
const runPromise = run({
244+
...baseRunOptions,
245+
handleEvent: (event) => {
246+
events.push(event)
247+
},
248+
handleStreamChunk: (chunk) => {
249+
streamChunks.push(chunk)
250+
},
251+
})
252+
253+
const handler = await waitForHandler()
254+
255+
await handler.options.onResponseChunk(
256+
responseChunk(handler, 'Root string '),
257+
)
258+
await handler.options.onResponseChunk(
259+
responseChunk(handler, {
260+
type: 'text',
261+
text: 'section complete',
262+
}),
263+
)
264+
await handler.options.onResponseChunk(
265+
responseChunk(handler, {
266+
type: 'finish',
267+
totalCost: 0,
268+
}),
269+
)
270+
271+
await resolvePrompt(handler)
272+
await runPromise
273+
274+
expect(streamChunks).toEqual(['Root string ', 'section complete'])
275+
276+
const textEvents = events.filter(
277+
(event): event is PrintModeEvent & { type: 'text' } =>
278+
event.type === 'text',
279+
)
280+
281+
expect(textEvents).toEqual([
282+
expect.objectContaining({
283+
type: 'text',
284+
text: 'Root string section complete',
285+
}),
286+
])
287+
})
288+
289+
test('keeps earlier text when new fragments are shorter than accumulated text', async () => {
290+
const events: PrintModeEvent[] = []
291+
const runPromise = run({
292+
...baseRunOptions,
293+
handleEvent: async (event) => {
294+
events.push(event)
295+
},
296+
})
297+
298+
const handler = await waitForHandler()
299+
300+
await handler.options.onResponseChunk(
301+
responseChunk(handler, {
302+
type: 'text',
303+
text: 'Intro line ',
304+
}),
305+
)
306+
await handler.options.onResponseChunk(
307+
responseChunk(handler, {
308+
type: 'text',
309+
text: 'continues',
310+
}),
311+
)
312+
await handler.options.onResponseChunk(
313+
responseChunk(handler, {
314+
type: 'text',
315+
text: ' and ends.<codebuff_tool_call>',
316+
}),
317+
)
318+
await handler.options.onResponseChunk(
319+
responseChunk(handler, {
320+
type: 'tool_call',
321+
toolCallId: 'tool-aggregate',
322+
toolName: 'example_tool',
323+
input: {},
324+
}),
325+
)
326+
await handler.options.onResponseChunk(
327+
responseChunk(handler, {
328+
type: 'finish',
329+
totalCost: 0,
330+
}),
331+
)
332+
333+
await resolvePrompt(handler)
334+
await runPromise
335+
336+
const textEvents = events.filter(
337+
(event): event is PrintModeEvent & { type: 'text' } =>
338+
event.type === 'text',
339+
)
340+
341+
expect(textEvents).toEqual([
342+
expect.objectContaining({
343+
type: 'text',
344+
text: 'Intro line continues and ends.',
345+
}),
346+
])
347+
})
348+
349+
test('flushes subagent text on subagent finish', async () => {
350+
const events: PrintModeEvent[] = []
351+
const runPromise = run({
352+
...baseRunOptions,
353+
handleEvent: async (event) => {
354+
events.push(event)
355+
},
356+
})
357+
358+
const handler = await waitForHandler()
359+
360+
await handler.options.onResponseChunk(
361+
responseChunk(handler, {
362+
type: 'text',
363+
agentId: 'agent-sub',
364+
text: 'Subagent output block',
365+
}),
366+
)
367+
await handler.options.onResponseChunk(
368+
responseChunk(handler, {
369+
type: 'subagent_finish',
370+
agentId: 'agent-sub',
371+
agentType: 'helper',
372+
displayName: 'Helper',
373+
onlyChild: false,
374+
}),
375+
)
376+
await handler.options.onResponseChunk(
377+
responseChunk(handler, {
378+
type: 'finish',
379+
totalCost: 0,
380+
}),
381+
)
382+
383+
await resolvePrompt(handler)
384+
await runPromise
385+
386+
const textEvents = events.filter(
387+
(event): event is PrintModeEvent & { type: 'text' } =>
388+
event.type === 'text',
389+
)
390+
391+
expect(textEvents).toContainEqual(
392+
expect.objectContaining({
393+
type: 'text',
394+
agentId: 'agent-sub',
395+
text: 'Subagent output block',
396+
}),
397+
)
398+
})
399+
400+
test('handles tool XML that spans multiple text chunks', async () => {
401+
const events: PrintModeEvent[] = []
402+
const runPromise = run({
403+
...baseRunOptions,
404+
handleEvent: (event) => {
405+
events.push(event)
406+
},
407+
})
408+
409+
const handler = await waitForHandler()
410+
411+
await handler.options.onResponseChunk(
412+
responseChunk(handler, {
413+
type: 'text',
414+
text: 'Before <codebuff_tool_call>{"x":1}',
415+
}),
416+
)
417+
await handler.options.onResponseChunk(
418+
responseChunk(handler, {
419+
type: 'text',
420+
text: '</codebuff_tool_call> after',
421+
}),
422+
)
423+
await handler.options.onResponseChunk(
424+
responseChunk(handler, {
425+
type: 'finish',
426+
totalCost: 0,
427+
}),
428+
)
429+
430+
await resolvePrompt(handler)
431+
await runPromise
432+
433+
const textEvents = events.filter(
434+
(event): event is PrintModeEvent & { type: 'text' } =>
435+
event.type === 'text',
436+
)
437+
438+
expect(textEvents).toEqual([
439+
expect.objectContaining({ text: 'Before' }),
440+
expect.objectContaining({ text: 'after' }),
441+
])
442+
})
443+
444+
test('trims surrounding newlines before emitting text', async () => {
445+
const events: PrintModeEvent[] = []
446+
const runPromise = run({
447+
...baseRunOptions,
448+
handleEvent: (event) => {
449+
events.push(event)
450+
},
451+
})
452+
453+
const handler = await waitForHandler()
454+
455+
await handler.options.onResponseChunk(
456+
responseChunk(handler, {
457+
type: 'text',
458+
text: '\nLine 1\nLine 2\n\n',
459+
}),
460+
)
461+
await handler.options.onResponseChunk(
462+
responseChunk(handler, {
463+
type: 'finish',
464+
totalCost: 0,
465+
}),
466+
)
467+
468+
await resolvePrompt(handler)
469+
await runPromise
470+
471+
const textEvents = events.filter(
472+
(event): event is PrintModeEvent & { type: 'text' } =>
473+
event.type === 'text',
474+
)
475+
476+
expect(textEvents).toEqual([
477+
expect.objectContaining({
478+
text: 'Line 1\nLine 2',
479+
}),
480+
])
481+
})
482+
483+
test('skips whitespace-only sections', async () => {
484+
const events: PrintModeEvent[] = []
485+
const runPromise = run({
486+
...baseRunOptions,
487+
handleEvent: (event) => {
488+
events.push(event)
489+
},
490+
})
491+
492+
const handler = await waitForHandler()
493+
494+
await handler.options.onResponseChunk(
495+
responseChunk(handler, {
496+
type: 'text',
497+
text: '\n\n',
498+
}),
499+
)
500+
await handler.options.onResponseChunk(
501+
responseChunk(handler, {
502+
type: 'finish',
503+
totalCost: 0,
504+
}),
505+
)
506+
507+
await resolvePrompt(handler)
508+
await runPromise
509+
510+
const textEvents = events.filter(
511+
(event): event is PrintModeEvent & { type: 'text' } =>
512+
event.type === 'text',
513+
)
514+
515+
expect(textEvents).toEqual([])
516+
})
517+
518+
test('flushes buffered text when finish clears residual tool XML state', async () => {
519+
const events: PrintModeEvent[] = []
520+
const streamChunks: string[] = []
521+
const runPromise = run({
522+
...baseRunOptions,
523+
handleEvent: (event) => {
524+
events.push(event)
525+
},
526+
handleStreamChunk: (chunk) => {
527+
streamChunks.push(chunk)
528+
},
529+
})
530+
531+
const handler = await waitForHandler()
532+
533+
await handler.options.onResponseChunk(
534+
responseChunk(handler, 'Streaming start '),
535+
)
536+
await handler.options.onResponseChunk(
537+
responseChunk(handler, 'continues before <codebuff_tool_call'),
538+
)
539+
await handler.options.onResponseChunk(
540+
responseChunk(handler, {
541+
type: 'finish',
542+
totalCost: 0,
543+
}),
544+
)
545+
546+
await resolvePrompt(handler)
547+
await runPromise
548+
549+
const textEvents = events.filter(
550+
(event): event is PrintModeEvent & { type: 'text' } =>
551+
event.type === 'text',
552+
)
553+
554+
expect(streamChunks).toEqual(['Streaming start ', 'continu'])
555+
556+
expect(textEvents).toEqual([
557+
expect.objectContaining({
558+
text: 'Streaming start continu',
559+
}),
560+
])
561+
})
562+
192563
test('splits root sections around tool events without duplication', async () => {
193564
const events: PrintModeEvent[] = []
194565
const runPromise = run({
@@ -335,9 +706,6 @@ describe('run() text emission', () => {
335706
event.type === 'text',
336707
)
337708

338-
expect(textEvents.map((event) => event.text)).toEqual([
339-
'Before ',
340-
' after',
341-
])
709+
expect(textEvents.map((event) => event.text)).toEqual(['Before', 'after'])
342710
})
343711
})

0 commit comments

Comments
 (0)