chore: import upstream snapshot with attribution

tests/test_engine/test_usage_tracker.py

@@ -0,0 +1,454 @@
+"""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"