diff --git a/backend/agents/stream_handler.py b/backend/agents/stream_handler.py index d8e2d63..8a0441a 100644 --- a/backend/agents/stream_handler.py +++ b/backend/agents/stream_handler.py @@ -34,19 +34,37 @@ class StreamEventHandler: """Observe stream chunks and drain AG-UI events. - Phases - ------ - 1. Construct → first ``drain()`` yields ``RUN_STARTED``. - 2. For each ``(chunk, metadata)`` from ``astream``: - ``observe(chunk, metadata)`` → ``drain()``. - 3. After stream exhausts: - ``finalize()`` → closes open blocks, yields ``RUN_FINISHED``. - Or ``error(message)`` → closes blocks, yields ``RUN_ERROR``. - - Only a single text message and a single reasoning block are supported - (typical for single ``create_agent()`` pipelines). Tool calls may be - concurrent (tracked independently by id). - """ + Two-phase ``observe()`` / ``drain()`` pattern decouples chunk processing + from event emission for clean testability. + + Reasoning-content strategies + ----------------------------- + LLM providers stream reasoning tokens in two incompatible ways: + + **A — Cumulative** (DeepSeek R1, some fine-tuned clones) + Each chunk carries the *full accumulated reasoning text* from the + start of the current reasoning block:: + + chunk 1: "Think step 1" + chunk 2: "Think step 1...Think step 2..." + + The handler diffs ``raw_n - raw_{n-1}`` via ``startswith`` and emits + only the new portion as the delta. + + **B — True delta** (OpenAI o-series, Gemini) + Each chunk carries *only new text*:: + + chunk 1: "Think step 1" + chunk 2: "...Think step 2..." + + The ``startswith`` diff typically fails (non-overlapping), so the + handler emits the incoming text as-is — which is already a delta. + + **Cross-block scenario (Pattern A)** + When the LLM opens a new reasoning block after a tool result, it may + re-iterate prior reasoning. The accumulated text is **not reset** on + block close, so the diff correctly strips old content from the new + block's first chunk.""" def __init__( self, @@ -69,8 +87,15 @@ def __init__( # Track the last known id per index to fill in the gap. self._last_tool_call_id_by_index: dict[int, str] = {} - # Accumulated reasoning content for delta computation - # (DeepSeek sends full accumulated text in each chunk) + # Last-seen raw reasoning text for the universal startswith diff. + # Handles both cumulative (DeepSeek) and delta (OpenAI) patterns: + # + # Cumulative: delta = raw - self._last_reasoning_content (startswith succeeds) + # Delta: delta = raw (startswith fails → emit as-is) + # + # Crucially NOT reset on _close_reasoning(): if the model re-iterates + # prior reasoning from an earlier block, we need this memory to + # subtract the old text from the new block's first chunk. self._last_reasoning_content: str = "" # Step counter for unique STEP_STARTED stepIds per reasoning block @@ -327,7 +352,13 @@ def _ensure_tool_open(self, tcc: ToolCallChunk, resolved_id: str | None = None) def _close_reasoning(self) -> None: if self._reasoning_open: self._reasoning_open = False - self._last_reasoning_content = "" + # NB: do NOT reset _last_reasoning_content here! + # DeepSeek (and some other models) include prior reasoning + # in subsequent reasoning blocks when they see their own + # previous output in the conversation context. If we reset + # the accumulated text, the diff against the next block's + # first chunk can't strip the old content, causing each + # reasoning block to grow by the sum of all prior blocks. self._current_reasoning_step_id = None self._pending.append( ReasoningMessageEndEvent(message_id=self._message_id), diff --git a/tests/agents/test_stream_handler.py b/tests/agents/test_stream_handler.py index 0d1be09..56b0e07 100644 --- a/tests/agents/test_stream_handler.py +++ b/tests/agents/test_stream_handler.py @@ -357,8 +357,13 @@ def test_accumulated_reasoning_deltas(self, handler): assert len(content_events3) == 1 assert content_events3[0].delta == "Let me read..." - def test_accumulated_reasoning_resets_on_close(self, handler): - """After reasoning closes, the accumulated text resets for the next block.""" + def test_cross_block_non_overlapping_delta(self, handler): + """Non-overlapping cross-block reasoning emits full new text. + + When the new reasoning block does NOT reiterate prior text + (Pattern B — OpenAI o-series, Gemini), the startswith diff + fails and the handler emits the incoming text as-is. + """ chunk1 = AIMessageChunk( content="", additional_kwargs={"reasoning_content": "First thought"}, @@ -371,7 +376,7 @@ def test_accumulated_reasoning_resets_on_close(self, handler): handler.observe(chunk2, {"langgraph_node": "agent"}) handler.drain() - # Start new reasoning block — should reset accumulated text + # New reasoning block with completely different content chunk3 = AIMessageChunk( content="", additional_kwargs={"reasoning_content": "Second thought"}, @@ -645,3 +650,117 @@ def test_error_closes_open_text_block(self): ends = [e for e in events if isinstance(e, TextMessageEndEvent)] assert len(ends) == 1 assert isinstance(events[-1], RunErrorEvent) + + +# ── Regression: cross-block accumulation (DeepSeek style) ──────────────── + + +class TestCrossBlockAccumulation: + """When DeepSeek re-iterates prior reasoning across tool-call boundaries. + + The handler must NOT reset ``_last_reasoning_content`` on block close, + otherwise the diff on the next block's first chunk emits old content as + a delta, producing ever-growing reasoning blocks:: + + Block 1: "foo" → delta "foo" ✓ + Tool result → close reasoning + Block 2: "foo bar" → delta " bar" ✓ (was: "foo bar" ❌) + Tool result → close reasoning + Block 3: "foo bar baz" → delta " baz" ✓ (was: "foo bar baz" ❌) + """ + + @pytest.fixture + def handler(self): + h = StreamEventHandler( + thread_id="th-1", run_id="run-1", message_id="msg-1", + ) + h.drain() # discard RUN_STARTED + return h + + def test_cross_block_accumulated_delta_strips_old_content(self, handler): + """Block 2 with repeated prefix emits only the new portion.""" + # Block 1: "foo" + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": "foo"}), + {"langgraph_node": "agent"}, + ) + handler.drain() + + # Close reasoning via tool result + handler.observe( + AIMessageChunk(content="", tool_call_chunks=[{"id": "c1", "name": "t", "args": "{}", "index": 0}]), + {"langgraph_node": "agent"}, + ) + handler.drain() + handler.observe(ToolMessage(content="ok", tool_call_id="c1"), {"langgraph_node": "tools"}) + handler.drain() + + # Block 2: DeepSeek re-iterates "foo" + adds " bar" + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": "foo bar"}), + {"langgraph_node": "agent"}, + ) + events = handler.drain() + content_events = [e for e in events if isinstance(e, ReasoningMessageContentEvent)] + assert len(content_events) == 1 + assert content_events[0].delta == " bar", ( + f"Expected delta ' bar', got {content_events[0].delta!r}" + ) + + def test_cross_block_three_blocks(self, handler): + """Three blocks: each emits only the new portion.""" + def _block(text: str) -> None: + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": text}), + {"langgraph_node": "agent"}, + ) + handler.drain() + # Close via tool result + handler.observe( + AIMessageChunk(content="", tool_call_chunks=[{"id": "c1", "name": "t", "args": "{}", "index": 0}]), + {"langgraph_node": "agent"}, + ) + handler.drain() + handler.observe(ToolMessage(content="ok", tool_call_id="c1"), {"langgraph_node": "tools"}) + handler.drain() + + _block("foo") + _block("foo bar") + _block("foo bar baz") + + # Fourth block: verify the chain stops growing + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": "foo bar baz qux"}), + {"langgraph_node": "agent"}, + ) + events = handler.drain() + content_events = [e for e in events if isinstance(e, ReasoningMessageContentEvent)] + assert len(content_events) == 1 + assert content_events[0].delta == " qux" + + def test_cross_block_delta_style_still_works(self, handler): + """OpenAI-style non-overlapping deltas still work across blocks.""" + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": "First block"}), + {"langgraph_node": "agent"}, + ) + handler.drain() + + # Close via tool result + handler.observe( + AIMessageChunk(content="", tool_call_chunks=[{"id": "c1", "name": "t", "args": "{}", "index": 0}]), + {"langgraph_node": "agent"}, + ) + handler.drain() + handler.observe(ToolMessage(content="ok", tool_call_id="c1"), {"langgraph_node": "tools"}) + handler.drain() + + # New block with completely different content (delta style) + handler.observe( + AIMessageChunk(content="", additional_kwargs={"reasoning_content": "Second block"}), + {"langgraph_node": "agent"}, + ) + events = handler.drain() + content_events = [e for e in events if isinstance(e, ReasoningMessageContentEvent)] + assert len(content_events) == 1 + assert content_events[0].delta == "Second block"