#![cfg(feature = "vertex")]
#![cfg_attr(docsrs, doc(cfg(feature = "vertex")))]
use std::sync::Arc;
use gcp_auth::TokenProvider;
use crate::auth::{RequestSigner, SignerResult};
const VERTEX_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const VERTEX_SCOPES: &[&str] = &[VERTEX_SCOPE];
#[derive(Clone)]
pub struct VertexCredentials {
inner: CredentialInner,
}
#[derive(Clone)]
enum CredentialInner {
Static(String),
Adc(Arc<dyn TokenProvider>),
}
impl VertexCredentials {
#[must_use]
pub fn from_token(token: impl Into<String>) -> Self {
Self {
inner: CredentialInner::Static(token.into()),
}
}
pub async fn from_adc() -> Result<Self, gcp_auth::Error> {
let provider = gcp_auth::provider().await?;
Ok(Self {
inner: CredentialInner::Adc(provider),
})
}
#[must_use]
pub fn from_env() -> Option<Self> {
if let Ok(token) = std::env::var("VERTEX_ACCESS_TOKEN") {
return Some(Self::from_token(token));
}
if std::env::var_os("GOOGLE_APPLICATION_CREDENTIALS").is_some() {
let handle = tokio::runtime::Handle::current();
return match handle.block_on(gcp_auth::provider()) {
Ok(provider) => Some(Self {
inner: CredentialInner::Adc(provider),
}),
Err(_) => None,
};
}
None
}
fn resolve_token(&self) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
match &self.inner {
CredentialInner::Static(t) => Ok(t.clone()),
CredentialInner::Adc(provider) => {
let handle = tokio::runtime::Handle::current();
let token = handle
.block_on(provider.token(VERTEX_SCOPES))
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?;
Ok(token.as_str().to_owned())
}
}
}
}
impl std::fmt::Debug for VertexCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self.inner {
CredentialInner::Static(t) => f
.debug_struct("VertexCredentials")
.field("kind", &"static-token")
.field("token", &format!("<redacted, {} chars>", t.len()))
.finish(),
CredentialInner::Adc(_) => f
.debug_struct("VertexCredentials")
.field("kind", &"adc")
.finish(),
}
}
}
#[derive(Debug, Clone)]
pub struct VertexSigner {
credentials: VertexCredentials,
}
impl VertexSigner {
#[must_use]
pub fn new(credentials: VertexCredentials) -> Self {
Self { credentials }
}
}
impl RequestSigner for VertexSigner {
fn sign(&self, request: &mut reqwest::Request) -> SignerResult {
request.headers_mut().remove("x-api-key");
let token = self.credentials.resolve_token()?;
let bearer = format!("Bearer {token}");
request
.headers_mut()
.insert("authorization", bearer.parse()?);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_request_with_api_key() -> reqwest::Request {
let client = reqwest::Client::new();
client
.post("https://us-east5-aiplatform.googleapis.com/v1/projects/my-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:rawPredict")
.header("x-api-key", "sk-ant-test-key")
.body(r#"{"messages":[{"role":"user","content":"hi"}]}"#)
.build()
.unwrap()
}
fn make_request_without_api_key() -> reqwest::Request {
let client = reqwest::Client::new();
client
.post("https://us-east5-aiplatform.googleapis.com/v1/projects/my-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:rawPredict")
.body(r#"{"messages":[{"role":"user","content":"hi"}]}"#)
.build()
.unwrap()
}
fn static_signer(token: &str) -> VertexSigner {
VertexSigner::new(VertexCredentials::from_token(token))
}
#[test]
fn sign_adds_authorization_bearer_header() {
let signer = static_signer("ya29.test-token");
let mut req = make_request_without_api_key();
signer.sign(&mut req).expect("sign succeeds");
let auth = req
.headers()
.get("authorization")
.expect("authorization header set by signer");
let auth_str = auth.to_str().expect("authorization is ASCII");
assert_eq!(
auth_str, "Bearer ya29.test-token",
"expected bearer prefix: {auth_str}"
);
}
#[test]
fn sign_removes_x_api_key_header() {
let signer = static_signer("ya29.test-token");
let mut req = make_request_with_api_key();
assert!(
req.headers().get("x-api-key").is_some(),
"x-api-key must be present before sign()"
);
signer.sign(&mut req).expect("sign succeeds");
assert!(
req.headers().get("x-api-key").is_none(),
"x-api-key must be removed after sign()"
);
}
#[test]
fn sign_sets_correct_bearer_format() {
let token = "ya29.c.long-token-value-here";
let signer = static_signer(token);
let mut req = make_request_without_api_key();
signer.sign(&mut req).expect("sign succeeds");
let auth = req
.headers()
.get("authorization")
.unwrap()
.to_str()
.unwrap();
assert!(auth.starts_with("Bearer "), "must start with 'Bearer '");
assert!(auth.contains(token), "must contain the token");
}
#[test]
fn credentials_redact_token_in_debug() {
let creds = VertexCredentials::from_token("ya29.very-secret-token");
let dbg = format!("{creds:?}");
assert!(!dbg.contains("very-secret-token"), "{dbg}");
assert!(dbg.contains("redacted"), "{dbg}");
}
#[test]
fn credentials_debug_shows_adc_kind_without_token() {
struct FakeProvider;
#[async_trait::async_trait]
impl TokenProvider for FakeProvider {
async fn token(
&self,
_scopes: &[&str],
) -> Result<Arc<gcp_auth::Token>, gcp_auth::Error> {
unimplemented!()
}
async fn project_id(&self) -> Result<Arc<str>, gcp_auth::Error> {
unimplemented!()
}
}
let creds = VertexCredentials {
inner: CredentialInner::Adc(Arc::new(FakeProvider)),
};
let dbg = format!("{creds:?}");
assert!(dbg.contains("adc"), "{dbg}");
}
#[test]
fn from_env_returns_none_when_no_vars_set() {
let has_token = std::env::var("VERTEX_ACCESS_TOKEN").is_ok();
let has_adc = std::env::var_os("GOOGLE_APPLICATION_CREDENTIALS").is_some();
if has_token || has_adc {
return;
}
let result = {
VertexCredentials::from_env()
};
assert!(
result.is_none(),
"expected None when env vars are absent: {result:?}"
);
}
#[test]
fn from_env_returns_static_creds_when_vertex_access_token_env_is_set() {
let creds = VertexCredentials::from_token("ya29.env-test-token");
assert!(
matches!(creds.inner, CredentialInner::Static(_)),
"from_token must yield a static credential"
);
}
}