mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-25 17:18:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6
hermes-agent/tests/hermes_state/test_aux_usage_accounting.py
Teknium 21dedb8586
fix(insights): include auxiliary usage in overview token totals (#65603) 
The overview's total_input/output/cache token counts summed only the
sessions counters (main-loop usage), while the per-model breakdown
already included auxiliary usage rows (task dimension from #65537) and
reconciled residuals. Result: hermes insights top-line totals
undercounted aux spend (compression summarizer, vision, titles) and
disagreed with the per-model table below them — the symptom reported
in #58592 and requested in #9979.

When the per-model breakdown is available, derive the overview token
totals from it (same pattern total_cost already used). Verified no
double-count across incremental CLI deltas, gateway absolute
overwrites, and aux rows.
2026-07-16 05:39:33 -07:00

388 lines
15 KiB
Python
Raw Blame History

"""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"}
class TestInsightsAuxTotals:
def test_overview_totals_include_aux_usage(self, db):
"""`hermes insights` overview must count aux tokens, not just the
sessions counters (issues #58592, #9979)."""
from agent.insights import InsightsEngine
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", "compression", model="glm-5",
billing_provider="openrouter", input_tokens=5000, output_tokens=500,
)
report = InsightsEngine(db).generate(days=30)
ov = report["overview"]
assert ov["total_input_tokens"] == 6000
assert ov["total_output_tokens"] == 600
models = {m["model"] for m in report["models"]}
assert {"main-model", "glm-5"} <= models
def test_overview_totals_not_double_counted_with_absolute_updates(self, db):
"""Gateway absolute overwrites + aux rows must not inflate totals."""
from agent.insights import InsightsEngine
db.create_session("s2", source="telegram")
db.update_token_counts(
"s2", input_tokens=2000, output_tokens=200,
model="main-model", billing_provider="nous", api_call_count=1,
)
db.update_token_counts(
"s2", input_tokens=2000, output_tokens=200,
model="main-model", billing_provider="nous",
absolute=True, api_call_count=1,
)
db.record_auxiliary_usage(
"s2", "title_generation", model="main-model",
billing_provider="nous", input_tokens=40, output_tokens=8,
)
report = InsightsEngine(db).generate(days=30)
ov = report["overview"]
assert ov["total_input_tokens"] == 2040
assert ov["total_output_tokens"] == 208
Reference in a new issue View git blame Copy permalink