fix reasoning accumulation

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),

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"