support new in/out token labels

routstr/payment/cost_calculation.py

@@ -1,4 +1,5 @@
 import math
+from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal
 
 from pydantic.v1 import BaseModel
 
@@ -157,11 +158,18 @@ async def calculate_cost(
                 },
             )
         try:
-            input_usd = _coerce_usd(
-                usage_data.get("cost_details", {}).get("input_cost", 0)
+            cost_details = usage_data.get("cost_details", {})
+            if not isinstance(cost_details, dict):
+                cost_details = {}
+            input_usd = _first_usd(
+                cost_details,
+                "input_cost",
+                "upstream_inference_prompt_cost",
             )
-            output_usd = _coerce_usd(
-                usage_data.get("cost_details", {}).get("output_cost", 0)
+            output_usd = _first_usd(
+                cost_details,
+                "output_cost",
+                "upstream_inference_completions_cost",
             )
             return _calculate_from_usd_cost(
                 usd_cost,
@@ -256,6 +264,15 @@ def _coerce_usd(value: object) -> float:
         return 0.0
 
 
+def _first_usd(source: dict, *fields: str) -> float:
+    """Return the first positive USD value among equivalent provider fields."""
+    for field in fields:
+        value = _coerce_usd(source.get(field))
+        if value > 0:
+            return value
+    return 0.0
+
+
 def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float:
     """Resolve USD cost with clear priority order.
 
@@ -263,7 +280,11 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float:
     """
     cost_details = usage_data.get("cost_details")
     if isinstance(cost_details, dict):
-        cost = _coerce_usd(cost_details.get("total_cost"))
+        cost = _first_usd(
+            cost_details,
+            "total_cost",
+            "upstream_inference_cost",
+        )
         if cost > 0:
             return cost
 
@@ -359,27 +380,66 @@ def _calculate_from_usd_cost(
 ) -> CostData:
     """Calculate cost from USD figures, deriving input/output split from tokens."""
     provider_fee = _resolve_provider_fee(response_data.get("model", ""))
-    usd_cost = usd_cost * provider_fee
-    input_usd = input_usd * provider_fee
-    output_usd = output_usd * provider_fee
-    sats_per_usd = 1.0 / sats_usd_price()
-    cost_in_sats = usd_cost * sats_per_usd
-    cost_in_msats = math.ceil(cost_in_sats * 1000)
+    fee_decimal = Decimal(str(provider_fee))
+    usd_cost_decimal = Decimal(str(usd_cost)) * fee_decimal
+    input_usd_decimal = Decimal(str(input_usd)) * fee_decimal
+    output_usd_decimal = Decimal(str(output_usd)) * fee_decimal
+    sats_usd_decimal = Decimal(str(sats_usd_price()))
 
-    if input_usd > 0 or output_usd > 0:
-        input_msats = int((input_usd * sats_per_usd) * 1000)
-        output_msats = int((output_usd * sats_per_usd) * 1000)
-    else:
-        effective_input_tokens = (
-            input_tokens + cache_read_tokens + cache_creation_tokens
+    usd_cost = float(usd_cost_decimal)
+    input_usd = float(input_usd_decimal)
+    output_usd = float(output_usd_decimal)
+    cost_in_sats = float(usd_cost_decimal / sats_usd_decimal)
+    cost_in_msats = int(
+        (usd_cost_decimal * Decimal(1000) / sats_usd_decimal).to_integral_value(
+            rounding=ROUND_CEILING
         )
-        total_tokens = effective_input_tokens + output_tokens
-        input_msats = (
-            int(cost_in_msats * effective_input_tokens / total_tokens)
-            if total_tokens > 0
-            else 0
+    )
+
+    if input_usd_decimal > 0 or output_usd_decimal > 0:
+        # The total is the authoritative billed amount. Allocating that integer
+        # total proportionally avoids losing sub-millisatoshi remainders when
+        # input and output components are each truncated independently.
+        component_usd = input_usd_decimal + output_usd_decimal
+        input_msats = int(
+            (
+                Decimal(cost_in_msats) * input_usd_decimal / component_usd
+            ).to_integral_value(rounding=ROUND_FLOOR)
         )
         output_msats = cost_in_msats - input_msats
+    else:
+        # Providers often report only a total USD cost. Derive the visible
+        # input/output split from the model's relative token prices; raw token
+        # counts alone are misleading when completion tokens cost more.
+        try:
+            pricing_rates = _get_pricing_rates(response_data)
+        except ValueError:
+            pricing_rates = None
+
+        if pricing_rates is None:
+            input_rate = float(settings.fixed_per_1k_input_tokens) * 1000.0
+            output_rate = float(settings.fixed_per_1k_output_tokens) * 1000.0
+            cache_read_rate = input_rate
+            cache_creation_rate = input_rate
+        else:
+            input_rate, output_rate, cache_read_rate, cache_creation_rate = (
+                pricing_rates
+            )
+
+        input_weight = (
+            input_tokens * input_rate
+            + cache_read_tokens * cache_read_rate
+            + cache_creation_tokens * cache_creation_rate
+        )
+        output_weight = output_tokens * output_rate
+        total_weight = input_weight + output_weight
+
+        if total_weight > 0:
+            input_msats = math.floor(cost_in_msats * input_weight / total_weight)
+            output_msats = cost_in_msats - input_msats
+        else:
+            input_msats = 0
+            output_msats = cost_in_msats
 
     logger.info(
         "Using cost from usage data/details",

routstr/upstream/base.py

@@ -1748,14 +1748,19 @@ class BaseUpstreamProvider:
                     total_cost = max(
                         total_cost,
                         _coerce_usd(cd.get("total_cost")),
+                        _coerce_usd(cd.get("upstream_inference_cost")),
                     )
                     input_cost = max(
                         input_cost,
                         _coerce_usd(cd.get("input_cost")),
+                        _coerce_usd(cd.get("upstream_inference_prompt_cost")),
                     )
                     output_cost = max(
                         output_cost,
                         _coerce_usd(cd.get("output_cost")),
+                        _coerce_usd(
+                            cd.get("upstream_inference_completions_cost")
+                        ),
                     )
                 for field in ("total_cost", "cost"):
                     total_cost = max(

routstr/upstream/messages_dispatch.py

@@ -317,6 +317,27 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent:
         total_cost += _coerce_float(usage.get("total_cost"))
         input_cost += _coerce_float(usage.get("input_cost"))
         output_cost += _coerce_float(usage.get("output_cost"))
+        cost_details = usage.get("cost_details")
+        if isinstance(cost_details, dict):
+            total_cost = max(
+                total_cost,
+                _coerce_float(cost_details.get("total_cost")),
+                _coerce_float(cost_details.get("upstream_inference_cost")),
+            )
+            input_cost = max(
+                input_cost,
+                _coerce_float(cost_details.get("input_cost")),
+                _coerce_float(
+                    cost_details.get("upstream_inference_prompt_cost")
+                ),
+            )
+            output_cost = max(
+                output_cost,
+                _coerce_float(cost_details.get("output_cost")),
+                _coerce_float(
+                    cost_details.get("upstream_inference_completions_cost")
+                ),
+            )
 
     msg_for_meta = event.get("message")
     if isinstance(msg_for_meta, dict):
@@ -340,14 +361,21 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent:
         total_cost = max(
             total_cost,
             _coerce_float(root_cost_details.get("total_cost")),
+            _coerce_float(root_cost_details.get("upstream_inference_cost")),
         )
         input_cost = max(
             input_cost,
             _coerce_float(root_cost_details.get("input_cost")),
+            _coerce_float(
+                root_cost_details.get("upstream_inference_prompt_cost")
+            ),
         )
         output_cost = max(
             output_cost,
             _coerce_float(root_cost_details.get("output_cost")),
+            _coerce_float(
+                root_cost_details.get("upstream_inference_completions_cost")
+            ),
         )
 
     event_type = str(event.get("type") or "")

tests/unit/test_cost_calculation_caching.py

@@ -5,7 +5,7 @@ edge cases, and billing accuracy.
 """
 
 import os
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 
@@ -465,6 +465,102 @@ async def test_cache_read_only_usd_cost_response_is_billed(
     assert result.cache_read_input_tokens == 1000
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("total_cost", "input_cost", "output_cost", "expected_msats"),
+    [
+        (0.000471, 0.00023451, 0.00023649, 9420),
+        (0.00000004, 0.00000002, 0.00000002, 1),
+    ],
+)
+async def test_small_usd_cost_components_sum_to_rounded_total(
+    total_cost: float,
+    input_cost: float,
+    output_cost: float,
+    expected_msats: int,
+) -> None:
+    """Small USD component costs must retain every billed millisatoshi."""
+    response = {
+        "model": "gpt-4",
+        "usage": {
+            "prompt_tokens": 1,
+            "completion_tokens": 1,
+            "cost_details": {
+                "total_cost": total_cost,
+                "input_cost": input_cost,
+                "output_cost": output_cost,
+            },
+        },
+    }
+
+    result = await calculate_cost(response, max_cost=100000)
+
+    assert isinstance(result, CostData)
+    assert result.total_msats == expected_msats
+    assert result.input_msats + result.output_msats == result.total_msats
+
+
+@pytest.mark.asyncio
+async def test_total_only_usd_cost_uses_model_prices_for_component_split(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A reported total is split by priced tokens, not raw token counts."""
+    monkeypatch.setattr(settings, "fixed_pricing", False)
+    response = {
+        "model": "z-ai/glm-5.2-20260616",
+        "usage": {
+            "prompt_tokens": 375,
+            "completion_tokens": 218,
+            "total_cost": 0.00039088,
+        },
+    }
+    pricing = Mock(
+        prompt=0.0001,
+        completion=0.001,
+        input_cache_read=0.0001,
+        input_cache_write=0.0001,
+    )
+    model = Mock(sats_pricing=pricing)
+
+    with patch("routstr.proxy.get_model_instance", return_value=model):
+        result = await calculate_cost(response, max_cost=100000)
+
+    assert isinstance(result, CostData)
+    assert result.total_msats == 7818
+    assert result.input_msats + result.output_msats == result.total_msats
+    assert result.input_msats == 1147
+    assert result.output_msats == 6671
+
+
+@pytest.mark.asyncio
+async def test_upstream_inference_cost_details_set_nonzero_components() -> None:
+    """OpenAI-compatible upstream inference aliases retain their exact split."""
+    response = {
+        "model": "z-ai/glm-5.2-20260616",
+        "usage": {
+            "prompt_tokens": 211,
+            "completion_tokens": 500,
+            "total_tokens": 711,
+            "cost": 0.00242155,
+            "cost_details": {
+                "upstream_inference_cost": 0.00242155,
+                "upstream_inference_prompt_cost": 0.00022155,
+                "upstream_inference_completions_cost": 0.0022,
+            },
+        },
+    }
+
+    result = await calculate_cost(response, max_cost=100000)
+
+    assert isinstance(result, CostData)
+    assert result.input_tokens == 211
+    assert result.output_tokens == 500
+    assert result.input_msats == 4431
+    assert result.output_msats == 44000
+    assert result.total_msats == 48431
+    assert result.input_msats + result.output_msats == result.total_msats
+
+
 # ============================================================================
 # Test 13: Missing Usage Block
 # ============================================================================

tests/unit/test_messages_dispatch_cost_accumulation.py

@@ -161,6 +161,31 @@ def test_cost_details_extracted_from_event_root() -> None:
     assert result.output_cost == 0.005
 
 
+@pytest.mark.unit
+def test_upstream_inference_cost_details_extracted_from_usage() -> None:
+    """Provider inference aliases inside usage preserve the USD split."""
+    event = {
+        "usage": {
+            "input_tokens": 211,
+            "output_tokens": 500,
+            "cost": 0.00242155,
+            "cost_details": {
+                "upstream_inference_cost": 0.00242155,
+                "upstream_inference_prompt_cost": 0.00022155,
+                "upstream_inference_completions_cost": 0.0022,
+            },
+        }
+    }
+
+    result = annotate_event(event, None)
+
+    assert result.input_tokens == 211
+    assert result.output_tokens == 500
+    assert result.total_cost == 0.00242155
+    assert result.input_cost == 0.00022155
+    assert result.output_cost == 0.0022
+
+
 # ============================================================================
 # Test 7: No Duplicated Dict Lookups
 # ============================================================================