fix(cost): capture OpenRouter streamed usage.cost (fixes $0 kimi-k3 cost)

strix/config/models.py

@@ -277,6 +277,51 @@ def _configure_litellm_compatibility() -> None:
     litellm.suppress_debug_info = True
 
     _register_litellm_cost_callback()
+    _install_openrouter_stream_cost_capture()
+
+
+def _install_openrouter_stream_cost_capture() -> None:
+    """Preserve OpenRouter's per-stream cost, which LiteLLM drops when streaming.
+
+    OpenRouter reports the real charge in ``usage.cost`` of the final stream
+    chunk, but LiteLLM rebuilds streamed responses from token-only fields and
+    discards it (its non-streamed path stashes the cost in hidden params; the
+    streaming path does not). Every scan streams, so without this the cost is
+    lost and Strix falls back to a cost-map estimate that is missing entirely
+    for new models (e.g. kimi-k3), reporting $0. Subclass the OpenRouter
+    streaming handler to record the cost keyed by response id so the cost
+    callback can recover the exact charge for the matching rebuilt response.
+    """
+    import litellm
+    from litellm.llms.openrouter.chat.transformation import (
+        OpenRouterChatCompletionStreamingHandler,
+        OpenrouterConfig,
+    )
+
+    from strix.report.state import remember_streamed_openrouter_cost
+
+    class _StrixOpenRouterStreamingHandler(OpenRouterChatCompletionStreamingHandler):
+        def chunk_parser(self, chunk: dict[str, Any]) -> Any:
+            stream = super().chunk_parser(chunk)
+            remember_streamed_openrouter_cost(
+                chunk.get("id") or getattr(stream, "id", None), chunk.get("usage")
+            )
+            return stream
+
+    class _StrixOpenrouterConfig(OpenrouterConfig):
+        def get_model_response_iterator(
+            self, streaming_response: Any, sync_stream: bool, json_mode: bool | None = False
+        ) -> Any:
+            return _StrixOpenRouterStreamingHandler(
+                streaming_response=streaming_response,
+                sync_stream=sync_stream,
+                json_mode=json_mode,
+            )
+
+    # LiteLLM's provider-config factory reads litellm.OpenrouterConfig at call
+    # time, so overriding the attribute is enough for the subclass to take
+    # effect. (type: ignore — mypy rejects reassigning a class attribute.)
+    litellm.OpenrouterConfig = _StrixOpenrouterConfig  # type: ignore[misc]
 
 
 _OPENROUTER_ATTRIBUTION_HEADERS = {

strix/report/state.py

@@ -1,6 +1,8 @@
 import json
 import logging
 import subprocess
+import threading
+from collections import OrderedDict
 from collections.abc import Callable
 from datetime import UTC, datetime
 from importlib.metadata import PackageNotFoundError, version
@@ -507,6 +509,63 @@ class ReportState:
         self._sync_llm_usage_record()
 
 
+# LiteLLM rebuilds streamed responses from token-only chunks and drops the
+# provider-reported ``usage.cost`` that OpenRouter sends in its final stream
+# chunk (unlike the non-streamed path, which stashes it in hidden params). Since
+# every scan streams, that cost never reaches the callback below. The OpenRouter
+# streaming handler (see strix.config.models) stashes the cost here keyed by the
+# response id so the callback can recover the exact charge for the matching
+# rebuilt response.
+_STREAMED_OPENROUTER_COST_LIMIT = 4096
+_streamed_openrouter_costs: OrderedDict[str, float] = OrderedDict()
+_streamed_openrouter_costs_lock = threading.Lock()
+
+
+def openrouter_stream_cost(usage: Any) -> float | None:
+    """Total OpenRouter-reported cost from a raw stream ``usage`` block, or None.
+
+    Non-BYOK responses bill everything to ``usage.cost``. BYOK responses put the
+    OpenRouter fee in ``usage.cost`` (often 0) and the provider charge in
+    ``usage.cost_details.upstream_inference_cost``, so BYOK totals sum the two.
+    """
+    if not isinstance(usage, dict):
+        return None
+    total = 0.0
+    cost = usage.get("cost")
+    if isinstance(cost, int | float) and cost > 0:
+        total += float(cost)
+    if bool(usage.get("is_byok")):
+        details = usage.get("cost_details")
+        upstream = details.get("upstream_inference_cost") if isinstance(details, dict) else None
+        if isinstance(upstream, int | float) and upstream > 0:
+            total += float(upstream)
+    return total if total > 0 else None
+
+
+def remember_streamed_openrouter_cost(response_id: Any, usage: Any) -> None:
+    """Record an OpenRouter stream's reported cost so the cost callback can read it."""
+    if not isinstance(response_id, str) or not response_id:
+        return
+    cost = openrouter_stream_cost(usage)
+    if cost is None:
+        return
+    with _streamed_openrouter_costs_lock:
+        _streamed_openrouter_costs[response_id] = cost
+        _streamed_openrouter_costs.move_to_end(response_id)
+        while len(_streamed_openrouter_costs) > _STREAMED_OPENROUTER_COST_LIMIT:
+            _streamed_openrouter_costs.popitem(last=False)
+
+
+def _take_streamed_openrouter_cost(completion_response: Any) -> float | None:
+    response_id = getattr(completion_response, "id", None)
+    if response_id is None and isinstance(completion_response, dict):
+        response_id = cast("dict[str, Any]", completion_response).get("id")
+    if not isinstance(response_id, str) or not response_id:
+        return None
+    with _streamed_openrouter_costs_lock:
+        return _streamed_openrouter_costs.pop(response_id, None)
+
+
 def litellm_cost_callback(
     kwargs: Any,
     completion_response: Any,
@@ -541,6 +600,11 @@ def litellm_cost_callback(
     if cost is None:
         cost = _usage_reported_cost(completion_response)
 
+    # Recover the exact OpenRouter cost the streaming handler stashed for this
+    # response — LiteLLM drops it from streamed usage, so nothing above sees it.
+    if cost is None:
+        cost = _take_streamed_openrouter_cost(completion_response)
+
     if cost is None:
         cost = _estimate_response_cost(kwargs, completion_response)
 

tests/test_cost_tracking.py

@@ -8,8 +8,21 @@ 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
+import strix.report.state as state_module
+from strix.config.models import (
+    _configure_litellm_compatibility,
+    _install_openrouter_stream_cost_capture,
+)
+from strix.report.state import (
+    litellm_cost_callback,
+    openrouter_stream_cost,
+    remember_streamed_openrouter_cost,
+)
+
+
+@pytest.fixture(autouse=True)
+def _clear_streamed_costs() -> None:
+    state_module._streamed_openrouter_costs.clear()
 
 
 def test_streaming_logging_stays_enabled_for_cost_callback() -> None:
@@ -151,3 +164,83 @@ def test_cost_callback_records_nothing_when_no_cost_available() -> None:
         litellm_cost_callback({"response_cost": None, "model": "x/y"}, response)
 
     report_state.record_observed_llm_cost.assert_not_called()
+
+
+def test_openrouter_stream_cost_extracts_plain_and_byok_totals() -> None:
+    assert openrouter_stream_cost({"cost": 0.003168}) == pytest.approx(0.003168)
+    assert openrouter_stream_cost(
+        {"cost": 0.01, "is_byok": True, "cost_details": {"upstream_inference_cost": 0.2}}
+    ) == pytest.approx(0.21)
+    # Upstream cost is only added for BYOK responses.
+    assert openrouter_stream_cost(
+        {"cost": 0.05, "is_byok": False, "cost_details": {"upstream_inference_cost": 0.04}}
+    ) == pytest.approx(0.05)
+    assert openrouter_stream_cost({"prompt_tokens": 10}) is None
+    assert openrouter_stream_cost(None) is None
+
+
+def test_cost_callback_recovers_streamed_openrouter_cost_by_response_id() -> None:
+    report_state = MagicMock()
+    remember_streamed_openrouter_cost("gen-abc", {"cost": 0.42})
+    # LiteLLM strips cost from the rebuilt streamed usage; only the id survives.
+    response = SimpleNamespace(id="gen-abc", usage=SimpleNamespace(cost=None), _hidden_params={})
+
+    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": "moonshotai/kimi-k3"}, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(0.42)
+    # The entry is consumed so a later response cannot double-count it.
+    assert "gen-abc" not in state_module._streamed_openrouter_costs
+
+
+def test_streamed_openrouter_cost_prefers_provider_report_over_estimate() -> None:
+    report_state = MagicMock()
+    remember_streamed_openrouter_cost("gen-xyz", {"cost": 0.9})
+    response = SimpleNamespace(
+        id="gen-xyz",
+        usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+        _hidden_params={},
+    )
+
+    with (
+        patch("strix.report.state.get_global_report_state", return_value=report_state),
+        patch("litellm.completion_cost", return_value=0.1) as estimate,
+    ):
+        litellm_cost_callback({"response_cost": None, "model": "moonshotai/kimi-k3"}, response)
+
+    report_state.record_observed_llm_cost.assert_called_once_with(0.9)
+    estimate.assert_not_called()
+
+
+def test_remember_streamed_openrouter_cost_evicts_oldest_over_limit() -> None:
+    limit = state_module._STREAMED_OPENROUTER_COST_LIMIT
+    for i in range(limit + 5):
+        remember_streamed_openrouter_cost(f"gen-{i}", {"cost": 0.001})
+
+    assert len(state_module._streamed_openrouter_costs) == limit
+    assert "gen-0" not in state_module._streamed_openrouter_costs
+    assert f"gen-{limit + 4}" in state_module._streamed_openrouter_costs
+
+
+def test_openrouter_stream_handler_records_cost() -> None:
+    _install_openrouter_stream_cost_capture()
+    handler_cls = (
+        litellm.OpenrouterConfig()
+        .get_model_response_iterator(streaming_response=iter([]), sync_stream=True)
+        .__class__
+    )
+
+    chunk = {
+        "id": "gen-stream",
+        "created": 1,
+        "model": "moonshotai/kimi-k3",
+        "choices": [{"index": 0, "delta": {"content": None}}],
+        "usage": {"prompt_tokens": 89, "completion_tokens": 138, "cost": 0.0035055},
+    }
+    handler = handler_cls(streaming_response=iter([]), sync_stream=True)
+    handler.chunk_parser(chunk)
+
+    assert state_module._streamed_openrouter_costs["gen-stream"] == pytest.approx(0.0035055)