From d6de546279ab2cacdac00be9d1a0bd301524fe19 Mon Sep 17 00:00:00 2001
From: 9qeklajc <tthefux@gmail.com>
Date: Sat, 11 Jul 2026 21:25:42 +0200
Subject: [PATCH] fix too low prices calculation

---
 routstr/payment/cost_calculation.py         | 50 +++++++++++----
 tests/unit/test_cost_calculation_caching.py | 69 ++++++++++++++++++++-
 2 files changed, 106 insertions(+), 13 deletions(-)

diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py
index 84f3847f..fc20a12e 100644
--- a/routstr/payment/cost_calculation.py
+++ b/routstr/payment/cost_calculation.py
@@ -367,19 +367,45 @@ def _calculate_from_usd_cost(
     cost_in_msats = math.ceil(cost_in_sats * 1000)
 
     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
-        )
-        total_tokens = effective_input_tokens + output_tokens
-        input_msats = (
-            int(cost_in_msats * effective_input_tokens / total_tokens)
-            if total_tokens > 0
-            else 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 + output_usd
+        input_msats = math.floor(cost_in_msats * input_usd / component_usd)
         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",
diff --git a/tests/unit/test_cost_calculation_caching.py b/tests/unit/test_cost_calculation_caching.py
index 93c68877..f9f53f58 100644
--- a/tests/unit/test_cost_calculation_caching.py
+++ b/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,73 @@ 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
+
+
 # ============================================================================
 # Test 13: Missing Usage Block
 # ============================================================================