diff --git a/agent/title_generator.py b/agent/title_generator.py index a84ab93af3..fdb0556e81 100644 --- a/agent/title_generator.py +++ b/agent/title_generator.py @@ -9,6 +9,9 @@ import json import logging import os import re +import threading +import time +import weakref from contextlib import suppress from typing import Any, Callable, Optional @@ -19,6 +22,18 @@ from agent.message_content import flatten_message_text logger = logging.getLogger(__name__) +# In-flight stage-2 upgrade threads. They bill their aux usage to the session from a daemon thread, +# so a process that reads the ledger right before exit (``-z --usage-file``) must be able to join +# them (bounded) instead of racing the write (#112848). +_UPGRADE_THREADS: "weakref.WeakSet[threading.Thread]" = weakref.WeakSet() + + +def wait_for_title_upgrades(timeout: float = 10.0) -> None: + """Bounded join of the auto-title threads still running; never raises.""" + deadline = time.monotonic() + timeout + for thread in list(_UPGRADE_THREADS): + thread.join(max(0.0, deadline - time.monotonic())) + # (task_name, exception) -> None; surfaces auxiliary failures so silent drops don't pile up as NULL titles. FailureCallback = Callable[[str, BaseException], None] # (title, source) -> None; source is the persisted provenance (``derived`` / ``llm``). Consumers paying a @@ -528,9 +543,11 @@ def maybe_auto_title( # profile whose turn this is: a bare Thread starts with an empty context and lands on the launch # profile under multiplex, titling X's session with the default profile's model and billing its key. from agent.memory_provider import spawn_context_thread - spawn_context_thread( + upgrade = spawn_context_thread( auto_title_session, name="auto-title", args=(session_db, session_id, user_message), kwargs=dict(failure_callback=failure_callback, main_runtime=main_runtime, title_callback=title_callback, runtime_validator=runtime_validator), - ).start() + ) + _UPGRADE_THREADS.add(upgrade) + upgrade.start() diff --git a/hermes_cli/oneshot.py b/hermes_cli/oneshot.py index 5c3603f2c4..6cfca65125 100644 --- a/hermes_cli/oneshot.py +++ b/hermes_cli/oneshot.py @@ -31,6 +31,55 @@ _USAGE_KEYS = ( "model", "provider", "session_id", "completed", ) +# Counters summed per auxiliary task (vision, compression, title_generation, ...) into the +# ``auxiliary`` block of the report. The main-loop keys above stay main-loop-only (backward +# compatible); ``total_including_auxiliary`` carries the grand total pipelines bill on (#112848). +_AUX_COUNTERS = ( + "api_calls", "input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens", + "reasoning_tokens", "estimated_cost_usd", +) + + +def _auxiliary_usage(session_db, session_id: Optional[str]) -> dict[str, dict]: + """Per-task aux usage recorded for *session_id*'s lineage (``{}`` without a store / session).""" + if session_db is None or not session_id: + return {} + try: + return session_db.auxiliary_usage_by_task(session_id) + except Exception: + logging.debug("oneshot: auxiliary usage read failed", exc_info=True) + return {} + + +def _attach_auxiliary_usage(result: dict, session_db, before: dict[str, dict]) -> None: + """Store this run's auxiliary usage on *result* as the delta against the pre-turn snapshot + (a resumed session already carries earlier runs' aux rows). Waits (bounded) for the auto-title + thread first: it bills from a daemon thread and can still be in flight when the turn returns.""" + from agent.title_generator import wait_for_title_upgrades + + wait_for_title_upgrades() + after = _auxiliary_usage(session_db, result.get("session_id")) + by_task: dict[str, dict] = {} + for task, counters in after.items(): + prior = before.get(task, {}) + delta = {key: (counters.get(key) or 0) - (prior.get(key) or 0) for key in _AUX_COUNTERS} + if any(delta.values()): + by_task[task] = delta + result["auxiliary_usage"] = by_task + + +def _auxiliary_report(report: dict, by_task: dict[str, dict]) -> None: + """Add the ``auxiliary`` breakdown and ``total_including_auxiliary`` to the ledger.""" + totals = {key: sum(t.get(key) or 0 for t in by_task.values()) for key in _AUX_COUNTERS} + totals["total_tokens"] = totals["input_tokens"] + totals["output_tokens"] + report["auxiliary"] = {**totals, "by_task": by_task} + main_cost = report.get("estimated_cost_usd") + report["total_including_auxiliary"] = { + "estimated_cost_usd": None if main_cost is None else main_cost + totals["estimated_cost_usd"], + "total_tokens": (report.get("total_tokens") or 0) + totals["total_tokens"], + "api_calls": (report.get("api_calls") or 0) + totals["api_calls"], + } + def _normalize_toolsets(toolsets: object = None) -> list[str] | None: """Split repeated/comma-separated toolset flags into a clean list (``None`` when empty).""" @@ -151,6 +200,8 @@ def _write_usage_file(path: Optional[str], result: dict, failure: Optional[str] report = {key: result.get(key) for key in _USAGE_KEYS} report["failed"] = bool(result.get("failed")) or failure is not None report["service_tier"] = result.get("service_tier") + if isinstance(result.get("auxiliary_usage"), dict): + _auxiliary_report(report, result["auxiliary_usage"]) if failure is not None: report["failure"] = failure out = Path(path).expanduser() @@ -225,6 +276,7 @@ def run_oneshot( skills=skills, resume=resume, reasoning=reasoning, + ledger=bool(usage_file), ) except BaseException as exc: # noqa: BLE001 # Capture anything escaping the agent (OSError from prompt_toolkit on a non-TTY pipe, @@ -418,9 +470,11 @@ def _run_agent( skills: object = None, resume: Optional[str] = None, reasoning: object = None, + ledger: bool = False, ) -> tuple[str, dict]: """Build an AIAgent exactly like a normal CLI chat turn, run one conversation, and return - ``(final_response, run_result)``. Imports are local to keep CLI startup cheap.""" + ``(final_response, run_result)``. Imports are local to keep CLI startup cheap. *ledger* (set when + ``--usage-file`` is requested) attaches this run's auxiliary usage to the result.""" from hermes_cli.config import load_config from hermes_cli.runtime_provider import resolve_runtime_provider from hermes_cli.tools_config import _get_platform_tools @@ -500,7 +554,10 @@ def _run_agent( agent.stream_delta_callback = None agent.tool_gen_callback = None + aux_before = _auxiliary_usage(session_db, resume_sid) if ledger else {} result = agent.run_conversation(prompt, conversation_history=conversation_history or None) + if ledger: + _attach_auxiliary_usage(result, session_db, aux_before) return (result.get("final_response") or "", result) finally: _close_agent(agent, session_db) diff --git a/hermes_state_usage.py b/hermes_state_usage.py index 5faf4f63e0..77c85ecb44 100644 --- a/hermes_state_usage.py +++ b/hermes_state_usage.py @@ -390,6 +390,29 @@ class SessionUsageMixin: self._insert_session_row(session_id, "unknown") self._execute_write(lambda conn: self._record_model_usage(conn, session_id, task=task, **usage)) + def auxiliary_usage_by_task(self, session_id: str) -> Dict[str, Dict[str, float]]: + """Per-task auxiliary usage (``task != ''``: vision, compression, title_generation, ...) summed + over the session's compression lineage. Aux calls bill to the id the turn STARTED with while + compression mints child ids mid-turn, so a single-id read misses rows (#112848).""" + if not session_id: + return {} + chain = self._session_lineage_root_to_tip(session_id) + rows = self._read_all( + f"""SELECT task, + COALESCE(SUM(api_call_count), 0) AS api_calls, + COALESCE(SUM(input_tokens), 0) AS input_tokens, + COALESCE(SUM(output_tokens), 0) AS output_tokens, + COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens, + COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens, + COALESCE(SUM(reasoning_tokens), 0) AS reasoning_tokens, + COALESCE(SUM(estimated_cost_usd), 0) AS estimated_cost_usd + FROM session_model_usage + WHERE session_id IN ({','.join('?' * len(chain))}) AND task != '' + GROUP BY task""", + chain, + ) + return {row["task"]: {k: row[k] for k in row.keys() if k != "task"} for row in rows} + def usage_totals(self, *, min_message_count: int = 1, include_archived: bool = False) -> Dict[str, float]: """Tokens and spend across the whole store (one scan), so the sidebar total does not shrink with paging. Spend prefers the billed figure over the estimate.""" diff --git a/tests/hermes_cli/test_oneshot_usage_file.py b/tests/hermes_cli/test_oneshot_usage_file.py index 48b3c79d2e..d582bc40d9 100644 --- a/tests/hermes_cli/test_oneshot_usage_file.py +++ b/tests/hermes_cli/test_oneshot_usage_file.py @@ -2,6 +2,8 @@ import json +import pytest + from hermes_cli.oneshot import _write_usage_file @@ -55,3 +57,66 @@ class TestWriteUsageFile: assert report["estimated_cost_usd"] is None + + +class TestAuxiliaryLedger: + """#112848: auxiliary LLM spend (title generation, vision, ...) recorded in session_model_usage + belongs in the pipeline ledger, additively — the main-loop keys stay main-loop-only.""" + + def test_aux_usage_is_a_separate_breakdown_and_main_keys_unchanged(self, tmp_path): + from hermes_cli.oneshot import _auxiliary_usage, _attach_auxiliary_usage + from hermes_state import SessionDB + + db = SessionDB(tmp_path / "state.db") + try: + db.create_session("root", "cli", model="main") + # Rows from an earlier run of a resumed session must not count toward this run. + db.record_auxiliary_usage("root", "vision", model="v", input_tokens=500, output_tokens=5, + estimated_cost_usd=0.5) + before = _auxiliary_usage(db, "root") + # This run: aux billed to the id the turn started with; compression mints a child id. + db.record_auxiliary_usage("root", "title_generation", model="t", input_tokens=40, output_tokens=8, + estimated_cost_usd=0.001) + db.create_session("child", "cli", model="main", parent_session_id="root") + result = _result(session_id="child") + _attach_auxiliary_usage(result, db, before) + finally: + db.close() + + path = tmp_path / "usage.json" + _write_usage_file(str(path), result) + report = json.loads(path.read_text()) + assert report["api_calls"] == 3 and report["total_tokens"] == 1250 # main loop untouched + assert report["auxiliary"]["by_task"] == { + "title_generation": {"api_calls": 1, "input_tokens": 40, "output_tokens": 8, "cache_read_tokens": 0, + "cache_write_tokens": 0, "reasoning_tokens": 0, "estimated_cost_usd": 0.001}, + } + assert report["auxiliary"]["total_tokens"] == 48 + assert report["total_including_auxiliary"] == { + "api_calls": 4, "total_tokens": 1298, "estimated_cost_usd": pytest.approx(0.1244), + } + + def test_waits_for_in_flight_title_thread(self, tmp_path): + import threading + import time + + from agent import title_generator + from hermes_cli.oneshot import _attach_auxiliary_usage + from hermes_state import SessionDB + + db = SessionDB(tmp_path / "state.db") + try: + db.create_session("s", "cli", model="main") + + def late_row(): + time.sleep(0.2) + db.record_auxiliary_usage("s", "title_generation", model="t", input_tokens=7, output_tokens=1) + + thread = threading.Thread(target=late_row, name="auto-title") + title_generator._UPGRADE_THREADS.add(thread) + thread.start() + result = _result(session_id="s") + _attach_auxiliary_usage(result, db, {}) + finally: + db.close() + assert result["auxiliary_usage"]["title_generation"]["api_calls"] == 1 diff --git a/website/docs/reference/cli-commands.md b/website/docs/reference/cli-commands.md index 47538ebe66..4d0cab7c54 100644 --- a/website/docs/reference/cli-commands.md +++ b/website/docs/reference/cli-commands.md @@ -239,11 +239,12 @@ Same agent, same tools, same skills — just strips every interactive / cosmetic #### `--usage-file` — JSON usage report for pipelines -`hermes -z "…" --usage-file /path/report.json` writes a machine-readable usage report after the run: `estimated_cost_usd`, `input_tokens` / `output_tokens` / `cache_read_tokens` / `cache_write_tokens` / `reasoning_tokens` / `total_tokens`, `api_calls`, `model`, `provider`, `session_id`, `service_tier`, and `completed` / `failed` flags. The report is written **even when the run fails**, so batch pipelines can always account for spend. It has no effect outside `-z`/`--oneshot`, and a broken usage write never masks the run's own outcome. +`hermes -z "…" --usage-file /path/report.json` writes a machine-readable usage report after the run: `estimated_cost_usd`, `input_tokens` / `output_tokens` / `cache_read_tokens` / `cache_write_tokens` / `reasoning_tokens` / `total_tokens`, `api_calls`, `model`, `provider`, `session_id`, `service_tier`, and `completed` / `failed` flags. Those top-level counters cover the **main agent loop** only. Auxiliary LLM calls made on the same run (title generation, vision, context compression, `web_extract`, background review, …) are reported separately under `auxiliary` — the same totals plus a per-task `by_task` map — and `total_including_auxiliary` (`estimated_cost_usd`, `total_tokens`, `api_calls`) is the grand total to bill on. The report is written **even when the run fails**, so batch pipelines can always account for spend. It has no effect outside `-z`/`--oneshot`, and a broken usage write never masks the run's own outcome. ```bash hermes -z "summarize this repo" --usage-file /tmp/usage.json -jq .estimated_cost_usd /tmp/usage.json +jq .total_including_auxiliary.estimated_cost_usd /tmp/usage.json +jq .auxiliary.by_task /tmp/usage.json # what did title generation / vision cost? ``` ## `hermes model`