From ed8ac914f963cad35586b921e47b6c9f4eda3785 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 17 Jul 2026 15:53:52 +0200 Subject: [PATCH 1/8] fix: thread the served model into bearer settlement pricing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Settlement previously re-derived pricing from the response's model string through the alias map, which resolves to the best-ranked candidate for that alias — not necessarily the provider/model that actually served the request. adjust_payment_for_tokens and calculate_cost now accept the routed Model and bill its pricing directly; the string lookup remains as a warning fallback for callers without routed identity (e.g. the generic streaming finalizer). Co-Authored-By: Claude Fable 5 --- routstr/auth.py | 12 +++- routstr/payment/cost_calculation.py | 24 ++++++- routstr/upstream/base.py | 27 ++++++++ tests/unit/test_settlement_identity.py | 88 ++++++++++++++++++++++++++ 4 files changed, 147 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_settlement_identity.py diff --git a/routstr/auth.py b/routstr/auth.py index 23ad6920..516ef29c 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -4,7 +4,7 @@ import math import random import time from datetime import datetime -from typing import Optional +from typing import TYPE_CHECKING, Optional from fastapi import HTTPException from sqlalchemy import case @@ -26,6 +26,9 @@ from .wallet import ( deserialize_token_from_string, ) +if TYPE_CHECKING: + from .payment.models import Model + logger = get_logger(__name__) payments_logger = get_logger("routstr.payments") @@ -771,12 +774,17 @@ async def adjust_payment_for_tokens( response_data: dict, session: AsyncSession, deducted_max_cost: int, + model_obj: "Model | None" = None, ) -> dict: """ Adjusts the payment based on token usage in the response. This is called after the initial payment and the upstream request is complete. Returns cost data to be included in the response. + ``model_obj`` is the model that actually served the request; it is passed + through to ``calculate_cost`` so billing uses the serving candidate's + pricing instead of re-deriving it from the response's model string. + The response's usage object is normalized with the default union parser in ``calculate_cost``. """ @@ -861,7 +869,7 @@ async def adjust_payment_for_tokens( extra={"error": str(e), "fee_msats": fee_msats}, ) - match await calculate_cost(response_data, deducted_max_cost): + match await calculate_cost(response_data, deducted_max_cost, model_obj): case MaxCostData() as cost: logger.debug( "Using max cost data (no token adjustment)", diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 06272d3a..aed83bd0 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -1,4 +1,5 @@ import math +from typing import TYPE_CHECKING from pydantic.v1 import BaseModel @@ -7,6 +8,9 @@ from ..core.settings import settings from .price import sats_usd_price from .usage import normalize_usage, parse_token_count +if TYPE_CHECKING: + from .models import Model + __all__ = [ "CostData", "CostDataError", @@ -66,12 +70,17 @@ def _empty_cost(cls: type[CostData] = CostData) -> CostData: async def calculate_cost( response_data: dict, max_cost: int, + model_obj: "Model | None" = None, ) -> CostData | MaxCostData | CostDataError: """Calculate the cost of an API request based on token usage. Args: response_data: Response data containing usage information max_cost: Maximum cost in millisats + model_obj: The model that actually served the request. When given, + its pricing is billed directly; without it, pricing is re-derived + from the response's model string via the alias map, which resolves + to the best-ranked candidate — not necessarily the serving one. Returns: Cost data or error information @@ -190,7 +199,7 @@ async def calculate_cost( # Fall back to token-based pricing try: - pricing_rates = _get_pricing_rates(response_data) + pricing_rates = _get_pricing_rates(response_data, model_obj) except ValueError as e: return CostDataError(message=str(e), code="pricing_error") @@ -307,9 +316,14 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: def _get_pricing_rates( response_data: dict, + model_obj: "Model | None" = None, ) -> tuple[float, float, float, float] | None: """Get configured rates, falling back to LiteLLM's model cost map. + The served ``model_obj`` (when the caller has it) is billed directly; + otherwise the response's model string is resolved through the alias map, + which yields the best-ranked candidate rather than the serving one. + Returns: (input_rate, output_rate, cache_read_rate, cache_write_rate). ``None`` means configured fixed pricing should be used by the caller. """ @@ -323,7 +337,13 @@ def _get_pricing_rates( from .models import litellm_cost_entry response_model = response_data.get("model", "") - model_obj = get_model_instance(response_model) + if model_obj is None: + logger.warning( + "Settling without routed model identity — re-deriving pricing " + "from the response's model string via the alias map", + extra={"response_model": response_model}, + ) + model_obj = get_model_instance(response_model) if model_obj and model_obj.sats_pricing: try: diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 306aca14..35996ac8 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -797,6 +797,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, background_tasks: BackgroundTasks, requested_model: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse: """Handle streaming chat completion responses with token usage tracking and cost adjustment. @@ -839,6 +840,7 @@ class BaseUpstreamProvider: {"model": last_model_seen or "unknown", "usage": None}, new_session, max_cost_for_model, + model_obj, ) usage_finalized = True except Exception: @@ -1011,6 +1013,7 @@ class BaseUpstreamProvider: adjustment_input, session, max_cost_for_model, + model_obj, ) usage_finalized = True except Exception as e: @@ -1107,6 +1110,7 @@ class BaseUpstreamProvider: session: AsyncSession, deducted_max_cost: int, requested_model: str | None = None, + model_obj: Model | None = None, ) -> Response: """Handle non-streaming chat completion responses with token usage tracking and cost adjustment. @@ -1153,6 +1157,7 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, + model_obj, ) await session.refresh(key) @@ -1248,6 +1253,7 @@ class BaseUpstreamProvider: key: ApiKey, max_cost_for_model: int, requested_model: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse: """Handle streaming Responses API responses with token usage tracking and cost adjustment. @@ -1291,6 +1297,7 @@ class BaseUpstreamProvider: {"model": last_model_seen or "unknown", "usage": None}, new_session, max_cost_for_model, + model_obj, ) usage_finalized = True except Exception: @@ -1420,6 +1427,7 @@ class BaseUpstreamProvider: adjustment_input, session, max_cost_for_model, + model_obj, ) usage_finalized = True except Exception as e: @@ -1538,6 +1546,7 @@ class BaseUpstreamProvider: session: AsyncSession, deducted_max_cost: int, requested_model: str | None = None, + model_obj: Model | None = None, ) -> Response: """Handle non-streaming Responses API responses with token usage tracking and cost adjustment. @@ -1587,6 +1596,7 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, + model_obj, ) await session.refresh(key) @@ -1720,6 +1730,7 @@ class BaseUpstreamProvider: key: ApiKey, max_cost_for_model: int, requested_model: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse: async def stream_with_cost( max_cost_for_model: int, @@ -1785,6 +1796,7 @@ class BaseUpstreamProvider: fallback, new_session, max_cost_for_model, + model_obj, ) usage_finalized = True return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() @@ -1936,6 +1948,7 @@ class BaseUpstreamProvider: combined_data, new_session, max_cost_for_model, + model_obj, ) self.inject_cost_metadata( @@ -1983,6 +1996,7 @@ class BaseUpstreamProvider: deducted_max_cost: int, path: str, requested_model: str | None = None, + model_obj: Model | None = None, ) -> Response: try: content = await response.aread() @@ -2007,6 +2021,7 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, + model_obj, ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2101,6 +2116,7 @@ class BaseUpstreamProvider: key, max_cost_for_model, requested_model, + model_obj, ) response_json = messages_dispatch.coerce_litellm_payload(result) @@ -2112,6 +2128,7 @@ class BaseUpstreamProvider: response_json, session, max_cost_for_model, + model_obj, ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2203,6 +2220,7 @@ class BaseUpstreamProvider: key: ApiKey, max_cost_for_model: int, requested_model: str | None, + model_obj: Model | None = None, ) -> StreamingResponse: """Re-emit a litellm Anthropic-event iterator as live SSE bytes with cost reconciliation appended at end of stream.""" @@ -2253,6 +2271,7 @@ class BaseUpstreamProvider: fallback, new_session, max_cost_for_model, + model_obj, ) usage_finalized = True return ( @@ -2326,6 +2345,7 @@ class BaseUpstreamProvider: combined_data, new_session, max_cost_for_model, + model_obj, ) self.inject_cost_metadata( combined_data, cost_data, fresh_key @@ -2671,6 +2691,7 @@ class BaseUpstreamProvider: key, max_cost_for_model, requested_model=original_model_id, + model_obj=model_obj, ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) @@ -2687,6 +2708,7 @@ class BaseUpstreamProvider: max_cost_for_model, path, requested_model=original_model_id, + model_obj=model_obj, ) finally: await response.aclose() @@ -2702,6 +2724,7 @@ class BaseUpstreamProvider: max_cost_for_model, path, requested_model=original_model_id, + model_obj=model_obj, ) finally: await response.aclose() @@ -2751,6 +2774,7 @@ class BaseUpstreamProvider: max_cost_for_model, background_tasks, requested_model=original_model_id, + model_obj=model_obj, ) result.background = background_tasks return result @@ -2764,6 +2788,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, requested_model=original_model_id, + model_obj=model_obj, ) finally: await response.aclose() @@ -3016,6 +3041,7 @@ class BaseUpstreamProvider: key, max_cost_for_model, requested_model=original_model_id, + model_obj=model_obj, ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) @@ -3031,6 +3057,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, requested_model=original_model_id, + model_obj=model_obj, ) finally: await response.aclose() diff --git a/tests/unit/test_settlement_identity.py b/tests/unit/test_settlement_identity.py new file mode 100644 index 00000000..b1300df8 --- /dev/null +++ b/tests/unit/test_settlement_identity.py @@ -0,0 +1,88 @@ +"""Settlement bills the model and provider fee that actually served. + +Covers ``calculate_cost``'s served-identity parameters: a passed ``model_obj`` +is billed directly instead of re-deriving pricing from the response's model +string through the alias map (which yields the best-ranked candidate, not the +serving one), and a passed ``provider_fee`` is applied on the USD-cost path +instead of the best-ranked provider's fee. The string/alias fallbacks remain +for callers without routed identity. +""" + +import os +from unittest.mock import patch + +import pytest + +os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") +os.environ.setdefault("UPSTREAM_API_KEY", "test") + +from routstr.payment.cost_calculation import CostData, calculate_cost +from routstr.payment.models import Architecture, Model, Pricing + + +def _make_model( + model_id: str, prompt_sats: float, completion_sats: float +) -> Model: + return Model( + id=model_id, + name=model_id, + created=0, + description="", + context_length=64000, + architecture=Architecture( + modality="text->text", + input_modalities=["text"], + output_modalities=["text"], + tokenizer="Other", + instruct_type=None, + ), + pricing=Pricing(prompt=prompt_sats, completion=completion_sats), + sats_pricing=Pricing(prompt=prompt_sats, completion=completion_sats), + ) + + +WINNER = _make_model("dual-model", 0.001, 0.002) +SERVED = _make_model("dual-model", 0.005, 0.010) + +RESPONSE = { + "model": "dual-model", + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + }, +} + + +@pytest.fixture(autouse=True) +def patch_sats_usd_price() -> None: # type: ignore[misc] + with patch( + "routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-4 + ): + yield + + +@pytest.mark.asyncio +async def test_served_model_pricing_wins_over_alias_lookup() -> None: + """With ``model_obj`` given, the alias map is not consulted for pricing.""" + with patch( + "routstr.proxy.get_model_instance", return_value=WINNER + ) as alias_lookup: + result = await calculate_cost( + dict(RESPONSE), max_cost=100_000, model_obj=SERVED + ) + + assert isinstance(result, CostData) + # 1000/1000 * 5000 + 500/1000 * 10000 msats at the SERVED model's rates. + assert result.total_msats == 10_000 + alias_lookup.assert_not_called() + + +@pytest.mark.asyncio +async def test_string_fallback_still_prices_without_model_obj() -> None: + """Callers without routed identity keep the alias-map string lookup.""" + with patch("routstr.proxy.get_model_instance", return_value=WINNER): + result = await calculate_cost(dict(RESPONSE), max_cost=100_000) + + assert isinstance(result, CostData) + assert result.total_msats == 2_000 From b76fa17f81c1f008c20986d7c47a2d1e24abd72b Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 17 Jul 2026 15:57:02 +0200 Subject: [PATCH 2/8] fix: thread the served model into x-cashu settlement pricing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit X-Cashu handlers do not rewrite the upstream's echoed model string, so get_x_cashu_cost previously priced whatever wire name the upstream reported — the most collapse-prone alias lookup of all. The routed Model is now threaded from forward_x_cashu_request through the chat and Responses handler chains (and the litellm messages path) into get_x_cashu_cost, so cost and refund are computed from the model that actually served. Co-Authored-By: Claude Fable 5 --- routstr/upstream/base.py | 41 +++++++++++++++++++++----- tests/unit/test_settlement_identity.py | 25 ++++++++++++++++ 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 35996ac8..89f283ed 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -2168,6 +2168,7 @@ class BaseUpstreamProvider: requested_model, mint, request_id, + model_obj, ) response_json = messages_dispatch.coerce_litellm_payload(result) @@ -2175,7 +2176,9 @@ class BaseUpstreamProvider: if requested_model and "model" in response_json: response_json["model"] = requested_model - cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + cost_data = await self.get_x_cashu_cost( + response_json, max_cost_for_model, model_obj + ) if cost_data and "usage" in response_json and isinstance( response_json["usage"], dict @@ -2383,6 +2386,7 @@ class BaseUpstreamProvider: requested_model: str | None, mint: str | None, request_id: str | None, + model_obj: Model | None = None, ) -> StreamingResponse: """Buffer a litellm stream end-to-end, compute cost, then replay. @@ -2474,7 +2478,7 @@ class BaseUpstreamProvider: } try: cost_data = await self.get_x_cashu_cost( - response_data, max_cost_for_model + response_data, max_cost_for_model, model_obj ) if cost_data: refund_amount = messages_dispatch.compute_refund( @@ -3234,13 +3238,19 @@ class BaseUpstreamProvider: ) async def get_x_cashu_cost( - self, response_data: dict, max_cost_for_model: int + self, + response_data: dict, + max_cost_for_model: int, + model_obj: Model | None = None, ) -> MaxCostData | CostData | None: """Calculate cost for X-Cashu payment based on response data. Args: response_data: Response data containing model and usage information max_cost_for_model: Maximum cost for the model + model_obj: The model that actually served the request; billed + directly instead of re-deriving pricing from the upstream's + echoed model string Returns: Cost data object (MaxCostData or CostData) or None if calculation fails @@ -3254,6 +3264,7 @@ class BaseUpstreamProvider: match await calculate_cost( response_data, max_cost_for_model, + model_obj, ): case MaxCostData() as cost: logger.debug( @@ -3398,6 +3409,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse: """Handle streaming response for X-Cashu payment, calculating refund if needed. @@ -3471,7 +3483,7 @@ class BaseUpstreamProvider: response_data = {"usage": usage_data, "model": model} try: cost_data = await self.get_x_cashu_cost( - response_data, max_cost_for_model + response_data, max_cost_for_model, model_obj ) if cost_data: if unit == "msat": @@ -3576,6 +3588,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> Response: """Handle non-streaming response for X-Cashu payment, calculating refund if needed. @@ -3597,7 +3610,9 @@ class BaseUpstreamProvider: try: response_json = json.loads(content_str) self._apply_provider_field(response_json) - cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + cost_data = await self.get_x_cashu_cost( + response_json, max_cost_for_model, model_obj + ) if cost_data and "usage" in response_json: response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 @@ -3726,6 +3741,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse | Response: """Handle chat completion response for X-Cashu payment, detecting streaming vs non-streaming. @@ -3769,6 +3785,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=request_id, + model_obj=model_obj, ) else: return await self.handle_x_cashu_non_streaming_response( @@ -3779,6 +3796,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=request_id, + model_obj=model_obj, ) except Exception as e: @@ -3964,6 +3982,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=getattr(request.state, "request_id", None), + model_obj=model_obj, ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) @@ -4256,6 +4275,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=getattr(request.state, "request_id", None), + model_obj=model_obj, ) background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) @@ -4306,6 +4326,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse | Response: """Handle Responses API completion response for X-Cashu payment. @@ -4350,6 +4371,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=request_id, + model_obj=model_obj, ) else: return await self.handle_x_cashu_non_streaming_responses_response( @@ -4360,6 +4382,7 @@ class BaseUpstreamProvider: max_cost_for_model, mint, request_id=request_id, + model_obj=model_obj, ) except Exception as e: @@ -4387,6 +4410,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> StreamingResponse: """Handle streaming Responses API response for X-Cashu payment. @@ -4445,7 +4469,7 @@ class BaseUpstreamProvider: response_data = {"usage": usage_data, "model": model} try: cost_data = await self.get_x_cashu_cost( - response_data, max_cost_for_model + response_data, max_cost_for_model, model_obj ) if cost_data: if unit == "msat": @@ -4551,6 +4575,7 @@ class BaseUpstreamProvider: max_cost_for_model: int, mint: str | None = None, request_id: str | None = None, + model_obj: Model | None = None, ) -> Response: """Handle non-streaming Responses API response for X-Cashu payment.""" logger.debug( @@ -4561,7 +4586,9 @@ class BaseUpstreamProvider: try: response_json = json.loads(content_str) self._apply_provider_field(response_json) - cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model) + cost_data = await self.get_x_cashu_cost( + response_json, max_cost_for_model, model_obj + ) if cost_data and "usage" in response_json: response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000 diff --git a/tests/unit/test_settlement_identity.py b/tests/unit/test_settlement_identity.py index b1300df8..67eb00c2 100644 --- a/tests/unit/test_settlement_identity.py +++ b/tests/unit/test_settlement_identity.py @@ -86,3 +86,28 @@ async def test_string_fallback_still_prices_without_model_obj() -> None: assert isinstance(result, CostData) assert result.total_msats == 2_000 + + +@pytest.mark.asyncio +async def test_x_cashu_cost_uses_served_model_not_upstream_echo() -> None: + """``get_x_cashu_cost`` bills the routed model, not the raw model echo. + + X-Cashu handlers do not rewrite the upstream's echoed model string, so + without the routed model the settle would look up whatever wire name the + upstream reported. With ``model_obj`` given, the echo must be irrelevant. + """ + from routstr.upstream import GenericUpstreamProvider + + provider = GenericUpstreamProvider("http://upstream.example", "key", 1.0) + response = dict(RESPONSE, model="totally-unknown-wire-name") + + with patch( + "routstr.proxy.get_model_instance", return_value=WINNER + ) as alias_lookup: + cost = await provider.get_x_cashu_cost( + response, max_cost_for_model=100_000, model_obj=SERVED + ) + + assert cost is not None + assert cost.total_msats == 10_000 + alias_lookup.assert_not_called() From 0aebfc6dbe9687e67c163c8b5121774cb12f62c5 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 17 Jul 2026 16:01:16 +0200 Subject: [PATCH 3/8] fix: bill the serving provider's fee on the USD-cost path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The USD-cost path (and the litellm pricing fallback) resolved the provider fee via get_provider_for_model(model_id)[0] — the best-ranked provider for the alias, not the one that served. Settlement callers in the upstream handlers now pass their own provider_fee through adjust_payment_for_tokens / get_x_cashu_cost into calculate_cost; the string-derived fallback remains for callers without a serving provider. Configured model pricing is unaffected (the fee is already baked into cached pricing). Co-Authored-By: Claude Fable 5 --- routstr/auth.py | 5 ++- routstr/payment/cost_calculation.py | 17 +++++++-- routstr/upstream/base.py | 13 +++++++ tests/unit/test_messages_litellm_dispatch.py | 14 ++++++- tests/unit/test_settlement_identity.py | 39 ++++++++++++++++++++ 5 files changed, 82 insertions(+), 6 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index 516ef29c..3dcd8660 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -775,6 +775,7 @@ async def adjust_payment_for_tokens( session: AsyncSession, deducted_max_cost: int, model_obj: "Model | None" = None, + provider_fee: float | None = None, ) -> dict: """ Adjusts the payment based on token usage in the response. @@ -869,7 +870,9 @@ async def adjust_payment_for_tokens( extra={"error": str(e), "fee_msats": fee_msats}, ) - match await calculate_cost(response_data, deducted_max_cost, model_obj): + match await calculate_cost( + response_data, deducted_max_cost, model_obj, provider_fee + ): case MaxCostData() as cost: logger.debug( "Using max cost data (no token adjustment)", diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index aed83bd0..b9d03d22 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -71,6 +71,7 @@ async def calculate_cost( response_data: dict, max_cost: int, model_obj: "Model | None" = None, + provider_fee: float | None = None, ) -> CostData | MaxCostData | CostDataError: """Calculate the cost of an API request based on token usage. @@ -81,6 +82,11 @@ async def calculate_cost( its pricing is billed directly; without it, pricing is re-derived from the response's model string via the alias map, which resolves to the best-ranked candidate — not necessarily the serving one. + provider_fee: The serving provider's fee multiplier, applied on the + USD-cost path and the litellm pricing fallback (configured model + pricing already carries the fee baked in). Without it, the fee is + re-derived from the response's model string, which yields the + best-ranked provider's fee. Returns: Cost data or error information @@ -186,6 +192,7 @@ async def calculate_cost( cache_creation_tokens, output_tokens, response_data, + provider_fee, ) except Exception as e: logger.warning( @@ -199,7 +206,7 @@ async def calculate_cost( # Fall back to token-based pricing try: - pricing_rates = _get_pricing_rates(response_data, model_obj) + pricing_rates = _get_pricing_rates(response_data, model_obj, provider_fee) except ValueError as e: return CostDataError(message=str(e), code="pricing_error") @@ -317,6 +324,7 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: def _get_pricing_rates( response_data: dict, model_obj: "Model | None" = None, + provider_fee: float | None = None, ) -> tuple[float, float, float, float] | None: """Get configured rates, falling back to LiteLLM's model cost map. @@ -380,7 +388,8 @@ def _get_pricing_rates( if input_usd <= 0 or output_usd <= 0: raise ValueError(f"Incomplete LiteLLM pricing for model: {pricing_model}") - provider_fee = _resolve_provider_fee(response_model) + if provider_fee is None: + provider_fee = _resolve_provider_fee(response_model) usd_per_sat = sats_usd_price() mspp_1k = input_usd * provider_fee * 1_000_000.0 / usd_per_sat mspc_1k = output_usd * provider_fee * 1_000_000.0 / usd_per_sat @@ -441,9 +450,11 @@ def _calculate_from_usd_cost( cache_creation_tokens: int, output_tokens: int, response_data: dict, + provider_fee: float | None = None, ) -> CostData: """Calculate cost from USD figures, deriving input/output split from tokens.""" - provider_fee = _resolve_provider_fee(response_data.get("model", "")) + if provider_fee is None: + 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 diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 89f283ed..64c03317 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -841,6 +841,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True except Exception: @@ -1014,6 +1015,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True except Exception as e: @@ -1158,6 +1160,7 @@ class BaseUpstreamProvider: session, deducted_max_cost, model_obj, + self.provider_fee, ) await session.refresh(key) @@ -1298,6 +1301,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True except Exception: @@ -1428,6 +1432,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True except Exception as e: @@ -1597,6 +1602,7 @@ class BaseUpstreamProvider: session, deducted_max_cost, model_obj, + self.provider_fee, ) await session.refresh(key) @@ -1797,6 +1803,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() @@ -1949,6 +1956,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) self.inject_cost_metadata( @@ -2022,6 +2030,7 @@ class BaseUpstreamProvider: session, deducted_max_cost, model_obj, + self.provider_fee, ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2129,6 +2138,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, model_obj, + self.provider_fee, ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2275,6 +2285,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) usage_finalized = True return ( @@ -2349,6 +2360,7 @@ class BaseUpstreamProvider: new_session, max_cost_for_model, model_obj, + self.provider_fee, ) self.inject_cost_metadata( combined_data, cost_data, fresh_key @@ -3265,6 +3277,7 @@ class BaseUpstreamProvider: response_data, max_cost_for_model, model_obj, + self.provider_fee, ): case MaxCostData() as cost: logger.debug( diff --git a/tests/unit/test_messages_litellm_dispatch.py b/tests/unit/test_messages_litellm_dispatch.py index d7f919a7..ab98c270 100644 --- a/tests/unit/test_messages_litellm_dispatch.py +++ b/tests/unit/test_messages_litellm_dispatch.py @@ -502,7 +502,12 @@ async def test_streaming_emits_sse_and_reconciles_cost_at_end() -> None: captured_cost_call: dict[str, Any] = {} async def fake_adjust( - fresh_key: Any, combined_data: Any, sess: Any, max_cost: int, usage: Any = None + fresh_key: Any, + combined_data: Any, + sess: Any, + max_cost: int, + model_obj: Any = None, + provider_fee: Any = None, ) -> dict: captured_cost_call["combined_data"] = combined_data captured_cost_call["max_cost"] = max_cost @@ -591,7 +596,12 @@ async def test_streaming_handles_iterator_yielding_raw_sse_bytes() -> None: captured: dict[str, Any] = {} async def fake_adjust( - fresh_key: Any, combined_data: Any, sess: Any, max_cost: int, usage: Any = None + fresh_key: Any, + combined_data: Any, + sess: Any, + max_cost: int, + model_obj: Any = None, + provider_fee: Any = None, ) -> dict: captured["combined_data"] = combined_data return fake_cost diff --git a/tests/unit/test_settlement_identity.py b/tests/unit/test_settlement_identity.py index 67eb00c2..6febb32d 100644 --- a/tests/unit/test_settlement_identity.py +++ b/tests/unit/test_settlement_identity.py @@ -88,6 +88,45 @@ async def test_string_fallback_still_prices_without_model_obj() -> None: assert result.total_msats == 2_000 +@pytest.mark.asyncio +async def test_usd_cost_path_applies_given_provider_fee() -> None: + """The USD-cost path bills the serving provider's fee when supplied.""" + from unittest.mock import Mock + + response = dict(RESPONSE) + response["usage"] = dict(RESPONSE["usage"], cost=0.001) # type: ignore[arg-type] + + best_ranked = Mock(provider_fee=1.0) + with patch( + "routstr.proxy.get_provider_for_model", return_value=[best_ranked] + ): + result = await calculate_cost( + response, max_cost=100_000, model_obj=SERVED, provider_fee=1.5 + ) + + assert isinstance(result, CostData) + # 0.001 USD * fee 1.5 / 0.0005 USD-per-sat = 3 sats = 3000 msats. + assert result.total_msats == 3_000 + + +@pytest.mark.asyncio +async def test_usd_cost_path_falls_back_to_best_ranked_fee() -> None: + """Without a supplied fee, the alias-map provider lookup still applies.""" + from unittest.mock import Mock + + response = dict(RESPONSE) + response["usage"] = dict(RESPONSE["usage"], cost=0.001) # type: ignore[arg-type] + + best_ranked = Mock(provider_fee=2.0) + with patch( + "routstr.proxy.get_provider_for_model", return_value=[best_ranked] + ): + result = await calculate_cost(response, max_cost=100_000) + + assert isinstance(result, CostData) + assert result.total_msats == 4_000 + + @pytest.mark.asyncio async def test_x_cashu_cost_uses_served_model_not_upstream_echo() -> None: """``get_x_cashu_cost`` bills the routed model, not the raw model echo. From 02cd2cfeec06042d40b6c98d5cbb4555eb5ce32e Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 17 Jul 2026 16:13:40 +0200 Subject: [PATCH 4/8] test: fix concurrent mock leak in overrun finalization tests Two tests entered patch("routstr.auth.calculate_cost", ...) inside concurrently gathered tasks. Interleaved patch exits restore in the wrong order, leaving the mock permanently installed for every later test in the session. Hoist the patch around the gather. Co-Authored-By: Claude Fable 5 --- .../test_balance_negative_on_cost_overrun.py | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/tests/integration/test_balance_negative_on_cost_overrun.py b/tests/integration/test_balance_negative_on_cost_overrun.py index a304d91d..1cf0f701 100644 --- a/tests/integration/test_balance_negative_on_cost_overrun.py +++ b/tests/integration/test_balance_negative_on_cost_overrun.py @@ -223,15 +223,18 @@ async def test_concurrent_cost_overruns_never_negative( async with create_session() as session: fresh_key = await session.get(ApiKey, key_hash) assert fresh_key is not None - with patch( - "routstr.auth.calculate_cost", - return_value=_cost_data(actual_token_cost), - ): - await adjust_payment_for_tokens( - fresh_key, response_data, session, deducted_max_cost - ) + await adjust_payment_for_tokens( + fresh_key, response_data, session, deducted_max_cost + ) - await asyncio.gather(*[finalize() for _ in range(n_requests)]) + # Patch once around the gather: entering the same patch target from + # concurrent tasks un-patches in the wrong order and leaks the mock into + # every later test in the session. + with patch( + "routstr.auth.calculate_cost", + return_value=_cost_data(actual_token_cost), + ): + await asyncio.gather(*[finalize() for _ in range(n_requests)]) async with create_session() as session: final_key = await session.get(ApiKey, key_hash) @@ -344,15 +347,18 @@ async def test_parallel_requests_no_free_inference( async with create_session() as session: fresh_key = await session.get(ApiKey, key_hash) assert fresh_key is not None - with patch( - "routstr.auth.calculate_cost", - return_value=_cost_data(actual_token_cost), - ): - await adjust_payment_for_tokens( - fresh_key, response_data, session, deducted_max_cost - ) + await adjust_payment_for_tokens( + fresh_key, response_data, session, deducted_max_cost + ) - await asyncio.gather(finalize(), finalize()) + # Patch once around the gather: entering the same patch target from two + # concurrent tasks un-patches in the wrong order and leaks the mock into + # every later test in the session. + with patch( + "routstr.auth.calculate_cost", + return_value=_cost_data(actual_token_cost), + ): + await asyncio.gather(finalize(), finalize()) async with create_session() as session: final_key = await session.get(ApiKey, key_hash) From df4d4c44e674539f798b69f2745f4ae4226af231 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Fri, 17 Jul 2026 16:13:40 +0200 Subject: [PATCH 5/8] fix: fail over with each candidate provider's own model The failover loop resolved a single Model for the request and reused it for every provider: a fallback provider was asked to serve the routing winner's model id and billed at the winner's pricing and fee. The alias map now keeps (model, provider) candidate pairs, the proxy rebinds both per attempt, and forwarding, max-cost echo, and settlement all use the candidate actually being tried. On a failover serve the response's model field now names the serving candidate's id. The unified candidate lookup also applies the version-suffix strip (-YYYYMMDD) that model resolution already had, so version-suffixed requests no longer resolve a model yet 400 with "no provider found". Co-Authored-By: Claude Fable 5 --- routstr/algorithm.py | 18 +- routstr/proxy.py | 86 ++-- tests/integration/test_failover_billing.py | 408 ++++++++++++++++++ .../integration/test_insufficient_balance.py | 6 +- tests/unit/test_algorithm.py | 10 +- tests/unit/test_stale_reservations.py | 7 +- tests/unit/test_upstream_rate_limit.py | 7 +- 7 files changed, 485 insertions(+), 57 deletions(-) create mode 100644 tests/integration/test_failover_billing.py diff --git a/routstr/algorithm.py b/routstr/algorithm.py index 777e5110..fbc5388e 100644 --- a/routstr/algorithm.py +++ b/routstr/algorithm.py @@ -89,7 +89,9 @@ def create_model_mappings( overrides_by_key: dict[tuple[str, int], tuple], disabled_model_keys: set[tuple[str, int]], ) -> tuple[ - dict[str, "Model"], dict[str, list["BaseUpstreamProvider"]], dict[str, "Model"] + dict[str, "Model"], + dict[str, list[tuple["Model", "BaseUpstreamProvider"]]], + dict[str, "Model"], ]: """Create optimal model mappings based on cost and provider preferences. @@ -97,7 +99,9 @@ def create_model_mappings( and creates three mappings based on cost optimization: 1. model_instances: alias -> Model (all model aliases mapped to their Model objects) - 2. provider_map: alias -> List[UpstreamProvider] (sorted list of providers for each alias) + 2. provider_map: alias -> List[(Model, UpstreamProvider)] (sorted candidate + list for each alias; each provider is paired with ITS OWN model so + failover can forward and bill the candidate that actually serves) 3. unique_models: base_id -> Model (unique models without provider prefixes) The algorithm: @@ -327,7 +331,7 @@ def create_model_mappings( # Sort candidates and build final maps model_instances: dict[str, "Model"] = {} - provider_map: dict[str, list["BaseUpstreamProvider"]] = {} + provider_map: dict[str, list[tuple["Model", "BaseUpstreamProvider"]]] = {} def alias_priority(model: "Model", alias: str) -> int: """Rank how strong the mapping of alias->model is. @@ -374,13 +378,13 @@ def create_model_mappings( best_model, best_provider = items[0] model_instances[alias] = best_model - provider_map[alias] = [p for _, p in items] + provider_map[alias] = list(items) # Log provider distribution (using top provider for stats) provider_counts: dict[str, int] = {} - for providers in provider_map.values(): - if providers: - provider = providers[0] + for candidate_list in provider_map.values(): + if candidate_list: + provider = candidate_list[0][1] provider_name = getattr(provider, "upstream_name", "unknown") provider_counts[provider_name] = provider_counts.get(provider_name, 0) + 1 diff --git a/routstr/proxy.py b/routstr/proxy.py index a5534e00..c7b8ecba 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -39,8 +39,8 @@ proxy_router = APIRouter() _upstreams: list[BaseUpstreamProvider] = [] _model_instances: dict[str, Model] = {} # All aliases -> Model _provider_map: dict[ - str, list[BaseUpstreamProvider] -] = {} # All aliases -> List[Provider] + str, list[tuple[Model, BaseUpstreamProvider]] +] = {} # All aliases -> sorted [(candidate Model, its Provider)] _unique_models: dict[str, Model] = {} # Unique model.id -> Model (no duplicates) @@ -72,32 +72,44 @@ def get_upstreams() -> list[BaseUpstreamProvider]: return _upstreams -def get_model_instance(model_id: str) -> Model | None: - """Get Model instance by ID from global cache.""" +def get_candidates( + model_id: str, +) -> list[tuple[Model, BaseUpstreamProvider]] | None: + """Get the sorted (model, provider) candidate list for a model ID. + + Each provider is paired with its own model for the alias, so routing can + forward and bill the candidate that actually serves. Version suffixes + (e.g. ``-20251222``) are stripped as a retry when the exact ID is + unknown, since upstreams may return a specific version of a base model + we track. + """ if not model_id: return None model_id_lower = model_id.lower() - # Try exact match first - if model := _model_instances.get(model_id_lower): - return model + if candidates := _provider_map.get(model_id_lower): + return candidates - # Try stripping common version suffixes (e.g., -20251222) - # This handles cases where upstream returns a specific version - # but we only track the base model name. import re base_model_id = re.sub(r"-\d{8}$", "", model_id_lower) if base_model_id != model_id_lower: - if model := _model_instances.get(base_model_id): - return model + if candidates := _provider_map.get(base_model_id): + return candidates return None +def get_model_instance(model_id: str) -> Model | None: + """Get the best-ranked Model instance for a model ID.""" + candidates = get_candidates(model_id) + return candidates[0][0] if candidates else None + + def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None: - """Get UpstreamProvider list for model ID from global cache.""" - return _provider_map.get(model_id.lower()) + """Get the sorted UpstreamProvider list for a model ID.""" + candidates = get_candidates(model_id) + return [provider for _, provider in candidates] if candidates else None def get_unique_models() -> list[Model]: @@ -283,25 +295,20 @@ async def proxy( "upstream_error", "All upstreams failed", 502, request=request ) - model_obj = get_model_instance(model_id) + candidates = get_candidates(model_id) - if not model_obj: + if not candidates: return create_error_response( "invalid_model", f"Model '{model_id}' not found", 400, request=request ) - upstreams = get_provider_for_model(model_id) - if not upstreams: - return create_error_response( - "invalid_model", - f"No provider found for model '{model_id}'", - 400, - request=request, - ) - if is_ehbp: - upstreams = [upstream for upstream in upstreams if upstream.supports_ehbp] - if not upstreams: + candidates = [ + (model, upstream) + for model, upstream in candidates + if upstream.supports_ehbp + ] + if not candidates: return create_error_response( "unsupported_request", f"No EHBP-capable provider found for model '{model_id}'", @@ -309,9 +316,10 @@ async def proxy( request=request, ) - # todo figure out cost calculation since fallback provider is usually not the same price - # Use first provider for initial checks/cost calculation - # primary_upstream = upstreams[0] + # Reserve/max-cost checks use the best-ranked candidate; the failover loop + # below rebinds (model_obj, upstream) per candidate so forwarding and + # settlement always use the model of the provider actually being tried. + model_obj = candidates[0][0] _max_cost_for_model = await get_max_cost_for_model( model=model_id, session=session, model_obj=model_obj @@ -326,7 +334,7 @@ async def proxy( if x_cashu := headers.get("x-cashu", None): last_error = None - for i, upstream in enumerate(upstreams): + for i, (model_obj, upstream) in enumerate(candidates): try: if is_ehbp: if not upstream.supports_ehbp: @@ -364,7 +372,7 @@ async def proxy( "status_code": e.status_code, }, ) - if i == len(upstreams) - 1: + if i == len(candidates) - 1: last_error = e continue @@ -391,12 +399,12 @@ async def proxy( logger.debug("Processing unauthenticated GET request", extra={"path": path}) last_error_response = None - for i, upstream in enumerate(upstreams): + for i, (_, upstream) in enumerate(candidates): try: headers = upstream.prepare_headers(dict(request.headers)) response = await upstream.forward_get_request(request, path, headers) - if response.status_code in [502, 429] and i < len(upstreams) - 1: + if response.status_code in [502, 429] and i < len(candidates) - 1: error_message = "" try: if hasattr(response, "body"): @@ -426,7 +434,7 @@ async def proxy( return response except UpstreamError as e: logger.warning(f"Upstream {upstream.provider_type} failed (GET): {e}") - if i == len(upstreams) - 1: + if i == len(candidates) - 1: last_error_response = create_upstream_error_response(e, request) continue return last_error_response or create_error_response( @@ -441,7 +449,7 @@ async def proxy( # the reactive retry can never loop unboundedly. already_stripped: set[str] = set() - for i, upstream in enumerate(upstreams): + for i, (model_obj, upstream) in enumerate(candidates): headers = upstream.prepare_headers(dict(request.headers)) try: @@ -542,7 +550,7 @@ async def proxy( if response.status_code != 200: # Check if we should retry (502 Upstream Error or 429 Rate Limit) should_retry = response.status_code in [502, 429, 400, 401, 403, 404] - if should_retry and i < len(upstreams) - 1: + if should_retry and i < len(candidates) - 1: error_message = "" try: if hasattr(response, "body"): @@ -622,12 +630,12 @@ async def proxy( "provider": upstream.provider_type, "model": model_id, "status_code": e.status_code, - "retry": i < len(upstreams) - 1, + "retry": i < len(candidates) - 1, }, ) # If this was the last provider - if i == len(upstreams) - 1: + if i == len(candidates) - 1: await revert_pay_for_request(key, session, max_cost_for_model) return create_upstream_error_response(e, request) diff --git a/tests/integration/test_failover_billing.py b/tests/integration/test_failover_billing.py new file mode 100644 index 00000000..10bdca94 --- /dev/null +++ b/tests/integration/test_failover_billing.py @@ -0,0 +1,408 @@ +"""Failover requests are billed and forwarded as the provider that served them. + +Covers the whole-system settlement path when two enabled providers expose the +same model under different spellings and prices: the routing winner fails with +a 502, the fallback provider serves, and the response must be billed at the +fallback's configured rate, carry the fallback's model id in the forwarded +request body, and echo the fallback's model id to the client. +""" + +import json +from typing import Any, AsyncGenerator +from unittest.mock import patch + +import httpx +import pytest +from httpx import AsyncClient + +from routstr.payment.models import Architecture, Model, Pricing +from routstr.proxy import refresh_model_maps +from routstr.upstream.base import BaseUpstreamProvider + +CHEAP_BASE_URL = "https://cheap.example.com/v1" +EXPENSIVE_BASE_URL = "https://expensive.example.com/v1" + + +def _make_model( + model_id: str, prompt_sats: float, completion_sats: float +) -> Model: + """Build a model whose USD and sats pricing rank consistently.""" + return Model( + id=model_id, + name=model_id, + created=1, + description="test model", + context_length=8192, + architecture=Architecture( + modality="text", + input_modalities=["text"], + output_modalities=["text"], + tokenizer="gpt", + instruct_type=None, + ), + pricing=Pricing( + prompt=prompt_sats, completion=completion_sats, max_cost=50.0 + ), + sats_pricing=Pricing( + prompt=prompt_sats, completion=completion_sats, max_cost=50.0 + ), + ) + + +class _StaticProvider(BaseUpstreamProvider): + """Upstream provider with a fixed model catalog and no remote refresh.""" + + def __init__( + self, base_url: str, api_key: str, fee: float, model: Model + ) -> None: + super().__init__(base_url, api_key, fee) + self.provider_type = "custom" + self._static_model = model + + def get_cached_models(self) -> list[Model]: + return [self._static_model] + + async def refresh_models_cache(self) -> None: + pass + + +async def _install_providers( + providers: list[_StaticProvider], +) -> AsyncGenerator[None, None]: + """Install providers into the routing maps, restoring the originals after.""" + from routstr import proxy + + original_upstreams = proxy.get_upstreams() + with patch("routstr.proxy._upstreams", providers): + await refresh_model_maps() + yield + with patch("routstr.proxy._upstreams", original_upstreams): + await refresh_model_maps() + + +@pytest.fixture +async def dual_provider_maps( + patched_db_engine: None, +) -> AsyncGenerator[tuple[_StaticProvider, _StaticProvider], None]: + """Two same-tail providers under different spellings and prices.""" + cheap = _StaticProvider( + CHEAP_BASE_URL, + "key-cheap", + 1.0, + _make_model("prova/dual-model", 0.001, 0.002), + ) + expensive = _StaticProvider( + EXPENSIVE_BASE_URL, + "key-expensive", + 1.0, + _make_model("provb/dual-model", 0.005, 0.010), + ) + async for _ in _install_providers([cheap, expensive]): + yield cheap, expensive + + +def _upstream_response(request: httpx.Request) -> httpx.Response: + """502 from the cheap (winning) provider; a served completion elsewhere.""" + if request.url.host == "cheap.example.com": + return httpx.Response( + 502, + content=json.dumps({"error": {"message": "bad gateway"}}).encode(), + headers={"content-type": "application/json"}, + ) + body = { + "id": "chatcmpl-served", + "object": "chat.completion", + "created": 1, + "model": "dual-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + }, + } + return httpx.Response( + 200, + content=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_failover_serve_billed_at_serving_providers_rate( + authenticated_client: AsyncClient, + dual_provider_maps: tuple[_StaticProvider, _StaticProvider], +) -> None: + """A fallback serve is billed at the fallback's price, not the winner's. + + The cheap provider ranks first for the shared tail; it 502s and the + expensive provider serves 1000 input + 500 output tokens. At the serving + provider's sats pricing (0.005/0.010 sats per token) that is 10_000 msats; + at the winner's (0.001/0.002) it would be 2_000 msats. + """ + sent_requests: list[httpx.Request] = [] + + # Patch the network transport (not AsyncClient.send) so the in-process + # ASGI test client is untouched and only the proxy's upstream hop is mocked. + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + sent_requests.append(request) + return _upstream_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + # cost_calculation binds sats_usd_price at import time, so the price + # patch in the app fixture does not reach it; patch its own binding. + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + payload = response.json() + + # Both providers were attempted, cheapest first. + assert [r.url.host for r in sent_requests] == [ + "cheap.example.com", + "expensive.example.com", + ] + + # The fallback must be asked for ITS OWN model spelling, not the winner's. + forwarded_body = json.loads(sent_requests[1].content) + assert forwarded_body["model"] == "provb/dual-model" + + # The response echo names the model that actually served. + assert payload["model"] == "provb/dual-model" + + # Billed at the serving provider's rate: 1000/1000*5000 + 500/1000*10000. + assert payload["cost"]["total_msats"] == 10_000 + + +@pytest.fixture +async def same_id_provider_maps( + patched_db_engine: None, +) -> AsyncGenerator[None, None]: + """Two providers exposing the IDENTICAL model id at different prices.""" + cheap = _StaticProvider( + CHEAP_BASE_URL, + "key-cheap", + 1.0, + _make_model("dual-model", 0.001, 0.002), + ) + expensive = _StaticProvider( + EXPENSIVE_BASE_URL, + "key-expensive", + 1.0, + _make_model("dual-model", 0.005, 0.010), + ) + async for _ in _install_providers([cheap, expensive]): + yield + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_same_id_failover_settles_at_serving_price( + authenticated_client: AsyncClient, + same_id_provider_maps: None, +) -> None: + """Settlement must not re-derive pricing from the response's model string. + + Both providers expose the exact same model id, so the forwarded body is + identical either way — the only observable difference is the settled + amount. The response's model string resolves to the alias winner (cheap), + but the expensive provider served, so the bill must be 10_000 msats, not + the winner's 2_000. + """ + sent_requests: list[httpx.Request] = [] + + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + sent_requests.append(request) + return _upstream_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + assert [r.url.host for r in sent_requests] == [ + "cheap.example.com", + "expensive.example.com", + ] + assert response.json()["cost"]["total_msats"] == 10_000 + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_version_suffixed_model_id_routes( + authenticated_client: AsyncClient, + same_id_provider_maps: None, +) -> None: + """A version-suffixed request (``…-YYYYMMDD``) routes to the base model. + + Model resolution stripped the suffix but the provider lookup did not, so + such requests resolved a model yet found no provider and 400'd. With the + unified candidate lookup the strip applies to both. + """ + + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + return _upstream_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model-20260101", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + + +@pytest.fixture +async def fee_split_provider_maps( + patched_db_engine: None, +) -> AsyncGenerator[None, None]: + """Same-tail providers whose fees differ; the serving one charges 1.5x.""" + cheap = _StaticProvider( + CHEAP_BASE_URL, + "key-cheap", + 1.0, + _make_model("dual-model", 0.001, 0.002), + ) + expensive = _StaticProvider( + EXPENSIVE_BASE_URL, + "key-expensive", + 1.5, + _make_model("dual-model", 0.005, 0.010), + ) + async for _ in _install_providers([cheap, expensive]): + yield + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_usd_cost_serve_carries_serving_providers_fee( + authenticated_client: AsyncClient, + fee_split_provider_maps: None, +) -> None: + """The USD-cost billing path applies the SERVING provider's fee. + + The upstream that serves reports ``usage.cost`` in USD, so billing goes + through the USD-cost path where the provider fee is applied explicitly. + The serving provider's fee is 1.5; the alias winner's is 1.0. At 0.001 USD + reported cost and 0.0005 USD/sat: 0.001 * 1.5 / 0.0005 = 3 sats = 3000 + msats (fee 1.0 would give 2000). + """ + sent_requests: list[httpx.Request] = [] + + def usd_cost_response(request: httpx.Request) -> httpx.Response: + if request.url.host == "cheap.example.com": + return httpx.Response( + 502, + content=json.dumps( + {"error": {"message": "bad gateway"}} + ).encode(), + headers={"content-type": "application/json"}, + ) + body = { + "id": "chatcmpl-usd", + "object": "chat.completion", + "created": 1, + "model": "dual-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + "cost": 0.001, + }, + } + return httpx.Response( + 200, + content=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + ) + + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + sent_requests.append(request) + return usd_cost_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + assert [r.url.host for r in sent_requests] == [ + "cheap.example.com", + "expensive.example.com", + ] + assert response.json()["cost"]["total_msats"] == 3_000 diff --git a/tests/integration/test_insufficient_balance.py b/tests/integration/test_insufficient_balance.py index 11860327..fdc04e2c 100644 --- a/tests/integration/test_insufficient_balance.py +++ b/tests/integration/test_insufficient_balance.py @@ -241,8 +241,10 @@ async def test_http_402_response_shape_on_insufficient_balance( mock_upstream.prepare_headers = MagicMock(return_value={}) with ( - patch("routstr.proxy.get_model_instance", return_value=mock_model), - patch("routstr.proxy.get_provider_for_model", return_value=[mock_upstream]), + patch( + "routstr.proxy.get_candidates", + return_value=[(mock_model, mock_upstream)], + ), # Patch where it is used (proxy imports it at module level) patch( "routstr.proxy.get_max_cost_for_model", diff --git a/tests/unit/test_algorithm.py b/tests/unit/test_algorithm.py index b3b457cc..21b418db 100644 --- a/tests/unit/test_algorithm.py +++ b/tests/unit/test_algorithm.py @@ -142,7 +142,7 @@ def test_create_model_mappings_includes_db_override_for_missing_cached_model( ) assert "azure/gpt-4o" in model_instances - assert provider_map["azure/gpt-4o"] == [provider] + assert [p for _, p in provider_map["azure/gpt-4o"]] == [provider] assert "gpt-4o" in unique_models @@ -186,7 +186,7 @@ def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type( disabled_model_keys=set(), ) - providers_for_alias = provider_map["azure/gpt-4o"] + providers_for_alias = [p for _, p in provider_map["azure/gpt-4o"]] assert provider_a in providers_for_alias assert provider_b in providers_for_alias assert len(providers_for_alias) == 2 @@ -226,8 +226,8 @@ def test_create_model_mappings_applies_override_only_to_matching_provider( disabled_model_keys=set(), ) - assert provider_map["provider-b-only"] == [provider_b] - assert set(provider_map["same-id"]) == {provider_a, provider_b} + assert [p for _, p in provider_map["provider-b-only"]] == [provider_b] + assert {p for _, p in provider_map["same-id"]} == {provider_a, provider_b} def test_create_model_mappings_disables_only_matching_provider() -> None: @@ -251,4 +251,4 @@ def test_create_model_mappings_disables_only_matching_provider() -> None: disabled_model_keys={("same-id", 2)}, ) - assert provider_map["same-id"] == [provider_a] + assert [p for _, p in provider_map["same-id"]] == [provider_a] diff --git a/tests/unit/test_stale_reservations.py b/tests/unit/test_stale_reservations.py index 8c7d7c8f..b0e4a88c 100644 --- a/tests/unit/test_stale_reservations.py +++ b/tests/unit/test_stale_reservations.py @@ -355,8 +355,11 @@ async def test_proxy_reverts_reservation_on_client_disconnect() -> None: revert_mock = AsyncMock(return_value=True) with ( - patch.object(proxy_module, "get_model_instance", return_value=MagicMock()), - patch.object(proxy_module, "get_provider_for_model", return_value=[upstream]), + patch.object( + proxy_module, + "get_candidates", + return_value=[(MagicMock(), upstream)], + ), patch.object( proxy_module, "get_max_cost_for_model", AsyncMock(return_value=1_000) ), diff --git a/tests/unit/test_upstream_rate_limit.py b/tests/unit/test_upstream_rate_limit.py index ea336c14..4290c32b 100644 --- a/tests/unit/test_upstream_rate_limit.py +++ b/tests/unit/test_upstream_rate_limit.py @@ -362,8 +362,11 @@ async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None: revert_mock = AsyncMock(return_value=True) with ( - patch.object(proxy_module, "get_model_instance", return_value=MagicMock()), - patch.object(proxy_module, "get_provider_for_model", return_value=[upstream]), + patch.object( + proxy_module, + "get_candidates", + return_value=[(MagicMock(), upstream)], + ), patch.object( proxy_module, "get_max_cost_for_model", AsyncMock(return_value=1_000) ), From 3f850c54f45fcf303ccd95a4f718cbf9b902d207 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Sat, 18 Jul 2026 07:22:00 +0200 Subject: [PATCH 6/8] fix: re-reserve the serving candidate's max cost on failover Admission and reservation were sized once to the best-ranked candidate's max cost while settlement bills the candidate that actually serves, so a failover to a pricier candidate could settle far beyond the admitted envelope and consume balance reserved by other in-flight requests. Before trying a fallback candidate, raise the reservation to its own envelope; reject candidates the key cannot cover, exactly as admission would have had they been ranked first. Co-Authored-By: Claude Fable 5 --- routstr/proxy.py | 24 ++++ tests/integration/test_failover_billing.py | 144 ++++++++++++++++++++- 2 files changed, 165 insertions(+), 3 deletions(-) diff --git a/routstr/proxy.py b/routstr/proxy.py index c7b8ecba..9ee5f93b 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -450,6 +450,30 @@ async def proxy( already_stripped: set[str] = set() for i, (model_obj, upstream) in enumerate(candidates): + if i > 0 and request_body_dict: + # The reservation was sized to the previous candidate's envelope; + # settlement bills the serving candidate, so a pricier fallback + # must be re-reserved at its own max cost before it is tried. A + # candidate whose envelope the key cannot cover is rejected, just + # as it would be had it been ranked first. + candidate_max = await get_max_cost_for_model( + model=model_id, session=session, model_obj=model_obj + ) + candidate_max = await calculate_discounted_max_cost( + candidate_max, request_body_dict, model_obj=model_obj + ) + candidate_max = max(candidate_max, settings.min_request_msat) + if candidate_max > max_cost_for_model: + await revert_pay_for_request(key, session, max_cost_for_model) + try: + await pay_for_request(key, candidate_max, session) + except HTTPException: + if i == len(candidates) - 1: + raise + await pay_for_request(key, max_cost_for_model, session) + continue + max_cost_for_model = candidate_max + headers = upstream.prepare_headers(dict(request.headers)) try: diff --git a/tests/integration/test_failover_billing.py b/tests/integration/test_failover_billing.py index 10bdca94..8633e075 100644 --- a/tests/integration/test_failover_billing.py +++ b/tests/integration/test_failover_billing.py @@ -24,7 +24,10 @@ EXPENSIVE_BASE_URL = "https://expensive.example.com/v1" def _make_model( - model_id: str, prompt_sats: float, completion_sats: float + model_id: str, + prompt_sats: float, + completion_sats: float, + max_cost: float = 50.0, ) -> Model: """Build a model whose USD and sats pricing rank consistently.""" return Model( @@ -41,10 +44,10 @@ def _make_model( instruct_type=None, ), pricing=Pricing( - prompt=prompt_sats, completion=completion_sats, max_cost=50.0 + prompt=prompt_sats, completion=completion_sats, max_cost=max_cost ), sats_pricing=Pricing( - prompt=prompt_sats, completion=completion_sats, max_cost=50.0 + prompt=prompt_sats, completion=completion_sats, max_cost=max_cost ), ) @@ -406,3 +409,138 @@ async def test_usd_cost_serve_carries_serving_providers_fee( "expensive.example.com", ] assert response.json()["cost"]["total_msats"] == 3_000 + + +@pytest.fixture +async def envelope_split_provider_maps( + patched_db_engine: None, +) -> AsyncGenerator[None, None]: + """Same-id providers where the fallback's max cost dwarfs the key balance.""" + cheap = _StaticProvider( + CHEAP_BASE_URL, + "key-cheap", + 1.0, + _make_model("dual-model", 0.001, 0.002, max_cost=50.0), + ) + expensive = _StaticProvider( + EXPENSIVE_BASE_URL, + "key-expensive", + 1.0, + _make_model("dual-model", 0.005, 0.010, max_cost=20_000.0), + ) + async for _ in _install_providers([cheap, expensive]): + yield + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_failover_beyond_balance_envelope_is_rejected( + authenticated_client: AsyncClient, + envelope_split_provider_maps: None, +) -> None: + """A fallback whose max-cost envelope exceeds the balance is not served. + + Admission and reservation are sized to the best-ranked candidate's max + cost. When that candidate fails and the next one's envelope exceeds the + key's balance, serving it could settle far beyond what admission allowed, + so the request must be rejected (as it would be if the pricier candidate + were ranked first) instead of forwarded. + """ + sent_requests: list[httpx.Request] = [] + + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + sent_requests.append(request) + return _upstream_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + # The 20_000-sat envelope exceeds the key's 10_000-sat balance: the + # fallback must be rejected before its upstream is ever contacted. + assert response.status_code == 402 + assert [r.url.host for r in sent_requests] == ["cheap.example.com"] + + +@pytest.fixture +async def raised_envelope_provider_maps( + patched_db_engine: None, +) -> AsyncGenerator[None, None]: + """Same-id providers where the fallback needs a larger, affordable reserve.""" + cheap = _StaticProvider( + CHEAP_BASE_URL, + "key-cheap", + 1.0, + _make_model("dual-model", 0.001, 0.002, max_cost=50.0), + ) + expensive = _StaticProvider( + EXPENSIVE_BASE_URL, + "key-expensive", + 1.0, + _make_model("dual-model", 0.005, 0.010, max_cost=100.0), + ) + async for _ in _install_providers([cheap, expensive]): + yield + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_failover_reserves_serving_candidates_envelope( + authenticated_client: AsyncClient, + raised_envelope_provider_maps: None, +) -> None: + """An affordable pricier fallback is re-reserved, served, and billed. + + The fallback's max cost (100 sats) exceeds the winner's (50 sats) but fits + the key's balance, so the reservation is raised to the serving candidate's + envelope and the request completes, billed at the serving rate with the + unused reserve refunded. + """ + sent_requests: list[httpx.Request] = [] + + async def fake_transport( + request: httpx.Request, *args: Any, **kwargs: Any + ) -> httpx.Response: + sent_requests.append(request) + return _upstream_response(request) + + with ( + patch( + "httpx.AsyncHTTPTransport.handle_async_request", + side_effect=fake_transport, + ), + patch( + "routstr.payment.cost_calculation.sats_usd_price", + return_value=0.0005, + ), + ): + response = await authenticated_client.post( + "/v1/chat/completions", + json={ + "model": "dual-model", + "messages": [{"role": "user", "content": "hello"}], + }, + ) + + assert response.status_code == 200 + assert [r.url.host for r in sent_requests] == [ + "cheap.example.com", + "expensive.example.com", + ] + assert response.json()["cost"]["total_msats"] == 10_000 From 7f5a0cf1ae2faaff81012d2aabe6bb5365f30598 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Sat, 18 Jul 2026 07:25:52 +0200 Subject: [PATCH 7/8] refactor: drop the dead model-instance alias map get_model_instance now derives from the candidate map, leaving the module-level alias map write-only; remove it so there is a single authoritative alias source. Co-Authored-By: Claude Fable 5 --- routstr/proxy.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/routstr/proxy.py b/routstr/proxy.py index 9ee5f93b..d0228c72 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -37,7 +37,6 @@ logger = get_logger(__name__) proxy_router = APIRouter() _upstreams: list[BaseUpstreamProvider] = [] -_model_instances: dict[str, Model] = {} # All aliases -> Model _provider_map: dict[ str, list[tuple[Model, BaseUpstreamProvider]] ] = {} # All aliases -> sorted [(candidate Model, its Provider)] @@ -149,7 +148,7 @@ async def refresh_model_maps() -> None: """Refresh global model and provider maps using the cost-based algorithm.""" from sqlalchemy.orm import selectinload - global _model_instances, _provider_map, _unique_models + global _provider_map, _unique_models async with create_session() as session: # Fetch all providers with their models in a single logical operation @@ -172,7 +171,7 @@ async def refresh_model_maps() -> None: else: disabled_model_keys.add(model_key) - _model_instances, _provider_map, _unique_models = create_model_mappings( + _, _provider_map, _unique_models = create_model_mappings( upstreams=_upstreams, overrides_by_key=overrides_by_key, disabled_model_keys=disabled_model_keys, From 50437a1cc6752b0f23167ab95ed3615de0eff5b4 Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Sat, 18 Jul 2026 07:25:52 +0200 Subject: [PATCH 8/8] refactor: require explicit settlement identity at the billing seams Make model_obj/provider_fee required (still nullable) on adjust_payment_for_tokens, get_x_cashu_cost and the private pricing helpers so a call site that fails to thread the served candidate is a type error instead of a silent fallback to alias-map re-derivation. calculate_cost keeps its defaults as the one documented fallback seam. Co-Authored-By: Claude Fable 5 --- routstr/auth.py | 4 ++-- routstr/payment/cost_calculation.py | 6 +++--- routstr/upstream/base.py | 6 +++++- .../test_balance_negative_on_cost_overrun.py | 20 +++++++++++++------ tests/integration/test_child_keys.py | 2 +- .../test_free_response_stale_reservation.py | 4 ++-- .../integration/test_reservation_lifecycle.py | 4 +++- tests/unit/test_coverage_base2.py | 4 ++-- 8 files changed, 32 insertions(+), 18 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index 3dcd8660..84a57b07 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -774,8 +774,8 @@ async def adjust_payment_for_tokens( response_data: dict, session: AsyncSession, deducted_max_cost: int, - model_obj: "Model | None" = None, - provider_fee: float | None = None, + model_obj: "Model | None", + provider_fee: float | None, ) -> dict: """ Adjusts the payment based on token usage in the response. diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index b9d03d22..37ac15d3 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -323,8 +323,8 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: def _get_pricing_rates( response_data: dict, - model_obj: "Model | None" = None, - provider_fee: float | None = None, + model_obj: "Model | None", + provider_fee: float | None, ) -> tuple[float, float, float, float] | None: """Get configured rates, falling back to LiteLLM's model cost map. @@ -450,7 +450,7 @@ def _calculate_from_usd_cost( cache_creation_tokens: int, output_tokens: int, response_data: dict, - provider_fee: float | None = None, + provider_fee: float | None, ) -> CostData: """Calculate cost from USD figures, deriving input/output split from tokens.""" if provider_fee is None: diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 64c03317..f355a1aa 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -1707,11 +1707,15 @@ class BaseUpstreamProvider: try: # Finalize with "unknown" model and no usage to release reservation/charge max cost + # (no routed identity here by design: the None usage settles at + # MaxCostData before any pricing lookup can happen). await adjust_payment_for_tokens( key, {"model": "unknown", "usage": None}, session, max_cost, + model_obj=None, + provider_fee=None, ) logger.debug( "Finalized generic streaming payment in background", @@ -3253,7 +3257,7 @@ class BaseUpstreamProvider: self, response_data: dict, max_cost_for_model: int, - model_obj: Model | None = None, + model_obj: Model | None, ) -> MaxCostData | CostData | None: """Calculate cost for X-Cashu payment based on response data. diff --git a/tests/integration/test_balance_negative_on_cost_overrun.py b/tests/integration/test_balance_negative_on_cost_overrun.py index 1cf0f701..e7b8fab8 100644 --- a/tests/integration/test_balance_negative_on_cost_overrun.py +++ b/tests/integration/test_balance_negative_on_cost_overrun.py @@ -82,7 +82,9 @@ async def test_balance_never_negative_when_cost_exceeds_reservation( "routstr.auth.calculate_cost", return_value=_cost_data(actual_token_cost), ): - await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost) + await adjust_payment_for_tokens( + key, response_data, integration_session, deducted_max_cost, None, None + ) await _refresh(integration_session, key) @@ -116,7 +118,9 @@ async def test_balance_floor_at_zero_on_overrun( "routstr.auth.calculate_cost", return_value=_cost_data(actual_token_cost), ): - await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost) + await adjust_payment_for_tokens( + key, response_data, integration_session, deducted_max_cost, None, None + ) await _refresh(integration_session, key) @@ -155,7 +159,9 @@ async def test_full_cost_charged_when_balance_sufficient_for_overrun( "routstr.auth.calculate_cost", return_value=_cost_data(actual_token_cost), ): - await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost) + await adjust_payment_for_tokens( + key, response_data, integration_session, deducted_max_cost, None, None + ) await _refresh(integration_session, key) @@ -224,7 +230,7 @@ async def test_concurrent_cost_overruns_never_negative( fresh_key = await session.get(ApiKey, key_hash) assert fresh_key is not None await adjust_payment_for_tokens( - fresh_key, response_data, session, deducted_max_cost + fresh_key, response_data, session, deducted_max_cost, None, None ) # Patch once around the gather: entering the same patch target from @@ -282,7 +288,9 @@ async def test_zero_free_balance_overrun_is_safe( "routstr.auth.calculate_cost", return_value=_cost_data(actual_token_cost), ): - await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost) + await adjust_payment_for_tokens( + key, response_data, integration_session, deducted_max_cost, None, None + ) await _refresh(integration_session, key) @@ -348,7 +356,7 @@ async def test_parallel_requests_no_free_inference( fresh_key = await session.get(ApiKey, key_hash) assert fresh_key is not None await adjust_payment_for_tokens( - fresh_key, response_data, session, deducted_max_cost + fresh_key, response_data, session, deducted_max_cost, None, None ) # Patch once around the gather: entering the same patch target from two diff --git a/tests/integration/test_child_keys.py b/tests/integration/test_child_keys.py index 4c141f80..5226d357 100644 --- a/tests/integration/test_child_keys.py +++ b/tests/integration/test_child_keys.py @@ -77,7 +77,7 @@ async def test_child_key_flow(integration_session: AsyncSession) -> None: try: adjustment = await adjust_payment_for_tokens( - child_key_db, response_data, integration_session, 500 + child_key_db, response_data, integration_session, 500, None, None ) assert adjustment["total_msats"] == 400 diff --git a/tests/integration/test_free_response_stale_reservation.py b/tests/integration/test_free_response_stale_reservation.py index 27f52681..00f3c924 100644 --- a/tests/integration/test_free_response_stale_reservation.py +++ b/tests/integration/test_free_response_stale_reservation.py @@ -58,7 +58,7 @@ async def test_overrun_charges_after_reservation_swept( return_value=_cost_data(actual_token_cost), ): await adjust_payment_for_tokens( - key, response_data, integration_session, deducted_max_cost + key, response_data, integration_session, deducted_max_cost, None, None ) await integration_session.refresh(key) @@ -129,7 +129,7 @@ async def test_free_response_path_closed_end_to_end( return_value=_cost_data(actual_token_cost), ): await adjust_payment_for_tokens( - key, response_data, session, deducted_max_cost + key, response_data, session, deducted_max_cost, None, None ) async with create_session() as session: diff --git a/tests/integration/test_reservation_lifecycle.py b/tests/integration/test_reservation_lifecycle.py index ee3fc875..1f60ce9d 100644 --- a/tests/integration/test_reservation_lifecycle.py +++ b/tests/integration/test_reservation_lifecycle.py @@ -120,7 +120,9 @@ async def test_finalise_releases_reservation_and_charges_balance( response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}} with patch("routstr.auth.calculate_cost", return_value=cost_data): - await adjust_payment_for_tokens(key, response_data, integration_session, cost) + await adjust_payment_for_tokens( + key, response_data, integration_session, cost, None, None + ) await integration_session.refresh(key) diff --git a/tests/unit/test_coverage_base2.py b/tests/unit/test_coverage_base2.py index ae7be5ad..1c6a7e46 100644 --- a/tests/unit/test_coverage_base2.py +++ b/tests/unit/test_coverage_base2.py @@ -146,7 +146,7 @@ def test_get_x_cashu_cost_with_usage() -> None: "usage": {"prompt_tokens": 100, "completion_tokens": 50}, } - result = p.get_x_cashu_cost(response_data, 100000) + result = p.get_x_cashu_cost(response_data, 100000, None) # Either returns None (needs more data) or a cost object assert result is not None @@ -157,7 +157,7 @@ def test_get_x_cashu_cost_no_usage() -> None: p = BaseUpstreamProvider("https://api.test.com", "sk-test") response_data = {"model": "gpt-4"} - result = p.get_x_cashu_cost(response_data, 100000) + result = p.get_x_cashu_cost(response_data, 100000, None) # Without usage, uses max_cost assert result is not None