i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

修复 Vertex AI 服务账号访问

Browse Source
This commit is contained in:
Codex
2026-05-10 15:51:38 +08:00
parent 7f101431c5
commit fb96771b56
17 changed files with 844 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/aether-provider-transport/src/conversion.rs
View File

@@ -15,8 +15,8 @@ use crate::policy::{
local_standard_transport_unsupported_reason_with_network,
};
use crate::vertex::{
is_vertex_api_key_transport_context,
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
is_vertex_api_key_transport_context, is_vertex_transport_context,
local_vertex_gemini_transport_unsupported_reason_with_network,
resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM,
};
use crate::GatewayProviderTransportSnapshot;
@@ -106,8 +106,8 @@ pub fn request_conversion_transport_unsupported_reason(
"claude:messages" => {
local_standard_transport_unsupported_reason_with_network(transport, "claude:messages")
}
"gemini:generate_content" if is_vertex_api_key_transport_context(transport) => {
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
"gemini:generate_content" if is_vertex_transport_context(transport) => {
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
}
"gemini:generate_content" => local_gemini_transport_unsupported_reason_with_network(
transport,

5
crates/aether-provider-transport/src/lib.rs
View File

@@ -103,7 +103,10 @@ pub use standard::{
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
};
pub use vertex::{is_vertex_api_key_transport_context, uses_vertex_api_key_query_auth};
pub use vertex::{
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
is_vertex_transport_context, uses_vertex_api_key_query_auth,
};
pub use video::{
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,

5
crates/aether-provider-transport/src/oauth_refresh/mod.rs
View File

@@ -19,6 +19,9 @@ use super::kiro::{
supports_local_kiro_request_auth_resolution, KiroOAuthRefreshAdapter, KiroRequestAuth,
};
use super::snapshot::GatewayProviderTransportSnapshot;
use super::vertex::{
supports_local_vertex_service_account_auth_resolution, VertexServiceAccountRefreshAdapter,
};
#[derive(Debug, Clone, PartialEq, Eq)]
#[allow(clippy::large_enum_variant)]
@@ -339,6 +342,7 @@ impl LocalOAuthRefreshCoordinator {
Self {
adapters: vec![
Arc::new(KiroOAuthRefreshAdapter::default()),
Arc::new(VertexServiceAccountRefreshAdapter),
Arc::new(GenericOAuthRefreshAdapter::default()),
],
cache: Mutex::new(BTreeMap::new()),
@@ -547,6 +551,7 @@ pub fn supports_local_oauth_request_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_kiro_request_auth_resolution(transport)
|| supports_local_vertex_service_account_auth_resolution(transport)
|| supports_local_generic_oauth_request_auth_resolution(transport)
}

48
crates/aether-provider-transport/src/request_url/mod.rs
View File

@@ -15,7 +15,8 @@ use crate::url::{
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
};
use crate::vertex::{
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
build_vertex_api_key_gemini_content_url, build_vertex_service_account_gemini_content_url,
resolve_local_vertex_api_key_query_auth, resolve_local_vertex_service_account_auth_config,
};
#[derive(Debug, Clone, Copy)]
@@ -256,6 +257,14 @@ fn build_transport_hook_url(
params.request_query,
);
}
if let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) {
return build_vertex_service_account_gemini_content_url(
params.mapped_model?,
params.upstream_is_stream,
&auth_config,
params.request_query,
);
}
}
if is_antigravity_provider_transport(transport) {
@@ -526,6 +535,43 @@ mod tests {
);
}
#[test]
fn uses_vertex_service_account_hook_before_default_gemini_url() {
let mut transport = sample_transport(
"vertex_ai",
"gemini:generate_content",
"https://aiplatform.googleapis.com",
None,
);
transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"client_email":"svc@example.iam.gserviceaccount.com",
"private_key":"TEST-PRIVATE-KEY",
"project_id":"demo-project"
}"#
.to_string(),
);
let url = build_transport_request_url(
&transport,
TransportRequestUrlParams {
provider_api_format: "gemini:generate_content",
mapped_model: Some("gemini-3.1-pro-preview"),
upstream_is_stream: false,
request_query: Some("foo=bar&beta=1"),
kiro_api_region: None,
},
)
.expect("vertex service account hook url");
assert_eq!(
url,
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent?foo=bar"
);
}
#[test]
fn builds_openai_responses_url_for_formal_format_name() {
let transport = sample_transport(

11
crates/aether-provider-transport/src/same_format_provider/mod.rs
View File

@@ -23,8 +23,8 @@ use crate::rules::{
};
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::vertex::{
is_vertex_api_key_transport_context,
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
is_vertex_service_account_transport_context, is_vertex_transport_context,
local_vertex_gemini_transport_unsupported_reason_with_network,
};
use crate::{build_transport_request_url, ensure_upstream_auth_header, TransportRequestUrlParams};
@@ -103,7 +103,7 @@ pub fn classify_same_format_provider_request_behavior(
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code");
let is_vertex = is_vertex_api_key_transport_context(transport);
let is_vertex = is_vertex_transport_context(transport);
let is_kiro = is_kiro_provider_transport(transport);
let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
transport.endpoint.config.as_ref(),
@@ -328,7 +328,7 @@ pub fn same_format_provider_transport_unsupported_reason(
} else if behavior.is_claude_code {
local_claude_code_transport_unsupported_reason_with_network(transport, api_format)
} else if behavior.is_vertex {
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
} else {
match family {
SameFormatProviderFamily::Standard => {
@@ -391,6 +391,9 @@ pub fn should_try_same_format_provider_oauth_auth(
behavior.is_kiro
|| matches!(family, SameFormatProviderFamily::Standard)
&& resolve_local_standard_auth(transport).is_none()
|| matches!(family, SameFormatProviderFamily::Gemini)
&& behavior.is_vertex
&& is_vertex_service_account_transport_context(transport)
|| matches!(family, SameFormatProviderFamily::Gemini)
&& !behavior.is_vertex
&& resolve_local_gemini_auth(transport).is_none()

348
crates/aether-provider-transport/src/vertex/auth.rs
View File

@@ -1,6 +1,29 @@
use std::collections::BTreeMap;
use async_trait::async_trait;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
use rsa::pkcs1::DecodeRsaPrivateKey;
use rsa::pkcs1v15::SigningKey;
use rsa::pkcs8::DecodePrivateKey;
use rsa::signature::{SignatureEncoding, Signer};
use rsa::RsaPrivateKey;
use serde_json::{json, Value};
use sha2::Sha256;
use url::form_urlencoded;
use super::super::oauth_refresh::{
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthRefreshAdapter,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
};
use super::super::snapshot::GatewayProviderTransportSnapshot;
pub const VERTEX_API_KEY_QUERY_PARAM: &str = "key";
pub const VERTEX_SERVICE_ACCOUNT_AUTH_HEADER: &str = "authorization";
pub const VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE: &str = "vertex_ai";
pub const GOOGLE_OAUTH_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const SERVICE_ACCOUNT_REFRESH_SKEW_SECS: u64 = 120;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VertexApiKeyQueryAuth {
@@ -8,6 +31,16 @@ pub struct VertexApiKeyQueryAuth {
pub value: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VertexServiceAccountAuthConfig {
pub client_email: String,
pub private_key: String,
pub project_id: String,
pub token_uri: String,
pub region: Option<String>,
pub model_regions: BTreeMap<String, String>,
}
pub fn resolve_local_vertex_api_key_query_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<VertexApiKeyQueryAuth> {
@@ -39,14 +72,274 @@ pub fn resolve_local_vertex_api_key_query_auth(
})
}
pub fn resolve_local_vertex_service_account_auth_config(
transport: &GatewayProviderTransportSnapshot,
) -> Option<VertexServiceAccountAuthConfig> {
if !super::is_vertex_service_account_transport_context(transport) {
return None;
}
parse_vertex_service_account_auth_config(transport.key.decrypted_auth_config.as_deref())
}
pub fn supports_local_vertex_service_account_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
resolve_local_vertex_service_account_auth_config(transport).is_some()
}
pub fn parse_vertex_service_account_auth_config(
raw: Option<&str>,
) -> Option<VertexServiceAccountAuthConfig> {
let raw = raw.map(str::trim).filter(|value| !value.is_empty())?;
let value: Value = serde_json::from_str(raw).ok()?;
parse_vertex_service_account_auth_config_value(&value)
}
fn parse_vertex_service_account_auth_config_value(
value: &Value,
) -> Option<VertexServiceAccountAuthConfig> {
let client_email = json_string(value.get("client_email"))?;
let private_key = json_string(value.get("private_key"))?;
let project_id = json_string(value.get("project_id"))?;
let token_uri =
json_string(value.get("token_uri")).unwrap_or_else(|| GOOGLE_OAUTH_TOKEN_URL.to_string());
let region = json_string(value.get("region"));
let model_regions = value
.get("model_regions")
.and_then(Value::as_object)
.map(|items| {
items
.iter()
.filter_map(|(model, region)| {
let model = model.trim();
let region = region.as_str()?.trim();
(!model.is_empty() && !region.is_empty())
.then(|| (model.to_string(), region.to_string()))
})
.collect::<BTreeMap<_, _>>()
})
.unwrap_or_default();
Some(VertexServiceAccountAuthConfig {
client_email,
private_key,
project_id,
token_uri,
region,
model_regions,
})
}
fn json_string(value: Option<&Value>) -> Option<String> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
#[derive(Debug, Clone, Default)]
pub struct VertexServiceAccountRefreshAdapter;
#[async_trait]
impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter {
fn provider_type(&self) -> &'static str {
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE
}
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
supports_local_vertex_service_account_auth_resolution(transport)
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
if !entry
.provider_type
.eq_ignore_ascii_case(VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE)
{
return None;
}
if service_account_token_expires_soon(entry.expires_at_unix_secs) {
return None;
}
let name = entry.auth_header_name.trim();
let value = entry.auth_header_value.trim();
if name.is_empty() || value.is_empty() {
return None;
}
Some(LocalResolvedOAuthRequestAuth::Header {
name: name.to_ascii_lowercase(),
value: value.to_string(),
})
}
fn resolve_without_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
None
}
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool {
supports_local_vertex_service_account_auth_resolution(transport)
&& entry
.and_then(|cached| self.resolve_cached(transport, cached))
.is_none()
}
async fn refresh(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) else {
return Ok(None);
};
let now = aether_oauth::core::current_unix_secs();
let assertion = build_vertex_service_account_assertion(&auth_config, now)?;
let body = form_urlencoded::Serializer::new(String::new())
.append_pair(
"grant_type",
"urn:ietf:params:oauth:grant-type:jwt-bearer",
)
.append_pair("assertion", &assertion)
.finish();
let response = executor
.execute(
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
transport,
&LocalOAuthHttpRequest {
request_id: "vertex_ai:service-account-token",
method: reqwest::Method::POST,
url: auth_config.token_uri.clone(),
headers: BTreeMap::from([(
"content-type".to_string(),
"application/x-www-form-urlencoded".to_string(),
)]),
json_body: None,
body_bytes: Some(body.into_bytes()),
},
)
.await?;
if response.status_code != 200 {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
status_code: response.status_code,
body_excerpt: body_excerpt(&response.body_text),
});
}
let body_json: Value =
serde_json::from_str(&response.body_text).map_err(|err| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
message: format!("vertex service account token response is not JSON: {err}"),
}
})?;
let access_token = json_string(body_json.get("access_token")).ok_or_else(|| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
message: "vertex service account token response missing access_token".to_string(),
}
})?;
let expires_in = body_json
.get("expires_in")
.and_then(Value::as_u64)
.unwrap_or(3600);
Ok(Some(CachedOAuthEntry {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE.to_string(),
auth_header_name: VERTEX_SERVICE_ACCOUNT_AUTH_HEADER.to_string(),
auth_header_value: format!("Bearer {access_token}"),
expires_at_unix_secs: Some(now.saturating_add(expires_in)),
metadata: Some(json!({
"project_id": auth_config.project_id,
"client_email": auth_config.client_email,
})),
}))
}
}
pub fn build_vertex_service_account_assertion(
auth_config: &VertexServiceAccountAuthConfig,
now_unix_secs: u64,
) -> Result<String, LocalOAuthRefreshError> {
let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
let payload = URL_SAFE_NO_PAD.encode(
serde_json::to_string(&json!({
"iss": auth_config.client_email,
"sub": auth_config.client_email,
"scope": GOOGLE_CLOUD_PLATFORM_SCOPE,
"aud": auth_config.token_uri,
"iat": now_unix_secs,
"exp": now_unix_secs.saturating_add(3600),
}))
.map_err(|err| LocalOAuthRefreshError::InvalidResponse {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
message: format!("vertex service account jwt payload encode failed: {err}"),
})?,
);
let message = format!("{header}.{payload}");
let private_key = decode_vertex_service_account_private_key(auth_config.private_key.as_str())?;
let signing_key = SigningKey::<Sha256>::new(private_key);
let signature = signing_key.sign(message.as_bytes());
Ok(format!(
"{message}.{}",
URL_SAFE_NO_PAD.encode(signature.to_bytes())
))
}
fn decode_vertex_service_account_private_key(
private_key_pem: &str,
) -> Result<RsaPrivateKey, LocalOAuthRefreshError> {
match RsaPrivateKey::from_pkcs8_pem(private_key_pem) {
Ok(private_key) => return Ok(private_key),
Err(pkcs8_err) => RsaPrivateKey::from_pkcs1_pem(private_key_pem).map_err(|pkcs1_err| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
message: format!(
"vertex service account private_key parse failed: pkcs8: {pkcs8_err}; pkcs1: {pkcs1_err}"
),
}
}),
}
}
fn service_account_token_expires_soon(expires_at_unix_secs: Option<u64>) -> bool {
expires_at_unix_secs
.map(|expires_at_unix_secs| {
aether_oauth::core::current_unix_secs()
>= expires_at_unix_secs.saturating_sub(SERVICE_ACCOUNT_REFRESH_SKEW_SECS)
})
.unwrap_or(true)
}
fn body_excerpt(value: &str) -> String {
value.chars().take(500).collect()
}
#[cfg(test)]
mod tests {
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use rsa::pkcs1::{EncodeRsaPrivateKey, LineEnding};
use rsa::rand_core::OsRng;
use rsa::RsaPrivateKey;
use super::{resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM};
use super::{
decode_vertex_service_account_private_key, parse_vertex_service_account_auth_config,
resolve_local_vertex_api_key_query_auth,
supports_local_vertex_service_account_auth_resolution, VERTEX_API_KEY_QUERY_PARAM,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
@@ -136,4 +429,57 @@ mod tests {
.expect("custom aiplatform transport should resolve");
assert_eq!(auth.value, "vertex-secret");
}
#[test]
fn parses_service_account_auth_config() {
let config = parse_vertex_service_account_auth_config(Some(
r#"{
"client_email":"svc@example.iam.gserviceaccount.com",
"private_key":"TEST-PRIVATE-KEY",
"project_id":"demo-project",
"region":"global",
"model_regions":{"gemini-2.0-flash":"us-central1"}
}"#,
))
.expect("service account config should parse");
assert_eq!(config.client_email, "svc@example.iam.gserviceaccount.com");
assert_eq!(config.project_id, "demo-project");
assert_eq!(config.region.as_deref(), Some("global"));
assert_eq!(
config.model_regions.get("gemini-2.0-flash").map(String::as_str),
Some("us-central1")
);
}
#[test]
fn supports_vertex_service_account_auth_resolution() {
let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"client_email":"svc@example.iam.gserviceaccount.com",
"private_key":"TEST-PRIVATE-KEY",
"project_id":"demo-project"
}"#
.to_string(),
);
assert!(supports_local_vertex_service_account_auth_resolution(
&transport
));
}
#[test]
fn decodes_pkcs1_service_account_private_key() {
let mut rng = OsRng;
let private_key = RsaPrivateKey::new(&mut rng, 1024)
.expect("test RSA private key should generate")
.to_pkcs1_pem(LineEnding::LF)
.expect("test RSA private key should encode as PKCS#1 PEM");
decode_vertex_service_account_private_key(private_key.as_str())
.expect("PKCS#1 private key should decode");
}
}

64
crates/aether-provider-transport/src/vertex/context.rs
View File

@@ -27,28 +27,36 @@ pub fn looks_like_vertex_ai_host(base_url: &str) -> bool {
}
pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
{
if is_vertex_provider_type(transport) {
return resolve_local_auth_type_for_transport_format(transport)
.eq_ignore_ascii_case("api_key");
}
if !looks_like_vertex_ai_host(&transport.endpoint.base_url) {
return false;
}
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
if !endpoint_api_format.starts_with("gemini:") && !endpoint_api_format.starts_with("claude:") {
if !is_vertex_host_format_context(transport) {
return false;
}
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key")
}
pub fn is_vertex_service_account_transport_context(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
if !is_vertex_provider_type(transport) && !is_vertex_host_format_context(transport) {
return false;
}
matches!(
resolve_local_auth_type_for_transport_format(transport).as_str(),
"service_account" | "vertex_ai"
)
}
pub fn is_vertex_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
is_vertex_api_key_transport_context(transport)
|| is_vertex_service_account_transport_context(transport)
}
pub fn uses_vertex_api_key_query_auth(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
@@ -60,11 +68,28 @@ pub fn uses_vertex_api_key_query_auth(
.starts_with("gemini:")
}
fn is_vertex_provider_type(transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
}
fn is_vertex_host_format_context(transport: &GatewayProviderTransportSnapshot) -> bool {
if !looks_like_vertex_ai_host(&transport.endpoint.base_url) {
return false;
}
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
endpoint_api_format.starts_with("gemini:") || endpoint_api_format.starts_with("claude:")
}
#[cfg(test)]
mod tests {
use super::{
is_vertex_api_key_transport_context, looks_like_vertex_ai_host,
uses_vertex_api_key_query_auth,
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
is_vertex_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -150,6 +175,17 @@ mod tests {
assert!(!is_vertex_api_key_transport_context(&transport));
}
#[test]
fn infers_vertex_service_account_context_for_fixed_provider() {
let mut transport = sample_transport();
transport.provider.provider_type = "vertex_ai".to_string();
transport.key.auth_type = "service_account".to_string();
assert!(is_vertex_service_account_transport_context(&transport));
assert!(is_vertex_transport_context(&transport));
assert!(!is_vertex_api_key_transport_context(&transport));
}
#[test]
fn detects_vertex_query_auth_usage_for_gemini_formats() {
let transport = sample_transport();

13
crates/aether-provider-transport/src/vertex/mod.rs
View File

@@ -4,20 +4,29 @@ mod policy;
mod url;
pub use auth::{
resolve_local_vertex_api_key_query_auth, VertexApiKeyQueryAuth, VERTEX_API_KEY_QUERY_PARAM,
parse_vertex_service_account_auth_config, resolve_local_vertex_api_key_query_auth,
resolve_local_vertex_service_account_auth_config,
supports_local_vertex_service_account_auth_resolution, VertexApiKeyQueryAuth,
VertexServiceAccountAuthConfig, VertexServiceAccountRefreshAdapter,
VERTEX_API_KEY_QUERY_PARAM, VERTEX_SERVICE_ACCOUNT_AUTH_HEADER,
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
};
pub use context::{
is_vertex_api_key_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
is_vertex_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
};
pub use policy::{
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
local_vertex_gemini_transport_unsupported_reason_with_network,
supports_local_vertex_api_key_gemini_transport,
supports_local_vertex_api_key_gemini_transport_with_network,
supports_local_vertex_api_key_imagen_transport,
supports_local_vertex_api_key_imagen_transport_with_network,
supports_local_vertex_gemini_transport_with_network,
};
pub use url::{
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
build_vertex_service_account_gemini_content_url, resolve_vertex_service_account_region,
VERTEX_API_KEY_BASE_URL,
};

64
crates/aether-provider-transport/src/vertex/policy.rs
View File

@@ -1,10 +1,13 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::{
body_rules_are_locally_supported, header_rules_are_locally_supported,
supports_local_oauth_request_auth_resolution,
resolve_transport_profile, transport_profile_is_configured,
transport_proxy_is_locally_supported,
};
use super::auth::resolve_local_vertex_api_key_query_auth;
use super::auth::{
resolve_local_vertex_api_key_query_auth, supports_local_vertex_service_account_auth_resolution,
};
fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> bool {
transport
@@ -17,6 +20,19 @@ fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> b
pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> Option<&'static str> {
local_vertex_gemini_transport_unsupported_reason_with_network_impl(transport, true)
}
pub fn local_vertex_gemini_transport_unsupported_reason_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> Option<&'static str> {
local_vertex_gemini_transport_unsupported_reason_with_network_impl(transport, false)
}
fn local_vertex_gemini_transport_unsupported_reason_with_network_impl(
transport: &GatewayProviderTransportSnapshot,
require_api_key: bool,
) -> Option<&'static str> {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return if !transport.provider.is_active {
@@ -41,7 +57,15 @@ pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
if !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref()) {
return Some("transport_body_rules_unsupported");
}
if resolve_local_vertex_api_key_query_auth(transport).is_none() {
let has_api_key_auth = resolve_local_vertex_api_key_query_auth(transport).is_some();
let has_service_account_auth =
supports_local_vertex_service_account_auth_resolution(transport)
&& supports_local_oauth_request_auth_resolution(transport);
if require_api_key {
if !has_api_key_auth {
return Some("transport_auth_unavailable");
}
} else if !has_api_key_auth && !has_service_account_auth {
return Some("transport_auth_unavailable");
}
if !transport_proxy_is_locally_supported(transport) {
@@ -71,6 +95,12 @@ pub fn supports_local_vertex_api_key_gemini_transport_with_network(
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport).is_none()
}
pub fn supports_local_vertex_gemini_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
local_vertex_gemini_transport_unsupported_reason_with_network(transport).is_none()
}
pub fn supports_local_vertex_api_key_imagen_transport(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
@@ -159,8 +189,10 @@ mod tests {
use super::{
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
local_vertex_gemini_transport_unsupported_reason_with_network,
supports_local_vertex_api_key_gemini_transport,
supports_local_vertex_api_key_gemini_transport_with_network,
supports_local_vertex_gemini_transport_with_network,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
@@ -242,13 +274,35 @@ mod tests {
}
#[test]
fn rejects_vertex_service_account_subset() {
fn rejects_vertex_service_account_from_api_key_subset() {
let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
assert!(!supports_local_vertex_api_key_gemini_transport(&transport));
}
#[test]
fn supports_vertex_service_account_gemini_transport_with_network() {
let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"client_email":"svc@example.iam.gserviceaccount.com",
"private_key":"TEST-PRIVATE-KEY",
"project_id":"demo-project"
}"#
.to_string(),
);
assert!(!supports_local_vertex_api_key_gemini_transport_with_network(
&transport
));
assert!(supports_local_vertex_gemini_transport_with_network(
&transport
));
}
#[test]
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
let mut transport = sample_transport();
@@ -272,5 +326,9 @@ mod tests {
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(&transport),
Some("transport_auth_unavailable")
);
assert_eq!(
local_vertex_gemini_transport_unsupported_reason_with_network(&transport),
Some("transport_auth_unavailable")
);
}
}

177
crates/aether-provider-transport/src/vertex/url.rs
View File

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use url::form_urlencoded;
use super::super::url::build_passthrough_path_url;
use super::auth::VertexServiceAccountAuthConfig;
pub const VERTEX_API_KEY_BASE_URL: &str = "https://aiplatform.googleapis.com";
@@ -24,6 +25,15 @@ pub fn build_vertex_api_key_imagen_content_url(
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
}
pub fn build_vertex_service_account_gemini_content_url(
model: &str,
stream: bool,
auth_config: &VertexServiceAccountAuthConfig,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_service_account_google_model_url(model, stream, auth_config, request_query)
}
fn build_vertex_api_key_google_model_url(
model: &str,
stream: bool,
@@ -46,6 +56,89 @@ fn build_vertex_api_key_google_model_url(
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
}
fn build_vertex_service_account_google_model_url(
model: &str,
stream: bool,
auth_config: &VertexServiceAccountAuthConfig,
request_query: Option<&str>,
) -> Option<String> {
let trimmed_model = model.trim();
let project_id = auth_config.project_id.trim();
if trimmed_model.is_empty() || project_id.is_empty() {
return None;
}
let region = resolve_vertex_service_account_region(trimmed_model, auth_config);
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
let base_url = if region == "global" {
VERTEX_API_KEY_BASE_URL.to_string()
} else {
format!("https://{region}-aiplatform.googleapis.com")
};
let path = format!(
"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{trimmed_model}:{action}"
);
let merged_query = build_vertex_service_account_query(request_query, stream);
build_passthrough_path_url(&base_url, &path, merged_query.as_deref(), &[])
}
pub fn resolve_vertex_service_account_region(
model: &str,
auth_config: &VertexServiceAccountAuthConfig,
) -> String {
let trimmed_model = model.trim();
if let Some(region) = auth_config
.model_regions
.get(trimmed_model)
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
return region.to_string();
}
if let Some(region) = default_vertex_model_region(trimmed_model) {
return region.to_string();
}
if let Some(region) = auth_config
.region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return region.to_string();
}
"global".to_string()
}
fn default_vertex_model_region(model: &str) -> Option<&'static str> {
if model.starts_with("gemini-3.") || model == "gemini-3-pro-image-preview" {
return Some("global");
}
if matches!(
model,
"gemini-2.0-flash"
| "gemini-2.0-flash-exp"
| "gemini-2.0-flash-001"
| "gemini-2.0-pro-exp"
| "gemini-2.0-flash-exp-image-generation"
| "gemini-1.5-pro"
| "gemini-1.5-pro-001"
| "gemini-1.5-pro-002"
| "gemini-1.5-flash"
| "gemini-1.5-flash-001"
| "gemini-1.5-flash-002"
| "imagen-3.0-generate-001"
| "imagen-3.0-fast-generate-001"
) {
return Some("us-central1");
}
None
}
fn build_vertex_api_key_query(
api_key: &str,
request_query: Option<&str>,
@@ -73,6 +166,29 @@ fn build_vertex_api_key_query(
}
}
fn build_vertex_service_account_query(request_query: Option<&str>, stream: bool) -> Option<String> {
let mut merged = BTreeMap::new();
merge_query_string(&mut merged, request_query);
merged.remove("beta");
merged.remove("key");
if stream {
merged
.entry("alt".to_string())
.or_insert_with(|| "sse".to_string());
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in merged {
serializer.append_pair(&key, &value);
}
let query = serializer.finish();
if query.is_empty() {
None
} else {
Some(query)
}
}
fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
return;
@@ -85,7 +201,13 @@ fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
#[cfg(test)]
mod tests {
use super::{build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url};
use std::collections::BTreeMap;
use super::{
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
build_vertex_service_account_gemini_content_url,
};
use crate::vertex::VertexServiceAccountAuthConfig;
#[test]
fn builds_vertex_gemini_api_key_stream_url() {
@@ -118,4 +240,57 @@ mod tests {
)
);
}
#[test]
fn builds_vertex_service_account_gemini_sync_url() {
let auth_config = VertexServiceAccountAuthConfig {
client_email: "svc@example.iam.gserviceaccount.com".to_string(),
private_key: "not-used".to_string(),
project_id: "demo-project".to_string(),
token_uri: "https://oauth2.googleapis.com/token".to_string(),
region: None,
model_regions: BTreeMap::new(),
};
assert_eq!(
build_vertex_service_account_gemini_content_url(
"gemini-3.1-pro-preview",
false,
&auth_config,
Some("foo=bar&beta=1&key=client-key")
)
.as_deref(),
Some(
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent?foo=bar"
)
);
}
#[test]
fn builds_vertex_service_account_gemini_stream_url_with_model_region_override() {
let auth_config = VertexServiceAccountAuthConfig {
client_email: "svc@example.iam.gserviceaccount.com".to_string(),
private_key: "not-used".to_string(),
project_id: "demo-project".to_string(),
token_uri: "https://oauth2.googleapis.com/token".to_string(),
region: Some("global".to_string()),
model_regions: BTreeMap::from([(
"gemini-2.0-flash".to_string(),
"us-central1".to_string(),
)]),
};
assert_eq!(
build_vertex_service_account_gemini_content_url(
"gemini-2.0-flash",
true,
&auth_config,
Some("foo=bar")
)
.as_deref(),
Some(
"https://us-central1-aiplatform.googleapis.com/v1/projects/demo-project/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent?alt=sse&foo=bar"
)
);
}
}