mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-17 14:42:06 +00:00
feat(agent): track per-model token usage for mid-session model switches
The `sessions` table records only the initial (model, billing_provider) for a session, so when a user switches models mid-session (via `/model` or programmatically) every token — including the switched model's — is attributed to the first model. Insights/billing reports then hide the cost of the new model entirely (e.g. a session that started on deepseek and switched to opus shows $0 for opus). Add a `session_model_usage` table keyed (session_id, model, billing_provider) that accumulates each per-API-call delta under the model active at the time of the call. `update_token_counts()` is the single chokepoint every per-call delta flows through (CLI, gateway, cron, delegated, codex), so recording there captures accurate attribution on every platform. Only the incremental path records — the gateway's `absolute=True` summary overwrite is skipped to avoid double-counting cumulative totals that can't be split per model. When a call omits the model, it falls back to the session's recorded model, matching the existing COALESCE-from-session summary behaviour. Insights `_compute_model_breakdown` now aggregates tokens and cost from `session_model_usage`, so a switched session splits correctly across models, with a defensive fallback to the per-session aggregate for any session lacking usage rows. A v17 migration backfills one usage row per existing token-bearing session from its aggregate totals (idempotent via INSERT OR IGNORE), validated lossless against a 1.3 GB production DB. Tests: per-model recording, mid-session split, model fallback, absolute no-double-count, v17 backfill, and an insights-level switch breakdown. Fixes #51607. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
022c4991fc
commit
cb7f6bbb2e
4 changed files with 438 additions and 20 deletions
166
hermes_state.py
166
hermes_state.py
|
|
@ -122,7 +122,7 @@ T = TypeVar("T")
|
|||
|
||||
DEFAULT_DB_PATH = get_hermes_home() / "state.db"
|
||||
|
||||
SCHEMA_VERSION = 19
|
||||
SCHEMA_VERSION = 20
|
||||
|
||||
# Cap on user-controlled FTS5 query input before regex/sanitizer processing.
|
||||
# Search queries do not need to be arbitrarily large, and bounding them keeps
|
||||
|
|
@ -768,6 +768,23 @@ CREATE TABLE IF NOT EXISTS messages (
|
|||
compacted INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS session_model_usage (
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id),
|
||||
model TEXT NOT NULL,
|
||||
billing_provider TEXT NOT NULL DEFAULT '',
|
||||
billing_base_url TEXT,
|
||||
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,
|
||||
first_seen REAL,
|
||||
last_seen REAL,
|
||||
PRIMARY KEY (session_id, model, billing_provider)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS state_meta (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT
|
||||
|
|
@ -794,6 +811,8 @@ CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id);
|
|||
CREATE INDEX IF NOT EXISTS idx_sessions_started ON sessions(started_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, timestamp);
|
||||
CREATE INDEX IF NOT EXISTS idx_compression_locks_expires ON compression_locks(expires_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_session_model_usage_session ON session_model_usage(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_session_model_usage_model ON session_model_usage(model);
|
||||
"""
|
||||
|
||||
# Indexes that reference columns added in later schema versions must be
|
||||
|
|
@ -1554,6 +1573,45 @@ class SessionDB:
|
|||
# means consumers fall back to sessions.json for those
|
||||
# rows until the gateway rewrites them.
|
||||
logger.debug("v18 gateway metadata backfill skipped: %s", exc)
|
||||
if current_version < 20:
|
||||
# v20: per-model usage attribution (issue #51607). Going
|
||||
# forward update_token_counts() records each API call into
|
||||
# session_model_usage keyed by the live model, but existing
|
||||
# sessions only have their aggregate totals on the sessions
|
||||
# row. Seed one usage row per historical session from those
|
||||
# aggregates so insights reads uniformly from the new table.
|
||||
# INSERT OR IGNORE keeps it idempotent: if newer code already
|
||||
# wrote a (session_id, model, provider) row for a session, the
|
||||
# PK conflict skips the stale aggregate rather than doubling it.
|
||||
try:
|
||||
cursor.execute(
|
||||
"""INSERT OR IGNORE INTO session_model_usage (
|
||||
session_id, model, billing_provider,
|
||||
billing_base_url, api_call_count, input_tokens,
|
||||
output_tokens, cache_read_tokens,
|
||||
cache_write_tokens, reasoning_tokens,
|
||||
estimated_cost_usd, first_seen, last_seen
|
||||
)
|
||||
SELECT id, COALESCE(model, 'unknown'),
|
||||
COALESCE(billing_provider, ''),
|
||||
billing_base_url,
|
||||
COALESCE(api_call_count, 0),
|
||||
COALESCE(input_tokens, 0),
|
||||
COALESCE(output_tokens, 0),
|
||||
COALESCE(cache_read_tokens, 0),
|
||||
COALESCE(cache_write_tokens, 0),
|
||||
COALESCE(reasoning_tokens, 0),
|
||||
COALESCE(estimated_cost_usd, 0),
|
||||
started_at, COALESCE(ended_at, started_at)
|
||||
FROM sessions
|
||||
WHERE COALESCE(input_tokens, 0)
|
||||
+ COALESCE(output_tokens, 0)
|
||||
+ COALESCE(cache_read_tokens, 0)
|
||||
+ COALESCE(cache_write_tokens, 0)
|
||||
+ COALESCE(reasoning_tokens, 0) > 0"""
|
||||
)
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
if current_version < SCHEMA_VERSION and fts_migrations_complete:
|
||||
cursor.execute(
|
||||
"UPDATE schema_version SET version = ?",
|
||||
|
|
@ -2500,10 +2558,116 @@ class SessionDB:
|
|||
api_call_count,
|
||||
session_id,
|
||||
)
|
||||
# Per-model usage attribution. ``update_token_counts`` is the single
|
||||
# chokepoint every per-API-call delta flows through (CLI, gateway, cron,
|
||||
# delegated runs — see conversation_loop / codex_runtime), and each call
|
||||
# carries the model/provider *active at the time of that call*. The
|
||||
# ``sessions`` row only keeps one (model, billing_provider) pair, so a
|
||||
# mid-session ``/model`` switch otherwise attributes every token to the
|
||||
# initial model (issue #51607). Recording the per-call delta into
|
||||
# session_model_usage keyed by the live model preserves an accurate
|
||||
# per-model breakdown regardless of how many times the user switches.
|
||||
#
|
||||
# Only the incremental path records here: the gateway also issues an
|
||||
# ``absolute=True`` call that overwrites the sessions summary totals
|
||||
# with the cached agent's cumulative figures — folding those in would
|
||||
# double-count, and cumulative totals can't be split back per model.
|
||||
record_model_usage = (not absolute) and (
|
||||
input_tokens or output_tokens or cache_read_tokens
|
||||
or cache_write_tokens or reasoning_tokens or api_call_count
|
||||
or estimated_cost_usd
|
||||
)
|
||||
|
||||
def _do(conn):
|
||||
conn.execute(sql, params)
|
||||
if record_model_usage:
|
||||
self._record_model_usage(
|
||||
conn,
|
||||
session_id,
|
||||
model=model,
|
||||
billing_provider=billing_provider,
|
||||
billing_base_url=billing_base_url,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
estimated_cost_usd=estimated_cost_usd,
|
||||
api_call_count=api_call_count,
|
||||
)
|
||||
self._execute_write(_do)
|
||||
|
||||
def _record_model_usage(
|
||||
self,
|
||||
conn,
|
||||
session_id: str,
|
||||
*,
|
||||
model: Optional[str],
|
||||
billing_provider: Optional[str],
|
||||
billing_base_url: Optional[str],
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_write_tokens: int,
|
||||
reasoning_tokens: int,
|
||||
estimated_cost_usd: Optional[float],
|
||||
api_call_count: int,
|
||||
) -> None:
|
||||
"""Accumulate a per-API-call usage delta into session_model_usage.
|
||||
|
||||
Runs inside the caller's write transaction (after the ``sessions``
|
||||
UPDATE) so the per-model rows stay consistent with the summary row.
|
||||
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.
|
||||
"""
|
||||
row = conn.execute(
|
||||
"SELECT model, billing_provider, billing_base_url "
|
||||
"FROM sessions WHERE id = ?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
sess_model = row["model"] if row is not None else None
|
||||
sess_provider = row["billing_provider"] if row is not None else None
|
||||
sess_base_url = row["billing_base_url"] 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
|
||||
now = time.time()
|
||||
conn.execute(
|
||||
"""INSERT INTO session_model_usage (
|
||||
session_id, model, billing_provider, billing_base_url,
|
||||
api_call_count, input_tokens, output_tokens,
|
||||
cache_read_tokens, cache_write_tokens, reasoning_tokens,
|
||||
estimated_cost_usd, first_seen, last_seen
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(session_id, model, billing_provider) DO UPDATE SET
|
||||
api_call_count = api_call_count + excluded.api_call_count,
|
||||
input_tokens = input_tokens + excluded.input_tokens,
|
||||
output_tokens = output_tokens + excluded.output_tokens,
|
||||
cache_read_tokens = cache_read_tokens + excluded.cache_read_tokens,
|
||||
cache_write_tokens = cache_write_tokens + excluded.cache_write_tokens,
|
||||
reasoning_tokens = reasoning_tokens + excluded.reasoning_tokens,
|
||||
estimated_cost_usd = estimated_cost_usd + excluded.estimated_cost_usd,
|
||||
billing_base_url = COALESCE(excluded.billing_base_url, billing_base_url),
|
||||
last_seen = excluded.last_seen""",
|
||||
(
|
||||
session_id,
|
||||
eff_model,
|
||||
eff_provider,
|
||||
eff_base_url,
|
||||
api_call_count or 0,
|
||||
input_tokens or 0,
|
||||
output_tokens or 0,
|
||||
cache_read_tokens or 0,
|
||||
cache_write_tokens or 0,
|
||||
reasoning_tokens or 0,
|
||||
float(estimated_cost_usd or 0.0),
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
|
||||
def ensure_session(
|
||||
self,
|
||||
session_id: str,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue