feat: convert StreamHandler to callback-based (remove Rich.Live)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

app/services/streaming.py

@@ -1,41 +1,56 @@
 """Streaming response handler — accumulates SSE chunks into a complete Message."""
 
-from collections.abc import AsyncIterator
+import time
-
+from collections.abc import AsyncIterator, Callable
-from rich.live import Live
-from rich.markdown import Markdown
-from rich.panel import Panel
 
 from app.models.config import DisplayConfig
 from app.models.message import Message
 from app.models.tool_call import ToolCall, ToolCallFunction
-from app.utils.logging import console, get_logger
+from app.utils.logging import get_logger
 from app.utils.token_counter import TokenUsage
 
 logger = get_logger(__name__)
 
+# Minimum interval between content update callbacks (seconds)
+_UPDATE_THROTTLE_INTERVAL = 0.1
+
 
 class StreamHandler:
-    """Processes an SSE chunk stream into a Rich live display and final Message.
+    """Processes an SSE chunk stream and produces a complete assistant Message.
 
-    Accumulates content deltas and tool call fragments, renders a live Markdown
+    Accumulates content deltas and tool call fragments. Notifies the UI via
-    panel during streaming, and produces a complete assistant Message on finish.
+    optional callbacks during streaming.
     """
 
     def __init__(self, display_config: DisplayConfig) -> None:
-        """Initialize the stream handler.
-
-        Args:
-            display_config: Display preferences (streaming toggle, etc.).
-        """
         self._display_config = display_config
         self._accumulated_content: str = ""
         self._accumulated_reasoning: str = ""
         self._tool_calls: dict[int, dict[str, str]] = {}
         self._usage: TokenUsage | None = None
+        self._on_content: Callable[[str], None] | None = None
+        self._on_thinking: Callable[[], None] | None = None
+        self._on_done: Callable[[], None] | None = None
+
+    def set_callbacks(
+        self,
+        on_content: Callable[[str], None] | None = None,
+        on_thinking: Callable[[], None] | None = None,
+        on_done: Callable[[], None] | None = None,
+    ) -> None:
+        """Set UI callbacks for streaming updates.
+
+        Args:
+            on_content: Called with accumulated content string (throttled to ~100ms).
+            on_thinking: Called once when first reasoning token arrives.
+            on_done: Called when streaming is complete.
+        """
+        self._on_content = on_content
+        self._on_thinking = on_thinking
+        self._on_done = on_done
 
     async def process_stream(self, chunk_iter: AsyncIterator[dict]) -> Message:
-        """Consume a chunk iterator, rendering live output and returning the final Message.
+        """Consume a chunk iterator and return the final Message.
 
         Args:
             chunk_iter: Async iterator of parsed SSE chunk dicts.
@@ -43,26 +58,42 @@ class StreamHandler:
         Returns:
             Complete assistant Message with accumulated content and tool calls.
         """
-        with Live(console=console, refresh_per_second=8) as live:
+        thinking_notified = False
-            async for chunk in chunk_iter:
+        last_update_time = 0.0
-                self._process_chunk(chunk)
 
-                # Show reasoning while waiting for content
+        async for chunk in chunk_iter:
-                display_text = self._accumulated_content
+            self._process_chunk(chunk)
-                if not display_text and self._accumulated_reasoning:
-                    display_text = "*thinking...*"
 
-                if display_text and self._display_config.stream_output:
+            if not self._display_config.stream_output:
-                    # Render inside the same Assistant panel used for final output
+                continue
-                    # so the live display and final frame are visually consistent
+
-                    live.update(
+            # Notify thinking once
-                        Panel(
+            if (
-                            Markdown(display_text),
+                not thinking_notified
-                            title="Assistant",
+                and not self._accumulated_content
-                            border_style="green",
+                and self._accumulated_reasoning
-                            expand=True,
+                and self._on_thinking is not None
-                        )
+            ):
-                    )
+                self._on_thinking()
+                thinking_notified = True
+
+            # Throttled content updates
+            if self._accumulated_content and self._on_content is not None:
+                now = time.monotonic()
+                if now - last_update_time >= _UPDATE_THROTTLE_INTERVAL:
+                    self._on_content(self._accumulated_content)
+                    last_update_time = now
+
+        # Final content update (ensures last chunk is shown)
+        if (
+            self._display_config.stream_output
+            and self._accumulated_content
+            and self._on_content is not None
+        ):
+            self._on_content(self._accumulated_content)
+
+        if self._on_done is not None:
+            self._on_done()
 
         tool_calls = self._build_tool_calls() or None
         return Message(
@@ -72,12 +103,7 @@ class StreamHandler:
         )
 
     def _process_chunk(self, chunk: dict) -> None:
-        """Extract content, tool calls, and usage from a single SSE chunk.
+        """Extract content, tool calls, and usage from a single SSE chunk."""
-
-        Args:
-            chunk: Parsed JSON dict from one SSE data line.
-        """
-        # Content delta
         choices = chunk.get("choices", [])
         if choices:
             delta = choices[0].get("delta", {})
@@ -86,12 +112,10 @@ class StreamHandler:
             if content_piece:
                 self._accumulated_content += content_piece
 
-            # Reasoning tokens (e.g. qwen3.5 thinking mode)
             reasoning_piece = delta.get("reasoning")
             if reasoning_piece:
                 self._accumulated_reasoning += reasoning_piece
 
-            # Tool call deltas (accumulated by index)
             for tc_delta in delta.get("tool_calls", []):
                 idx = tc_delta.get("index", 0)
                 if idx not in self._tool_calls:
@@ -109,7 +133,6 @@ class StreamHandler:
                 if func.get("arguments"):
                     entry["arguments"] += func["arguments"]
 
-        # Token usage (typically in the final chunk)
         usage_data = chunk.get("usage")
         if usage_data:
             self._usage = TokenUsage(
@@ -119,11 +142,7 @@ class StreamHandler:
             )
 
     def _build_tool_calls(self) -> list[ToolCall]:
-        """Convert accumulated tool call fragments into sorted ToolCall list.
+        """Convert accumulated tool call fragments into sorted ToolCall list."""
-
-        Returns:
-            List of ToolCall objects sorted by stream index.
-        """
         if not self._tool_calls:
             return []
 
@@ -143,14 +162,7 @@ class StreamHandler:
         return result
 
     def get_partial_message(self) -> Message | None:
-        """Return whatever content/tool_calls have been accumulated so far.
+        """Return whatever content/tool_calls have been accumulated so far."""
-
-        Useful for mid-stream interruption recovery — returns None if nothing
-        has been accumulated yet.
-
-        Returns:
-            Partial assistant Message, or None if no content accumulated.
-        """
         tool_calls = self._build_tool_calls() or None
         if not self._accumulated_content and not tool_calls:
             return None
@@ -176,3 +188,6 @@ class StreamHandler:
         self._accumulated_reasoning = ""
         self._tool_calls.clear()
         self._usage = None
+        self._on_content = None
+        self._on_thinking = None
+        self._on_done = None

tests/unit/test_streaming.py

@@ -0,0 +1,111 @@
+"""Tests for callback-based StreamHandler."""
+
+import asyncio
+from unittest.mock import MagicMock, call
+
+import pytest
+
+from app.models.config import DisplayConfig
+from app.services.streaming import StreamHandler
+
+
+def _make_chunk(content: str | None = None, reasoning: str | None = None) -> dict:
+    """Helper to create a fake SSE chunk."""
+    delta: dict = {}
+    if content is not None:
+        delta["content"] = content
+    if reasoning is not None:
+        delta["reasoning"] = reasoning
+    return {"choices": [{"delta": delta}]}
+
+
+def _make_tool_call_chunk(index: int, tc_id: str = "", name: str = "", args: str = "") -> dict:
+    """Helper to create a fake tool call chunk."""
+    tc_delta: dict = {"index": index}
+    if tc_id:
+        tc_delta["id"] = tc_id
+    func: dict = {}
+    if name:
+        func["name"] = name
+    if args:
+        func["arguments"] = args
+    if func:
+        tc_delta["function"] = func
+    return {"choices": [{"delta": {"tool_calls": [tc_delta]}}]}
+
+
+async def _async_iter(items: list[dict]):
+    for item in items:
+        yield item
+
+
+class TestStreamHandlerCallbacks:
+    @pytest.mark.asyncio
+    async def test_on_content_called_with_accumulated_text(self) -> None:
+        handler = StreamHandler(DisplayConfig(stream_output=True))
+        on_content = MagicMock()
+        on_thinking = MagicMock()
+        on_done = MagicMock()
+        handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
+
+        chunks = [_make_chunk(content="Hello"), _make_chunk(content=" world")]
+        msg = await handler.process_stream(_async_iter(chunks))
+
+        assert msg.content == "Hello world"
+        assert on_content.call_count >= 1
+        # Last call should have full accumulated content
+        last_content = on_content.call_args_list[-1][0][0]
+        assert last_content == "Hello world"
+        on_done.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_on_thinking_called_for_reasoning(self) -> None:
+        handler = StreamHandler(DisplayConfig(stream_output=True))
+        on_content = MagicMock()
+        on_thinking = MagicMock()
+        on_done = MagicMock()
+        handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
+
+        chunks = [_make_chunk(reasoning="let me think")]
+        msg = await handler.process_stream(_async_iter(chunks))
+
+        on_thinking.assert_called_once()
+        on_content.assert_not_called()
+        on_done.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_display_callbacks_skip_when_stream_output_disabled(self) -> None:
+        """on_content and on_thinking are suppressed, but on_done always fires."""
+        handler = StreamHandler(DisplayConfig(stream_output=False))
+        on_content = MagicMock()
+        on_thinking = MagicMock()
+        on_done = MagicMock()
+        handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
+
+        chunks = [_make_chunk(content="Hello")]
+        msg = await handler.process_stream(_async_iter(chunks))
+
+        assert msg.content == "Hello"
+        on_content.assert_not_called()
+        on_thinking.assert_not_called()
+        on_done.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_no_callbacks_by_default(self) -> None:
+        """process_stream works without set_callbacks (backward compat)."""
+        handler = StreamHandler(DisplayConfig(stream_output=True))
+        chunks = [_make_chunk(content="Hello")]
+        msg = await handler.process_stream(_async_iter(chunks))
+        assert msg.content == "Hello"
+
+    @pytest.mark.asyncio
+    async def test_tool_calls_still_accumulated(self) -> None:
+        handler = StreamHandler(DisplayConfig(stream_output=True))
+        chunks = [
+            _make_tool_call_chunk(0, tc_id="call_1", name="read_file"),
+            _make_tool_call_chunk(0, args='{"path": "foo.py"}'),
+        ]
+        msg = await handler.process_stream(_async_iter(chunks))
+        assert msg.tool_calls is not None
+        assert len(msg.tool_calls) == 1
+        assert msg.tool_calls[0].function.name == "read_file"