feat: configure auth channel mismatch formats

This commit is contained in:
fawney19
2026-05-03 00:49:22 +08:00
parent 3a770306cc
commit e3ea2d1451
63 changed files with 585 additions and 39 deletions

View File

@@ -50,6 +50,7 @@ struct GatewayLocalCandidateResolutionPort<'a> {
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&'a serde_json::Value>,
sticky_session_token: Option<&'a str>,
request_auth_channel: Option<&'a str>,
}
#[async_trait]
@@ -86,6 +87,11 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
transport: &Self::Transport,
requested_model: Option<&str>,
) -> Option<&'static str> {
if let Some(skip_reason) =
candidate_auth_channel_skip_reason(transport, self.request_auth_channel)
{
return Some(skip_reason);
}
candidate_common_transport_skip_reason(
transport,
candidate_transport_policy_facts(candidate),
@@ -169,6 +175,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -182,6 +189,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
auth_snapshot,
required_capabilities,
sticky_session_token,
request_auth_channel,
AiCandidateResolutionMode::Standard,
)
.await
@@ -195,6 +203,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -208,6 +217,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
auth_snapshot,
required_capabilities,
sticky_session_token,
request_auth_channel,
AiCandidateResolutionMode::WithoutTransportPairGate,
)
.await
@@ -221,6 +231,7 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -232,6 +243,7 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
auth_snapshot,
required_capabilities,
sticky_session_token,
request_auth_channel,
};
let request = AiCandidateResolutionRequest {
@@ -257,6 +269,87 @@ fn candidate_transport_policy_facts(
}
}
fn candidate_auth_channel_skip_reason(
transport: &GatewayProviderTransportSnapshot,
request_auth_channel: Option<&str>,
) -> Option<&'static str> {
let request_auth_channel = normalize_request_auth_channel(request_auth_channel?)?;
let upstream_auth_channel = resolve_transport_request_auth_channel(transport)?;
if request_auth_channel == upstream_auth_channel
|| allow_auth_channel_mismatch_for_format(transport)
{
None
} else {
Some("auth_channel_mismatch")
}
}
fn normalize_request_auth_channel(value: &str) -> Option<&'static str> {
match value.trim().to_ascii_lowercase().as_str() {
"api_key" | "api-key" | "apikey" => Some("api_key"),
"bearer_like" | "bearer-like" | "bearer" | "oauth" => Some("bearer_like"),
_ => None,
}
}
fn resolve_transport_request_auth_channel(
transport: &GatewayProviderTransportSnapshot,
) -> Option<&'static str> {
let auth_type = resolve_transport_auth_type_for_endpoint_format(transport);
match auth_type.as_str() {
"api_key" => Some("api_key"),
"bearer" => Some("bearer_like"),
"oauth" if provider_uses_bearer_like_oauth(&transport.provider.provider_type) => {
Some("bearer_like")
}
_ => None,
}
}
fn resolve_transport_auth_type_for_endpoint_format(
transport: &GatewayProviderTransportSnapshot,
) -> String {
let default_auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
transport
.key
.auth_type_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|overrides| {
overrides
.get(&api_format)
.or_else(|| overrides.get(transport.endpoint.api_format.trim()))
})
.and_then(serde_json::Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.filter(|value| matches!(value.as_str(), "api_key" | "bearer"))
.unwrap_or(default_auth_type)
}
fn provider_uses_bearer_like_oauth(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "gemini_cli" | "antigravity" | "kiro"
)
}
fn allow_auth_channel_mismatch_for_format(transport: &GatewayProviderTransportSnapshot) -> bool {
let api_format = crate::ai_serving::normalize_api_format_alias(&transport.endpoint.api_format);
transport
.key
.allow_auth_channel_mismatch_formats
.as_ref()
.and_then(serde_json::Value::as_array)
.is_some_and(|items| {
items
.iter()
.filter_map(serde_json::Value::as_str)
.any(|item| crate::ai_serving::normalize_api_format_alias(item) == api_format)
})
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
@@ -285,3 +378,102 @@ pub(crate) async fn read_candidate_transport_snapshot(
}
}
}
#[cfg(test)]
mod tests {
use super::candidate_auth_channel_skip_reason;
use crate::ai_serving::GatewayProviderTransportSnapshot;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
use serde_json::json;
fn sample_transport(auth_type: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".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: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:messages".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("messages".to_string()),
is_active: true,
base_url: "https://example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: auth_type.to_string(),
is_active: true,
api_formats: Some(vec!["claude:messages".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: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn auth_channel_gate_skips_mismatched_raw_secret_auth() {
let transport = sample_transport("bearer");
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
);
}
#[test]
fn auth_channel_gate_allows_explicit_mismatch_format() {
let mut transport = sample_transport("bearer");
transport.key.allow_auth_channel_mismatch_formats = Some(json!(["claude:messages"]));
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
None
);
}
#[test]
fn auth_channel_gate_treats_cli_oauth_provider_as_bearer_like() {
let mut transport = sample_transport("oauth");
transport.provider.provider_type = "claude_code".to_string();
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("bearer_like")),
None
);
assert_eq!(
candidate_auth_channel_skip_reason(&transport, Some("api_key")),
Some("auth_channel_mismatch")
);
}
}