Skip to content

Commit d931aa3

Browse files
brandonkachenbrandonkachen
committed
Handle string response chunks in sdk
1 parent 033722f commit d931aa3

File tree

2 files changed

+30
-163
lines changed

2 files changed

+30
-163
lines changed

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

Lines changed: 18 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -204,17 +204,9 @@ describe('run() text emission', () => {
204204

205205
const handler = await waitForHandler()
206206

207+
await handler.options.onResponseChunk(responseChunk(handler, 'Hello '))
207208
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-
}),
209+
responseChunk(handler, 'Hello world'),
218210
)
219211
await handler.options.onResponseChunk(
220212
responseChunk(handler, {
@@ -237,7 +229,7 @@ describe('run() text emission', () => {
237229
])
238230
})
239231

240-
test('emits combined text when raw string and structured chunks interleave', async () => {
232+
test('emits combined text when consecutive string chunks overlap', async () => {
241233
const events: PrintModeEvent[] = []
242234
const streamChunks: string[] = []
243235
const runPromise = run({
@@ -256,10 +248,7 @@ describe('run() text emission', () => {
256248
responseChunk(handler, 'Root string '),
257249
)
258250
await handler.options.onResponseChunk(
259-
responseChunk(handler, {
260-
type: 'text',
261-
text: 'section complete',
262-
}),
251+
responseChunk(handler, 'Root string section complete'),
263252
)
264253
await handler.options.onResponseChunk(
265254
responseChunk(handler, {
@@ -298,22 +287,13 @@ describe('run() text emission', () => {
298287
const handler = await waitForHandler()
299288

300289
await handler.options.onResponseChunk(
301-
responseChunk(handler, {
302-
type: 'text',
303-
text: 'Intro line ',
304-
}),
290+
responseChunk(handler, 'Intro line '),
305291
)
306292
await handler.options.onResponseChunk(
307-
responseChunk(handler, {
308-
type: 'text',
309-
text: 'continues',
310-
}),
293+
responseChunk(handler, 'continues'),
311294
)
312295
await handler.options.onResponseChunk(
313-
responseChunk(handler, {
314-
type: 'text',
315-
text: ' and ends.<codebuff_tool_call>',
316-
}),
296+
responseChunk(handler, ' and ends.<codebuff_tool_call>'),
317297
)
318298
await handler.options.onResponseChunk(
319299
responseChunk(handler, {
@@ -409,16 +389,10 @@ describe('run() text emission', () => {
409389
const handler = await waitForHandler()
410390

411391
await handler.options.onResponseChunk(
412-
responseChunk(handler, {
413-
type: 'text',
414-
text: 'Before <codebuff_tool_call>{"x":1}',
415-
}),
392+
responseChunk(handler, 'Before <codebuff_tool_call>{"x":1}'),
416393
)
417394
await handler.options.onResponseChunk(
418-
responseChunk(handler, {
419-
type: 'text',
420-
text: '</codebuff_tool_call> after',
421-
}),
395+
responseChunk(handler, '</codebuff_tool_call> after'),
422396
)
423397
await handler.options.onResponseChunk(
424398
responseChunk(handler, {
@@ -436,8 +410,7 @@ describe('run() text emission', () => {
436410
)
437411

438412
expect(textEvents).toEqual([
439-
expect.objectContaining({ text: 'Before' }),
440-
expect.objectContaining({ text: 'after' }),
413+
expect.objectContaining({ text: 'Before after' }),
441414
])
442415
})
443416

@@ -453,10 +426,7 @@ describe('run() text emission', () => {
453426
const handler = await waitForHandler()
454427

455428
await handler.options.onResponseChunk(
456-
responseChunk(handler, {
457-
type: 'text',
458-
text: '\nLine 1\nLine 2\n\n',
459-
}),
429+
responseChunk(handler, '\nLine 1\nLine 2\n\n'),
460430
)
461431
await handler.options.onResponseChunk(
462432
responseChunk(handler, {
@@ -491,12 +461,7 @@ describe('run() text emission', () => {
491461

492462
const handler = await waitForHandler()
493463

494-
await handler.options.onResponseChunk(
495-
responseChunk(handler, {
496-
type: 'text',
497-
text: '\n\n',
498-
}),
499-
)
464+
await handler.options.onResponseChunk(responseChunk(handler, '\n\n'))
500465
await handler.options.onResponseChunk(
501466
responseChunk(handler, {
502467
type: 'finish',
@@ -572,10 +537,7 @@ describe('run() text emission', () => {
572537
const handler = await waitForHandler()
573538

574539
await handler.options.onResponseChunk(
575-
responseChunk(handler, {
576-
type: 'text',
577-
text: 'First section',
578-
}),
540+
responseChunk(handler, 'First section'),
579541
)
580542
await handler.options.onResponseChunk(
581543
responseChunk(handler, {
@@ -586,10 +548,7 @@ describe('run() text emission', () => {
586548
}),
587549
)
588550
await handler.options.onResponseChunk(
589-
responseChunk(handler, {
590-
type: 'text',
591-
text: 'Second section',
592-
}),
551+
responseChunk(handler, 'Second section'),
593552
)
594553
await handler.options.onResponseChunk(
595554
responseChunk(handler, {
@@ -674,22 +633,13 @@ describe('run() text emission', () => {
674633
const handler = await waitForHandler()
675634

676635
await handler.options.onResponseChunk(
677-
responseChunk(handler, {
678-
type: 'text',
679-
text: 'Before <codebuff_tool_call>{"a":1}',
680-
}),
636+
responseChunk(handler, 'Before <codebuff_tool_call>{"a":1}'),
681637
)
682638
await handler.options.onResponseChunk(
683-
responseChunk(handler, {
684-
type: 'text',
685-
text: '</codebuff_tool_call>',
686-
}),
639+
responseChunk(handler, '</codebuff_tool_call>'),
687640
)
688641
await handler.options.onResponseChunk(
689-
responseChunk(handler, {
690-
type: 'text',
691-
text: ' after',
692-
}),
642+
responseChunk(handler, ' after'),
693643
)
694644
await handler.options.onResponseChunk(
695645
responseChunk(handler, {
@@ -706,6 +656,6 @@ describe('run() text emission', () => {
706656
event.type === 'text',
707657
)
708658

709-
expect(textEvents.map((event) => event.text)).toEqual(['Before', 'after'])
659+
expect(textEvents.map((event) => event.text)).toEqual(['Before after'])
710660
})
711661
})

‎sdk/src/run.ts‎

Lines changed: 12 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,6 @@ import type { SessionState } from '../../common/src/types/session-state'
4444
import type { Source } from '../../common/src/types/source'
4545
import type { CodebuffFileSystem } from '@codebuff/common/types/filesystem'
4646

47-
type TextPrintEvent = Extract<PrintModeEvent, { type: 'text' }>
48-
4947
export type CodebuffClientOptions = {
5048
apiKey?: string
5149

@@ -140,23 +138,12 @@ export async function run({
140138
const ROOT_AGENT_KEY = '__root__'
141139

142140
const streamFilterState = createToolXmlFilterState()
143-
const textFilterStates = new Map<string, ToolXmlFilterState>()
144141
const textAccumulator = new Map<string, string>()
145142
const lastStreamedTextByAgent = new Map<string, string>()
146-
const lastTextEventByAgent = new Map<string, TextPrintEvent>()
147143
const sectionStartIndexByAgent = new Map<string, number>()
148144

149145
const subagentFilterStates = new Map<string, ToolXmlFilterState>()
150146

151-
const getTextFilterState = (agentKey: string): ToolXmlFilterState => {
152-
let state = textFilterStates.get(agentKey)
153-
if (!state) {
154-
state = createToolXmlFilterState()
155-
textFilterStates.set(agentKey, state)
156-
}
157-
return state
158-
}
159-
160147
const getSubagentFilterState = (agentId: string): ToolXmlFilterState => {
161148
let state = subagentFilterStates.get(agentId)
162149
if (!state) {
@@ -263,30 +250,15 @@ export async function run({
263250
}
264251

265252
const eventAgentId = resolveAgentId(agentKey, agentIdHint)
266-
const lastChunk = lastTextEventByAgent.get(agentKey)
267-
268-
let eventPayload: PrintModeEvent
269-
if (lastChunk) {
270-
eventPayload = { ...lastChunk, text: trimmedText }
271-
272-
if (
273-
eventAgentId &&
274-
(!('agentId' in eventPayload) ||
275-
(eventPayload as { agentId?: string | null }).agentId == null)
276-
) {
277-
const eventWithAgent = eventPayload as { agentId?: string }
278-
eventWithAgent.agentId = eventAgentId
279-
}
280-
} else {
281-
eventPayload = {
282-
type: 'text',
283-
text: trimmedText,
284-
} as PrintModeEvent
285-
286-
if (eventAgentId) {
287-
const eventWithAgent = eventPayload as { agentId?: string }
288-
eventWithAgent.agentId = eventAgentId
289-
}
253+
254+
const eventPayload = {
255+
type: 'text',
256+
text: trimmedText,
257+
} as PrintModeEvent
258+
259+
if (eventAgentId) {
260+
const eventWithAgent = eventPayload as { agentId?: string }
261+
eventWithAgent.agentId = eventAgentId
290262
}
291263

292264
await handleEvent?.(eventPayload)
@@ -312,46 +284,18 @@ export async function run({
312284
agentKey: string,
313285
eventAgentId?: string,
314286
): Promise<void> => {
315-
const state = textFilterStates.get(agentKey)
316-
let pending = ''
317-
318-
if (state) {
319-
const { text: pendingText } = filterToolXmlFromText(
320-
state,
321-
'',
322-
MAX_TOOL_XML_BUFFER,
323-
)
324-
pending = pendingText
325-
326-
if (state.buffer && !state.buffer.includes('<')) {
327-
pending += state.buffer
328-
}
329-
330-
state.buffer = ''
331-
state.activeTag = null
332-
333-
textFilterStates.delete(agentKey)
334-
} else {
335-
ensureSectionStart(agentKey)
336-
}
337-
338-
let nextFullText = textAccumulator.get(agentKey) ?? ''
339287
ensureSectionStart(agentKey)
340288

341-
if (pending) {
342-
nextFullText = accumulateText(agentKey, pending)
343-
if (agentKey === ROOT_AGENT_KEY) {
344-
await emitStreamDelta(agentKey, nextFullText)
345-
}
289+
const nextFullText = textAccumulator.get(agentKey) ?? ''
290+
if (agentKey === ROOT_AGENT_KEY && nextFullText) {
291+
await emitStreamDelta(agentKey, nextFullText)
346292
}
347293

348294
await emitPendingSection(agentKey, eventAgentId)
349295

350296
textAccumulator.delete(agentKey)
351297
lastStreamedTextByAgent.delete(agentKey)
352298
sectionStartIndexByAgent.delete(agentKey)
353-
354-
lastTextEventByAgent.delete(agentKey)
355299
}
356300

357301
const flushSubagentState = async (
@@ -430,33 +374,6 @@ export async function run({
430374
const nextFullText = accumulateText(ROOT_AGENT_KEY, sanitized)
431375
await emitStreamDelta(ROOT_AGENT_KEY, nextFullText)
432376
}
433-
} else if (chunk.type === 'text') {
434-
const agentKey = chunk.agentId ?? ROOT_AGENT_KEY
435-
const state = getTextFilterState(agentKey)
436-
lastTextEventByAgent.set(agentKey, { ...chunk })
437-
ensureSectionStart(agentKey)
438-
const { text: sanitized } = filterToolXmlFromText(
439-
state,
440-
chunk.text,
441-
MAX_TOOL_XML_BUFFER,
442-
)
443-
444-
if (sanitized) {
445-
const nextFullText = accumulateText(agentKey, sanitized)
446-
if (agentKey === ROOT_AGENT_KEY) {
447-
await emitStreamDelta(agentKey, nextFullText)
448-
}
449-
}
450-
451-
const fullText = textAccumulator.get(agentKey) ?? ''
452-
const startIndex =
453-
sectionStartIndexByAgent.get(agentKey) ?? fullText.length
454-
const shouldFlushForToolXml =
455-
state.activeTag != null && startIndex < fullText.length
456-
457-
if (shouldFlushForToolXml) {
458-
await emitPendingSection(agentKey, chunk.agentId)
459-
}
460377
} else {
461378
const chunkType = chunk.type as string
462379

0 commit comments

Comments
 (0)