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

fix(gateway): harden Gemini endpoint routing

Browse Source
This commit is contained in:
MMEXA
2026-05-18 00:53:34 +00:00
parent 81ff375bfd
commit b004a02e4a
20 changed files with 1348 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,8 +24,9 @@ pub use policy::{
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,
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
build_vertex_api_key_imagen_content_url, build_vertex_service_account_gemini_content_url,
build_vertex_service_account_gemini_embedding_url, resolve_vertex_service_account_region,
VERTEX_API_KEY_BASE_URL,
};

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

@@ -42,9 +42,12 @@ fn local_vertex_gemini_transport_unsupported_reason_with_network_impl(
Some("key_inactive")
};
}
if aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format)
!= "gemini:generate_content"
{
let endpoint_api_format =
aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format);
if !matches!(
endpoint_api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
) {
return Some("transport_api_format_mismatch");
}
if !is_vertex_transport_family(transport) {
@@ -299,6 +302,28 @@ mod tests {
));
}
#[test]
fn supports_vertex_service_account_gemini_embedding_transport_with_network() {
let mut transport = sample_transport();
transport.endpoint.api_format = "gemini:embedding".to_string();
transport.endpoint.endpoint_kind = Some("embedding".to_string());
transport.key.api_formats = Some(vec!["gemini:embedding".to_string()]);
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_gemini_transport_with_network(
&transport
));
}
#[test]
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
let mut transport = sample_transport();

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

@@ -13,7 +13,12 @@ pub fn build_vertex_api_key_gemini_content_url(
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
build_vertex_api_key_google_model_url(model, action, stream, api_key, request_query)
}
pub fn build_vertex_api_key_imagen_content_url(
@@ -22,7 +27,20 @@ pub fn build_vertex_api_key_imagen_content_url(
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
build_vertex_api_key_google_model_url(model, action, stream, api_key, request_query)
}
pub fn build_vertex_api_key_gemini_embedding_url(
model: &str,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_api_key_google_model_url(model, "embedContent", false, api_key, request_query)
}
pub fn build_vertex_service_account_gemini_content_url(
@@ -31,56 +49,69 @@ pub fn build_vertex_service_account_gemini_content_url(
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,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
let trimmed_model = model.trim();
let trimmed_api_key = api_key.trim();
if trimmed_model.is_empty() || trimmed_api_key.is_empty() {
return None;
}
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
let path = format!("/v1/publishers/google/models/{trimmed_model}:{action}");
build_vertex_service_account_google_model_url(model, action, stream, auth_config, request_query)
}
pub fn build_vertex_service_account_gemini_embedding_url(
model: &str,
auth_config: &VertexServiceAccountAuthConfig,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_service_account_google_model_url(
model,
"embedContent",
false,
auth_config,
request_query,
)
}
fn build_vertex_api_key_google_model_url(
model: &str,
action: &str,
stream: bool,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
let trimmed_model = model.trim();
let trimmed_action = action.trim();
let trimmed_api_key = api_key.trim();
if trimmed_model.is_empty() || trimmed_action.is_empty() || trimmed_api_key.is_empty() {
return None;
}
let path = format!("/v1/publishers/google/models/{trimmed_model}:{trimmed_action}");
let merged_query = build_vertex_api_key_query(trimmed_api_key, request_query, stream);
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
}
fn build_vertex_service_account_google_model_url(
model: &str,
action: &str,
stream: bool,
auth_config: &VertexServiceAccountAuthConfig,
request_query: Option<&str>,
) -> Option<String> {
let trimmed_model = model.trim();
let trimmed_action = action.trim();
let project_id = auth_config.project_id.trim();
if trimmed_model.is_empty() || project_id.is_empty() {
if trimmed_model.is_empty() || trimmed_action.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}"
"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{trimmed_model}:{trimmed_action}"
);
let merged_query = build_vertex_service_account_query(request_query, stream);
build_passthrough_path_url(&base_url, &path, merged_query.as_deref(), &[])