From c3997cdb35c9bee6e9f655977eb20f6bfd640033 Mon Sep 17 00:00:00 2001
From: Sonai Biswas <79517696+Sonai124@users.noreply.github.com>
Date: Fri, 3 Jul 2026 06:10:28 +0200
Subject: [PATCH] fix: report cost for streamed OpenRouter calls (#634)

* fix: capture cost for streamed LiteLLM responses

* docs: note LiteLLM streaming metadata callbacks
---
 docs/usage/cli.mdx          |  5 +++++
 strix/config/models.py      |  6 +++--
 strix/report/state.py       | 14 +++++++++++-
 tests/test_cost_tracking.py | 44 +++++++++++++++++++++++++++++++++++++
 4 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_cost_tracking.py

diff --git a/docs/usage/cli.mdx b/docs/usage/cli.mdx
index 58d3e99..0827046 100644
--- a/docs/usage/cli.mdx
+++ b/docs/usage/cli.mdx
@@ -73,6 +73,11 @@ strix --target <target> [options]
     concurrently).
   - Cost is a best-effort estimate derived from token usage and model pricing;
     providers that do not expose priced usage may under-count.
+  - For LiteLLM-routed models, Strix enables streaming success callbacks to
+    capture provider-reported cost. Message content remains excluded, but
+    third-party LiteLLM callbacks configured in the same process can receive
+    other streaming metadata such as model names, request IDs, and token
+    counts.
 </ParamField>
 
 ## Examples
diff --git a/strix/config/models.py b/strix/config/models.py
index 7f6227f..3cca9fb 100644
--- a/strix/config/models.py
+++ b/strix/config/models.py
@@ -97,13 +97,15 @@ def _mirror_api_key_to_provider_env(model_name: str | None, api_key: str) -> Non
 
 
 def _configure_litellm_compatibility() -> None:
-    """Enable LiteLLM's permissive param handling and disable its callbacks."""
+    """Apply LiteLLM compatibility, privacy, and callback settings."""
     import litellm
 
     litellm.drop_params = True
     litellm.modify_params = True
     litellm.turn_off_message_logging = True
-    litellm.disable_streaming_logging = True
+    # Strix uses LiteLLM's success callback to capture provider-reported cost.
+    # Disabling streaming logging also disables that callback for streamed calls.
+    litellm.disable_streaming_logging = False
     litellm.suppress_debug_info = True
 
     _register_litellm_cost_callback()
diff --git a/strix/report/state.py b/strix/report/state.py
index 6f626c1..809aa09 100644
--- a/strix/report/state.py
+++ b/strix/report/state.py
@@ -3,7 +3,7 @@ import logging
 from collections.abc import Callable
 from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, cast
 from uuid import uuid4
 
 from agents.usage import Usage
@@ -383,6 +383,18 @@ def litellm_cost_callback(
             if value is not None and value > 0:
                 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)
+
     if cost is None or cost <= 0:
         return
     report_state = get_global_report_state()
diff --git a/tests/test_cost_tracking.py b/tests/test_cost_tracking.py
new file mode 100644
index 0000000..fdbf7c1
--- /dev/null
+++ b/tests/test_cost_tracking.py
@@ -0,0 +1,44 @@
+"""Tests for provider-reported LLM cost capture."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import litellm
+
+from strix.config.models import _configure_litellm_compatibility
+from strix.report.state import litellm_cost_callback
+
+
+def test_streaming_logging_stays_enabled_for_cost_callback() -> None:
+    with (
+        patch.object(litellm, "disable_streaming_logging", new=True),
+        patch("strix.config.models._register_litellm_cost_callback") as register,
+    ):
+        _configure_litellm_compatibility()
+        assert litellm.disable_streaming_logging is False
+        register.assert_called_once_with()
+
+
+def test_cost_callback_reads_openrouter_stream_usage_cost() -> None:
+    report_state = MagicMock()
+    response = SimpleNamespace(
+        usage=SimpleNamespace(cost=1.2345),
+        _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(1.2345)
+
+
+def test_cost_callback_reads_usage_cost_from_mapping_response() -> None:
+    report_state = MagicMock()
+    response = {"usage": {"cost": 0.125}}
+
+    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.125)