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) ),