diff --git a/agent/aux_accounting.py b/agent/aux_accounting.py new file mode 100644 index 000000000000..6f851a642ecd --- /dev/null +++ b/agent/aux_accounting.py @@ -0,0 +1,138 @@ +"""Ambient session-accounting context for auxiliary LLM calls. + +Auxiliary calls (vision, compression, title generation, web_extract, +session_search, ...) funnel through ``agent.auxiliary_client`` which has no +session handle — so their token usage was historically discarded, leaving +dashboard analytics blind to aux model spend (issue #23270). + +Instead of threading ``session_db``/``session_id`` parameters through every +aux call site, the agent loop publishes them here (mirroring the Nous Portal +conversation context in ``agent.portal_tags``) and the auxiliary client +records usage at its single response-validation chokepoint. + +ContextVar semantics give us the right isolation for free: + +* concurrent agents in one process (gateway sessions, delegate subagents) + never see each other's accounting context; +* worker threads spawned via ``tools.thread_context.propagate_context_to_thread`` + (MoA fan-out, background review) inherit the parent turn's context; +* asyncio tasks inherit the context of the code that created them. + +MoA reference/aggregator slots are explicitly EXCLUDED from recording: +``agent/conversation_loop.py`` already folds MoA advisor usage and cost into +the main loop's ``update_token_counts`` delta, so recording them here would +double-count (see ``_EXCLUDED_TASKS``). +""" + +from __future__ import annotations + +import logging +from contextvars import ContextVar +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +# (session_db, session_id) for the active agent turn, or None outside one. +_accounting: ContextVar[Optional[tuple]] = ContextVar( + "aux_accounting_context", default=None +) + +# Aux tasks whose usage is already accounted by the main loop — recording +# them here would double-count. MoA advisor/aggregator usage is folded into +# conversation_loop's update_token_counts delta (tokens AND cost). +_EXCLUDED_TASKS = frozenset({"moa_reference", "moa_aggregator"}) + + +def set_accounting_context(session_db: Any, session_id: Optional[str]): + """Publish the active session's accounting handles for aux usage recording. + + Called by the agent loop at turn entry. Returns the ContextVar token so + callers can ``reset_accounting_context(token)`` on turn exit. Publishing + ``None`` handles (no DB / no session id) clears the context. + """ + if session_db is None or not session_id: + return _accounting.set(None) + return _accounting.set((session_db, session_id)) + + +def reset_accounting_context(token) -> None: + """Restore the previous accounting context (pair with ``set_...``).""" + try: + _accounting.reset(token) + except Exception: + _accounting.set(None) + + +def get_accounting_context() -> Optional[tuple]: + """Return ``(session_db, session_id)`` for the active turn, or ``None``.""" + return _accounting.get() + + +def record_aux_usage( + response: Any, + task: Optional[str], + *, + provider: Optional[str] = None, + base_url: Optional[str] = None, +) -> None: + """Record an auxiliary response's token usage against the ambient session. + + Called from the auxiliary client's response-validation chokepoint. Strictly + best-effort: any failure is swallowed (accounting must never break an aux + call). No-ops when: + + * no accounting context is published (call is outside any agent turn), + * the task is main-loop-accounted (MoA slots — see ``_EXCLUDED_TASKS``), + * the response carries no usage object. + + The model is read from ``response.model`` (accurate even after the aux + client's provider-fallback chains); *provider*/*base_url* reflect the + originally-resolved route and are best-effort. + """ + try: + if not task or task in _EXCLUDED_TASKS: + return + ctx = _accounting.get() + if ctx is None: + return + session_db, session_id = ctx + raw_usage = getattr(response, "usage", None) + if raw_usage is None: + return + + from agent.usage_pricing import estimate_usage_cost, normalize_usage + + usage = normalize_usage(raw_usage, provider=provider) + if not ( + usage.input_tokens or usage.output_tokens + or usage.cache_read_tokens or usage.cache_write_tokens + or usage.reasoning_tokens + ): + return + + model = str(getattr(response, "model", "") or "") or "unknown" + estimated_cost = None + try: + cost = estimate_usage_cost( + model, usage, provider=provider, base_url=base_url + ) + if cost.amount_usd is not None: + estimated_cost = float(cost.amount_usd) + except Exception: + logger.debug("Aux usage cost estimation failed", exc_info=True) + + session_db.record_auxiliary_usage( + session_id, + task, + model=model, + billing_provider=provider, + billing_base_url=base_url, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + cache_read_tokens=usage.cache_read_tokens, + cache_write_tokens=usage.cache_write_tokens, + reasoning_tokens=usage.reasoning_tokens, + estimated_cost_usd=estimated_cost, + ) + except Exception: + logger.debug("Aux usage recording failed (non-fatal)", exc_info=True) diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 54c7e4678830..bacd3e7749ff 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -6626,7 +6626,12 @@ def _build_call_kwargs( return kwargs -def _validate_llm_response(response: Any, task: str = None) -> Any: +def _validate_llm_response( + response: Any, + task: Optional[str] = None, + provider: Optional[str] = None, + base_url: Optional[str] = None, +) -> Any: """Validate that an LLM response has the expected .choices[0].message shape. Fails fast with a clear error instead of letting malformed payloads @@ -6634,11 +6639,21 @@ def _validate_llm_response(response: Any, task: str = None) -> Any: AttributeError (e.g. "'str' object has no attribute 'choices'"). See #7264. + + Also the single accounting chokepoint for auxiliary usage: every + successful non-streaming aux response passes through here exactly once, + so token usage is recorded against the ambient session context published + by the agent loop (``agent.aux_accounting``, issue #23270). Recording is + best-effort and never affects validation. *provider*/*base_url* are + optional accounting hints — fallback-path calls omit them and the row + keeps the model (read from the response itself) with an empty route. """ if response is None: raise RuntimeError( f"Auxiliary {task or 'call'}: LLM returned None response" ) + from agent.aux_accounting import record_aux_usage + record_aux_usage(response, task, provider=provider, base_url=base_url) # Allow SimpleNamespace responses from adapters (CodexAuxiliaryClient, # AnthropicAuxiliaryClient) — they have .choices[0].message. try: @@ -6905,7 +6920,8 @@ def call_llm( # for the transient retry every auxiliary task shares. (PR #16587) try: return _validate_llm_response( - client.chat.completions.create(**kwargs), task) + client.chat.completions.create(**kwargs), task, + provider=resolved_provider, base_url=_base_info) except Exception as transient_err: if not _is_transient_transport_error(transient_err): raise @@ -7484,7 +7500,8 @@ async def async_call_llm( # for the rationale. (PR #16587) try: return _validate_llm_response( - await client.chat.completions.create(**kwargs), task) + await client.chat.completions.create(**kwargs), task, + provider=resolved_provider, base_url=_client_base) except Exception as transient_err: if not _is_transient_transport_error(transient_err): raise diff --git a/agent/title_generator.py b/agent/title_generator.py index 840b56aaf09b..2270d75297fd 100644 --- a/agent/title_generator.py +++ b/agent/title_generator.py @@ -151,6 +151,7 @@ def auto_title_session( # ``conversation=`` Portal tag as the turn it titles. Root-of-lineage for # consistency with the agent loop (a no-op on first exchange, where # titling happens, but correct if this ever runs on a continuation). + from agent.aux_accounting import set_accounting_context from agent.portal_tags import set_conversation_context conversation_id = session_id @@ -159,6 +160,9 @@ def auto_title_session( except Exception: pass set_conversation_context(conversation_id) + # Same for the accounting context, so the title call's token usage is + # recorded against this session (task='title_generation', #23270). + set_accounting_context(session_db, session_id) title = generate_title( user_message, assistant_response, failure_callback=failure_callback, main_runtime=main_runtime diff --git a/hermes_cli/web_server.py b/hermes_cli/web_server.py index ac40ff8eba4e..38487c14ee2c 100644 --- a/hermes_cli/web_server.py +++ b/hermes_cli/web_server.py @@ -14297,6 +14297,115 @@ async def update_config_raw(body: RawConfigUpdate, profile: Optional[str] = None # --------------------------------------------------------------------------- +def _aux_usage_rows(db, cutoff: float) -> List[Dict[str, Any]]: + """Per-(model, task) auxiliary usage within the window (issue #23270). + + Reads the task-dimension rows (task != '') that record_auxiliary_usage + writes into session_model_usage. Returns [] when the table predates the + task column (older DB opened read-only by newer code). + """ + try: + cur = db._conn.execute(""" + SELECT u.model, + u.task, + u.billing_provider, + SUM(u.input_tokens) as input_tokens, + SUM(u.output_tokens) as output_tokens, + SUM(u.cache_read_tokens) as cache_read_tokens, + SUM(u.reasoning_tokens) as reasoning_tokens, + COALESCE(SUM(u.estimated_cost_usd), 0) as estimated_cost, + COUNT(DISTINCT u.session_id) as sessions, + SUM(COALESCE(u.api_call_count, 0)) as api_calls, + MAX(u.last_seen) as last_used_at + FROM session_model_usage u + JOIN sessions s ON s.id = u.session_id + WHERE s.started_at > ? AND u.task != '' + GROUP BY u.model, u.task, u.billing_provider + ORDER BY SUM(u.input_tokens) + SUM(u.output_tokens) DESC + """, (cutoff,)) + return [dict(r) for r in cur.fetchall()] + except Exception: + # Table predates the task column (older DB opened by newer code) — + # aux breakdown is simply unavailable. + return [] + + +def _merge_aux_into_by_model( + by_model: List[Dict[str, Any]], aux_rows: List[Dict[str, Any]] +) -> List[Dict[str, Any]]: + """Fold aux usage rows into the sessions-derived per-model list. + + Aux usage lives only in session_model_usage (never in the sessions + counters), so adding it here cannot double-count. Models that ONLY + appear via aux calls (e.g. a dedicated vision model) get their own + entry — previously they were entirely invisible. + """ + if not aux_rows: + return by_model + merged: Dict[str, Dict[str, Any]] = {} + for row in by_model: + merged[row.get("model") or "unknown"] = row + for aux in aux_rows: + model = aux.get("model") or "unknown" + target = merged.get(model) + if target is None: + target = { + "model": model, + "input_tokens": 0, + "output_tokens": 0, + "estimated_cost": 0, + "sessions": 0, + "api_calls": 0, + } + merged[model] = target + target["input_tokens"] = (target.get("input_tokens") or 0) + (aux.get("input_tokens") or 0) + target["output_tokens"] = (target.get("output_tokens") or 0) + (aux.get("output_tokens") or 0) + target["estimated_cost"] = (target.get("estimated_cost") or 0) + (aux.get("estimated_cost") or 0) + target["api_calls"] = (target.get("api_calls") or 0) + (aux.get("api_calls") or 0) + tasks = target.setdefault("aux_tasks", []) + tasks.append({ + "task": aux.get("task") or "", + "input_tokens": aux.get("input_tokens") or 0, + "output_tokens": aux.get("output_tokens") or 0, + "estimated_cost": aux.get("estimated_cost") or 0, + "api_calls": aux.get("api_calls") or 0, + }) + result = list(merged.values()) + result.sort( + key=lambda r: (r.get("input_tokens") or 0) + (r.get("output_tokens") or 0), + reverse=True, + ) + return result + + +def _aux_task_summary(aux_rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Aggregate aux usage rows across models into a per-task summary.""" + by_task: Dict[str, Dict[str, Any]] = {} + for aux in aux_rows: + task = aux.get("task") or "" + d = by_task.setdefault(task, { + "task": task, + "input_tokens": 0, + "output_tokens": 0, + "estimated_cost": 0, + "api_calls": 0, + "models": [], + }) + d["input_tokens"] += aux.get("input_tokens") or 0 + d["output_tokens"] += aux.get("output_tokens") or 0 + d["estimated_cost"] += aux.get("estimated_cost") or 0 + d["api_calls"] += aux.get("api_calls") or 0 + model = aux.get("model") or "unknown" + if model not in d["models"]: + d["models"].append(model) + result = list(by_task.values()) + result.sort( + key=lambda r: (r.get("input_tokens") or 0) + (r.get("output_tokens") or 0), + reverse=True, + ) + return result + + @app.get("/api/analytics/usage") async def get_usage_analytics(days: int = 30, profile: Optional[str] = None): from agent.insights import InsightsEngine @@ -14331,6 +14440,14 @@ async def get_usage_analytics(days: int = 30, profile: Optional[str] = None): """, (cutoff,)) by_model = [dict(r) for r in cur2.fetchall()] + # Fold in auxiliary usage (vision, compression, title_generation, ...) + # recorded per (model, task) in session_model_usage. Aux calls never + # touch the sessions counters, so this is add-only — no double count. + # Without it the models list shows only the main agent model even when + # aux models are actively burning tokens (issue #23270). + aux_rows = _aux_usage_rows(db, cutoff) + by_model = _merge_aux_into_by_model(by_model, aux_rows) + cur3 = db._conn.execute(""" SELECT SUM(input_tokens) as total_input, SUM(output_tokens) as total_output, @@ -14357,6 +14474,9 @@ async def get_usage_analytics(days: int = 30, profile: Optional[str] = None): return { "daily": daily, "by_model": by_model, + # Aux-task summary across models (vision, compression, ...). Lets + # the dashboard answer "what is compression costing me" directly. + "by_task": _aux_task_summary(aux_rows), "totals": totals, "period_days": days, "skills": skills, @@ -14399,6 +14519,28 @@ async def get_models_analytics(days: int = 30, profile: Optional[str] = None): """, (cutoff,)) raw_rows = [dict(r) for r in cur.fetchall()] + # Add auxiliary usage as (model, provider) rows so aux-only models + # (dedicated vision/compression models) appear on the Models page + # instead of being invisible (issue #23270). Keyed by + # model+billing_provider to match the GROUP BY above. + for aux in _aux_usage_rows(db, cutoff): + raw_rows.append({ + "model": aux.get("model") or "unknown", + "billing_provider": aux.get("billing_provider") or "", + "input_tokens": aux.get("input_tokens") or 0, + "output_tokens": aux.get("output_tokens") or 0, + "cache_read_tokens": aux.get("cache_read_tokens") or 0, + "reasoning_tokens": aux.get("reasoning_tokens") or 0, + "estimated_cost": aux.get("estimated_cost") or 0, + "actual_cost": 0, + "sessions": aux.get("sessions") or 0, + "api_calls": aux.get("api_calls") or 0, + "tool_calls": 0, + "last_used_at": aux.get("last_used_at"), + "avg_tokens_per_session": 0, + "aux_task": aux.get("task") or "", + }) + # Session rows can be created before the first billable provider call # finishes. If that early row records only the model name, and a later # row for the same model has real accounting + billing_provider, the diff --git a/hermes_state.py b/hermes_state.py index e5a517c50534..74d2d6df50b3 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -140,7 +140,7 @@ T = TypeVar("T") DEFAULT_DB_PATH = get_hermes_home() / "state.db" -SCHEMA_VERSION = 21 +SCHEMA_VERSION = 22 # Cap on user-controlled FTS5 query input before regex/sanitizer processing. # Search queries do not need to be arbitrarily large, and bounding them keeps @@ -826,6 +826,7 @@ CREATE TABLE IF NOT EXISTS session_model_usage ( billing_provider TEXT NOT NULL DEFAULT '', billing_base_url TEXT NOT NULL DEFAULT '', billing_mode TEXT NOT NULL DEFAULT '', + task TEXT NOT NULL DEFAULT '', api_call_count INTEGER NOT NULL DEFAULT 0, input_tokens INTEGER NOT NULL DEFAULT 0, output_tokens INTEGER NOT NULL DEFAULT 0, @@ -838,7 +839,7 @@ CREATE TABLE IF NOT EXISTS session_model_usage ( cost_source TEXT, first_seen REAL, last_seen REAL, - PRIMARY KEY (session_id, model, billing_provider, billing_base_url, billing_mode) + PRIMARY KEY (session_id, model, billing_provider, billing_base_url, billing_mode, task) ); CREATE TABLE IF NOT EXISTS state_meta ( @@ -1704,6 +1705,72 @@ class SessionDB: ) except sqlite3.OperationalError: pass + if current_version < 22: + # v22: task-dimension usage attribution (issue #23270). + # session_model_usage gains a ``task`` column ('' = main agent + # loop; 'vision'/'compression'/'title_generation'/... = + # auxiliary calls) so aux model spend is visible in analytics. + # The column participates in the PRIMARY KEY and SQLite cannot + # ALTER a PK, so rebuild the table. The reconciler will have + # already ADDed the plain column on legacy DBs (harmless); + # the rebuild bakes it into the PK properly. Existing rows are + # main-loop accounting by definition → task=''. + try: + legacy_pk = cursor.execute( + "SELECT COUNT(*) FROM pragma_table_info('session_model_usage') " + "WHERE name = 'task' AND pk > 0" + ).fetchone()[0] + if not legacy_pk: + cursor.execute("ALTER TABLE session_model_usage RENAME TO session_model_usage_v21") + cursor.execute( + """CREATE TABLE session_model_usage ( + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + model TEXT NOT NULL, + billing_provider TEXT NOT NULL DEFAULT '', + billing_base_url TEXT NOT NULL DEFAULT '', + billing_mode TEXT NOT NULL DEFAULT '', + task TEXT NOT NULL DEFAULT '', + api_call_count INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + estimated_cost_usd REAL NOT NULL DEFAULT 0, + actual_cost_usd REAL NOT NULL DEFAULT 0, + cost_status TEXT, + cost_source TEXT, + first_seen REAL, + last_seen REAL, + PRIMARY KEY (session_id, model, billing_provider, billing_base_url, billing_mode, task) + )""" + ) + cursor.execute( + """INSERT INTO session_model_usage ( + session_id, model, billing_provider, billing_base_url, + billing_mode, task, api_call_count, input_tokens, + output_tokens, cache_read_tokens, cache_write_tokens, + reasoning_tokens, estimated_cost_usd, actual_cost_usd, + cost_status, cost_source, first_seen, last_seen + ) + SELECT session_id, model, billing_provider, billing_base_url, + billing_mode, '', api_call_count, input_tokens, + output_tokens, cache_read_tokens, cache_write_tokens, + reasoning_tokens, estimated_cost_usd, actual_cost_usd, + cost_status, cost_source, first_seen, last_seen + FROM session_model_usage_v21""" + ) + cursor.execute("DROP TABLE session_model_usage_v21") + cursor.execute( + "CREATE INDEX IF NOT EXISTS idx_session_model_usage_session " + "ON session_model_usage(session_id)" + ) + cursor.execute( + "CREATE INDEX IF NOT EXISTS idx_session_model_usage_model " + "ON session_model_usage(model)" + ) + except sqlite3.OperationalError as exc: + logger.debug("v22 session_model_usage rebuild skipped: %s", exc) if current_version < SCHEMA_VERSION and fts_migrations_complete: cursor.execute( "UPDATE schema_version SET version = ?", @@ -2786,6 +2853,7 @@ class SessionDB: cost_status: Optional[str], cost_source: Optional[str], api_call_count: int, + task: str = "", ) -> None: """Accumulate a per-API-call usage delta into session_model_usage. @@ -2794,6 +2862,11 @@ class SessionDB: When the caller omits the model/provider (some paths only pass token deltas), fall back to the values already recorded on the session row — the same COALESCE-from-session behaviour the summary update uses. + + ``task`` distinguishes what kind of work consumed the tokens: + ``''`` (empty) is the main agent loop; auxiliary calls record their + task name (``vision``, ``compression``, ``title_generation``, ...) + via :meth:`record_auxiliary_usage` (issue #23270). """ row = conn.execute( "SELECT model, billing_provider, billing_base_url, billing_mode " @@ -2805,20 +2878,30 @@ class SessionDB: sess_base_url = row["billing_base_url"] if row is not None else None sess_billing_mode = row["billing_mode"] if row is not None else None - eff_model = model or sess_model or "unknown" - eff_provider = billing_provider or sess_provider or "" - eff_base_url = billing_base_url or sess_base_url or "" - eff_billing_mode = billing_mode or sess_billing_mode or "" + # Aux-task rows (task != '') must NOT inherit the session's main-loop + # route: an aux call may use a completely different provider/model + # (vision on gemini while the main loop runs anthropic). Missing info + # stays 'unknown'/empty rather than borrowing a misleading route. + if task: + eff_model = model or "unknown" + eff_provider = billing_provider or "" + eff_base_url = billing_base_url or "" + eff_billing_mode = billing_mode or "" + else: + eff_model = model or sess_model or "unknown" + eff_provider = billing_provider or sess_provider or "" + eff_base_url = billing_base_url or sess_base_url or "" + eff_billing_mode = billing_mode or sess_billing_mode or "" now = time.time() conn.execute( """INSERT INTO session_model_usage ( session_id, model, billing_provider, billing_base_url, billing_mode, - api_call_count, input_tokens, output_tokens, + task, api_call_count, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, estimated_cost_usd, actual_cost_usd, cost_status, cost_source, first_seen, last_seen - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(session_id, model, billing_provider, billing_base_url, billing_mode) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(session_id, model, billing_provider, billing_base_url, billing_mode, task) DO UPDATE SET api_call_count = api_call_count + excluded.api_call_count, input_tokens = input_tokens + excluded.input_tokens, @@ -2837,6 +2920,7 @@ class SessionDB: eff_provider, eff_base_url, eff_billing_mode, + task or "", api_call_count or 0, input_tokens or 0, output_tokens or 0, @@ -2863,6 +2947,65 @@ class SessionDB: self._insert_session_row(session_id, source, model=model, **kwargs) return session_id + def record_auxiliary_usage( + self, + session_id: str, + task: str, + *, + model: Optional[str] = None, + billing_provider: Optional[str] = None, + billing_base_url: Optional[str] = None, + input_tokens: int = 0, + output_tokens: int = 0, + cache_read_tokens: int = 0, + cache_write_tokens: int = 0, + reasoning_tokens: int = 0, + estimated_cost_usd: Optional[float] = None, + ) -> None: + """Record an auxiliary LLM call's usage against *session_id* (issue #23270). + + Auxiliary calls (vision, compression, title_generation, web_extract, + session_search, ...) historically discarded their usage, leaving the + dashboard's per-model analytics blind to aux model spend. This writes + a per-(model, provider, task) delta into ``session_model_usage`` — + the same table the main loop's ``update_token_counts`` feeds — WITHOUT + touching the ``sessions`` summary row. That separation is deliberate: + the gateway overwrites session counters with absolute main-loop totals, + so folding aux tokens into the summary row would either be clobbered + or double-counted. Insights/analytics read the union of both. + + Best-effort by contract: callers must never fail an aux call because + accounting failed. + """ + if not session_id or not task: + return + # FK on session_model_usage.session_id → sessions.id: ensure the row + # exists (same INSERT OR IGNORE guard update_token_counts uses — the + # initial create_session() can fail under concurrent SQLite locking). + self._insert_session_row(session_id, "unknown") + + def _do(conn): + self._record_model_usage( + conn, + session_id, + model=model, + billing_provider=billing_provider, + billing_base_url=billing_base_url, + billing_mode=None, + input_tokens=input_tokens or 0, + output_tokens=output_tokens or 0, + cache_read_tokens=cache_read_tokens or 0, + cache_write_tokens=cache_write_tokens or 0, + reasoning_tokens=reasoning_tokens or 0, + estimated_cost_usd=estimated_cost_usd, + actual_cost_usd=None, + cost_status=None, + cost_source=None, + api_call_count=1, + task=task, + ) + self._execute_write(_do) + def prune_empty_ghost_sessions(self, sessions_dir: "Optional[Path]" = None) -> int: """Remove empty TUI ghost sessions (no messages, no title, >24hr old).""" cutoff = time.time() - 86400 # Only sessions older than 24 hours diff --git a/run_agent.py b/run_agent.py index 3ba625d9cca8..993ff8032a66 100644 --- a/run_agent.py +++ b/run_agent.py @@ -5920,6 +5920,10 @@ class AIAgent: moa_config: Optional[dict[str, Any]] = None, ) -> Dict[str, Any]: """Forwarder — see ``agent.conversation_loop.run_conversation``.""" + from agent.aux_accounting import ( + reset_accounting_context, + set_accounting_context, + ) from agent.conversation_loop import run_conversation from agent.portal_tags import ( reset_conversation_context, @@ -5931,6 +5935,13 @@ class AIAgent: # (which copy this Context into their thread) — inherits the # ``conversation=`` tag with zero per-call-site plumbing. token = set_conversation_context(self._conversation_root_id()) + # Publish the session accounting handles the same way so auxiliary + # calls record their token usage into session_model_usage (task + # dimension) — the fix for aux spend being invisible in analytics + # (issue #23270). + acct_token = set_accounting_context( + getattr(self, "_session_db", None), getattr(self, "session_id", None) + ) try: return run_conversation( self, @@ -5944,6 +5955,7 @@ class AIAgent: moa_config=moa_config, ) finally: + reset_accounting_context(acct_token) reset_conversation_context(token) def chat(self, message: str, stream_callback: Optional[callable] = None) -> str: diff --git a/tests/agent/test_auxiliary_client.py b/tests/agent/test_auxiliary_client.py index c426a3815375..0fd98e9f4206 100644 --- a/tests/agent/test_auxiliary_client.py +++ b/tests/agent/test_auxiliary_client.py @@ -1474,7 +1474,7 @@ class TestAuxiliaryPoolAwareness: patch("agent.auxiliary_client._resolve_task_provider_model", return_value=("nous", "nous-model", None, None, None)), patch("agent.auxiliary_client._get_cached_client", return_value=(stale_client, "nous-model")), patch("agent.auxiliary_client.OpenAI", return_value=fresh_client), - patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task: resp), + patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._resolve_nous_runtime_api", return_value=("fresh-agent-key", "https://inference-api.nousresearch.com/v1")), ): result = call_llm( @@ -1506,7 +1506,7 @@ class TestAuxiliaryPoolAwareness: patch("agent.auxiliary_client._resolve_task_provider_model", return_value=("nous", "nous-model", None, None, None)), patch("agent.auxiliary_client._get_cached_client", return_value=(stale_client, "nous-model")), patch("agent.auxiliary_client.OpenAI", return_value=fresh_client), - patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task: resp), + patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._resolve_nous_runtime_api", return_value=("fresh-agent-key", "https://inference-api.nousresearch.com/v1")), patch( "hermes_cli.nous_account.get_nous_portal_account_info", @@ -1544,7 +1544,7 @@ class TestAuxiliaryPoolAwareness: patch("agent.auxiliary_client._resolve_task_provider_model", return_value=("nous", "nous-model", None, None, None)), patch("agent.auxiliary_client._get_cached_client", return_value=(stale_client, "nous-model")), patch("agent.auxiliary_client._to_async_client", return_value=(fresh_async_client, "nous-model")), - patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task: resp), + patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._resolve_nous_runtime_api", return_value=("fresh-agent-key", "https://inference-api.nousresearch.com/v1")), ): result = await async_call_llm( @@ -1577,7 +1577,7 @@ class TestAuxiliaryPoolAwareness: patch("agent.auxiliary_client._resolve_task_provider_model", return_value=("nous", "nous-model", None, None, None)), patch("agent.auxiliary_client._get_cached_client", return_value=(stale_client, "nous-model")), patch("agent.auxiliary_client._to_async_client", return_value=(fresh_async_client, "nous-model")), - patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task: resp), + patch("agent.auxiliary_client._validate_llm_response", side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._resolve_nous_runtime_api", return_value=("fresh-agent-key", "https://inference-api.nousresearch.com/v1")), patch( "hermes_cli.nous_account.get_nous_portal_account_info", @@ -2760,7 +2760,7 @@ class TestTransientTransportRetry: ), patch( "agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp, + side_effect=lambda resp, _task, **_kw: resp, ), ) diff --git a/tests/agent/test_auxiliary_compression_timeout_floor.py b/tests/agent/test_auxiliary_compression_timeout_floor.py index 64209ce80689..b97b4ab11d84 100644 --- a/tests/agent/test_auxiliary_compression_timeout_floor.py +++ b/tests/agent/test_auxiliary_compression_timeout_floor.py @@ -65,7 +65,7 @@ def _patches(client, *, task_timeout): patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._get_task_timeout", return_value=task_timeout), ) diff --git a/tests/agent/test_unsupported_parameter_retry.py b/tests/agent/test_unsupported_parameter_retry.py index d8f9e53c4267..288906e36705 100644 --- a/tests/agent/test_unsupported_parameter_retry.py +++ b/tests/agent/test_unsupported_parameter_retry.py @@ -102,7 +102,7 @@ class TestMaxTokensRetryHardening: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), ): with pytest.raises(RuntimeError): call_llm( @@ -129,7 +129,7 @@ class TestMaxTokensRetryHardening: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), ): with pytest.raises(RuntimeError): await async_call_llm( diff --git a/tests/agent/test_unsupported_temperature_retry.py b/tests/agent/test_unsupported_temperature_retry.py index 4d2ebb980e30..fd9c72ef97f0 100644 --- a/tests/agent/test_unsupported_temperature_retry.py +++ b/tests/agent/test_unsupported_temperature_retry.py @@ -97,7 +97,7 @@ class TestCallLlmUnsupportedTemperatureRetry: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), ): result = call_llm( task="compression", @@ -135,7 +135,7 @@ class TestCallLlmUnsupportedTemperatureRetry: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._try_payment_fallback", return_value=None), ): @@ -165,7 +165,7 @@ class TestCallLlmUnsupportedTemperatureRetry: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._try_payment_fallback", return_value=None), ): @@ -197,7 +197,7 @@ class TestAsyncCallLlmUnsupportedTemperatureRetry: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), ): result = await async_call_llm( task="session_search", @@ -232,7 +232,7 @@ class TestAsyncCallLlmUnsupportedTemperatureRetry: patch("agent.auxiliary_client._get_cached_client", return_value=(client, "gpt-5.5")), patch("agent.auxiliary_client._validate_llm_response", - side_effect=lambda resp, _task: resp), + side_effect=lambda resp, _task, **_kw: resp), patch("agent.auxiliary_client._try_payment_fallback", return_value=None), ): diff --git a/tests/hermes_state/test_aux_usage_accounting.py b/tests/hermes_state/test_aux_usage_accounting.py new file mode 100644 index 000000000000..1797e1a468ef --- /dev/null +++ b/tests/hermes_state/test_aux_usage_accounting.py @@ -0,0 +1,342 @@ +"""Tests for auxiliary usage accounting (issue #23270). + +Auxiliary LLM calls (vision, compression, title_generation, ...) record +their token usage into session_model_usage with a ``task`` dimension via +the ambient accounting context (agent/aux_accounting.py), making aux model +spend visible in analytics. +""" +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from hermes_state import SessionDB + + +@pytest.fixture +def db(tmp_path): + return SessionDB(tmp_path / "state.db") + + +def _mk_response(model="aux-model", prompt=100, completion=20): + return SimpleNamespace( + model=model, + usage=SimpleNamespace( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=prompt + completion, + ), + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))], + ) + + +def _usage_rows(db, session_id): + with db._lock: + rows = db._conn.execute( + "SELECT * FROM session_model_usage WHERE session_id = ? ORDER BY task", + (session_id,), + ).fetchall() + return [dict(r) for r in rows] + + +class TestRecordAuxiliaryUsage: + def test_records_task_row(self, db): + db.create_session("s1", source="cli") + db.record_auxiliary_usage( + "s1", "vision", model="gemini-3-flash", + billing_provider="gemini", input_tokens=500, output_tokens=50, + ) + rows = _usage_rows(db, "s1") + assert len(rows) == 1 + r = rows[0] + assert r["task"] == "vision" + assert r["model"] == "gemini-3-flash" + assert r["billing_provider"] == "gemini" + assert r["input_tokens"] == 500 + assert r["output_tokens"] == 50 + assert r["api_call_count"] == 1 + + def test_accumulates_same_task_and_model(self, db): + db.create_session("s1", source="cli") + for _ in range(3): + db.record_auxiliary_usage( + "s1", "compression", model="glm-5", input_tokens=1000, output_tokens=100, + ) + rows = _usage_rows(db, "s1") + assert len(rows) == 1 + assert rows[0]["input_tokens"] == 3000 + assert rows[0]["api_call_count"] == 3 + + def test_task_rows_do_not_touch_session_counters(self, db): + """Aux usage must NOT increment sessions.input_tokens — the gateway + overwrites those with absolute main-loop totals.""" + db.create_session("s1", source="cli") + db.record_auxiliary_usage("s1", "vision", model="m", input_tokens=999) + sess = db.get_session("s1") + assert (sess.get("input_tokens") or 0) == 0 + + def test_task_row_does_not_inherit_session_route(self, db): + """An aux call on a different provider must not borrow the session's + main-loop model/provider.""" + db.create_session("s1", source="cli", model="anthropic/claude-opus-4.6") + db.update_token_counts( + "s1", input_tokens=10, model="anthropic/claude-opus-4.6", + billing_provider="anthropic", api_call_count=1, + ) + db.record_auxiliary_usage("s1", "vision", input_tokens=5) # no model given + rows = {r["task"]: r for r in _usage_rows(db, "s1")} + assert rows["vision"]["model"] == "unknown" + assert rows["vision"]["billing_provider"] == "" + # main-loop row unaffected + assert rows[""]["model"] == "anthropic/claude-opus-4.6" + + def test_main_loop_and_aux_rows_coexist(self, db): + db.create_session("s1", source="cli") + db.update_token_counts( + "s1", input_tokens=100, output_tokens=10, + model="main-model", billing_provider="nous", api_call_count=1, + ) + db.record_auxiliary_usage( + "s1", "title_generation", model="main-model", + billing_provider="nous", input_tokens=40, output_tokens=8, + ) + rows = _usage_rows(db, "s1") + tasks = sorted(r["task"] for r in rows) + assert tasks == ["", "title_generation"] + + def test_noop_without_session_or_task(self, db): + db.record_auxiliary_usage("", "vision", input_tokens=5) + db.create_session("s1", source="cli") + db.record_auxiliary_usage("s1", "", input_tokens=5) + assert _usage_rows(db, "s1") == [] + + def test_creates_session_row_if_missing(self, db): + """FK safety: recording against a not-yet-created session must not fail.""" + db.record_auxiliary_usage("ghost", "vision", model="m", input_tokens=5) + rows = _usage_rows(db, "ghost") + assert len(rows) == 1 + + +class TestSchemaMigrationV22: + def test_v21_db_migrates_with_existing_rows(self, tmp_path): + """A legacy DB with pre-task rows migrates: rows preserved, task=''.""" + import sqlite3 as _sq + db = SessionDB(tmp_path / "state.db") + db.create_session("legacy", source="cli") + db.update_token_counts( + "legacy", input_tokens=42, model="old-model", + billing_provider="openrouter", api_call_count=1, + ) + db.close() + + # Rebuild the legacy (v21) table shape: task column absent. + conn = _sq.connect(tmp_path / "state.db") + conn.executescript(""" + CREATE TABLE smu_old AS SELECT session_id, model, billing_provider, + billing_base_url, billing_mode, api_call_count, input_tokens, + output_tokens, cache_read_tokens, cache_write_tokens, + reasoning_tokens, estimated_cost_usd, actual_cost_usd, + cost_status, cost_source, first_seen, last_seen + FROM session_model_usage; + DROP TABLE session_model_usage; + CREATE TABLE session_model_usage ( + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + model TEXT NOT NULL, + billing_provider TEXT NOT NULL DEFAULT '', + billing_base_url TEXT NOT NULL DEFAULT '', + billing_mode TEXT NOT NULL DEFAULT '', + api_call_count INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + estimated_cost_usd REAL NOT NULL DEFAULT 0, + actual_cost_usd REAL NOT NULL DEFAULT 0, + cost_status TEXT, + cost_source TEXT, + first_seen REAL, + last_seen REAL, + PRIMARY KEY (session_id, model, billing_provider, billing_base_url, billing_mode) + ); + INSERT INTO session_model_usage SELECT * FROM smu_old; + DROP TABLE smu_old; + UPDATE schema_version SET version = 21; + """) + conn.commit() + conn.close() + + # Reopen: v22 migration rebuilds the table with task in the PK. + db2 = SessionDB(tmp_path / "state.db") + with db2._lock: + pk_cols = [ + r[1] for r in db2._conn.execute( + "SELECT * FROM pragma_table_info('session_model_usage') WHERE pk > 0" + ).fetchall() + ] + row = db2._conn.execute( + "SELECT task, input_tokens FROM session_model_usage WHERE session_id = 'legacy'" + ).fetchone() + assert "task" in pk_cols + assert row is not None + assert row[0] == "" # legacy rows are main-loop accounting + assert row[1] == 42 + # And the new dimension works post-migration. + db2.record_auxiliary_usage("legacy", "vision", model="v", input_tokens=1) + db2.close() + + +class TestAmbientAccountingContext: + def test_record_aux_usage_writes_through_context(self, db): + from agent.aux_accounting import ( + record_aux_usage, + reset_accounting_context, + set_accounting_context, + ) + + db.create_session("s1", source="cli") + token = set_accounting_context(db, "s1") + try: + record_aux_usage(_mk_response(model="aux-m"), "vision", provider="gemini") + finally: + reset_accounting_context(token) + rows = _usage_rows(db, "s1") + assert len(rows) == 1 + assert rows[0]["task"] == "vision" + assert rows[0]["model"] == "aux-m" + assert rows[0]["input_tokens"] == 100 + assert rows[0]["output_tokens"] == 20 + + def test_noop_outside_context(self, db): + from agent.aux_accounting import record_aux_usage + + db.create_session("s1", source="cli") + record_aux_usage(_mk_response(), "vision") + assert _usage_rows(db, "s1") == [] + + def test_moa_tasks_excluded(self, db): + """MoA advisor usage is already folded into the main-loop delta by + conversation_loop — recording it here would double-count.""" + from agent.aux_accounting import ( + record_aux_usage, + reset_accounting_context, + set_accounting_context, + ) + + db.create_session("s1", source="cli") + token = set_accounting_context(db, "s1") + try: + record_aux_usage(_mk_response(), "moa_reference") + record_aux_usage(_mk_response(), "moa_aggregator") + finally: + reset_accounting_context(token) + assert _usage_rows(db, "s1") == [] + + def test_no_usage_object_is_noop(self, db): + from agent.aux_accounting import ( + record_aux_usage, + reset_accounting_context, + set_accounting_context, + ) + + db.create_session("s1", source="cli") + resp = SimpleNamespace(model="m", choices=[]) + token = set_accounting_context(db, "s1") + try: + record_aux_usage(resp, "vision") + finally: + reset_accounting_context(token) + assert _usage_rows(db, "s1") == [] + + def test_recording_failure_never_raises(self, db): + from agent.aux_accounting import ( + record_aux_usage, + reset_accounting_context, + set_accounting_context, + ) + + class ExplodingDB: + def record_auxiliary_usage(self, *a, **kw): + raise RuntimeError("disk full") + + token = set_accounting_context(ExplodingDB(), "s1") + try: + record_aux_usage(_mk_response(), "vision") # must not raise + finally: + reset_accounting_context(token) + + def test_validate_llm_response_records(self, db): + """The aux client's validation chokepoint feeds the recorder.""" + from agent.aux_accounting import ( + reset_accounting_context, + set_accounting_context, + ) + from agent.auxiliary_client import _validate_llm_response + + db.create_session("s1", source="cli") + token = set_accounting_context(db, "s1") + try: + out = _validate_llm_response(_mk_response(), "web_extract", provider="openrouter") + finally: + reset_accounting_context(token) + assert out is not None + rows = _usage_rows(db, "s1") + assert len(rows) == 1 + assert rows[0]["task"] == "web_extract" + assert rows[0]["billing_provider"] == "openrouter" + + def test_context_isolated_between_copied_contexts(self, db): + import contextvars + + from agent.aux_accounting import get_accounting_context, set_accounting_context + + def _set_and_get(sid): + set_accounting_context(db, sid) + return get_accounting_context()[1] + + a = contextvars.copy_context().run(_set_and_get, "agent-a") + b = contextvars.copy_context().run(_set_and_get, "agent-b") + assert (a, b) == ("agent-a", "agent-b") + assert get_accounting_context() is None + + +class TestAnalyticsAuxRows: + def test_aux_usage_rows_and_merge(self, db): + from hermes_cli.web_server import ( + _aux_task_summary, + _aux_usage_rows, + _merge_aux_into_by_model, + ) + + db.create_session("s1", source="cli") + db.update_token_counts( + "s1", input_tokens=1000, output_tokens=100, + model="main-model", billing_provider="nous", api_call_count=1, + ) + db.record_auxiliary_usage( + "s1", "vision", model="vision-model", + billing_provider="gemini", input_tokens=300, output_tokens=30, + ) + db.record_auxiliary_usage( + "s1", "compression", model="main-model", + billing_provider="nous", input_tokens=200, output_tokens=20, + ) + + aux = _aux_usage_rows(db, cutoff=0) + assert {r["task"] for r in aux} == {"vision", "compression"} + + by_model = [{ + "model": "main-model", "input_tokens": 1000, "output_tokens": 100, + "estimated_cost": 0, "sessions": 1, "api_calls": 1, + }] + merged = _merge_aux_into_by_model(by_model, aux) + by_name = {r["model"]: r for r in merged} + # vision-only model surfaces as its own entry + assert "vision-model" in by_name + assert by_name["vision-model"]["input_tokens"] == 300 + # compression folded into the main model's totals + assert by_name["main-model"]["input_tokens"] == 1200 + assert by_name["main-model"]["api_calls"] == 2 + + tasks = _aux_task_summary(aux) + assert {t["task"] for t in tasks} == {"vision", "compression"} diff --git a/tests/run_agent/test_moa_streaming.py b/tests/run_agent/test_moa_streaming.py index a61780aca088..03d4c2f5e420 100644 --- a/tests/run_agent/test_moa_streaming.py +++ b/tests/run_agent/test_moa_streaming.py @@ -207,7 +207,7 @@ def test_call_llm_non_stream_still_validates(monkeypatch): validated = {"called": False} - def _validate(resp, task): + def _validate(resp, task, provider=None, base_url=None): validated["called"] = True return resp