修复 Vertex AI 服务账号访问

This commit is contained in:
Codex
2026-05-10 15:51:38 +08:00
parent 7f101431c5
commit 7e804c408f
20 changed files with 851 additions and 42 deletions

2
Cargo.lock generated
View File

@@ -280,9 +280,11 @@ dependencies = [
"aether-video-tasks-core", "aether-video-tasks-core",
"async-trait", "async-trait",
"axum", "axum",
"base64 0.22.1",
"http", "http",
"regex", "regex",
"reqwest", "reqwest",
"rsa",
"serde", "serde",
"serde_json", "serde_json",
"sha2", "sha2",

View File

@@ -6,7 +6,9 @@ use crate::ai_serving::planner::candidate_preparation::{
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate; use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata; use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_serving::transport::kiro::KiroRequestAuth; use crate::ai_serving::transport::kiro::KiroRequestAuth;
use crate::ai_serving::transport::vertex::resolve_local_vertex_api_key_query_auth; use crate::ai_serving::transport::vertex::{
is_vertex_api_key_transport_context, resolve_local_vertex_api_key_query_auth,
};
use crate::ai_serving::transport::SameFormatProviderRequestBehavior; use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
use crate::ai_serving::{ use crate::ai_serving::{
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
@@ -134,7 +136,10 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
return None; return None;
} }
}; };
if behavior.is_vertex && vertex_query_auth.is_none() { if behavior.is_vertex
&& is_vertex_api_key_transport_context(&transport)
&& vertex_query_auth.is_none()
{
super::super::payload::mark_skipped_local_same_format_provider_candidate( super::super::payload::mark_skipped_local_same_format_provider_candidate(
state, state,
input, input,

View File

@@ -362,8 +362,8 @@ fn provider_query_transport_supports_standard_test_execution(
) )
} }
"gemini:generate_content" => { "gemini:generate_content" => {
if crate::provider_transport::is_vertex_api_key_transport_context(transport) { if crate::provider_transport::is_vertex_transport_context(transport) {
aether_provider_transport::vertex::supports_local_vertex_api_key_gemini_transport_with_network(transport) aether_provider_transport::vertex::supports_local_vertex_gemini_transport_with_network(transport)
} else { } else {
state.supports_local_gemini_transport_with_network(transport, api_format) state.supports_local_gemini_transport_with_network(transport, api_format)
} }

View File

@@ -194,7 +194,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
} }
let oauth_auth = match format_value.as_str() { let oauth_auth = match format_value.as_str() {
"openai:chat" | "claude:messages" => { "openai:chat" | "claude:messages" | "gemini:generate_content" => {
match state.resolve_local_oauth_request_auth(&transport).await { match state.resolve_local_oauth_request_auth(&transport).await {
Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header { Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header {
name, name,
@@ -217,16 +217,26 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
} }
"gemini:generate_content" => { "gemini:generate_content" => {
crate::provider_transport::auth::resolve_local_gemini_auth(&transport) crate::provider_transport::auth::resolve_local_gemini_auth(&transport)
.or(oauth_auth.clone())
} }
_ => None, _ => None,
}; };
let Some((auth_header, auth_value)) = auth else {
return None;
};
let uses_vertex_query_auth = crate::provider_transport::uses_vertex_api_key_query_auth( let uses_vertex_query_auth = crate::provider_transport::uses_vertex_api_key_query_auth(
&transport, &transport,
format_value.as_str(), format_value.as_str(),
); );
let vertex_query_auth = if uses_vertex_query_auth {
crate::provider_transport::vertex::resolve_local_vertex_api_key_query_auth(&transport)
} else {
None
};
let (auth_header, auth_value) = match auth {
Some((auth_header, auth_value)) => (auth_header, auth_value),
None if uses_vertex_query_auth && vertex_query_auth.is_some() => {
(String::new(), String::new())
}
None => return None,
};
let upstream_url = crate::provider_transport::build_transport_request_url( let upstream_url = crate::provider_transport::build_transport_request_url(
&transport, &transport,

View File

@@ -57,6 +57,17 @@ fn oauth_access_token_expired(expires_at_unix_secs: Option<u64>, now_unix_secs:
expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs) expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
} }
fn local_oauth_refresh_entry_should_stay_memory_only(
transport: &provider_transport::GatewayProviderTransportSnapshot,
entry: &provider_transport::CachedOAuthEntry,
) -> bool {
entry
.provider_type
.trim()
.eq_ignore_ascii_case(provider_transport::vertex::VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE)
&& provider_transport::is_vertex_service_account_transport_context(transport)
}
fn oauth_auth_config_refresh_token_fingerprint(auth_config: Option<&str>) -> Option<String> { fn oauth_auth_config_refresh_token_fingerprint(auth_config: Option<&str>) -> Option<String> {
let parsed = auth_config let parsed = auth_config
.map(str::trim) .map(str::trim)
@@ -1094,6 +1105,17 @@ impl AppState {
return Ok(()); return Ok(());
} }
if local_oauth_refresh_entry_should_stay_memory_only(transport, entry) {
tracing::info!(
key_id = %key_id,
provider_id = %transport.provider.id,
provider_type = %transport.provider.provider_type,
expires_at_unix_secs = ?entry.expires_at_unix_secs,
"gateway local oauth refresh entry kept in memory only"
);
return Ok(());
}
let Some(encryption_key) = self.data.encryption_key() else { let Some(encryption_key) = self.data.encryption_key() else {
return Ok(()); return Ok(());
}; };
@@ -1703,4 +1725,71 @@ mod tests {
Some("[OAUTH_EXPIRED] access token invalid".to_string()), Some("[OAUTH_EXPIRED] access token invalid".to_string()),
); );
} }
#[test]
fn vertex_service_account_refresh_entry_stays_memory_only() {
let transport = crate::provider_transport::GatewayProviderTransportSnapshot {
provider: crate::provider_transport::snapshot::GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Vertex".to_string(),
provider_type: "vertex_ai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: crate::provider_transport::snapshot::GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "gemini:generate_content".to_string(),
api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://aiplatform.googleapis.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: crate::provider_transport::snapshot::GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "Gemini".to_string(),
auth_type: "service_account".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some("{\"project_id\":\"demo\"}".to_string()),
},
};
let entry = crate::provider_transport::CachedOAuthEntry {
provider_type: "vertex_ai".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer access-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
};
assert!(super::local_oauth_refresh_entry_should_stay_memory_only(
&transport, &entry
));
}
} }

View File

@@ -0,0 +1,2 @@
-- Historical compatibility no-op.
-- The auth module tables are already present in the MySQL baseline.

View File

@@ -0,0 +1,3 @@
-- Historical compatibility no-op.
-- The auth module tables are already present in the squashed baseline and
-- 20260505000000_sync_core_export_columns for existing Postgres deployments.

View File

@@ -0,0 +1,2 @@
-- Historical compatibility no-op.
-- The auth module tables are already present in the SQLite baseline.

View File

@@ -8,6 +8,7 @@ use aether_provider_transport::{
}; };
use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD}; use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD};
use base64::Engine as _; use base64::Engine as _;
use rsa::pkcs1::DecodeRsaPrivateKey;
use rsa::pkcs1v15::SigningKey; use rsa::pkcs1v15::SigningKey;
use rsa::pkcs8::DecodePrivateKey; use rsa::pkcs8::DecodePrivateKey;
use rsa::signature::{SignatureEncoding, Signer}; use rsa::signature::{SignatureEncoding, Signer};
@@ -675,8 +676,7 @@ fn build_vertex_service_account_assertion(
.map_err(|err| format!("vertex_ai(service_account): jwt payload encode failed: {err}"))?, .map_err(|err| format!("vertex_ai(service_account): jwt payload encode failed: {err}"))?,
); );
let message = format!("{header}.{payload}"); let message = format!("{header}.{payload}");
let private_key = RsaPrivateKey::from_pkcs8_pem(private_key_pem) let private_key = decode_vertex_service_account_private_key(private_key_pem)?;
.map_err(|err| format!("vertex_ai(service_account): private_key parse failed: {err}"))?;
let signing_key = SigningKey::<Sha256>::new(private_key); let signing_key = SigningKey::<Sha256>::new(private_key);
let signature = signing_key.sign(message.as_bytes()); let signature = signing_key.sign(message.as_bytes());
Ok(format!( Ok(format!(
@@ -685,6 +685,19 @@ fn build_vertex_service_account_assertion(
)) ))
} }
fn decode_vertex_service_account_private_key(
private_key_pem: &str,
) -> Result<RsaPrivateKey, String> {
match RsaPrivateKey::from_pkcs8_pem(private_key_pem) {
Ok(private_key) => Ok(private_key),
Err(pkcs8_err) => RsaPrivateKey::from_pkcs1_pem(private_key_pem).map_err(|pkcs1_err| {
format!(
"vertex_ai(service_account): private_key parse failed: pkcs8: {pkcs8_err}; pkcs1: {pkcs1_err}"
)
}),
}
}
fn execution_result_json_body(result: &ExecutionResult) -> Result<Value, String> { fn execution_result_json_body(result: &ExecutionResult) -> Result<Value, String> {
if result.status_code != 200 { if result.status_code != 200 {
return Err(execution_result_error_message(result)); return Err(execution_result_error_message(result));

View File

@@ -15,12 +15,14 @@ aether-oauth.workspace = true
aether-runtime-state.workspace = true aether-runtime-state.workspace = true
aether-video-tasks-core.workspace = true aether-video-tasks-core.workspace = true
async-trait.workspace = true async-trait.workspace = true
base64.workspace = true
http.workspace = true http.workspace = true
regex.workspace = true regex.workspace = true
reqwest.workspace = true reqwest.workspace = true
rsa = "0.9.10"
serde.workspace = true serde.workspace = true
serde_json.workspace = true serde_json.workspace = true
sha2.workspace = true sha2 = { workspace = true, features = ["oid"] }
thiserror.workspace = true thiserror.workspace = true
tokio.workspace = true tokio.workspace = true
tracing.workspace = true tracing.workspace = true

View File

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

View File

@@ -103,7 +103,10 @@ pub use standard::{
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput, StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput, 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::{ pub use video::{
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url, build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport, reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,

View File

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

View File

@@ -15,7 +15,8 @@ use crate::url::{
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path, build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
}; };
use crate::vertex::{ 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)] #[derive(Debug, Clone, Copy)]
@@ -256,6 +257,14 @@ fn build_transport_hook_url(
params.request_query, 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) { 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] #[test]
fn builds_openai_responses_url_for_formal_format_name() { fn builds_openai_responses_url_for_formal_format_name() {
let transport = sample_transport( let transport = sample_transport(

View File

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

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; use super::super::snapshot::GatewayProviderTransportSnapshot;
pub const VERTEX_API_KEY_QUERY_PARAM: &str = "key"; 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)] #[derive(Debug, Clone, PartialEq, Eq)]
pub struct VertexApiKeyQueryAuth { pub struct VertexApiKeyQueryAuth {
@@ -8,6 +31,16 @@ pub struct VertexApiKeyQueryAuth {
pub value: String, 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( pub fn resolve_local_vertex_api_key_query_auth(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
) -> Option<VertexApiKeyQueryAuth> { ) -> 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)] #[cfg(test)]
mod tests { mod tests {
use super::super::super::snapshot::{ use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, 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 { fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot { GatewayProviderTransportSnapshot {
@@ -136,4 +429,57 @@ mod tests {
.expect("custom aiplatform transport should resolve"); .expect("custom aiplatform transport should resolve");
assert_eq!(auth.value, "vertex-secret"); 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");
}
} }

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 { pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
if transport if is_vertex_provider_type(transport) {
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
{
return resolve_local_auth_type_for_transport_format(transport) return resolve_local_auth_type_for_transport_format(transport)
.eq_ignore_ascii_case("api_key"); .eq_ignore_ascii_case("api_key");
} }
if !looks_like_vertex_ai_host(&transport.endpoint.base_url) { if !is_vertex_host_format_context(transport) {
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:") {
return false; return false;
} }
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key") 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( pub fn uses_vertex_api_key_query_auth(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str, provider_api_format: &str,
@@ -60,11 +68,28 @@ pub fn uses_vertex_api_key_query_auth(
.starts_with("gemini:") .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)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
is_vertex_api_key_transport_context, looks_like_vertex_ai_host, is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
uses_vertex_api_key_query_auth, is_vertex_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
}; };
use crate::snapshot::{ use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -150,6 +175,17 @@ mod tests {
assert!(!is_vertex_api_key_transport_context(&transport)); 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] #[test]
fn detects_vertex_query_auth_usage_for_gemini_formats() { fn detects_vertex_query_auth_usage_for_gemini_formats() {
let transport = sample_transport(); let transport = sample_transport();

View File

@@ -4,20 +4,29 @@ mod policy;
mod url; mod url;
pub use auth::{ 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::{ 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::{ pub use policy::{
local_vertex_api_key_gemini_transport_unsupported_reason_with_network, 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,
supports_local_vertex_api_key_gemini_transport_with_network, supports_local_vertex_api_key_gemini_transport_with_network,
supports_local_vertex_api_key_imagen_transport, supports_local_vertex_api_key_imagen_transport,
supports_local_vertex_api_key_imagen_transport_with_network, supports_local_vertex_api_key_imagen_transport_with_network,
supports_local_vertex_gemini_transport_with_network,
}; };
pub use url::{ pub use url::{
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_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, VERTEX_API_KEY_BASE_URL,
}; };

View File

@@ -1,10 +1,13 @@
use super::super::snapshot::GatewayProviderTransportSnapshot; use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::{ use super::super::{
body_rules_are_locally_supported, header_rules_are_locally_supported, body_rules_are_locally_supported, header_rules_are_locally_supported,
supports_local_oauth_request_auth_resolution,
resolve_transport_profile, transport_profile_is_configured, resolve_transport_profile, transport_profile_is_configured,
transport_proxy_is_locally_supported, 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 { fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> bool {
transport 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( pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
transport: &GatewayProviderTransportSnapshot, 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> { ) -> Option<&'static str> {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active { if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return if !transport.provider.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()) { if !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref()) {
return Some("transport_body_rules_unsupported"); 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"); return Some("transport_auth_unavailable");
} }
if !transport_proxy_is_locally_supported(transport) { 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() 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( pub fn supports_local_vertex_api_key_imagen_transport(
transport: &GatewayProviderTransportSnapshot, transport: &GatewayProviderTransportSnapshot,
) -> bool { ) -> bool {
@@ -159,8 +189,10 @@ mod tests {
use super::{ use super::{
local_vertex_api_key_gemini_transport_unsupported_reason_with_network, 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,
supports_local_vertex_api_key_gemini_transport_with_network, supports_local_vertex_api_key_gemini_transport_with_network,
supports_local_vertex_gemini_transport_with_network,
}; };
fn sample_transport() -> GatewayProviderTransportSnapshot { fn sample_transport() -> GatewayProviderTransportSnapshot {
@@ -242,13 +274,35 @@ mod tests {
} }
#[test] #[test]
fn rejects_vertex_service_account_subset() { fn rejects_vertex_service_account_from_api_key_subset() {
let mut transport = sample_transport(); let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string(); transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string()); transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
assert!(!supports_local_vertex_api_key_gemini_transport(&transport)); 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] #[test]
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() { fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
let mut transport = sample_transport(); let mut transport = sample_transport();
@@ -272,5 +326,9 @@ mod tests {
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(&transport), local_vertex_api_key_gemini_transport_unsupported_reason_with_network(&transport),
Some("transport_auth_unavailable") Some("transport_auth_unavailable")
); );
assert_eq!(
local_vertex_gemini_transport_unsupported_reason_with_network(&transport),
Some("transport_auth_unavailable")
);
} }
} }

View File

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use url::form_urlencoded; use url::form_urlencoded;
use super::super::url::build_passthrough_path_url; use super::super::url::build_passthrough_path_url;
use super::auth::VertexServiceAccountAuthConfig;
pub const VERTEX_API_KEY_BASE_URL: &str = "https://aiplatform.googleapis.com"; 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) 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( fn build_vertex_api_key_google_model_url(
model: &str, model: &str,
stream: bool, 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(), &[]) 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( fn build_vertex_api_key_query(
api_key: &str, api_key: &str,
request_query: Option<&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>) { 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 { let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
return; return;
@@ -85,7 +201,13 @@ fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
#[cfg(test)] #[cfg(test)]
mod tests { 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] #[test]
fn builds_vertex_gemini_api_key_stream_url() { 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"
)
);
}
} }