From 892aed61cc93c1c16b1a717d85f944f809c3e027 Mon Sep 17 00:00:00 2001 From: redshift <213178690+1ftredsh@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:01:13 +0800 Subject: [PATCH] Fix mypy: move type: ignore onto ASGITransport line --- docs/ehbp-proxy-support.md | 64 ++++++++----------- docs/tinfoil-direct-integration.md | 10 +-- routstr/proxy.py | 25 +++++--- routstr/upstream/tinfoil.py | 2 +- routstr/upstream/tinfoil_trailer.py | 30 ++++++++- .../test_proxy_tinfoil_attestation_routing.py | 54 ++++++++++++++-- tests/unit/test_tinfoil_trailer.py | 48 ++++++++++++++ 7 files changed, 176 insertions(+), 57 deletions(-) diff --git a/docs/ehbp-proxy-support.md b/docs/ehbp-proxy-support.md index 05b5aad3..f9edc6e0 100644 --- a/docs/ehbp-proxy-support.md +++ b/docs/ehbp-proxy-support.md @@ -53,25 +53,22 @@ The actual EHBP forwarding logic does **not** live in `base.py`. ### `routstr/upstream/ehbp.py` -Contains the shared opaque EHBP transport helpers: +Contains the shared opaque EHBP transport and billing helpers: - `EHBPForwardingTarget` — provider-specific target URL plus extra headers -- `forward_ehbp_request()` — forwards the raw encrypted body to an EHBP-capable - provider, streams the encrypted response back untouched, and finalizes bearer - billing at max cost because usage is encrypted -- `forward_ehbp_x_cashu_request()` — redeems the Cashu token, forwards raw, - refunds the full token on upstream failure, and refunds any value above - `max_cost_for_model` on success +- `forward_ehbp_request()` — forwards the encrypted body, captures Tinfoil + usage from a response header or streaming HTTP trailer, and finalizes bearer + billing at actual cost (falling back to max cost when usage is unavailable) +- `forward_ehbp_x_cashu_request()` — redeems the Cashu token, refunds the full + token on upstream failure, and refunds the difference between the redeemed + amount and actual cost (or max cost when usage is unavailable) -### `routstr/upstream/ppqai.py` +### Provider support -- Sets `supports_ehbp = True`. -- Implements `get_ehbp_forwarding_target()` to forward to - `https://api.ppq.ai/private/v1/...` — the PPQ.AI enclave endpoint that - understands EHBP and returns the `Ehbp-Response-Nonce` header. -- Adds `X-Private-Model` with the model's `forwarded_model_id` (e.g. - `private/kimi-k2-6`). PPQ.AI's billing layer needs this since it can't - decrypt the body. +EHBP is currently enabled only for `TinfoilUpstreamProvider`. It forwards to +Tinfoil's attested enclave and requests `X-Tinfoil-Usage-Metrics` for billing. +PPQ.AI retains its private-target implementation, but `supports_ehbp = False` +until it has a provider-specific trusted usage/model-binding strategy. ## Why it's done this way @@ -84,10 +81,10 @@ The proxy is a **blind relay** for EHBP requests. It cannot decrypt the body 4. Pass through EHBP protocol headers (`Ehbp-Encapsulated-Key` on request, `Ehbp-Response-Nonce` on response) -Cost tracking happens at the proxy level using `max_cost_for_model` from the -model registry. Because EHBP responses are encrypted, Routstr cannot reconcile -against token usage. Bearer requests reserve and then finalize max-cost billing; -X-Cashu requests redeem the token and refund any amount above max cost. +Cost tracking happens at the proxy level. Routstr reserves or redeems up to +`max_cost_for_model`, then Tinfoil's out-of-band usage header/trailer allows it +to finalize at actual token cost. If trusted usage is missing or invalid, the +proxy safely falls back to max-cost billing. ## End-to-end flow @@ -125,7 +122,7 @@ Three parties see three different model IDs: | Party | Header/Body | Value | Source | |---|---|---|---| | Routstr proxy | `X-Routstr-Model` header | `tinfoil-kimi-k2-6` | SDK sends full caller-facing id | -| PPQ.AI billing | `X-Private-Model` header | `private/kimi-k2-6` | Proxy sends `forwarded_model_id` | +| Tinfoil usage metrics | `model` field | `kimi-k2-6` | Enclave reports the model actually served | | Tinfoil enclave | `body.model` (encrypted) | `kimi-k2-6` | SDK strips `tinfoil-` prefix before encryption | ## Implementation status @@ -135,24 +132,19 @@ implements the direct blind-upstream pattern described above. The shared EHBP helpers in `routstr/upstream/ehbp.py` were extended to: - Request usage metrics via `X-Tinfoil-Request-Usage-Metrics: true`. -- Parse `X-Tinfoil-Usage-Metrics` from the response header (non-streaming). -- Override the forwarding URL with `X-Tinfoil-Enclave-Url` when the SDK sends it. -- Finalize bearer billing with actual token cost via `adjust_payment_for_tokens`. +- Parse `X-Tinfoil-Usage-Metrics` from the response header (non-streaming) or + HTTP trailer (streaming). +- Override the forwarding URL with a validated `X-Tinfoil-Enclave-Url` when the + SDK sends it. +- Finalize bearer billing with the dedicated EHBP actual-cost finalizer. - Compute X-Cashu refunds from actual cost instead of max cost. See `docs/tinfoil-direct-integration.md` for the full implementation notes. -## Not yet tested +## Verification status -These changes were written without integration testing due to the complexity -of the full stack (SDK + proxy + PPQ.AI enclave + Cashu mint). Needs end-to-end -verification with a real `tinfoil-*` model request. - -Important assumptions to verify: - -- PPQ.AI accepts `/private/v1/...` with `X-Private-Model`. -- PPQ.AI enforces consistency between `X-Private-Model` and the encrypted - `body.model`, otherwise a malicious client could understate - `X-Routstr-Model` for billing. -- SDK behavior on non-2xx proxy-generated errors that do not carry - `Ehbp-Response-Nonce`. +Unit coverage includes usage parsing, target validation, HTTP trailer capture, +response-size limits, and bearer payment finalization. End-to-end requests have +verified both non-streaming usage headers and streaming usage trailers against +Tinfoil. SDK behavior on proxy-generated non-2xx responses without an +`Ehbp-Response-Nonce` still merits explicit end-to-end coverage. diff --git a/docs/tinfoil-direct-integration.md b/docs/tinfoil-direct-integration.md index 80d104c1..c647699a 100644 --- a/docs/tinfoil-direct-integration.md +++ b/docs/tinfoil-direct-integration.md @@ -485,9 +485,9 @@ back to the requested model's pricing. - ~~End-to-end test with a real Tinfoil SDK client against a Routstr node with `TINFOIL_API_KEY` set.~~ Verified: both non-streaming (header) and streaming (trailer) responses include `model=`. -- Streaming requests: usage is delivered as an HTTP trailer. Currently the - bearer path finalizes max-cost before streaming begins. Supporting streaming - usage would require buffering the response (for X-Cashu) or a deferred - finalization (for bearer). +- Streaming trailer capture is implemented by buffering the encrypted response + in `forward_with_trailer()` and then using the dedicated EHBP payment + finalizers for bearer and X-Cashu requests. This provides actual-cost billing + today, at the cost of full time-to-last-byte latency for streaming responses. - Whether Tinfoil's `/v1/responses` endpoint also returns usage metrics - headers. + headers or trailers. diff --git a/routstr/proxy.py b/routstr/proxy.py index fc67797b..a5534e00 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -106,8 +106,13 @@ def get_unique_models() -> list[Model]: def _is_tinfoil_attestation_path(path: str) -> bool: - """Return True for Tinfoil attestation-bundle proxy paths.""" - return path in {"attestation", "tee/attestation"} + """Return True for exact Tinfoil attestation routes, with optional slash.""" + return path in { + "attestation", + "attestation/", + "tee/attestation", + "tee/attestation/", + } def _select_unauthenticated_get_upstreams( @@ -234,13 +239,10 @@ async def proxy( else: model_id = request_body_dict.get("model", "unknown") - # /tee/* and /attestation GET requests don't map to models — forward - # without model/cost/auth lookups. Tinfoil attestation paths are routed - # only to Tinfoil providers so an unrelated upstream's 404 cannot - # short-circuit before the attestation proxy is tried. - if request.method == "GET" and ( - path.startswith("tee/") or path.startswith("attestation") - ): + # Exact Tinfoil attestation GET routes don't map to models — forward + # without model/cost/auth lookups. Do not prefix-match here: paths such as + # /attestationjunk must continue through normal authentication. + if request.method == "GET" and _is_tinfoil_attestation_path(path): selected_upstreams = _select_unauthenticated_get_upstreams(path, _upstreams) if not selected_upstreams: return create_error_response( @@ -255,7 +257,10 @@ async def proxy( 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(selected_upstreams) - 1: + if ( + response.status_code in [502, 429] + and i < len(selected_upstreams) - 1 + ): logger.warning( "Upstream %s returned %s for unauthenticated GET %s, trying next", upstream.provider_type, diff --git a/routstr/upstream/tinfoil.py b/routstr/upstream/tinfoil.py index 795ae904..6eeb1147 100644 --- a/routstr/upstream/tinfoil.py +++ b/routstr/upstream/tinfoil.py @@ -116,7 +116,7 @@ class TinfoilUpstreamProvider(BaseUpstreamProvider): EHBP-only header used for encrypted POST requests and is not honored for unencrypted GET requests. """ - clean_path = path.removeprefix("tee/") + clean_path = path.removeprefix("tee/").rstrip("/") if clean_path == "attestation": return await self._proxy_attestation(headers) return await super().forward_get_request(request, path, headers) diff --git a/routstr/upstream/tinfoil_trailer.py b/routstr/upstream/tinfoil_trailer.py index abffd6ba..0357864f 100644 --- a/routstr/upstream/tinfoil_trailer.py +++ b/routstr/upstream/tinfoil_trailer.py @@ -27,6 +27,16 @@ _READ_BUFSIZE = 65536 _DEFAULT_TIMEOUT_SECONDS = 30.0 _DEFAULT_CLOSE_TIMEOUT_SECONDS = 1.0 _DEFAULT_MAX_RESPONSE_BYTES = 25 * 1024 * 1024 +_HOP_BY_HOP_HEADERS = { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} @dataclass @@ -47,6 +57,19 @@ def _get_header(headers: list[tuple[str, str]], name: str) -> str | None: return None +def _strip_hop_by_hop_headers(headers: dict[str, str]) -> dict[str, str]: + """Remove connection-specific headers before serializing a new request.""" + connection_tokens: set[str] = set() + for key, value in headers.items(): + if key.lower() == "connection": + connection_tokens.update( + token.strip().lower() for token in value.split(",") if token.strip() + ) + + excluded = _HOP_BY_HOP_HEADERS | connection_tokens + return {key: value for key, value in headers.items() if key.lower() not in excluded} + + async def forward_with_trailer( *, method: str, @@ -71,6 +94,11 @@ async def forward_with_trailer( if parsed.query: path = f"{path}?{parsed.query}" + # FastAPI has already decoded the incoming request body. Do not carry the + # original connection's framing or other hop-by-hop metadata into the new + # upstream connection. + headers = _strip_hop_by_hop_headers(headers) + ssl_ctx = ssl.create_default_context() reader, writer = await asyncio.wait_for( asyncio.open_connection(host, port, ssl=ssl_ctx), @@ -86,7 +114,7 @@ async def forward_with_trailer( header_lines.append("Connection: close") for key, value in headers.items(): - if key.lower() in ("host", "connection"): + if key.lower() == "host": continue header_lines.append(f"{key}: {value}") diff --git a/tests/unit/test_proxy_tinfoil_attestation_routing.py b/tests/unit/test_proxy_tinfoil_attestation_routing.py index 35b618ca..b6ba2f88 100644 --- a/tests/unit/test_proxy_tinfoil_attestation_routing.py +++ b/tests/unit/test_proxy_tinfoil_attestation_routing.py @@ -38,7 +38,8 @@ async def test_attestation_get_routes_directly_to_tinfoil_provider( monkeypatch.setattr(proxy_module, "_upstreams", [non_tinfoil, tinfoil]) async with AsyncClient( - transport=ASGITransport(app=proxy_app), base_url="http://test" # type: ignore[arg-type] + transport=ASGITransport(app=proxy_app), # type: ignore[arg-type] + base_url="http://test", ) as client: response = await client.get("/attestation") @@ -69,7 +70,8 @@ async def test_tee_attestation_get_routes_directly_to_tinfoil_provider( monkeypatch.setattr(proxy_module, "_upstreams", [non_tinfoil, tinfoil]) async with AsyncClient( - transport=ASGITransport(app=proxy_app), base_url="http://test" # type: ignore[arg-type] + transport=ASGITransport(app=proxy_app), # type: ignore[arg-type] + base_url="http://test", ) as client: response = await client.get("/tee/attestation") @@ -79,6 +81,48 @@ async def test_tee_attestation_get_routes_directly_to_tinfoil_provider( tinfoil.forward_get_request.assert_awaited_once() +@pytest.mark.parametrize("path", ["attestation/", "tee/attestation/"]) +@pytest.mark.asyncio +async def test_attestation_trailing_slash_routes_directly_to_tinfoil( + monkeypatch: pytest.MonkeyPatch, proxy_app: FastAPI, path: str +) -> None: + tinfoil = MagicMock() + tinfoil.provider_type = "tinfoil" + tinfoil.prepare_headers = MagicMock(return_value={}) + tinfoil.forward_get_request = AsyncMock(return_value=Response(status_code=200)) + monkeypatch.setattr(proxy_module, "_upstreams", [tinfoil]) + + async with AsyncClient( + transport=ASGITransport(app=proxy_app), # type: ignore[arg-type] + base_url="http://test", + ) as client: + response = await client.get(f"/{path}") + + assert response.status_code == 200 + tinfoil.forward_get_request.assert_awaited_once() + + +@pytest.mark.parametrize("path", ["attestation/foo", "attestationjunk"]) +@pytest.mark.asyncio +async def test_non_attestation_prefix_does_not_bypass_authentication( + monkeypatch: pytest.MonkeyPatch, proxy_app: FastAPI, path: str +) -> None: + tinfoil = MagicMock() + tinfoil.provider_type = "tinfoil" + tinfoil.forward_get_request = AsyncMock() + monkeypatch.setattr(proxy_module, "_upstreams", [tinfoil]) + + async with AsyncClient( + transport=ASGITransport(app=proxy_app), # type: ignore[arg-type] + base_url="http://test", + ) as client: + response = await client.get(f"/{path}") + + assert response.status_code == 400 + assert response.json()["error"]["type"] == "invalid_model" + tinfoil.forward_get_request.assert_not_awaited() + + def test_attestation_upstream_selection_is_tinfoil_only() -> None: non_tinfoil = MagicMock(provider_type="openai") tinfoil = MagicMock(provider_type="tinfoil") @@ -90,6 +134,8 @@ def test_attestation_upstream_selection_is_tinfoil_only() -> None: "tee/attestation", [non_tinfoil, tinfoil] ) == [tinfoil] assert proxy_module._select_unauthenticated_get_upstreams( - "tee/other", [non_tinfoil, tinfoil] + "attestation/", [non_tinfoil, tinfoil] + ) == [tinfoil] + assert proxy_module._select_unauthenticated_get_upstreams( + "attestationjunk", [non_tinfoil, tinfoil] ) == [non_tinfoil, tinfoil] - diff --git a/tests/unit/test_tinfoil_trailer.py b/tests/unit/test_tinfoil_trailer.py index d8319375..ef1c96f1 100644 --- a/tests/unit/test_tinfoil_trailer.py +++ b/tests/unit/test_tinfoil_trailer.py @@ -66,6 +66,54 @@ async def test_forward_with_trailer_captures_usage_trailer( writer.wait_closed.assert_awaited_once() +@pytest.mark.asyncio +async def test_forward_with_trailer_strips_hop_by_hop_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok" + reader = FakeReader([response]) + writer = FakeWriter() + monkeypatch.setattr( + "routstr.upstream.tinfoil_trailer.asyncio.open_connection", + AsyncMock(return_value=(reader, writer)), + ) + + await forward_with_trailer( + method="POST", + url="https://enclave.tinfoil.sh/v1/chat/completions", + headers={ + "Authorization": "Bearer upstream", + "Connection": "keep-alive, X-Client-Hop", + "Keep-Alive": "timeout=5", + "Proxy-Authenticate": "Basic", + "Proxy-Authorization": "Basic secret", + "TE": "trailers", + "Trailer": "X-Usage", + "Transfer-Encoding": "chunked", + "Upgrade": "websocket", + "X-Client-Hop": "remove-me", + "X-End-To-End": "preserve-me", + }, + body=b"opaque", + ) + + serialized_headers = writer.written.split(b"\r\n\r\n", 1)[0].lower() + for name in ( + b"keep-alive", + b"proxy-authenticate", + b"proxy-authorization", + b"te:", + b"trailer:", + b"transfer-encoding", + b"upgrade:", + b"x-client-hop", + ): + assert name not in serialized_headers + assert b"connection: close" in serialized_headers + assert b"content-length: 6" in serialized_headers + assert b"x-end-to-end: preserve-me" in serialized_headers + + @pytest.mark.asyncio async def test_forward_with_trailer_enforces_response_size_limit( monkeypatch: pytest.MonkeyPatch,