From 19a92568b704b7e764081946e17fd63db804eed0 Mon Sep 17 00:00:00 2001 From: redshift <213178690+1ftredsh@users.noreply.github.com> Date: Mon, 17 Aug 2026 08:24:48 +0100 Subject: [PATCH] feat(ehbp): surface tinfoil max_tokens via header for balance discounting EHBP/encrypted tinfoil requests have an opaque, HPKE-sealed body, so the proxy cannot read the client's completion cap to size the required balance. The proxy previously required the full, context-window-sized max_completion_cost regardless of what the client requested, causing 402 insufficient-balance errors for users who capped their tokens. Add client-supplied `X-Routstr-Max-Tokens` (mirroring `X-Routstr-Model`) and thread it into `calculate_discounted_max_cost` so the required balance scales with the requested completion cap without decrypting the body: - proxy reads the header on EHBP requests and passes it as the discount override (with `max_completion_tokens` body fallback for plain requests). - a present header below the 64k enclave floor (or unparseable/<=0) is rejected with a 400 invalid_request, matching Tinfoil's minimum cap. - the header is stripped before being forwarded to the enclave. --- .gitignore | 3 + routstr/payment/helpers.py | 16 ++- routstr/proxy.py | 67 +++++++++- routstr/upstream/ehbp.py | 5 + tests/integration/conftest.py | 1 + tests/unit/test_coverage_proxy.py | 172 +++++++++++++++++++++++++ tests/unit/test_payment_helpers.py | 92 +++++++++++++ tests/unit/test_tinfoil_integration.py | 3 + 8 files changed, 355 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 903d526e..35911ef7 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,6 @@ proof_backups *.todo ui_out .worktrees + +# macOS +.DS_Store diff --git a/routstr/payment/helpers.py b/routstr/payment/helpers.py index 702ab527..84fda539 100644 --- a/routstr/payment/helpers.py +++ b/routstr/payment/helpers.py @@ -155,8 +155,15 @@ async def calculate_discounted_max_cost( max_cost_for_model: int, body: dict, model_obj: Any | None = None, + max_tokens: int | None = None, ) -> int: - """Calculate the discounted max cost for a request using model pricing when available.""" + """Calculate the discounted max cost for a request using model pricing when available. + + ``max_tokens`` is an explicit override used for opaque-body requests (e.g. + EHBP/encrypted tinfoil calls) where the proxy cannot read ``max_tokens`` + from the request body. When ``None`` it falls back to the body's + ``max_tokens`` (then ``max_completion_tokens``). + """ if settings.fixed_pricing: return max_cost_for_model @@ -216,7 +223,12 @@ async def calculate_discounted_max_cost( if estimated_prompt_delta_sats > 0: adjusted = adjusted - math.floor(estimated_prompt_delta_sats * 1000) - max_tokens_raw = body.get("max_tokens", None) + if max_tokens is not None: + max_tokens_raw = max_tokens + else: + max_tokens_raw = body.get("max_tokens", None) + if max_tokens_raw is None: + max_tokens_raw = body.get("max_completion_tokens", None) if max_tokens_raw is not None: try: max_tokens_int = int(max_tokens_raw) diff --git a/routstr/proxy.py b/routstr/proxy.py index 4542d008..e94c4bcc 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -273,9 +273,22 @@ async def _proxy( # raw encrypted body to the upstream's /private/ endpoint and stream the # encrypted response back untouched — the SDK's SecureClient decrypts it. is_ehbp = "ehbp-encapsulated-key" in headers + ehbp_max_tokens: int | None = None if is_ehbp: + # The body is HPKE-sealed/opaque, so the SDK also surfaces the client's + # completion cap ("max_tokens") via X-Routstr-Max-Tokens. This lets the + # proxy size the required balance down without decrypting the body. + # A present header below the enclave's floor is rejected outright. request_body_dict = {} model_id = headers.get("x-routstr-model", "") + try: + ehbp_max_tokens = _validated_ehbp_max_tokens( + headers.get("x-routstr-max-tokens", "") + ) + except ValueError as exc: + return create_error_response( + "invalid_request", str(exc), 400, request=request + ) if not model_id: return create_error_response( "invalid_request", @@ -364,7 +377,10 @@ async def _proxy( model=model_id, session=session, model_obj=model_obj ) max_cost_for_model = await calculate_discounted_max_cost( - _max_cost_for_model, request_body_dict, model_obj=model_obj + _max_cost_for_model, + request_body_dict, + model_obj=model_obj, + max_tokens=ehbp_max_tokens, ) check_token_balance(headers, request_body_dict, max_cost_for_model) @@ -502,7 +518,10 @@ async def _proxy( 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, + request_body_dict, + model_obj=model_obj, + max_tokens=ehbp_max_tokens, ) if candidate_max > max_cost_for_model: await revert_pay_for_request( @@ -831,6 +850,50 @@ async def get_bearer_token_key( raise +def _parse_ehbp_max_tokens(raw: str) -> int | None: + """Parse the ``X-Routstr-Max-Tokens`` header into an integer (or ``None``). + + ``None`` is returned for missing, non-integer, and non-positive values. + Range validation (e.g. a Tinfoil minimum completion cap) is the caller's + responsibility via :func:`_validated_ehbp_max_tokens`. + """ + if not raw: + return None + try: + value = int(raw) + except (TypeError, ValueError): + return None + if value <= 0: + return None + return value + + +# Tinfoil enclaves reject completion caps below this floor, so a client that +# sends a lower ``X-Routstr-Max-Tokens`` is rejected rather than silently +# priced off the (unreachable) smaller cap. +_EHBP_MIN_MAX_TOKENS = 64_000 + + +def _validated_ehbp_max_tokens( + raw: str, min_tokens: int = _EHBP_MIN_MAX_TOKENS +) -> int | None: + """Parse and range-check an ``X-Routstr-Max-Tokens`` header value. + + Returns ``None`` when the header is absent (no constraint), or the parsed + integer when it is present and at least ``min_tokens``. Raises + ``ValueError`` when the header is present but invalid or below + ``min_tokens`` — callers should turn that into a 400. + """ + if not raw: + return None + value = _parse_ehbp_max_tokens(raw) + if value is None or value < min_tokens: + raise ValueError( + f"X-Routstr-Max-Tokens must be an integer of at least {min_tokens}" + ) + return value + + def extract_model_from_responses_request(request_body_dict: dict[str, Any]) -> str: if model := request_body_dict.get("model"): return model diff --git a/routstr/upstream/ehbp.py b/routstr/upstream/ehbp.py index c70a3b68..357de867 100644 --- a/routstr/upstream/ehbp.py +++ b/routstr/upstream/ehbp.py @@ -52,6 +52,10 @@ logger = get_logger(__name__) _ENCLAVE_URL_HEADER = "X-Tinfoil-Enclave-Url" _REQUEST_USAGE_HEADER = "X-Tinfoil-Request-Usage-Metrics" _RESPONSE_USAGE_HEADER = "X-Tinfoil-Usage-Metrics" +# Client supplies the completion cap via this header on EHBP (opaque-body) +# requests so the proxy can size the required balance without decrypting the +# body. It is stripped before the request is forwarded to the enclave. +_MAX_TOKENS_HEADER = "X-Routstr-Max-Tokens" _TINFOIL_PROVIDER_TYPE = "tinfoil" _TINFOIL_ALLOWED_ENCLAVE_HOST_SUFFIX = ".tinfoil.sh" _TINFOIL_ALLOWED_ENCLAVE_HOSTS = frozenset({"tinfoil.sh"}) @@ -68,6 +72,7 @@ def _normalize_upstream_model_id(model_id: str | None) -> str: _PROXY_ONLY_HEADERS = frozenset( { "x-routstr-model", + _MAX_TOKENS_HEADER.lower(), "x-tinfoil-enclave-url", "x-tinfoil-request-usage-metrics", } diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index aa10a81c..10293793 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -531,6 +531,7 @@ async def integration_app( max_cost_for_model: int, body: dict, model_obj: Any = None, + max_tokens: int | None = None, ) -> int: return max_cost_for_model diff --git a/tests/unit/test_coverage_proxy.py b/tests/unit/test_coverage_proxy.py index 0e834dc5..4722d272 100644 --- a/tests/unit/test_coverage_proxy.py +++ b/tests/unit/test_coverage_proxy.py @@ -7,6 +7,8 @@ import json import pytest from fastapi import HTTPException +from fastapi.responses import Response +from unittest.mock import AsyncMock, MagicMock, patch # =========================================================================== # parse_request_body_json @@ -64,6 +66,176 @@ def test_parse_json_rejects_non_integer_max_tokens() -> None: assert exc_info.value.status_code == 400 +# =========================================================================== +# _parse_ehbp_max_tokens +# =========================================================================== + +def test_parse_ehbp_max_tokens_valid() -> None: + """A numeric header value is parsed to an int.""" + from routstr.proxy import _parse_ehbp_max_tokens + + assert _parse_ehbp_max_tokens("1024") == 1024 + assert _parse_ehbp_max_tokens("0") is None + assert _parse_ehbp_max_tokens("-5") is None + assert _parse_ehbp_max_tokens("") is None + assert _parse_ehbp_max_tokens(None) is None + assert _parse_ehbp_max_tokens("abc") is None + assert _parse_ehbp_max_tokens("12.5") is None + + +def test_validated_ehbp_max_tokens_range() -> None: + """The validated helper enforces the 64000-token floor.""" + from routstr.proxy import ( + _EHBP_MIN_MAX_TOKENS, + _validated_ehbp_max_tokens, + ) + + # Absent header: no constraint. + assert _validated_ehbp_max_tokens("") is None + assert _validated_ehbp_max_tokens(None) is None + + # At-or-above the floor: parsed value returned. + assert _validated_ehbp_max_tokens(str(_EHBP_MIN_MAX_TOKENS)) == _EHBP_MIN_MAX_TOKENS + assert _validated_ehbp_max_tokens("100000") == 100000 + + # Below the floor / invalid: rejected. + with pytest.raises(ValueError): + _validated_ehbp_max_tokens(str(_EHBP_MIN_MAX_TOKENS - 1)) + with pytest.raises(ValueError): + _validated_ehbp_max_tokens("0") + with pytest.raises(ValueError): + _validated_ehbp_max_tokens("abc") + + +def _ehbp_request(headers: dict) -> MagicMock: + request = MagicMock() + request.method = "POST" + request.headers = { + "authorization": "Bearer test-key", + "ehbp-encapsulated-key": "abc123", + "x-routstr-model": "tinfoil-llama3-3-70b", + **headers, + } + request.body = AsyncMock(return_value=b"sealed-opaque-body") + request.state.request_id = "ehbp-max-tokens-test" + return request + + +async def test_ehbp_max_tokens_below_min_returns_400() -> None: + """A present header below the floor yields a 400 before any DB access.""" + from routstr.proxy import _proxy + + response = await _proxy( + _ehbp_request({"x-routstr-max-tokens": "1024"}), + "v1/chat/completions", + MagicMock(), + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert body["error"]["type"] == "invalid_request" + assert "at least 64000" in body["error"]["message"] + + +async def test_ehbp_max_tokens_invalid_returns_400() -> None: + """A present but unparseable header value returns 400 as well.""" + from routstr.proxy import _proxy + + response = await _proxy( + _ehbp_request({"x-routstr-max-tokens": "not-a-number"}), + "v1/chat/completions", + MagicMock(), + ) + + assert response.status_code == 400 + body = json.loads(response.body) + assert "at least 64000" in body["error"]["message"] + + +async def test_ehbp_max_tokens_flows_into_discount() -> None: + """A valid header cap is passed through to the cost discount.""" + from routstr.proxy import _proxy + + model = MagicMock() + upstream = MagicMock() + key = MagicMock() + seen: dict[str, object] = {} + + async def fake_discount( + max_cost_for_model: int, + body: dict, + model_obj: object = None, + max_tokens: int | None = None, + ) -> int: + seen["max_tokens"] = max_tokens + seen["body"] = body + return max_cost_for_model + + with ( + patch("routstr.proxy.get_candidates", return_value=[(model, upstream)]), + patch("routstr.proxy.get_max_cost_for_model", AsyncMock(return_value=100_000)), + patch( + "routstr.proxy.calculate_discounted_max_cost", + side_effect=fake_discount, + ), + patch("routstr.proxy.check_token_balance"), + patch("routstr.proxy.get_bearer_token_key", AsyncMock(return_value=key)), + patch("routstr.proxy.pay_for_request", AsyncMock()), + patch("routstr.proxy.get_reservation_snapshot", AsyncMock(return_value=None)), + patch("routstr.proxy.forward_ehbp_request", AsyncMock(return_value=Response(status_code=200))), + ): + response = await _proxy( + _ehbp_request({"x-routstr-max-tokens": "100000"}), + "v1/chat/completions", + MagicMock(), + ) + + assert response.status_code == 200 + assert seen["max_tokens"] == 100000 + assert seen["body"] == {} # opaque body stays opaque + + +async def test_ehbp_max_tokens_absent_is_allowed() -> None: + """No header still works (backwards compatible with older SDKs).""" + from routstr.proxy import _proxy + + model = MagicMock() + upstream = MagicMock() + key = MagicMock() + seen: dict[str, object] = {} + + async def fake_discount( + max_cost_for_model: int, + body: dict, + model_obj: object = None, + max_tokens: int | None = None, + ) -> int: + seen["max_tokens"] = max_tokens + return max_cost_for_model + + with ( + patch("routstr.proxy.get_candidates", return_value=[(model, upstream)]), + patch("routstr.proxy.get_max_cost_for_model", AsyncMock(return_value=100_000)), + patch( + "routstr.proxy.calculate_discounted_max_cost", + side_effect=fake_discount, + ), + patch("routstr.proxy.check_token_balance"), + patch("routstr.proxy.get_bearer_token_key", AsyncMock(return_value=key)), + patch("routstr.proxy.pay_for_request", AsyncMock()), + patch("routstr.proxy.get_reservation_snapshot", AsyncMock(return_value=None)), + patch("routstr.proxy.forward_ehbp_request", AsyncMock(return_value=Response(status_code=200))), + ): + response = await _proxy( + _ehbp_request({}), + "v1/chat/completions", + MagicMock(), + ) + + assert response.status_code == 200 + assert seen["max_tokens"] is None + + # =========================================================================== # extract_model_from_responses_request # =========================================================================== diff --git a/tests/unit/test_payment_helpers.py b/tests/unit/test_payment_helpers.py index 6809d94c..2151233c 100644 --- a/tests/unit/test_payment_helpers.py +++ b/tests/unit/test_payment_helpers.py @@ -155,3 +155,95 @@ async def test_discounted_max_cost_floors_at_min_request_msat() -> None: cost = await calculate_discounted_max_cost(150_000, body, model_obj) assert cost == 1000 + + +async def test_discounted_max_cost_honors_explicit_max_tokens_override() -> None: + """The ``max_tokens`` override parameter (opaque/EHBP body) is honored. + + Body is empty (encrypted) so only the override can drive the discount. + completion max budget = 100 sats; 80k tokens * 0.001 = 80 sats used => the + required balance is trimmed from 100 sats to 80 sats (80000 msats). + """ + from routstr.payment.helpers import calculate_discounted_max_cost + + pricing = Mock() + pricing.prompt = 0.001 + pricing.completion = 0.001 + pricing.max_prompt_cost = 0.0 # no prompt-side discount to keep math clean + pricing.max_completion_cost = 100.0 + + model_obj = Mock() + model_obj.sats_pricing = pricing + model_obj.top_provider = None + model_obj.context_length = None + + with ( + patch.object(settings, "fixed_pricing", False), + patch.object(settings, "tolerance_percentage", 0), + patch.object(settings, "min_request_msat", 1000), + ): + cost = await calculate_discounted_max_cost( + 100_000, {}, model_obj, max_tokens=80_000 + ) + + # full max_cost (100 sats = 100000 msats) minus completion discount + # (100 - 80 = 20 sats = 20000 msats) => 80000 msats + assert cost == 80_000 + + +async def test_discounted_max_cost_body_max_completion_tokens_fallback() -> None: + """Body ``max_completion_tokens`` is honoured when no override is set.""" + from routstr.payment.helpers import calculate_discounted_max_cost + + pricing = Mock() + pricing.prompt = 0.001 + pricing.completion = 0.001 + pricing.max_prompt_cost = 0.0 + pricing.max_completion_cost = 100.0 + + model_obj = Mock() + model_obj.sats_pricing = pricing + model_obj.top_provider = None + model_obj.context_length = None + + body = {"max_completion_tokens": 80_000} + + with ( + patch.object(settings, "fixed_pricing", False), + patch.object(settings, "tolerance_percentage", 0), + patch.object(settings, "min_request_msat", 1000), + ): + cost = await calculate_discounted_max_cost(100_000, body, model_obj) + + assert cost == 80_000 + + +async def test_discounted_max_cost_override_takes_precedence_over_body() -> None: + """An explicit override wins over any body max_tokens value.""" + from routstr.payment.helpers import calculate_discounted_max_cost + + pricing = Mock() + pricing.prompt = 0.001 + pricing.completion = 0.001 + pricing.max_prompt_cost = 0.0 + pricing.max_completion_cost = 100.0 + + model_obj = Mock() + model_obj.sats_pricing = pricing + model_obj.top_provider = None + model_obj.context_length = None + + # Body says 10 tokens (would floor to min_request_msat = 1000); the + # 80000-token override must win, so the result is 80000 msats. + body = {"max_tokens": 10} + + with ( + patch.object(settings, "fixed_pricing", False), + patch.object(settings, "tolerance_percentage", 0), + patch.object(settings, "min_request_msat", 1000), + ): + cost = await calculate_discounted_max_cost( + 100_000, body, model_obj, max_tokens=80_000 + ) + + assert cost == 80_000 diff --git a/tests/unit/test_tinfoil_integration.py b/tests/unit/test_tinfoil_integration.py index 9435a851..c4af5490 100644 --- a/tests/unit/test_tinfoil_integration.py +++ b/tests/unit/test_tinfoil_integration.py @@ -126,6 +126,7 @@ class TestStripProxyHeaders: def test_strips_all_proxy_only(self) -> None: headers = { "x-routstr-model": "tinfoil-llama3-3-70b", + "x-routstr-max-tokens": "1024", "X-Tinfoil-Enclave-Url": "https://inference.tinfoil.sh", "X-Tinfoil-Request-Usage-Metrics": "true", "Authorization": "Bearer secret", @@ -133,6 +134,7 @@ class TestStripProxyHeaders: } clean = _strip_proxy_headers(headers) assert "x-routstr-model" not in clean + assert "x-routstr-max-tokens" not in clean assert "X-Tinfoil-Enclave-Url" not in clean assert "X-Tinfoil-Request-Usage-Metrics" not in clean assert clean["Authorization"] == "Bearer secret" @@ -141,6 +143,7 @@ class TestStripProxyHeaders: def test_all_proxy_only_headers_covered(self) -> None: assert _PROXY_ONLY_HEADERS == { "x-routstr-model", + "x-routstr-max-tokens", "x-tinfoil-enclave-url", "x-tinfoil-request-usage-metrics", }