Uh oh!
There was an error while loading. Please reload this page.
feat(provider-vertex): wire Gemini chat via pre-minted OAuth2 bearer (D5.2.a) - #321
Merged
Conversation
…(D5.2.a, #302 Phase E) Replaces the skeleton's `BridgeError::Config("not yet implemented")` stubs in `aisix-provider-vertex` with real Vertex AI dispatch for the `google` publisher (Gemini chat). Other Vertex publishers (`anthropic.*`, `meta.*`, `mistral.*`, `ai21.*`) surface a clear publisher-named "not yet implemented — D5.3/D5.4" error. `chat_stream()` returns "streaming not yet implemented — D5.2.b" for all publishers. ## Wire shape pinned (Gemini on Vertex) Per https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini: - URL: `POST https://<region>-aiplatform.googleapis.com/v1/projects/<project>/locations/<region>/publishers/google/models/<model>:generateContent` - Body: Gemini's `generateContent` JSON shape: 1. `contents[].role` is `"user"` or `"model"` (NOT `"assistant"`) 2. `contents[].parts[].text` (not `content`) 3. `systemInstruction` is top-level — system messages do NOT appear in `contents[]` (Gemini 400s otherwise) 4. `generationConfig.temperature` / `topP` / `maxOutputTokens` (camelCase, emitted only when set) - Auth: `Authorization: Bearer <access_token>` where `access_token` comes from the operator-supplied pre-minted GCP OAuth2 token ## Credentials convention `ProviderKey.secret` is a JSON-encoded `{access_token, project, region}` blob. The `access_token` is operator-managed — they refresh it before the GCP TTL (~1 hour) and re-PUT the ProviderKey. D5.1 follow-up will add in-process token mint via `yup-oauth2` / `gcp_auth` from a `service_account_json` field. `project`, `region`, and `upstream_id` (model name from `ctx.model.model_name`) are validated for URL-control chars before being interpolated into the path: `/`, `?`, `#`, whitespace, `..` all rejected at registration time so a malicious model_name can't redirect dispatch or corrupt metrics labels. ## Audit lessons applied proactively This PR was written after the D6 audit (#319) and D7.2.a audit (#320) returned their reports. The lessons baked in from the start: - `bridge.chat()` is exercised end-to-end via wiremock + a `#[cfg(test)] api_base_override` seam — credentials, URL stitching, body shaping all run normally; only the destination host is different. - 4xx error body redacted to canned status-keyed phrases (does NOT echo Vertex's `Permission denied on project my-proj-prod-123` envelope into customer-visible errors). - `VertexSecret::parse` error message uses generic JSON-shape hint; does NOT echo raw secret bytes (serde error messages can leak partial content). Pinned by `vertex_secret_error_does_not_leak_secret_content`. - `BridgeContext.deadline` threaded into `chat()` via `with_deadline`. - Header-injection guard on `access_token` value (`HeaderValue::from_str` rejects newlines / NULs); same for `request_id`. ## Test coverage 38 unit tests, all passing: ### Publisher resolution (8 tests, preserved from skeleton) gemini-*, claude-*, meta/llama-*, llama*, mistral-*, codestral-*, jamba-*; case-insensitive on model name; URL segment for each publisher; `Meta::url_segment() == None` (Llama uses OpenAPI shim). ### `VertexSecret` parsing (4 tests) Full form, empty rejected, non-JSON rejected with generic shape error, raw secret bytes do NOT leak into error messages. ### URL token validation (3 tests) Canonical IDs pass (`my-proj-prod-123`, `us-central1`, `europe-west4`, `gemini-1.5-pro`); URL-injection chars (`/`, `?`, `#`, `\n`, `..`) rejected; empty rejected. ### Gemini request body translation (5 tests) User turn → role=user; assistant turn → role=model (NOT "assistant"); system messages → top-level `systemInstruction` (NOT in contents[]); multiple system messages concatenated with `\n\n`; `generationConfig` only emitted when at least one field is set. ### Gemini response translation (4 tests) Text + STOP finishReason → ChatResponse + FinishReason::Stop; MAX_TOKENS → FinishReason::Length; safety reasons (SAFETY / RECITATION / BLOCKLIST / PROHIBITED_CONTENT / SPII) → ContentFilter; missing usageMetadata → 0 tokens. ### Pre-dispatch validation (6 tests) Unknown publisher, non-Google publisher named in error, invalid secret, missing model_name, `chat_ignores_req_model_and_uses_ctx_model_name` (D6 audit HIGH-1 regression carried over), `chat_stream` not-implemented error references D5.2.b. ### Bridge dispatch via `bridge.chat()` end-to-end against wiremock (8 tests) - URL path exactly `/v1/projects/my-proj/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent` - Body uses Gemini wire shape (camelCase, no `model` field, no `stream`) - System role lifted to top-level `systemInstruction` - Assistant role uses `model`, NOT `assistant` - Authorization header carries `Bearer <access_token>` verbatim - 4xx upstream error body redacted (does NOT echo project id) - Path injection in project field rejected before URL stitching - MAX_TOKENS finishReason maps to FinishReason::Length ## References - Vertex AI REST API — https://cloud.google.com/vertex-ai/docs/reference/rest - Gemini generateContent — https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini - Gemini finishReason enum — https://ai.google.dev/api/generate-content#FinishReason - Google Gemini Python SDK — https://github.com/google-gemini/generative-ai-python
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThis PR transforms the Vertex provider bridge from a skeleton into a fully functional Gemini chat implementation, including HTTP client integration, publisher dispatch, Gemini request/response translation, comprehensive security validation, and integration tests with mock HTTP endpoints. ChangesVertex Gemini Bridge
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Note 🎁 Summarized by CodeRabbit FreeYour organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above. Comment |
… + LOW-4) PR #321 audit surfaced 1 MEDIUM + 4 LOW; this commit addresses the MEDIUM + 2 of the 4 LOW items that have concrete code fixes. LOW-1 (generationConfig missing topK/stopSequences) and LOW-2 (usageMetadata missing cachedContentTokenCount/thoughtsTokenCount) deferred as scope-extension follow-ups. ## MEDIUM (fixed) MEDIUM-1 — `BridgeError::Config` no longer formats the underlying `InvalidHeaderValue` Display into the customer-visible message. The `http` crate's current Display impl is opaque, but it's an implementation detail; a future change including the offending byte position would leak partial bearer-token content. Since the bytes being validated ARE the customer's bearer token, defense in depth wins. Pinned by two new tests (`header_invalid_access_token_error_does_not_leak_bytes` + `header_invalid_request_id_error_does_not_leak_bytes`). ## LOW (fixed) LOW-3 — `finishReason` IMAGE_SAFETY and LANGUAGE now map to `FinishReason::ContentFilter` (previously fell through to Stop, misleading tracing — a customer's dashboard would show a successful "stop" when Google in fact filtered the response). Added to the existing safety-finish-reason loop test. LOW-4 — `chat_gemini` now fails fast with a clear Config error when `contents[]` ends up empty (system-only messages, since system lifts to top-level `systemInstruction`). Gemini's schema requires `contents` non-empty; without this guard, Vertex 400s with a generic envelope and the customer sees the canned "upstream returned 400" phrase. Pinned by `chat_gemini_with_system_only_messages_fails_fast` using `Mock::expect(0)` to catch a regression that leaks the request to dispatch. ## Deferred (justified) LOW-1 — `generationConfig` missing `topK` / `stopSequences` / `responseMimeType` / `candidateCount` / `seed` / `thinkingConfig`. Out of scope for the minimal D5.2.a; will land alongside the `ChatFormat.extra` pass-through follow-up that lifts these fields uniformly across all bridges. LOW-2 — `usageMetadata` missing `cachedContentTokenCount` and `thoughtsTokenCount`. Needs `UsageStats` schema additions (cross-bridge change) — will land alongside D5.5 metrics work. ## Verification cargo test -p aisix-provider-vertex → 41 passed (was 38; +3 audit regression tests) cargo clippy --workspace --all-targets -- -D warnings → clean cargo fmt --check → clean
moonming
commented
May 17, 2026
MemberAuthor
Audit follow-up pushed (67bad57): Audit returned 0 HIGH + 1 MEDIUM + 4 LOW. This commit addresses MEDIUM-1 + 2 of the 4 LOW items. MEDIUM (fixed)
LOW (fixed)
Deferred (justified)
Verification |
Uh oh!
There was an error while loading. Please reload this page.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces the skeleton's "not yet implemented" stubs in
aisix-provider-vertexwith real Vertex AI dispatch for thegooglepublisher (Gemini chat). Other Vertex publishers + streaming surface clear publisher-named follow-up errors.Scope: deliberately minimal —
chat()forgoogleonly, with credentials supplied as a pre-minted OAuth2 access token (operator manages refresh). Mirrors the D6 #319 / D7.2.a #320 pattern.Wire shape (Gemini on Vertex)
Per https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/gemini:
POST https://<region>-aiplatform.googleapis.com/v1/projects/<project>/locations/<region>/publishers/google/models/<model>:generateContentgenerateContentJSON:contents[].roleis\"user\"or\"model\"(NOT\"assistant\")contents[].parts[].text(notcontent)systemInstructionis top-level — system messages do NOT appear incontents[]generationConfig.temperature/topP/maxOutputTokens(camelCase, emitted only when set)Authorization: Bearer <access_token>(pre-minted OAuth2 token from operator)Credentials convention
ProviderKey.secretis JSON-encoded{access_token, project, region}. The operator manages token refresh (GCP TTL ~1 hour). D5.1 follow-up adds in-process token mint viaservice_account_json(e.g.yup-oauth2/gcp_auth).project,region, andupstream_idare validated for URL-control chars (/,?,#,\\n,..) before being interpolated into the path — defense in depth against malicious model_name redirecting dispatch or corrupting metrics labels.Audit lessons applied proactively
Written after the D6 audit (#319) and D7.2.a audit (#320) results came back. The lessons baked in from the start:
bridge.chat()exercised end-to-end via wiremock + a#[cfg(test)] api_base_overrideseamVertexSecret::parseerror uses generic JSON-shape hint; does NOT echo raw secret bytesBridgeContext.deadlinethreaded intochat()viawith_deadlineaccess_tokenvalue (HeaderValue::from_strrejects newlines)References
Test plan
38 unit tests, all passing:
Publisher resolution (8 tests, preserved from skeleton)
gemini-, claude-, meta/llama-, llama, mistral-, codestral-, jamba-*; case-insensitive; URL segment for each;
Meta::url_segment() == None.VertexSecretparsing (4 tests)Full form, empty rejected, non-JSON rejected, raw secret bytes do NOT leak into error messages (pinned by
vertex_secret_error_does_not_leak_secret_content).URL token validation (3 tests)
Canonical IDs pass; URL-injection chars (
/,?,#,\\n,..) rejected; empty rejected.Gemini request body translation (5 tests)
User turn → role=user; assistant → role=model (NOT "assistant"); system → top-level
systemInstruction(NOT in contents[]); multiple system messages concatenated;generationConfigonly emitted when set.Gemini response translation (4 tests)
STOP → Stop; MAX_TOKENS → Length; safety reasons (SAFETY / RECITATION / BLOCKLIST / PROHIBITED_CONTENT / SPII) → ContentFilter; missing usageMetadata.
Pre-dispatch validation (6 tests)
Unknown publisher, non-Google publisher named in error, invalid secret, missing model_name,
chat_ignores_req_model_and_uses_ctx_model_name,chat_streamreferences D5.2.b.Bridge dispatch via
bridge.chat()end-to-end against wiremock (8 tests)/v1/projects/my-proj/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContentmodel/stream)systemInstructionmodel, NOTassistantBearer <access_token>verbatimTest plan TODO
cargo test -p aisix-provider-vertex→ 38/38 passcargo clippy --workspace --all-targets -- -D warningscleancargo fmt --checkcleancargo test --workspacecargo clippycargo fmt --checkSummary by CodeRabbit
New Features
Tests
Documentation