From f4c382e5bacdaa6151ad4d2e03f5a8563eb08b4b Mon Sep 17 00:00:00 2001 From: Fsocietyhhh <1211904451@qq.com> Date: Wed, 22 Jul 2026 16:56:47 +0800 Subject: [PATCH 1/2] feat: add native Gemini proxy passthrough --- CHANGELOG.md | 27 ++++ README.md | 57 ++++++++- blockrun_litellm/__init__.py | 2 +- blockrun_litellm/proxy.py | 73 ++++++++++- pyproject.toml | 2 +- tests/test_proxy_gemini_native.py | 201 ++++++++++++++++++++++++++++++ tests/test_stream_usage.py | 10 +- 7 files changed, 362 insertions(+), 10 deletions(-) create mode 100644 tests/test_proxy_gemini_native.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e320bf6..eaf0bbe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,32 @@ # Changelog +## 0.8.0 — 2026-07-22 + +### Added + +- **Gemini's native `v1beta` protocol now works through the local sidecar.** + `POST /v1beta/models/{model}:generateContent` and + `POST /v1beta/models/{model}:streamGenerateContent` are forwarded verbatim + to BlockRun while the sidecar handles Base or Solana x402 payment locally. + This lets a customer keep Gemini-native request/response bodies and SSE + frames instead of translating through OpenAI Chat Completions. + +- Streaming is selected from the Gemini URL method, matching Google's wire + protocol (there is no OpenAI-style `"stream": true` field to inspect). + Upstream errors keep their native HTTP status and Gemini error envelope. + +- Client Google credentials are never forwarded. Google SDKs may require a + dummy API key and attach it as `x-goog-api-key` or `?key=...`; the sidecar + removes both because BlockRun authenticates with the local x402 wallet and + the gateway owns the upstream Google credential. + +### Compatibility + +- Existing OpenAI, Anthropic, image, video, and audio routes are unchanged. +- Native Gemini is available in sidecar/proxy mode. The in-process LiteLLM + custom provider remains OpenAI-shaped because it returns LiteLLM response + objects rather than exposing an HTTP protocol surface. + ## 0.7.6 — 2026-07-16 ### Fixed diff --git a/README.md b/README.md index 121228c..c25de4b 100644 --- a/README.md +++ b/README.md @@ -361,6 +361,8 @@ curl http://localhost:4001/v1/chat/completions \ | Method | Path | Notes | |---|---|---| | `POST` | `/v1/chat/completions` | OpenAI Chat Completions. `stream=True` returns `text/event-stream`; otherwise JSON. | +| `POST` | `/v1beta/models/{model}:generateContent` | Native Gemini JSON request and response, with automatic x402 payment. | +| `POST` | `/v1beta/models/{model}:streamGenerateContent` | Native Gemini SSE, with automatic x402 payment. | | `POST` | `/v1/responses` | OpenAI Responses API, bridged onto Chat Completions (`input`→`messages`, `output`/`response.*` SSE out). Text-in/text-out; for advanced tool/state flows use `/v1/chat/completions`. | | `POST` | `/v1/images/generations` | OpenAI Image Generations. Accepts `prompt`, `model`, `size`, `n`, and `quality` (Solana only — see below). | | `POST` | `/v1/images/edits` | OpenAI-compatible image editing. Accepts JSON data URIs or multipart `image`/`image[]`; supports multiple source images, `mask`, and `quality` (Solana only). `/v1/images/image2image` is an alias. | @@ -375,6 +377,59 @@ curl http://localhost:4001/v1/chat/completions \ | `GET` | `/healthz` | Liveness probe (no upstream call) | | `GET` | `/docs` | Auto-generated Swagger UI | +### 2e. Native Gemini protocol + +Native Gemini calls use the sidecar root (`http://localhost:4001`), not the +OpenAI `/v1` base. The sidecar preserves Gemini request/response JSON and SSE +frames and adds the x402 payment signature using the configured wallet. + +```bash +# Non-streaming +curl http://localhost:4001/v1beta/models/gemini-2.5-flash:generateContent \ + -H "Content-Type: application/json" \ + -d '{"contents":[{"role":"user","parts":[{"text":"Reply with native-ok"}]}]}' + +# Streaming +curl -N 'http://localhost:4001/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse' \ + -H "Content-Type: application/json" \ + -d '{"contents":[{"role":"user","parts":[{"text":"Reply with stream-ok"}]}]}' +``` + +The official Google Gen AI Python SDK can use the same native protocol. Keep +the SDK itself; point its custom base URL at the sidecar and use any non-empty +placeholder API key (the sidecar removes it before forwarding): + +```python +from google import genai +from google.genai import types + +client = genai.Client( + api_key="unused", + http_options=types.HttpOptions(base_url="http://localhost:4001"), +) + +response = client.models.generate_content( + model="gemini-2.5-flash-lite", + contents="Reply with native-ok", +) +print(response.text) + +for chunk in client.models.generate_content_stream( + model="gemini-2.5-flash-lite", + contents="Reply with stream-ok", +): + print(chunk.text or "", end="") +``` + +For Solana, start the sidecar with +`BLOCKRUN_API_URL=https://sol.blockrun.ai/api` and a `SOLANA_WALLET_KEY` (or +the wallet already stored at `~/.blockrun/.solana-session`). A client-supplied +Google API key is ignored; the local wallet is the authentication and payment +mechanism. + +This is a sidecar HTTP feature. `litellm.completion(...)` remains the +OpenAI-compatible interface and continues to use `/v1/chat/completions`. + #### Image request limits `n` is bounded to **1–10** locally. The Base `image2image` gateway schema has no @@ -403,7 +458,7 @@ otherwise degrades silently to text-to-video and still bills you. Reference-to-video (`reference_videos`/`reference_audios`) is **not** supported: both gateways currently gate it off and would return 503. -### 2e. Image generation +### 2f. Image generation ```python from openai import OpenAI diff --git a/blockrun_litellm/__init__.py b/blockrun_litellm/__init__.py index 55e72d5..faf2a32 100644 --- a/blockrun_litellm/__init__.py +++ b/blockrun_litellm/__init__.py @@ -26,4 +26,4 @@ from blockrun_litellm.provider import BlockRunLLM, register __all__ = ["BlockRunLLM", "register", "enable_local_logging"] -__version__ = "0.7.6" +__version__ = "0.8.0" diff --git a/blockrun_litellm/proxy.py b/blockrun_litellm/proxy.py index 8c6e71c..820cb15 100644 --- a/blockrun_litellm/proxy.py +++ b/blockrun_litellm/proxy.py @@ -11,6 +11,8 @@ - ``POST /v1/messages`` — native Anthropic Messages (Claude Code, Anthropic SDK); verbatim x402-signed passthrough — tools/thinking preserved - ``POST /v1/messages/count_tokens`` — Anthropic token counting (passthrough) +- ``POST /v1beta/models/{model}:generateContent`` — native Gemini JSON +- ``POST /v1beta/models/{model}:streamGenerateContent`` — native Gemini SSE - ``POST /v1/images/generations`` — OpenAI Image Generations (DALL-E compatible) - ``POST /v1/videos/generations`` — video generation (xai/grok-imagine-video, Seedance); async submit+poll, settles only on completion @@ -386,6 +388,17 @@ def _openai_fwd_headers(request: Request) -> Dict[str, str]: return {"content-type": request.headers.get("content-type", "application/json")} +def _gemini_fwd_headers(request: Request) -> Dict[str, str]: + """Headers for native Gemini requests. + + Google SDKs require an ``api_key`` even when pointed at a custom base URL + and may send it as ``x-goog-api-key`` or ``?key=...``. The sidecar uses + the local x402 wallet instead, so neither credential is forwarded. Keep + only the content type; the BlockRun gateway supplies its own Google key. + """ + return {"content-type": request.headers.get("content-type", "application/json")} + + def _open_upstream_stream( client: httpx.Client, target: str, raw: bytes, headers: Dict[str, str] ) -> httpx.Response: @@ -511,7 +524,14 @@ def _body_model(raw: bytes) -> Optional[str]: async def _forward_passthrough( - request: Request, path: str, headers: Dict[str, str], *, allow_stream: bool + request: Request, + path: str, + headers: Dict[str, str], + *, + allow_stream: bool, + stream_override: Optional[bool] = None, + model_override: Optional[str] = None, + forward_query: bool = True, ) -> Response: """Verbatim x402-signed passthrough to a BlockRun native endpoint. @@ -528,18 +548,18 @@ async def _forward_passthrough( api_url = _resolve_api_url() raw = await request.body() client = _messages_client(api_url) - qs = request.url.query + qs = request.url.query if forward_query else "" target = f"{api_url}{path}" + (f"?{qs}" if qs else "") _t0 = time.monotonic() - model = _body_model(raw) + model = model_override or _body_model(raw) req_id = request.headers.get("x-request-id") def _latency_ms() -> float: return (time.monotonic() - _t0) * 1000.0 - wants_stream = False - if allow_stream: + wants_stream = bool(stream_override) if stream_override is not None else False + if allow_stream and stream_override is None: try: wants_stream = bool(_json.loads(raw or b"{}").get("stream")) except Exception: @@ -665,6 +685,49 @@ def _post(): ) +# --------------------------------------------------------------------------- +# Native Gemini /v1beta passthrough (Google GenAI SDK, raw REST clients, ...) +# --------------------------------------------------------------------------- +# Gemini selects streaming in the URL method, not with a JSON ``stream`` field. +# The request and response bodies stay byte-for-byte native; only x402 payment +# signing is added by the shared transport. A Google SDK's dummy ``key`` query +# and x-goog-api-key header are deliberately dropped before the gateway hop. + +_GEMINI_METHODS = {"generateContent", "streamGenerateContent"} + + +def _gemini_error(status: int, message: str, code: str = "INVALID_ARGUMENT") -> JSONResponse: + return JSONResponse( + status_code=status, + content={"error": {"code": status, "message": message, "status": code}}, + ) + + +@app.post("/v1beta/models/{model_and_method:path}", dependencies=[Depends(_require_token)]) +async def gemini_native(request: Request, model_and_method: str) -> Any: + raw_model, separator, method = model_and_method.rpartition(":") + if not separator or not raw_model.strip() or method not in _GEMINI_METHODS: + return _gemini_error( + 400, + "This endpoint only supports generateContent and streamGenerateContent.", + ) + + model = raw_model.strip().removeprefix("google/") + return await _forward_passthrough( + request, + f"/v1beta/models/{model_and_method}", + _gemini_fwd_headers(request), + allow_stream=True, + stream_override=method == "streamGenerateContent", + model_override=f"google/{model}", + # Google SDKs append ?key=... and ?alt=sse. The method already tells + # the gateway whether to stream, so forwarding either is unnecessary; + # dropping the query also prevents a real Google key from reaching + # gateway access logs when a caller accidentally configured one. + forward_query=False, + ) + + @app.post("/v1/messages", dependencies=[Depends(_require_token)]) async def anthropic_messages(request: Request) -> Any: return await _forward_passthrough( diff --git a/pyproject.toml b/pyproject.toml index bce67e4..c26fe03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "blockrun-litellm" -version = "0.7.6" +version = "0.8.0" description = "LiteLLM adapter for BlockRun — call x402-paid AI models via LiteLLM (custom provider or local OpenAI-compatible proxy)" readme = "README.md" license = "MIT" diff --git a/tests/test_proxy_gemini_native.py b/tests/test_proxy_gemini_native.py new file mode 100644 index 0000000..4d623fe --- /dev/null +++ b/tests/test_proxy_gemini_native.py @@ -0,0 +1,201 @@ +"""Native Gemini protocol coverage for the x402 sidecar. + +The sidecar must preserve Gemini JSON/SSE bytes while adding wallet payment. +Streaming is selected by ``:streamGenerateContent`` in the URL, not by a +``stream`` field in the request body. +""" + +from __future__ import annotations + +import contextlib + +import httpx +import pytest +from fastapi.testclient import TestClient + +import blockrun_litellm.proxy as proxy + + +class _BytesStream(httpx.SyncByteStream): + def __init__(self, data: bytes): + self._data = data + + def __iter__(self): + yield self._data + + def close(self) -> None: + pass + + +class _FakeClient: + def __init__(self, response: httpx.Response): + self._response = response + self.requests: list[httpx.Request] = [] + self.posts = 0 + self.stream_sends = 0 + + def build_request(self, method, url, content=None, headers=None): + request = httpx.Request(method, url, content=content, headers=headers) + self.requests.append(request) + return request + + def send(self, request, stream=False): + self.stream_sends += 1 + return self._response + + def post(self, url, content=None, headers=None): + self.posts += 1 + request = httpx.Request("POST", url, content=content, headers=headers) + self.requests.append(request) + return self._response + + +@pytest.fixture +def client() -> TestClient: + return TestClient(proxy.app) + + +@pytest.fixture +def count_semaphore(monkeypatch: pytest.MonkeyPatch): + calls = {"acquired": 0} + + @contextlib.asynccontextmanager + async def _spy(): + calls["acquired"] += 1 + yield + + monkeypatch.setattr(proxy, "_get_semaphore", _spy) + return calls + + +def _patch_client( + monkeypatch: pytest.MonkeyPatch, response: httpx.Response +) -> _FakeClient: + fake = _FakeClient(response) + monkeypatch.setattr(proxy, "_messages_client", lambda api_url: fake) + return fake + + +def test_generate_content_is_verbatim_and_drops_google_credentials( + monkeypatch, client, count_semaphore +): + upstream_body = b'{"candidates":[{"content":{"parts":[{"text":"native-ok"}]}}]}' + fake = _patch_client( + monkeypatch, + httpx.Response(200, headers={"content-type": "application/json"}, content=upstream_body), + ) + raw = b'{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}' + + response = client.post( + "/v1beta/models/gemini-2.5-flash:generateContent?key=client-secret", + content=raw, + headers={ + "content-type": "application/json", + "x-goog-api-key": "client-secret", + "authorization": "Bearer proxy-or-google-token", + }, + ) + + assert response.status_code == 200 + assert response.content == upstream_body + assert fake.posts == 1 and fake.stream_sends == 0 + assert count_semaphore["acquired"] == 1 + forwarded = fake.requests[0] + assert str(forwarded.url) == ( + "https://blockrun.ai/api/v1beta/models/gemini-2.5-flash:generateContent" + ) + assert forwarded.content == raw + assert "x-goog-api-key" not in forwarded.headers + assert "authorization" not in forwarded.headers + + +def test_stream_generate_content_uses_url_method_and_preserves_sse( + monkeypatch, client, count_semaphore +): + sse = ( + b'data: {"candidates":[{"content":{"parts":[{"text":"stream-ok"}]}}]}\n\n' + b'data: {"usageMetadata":{"promptTokenCount":3,"candidatesTokenCount":2}}\n\n' + ) + fake = _patch_client( + monkeypatch, + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=_BytesStream(sse), + ), + ) + + response = client.post( + "/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse&key=dummy", + json={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + ) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert response.content == sse + assert fake.stream_sends == 1 and fake.posts == 0 + assert count_semaphore["acquired"] == 1 + assert str(fake.requests[0].url) == ( + "https://blockrun.ai/api/v1beta/models/" + "gemini-2.5-flash:streamGenerateContent" + ) + + +def test_streaming_upstream_error_keeps_native_status(monkeypatch, client, count_semaphore): + error = b'{"error":{"code":429,"message":"quota","status":"RESOURCE_EXHAUSTED"}}' + fake = _patch_client( + monkeypatch, + httpx.Response(429, headers={"content-type": "application/json"}, content=error), + ) + + response = client.post( + "/v1beta/models/gemini-2.5-flash:streamGenerateContent", + json={"contents": []}, + ) + + assert response.status_code == 429 + assert response.content == error + assert fake.stream_sends == 1 + assert count_semaphore["acquired"] == 1 + + +def test_google_prefixed_model_is_forwarded_and_logged(monkeypatch, client, count_semaphore): + fake = _patch_client( + monkeypatch, + httpx.Response(200, headers={"content-type": "application/json"}, content=b"{}"), + ) + logged = [] + monkeypatch.setattr(proxy._logger, "log_proxy_call", lambda **kwargs: logged.append(kwargs)) + + response = client.post( + "/v1beta/models/google/gemini-3.1-pro:generateContent", + json={"contents": []}, + ) + + assert response.status_code == 200 + assert str(fake.requests[0].url).endswith( + "/v1beta/models/google/gemini-3.1-pro:generateContent" + ) + assert logged[0]["model"] == "google/gemini-3.1-pro" + assert logged[0]["path"] == ( + "/v1beta/models/google/gemini-3.1-pro:generateContent" + ) + + +@pytest.mark.parametrize( + "path", + [ + "/v1beta/models/gemini-2.5-flash:countTokens", + "/v1beta/models/gemini-2.5-flash", + ], +) +def test_unsupported_native_method_fails_before_payment(monkeypatch, client, path): + def _must_not_create_client(api_url): + raise AssertionError("wallet client must not be created") + + monkeypatch.setattr(proxy, "_messages_client", _must_not_create_client) + response = client.post(path, json={"contents": []}) + + assert response.status_code == 400 + assert response.json()["error"]["status"] == "INVALID_ARGUMENT" + diff --git a/tests/test_stream_usage.py b/tests/test_stream_usage.py index d11a1e3..0936f76 100644 --- a/tests/test_stream_usage.py +++ b/tests/test_stream_usage.py @@ -155,8 +155,9 @@ def test_usage_survives_custom_stream_wrapper() -> None: def test_usage_dropped_without_provider_specific_fields_key() -> None: """Documents WHY the key contract matters: strip the key from the usage - frame (simulating a 'simplified' choice-less return) and LiteLLM falls - back to tokenizer estimates instead of the gateway's real counts.""" + frame (simulating a 'simplified' choice-less return). Older LiteLLM falls + back to tokenizer estimates; newer releases accept the usage frame without + the key, in which case this historical compatibility guard is unnecessary.""" from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import custom_llm_setup @@ -212,6 +213,11 @@ def _stripped(c: ChatCompletionChunk): got = ( built.usage.prompt_tokens if built is not None and built.usage else None ) + if got == 100: + pytest.skip( + "This LiteLLM release accepts post-finish usage frames without " + "provider_specific_fields; the compatibility guard is no longer required." + ) assert got != 100, ( "LiteLLM consumed the usage frame even without the " "provider_specific_fields key — its post-finish guard changed, so " From fbcc6142708ddf4345a193eb55174c661b315ec6 Mon Sep 17 00:00:00 2001 From: 1bcMax Date: Wed, 22 Jul 2026 13:52:33 -0500 Subject: [PATCH 2/2] fix(gemini): harden native passthrough after review - sanitize the model id (strip models//google/ prefixes, reject slashes) and rebuild the forwarded path from clean parts, so no crafted model name can move the signed upstream target outside /v1beta/models/ (defense-in-depth; gateway already backstops, but the wallet shouldn't rely on it) - log every local 400 reject, honoring the 0.7.6 every-exit-logs invariant - document the Base enterprise/allowlist gate (403 for retail) in README + CHANGELOG, point retail callers to Solana or /v1/chat/completions - de-duplicate _gemini_fwd_headers, refresh _forward_passthrough docstring for the new kwargs, simplify the stream_override ternary - tests: isolate BLOCKRUN_API_URL, add slash-rejection, blank-model, dot-segment, and proxy-token denial coverage; assert the reject logs --- CHANGELOG.md | 14 +++++ README.md | 16 ++++++ blockrun_litellm/proxy.py | 91 +++++++++++++++++++++++++------ tests/test_proxy_gemini_native.py | 81 ++++++++++++++++++++++++--- 4 files changed, 177 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eaf0bbe..911923c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,20 @@ removes both because BlockRun authenticates with the local x402 wallet and the gateway owns the upstream Google credential. +- The model segment is sanitized before forwarding: `models/`/`google/` + prefixes are stripped and the path is rebuilt from the clean id, so a crafted + model name cannot move the signed upstream target outside `/v1beta/models/`. + Unsupported methods and invalid model ids are rejected with a `400` — and, per + the 0.7.6 every-exit-logs guarantee, still write an audit row. + +### Availability + +- Native Gemini passthrough is **enterprise/allowlist-only on the Base gateway** + (non-allowlisted payers receive `403 PERMISSION_DENIED` and are not charged); + it is generally available on Solana (`sol.blockrun.ai`). Retail callers on + Base should use `/v1/chat/completions` with `model="google/gemini-3.1-pro"`. + Only `generateContent`/`streamGenerateContent` are proxied. + ### Compatibility - Existing OpenAI, Anthropic, image, video, and audio routes are unchanged. diff --git a/README.md b/README.md index c25de4b..337c22b 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,22 @@ the wallet already stored at `~/.blockrun/.solana-session`). A client-supplied Google API key is ignored; the local wallet is the authentication and payment mechanism. +**Availability.** Native Gemini passthrough on the **Base** gateway (the +default, `https://blockrun.ai/api`) is limited to enterprise/allowlisted +wallets — other payers get `403 PERMISSION_DENIED` and are **not** charged. On +**Solana** (`https://sol.blockrun.ai/api`) it is generally available. If your +wallet is not on the Base allowlist, use Solana for native Gemini, or call the +OpenAI-format `/v1/chat/completions` route with `model="google/gemini-3.1-pro"`, +which works on both chains. Only `generateContent` and `streamGenerateContent` +are proxied; `countTokens`, model list/get, and `embedContent` are not served by +the gateway. Streaming always returns SSE regardless of `alt=` (the client query +string is dropped and the gateway sets `alt=sse` itself), so raw REST clients +must parse SSE, not the chunked-JSON-array form. + +If you set `BLOCKRUN_PROXY_TOKEN` to guard the sidecar, the Google SDK cannot +send it (it only sends `x-goog-api-key`), so pass it explicitly as a Bearer +header: `types.HttpOptions(base_url=..., headers={"Authorization": "Bearer "})`. + This is a sidecar HTTP feature. `litellm.completion(...)` remains the OpenAI-compatible interface and continues to use `/v1/chat/completions`. diff --git a/blockrun_litellm/proxy.py b/blockrun_litellm/proxy.py index 820cb15..298d0de 100644 --- a/blockrun_litellm/proxy.py +++ b/blockrun_litellm/proxy.py @@ -50,6 +50,7 @@ import logging import math import os +import re import threading import time import uuid @@ -395,8 +396,12 @@ def _gemini_fwd_headers(request: Request) -> Dict[str, str]: and may send it as ``x-goog-api-key`` or ``?key=...``. The sidecar uses the local x402 wallet instead, so neither credential is forwarded. Keep only the content type; the BlockRun gateway supplies its own Google key. + + Same shape as ``_openai_fwd_headers`` today; delegate so a future change to + the forwarded-header policy can't silently apply to one protocol and not + the other. """ - return {"content-type": request.headers.get("content-type", "application/json")} + return _openai_fwd_headers(request) def _open_upstream_stream( @@ -535,15 +540,28 @@ async def _forward_passthrough( ) -> Response: """Verbatim x402-signed passthrough to a BlockRun native endpoint. - Shared by ``/v1/messages``, ``/v1/messages/count_tokens`` and - ``/v1/chat/completions``. The body is forwarded byte-for-byte — no - translation, no SDK typed parsing — so OpenAI clients (incl. Codex with - wire_api=chat) and Claude Code stay off the SDK's ``chat_completion_stream`` - path, the source of the streamed-tool-call ``'dict' object has no attribute - 'delta'`` crash. The upstream call is gated by ``_get_semaphore()`` like - every other paid route. For streaming we open the upstream first to learn its - real status, then either return a real error ``Response`` (preserving the - 4xx/5xx the client must see) or stream the body. + Shared by ``/v1/messages``, ``/v1/messages/count_tokens``, + ``/v1/chat/completions`` and the native Gemini ``/v1beta/models/...`` route. + The body is forwarded byte-for-byte — no translation, no SDK typed parsing — + so OpenAI clients (incl. Codex with wire_api=chat) and Claude Code stay off + the SDK's ``chat_completion_stream`` path, the source of the + streamed-tool-call ``'dict' object has no attribute 'delta'`` crash. The + upstream call is gated by ``_get_semaphore()`` like every other paid route. + For streaming we open the upstream first to learn its real status, then + either return a real error ``Response`` (preserving the 4xx/5xx the client + must see) or stream the body. + + Keyword-only overrides let a caller adapt the shared path to a protocol whose + conventions differ from the OpenAI body: + + * ``stream_override`` — when not ``None``, decides streaming directly instead + of sniffing a ``"stream"`` field in the body (Gemini encodes streaming in + the URL method, not the body). ``allow_stream`` is then irrelevant. + * ``model_override`` — the model string used for the audit log, when it can't + be read from the body (e.g. Gemini carries it in the URL). + * ``forward_query`` — ``False`` drops the client query string before + forwarding (Gemini strips ``?key=``/``?alt=sse``; the gateway re-derives + both server-side). """ api_url = _resolve_api_url() raw = await request.body() @@ -558,7 +576,7 @@ async def _forward_passthrough( def _latency_ms() -> float: return (time.monotonic() - _t0) * 1000.0 - wants_stream = bool(stream_override) if stream_override is not None else False + wants_stream = bool(stream_override) if allow_stream and stream_override is None: try: wants_stream = bool(_json.loads(raw or b"{}").get("stream")) @@ -695,6 +713,13 @@ def _post(): _GEMINI_METHODS = {"generateContent", "streamGenerateContent"} +# A native Gemini model id, once the id-only ``models/``/``google/`` prefixes are +# stripped: letters, digits, dot, dash, underscore. Deliberately no slash and no +# dot-segment — those are how a crafted model name (``../../v1/chat/completions``) +# escapes the ``/v1beta/models/`` prefix after httpx normalizes the URL, so we +# reject them locally before the wallet ever signs a payment. +_GEMINI_MODEL_RE = re.compile(r"[A-Za-z0-9._-]+") + def _gemini_error(status: int, message: str, code: str = "INVALID_ARGUMENT") -> JSONResponse: return JSONResponse( @@ -703,27 +728,57 @@ def _gemini_error(status: int, message: str, code: str = "INVALID_ARGUMENT") -> ) +def _log_gemini_reject(path: str, req_id: Optional[str]) -> None: + """Log a local, pre-payment 400 the way the media routes do. + + The 0.7.6 invariant is that every exit logs. This reject never reaches the + gateway, so it's genuinely free (no ``settlement_status`` flag), but it still + gets a row — a call that leaves no trace is invisible to reconciliation. + """ + _logger.log_proxy_call( + model=None, + path=path, + stream=False, + http_status=400, + cost_usd=None, + settlement=None, + latency_ms=0.0, + request_id=req_id, + ) + + @app.post("/v1beta/models/{model_and_method:path}", dependencies=[Depends(_require_token)]) async def gemini_native(request: Request, model_and_method: str) -> Any: + req_id = request.headers.get("x-request-id") + log_path = f"/v1beta/models/{model_and_method}" raw_model, separator, method = model_and_method.rpartition(":") - if not separator or not raw_model.strip() or method not in _GEMINI_METHODS: + if not separator or method not in _GEMINI_METHODS: + _log_gemini_reject(log_path, req_id) return _gemini_error( 400, "This endpoint only supports generateContent and streamGenerateContent.", ) - model = raw_model.strip().removeprefix("google/") + # Strip the id-only prefixes the gateway also accepts, then require a clean + # model token. The forwarded path is rebuilt from the sanitized parts — never + # the raw segment — so no slash or ``..`` in the model can move the signed + # upstream target outside ``/v1beta/models/``. + model = raw_model.strip().removeprefix("models/").removeprefix("google/") + if not _GEMINI_MODEL_RE.fullmatch(model): + _log_gemini_reject(log_path, req_id) + return _gemini_error(400, f"Invalid Gemini model id: {raw_model.strip()!r}.") + return await _forward_passthrough( request, - f"/v1beta/models/{model_and_method}", + f"/v1beta/models/{model}:{method}", _gemini_fwd_headers(request), allow_stream=True, stream_override=method == "streamGenerateContent", model_override=f"google/{model}", - # Google SDKs append ?key=... and ?alt=sse. The method already tells - # the gateway whether to stream, so forwarding either is unnecessary; - # dropping the query also prevents a real Google key from reaching - # gateway access logs when a caller accidentally configured one. + # Google SDKs append ?key=... and ?alt=sse. The BlockRun gateway derives + # streaming from the URL method and re-adds ?alt=sse itself on the real + # Google call, so it ignores the client query entirely; dropping it here + # loses nothing and keeps a stray real Google key out of gateway logs. forward_query=False, ) diff --git a/tests/test_proxy_gemini_native.py b/tests/test_proxy_gemini_native.py index 4d623fe..c1d4dc5 100644 --- a/tests/test_proxy_gemini_native.py +++ b/tests/test_proxy_gemini_native.py @@ -51,7 +51,11 @@ def post(self, url, content=None, headers=None): @pytest.fixture -def client() -> TestClient: +def client(monkeypatch: pytest.MonkeyPatch) -> TestClient: + # _resolve_api_url() reads BLOCKRUN_API_URL first; a developer or CI job with + # it exported (the documented Solana workflow) would otherwise flip the + # absolute-URL assertions below. Pin the Base default for these tests. + monkeypatch.delenv("BLOCKRUN_API_URL", raising=False) return TestClient(proxy.app) @@ -159,7 +163,11 @@ def test_streaming_upstream_error_keeps_native_status(monkeypatch, client, count assert count_semaphore["acquired"] == 1 -def test_google_prefixed_model_is_forwarded_and_logged(monkeypatch, client, count_semaphore): +def test_google_prefixed_model_is_stripped_before_forwarding(monkeypatch, client, count_semaphore): + """A ``google/`` (or ``models/``) prefix is an id-only convenience the gateway + also strips. We rebuild the forwarded path from the sanitized model so the + upstream sees the bare id, and the audit log canonicalizes to ``google/``. + """ fake = _patch_client( monkeypatch, httpx.Response(200, headers={"content-type": "application/json"}, content=b"{}"), @@ -173,13 +181,12 @@ def test_google_prefixed_model_is_forwarded_and_logged(monkeypatch, client, coun ) assert response.status_code == 200 + # Path is rebuilt from the sanitized model — the ``google/`` prefix is gone. assert str(fake.requests[0].url).endswith( - "/v1beta/models/google/gemini-3.1-pro:generateContent" + "/v1beta/models/gemini-3.1-pro:generateContent" ) assert logged[0]["model"] == "google/gemini-3.1-pro" - assert logged[0]["path"] == ( - "/v1beta/models/google/gemini-3.1-pro:generateContent" - ) + assert logged[0]["path"] == "/v1beta/models/gemini-3.1-pro:generateContent" @pytest.mark.parametrize( @@ -187,15 +194,75 @@ def test_google_prefixed_model_is_forwarded_and_logged(monkeypatch, client, coun [ "/v1beta/models/gemini-2.5-flash:countTokens", "/v1beta/models/gemini-2.5-flash", + # Blank / whitespace-only model — the model token is empty after stripping. + "/v1beta/models/:generateContent", + "/v1beta/models/%20:streamGenerateContent", + # A slash in the model reaches the handler (the :path converter keeps it) + # and must be rejected before payment — this is the escape the sanitizer + # closes. httpx would otherwise normalize it out of /v1beta/models/. + "/v1beta/models/foo/bar:generateContent", + "/v1beta/models/google/gemini/escape:generateContent", ], ) -def test_unsupported_native_method_fails_before_payment(monkeypatch, client, path): +def test_unsupported_or_unsafe_native_method_fails_before_payment(monkeypatch, client, path): def _must_not_create_client(api_url): raise AssertionError("wallet client must not be created") monkeypatch.setattr(proxy, "_messages_client", _must_not_create_client) + logged = [] + monkeypatch.setattr(proxy._logger, "log_proxy_call", lambda **kwargs: logged.append(kwargs)) + response = client.post(path, json={"contents": []}) assert response.status_code == 400 assert response.json()["error"]["status"] == "INVALID_ARGUMENT" + # 0.7.6 invariant: every exit logs, even a local pre-payment reject. + assert logged and logged[0]["http_status"] == 400 + assert "settlement_status" not in logged[0], "local reject never reached the gateway" + + +@pytest.mark.parametrize( + "path", + [ + "/v1beta/models/../../../v1/chat/completions:generateContent", + "/v1beta/models/google/../../v1/videos/generations:generateContent", + ], +) +def test_dot_segment_paths_never_reach_the_wallet(monkeypatch, client, path): + """Literal ``..`` segments are normalized away by the HTTP layer before they + reach the handler (404), so they never touch the wallet. Belt-and-suspenders + alongside the in-handler slash rejection above.""" + monkeypatch.setattr( + proxy, + "_messages_client", + lambda api_url: (_ for _ in ()).throw(AssertionError("wallet client must not be created")), + ) + response = client.post(path, json={"contents": []}) + assert response.status_code in (400, 404) + + +def test_proxy_token_gate_denies_gemini_route(monkeypatch, client): + """With BLOCKRUN_PROXY_TOKEN set, the route requires a Bearer token. The Google + SDK's own credential (x-goog-api-key) must not satisfy the gate, and no wallet + client is created for a denied request.""" + monkeypatch.setenv("BLOCKRUN_PROXY_TOKEN", "sekrit") + monkeypatch.setattr( + proxy, + "_messages_client", + lambda api_url: (_ for _ in ()).throw(AssertionError("must not reach wallet")), + ) + + denied = client.post( + "/v1beta/models/gemini-2.5-flash:generateContent", + json={"contents": []}, + headers={"x-goog-api-key": "sekrit"}, + ) + assert denied.status_code == 401 + + wrong_bearer = client.post( + "/v1beta/models/gemini-2.5-flash:generateContent", + json={"contents": []}, + headers={"authorization": "Bearer nope"}, + ) + assert wrong_bearer.status_code == 401