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

fix(provider): 将rust分支的gemini cli端点行为对齐到python分支 (#321)

Browse Source 
* fix(provider): 对齐 Vertex/Gemini 上游发包与 Python master

- provider-transport: 为 custom+aiplatform 推断 Vertex API key 上下文并统一 URL 构建顺序，复用共享 request_url 构建最终上游地址
- ai-pipeline/gateway: Vertex Gemini 路径改为仅使用 URL query key，不再向上游附带 x-goog-api-key header；同步对齐 standard/admin/test-connection/runtime miss 摘要中的最终 URL
- gemini conversion: 按 Python master 输出 Gemini 请求体，补齐 system_instruction / generation_config / tool_config / function_declarations 形态，并移植 Gemini schema 清洗逻辑
- scheduler/executor: 将最终 upstream_url、mapped_model、key_name 写入候选 extra_data，运行时 miss 诊断优先展示真实展开后的上游 URL 便于服务器排障

* fix(provider): 修复 Vertex provider 测试与本地调度链路

* fix(provider): 对齐 Vertex 本地执行与 Rust CI
This commit is contained in:
Entropy.Xu
2026-04-23 23:01:06 +08:00
committed by GitHub
parent ccec46eddd
commit 0f94f92c37
29 changed files with 1879 additions and 497 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
apps/aether-gateway/src/ai_pipeline/mod.rs
View File

@@ -48,6 +48,26 @@ pub(crate) use self::pure::*;
pub(crate) use crate::control::GatewayControlDecision;
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
pub(crate) fn build_provider_transport_request_url(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
mapped_model: Option<&str>,
upstream_is_stream: bool,
request_query: Option<&str>,
kiro_api_region: Option<&str>,
) -> Option<String> {
crate::provider_transport::build_transport_request_url(
transport,
crate::provider_transport::TransportRequestUrlParams {
provider_api_format,
mapped_model,
upstream_is_stream,
request_query,
kiro_api_region,
},
)
}
pub(crate) async fn resolve_execution_runtime_auth_context(
state: &AppState,
decision: &GatewayControlDecision,

11
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/request/policy.rs
View File

@@ -6,7 +6,10 @@ use crate::ai_pipeline::transport::policy::{
local_gemini_transport_unsupported_reason_with_network,
local_standard_transport_unsupported_reason_with_network,
};
use crate::ai_pipeline::transport::vertex::local_vertex_api_key_gemini_transport_unsupported_reason_with_network;
use crate::ai_pipeline::transport::vertex::{
is_vertex_api_key_transport_context,
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
};
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use super::super::LocalSameFormatProviderFamily;
@@ -34,11 +37,7 @@ pub(super) fn classify_same_format_provider_request_behavior(
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code");
let is_vertex = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai");
let is_vertex = is_vertex_api_key_transport_context(transport);
let is_kiro = transport
.provider
.provider_type

135
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/request/url.rs
View File

@@ -1,16 +1,6 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use super::super::{
build_antigravity_v1internal_url, build_claude_code_messages_url, build_claude_messages_url,
build_gemini_content_url, build_kiro_generate_assistant_response_url,
build_passthrough_path_url, build_vertex_api_key_gemini_content_url,
resolve_local_vertex_api_key_query_auth, AntigravityRequestUrlAction,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
use super::super::LocalSameFormatProviderSpec;
pub(crate) fn build_same_format_upstream_url(
parts: &http::request::Parts,
@@ -20,119 +10,16 @@ pub(crate) fn build_same_format_upstream_url(
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_pipeline::transport::kiro::KiroRequestAuth>,
) -> Option<String> {
if let Some(kiro_auth) = kiro_auth {
return build_kiro_generate_assistant_response_url(
&transport.endpoint.base_url,
parts.uri.query(),
Some(kiro_auth.auth_config.effective_api_region()),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code")
{
return Some(build_claude_code_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
));
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("vertex_ai")
{
let auth = resolve_local_vertex_api_key_query_auth(transport)?;
return build_vertex_api_key_gemini_content_url(
mapped_model,
upstream_is_stream,
&auth.value,
parts.uri.query(),
);
}
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("antigravity")
{
let query = parts.uri.query().map(|query| {
form_urlencoded::parse(query.as_bytes())
.into_owned()
.collect::<BTreeMap<String, String>>()
});
return build_antigravity_v1internal_url(
&transport.endpoint.base_url,
if upstream_is_stream {
AntigravityRequestUrlAction::StreamGenerateContent
} else {
AntigravityRequestUrlAction::GenerateContent
},
query.as_ref(),
);
}
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(path) = custom_path {
let blocked_keys = match spec.family {
LocalSameFormatProviderFamily::Standard => &[][..],
LocalSameFormatProviderFamily::Gemini => &["key"][..],
};
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
blocked_keys,
)?;
return Some(maybe_add_gemini_stream_alt_sse(url, spec));
}
let url = match spec.family {
LocalSameFormatProviderFamily::Standard => Some(build_claude_messages_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
LocalSameFormatProviderFamily::Gemini => build_gemini_content_url(
&transport.endpoint.base_url,
mapped_model,
spec.require_streaming,
parts.uri.query(),
),
}?;
Some(maybe_add_gemini_stream_alt_sse(url, spec))
maybe_add_gemini_stream_alt_sse(crate::ai_pipeline::build_provider_transport_request_url(
transport,
spec.api_format,
Some(mapped_model),
upstream_is_stream,
parts.uri.query(),
kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
))
}
fn maybe_add_gemini_stream_alt_sse(
upstream_url: String,
spec: LocalSameFormatProviderSpec,
) -> String {
if spec.family != LocalSameFormatProviderFamily::Gemini || !spec.require_streaming {
return upstream_url;
}
let has_alt = upstream_url
.split_once('?')
.map(|(_, query)| {
form_urlencoded::parse(query.as_bytes())
.any(|(key, _)| key.as_ref().eq_ignore_ascii_case("alt"))
})
.unwrap_or(false);
if has_alt {
return upstream_url;
}
if upstream_url.contains('?') {
format!("{upstream_url}&alt=sse")
} else {
format!("{upstream_url}?alt=sse")
}
fn maybe_add_gemini_stream_alt_sse(url: Option<String>) -> Option<String> {
url
}

32
apps/aether-gateway/src/ai_pipeline/planner/standard/family/request.rs
View File

@@ -13,6 +13,7 @@ use crate::ai_pipeline::transport::apply_local_header_rules;
use crate::ai_pipeline::transport::auth::{
build_claude_passthrough_headers, build_openai_passthrough_headers, ensure_upstream_auth_header,
};
use crate::ai_pipeline::transport::vertex::uses_vertex_api_key_query_auth;
use crate::ai_pipeline::GatewayProviderTransportSnapshot;
use crate::AppState;
@@ -155,6 +156,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return None;
}
};
let uses_vertex_query_auth = uses_vertex_api_key_query_auth(transport, provider_api_format);
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
@@ -173,10 +175,15 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
Some("application/json"),
)
};
let protected_headers = if uses_vertex_query_auth {
&["content-type"][..]
} else {
&[prepared_candidate.auth_header.as_str(), "content-type"][..]
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&prepared_candidate.auth_header, "content-type"],
protected_headers,
&provider_request_body,
Some(body_json),
) {
@@ -201,11 +208,20 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
let (auth_header, auth_value) = if uses_vertex_query_auth {
provider_request_headers.remove("x-goog-api-key");
(String::new(), String::new())
} else {
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
(
prepared_candidate.auth_header.clone(),
prepared_candidate.auth_value.clone(),
)
};
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
@@ -213,8 +229,8 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
}
Some(LocalStandardCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
auth_header,
auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body,

20
apps/aether-gateway/src/ai_pipeline/planner/standard/normalize/chat.rs
View File

@@ -3,8 +3,8 @@ use serde_json::Value;
use crate::ai_pipeline::conversion::{request_conversion_kind, RequestConversionKind};
use crate::ai_pipeline::transport::apply_local_body_rules;
use crate::ai_pipeline::transport::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
build_claude_messages_url, build_openai_chat_url, build_openai_cli_url,
build_passthrough_path_url,
};
use crate::ai_pipeline::{
apply_codex_openai_cli_special_body_edits, apply_openai_compact_special_body_edits,
@@ -102,12 +102,16 @@ pub(crate) fn build_cross_format_openai_chat_upstream_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
RequestConversionKind::ToGeminiStandard => build_gemini_content_url(
&transport.endpoint.base_url,
mapped_model,
upstream_is_stream,
parts.uri.query(),
),
RequestConversionKind::ToGeminiStandard => {
crate::ai_pipeline::build_provider_transport_request_url(
transport,
provider_api_format,
Some(mapped_model),
upstream_is_stream,
parts.uri.query(),
None,
)
}
RequestConversionKind::ToOpenAIFamilyCli => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),

20
apps/aether-gateway/src/ai_pipeline/planner/standard/normalize/cli.rs
View File

@@ -9,8 +9,8 @@ use crate::ai_pipeline::transport::antigravity::{
};
use crate::ai_pipeline::transport::apply_local_body_rules;
use crate::ai_pipeline::transport::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
build_claude_messages_url, build_openai_chat_url, build_openai_cli_url,
build_passthrough_path_url,
};
use crate::ai_pipeline::{
apply_codex_openai_cli_special_body_edits, apply_openai_compact_special_body_edits,
@@ -155,12 +155,16 @@ pub(crate) fn build_cross_format_openai_cli_upstream_url(
&transport.endpoint.base_url,
parts.uri.query(),
)),
RequestConversionKind::ToGeminiStandard => build_gemini_content_url(
&transport.endpoint.base_url,
mapped_model,
upstream_is_stream,
parts.uri.query(),
),
RequestConversionKind::ToGeminiStandard => {
crate::ai_pipeline::build_provider_transport_request_url(
transport,
provider_api_format,
Some(mapped_model),
upstream_is_stream,
parts.uri.query(),
None,
)
}
},
}
}

33
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/request.rs
View File

@@ -21,6 +21,7 @@ use crate::ai_pipeline::transport::auth::{
resolve_local_openai_bearer_auth,
};
use crate::ai_pipeline::transport::local_openai_chat_transport_unsupported_reason;
use crate::ai_pipeline::transport::vertex::uses_vertex_api_key_query_auth;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy, GatewayProviderTransportSnapshot};
use crate::AppState;
@@ -301,6 +302,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
.await;
return None;
};
let uses_vertex_query_auth =
uses_vertex_api_key_query_auth(transport, provider_api_format.as_str());
let mut provider_request_headers = if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
@@ -319,10 +322,15 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
Some("application/json"),
)
};
let protected_headers = if uses_vertex_query_auth {
&["content-type"][..]
} else {
&[prepared_candidate.auth_header.as_str(), "content-type"][..]
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&prepared_candidate.auth_header, "content-type"],
protected_headers,
&provider_request_body,
Some(body_json),
) {
@@ -347,11 +355,20 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
let (auth_header, auth_value) = if uses_vertex_query_auth {
provider_request_headers.remove("x-goog-api-key");
(String::new(), String::new())
} else {
ensure_upstream_auth_header(
&mut provider_request_headers,
&prepared_candidate.auth_header,
&prepared_candidate.auth_value,
);
(
prepared_candidate.auth_header.clone(),
prepared_candidate.auth_value.clone(),
)
};
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())
@@ -365,8 +382,8 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
};
Some(LocalOpenAiChatCandidatePayloadParts {
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
auth_header,
auth_value,
mapped_model: prepared_candidate.mapped_model,
provider_api_format,
provider_request_body,

17
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/request.rs
View File

@@ -28,6 +28,7 @@ use crate::ai_pipeline::transport::auth::{
resolve_local_openai_bearer_auth, resolve_local_standard_auth,
};
use crate::ai_pipeline::transport::local_standard_transport_unsupported_reason_with_network;
use crate::ai_pipeline::transport::vertex::uses_vertex_api_key_query_auth;
use crate::ai_pipeline::{ConversionMode, ExecutionStrategy};
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use crate::AppState;
@@ -267,6 +268,7 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
.await;
return None;
};
let uses_vertex_query_auth = uses_vertex_api_key_query_auth(transport, provider_api_format);
let extra_headers = antigravity_auth
.as_ref()
@@ -297,10 +299,15 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
Some("application/json"),
)
};
let protected_headers = if uses_vertex_query_auth {
&["content-type"][..]
} else {
&[auth_header.as_str(), "content-type"][..]
};
if !apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[&auth_header, "content-type"],
protected_headers,
&provider_request_body,
Some(body_json),
) {
@@ -325,7 +332,13 @@ pub(crate) async fn resolve_local_openai_cli_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
let (auth_header, auth_value) = if uses_vertex_query_auth {
provider_request_headers.remove("x-goog-api-key");
(String::new(), String::new())
} else {
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
(auth_header, auth_value)
};
if upstream_is_stream {
provider_request_headers
.entry("accept".to_string())

61
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/plan_builders/stream.rs
View File

@@ -31,12 +31,11 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
return Ok(None);
};
let Some(auth_header) = take_non_empty_string(&mut payload.auth_header) else {
let auth_header = take_non_empty_string(&mut payload.auth_header);
let auth_value = take_non_empty_string(&mut payload.auth_value);
if auth_header.is_some() != auth_value.is_some() {
return Ok(None);
};
let Some(auth_value) = take_non_empty_string(&mut payload.auth_value) else {
return Ok(None);
};
}
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
return Ok(None);
};
@@ -86,27 +85,37 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
let existing_provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers = if existing_provider_request_headers.is_empty() {
if provider_api_format == client_api_format {
build_complete_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
if let (Some(auth_header), Some(auth_value)) =
(auth_header.as_deref(), auth_value.as_deref())
{
if provider_api_format == client_api_format {
build_complete_passthrough_headers_with_auth(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
}
} else {
build_openai_passthrough_headers(
crate::ai_pipeline::transport::auth::build_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
@@ -114,7 +123,9 @@ pub(crate) fn build_openai_chat_stream_plan_from_decision(
} else {
existing_provider_request_headers
};
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
}
provider_request_headers.insert("accept".to_string(), "text/event-stream".to_string());
let content_type = payload
.content_type

61
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/plan_builders/sync.rs
View File

@@ -30,12 +30,11 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
let Some(key_id) = take_non_empty_string(&mut payload.key_id) else {
return Ok(None);
};
let Some(auth_header) = take_non_empty_string(&mut payload.auth_header) else {
let auth_header = take_non_empty_string(&mut payload.auth_header);
let auth_value = take_non_empty_string(&mut payload.auth_value);
if auth_header.is_some() != auth_value.is_some() {
return Ok(None);
};
let Some(auth_value) = take_non_empty_string(&mut payload.auth_value) else {
return Ok(None);
};
}
let Some(provider_api_format) = take_non_empty_string(&mut payload.provider_api_format) else {
return Ok(None);
};
@@ -86,27 +85,37 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
let existing_provider_request_headers = std::mem::take(&mut payload.provider_request_headers);
let extra_headers = std::mem::take(&mut payload.extra_headers);
let mut provider_request_headers = if existing_provider_request_headers.is_empty() {
if provider_api_format == client_api_format {
build_complete_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
if let (Some(auth_header), Some(auth_value)) =
(auth_header.as_deref(), auth_value.as_deref())
{
if provider_api_format == client_api_format {
build_complete_passthrough_headers_with_auth(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else if provider_api_format.starts_with("claude:") {
build_claude_passthrough_headers(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
} else {
build_openai_passthrough_headers(
&parts.headers,
auth_header,
auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
}
} else {
build_openai_passthrough_headers(
crate::ai_pipeline::transport::auth::build_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
&extra_headers,
payload.content_type.as_deref(),
)
@@ -114,7 +123,9 @@ pub(crate) fn build_openai_chat_sync_plan_from_decision(
} else {
existing_provider_request_headers
};
ensure_upstream_auth_header(&mut provider_request_headers, &auth_header, &auth_value);
if let (Some(auth_header), Some(auth_value)) = (auth_header.as_deref(), auth_value.as_deref()) {
ensure_upstream_auth_header(&mut provider_request_headers, auth_header, auth_value);
}
if payload.upstream_is_stream {
provider_request_headers
.entry("accept".to_string())

10
apps/aether-gateway/src/executor/outcome.rs
View File

@@ -619,8 +619,9 @@ async fn load_runtime_miss_candidate_contexts(
&candidate,
"selected_provider_model_name",
),
endpoint_url: endpoint
.and_then(|value| build_runtime_miss_candidate_endpoint_url(value, decision)),
endpoint_url: endpoint.and_then(|value| {
build_runtime_miss_candidate_endpoint_url(&candidate, value, decision)
}),
candidate,
}
})
@@ -671,9 +672,14 @@ fn candidate_extra_data_string(candidate: &StoredRequestCandidate, key: &str) ->
}
fn build_runtime_miss_candidate_endpoint_url(
candidate: &StoredRequestCandidate,
endpoint: &StoredProviderCatalogEndpoint,
decision: Option<&GatewayControlDecision>,
) -> Option<String> {
if let Some(upstream_url) = candidate_extra_data_string(candidate, "upstream_url") {
return Some(upstream_url);
}
let path = endpoint
.custom_path
.as_deref()

183
apps/aether-gateway/src/handlers/admin/provider/query/models.rs
View File

@@ -308,7 +308,11 @@ fn provider_query_transport_supports_standard_test_execution(
)
}
"gemini:chat" | "gemini:cli" => {
state.supports_local_gemini_transport_with_network(transport, api_format)
if crate::provider_transport::is_vertex_api_key_transport_context(transport) {
aether_provider_transport::vertex::supports_local_vertex_api_key_gemini_transport_with_network(transport)
} else {
state.supports_local_gemini_transport_with_network(transport, api_format)
}
}
_ => false,
}
@@ -1018,6 +1022,13 @@ async fn provider_query_execute_standard_test_candidate(
}
};
let uses_vertex_query_auth =
crate::provider_transport::uses_vertex_api_key_query_auth(&transport, provider_api_format);
let vertex_query_auth = if uses_vertex_query_auth {
aether_provider_transport::vertex::resolve_local_vertex_api_key_query_auth(&transport)
} else {
None
};
let oauth_auth = match provider_api_format {
"openai:chat" | "openai:cli" | "claude:chat" | "claude:cli" | "gemini:chat"
| "gemini:cli" => state.resolve_local_oauth_header_auth(&transport).await?,
@@ -1031,14 +1042,24 @@ async fn provider_query_execute_standard_test_candidate(
"claude:chat" | "claude:cli" => {
crate::provider_transport::auth::resolve_local_standard_auth(&transport).or(oauth_auth)
}
"gemini:chat" | "gemini:cli" => state.resolve_local_gemini_auth(&transport).or(oauth_auth),
"gemini:chat" | "gemini:cli" => {
if uses_vertex_query_auth {
oauth_auth
} else {
state.resolve_local_gemini_auth(&transport).or(oauth_auth)
}
}
_ => None,
};
let Some((auth_header, auth_value)) = auth else {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
format!("Provider auth is unavailable for {provider_api_format}"),
));
let (auth_header, auth_value) = match auth {
Some((auth_header, auth_value)) => (Some(auth_header), Some(auth_value)),
None if uses_vertex_query_auth && vertex_query_auth.is_some() => (None, None),
None => {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
format!("Provider auth is unavailable for {provider_api_format}"),
));
}
};
let mut synthetic_request = http::Request::builder()
@@ -1048,91 +1069,16 @@ async fn provider_query_execute_standard_test_candidate(
*synthetic_request.headers_mut() = provider_query_extract_request_headers(payload);
let (parts, _) = synthetic_request.into_parts();
let request_url = match provider_api_format {
"openai:chat" => {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match custom_path {
Some(path) => state.build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
&[],
),
None => Some(
state.build_openai_chat_url(&transport.endpoint.base_url, parts.uri.query()),
),
}
}
"claude:chat" | "claude:cli" => {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match custom_path {
Some(path) => state.build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
&[],
),
None => Some(
state
.build_claude_messages_url(&transport.endpoint.base_url, parts.uri.query()),
),
}
}
"gemini:chat" | "gemini:cli" => {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match custom_path {
Some(path) => state.build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
&["key"],
),
None => state.build_gemini_content_url(
&transport.endpoint.base_url,
&candidate.effective_model,
false,
parts.uri.query(),
),
}
}
"openai:cli" => {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match custom_path {
Some(path) => state.build_passthrough_path_url(
&transport.endpoint.base_url,
path,
parts.uri.query(),
&[],
),
None => Some(build_openai_cli_url(
&transport.endpoint.base_url,
parts.uri.query(),
false,
)),
}
}
_ => None,
};
let request_url = crate::provider_transport::build_transport_request_url(
&transport,
crate::provider_transport::TransportRequestUrlParams {
provider_api_format,
mapped_model: Some(candidate.effective_model.as_str()),
upstream_is_stream: false,
request_query: parts.uri.query(),
kiro_api_region: None,
},
);
let Some(request_url) = request_url else {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
@@ -1144,8 +1090,8 @@ async fn provider_query_execute_standard_test_candidate(
"claude:chat" | "claude:cli" => {
crate::provider_transport::auth::build_claude_passthrough_headers(
&parts.headers,
&auth_header,
&auth_value,
auth_header.as_deref().unwrap_or_default(),
auth_value.as_deref().unwrap_or_default(),
&BTreeMap::new(),
Some("application/json"),
)
@@ -1153,26 +1099,41 @@ async fn provider_query_execute_standard_test_candidate(
"openai:cli" => {
crate::provider_transport::auth::build_complete_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
auth_header.as_deref().unwrap_or_default(),
auth_value.as_deref().unwrap_or_default(),
&BTreeMap::new(),
Some("application/json"),
)
}
_ => state.build_passthrough_headers_with_auth(
&parts.headers,
&auth_header,
&auth_value,
&BTreeMap::new(),
),
_ => match (auth_header.as_deref(), auth_value.as_deref()) {
(Some(auth_header), Some(auth_value)) => state.build_passthrough_headers_with_auth(
&parts.headers,
auth_header,
auth_value,
&BTreeMap::new(),
),
_ => crate::provider_transport::auth::build_passthrough_headers(
&parts.headers,
&BTreeMap::new(),
Some("application/json"),
),
},
};
if uses_vertex_query_auth {
request_headers.remove("x-goog-api-key");
}
request_headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
let protected_headers = if uses_vertex_query_auth {
vec!["content-type"]
} else {
vec![auth_header.as_deref().unwrap_or_default(), "content-type"]
};
if !state.apply_local_header_rules(
&mut request_headers,
transport.endpoint.header_rules.as_ref(),
&[auth_header.as_str(), "content-type"],
&protected_headers,
&provider_request_body,
Some(&request_body),
) {
@@ -1200,11 +1161,17 @@ async fn provider_query_execute_standard_test_candidate(
transport.key.decrypted_auth_config.as_deref(),
);
}
crate::provider_transport::ensure_upstream_auth_header(
&mut request_headers,
&auth_header,
&auth_value,
);
if !uses_vertex_query_auth {
if let (Some(auth_header), Some(auth_value)) =
(auth_header.as_deref(), auth_value.as_deref())
{
crate::provider_transport::ensure_upstream_auth_header(
&mut request_headers,
auth_header,
auth_value,
);
}
}
let plan = ExecutionPlan {
request_id: trace_id.to_string(),

79
apps/aether-gateway/src/handlers/public/support/test_connection/route.rs
View File

@@ -219,50 +219,21 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
let Some((auth_header, auth_value)) = auth else {
return None;
};
let uses_vertex_query_auth = crate::provider_transport::uses_vertex_api_key_query_auth(
&transport,
format_value.as_str(),
);
let upstream_url = {
let custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match (format_value.as_str(), custom_path) {
("openai:chat", Some(path)) | ("claude:chat", Some(path)) => {
crate::provider_transport::url::build_passthrough_path_url(
&transport.endpoint.base_url,
path,
None,
&[],
)
}
("gemini:chat", Some(path)) => {
crate::provider_transport::url::build_passthrough_path_url(
&transport.endpoint.base_url,
path,
None,
&["key"],
)
}
("openai:chat", None) => Some(crate::provider_transport::url::build_openai_chat_url(
&transport.endpoint.base_url,
None,
)),
("claude:chat", None) => {
Some(crate::provider_transport::url::build_claude_messages_url(
&transport.endpoint.base_url,
None,
))
}
("gemini:chat", None) => crate::provider_transport::url::build_gemini_content_url(
&transport.endpoint.base_url,
&model,
false,
None,
),
_ => None,
}
};
let upstream_url = crate::provider_transport::build_transport_request_url(
&transport,
crate::provider_transport::TransportRequestUrlParams {
provider_api_format: format_value.as_str(),
mapped_model: Some(model.as_str()),
upstream_is_stream: false,
request_query: None,
kiro_api_region: None,
},
);
let Some(upstream_url) = upstream_url else {
return None;
};
@@ -271,20 +242,30 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
("content-type".to_string(), "application/json".to_string()),
(auth_header.clone(), auth_value.clone()),
]);
if uses_vertex_query_auth {
provider_request_headers.remove("x-goog-api-key");
}
let protected_headers = if uses_vertex_query_auth {
&["content-type"][..]
} else {
&[auth_header.as_str(), "content-type"][..]
};
if !crate::provider_transport::apply_local_header_rules(
&mut provider_request_headers,
transport.endpoint.header_rules.as_ref(),
&[auth_header.as_str(), "content-type"],
protected_headers,
&provider_request_body,
None,
) {
return None;
}
crate::provider_transport::ensure_upstream_auth_header(
&mut provider_request_headers,
&auth_header,
&auth_value,
);
if !uses_vertex_query_auth {
crate::provider_transport::ensure_upstream_auth_header(
&mut provider_request_headers,
&auth_header,
&auth_value,
);
}
let mut upstream_request = state.client.post(&upstream_url);
for (name, value) in &provider_request_headers {