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 a8491330fa
commit 90063a7cfe
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
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.
+4
View File
@@ -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",
]
@@ -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.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",
]
+19 -1
View File
@@ -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):