From bef2d0eceb948e4b34f7ebe99f3d6a8b65601576 Mon Sep 17 00:00:00 2001 From: Lavindeep Dhillon Date: Sat, 27 Jun 2026 18:02:14 -0400 Subject: [PATCH] feat(ui): plumb reasoning stream + output-token count to the UI seam Add an on_thinking callback through OllamaClient.chat/_stream_chat and the LLMClient/RuntimeUI protocols, mirroring the existing on_token path, and parse Ollama's eval_count into Message.output_tokens. The runtime accumulates output tokens per turn and surfaces the total in TurnStats. The current TerminalUI consumes the new hook as a no-op, so a default session stays byte-identical; a later full-screen UI renders the live reasoning readout. --- docs/DESIGN.md | 2 +- shellpilot/cli/terminal.py | 8 +++ shellpilot/llm/client.py | 1 + shellpilot/llm/messages.py | 7 +- shellpilot/llm/ollama.py | 19 +++-- shellpilot/runtime/conversation.py | 5 ++ shellpilot/runtime/events.py | 3 + tests/fakes/fake_llm.py | 4 ++ tests/fakes/fake_ui.py | 4 ++ tests/test_conversation.py | 25 +++++++ tests/test_ollama_chat.py | 109 +++++++++++++++++++++++++++++ tests/test_terminal_ui.py | 18 +++++ 12 files changed, 198 insertions(+), 7 deletions(-) diff --git a/docs/DESIGN.md b/docs/DESIGN.md index 8cc4af5..7fe1235 100644 --- a/docs/DESIGN.md +++ b/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 diff --git a/shellpilot/cli/terminal.py b/shellpilot/cli/terminal.py index 2e8e4de..0003351 100644 --- a/shellpilot/cli/terminal.py +++ b/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]") diff --git a/shellpilot/llm/client.py b/shellpilot/llm/client.py index f964e18..57d6ee9 100644 --- a/shellpilot/llm/client.py +++ b/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: ... diff --git a/shellpilot/llm/messages.py b/shellpilot/llm/messages.py index 860806d..e8286a1 100644 --- a/shellpilot/llm/messages.py +++ b/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) diff --git a/shellpilot/llm/ollama.py b/shellpilot/llm/ollama.py index 1bbe006..4062658 100644 --- a/shellpilot/llm/ollama.py +++ b/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, ) diff --git a/shellpilot/runtime/conversation.py b/shellpilot/runtime/conversation.py index f3471b0..459744b 100644 --- a/shellpilot/runtime/conversation.py +++ b/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() diff --git a/shellpilot/runtime/events.py b/shellpilot/runtime/events.py index c7d2c5f..3bc7098 100644 --- a/shellpilot/runtime/events.py +++ b/shellpilot/runtime/events.py @@ -18,6 +18,7 @@ class TurnStats: context_tokens: int context_pct: int warn: bool + output_tokens: int class RuntimeUI(Protocol): @@ -25,6 +26,8 @@ class RuntimeUI(Protocol): 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).""" ... diff --git a/tests/fakes/fake_llm.py b/tests/fakes/fake_llm.py index 18c7be9..c505a26 100644 --- a/tests/fakes/fake_llm.py +++ b/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: diff --git a/tests/fakes/fake_ui.py b/tests/fakes/fake_ui.py index a9d9a17..a484124 100644 --- a/tests/fakes/fake_ui.py +++ b/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 diff --git a/tests/test_conversation.py b/tests/test_conversation.py index d021cc8..9e6b07b 100644 --- a/tests/test_conversation.py +++ b/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 diff --git a/tests/test_ollama_chat.py b/tests/test_ollama_chat.py index ddfa873..c487f27 100644 --- a/tests/test_ollama_chat.py +++ b/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 diff --git a/tests/test_terminal_ui.py b/tests/test_terminal_ui.py index 4e0ff84..6860fb5 100644 --- a/tests/test_terminal_ui.py +++ b/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 == ""