fix(agent-runtime): release an unclosed think-tag open as text at flush

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<string, string>) => {
         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<Message>([
       ...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,
     ])

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(['<think>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 <think> ', '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. <think>plan it')).toEqual([
+      { type: 'text', text: 'Answer part one. ' },
+    ])
+    expect(stream.push('</think>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 <think>${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('</think>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(`<think>${long}`)).toEqual([
+      { type: 'reasoning', text: long },
+    ])
+    expect(stream.push(' still going')).toEqual([
+      { type: 'reasoning', text: ' still going' },
+    ])
+    expect(stream.push('</think>answer')).toEqual([
+      { type: 'text', text: 'answer' },
+    ])
+  })
+
+  it('flushes a clean-lane unclosed block tail as reasoning', () => {
+    const stream = new ThinkTagStream()
+    expect(stream.push('<think>thought</thi')).toEqual([
+      { type: 'reasoning', text: 'thought' },
+    ])
+    expect(stream.flush()).toEqual([{ type: 'reasoning', text: '</thi' }])
+  })
+
+  it('keeps the implicit head release as text at the shared budget', () => {
+    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('</think>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 <think> 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.</think>Real answer.'],
+      [
+        'Ключевая зацепка: the bundle knows the type.',
+        'Do that.</think>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('</think>tail')).toEqual([{ type: 'text', text: 'tail' }])
+    expect(stream.push('</think>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 </think>?' }] },
-        { role: 'assistant', content: [{ type: 'reasoning', text: '</think>' }] },
+        {
+          role: 'assistant',
+          content: [{ type: 'reasoning', text: '</think>' }],
+        },
       ]),
     ).toBe(false)
   })
@@ -198,7 +298,9 @@ describe('stripThinkScaffolding', () => {
 
   it('leaves surrounding whitespace alone, unlike stripThinkTags', () => {
     expect(stripThinkScaffolding('  spaced  ')).toBe('  spaced  ')
-    expect(stripThinkScaffolding('a\n\n<think>x</think>\n\nb')).toBe('a\n\n\n\nb')
+    expect(stripThinkScaffolding('a\n\n<think>x</think>\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('<think>x</think>answer')])).toBe(
+      false,
+    )
+  })
+
+  it('is true when the last assistant turn leaves an open unclosed', () => {
+    expect(
+      historyHasUnclosedOpen([assistant('answer<think>truncated 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 <think> and </think> 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('answer<think>quoted once'),
+        assistant('<think>x</think>clean answer'),
+      ]),
+    ).toBe(false)
+  })
+
+  it('ignores user messages and reasoning parts', () => {
+    expect(
+      historyHasUnclosedOpen([
+        { role: 'user', content: [{ type: 'text', text: 'why <think>?' }] },
+        {
+          role: 'assistant',
+          content: [{ type: 'reasoning', text: '<think>native' }],
+        },
+      ]),
+    ).toBe(false)
+  })
+})