From b6ca4fc856c8e8bc3bc922ed48957af853fa3dfa Mon Sep 17 00:00:00 2001 From: RelaxJonh <92573950+RelaxJonh@users.noreply.github.com> Date: Fri, 31 Jul 2026 21:59:20 -0700 Subject: [PATCH] fix(state): heal session_model_usage PK unconditionally to restore token/cost accounting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- hermes_state_schema.py | 126 +++++++++++++ .../state/test_session_model_usage_pk_heal.py | 170 ++++++++++++++++++ 2 files changed, 296 insertions(+) create mode 100644 tests/state/test_session_model_usage_pk_heal.py diff --git a/hermes_state_schema.py b/hermes_state_schema.py index 80a4494337..15513a7e73 100644 --- a/hermes_state_schema.py +++ b/hermes_state_schema.py @@ -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 diff --git a/tests/state/test_session_model_usage_pk_heal.py b/tests/state/test_session_model_usage_pk_heal.py new file mode 100644 index 0000000000..5cc6a568d8 --- /dev/null +++ b/tests/state/test_session_model_usage_pk_heal.py @@ -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()