From 9911ce9dcd97fab0d22bd9894009cf344aba2bf7 Mon Sep 17 00:00:00 2001
From: 9qeklajc <9qeklajc>
Date: Tue, 14 Oct 2025 00:02:04 +0200
Subject: [PATCH] fix different max cost calculation

---
 routstr/core/settings.py           |  1 +
 routstr/payment/helpers.py         | 26 +++++----
 routstr/payment/models.py          | 37 ++++++++++++
 tests/unit/test_payment_helpers.py | 94 +++++++++++++++++++++++++++++-
 4 files changed, 146 insertions(+), 12 deletions(-)

diff --git a/routstr/core/settings.py b/routstr/core/settings.py
index 35560a4c..efaa8484 100644
--- a/routstr/core/settings.py
+++ b/routstr/core/settings.py
@@ -53,6 +53,7 @@ class Settings(BaseSettings):
     tolerance_percentage: float = Field(default=1.0, env="TOLERANCE_PERCENTAGE")
     # Minimum per-request charge in millisatoshis when model pricing is free/zero
     min_request_msat: int = Field(default=1, env="MIN_REQUEST_MSAT")
+    cashu_mint_fee_msat: int = Field(default=60, env="CASHU_MINT_FEE_MSAT")
 
     # Network
     cors_origins: list[str] = Field(default_factory=lambda: ["*"], env="CORS_ORIGINS")
diff --git a/routstr/payment/helpers.py b/routstr/payment/helpers.py
index 6dc4b8ff..2064727c 100644
--- a/routstr/payment/helpers.py
+++ b/routstr/payment/helpers.py
@@ -11,7 +11,7 @@ from ..core import get_logger
 from ..core.db import ModelRow
 from ..core.settings import settings
 from ..wallet import deserialize_token_from_string
-from .models import Pricing
+from .models import Pricing, compute_effective_max_cost_msats
 
 logger = get_logger(__name__)
 
@@ -72,7 +72,8 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
         token_obj.amount if token_obj.unit == "msat" else token_obj.amount * 1000
     )
 
-    if max_cost_for_model > amount_msat:
+    fee_buffer = getattr(settings, "cashu_mint_fee_msat", 60)
+    if max_cost_for_model > amount_msat + fee_buffer:
         raise HTTPException(
             status_code=413,
             detail={
@@ -133,16 +134,17 @@ async def get_max_cost_for_model(
     row = await session.get(ModelRow, model)
     if row and row.sats_pricing:
         try:
-            sats = Pricing(**json.loads(row.sats_pricing))  # type: ignore
-            max_cost = sats.max_cost * 1000 * (1 - settings.tolerance_percentage / 100)
-            logger.debug(
-                "Found model-specific max cost",
-                extra={"model": model, "max_cost_msats": max_cost},
-            )
-            calculated_msats = int(max_cost)
-            return max(settings.min_request_msat, calculated_msats)
+            sats_dict = json.loads(row.sats_pricing)
         except Exception:
-            pass
+            sats_dict = None
+        if isinstance(sats_dict, dict):
+            effective_msats = compute_effective_max_cost_msats(sats_dict)
+            if effective_msats > 0:
+                logger.debug(
+                    "Found model-specific max cost",
+                    extra={"model": model, "max_cost_msats": effective_msats},
+                )
+                return effective_msats
 
     logger.warning(
         "Model pricing not found, using fixed cost",
@@ -201,6 +203,8 @@ async def calculate_discounted_max_cost(
             else:
                 adjusted = adjusted + math.ceil(-estimated_completion_delta_sats * 1000)
 
+    adjusted = min(max_cost_for_model, adjusted)
+
     logger.debug(
         "Discounted max cost computed",
         extra={
diff --git a/routstr/payment/models.py b/routstr/payment/models.py
index c064a8df..38e18505 100644
--- a/routstr/payment/models.py
+++ b/routstr/payment/models.py
@@ -58,6 +58,39 @@ class Model(BaseModel):
     top_provider: TopProvider | None = None
 
 
+def compute_effective_max_cost_msats(
+    pricing: Pricing | dict | None,
+) -> int:
+    if pricing is None:
+        try:
+            return max(1, int(settings.min_request_msat))
+        except Exception:
+            return 1
+
+    pricing_obj = (
+        pricing if isinstance(pricing, Pricing) else Pricing.parse_obj(pricing)  # type: ignore[arg-type]
+    )
+
+    try:
+        tolerance = float(getattr(settings, "tolerance_percentage", 0.0))
+    except Exception:
+        tolerance = 0.0
+
+    tolerance_factor = max(0.0, 1.0 - tolerance / 100.0)
+
+    try:
+        min_request_msat = max(1, int(settings.min_request_msat))
+    except Exception:
+        min_request_msat = 1
+
+    base_msats = int(float(pricing_obj.max_cost or 0.0) * 1000.0 * tolerance_factor)
+
+    if base_msats <= 0:
+        return min_request_msat
+
+    return max(min_request_msat, base_msats)
+
+
 def fetch_openrouter_models(source_filter: str | None = None) -> list[dict]:
     """Fetches model information from OpenRouter API."""
     base_url = "https://openrouter.ai/api/v1"
@@ -177,6 +210,10 @@ def _row_to_model(row: ModelRow) -> Model:
     except Exception:
         pass
 
+    if isinstance(sats_pricing, dict):
+        effective_msats = compute_effective_max_cost_msats(sats_pricing)
+        sats_pricing["max_cost"] = effective_msats / 1000.0
+
     return Model(
         id=row.id,
         name=row.name,
diff --git a/tests/unit/test_payment_helpers.py b/tests/unit/test_payment_helpers.py
index aa1b0fc2..ad757e89 100644
--- a/tests/unit/test_payment_helpers.py
+++ b/tests/unit/test_payment_helpers.py
@@ -1,12 +1,20 @@
 import os
 from unittest.mock import AsyncMock, Mock, patch
 
+import pytest
+from fastapi import HTTPException
+
 # Set required env vars before importing
 os.environ["UPSTREAM_BASE_URL"] = "http://test"
 os.environ["UPSTREAM_API_KEY"] = "test"
 
 from routstr.core.settings import settings  # noqa: E402
-from routstr.payment.helpers import get_max_cost_for_model  # noqa: E402
+from routstr.payment.helpers import (  # noqa: E402
+    calculate_discounted_max_cost,
+    check_token_balance,
+    get_max_cost_for_model,
+)
+from routstr.payment.models import Pricing  # noqa: E402
 
 
 async def test_get_max_cost_for_model_known() -> None:
@@ -73,3 +81,87 @@ async def test_get_max_cost_for_model_tolerance() -> None:
         with patch.object(settings, "tolerance_percentage", 10):
             cost = await get_max_cost_for_model("gpt-4", session=mock_session)
             assert cost == 450000  # 500 sats * 1000 * 0.9 = 450000
+
+
+async def test_calculate_discounted_max_cost_no_session() -> None:
+    with patch.object(settings, "fixed_pricing", False):
+        result = await calculate_discounted_max_cost(123, {}, session=None)
+        assert result == 123
+
+
+async def test_calculate_discounted_max_cost_clamped(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    mock_session = AsyncMock()
+    mock_exec_result = Mock()
+    mock_exec_result.all = Mock(return_value=[("gpt-4",)])
+    mock_session.exec.return_value = mock_exec_result
+
+    pricing = {
+        "prompt": 0.0,
+        "completion": 0.0,
+        "request": 0.0,
+        "image": 0.0,
+        "web_search": 0.0,
+        "internal_reasoning": 0.0,
+        "max_prompt_cost": 100.0,
+        "max_completion_cost": 200.0,
+        "max_cost": 300.0,
+    }
+
+    async def mock_get_model_cost_info(*args, **kwargs):  # type: ignore
+        return Pricing(**pricing)
+
+    monkeypatch.setattr(
+        "routstr.payment.helpers.get_model_cost_info", mock_get_model_cost_info
+    )
+
+    with patch.object(settings, "fixed_pricing", False):
+        with patch.object(settings, "tolerance_percentage", 0):
+            result = await calculate_discounted_max_cost(
+                320000,
+                {"model": "gpt-4", "max_tokens": 10, "messages": ["one"]},
+                session=mock_session,
+            )
+            assert 0 <= result <= 320000
+
+
+async def test_check_token_balance_with_fee_buffer(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class Token:
+        unit = "msat"
+        amount = 320400
+
+    def mock_deserialize(_value: str) -> Token:
+        return Token()
+
+    monkeypatch.setattr(
+        "routstr.payment.helpers.deserialize_token_from_string", mock_deserialize
+    )
+    with patch.object(settings, "cashu_mint_fee_msat", 100):
+        headers = {"x-cashu": "token"}
+        body = {"model": "gpt-4"}
+        check_token_balance(headers, body, max_cost_for_model=320450)
+
+
+def test_check_token_balance_insufficient(monkeypatch: pytest.MonkeyPatch) -> None:
+    class Token:
+        unit = "msat"
+        amount = 320200
+
+    def mock_deserialize(_value: str) -> Token:
+        return Token()
+
+    monkeypatch.setattr(
+        "routstr.payment.helpers.deserialize_token_from_string", mock_deserialize
+    )
+    with patch.object(settings, "cashu_mint_fee_msat", 100):
+        headers = {"x-cashu": "token"}
+        body = {"model": "gpt-4"}
+        with pytest.raises(HTTPException) as exc:
+            check_token_balance(headers, body, max_cost_for_model=320450)
+        assert exc.value.status_code == 413
+        detail = exc.value.detail  # type: ignore[assignment]
+        assert isinstance(detail, dict)
+        assert detail.get("type") == "minimum_balance_required"