From 63014a487d5f3c8eb645efdae0fb9008863e5f04 Mon Sep 17 00:00:00 2001 From: h Date: Tue, 1 Sep 2026 15:37:53 +0200 Subject: [PATCH] fix(storage,api): usage cost as per-turn delta of the sdk running total, backfill --- src/beaver_gateway/frontends/api/frontend.py | 7 +- src/beaver_gateway/storage/__init__.py | 4 + src/beaver_gateway/storage/backfill_usage.py | 37 +++++ src/beaver_gateway/storage/db.py | 136 ++++++++++++++++++- src/beaver_gateway/storage/models.py | 20 ++- tests/test_usage_deltas.py | 105 ++++++++++++++ 6 files changed, 304 insertions(+), 5 deletions(-) create mode 100644 src/beaver_gateway/storage/backfill_usage.py create mode 100644 tests/test_usage_deltas.py diff --git a/src/beaver_gateway/frontends/api/frontend.py b/src/beaver_gateway/frontends/api/frontend.py index 5906edc..cc46705 100644 --- a/src/beaver_gateway/frontends/api/frontend.py +++ b/src/beaver_gateway/frontends/api/frontend.py @@ -6,9 +6,10 @@ write goes through ``core/conversations``; the frontend only shapes JSON. bus as SSE with the same keepalive the markdown frontend uses, so a proxy never sees a silent socket. -Usage figures come from the ``usage`` table (one row per turn, API-price -``cost_usd`` and per-model ``model_usage`` from the SDK's -``ResultMessage``); subscription quotas come from ``rate_limits`` +Usage figures come from the ``usage`` table (one row per turn; API-price +``cost_usd`` and per-model ``model_usage`` are per-turn deltas of the SDK's +cumulative ``ResultMessage`` counters, see ``storage.append_usage``); +subscription quotas come from ``rate_limits`` (``RateLimitEvent``). The quota covers the whole subscription, so ``/api/limits`` puts the gateway's own spend for the window next to it for calibration by eye. diff --git a/src/beaver_gateway/storage/__init__.py b/src/beaver_gateway/storage/__init__.py index 4782f34..7477011 100644 --- a/src/beaver_gateway/storage/__init__.py +++ b/src/beaver_gateway/storage/__init__.py @@ -9,12 +9,14 @@ from beaver_gateway.storage.db import ( Database, append_audit, append_usage, + backfill_usage_deltas, create_token, list_active_tokens, list_audit_records, list_tokens, revoke_token, touch_token, + usage_deltas, ) from beaver_gateway.storage.models import ( AuditLog, @@ -45,10 +47,12 @@ __all__ = [ "Usage", "append_audit", "append_usage", + "backfill_usage_deltas", "create_token", "list_active_tokens", "list_audit_records", "list_tokens", "revoke_token", "touch_token", + "usage_deltas", ] diff --git a/src/beaver_gateway/storage/backfill_usage.py b/src/beaver_gateway/storage/backfill_usage.py new file mode 100644 index 0000000..dc0172c --- /dev/null +++ b/src/beaver_gateway/storage/backfill_usage.py @@ -0,0 +1,37 @@ +"""Rewrite ``usage.cost_usd`` / ``model_usage`` as per-turn deltas. + +Until 2026-09-01 the gateway stored ``ResultMessage.total_cost_usd`` verbatim, +which is cumulative for the claude process - summing the column overstated a +day by an order of magnitude. Run once after deploying the delta-aware +:func:`beaver_gateway.storage.append_usage`:: + + python -m beaver_gateway.storage.backfill_usage + +Reads ``DATABASE_URL`` like the gateway (``.env`` included); safe to rerun. +""" + +from __future__ import annotations + +import asyncio + +from dotenv import load_dotenv + +from beaver_gateway.settings import Settings +from beaver_gateway.storage.db import Database, backfill_usage_deltas + + +async def _main() -> None: + load_dotenv(override=False) + settings = Settings() # ty: ignore[missing-argument] + db = Database(settings.database_url) + try: + await db.create_all() + async with db.session() as session: + changed = await backfill_usage_deltas(session) + finally: + await db.dispose() + print(f"usage rows rewritten: {changed}") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/src/beaver_gateway/storage/db.py b/src/beaver_gateway/storage/db.py index b51dcfe..39f3e32 100644 --- a/src/beaver_gateway/storage/db.py +++ b/src/beaver_gateway/storage/db.py @@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any from sqlalchemy import inspect, text from sqlalchemy.ext.asyncio import create_async_engine -from sqlmodel import SQLModel, select +from sqlmodel import SQLModel, col, select from sqlmodel.ext.asyncio.session import AsyncSession from beaver_gateway.storage.models import AuditLog, Token, Usage @@ -227,18 +227,152 @@ async def list_audit_records( async def append_usage(session: AsyncSession, row: Usage) -> None: + """Persist a turn's usage with ``cost_usd`` / ``model_usage`` as per-turn deltas. + + ``row`` arrives with the SDK's cumulative ``total_cost_usd`` and + ``model_usage`` in ``cost_usd`` / ``model_usage``; they move to the + ``*_total`` columns and the stored values become the difference against + the previous row of the same session (see :func:`usage_deltas`). + """ + prev = await _last_usage(session, row.session_id) + usage_deltas(row, prev) session.add(row) await session.commit() +async def _last_usage(session: AsyncSession, session_id: str | None) -> Usage | None: + if not session_id: + return None + stmt = ( + select(Usage) + .where(col(Usage.session_id) == session_id) + .order_by(col(Usage.ts).desc(), col(Usage.id).desc()) + .limit(1) + ) + return (await session.exec(stmt)).first() + + +# Counters inside ``ResultMessage.model_usage[model]``; everything else there +# (provider, costBasis, contextWindow, maxOutputTokens, ...) is a constant. +_MODEL_USAGE_COUNTERS = ( + "inputTokens", + "outputTokens", + "cacheReadInputTokens", + "cacheCreationInputTokens", + "webSearchRequests", + "costUSD", +) + + +def usage_deltas(row: Usage, prev: Usage | None) -> None: + """Turn the SDK's cumulative cost counters on ``row`` into per-turn deltas. + + ``total_cost_usd`` and ``model_usage`` in a ``ResultMessage`` are running + totals of the claude process: they grow with every turn of the session and + start over when the process is respawned (rotation, resume after a + restart). A row fresh from the SDK (``cost_usd_total is None``) first moves + the raw values into ``cost_usd_total`` / ``model_usage_total``; then + ``cost_usd`` / ``model_usage`` become ``current - previous`` against + ``prev`` (the last row of the same session), or the raw value when there + is no previous row or the counter went backwards - a fresh process. + Idempotent on rows that already carry their totals, which is what the + backfill relies on. + """ + if row.cost_usd_total is None and row.model_usage_total is None: + row.cost_usd_total = row.cost_usd + row.model_usage_total = row.model_usage + prev_cost = prev.cost_usd_total if prev is not None else None + prev_models = prev.model_usage_total if prev is not None else None + row.cost_usd = _counter_delta(row.cost_usd_total, prev_cost) + row.model_usage = _model_usage_delta(row.model_usage_total, prev_models) + + +def _counter_delta(current: float | None, previous: float | None) -> float | None: + if current is None: + return None + if previous is None or current < previous: + return current + return round(current - previous, 6) + + +def _model_usage_delta( + current: dict[str, Any] | None, previous: dict[str, Any] | None +) -> dict[str, Any] | None: + if current is None: + return None + out: dict[str, Any] = {} + for model, fields in current.items(): + if not isinstance(fields, dict): + out[model] = fields + continue + before = (previous or {}).get(model) + before = before if isinstance(before, dict) else {} + delta = dict(fields) + for key in _MODEL_USAGE_COUNTERS: + if key not in fields: + continue + cur, prev = fields.get(key), before.get(key) + if not isinstance(cur, int | float) or isinstance(cur, bool): + continue + prev_num = ( + prev + if isinstance(prev, int | float) and not isinstance(prev, bool) + else None + ) + value = _counter_delta(cur, prev_num) + delta[key] = ( + round(value or 0, 6) if isinstance(cur, float) else int(value or 0) + ) + out[model] = delta + return out + + +async def backfill_usage_deltas(session: AsyncSession) -> int: + """Rewrite ``cost_usd`` / ``model_usage`` of every row as per-turn deltas. + + Rows written before the ``*_total`` columns existed hold the SDK's + cumulative values in ``cost_usd`` / ``model_usage``; they are moved to the + totals first. Rows are replayed per session in time order. Idempotent: + a second run rewrites nothing. Returns the number of rows changed. + """ + stmt = select(Usage).order_by(col(Usage.session_id), col(Usage.ts), col(Usage.id)) + rows = (await session.exec(stmt)).all() + last_by_session: dict[str, Usage] = {} + changed = 0 + for row in rows: + before = ( + row.cost_usd, + row.model_usage, + row.cost_usd_total, + row.model_usage_total, + ) + prev = last_by_session.get(row.session_id) if row.session_id else None + usage_deltas(row, prev) + after = ( + row.cost_usd, + row.model_usage, + row.cost_usd_total, + row.model_usage_total, + ) + if after != before: + session.add(row) + changed += 1 + if row.session_id: + last_by_session[row.session_id] = row + await session.commit() + return changed + + __all__ = [ "Database", "append_audit", "append_usage", + "backfill_usage_deltas", "create_token", "list_active_tokens", "list_audit_records", "list_tokens", "revoke_token", "touch_token", + "usage_deltas", ] diff --git a/src/beaver_gateway/storage/models.py b/src/beaver_gateway/storage/models.py index 0d8eb36..c9ce2bb 100644 --- a/src/beaver_gateway/storage/models.py +++ b/src/beaver_gateway/storage/models.py @@ -287,7 +287,17 @@ class TranscriptEntry(SQLModel, table=True): class Usage(SQLModel, table=True): - """Per-turn token accounting from ``ResultMessage.usage``.""" + """Per-turn token accounting from ``ResultMessage``. + + Token columns come from ``ResultMessage.usage`` and are sums over the + turn's API calls. ``cost_usd`` and ``model_usage`` are **per-turn deltas**: + the SDK reports ``total_cost_usd`` and ``model_usage`` cumulatively for + the claude process (growing across the turns of a session, reset when + the process is respawned), so :func:`append_usage` stores the raw values + in ``cost_usd_total`` / ``model_usage_total`` and the difference against + the previous row of the same session here. Summing ``cost_usd`` over rows + is therefore meaningful; summing ``cost_usd_total`` is not. + """ __tablename__ = "usage" @@ -311,6 +321,14 @@ class Usage(SQLModel, table=True): default=None, sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True), ) + """Per-model tokens, cost and web searches - per-turn delta (see class doc).""" + cost_usd_total: float | None = Field(default=None) + """``ResultMessage.total_cost_usd`` verbatim: cumulative for the claude process.""" + model_usage_total: dict[str, Any] | None = Field( + default=None, + sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True), + ) + """``ResultMessage.model_usage`` verbatim: cumulative for the claude process.""" class RateLimit(SQLModel, table=True): diff --git a/tests/test_usage_deltas.py b/tests/test_usage_deltas.py new file mode 100644 index 0000000..f47dcc7 --- /dev/null +++ b/tests/test_usage_deltas.py @@ -0,0 +1,105 @@ +"""``cost_usd`` / ``model_usage`` are per-turn deltas of the SDK's running totals.""" + +import tempfile +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any + +from sqlmodel import col, select + +from beaver_gateway.storage import Database, Usage, append_usage, backfill_usage_deltas + + +def _mu(cost: float, out: int, cache_read: int) -> dict[str, Any]: + return { + "claude-opus-5": { + "costUSD": cost, + "inputTokens": 1, + "outputTokens": out, + "cacheReadInputTokens": cache_read, + "cacheCreationInputTokens": 0, + "webSearchRequests": 0, + "contextWindow": 1_000_000, + "provider": "firstParty", + } + } + + +def _row( + ts: datetime, sid: str | None, cost: float, out: int, cache_read: int +) -> Usage: + return Usage( + ts=ts, + agent_name="d", + session_id=sid, + model="claude-opus-5", + output_tokens=out, + cache_read_tokens=cache_read, + cost_usd=cost, + model_usage=_mu(cost, out, cache_read), + ) + + +async def _db() -> Database: + path = Path(tempfile.mkdtemp(prefix="beaver-usage-")) / "u.db" + db = Database(f"sqlite:///{path}") + await db.create_all() + return db + + +async def _all(db: Database) -> list[Usage]: + async with db.session() as session: + return list((await session.exec(select(Usage).order_by(col(Usage.id)))).all()) + + +async def test_append_usage_stores_deltas_and_keeps_totals() -> None: + db = await _db() + t0 = datetime.now(UTC).replace(tzinfo=None) + async with db.session() as session: + await append_usage(session, _row(t0, "s1", 0.5, 100, 1000)) + await append_usage( + session, _row(t0 + timedelta(minutes=1), "s1", 0.8, 160, 2500) + ) + # the claude process was respawned: the counters start over + await append_usage(session, _row(t0 + timedelta(minutes=2), "s1", 0.2, 30, 400)) + # another session is a separate counter + await append_usage(session, _row(t0 + timedelta(minutes=3), "s2", 0.7, 50, 900)) + # no session id - nothing to diff against + await append_usage(session, _row(t0 + timedelta(minutes=4), None, 0.3, 10, 100)) + rows = await _all(db) + assert [r.cost_usd for r in rows] == [0.5, 0.3, 0.2, 0.7, 0.3] + assert [r.cost_usd_total for r in rows] == [0.5, 0.8, 0.2, 0.7, 0.3] + second = rows[1].model_usage + assert second is not None + opus = second["claude-opus-5"] + assert (opus["costUSD"], opus["outputTokens"], opus["cacheReadInputTokens"]) == ( + 0.3, + 60, + 1500, + ) + assert (opus["contextWindow"], opus["provider"]) == (1_000_000, "firstParty") + assert rows[1].model_usage_total == _mu(0.8, 160, 2500) + assert rows[2].model_usage is not None + assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 30 + await db.dispose() + + +async def test_backfill_rewrites_legacy_cumulative_rows_once() -> None: + db = await _db() + t0 = datetime.now(UTC).replace(tzinfo=None) + async with db.session() as session: + # legacy rows: the raw running total sits in cost_usd / model_usage + for i, cost in enumerate((1.0, 1.5, 2.5, 0.4)): + session.add(_row(t0 + timedelta(minutes=i), "s1", cost, 10 * (i + 1), 100)) + session.add(_row(t0 + timedelta(minutes=9), "s9", 0.9, 5, 50)) + await session.commit() + async with db.session() as session: + assert await backfill_usage_deltas(session) == 5 + rows = await _all(db) + assert [r.cost_usd for r in rows] == [1.0, 0.5, 1.0, 0.4, 0.9] + assert [r.cost_usd_total for r in rows] == [1.0, 1.5, 2.5, 0.4, 0.9] + assert rows[2].model_usage is not None + assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 10 + async with db.session() as session: + assert await backfill_usage_deltas(session) == 0 + await db.dispose()