fix(gateway): harden Gemini endpoint routing

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

View File

@@ -66,6 +66,10 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
extensions: Default::default(),
});
}
outputs.retain(gemini_response_output_has_visible_content);
if outputs.is_empty() {
return None;
}
let content = outputs
.first()
.map(|output| output.content.clone())
@@ -108,6 +112,18 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
Some(canonical)
}
fn gemini_response_output_has_visible_content(output: &CanonicalResponseOutput) -> bool {
output.content.iter().any(|block| match block {
CanonicalContentBlock::Text { text, .. } => !text.trim().is_empty(),
CanonicalContentBlock::ToolUse { .. }
| CanonicalContentBlock::ToolResult { .. }
| CanonicalContentBlock::Image { .. }
| CanonicalContentBlock::File { .. }
| CanonicalContentBlock::Audio { .. } => true,
CanonicalContentBlock::Thinking { .. } | CanonicalContentBlock::Unknown { .. } => false,
})
}
pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value) -> Option<Value> {
let mut response = canonical_to_gemini_response(canonical, report_context)?;
if let Some(object) = response.as_object_mut() {
@@ -353,3 +369,72 @@ fn canonical_usage_to_gemini_usage_metadata(usage: &CanonicalUsage) -> Value {
}
Value::Object(out)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::CanonicalContentBlock;
#[test]
fn gemini_response_without_visible_parts_is_not_success() {
let body = json!({
"candidates": [{
"content": {"role": "model"},
"finishReason": "MAX_TOKENS"
}],
"usageMetadata": {
"promptTokenCount": 8,
"candidatesTokenCount": 1,
"thoughtsTokenCount": 25,
"totalTokenCount": 34
},
"modelVersion": "gemini-3-flash-preview",
"responseId": "resp-empty"
});
assert!(from_raw(&body).is_none());
}
#[test]
fn gemini_response_with_only_thought_parts_is_not_success() {
let body = json!({
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "hidden plan", "thought": true}]
},
"finishReason": "MAX_TOKENS"
}],
"modelVersion": "gemini-3-flash-preview",
"responseId": "resp-thought-only"
});
assert!(from_raw(&body).is_none());
}
#[test]
fn gemini_response_with_function_call_is_visible_output() {
let body = json!({
"candidates": [{
"content": {
"role": "model",
"parts": [{
"functionCall": {
"name": "lookup",
"args": {"query": "weather"}
}
}]
},
"finishReason": "STOP"
}],
"modelVersion": "gemini-3-flash-preview",
"responseId": "resp-tool"
});
let canonical = from_raw(&body).expect("function call should be visible output");
assert!(matches!(
canonical.content.first(),
Some(CanonicalContentBlock::ToolUse { name, .. }) if name == "lookup"
));
}
}

View File

@@ -317,11 +317,15 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",

View File

@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",

View File

@@ -113,11 +113,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($3) = 'gemini:generate_content'
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -296,11 +296,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($4) = 'gemini:generate_content'
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -478,11 +478,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($6) = 'gemini:generate_content'
AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -1287,6 +1287,22 @@ mod tests {
}
}
#[test]
fn candidate_selection_sql_allows_vertex_embedding_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'"));
assert!(sql.contains("gemini:embedding"));
assert!(sql.contains("gemini:generate_content"));
assert!(sql.contains("claude:messages"));
}
}
#[test]
fn requested_model_selection_page_sql_adds_limit_and_offset() {
let sql = requested_model_selection_page_sql();

View File

@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",

View File

@@ -423,7 +423,7 @@ mod tests {
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
request_conversion_direct_auth, request_conversion_enabled_for_transport,
request_conversion_transport_supported, request_pair_allowed_for_transport,
CandidateTransportPolicyFacts,
request_pair_direct_auth, CandidateTransportPolicyFacts,
};
use aether_ai_formats::formats::matrix::RequestConversionKind;
use serde_json::json;
@@ -593,6 +593,21 @@ mod tests {
);
}
#[test]
fn vertex_gemini_embedding_transport_supports_openai_embedding_conversion() {
let transport = transport_snapshot("vertex_ai", "gemini:embedding", "api_key", true, None);
assert!(request_pair_allowed_for_transport(
&transport,
"openai:embedding",
"gemini:embedding"
));
assert_eq!(
request_pair_direct_auth(&transport, "gemini:embedding"),
Some(("key".to_string(), "secret".to_string()))
);
}
#[test]
fn kiro_claude_messages_transport_supports_cross_format_conversion_via_envelope() {
let transport = transport_snapshot("kiro", "claude:messages", "bearer", true, None);

View File

@@ -47,6 +47,7 @@ pub enum ProviderApiFormatInheritance {
None,
OAuth,
OAuthOrBearer,
OAuthOrServiceAccount,
OAuthOrConfiguredBearer,
}
@@ -61,6 +62,9 @@ impl ProviderApiFormatInheritance {
Self::None => false,
Self::OAuth => auth_type == "oauth",
Self::OAuthOrBearer => auth_type == "oauth" || auth_type == "bearer",
Self::OAuthOrServiceAccount => {
auth_type == "oauth" || auth_type == "service_account" || auth_type == "vertex_ai"
}
Self::OAuthOrConfiguredBearer => {
auth_type == "oauth"
|| auth_type == "bearer"
@@ -221,11 +225,12 @@ const GEMINI_CLI_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
};
const VERTEX_AI_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
fixed_provider: true,
api_format_inheritance: ProviderApiFormatInheritance::OAuth,
api_format_inheritance: ProviderApiFormatInheritance::OAuthOrServiceAccount,
enable_format_conversion_by_default: true,
supports_model_fetch: false,
supports_local_openai_chat_transport: false,
supports_local_same_format_transport: false,
local_embedding_support: ProviderLocalEmbeddingSupport::Gemini,
..STANDARD_RUNTIME_POLICY
};
const ANTIGRAVITY_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
@@ -329,6 +334,12 @@ const VERTEX_AI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTe
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "gemini:embedding",
api_format: "gemini:embedding",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "claude:messages",
api_format: "claude:messages",
@@ -613,6 +624,11 @@ mod tests {
"bearer",
Some("{}")
));
assert!(fixed_provider_key_inherits_api_formats(
"vertex_ai",
"service_account",
None
));
assert!(!fixed_provider_key_inherits_api_formats(
"kiro", "bearer", None
));
@@ -681,6 +697,7 @@ mod tests {
("custom", "openai:embedding"),
("gemini", "gemini:embedding"),
("google", "gemini:embedding"),
("vertex_ai", "gemini:embedding"),
("jina", "jina:embedding"),
("doubao", "doubao:embedding"),
("volcengine", "doubao:embedding"),
@@ -694,6 +711,7 @@ mod tests {
for (provider_type, api_format) in [
("openai", "gemini:embedding"),
("gemini", "openai:embedding"),
("vertex_ai", "openai:embedding"),
("jina", "doubao:embedding"),
("doubao", "jina:embedding"),
("claude_code", "openai:embedding"),
@@ -710,4 +728,28 @@ mod tests {
"GEMINI:EMBEDDING"
));
}
#[test]
fn vertex_fixed_provider_template_includes_gemini_embedding_endpoint() {
let template =
fixed_provider_template("vertex_ai").expect("vertex_ai template should exist");
assert_eq!(
template
.endpoints
.iter()
.map(|item| item.api_format)
.collect::<Vec<_>>(),
vec![
"gemini:generate_content",
"gemini:embedding",
"claude:messages",
]
);
assert!(
fixed_provider_endpoint_template_by_api_format("vertex_ai", "gemini:embedding")
.is_some()
);
}
}

View File

@@ -16,7 +16,9 @@ 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, build_vertex_service_account_gemini_content_url,
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
build_vertex_service_account_gemini_content_url,
build_vertex_service_account_gemini_embedding_url, is_vertex_transport_context,
resolve_local_vertex_api_key_query_auth, resolve_local_vertex_service_account_auth_config,
};
@@ -62,13 +64,20 @@ fn build_transport_request_url_inner(
params: TransportRequestUrlParams<'_>,
gemini_embedding_batch: bool,
) -> Option<String> {
let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase();
let normalized_provider_api_format =
aether_ai_formats::normalize_api_format_alias(&provider_api_format);
if normalized_provider_api_format == "gemini:embedding"
&& gemini_embedding_batch
&& is_vertex_transport_context(transport)
{
return None;
}
if let Some(url) = build_transport_hook_url(transport, params) {
return Some(url);
}
let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase();
let normalized_provider_api_format =
aether_ai_formats::normalize_api_format_alias(&provider_api_format);
let custom_path = transport
.endpoint
.custom_path
@@ -281,25 +290,42 @@ fn build_transport_hook_url(
));
}
if aether_ai_formats::normalize_api_format_alias(params.provider_api_format)
== "gemini:generate_content"
{
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
return build_vertex_api_key_gemini_content_url(
params.mapped_model?,
params.upstream_is_stream,
&auth.value,
params.request_query,
);
match aether_ai_formats::normalize_api_format_alias(params.provider_api_format).as_str() {
"gemini:generate_content" => {
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
return build_vertex_api_key_gemini_content_url(
params.mapped_model?,
params.upstream_is_stream,
&auth.value,
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 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,
);
"gemini:embedding" => {
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
return build_vertex_api_key_gemini_embedding_url(
params.mapped_model?,
&auth.value,
params.request_query,
);
}
if let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) {
return build_vertex_service_account_gemini_embedding_url(
params.mapped_model?,
&auth_config,
params.request_query,
);
}
}
_ => {}
}
if is_antigravity_provider_transport(transport) {
@@ -629,6 +655,91 @@ mod tests {
);
}
#[test]
fn uses_vertex_service_account_hook_for_gemini_embedding_url() {
let mut transport = sample_transport(
"vertex_ai",
"gemini:embedding",
"https://aiplatform.googleapis.com",
None,
);
transport.endpoint.endpoint_kind = Some("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(),
);
let provider_request_body = json!({
"content": {"parts": [{"text": "hello"}]}
});
let url = build_transport_request_url_for_request_body(
&transport,
TransportRequestUrlParams {
provider_api_format: "gemini:embedding",
mapped_model: Some("gemini-embedding-2"),
upstream_is_stream: false,
request_query: Some("foo=bar&beta=1"),
kiro_api_region: None,
},
Some(&provider_request_body),
)
.expect("vertex embedding service account hook url");
assert_eq!(
url,
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-embedding-2:embedContent?foo=bar"
);
}
#[test]
fn vertex_gemini_embedding_batch_request_does_not_use_gemini_api_batch_endpoint() {
let mut transport = sample_transport(
"vertex_ai",
"gemini:embedding",
"https://aiplatform.googleapis.com",
None,
);
transport.endpoint.endpoint_kind = Some("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(),
);
let batch_body = json!({
"requests": [
{
"model": "models/gemini-embedding-2",
"content": {"parts": [{"text": "alpha"}]}
}
]
});
assert!(build_transport_request_url_for_request_body(
&transport,
TransportRequestUrlParams {
provider_api_format: "gemini:embedding",
mapped_model: Some("gemini-embedding-2"),
upstream_is_stream: false,
request_query: None,
kiro_api_region: None,
},
Some(&batch_body),
)
.is_none());
}
#[test]
fn builds_openai_responses_url_for_formal_format_name() {
let transport = sample_transport(

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,
};

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();

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(), &[])