diff --git a/packages/agent-runtime/src/tools/stream-parser.ts b/packages/agent-runtime/src/tools/stream-parser.ts index edd0e60bde..69bccde3b3 100644 --- a/packages/agent-runtime/src/tools/stream-parser.ts +++ b/packages/agent-runtime/src/tools/stream-parser.ts @@ -2,10 +2,7 @@ import { toolNames } from '@codebuff/common/tools/constants' import { buildArray } from '@codebuff/common/util/array' import { STREAM_RECOVERY_EVENT } from '@codebuff/common/util/axiom-only-log' import { AbortError } from '@codebuff/common/util/error' -import { - assistantMessage, - userMessage, -} from '@codebuff/common/util/messages' +import { assistantMessage, userMessage } from '@codebuff/common/util/messages' import { generateCompactId } from '@codebuff/common/util/string' import { processStreamWithTools } from '../tool-stream-parser' @@ -18,11 +15,11 @@ import { } from './tool-executor' import { withSystemTags } from '../util/messages' import { + historyHasUnclosedOpen, historyLeaksThinkTags, stripThinkScaffolding, ThinkTagStream, } from '../util/think-tag-stream' - import type { CustomToolCall, ExecuteToolCallParams } from './tool-executor' import type { ThinkStreamSegment } from '../util/think-tag-stream' import type { AgentTemplate } from '../templates/types' @@ -156,9 +153,7 @@ export async function processStream( > & ParamsExcluding< typeof processStreamWithTools, - | 'processors' - | 'defaultProcessor' - | 'executeXmlToolCall' + 'processors' | 'defaultProcessor' | 'executeXmlToolCall' >, ) { const { @@ -180,10 +175,11 @@ export async function processStream( // Reasoning that a lane failed to put in its native field arrives here as // ordinary text, tags and all. Reclassify it before it reaches a surface, so // the thinking box is the only place a chain of thought is ever rendered. - // See util/think-tag-stream.ts for the three shapes and why the implicit-open - // rule is armed from the history rather than from a model id. + // See util/think-tag-stream.ts for the three shapes and why the hold rules + // are armed from the history rather than from a model id. const thinkTagStream = new ThinkTagStream({ implicitOpen: historyLeaksThinkTags(agentState.messageHistory), + holdExplicitOpens: historyHasUnclosedOpen(agentState.messageHistory), }) const emitThinkSegments = (segments: ThinkStreamSegment[]): void => { for (const segment of segments) { @@ -205,7 +201,8 @@ export async function processStream( const toolResults: ToolMessage[] = [] const toolResultsToAddToMessageHistory: ToolMessage[] = [] const toolCalls: (CodebuffToolCall | CustomToolCall)[] = [] - const toolCallsToAddToMessageHistory: (CodebuffToolCall | CustomToolCall)[] = [] + const toolCallsToAddToMessageHistory: (CodebuffToolCall | CustomToolCall)[] = + [] const assistantMessages: Message[] = [] // Inline agents replace the parent's history with their result. Track which // current-step messages they inherited so finalization does not append them @@ -255,7 +252,7 @@ export async function processStream( function createToolExecutionCallback(toolName: string, isXmlMode: boolean) { const responseHandler = createResponseHandler() return { - onTagStart: () => { }, + onTagStart: () => {}, onTagEnd: async (_: string, input: Record) => { if (signal.aborted) { return @@ -266,10 +263,10 @@ export async function processStream( // Check if this is an agent tool call that should be transformed to spawn_agents const transformed = !isNativeTool ? tryTransformAgentToolCall({ - toolName, - input, - spawnableAgents: agentTemplate.spawnableAgents, - }) + toolName, + input, + spawnableAgents: agentTemplate.spawnableAgents, + }) : null const isSpawnCall = Boolean(transformed) || @@ -629,17 +626,18 @@ export async function processStream( const completedToolCallIds = new Set( toolResultsToAddToMessageHistory.map((r) => r.toolCallId), ) - const filteredToolCalls = - toolCallsToAddToMessageHistory.filter((tc) => - completedToolCallIds.has(tc.toolCallId), - ) + const filteredToolCalls = toolCallsToAddToMessageHistory.filter((tc) => + completedToolCallIds.has(tc.toolCallId), + ) agentState.messageHistory = buildArray([ ...agentState.messageHistory, ...assistantMessages.filter( (message) => !claimedByInlineAgent.has(message), ), - ...filteredToolCalls.map((toolCall) => assistantMessage({ ...toolCall, type: 'tool-call' })), + ...filteredToolCalls.map((toolCall) => + assistantMessage({ ...toolCall, type: 'tool-call' }), + ), ...toolResultsToAddToMessageHistory, ...errorMessages, ]) diff --git a/packages/agent-runtime/src/util/__tests__/think-tag-stream.test.ts b/packages/agent-runtime/src/util/__tests__/think-tag-stream.test.ts index 702f5fe3f7..f62a8eacd9 100644 --- a/packages/agent-runtime/src/util/__tests__/think-tag-stream.test.ts +++ b/packages/agent-runtime/src/util/__tests__/think-tag-stream.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'bun:test' import { + historyHasUnclosedOpen, historyLeaksThinkTags, IMPLICIT_OPEN_BUDGET_CHARS, stripThinkScaffolding, @@ -12,7 +13,7 @@ import type { ThinkStreamSegment } from '../think-tag-stream' /** Feed the deltas one at a time, then flush — the shape a real stream has. */ function run( deltas: string[], - options?: { implicitOpen?: boolean }, + options?: { implicitOpen?: boolean; holdExplicitOpens?: boolean }, ): ThinkStreamSegment[] { const stream = new ThinkTagStream(options) const out: ThinkStreamSegment[] = [] @@ -59,13 +60,104 @@ describe('ThinkTagStream — paired tags', () => { expect(joined(run(['ends with <']), 'text')).toBe('ends with <') }) - it('treats an unclosed open tag as reasoning through end of stream', () => { + it('streams an unclosed open as reasoning on a clean lane, live like rule 1', () => { + // Default lane: a bare open is near-certainly a real block (truncated + // thought). Streaming it live keeps the zero-latency main behavior — + // the review's concern was exactly this trace freezing until a close. const out = run(['truncated thou', 'ght']) expect(joined(out, 'reasoning')).toBe('truncated thought') expect(joined(out, 'text')).toBe('') }) -}) + it('keeps an answer quoted around a prose open tag as text when armed', () => { + // History-proven lane (an open was left unclosed before): the hold is + // armed, so a quoted tag cannot swallow the answer into the thinking box. + const out = run( + ['Write ', 'like this in your docs. The answer continues here.'], + { holdExplicitOpens: true }, + ) + expect(joined(out, 'reasoning')).toBe('') + expect(joined(out, 'text')).toBe( + 'Write like this in your docs. The answer continues here.', + ) + }) + + it('holds an explicit open only until its close arrives, when armed', () => { + const stream = new ThinkTagStream({ holdExplicitOpens: true }) + expect(stream.push('Answer part one. plan it')).toEqual([ + { type: 'text', text: 'Answer part one. ' }, + ]) + expect(stream.push('Here is the answer.')).toEqual([ + { type: 'reasoning', text: 'plan it' }, + { type: 'text', text: 'Here is the answer.' }, + ]) + }) + + it('releases an armed hold past the budget as text, and the rest streams live', () => { + // Same give-up as the implicit head: past the budget the step is + // answering, not thinking — a quoted tag with a long answer after it + // must not be swallowed (issue #1155, bug 2). + const stream = new ThinkTagStream({ holdExplicitOpens: true }) + const long = 'x'.repeat(IMPLICIT_OPEN_BUDGET_CHARS) + expect(stream.push(`Answer ${long}`)).toEqual([ + { type: 'text', text: 'Answer ' }, + { type: 'text', text: long }, + ]) + expect(stream.push(' still answering')).toEqual([ + { type: 'text', text: ' still answering' }, + ]) + // Disarmed: a later marker is stripped, not treated as a close. + expect(stream.push('tail')).toEqual([ + { type: 'text', text: 'tail' }, + ]) + }) + + it('streams a long well-formed trace per-delta on a clean lane, never buffered', () => { + // The review's measurement: a DeepSeek-R1-style trace must not wait for + // its close. One push in, everything so far is already out. + const stream = new ThinkTagStream() + const long = 'x'.repeat(5000) + expect(stream.push(`${long}`)).toEqual([ + { type: 'reasoning', text: long }, + ]) + expect(stream.push(' still going')).toEqual([ + { type: 'reasoning', text: ' still going' }, + ]) + expect(stream.push('answer')).toEqual([ + { type: 'text', text: 'answer' }, + ]) + }) + + it('flushes a clean-lane unclosed block tail as reasoning', () => { + const stream = new ThinkTagStream() + expect(stream.push('thought { + const stream = new ThinkTagStream({ implicitOpen: true }) + const long = 'x'.repeat(IMPLICIT_OPEN_BUDGET_CHARS) + expect(stream.push(long)).toEqual([{ type: 'text', text: long }]) + expect(stream.push('tail')).toEqual([ + { type: 'text', text: 'tail' }, + ]) + }) + + it('releases an explicit open on a native reasoning chunk as text, when armed', () => { + const stream = new ThinkTagStream({ holdExplicitOpens: true }) + expect(stream.push('A quoted open')).toEqual([ + { type: 'text', text: 'A ' }, + ]) + expect(stream.disarmImplicitOpen()).toEqual([ + { type: 'text', text: ' quoted open' }, + ]) + expect(stream.push(' and the answer continues')).toEqual([ + { type: 'text', text: ' and the answer continues' }, + ]) + }) +}) describe('ThinkTagStream — orphan close, not armed', () => { // The default for every non-leaking model: the prose is not reclassified // (it already streamed), but the bare marker must never reach a transcript. @@ -86,7 +178,10 @@ describe('ThinkTagStream — orphan close, not armed', () => { describe('ThinkTagStream — orphan close, armed', () => { it('reclassifies the head as reasoning once the marker lands', () => { const out = run( - ['Ключевая зацепка: the bundle knows the type.', 'Do that.Real answer.'], + [ + 'Ключевая зацепка: the bundle knows the type.', + 'Do that.Real answer.', + ], { implicitOpen: true }, ) expect(out).toEqual([ @@ -133,7 +228,9 @@ describe('ThinkTagStream — orphan close, armed', () => { { type: 'text', text: ' and more' }, ]) // Disarmed: a later marker is stripped, not treated as a close. - expect(stream.push('tail')).toEqual([{ type: 'text', text: 'tail' }]) + expect(stream.push('tail')).toEqual([ + { type: 'text', text: 'tail' }, + ]) }) it('disarms on a native reasoning chunk and releases the head as text', () => { @@ -183,7 +280,10 @@ describe('historyLeaksThinkTags', () => { expect( historyLeaksThinkTags([ { role: 'user', content: [{ type: 'text', text: 'why ?' }] }, - { role: 'assistant', content: [{ type: 'reasoning', text: '' }] }, + { + role: 'assistant', + content: [{ type: 'reasoning', text: '' }], + }, ]), ).toBe(false) }) @@ -198,7 +298,9 @@ describe('stripThinkScaffolding', () => { it('leaves surrounding whitespace alone, unlike stripThinkTags', () => { expect(stripThinkScaffolding(' spaced ')).toBe(' spaced ') - expect(stripThinkScaffolding('a\n\nx\n\nb')).toBe('a\n\n\n\nb') + expect(stripThinkScaffolding('a\n\nx\n\nb')).toBe( + 'a\n\n\n\nb', + ) }) }) @@ -215,3 +317,55 @@ describe('historyLeaksThinkTags — head window', () => { ).toBe(false) }) }) + +describe('historyHasUnclosedOpen', () => { + const assistant = (text: string) => ({ + role: 'assistant', + content: [{ type: 'text', text }], + }) + + it('is false for a clean or properly paired last turn', () => { + expect(historyHasUnclosedOpen([])).toBe(false) + expect(historyHasUnclosedOpen([assistant('xanswer')])).toBe( + false, + ) + }) + + it('is true when the last assistant turn leaves an open unclosed', () => { + expect( + historyHasUnclosedOpen([assistant('answertruncated thought')]), + ).toBe(true) + }) + + it('is false when both tags are quoted in prose', () => { + // Docs quoting the pair close after they open — not evidence of a leak. + expect( + historyHasUnclosedOpen([ + assistant('use and to mark reasoning'), + ]), + ).toBe(false) + }) + + it('heals: only the last assistant turn counts', () => { + // A one-off quoted open arms exactly the next step; a clean reply after + // it disarms, so later genuine traces are never held. + expect( + historyHasUnclosedOpen([ + assistant('answerquoted once'), + assistant('xclean answer'), + ]), + ).toBe(false) + }) + + it('ignores user messages and reasoning parts', () => { + expect( + historyHasUnclosedOpen([ + { role: 'user', content: [{ type: 'text', text: 'why ?' }] }, + { + role: 'assistant', + content: [{ type: 'reasoning', text: 'native' }], + }, + ]), + ).toBe(false) + }) +}) diff --git a/packages/agent-runtime/src/util/think-tag-stream.ts b/packages/agent-runtime/src/util/think-tag-stream.ts index 86b37af3da..8b2ebb71de 100644 --- a/packages/agent-runtime/src/util/think-tag-stream.ts +++ b/packages/agent-runtime/src/util/think-tag-stream.ts @@ -22,8 +22,20 @@ * * 1. `` — paired tags. Content between them is reasoning. * Unambiguous, free, always on. - * 2. A bare `` that never closes (a truncated thought). Everything - * after it is reasoning. + * 2. A bare `` that never closes. Two different events produce this: + * a thought truncated mid-stream, or the tag written as PROSE — docs that + * mention the tag, a quoted template, a lane with a broken chat template. + * The two are indistinguishable while the deltas arrive. On a lane whose + * history shows no unclosed open — the majority, every model that pairs + * its tags — the block streams live as reasoning, exactly like rule 1: + * buffering a genuine chain of thought until its close would freeze the + * thinking box for the whole trace. When {@link historyHasUnclosedOpen} + * proves the lane leaves opens unclosed, + * {@link ThinkTagStreamOptions.holdExplicitOpens} holds from the open + * instead: a close settles it as reasoning, and the budget or the end of + * the step releases it as text, so an answer is delayed but never + * swallowed. The cost of arming from history is one step of lag on a + * lane that leaks for the first time — the same accepted gap rule 3 has. * 3. An orphan `` with no open tag — the DeepSeek shape above, where * the open tag was consumed by the chat template's prefill. The text * BEFORE it is reasoning, but by the time the marker arrives that text has @@ -58,14 +70,30 @@ export interface ThinkTagStreamOptions { * buffer as text. */ implicitOpen?: boolean + + /** + * Hold from every explicit `` open until a close, the + * {@link IMPLICIT_OPEN_BUDGET_CHARS}, or the end of the step settles it. + * + * Off — the default — keeps rule 1 free: a paired block streams live with + * zero buffering, and the only cost is that a prose-quoted open on a lane + * with no history swallows that one step into the thinking box. Arm it from + * {@link historyHasUnclosedOpen}, never unconditionally: on a lane proven + * to leave opens unclosed the hold is protection, but on a clean lane it + * would delay every genuine trace for nothing. + */ + holdExplicitOpens?: boolean } /** - * How much leading content to hold while waiting for an orphan ``. + * How much leading content to hold while waiting for a settling marker. * * A leaked chain of thought runs well past this, so the cap is not there to - * fit one — it bounds the wrong case. If the marker has not arrived by here + * fit one — it bounds the wrong case. If the close has not arrived by here * the step is answering, not thinking, and the buffer is released as text. + * It bounds the implicit head ({@link ThinkTagStreamOptions.implicitOpen}) + * and, when armed, every explicit open + * ({@link ThinkTagStreamOptions.holdExplicitOpens}) alike. */ export const IMPLICIT_OPEN_BUDGET_CHARS = 4000 @@ -142,6 +170,38 @@ export function historyLeaksThinkTags( } return false } +/** + * True when the most recent assistant turn ended inside an explicit think + * block — its text contains an open with no close after it. + * + * The arming signal for {@link ThinkTagStreamOptions.holdExplicitOpens}, + * symmetric with {@link historyLeaksThinkTags}: the same lane decision the + * history is kept for. A lane that pairs its tags cleanly never arms, so + * rule 1 stays free; a lane proven to leave an open unclosed arms the next + * step so a prose-quoted tag cannot swallow the answer. + */ +export function historyHasUnclosedOpen( + messages: readonly { role: string; content: unknown }[], +): boolean { + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i] + if (message.role !== 'assistant') continue + if (!Array.isArray(message.content)) return false + let text = '' + for (const part of message.content) { + if ( + part && + typeof part === 'object' && + (part as { type?: unknown }).type === 'text' && + typeof (part as { text?: unknown }).text === 'string' + ) { + text += (part as { text: string }).text + } + } + return text.lastIndexOf(OPEN_TAG) > text.lastIndexOf(CLOSE_TAG) + } + return false +} /** * Incremental classifier over one step's content stream. @@ -153,27 +213,40 @@ export function historyLeaksThinkTags( export class ThinkTagStream { /** Trailing bytes withheld because they may be the start of a tag. */ private partial = '' - /** Leading content withheld while `implicitOpen` is still undecided. */ + /** Content withheld while the classification of the open block is still + * undecided: the implicit head, or everything since an explicit open on an + * armed lane. */ private held = '' private implicitOpen: boolean private inThinkBlock: boolean + /** True while `held` is waiting for the close that settles its + * classification. Armed by construction (implicitOpen) and by an explicit + * open when {@link ThinkTagStreamOptions.holdExplicitOpens} is set; + * cleared by the close ({@link confirmOpenHold}) or by giving up + * ({@link abandonOpenHold}) — budget, a native reasoning chunk, or flush. */ + private holdingForOpen: boolean + private holdExplicitOpens: boolean constructor(options: ThinkTagStreamOptions = {}) { this.implicitOpen = options.implicitOpen ?? false + this.holdExplicitOpens = options.holdExplicitOpens ?? false this.inThinkBlock = this.implicitOpen + this.holdingForOpen = this.implicitOpen } /** - * Give up on `implicitOpen` and release anything held as text. + * Give up on the hold — the implicit head's, or an explicit open's when + * {@link ThinkTagStreamOptions.holdExplicitOpens} armed it — and release + * anything held as text. * - * Called when the step turns out not to be leaking after all. The strongest - * such signal is a native reasoning chunk: a lane that populates - * `reasoning_content` is by definition not putting the thought in `content`, - * so whatever is in `content` is the answer. + * Called when the step turns out not to be thinking in `content` after all. + * The strongest such signal is a native reasoning chunk: a lane that + * populates `reasoning_content` is by definition not putting the thought in + * `content`, so whatever is in `content` is the answer. */ disarmImplicitOpen(): ThinkStreamSegment[] { - if (!this.implicitOpen) return [] - return this.abandonImplicitOpen() + if (!this.holdingForOpen) return [] + return this.abandonOpenHold() } push(chunk: string): ThinkStreamSegment[] { @@ -189,10 +262,11 @@ export class ThinkTagStream { this.addReasoning(segments, buffer.slice(0, closeIdx)) buffer = buffer.slice(closeIdx + CLOSE_TAG.length) this.inThinkBlock = false - // The close the implicit block was waiting for: everything held is - // confirmed reasoning. It can only happen once — a later orphan close - // is an ordinary stray marker and is stripped below. - this.confirmImplicitOpen(segments) + // The close the hold was waiting for: everything held is confirmed + // reasoning — whether the block was opened implicitly (a leaked chain + // of thought) or explicitly. A later orphan close is an ordinary + // stray marker and is stripped below. + this.confirmOpenHold(segments) continue } @@ -203,6 +277,12 @@ export class ThinkTagStream { this.addText(segments, buffer.slice(0, openIdx)) buffer = buffer.slice(openIdx + OPEN_TAG.length) this.inThinkBlock = true + // An open tag could be a block the model is thinking in, or the tag + // quoted as prose. On a clean lane it is near-certainly the former — + // rule 1 streams it live, free. Only a lane the history has proven + // to leave opens unclosed holds from here, so the answer behind a + // quoted tag is delayed rather than swallowed. + if (this.holdExplicitOpens) this.holdingForOpen = true continue } // Orphan close with nothing to close: drop the marker so it cannot reach @@ -220,23 +300,28 @@ export class ThinkTagStream { } /** Emit everything withheld. A partial tag that never completed was always - * just text, and content held for an orphan close that never came is the - * answer — releasing both here is what makes the speculation lossless. */ + * just text — except at the tail of an unclosed block on a clean lane, + * where it belongs with the reasoning it would have ended. Content held + * for a close that never came is the answer: releasing it here is what + * makes the hold lossless. An unclosed open on an armed lane is treated + * exactly like an orphan close — the marker is scaffolding, the text + * around it is the answer. */ flush(): ThinkStreamSegment[] { const segments: ThinkStreamSegment[] = [] + if (this.holdingForOpen) segments.push(...this.abandonOpenHold()) const trailing = this.partial this.partial = '' if (trailing) { if (this.inThinkBlock) this.addReasoning(segments, trailing) else this.addText(segments, trailing) } - if (this.implicitOpen) segments.push(...this.abandonImplicitOpen()) return segments } - /** The orphan close arrived: what was held was reasoning after all. */ - private confirmImplicitOpen(segments: ThinkStreamSegment[]): void { - if (!this.implicitOpen) return + /** The close arrived: what was held was reasoning after all. */ + private confirmOpenHold(segments: ThinkStreamSegment[]): void { + if (!this.holdingForOpen) return + this.holdingForOpen = false this.implicitOpen = false const held = this.held this.held = '' @@ -244,7 +329,8 @@ export class ThinkTagStream { } /** No close is coming: what was held was the answer. */ - private abandonImplicitOpen(): ThinkStreamSegment[] { + private abandonOpenHold(): ThinkStreamSegment[] { + this.holdingForOpen = false this.implicitOpen = false this.inThinkBlock = false const held = this.held @@ -252,23 +338,21 @@ export class ThinkTagStream { return held ? [{ type: 'text', text: held }] : [] } - private addReasoning( - segments: ThinkStreamSegment[], - text: string, - ): void { + private addReasoning(segments: ThinkStreamSegment[], text: string): void { // A nested/duplicated open tag inside a block is scaffolding, never thought. const cleaned = text.split(OPEN_TAG).join('') if (!cleaned) return - if (!this.implicitOpen) { + if (!this.holdingForOpen) { push(segments, 'reasoning', cleaned) return } - // Still undecided: this is reasoning only if an orphan close confirms it, - // so hold rather than send. Past the budget the step is answering, not + // Undecided: reasoning only if a close confirms it, so hold rather than + // send. A held chain of thought — implicit head or armed explicit open — + // runs well past the budget, so past it the step is answering, not // thinking, and the hold is released as text. this.held += cleaned if (this.held.length >= IMPLICIT_OPEN_BUDGET_CHARS) { - segments.push(...this.abandonImplicitOpen()) + segments.push(...this.abandonOpenHold()) } }