feat(rust): port VertexBase auth (Google credential resolution + token) to litellm-core as a base provider - #33906
Conversation
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Greptile SummaryThis PR ports the Python
Confidence Score: 3/5Safe to merge for now since the module is behind a default-off feature flag with no live route, but the cache error handling in The credential classifier and project-ID resolution are correct and well-tested. The process-wide singleton cache propagates Mutex lock-poisoning as a hard auth error at both the lookup and the store step — if the lock is ever poisoned, every subsequent resolve_token call fails permanently without attempting a network fetch, and a successfully-fetched token is silently discarded at the store step. litellm-rust/crates/core/src/providers/vertex_ai/vertex_base.rs — specifically the
|
| Filename | Overview |
|---|---|
| litellm-rust/crates/core/src/providers/vertex_ai/vertex_base.rs | New Vertex AI base provider for Google credential resolution and token minting. The token fetch is not resilient to cache errors: a poisoned cache lock discards a successfully-fetched token and permanently blocks all future auth. |
| litellm-rust/crates/core/Cargo.toml | Adds google-cloud-auth = "=1.12.0" pinned as an optional dep behind the new vertex-auth feature; existing workspace deps like sha2 are already present. Clean addition. |
| litellm-rust/crates/core/src/providers/vertex_ai/constants.rs | New constants file defining all Google/Vertex environment variable names and the cloud-platform scope. Values are correct. |
| litellm-rust/crates/core/src/providers/vertex_ai/mod.rs | Exposes vertex_base and constants modules behind #[cfg(feature = "vertex-auth")] guards. Correct and minimal. |
Reviews (1): Last reviewed commit: "feat(rust): add Vertex auth base provide..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| None => { | ||
| let provider = build_credentials(source, credentials)?; | ||
| let access_token = provider.access_token().await.map_err(|error| { | ||
| CoreError::Auth(format!("Google token resolution failed: {error}")) | ||
| })?; | ||
| store_credentials(cache_key, provider.clone())?; | ||
| return Ok(VertexToken { | ||
| access_token: access_token.token, | ||
| project_id, | ||
| }); | ||
| } | ||
| }; | ||
| let access_token = provider | ||
| .access_token() | ||
| .await | ||
| .map_err(|error| CoreError::Auth(format!("Google token resolution failed: {error}")))?; | ||
| Ok(VertexToken { | ||
| access_token: access_token.token, | ||
| project_id, | ||
| }) |
There was a problem hiding this comment.
store_credentials(cache_key, provider.clone())? propagates a cache write failure (e.g. a poisoned Mutex) after access_token was already successfully fetched. The caller receives an auth error and must retry even though a perfectly good token was in hand. Worse, cached_credentials(&cache_key)? on line 137 fails at the lookup step too, meaning a poisoned lock prevents any credential building at all—every subsequent resolve_token call fails immediately without attempting a network fetch, which is unrecoverable without a process restart. Both helpers should treat lock-poisoning as a non-fatal cache miss rather than a hard error: lock().unwrap_or_else(|e| e.into_inner()), or simply let _ = store_credentials(...) so a valid token is always returned if the fetch succeeds.
| project_id.ok_or(CoreError::MissingField("project_id")) | ||
| } | ||
|
|
||
| fn cache_key(credentials: Option<&Value>, project_id: &str) -> String { | ||
| let mut hasher = Sha256::new(); | ||
| if let Some(credentials) = credentials { | ||
| hasher.update(credentials.to_string()); | ||
| } | ||
| hasher.update(project_id); | ||
| format!("{:x}", hasher.finalize()) | ||
| } | ||
|
|
||
| fn cached_credentials(key: &str) -> CoreResult<Option<AccessTokenCredentials>> { | ||
| let cache = VERTEX_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); | ||
| let mut cache = cache | ||
| .lock() | ||
| .map_err(|_| CoreError::Auth("Vertex credential cache lock poisoned".to_string()))?; | ||
| Ok(cache.get_cache(key)) | ||
| } |
There was a problem hiding this comment.
Clone semantics of AccessTokenCredentials determine cache effectiveness
InMemoryCache::get_cache returns a clone of the stored AccessTokenCredentials. If Clone on that type is a value copy (not Arc-backed shared state), each caller gets an independent snapshot of the credentials as they existed at storage time. After the internal token expires, every caller would independently re-fetch from Google rather than sharing a single refresh. Given the InMemoryCache default TTL is 10 minutes (vs. a 1-hour token lifetime), the cache's primary benefit reduces to avoiding credential-object rebuilds rather than avoiding token network calls. If the google-cloud-auth crate's AccessTokenCredentials uses Arc internally for its token cache, this is a non-issue—but that contract is implicit. A clarifying comment noting the assumption (or preferring an Arc<AccessTokenCredentials> in the cache) would make the intended behavior explicit.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| fn classifier_covers_authorized_service_and_default_sources() { | ||
| assert_eq!( | ||
| classify_auth(Some(&json!({"type": "authorized_user"}))), | ||
| VertexAuthSource::AuthorizedUser | ||
| ); | ||
| assert_eq!( | ||
| classify_auth(Some(&json!({"type": "service_account"}))), | ||
| VertexAuthSource::ServiceAccount | ||
| ); | ||
| assert_eq!( | ||
| classify_auth(Some(&json!({"client_email": "test@example.com"}))), | ||
| VertexAuthSource::ServiceAccount | ||
| ); | ||
| assert_eq!(classify_auth(None), VertexAuthSource::DefaultAdc); | ||
| } | ||
|
|
||
| #[test] | ||
| fn credential_input_reads_files_inline_json_and_objects() { | ||
| let path = std::env::temp_dir().join(format!("vertex-auth-{}.json", std::process::id())); | ||
| std::fs::write(&path, r#"{"type":"authorized_user"}"#).unwrap(); | ||
| assert_eq!( |
There was a problem hiding this comment.
std::fs::remove_file(path).unwrap() is called unconditionally after the first assertion, but if the assert_eq! on line 304 panics the file is left in the temp directory. Wrapping the body in a scope with a defer-style guard (or using a Drop wrapper) would ensure cleanup even on assertion failure.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
2 participants
Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
This is PR 2 of the Rust migration's auth de-risking; PR 1 (#33888, merged) did the same for Bedrock's
boto3/botocoreauth. The question here is whether thegoogle-authauth surface that the Python Vertex path depends on has a maintained Rust equivalent that behaves the same, so the Rust gateway can resolve Google credentials and mint access tokens natively instead of calling back into Python. It does: Google's officialgoogle-cloud-authcrate (fromgoogleapis/google-cloud-rust)Nothing is wired into a live route yet and the whole module is behind a default-off
vertex-authcargo feature, so there is no proxy endpoint to curl at this stage; core still builds and clippies clean with the feature off. A live token call is not included because there are no Google credentials provisioned in this environment, unlike the AWS keys used for PR 1. The proof at this stage is the pure credential-source classifier and the credential/project resolution, pinned by unit tests that assert the exact Pythonload_authdispatch order and the project-id precedence; both fail if the branch order or precedence driftsCaptured at commit e0f2e5e
Type
🆕 New Feature
Changes
Ports the auth surface of Python's
VertexBase(litellm/llms/vertex_ai/vertex_llm_base.py) into the Rust workspace as a base provider atlitellm-core/src/providers/vertex_ai/vertex_base.rs, so the Rust Vertex path can resolve Google credentials and obtain an access token natively. It is gated behind a new, default-offvertex-authcargo feature onlitellm-core; thegoogle-cloud-authdependency is optional, pinned to=1.12.0, and rustls-only (default-features = false,default-rustls-provider), so no OpenSSL or native-tls is linkedCredential-source selection is a pure function over the parsed credentials JSON, preserving the same predicate order as
load_auth, so it is unit-testable without the networkCredential input resolution mirrors Python: a string that names an existing file is read and parsed as JSON, otherwise the string is parsed as inline JSON, and an already-parsed object is passed through; empty or whitespace-only credentials are treated as absent, and non-object or malformed JSON returns a typed error rather than panicking. Project id follows Python's precedence: an explicit project wins, then the credential's
project_id, then itsquota_project_id, thenGOOGLE_CLOUD_PROJECT, thenGOOGLE_CLOUD_QUOTA_PROJECT. Tokens are requested with thecloud-platformscopeAll six sources are handled by
google-cloud-authwith no silent downgrade or faked parity: service account, authorized user, default ADC, and the threeexternal_accountworkload-identity variants (AWS-sourced, executable/pluggable, and identity-pool), the last three built through the crate'sexternal_accountbuilder which dispatches on the credential JSON internallyCaching reuses the shared
litellm-coreInMemoryCache(added in PR 1) through a process-wideOnceLock<Mutex<InMemoryCache<AccessTokenCredentials>>>, keyed by a SHA-256 hash of the resolved credential JSON plus project id so raw credential material is never used as a key or logged. Python's sync/async single-flight and background-refresh machinery is intentionally not reimplemented: it exists to work aroundgoogle-auth's syncrefresh()not being concurrency-safe, whereas the RustAccessTokenCredentialsare async and own their token cache and refresh internallyPlacement note for reviewers: reading credentials and minting tokens is auth work that the core-purity guidance in
litellm-rust/CLAUDE.mdnormally keeps out oflitellm-core. Putting the base provider here, mirroring how Python's base provider owns auth and matching the merged Bedrock PR, is a deliberate, directed decision; the module carries a short note to that effect and the purity guidance will be reconciled separately rather than silentlyFinal Attestation
Link to Devin session: https://app.devin.ai/sessions/bdf99b89a6584cc0989b5605765aec03
Requested by: @ishaan-berri