feat(agent): track per-model token usage for mid-session model switches

The `sessions` table records only the initial (model, billing_provider)
for a session, so when a user switches models mid-session (via `/model`
or programmatically) every token — including the switched model's — is
attributed to the first model. Insights/billing reports then hide the
cost of the new model entirely (e.g. a session that started on deepseek
and switched to opus shows $0 for opus).

Add a `session_model_usage` table keyed (session_id, model,
billing_provider) that accumulates each per-API-call delta under the
model active at the time of the call. `update_token_counts()` is the
single chokepoint every per-call delta flows through (CLI, gateway,
cron, delegated, codex), so recording there captures accurate
attribution on every platform. Only the incremental path records — the
gateway's `absolute=True` summary overwrite is skipped to avoid
double-counting cumulative totals that can't be split per model. When a
call omits the model, it falls back to the session's recorded model,
matching the existing COALESCE-from-session summary behaviour.

Insights `_compute_model_breakdown` now aggregates tokens and cost from
`session_model_usage`, so a switched session splits correctly across
models, with a defensive fallback to the per-session aggregate for any
session lacking usage rows. A v17 migration backfills one usage row per
existing token-bearing session from its aggregate totals (idempotent via
INSERT OR IGNORE), validated lossless against a 1.3 GB production DB.

Tests: per-model recording, mid-session split, model fallback, absolute
no-double-count, v17 backfill, and an insights-level switch breakdown.

Fixes #51607.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Thomas Connally 2026-06-23 21:23:43 -05:00 committed by Teknium
parent 022c4991fc
commit cb7f6bbb2e
4 changed files with 438 additions and 20 deletions

View file

@ -17,6 +17,7 @@ Usage:
"""
import json
import sqlite3
import time
from collections import Counter, defaultdict
from datetime import datetime
@ -142,7 +143,7 @@ class InsightsEngine:
# Compute insights
overview = self._compute_overview(sessions, message_stats)
models = self._compute_model_breakdown(sessions)
models = self._compute_model_breakdown(sessions, cutoff, source)
platforms = self._compute_platform_breakdown(sessions)
tools = self._compute_tool_breakdown(tool_usage)
skills = self._compute_skill_breakdown(skill_usage)
@ -473,39 +474,140 @@ class InsightsEngine:
"included_cost_sessions": included_cost_sessions,
}
def _compute_model_breakdown(self, sessions: List[Dict]) -> List[Dict]:
"""Break down usage by model."""
_GET_MODEL_USAGE_WITH_SOURCE = (
"SELECT u.session_id, u.model, u.billing_provider, u.billing_base_url,"
" u.api_call_count, u.input_tokens, u.output_tokens,"
" u.cache_read_tokens, u.cache_write_tokens, u.reasoning_tokens,"
" u.estimated_cost_usd"
" FROM session_model_usage u"
" JOIN sessions s ON s.id = u.session_id"
" WHERE s.started_at >= ? AND s.source = ?"
)
_GET_MODEL_USAGE_ALL = (
"SELECT u.session_id, u.model, u.billing_provider, u.billing_base_url,"
" u.api_call_count, u.input_tokens, u.output_tokens,"
" u.cache_read_tokens, u.cache_write_tokens, u.reasoning_tokens,"
" u.estimated_cost_usd"
" FROM session_model_usage u"
" JOIN sessions s ON s.id = u.session_id"
" WHERE s.started_at >= ?"
)
def _get_model_usage(self, cutoff: float, source: str = None) -> List[Dict]:
"""Fetch per-model usage rows within the window (issue #51607).
Returns an empty list when the table is missing (e.g. a DB opened by
older code that never created it) so the caller can fall back to the
per-session aggregate.
"""
try:
if source:
cursor = self._conn.execute(
self._GET_MODEL_USAGE_WITH_SOURCE, (cutoff, source)
)
else:
cursor = self._conn.execute(self._GET_MODEL_USAGE_ALL, (cutoff,))
return [dict(row) for row in cursor.fetchall()]
except sqlite3.OperationalError:
return []
def _compute_model_breakdown(
self, sessions: List[Dict], cutoff: float, source: str = None
) -> List[Dict]:
"""Break down token usage and cost by model.
Tokens and cost are attributed per model from session_model_usage, so a
session that switched models mid-flight (via ``/model``) splits across
every model it used instead of dumping everything on the initial model
(issue #51607). Sessions without per-model rows — e.g. data written
before this table existed and not yet backfilled fall back to their
single recorded (model, billing_provider) aggregate so nothing is lost.
Tool calls aren't tied to a specific API invocation, so they stay
attributed to the session's recorded model.
"""
model_data = defaultdict(lambda: {
"sessions": 0, "input_tokens": 0, "output_tokens": 0,
"sessions": set(), "input_tokens": 0, "output_tokens": 0,
"cache_read_tokens": 0, "cache_write_tokens": 0,
"total_tokens": 0, "tool_calls": 0, "cost": 0.0,
"reasoning_tokens": 0, "total_tokens": 0, "api_calls": 0,
"tool_calls": 0, "cost": 0.0,
})
for s in sessions:
model = s.get("model") or "unknown"
def _accumulate(model, provider, base_url, session_id, inp, out,
cache_read, cache_write, reasoning):
model = model or "unknown"
# Normalize: strip provider prefix for display
display_model = model.split("/")[-1] if "/" in model else model
d = model_data[display_model]
d["sessions"] += 1
inp = s.get("input_tokens") or 0
out = s.get("output_tokens") or 0
cache_read = s.get("cache_read_tokens") or 0
cache_write = s.get("cache_write_tokens") or 0
d["sessions"].add(session_id)
d["input_tokens"] += inp
d["output_tokens"] += out
d["cache_read_tokens"] += cache_read
d["cache_write_tokens"] += cache_write
d["reasoning_tokens"] += reasoning
d["total_tokens"] += inp + out + cache_read + cache_write
d["tool_calls"] += s.get("tool_call_count") or 0
estimate, status = _estimate_cost(s)
estimate, status = _estimate_cost(
model, inp, out,
cache_read_tokens=cache_read, cache_write_tokens=cache_write,
provider=provider or None, base_url=base_url,
)
d["cost"] += estimate
d["has_pricing"] = has_known_pricing(model, s.get("billing_provider"), s.get("billing_base_url"))
d["cost_status"] = status
if has_known_pricing(model, provider or None, base_url):
d["has_pricing"] = True
else:
d.setdefault("has_pricing", False)
return display_model
result = [
{"model": model, **data}
for model, data in model_data.items()
]
usage_rows = self._get_model_usage(cutoff, source)
covered: set = set()
for r in usage_rows:
covered.add(r["session_id"])
d = _accumulate(
r["model"], r["billing_provider"], r.get("billing_base_url"),
r["session_id"], r["input_tokens"] or 0, r["output_tokens"] or 0,
r["cache_read_tokens"] or 0, r["cache_write_tokens"] or 0,
r["reasoning_tokens"] or 0,
)
model_data[d]["api_calls"] += r["api_call_count"] or 0
# Fallback for sessions with token totals but no per-model rows
# (legacy data not covered by the v17 backfill). Attribute their
# aggregate to the single recorded model so totals never regress.
for s in sessions:
if s["id"] in covered:
continue
inp = s.get("input_tokens") or 0
out = s.get("output_tokens") or 0
cache_read = s.get("cache_read_tokens") or 0
cache_write = s.get("cache_write_tokens") or 0
if not (inp or out or cache_read or cache_write):
continue
_accumulate(
s.get("model"), s.get("billing_provider"),
s.get("billing_base_url"), s["id"],
inp, out, cache_read, cache_write, 0,
)
# Tool calls are attributed by the session's recorded model.
for s in sessions:
tool_calls = s.get("tool_call_count") or 0
if not tool_calls:
continue
model = s.get("model") or "unknown"
display_model = model.split("/")[-1] if "/" in model else model
model_data[display_model]["tool_calls"] += tool_calls
result = []
for model, data in model_data.items():
entry = {"model": model, **data}
entry["sessions"] = len(data["sessions"])
# Models that surfaced only via tool-call attribution (no token
# rows) won't have these set by _accumulate — default them so the
# output shape is uniform for downstream/JSON consumers.
entry.setdefault("has_pricing", False)
entry.setdefault("cost_status", "unknown")
result.append(entry)
# Sort by tokens first, fall back to session count when tokens are 0
result.sort(key=lambda x: (x["total_tokens"], x["sessions"]), reverse=True)
return result