Merge 37d03fa11a into 851743adcf

strix/report/usage.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import Any, cast
 
 from agents.usage import Usage, deserialize_usage, serialize_usage
 from pydantic import BaseModel, TypeAdapter, ValidationError
@@ -58,13 +58,21 @@ _PROVIDER_USAGE = TypeAdapter(dict[str, ProviderUsage])
 
 
 class LLMUsageLedger:
-    """Aggregate SDK ``Usage`` objects and attach best-effort cost estimates."""
+    """Aggregate SDK ``Usage`` objects and attach best-effort cost estimates.
+
+    ``agents`` remains the public aggregate view in ``run.json``. The optional,
+    versioned ``model_usage`` section preserves proven agent/model ownership for
+    resumes; records without it hydrate as one legacy bucket per agent.
+    """
 
     def __init__(self) -> None:
         self._total_usage = Usage()
         self._agent_usage: dict[str, Usage] = {}
         self._agent_metadata: dict[str, dict[str, str]] = {}
         self._agent_estimated_cost: dict[str, float] = {}
+        self._model_usage: dict[tuple[str, str | None], Usage] = {}
+        self._model_fallback_cost: dict[tuple[str, str | None], float] = {}
+        self._unattributed_cost = 0.0
         self._observed_cost = 0.0
         self._estimated_cost = 0.0
         self._has_observed_cost = False
@@ -88,22 +96,34 @@ class LLMUsageLedger:
             return False
 
         normalized_agent_id = str(agent_id or "unknown")
+        normalized_model = (model or "").strip() or None
         self._total_usage.add(usage)
         self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage)
+        model_key = (normalized_agent_id, normalized_model)
+        self._model_usage.setdefault(model_key, Usage()).add(usage)
 
         metadata = self._agent_metadata.setdefault(normalized_agent_id, {})
         if agent_name:
             metadata["agent_name"] = agent_name
-        if model:
-            metadata["model"] = model
+        if normalized_model:
+            metadata["model"] = normalized_model
 
         if not self.zero_cost:
-            estimated = _estimate_litellm_cost(usage, model)
+            estimated = _estimate_litellm_cost(usage, normalized_model)
             if estimated:
                 self._estimated_cost += estimated
                 self._agent_estimated_cost[normalized_agent_id] = (
                     self._agent_estimated_cost.get(normalized_agent_id, 0.0) + estimated
                 )
+            configured_model, prices = _configured_prices()
+            fallback = (
+                _configured_usage_cost(usage, prices)
+                if prices is not None and normalized_model == configured_model
+                else estimated or 0.0
+            )
+            self._model_fallback_cost[model_key] = _round_cost(
+                self._model_fallback_cost.get(model_key, 0.0) + fallback
+            )
 
         return True
 
@@ -148,22 +168,70 @@ class LLMUsageLedger:
         model, prices = _configured_prices()
         if prices is not None:
             return _round_cost(
-                sum(self._configured_agent_cost(aid, model, prices) for aid in self._agent_usage)
+                self._unattributed_cost
+                + sum(self._configured_model_cost(key, model, prices) for key in self._model_usage)
             )
         return _round_cost(self._observed_cost if self._has_observed_cost else self._estimated_cost)
 
     def _configured_agent_cost(self, agent_id: str, model: str | None, prices: Prices) -> float:
-        if self._agent_metadata.get(agent_id, {}).get("model") == model:
-            return _configured_usage_cost(self._agent_usage[agent_id], prices)
-        return self._agent_estimated_cost.get(agent_id, 0.0)
+        return _round_cost(
+            sum(
+                self._configured_model_cost(key, model, prices)
+                for key in self._model_usage
+                if key[0] == agent_id
+            )
+        )
+
+    def _configured_model_cost(
+        self,
+        key: tuple[str, str | None],
+        model: str | None,
+        prices: Prices,
+    ) -> float:
+        if key[1] == model:
+            return _configured_usage_cost(self._model_usage[key], prices)
+        return self._model_fallback_cost.get(key, 0.0)
+
+    def _model_costs_for_record(
+        self,
+        total_cost: float,
+        model: str | None,
+        prices: Prices | None,
+    ) -> dict[tuple[str, str | None], float]:
+        if self.zero_cost:
+            return dict.fromkeys(self._model_usage, 0.0)
+        if prices is not None:
+            return {
+                key: self._configured_model_cost(key, model, prices) for key in self._model_usage
+            }
+        if not self._has_observed_cost:
+            return {key: self._model_fallback_cost.get(key, 0.0) for key in self._model_usage}
+
+        model_tokens = {
+            key: _resolve_total_tokens(usage) for key, usage in self._model_usage.items()
+        }
+        total_tokens = sum(model_tokens.values())
+        attributable_cost = max(0.0, total_cost - self._unattributed_cost)
+        return {
+            key: attributable_cost * (tokens / total_tokens) if total_tokens else 0.0
+            for key, tokens in model_tokens.items()
+        }
 
     def to_record(self) -> dict[str, Any]:
-        record = serialize_usage(self._total_usage)
+        record: dict[str, Any] = serialize_usage(self._total_usage)
         model, prices = _configured_prices()
         total_cost = self.total_cost
+        model_costs = self._model_costs_for_record(total_cost, model, prices)
         record["cost"] = total_cost
         record["providers"] = {name: tally.model_dump() for name, tally in self._providers.items()}
-        record["agents"] = []
+        agent_records: list[dict[str, Any]] = []
+        model_entries: list[dict[str, Any]] = []
+        record["agents"] = agent_records
+        record["model_usage"] = {
+            "version": 1,
+            "unattributed_cost": self._unattributed_cost,
+            "entries": model_entries,
+        }
 
         agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()}
         total_tokens = sum(agent_tokens.values())
@@ -178,7 +246,7 @@ class LLMUsageLedger:
                 else 0.0
             )
 
-            agent_record = serialize_usage(usage)
+            agent_record: dict[str, Any] = serialize_usage(usage)
             agent_record.update(
                 {
                     "agent_id": agent_id,
@@ -187,22 +255,29 @@ class LLMUsageLedger:
                     "cost": _round_cost(agent_cost),
                 }
             )
-            record["agents"].append(agent_record)
+            agent_records.append(agent_record)
+
+        for (agent_id, usage_model), usage in sorted(
+            self._model_usage.items(), key=lambda item: (item[0][0], item[0][1] or "")
+        ):
+            model_record: dict[str, Any] = serialize_usage(usage)
+            model_record.update(
+                {
+                    "agent_id": agent_id,
+                    "model": usage_model,
+                    "cost": _round_cost(model_costs[(agent_id, usage_model)]),
+                }
+            )
+            model_entries.append(model_record)
 
         return record
 
     def hydrate(self, raw_usage: Any) -> None:
-        self._total_usage = Usage()
-        self._agent_usage.clear()
-        self._agent_metadata.clear()
-        self._agent_estimated_cost.clear()
-        self._observed_cost = 0.0
-        self._estimated_cost = 0.0
-        self._has_observed_cost = False
-        self._providers = {}
+        self._reset_for_hydration()
 
         if not isinstance(raw_usage, dict):
             return
+        raw_usage = cast("dict[str, Any]", raw_usage)
 
         try:
             self._providers = _PROVIDER_USAGE.validate_python(raw_usage.get("providers") or {})
@@ -241,6 +316,107 @@ class LLMUsageLedger:
                 metadata["model"] = model
             self._agent_metadata[agent_id] = metadata
 
+        if not self._hydrate_model_usage(raw_usage.get("model_usage")):
+            # Legacy records can prove only their aggregate agent/model snapshot.
+            # Treat it as one bucket; never invent a historical model split.
+            for agent_id, usage in self._agent_usage.items():
+                model = self._agent_metadata.get(agent_id, {}).get("model")
+                key = (agent_id, model)
+                model_usage = Usage()
+                model_usage.add(usage)
+                self._model_usage[key] = model_usage
+                self._model_fallback_cost[key] = self._agent_estimated_cost.get(agent_id, 0.0)
+
+            attributed_cost = sum(self._model_fallback_cost.values())
+            self._unattributed_cost = _round_cost(max(0.0, persisted_cost - attributed_cost))
+
+    def _reset_for_hydration(self) -> None:
+        self._total_usage = Usage()
+        self._agent_usage.clear()
+        self._agent_metadata.clear()
+        self._agent_estimated_cost.clear()
+        self._model_usage.clear()
+        self._model_fallback_cost.clear()
+        self._unattributed_cost = 0.0
+        self._observed_cost = 0.0
+        self._estimated_cost = 0.0
+        self._has_observed_cost = False
+        self._providers = {}
+
+    def _hydrate_model_usage(self, raw_model_usage: Any) -> bool:
+        if raw_model_usage is None:
+            return False
+        if not isinstance(raw_model_usage, dict):
+            logger.warning("Ignoring invalid model-specific llm_usage in run.json")
+            return False
+        if raw_model_usage.get("version") != 1:
+            logger.warning("Ignoring unsupported model-specific llm_usage version in run.json")
+            return False
+        try:
+            model_usage, model_costs, unattributed_cost = _deserialize_model_usage(
+                raw_model_usage, self._agent_usage
+            )
+        except Exception:  # noqa: BLE001 - malformed legacy records fall back safely
+            logger.warning("Ignoring invalid model-specific llm_usage in run.json")
+            return False
+
+        self._model_usage = model_usage
+        self._model_fallback_cost = model_costs
+        self._unattributed_cost = unattributed_cost
+        return True
+
+
+def _deserialize_model_usage(
+    raw_model_usage: Any,
+    agent_usage: dict[str, Usage],
+) -> tuple[
+    dict[tuple[str, str | None], Usage],
+    dict[tuple[str, str | None], float],
+    float,
+]:
+    if not isinstance(raw_model_usage, dict):
+        raise TypeError
+    raw_model_usage = cast("dict[str, Any]", raw_model_usage)
+    if raw_model_usage.get("version") != 1:
+        raise ValueError
+    raw_entries = raw_model_usage.get("entries")
+    if not isinstance(raw_entries, list):
+        raise TypeError
+    typed_entries = cast("list[Any]", raw_entries)  # type: ignore[redundant-cast]
+
+    model_usage: dict[tuple[str, str | None], Usage] = {}
+    model_costs: dict[tuple[str, str | None], float] = {}
+    for raw_entry in typed_entries:
+        if not isinstance(raw_entry, dict):
+            raise TypeError
+        raw_entry = cast("dict[str, Any]", raw_entry)
+        agent_id = str(raw_entry.get("agent_id") or "").strip()
+        raw_model = raw_entry.get("model")
+        if not agent_id or agent_id not in agent_usage:
+            raise ValueError
+        if raw_model is not None and not isinstance(raw_model, str):
+            raise TypeError
+        model = (raw_model or "").strip() or None
+        key = (agent_id, model)
+        if key in model_usage:
+            raise ValueError
+        model_usage[key] = deserialize_usage(raw_entry)
+        model_costs[key] = _float_or_zero(raw_entry.get("cost"))
+
+    for agent_id, aggregate in agent_usage.items():
+        combined = Usage()
+        for key, usage in model_usage.items():
+            if key[0] == agent_id:
+                combined.add(usage)
+        if _usage_signature(combined) != _usage_signature(aggregate):
+            raise ValueError
+
+    return (
+        model_usage,
+        model_costs,
+        _float_or_zero(raw_model_usage.get("unattributed_cost")),
+    )
+
 
 def _resolve_total_tokens(usage: Usage) -> int:
     total = max(0, int(usage.total_tokens or 0))
@@ -261,6 +437,17 @@ def _usage_has_activity(usage: Usage) -> bool:
     )
 
 
+def _usage_signature(usage: Usage) -> tuple[int, int, int, int, int]:
+    details = getattr(usage, "input_tokens_details", None)
+    return (
+        _int_or_zero(getattr(usage, "requests", 0)),
+        _int_or_zero(getattr(usage, "input_tokens", 0)),
+        _int_or_zero(getattr(usage, "output_tokens", 0)),
+        _int_or_zero(getattr(usage, "total_tokens", 0)),
+        _int_or_zero(getattr(details, "cached_tokens", 0)),
+    )
+
+
 def _estimate_litellm_cost(usage: Usage, model: str | None) -> float | None:
     litellm_model = _litellm_model_name(model)
     if not litellm_model:

tests/test_hooks.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from contextlib import nullcontext
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any
 from unittest.mock import MagicMock, patch
@@ -125,6 +126,46 @@ async def test_manual_rates_make_unpriced_model_trip_budget(
         await hooks.on_llm_end(_make_context(), MagicMock(), response)
 
 
+@pytest.mark.parametrize(("budget", "raises"), [(0.01, True), (0.0101, False)])
+@pytest.mark.asyncio
+async def test_resumed_model_switch_uses_correct_live_budget(
+    budget: float,
+    raises: bool,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    monkeypatch.setattr(loader, "_override", tmp_path / "missing.json")
+    monkeypatch.setenv("STRIX_LLM", "gateway/model-a")
+    monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "2")
+    monkeypatch.setattr(loader, "_cached", None)
+    original = LLMUsageLedger()
+    original.record(agent_id="root", usage=Usage(input_tokens=1000), model="gateway/model-a")
+
+    monkeypatch.setenv("STRIX_LLM", "gateway/model-b")
+    monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "8")
+    monkeypatch.setattr(loader, "_cached", None)
+    resumed = LLMUsageLedger()
+    resumed.hydrate(original.to_record())
+
+    class ReportState:
+        def record_sdk_usage(self, **kwargs: Any) -> None:
+            resumed.record(**kwargs)
+
+        def get_total_llm_cost(self) -> float:
+            return resumed.total_cost
+
+    hooks = ReportUsageHooks(model="gateway/model-b", max_budget_usd=budget)
+    response: Any = SimpleNamespace(usage=Usage(input_tokens=1000))
+    expectation = pytest.raises(BudgetExceededError) if raises else nullcontext()
+    with (
+        patch("strix.core.hooks.get_global_report_state", return_value=ReportState()),
+        expectation,
+    ):
+        await hooks.on_llm_end(_make_context(agent_id="root"), MagicMock(), response)
+
+    assert resumed.total_cost == 0.01
+
+
 @pytest.mark.asyncio
 async def test_budget_check_uses_live_cost_accessor() -> None:
     # The check must read the live ledger, not the persisted run-record snapshot,

tests/test_pricing.py

@@ -6,12 +6,14 @@ from unittest.mock import patch
 
 import litellm
 import pytest
-from agents.usage import Usage
+from agents.usage import Usage, serialize_usage
 from pydantic import ValidationError
 
 from strix.config import loader
 from strix.config.settings import LlmSettings
+from strix.core.hooks import recomputed_budget_flags
 from strix.report.pricing import resolve_litellm_model
+from strix.report.state import ReportState
 from strix.report.usage import LLMUsageLedger
 
 
@@ -45,6 +47,27 @@ def _usage(input_tokens: int, output_tokens: int = 0, cached_tokens: int = 0) ->
     return usage
 
 
+def _configure_manual_model(
+    monkeypatch: pytest.MonkeyPatch,
+    model: str,
+    *,
+    input_rate: float,
+    output_rate: float | None = None,
+    cached_rate: float | None = None,
+) -> None:
+    monkeypatch.setenv("STRIX_LLM", model)
+    monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", str(input_rate))
+    if output_rate is None:
+        monkeypatch.delenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", raising=False)
+    else:
+        monkeypatch.setenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", str(output_rate))
+    if cached_rate is None:
+        monkeypatch.delenv("STRIX_LLM_CACHED_INPUT_COST_PER_MTOK", raising=False)
+    else:
+        monkeypatch.setenv("STRIX_LLM_CACHED_INPUT_COST_PER_MTOK", str(cached_rate))
+    monkeypatch.setattr(loader, "_cached", None)
+
+
 def test_resolves_common_bare_model_names() -> None:
     resolve_litellm_model.cache_clear()
     assert resolve_litellm_model("deepseek-v4-flash") == "deepseek/deepseek-v4-flash"
@@ -176,6 +199,7 @@ def test_hydrated_usage_uses_configured_rates_not_persisted_cost(
     original = LLMUsageLedger()
     original.record(agent_id="a", usage=_usage(1000, 200), model="gateway/unpriced-model")
     persisted = original.to_record()
+    persisted.pop("model_usage")
     persisted["cost"] = 42.0
     persisted["agents"][0]["cost"] = 42.0
     hydrated = LLMUsageLedger()
@@ -185,6 +209,208 @@ def test_hydrated_usage_uses_configured_rates_not_persisted_cost(
     assert hydrated.total_cost == 0.0036
 
 
+@pytest.mark.parametrize(
+    ("model_b_rate", "expected", "budget", "expected_flags"),
+    [
+        (8.0, 0.01, 0.012, (False, False)),
+        (0.5, 0.0025, 0.0023, (True, True)),
+    ],
+)
+def test_resumed_agent_model_switch_preserves_historical_cost(
+    model_b_rate: float,
+    expected: float,
+    budget: float,
+    expected_flags: tuple[bool, bool],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    original = LLMUsageLedger()
+    original.record(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=model_b_rate)
+    resumed = LLMUsageLedger()
+    resumed.hydrate(original.to_record())
+    resumed.record(agent_id="root", usage=_usage(1000), model="gateway/model-b")
+
+    assert resumed.total_cost == expected
+    assert resumed.to_record()["cost"] == expected
+    assert recomputed_budget_flags(expected, budget, interactive=False) == expected_flags
+
+
+def test_repeated_model_switches_keep_each_models_usage(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    first = LLMUsageLedger()
+    first.record(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=8)
+    second = LLMUsageLedger()
+    second.hydrate(first.to_record())
+    second.record(agent_id="root", usage=_usage(1000), model="gateway/model-b")
+
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    third = LLMUsageLedger()
+    third.hydrate(second.to_record())
+    third.record(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+
+    assert third.total_cost == 0.012
+    record = third.to_record()
+    assert len(record["agents"]) == 1
+    assert record["agents"][0]["input_tokens"] == 3000
+    assert record["model_usage"]["version"] == 1
+    assert len(record["model_usage"]["entries"]) == 2
+
+
+def test_model_usage_round_trips_for_root_child_and_dedupe(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    original = LLMUsageLedger()
+    for agent_id in ("root", "child"):
+        original.record(agent_id=agent_id, usage=_usage(1000), model="gateway/model-a")
+    with patch("litellm.completion_cost", return_value=0.42):
+        original.record(agent_id="dedupe", usage=_usage(500), model="deepseek/deepseek-v4-flash")
+
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=8)
+    resumed = LLMUsageLedger()
+    resumed.hydrate(original.to_record())
+    for agent_id in ("root", "child"):
+        resumed.record(agent_id=agent_id, usage=_usage(1000), model="gateway/model-b")
+
+    assert resumed.total_cost == 0.44
+    persisted = resumed.to_record()
+    hydrated = LLMUsageLedger()
+    hydrated.hydrate(persisted)
+    assert hydrated.to_record() == persisted
+    assert {entry["agent_id"] for entry in persisted["agents"]} == {
+        "root",
+        "child",
+        "dedupe",
+    }
+
+
+def test_model_switch_uses_matching_cached_token_rates(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=5, cached_rate=1)
+    original = LLMUsageLedger()
+    original.record(
+        agent_id="root",
+        usage=_usage(10, cached_tokens=4),
+        model="gateway/model-a",
+    )
+
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=10, cached_rate=2)
+    resumed = LLMUsageLedger()
+    resumed.hydrate(original.to_record())
+    resumed.record(
+        agent_id="root",
+        usage=_usage(10, cached_tokens=4),
+        model="gateway/model-b",
+    )
+
+    assert resumed.total_cost == 0.000102
+
+
+def test_legacy_agent_record_freezes_cost_when_model_changes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    legacy_usage = serialize_usage(_usage(1000))
+    legacy_usage.update(
+        {
+            "cost": 0.002,
+            "agents": [
+                {
+                    **serialize_usage(_usage(1000)),
+                    "agent_id": "root",
+                    "agent_name": "root",
+                    "model": "gateway/model-a",
+                    "cost": 0.002,
+                }
+            ],
+        }
+    )
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=8)
+    resumed = LLMUsageLedger()
+
+    resumed.hydrate(legacy_usage)
+    resumed.record(agent_id="root", usage=_usage(1000), model="gateway/model-b")
+
+    assert resumed.total_cost == 0.01
+
+
+def test_legacy_same_model_resume_reprices_without_double_counting(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    legacy_usage = {
+        **serialize_usage(_usage(1000)),
+        "cost": 42.0,
+        "agents": [
+            {
+                **serialize_usage(_usage(1000)),
+                "agent_id": "root",
+                "agent_name": "root",
+                "model": "gateway/model-a",
+                "cost": 42.0,
+            }
+        ],
+    }
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    resumed = LLMUsageLedger()
+
+    resumed.hydrate(legacy_usage)
+    resumed.record(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+
+    record = resumed.to_record()
+    assert resumed.total_cost == 0.004
+    assert record["input_tokens"] == 2000
+    assert record["model_usage"]["entries"][0]["input_tokens"] == 2000
+
+
+def test_observed_cost_precedence_survives_model_switch() -> None:
+    ledger = LLMUsageLedger()
+    with patch("litellm.completion_cost", side_effect=[0.42, 0.17]):
+        ledger.record(agent_id="root", usage=_usage(1000), model="deepseek/deepseek-v4-flash")
+        ledger.record(agent_id="root", usage=_usage(1000), model="openai/gpt-4o")
+    ledger.record_observed_cost(0.25)
+
+    assert ledger.total_cost == 0.25
+
+
+def test_zero_cost_model_switch_remains_free(monkeypatch: pytest.MonkeyPatch) -> None:
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    ledger = LLMUsageLedger()
+    ledger.zero_cost = True
+    ledger.record(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=8)
+    ledger.record(agent_id="root", usage=_usage(1000), model="gateway/model-b")
+
+    assert ledger.total_cost == 0.0
+
+
+def test_run_record_and_live_budget_agree_after_model_switch(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    _configure_manual_model(monkeypatch, "gateway/model-a", input_rate=2)
+    initial = ReportState(run_name="model-switch")
+    initial.record_sdk_usage(agent_id="root", usage=_usage(1000), model="gateway/model-a")
+
+    _configure_manual_model(monkeypatch, "gateway/model-b", input_rate=8)
+    resumed = ReportState(run_name="model-switch")
+    resumed.hydrate_from_run_dir()
+    resumed.record_sdk_usage(agent_id="root", usage=_usage(1000), model="gateway/model-b")
+
+    persisted = json.loads((resumed.get_run_dir() / "run.json").read_text())
+    assert resumed.get_total_llm_cost() == 0.01
+    assert persisted["llm_usage"]["cost"] == 0.01
+    assert recomputed_budget_flags(resumed.get_total_llm_cost(), 0.012, interactive=False) == (
+        False,
+        False,
+    )
+
+
 @pytest.mark.parametrize("value", ["invalid", "-1", "nan", "inf"])
 def test_invalid_manual_rate_is_rejected(value: str, monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", value)