diff --git a/src/repl/core.py b/src/repl/core.py index 4cfd73182..edbb862fe 100644 --- a/src/repl/core.py +++ b/src/repl/core.py @@ -685,11 +685,24 @@ def _bottom_toolbar(self): from src.utils.advisor import format_advisor_status advisor_seg = format_advisor_status(self.provider, model) advisor_part = f" {advisor_seg} ·" if advisor_seg else "" + # Advisor token counts — accumulated on the ToolContext + # by ``src/tool_system/tools/advisor.py`` per consultation. + # Surface them next to the worker's counts so the user can + # see how much of the spend went to the reviewer model. + # Hidden when zero so the toolbar stays compact for users + # who haven't enabled the advisor. + adv_in = int(getattr(self.tool_context, "advisor_input_tokens", 0) or 0) + adv_out = int(getattr(self.tool_context, "advisor_output_tokens", 0) or 0) + advisor_tokens = ( + f" (advisor: {adv_in} in / {adv_out} out)" + if (adv_in or adv_out) else "" + ) return ( f" {provider} · {model} · {cwd} ·{advisor_part} " f"turns: {self._stats_turns} · " f"tokens: {self._stats_input_tokens} in / " - f"{self._stats_output_tokens} out " + f"{self._stats_output_tokens} out" + f"{advisor_tokens} " ) except Exception: # Never let the toolbar break the input prompt. diff --git a/src/tool_system/context.py b/src/tool_system/context.py index e5401c51b..71bb857e8 100644 --- a/src/tool_system/context.py +++ b/src/tool_system/context.py @@ -104,6 +104,15 @@ class ToolContext: # block without I/O). tool_result_chars_so_far: int = 0 _aggregate_lock: threading.Lock = field(default_factory=threading.Lock) + # Session-cumulative tokens spent on client-side advisor calls. + # ``src/tool_system/tools/advisor.py`` accumulates here on every + # consultation; the REPL bottom_toolbar + TUI StatusLine read + # them to display ``advisor: /`` next to the worker's + # token counts. Distinct from ``tool_result_chars_so_far`` (which + # is a per-message budget tied to API-result persistence) — these + # are per-session totals for UI display. + advisor_input_tokens: int = 0 + advisor_output_tokens: int = 0 # Chapter-10 / Chunk F / WI-6.1 — agent-name registry. Maps the # human-readable ``name`` (passed via Agent({name: "researcher"})) # to the random ``agent_id`` returned by the spawn. SendMessage diff --git a/src/tool_system/tools/advisor.py b/src/tool_system/tools/advisor.py index 805808702..d36a8aded 100644 --- a/src/tool_system/tools/advisor.py +++ b/src/tool_system/tools/advisor.py @@ -63,12 +63,26 @@ def _advisor_call(tool_input: dict[str, Any], context: ToolContext) -> ToolResul # the proxy and hit api.anthropic.com directly. main_provider = getattr(context, "_active_provider", None) - ok, text = execute_client_advisor( + ok, text, usage = execute_client_advisor( advisor_model, forwarded, abort_signal=abort_signal, main_provider=main_provider, ) + + # Accumulate the advisor's tokens onto the ToolContext so the + # REPL / TUI status bar can display them next to the worker's. + # Counters are added lazily via getattr so a long-lived context + # that was constructed before this field existed still works. + try: + prior_in = int(getattr(context, "advisor_input_tokens", 0) or 0) + prior_out = int(getattr(context, "advisor_output_tokens", 0) or 0) + context.advisor_input_tokens = prior_in + int(usage.get("input_tokens", 0)) + context.advisor_output_tokens = prior_out + int(usage.get("output_tokens", 0)) + except Exception: + # Status-bar bookkeeping must never break the tool result. + pass + return ToolResult( name="advisor", output=text, diff --git a/src/tui/agent_bridge.py b/src/tui/agent_bridge.py index a9653b134..437643300 100644 --- a/src/tui/agent_bridge.py +++ b/src/tui/agent_bridge.py @@ -317,6 +317,27 @@ def _persist(msg: Any) -> None: ) except Exception: pass + # Mirror the client-side advisor token counts from tool_context + # into ``state.usage`` so the StatusLine widget can surface them + # next to the worker tokens. ``tool_context.advisor_*`` is + # accumulated by ``AdvisorTool._advisor_call`` (PR #190) on every + # client-side consultation; this is the TUI-side parity of the + # legacy REPL's ``_bottom_toolbar`` reading the same ctx fields. + # NOTE: server-side advisor attribution (Anthropic API's + # ``usage.iterations[]`` with ``type="advisor_message"``) is NOT + # currently surfaced — the Python SDK 0.88.0 we depend on only + # defines ``"message"`` and ``"compaction"`` discriminators + # (see anthropic/types/beta/beta_message_iteration_usage.py). + # Add a separate accumulator + parser when the SDK gains the + # advisor discriminator OR when we move to direct HTTP parsing. + try: + adv_in = int(getattr(self._tool_context, "advisor_input_tokens", 0) or 0) + adv_out = int(getattr(self._tool_context, "advisor_output_tokens", 0) or 0) + if adv_in or adv_out: + self._state.usage["advisor_input_tokens"] = adv_in + self._state.usage["advisor_output_tokens"] = adv_out + except Exception: + pass self._post( AgentRunFinished( response_text=result.response_text, diff --git a/src/tui/widgets/status_line.py b/src/tui/widgets/status_line.py index 5d129a75b..7eaba0813 100644 --- a/src/tui/widgets/status_line.py +++ b/src/tui/widgets/status_line.py @@ -168,6 +168,18 @@ def _compose_text(self, *, spinner: str) -> Text: total = state.usage.get("input_tokens", 0) + state.usage.get("output_tokens", 0) if total: right_bits.append(f"tokens {total}") + # Advisor token segment — appears next to worker tokens + # whenever the advisor has been consulted this session. + # ``state.usage["advisor_*"]`` is mirrored from + # ``tool_context.advisor_*`` by ``agent_bridge.py`` after + # each run; the underlying ctx counter is accumulated by + # ``AdvisorTool._advisor_call`` on every consultation. + # Hidden when zero so the bar stays compact for users who + # haven't enabled the advisor yet. + adv_in = state.usage.get("advisor_input_tokens", 0) + adv_out = state.usage.get("advisor_output_tokens", 0) + if adv_in or adv_out: + right_bits.append(f"advisor {adv_in}/{adv_out}") right = " · ".join(right_bits) return Text(f"{left} {middle} {cwd} {right}") diff --git a/src/utils/advisor.py b/src/utils/advisor.py index 230003528..73d4cf9a0 100644 --- a/src/utils/advisor.py +++ b/src/utils/advisor.py @@ -694,12 +694,16 @@ def execute_client_advisor( *, abort_signal: Any = None, main_provider: Any = None, -) -> tuple[bool, str]: +) -> tuple[bool, str, dict[str, int]]: """Run one client-side advisor consultation. - Returns ``(ok, text)``: when ``ok`` is True, ``text`` is the + Returns ``(ok, text, usage)``: when ``ok`` is True, ``text`` is the advisor's advice; when False, ``text`` is a short error message suitable for surfacing as a tool_result with ``is_error=True``. + ``usage`` is a dict with ``input_tokens`` / ``output_tokens`` keys + (zero-filled on failure paths). The caller accumulates these into + a session-level counter so the status bar can show advisor token + spend separately from the worker's. Provider routing: * If ``main_provider`` was passed AND it has a custom base URL @@ -715,11 +719,12 @@ def execute_client_advisor( case (e.g. Anthropic main + Gemini advisor). Network failures, model errors, and missing-config conditions are - all caught and surfaced as ``(False, "...")`` rather than raised — - a tool that throws inside dispatch kills the turn, but a failed - advisor consultation should just leave the worker model uninformed - and let it continue. + all caught and surfaced as ``(False, "...", {0,0})`` rather than + raised — a tool that throws inside dispatch kills the turn, but + a failed advisor consultation should just leave the worker model + uninformed and let it continue. """ + _zero_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0} try: from src.config import get_provider_config from src.providers import get_provider_class @@ -734,7 +739,7 @@ def execute_client_advisor( else: provider_name = infer_provider_for_model(advisor_model) if provider_name is None: - return (False, f"Advisor unavailable: cannot route model {advisor_model!r} to a known provider.") + return (False, f"Advisor unavailable: cannot route model {advisor_model!r} to a known provider.", _zero_usage) provider_cls = get_provider_class(provider_name) cfg_raw = dict(get_provider_config(provider_name)) @@ -750,7 +755,7 @@ def execute_client_advisor( model=advisor_model, ) except Exception as e: # noqa: BLE001 — surface as advisor failure - return (False, f"Advisor unavailable: failed to construct provider for {advisor_model!r}: {e}") + return (False, f"Advisor unavailable: failed to construct provider for {advisor_model!r}: {e}", _zero_usage) # System-prompt delivery is provider-specific: # * Anthropic-shaped providers (AnthropicProvider / MinimaxProvider) @@ -804,12 +809,21 @@ def execute_client_advisor( # we can't pass it portably. response = provider.chat(request_messages, **call_kwargs) except Exception as e: # noqa: BLE001 — surface as advisor failure - return (False, f"Advisor unavailable: {type(e).__name__}: {e}") + return (False, f"Advisor unavailable: {type(e).__name__}: {e}", _zero_usage) + + # Pull token counts off the ChatResponse for the session + # accumulator. Defaults to zero when the provider didn't return + # a usage dict (some mocks / older providers). + raw_usage = getattr(response, "usage", None) or {} + usage: dict[str, int] = { + "input_tokens": int(raw_usage.get("input_tokens", 0) or 0), + "output_tokens": int(raw_usage.get("output_tokens", 0) or 0), + } text = getattr(response, "content", None) or "" if not isinstance(text, str) or not text.strip(): - return (False, "Advisor returned no text content.") - return (True, text) + return (False, "Advisor returned no text content.", usage) + return (True, text, usage) __all__ = [ diff --git a/tests/test_advisor_client_side.py b/tests/test_advisor_client_side.py index 160941406..ba1076d32 100644 --- a/tests/test_advisor_client_side.py +++ b/tests/test_advisor_client_side.py @@ -371,7 +371,7 @@ def test_returns_text_on_success_anthropic(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "test"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "claude-opus-4-6", [{"role": "user", "content": "hi"}] ) self.assertTrue(ok) @@ -393,7 +393,7 @@ def test_openai_shape_gets_system_as_first_message(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "test"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "gpt-5.4", [{"role": "user", "content": "hi"}] ) self.assertTrue(ok) @@ -407,7 +407,7 @@ def test_openai_shape_gets_system_as_first_message(self) -> None: self.assertEqual(forwarded_messages[1]["role"], "user") def test_returns_error_when_model_unroutable(self) -> None: - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "totally-fake-zzz-9999", [{"role": "user", "content": "hi"}] ) self.assertFalse(ok) @@ -424,7 +424,7 @@ def test_returns_error_when_provider_raises(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "test"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "claude-opus-4-6", [{"role": "user", "content": "hi"}] ) self.assertFalse(ok) @@ -438,7 +438,7 @@ def test_returns_error_when_response_empty(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "test"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "claude-opus-4-6", [{"role": "user", "content": "hi"}] ) self.assertFalse(ok) @@ -476,7 +476,7 @@ def _fake_openai_init(**kwargs: Any) -> Any: return_value={"api_key": "k", "base_url": "https://litellm.singula.ai"}, ): with patch.object(OpenAIProvider, "__new__", lambda cls, **kw: _fake_openai_init(**kw)): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "claude-opus-4-7", [{"role": "user", "content": "hi"}], main_provider=main_provider, @@ -510,7 +510,7 @@ def test_uses_inferred_provider_when_main_is_not_proxy(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "k"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "gemini-2.5-pro", [{"role": "user", "content": "hi"}], main_provider=main_provider, @@ -533,7 +533,7 @@ def test_falls_back_to_chat_when_stream_unimplemented(self) -> None: with patch( "src.config.get_provider_config", return_value={"api_key": "test"} ): - ok, text = execute_client_advisor( + ok, text, _usage = execute_client_advisor( "claude-opus-4-6", [{"role": "user", "content": "hi"}] ) self.assertTrue(ok) @@ -569,7 +569,7 @@ def test_call_returns_advice_text(self) -> None: with patch("src.settings.settings.get_settings", return_value=fake_settings): with patch( "src.utils.advisor.execute_client_advisor", - return_value=(True, "use the foo pattern"), + return_value=(True, "use the foo pattern", {"input_tokens": 0, "output_tokens": 0}), ): result = AdvisorTool.call({}, ctx) self.assertEqual(result.name, "advisor") @@ -584,12 +584,38 @@ def test_call_marks_error_when_execute_fails(self) -> None: with patch("src.settings.settings.get_settings", return_value=fake_settings): with patch( "src.utils.advisor.execute_client_advisor", - return_value=(False, "Advisor unavailable: foo"), + return_value=(False, "Advisor unavailable: foo", {"input_tokens": 0, "output_tokens": 0}), ): result = AdvisorTool.call({}, ctx) self.assertTrue(result.is_error) self.assertIn("unavailable", result.output) + def test_call_accumulates_advisor_tokens_onto_ctx(self) -> None: + """The dispatcher writes ``advisor_input_tokens`` / + ``advisor_output_tokens`` onto the ToolContext on every + consultation so the REPL/TUI status bar can surface them next + to the worker's totals. Multiple calls accumulate (not replace).""" + ctx = ToolContext(workspace_root=self._tmpdir) + ctx.messages = [{"role": "user", "content": "task"}] + fake_settings = MagicMock() + fake_settings.advisor_model = "claude-opus-4-6" + with patch("src.settings.settings.get_settings", return_value=fake_settings): + with patch( + "src.utils.advisor.execute_client_advisor", + return_value=(True, "advice 1", {"input_tokens": 100, "output_tokens": 50}), + ): + AdvisorTool.call({}, ctx) + self.assertEqual(ctx.advisor_input_tokens, 100) + self.assertEqual(ctx.advisor_output_tokens, 50) + # Second call ACCUMULATES on top of the first. + with patch( + "src.utils.advisor.execute_client_advisor", + return_value=(True, "advice 2", {"input_tokens": 30, "output_tokens": 10}), + ): + AdvisorTool.call({}, ctx) + self.assertEqual(ctx.advisor_input_tokens, 130) + self.assertEqual(ctx.advisor_output_tokens, 60) + def test_call_when_no_model_configured_returns_error(self) -> None: ctx = ToolContext(workspace_root=self._tmpdir) fake_settings = MagicMock()