feat(ui): plumb reasoning stream + output-token count to the UI seam

docs/DESIGN.md

@@ -2405,7 +2405,7 @@ Most rows here concern the v2 memory system. The prompt-injection and secret row
 | Model emits malformed tool call | Reject tool call, show compact error, and allow one retry. |
 | Model loops tool calls | Enforce turn/tool budgets and replan or stop. |
 | Reasoning mode unavailable | Per-model fallback: add the model to `_no_think`; retry once without `think`; other models keep sending `think`. `_reasoning` (the config-level flag) is never mutated. |
-| Reasoning-only turn (think text, empty content) | The streamed `thinking` field is now accumulated and captured on the reply message (it is never echoed back to the API and never rendered), so a turn that reasons and then emits nothing is observable in the audit log rather than silently empty; the runtime nudges such a reply once the model has already run a tool (section 10.4). |
+| Reasoning-only turn (think text, empty content) | The streamed `thinking` field is accumulated and captured on the reply message. It is never echoed back to the API; it may be surfaced live to the UI when a consumer is wired via `on_thinking` (the default `TerminalUI` is a no-op). A turn that reasons then emits nothing is observable in the audit log; the runtime nudges such a reply once the model has already run a tool (section 10.4). |
 | Stream ends without `done` sentinel | Rejected as incomplete: `_stream_chat` raises `OllamaResponseError` after the `with` block closes without seeing a chunk whose `done` field is truthy. Legitimate early stops (`done_reason = "length"` or `"load"`) still carry the sentinel and are accepted normally. |
 
 ### 24.7 Privacy And Log Edge Cases

shellpilot/cli/terminal.py

@@ -213,6 +213,14 @@ def stream_token(self, token: str) -> None:
         self._spinner.stop()
         self._stream.feed(token)
 
+    def stream_thinking(self, text: str) -> None:
+        """No-op: the current terminal UI does not surface reasoning text.
+
+        A later full-screen UI consumes this hook; keeping it silent here keeps
+        a default session byte-identical.
+        """
+        return
+
     def show_status(self, text: str) -> None:
         self._console.print(f"[sp.dim]{escape(_sanitize_line(text))}[/sp.dim]")
 

shellpilot/llm/client.py

@@ -23,6 +23,7 @@ def chat(
         num_ctx: int,
         options: dict[str, Any] | None = None,
         on_token: TokenCallback | None = None,
+        on_thinking: TokenCallback | None = None,
     ) -> Message: ...
 
     def health(self) -> bool: ...

shellpilot/llm/messages.py

@@ -44,9 +44,12 @@ class Message:
     tool_calls: tuple[ToolCall, ...] = ()
     images: tuple[ImageRef, ...] = ()
     # Streamed reasoning text from a think-capable model, captured for audit
-    # observability only. Never echoed back to the API (see ollama._encode_message)
-    # and never rendered; kept default-empty so existing constructions stand.
+    # observability only. Never echoed back to the API (see ollama._encode_message);
+    # may be surfaced live to the UI when a consumer is wired (see on_thinking).
+    # Kept default-empty so existing constructions stand.
     thinking: str = ""
+    # Total output tokens for this call, from Ollama's `eval_count`; 0 if absent.
+    output_tokens: int = 0
 
 
 @dataclass(frozen=True)

shellpilot/llm/ollama.py

@@ -198,6 +198,7 @@ def chat(
         num_ctx: int,
         options: dict[str, Any] | None = None,
         on_token: Callable[[str], None] | None = None,
+        on_thinking: Callable[[str], None] | None = None,
     ) -> Message:
         """Stream one chat completion; num_ctx is set explicitly on every request.
 
@@ -215,23 +216,27 @@ def chat(
         if self._reasoning and model not in self._no_think:
             payload["think"] = True
         try:
-            return self._stream_chat(payload, on_token)
+            return self._stream_chat(payload, on_token, on_thinking)
         except OllamaResponseError as exc:
             # Reasoning mode unavailable for this model: cache and retry once without
             # think (design section 24.6). Other models are not affected.
             if self._reasoning and model not in self._no_think and "think" in str(exc).lower():
                 self._no_think.add(model)
                 payload.pop("think", None)
-                return self._stream_chat(payload, on_token)
+                return self._stream_chat(payload, on_token, on_thinking)
             raise
 
     def _stream_chat(
-        self, payload: dict[str, Any], on_token: Callable[[str], None] | None
+        self,
+        payload: dict[str, Any],
+        on_token: Callable[[str], None] | None,
+        on_thinking: Callable[[str], None] | None = None,
     ) -> Message:
         content_parts: list[str] = []
         thinking_parts: list[str] = []
         tool_calls: list[ToolCall] = []
         done_seen = False
+        output_tokens = 0
         try:
             with self._client.stream("POST", "/api/chat", json=payload) as response:
                 if response.status_code >= 400:
@@ -250,6 +255,8 @@ def _stream_chat(
                         raise OllamaResponseError(str(chunk["error"]))
                     if chunk.get("done"):
                         done_seen = True
+                        raw_count = chunk.get("eval_count")
+                        output_tokens = raw_count if isinstance(raw_count, int) else 0
                     message = chunk.get("message") or {}
                     if not isinstance(message, dict):
                         raise OllamaResponseError(f"unexpected stream message shape: {line[:200]}")
@@ -260,10 +267,13 @@ def _stream_chat(
                             on_token(token)
                     # Reasoning text streams in a separate field; capture it so a
                     # reasoning-only turn is observable instead of silently empty
-                    # (design section 24.6). It is never streamed to the UI.
+                    # (design section 24.6). May be streamed to the UI via on_thinking
+                    # when a consumer is wired; never echoed back to the API.
                     thinking = message.get("thinking") or ""
                     if thinking:
                         thinking_parts.append(thinking)
+                        if on_thinking is not None:
+                            on_thinking(thinking)
                     raw_calls = message.get("tool_calls") or []
                     if not isinstance(raw_calls, list):
                         raise OllamaResponseError(f"unexpected tool_calls shape: {line[:200]}")
@@ -280,6 +290,7 @@ def _stream_chat(
             content="".join(content_parts),
             tool_calls=tuple(tool_calls),
             thinking="".join(thinking_parts),
+            output_tokens=output_tokens,
         )
 
 

shellpilot/runtime/conversation.py

@@ -155,6 +155,7 @@ def __init__(
         self._staged_tool_images: list[ImageRef] = []
         self._last_user_text = ""
         self._last_failure_signature: str | None = None
+        self._turn_output_tokens: int = 0
         self.snapshots = SnapshotStore()
         self.recent_diffs: list[str] = []
         self.plan_manager = PlanManager(workspace, settings.runtime.security_profile)
@@ -481,6 +482,7 @@ def run_turn(self, text: str, *, images: Sequence[ImageRef] = ()) -> str:
         # Belt-and-braces: a stale stage left by an aborted prior turn must not
         # attach to this unrelated turn's first message.
         self._staged_tool_images.clear()
+        self._turn_output_tokens = 0
         if estimate_tokens(text) > self.budget.max_user_message_tokens:
             self._ui.show_status(
                 "Message too large for the model context "
@@ -526,6 +528,7 @@ def _turn_stats(self, elapsed_s: float) -> TurnStats:
             context_tokens=used,
             context_pct=pct,
             warn=used > self.budget.compact_at_tokens,
+            output_tokens=self._turn_output_tokens,
         )
 
     def _pending_plan_step(self) -> tuple[int, str] | None:
@@ -610,9 +613,11 @@ def _tool_loop(self) -> Message:
                     num_ctx=self.budget.model_context_tokens,
                     options=self._settings.model.options,
                     on_token=self._ui.stream_token,
+                    on_thinking=self._ui.stream_thinking,
                 )
             finally:
                 self._ui.end_response()
+            self._turn_output_tokens += reply.output_tokens
             self._record(reply)
             if not reply.tool_calls:
                 pending = self._pending_plan_step()

shellpilot/runtime/events.py

@@ -18,13 +18,16 @@ class TurnStats:
     context_tokens: int
     context_pct: int
     warn: bool
+    output_tokens: int
 
 
 class RuntimeUI(Protocol):
     """What the conversation runtime needs from a user interface."""
 
     def stream_token(self, token: str) -> None: ...
 
+    def stream_thinking(self, text: str) -> None: ...
+
     def begin_response(self) -> None:
         """The runtime is about to call the model (start a waiting indicator)."""
         ...

tests/fakes/fake_llm.py

@@ -60,6 +60,7 @@ def chat(
         num_ctx: int,
         options: dict[str, Any] | None = None,
         on_token: TokenCallback | None = None,
+        on_thinking: TokenCallback | None = None,
     ) -> Message:
         self.calls.append(
             RecordedCall(
@@ -76,6 +77,9 @@ def chat(
         if on_token is not None and reply.content:
             for chunk in _chunks(reply.content, 8):
                 on_token(chunk)
+        if on_thinking is not None and reply.thinking:
+            for chunk in _chunks(reply.thinking, 8):
+                on_thinking(chunk)
         return reply
 
     def health(self) -> bool:

tests/fakes/fake_ui.py

@@ -16,6 +16,7 @@
 @dataclass
 class FakeUI:
     tokens: list[str] = field(default_factory=list)
+    thinking_fragments: list[str] = field(default_factory=list)
     began: int = 0
     ended: int = 0
     turn_stats: list[TurnStats] = field(default_factory=list)
@@ -36,6 +37,9 @@ class FakeUI:
     def stream_token(self, token: str) -> None:
         self.tokens.append(token)
 
+    def stream_thinking(self, text: str) -> None:
+        self.thinking_fragments.append(text)
+
     def begin_response(self) -> None:
         self.began += 1
 

tests/test_conversation.py

@@ -1194,3 +1194,28 @@ def test_build_system_prompt_egressing_flag() -> None:
     assert "no independent network access" not in remote
     assert "entirely on this machine" not in remote
     assert "leaves this device" in remote
+
+
+# ---------------------------------------------------------------------------
+# output_tokens accumulation across multi-call tool loop (thinking-stream plumbing)
+# ---------------------------------------------------------------------------
+
+
+def test_turn_stats_output_tokens_accumulated(tmp_path: Path) -> None:
+    """TurnStats.output_tokens is the sum of output_tokens across all chat() calls in a turn.
+
+    Uses two empty-content replies to trigger the empty-reply nudge path (the
+    only multi-call-per-turn route that needs no tool registry or approvals).
+    """
+    import dataclasses
+
+    from tests.fakes.fake_llm import answer
+
+    first = dataclasses.replace(answer(""), output_tokens=100)  # empty → triggers nudge
+    second = dataclasses.replace(answer("done"), output_tokens=50)  # real answer
+    fake = FakeLLM(script=[first, second])
+    ui = FakeUI()
+    runtime = make_runtime(fake, ui, tmp_path)
+    runtime.run_turn("do something")
+    assert len(ui.turn_stats) == 1
+    assert ui.turn_stats[0].output_tokens == 150

tests/test_ollama_chat.py

@@ -245,3 +245,112 @@ def handler(request: httpx.Request) -> httpx.Response:
 
     client = OllamaClient(transport=httpx.MockTransport(handler))
     assert client.model_context_length("missing") is None
+
+
+# ---------------------------------------------------------------------------
+# on_thinking callback and eval_count → output_tokens (thinking-stream plumbing)
+# ---------------------------------------------------------------------------
+
+
+def test_on_thinking_fires_callback() -> None:
+    """Thinking fragments forwarded to on_thinking in order; on_token still fires for content."""
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=stream_body(
+                {"message": {"role": "assistant", "thinking": "Hmm ", "content": ""}},
+                {"message": {"role": "assistant", "thinking": "let me think.", "content": ""}},
+                {
+                    "message": {"role": "assistant", "thinking": "", "content": "Answer"},
+                    "done": False,
+                },
+                {"message": {"role": "assistant", "content": ""}, "done": True},
+            ),
+        )
+
+    thinking_chunks: list[str] = []
+    content_tokens: list[str] = []
+    client = OllamaClient(transport=httpx.MockTransport(handler))
+    reply = client.chat(
+        "gemma4:e4b",
+        [user("hi")],
+        num_ctx=4096,
+        on_token=content_tokens.append,
+        on_thinking=thinking_chunks.append,
+    )
+    assert thinking_chunks == ["Hmm ", "let me think."]
+    assert content_tokens == ["Answer"]
+    assert reply.thinking == "Hmm let me think."
+    assert reply.content == "Answer"
+
+
+def test_on_thinking_none_does_not_raise() -> None:
+    """on_thinking=None (the default) leaves the existing stream behavior unchanged."""
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=stream_body(
+                {
+                    "message": {"role": "assistant", "thinking": "thinking...", "content": ""},
+                    "done": True,
+                },
+            ),
+        )
+
+    client = OllamaClient(transport=httpx.MockTransport(handler))
+    reply = client.chat("gemma4:e4b", [user("hi")], num_ctx=4096)
+    assert reply.thinking == "thinking..."
+
+
+def test_eval_count_sets_output_tokens() -> None:
+    """A done chunk carrying eval_count sets Message.output_tokens to that integer."""
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=stream_body(
+                {"message": {"role": "assistant", "content": "hi"}, "done": True, "eval_count": 42},
+            ),
+        )
+
+    client = OllamaClient(transport=httpx.MockTransport(handler))
+    reply = client.chat("gemma4:e4b", [user("hello")], num_ctx=4096)
+    assert reply.output_tokens == 42
+
+
+def test_eval_count_absent_gives_zero() -> None:
+    """A done chunk without eval_count yields output_tokens == 0."""
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=stream_body(
+                {"message": {"role": "assistant", "content": "hi"}, "done": True},
+            ),
+        )
+
+    client = OllamaClient(transport=httpx.MockTransport(handler))
+    reply = client.chat("gemma4:e4b", [user("hello")], num_ctx=4096)
+    assert reply.output_tokens == 0
+
+
+def test_eval_count_non_int_gives_zero() -> None:
+    """A non-int eval_count in the done chunk yields output_tokens == 0 (no raise)."""
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=stream_body(
+                {
+                    "message": {"role": "assistant", "content": "hi"},
+                    "done": True,
+                    "eval_count": "not-an-int",
+                },
+            ),
+        )
+
+    client = OllamaClient(transport=httpx.MockTransport(handler))
+    reply = client.chat("gemma4:e4b", [user("hello")], num_ctx=4096)
+    assert reply.output_tokens == 0

tests/test_terminal_ui.py

@@ -938,3 +938,21 @@ def read(self, context: object) -> str:
         json.loads(line) for line in (fake_paths.state_dir / "audit.jsonl").read_text().splitlines()
     ]
     assert any(e["event"] == "session_resume" for e in events), "resume load must still happen"
+
+
+# ---------------------------------------------------------------------------
+# stream_thinking no-op on TerminalUI (thinking-stream plumbing)
+# ---------------------------------------------------------------------------
+
+
+def test_stream_thinking_is_noop() -> None:
+    """TerminalUI.stream_thinking produces no console output.
+
+    Default session stays byte-identical: thinking is not echoed to the terminal.
+    """
+    console = make_console()
+    ui = TerminalUI(console, glyphs=GLYPHS, spinner=False)
+    console.begin_capture()
+    ui.stream_thinking("I am reasoning about this…")
+    captured = console.end_capture()
+    assert captured == ""