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

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",
 ]