fix(report): restore cost tracking for OpenRouter and other LiteLLM-routed models (#801)

strix/core/hooks.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import math
 from typing import TYPE_CHECKING, Any
 
 from agents.lifecycle import RunHooks
@@ -27,8 +28,6 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]):
     """Persist SDK-native usage after every model response."""
 
     def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None:
-        import math
-
         if max_budget_usd is not None and (
             not math.isfinite(max_budget_usd) or max_budget_usd <= 0
         ):

strix/report/state.py

@@ -534,16 +534,10 @@ def litellm_cost_callback(
                 cost = value
 
     if cost is None:
-        usage: Any = getattr(completion_response, "usage", None)
-        if usage is None and isinstance(completion_response, dict):
-            usage = cast("dict[str, Any]", completion_response).get("usage")
-        usage_cost: Any
-        if isinstance(usage, dict):
-            usage_cost = cast("dict[str, Any]", usage).get("cost")
-        else:
-            usage_cost = getattr(usage, "cost", None)
-        if isinstance(usage_cost, int | float) and usage_cost > 0:
-            cost = float(usage_cost)
+        cost = _usage_reported_cost(completion_response)
+
+    if cost is None:
+        cost = _estimate_response_cost(kwargs, completion_response)
 
     if cost is None or cost <= 0:
         return
@@ -554,3 +548,101 @@ def litellm_cost_callback(
         report_state.record_observed_llm_cost(cost)
     except Exception:
         logger.exception("Failed to record observed LiteLLM cost")
+
+
+def _usage_reported_cost(completion_response: Any) -> float | None:
+    """Provider-reported cost from the ``usage`` block (e.g. OpenRouter).
+
+    Non-BYOK responses charge everything to ``usage.cost``. BYOK responses
+    charge only the OpenRouter fee to ``usage.cost`` (often 0) and report the
+    provider charge in ``usage.cost_details.upstream_inference_cost``, so the
+    true BYOK total is the sum of the two.
+    """
+    usage: Any = getattr(completion_response, "usage", None)
+    if usage is None and isinstance(completion_response, dict):
+        usage = cast("dict[str, Any]", completion_response).get("usage")
+    if usage is None:
+        return None
+
+    def _field(container: Any, name: str) -> Any:
+        if isinstance(container, dict):
+            return cast("dict[str, Any]", container).get(name)
+        return getattr(container, name, None)
+
+    total = 0.0
+    usage_cost = _field(usage, "cost")
+    if isinstance(usage_cost, int | float) and usage_cost > 0:
+        total += float(usage_cost)
+
+    if bool(_field(usage, "is_byok")):
+        upstream = _field(_field(usage, "cost_details"), "upstream_inference_cost")
+        if isinstance(upstream, int | float) and upstream > 0:
+            total += float(upstream)
+
+    return total if total > 0 else None
+
+
+def _estimate_response_cost(kwargs: Any, completion_response: Any) -> float | None:
+    """Best-effort LiteLLM cost-map estimate when no provider-reported cost exists.
+
+    LiteLLM strips provider cost fields when rebuilding streamed responses and
+    returns no ``response_cost`` for models missing from its cost map, so try
+    the provider-prefixed name, the raw name, and the bare model name.
+    """
+    from litellm import completion_cost
+
+    model = kwargs.get("model") if isinstance(kwargs, dict) else None
+    if not isinstance(model, str) or not model:
+        if isinstance(completion_response, dict):
+            model = cast("dict[str, Any]", completion_response).get("model")
+        else:
+            model = getattr(completion_response, "model", None)
+    if not isinstance(model, str) or not model:
+        return None
+
+    provider = None
+    litellm_params = kwargs.get("litellm_params") if isinstance(kwargs, dict) else None
+    if isinstance(litellm_params, dict):
+        provider = litellm_params.get("custom_llm_provider")
+
+    usage_payload = _usage_payload(completion_response)
+    if usage_payload is None:
+        return None
+
+    candidates: list[str] = []
+    if isinstance(provider, str) and provider and not model.startswith(f"{provider}/"):
+        candidates.append(f"{provider}/{model}")
+    candidates.append(model)
+    if "/" in model:
+        candidates.append(model.rsplit("/", 1)[-1])
+
+    for candidate in candidates:
+        try:
+            value = completion_cost(
+                completion_response={"model": candidate, "usage": usage_payload},
+                model=candidate,
+            )
+        except Exception:  # nosec B112  # noqa: BLE001, S112
+            continue
+        if isinstance(value, int | float) and value > 0:
+            return float(value)
+    return None
+
+
+def _usage_payload(completion_response: Any) -> dict[str, Any] | None:
+    """Token counts as a plain dict, detached from the response's provider metadata."""
+    usage: Any = getattr(completion_response, "usage", None)
+    if usage is None and isinstance(completion_response, dict):
+        usage = cast("dict[str, Any]", completion_response).get("usage")
+    if usage is None:
+        return None
+    if hasattr(usage, "model_dump"):
+        usage = usage.model_dump()
+    if not isinstance(usage, dict):
+        return None
+    payload = cast("dict[str, Any]", usage)
+    if not payload.get("total_tokens") and not (
+        payload.get("prompt_tokens") or payload.get("completion_tokens")
+    ):
+        return None
+    return payload

tests/test_cost_tracking.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import litellm
+import pytest
 
 from strix.config.models import _configure_litellm_compatibility
 from strix.report.state import litellm_cost_callback
@@ -42,3 +43,111 @@ def test_cost_callback_reads_usage_cost_from_mapping_response() -> None:
         litellm_cost_callback({}, response)
 
     report_state.record_observed_llm_cost.assert_called_once_with(0.125)
+
+
+def test_cost_callback_reads_byok_upstream_inference_cost() -> None:
+    report_state = MagicMock()
+    response = SimpleNamespace(
+        usage=SimpleNamespace(
+            cost=0,
+            is_byok=True,
+            cost_details=SimpleNamespace(upstream_inference_cost=6.75e-06),
+        ),
+        _hidden_params={},
+    )
+
+    with patch("strix.report.state.get_global_report_state", return_value=report_state):
+        litellm_cost_callback({"response_cost": None}, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(6.75e-06)
+
+
+def test_cost_callback_sums_usage_cost_and_upstream_inference_cost() -> None:
+    report_state = MagicMock()
+    response = {
+        "usage": {
+            "cost": 0.01,
+            "is_byok": True,
+            "cost_details": {"upstream_inference_cost": 0.2},
+        }
+    }
+
+    with patch("strix.report.state.get_global_report_state", return_value=report_state):
+        litellm_cost_callback({}, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(pytest.approx(0.21))
+
+
+def test_cost_callback_ignores_upstream_cost_for_non_byok_responses() -> None:
+    report_state = MagicMock()
+    response = {
+        "usage": {
+            "cost": 0.05,
+            "is_byok": False,
+            "cost_details": {"upstream_inference_cost": 0.04},
+        }
+    }
+
+    with patch("strix.report.state.get_global_report_state", return_value=report_state):
+        litellm_cost_callback({}, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(0.05)
+
+
+def test_cost_callback_estimates_cost_with_provider_prefixed_model() -> None:
+    report_state = MagicMock()
+    response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
+    kwargs = {
+        "response_cost": None,
+        "model": "anthropic/claude-sonnet-4.5",
+        "litellm_params": {"custom_llm_provider": "openrouter"},
+    }
+
+    def fake_completion_cost(**kwargs: object) -> float:
+        if kwargs["model"] == "openrouter/anthropic/claude-sonnet-4.5":
+            return 0.5
+        raise ValueError(kwargs["model"])
+
+    with (
+        patch("strix.report.state.get_global_report_state", return_value=report_state),
+        patch("litellm.completion_cost", side_effect=fake_completion_cost),
+    ):
+        litellm_cost_callback(kwargs, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(0.5)
+
+
+def test_cost_callback_estimates_cost_with_bare_model_fallback() -> None:
+    report_state = MagicMock()
+    response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
+    kwargs = {
+        "response_cost": None,
+        "model": "openai/gpt-4o-mini",
+        "litellm_params": {"custom_llm_provider": "openrouter"},
+    }
+
+    def fake_completion_cost(**kwargs: object) -> float:
+        if kwargs["model"] == "gpt-4o-mini":
+            return 0.025
+        raise ValueError(kwargs["model"])
+
+    with (
+        patch("strix.report.state.get_global_report_state", return_value=report_state),
+        patch("litellm.completion_cost", side_effect=fake_completion_cost),
+    ):
+        litellm_cost_callback(kwargs, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(0.025)
+
+
+def test_cost_callback_records_nothing_when_no_cost_available() -> None:
+    report_state = MagicMock()
+    response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
+
+    with (
+        patch("strix.report.state.get_global_report_state", return_value=report_state),
+        patch("litellm.completion_cost", side_effect=ValueError("unknown model")),
+    ):
+        litellm_cost_callback({"response_cost": None, "model": "x/y"}, response)
+
+    report_state.record_observed_llm_cost.assert_not_called()