diff --git a/tests/test_coord_rich_ws_state_payload.py b/tests/test_coord_rich_ws_state_payload.py index cbe9ba57..1deea4b9 100644 --- a/tests/test_coord_rich_ws_state_payload.py +++ b/tests/test_coord_rich_ws_state_payload.py @@ -73,6 +73,47 @@ def test_coord_on_status_persists_usage_event() -> None: assert kwargs["completion_tokens"] == 3 +def test_coord_on_aux_usage_persists_usage_event() -> None: + """Auxiliary LLM calls (plan/task sub-agents, compaction, web-fetch + summarisation, title gen) bypass ``on_status`` entirely. ``on_aux_usage`` + is what gets their token spend onto the governance dashboard.""" + storage = MagicMock() + ui = ConsoleCoordinatorUI(ws_id="coord-ws", user_id="u1") + with _patch_get_storage(storage): + ui.on_aux_usage( + { + "prompt_tokens": 500, + "completion_tokens": 40, + "cache_creation_tokens": 12, + "cache_read_tokens": 8, + "model": "plan-model", + } + ) + storage.record_usage_event.assert_called_once() + kwargs = storage.record_usage_event.call_args.kwargs + assert kwargs["ws_id"] == "coord-ws" + assert kwargs["user_id"] == "u1" + assert kwargs["model"] == "plan-model" + assert kwargs["prompt_tokens"] == 500 + assert kwargs["completion_tokens"] == 40 + assert kwargs["cache_creation_tokens"] == 12 + assert kwargs["cache_read_tokens"] == 8 + # Tools a sub-agent calls internally are its own tally, not this ws's. + assert kwargs["tool_calls_count"] == 0 + + +def test_coord_on_aux_usage_leaves_live_counters_untouched() -> None: + """Unlike ``on_status``, ``on_aux_usage`` must NOT fold the auxiliary + prompt into the live per-ws context gauge — that tracks the main + conversation's window, and an agent/compaction prompt isn't it.""" + ui = ConsoleCoordinatorUI(ws_id="coord-ws", user_id="u1") + with _patch_get_storage(MagicMock()): + ui.on_aux_usage({"prompt_tokens": 9999, "completion_tokens": 9999}) + assert ui._ws_prompt_tokens == 0 + assert ui._ws_completion_tokens == 0 + assert ui._ws_context_ratio == 0.0 + + def test_coord_on_content_token_accumulates() -> None: """Pre-lift coord ``on_content_token`` only enqueued; lift turns it into the same per-ws accumulator WebUI uses so the collector @@ -580,6 +621,37 @@ def test_webui_on_status_still_records_prometheus_metrics() -> None: WebUI._global_queue = None +def test_webui_on_aux_usage_records_prometheus_metrics() -> None: + """WebUI.on_aux_usage must feed ``_metrics.record_*`` so auxiliary + (sub-agent / compaction / utility) tokens land in + ``turnstone_tokens_total``, not just main-loop turns. Regression guard + mirroring ``test_webui_on_status_still_records_prometheus_metrics`` — + without it, a refactor dropping the override would silently stop + counting aux tokens with nothing failing.""" + import queue + + from turnstone.server import WebUI + + WebUI._global_queue = queue.Queue() + try: + ui = WebUI(ws_id="ws-int", user_id="u1") + with patch("turnstone.server._metrics") as mock_metrics, _patch_get_storage(MagicMock()): + ui.on_aux_usage( + { + "prompt_tokens": 64, + "completion_tokens": 8, + "cache_creation_tokens": 3, + "cache_read_tokens": 5, + } + ) + mock_metrics.record_tokens.assert_called_once_with(64, 8) + mock_metrics.record_cache_tokens.assert_called_once_with(3, 5) + # Unlike on_status, aux usage must NOT touch the context-ratio gauge. + mock_metrics.record_context_ratio.assert_not_called() + finally: + WebUI._global_queue = None + + def test_webui_on_tool_result_still_records_prometheus_tool_call() -> None: """Same as above for ``on_tool_result``.""" import queue diff --git a/tests/test_session.py b/tests/test_session.py index fc445077..00b6c570 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -5854,3 +5854,104 @@ class TestSearchCaptureStreaming: from turnstone.core.session import _SEARCH_STDERR_CAP assert len(stderr) <= _SEARCH_STDERR_CAP + + +# --------------------------------------------------------------------------- +# Auxiliary-usage accounting — non-streaming LLM calls (title gen, +# compaction, web-fetch summarisation, plan/task sub-agents) bypass the +# streaming on_status path; _record_aux_usage routes their usage to the +# UI's on_aux_usage hook so it still reaches the governance dashboard. +# --------------------------------------------------------------------------- + + +class _AuxRecordingUI(NullUI): + """NullUI plus the on_aux_usage hook, capturing each recorded dict.""" + + def __init__(self) -> None: + self.aux_calls: list[dict[str, Any]] = [] + + def on_aux_usage(self, usage): + self.aux_calls.append(usage) + + +def test_utility_completion_records_aux_usage(): + """A utility completion's token usage is routed to on_aux_usage with the + fields mapped from the provider's UsageInfo and the session model.""" + from turnstone.core.providers._protocol import ( + CompletionResult, + ModelCapabilities, + UsageInfo, + ) + + ui = _AuxRecordingUI() + session = _make_session(ui=ui) + session._provider = MagicMock() + session._provider.get_capabilities.return_value = ModelCapabilities() + session._provider.create_completion.return_value = CompletionResult( + content="A Generated Title", + usage=UsageInfo( + prompt_tokens=120, + completion_tokens=8, + total_tokens=128, + cache_creation_tokens=4, + cache_read_tokens=16, + ), + ) + + session._utility_completion([{"role": "user", "content": "hi"}]) + + assert len(ui.aux_calls) == 1 + rec = ui.aux_calls[0] + assert rec["prompt_tokens"] == 120 + assert rec["completion_tokens"] == 8 + assert rec["cache_creation_tokens"] == 4 + assert rec["cache_read_tokens"] == 16 + assert rec["model"] == "test-model" + + +def test_record_aux_usage_skips_when_usage_missing(): + """A provider that reports no usage object must not emit a phantom + zero-token row.""" + from turnstone.core.providers._protocol import CompletionResult + + ui = _AuxRecordingUI() + session = _make_session(ui=ui) + session._record_aux_usage(CompletionResult(content="x", usage=None)) + assert ui.aux_calls == [] + + +def test_record_aux_usage_noop_without_ui_hook(): + """Minimal UI stubs predating on_aux_usage (e.g. NullUI) must not crash + a title-gen or sub-agent turn — recording silently no-ops.""" + from turnstone.core.providers._protocol import CompletionResult, UsageInfo + + session = _make_session(ui=NullUI()) # NullUI has no on_aux_usage + session._record_aux_usage( + CompletionResult( + content="x", + usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + ) # no exception raised == pass + + +def test_record_aux_usage_attributes_explicit_model(): + """Sub-agent turns record under the agent's OWN model — session.py's + _api_call passes model=agent_model so plan/task spend attributes to the + sub-agent's model, not the coordinating session's. Verify the override + reaches on_aux_usage rather than defaulting to self.model.""" + from turnstone.core.providers._protocol import CompletionResult, UsageInfo + + ui = _AuxRecordingUI() + session = _make_session(ui=ui) # session model == "test-model" + session._record_aux_usage( + CompletionResult( + content="plan output", + usage=UsageInfo(prompt_tokens=900, completion_tokens=60, total_tokens=960), + ), + model="plan-model-xyz", + ) + + assert len(ui.aux_calls) == 1 + # The explicit agent model wins over the session default. + assert ui.aux_calls[0]["model"] == "plan-model-xyz" + assert ui.aux_calls[0]["prompt_tokens"] == 900 diff --git a/turnstone/console/static/governance.js b/turnstone/console/static/governance.js index 33514d16..570d45c4 100644 --- a/turnstone/console/static/governance.js +++ b/turnstone/console/static/governance.js @@ -2462,21 +2462,21 @@ function loadGovUsage() { else since = new Date(now - 7 * 24 * 60 * 60 * 1000); const sinceStr = since.toISOString().slice(0, 19); - // Fetch summary + breakdown in parallel - const summaryUrl = - "/v1/api/admin/usage?since=" + encodeURIComponent(sinceStr); - const breakdownUrl = summaryUrl + "&group_by=" + _govUsageGroupBy; + // A single request returns both the window-total `summary` (one SUM + // row, computed independent of group_by) and the grouped `breakdown` + // rows — the readout cards read the former, the bar chart the latter. + const url = + "/v1/api/admin/usage?since=" + + encodeURIComponent(sinceStr) + + "&group_by=" + + _govUsageGroupBy; - Promise.all([ - authFetch(summaryUrl).then(function (r) { + authFetch(url) + .then(function (r) { return r.json(); - }), - authFetch(breakdownUrl).then(function (r) { - return r.json(); - }), - ]) - .then(function (results) { - _renderGovUsage(results[0], results[1]); + }) + .then(function (data) { + _renderGovUsage(data); }) .catch(function () { setSafeHtml( @@ -2486,9 +2486,11 @@ function loadGovUsage() { }); } -function _renderGovUsage(summary, breakdown) { +function _renderGovUsage(data) { const container = document.getElementById("admin-usage-content"); - const s = (summary.breakdown && summary.breakdown[0]) || {}; + // Cards show the window total (the SUM row), NOT breakdown[0] — the + // latter is the oldest bucket and silently understated every headline. + const s = (data.summary && data.summary[0]) || {}; const prompt = s.prompt_tokens || 0; const completion = s.completion_tokens || 0; const total = prompt + completion; @@ -2529,7 +2531,7 @@ function _renderGovUsage(summary, breakdown) { ""; // Bar chart breakdown - const items = breakdown.breakdown || []; + const items = data.breakdown || []; if (items.length) { let maxVal = 0; for (let i = 0; i < items.length; i++) { diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 33af1c81..14a13d64 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -3100,7 +3100,7 @@ class ChatSession: caps = self._get_capabilities() clamped = min(max_tokens, caps.max_output_tokens) if caps.max_output_tokens else max_tokens messages = self._maybe_attach_vllm_chat_reasoning(messages, self._provider) - return self._provider.create_completion( + result = self._provider.create_completion( client=self.client, model=self.model, messages=messages, @@ -3111,6 +3111,44 @@ class ChatSession: capabilities=caps, replay_reasoning_to_model=self._resolve_replay_reasoning_to_model(caps=caps), ) + # Utility completions (title gen, compaction, web-fetch extraction) + # bypass the streaming on_status path — record their usage so the + # governance dashboard reflects this spend. + self._record_aux_usage(result) + return result + + def _record_aux_usage(self, result: CompletionResult, *, model: str | None = None) -> None: + """Persist token usage for a non-streaming auxiliary completion. + + Title generation, compaction, web-fetch summarisation, and + plan/task sub-agents all run via ``create_completion`` and bypass + the streaming ``on_status`` accounting path; without this their + spend never reaches the usage dashboard. Delegates to the UI's + ``on_aux_usage`` hook (which owns the storage write + any node + metrics), mirroring how ``_print_status_line`` routes main-loop + usage through ``on_status``. + + ``model`` defaults to the session model (utility calls share it); + sub-agent callers pass the agent's own model so per-model + attribution stays accurate. The hook is looked up defensively — + minimal UI stubs (some tests, replay shims) predate it and should + skip recording rather than crash a title-gen or sub-agent turn. + """ + u = result.usage + if u is None: + return + record = getattr(self.ui, "on_aux_usage", None) + if record is None: + return + record( + { + "prompt_tokens": u.prompt_tokens, + "completion_tokens": u.completion_tokens, + "cache_creation_tokens": u.cache_creation_tokens, + "cache_read_tokens": u.cache_read_tokens, + "model": model or self.model, + } + ) # -- tool search helpers -------------------------------------------------- @@ -11025,7 +11063,7 @@ class ChatSession: last_err: Exception | None = None for attempt in range(self._MAX_RETRIES + 1): try: - return agent_provider.create_completion( + agent_result = agent_provider.create_completion( client=agent_client, model=agent_model, messages=messages, @@ -11039,6 +11077,11 @@ class ChatSession: agent_alias, caps=agent_caps ), ) + # Sub-agent turns bypass on_status — record per-turn so + # plan/task spend is visible in the dashboard, attributed + # to the agent's own model. + self._record_aux_usage(agent_result, model=agent_model) + return agent_result except Exception as e: ename = type(e).__name__ if ( diff --git a/turnstone/core/session_ui_base.py b/turnstone/core/session_ui_base.py index 21a0541e..5cf1bf63 100644 --- a/turnstone/core/session_ui_base.py +++ b/turnstone/core/session_ui_base.py @@ -1930,6 +1930,66 @@ class SessionUIBase: "turn_count": turn_count, } ) + self._write_usage_row( + model=usage.get("model", ""), + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + tool_calls_count=tool_count, + cache_creation_tokens=cache_creation, + cache_read_tokens=cache_read, + ) + + def on_aux_usage(self, usage: dict[str, Any]) -> None: + """Persist a ``usage_event`` for an auxiliary (non-main-loop) LLM call. + + Title generation, conversation compaction, web-fetch + summarisation, and plan/task sub-agents all run through the + provider's non-streaming ``create_completion`` and never reach + :meth:`on_status` — so without this their tokens are invisible to + the governance usage dashboard, undercounting real consumption + (potentially by a large factor for agent-heavy workstreams). + + Deliberately narrower than :meth:`on_status`: it records ONLY the + storage row. It does NOT enqueue a ``status`` UI event, advance + ``_ws_prompt_tokens`` / ``_ws_context_ratio`` (an auxiliary + prompt is not the main conversation's live context — folding it + in would make the status-bar context gauge lie), or touch the + tool-call delta counter. ``tool_calls_count`` is recorded as 0: + any tools a sub-agent or judge calls internally are its own + concern, not part of this workstream's surfaced tool tally. + + ``WebUI`` overrides this to also feed the node Prometheus token + counters; call ``super().on_aux_usage(...)`` to keep the + ``usage_event`` row consistent. + """ + self._write_usage_row( + model=usage.get("model", ""), + prompt_tokens=usage.get("prompt_tokens", 0), + completion_tokens=usage.get("completion_tokens", 0), + tool_calls_count=0, + cache_creation_tokens=usage.get("cache_creation_tokens", 0), + cache_read_tokens=usage.get("cache_read_tokens", 0), + ) + + def _write_usage_row( + self, + *, + model: str, + prompt_tokens: int, + completion_tokens: int, + tool_calls_count: int, + cache_creation_tokens: int, + cache_read_tokens: int, + ) -> None: + """Persist one ``usage_event`` row, swallowing + logging storage + errors so a reporting-sink failure never breaks a live turn. + + Shared by :meth:`on_status` (main-loop turns, real tool-call + delta) and :meth:`on_aux_usage` (auxiliary calls, with + ``tool_calls_count=0``) so the get-storage / record / except-log + shape lives in one place and can't drift if the + ``record_usage_event`` signature changes. + """ try: from turnstone.core.storage._registry import get_storage @@ -1940,12 +2000,12 @@ class SessionUIBase: user_id=self._user_id, ws_id=self.ws_id, node_id="", - model=usage.get("model", ""), + model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, - tool_calls_count=tool_count, - cache_creation_tokens=cache_creation, - cache_read_tokens=cache_read, + tool_calls_count=tool_calls_count, + cache_creation_tokens=cache_creation_tokens, + cache_read_tokens=cache_read_tokens, ) except Exception: log.warning("Failed to record usage event", exc_info=True) diff --git a/turnstone/server.py b/turnstone/server.py index a5693301..404dff76 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -354,6 +354,22 @@ class WebUI(SessionUIBase): _metrics.record_context_ratio(total_tok / context_window if context_window > 0 else 0.0) super().on_status(usage, context_window, effort) + def on_aux_usage(self, usage: dict[str, Any]) -> None: + """Feed node Prometheus token counters for auxiliary LLM calls. + + Mirrors :meth:`on_status`' token-metric writes (minus + context-ratio — an auxiliary prompt is not the main context + window) so ``turnstone_tokens_total`` reflects sub-agent, + compaction, and utility spend, not just main-loop turns. The + ``usage_event`` storage row is inherited from + :meth:`SessionUIBase.on_aux_usage`. + """ + _metrics.record_tokens(usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0)) + _metrics.record_cache_tokens( + usage.get("cache_creation_tokens", 0), usage.get("cache_read_tokens", 0) + ) + super().on_aux_usage(usage) + def on_plan_review(self, content: str) -> str: self._plan_event.clear() self._pending_plan_review = {"type": "plan_review", "content": content}