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

Merge remote-tracking branch 'origin/main' into codex/provider-policy-hardening

Browse Source
This commit is contained in:
elky
2026-09-05 16:21:10 +08:00
parent d5f34b2ee2 882bb43125
commit 9d7a0665c0
1228 changed files with 241230 additions and 31814 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/adaptation/private_envelope/sync.rs
+3 -3
View File
@@ -1,3 +1,4 @@
use aether_usage_runtime::decode_internal_report_body_base64;
use base64::Engine as _;
use serde_json::Value;
@@ -43,9 +44,8 @@ pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
}
if let Some(body_base64) = payload.body_base64.as_deref() {
let body_bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_bytes =
decode_internal_report_body_base64(body_base64).map_err(GatewayError::Internal)?;
let Some(normalized_bytes) =
normalize_provider_private_stream_bytes(report_context, &body_bytes)?
else {
apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs
+8 -8
View File
@@ -11,10 +11,10 @@ use super::{
convert_gemini_chat_response_to_openai_chat, convert_gemini_response_to_openai_responses,
maybe_build_local_core_sync_finalize_response,
};
use crate::ai_serving::GatewayControlDecision;
use crate::ai_serving::{
convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat,
convert_openai_responses_response_to_openai_chat, openai_responses_message_item_id,
GatewayControlDecision,
};
use crate::usage::GatewaySyncReportRequest;
@@ -192,7 +192,7 @@ fn aggregates_openai_responses_stream_completed_event_to_final_response() {
"output_text": "Hello",
"output": [{
"type": "message",
"id": "resp_123_msg",
"id": openai_responses_message_item_id("resp_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -843,7 +843,7 @@ fn converts_claude_cli_response_to_openai_responses_response() {
"output_text": "Hello Claude CLI",
"output": [{
"type": "message",
"id": "msg_cli_123_msg",
"id": openai_responses_message_item_id("msg_cli_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -907,7 +907,7 @@ fn converts_claude_cli_tool_use_to_openai_responses_function_call() {
"output": [
{
"type": "message",
"id": "msg_cli_tool_123_msg",
"id": openai_responses_message_item_id("msg_cli_tool_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -977,7 +977,7 @@ fn converts_gemini_cli_response_to_openai_responses_response() {
"output_text": "Hello Gemini CLI",
"output": [{
"type": "message",
"id": "resp_cli_123_msg",
"id": openai_responses_message_item_id("resp_cli_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -1046,7 +1046,7 @@ fn converts_gemini_cli_function_call_to_openai_responses_function_call() {
"output": [
{
"type": "message",
"id": "resp_cli_tool_123_msg",
"id": openai_responses_message_item_id("resp_cli_tool_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -1252,7 +1252,7 @@ fn local_finalize_handles_openai_responses_openai_family_sync_response_even_when
"model": "gpt-5",
"output": [{
"type": "message",
"id": "resp_cli_family_123_msg",
"id": openai_responses_message_item_id("resp_cli_family_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+83 -123
View File
@@ -9,7 +9,7 @@ use aether_ai_serving::{
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
use aether_routing_core::{
rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts,
RoutingCandidateTrace, RoutingDecisionTrace,
RoutingCandidateTrace, RoutingDecisionTrace, RoutingExecutionPolicy,
};
use aether_scheduler_core::{
ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome,
@@ -47,13 +47,12 @@ use crate::cache::{
use crate::clock::current_unix_ms;
use crate::dispatch::refs::dispatch_ref_for_local_candidate;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
use crate::orchestration::{ExecutionAttemptIdentity, POOL_KEY_RETRY_INDEX_STRIDE};
use crate::scheduler::candidate::is_auth_api_key_concurrency_limit_skip_reason;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
const AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET: Duration = Duration::from_millis(100);
const AUTH_API_KEY_CONCURRENCY_RETRY_DELAY: Duration = Duration::from_millis(10);
@@ -80,6 +79,13 @@ type DecorateSkippedCandidateFn<'a> = Arc<
pub(crate) trait LocalExecutionAttemptSource<T>: Send {
async fn next_execution_attempt(&mut self) -> Result<Option<T>, GatewayError>;
/// Returns the request-scoped execution behaviour selected by routing.
/// Execution wrappers use this snapshot before consuming the first
/// attempt, avoiding a second lookup against mutable system settings.
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
None
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<T>, GatewayError>;
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError>;
@@ -481,10 +487,6 @@ where
type ExtraData = Value;
type Error = Infallible;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
local_attempt_slot_count(&candidate.transport)
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
available_candidate_extra_data_with_dispatch_ref(candidate, &self.build_extra_data)
}
@@ -1242,9 +1244,7 @@ async fn scheduler_cache_affinity_enabled(
state: PlannerAppState<'_>,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> bool {
scheduler_ordering_config_for_routing_policy(state, routing_policy)
.await
.scheduling_mode
scheduler_ordering_config_for_routing_policy(routing_policy).scheduling_mode
== SchedulerSchedulingMode::CacheAffinity
}
@@ -1610,7 +1610,8 @@ async fn persist_available_local_execution_candidate_at_index<F>(
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
// Exactly one attempt is materialized per candidate; same-key retries are
// derived lazily by the attempt loop after a failure.
let extra_data = ai_candidate_extra_data_with_ranking(
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
candidate.ranking.as_ref(),
@@ -1625,53 +1626,34 @@ where
Some(candidate_index),
extra_data,
);
let should_persist = should_persist_available_local_candidate(&candidate);
let mut attempts = Vec::with_capacity(attempt_slots as usize);
let mut owned_candidate = Some(candidate);
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = if should_persist_available_local_candidate(&candidate) {
state
.persist_available_local_candidate(
trace_id,
context.user_id,
context.api_key_id,
&candidate.candidate,
candidate_index,
retry_index,
generated_candidate_id.as_str(),
context.required_capabilities,
extra_data,
current_unix_ms(),
context.error_context,
)
.await
} else {
generated_candidate_id
};
for retry_index in 0..attempt_slots {
let candidate_ref = owned_candidate
.as_ref()
.expect("candidate should remain available until final retry");
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = if should_persist {
state
.persist_available_local_candidate(
trace_id,
context.user_id,
context.api_key_id,
&candidate_ref.candidate,
candidate_index,
effective_retry_index(retry_index, candidate_ref.orchestration.pool_key_index),
generated_candidate_id.as_str(),
context.required_capabilities,
extra_data.clone(),
current_unix_ms(),
context.error_context,
)
.await
} else {
generated_candidate_id
};
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
candidate_ref.clone()
};
let retry_index =
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
attempts.push(LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id,
});
}
attempts
vec![LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id,
}]
}
fn available_candidate_extra_data_with_dispatch_ref<F>(
@@ -1840,6 +1822,7 @@ fn routing_trace_for_candidate(
CandidateKind::Provider => Some(candidate.key_id.clone()),
CandidateKind::PoolGroup => None,
},
api_format: Some(candidate.endpoint_api_format.clone()),
provider_priority: candidate.provider_priority,
key_priority: candidate
.key_global_priority_for_format
@@ -1923,32 +1906,15 @@ fn build_unpersisted_local_execution_candidate_attempts(
candidate: EligibleLocalExecutionCandidate,
candidate_index: u32,
) -> VecDeque<LocalExecutionCandidateAttempt> {
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
let mut attempts = VecDeque::with_capacity(attempt_slots as usize);
let mut owned_candidate = Some(candidate);
for retry_index in 0..attempt_slots {
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
owned_candidate
.as_ref()
.expect("candidate should remain available until final retry")
.clone()
};
let retry_index =
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
attempts.push_back(LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id: Uuid::new_v4().to_string(),
});
}
attempts
// One attempt per candidate; same-key retries are derived lazily by the
// attempt loop after a failure.
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
VecDeque::from([LocalExecutionCandidateAttempt {
eligible: candidate,
candidate_index,
retry_index,
candidate_id: Uuid::new_v4().to_string(),
}])
}
async fn persist_pool_group_exhaustion_skipped_candidate(
@@ -2276,6 +2242,8 @@ mod tests {
pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch: None,
// These tests cover persistence shape, not same-key retries.
sticky_key_attempts: Some(1),
},
ranking: None,
}
@@ -2357,16 +2325,11 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 2);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
assert!(stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.is_none());
}
#[test]
@@ -2514,14 +2477,23 @@ mod tests {
assert!(should_cache_resolved_candidate_page(&cursor));
let fixed_order_app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::disabled().with_system_config_values_for_tests([(
"scheduling_mode".to_string(),
json!("fixed_order"),
)]),
);
let fixed_order_app = AppState::new().expect("state should build");
let fixed_order_policy = ResolvedRoutingPolicy {
group_id: Some("routing-group-fixed-order".to_string()),
group_version: Some(1),
selection_source: "test".to_string(),
requested_model: "gpt-5".to_string(),
resolved_model: "gpt-5".to_string(),
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
matched_rules: Vec::new(),
};
let mut page_cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&fixed_order_app),
&model_directive_policy,
@@ -2531,7 +2503,7 @@ mod tests {
true,
None,
&auth_snapshot,
None,
Some(&fixed_order_policy),
None,
None,
false,
@@ -2549,7 +2521,7 @@ mod tests {
auth_snapshot,
client_session_affinity: None,
required_capabilities: None,
routing_policy: None,
routing_policy: Some(fixed_order_policy),
sticky_session_token: None,
request_auth_channel: None,
skipped_user_id: "user-1".to_string(),
@@ -2647,16 +2619,11 @@ mod tests {
);
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[1].candidate_index, 1);
assert_eq!(
stored[1]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
assert!(stored[1]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.is_none());
}
#[test]
@@ -2726,7 +2693,7 @@ mod tests {
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("ranking metadata should persist as object extra data");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert!(extra_data.get("existing").is_none());
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
@@ -2739,14 +2706,7 @@ mod tests {
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
assert_eq!(
extra_data
.get("dispatch_ref")
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("ranked-key"))
);
assert!(extra_data.get("dispatch_ref").is_none());
}
#[tokio::test]
@@ -3084,7 +3044,7 @@ mod tests {
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("skipped ranking metadata should persist");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert!(extra_data.get("existing").is_none());
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs
+10 -2
View File
@@ -278,13 +278,21 @@ mod tests {
assert_eq!(metadata["transport_diagnostics"]["provider_type"], "codex");
assert_eq!(
metadata["transport_diagnostics"]["fingerprint"]["transport_profile"]["profile_id"],
"chrome_136"
metadata["transport_diagnostics"]["key_fingerprint_configured"],
Value::Bool(true)
);
assert_eq!(
metadata["transport_diagnostics"]["key_transport_profile_configured"],
Value::Bool(true)
);
assert_eq!(
metadata["transport_diagnostics"]["resolved_transport_profile_id"],
"chrome_136"
);
assert_eq!(
metadata["transport_diagnostics"]["resolved_transport_profile"]["profile_id"],
"chrome_136"
);
assert_eq!(
metadata["transport_diagnostics"]["request_pair"]["conversion_enabled"],
Value::Bool(true)
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
+113 -86
View File
@@ -3,17 +3,14 @@ use aether_ai_serving::{
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
AiRankingSchedulingMode,
};
use aether_routing_core::{ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode};
use aether_routing_core::ResolvedRoutingPolicy;
use async_trait::async_trait;
use tokio::sync::Mutex;
use tracing::warn;
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::scheduler::config::{
read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use aether_scheduler_core::{
matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
@@ -133,7 +130,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
required_capabilities: Option<&serde_json::Value>,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> Vec<EligibleLocalExecutionCandidate> {
let ordering_config = scheduler_ordering_config_for_routing_policy(state, routing_policy).await;
let ordering_config = scheduler_ordering_config_for_routing_policy(routing_policy);
let port = GatewayLocalCandidateRankingPort {
state,
requested_model,
@@ -184,35 +181,24 @@ fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedul
}
}
pub(crate) async fn scheduler_ordering_config_for_routing_policy(
state: PlannerAppState<'_>,
/// Return the immutable scheduler snapshot carried by a resolved routing
/// policy. A missing policy is a programming error in production request
/// paths; unit tests may use the scheduler default for isolated ranking tests.
pub(crate) fn scheduler_ordering_config_for_routing_policy(
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> SchedulerOrderingConfig {
let system_config = read_scheduler_ordering_config_or_default(state).await;
match routing_policy {
Some(policy) => {
let mut config = scheduler_ordering_config_from_routing_policy(policy);
config.keep_priority_on_conversion |= system_config.keep_priority_on_conversion;
config
Some(policy) => SchedulerOrderingConfig::from_routing_policy(policy),
None => {
#[cfg(test)]
{
SchedulerOrderingConfig::default()
}
#[cfg(not(test))]
{
panic!("resolved routing policy is required before candidate scheduling")
}
}
None => system_config,
}
}
fn scheduler_ordering_config_from_routing_policy(
policy: &ResolvedRoutingPolicy,
) -> SchedulerOrderingConfig {
SchedulerOrderingConfig {
priority_mode: match policy.priority_mode {
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
},
scheduling_mode: match policy.scheduling_mode {
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
},
keep_priority_on_conversion: policy.keep_priority_on_conversion,
}
}
@@ -231,37 +217,32 @@ fn routing_overlaid_candidate(
let overlaid_key_priority = match kind {
LocalExecutionCandidateKind::SingleKey => policy
.ranking_overlay
.key_priority_overrides
.get(candidate.key_id.as_str()),
.key_priority_override_matching_format(candidate.key_id.as_str(), |format| {
crate::ai_serving::api_format_alias_matches(
format,
candidate.endpoint_api_format.as_str(),
)
})
.or_else(|| {
policy
.ranking_overlay
.key_priority_overrides
.get(candidate.key_id.as_str())
.copied()
}),
LocalExecutionCandidateKind::PoolGroup => policy
.ranking_overlay
.pool_priority_overrides
.get(candidate.provider_id.as_str()),
.get(candidate.provider_id.as_str())
.copied(),
};
if let Some(overlaid_key_priority) = overlaid_key_priority.copied() {
if let Some(overlaid_key_priority) = overlaid_key_priority {
overlaid.key_internal_priority = overlaid_key_priority;
overlaid.key_global_priority_for_format = Some(overlaid_key_priority);
}
overlaid
}
async fn read_scheduler_ordering_config_or_default(
state: PlannerAppState<'_>,
) -> SchedulerOrderingConfig {
match read_scheduler_ordering_config(state.app()).await {
Ok(config) => config,
Err(error) => {
warn!(
event_name = "planner_scheduler_ordering_config_load_failed",
log_type = "event",
error = ?error,
"failed to load scheduler ordering config while ranking local execution candidates"
);
SchedulerOrderingConfig::default()
}
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
@@ -270,10 +251,17 @@ mod tests {
use aether_ai_serving::{
ai_ranking_context, build_ai_rankable_candidate, AiRankableCandidateParts,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::{
provider_catalog::InMemoryProviderCatalogReadRepository,
routing_profiles::InMemoryRoutingGroupRepository,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::routing_profiles::{
CreateRoutingGroupRecord, RoutingGroupWriteRepository,
};
use aether_scheduler_core::{
apply_scheduler_candidate_ranking,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
@@ -303,7 +291,11 @@ mod tests {
required_capabilities: Option<&serde_json::Value>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let ordering_config = super::read_scheduler_ordering_config_or_default(state).await;
let ordering_config =
crate::scheduler::config::read_system_default_routing_ordering_config(state.app())
.await
.expect("routing strategy should load")
.unwrap_or_default();
let mut candidates = candidates;
let mut rankables = Vec::with_capacity(candidates.len());
let mut ordering_cache = CandidateTransportRankingFactsCache::default();
@@ -378,6 +370,8 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: aether_routing_core::RankingOverlay::default(),
mutation_plan: Default::default(),
pool_policy_overrides: BTreeMap::new(),
@@ -396,7 +390,7 @@ mod tests {
}
#[tokio::test]
async fn routing_policy_inherits_global_conversion_priority_override() {
async fn routing_policy_ignores_legacy_global_conversion_priority_override() {
let data_state = GatewayDataState::default().with_system_config_values_for_tests([(
"keep_priority_on_conversion".to_string(),
json!(true),
@@ -413,23 +407,24 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
matched_rules: Vec::new(),
};
let ordering = super::scheduler_ordering_config_for_routing_policy(
PlannerAppState::new(&state),
Some(&policy),
)
.await;
let ordering = super::scheduler_ordering_config_for_routing_policy(Some(&policy));
assert_eq!(
ordering.scheduling_mode,
crate::scheduler::config::SchedulerSchedulingMode::FixedOrder
);
assert!(ordering.keep_priority_on_conversion);
assert!(
!ordering.keep_priority_on_conversion,
"a resolved routing policy must not inherit the legacy system-config flag"
);
}
#[test]
@@ -447,6 +442,8 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::GlobalKey,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: aether_routing_core::RankingOverlay {
pool_priority_overrides: BTreeMap::from([("provider-1".to_string(), 4)]),
key_priority_overrides: BTreeMap::from([("representative-key".to_string(), 1)]),
@@ -570,6 +567,15 @@ mod tests {
api_formats: Option<serde_json::Value>,
allowed_models: Option<serde_json::Value>,
) -> StoredProviderCatalogKey {
let credential_state = AppState::new()
.expect("credential state should build")
.with_data_state_for_tests(
GatewayDataState::disabled()
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let encrypted_api_key = credential_state
.seal_provider_catalog_key_api_key(provider_id, id, "plain-upstream-key")
.expect("api key should encrypt");
StoredProviderCatalogKey::new(
id.to_string(),
provider_id.to_string(),
@@ -581,7 +587,7 @@ mod tests {
.expect("key should build")
.with_transport_fields(
api_formats,
"plain-upstream-key".to_string(),
encrypted_api_key,
None,
None,
Some(json!({"openai:chat": 1})),
@@ -695,7 +701,7 @@ mod tests {
let observed_at_unix_secs = current_unix_secs();
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![
("provider_priority_mode".to_string(), json!("provider")),
@@ -704,6 +710,7 @@ mod tests {
serde_json::to_value(TunnelAttachmentRecord {
gateway_instance_id: "gateway-b".to_string(),
relay_base_url: "http://gateway-b:8080".to_string(),
tunnel_generation: "test-generation-remote".to_string(),
conn_count: 1,
observed_at_unix_secs,
})
@@ -714,6 +721,7 @@ mod tests {
serde_json::to_value(TunnelAttachmentRecord {
gateway_instance_id: "gateway-a".to_string(),
relay_base_url: "http://gateway-a:8080".to_string(),
tunnel_generation: "test-generation-local".to_string(),
conn_count: 1,
observed_at_unix_secs,
})
@@ -772,7 +780,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -825,7 +833,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
@@ -882,7 +890,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -918,7 +926,8 @@ mod tests {
}
#[tokio::test]
async fn local_execution_ranking_keeps_cross_format_priority_when_global_override_is_enabled() {
async fn local_execution_ranking_keeps_cross_format_priority_when_strategy_override_is_enabled()
{
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider_with_options("provider-same", false, 10),
@@ -933,14 +942,32 @@ mod tests {
sample_key_for_provider("provider-cross", "key-cross", ""),
],
);
let routing_repository = std::sync::Arc::new(InMemoryRoutingGroupRepository::default());
routing_repository
.create_routing_group(CreateRoutingGroupRecord {
id: "strategy-default".to_string(),
name: "strategy-default".to_string(),
description: None,
enabled: true,
is_system_default: true,
sort_order: 0,
config_json: json!({
"default_policy": {
"keep_priority_on_conversion": true
}
}),
version: 1,
created_at: 1,
updated_at: 1,
published_at: None,
})
.await
.expect("routing strategy should be created");
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"keep_priority_on_conversion".to_string(),
json!(true),
)]);
.with_routing_group_repository_for_tests(routing_repository);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
@@ -1000,7 +1027,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
@@ -1066,7 +1093,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1119,7 +1146,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1193,7 +1220,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1273,7 +1300,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1349,7 +1376,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1416,7 +1443,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1499,7 +1526,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1564,7 +1591,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1653,7 +1680,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1739,7 +1766,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1836,7 +1863,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -1941,7 +1968,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
@@ -2035,7 +2062,7 @@ mod tests {
);
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
std::sync::Arc::new(provider_catalog),
"development-key",
DEVELOPMENT_ENCRYPTION_KEY,
);
let state = AppState::new()
.expect("state should build")
apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs
+12 -1
View File
@@ -23,7 +23,9 @@ use crate::ai_serving::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::stage_metrics::observe_gateway_stage_ms;
use super::candidate_ranking::rank_eligible_local_execution_candidates;
use super::candidate_ranking::{
rank_eligible_local_execution_candidates, scheduler_ordering_config_for_routing_policy,
};
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
@@ -378,8 +380,17 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
"candidate_resolution_core",
started_at.elapsed().as_millis() as u64,
);
let sticky_key_attempts = if outcome.eligible_candidates.is_empty() {
None
} else {
Some(
scheduler_ordering_config_for_routing_policy(routing_policy)
.sticky_key_attempts,
)
};
for candidate in &mut outcome.eligible_candidates {
candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch);
candidate.orchestration.sticky_key_attempts = sticky_key_attempts;
}
(outcome.eligible_candidates, outcome.skipped_candidates)
}
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
+33 -11
View File
@@ -174,6 +174,9 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
self.ranking_seed,
false,
self.request_operation,
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(
self.routing_policy,
),
)
.await?;
@@ -425,11 +428,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
);
let ordering_config =
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(
state,
routing_policy,
)
.await;
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(routing_policy);
Self {
state,
@@ -1291,6 +1290,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.then_some(self.client_session_affinity.as_ref())
.flatten(),
self.ranking_seed,
self.ordering_config,
)
.await?;
let skipped_candidates = skipped_candidates
@@ -1473,6 +1473,7 @@ mod tests {
use super::*;
use crate::data::GatewayDataState;
use crate::AppState;
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::DataLayerError;
@@ -1884,6 +1885,8 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -1947,6 +1950,8 @@ mod tests {
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -2170,6 +2175,19 @@ mod tests {
None,
)
.expect("endpoint transport should build");
let credential_state = AppState::new()
.expect("credential state should build")
.with_data_state_for_tests(
GatewayDataState::disabled()
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let encrypted_api_key = credential_state
.seal_provider_catalog_key_api_key(
row.provider_id.as_str(),
row.key_id.as_str(),
"plain-upstream-key",
)
.expect("api key should encrypt");
let key = StoredProviderCatalogKey::new(
row.key_id.clone(),
row.provider_id.clone(),
@@ -2181,7 +2199,7 @@ mod tests {
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!([row.endpoint_api_format.clone()])),
"plain-upstream-key".to_string(),
encrypted_api_key,
None,
None,
None,
@@ -2536,7 +2554,7 @@ mod tests {
provider_repository,
candidate_repository,
)
.with_encryption_key_for_tests("development-key");
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
@@ -2656,15 +2674,17 @@ mod tests {
provider_repository,
candidate_repository,
)
.with_encryption_key_for_tests("development-key")
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY)
// Legacy keys deliberately disagree with the routing policy: the
// resolved policy must be the only source of scheduler ordering.
.with_system_config_values_for_tests([
(
"scheduling_mode".to_string(),
serde_json::json!("fixed_order"),
serde_json::json!("cache_affinity"),
),
(
"keep_priority_on_conversion".to_string(),
serde_json::json!(true),
serde_json::json!(false),
),
]);
let app = AppState::new()
@@ -2681,7 +2701,9 @@ mod tests {
resolved_model: "gpt-5.4-mini".to_string(),
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
keep_priority_on_conversion: false,
keep_priority_on_conversion: true,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
execution_policy: Default::default(),
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
apps/aether-gateway/src/ai_serving/planner/decision/stream.rs
+31 -9
View File
@@ -11,6 +11,7 @@ use crate::ai_serving::planner::route::{
is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
};
use crate::ai_serving::{resolve_decision_execution_runtime_auth_context, GatewayControlDecision};
use crate::state::VideoTaskRouteAccess;
use crate::{AiExecutionDecision, AppState, GatewayError};
pub(crate) async fn maybe_build_stream_decision_payload(
@@ -155,16 +156,37 @@ async fn maybe_build_local_video_task_content_stream_decision_payload(
return Ok(None);
}
let _ = state
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
.await?;
let Some(user_id) = decision
.auth_context
.as_ref()
.filter(|auth_context| auth_context.access_allowed)
.map(|auth_context| auth_context.user_id.trim())
.filter(|value| !value.is_empty())
else {
return Err(crate::video_tasks::not_found_error());
};
if state
.hydrate_video_task_for_route_for_user(
decision.route_family.as_deref(),
parts.uri.path(),
user_id,
)
.await?
!= VideoTaskRouteAccess::Allowed
{
return Err(crate::video_tasks::not_found_error());
}
let Some(action) = state.video_tasks.prepare_openai_content_stream_action(
parts.uri.path(),
parts.uri.query(),
trace_id,
) else {
return Ok(None);
let Some(action) = state
.video_tasks
.prepare_openai_content_stream_action_for_user(
parts.uri.path(),
parts.uri.query(),
trace_id,
user_id,
)
else {
return Err(crate::video_tasks::not_found_error());
};
let crate::video_tasks::LocalVideoTaskContentAction::StreamPlan(plan) = action else {
apps/aether-gateway/src/ai_serving/planner/decision/sync.rs
+19 -9
View File
@@ -16,6 +16,7 @@ use crate::ai_serving::{
build_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
GatewayControlDecision,
};
use crate::state::VideoTaskRouteAccess;
use crate::{AiExecutionDecision, AppState, GatewayError};
pub(crate) async fn maybe_build_sync_decision_payload(
@@ -191,10 +192,6 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
return Ok(None);
}
let _ = state
.hydrate_video_task_for_route(decision.route_family.as_deref(), parts.uri.path())
.await?;
let auth_context = resolve_execution_runtime_auth_context(
state,
decision,
@@ -204,16 +201,30 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
)
.await?;
let Some(auth_context) = auth_context else {
return Ok(None);
return Err(crate::video_tasks::not_found_error());
};
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan(
if !auth_context.access_allowed || auth_context.user_id.trim().is_empty() {
return Err(crate::video_tasks::not_found_error());
}
if state
.hydrate_video_task_for_route_for_user(
decision.route_family.as_deref(),
parts.uri.path(),
&auth_context.user_id,
)
.await?
!= VideoTaskRouteAccess::Allowed
{
return Err(crate::video_tasks::not_found_error());
}
let Some(follow_up) = state.video_tasks.prepare_follow_up_sync_plan_for_user(
plan_kind,
parts.uri.path(),
Some(body_json),
Some(&auth_context),
trace_id,
) else {
return Ok(None);
return Err(crate::video_tasks::not_found_error());
};
let aether_video_tasks_core::LocalVideoTaskFollowUpPlan {
@@ -236,8 +247,7 @@ async fn maybe_build_local_video_task_follow_up_sync_decision_payload(
downstream_path = %parts.uri.path(),
provider_api_format = %plan.provider_api_format,
client_api_format = %plan.client_api_format,
upstream_base_url = ?upstream_base_url,
upstream_url = %plan.url,
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
"gateway built local video follow-up sync decision payload"
);
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
+90 -25
View File
@@ -37,6 +37,10 @@ const ROUTING_GROUP_SELECTION_CACHE_TTL: Duration = Duration::from_secs(30);
const ROUTING_GROUP_SELECTION_CACHE_STALE_TTL: Duration = Duration::from_secs(120);
const CODEX_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
const CODEX_FEDRAMP_HEADER: &str = "x-openai-fedramp";
const INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE: &str =
"routing provider request violates provider contract";
const INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE: &str =
"invalid provider request headers in routing mutation";
#[derive(Debug, Clone)]
pub(crate) struct ResolvedLocalDecisionAuthInput {
@@ -292,6 +296,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
crate::ai_serving::openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
provider_model.as_str(),
)
})
.unwrap_or_default();
@@ -311,10 +316,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
)
}
}
.map_err(|violation| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("routing provider_request violates provider contract: {violation:?}"),
})?;
.map_err(|_| invalid_routing_provider_contract())?;
}
let provider_model = provider_request_body
.get("model")
@@ -649,21 +651,17 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
GatewayRoutingSelectionError::NotFound(explicit_group.unwrap_or_default()),
));
}
None
return Err(routing_selection_error(
GatewayRoutingSelectionError::NoDefault,
));
}
};
let Some((group_id, group_version, group_config_json, selection_source)) = selected_group
else {
input.client_session_affinity = client_session_affinity_from_api_request(
client_api_format,
&parts.headers,
Some(body_json),
);
input.routing_policy = None;
input.routing_trace_seed = None;
input.routing_context = None;
return Ok(());
return Err(routing_selection_error(
GatewayRoutingSelectionError::NoDefault,
));
};
if try_attach_static_default_routing_policy_to_input(
@@ -887,10 +885,36 @@ fn routing_selection_error(error: GatewayRoutingSelectionError) -> GatewayError
GatewayRoutingSelectionError::Repository(message) => {
GatewayError::Internal(format!("routing group repository lookup failed: {message}"))
}
error => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: error.to_string(),
GatewayRoutingSelectionError::NoDefault => GatewayError::Client {
status: StatusCode::SERVICE_UNAVAILABLE,
message: "no enabled routing strategy is configured for this request".to_string(),
},
GatewayRoutingSelectionError::NotFound(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: "requested routing group was not found".to_string(),
},
GatewayRoutingSelectionError::Disabled(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: "requested routing group is not enabled".to_string(),
},
GatewayRoutingSelectionError::Forbidden(_) => GatewayError::Client {
status: StatusCode::FORBIDDEN,
message: "requested routing group is not allowed for this principal".to_string(),
},
}
}
fn invalid_routing_provider_contract() -> GatewayError {
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE.to_string(),
}
}
fn invalid_routing_provider_headers() -> GatewayError {
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE.to_string(),
}
}
@@ -945,14 +969,9 @@ fn btree_headers_to_header_map(
) -> Result<HeaderMap, GatewayError> {
let mut output = HeaderMap::new();
for (name, value) in headers {
let name = HeaderName::from_bytes(name.as_bytes()).map_err(|err| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("invalid provider request header name in routing mutation: {err}"),
})?;
let value = HeaderValue::from_str(value).map_err(|err| GatewayError::Client {
status: StatusCode::BAD_REQUEST,
message: format!("invalid provider request header value in routing mutation: {err}"),
})?;
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(|_| invalid_routing_provider_headers())?;
let value = HeaderValue::from_str(value).map_err(|_| invalid_routing_provider_headers())?;
output.insert(name, value);
}
Ok(output)
@@ -1097,6 +1116,7 @@ fn ensure_report_context_routing_trace(
endpoint_id: decision.endpoint_id.clone().unwrap_or_default(),
model_id,
key_id,
api_format: decision.provider_api_format.clone(),
provider_priority,
key_priority,
},
@@ -1174,6 +1194,50 @@ mod tests {
}
}
#[test]
fn routing_selection_errors_do_not_echo_explicit_group() {
let secret = "private-group?token=Bearer-secret";
for error in [
GatewayRoutingSelectionError::NotFound(secret.to_string()),
GatewayRoutingSelectionError::Disabled(secret.to_string()),
GatewayRoutingSelectionError::Forbidden(secret.to_string()),
] {
let error = routing_selection_error(error);
assert!(matches!(
error,
GatewayError::Client {
status: StatusCode::FORBIDDEN,
ref message,
} if !message.contains(secret)
));
}
}
#[test]
fn routing_provider_errors_do_not_echo_dynamic_details() {
let secret = "https://internal.example/?token=Bearer-secret";
let contract_error = invalid_routing_provider_contract();
let header_error = btree_headers_to_header_map(&BTreeMap::from([(
format!("Authorization: {secret}"),
secret.to_string(),
)]))
.expect_err("invalid header should fail");
for (error, expected_message) in [
(contract_error, INVALID_ROUTING_PROVIDER_CONTRACT_MESSAGE),
(header_error, INVALID_ROUTING_PROVIDER_HEADERS_MESSAGE),
] {
assert!(matches!(
error,
GatewayError::Client {
status: StatusCode::BAD_REQUEST,
ref message,
} if message == expected_message && !message.contains(secret)
));
}
}
#[tokio::test]
async fn explicit_auth_snapshot_override_does_not_fall_back_to_the_planner_cache() {
// AppState::new has no auth snapshot repository. Without the explicit
@@ -1230,6 +1294,7 @@ mod tests {
description: None,
enabled: true,
is_system_default: false,
sort_order: 0,
config_json: json!({}),
version: 1,
created_at: 1,
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs
+6
View File
@@ -140,6 +140,9 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -246,6 +249,9 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await?;
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
+1
View File
@@ -203,6 +203,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
+1
View File
@@ -183,6 +183,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
prepared.transport.provider.provider_type.as_str(),
prepared.transport.endpoint.base_url.as_str(),
prepared.mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs
+15
View File
@@ -26,6 +26,7 @@ use super::{
LocalSameFormatProviderCandidateAttemptSource, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderSpec,
};
use aether_routing_core::RoutingExecutionPolicy;
pub(crate) struct LocalSameFormatProviderSyncAttemptSource<'a> {
state: &'a AppState,
@@ -189,6 +190,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
@@ -234,6 +242,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncA
impl LocalExecutionAttemptSource<AiStreamAttempt>
for LocalSameFormatProviderStreamAttemptSource<'_>
{
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_stream_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/report_context.rs
+68 -10
View File
@@ -4,7 +4,7 @@ use aether_ai_serving::{
build_ai_execution_report_context,
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
AiExecutionReportContextParts, AiRequestOrigin,
AiExecutionReportContextParts, AiRequestOrigin, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
};
use aether_routing_core::ResolvedRoutingPolicy;
use aether_runtime_state::RuntimeLockLease;
@@ -21,7 +21,8 @@ use crate::client_session_affinity::{
};
use crate::orchestration::{
insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity,
ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
ROUTING_EXECUTION_POLICY_REPORT_FIELD, ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD,
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
};
use crate::scheduler::affinity::insert_scheduler_affinity_policy_report_context_field;
@@ -59,6 +60,9 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
pub(crate) routing_policy: Option<&'a ResolvedRoutingPolicy>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
/// Routing policy sticky-key attempt budget; read back by the attempt
/// loop to derive same-key retries lazily.
pub(crate) sticky_key_attempts: Option<u32>,
pub(crate) client_requested_stream: bool,
pub(crate) upstream_is_stream: bool,
pub(crate) has_envelope: bool,
@@ -72,10 +76,12 @@ pub(crate) fn build_local_execution_report_context(
let RequestOrigin {
client_ip,
user_agent,
forwarded_headers_trusted,
} = parts
.request_origin
.unwrap_or_else(|| request_origin_from_headers(parts.original_headers));
let original_headers = crate::ai_serving::collect_control_headers(parts.original_headers);
let original_headers =
collect_report_context_original_headers(parts.original_headers, forwarded_headers_trusted);
let original_request_body = crate::ai_serving::build_report_context_original_request_echo(
parts.original_request_body_json,
parts.original_request_body_base64,
@@ -102,13 +108,20 @@ pub(crate) fn build_local_execution_report_context(
value,
);
}
if let Some(incoming_tls) =
crate::ai_serving::tls_fingerprint_from_headers(parts.original_headers)
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
if forwarded_headers_trusted {
if let Some(incoming_tls) =
crate::ai_serving::tls_fingerprint_from_headers(parts.original_headers)
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
}
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
insert_scheduler_affinity_policy_report_context_field(&mut extra_fields, parts.routing_policy);
if let Some(policy) = parts.routing_policy {
if let Ok(value) = serde_json::to_value(policy.execution_policy) {
extra_fields.insert(ROUTING_EXECUTION_POLICY_REPORT_FIELD.to_string(), value);
}
}
if let Some(override_policy) = parts
.routing_policy
.and_then(|policy| policy.pool_policy_overrides.get(parts.provider_id))
@@ -124,6 +137,12 @@ pub(crate) fn build_local_execution_report_context(
Value::Number(epoch.into()),
);
}
if let Some(sticky_key_attempts) = parts.sticky_key_attempts {
extra_fields.insert(
STICKY_KEY_ATTEMPTS_REPORT_FIELD.to_string(),
Value::Number(sticky_key_attempts.into()),
);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
@@ -174,6 +193,17 @@ pub(crate) fn build_local_execution_report_context(
})
}
fn collect_report_context_original_headers(
headers: &http::HeaderMap,
forwarded_headers_trusted: bool,
) -> BTreeMap<String, String> {
let mut collected = crate::ai_serving::collect_control_headers(headers);
if !forwarded_headers_trusted {
collected.retain(|name, _| !name.starts_with("x-aether-tls-"));
}
collected
}
fn insert_request_path_fields(
extra_fields: &mut Map<String, Value>,
request_path: Option<&str>,
@@ -243,8 +273,8 @@ mod tests {
use serde_json::{json, Map, Value};
use super::{
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
LocalExecutionReportContextParts,
build_local_execution_report_context, collect_report_context_original_headers,
provider_stream_event_api_format_for_provider_type, LocalExecutionReportContextParts,
};
use crate::ai_serving::ExecutionRuntimeAuthContext;
use crate::ai_serving::RequestOrigin;
@@ -274,6 +304,26 @@ mod tests {
);
}
#[test]
fn untrusted_tls_forwarding_headers_are_excluded_from_report_context() {
let mut headers = http::HeaderMap::new();
headers.insert("x-aether-tls-ja3", "spoofed-ja3".parse().unwrap());
headers.insert(http::header::USER_AGENT, "test-client".parse().unwrap());
let untrusted = collect_report_context_original_headers(&headers, false);
assert!(!untrusted.contains_key("x-aether-tls-ja3"));
assert_eq!(
untrusted.get("user-agent").map(String::as_str),
Some("test-client")
);
let trusted = collect_report_context_original_headers(&headers, true);
assert_eq!(
trusted.get("x-aether-tls-ja3").map(String::as_str),
Some("spoofed-ja3")
);
}
#[test]
fn local_execution_report_context_records_request_origin_and_session_affinity() {
let auth_context = ExecutionRuntimeAuthContext {
@@ -324,12 +374,14 @@ mod tests {
request_origin: Some(RequestOrigin {
client_ip: Some("203.0.113.8".to_string()),
user_agent: Some("Claude-Code/1.0".to_string()),
forwarded_headers_trusted: false,
}),
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: Some(&client_session_affinity),
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
@@ -413,6 +465,7 @@ mod tests {
client_session_affinity: None,
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
@@ -474,12 +527,17 @@ mod tests {
original_headers: &original_headers,
request_path: None,
request_query_string: None,
request_origin: None,
request_origin: Some(RequestOrigin {
client_ip: None,
user_agent: None,
forwarded_headers_trusted: true,
}),
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: None,
routing_policy: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/files.rs
+15
View File
@@ -17,6 +17,7 @@ use crate::ai_serving::{
resolve_gemini_files_sync_spec as resolve_sync_spec, LocalGeminiFilesSpec,
};
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_routing_core::RoutingExecutionPolicy;
use self::decision::maybe_build_local_gemini_files_decision_payload_for_candidate;
use self::support::{
@@ -174,6 +175,13 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalGeminiFilesSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
@@ -212,6 +220,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalGeminiFilesSyncAttemptS
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalGeminiFilesStreamAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_stream_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs
+1
View File
@@ -109,6 +109,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream: spec_metadata.require_streaming,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/files/request.rs
+85 -1
View File
@@ -8,7 +8,10 @@ use crate::ai_serving::transport::{
GeminiFilesRequestBodyError,
};
use crate::ai_serving::GEMINI_FILES_UPLOAD_PLAN_KIND;
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::ai_serving::{
CandidateFailureDiagnostic, GatewayProviderTransportSnapshot, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
};
use crate::AppState;
use super::support::{
@@ -47,6 +50,26 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
let transport = &attempt.eligible.transport;
let effective_headers = input.effective_headers(&parts.headers);
if matches!(
spec_metadata.decision_kind,
GEMINI_FILES_GET_PLAN_KIND
| GEMINI_FILES_DELETE_PLAN_KIND
| GEMINI_FILES_DOWNLOAD_PLAN_KIND
) && !candidate_matches_owned_gemini_file_mapping(state, parts, input, attempt).await
{
mark_skipped_local_gemini_files_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"gemini_file_mapping_mismatch",
)
.await;
return None;
}
if let Some(skip_reason) =
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
{
@@ -191,3 +214,64 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
file_name,
})
}
async fn candidate_matches_owned_gemini_file_mapping(
state: &AppState,
parts: &http::request::Parts,
input: &LocalGeminiFilesDecisionInput,
attempt: &LocalGeminiFilesCandidateAttempt,
) -> bool {
let Some(file_name) = normalize_gemini_file_name_from_path(parts.uri.path()) else {
return false;
};
let user_id = input.auth_context.user_id.trim();
if user_id.is_empty() || !state.has_gemini_file_mapping_data_reader() {
return false;
}
let Ok(Some(mapping)) = state
.find_active_gemini_file_mapping_for_owner(
file_name.as_str(),
&attempt.eligible.transport.key.id,
user_id,
crate::clock::current_unix_secs(),
)
.await
else {
return false;
};
mapping.user_id.as_deref().map(str::trim) == Some(user_id)
&& mapping.key_id == attempt.eligible.transport.key.id
}
pub(crate) fn normalize_gemini_file_name_from_path(path: &str) -> Option<String> {
let suffix = path.strip_prefix("/v1beta/files/")?.trim_matches('/');
let suffix = suffix.strip_suffix(":download").unwrap_or(suffix).trim();
let suffix = suffix.strip_prefix("files/").unwrap_or(suffix).trim();
if suffix.is_empty() || suffix.contains('/') {
return None;
}
Some(format!("files/{suffix}"))
}
#[cfg(test)]
mod tests {
use super::normalize_gemini_file_name_from_path;
#[test]
fn normalizes_supported_gemini_file_object_paths() {
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/file-123"),
Some("files/file-123".to_string())
);
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/file-123:download"),
Some("files/file-123".to_string())
);
assert_eq!(
normalize_gemini_file_name_from_path("/v1beta/files/files/abc-123"),
Some("files/abc-123".to_string())
);
assert_eq!(normalize_gemini_file_name_from_path("/v1beta/files"), None);
}
}
apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs
+6
View File
@@ -108,6 +108,9 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -181,6 +184,9 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
+15
View File
@@ -19,6 +19,7 @@ use crate::ai_serving::{
resolve_local_image_sync_spec as resolve_sync_spec,
};
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_routing_core::RoutingExecutionPolicy;
use self::decision::maybe_build_local_openai_image_decision_payload_for_candidate;
use self::support::{
@@ -252,6 +253,13 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiImageSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
@@ -290,6 +298,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiImageSyncAttemptS
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiImageStreamAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_stream_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs
+1
View File
@@ -122,6 +122,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs
+20
View File
@@ -9,6 +9,7 @@ use crate::ai_serving::planner::candidate_preparation::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::pure::normalize_openai_image_request_with_options;
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
@@ -338,6 +339,25 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
let candidate = &attempt.eligible.candidate;
let transport = &attempt.eligible.transport;
let provider_api_format = "gemini:generate_content";
// The gemini:generate_content URL hook rewrites an Antigravity endpoint to
// /v1internal:, and this image path has no v1internal envelope to match it.
// Skip the candidate instead of posting a bare Gemini body that upstream
// would only reject.
if is_antigravity_provider_transport(transport) {
mark_skipped_local_openai_image_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
let effective_headers = input.effective_headers(&parts.headers);
let prepared_candidate = match prepare_header_authenticated_candidate(
apps/aether-gateway/src/ai_serving/planner/specialized/image/support.rs
+6
View File
@@ -127,6 +127,9 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await
{
@@ -201,6 +204,9 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await
{
apps/aether-gateway/src/ai_serving/planner/specialized/video.rs
+8
View File
@@ -16,6 +16,7 @@ use crate::ai_serving::{
LocalVideoCreateSpec,
};
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_routing_core::RoutingExecutionPolicy;
use self::decision::maybe_build_local_video_create_decision_payload_for_candidate;
use self::support::{
@@ -104,6 +105,13 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalVideoCreateSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs
+1
View File
@@ -90,6 +90,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs
+6
View File
@@ -133,6 +133,9 @@ pub(super) async fn list_local_video_create_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await
{
@@ -190,6 +193,9 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
input.routing_policy.as_ref(),
),
)
.await
{
apps/aether-gateway/src/ai_serving/planner/standard/codex/tests.rs
+4 -4
View File
@@ -505,7 +505,7 @@ fn projects_uuid_prompt_cache_identity_into_missing_session_headers() {
assert_eq!(headers.get("x-client-request-id"), None);
assert_eq!(
headers.get("user-agent"),
Some(&"codex_cli_rs/0.144.1".to_string())
Some(&"codex_cli_rs/0.153.3".to_string())
);
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
assert!(!headers.contains_key("version"));
@@ -615,7 +615,7 @@ fn injects_only_codex_client_headers_for_images_requests() {
);
assert_eq!(
headers.get("user-agent"),
Some(&"codex_cli_rs/0.144.1".to_string())
Some(&"codex_cli_rs/0.153.3".to_string())
);
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
assert!(!headers.contains_key("version"));
@@ -699,7 +699,7 @@ fn preserves_client_context_headers_and_enforces_codex_provider_identity() {
);
assert_eq!(
headers.get("user-agent"),
Some(&"codex_cli_rs/0.144.1".to_string())
Some(&"codex_cli_rs/0.153.3".to_string())
);
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
assert_eq!(
@@ -763,7 +763,7 @@ fn compact_projects_uuid_prompt_cache_identity_into_session_headers() {
assert_eq!(headers.get("x-client-request-id"), None);
assert_eq!(
headers.get("user-agent"),
Some(&"codex_cli_rs/0.144.1".to_string())
Some(&"codex_cli_rs/0.153.3".to_string())
);
assert_eq!(headers.get("originator"), Some(&"codex_cli_rs".to_string()));
assert!(!headers.contains_key("version"));
apps/aether-gateway/src/ai_serving/planner/standard/deepseek.rs
+107 -8
View File
@@ -15,11 +15,25 @@ pub(crate) fn is_deepseek_provider(provider_type: &str, base_url: &str) -> bool
host == "deepseek.com" || host.ends_with(".deepseek.com")
}
fn is_deepseek_model(provider_model: &str) -> bool {
let provider_model = provider_model.trim().to_ascii_lowercase();
let leaf = provider_model
.rsplit(['/', ':'])
.next()
.unwrap_or(provider_model.as_str());
leaf == "deepseek" || leaf.starts_with("deepseek-") || leaf.starts_with("deepseek_")
}
fn is_deepseek_upstream(provider_type: &str, base_url: &str, provider_model: &str) -> bool {
is_deepseek_provider(provider_type, base_url) || is_deepseek_model(provider_model)
}
pub(crate) fn openai_responses_reasoning_replay_policy(
provider_type: &str,
base_url: &str,
provider_model: &str,
) -> crate::ai_serving::OpenAiResponsesReasoningReplayPolicy {
if is_deepseek_provider(provider_type, base_url) {
if is_deepseek_upstream(provider_type, base_url, provider_model) {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
} else {
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
@@ -33,7 +47,11 @@ pub(crate) fn apply_deepseek_tool_call_thinking_compat(
provider_api_format: &str,
original_request_body: Option<&Value>,
) {
if !is_deepseek_provider(provider_type, base_url) {
let provider_model = provider_request_body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default();
if !is_deepseek_upstream(provider_type, base_url, provider_model) {
return;
}
@@ -302,11 +320,35 @@ mod tests {
));
assert!(!is_deepseek_provider("custom", "ftp://api.deepseek.com/v1"));
assert_eq!(
openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1"),
openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
);
assert_eq!(
openai_responses_reasoning_replay_policy("openai", "https://api.openai.com/v1"),
openai_responses_reasoning_replay_policy(
"openai",
"https://api.openai.com/v1",
"gpt-5.6-sol",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"custom",
"https://api.b.ai/v1",
"deepseek-v4-flash",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::DeepSeekOpaque
);
assert_eq!(
openai_responses_reasoning_replay_policy(
"custom",
"https://api.b.ai/v1",
"not-deepseek-compatible",
),
crate::ai_serving::OpenAiResponsesReasoningReplayPolicy::OpenAiItemIds
);
}
@@ -330,8 +372,11 @@ mod tests {
"input": reasoning_items.clone(),
"future_request_field": {"preserve": true}
});
let replay_policy =
openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1");
let replay_policy = openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
);
let mut provider_body = crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers_and_reasoning_replay_policy(
&request,
"openai:responses",
@@ -373,7 +418,11 @@ mod tests {
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy(
&mut deepseek,
"openai:responses",
openai_responses_reasoning_replay_policy("custom", "https://api.deepseek.com/v1"),
openai_responses_reasoning_replay_policy(
"custom",
"https://api.deepseek.com/v1",
"deepseek-v4-flash",
),
),
0
);
@@ -383,7 +432,11 @@ mod tests {
crate::ai_serving::strip_incompatible_openai_responses_reasoning_items_with_policy(
&mut openai,
"openai:responses",
openai_responses_reasoning_replay_policy("openai", "https://api.openai.com/v1"),
openai_responses_reasoning_replay_policy(
"openai",
"https://api.openai.com/v1",
"gpt-5.6-sol",
),
),
66
);
@@ -417,6 +470,52 @@ mod tests {
assert_eq!(body["messages"][1]["reasoning_content"], "");
}
#[test]
fn custom_relay_deepseek_model_adds_chat_thinking_compat() {
let mut body = json!({
"model": "deepseek-v4-flash",
"messages": [
{"role": "user", "content": "inspect the repository"},
{"role": "assistant", "content": null, "tool_calls": [{
"id": "call_1",
"type": "function",
"function": {"name": "inspect", "arguments": "{}"}
}]},
{"role": "tool", "tool_call_id": "call_1", "content": "done"}
]
});
apply_deepseek_tool_call_thinking_compat(
&mut body,
"custom",
"https://api.b.ai/v1",
"openai:chat",
None,
);
assert_eq!(body["thinking"]["type"], "enabled");
assert_eq!(body["messages"][1]["reasoning_content"], "");
}
#[test]
fn custom_relay_non_deepseek_model_is_not_rewritten() {
let original = json!({
"model": "not-deepseek-compatible",
"messages": [{"role": "assistant", "content": "done"}]
});
let mut body = original.clone();
apply_deepseek_tool_call_thinking_compat(
&mut body,
"custom",
"https://api.b.ai/v1",
"openai:chat",
None,
);
assert_eq!(body, original);
}
#[test]
fn openai_chat_deepseek_honors_disabled_thinking() {
let original = json!({"reasoning_effort": "none"});
apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs
+15
View File
@@ -18,6 +18,7 @@ use crate::ai_serving::planner::spec_metadata::{
};
use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use aether_routing_core::RoutingExecutionPolicy;
use super::candidates::{
build_local_standard_candidate_attempt_source, resolve_local_standard_decision_input,
@@ -177,6 +178,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
@@ -220,6 +228,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSour
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalStandardStreamAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_stream_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
+1
View File
@@ -142,6 +142,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
+169 -1
View File
@@ -4,6 +4,10 @@ use std::sync::Arc;
use aether_contracts::ResolvedTransportProfile;
use serde_json::Value;
use crate::ai_serving::planner::antigravity::{
build_antigravity_v1internal_provider_request, AntigravityV1InternalRequestError,
AntigravityV1InternalRequestInput, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
};
use crate::ai_serving::planner::candidate_preparation::{
prepare_header_authenticated_candidate, prepare_header_authenticated_candidate_from_auth,
OauthPreparationContext,
@@ -26,6 +30,7 @@ use crate::ai_serving::planner::standard::{
openai_provider_request_contract_failure_extra_data, openai_responses_reasoning_replay_policy,
request_body_build_failure_extra_data, request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::is_antigravity_provider_transport;
use crate::ai_serving::transport::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body,
is_kiro_claude_messages_transport, KiroProviderHeadersInput, KiroRequestAuth,
@@ -587,7 +592,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
state,
body_json,
spec_metadata.api_format,
provider_api_format,
@@ -597,6 +602,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
prepared_candidate.mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
@@ -836,6 +842,29 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await);
}
if normalized_provider_api_format == "gemini:generate_content"
&& is_antigravity_provider_transport(transport)
{
return Ok(build_antigravity_cross_format_payload_parts(
state,
parts,
trace_id,
body_json,
input,
attempt,
transport,
spec_metadata.api_format,
provider_api_format,
prepared_candidate.mapped_model,
prepared_candidate.auth_header,
prepared_candidate.auth_value,
provider_request_body,
upstream_is_stream,
redaction.redacted,
)
.await);
}
if normalized_provider_api_format == "gemini:generate_content"
&& is_gemini_cli_provider_transport(transport)
{
@@ -962,6 +991,145 @@ fn apply_transport_request_body_semantics(
)
}
#[allow(clippy::too_many_arguments)]
async fn build_antigravity_cross_format_payload_parts(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
original_body_json: &serde_json::Value,
input: &LocalStandardDecisionInput,
attempt: &LocalStandardCandidateAttempt,
transport: &Arc<GatewayProviderTransportSnapshot>,
client_api_format: &str,
provider_api_format: &str,
mapped_model: String,
auth_header: String,
auth_value: String,
gemini_request_body: Value,
upstream_is_stream: bool,
request_redacted: bool,
) -> Option<LocalStandardCandidatePayloadParts> {
let candidate = &attempt.eligible.candidate;
let effective_headers = input.effective_headers(&parts.headers);
let resolved =
match build_antigravity_v1internal_provider_request(AntigravityV1InternalRequestInput {
state,
parts,
transport,
trace_id,
mapped_model: &mapped_model,
provider_api_format,
auth_header: &auth_header,
auth_value: &auth_value,
request_headers: effective_headers,
original_request_body: original_body_json,
gemini_request_body: &gemini_request_body,
upstream_is_stream,
same_format: false,
})
.await
{
Ok(resolved) => resolved,
Err(AntigravityV1InternalRequestError::TransportUnsupported) => {
mark_skipped_local_standard_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_unsupported",
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::EnvelopeUnsupported) => {
mark_skipped_local_standard_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
original_body_json,
client_api_format,
provider_api_format,
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::UpstreamUrlUnavailable) => {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"upstream_url_missing",
CandidateFailureDiagnostic::upstream_url_missing(
client_api_format,
provider_api_format,
"standard_family_antigravity_url",
),
)
.await;
return None;
}
Err(AntigravityV1InternalRequestError::HeaderRulesApplyFailed) => {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_header_rules_apply_failed",
CandidateFailureDiagnostic::header_rules_apply_failed(
client_api_format,
provider_api_format,
"standard_family_antigravity_headers",
),
)
.await;
return None;
}
};
let mut provider_request_headers = resolved.headers.headers;
apply_codex_openai_special_headers(
&mut provider_request_headers,
&resolved.body,
effective_headers,
resolved.transport.provider.provider_type.as_str(),
provider_api_format,
Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(),
);
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
);
Some(LocalStandardCandidatePayloadParts {
auth_header: resolved.headers.auth_header,
auth_value: resolved.headers.auth_value,
mapped_model,
provider_api_format: provider_api_format.to_string(),
provider_request_body: resolved.body,
provider_request_headers,
upstream_url: resolved.upstream_url,
upstream_is_stream,
envelope_name: Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME),
transport: resolved.transport,
transport_profile: None,
request_redacted,
})
}
#[allow(clippy::too_many_arguments)]
async fn build_gemini_cli_cross_format_payload_parts(
state: &AppState,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
+1
View File
@@ -195,6 +195,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+3 -2
View File
@@ -159,6 +159,7 @@ fn finalize_openai_chat_provider_request_body(
openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model,
),
)
.err()
@@ -2740,7 +2741,7 @@ mod tests {
.provider_request_headers
.get("x-client-version")
.map(String::as_str),
Some("1.2.3")
Some("4.3.0")
);
assert_eq!(
payload
@@ -2760,7 +2761,7 @@ mod tests {
assert_eq!(payload.provider_request_body["model"], "gemini-2.5-pro");
assert_eq!(
payload.provider_request_body["userAgent"],
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)"
"vscode/1.X.X (Antigravity/4.3.0)"
);
assert_eq!(payload.provider_request_body["requestType"], "agent");
assert!(payload.provider_request_body.get("contents").is_none());
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
+8
View File
@@ -1,3 +1,4 @@
use aether_routing_core::RoutingExecutionPolicy;
use async_trait::async_trait;
use std::collections::VecDeque;
use tracing::warn;
@@ -119,6 +120,13 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
let select_started_at = std::time::Instant::now();
let selected = self.next_execution_attempt_with_target_select().await?;
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
+8
View File
@@ -1,3 +1,4 @@
use aether_routing_core::RoutingExecutionPolicy;
use async_trait::async_trait;
use tracing::warn;
@@ -92,6 +93,13 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiChatSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/stream.rs
+12 -27
View File
@@ -1,7 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use tracing::debug;
use super::super::{
augment_sync_report_context, build_ai_execution_plan_from_decision,
@@ -10,7 +9,6 @@ use super::super::{
AiStreamAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::provider_adaptation_requires_eventstream_accept;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
@@ -157,21 +155,16 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
return Ok(None);
};
let (url, url_source) = if let Some(upstream_url) =
take_non_empty_string(&mut payload.upstream_url)
{
(upstream_url, "upstream_url")
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
upstream_url
} else {
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
return Ok(None);
};
(
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
),
"upstream_base_url",
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
)
};
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
@@ -238,16 +231,7 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
tracing::debug!(
event_name = "local_openai_responses_stream_plan_built",
log_type = "debug",
request_id = %plan.request_id,
@@ -255,12 +239,13 @@ pub(crate) fn build_openai_responses_stream_plan_from_decision(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
downstream_path = %parts.uri.path(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = effective_upstream_is_stream,
apps/aether-gateway/src/ai_serving/planner/standard/openai/plan_builders/sync.rs
+12 -27
View File
@@ -1,7 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::RequestBody;
use tracing::debug;
use super::super::{
augment_sync_report_context, build_ai_execution_plan_from_decision,
@@ -10,7 +9,6 @@ use super::super::{
AiSyncAttempt,
};
use crate::ai_serving::planner::common::enforce_provider_body_stream_policy;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::transport::{
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, StandardPlanFallbackAcceptPolicy,
@@ -142,21 +140,16 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
let Some(auth_pair) = take_ai_upstream_auth_pair(&mut payload) else {
return Ok(None);
};
let (url, url_source) = if let Some(upstream_url) =
take_non_empty_string(&mut payload.upstream_url)
{
(upstream_url, "upstream_url")
let url = if let Some(upstream_url) = take_non_empty_string(&mut payload.upstream_url) {
upstream_url
} else {
let Some(upstream_base_url) = take_non_empty_string(&mut payload.upstream_base_url) else {
return Ok(None);
};
(
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
),
"upstream_base_url",
build_standard_plan_fallback_openai_responses_url(
&upstream_base_url,
parts.uri.query(),
compact,
)
};
let Some(provider_request_body_value) = payload.provider_request_body.take() else {
@@ -205,16 +198,7 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_decision_upstream_base_url = payload
.upstream_base_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_decision_upstream_url = payload
.upstream_url
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_plan_url = sanitize_upstream_url_for_log(plan.url.as_str());
debug!(
tracing::debug!(
event_name = "local_openai_responses_sync_plan_built",
log_type = "debug",
request_id = %plan.request_id,
@@ -222,12 +206,13 @@ pub(crate) fn build_openai_responses_sync_plan_from_decision(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
downstream_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&plan.url),
downstream_path = %parts.uri.path(),
downstream_query = ?log_downstream_query,
url_source,
decision_upstream_base_url = ?log_decision_upstream_base_url,
decision_upstream_url = ?log_decision_upstream_url,
plan_url = %log_plan_url,
client_api_format = %plan.client_api_format,
provider_api_format = %plan.provider_api_format,
upstream_is_stream = payload.upstream_is_stream,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
+6 -5
View File
@@ -3,7 +3,6 @@ use tracing::debug;
use crate::ai_serving::build_request_trace_proxy_value;
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision_with_websocket_mode;
use crate::ai_serving::planner::redaction::sanitize_upstream_url_for_log;
use crate::ai_serving::planner::report_context::{
build_local_execution_report_context, insert_native_client_envelope_name,
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
@@ -184,6 +183,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_session_affinity: input.client_session_affinity.as_ref(),
routing_policy: input.routing_policy.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)
@@ -204,12 +204,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
&resolved.transport,
);
let log_base_url = sanitize_upstream_url_for_log(resolved.transport.endpoint.base_url.as_str());
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(resolved.upstream_url.as_str());
debug!(
event_name = "local_openai_responses_decision_payload_built",
log_type = "debug",
@@ -226,9 +224,12 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_api_format = spec_metadata.api_format,
provider_api_format = %resolved.provider_api_format,
request_path = %parts.uri.path(),
request_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&resolved.upstream_url),
request_query = ?log_request_query,
upstream_base_url = %log_base_url,
upstream_url = %log_upstream_url,
upstream_is_stream = resolved.upstream_is_stream,
has_envelope = resolved.envelope_name.is_some(),
"gateway built local openai responses decision payload"
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
+8 -15
View File
@@ -24,7 +24,6 @@ use crate::ai_serving::planner::gemini_cli::{
};
use crate::ai_serving::planner::redaction::{
request_identity_response_encoding_when_redacted, resolve_provider_chat_pii_redaction,
sanitize_upstream_url_for_log,
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::standard::{
@@ -428,7 +427,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
state,
body_json,
spec_metadata.api_format,
provider_api_format,
@@ -438,6 +437,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
let reasoning_replay_policy = openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model.as_str(),
);
let redaction = resolve_provider_chat_pii_redaction(
state,
@@ -866,17 +866,10 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, provider_api_format);
let log_base_url = sanitize_upstream_url_for_log(transport.endpoint.base_url.as_str());
let log_custom_path = transport
.endpoint
.custom_path
.as_deref()
.map(sanitize_upstream_url_for_log);
let log_request_query = parts
.uri
.query()
.and_then(crate::ai_serving::api::sanitize_request_query_string);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_upstream_url_resolved",
@@ -892,12 +885,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts_with_
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
base_url = %log_base_url,
custom_path = ?log_custom_path,
request_path_and_query = %crate::ai_serving::pure::sanitize_request_path_and_query(
parts.uri.path(),
parts.uri.query(),
).unwrap_or_else(|| "/".to_string()),
upstream_origin = %crate::handlers::shared::security_log_url_origin(&upstream_url),
request_path = %parts.uri.path(),
request_query = ?log_request_query,
mapped_model = %mapped_model,
upstream_url = %log_upstream_url,
upstream_is_stream,
"gateway resolved local openai responses upstream url"
);
@@ -2010,8 +2005,6 @@ async fn build_kiro_openai_responses_payload_parts(
};
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(client_api_format, provider_api_format);
let log_upstream_url = sanitize_upstream_url_for_log(upstream_url.as_str());
debug!(
event_name = "local_openai_responses_kiro_upstream_url_resolved",
log_type = "debug",
@@ -2026,7 +2019,7 @@ async fn build_kiro_openai_responses_payload_parts(
provider_api_format = %provider_api_format,
execution_strategy = execution_strategy.as_str(),
conversion_mode = conversion_mode.as_str(),
upstream_url = %log_upstream_url,
upstream_origin = %crate::handlers::shared::security_log_url_origin(&upstream_url),
upstream_is_stream,
"gateway resolved local openai responses kiro upstream url"
);
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
+1
View File
@@ -1005,6 +1005,7 @@ pub(crate) async fn maybe_build_responses_websocket_decision(
reasoning_replay_policy: openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
mapped_model.as_str(),
),
model_directive_patch: input
.model_directive_policy
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs
+15
View File
@@ -1,3 +1,4 @@
use aether_routing_core::RoutingExecutionPolicy;
use async_trait::async_trait;
use tracing::warn;
@@ -161,6 +162,13 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_sync_attempt(attempt).await? {
@@ -204,6 +212,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAtte
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiResponsesStreamAttemptSource<'_> {
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
self.input
.routing_policy
.as_ref()
.map(|policy| policy.execution_policy)
}
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await? {
match self.build_stream_attempt(attempt).await? {
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
+18
View File
@@ -8,9 +8,13 @@ use crate::constants::{
API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
};
use crate::scheduler::candidate::SchedulerSkippedCandidate;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::GatewayError;
impl<'a> PlannerAppState<'a> {
/// `ordering_config` is the immutable scheduler snapshot derived from the
/// request's resolved routing policy.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates(
self,
api_format: &str,
@@ -21,6 +25,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: SchedulerOrderingConfig,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates(
self.app().data.as_ref(),
@@ -33,10 +38,12 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity,
now_unix_secs,
enable_model_directives,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
self,
api_format: &str,
@@ -47,6 +54,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: SchedulerOrderingConfig,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -64,10 +72,12 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs,
enable_model_directives,
None,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_operation(
self,
api_format: &str,
@@ -79,6 +89,7 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: SchedulerOrderingConfig,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -103,6 +114,7 @@ impl<'a> PlannerAppState<'a> {
attempt_now_unix_secs,
enable_model_directives,
request_operation,
ordering_config,
)
.await?;
@@ -123,6 +135,7 @@ impl<'a> PlannerAppState<'a> {
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
self,
api_format: &str,
@@ -132,6 +145,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: SchedulerOrderingConfig,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -148,10 +162,12 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
self,
candidate_api_format: &str,
@@ -160,6 +176,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: SchedulerOrderingConfig,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
@@ -176,6 +193,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot,
client_session_affinity,
attempt_now_unix_secs,
ordering_config,
)
.await?;
apps/aether-gateway/src/ai_serving/pure/mod.rs
+1
View File
@@ -177,6 +177,7 @@ pub(crate) use aether_ai_formats::{
api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream,
api_format_permission_covers, codex_responses_lite_tool_is_client_executed,
intersect_api_format_allowed_lists, is_embedding_api_format, is_rerank_api_format,
normalize_openai_responses_message_item_ids, openai_responses_message_item_id,
openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy, ApiOperation, ClientSurface,
apps/aether-gateway/src/ai_serving/response_history.rs
+160 -8
View File
@@ -2,14 +2,15 @@ use crate::ai_serving::{
hydrate_response_history, normalize_api_format_alias, record_converted_response_history,
response_history_is_loaded, response_history_storage_key, ResponseHistoryRecord,
};
use aether_runtime_state::RuntimeState;
use serde_json::Value;
use tracing::warn;
use crate::GatewayError;
use crate::{AppState, GatewayError};
const RESPONSE_HISTORY_SECRET_PURPOSE: &str = "openai-response-history";
pub(crate) async fn hydrate_openai_response_history(
runtime_state: &RuntimeState,
state: &AppState,
request: &Value,
client_api_format: &str,
provider_api_format: &str,
@@ -33,6 +34,7 @@ pub(crate) async fn hydrate_openai_response_history(
}
let storage_key = response_history_storage_key(previous_response_id, Some(history_scope));
let runtime_state = state.runtime_state();
let payload = runtime_state.kv_get(&storage_key).await.map_err(|error| {
warn!(
event_name = "openai_response_history_read_failed",
@@ -46,8 +48,24 @@ pub(crate) async fn hydrate_openai_response_history(
let Some(payload) = payload else {
return Ok(());
};
let Some(payload) = crate::handlers::shared::open_runtime_secret_payload(
state,
RESPONSE_HISTORY_SECRET_PURPOSE,
&payload,
) else {
let _ = runtime_state.kv_delete(&storage_key).await;
warn!(
event_name = "openai_response_history_decryption_failed",
log_type = "ops",
backend = runtime_state.backend_kind().as_str(),
"gateway rejected undecryptable shared OpenAI response history"
);
return Err(GatewayError::Internal(
"OpenAI response history decryption failed".to_string(),
));
};
if let Err(error) =
hydrate_response_history(previous_response_id, Some(history_scope), &payload)
hydrate_response_history(previous_response_id, Some(history_scope), payload.as_str())
{
let _ = runtime_state.kv_delete(&storage_key).await;
warn!(
@@ -65,11 +83,25 @@ pub(crate) async fn hydrate_openai_response_history(
}
pub(crate) async fn persist_response_history_record(
runtime_state: &RuntimeState,
state: &AppState,
record: ResponseHistoryRecord,
) {
let runtime_state = state.runtime_state();
let Some(sealed_payload) = crate::handlers::shared::seal_runtime_secret_payload(
state,
RESPONSE_HISTORY_SECRET_PURPOSE,
&record.payload,
) else {
warn!(
event_name = "openai_response_history_encryption_unavailable",
log_type = "ops",
backend = runtime_state.backend_kind().as_str(),
"gateway refused to persist unencrypted OpenAI response history"
);
return;
};
if let Err(error) = runtime_state
.kv_set(&record.storage_key, record.payload, Some(record.ttl))
.kv_set(&record.storage_key, sealed_payload, Some(record.ttl))
.await
{
warn!(
@@ -83,7 +115,7 @@ pub(crate) async fn persist_response_history_record(
}
pub(crate) async fn persist_converted_response_history(
runtime_state: &RuntimeState,
state: &AppState,
report_context: &Value,
response: Option<&Value>,
) {
@@ -91,6 +123,126 @@ pub(crate) async fn persist_converted_response_history(
return;
};
if let Some(record) = record_converted_response_history(report_context, response) {
persist_response_history_record(runtime_state, record).await;
persist_response_history_record(state, record).await;
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
use serde_json::json;
use sha2::{Digest, Sha256};
use super::{
hydrate_openai_response_history, persist_response_history_record, ResponseHistoryRecord,
};
use crate::{ai_serving::response_history_storage_key, data::GatewayDataState, AppState};
fn response_history_test_state() -> AppState {
AppState::new()
.expect("test state should build")
.with_data_state_for_tests(
GatewayDataState::disabled()
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_runtime_state(Arc::new(RuntimeState::memory(
MemoryRuntimeStateConfig::default(),
)))
}
fn response_history_payload(response_id: &str, scope: &str, marker: &str) -> String {
let expires_at_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
.saturating_add(3600);
json!({
"version": 1,
"response_id": response_id,
"scope_fingerprint": format!("{:x}", Sha256::digest(scope.trim().as_bytes())),
"expires_at_unix_secs": expires_at_unix_secs,
"transcript": [{"type": "message", "content": marker}],
})
.to_string()
}
#[tokio::test]
async fn response_history_is_encrypted_at_rest_and_hydrates() {
let state = response_history_test_state();
let response_id = "resp_gateway_encrypted_history_v1";
let scope = "response-history-encrypted-scope";
let marker = "private-response-history-marker";
let storage_key = response_history_storage_key(response_id, Some(scope));
let payload = response_history_payload(response_id, scope, marker);
persist_response_history_record(
&state,
ResponseHistoryRecord {
storage_key: storage_key.clone(),
payload,
ttl: Duration::from_secs(6 * 60 * 60),
},
)
.await;
let stored = state
.runtime_kv_get(&storage_key)
.await
.expect("history lookup should succeed")
.expect("history should be persisted");
assert!(crate::handlers::shared::runtime_secret_payload_is_sealed(
&stored
));
assert!(!stored.contains(marker));
hydrate_openai_response_history(
&state,
&json!({"previous_response_id": response_id}),
"openai:responses",
"openai:chat",
scope,
)
.await
.expect("encrypted history should hydrate");
assert!(crate::ai_serving::response_history_is_loaded(
response_id,
Some(scope)
));
}
#[tokio::test]
async fn response_history_reader_rejects_and_deletes_legacy_plaintext() {
let state = response_history_test_state();
let response_id = "resp_gateway_legacy_history_v1";
let scope = "response-history-legacy-scope";
let storage_key = response_history_storage_key(response_id, Some(scope));
let payload = response_history_payload(response_id, scope, "legacy-private-history");
state
.runtime_kv_setex(&storage_key, &payload, 6 * 60 * 60)
.await
.expect("legacy history should store");
let result = hydrate_openai_response_history(
&state,
&json!({"previous_response_id": response_id}),
"openai:responses",
"openai:chat",
scope,
)
.await;
assert!(result.is_err());
assert!(!crate::ai_serving::response_history_is_loaded(
response_id,
Some(scope)
));
assert!(state
.runtime_kv_get(&storage_key)
.await
.expect("history lookup should succeed")
.is_none());
}
}