"""Unit tests for SessionUsage / UsageTracker cache token accumulation.""" import pytest from opensquilla.engine.usage import ModelUsage, SessionUsage, UsageTracker def test_session_usage_accumulates_cache_tokens() -> None: usage = SessionUsage(model_id="claude-opus-4-7") usage.add(1000, 50, "claude-opus-4-7", cache_read_tokens=500, cache_write_tokens=100) usage.add(2000, 80, "claude-opus-4-7", cache_read_tokens=300, cache_write_tokens=40) assert usage.input_tokens == 3000 assert usage.output_tokens == 130 assert usage.cache_read_tokens == 800 assert usage.cache_write_tokens == 140 def test_session_usage_per_model_breakdown_isolates_cache_tokens() -> None: usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", cache_read_tokens=500, cache_write_tokens=100) usage.add(2000, 80, "deepseek-v4-pro", cache_read_tokens=300, cache_write_tokens=40) assert usage._per_model is not None opus = usage._per_model["claude-opus-4-7"] deepseek = usage._per_model["deepseek-v4-pro"] assert opus.cache_read_tokens == 500 assert opus.cache_write_tokens == 100 assert deepseek.cache_read_tokens == 300 assert deepseek.cache_write_tokens == 40 def test_usage_tracker_session_total_equals_sum_across_providers() -> None: # R9c guard: add() is synchronous, so a session's totals are the exact # arithmetic sum of every accumulation, including turns billed to # different providers/models (no interleaving / lost-update race). tracker = UsageTracker() sk = "agent:main:s1" tracker.add(sk, 1000, 100, model_id="anthropic/m1", provider="anthropic") tracker.add(sk, 2000, 200, model_id="deepseek/m2", provider="deepseek") tracker.add(sk, 500, 50, model_id="anthropic/m1", provider="anthropic") usage = tracker.get(sk) assert usage is not None assert usage.input_tokens == 3500 assert usage.output_tokens == 350 assert usage._per_model is not None assert usage._per_model["anthropic/m1"].input_tokens == 1500 assert usage._per_model["deepseek/m2"].input_tokens == 2000 def test_session_usage_add_default_cache_zero() -> None: """Existing positional callers (no cache kwargs) still work; cache fields stay at 0.""" usage = SessionUsage(model_id="claude-opus-4-7") usage.add(1000, 50, "claude-opus-4-7") assert usage.input_tokens == 1000 assert usage.cache_read_tokens == 0 assert usage.cache_write_tokens == 0 def test_usage_tracker_add_passes_cache_to_session() -> None: tracker = UsageTracker() tracker.add( "session-a", input_tokens=1000, output_tokens=50, model_id="claude-opus-4-7", cache_read_tokens=200, cache_write_tokens=80, ) tracker.add( "session-a", input_tokens=500, output_tokens=20, model_id="claude-opus-4-7", cache_read_tokens=100, cache_write_tokens=40, ) usage = tracker.get("session-a") assert usage is not None assert usage.cache_read_tokens == 300 assert usage.cache_write_tokens == 120 def test_usage_tracker_isolates_sessions() -> None: tracker = UsageTracker() tracker.add( "session-a", input_tokens=100, output_tokens=10, model_id="m", cache_read_tokens=50, cache_write_tokens=5, ) tracker.add( "session-b", input_tokens=200, output_tokens=20, model_id="m", cache_read_tokens=70, cache_write_tokens=15, ) a = tracker.get("session-a") b = tracker.get("session-b") assert a is not None and b is not None assert a.cache_read_tokens == 50 assert a.cache_write_tokens == 5 assert b.cache_read_tokens == 70 assert b.cache_write_tokens == 15 # --------------------------------------------------------------------------- # Positional dataclass construction safety. # --------------------------------------------------------------------------- def test_session_usage_positional_construction_does_not_shift_fields() -> None: """Regression: SessionUsage(1, 2, "claude-opus-4-7") must keep the third positional arg in model_id, not in a cache field. Asserts the new cache counters were appended at the *end* of the dataclass.""" usage = SessionUsage(1000, 50, "claude-opus-4-7") assert usage.input_tokens == 1000 assert usage.output_tokens == 50 assert usage.model_id == "claude-opus-4-7" assert usage.cache_read_tokens == 0 assert usage.cache_write_tokens == 0 def test_model_usage_positional_construction_keeps_model_id_first() -> None: """ModelUsage(model_id, in, out) — sanity check for positional callers.""" mu = ModelUsage("claude-opus-4-7", 1000, 50) assert mu.model_id == "claude-opus-4-7" assert mu.input_tokens == 1000 assert mu.output_tokens == 50 assert mu.cache_read_tokens == 0 assert mu.cache_write_tokens == 0 # --------------------------------------------------------------------------- # model_breakdown must surface cache fields. # --------------------------------------------------------------------------- def test_model_breakdown_serializes_cache_fields_for_per_model_path() -> None: """Multi-model session: each breakdown entry must carry cache R/W counters so the UI's modelBreakdown column can show per-model cache split.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", cache_read_tokens=500, cache_write_tokens=100) usage.add(2000, 80, "deepseek-v4-pro", cache_read_tokens=300, cache_write_tokens=40) breakdown = usage.model_breakdown assert len(breakdown) == 2 by_model = {row["model"]: row for row in breakdown} assert by_model["claude-opus-4-7"]["cacheReadTokens"] == 500 assert by_model["claude-opus-4-7"]["cacheWriteTokens"] == 100 assert by_model["deepseek-v4-pro"]["cacheReadTokens"] == 300 assert by_model["deepseek-v4-pro"]["cacheWriteTokens"] == 40 def test_model_breakdown_serializes_cache_fields_for_single_model_path() -> None: """Single-model session (no per-model dict yet): the synthesized one-row breakdown must also carry cache R/W counters.""" usage = SessionUsage(model_id="claude-opus-4-7") # Direct field mutation — exercises the no-_per_model branch in model_breakdown. usage.input_tokens = 1000 usage.output_tokens = 50 usage.cache_read_tokens = 200 usage.cache_write_tokens = 80 [row] = usage.model_breakdown assert row["model"] == "claude-opus-4-7" assert row["cacheReadTokens"] == 200 assert row["cacheWriteTokens"] == 80 # --------------------------------------------------------------------------- # Per-model billed_cost: real provider-billed amounts flow through ModelUsage # so the breakdown can show actual numbers instead of relying on the rpc_usage # pro-rate workaround. # --------------------------------------------------------------------------- def test_model_usage_accumulates_billed_cost() -> None: """ModelUsage.billed_cost accumulates across SessionUsage.add(...) calls for the same model. Mirrors the agent.py raw_ev loop where one model can receive multiple provider calls within a turn.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.0500) usage.add(2000, 80, "claude-opus-4-7", billed_cost=0.0125) assert usage._per_model is not None opus = usage._per_model["claude-opus-4-7"] assert opus.billed_cost == pytest.approx(0.0625) def test_session_usage_per_model_billed_isolates_by_model() -> None: """Multi-model session: each model's billed_cost stays independent.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.1254) usage.add(9000, 100, "z-ai/glm-5.1", billed_cost=0.0111) assert usage._per_model is not None assert usage._per_model["claude-opus-4-7"].billed_cost == pytest.approx(0.1254) assert usage._per_model["z-ai/glm-5.1"].billed_cost == pytest.approx(0.0111) def test_model_breakdown_uses_billed_when_present() -> None: """When provider returned real billed costs, model_breakdown must surface those (not the cache-blind pricing-table estimate). Source must flip to 'provider_billed' so the UI renders the solid 'Actual' chip without triggering the pro-rate disclosure notice.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.1254) usage.add(9000, 100, "z-ai/glm-5.1", billed_cost=0.0111) breakdown = usage.model_breakdown by_model = {row["model"]: row for row in breakdown} assert by_model["claude-opus-4-7"]["costUsd"] == pytest.approx(0.1254) assert by_model["claude-opus-4-7"]["billedCostUsd"] == pytest.approx(0.1254) assert by_model["claude-opus-4-7"]["costSource"] == "provider_billed" assert by_model["z-ai/glm-5.1"]["costUsd"] == pytest.approx(0.0111) assert by_model["z-ai/glm-5.1"]["costSource"] == "provider_billed" # Sum of per-model billed equals total session billed (the property the # WebUI relies on to render the "Actual" badge without pro-rating). total = sum(row["costUsd"] for row in breakdown) assert total == pytest.approx(0.1365) def test_model_breakdown_falls_back_to_estimate_when_no_billed() -> None: """When billed_cost is 0 (provider unavailable / estimate-only path), model_breakdown reverts to the local pricing-table estimate and tags the source as 'opensquilla_estimate'.""" usage = SessionUsage() # Use a model that has a stable price entry. usage.add(1000, 50, "claude-opus-4-7") # no billed_cost kwarg → defaults to 0.0 [row] = usage.model_breakdown assert row["billedCostUsd"] == 0.0 assert row["costSource"] == "opensquilla_estimate" # costUsd here is the pricing-table estimate (non-zero for known models). assert row["costUsd"] > 0 def test_model_breakdown_mixed_billed_and_estimate() -> None: """If only one of two models has a billed cost, each row's source reflects its own situation — one provider_billed, one opensquilla_estimate. Total is mixed but the rpc_usage layer handles that aggregation separately.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) usage.add(2000, 80, "deepseek-v4-pro") # no billed by_model = {row["model"]: row for row in usage.model_breakdown} assert by_model["claude-opus-4-7"]["costSource"] == "provider_billed" assert by_model["deepseek-v4-pro"]["costSource"] == "opensquilla_estimate" def test_usage_tracker_add_forwards_billed_cost() -> None: """UsageTracker.add must thread billed_cost through to SessionUsage — this is the surface agent.py uses, so a regression here breaks the whole provider-billed cost path.""" tracker = UsageTracker() tracker.add( "session-a", input_tokens=1000, output_tokens=50, model_id="claude-opus-4-7", billed_cost=0.05, ) tracker.add( "session-a", input_tokens=500, output_tokens=20, model_id="claude-opus-4-7", billed_cost=0.025, ) usage = tracker.get("session-a") assert usage is not None assert usage._per_model is not None assert usage._per_model["claude-opus-4-7"].billed_cost == pytest.approx(0.075) def test_session_billed_cost_aggregates_across_models() -> None: """SessionUsage.billed_cost sums per-model real billed totals. Used by rpc_usage._tracker_rows to decide the row-level cost_source.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.1254) usage.add(9000, 100, "z-ai/glm-5.1", billed_cost=0.0111) assert usage.billed_cost == pytest.approx(0.1365) def test_session_cost_source_provider_billed_when_all_models_billed() -> None: """Every model has real billed → row source = provider_billed.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) usage.add(2000, 80, "deepseek-v4-pro", billed_cost=0.01) assert usage.cost_source == "provider_billed" def test_session_cost_source_mixed_when_only_some_models_billed() -> None: """Some models billed, some estimate-only → source = mixed. Prevents the row from claiming 'Actual' when half its breakdown is estimated.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) usage.add(2000, 80, "deepseek-v4-pro") # no billed assert usage.cost_source == "mixed" def test_session_cost_source_estimate_when_no_models_billed() -> None: """No billed data anywhere → source = opensquilla_estimate (the estimate-only path; preserved for sessions where provider returned no cost on any call).""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7") usage.add(2000, 80, "deepseek-v4-pro") assert usage.cost_source == "opensquilla_estimate" assert usage.billed_cost == 0.0 def test_session_total_cost_uses_billed_when_available() -> None: """All-billed session: total_cost equals billed_cost (no double-count).""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) usage.add(2000, 80, "deepseek-v4-pro", billed_cost=0.01) assert usage.total_cost == pytest.approx(0.06) assert usage.billed_cost == pytest.approx(0.06) def test_session_total_cost_mixes_billed_and_estimate() -> None: """Mixed session: total_cost = billed (for billed models) + estimate (for unbilled). This is the key property that prevents the row from under-reporting the unbilled portion of cost. Without this, _tracker_rows would set row.cost_usd to billed_cost only ($0.05) while the breakdown's unbilled item still contributes its estimate ($0.03) — the row would show $0.05 but the breakdown sum would be $0.08, a visible inconsistency.""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) # deepseek-v4-pro: no billed; pricing-table estimate kicks in usage.add(2000, 80, "deepseek-v4-pro") # billed_cost only counts billed models billed = usage.billed_cost assert billed == pytest.approx(0.05) # cost (estimate) sums both models' pricing-table estimates estimate = usage.cost assert estimate > 0 # both models have pricing entries # total_cost: billed for first + estimate for second deepseek_estimate = usage._per_model["deepseek-v4-pro"].cost expected_total = 0.05 + deepseek_estimate assert usage.total_cost == pytest.approx(expected_total) # And total_cost equals the breakdown sum (key invariant). The breakdown # serializer rounds to 6 decimals so we allow a 1e-6 tolerance against # the un-rounded total_cost; user-facing display rounds to <=4 decimals. breakdown_sum = sum(item["costUsd"] for item in usage.model_breakdown) assert breakdown_sum == pytest.approx(usage.total_cost, abs=1e-6) def test_session_total_cost_same_model_mixed_calls_adds_estimate( monkeypatch: pytest.MonkeyPatch, ) -> None: """One model, one billed call plus one unbilled call: total_cost must be billed + estimate-of-unbilled, not the old billed-only collapse. The breakdown row flips to 'mixed' and still sums to total_cost.""" monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0") usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7", billed_cost=0.05) usage.add(2000, 80, "claude-opus-4-7") # provider returned no cost # $15/M in + $75/M out (static table): (2000 * 15 + 80 * 75) / 1e6 = 0.036 assert usage.total_cost == pytest.approx(0.05 + 0.036) [row] = usage.model_breakdown assert row["costSource"] == "mixed" assert row["billedCostUsd"] == pytest.approx(0.05) assert row["costUsd"] == pytest.approx(usage.total_cost, abs=1e-6) def test_session_total_cost_falls_back_to_estimate_when_no_billed() -> None: """No billed anywhere → total_cost == cost (pure estimate).""" usage = SessionUsage() usage.add(1000, 50, "claude-opus-4-7") usage.add(2000, 80, "deepseek-v4-pro") assert usage.total_cost == pytest.approx(usage.cost) def test_session_total_cost_empty_session_is_zero() -> None: """Defensive: empty session returns 0 (falls through to .cost which is 0).""" usage = SessionUsage() assert usage.total_cost == 0.0 def test_session_billed_cost_zero_when_no_per_model() -> None: """Defensive: empty session (no add() calls) returns 0 billed and estimate cost source. Guards against AttributeError when _per_model is None.""" usage = SessionUsage() assert usage.billed_cost == 0.0 assert usage.cost_source == "opensquilla_estimate" def test_model_usage_positional_construction_keeps_billed_cost_default() -> None: """Regression: the new billed_cost field MUST be appended at the end of ModelUsage so positional callers (ModelUsage(model_id, in, out, ...)) keep aligning. Existing test (around line 112) already covers cache_read/write at the tail; this extends to billed_cost.""" mu = ModelUsage("claude-opus-4-7", 1000, 50, 200, 80) assert mu.model_id == "claude-opus-4-7" assert mu.input_tokens == 1000 assert mu.output_tokens == 50 assert mu.cache_read_tokens == 200 assert mu.cache_write_tokens == 80 assert mu.billed_cost == 0.0 # default at the tail def test_local_provider_estimates_as_free() -> None: usage = SessionUsage() usage.add(10_000, 500, "qwen3:4b", provider="ollama") assert usage.cost == 0.0 assert usage._per_model is not None assert usage._per_model["qwen3:4b"].cost == 0.0 breakdown = usage.model_breakdown[0] assert breakdown["costUsd"] == 0.0 assert breakdown["billedCostUsd"] == 0.0 def test_cloud_provider_cost_unchanged() -> None: usage = SessionUsage() usage.add(1_000_000, 0, "claude-sonnet-4", provider="anthropic") assert usage.cost == pytest.approx(3.0) def test_usage_tracker_threads_provider_to_per_model() -> None: tracker = UsageTracker() tracker.add("s", 8000, 300, "gemma3:4b", provider="ollama") sess = tracker._sessions["s"] assert sess.cost == 0.0 assert sess._per_model is not None assert sess._per_model["gemma3:4b"].provider == "ollama"