From 37d03fa11aeea2c3cd8ae0594001ec4bd9df395b Mon Sep 17 00:00:00 2001 From: Muzzy5150 Date: Sat, 31 Oct 2026 15:53:38 -0700 Subject: [PATCH] fix(report): preserve usage pricing across model resumes --- strix/report/usage.py | 229 ++++++++++++++++++++++++++++++++++++++---- tests/test_hooks.py | 41 ++++++++ tests/test_pricing.py | 228 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 476 insertions(+), 22 deletions(-) diff --git a/strix/report/usage.py b/strix/report/usage.py index 3f3f10d01..c8bbb3e22 100644 --- a/strix/report/usage.py +++ b/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: diff --git a/tests/test_hooks.py b/tests/test_hooks.py index b13b2f310..74ffda2ef 100644 --- a/tests/test_hooks.py +++ b/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, diff --git a/tests/test_pricing.py b/tests/test_pricing.py index ee1130e66..3ad7dd46f 100644 --- a/tests/test_pricing.py +++ b/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)