fix(insights): harden per-route usage attribution

Preserve deletability, route identity, stored costs, aggregate reconciliation,
and zero-usage Codex route accounting on top of the salvaged per-model usage
work.
This commit is contained in:
teknium1 2026-07-11 04:30:08 -07:00 committed by Teknium
parent d14006ead2
commit 0d63c23f36
5 changed files with 245 additions and 49 deletions

View file

@ -120,6 +120,9 @@ def _record_codex_app_server_usage(agent, turn) -> dict[str, Any]:
agent._session_db.update_token_counts(
agent.session_id,
model=agent.model,
billing_provider=agent.provider,
billing_base_url=agent.base_url,
billing_mode="subscription_included",
api_call_count=1,
)
except Exception as exc:

View file

@ -142,8 +142,8 @@ class InsightsEngine:
}
# Compute insights
overview = self._compute_overview(sessions, message_stats)
models = self._compute_model_breakdown(sessions, cutoff, source)
overview = self._compute_overview(sessions, message_stats, models)
platforms = self._compute_platform_breakdown(sessions)
tools = self._compute_tool_breakdown(tool_usage)
skills = self._compute_skill_breakdown(skill_usage)
@ -173,7 +173,7 @@ class InsightsEngine:
"message_count, tool_call_count, input_tokens, output_tokens, "
"cache_read_tokens, cache_write_tokens, billing_provider, "
"billing_base_url, billing_mode, estimated_cost_usd, "
"actual_cost_usd, cost_status, cost_source")
"actual_cost_usd, cost_status, cost_source, api_call_count")
# Pre-computed query strings — f-string evaluated once at class definition,
# not at runtime, so no user-controlled value can alter the query structure.
@ -400,7 +400,12 @@ class InsightsEngine:
# Computation
# =========================================================================
def _compute_overview(self, sessions: List[Dict], message_stats: Dict) -> Dict:
def _compute_overview(
self,
sessions: List[Dict],
message_stats: Dict,
models: Optional[List[Dict]] = None,
) -> Dict:
"""Compute high-level overview statistics."""
total_input = sum(s.get("input_tokens") or 0 for s in sessions)
total_output = sum(s.get("output_tokens") or 0 for s in sessions)
@ -432,6 +437,9 @@ class InsightsEngine:
else:
models_without_pricing.add(display)
if models:
total_cost = sum(float(m.get("cost") or 0.0) for m in models)
# Session duration stats (guard against negative durations from clock drift)
durations = []
for s in sessions:
@ -478,7 +486,8 @@ class InsightsEngine:
"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"
" u.estimated_cost_usd, u.actual_cost_usd, u.cost_status,"
" u.cost_source, u.billing_mode"
" FROM session_model_usage u"
" JOIN sessions s ON s.id = u.session_id"
" WHERE s.started_at >= ? AND s.source = ?"
@ -487,7 +496,8 @@ class InsightsEngine:
"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"
" u.estimated_cost_usd, u.actual_cost_usd, u.cost_status,"
" u.cost_source, u.billing_mode"
" FROM session_model_usage u"
" JOIN sessions s ON s.id = u.session_id"
" WHERE s.started_at >= ?"
@ -530,15 +540,16 @@ class InsightsEngine:
"sessions": set(), "input_tokens": 0, "output_tokens": 0,
"cache_read_tokens": 0, "cache_write_tokens": 0,
"reasoning_tokens": 0, "total_tokens": 0, "api_calls": 0,
"tool_calls": 0, "cost": 0.0,
"tool_calls": 0, "cost": 0.0, "actual_cost": 0.0,
})
def _accumulate(model, provider, base_url, session_id, inp, out,
cache_read, cache_write, reasoning):
cache_read, cache_write, reasoning, *,
stored_cost=None, actual_cost=None, cost_status=None):
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: Dict[str, Any] = model_data[display_model]
d["sessions"].add(session_id)
d["input_tokens"] += inp
d["output_tokens"] += out
@ -546,12 +557,17 @@ class InsightsEngine:
d["cache_write_tokens"] += cache_write
d["reasoning_tokens"] += reasoning
d["total_tokens"] += inp + out + cache_read + cache_write
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,
)
if stored_cost is None:
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,
)
else:
estimate = float(stored_cost or 0.0)
status = cost_status or "unknown"
d["cost"] += estimate
d["actual_cost"] += float(actual_cost or 0.0)
d["cost_status"] = status
if has_known_pricing(model, provider or None, base_url):
d["has_pricing"] = True
@ -560,34 +576,75 @@ class InsightsEngine:
return display_model
usage_rows = self._get_model_usage(cutoff, source)
covered: set = set()
usage_totals = defaultdict(lambda: {
"input_tokens": 0, "output_tokens": 0, "cache_read_tokens": 0,
"cache_write_tokens": 0, "reasoning_tokens": 0,
"api_call_count": 0, "estimated_cost_usd": 0.0,
"actual_cost_usd": 0.0,
})
for r in usage_rows:
covered.add(r["session_id"])
totals: Dict[str, Any] = usage_totals[r["session_id"]]
for key in (
"input_tokens", "output_tokens", "cache_read_tokens",
"cache_write_tokens", "reasoning_tokens", "api_call_count",
):
totals[key] += r[key] or 0
totals["estimated_cost_usd"] += r["estimated_cost_usd"] or 0.0
totals["actual_cost_usd"] += r["actual_cost_usd"] or 0.0
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,
stored_cost=(
r["estimated_cost_usd"]
if r.get("cost_status") or r.get("cost_source")
else None
),
actual_cost=r["actual_cost_usd"],
cost_status=r.get("cost_status"),
)
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.
# Reconcile against the aggregate row. This covers legacy sessions,
# interrupted migrations, and absolute cumulative updates without
# double-counting already-attributed route deltas.
for s in sessions:
if s["id"] in covered:
totals = usage_totals[s["id"]]
inp = max(0, (s.get("input_tokens") or 0) - totals["input_tokens"])
out = max(0, (s.get("output_tokens") or 0) - totals["output_tokens"])
cache_read = max(
0, (s.get("cache_read_tokens") or 0) - totals["cache_read_tokens"]
)
cache_write = max(
0, (s.get("cache_write_tokens") or 0) - totals["cache_write_tokens"]
)
residual_cost = max(
0.0, float(s.get("estimated_cost_usd") or 0.0)
- totals["estimated_cost_usd"],
)
residual_actual = max(
0.0, float(s.get("actual_cost_usd") or 0.0)
- totals["actual_cost_usd"],
)
residual_calls = max(
0, (s.get("api_call_count") or 0) - totals["api_call_count"]
)
if not (
inp or out or cache_read or cache_write or residual_cost
or residual_actual or residual_calls
):
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(
d = _accumulate(
s.get("model"), s.get("billing_provider"),
s.get("billing_base_url"), s["id"],
inp, out, cache_read, cache_write, 0,
stored_cost=residual_cost,
actual_cost=residual_actual,
cost_status=s.get("cost_status"),
)
residual_bucket: Dict[str, Any] = model_data[d]
residual_bucket["api_calls"] += residual_calls
# Tool calls are attributed by the session's recorded model.
for s in sessions:

View file

@ -769,10 +769,11 @@ CREATE TABLE IF NOT EXISTS messages (
);
CREATE TABLE IF NOT EXISTS session_model_usage (
session_id TEXT NOT NULL REFERENCES sessions(id),
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
model TEXT NOT NULL,
billing_provider TEXT NOT NULL DEFAULT '',
billing_base_url TEXT,
billing_base_url TEXT NOT NULL DEFAULT '',
billing_mode TEXT NOT NULL DEFAULT '',
api_call_count INTEGER NOT NULL DEFAULT 0,
input_tokens INTEGER NOT NULL DEFAULT 0,
output_tokens INTEGER NOT NULL DEFAULT 0,
@ -780,9 +781,12 @@ CREATE TABLE IF NOT EXISTS session_model_usage (
cache_write_tokens INTEGER NOT NULL DEFAULT 0,
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
estimated_cost_usd REAL NOT NULL DEFAULT 0,
actual_cost_usd REAL NOT NULL DEFAULT 0,
cost_status TEXT,
cost_source TEXT,
first_seen REAL,
last_seen REAL,
PRIMARY KEY (session_id, model, billing_provider)
PRIMARY KEY (session_id, model, billing_provider, billing_base_url, billing_mode)
);
CREATE TABLE IF NOT EXISTS state_meta (
@ -1587,14 +1591,17 @@ class SessionDB:
cursor.execute(
"""INSERT OR IGNORE INTO session_model_usage (
session_id, model, billing_provider,
billing_base_url, api_call_count, input_tokens,
billing_base_url, billing_mode,
api_call_count, input_tokens,
output_tokens, cache_read_tokens,
cache_write_tokens, reasoning_tokens,
estimated_cost_usd, first_seen, last_seen
estimated_cost_usd, actual_cost_usd,
cost_status, cost_source, first_seen, last_seen
)
SELECT id, COALESCE(model, 'unknown'),
COALESCE(billing_provider, ''),
billing_base_url,
COALESCE(billing_base_url, ''),
COALESCE(billing_mode, ''),
COALESCE(api_call_count, 0),
COALESCE(input_tokens, 0),
COALESCE(output_tokens, 0),
@ -1602,6 +1609,8 @@ class SessionDB:
COALESCE(cache_write_tokens, 0),
COALESCE(reasoning_tokens, 0),
COALESCE(estimated_cost_usd, 0),
COALESCE(actual_cost_usd, 0),
cost_status, cost_source,
started_at, COALESCE(ended_at, started_at)
FROM sessions
WHERE COALESCE(input_tokens, 0)
@ -2539,6 +2548,11 @@ class SessionDB:
model = COALESCE(model, ?),
api_call_count = COALESCE(api_call_count, 0) + ?
WHERE id = ?"""
has_accounted_usage = bool(
input_tokens or output_tokens or cache_read_tokens
or cache_write_tokens or reasoning_tokens or api_call_count
or estimated_cost_usd or actual_cost_usd
)
params = (
input_tokens,
output_tokens,
@ -2551,10 +2565,10 @@ class SessionDB:
cost_status,
cost_source,
pricing_version,
billing_provider,
billing_base_url,
billing_mode,
model,
billing_provider if has_accounted_usage else None,
billing_base_url if has_accounted_usage else None,
billing_mode if has_accounted_usage else None,
model if has_accounted_usage else None,
api_call_count,
session_id,
)
@ -2568,10 +2582,9 @@ class SessionDB:
# session_model_usage keyed by the live model preserves an accurate
# per-model breakdown regardless of how many times the user switches.
#
# Only the incremental path records here: the gateway also issues an
# ``absolute=True`` call that overwrites the sessions summary totals
# with the cached agent's cumulative figures — folding those in would
# double-count, and cumulative totals can't be split back per model.
# Only the incremental path records here. Absolute cumulative updates
# cannot be split back into routes; Insights reconciles any positive
# residual against the aggregate session row instead.
record_model_usage = (not absolute) and (
input_tokens or output_tokens or cache_read_tokens
or cache_write_tokens or reasoning_tokens or api_call_count
@ -2593,6 +2606,7 @@ class SessionDB:
# legacy row: one row cannot represent mixed-provider usage.
first_accounted_route = (
existing_api_calls == 0
and has_accounted_usage
and bool(model)
and bool(billing_provider)
and (existing_model != model or existing_provider != billing_provider)
@ -2601,8 +2615,7 @@ class SessionDB:
conn.execute(
"""UPDATE sessions
SET model = ?, billing_provider = ?,
billing_base_url = COALESCE(?, billing_base_url),
billing_mode = COALESCE(?, billing_mode)
billing_base_url = ?, billing_mode = ?
WHERE id = ?""",
(model, billing_provider, billing_base_url, billing_mode, session_id),
)
@ -2614,12 +2627,16 @@ class SessionDB:
model=model,
billing_provider=billing_provider,
billing_base_url=billing_base_url,
billing_mode=billing_mode,
input_tokens=input_tokens,
output_tokens=output_tokens,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
reasoning_tokens=reasoning_tokens,
estimated_cost_usd=estimated_cost_usd,
actual_cost_usd=actual_cost_usd,
cost_status=cost_status,
cost_source=cost_source,
api_call_count=api_call_count,
)
self._execute_write(_do)
@ -2632,12 +2649,16 @@ class SessionDB:
model: Optional[str],
billing_provider: Optional[str],
billing_base_url: Optional[str],
billing_mode: Optional[str],
input_tokens: int,
output_tokens: int,
cache_read_tokens: int,
cache_write_tokens: int,
reasoning_tokens: int,
estimated_cost_usd: Optional[float],
actual_cost_usd: Optional[float],
cost_status: Optional[str],
cost_source: Optional[str],
api_call_count: int,
) -> None:
"""Accumulate a per-API-call usage delta into session_model_usage.
@ -2649,26 +2670,30 @@ class SessionDB:
the same COALESCE-from-session behaviour the summary update uses.
"""
row = conn.execute(
"SELECT model, billing_provider, billing_base_url "
"SELECT model, billing_provider, billing_base_url, billing_mode "
"FROM sessions WHERE id = ?",
(session_id,),
).fetchone()
sess_model = row["model"] if row is not None else None
sess_provider = row["billing_provider"] if row is not None else None
sess_base_url = row["billing_base_url"] if row is not None else None
sess_billing_mode = row["billing_mode"] if row is not None else None
eff_model = model or sess_model or "unknown"
eff_provider = billing_provider or sess_provider or ""
eff_base_url = billing_base_url or sess_base_url
eff_base_url = billing_base_url or sess_base_url or ""
eff_billing_mode = billing_mode or sess_billing_mode or ""
now = time.time()
conn.execute(
"""INSERT INTO session_model_usage (
session_id, model, billing_provider, billing_base_url,
session_id, model, billing_provider, billing_base_url, billing_mode,
api_call_count, input_tokens, output_tokens,
cache_read_tokens, cache_write_tokens, reasoning_tokens,
estimated_cost_usd, first_seen, last_seen
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(session_id, model, billing_provider) DO UPDATE SET
estimated_cost_usd, actual_cost_usd, cost_status, cost_source,
first_seen, last_seen
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(session_id, model, billing_provider, billing_base_url, billing_mode)
DO UPDATE SET
api_call_count = api_call_count + excluded.api_call_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
@ -2676,13 +2701,16 @@ class SessionDB:
cache_write_tokens = cache_write_tokens + excluded.cache_write_tokens,
reasoning_tokens = reasoning_tokens + excluded.reasoning_tokens,
estimated_cost_usd = estimated_cost_usd + excluded.estimated_cost_usd,
billing_base_url = COALESCE(excluded.billing_base_url, billing_base_url),
actual_cost_usd = actual_cost_usd + excluded.actual_cost_usd,
cost_status = COALESCE(excluded.cost_status, cost_status),
cost_source = COALESCE(excluded.cost_source, cost_source),
last_seen = excluded.last_seen""",
(
session_id,
eff_model,
eff_provider,
eff_base_url,
eff_billing_mode,
api_call_count or 0,
input_tokens or 0,
output_tokens or 0,
@ -2690,6 +2718,9 @@ class SessionDB:
cache_write_tokens or 0,
reasoning_tokens or 0,
float(estimated_cost_usd or 0.0),
float(actual_cost_usd or 0.0),
cost_status,
cost_source,
now,
now,
),

View file

@ -349,6 +349,47 @@ class TestInsightsPopulated:
assert models["deepseek-v4-pro"]["total_tokens"] == 48000
assert models["claude-opus-4.8"]["total_tokens"] == 54000
def test_partial_per_model_rows_preserve_session_totals(self, db):
"""A partial rolling-upgrade row must not hide aggregate residuals."""
db.create_session(session_id="partial", source="cli", model="gpt-4o")
db.update_token_counts(
"partial", input_tokens=100, output_tokens=20,
model="gpt-4o", billing_provider="openai", api_call_count=1,
)
db.update_token_counts(
"partial", input_tokens=1000, output_tokens=200,
model="gpt-4o", billing_provider="openai", api_call_count=10,
absolute=True,
)
report = InsightsEngine(db).generate(days=30)
model = next(m for m in report["models"] if m["model"] == "gpt-4o")
assert model["input_tokens"] == 1000
assert model["output_tokens"] == 200
assert model["api_calls"] == 10
assert sum(m["total_tokens"] for m in report["models"]) == \
report["overview"]["total_tokens"]
def test_overview_cost_matches_per_model_stored_cost(self, db):
db.create_session(session_id="cost", source="cli", model="model-a")
db.update_token_counts(
"cost", input_tokens=10, model="model-a", billing_provider="custom",
estimated_cost_usd=1.25, actual_cost_usd=1.0,
cost_status="estimated", cost_source="provider", api_call_count=1,
)
db.update_session_model("cost", "model-b")
db.update_session_billing_route("cost", provider="custom-b", base_url=None)
db.update_token_counts(
"cost", input_tokens=20, model="model-b", billing_provider="custom-b",
estimated_cost_usd=2.5, actual_cost_usd=2.0,
cost_status="estimated", cost_source="provider", api_call_count=1,
)
report = InsightsEngine(db).generate(days=30)
assert sum(m["cost"] for m in report["models"]) == pytest.approx(3.75)
assert report["overview"]["estimated_cost"] == pytest.approx(3.75)
assert report["overview"]["actual_cost"] == pytest.approx(3.0)
def test_platform_breakdown(self, populated_db):
engine = InsightsEngine(populated_db)
report = engine.generate(days=30)

View file

@ -480,6 +480,58 @@ class TestSessionLifecycle:
).fetchone()
assert rows["n"] == 0
def test_per_model_usage_keeps_distinct_billing_routes(self, db):
"""The same model through distinct billing routes must not collapse."""
db.create_session(session_id="routes", source="cli", model="shared-model")
db.update_token_counts(
"routes", input_tokens=10, model="shared-model",
billing_provider="custom", billing_base_url="https://one.example/v1",
billing_mode="api_key", estimated_cost_usd=0.01, api_call_count=1,
)
db.update_token_counts(
"routes", input_tokens=20, model="shared-model",
billing_provider="custom", billing_base_url="https://two.example/v1",
billing_mode="subscription_included", estimated_cost_usd=0.0,
cost_status="included", api_call_count=1,
)
rows = db._conn.execute(
"SELECT billing_base_url, billing_mode, input_tokens "
"FROM session_model_usage WHERE session_id = 'routes' "
"ORDER BY billing_base_url"
).fetchall()
assert [(r["billing_base_url"], r["billing_mode"], r["input_tokens"])
for r in rows] == [
("https://one.example/v1", "api_key", 10),
("https://two.example/v1", "subscription_included", 20),
]
def test_metadata_only_update_does_not_replace_requested_route(self, db):
db.create_session(session_id="metadata", source="cli", model="primary")
db.update_token_counts(
"metadata", model="fallback", billing_provider="fallback-provider",
api_call_count=0,
)
row = db.get_session("metadata")
assert row["model"] == "primary"
assert row["billing_provider"] is None
def test_first_accounted_route_replaces_all_route_fields_atomically(self, db):
db.create_session(session_id="route", source="cli", model="primary")
db.update_session_billing_route(
"route", provider="primary-provider",
base_url="https://primary.example/v1", billing_mode="api_key",
)
db.update_token_counts(
"route", model="fallback", billing_provider="fallback-provider",
billing_base_url=None, billing_mode=None, api_call_count=1,
)
row = db.get_session("route")
assert row["model"] == "fallback"
assert row["billing_provider"] == "fallback-provider"
assert row["billing_base_url"] is None
assert row["billing_mode"] is None
def test_v17_backfill_seeds_existing_session_usage(self, tmp_path):
"""A DB upgraded from <17 seeds one usage row per historical session
from its aggregate totals, so insights read uniformly from the table.
@ -2054,6 +2106,18 @@ class TestDeleteAndExport:
assert db.get_session("s1") is None
assert db.message_count(session_id="s1") == 0
def test_delete_session_cascades_per_model_usage(self, db):
db.create_session(session_id="usage", source="cli", model="gpt-5")
db.update_token_counts(
"usage", input_tokens=10, model="gpt-5",
billing_provider="openai", api_call_count=1,
)
assert db.delete_session("usage") is True
count = db._conn.execute(
"SELECT COUNT(*) FROM session_model_usage WHERE session_id = 'usage'"
).fetchone()[0]
assert count == 0
def test_delete_nonexistent(self, db):
assert db.delete_session("nope") is False