fix(storage,api): usage cost as per-turn delta of the sdk running total, backfill

This commit is contained in:
hh
2026-09-01 15:37:53 +02:00
parent 3406e2e178
commit 63014a487d
6 changed files with 304 additions and 5 deletions
+4 -3
View File
@@ -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 bus as SSE with the same keepalive the markdown frontend uses, so a proxy
never sees a silent socket. never sees a silent socket.
Usage figures come from the ``usage`` table (one row per turn, API-price Usage figures come from the ``usage`` table (one row per turn; API-price
``cost_usd`` and per-model ``model_usage`` from the SDK's ``cost_usd`` and per-model ``model_usage`` are per-turn deltas of the SDK's
``ResultMessage``); subscription quotas come from ``rate_limits`` cumulative ``ResultMessage`` counters, see ``storage.append_usage``);
subscription quotas come from ``rate_limits``
(``RateLimitEvent``). The quota covers the whole subscription, so (``RateLimitEvent``). The quota covers the whole subscription, so
``/api/limits`` puts the gateway's own spend for the window next to it ``/api/limits`` puts the gateway's own spend for the window next to it
for calibration by eye. for calibration by eye.
+4
View File
@@ -9,12 +9,14 @@ from beaver_gateway.storage.db import (
Database, Database,
append_audit, append_audit,
append_usage, append_usage,
backfill_usage_deltas,
create_token, create_token,
list_active_tokens, list_active_tokens,
list_audit_records, list_audit_records,
list_tokens, list_tokens,
revoke_token, revoke_token,
touch_token, touch_token,
usage_deltas,
) )
from beaver_gateway.storage.models import ( from beaver_gateway.storage.models import (
AuditLog, AuditLog,
@@ -45,10 +47,12 @@ __all__ = [
"Usage", "Usage",
"append_audit", "append_audit",
"append_usage", "append_usage",
"backfill_usage_deltas",
"create_token", "create_token",
"list_active_tokens", "list_active_tokens",
"list_audit_records", "list_audit_records",
"list_tokens", "list_tokens",
"revoke_token", "revoke_token",
"touch_token", "touch_token",
"usage_deltas",
] ]
@@ -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())
+135 -1
View File
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any
from sqlalchemy import inspect, text from sqlalchemy import inspect, text
from sqlalchemy.ext.asyncio import create_async_engine 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 sqlmodel.ext.asyncio.session import AsyncSession
from beaver_gateway.storage.models import AuditLog, Token, Usage 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: 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) session.add(row)
await session.commit() 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__ = [ __all__ = [
"Database", "Database",
"append_audit", "append_audit",
"append_usage", "append_usage",
"backfill_usage_deltas",
"create_token", "create_token",
"list_active_tokens", "list_active_tokens",
"list_audit_records", "list_audit_records",
"list_tokens", "list_tokens",
"revoke_token", "revoke_token",
"touch_token", "touch_token",
"usage_deltas",
] ]
+19 -1
View File
@@ -287,7 +287,17 @@ class TranscriptEntry(SQLModel, table=True):
class Usage(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" __tablename__ = "usage"
@@ -311,6 +321,14 @@ class Usage(SQLModel, table=True):
default=None, default=None,
sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True), 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): class RateLimit(SQLModel, table=True):
+105
View File
@@ -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()