hermes-admin/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 28 Packages Projects Releases Wiki Activity

fix(state): heal session_model_usage PK unconditionally to restore token/cost accounting

Browse Source 
Installs whose state.db reached schema_version >= 22 before the task
dimension was added carry a 5-column PRIMARY KEY on
session_model_usage. The column reconciler ADDs task as a bare
nullable, but SQLite cannot ALTER a primary key, and the version-gated
v22 rebuild is unreachable (current_version < 22 already false), so
the composite 6-column key never lands. Every upsert in
_record_model_usage then fails with 'ON CONFLICT clause does not match
any PRIMARY KEY or UNIQUE constraint', aborting the enclosing write
transaction — token/cost accounting permanently dead (#73823).

Add an idempotent _heal_session_model_usage_pk() modeled on
_heal_gateway_routing_pk(), run unconditionally from _init_schema on
every open. Salvaged from #73838 with fix-ups:

- ported to SessionSchemaMixin in hermes_state_schema.py (the schema
  code moved out of hermes_state.py in 21c7ae8563; the PR targeted the
  old location)
- rebuild wrapped in a PRAGMA foreign_keys=OFF/ON window: the
  connection enables FKs before _init_schema and OR IGNORE does NOT
  suppress FK violations, so a single orphaned usage row (session
  pruned while accounting was broken) would have aborted the heal
- COALESCE('') on the nullable reconciler-added task column (and the
  billing columns) during the copy
- stale-v22+ regression tests: rebuilt PK + restored upsert, orphan
  rows survive the FK window, healthy-DB no-op, no legacy leftover

Fixes #73823
This commit is contained in:
RelaxJonh
2026-07-31 21:59:20 -07:00
committed by Teknium
parent 14eca89779
commit b6ca4fc856
2 changed files with 296 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
hermes_state_schema.py
View File

@@ -290,6 +290,126 @@ class SessionSchemaMixin:
)
cursor.execute("DROP TABLE gateway_routing_legacy_pk")
def _heal_session_model_usage_pk(self, cursor: sqlite3.Cursor) -> None:
"""Rebuild ``session_model_usage`` when its PRIMARY KEY lacks ``task``.
Installs whose ``state.db`` reached ``schema_version >= 22`` before
the ``task`` dimension was added carry a 5-column PRIMARY KEY
``(session_id, model, billing_provider, billing_base_url,
billing_mode)``. ``_reconcile_columns()`` ADDs the ``task`` column
as a bare nullable, but SQLite cannot ALTER a primary key, so the
shipped composite 6-column key never lands. The version-gated v22
rebuild is unreachable on those installs (``current_version < 22``
is already false), so every upsert in ``_record_model_usage()``
fails with "ON CONFLICT clause does not match any PRIMARY KEY or
UNIQUE constraint" — aborting the enclosing write transaction and
silently zeroing all token *and* cost accounting (#73823).
Idempotent; runs unconditionally on every open, same pattern as
:meth:`_heal_gateway_routing_pk` above. On healthy databases the
PRAGMA check short-circuits and this is a no-op.
"""
try:
rows = cursor.execute(
'PRAGMA table_info("session_model_usage")'
).fetchall()
except sqlite3.OperationalError:
return
if not rows:
# Table doesn't exist yet — SCHEMA_SQL creates it correctly.
return
def _col(row, idx, name):
return row[idx] if isinstance(row, (tuple, list)) else row[name]
pk_cols = {
_col(r, 1, "name") for r in rows if _col(r, 5, "pk")
}
if "task" in pk_cols:
# task is already in the PK — healthy.
return
logger.info(
"session_model_usage has legacy primary key %r (missing task); "
"rebuilding with composite 6-column key",
sorted(pk_cols),
)
# FK-off window: the connection enables PRAGMA foreign_keys=ON
# before _init_schema runs, and session_model_usage.session_id
# REFERENCES sessions(id). INSERT OR IGNORE does NOT suppress
# foreign-key violations (OR IGNORE only covers uniqueness/NOT
# NULL conflicts), so an orphaned usage row — possible after a
# partial prune while accounting was broken — would abort the
# whole rebuild. Disable FK enforcement for the copy and restore
# it afterwards. PRAGMA foreign_keys is a no-op inside a
# transaction, which is fine here: _init_schema runs on an
# isolation_level=None connection with no transaction open.
cursor.execute("PRAGMA foreign_keys=OFF")
try:
cursor.execute(
"ALTER TABLE session_model_usage "
"RENAME TO session_model_usage_legacy_pk"
)
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)
)"""
)
# OR IGNORE: while the PK was wrong the reconciler may have left
# ``task`` NULL on old rows; COALESCE to '' can theoretically
# collide with a genuine ''-task row — keep the first, drop the
# duplicate rather than fail the heal.
cursor.execute(
"""INSERT OR IGNORE 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,
COALESCE(billing_provider, ''),
COALESCE(billing_base_url, ''),
COALESCE(billing_mode, ''),
COALESCE(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
FROM session_model_usage_legacy_pk"""
)
cursor.execute("DROP TABLE session_model_usage_legacy_pk")
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("session_model_usage PK heal skipped: %s", exc)
finally:
cursor.execute("PRAGMA foreign_keys=ON")
def _init_schema(self):
"""Create tables and FTS if they don't exist, reconcile columns.
@@ -319,6 +439,12 @@ class SessionSchemaMixin:
# the one table-shape repair reconciliation can't express.
self._heal_gateway_routing_pk(cursor)
# Rebuild session_model_usage if its PRIMARY KEY lacks the ``task``
# column (5-column PK on installs already at v22+ when the column
# landed — the version-gated rebuild is unreachable there, #73823).
# Same PK-rebuild constraint as gateway_routing above.
self._heal_session_model_usage_pk(cursor)
# Indexes that reference reconciler-added columns must be created
# AFTER _reconcile_columns runs — declaring them in SCHEMA_SQL
# makes the initial executescript fail on legacy DBs (the index's

170
tests/state/test_session_model_usage_pk_heal.py Normal file
View File

@@ -0,0 +1,170 @@
"""Unconditional session_model_usage PK heal (#73823, salvage of #73838).
Installs whose state.db reached ``schema_version >= 22`` before the
``task`` dimension was added carry a 5-column PRIMARY KEY on
``session_model_usage``. The column reconciler ADDs ``task`` as a bare
nullable, but SQLite cannot ALTER a primary key, and the version-gated
v22 rebuild is unreachable (``current_version < 22`` already false), so
the composite 6-column key never lands. Every upsert in
``_record_model_usage`` then fails with "ON CONFLICT clause does not
match any PRIMARY KEY or UNIQUE constraint", aborting the enclosing
write transaction — token/cost accounting permanently dead.
``_heal_session_model_usage_pk`` runs unconditionally on every open
(same pattern as ``_heal_gateway_routing_pk``) and rebuilds the table
once, inside an FK-off window (OR IGNORE does not suppress FK
violations and the connection enables foreign_keys before init).
"""
import sqlite3
import pytest
from hermes_state import SessionDB
from hermes_state_common import SCHEMA_VERSION
LEGACY_SQL = """
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)
)
"""
def _make_stale_v22_db(tmp_path, usage_rows=(), sessions=("s1",)):
"""Build a state.db in the stale-v22+ shape: current schema everywhere,
but session_model_usage carrying the legacy 5-column PK with ``task``
reconciler-appended OUTSIDE the key, and schema_version already at
current — so the version-gated v22 rebuild can never run."""
db_path = tmp_path / "state.db"
# Born-current DB for everything else...
db = SessionDB(db_path=db_path)
for sid in sessions:
db.create_session(sid, "cli")
db.close()
# ...then regress session_model_usage to the legacy shape.
conn = sqlite3.connect(db_path)
conn.execute("DROP TABLE session_model_usage")
conn.execute(LEGACY_SQL)
# Mimic the column reconciler: task appended OUTSIDE the primary key.
conn.execute('ALTER TABLE session_model_usage ADD COLUMN "task" TEXT')
conn.executemany(
"INSERT INTO session_model_usage "
"(session_id, model, input_tokens, output_tokens) VALUES (?, ?, ?, ?)",
list(usage_rows),
)
# Version already current: proves the heal does not depend on the gate.
conn.execute("UPDATE schema_version SET version = ?", (SCHEMA_VERSION,))
conn.commit()
conn.close()
return db_path
def _pk_cols(db):
rows = db._conn.execute(
'PRAGMA table_info("session_model_usage")'
).fetchall()
return sorted(r["name"] for r in rows if r["pk"])
class TestSessionModelUsagePkHeal:
def test_stale_v22_pk_rebuilt_and_accounting_restored(self, tmp_path):
"""The broken-PK table is rebuilt on open even though schema_version
is already current, and the usage upsert works again."""
db_path = _make_stale_v22_db(
tmp_path, usage_rows=[("s1", "m-old", 10, 20)]
)
db = SessionDB(db_path=db_path)
try:
assert "task" in _pk_cols(db)
# Existing rows survive the rebuild (task backfilled to '').
row = db._conn.execute(
"SELECT task, input_tokens FROM session_model_usage "
"WHERE session_id='s1' AND model='m-old'"
).fetchone()
assert row is not None
assert row["task"] == ""
assert row["input_tokens"] == 10
# The killed write path works again: the upsert used to abort
# the whole transaction with an ON CONFLICT mismatch.
db.update_token_counts(
"s1", input_tokens=5, output_tokens=7,
model="m-new", billing_provider="p", api_call_count=1,
)
row = db._conn.execute(
"SELECT input_tokens FROM session_model_usage "
"WHERE session_id='s1' AND model='m-new'"
).fetchone()
assert row is not None and row["input_tokens"] == 5
finally:
db.close()
def test_orphan_rows_survive_fk_enforcement(self, tmp_path):
"""The rebuild copies rows inside an FK-off window: an orphaned
usage row (session pruned while accounting was broken) must not
abort the heal — OR IGNORE does NOT suppress FK violations."""
db_path = _make_stale_v22_db(
tmp_path,
usage_rows=[("s1", "m1", 1, 1), ("ghost-session", "m1", 2, 2)],
)
db = SessionDB(db_path=db_path)
try:
assert "task" in _pk_cols(db)
rows = db._conn.execute(
"SELECT session_id FROM session_model_usage ORDER BY session_id"
).fetchall()
assert [r["session_id"] for r in rows] == ["ghost-session", "s1"]
# FK enforcement is restored after the heal window.
assert db._conn.execute("PRAGMA foreign_keys").fetchone()[0] == 1
finally:
db.close()
def test_healthy_db_is_a_noop(self, tmp_path):
"""A DB born with the composite PK is left untouched (idempotence)."""
db = SessionDB(db_path=tmp_path / "state.db")
try:
db.create_session("s1", "cli")
db.update_token_counts(
"s1", input_tokens=3, model="m", billing_provider="p",
api_call_count=1,
)
assert "task" in _pk_cols(db)
# Re-running the heal directly is a no-op.
cur = db._conn.cursor()
db._heal_session_model_usage_pk(cur)
row = db._conn.execute(
"SELECT input_tokens FROM session_model_usage "
"WHERE session_id='s1'"
).fetchone()
assert row is not None and row["input_tokens"] == 3
finally:
db.close()
def test_no_legacy_leftover_table(self, tmp_path):
"""The rename-copy-drop leaves no *_legacy_pk residue behind."""
db_path = _make_stale_v22_db(tmp_path, usage_rows=[("s1", "m1", 1, 1)])
db = SessionDB(db_path=db_path)
try:
left = db._conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name='session_model_usage_legacy_pk'"
).fetchone()
assert left is None
finally:
db.close()