diff --git a/cli/src/utils/__tests__/block-operations-streaming.test.ts b/cli/src/utils/__tests__/block-operations-streaming.test.ts new file mode 100644 index 0000000000..5bec86b451 --- /dev/null +++ b/cli/src/utils/__tests__/block-operations-streaming.test.ts @@ -0,0 +1,130 @@ +import { describe, expect, test } from 'bun:test' + +import { appendTextToRootStream } from '../block-operations' + +import type { ContentBlock, TextContentBlock } from '../../types/chat' + +/** + * The text stream path only scans the seam between the last block and the new + * delta when looking for a partial think tag (see block-operations.ts), so it + * must still notice a tag split at any point inside the tag syntax, including + * the widest partial prefix it has to recognise (` { + let blocks: ContentBlock[] = [] + for (let i = 0; i < text.length; i += chunkSize) { + blocks = appendTextToRootStream(blocks, { + type: 'text', + text: text.slice(i, i + chunkSize), + }) + } + return blocks +} + +const shape = (blocks: ContentBlock[]) => + blocks.map((block) => { + const text = block as TextContentBlock + return { + type: block.type, + textType: text.textType, + content: text.content, + thinkingOpen: text.thinkingOpen, + } + }) + +const textOf = (blocks: ContentBlock[]) => + blocks.map((block) => (block as TextContentBlock).content ?? '').join('') + +describe('think tags split across delta boundaries', () => { + test('opening tag split across two deltas opens a thinking block', () => { + const first = appendTextToRootStream( + [{ type: 'text', content: 'a', textType: 'text' }], + { type: 'text', text: ' b' }) + + expect(shape(second)).toEqual([ + { type: 'text', textType: 'text', content: 'a ', thinkingOpen: undefined }, + { type: 'text', textType: 'reasoning', content: 'b', thinkingOpen: true }, + ]) + }) + + test('closing tag split across two deltas closes the thinking block', () => { + const first = appendTextToRootStream([], { + type: 'text', + text: 'a b' }) + + expect(shape(second)).toEqual([ + { type: 'text', textType: 'reasoning', content: 'a', thinkingOpen: false }, + { type: 'text', textType: 'text', content: ' b', thinkingOpen: undefined }, + ]) + }) + + test('an empty think pair produces no blocks, however it is chunked', () => { + for (const chunkSize of [1, 2, 7, 8, 100]) { + expect(stream('', chunkSize)).toEqual([]) + } + }) + + test('the widest partial close prefix is still recognised', () => { + // `` arrives on its own. + const first = appendTextToRootStream([], { + type: 'text', + text: 'thought' }) + + expect(shape(second)).toEqual([ + { + type: 'text', + textType: 'reasoning', + content: 'thought', + thinkingOpen: false, + }, + ]) + }) + + test('ordinary less-than at the end of a delta is kept as text', () => { + const first = appendTextToRootStream([], { type: 'text', text: 'a <' }) + const second = appendTextToRootStream(first, { type: 'text', text: ' b' }) + + expect(shape(second)).toEqual([ + { + type: 'text', + textType: 'text', + content: 'a < b', + thinkingOpen: undefined, + }, + ]) + }) +}) + +describe('long streams stay a single block', () => { + test('5000 characters without tags accumulate into one block', () => { + const text = 'abcdefghij'.repeat(500) + const blocks = stream(text, 7) + + expect(blocks).toHaveLength(1) + expect((blocks[0] as TextContentBlock).content).toBe(text) + }) + + test('a tag completed after a long prefix is still split out', () => { + const prefix = 'x'.repeat(10_000) + const blocks = stream(`${prefix}thoughttail`, 64) + + // The prefix must stay intact and keep its length: no truncation at the seam. + expect(textOf(blocks)).toBe(`${prefix}thoughttail`) + expect(blocks.some((b) => (b as TextContentBlock).textType === 'reasoning')) + .toBe(true) + }) + + test('a large block with a less-than keeps every character', () => { + const text = `${'y'.repeat(10_000)} < not a tag` + const blocks = stream(text, 33) + + expect(textOf(blocks)).toBe(text) + }) +}) diff --git a/cli/src/utils/block-operations.ts b/cli/src/utils/block-operations.ts index 86396507ac..789dc42220 100644 --- a/cli/src/utils/block-operations.ts +++ b/cli/src/utils/block-operations.ts @@ -170,14 +170,27 @@ const appendTextWithThinkParsingToBlocks = ( lastBlock?.type === 'text' && lastBlock.textType === 'text' ) { - const combinedText = lastBlock.content + text + // A think tag can only be *completed* by this delta, at the seam between + // the block's tail and `text`; a tag already sitting inside the block was + // split out when it arrived. So scan the seam first and only fall back to + // the whole block when a '<' could still be mid-tag. Concatenating and + // scanning the full block per delta made this O(block) per chunk, i.e. + // quadratic over a long reply: a 470KB stream spent ~430ms here. + const seamLength = Math.max(THINK_OPEN_TAG.length, THINK_CLOSE_TAG.length) - 1 + const seam = lastBlock.content.slice( + Math.max(0, lastBlock.content.length - seamLength), + ) - if ( - combinedText.includes(THINK_OPEN_TAG) || - combinedText.includes(THINK_CLOSE_TAG) - ) { - textToParse = combinedText - nextBlocks.pop() + if (seam.includes('<') || text.includes('<')) { + const combinedText = lastBlock.content + text + + if ( + combinedText.includes(THINK_OPEN_TAG) || + combinedText.includes(THINK_CLOSE_TAG) + ) { + textToParse = combinedText + nextBlocks.pop() + } } }