feat(vertex): in-process SA JSON → JWT → OAuth + token cache (#302 Phase E D5.1) - #387
Merged
Conversation
…ase E D5.1) Closes the second half of #302 Phase E productionization (Step 2a streaming merged as ai-gateway#386). Before this PR, `ProviderKey .secret` required a pre-minted GCP OAuth2 `access_token` — operator was responsible for refresh (GCP token TTL ~1h), which is awkward for production-scale deployments with many provider keys. This PR extends `VertexSecret` to accept EITHER mode, with mutual- exclusion validation at parse time: - **Pre-minted token** (backward-compat): set `access_token`. The bridge uses it verbatim. Existing deployments are unaffected. - **In-process SA mint** (new): set `service_account_json` to the full GCP SA JSON key. The bridge signs a JWT with the SA's RSA `private_key`, exchanges it for an OAuth2 access token via the SA's `token_uri` (JWT-bearer assertion grant per RFC 7523), and caches it in-process keyed by SA `client_email` with TTL refresh ~60s before the upstream-reported expiry. ## New module: `token_mint` `token_mint::TokenMinter` owns the cache + mint pipeline: - `ServiceAccountKey` — minimum SA JSON fields needed for minting (type, private_key, client_email, token_uri). Other on-disk SA fields (project_id, private_key_id, client_id, auth_uri, *_x509_cert_url) deserialize-skipped; operators can paste the whole SA JSON verbatim. - `validate()` — cheap shape checks (type == "service_account", private_key starts with -----BEGIN, client_email + token_uri non-empty). Run at parse time so the operator gets an actionable error before the first chat. - JWT claims: `{iss=client_email, scope=cloud-platform, aud=token_uri, iat=now, exp=now+3600}` per Google's documented flow. Signed RS256 via `jsonwebtoken` crate. - Token endpoint POST: form-encoded `grant_type=urn:ietf:params :oauth:grant-type:jwt-bearer` + `assertion=<jwt>`. Response is `{access_token, expires_in}`; `token_type` always "Bearer" so discarded. - Cache: `Arc<RwLock<HashMap<String, CachedToken>>>` keyed by `client_email`. Read-lock on the common cache-hit path; write- lock only on mint. Cache stores `expires_at = now + expires_in - 60s` to refresh proactively. - Endpoint error handling: 4xx/5xx from token endpoint becomes `BridgeError::Config` with body truncated to 500 chars (GCP returns operator-actionable OAuth-shape errors like `invalid_grant`; doesn't echo the SA private key). ## Bridge wiring `VertexBridge` now holds `Arc<TokenMinter>`. Both `chat_gemini` and `chat_gemini_stream` resolve the bearer via `creds.resolve_access_token(&minter).await?` BEFORE entering the deadline-wrapped request future, so token-mint failures surface as direct `Err` returns from `chat()` / `chat_stream()` rather than being yielded mid-stream. The two streams test seams (`with_api_base_override` for Vertex host, `with_token_endpoint_override` for OAuth host) are independent — a wiremock test can override either or both, e.g. exercise the full SA mint + chat round-trip end-to-end against two separate wiremock servers. ## Tests (10 new across token_mint + bridge, all wiremock-based) token_mint::tests: - mint_signs_jwt_with_correct_claims_and_posts_to_token_uri: captures the assertion JWT, decodes against the matching public key, asserts iss/scope/aud/iat/exp claims have the expected values + relationships - get_token_caches_within_ttl_and_only_calls_endpoint_once: 3 calls, wiremock `.expect(1)` enforces single mint - token_endpoint_5xx_surfaces_as_config_error_with_body_truncated: 503 + arbitrary body → Config error with HTTP status + body - invalid_pem_surfaces_clear_error_before_endpoint_call: bad PEM → Config error, `.expect(0)` on endpoint enforces no network call - wrong_type_field_rejected_before_endpoint_call: type="external_account" → Config error before network call bridge::tests: - vertex_secret_accepts_service_account_json_path: SA-only secret parses cleanly + SA validate runs - vertex_secret_rejects_both_credential_modes_set: mutual exclusion - vertex_secret_rejects_neither_credential_mode_set: mutual exclusion - vertex_secret_rejects_empty_access_token_string: edge case where operator pastes "" rather than omitting the field - chat_gemini_sa_path_mints_token_and_forwards_as_bearer: full end-to-end pipeline — SA JSON in secret → token minted via wiremock OAuth endpoint → minted token forwarded as Authorization Bearer on Vertex chat call. Asserts the OAuth endpoint is hit exactly ONCE across two chats (cache hit on the second). Plus the existing `vertex_secret_parses_full_form` is updated for the new `Option<String>` shape. ## Test fixtures Embedded 2048-bit RSA test keypair at `crates/aisix-provider-vertex/test-fixtures/test_sa_{private,public}.pem`. Generated via `openssl genpkey -algorithm RSA -pkeyopt rsa_keygen_bits:2048` + matching pubout. Deterministic so JWT byte sequences reproduce across runs. NOT a real GCP key — safe to commit. `cargo test -p aisix-provider-vertex` → 60/60 PASS. `cargo clippy -p aisix-provider-vertex --all-targets -- -D warnings` clean. `cargo fmt --all` applied. ## Dependency added `jsonwebtoken = "9"` — actively maintained, audited via cargo-audit, no transitive heavy deps. Alternatives considered + rejected: `yup-oauth2` (much larger surface, includes a full OAuth client we don't need) and `gcp_auth` (pulls hyper directly, doesn't compose with the existing reqwest::Client). ## References - GCP SA OAuth flow: https://developers.google.com/identity/protocols/oauth2/service-account#authorizingrequests - JWT Bearer grant (RFC 7523): https://www.rfc-editor.org/rfc/rfc7523 - Standard SA JSON shape: emitted by `gcloud iam service-accounts keys create`. Field set confirmed against the python-aiplatform SDK's `google.oauth2.service_account.Credentials.from_service_account_info()` consumer signature. ## Unblocks AC.11 partial in api7/AISIX-Cloud#302: "Vertex 上 Gemini / Claude / Llama 都通过单一 google-vertex catalog 入口能 stream" — Gemini half is now fully done (chat + stream wired by ai-gateway#386, SA OAuth wired by this PR). Anthropic + Llama publishers remain blocked on D5.3 / D5.4 per the roadmap. ## Out of scope - Live-mode e2e against a real GCP project (the cp-api/dashboard matrix spec is a separate AISIX-Cloud PR — gated on #481 wiring mock-vertex into compose first) - TTL-expiry behavior under simulated time advance (clock-mocking infra not yet established in the crate; manual sanity-checked the expires_at arithmetic; covered indirectly by the cache-hit test)
|
Warning Review limit reached
Your plan currently allows 3 reviews/hour. Refill in 13 minutes and 18 seconds. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more review capacity refills, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than trial, open-source, and free plans. In all cases, review capacity refills continuously over time. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Free Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds in-process GCP OAuth2 token minting to the Vertex provider, enabling automatic JWT-based access token generation from service-account JSON keys. The bridge now supports both pre-minted tokens and full service-account JSON secrets, with per-bridge caching and TTL-aware token refresh. ChangesVertex OAuth2 token minting
Sequence DiagramssequenceDiagram
participant VertexBridge
participant TokenMinter
participant Cache
participant TokenEndpoint
VertexBridge->>TokenMinter: get_token(sa)
TokenMinter->>Cache: read cached token for client_email
alt Token valid and not expired
Cache-->>TokenMinter: return cached token
else Token missing or expired
TokenMinter->>TokenMinter: build JWT claims (iss, scope, aud, iat, exp)
TokenMinter->>TokenMinter: sign with RS256 private key
TokenMinter->>TokenEndpoint: POST grant_type=jwt-bearer
TokenEndpoint-->>TokenMinter: access_token, expires_in
TokenMinter->>Cache: store token with computed TTL
Cache-->>TokenMinter: stored
end
TokenMinter-->>VertexBridge: access_token
sequenceDiagram
participant Client
participant VertexBridge
participant VertexSecret
participant TokenMinter
participant VertexAPI
Client->>VertexBridge: chat_gemini(creds)
VertexBridge->>VertexSecret: resolve_access_token(&token_minter)
alt access_token mode
VertexSecret-->>VertexBridge: return pre-minted token
else service_account_json mode
VertexSecret->>TokenMinter: get_token(sa)
TokenMinter-->>VertexSecret: minted token
VertexSecret-->>VertexBridge: minted token
end
VertexBridge->>VertexAPI: POST Authorization: Bearer {token}
VertexAPI-->>VertexBridge: response
🎯 4 (Complex) | ⏱️ ~45 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 |
There was a problem hiding this comment.
Pull request overview
Implements in-process GCP service-account JWT-bearer token minting for the Vertex provider, enabling ProviderKey.secret to use either a pre-minted access_token (backward compatible) or a full service_account_json, with an in-process token cache to avoid hourly operator refresh.
Changes:
- Add
token_mintmodule to sign JWTs (RS256), exchange them for OAuth2 access tokens, and cache tokens keyed byclient_email. - Extend
VertexSecretparsing/validation to enforce mutual exclusion betweenaccess_tokenandservice_account_json, and wire token resolution into both chat and streaming requests. - Add wiremock-based tests and RSA PEM fixtures; add
jsonwebtokendependency.
Reviewed changes
Copilot reviewed 6 out of 7 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| crates/aisix-provider-vertex/test-fixtures/test_sa_public.pem | Adds deterministic RSA public key fixture for JWT verification in tests. |
| crates/aisix-provider-vertex/test-fixtures/test_sa_private.pem | Adds deterministic RSA private key fixture for signing JWTs in tests. |
| crates/aisix-provider-vertex/src/token_mint.rs | New token minting + caching implementation and unit tests. |
| crates/aisix-provider-vertex/src/lib.rs | Registers the new token_mint module and updates crate status docs. |
| crates/aisix-provider-vertex/src/bridge.rs | Adds SA-vs-token secret parsing, token resolution, and test seams for OAuth endpoint override. |
| crates/aisix-provider-vertex/Cargo.toml | Adds jsonwebtoken dependency. |
| Cargo.lock | Locks new dependency versions (jsonwebtoken, simple_asn1, etc.). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+176
to
+177
| // Cache miss or expired — mint fresh under write-lock. | ||
| let (access_token, expires_in_secs) = self.mint(sa).await?; |
Comment on lines
+238
to
+242
| let body = resp.text().await.unwrap_or_default(); | ||
| let truncated: String = body.chars().take(500).collect(); | ||
| return Err(BridgeError::Config(format!( | ||
| "vertex token mint upstream returned HTTP {status}: {truncated}" | ||
| ))); |
Comment on lines
+154
to
+160
| /// Test-only seam: replace the SA's `token_uri` host with this | ||
| /// URL. JWT contents are unchanged so the assertion shape is | ||
| /// still verifiable end-to-end. | ||
| #[cfg(test)] | ||
| pub(crate) fn with_token_endpoint_override(mut self, url: impl Into<String>) -> Self { | ||
| self.token_endpoint_override = Some(url.into()); | ||
| self |
Comment on lines
+248
to
+252
| if parsed.access_token.as_deref().is_some_and(str::is_empty) { | ||
| return Err(BridgeError::Config( | ||
| "vertex provider_key.secret.access_token is empty".into(), | ||
| )); | ||
| } |
Comment on lines
+510
to
515
| // Resolve bearer: pre-minted token verbatim, or mint+cache | ||
| // via the in-process token minter from SA JSON. Failure | ||
| // surfaces as a Config error (operator-actionable). | ||
| let access_token = creds.resolve_access_token(&self.token_minter).await?; | ||
| let headers = build_request_headers(&access_token, &ctx.request_id)?; | ||
| let client = self.client.clone(); |
Comment on lines
+585
to
591
| // Resolve bearer (pre-minted OR minted-from-SA) BEFORE | ||
| // entering the stream future so token-mint errors surface | ||
| // as a direct Err return rather than being yielded mid-stream. | ||
| let access_token = creds.resolve_access_token(&self.token_minter).await?; | ||
| let headers = build_request_headers(&access_token, &ctx.request_id)?; | ||
| let client = self.client.clone(); | ||
| let started = Instant::now(); |
Comment on lines
+85
to
+91
| /// Test-only seam: replace the SA `token_uri` host on the | ||
| /// internal token minter. Used by the SA-flow tests to redirect | ||
| /// JWT-bearer assertions to a wiremock endpoint. | ||
| #[cfg(test)] | ||
| pub(crate) fn with_token_endpoint_override(mut self, url: impl Into<String>) -> Self { | ||
| let new_minter = TokenMinter::new(self.client.clone()).with_token_endpoint_override(url); | ||
| self.token_minter = Arc::new(new_minter); |
…IUM on #387) Audit-aigw-387-vertex-oauth flagged that all non-2xx responses from the GCP token endpoint were returned as BridgeError::Config (HTTP 500 "operator must fix"). A flapping GCP OAuth backend would surface to the customer as an unconditional 500 with no retry hint — wrong classification for a transient upstream failure. Fix per audit's suggested code: - 5xx → BridgeError::upstream_status_with_retry_after(status, msg, retry_after) Surfaces the upstream status (cooldown layer maps to 502), and forwards GCP's Retry-After header so a token-endpoint backoff hint flows end-to-end through the cooldown layer rather than being silently dropped. - 4xx → BridgeError::Config (unchanged) invalid_grant / clock skew / bad SA ARE operator-actionable; Config (500) remains correct. Reads Retry-After BEFORE consuming the body via resp.text() so the header survives the body read. Test updates: - token_endpoint_5xx_surfaces_as_config_error_with_body_truncated → token_endpoint_5xx_surfaces_as_upstream_status_with_retry_hint (rename + reshape; asserts status=503, retry_after=Some(30s), message contains body) - New: token_endpoint_4xx_surfaces_as_config_error — pins the 4xx-still-Config branch with an invalid_grant fixture so a future regression that broadened the upstream_status mapping surfaces here `cargo test -p aisix-provider-vertex` → 61/61 PASS (was 60; +1 new 4xx test). `cargo clippy -p aisix-provider-vertex --all-targets -- -D warnings` clean. Audit LOWs (non-blocking, deferred per CLAUDE.md §8): - LOW: no e2e for chat_gemini_stream SA path — both code paths call resolve_access_token at the same lifecycle point; a regression would have to diverge both sites - LOW: cache-miss race on first-call burst — acceptable (worst case one extra mint; spec doesn't require single-flight) - LOW: expires_in:0 edge case — saturating_sub handles cleanly
Collaborator
Author
Audit response — addressedIndependent audit-aigw-387-vertex-oauth returned with 1 MEDIUM + 3 LOWs. Addressed in commit `fbd8d95`: MEDIUM — token-endpoint 5xx misclassified as Config (HTTP 500 operator-must-fix) ✅ fixed in codeCritical correctness bug surfaced by the audit. A flapping GCP OAuth backend would have hit the customer with unconditional 500s with no retry hint. Fixed per audit's suggested code: ```rust Reads `Retry-After` BEFORE consuming the body so the header survives the `resp.text()` call and flows through to the cooldown layer. Tests updated:
LOWs (non-blocking, deferred per CLAUDE.md §8)
Result`cargo test -p aisix-provider-vertex` → 61/61 PASS (was 60; +1 new 4xx test). All merge-blocking findings addressed. Awaiting fresh CI green. |
Merged
5 tasks
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 free
to 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
Closes the second half of #302 Phase E productionization (Step 2a streaming merged as #386). Before this PR, `ProviderKey.secret` required a pre-minted GCP OAuth2 `access_token` — operator was responsible for refresh (GCP token TTL ~1h), awkward for production-scale deployments.
`VertexSecret` now accepts EITHER mode with mutual-exclusion validation at parse time:
New module: `token_mint`
`TokenMinter` owns the cache + mint pipeline:
Bridge wiring
`VertexBridge` holds `Arc`. `chat_gemini` and `chat_gemini_stream` resolve the bearer via `creds.resolve_access_token(&minter).await?` BEFORE entering the deadline-wrapped request future, so token-mint failures surface as direct `Err` returns rather than yielded mid-stream.
Two independent test seams:
A wiremock test can exercise the full SA mint + chat round-trip against two separate wiremock servers.
Tests (10 new across token_mint + bridge, all wiremock-based)
`token_mint::tests` (5):
`bridge::tests` (5 new + 1 updated):
Test fixtures
Embedded 2048-bit RSA test keypair at `crates/aisix-provider-vertex/test-fixtures/test_sa_{private,public}.pem`. Generated via `openssl genpkey -algorithm RSA -pkeyopt rsa_keygen_bits:2048`. NOT a real GCP key — safe to commit.
Test plan
Dependency added
`jsonwebtoken = "9"` — actively maintained, audited via cargo-audit, no transitive heavy deps. Alternatives `yup-oauth2` (much larger surface) and `gcp_auth` (pulls hyper directly) considered + rejected.
References (per CLAUDE.md §7)
Unblocks
AC.11 in api7/AISIX-Cloud#302: "Vertex 上 Gemini / Claude / Llama 都通过单一 google-vertex catalog 入口能 stream" — Gemini half is now fully done (chat + stream wired by #386, SA OAuth wired by this PR). Anthropic + Llama publishers remain blocked on D5.3 / D5.4 per the roadmap.
Out of scope
Summary by CodeRabbit