fix: include cache-write tokens in quota audit (#577)

[DhruvGarg111]

Fixes the quota-audit alternative-cost calculation by carrying cache-write tokens through the session aggregate and including them in alternative pricing. Adds regression coverage for cache-read and cache-write pricing.

Closes #577

Summary

• Carry cache_write_tokens through the per-session aggregate and price them on the alternative-model side, closing the last _alt_unit_cost call site that omitted them.
• Split the counterfactual into a per-original-model aggregate (_ModelCostAgg), so a mixed-model session no longer prices every model's tokens at the dominant model's downgrade.
• Add regression coverage for both.

Root cause

The actual side of the counterfactual reads the stored cost_usd, which calculate_cost computes across all four token classes. The alternative side hand-rolled its arithmetic over three, silently dropping cache-write. The two sides were therefore measuring different things, and the difference between them — the headline recoverable figure — was inflated by the full cache-write cost of the candidate.

Before / after headline figure

One example window: 10 cache-heavy Opus sessions, each 20k input / 2k output / 150k cache-read / 40k cache-write, claude-opus-4-8 → claude-sonnet-5, priced at 2026-08-15.

	alternative cost	recoverable headline
before (cache-write dropped from the alt side)	$0.90	$3.85
after (cache-write priced on the alt side)	$1.90	$2.85

Actual window cost is $4.75 in both. The fix removes $1.00 of overstatement, 26% of the old headline, on a window shaped to make cache-write matter. The correction always moves the number down: the alt side was understated, so the implied saving was overstated. A less cache-heavy window moves less.

What's also in this PR

The per-model split goes beyond what #577 asked for (an accumulator on _SessionAgg). It is a genuine correctness fix in the same function, surfaced by Greptile's P1 on the first commit, and it describes a defect that already existed on main rather than one this PR introduced: agg.new_input/output/reread were summed across every premium model in a session and priced at dominant_model()'s alternative.

It carries one behaviour change worth naming: a session whose dominant model had no downgrade used to be dropped whole from both actual_cost and alternative_cost. Each model is now included or excluded independently. Both sides move together, so the pair stays coherent for the downstream consumer in cli/cmd_quota_audit.py.

Tests / Verification

• test_api_counterfactual_prices_mixed_models_per_model — fails against origin/main, pinning the per-model split.
• Cache-read / cache-write pricing coverage — fails with the product file reverted.
• Both new tests verified load-bearing by reverting model_downgrade.py to origin/main.
• Full CI matrix green (3.10 / 3.11 / 3.12), lint, test-ts, version-lockstep.
• The earlier red runs on this branch were main's expired claude-sonnet-5 introductory rate, unrelated to this PR and fixed in fix(tests): pin the pricing instant so a dated rate can't break main #741.

What's NOT in this PR

• The quota-weighting question. TurnComposition.quota_weighted_tokens weights cache-write at 1.0, and this PR does not touch it — only the implied-dollar counterfactual. Anthropic's docs state cache hits are not deducted against rate limits but say nothing explicit about cache creation, so the 1.0 weight is plausible and unverified. Out of scope here.
• The 5m vs 1h cache-write TTL distinction, which models.toml does not model at all. Pre-existing.
• _SessionAgg's new_input/output/reread/cost are now duplicated by pricing_by_model and survive only to feed _example_for. Left as is; two aggregations of the same turns can drift, so worth revisiting if that function is touched again.

---

Body expanded by the maintainer at merge time to record the before/after figure required by #577's acceptance criteria and to describe the second commit. The code is entirely @DhruvGarg111's.

[greptile-apps]

Greptile Summary

The PR carries cache-write tokens into quota-audit counterfactual pricing and splits mixed-model sessions by original model so each group uses its own downgrade.

• Adds per-model token and observed-cost aggregates.
• Prices all four token classes on the alternative side.
• Adds regression tests for cache-write and mixed-model pricing.

Confidence Score: 4/5

The PR is not yet safe to merge because mixed-model audits can still expose incomplete actual and alternative dollar totals as a complete comparison.

The per-model split fixes mixed-model downgrade selection, but groups lacking a provider, downgrade, or alternative price are still skipped together with their known observed cost, while downstream rendering provides no partial-coverage indication.

Files Needing Attention: tokenjam/core/optimize/analyzers/model_downgrade.py

Important Files Changed

Filename	Overview
tokenjam/core/optimize/analyzers/model_downgrade.py	Adds four-bucket per-model counterfactual aggregation, but unpriceable groups still cause incomplete dollar totals.
tests/unit/test_opus_quota_audit.py	Adds focused regression coverage for cache-write alternative pricing and fully priceable mixed-model sessions.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Flagged premium turns] --> B[Group by provider and original model]
  B --> C[Accumulate input, output, cache-read, cache-write, and observed cost]
  C --> D[Resolve each model's downgrade]
  D --> E[Price alternative token mix]
  E --> F[Aggregate actual and alternative audit totals]

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into fix/577-cache-w..." | Re-trigger Greptile}

[DhruvGarg111]

Addressed the mixed-model pricing finding. Quota-audit counterfactuals now aggregate token classes by original model and price each group at its own downgrade; cache-write tokens are passed explicitly. Added mixed-model regression coverage.

Validation: 25 quota-audit tests and 53 related optimize tests passed; Ruff and mypy passed.

[greptile-apps]

Partial costs appear complete

When a flagged mixed-model session contains one priceable group and another premium-model group with a missing provider or no matching downgrade, this branch omits the unresolved group's observed and alternative costs while its tokens remain in the audit population. The CLI and API consequently present incomplete dollar totals as the actual billed amount and counterfactual cost without indicating partial pricing coverage.

Knowledge Base Used:

• Usage data lifecycle
• Optimization analysis and recommendations

Prompt To Fix With AI

This is a comment left during a code review.
Path: tokenjam/core/optimize/analyzers/model_downgrade.py
Line: 1170-1172

Comment:
**Partial costs appear complete**

When a flagged mixed-model session contains one priceable group and another premium-model group with a missing provider or no matching downgrade, this branch omits the unresolved group's observed and alternative costs while its tokens remain in the audit population. The CLI and API consequently present incomplete dollar totals as the actual billed amount and counterfactual cost without indicating partial pricing coverage.

**Knowledge Base Used:**
- [Usage data lifecycle](https://app.greptile.com/metabldr/-/custom-context/knowledge-base/metabuilder-labs/tokenjam/-/docs/usage-data-lifecycle.md)
- [Optimization analysis and recommendations](https://app.greptile.com/metabldr/-/custom-context/knowledge-base/metabuilder-labs/tokenjam/-/docs/optimization-analysis.md)

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[anilmurty]

@DhruvGarg111 thanks for this — the arithmetic is right and I want to land it. Two small asks on the PR body first, plus context on the CI you've been fighting.

The red CI was never yours. claude-sonnet-5's introductory pricing expired on 2026-09-01 and five tests were reading the live rate, so main itself went red that day and every open PR inherited it. Fixed in #741. I've updated your branch onto it and re-approved the workflow, so the run in flight should be clean.

I traced the cost path before approving, since "add a token type to a cost figure" is the kind of change that goes wrong quietly:

• Rate weighting is correct. _alt_unit_cost → span_pricing.price_span → calculate_cost, which applies cache_write_per_mtok as its own term rather than folding cache-write in with input.
• No double-count. TurnComposition's four buckets are disjoint, and input_tokens is exclusive of cache creation, so the fourth term adds no overlap.
• The direction is the honest one. The actual side was already four-type via the stored cost_usd while the alt side priced three, so the alt was understated and the implied saving overstated. This moves the headline number down. That's the right way for a correction like this to run.
• at=pricing_agg.first_turn_at keeps the original traffic's instant, which is what the span-pricing contract asks for.

Greptile's P1 about mixed-model cache pricing was real, and worth saying clearly: it described a defect that already existed on main, not one you introduced — agg.new_input/output/reread were summed across all premium models and priced at the dominant model's alternative. Your pricing_by_model split fixes it properly, and test_api_counterfactual_prices_mixed_models_per_model fails against main, so it's pinned. Both new tests are load-bearing; I verified by reverting the product file.

Two asks, both on the body:

1. Issue downsize headline savings overstated because the alternative-model price omits cache write tokens #577's third "Done when" isn't met yet — it asks the PR to state the before/after headline recoverable-savings figure for at least one example window. Since this changes a number a user reads, that figure is the part a reviewer can't reconstruct. Please add it.
2. The body still describes only your first commit. The per-model split in commit 2 is a bigger change than the issue asked for — a genuine correctness fix in the same function, but it goes beyond "add a cache_write accumulator", and the next person reading this PR won't know it's there. Please describe it, and note the behaviour change it carries: a session whose dominant model had no downgrade used to be dropped whole from both sides, and now each model is included or excluded independently.

One observation, no action needed: _SessionAgg's new_input/output/reread/cost are now duplicated by pricing_by_model and survive only to feed _example_for. Two aggregations of the same turns can drift apart later. Fine as is; worth a comment if you touch it again.

Once the body has those two items and CI is green, this is ready.

[anilmurty]

Thanks @DhruvGarg111 — merging. The arithmetic is right where it matters: cache-write is priced as its own term rather than folded into input, the four token buckets stay disjoint so nothing double-counts, and at=pricing_agg.first_turn_at keeps the counterfactual on the original traffic's instant. The correction moves the headline figure down, which is the honest direction — the alt side was understated, so the implied saving was overstated.

I expanded the PR body at merge time to record the before/after figure #577 asks for (a cache-heavy example window: $3.85 → $2.85 recoverable, 26% of overstatement removed) and to describe the per-model split from your second commit, which the body did not mention. The code is entirely yours.

Triaging Greptile's two P1s rather than merging past them:

1. "Mixed-model cache pricing is wrong" — valid against your first commit, and it described a defect that already existed on main. Your pricing_by_model split fixes it, and test_api_counterfactual_prices_mixed_models_per_model fails against main, so it is pinned. Resolved.

2. "Partial costs appear complete" — a real observation, not a blocker here, and not introduced by this PR. Both totals exclude the same groups and move together, so the delta is a fair comparison and "Same work" scopes the sentence. The missing piece is a coverage indication, and that predates you: before the per-model split, a session whose dominant model had no downgrade was dropped whole from both sides, so actual_cost_usd was never the user's full spend. Your change covers strictly more real work than before. Filed as #742 rather than dropped.

The red CI you were fighting was main's expired claude-sonnet-5 introductory rate, fixed in #741 — nothing to do with your work.

[DhruvGarg111]

Thanks @anilmurty , by the way i am looking forward to contributing more to your repo.