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

@@ -43,3 +43,6 @@ proof_backups
 *.todo
 ui_out
 .worktrees
+
+# macOS
+.DS_Store

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)

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

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",
     }

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
 

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
 # ===========================================================================

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

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",
         }