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

Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota

Browse Source
This commit is contained in:
ZheFox
2026-09-03 10:46:49 +08:00
parent 45c840b8d3 cae9aa4134
commit f822df6cce
57 changed files with 1850 additions and 696 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+42 -79
View File
@@ -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);
@@ -481,10 +480,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)
}
@@ -1610,7 +1605,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 +1621,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 +1817,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 +1901,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 +2237,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,
}
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
+30 -31
View File
@@ -3,7 +3,7 @@ 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;
@@ -184,35 +184,16 @@ fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedul
}
}
/// Ordering config for a request. A resolved routing policy is authoritative
/// and is never merged with legacy system-config values; without a policy the
/// effective default (system-default routing group, then legacy keys) applies.
pub(crate) async fn scheduler_ordering_config_for_routing_policy(
state: PlannerAppState<'_>,
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
}
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,
Some(policy) => SchedulerOrderingConfig::from_routing_policy(policy),
None => read_scheduler_ordering_config_or_default(state).await,
}
}
@@ -231,14 +212,26 @@ 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);
}
@@ -378,6 +371,7 @@ 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,
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,6 +407,7 @@ 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,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -429,7 +424,10 @@ mod tests {
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 +445,7 @@ 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,
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)]),
apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs
+13 -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,18 @@ 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(state, routing_policy)
.await
.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
+13 -3
View File
@@ -174,6 +174,8 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
self.ranking_seed,
false,
self.request_operation,
self.routing_policy
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
@@ -1291,6 +1293,9 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
.then_some(self.client_session_affinity.as_ref())
.flatten(),
self.ranking_seed,
self.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let skipped_candidates = skipped_candidates
@@ -1884,6 +1889,7 @@ 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,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -1947,6 +1953,7 @@ 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,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
@@ -2657,14 +2664,16 @@ mod tests {
candidate_repository,
)
.with_encryption_key_for_tests("development-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 +2690,8 @@ 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,
ranking_overlay: Default::default(),
mutation_plan: Default::default(),
pool_policy_overrides: Default::default(),
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
+1
View File
@@ -1073,6 +1073,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,
},
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs
+9
View File
@@ -26,6 +26,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_api_request;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
use super::{
@@ -140,6 +141,10 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -246,6 +251,10 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
current_unix_secs(),
false,
spec.operation.map(|operation| operation.as_str()),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.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/report_context.rs
+13 -1
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;
@@ -59,6 +59,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,
@@ -124,6 +127,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,
@@ -330,6 +339,7 @@ mod tests {
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 +423,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,
@@ -480,6 +491,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: false,
has_envelope: false,
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/support.rs
+9
View File
@@ -26,6 +26,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
@@ -108,6 +109,10 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
current_unix_secs(),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
let outcome = materialize_local_execution_candidates_with_serving(
@@ -181,6 +186,10 @@ 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(),
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
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/support.rs
+9
View File
@@ -27,6 +27,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
@@ -127,6 +128,10 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
@@ -201,6 +206,10 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.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
+9
View File
@@ -29,6 +29,7 @@ use crate::ai_serving::{
};
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::clock::current_unix_secs;
use crate::scheduler::config::SchedulerOrderingConfig;
use crate::{AppState, GatewayError};
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalVideoCreateCandidateAttempt;
@@ -133,6 +134,10 @@ pub(super) async fn list_local_video_create_candidate_attempts(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.await
{
@@ -190,6 +195,10 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
input.client_session_affinity.as_ref(),
current_unix_secs(),
false,
input
.routing_policy
.as_ref()
.map(SchedulerOrderingConfig::from_routing_policy),
)
.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/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/responses/decision/payload.rs
+1
View File
@@ -184,6 +184,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)
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
+19
View File
@@ -8,9 +8,14 @@ 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 request's routing-policy derived scheduler
/// config (see `SchedulerOrderingConfig::from_routing_policy`). `None`
/// falls back to the runtime default.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates(
self,
api_format: &str,
@@ -21,6 +26,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
crate::scheduler::candidate::list_selectable_candidates(
self.app().data.as_ref(),
@@ -33,10 +39,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 +55,7 @@ impl<'a> PlannerAppState<'a> {
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -64,10 +73,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 +90,7 @@ impl<'a> PlannerAppState<'a> {
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -103,6 +115,7 @@ impl<'a> PlannerAppState<'a> {
attempt_now_unix_secs,
enable_model_directives,
request_operation,
ordering_config,
)
.await?;
@@ -123,6 +136,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 +146,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -148,10 +163,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 +177,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<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 +194,7 @@ impl<'a> PlannerAppState<'a> {
auth_snapshot,
client_session_affinity,
attempt_now_unix_secs,
ordering_config,
)
.await?;
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+26 -12
View File
@@ -756,19 +756,28 @@ impl<'a> PoolKeyCursor<'a> {
return None;
}
};
let api_format = self.group.candidate.endpoint_api_format.as_str();
rows.sort_by(|left, right| {
let left_priority = self
.routing_overlay
.as_ref()
.map_or(left.key_internal_priority, |overlay| {
overlay.key_priority(&left.key_id, left.key_internal_priority)
});
let right_priority = self
.routing_overlay
.as_ref()
.map_or(right.key_internal_priority, |overlay| {
overlay.key_priority(&right.key_id, right.key_internal_priority)
});
let left_priority = self.routing_overlay.as_ref().map_or(
left.key_internal_priority,
|overlay| {
overlay.key_priority_for_format(
&left.key_id,
api_format,
left.key_internal_priority,
)
},
);
let right_priority = self.routing_overlay.as_ref().map_or(
right.key_internal_priority,
|overlay| {
overlay.key_priority_for_format(
&right.key_id,
api_format,
right.key_internal_priority,
)
},
);
left_priority
.cmp(&right_priority)
.then(left.key_id.cmp(&right.key_id))
@@ -1987,11 +1996,13 @@ fn apply_pool_orchestration(
orchestration: PoolCandidateOrchestration,
) -> EligibleLocalExecutionCandidate {
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
let sticky_key_attempts = candidate.orchestration.sticky_key_attempts;
candidate.orchestration = LocalExecutionCandidateMetadata {
candidate_group_id: orchestration.candidate_group_id,
pool_key_index: orchestration.pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch,
sticky_key_attempts,
};
candidate
}
@@ -2244,6 +2255,7 @@ mod tests {
pool_key_index: Some(0),
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
}
);
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
@@ -2261,6 +2273,7 @@ mod tests {
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
}
);
}
@@ -5157,6 +5170,7 @@ mod tests {
priority_mode: RoutingSetPriorityMode::Provider,
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
ranking_overlay: RankingOverlay {
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
..RankingOverlay::default()
apps/aether-gateway/src/dispatch/refs.rs
+1
View File
@@ -208,6 +208,7 @@ mod tests {
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
sticky_key_attempts: None,
},
ranking: None,
}
apps/aether-gateway/src/executor/candidate_loop.rs
+33 -9
View File
@@ -252,6 +252,10 @@ where
Ok(())
}
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
}
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
record_provider_transfer_attempt_failed(
self.state,
@@ -788,18 +792,31 @@ where
{
let mut last_attempted = None;
let mut fallback_response = None;
// A same-key retry derived after a candidate-scoped failure runs before
// the source is asked for the next candidate.
let mut pending_same_key_retry: Option<Attempt> = None;
loop {
let next_started_at = std::time::Instant::now();
let next_attempt =
next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout)
let attempt = match pending_same_key_retry.take() {
Some(attempt) => attempt,
None => {
let next_started_at = std::time::Instant::now();
let next_attempt = next_execution_attempt_with_timeout(
source,
trace_id,
plan_kind,
planning_timeout,
)
.await?;
observe_gateway_stage_ms(
"stream_candidate_next",
next_started_at.elapsed().as_millis() as u64,
);
let Some(attempt) = next_attempt else {
break;
observe_gateway_stage_ms(
"stream_candidate_next",
next_started_at.elapsed().as_millis() as u64,
);
let Some(attempt) = next_attempt else {
break;
};
attempt
}
};
if port.should_skip_attempt(&attempt).await? {
let provider_id = attempt.execution_plan().provider_id.clone();
@@ -839,6 +856,9 @@ where
if attempt_fallback_response.is_some() {
fallback_response = attempt_fallback_response;
}
if scope == AiAttemptRetryScope::Candidate {
pending_same_key_retry = port.next_same_key_retry(&attempt).await?;
}
apply_attempt_retry_scope(source, &attempt, scope).await?;
}
}
@@ -952,6 +972,10 @@ where
Ok(())
}
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
}
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
record_provider_transfer_attempt_failed(
self.state,
apps/aether-gateway/src/executor/orchestration.rs
+105
View File
@@ -1663,6 +1663,22 @@ mod tests {
candidate_index: u32,
endpoint_id: &str,
candidate_id: &str,
) -> AiSyncAttempt {
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
candidate_index,
endpoint_id,
candidate_id,
1,
)
}
/// `sticky_key_attempts` is pinned so these tests exercise candidate
/// failover; the default same-key retry is covered separately.
fn test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
candidate_index: u32,
endpoint_id: &str,
candidate_id: &str,
sticky_key_attempts: u32,
) -> AiSyncAttempt {
AiSyncAttempt {
plan: test_openai_image_heartbeat_plan(endpoint_id, candidate_id),
@@ -1670,6 +1686,7 @@ mod tests {
report_context: Some(json!({
"candidate_index": candidate_index,
"retry_index": 0,
"sticky_key_attempts": sticky_key_attempts,
})),
}
}
@@ -1828,6 +1845,9 @@ mod tests {
report_context: Some(json!({
"candidate_index": candidate_index,
"retry_index": 0,
// Pin to a single attempt so this helper exercises candidate
// failover rather than the default same-key retry.
"sticky_key_attempts": 1,
"client_api_format": client_api_format,
"provider_api_format": client_api_format,
})),
@@ -1980,6 +2000,90 @@ mod tests {
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
}
#[tokio::test]
async fn openai_image_sync_heartbeat_retries_sticky_key_lazily_before_failover() {
let seen_plans = Arc::new(std::sync::Mutex::new(Vec::<(String, Option<String>)>::new()));
let seen_plans_for_override = Arc::clone(&seen_plans);
let state = AppState::new()
.expect("state should build")
.with_execution_runtime_sync_override_for_tests(move |plan| {
seen_plans_for_override
.lock()
.expect("mutex should lock")
.push((plan.endpoint_id.clone(), plan.candidate_id.clone()));
if plan.endpoint_id == "endpoint-retry" {
Ok(test_openai_image_execution_result(
plan,
StatusCode::TOO_MANY_REQUESTS.as_u16(),
json!({"error": {"message": "retry this candidate"}}),
))
} else {
Ok(test_openai_image_execution_result(
plan,
StatusCode::OK.as_u16(),
json!({"data": [{"b64_json": "second-candidate"}]}),
))
}
});
// Three total attempts on the sticky key; only one attempt is
// materialized up front, the other two are derived after each failure.
let attempts = vec![
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
0,
"endpoint-retry",
"candidate-retry",
3,
),
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
1,
"endpoint-success",
"candidate-success",
3,
),
];
let outcome = execute_openai_image_sync_heartbeat_attempts(
state,
"/v1/images/generations".to_string(),
"trace-image-heartbeat-sticky-retry".to_string(),
test_openai_image_heartbeat_decision(),
TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(),
attempts,
ProviderTransferTracker::default(),
Instant::now(),
)
.await
.expect("heartbeat attempts should execute");
let LocalExecutionRequestOutcome::Responded(response) = outcome else {
panic!("second candidate should return a response");
};
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
let seen_plans = seen_plans.lock().expect("mutex should lock").clone();
assert_eq!(
seen_plans
.iter()
.map(|(endpoint_id, _)| endpoint_id.as_str())
.collect::<Vec<_>>(),
[
"endpoint-retry",
"endpoint-retry",
"endpoint-retry",
"endpoint-success"
]
);
let sticky_candidate_ids = seen_plans[..3]
.iter()
.map(|(_, candidate_id)| candidate_id.clone())
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
sticky_candidate_ids.len(),
3,
"each derived same-key retry must carry a fresh candidate id"
);
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
}
#[tokio::test]
async fn openai_image_sync_heartbeat_honors_provider_transfer_limit() {
let call_count = Arc::new(AtomicUsize::new(0));
@@ -2013,6 +2117,7 @@ mod tests {
attempt.report_context = Some(json!({
"candidate_index": index,
"retry_index": 0,
"sticky_key_attempts": 1,
"local_failover_policy": {
"max_transfer_count": 1,
"max_transfer_timeout_seconds": 0
apps/aether-gateway/src/handlers/admin/model/routing.rs
+7 -20
View File
@@ -67,27 +67,14 @@ pub(crate) async fn build_admin_global_model_routing_payload(
.push(key);
}
let scheduling_mode = state
.read_system_config_json_value("scheduling_mode")
// Effective default scheduling: system-default routing group first, then
// legacy system-config keys.
let ordering_config = crate::scheduler::config::read_scheduler_ordering_config(state.app())
.await
.ok()
.flatten()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "cache_affinity".to_string());
let priority_mode = state
.read_system_config_json_value("provider_priority_mode")
.await
.ok()
.flatten()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "provider".to_string());
let keep_priority_on_conversion = state
.read_system_config_json_value("keep_priority_on_conversion")
.await
.ok()
.flatten()
.and_then(|value| value.as_bool())
.unwrap_or(false);
.unwrap_or_default();
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
let priority_mode = ordering_config.priority_mode_str().to_string();
let keep_priority_on_conversion = ordering_config.keep_priority_on_conversion;
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_store.rs
+4 -10
View File
@@ -264,16 +264,10 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
pub(super) async fn build_admin_monitoring_cache_snapshot(
state: &AdminAppState<'_>,
) -> Result<AdminMonitoringCacheSnapshot, GatewayError> {
let scheduling_mode = state
.read_system_config_json_value("scheduling_mode")
.await?
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "cache_affinity".to_string());
let provider_priority_mode = state
.read_system_config_json_value("provider_priority_mode")
.await?
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "provider".to_string());
let ordering_config =
crate::scheduler::config::read_scheduler_ordering_config(state.app()).await?;
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
let provider_priority_mode = ordering_config.priority_mode_str().to_string();
let now = chrono::Utc::now();
let usage_summary = if state.has_usage_data_reader() {
apps/aether-gateway/src/main.rs
+16
View File
@@ -2096,6 +2096,22 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
);
}
state.bootstrap_admin_from_env().await?;
match state.ensure_system_default_routing_group().await {
Ok(Some(group)) => {
info!(
group_id = %group.id,
group_name = %group.name,
"created system default routing group from legacy scheduler config"
);
}
Ok(None) => {}
Err(err) => {
warn!(
error = %err,
"failed to bootstrap system default routing group; scheduler falls back to legacy system config"
);
}
}
match state.prewarm_chat_pii_redaction_runtime_config().await {
Ok(enabled) => {
info!(
apps/aether-gateway/src/orchestration/attempt.rs
+180 -302
View File
@@ -1,8 +1,9 @@
use aether_ai_serving::{AiExecutionAttempt, STICKY_KEY_ATTEMPTS_REPORT_FIELD};
use aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS;
use aether_runtime_state::RuntimeLockLease;
use aether_scheduler_core::parse_request_candidate_report_context;
use serde_json::Value;
use crate::provider_transport::GatewayProviderTransportSnapshot;
use uuid::Uuid;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ExecutionAttemptIdentity {
@@ -32,6 +33,9 @@ pub(crate) struct LocalExecutionCandidateMetadata {
pub(crate) pool_key_index: Option<u32>,
pub(crate) pool_key_lease: Option<RuntimeLockLease>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
/// Routing-policy `sticky_key_attempts` in effect for this request. `None`
/// means the policy default applies.
pub(crate) sticky_key_attempts: Option<u32>,
}
pub(crate) const SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD: &str = "scheduler_affinity_epoch";
@@ -42,6 +46,11 @@ pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token
pub(crate) const POOL_KEY_LEASE_FENCING_REPORT_FIELD: &str = "pool_key_lease_fencing_token";
pub(crate) const POOL_KEY_LEASE_TTL_MS_REPORT_FIELD: &str = "pool_key_lease_ttl_ms";
/// Pool-expanded keys encode `pool_key_index * STRIDE + retry_index` into the
/// persisted `retry_index` so a pool group's keys stay ordered in one candidate
/// slot. Same-key retries on a pool key are therefore bounded by the stride.
pub(crate) const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
pub(crate) fn attempt_identity_from_report_context(
report_context: Option<&Value>,
) -> Option<ExecutionAttemptIdentity> {
@@ -72,6 +81,10 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
scheduler_affinity_epoch: report_context
.and_then(|value| value.get(SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD))
.and_then(Value::as_u64),
sticky_key_attempts: report_context
.and_then(|value| value.get(STICKY_KEY_ATTEMPTS_REPORT_FIELD))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
}
}
@@ -140,58 +153,55 @@ fn pool_key_lease_from_report_context(report_context: Option<&Value>) -> Option<
})
}
pub(crate) fn build_local_attempt_identities(
candidate_index: u32,
transport: &GatewayProviderTransportSnapshot,
) -> Vec<ExecutionAttemptIdentity> {
let attempt_slots = local_attempt_slot_count(transport);
(0..attempt_slots)
.map(|retry_index| ExecutionAttemptIdentity::new(candidate_index, retry_index))
.collect()
/// Retry index of the next same-key attempt, or `None` when the sticky-key
/// budget for this candidate is used up.
///
/// Only the first-ranked candidate (index `0`, the cache-affinity sticky key)
/// is retried on the same key; every later candidate gets exactly one attempt
/// so that once failover has started it keeps advancing. `sticky_key_attempts`
/// is the *total* attempt count on that key: `2` means one retry, `0` and `1`
/// mean none. There is no upper bound: attempts are derived one at a time
/// after each failure, never materialized ahead of time.
///
/// Inside a pool group only the first key (`pool_key_index == 0`) is treated
/// as sticky, and its retries stay below `POOL_KEY_RETRY_INDEX_STRIDE` so the
/// encoded retry index never collides with the next pool key.
pub(crate) fn next_same_key_retry_index(
identity: ExecutionAttemptIdentity,
sticky_key_attempts: Option<u32>,
) -> Option<u32> {
if identity.candidate_index != 0 {
return None;
}
let pool_limit = match identity.pool_key_index {
None => u32::MAX,
Some(0) => POOL_KEY_RETRY_INDEX_STRIDE,
Some(_) => return None,
};
let budget = sticky_key_attempts.unwrap_or(DEFAULT_STICKY_KEY_ATTEMPTS);
let attempts_so_far = identity.retry_index.checked_add(1)?;
if attempts_so_far >= budget || attempts_so_far >= pool_limit {
return None;
}
Some(attempts_so_far)
}
pub(crate) fn local_attempt_slot_count(transport: &GatewayProviderTransportSnapshot) -> u32 {
local_attempt_slots_from_transport(transport).unwrap_or(1)
}
/// For endpoint/provider table fields, `2` is the legacy admin default and is
/// treated as "not explicitly configured" so existing local-execution behaviour
/// (one attempt slot per candidate) stays unchanged. Values `0`, `1`, and `>2`
/// are treated as explicit.
const LEGACY_DEFAULT_MAX_RETRIES: u32 = 2;
/// Upper bound on local attempt slots. This is intentionally stricter than
/// admin max_retries validation to prevent unbounded pre-materialization from
/// arbitrarily large JSON config values.
const MAX_LOCAL_ATTEMPT_SLOTS: u32 = 99;
fn local_attempt_slots_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<u32> {
let rules = transport
.provider
.config
.as_ref()
.and_then(|config| config.get("failover_rules"))
.and_then(Value::as_object);
rules
.and_then(|value| value.get("max_retries"))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.or_else(|| {
transport
.endpoint
.max_retries
.and_then(|value| u32::try_from(value).ok())
.filter(|&value| value != LEGACY_DEFAULT_MAX_RETRIES)
})
.or_else(|| {
transport
.provider
.max_retries
.and_then(|value| u32::try_from(value).ok())
.filter(|&value| value != LEGACY_DEFAULT_MAX_RETRIES)
})
.map(|value| value.clamp(1, MAX_LOCAL_ATTEMPT_SLOTS))
/// Derive the next same-key attempt for `attempt` after a candidate-scoped
/// failure, reading the attempt identity and sticky budget from its report
/// context. Returns `None` when no further same-key retry is allowed.
pub(crate) fn next_same_key_retry_attempt<A: AiExecutionAttempt>(attempt: &A) -> Option<A> {
let owned_report_context = attempt
.report_context_ref()
.is_none()
.then(|| attempt.report_context())
.flatten();
let report_context = attempt
.report_context_ref()
.or(owned_report_context.as_ref());
let identity = attempt_identity_from_report_context(report_context)?;
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
let retry_index = next_same_key_retry_index(identity, metadata.sticky_key_attempts)?;
attempt.with_same_key_retry(retry_index, Uuid::new_v4().to_string())
}
#[cfg(test)]
@@ -199,277 +209,143 @@ mod tests {
use serde_json::json;
use super::{
attempt_identity_from_report_context, build_local_attempt_identities,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
};
use crate::provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
attempt_identity_from_report_context,
local_execution_candidate_metadata_from_report_context, next_same_key_retry_attempt,
next_same_key_retry_index, ExecutionAttemptIdentity, LocalExecutionCandidateMetadata,
POOL_KEY_RETRY_INDEX_STRIDE,
};
use aether_ai_serving::{AiExecutionAttempt, AiSyncAttempt};
use aether_runtime_state::RuntimeLockLease;
fn sample_transport(
provider_max_retries: Option<i32>,
endpoint_max_retries: Option<i32>,
provider_config: Option<serde_json::Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "OpenAI".to_string(),
provider_type: "llm".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: provider_max_retries,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: provider_config,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:chat".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://example.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: endpoint_max_retries,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "primary".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
#[test]
fn first_candidate_defaults_to_one_same_key_retry() {
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), None),
Some(1)
);
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 1), None),
None
);
}
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
#[test]
fn first_candidate_uses_policy_sticky_key_attempts_without_upper_bound() {
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 2), Some(3)),
None
);
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 1), Some(3)),
Some(2)
);
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 4_999), Some(10_000)),
Some(5_000)
);
}
#[test]
fn zero_and_one_mean_single_attempt() {
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), Some(0)),
None
);
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), Some(1)),
None
);
}
#[test]
fn failover_candidates_never_retry_on_the_same_key() {
for candidate_index in 1..5 {
assert_eq!(
next_same_key_retry_index(ExecutionAttemptIdentity::new(candidate_index, 0), None),
None
);
assert_eq!(
next_same_key_retry_index(
ExecutionAttemptIdentity::new(candidate_index, 0),
Some(50)
),
None
);
}
}
#[test]
fn build_local_attempt_identities_defaults_to_single_attempt() {
let identities = build_local_attempt_identities(3, &sample_transport(None, None, None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
}
#[test]
fn build_local_attempt_identities_prefer_failover_rules_over_endpoint_and_provider() {
let identities = build_local_attempt_identities(
1,
&sample_transport(
Some(5),
Some(4),
Some(json!({
"failover_rules": {
"max_retries": 2
}
})),
),
);
fn pool_groups_only_retry_their_first_key_within_the_stride() {
let first_pool_key = ExecutionAttemptIdentity::new(0, 0).with_pool_key_index(Some(0));
assert_eq!(next_same_key_retry_index(first_pool_key, Some(3)), Some(1));
let at_stride_limit = ExecutionAttemptIdentity::new(0, POOL_KEY_RETRY_INDEX_STRIDE - 1)
.with_pool_key_index(Some(0));
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(1, 0),
ExecutionAttemptIdentity::new(1, 1),
]
next_same_key_retry_index(at_stride_limit, Some(10_000)),
None
);
}
#[test]
fn build_local_attempt_identities_falls_back_to_endpoint_max_retries() {
let identities =
build_local_attempt_identities(2, &sample_transport(Some(5), Some(3), None));
let second_pool_key = ExecutionAttemptIdentity::new(0, POOL_KEY_RETRY_INDEX_STRIDE)
.with_pool_key_index(Some(1));
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(2, 0),
ExecutionAttemptIdentity::new(2, 1),
ExecutionAttemptIdentity::new(2, 2),
]
next_same_key_retry_index(second_pool_key, Some(10_000)),
None
);
}
#[test]
fn build_local_attempt_identities_falls_back_to_provider_max_retries() {
let identities = build_local_attempt_identities(0, &sample_transport(Some(4), None, None));
fn next_same_key_retry_attempt_rewrites_candidate_id_and_retry_index() {
let attempt = AiSyncAttempt {
plan: aether_contracts::ExecutionPlan {
request_id: "trace-1".to_string(),
candidate_id: Some("candidate-a".to_string()),
provider_name: None,
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com".to_string(),
headers: Default::default(),
content_type: None,
content_encoding: None,
body: aether_contracts::RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: None,
proxy: None,
transport_profile: None,
timeouts: None,
},
report_kind: None,
report_context: Some(json!({
"candidate_id": "candidate-a",
"candidate_index": 0,
"retry_index": 0,
"sticky_key_attempts": 2,
})),
};
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(0, 0),
ExecutionAttemptIdentity::new(0, 1),
ExecutionAttemptIdentity::new(0, 2),
ExecutionAttemptIdentity::new(0, 3),
]
let retry = next_same_key_retry_attempt(&attempt).expect("one same-key retry remains");
let retry_candidate_id = retry.plan.candidate_id.clone().expect("fresh candidate id");
assert_ne!(retry_candidate_id, "candidate-a");
assert_eq!(retry.plan.key_id, "key-1");
let context = retry.report_context_ref().expect("context retained");
assert_eq!(context["candidate_id"], json!(retry_candidate_id));
assert_eq!(context["retry_index"], json!(1));
assert_eq!(context["candidate_index"], json!(0));
assert!(
next_same_key_retry_attempt(&retry).is_none(),
"budget of 2 attempts is exhausted after one retry"
);
}
#[test]
fn build_local_attempt_identities_endpoint_overrides_provider() {
let identities =
build_local_attempt_identities(7, &sample_transport(Some(10), Some(3), None));
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(7, 0),
ExecutionAttemptIdentity::new(7, 1),
ExecutionAttemptIdentity::new(7, 2),
]
);
}
#[test]
fn build_local_attempt_identities_default_two_treated_as_unset() {
let identities =
build_local_attempt_identities(5, &sample_transport(Some(2), Some(2), None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(5, 0)]);
}
#[test]
fn build_local_attempt_identities_endpoint_two_falls_back_to_provider_ten() {
let identities =
build_local_attempt_identities(1, &sample_transport(Some(10), Some(2), None));
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(1, 0),
ExecutionAttemptIdentity::new(1, 1),
ExecutionAttemptIdentity::new(1, 2),
ExecutionAttemptIdentity::new(1, 3),
ExecutionAttemptIdentity::new(1, 4),
ExecutionAttemptIdentity::new(1, 5),
ExecutionAttemptIdentity::new(1, 6),
ExecutionAttemptIdentity::new(1, 7),
ExecutionAttemptIdentity::new(1, 8),
ExecutionAttemptIdentity::new(1, 9),
]
);
}
#[test]
fn build_local_attempt_identities_failover_rules_zero_produces_one_slot() {
let identities = build_local_attempt_identities(
1,
&sample_transport(
Some(5),
Some(4),
Some(json!({
"failover_rules": {
"max_retries": 0
}
})),
),
);
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(1, 0)]);
}
#[test]
fn build_local_attempt_identities_endpoint_zero_produces_one_slot() {
let identities =
build_local_attempt_identities(3, &sample_transport(Some(5), Some(0), None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
}
#[test]
fn build_local_attempt_identities_provider_zero_produces_one_slot() {
let identities = build_local_attempt_identities(3, &sample_transport(Some(0), None, None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
}
#[test]
fn build_local_attempt_identities_provider_ten_creates_ten_slots() {
let identities = build_local_attempt_identities(2, &sample_transport(Some(10), None, None));
assert_eq!(identities.len(), 10);
assert_eq!(identities[0], ExecutionAttemptIdentity::new(2, 0));
assert_eq!(identities[9], ExecutionAttemptIdentity::new(2, 9));
}
#[test]
fn build_local_attempt_identities_failover_rules_over_limit_clamped_to_max() {
let identities = build_local_attempt_identities(
0,
&sample_transport(
Some(3),
Some(5),
Some(json!({
"failover_rules": {
"max_retries": 1000
}
})),
),
);
assert_eq!(identities.len(), 99);
}
#[test]
fn build_local_attempt_identities_failover_rules_u32_max_clamped_to_max() {
let identities = build_local_attempt_identities(
0,
&sample_transport(
None,
None,
Some(json!({
"failover_rules": {
"max_retries": u32::MAX
}
})),
),
);
assert_eq!(identities.len(), 99);
}
#[test]
fn build_local_attempt_identities_endpoint_over_limit_clamped_to_max() {
let identities =
build_local_attempt_identities(0, &sample_transport(None, Some(2000), None));
assert_eq!(identities.len(), 99);
}
#[test]
fn build_local_attempt_identities_provider_over_limit_clamped_to_max() {
let identities =
build_local_attempt_identities(0, &sample_transport(Some(5000), None, None));
assert_eq!(identities.len(), 99);
}
#[test]
fn parse_attempt_identity_from_report_context_reads_candidate_and_retry_indices() {
let identity = attempt_identity_from_report_context(Some(&json!({
@@ -499,6 +375,7 @@ mod tests {
"pool_key_lease_token": "gateway-1:token-1",
"pool_key_lease_fencing_token": 7,
"pool_key_lease_ttl_ms": 900000,
"sticky_key_attempts": 3,
})));
assert_eq!(
@@ -514,6 +391,7 @@ mod tests {
ttl_ms: 900000,
}),
scheduler_affinity_epoch: None,
sticky_key_attempts: Some(3),
}
);
}
apps/aether-gateway/src/orchestration/mod.rs
+4 -5
View File
@@ -20,11 +20,10 @@ pub(crate) use self::adaptive::{
LocalAdaptiveRateLimitProjection, LocalAdaptiveSuccessProjection,
};
pub(crate) use self::attempt::{
attempt_identity_from_report_context, build_local_attempt_identities,
insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata, ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD,
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
attempt_identity_from_report_context, insert_pool_key_lease_report_context_fields,
local_execution_candidate_metadata_from_report_context, next_same_key_retry_attempt,
ExecutionAttemptIdentity, LocalExecutionCandidateMetadata, POOL_KEY_RETRY_INDEX_STRIDE,
ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
};
pub(crate) use self::classifier::{
classify_anthropic_failure_disposition, classify_failure_disposition, classify_local_failover,
apps/aether-gateway/src/routing/resolver.rs
+21 -16
View File
@@ -1,7 +1,7 @@
use aether_routing_core::{
resolve_routing_policy, MutationPlan, RankingOverlay, ResolvedRoutingPolicy,
RoutingGroupConfig, RoutingPolicyInput, RoutingRulePhase, RoutingSchedulingMode,
RoutingSetPriorityMode,
RoutingDefaultPolicy, RoutingGroupConfig, RoutingPolicyInput, RoutingRulePhase,
RoutingSchedulingMode, RoutingSetPriorityMode, DEFAULT_STICKY_KEY_ATTEMPTS,
};
use http::StatusCode;
use serde_json::Value;
@@ -81,9 +81,7 @@ pub(crate) fn resolve_gateway_routing_policy(
pub(crate) fn resolve_gateway_static_default_routing_policy(
input: GatewayStaticRoutingPolicyInput<'_>,
) -> Result<Option<ResolvedRoutingPolicy>, GatewayError> {
let Some((priority_mode, scheduling_mode, keep_priority_on_conversion)) =
static_default_policy_fields(input.group_config_json)?
else {
let Some(default_policy) = static_default_policy_fields(input.group_config_json)? else {
return Ok(None);
};
@@ -93,9 +91,10 @@ pub(crate) fn resolve_gateway_static_default_routing_policy(
selection_source: input.selection_source.to_string(),
requested_model: input.requested_model.to_string(),
resolved_model: input.resolved_model.to_string(),
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
priority_mode: default_policy.priority_mode,
scheduling_mode: default_policy.scheduling_mode,
keep_priority_on_conversion: default_policy.keep_priority_on_conversion,
sticky_key_attempts: default_policy.sticky_key_attempts,
ranking_overlay: RankingOverlay::default(),
mutation_plan: MutationPlan::default(),
pool_policy_overrides: BTreeMap::new(),
@@ -105,7 +104,7 @@ pub(crate) fn resolve_gateway_static_default_routing_policy(
fn static_default_policy_fields(
config_json: &Value,
) -> Result<Option<(RoutingSetPriorityMode, RoutingSchedulingMode, bool)>, GatewayError> {
) -> Result<Option<RoutingDefaultPolicy>, GatewayError> {
let Some(object) = config_json.as_object() else {
return Ok(None);
};
@@ -117,11 +116,7 @@ fn static_default_policy_fields(
}
let Some(default_policy) = object.get("default_policy") else {
return Ok(Some((
RoutingSetPriorityMode::default(),
RoutingSchedulingMode::default(),
false,
)));
return Ok(Some(RoutingDefaultPolicy::default()));
};
let Some(default_policy) = default_policy.as_object() else {
return Ok(None);
@@ -141,12 +136,22 @@ fn static_default_policy_fields(
})?,
None => false,
};
let sticky_key_attempts = match default_policy.get("sticky_key_attempts") {
Some(value) => value
.as_u64()
.and_then(|value| u32::try_from(value).ok())
.ok_or_else(|| {
invalid_routing_group_config("sticky_key_attempts must be a non-negative integer")
})?,
None => DEFAULT_STICKY_KEY_ATTEMPTS,
};
Ok(Some((
Ok(Some(RoutingDefaultPolicy {
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
)))
sticky_key_attempts,
}))
}
fn routing_array_field_is_missing_or_empty(
apps/aether-gateway/src/scheduler/candidate/mod.rs
+41 -18
View File
@@ -1,7 +1,9 @@
use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons_and_ordering,
collect_selectable_enumerated_candidates_with_skip_reasons,
resolve_preselection_ordering_config,
};
use super::config::SchedulerOrderingConfig;
use super::state::SchedulerRuntimeState;
mod affinity;
@@ -53,6 +55,10 @@ enum RequiredCapabilityMatchMode {
Exclusive,
}
/// `ordering_config` carries the request's routing-policy derived scheduler
/// config. `None` falls back to the runtime default (system-default routing
/// group, then legacy system-config keys).
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -64,6 +70,7 @@ pub(crate) async fn list_selectable_candidates(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates(
selection_row_source,
@@ -76,6 +83,7 @@ pub(crate) async fn list_selectable_candidates(
client_session_affinity,
now_unix_secs,
enable_model_directives,
ordering_config,
)
.await
}
@@ -87,6 +95,7 @@ pub(crate) fn is_exact_all_skipped_by_auth_limit(
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -98,6 +107,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -105,7 +115,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
),
GatewayError,
> {
collect_selectable_candidates_with_skip_reasons(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
@@ -117,10 +127,12 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
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(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -133,6 +145,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -140,7 +153,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
),
GatewayError,
> {
collect_selectable_candidates_with_skip_reasons(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
@@ -152,6 +165,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
now_unix_secs,
enable_model_directives,
request_operation,
ordering_config,
)
.await
}
@@ -166,6 +180,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -173,7 +188,8 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
),
GatewayError,
> {
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let ordering_config =
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
let priority_affinity_key = selection::scheduling_priority_affinity_key(
auth_snapshot,
client_session_affinity,
@@ -194,6 +210,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -203,6 +220,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
@@ -214,12 +232,14 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
)
.await?
.0,
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -229,6 +249,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
if normalized_api_format.is_empty() {
@@ -262,20 +283,22 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
for global_model_name in model_names {
let (candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
client_session_affinity,
now_unix_secs,
false,
None,
)
.await?;
let (candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
client_session_affinity,
now_unix_secs,
false,
None,
ordering_config,
)
.await?;
all_attempts_blocked_by_auth_limit &=
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
match capability_mode {
apps/aether-gateway/src/scheduler/candidate/selection.rs
+78 -17
View File
@@ -1,7 +1,7 @@
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use crate::GatewayError;
use aether_scheduler_core::ClientSessionAffinity;
@@ -115,6 +115,7 @@ pub(super) async fn select_minimal_candidate(
Ok(selected)
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -126,24 +127,32 @@ pub(super) async fn collect_selectable_candidates(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(collect_selectable_candidates_with_skip_reasons(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
None,
Ok(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
None,
ordering_config,
)
.await?
.0,
)
.await?
.0)
}
/// Legacy-shaped entrypoint that resolves the ordering config from the
/// runtime state. Prefer `collect_selectable_candidates_with_skip_reasons_and_ordering`
/// and pass the request's routing-policy config explicitly.
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates_with_skip_reasons(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -163,7 +172,59 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
),
GatewayError,
> {
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
request_operation,
None,
)
.await
}
/// Resolve the ordering config for a preselection pass: the routing-policy
/// derived config wins when the caller has one; otherwise fall back to the
/// runtime default (system-default routing group, then legacy keys).
pub(super) async fn resolve_preselection_ordering_config(
runtime_state: &impl SchedulerRuntimeState,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<SchedulerOrderingConfig, GatewayError> {
match ordering_config {
Some(config) => Ok(config),
None => runtime_state.read_scheduler_ordering_config().await,
}
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
),
GatewayError,
> {
let ordering_config =
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
let priority_affinity_key = scheduling_priority_affinity_key(
auth_snapshot,
client_session_affinity,
@@ -205,7 +266,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
ordering_config: SchedulerOrderingConfig,
priority_affinity_key: Option<&str>,
) -> Result<
(
apps/aether-gateway/src/scheduler/candidate/tests/required_capability.rs
+4
View File
@@ -65,6 +65,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
None,
None,
100,
None,
)
.await
.expect("selection should succeed");
@@ -120,6 +121,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
None,
None,
100,
None,
)
.await
.expect("selection should succeed");
@@ -196,6 +198,7 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
Some(&auth_snapshot),
Some(&client_session_affinity),
100,
None,
)
.await
.expect("selection should succeed");
@@ -273,6 +276,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
Some(&auth_snapshot),
None,
100,
None,
)
.await
.expect("selection should succeed");
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
+1
View File
@@ -75,6 +75,7 @@ async fn collect_selectable_candidates(
None,
now_unix_secs,
false,
None,
)
.await
}
apps/aether-gateway/src/scheduler/config.rs
+303
View File
@@ -1,4 +1,10 @@
use aether_data_contracts::repository::routing_profiles::RoutingGroupLookupKey;
use aether_routing_core::{
ResolvedRoutingPolicy, RoutingDefaultPolicy, RoutingSchedulingMode, RoutingSetPriorityMode,
DEFAULT_STICKY_KEY_ATTEMPTS,
};
use aether_scheduler_core::SchedulerPriorityMode;
use tracing::warn;
use crate::{AppState, GatewayError};
@@ -10,11 +16,30 @@ pub(crate) enum SchedulerSchedulingMode {
LoadBalance,
}
impl SchedulerSchedulingMode {
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::FixedOrder => "fixed_order",
Self::CacheAffinity => "cache_affinity",
Self::LoadBalance => "load_balance",
}
}
}
pub(crate) fn scheduler_priority_mode_as_str(mode: SchedulerPriorityMode) -> &'static str {
match mode {
SchedulerPriorityMode::Provider => "provider",
SchedulerPriorityMode::GlobalKey => "global_key",
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct SchedulerOrderingConfig {
pub(crate) priority_mode: SchedulerPriorityMode,
pub(crate) scheduling_mode: SchedulerSchedulingMode,
pub(crate) keep_priority_on_conversion: bool,
/// Total attempts on the first-ranked (sticky) candidate before failover.
pub(crate) sticky_key_attempts: u32,
}
impl Default for SchedulerOrderingConfig {
@@ -23,10 +48,72 @@ impl Default for SchedulerOrderingConfig {
priority_mode: SchedulerPriorityMode::Provider,
scheduling_mode: SchedulerSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
}
}
}
impl SchedulerOrderingConfig {
/// Ordering config derived from a resolved routing policy. The policy is
/// the single source of truth: no legacy system-config value is merged in.
pub(crate) fn from_routing_policy(policy: &ResolvedRoutingPolicy) -> Self {
Self {
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
keep_priority_on_conversion: policy.keep_priority_on_conversion,
sticky_key_attempts: policy.sticky_key_attempts,
}
}
pub(crate) fn from_routing_default_policy(policy: &RoutingDefaultPolicy) -> Self {
Self {
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
keep_priority_on_conversion: policy.keep_priority_on_conversion,
sticky_key_attempts: policy.sticky_key_attempts,
}
}
pub(crate) fn to_routing_default_policy(self) -> RoutingDefaultPolicy {
RoutingDefaultPolicy {
priority_mode: match self.priority_mode {
SchedulerPriorityMode::Provider => RoutingSetPriorityMode::Provider,
SchedulerPriorityMode::GlobalKey => RoutingSetPriorityMode::GlobalKey,
},
scheduling_mode: match self.scheduling_mode {
SchedulerSchedulingMode::FixedOrder => RoutingSchedulingMode::FixedOrder,
SchedulerSchedulingMode::CacheAffinity => RoutingSchedulingMode::CacheAffinity,
SchedulerSchedulingMode::LoadBalance => RoutingSchedulingMode::LoadBalance,
},
keep_priority_on_conversion: self.keep_priority_on_conversion,
sticky_key_attempts: self.sticky_key_attempts,
}
}
pub(crate) fn priority_mode_str(self) -> &'static str {
scheduler_priority_mode_as_str(self.priority_mode)
}
pub(crate) fn scheduling_mode_str(self) -> &'static str {
self.scheduling_mode.as_str()
}
}
fn scheduler_priority_mode_from_routing(mode: RoutingSetPriorityMode) -> SchedulerPriorityMode {
match mode {
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
}
}
fn scheduler_scheduling_mode_from_routing(mode: RoutingSchedulingMode) -> SchedulerSchedulingMode {
match mode {
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
}
}
pub(crate) fn parse_scheduler_priority_mode(
value: Option<&serde_json::Value>,
) -> SchedulerPriorityMode {
@@ -62,8 +149,62 @@ pub(crate) fn parse_scheduler_scheduling_mode(
}
}
/// Effective scheduler ordering config for requests that carry no resolved
/// routing policy.
///
/// Resolution order:
/// 1. the enabled system-default routing group's `default_policy`;
/// 2. the legacy system-config keys (`provider_priority_mode`,
/// `scheduling_mode`, `keep_priority_on_conversion`).
///
/// Step 2 only exists so deployments that never created a routing group keep
/// their behaviour; once the legacy keys are removed this function collapses
/// to step 1 plus `SchedulerOrderingConfig::default()`.
pub(crate) async fn read_scheduler_ordering_config(
state: &AppState,
) -> Result<SchedulerOrderingConfig, GatewayError> {
if let Some(config) = read_system_default_routing_ordering_config(state).await? {
return Ok(config);
}
read_legacy_scheduler_ordering_config(state).await
}
/// Ordering config from the enabled system-default routing group, if any.
pub(crate) async fn read_system_default_routing_ordering_config(
state: &AppState,
) -> Result<Option<SchedulerOrderingConfig>, GatewayError> {
let Some(group) = state
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await?
.filter(|group| group.enabled)
else {
return Ok(None);
};
let default_policy = match group.config_json.get("default_policy") {
None | Some(serde_json::Value::Null) => RoutingDefaultPolicy::default(),
Some(value) => match serde_json::from_value::<RoutingDefaultPolicy>(value.clone()) {
Ok(policy) => policy,
Err(error) => {
warn!(
event_name = "scheduler_system_default_routing_policy_invalid",
log_type = "event",
group_id = %group.id,
error = %error,
"system default routing group has an invalid default_policy; ignoring it"
);
return Ok(None);
}
},
};
Ok(Some(SchedulerOrderingConfig::from_routing_default_policy(
&default_policy,
)))
}
/// Legacy system-config based ordering config. Kept only as a migration
/// fallback; see `read_scheduler_ordering_config`.
pub(crate) async fn read_legacy_scheduler_ordering_config(
state: &AppState,
) -> Result<SchedulerOrderingConfig, GatewayError> {
let priority_mode = parse_scheduler_priority_mode(
state
@@ -87,5 +228,167 @@ pub(crate) async fn read_scheduler_ordering_config(
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
// Legacy config never carried a sticky-key setting; use the routing default.
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
})
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
use aether_data_contracts::repository::routing_profiles::{
CreateRoutingGroupRecord, RoutingGroupLookupKey, RoutingGroupReadRepository,
RoutingGroupWriteRepository,
};
use serde_json::json;
use super::*;
use crate::data::GatewayDataState;
fn legacy_values() -> [(String, serde_json::Value); 3] {
[
("provider_priority_mode".to_string(), json!("global_key")),
("scheduling_mode".to_string(), json!("load_balance")),
("keep_priority_on_conversion".to_string(), json!(true)),
]
}
async fn create_system_default(
repository: &InMemoryRoutingGroupRepository,
enabled: bool,
config_json: serde_json::Value,
) {
repository
.create_routing_group(CreateRoutingGroupRecord {
id: "system-default".to_string(),
name: "system-default".to_string(),
description: None,
enabled,
is_system_default: true,
config_json,
version: 1,
created_at: 1,
updated_at: 1,
published_at: None,
})
.await
.unwrap();
}
#[tokio::test]
async fn system_default_routing_group_overrides_legacy_keys() {
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
create_system_default(
&repository,
true,
json!({
"default_policy": {
"priority_mode": "provider",
"scheduling_mode": "fixed_order",
"keep_priority_on_conversion": false
}
}),
)
.await;
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::disabled()
.with_system_config_values_for_tests(legacy_values())
.with_routing_group_repository_for_tests(repository),
);
let config = read_scheduler_ordering_config(&state).await.unwrap();
assert_eq!(config.priority_mode, SchedulerPriorityMode::Provider);
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::FixedOrder);
assert!(!config.keep_priority_on_conversion);
}
#[tokio::test]
async fn missing_default_policy_in_system_default_group_uses_routing_defaults() {
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
create_system_default(&repository, true, json!({})).await;
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::disabled()
.with_system_config_values_for_tests(legacy_values())
.with_routing_group_repository_for_tests(repository),
);
let config = read_scheduler_ordering_config(&state).await.unwrap();
assert_eq!(config, SchedulerOrderingConfig::default());
}
#[tokio::test]
async fn disabled_or_missing_system_default_group_falls_back_to_legacy_keys() {
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
create_system_default(
&repository,
false,
json!({"default_policy": {"scheduling_mode": "fixed_order"}}),
)
.await;
let with_disabled_group = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::disabled()
.with_system_config_values_for_tests(legacy_values())
.with_routing_group_repository_for_tests(repository),
);
let without_repository = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::disabled().with_system_config_values_for_tests(legacy_values()),
);
for state in [with_disabled_group, without_repository] {
let config = read_scheduler_ordering_config(&state).await.unwrap();
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
assert!(config.keep_priority_on_conversion);
}
}
#[tokio::test]
async fn bootstrap_creates_system_default_group_from_legacy_keys_once() {
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
let state = AppState::new().unwrap().with_data_state_for_tests(
GatewayDataState::disabled()
.with_system_config_values_for_tests(legacy_values())
.with_routing_group_repository_for_tests(repository.clone()),
);
let created = state
.ensure_system_default_routing_group_inner()
.await
.unwrap()
.expect("first bootstrap should create the system default group");
assert!(created.enabled);
assert!(created.is_system_default);
assert_eq!(
created.config_json["default_policy"],
json!({
"priority_mode": "global_key",
"scheduling_mode": "load_balance",
"keep_priority_on_conversion": true,
"sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS
})
);
let second = state
.ensure_system_default_routing_group_inner()
.await
.unwrap();
assert!(second.is_none(), "bootstrap must be idempotent");
assert_eq!(
repository
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await
.unwrap()
.map(|group| group.id),
Some(created.id)
);
let config = read_scheduler_ordering_config(&state).await.unwrap();
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
assert!(config.keep_priority_on_conversion);
}
}
apps/aether-gateway/src/state/routing_profiles.rs
+78
View File
@@ -4,11 +4,89 @@ use aether_data_contracts::repository::routing_profiles::{
StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion,
UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
};
use aether_routing_core::RoutingGroupConfig;
use std::sync::Arc;
use tracing::warn;
use super::{AppState, GatewayError};
const BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME: &str = "system-default";
impl AppState {
/// Make sure an enabled system-default routing group exists.
///
/// When none exists, one is created from the legacy scheduler system-config
/// keys so that removing those keys later does not change behaviour. Returns
/// the created group, or `None` when nothing had to be created (no routing
/// storage, no writer, or a system default already exists).
pub async fn ensure_system_default_routing_group(
&self,
) -> Result<Option<StoredRoutingGroup>, std::io::Error> {
self.ensure_system_default_routing_group_inner()
.await
.map_err(|err| std::io::Error::other(format!("{err:?}")))
}
pub(crate) async fn ensure_system_default_routing_group_inner(
&self,
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
if !self.has_routing_group_data_reader() {
return Ok(None);
}
if self
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
.await?
.is_some()
{
return Ok(None);
}
if !self.has_routing_group_data_writer() {
warn!(
event_name = "routing_system_default_bootstrap_skipped",
log_type = "event",
"no system default routing group exists and routing storage is read-only; scheduler falls back to legacy system config"
);
return Ok(None);
}
let legacy = crate::scheduler::config::read_legacy_scheduler_ordering_config(self).await?;
let config = RoutingGroupConfig {
default_policy: legacy.to_routing_default_policy(),
..RoutingGroupConfig::default()
};
let config_json = serde_json::to_value(config)
.map_err(|err| GatewayError::Internal(format!("serialize routing config: {err}")))?;
let name = if self
.find_routing_group(RoutingGroupLookupKey::Name(
BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME,
))
.await?
.is_some()
{
format!(
"{BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME}-{}",
&uuid::Uuid::new_v4().simple().to_string()[..8]
)
} else {
BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME.to_string()
};
let now = crate::clock::current_unix_secs() as i64;
self.create_routing_group(CreateRoutingGroupRecord {
id: uuid::Uuid::new_v4().to_string(),
name,
description: Some("自动从旧版调度配置迁移生成的系统默认策略".to_string()),
enabled: true,
is_system_default: true,
config_json,
version: 1,
created_at: now,
updated_at: now,
published_at: Some(now),
})
.await
}
pub(crate) fn has_routing_group_data_reader(&self) -> bool {
self.data.has_routing_group_reader()
}
apps/aether-gateway/src/tests/ai_execute/stream/decision.rs
+26 -5
View File
@@ -2220,7 +2220,9 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
.to_string(),
});
let frames = if attempt == 1 {
// The primary key gets two attempts under the default
// sticky_key_attempts; both must fail to reach the backup.
let frames = if attempt <= 2 {
concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":429,\"headers\":{\"content-type\":\"application/json\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"error\\\":{\\\"message\\\":\\\"rate limited\\\",\\\"type\\\":\\\"rate_limit_error\\\"}}\"}}\n",
@@ -2362,14 +2364,25 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
.lock()
.expect("mutex should lock")
.len()
>= 2
>= 3
})
.await;
let seen_execution_runtime_requests = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(seen_execution_runtime_requests.len(), 2);
// Default sticky_key_attempts is 2: the primary key is retried once on
// the same key, then failover moves to the backup with a single attempt.
assert_eq!(seen_execution_runtime_requests.len(), 3);
assert_eq!(
seen_execution_runtime_requests
.iter()
.filter(|request| {
request.url == "https://api.openai.primary.example/chat/completions"
})
.count(),
2
);
let primary_request = seen_execution_runtime_requests
.iter()
.find(|request| request.url == "https://api.openai.primary.example/chat/completions")
@@ -2404,7 +2417,15 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
.list_by_request_id("trace-openai-chat-local-stream-failover-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 2);
assert_eq!(stored_candidates.len(), 3);
assert_eq!(
stored_candidates
.iter()
.filter(|candidate| candidate.status == RequestCandidateStatus::Failed)
.count(),
2,
"both sticky-key attempts on the primary should be recorded as failed"
);
let failed_candidate = stored_candidates
.iter()
.find(|candidate| {
@@ -2465,7 +2486,7 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
assert_eq!(
execution_runtime_hits.load(std::sync::atomic::Ordering::SeqCst),
2
3
);
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs
+54 -45
View File
@@ -112,7 +112,7 @@ async fn gateway_skips_unsupported_local_openai_chat_sync_candidate_before_tryin
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -134,7 +134,7 @@ async fn gateway_skips_unsupported_local_openai_chat_sync_candidate_before_tryin
"https://api.openai.skip.example".to_string(),
None,
None,
Some(2),
Some(1),
None,
None,
None,
@@ -520,7 +520,7 @@ async fn gateway_surfaces_local_execution_runtime_miss_reason_when_all_openai_ch
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -542,7 +542,7 @@ async fn gateway_surfaces_local_execution_runtime_miss_reason_when_all_openai_ch
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(2),
Some(1),
None,
None,
None,
@@ -802,7 +802,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -828,7 +828,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
base_url.to_string(),
None,
None,
Some(2),
Some(1),
None,
None,
None,
@@ -970,7 +970,9 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
.to_string(),
});
if attempt == 1 {
// The primary key gets two attempts under the default
// sticky_key_attempts; both must fail to reach the backup.
if attempt <= 2 {
return Json(json!({
"request_id": "trace-openai-chat-local-failover-123",
"status_code": 401,
@@ -1125,56 +1127,63 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
.lock()
.expect("mutex should lock")
.clone();
assert_eq!(seen_execution_runtime_requests.len(), 2);
// Default sticky_key_attempts is 2: the primary key is retried once on
// the same key, then failover moves to the backup with a single attempt.
assert_eq!(seen_execution_runtime_requests.len(), 3);
for primary_request in &seen_execution_runtime_requests[..2] {
assert_eq!(
primary_request.trace_id,
"trace-openai-chat-local-failover-123"
);
assert_eq!(
primary_request.url,
"https://api.openai.primary.example/chat/completions"
);
assert_eq!(
primary_request.authorization,
"Bearer sk-upstream-openai-primary"
);
}
assert_eq!(
seen_execution_runtime_requests[0].trace_id,
"trace-openai-chat-local-failover-123"
);
assert_eq!(
seen_execution_runtime_requests[0].url,
"https://api.openai.primary.example/chat/completions"
);
assert_eq!(
seen_execution_runtime_requests[0].authorization,
"Bearer sk-upstream-openai-primary"
);
assert_eq!(
seen_execution_runtime_requests[1].url,
seen_execution_runtime_requests[2].url,
"https://api.openai.backup.example/chat/completions"
);
assert_eq!(
seen_execution_runtime_requests[1].model,
seen_execution_runtime_requests[2].model,
"gpt-5-upstream-backup"
);
assert_eq!(
seen_execution_runtime_requests[1].authorization,
seen_execution_runtime_requests[2].authorization,
"Bearer sk-upstream-openai-backup"
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-failover-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 2);
assert_eq!(stored_candidates[0].candidate_index, 0);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
assert_eq!(stored_candidates[0].status_code, Some(401));
assert_eq!(
stored_candidates[0].error_message.as_deref(),
Some("invalid auth token")
);
let failed_upstream_response = stored_candidates[0]
.extra_data
.as_ref()
.and_then(|value| value.get("upstream_response"))
.expect("failed candidate should keep its upstream response");
assert_eq!(failed_upstream_response["status_code"], json!(401));
assert_eq!(
failed_upstream_response["body"]["error"]["message"],
json!("invalid auth token")
);
assert_eq!(stored_candidates[1].candidate_index, 1);
assert_eq!(stored_candidates[1].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates[1].status_code, Some(200));
assert_eq!(stored_candidates.len(), 3);
for (retry_index, failed_candidate) in stored_candidates[..2].iter().enumerate() {
assert_eq!(failed_candidate.candidate_index, 0);
assert_eq!(failed_candidate.retry_index, retry_index as u32);
assert_eq!(failed_candidate.status, RequestCandidateStatus::Failed);
assert_eq!(failed_candidate.status_code, Some(401));
assert_eq!(
failed_candidate.error_message.as_deref(),
Some("invalid auth token")
);
let failed_upstream_response = failed_candidate
.extra_data
.as_ref()
.and_then(|value| value.get("upstream_response"))
.expect("failed candidate should keep its upstream response");
assert_eq!(failed_upstream_response["status_code"], json!(401));
assert_eq!(
failed_upstream_response["body"]["error"]["message"],
json!("invalid auth token")
);
}
assert_eq!(stored_candidates[2].candidate_index, 1);
assert_eq!(stored_candidates[2].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates[2].status_code, Some(200));
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
@@ -1184,7 +1193,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
assert_eq!(
*execution_runtime_hits.lock().expect("mutex should lock"),
2
3
);
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
apps/aether-gateway/src/tests/ai_execute/sync/search.rs
+12 -10
View File
@@ -140,7 +140,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
false,
false,
None,
Some(2),
Some(1),
None,
Some(900.0),
None,
@@ -162,7 +162,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
"https://chatgpt.com/backend-api/codex".to_string(),
None,
None,
Some(2),
Some(1),
None,
None,
None,
@@ -570,9 +570,12 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
.filter(|plan| plan["request_id"] == "trace-search-failover-1")
.map(|plan| plan["provider_id"].clone())
.collect::<Vec<_>>();
// Default sticky_key_attempts is 2: the first provider is retried once on
// the same key before failover advances to the second provider.
assert_eq!(
failover_plans,
vec![
json!("provider-codex-search-1"),
json!("provider-codex-search-1"),
json!("provider-codex-search-2")
]
@@ -581,17 +584,16 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
.list_by_request_id("trace-search-failover-1")
.await
.expect("failover request candidates should read");
assert_eq!(failover_candidates.len(), 2);
assert_eq!(failover_candidates.len(), 3);
for failed_candidate in &failover_candidates[..2] {
assert_eq!(failed_candidate.status, RequestCandidateStatus::Failed);
assert_eq!(failed_candidate.status_code, Some(500));
}
assert_eq!(
failover_candidates[0].status,
RequestCandidateStatus::Failed
);
assert_eq!(failover_candidates[0].status_code, Some(500));
assert_eq!(
failover_candidates[1].status,
failover_candidates[2].status,
RequestCandidateStatus::Success
);
assert_eq!(failover_candidates[1].status_code, Some(200));
assert_eq!(failover_candidates[2].status_code, Some(200));
gateway_handle.abort();
execution_runtime_handle.abort();
apps/aether-gateway/src/tests/architecture/ai_serving.rs
-1
View File
@@ -1714,7 +1714,6 @@ fn ai_serving_candidate_materialization_owns_affinity_and_candidate_runtime_pers
"pub fn ai_should_persist_available_candidate_for_pool_key",
"pub fn ai_should_persist_skipped_candidate_for_pool_membership",
"pub fn ai_candidate_extra_data_with_ranking",
"attempt_slot_count",
"should_persist_available_candidate",
"persist_available_candidate",
"build_attempt",
apps/aether-gateway/src/tests/usage.rs
+2 -2
View File
@@ -122,7 +122,7 @@ pub(super) fn sample_local_openai_provider() -> StoredProviderCatalogProvider {
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -144,7 +144,7 @@ pub(super) fn sample_local_openai_endpoint() -> StoredProviderCatalogEndpoint {
"https://api.openai.example/v1".to_string(),
None,
None,
Some(2),
Some(1),
None,
None,
None,
apps/aether-gateway/src/tests/usage/local.rs
+22 -17
View File
@@ -879,9 +879,13 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
.list_by_request_id("trace-openai-chat-local-report-sync-failure-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
assert_eq!(stored_candidates[0].status_code, Some(503));
// The only candidate is the sticky first key: the default policy retries
// it once on the same key before the request is exhausted.
assert_eq!(stored_candidates.len(), 2);
for candidate in &stored_candidates {
assert_eq!(candidate.status, RequestCandidateStatus::Failed);
assert_eq!(candidate.status_code, Some(503));
}
}
#[test]
@@ -957,7 +961,8 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
let response = send_request(gateway, request).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 1);
// The sticky first key is retried once on the same key before exhaustion.
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 2);
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
@@ -982,15 +987,15 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
.list_by_request_id("trace-openai-chat-local-transport-unavailable-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
assert!(stored_candidates[0]
.latency_ms
.is_some_and(|value| value >= 5));
assert_eq!(
stored_candidates[0].error_type.as_deref(),
Some("execution_runtime_unavailable")
);
assert_eq!(stored_candidates.len(), 2);
for candidate in &stored_candidates {
assert_eq!(candidate.status, RequestCandidateStatus::Failed);
assert!(candidate.latency_ms.is_some_and(|value| value >= 5));
assert_eq!(
candidate.error_type.as_deref(),
Some("execution_runtime_unavailable")
);
}
}
#[test]
@@ -1786,7 +1791,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -1808,7 +1813,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
"https://right.codes/codex".to_string(),
None,
None,
Some(2),
Some(1),
Some("/v1/messages".to_string()),
None,
None,
@@ -2111,7 +2116,7 @@ fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude
false,
false,
None,
Some(2),
Some(1),
None,
Some(20.0),
None,
@@ -2133,7 +2138,7 @@ fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude
"https://right.codes/codex".to_string(),
None,
None,
Some(2),
Some(1),
Some("/v1/messages".to_string()),
None,
None,
crates/aether-ai/serving/src/attempt_loop.rs
+88 -2
View File
@@ -13,8 +13,23 @@ pub trait AiExecutionAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
None
}
/// Re-issue this attempt against the same key as a fresh attempt with the
/// given retry index and candidate id. Attempt types that cannot be
/// re-issued return `None`, which disables same-key retries for them.
fn with_same_key_retry(&self, _retry_index: u32, _candidate_id: String) -> Option<Self>
where
Self: Sized,
{
None
}
}
/// Report-context field carrying the routing policy's sticky-key attempt
/// budget for the request, so the attempt loop can derive same-key retries
/// lazily instead of pre-materializing them.
pub const STICKY_KEY_ATTEMPTS_REPORT_FIELD: &str = "sticky_key_attempts";
#[derive(Debug)]
pub enum AiAttemptLoopOutcome<Response, Exhaustion> {
Responded(Response),
@@ -83,6 +98,17 @@ where
Ok(())
}
/// After `attempt` failed with candidate scope, return the next attempt on
/// the same key, or `None` once the sticky-key budget is used up. Retries
/// are derived here on demand so no attempt is materialized before it is
/// actually needed.
async fn next_same_key_retry(
&self,
_attempt: &Attempt,
) -> Result<Option<Attempt>, Self::Error> {
Ok(None)
}
async fn mark_unused_attempts(&self, attempts: Vec<Attempt>) -> Result<(), Self::Error>;
async fn build_exhaustion(
@@ -101,11 +127,15 @@ where
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
let mut remaining = attempts.into_iter();
let mut pending_same_key_retry: Option<Attempt> = None;
let mut last_attempted = None;
let mut retry_filters: Vec<AiAttemptRetryFilter> = Vec::new();
let mut fallback_response = None;
while let Some(attempt) = remaining.next() {
loop {
let Some(attempt) = pending_same_key_retry.take().or_else(|| remaining.next()) else {
break;
};
if retry_filters.iter().any(|filter| filter.matches(&attempt))
|| port.should_skip_attempt(&attempt).await?
{
@@ -133,7 +163,9 @@ where
if attempt_fallback_response.is_some() {
fallback_response = attempt_fallback_response;
}
if scope != AiAttemptRetryScope::Candidate {
if scope == AiAttemptRetryScope::Candidate {
pending_same_key_retry = port.next_same_key_retry(&attempt).await?;
} else {
retry_filters.push(AiAttemptRetryFilter::new(&attempt, scope));
}
}
@@ -188,6 +220,32 @@ impl AiAttemptRetryFilter {
}
}
/// Clone `plan`/`report_context` for a same-key retry: only the candidate id
/// and retry index change, everything else (url, headers, body) is reused.
fn same_key_retry_parts(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
retry_index: u32,
candidate_id: String,
) -> (aether_contracts::ExecutionPlan, Option<serde_json::Value>) {
let mut plan = plan.clone();
plan.candidate_id = Some(candidate_id.clone());
let report_context = report_context.cloned().map(|mut value| {
if let Some(object) = value.as_object_mut() {
object.insert(
"candidate_id".to_string(),
serde_json::Value::String(candidate_id),
);
object.insert(
"retry_index".to_string(),
serde_json::Value::Number(retry_index.into()),
);
}
value
});
(plan, report_context)
}
impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
@@ -204,6 +262,20 @@ impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
self.report_context.as_ref()
}
fn with_same_key_retry(&self, retry_index: u32, candidate_id: String) -> Option<Self> {
let (plan, report_context) = same_key_retry_parts(
&self.plan,
self.report_context.as_ref(),
retry_index,
candidate_id,
);
Some(Self {
plan,
report_kind: self.report_kind.clone(),
report_context,
})
}
}
impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
@@ -222,6 +294,20 @@ impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
self.report_context.as_ref()
}
fn with_same_key_retry(&self, retry_index: u32, candidate_id: String) -> Option<Self> {
let (plan, report_context) = same_key_retry_parts(
&self.plan,
self.report_context.as_ref(),
retry_index,
candidate_id,
);
Some(Self {
plan,
report_kind: self.report_kind.clone(),
report_context,
})
}
}
#[cfg(test)]
crates/aether-ai/serving/src/candidate_persistence.rs
+22 -60
View File
@@ -9,8 +9,6 @@ pub trait AiAvailableCandidatePersistencePort: Send + Sync {
type ExtraData: Clone + Send + Sync;
type Error: Send;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32;
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData>;
fn generate_candidate_id(&self) -> String;
@@ -42,50 +40,28 @@ pub async fn run_ai_available_candidate_persistence<Port>(
where
Port: AiAvailableCandidatePersistencePort,
{
let total_attempts = candidates
.iter()
.map(|candidate| port.attempt_slot_count(candidate) as usize)
.sum();
let mut materialized = Vec::with_capacity(total_attempts);
// One attempt per candidate. Same-key retries are derived lazily by the
// attempt loop (`AiAttemptLoopPort::next_same_key_retry`) after a failure,
// so the sticky-key budget never inflates up-front materialization.
let mut materialized = Vec::with_capacity(candidates.len());
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_index = candidate_index as u32;
let attempt_slots = port.attempt_slot_count(&candidate).max(1);
let extra_data = port.build_extra_data(&candidate);
let mut owned_candidate = Some(candidate);
for retry_index in 0..attempt_slots {
let candidate = owned_candidate
.as_ref()
.expect("candidate should remain available until final retry");
let generated_candidate_id = port.generate_candidate_id();
let candidate_id = if port.should_persist_available_candidate(candidate) {
port.persist_available_candidate(
candidate,
candidate_index,
retry_index,
generated_candidate_id.as_str(),
extra_data.clone(),
)
.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.clone()
};
materialized.push(port.build_attempt(
candidate,
let generated_candidate_id = port.generate_candidate_id();
let candidate_id = if port.should_persist_available_candidate(&candidate) {
port.persist_available_candidate(
&candidate,
candidate_index,
retry_index,
candidate_id,
));
}
0,
generated_candidate_id.as_str(),
extra_data,
)
.await?
} else {
generated_candidate_id
};
materialized.push(port.build_attempt(candidate, candidate_index, 0, candidate_id));
}
Ok(materialized)
@@ -177,7 +153,6 @@ mod tests {
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestCandidate {
id: &'static str,
attempt_slots: u32,
persist: bool,
}
@@ -216,10 +191,6 @@ mod tests {
type ExtraData = String;
type Error = std::convert::Infallible;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
candidate.attempt_slots
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
Some(format!("extra:{}", candidate.id))
}
@@ -299,7 +270,7 @@ mod tests {
}
#[tokio::test]
async fn available_persistence_expands_candidates_into_retry_attempts() {
async fn available_persistence_materializes_one_attempt_per_candidate() {
let port = TestPort::default();
let attempts = run_ai_available_candidate_persistence(
@@ -307,12 +278,10 @@ mod tests {
vec![
TestCandidate {
id: "a",
attempt_slots: 2,
persist: true,
},
TestCandidate {
id: "b",
attempt_slots: 1,
persist: false,
},
],
@@ -320,6 +289,8 @@ mod tests {
.await
.unwrap();
// Same-key retries are never pre-materialized; the attempt loop
// derives them on demand after a failure.
assert_eq!(
attempts,
[
@@ -329,26 +300,17 @@ mod tests {
retry_index: 0,
candidate_id: "stored-candidate-1".to_string(),
},
TestAttempt {
id: "a",
candidate_index: 0,
retry_index: 1,
candidate_id: "stored-candidate-2".to_string(),
},
TestAttempt {
id: "b",
candidate_index: 1,
retry_index: 0,
candidate_id: "candidate-3".to_string(),
candidate_id: "candidate-2".to_string(),
},
]
);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"available:a:0:0:candidate-1:extra:a",
"available:a:0:1:candidate-2:extra:a",
]
["available:a:0:0:candidate-1:extra:a"]
);
}
crates/aether-ai/serving/src/lib.rs
+1 -1
View File
@@ -54,7 +54,7 @@ pub use aether_pool_core::{
};
pub use attempt_loop::{
run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopOutcome, AiAttemptLoopPort,
AiAttemptRetryScope, AiExecutionAttempt,
AiAttemptRetryScope, AiExecutionAttempt, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
};
pub use attempt_plan::{
build_ai_execution_decision_from_plan, build_ai_execution_plan_from_decision,
crates/aether-routing-core/src/actions.rs
+6
View File
@@ -73,6 +73,8 @@ pub enum RoutingAction {
priority_mode: Option<RoutingSetPriorityMode>,
scheduling_mode: Option<RoutingSchedulingMode>,
keep_priority_on_conversion: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
sticky_key_attempts: Option<u32>,
},
SetProviderPriority {
provider_id: String,
@@ -81,6 +83,10 @@ pub enum RoutingAction {
SetKeyPriority {
key_id: String,
priority: i32,
/// When set, the override only applies to candidates served through
/// this API format; otherwise it applies to the key on every format.
#[serde(default, skip_serializing_if = "Option::is_none")]
api_format: Option<String>,
},
JsonPatchBody {
patch: Vec<RoutingJsonPatchOperation>,
crates/aether-routing-core/src/lib.rs
+3 -3
View File
@@ -13,9 +13,9 @@ pub use actions::{
};
pub use conditions::{RoutingCondition, RoutingConditionContext, RoutingConditionOp};
pub use model::{
RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord,
RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride, RoutingRule,
RoutingSchedulingPreset,
RoutingDefaultPolicy, RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig,
RoutingGroupRecord, RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride,
RoutingRule, RoutingSchedulingPreset, DEFAULT_STICKY_KEY_ATTEMPTS,
};
pub use mutations::{
apply_json_patch_operations, validate_header_patch, validate_json_patch_operations,
crates/aether-routing-core/src/model.rs
+32 -1
View File
@@ -23,7 +23,11 @@ pub struct RoutingPoolPolicyOverride {
pub scheduling_presets: Vec<RoutingSchedulingPreset>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
/// Default number of attempts on the first-ranked (sticky) candidate before
/// failing over: one retry on the same key.
pub const DEFAULT_STICKY_KEY_ATTEMPTS: u32 = 2;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingDefaultPolicy {
#[serde(default)]
pub priority_mode: RoutingSetPriorityMode,
@@ -31,6 +35,26 @@ pub struct RoutingDefaultPolicy {
pub scheduling_mode: RoutingSchedulingMode,
#[serde(default)]
pub keep_priority_on_conversion: bool,
/// Total attempts on the first-ranked candidate before moving on. Later
/// candidates always get a single attempt so failover keeps advancing.
/// `0` and `1` both mean no same-key retry.
#[serde(default = "default_sticky_key_attempts")]
pub sticky_key_attempts: u32,
}
impl Default for RoutingDefaultPolicy {
fn default() -> Self {
Self {
priority_mode: RoutingSetPriorityMode::default(),
scheduling_mode: RoutingSchedulingMode::default(),
keep_priority_on_conversion: false,
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
}
}
}
fn default_sticky_key_attempts() -> u32 {
DEFAULT_STICKY_KEY_ATTEMPTS
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
@@ -44,6 +68,13 @@ pub struct RoutingModelPolicy {
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
/// Key priority overrides scoped to one API format: `api_format -> key_id -> priority`.
///
/// A key can serve several API formats and legacy `global_priority_by_format`
/// ranks it independently per format. Entries here take precedence over
/// `key_priority_overrides` when the candidate format matches.
#[serde(default)]
pub key_priority_overrides_by_format: BTreeMap<String, BTreeMap<String, i32>>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
crates/aether-routing-core/src/policy.rs
+113 -6
View File
@@ -57,6 +57,9 @@ pub struct ResolvedRoutingPolicy {
pub priority_mode: RoutingSetPriorityMode,
pub scheduling_mode: RoutingSchedulingMode,
pub keep_priority_on_conversion: bool,
/// See `RoutingDefaultPolicy::sticky_key_attempts`.
#[serde(default = "default_sticky_key_attempts")]
pub sticky_key_attempts: u32,
pub ranking_overlay: RankingOverlay,
pub mutation_plan: MutationPlan,
#[serde(default)]
@@ -89,6 +92,7 @@ pub fn resolve_routing_policy(
priority_mode: config.default_policy.priority_mode,
scheduling_mode: config.default_policy.scheduling_mode,
keep_priority_on_conversion: config.default_policy.keep_priority_on_conversion,
sticky_key_attempts: config.default_policy.sticky_key_attempts,
ranking_overlay: RankingOverlay::default(),
mutation_plan: MutationPlan::default(),
pool_policy_overrides: BTreeMap::new(),
@@ -163,6 +167,13 @@ fn apply_model_policy(policy: &mut ResolvedRoutingPolicy, model_policy: &Routing
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
for (api_format, overrides) in &model_policy.key_priority_overrides_by_format {
for (key_id, priority) in overrides {
policy
.ranking_overlay
.insert_key_priority_override_for_format(api_format, key_id.clone(), *priority);
}
}
policy.ranking_overlay.pool_priority_overrides.extend(
model_policy
.pool_priority_overrides
@@ -198,6 +209,7 @@ fn apply_action(
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
sticky_key_attempts,
} => {
if let Some(priority_mode) = priority_mode {
policy.priority_mode = *priority_mode;
@@ -208,6 +220,9 @@ fn apply_action(
if let Some(keep_priority_on_conversion) = keep_priority_on_conversion {
policy.keep_priority_on_conversion = *keep_priority_on_conversion;
}
if let Some(sticky_key_attempts) = sticky_key_attempts {
policy.sticky_key_attempts = *sticky_key_attempts;
}
}
RoutingAction::SetProviderPriority {
provider_id,
@@ -218,12 +233,27 @@ fn apply_action(
.provider_priority_overrides
.insert(provider_id.clone(), *priority);
}
RoutingAction::SetKeyPriority { key_id, priority } => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
RoutingAction::SetKeyPriority {
key_id,
priority,
api_format,
} => match api_format
.as_deref()
.map(str::trim)
.filter(|f| !f.is_empty())
{
Some(api_format) => {
policy
.ranking_overlay
.insert_key_priority_override_for_format(api_format, key_id.clone(), *priority);
}
None => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
},
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
@@ -260,6 +290,10 @@ fn model_allowed(patterns: &[String], requested_model: &str) -> bool {
.any(|pattern| model_pattern_matches(pattern, requested_model))
}
fn default_sticky_key_attempts() -> u32 {
crate::model::DEFAULT_STICKY_KEY_ATTEMPTS
}
fn model_pattern_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.trim();
if pattern == "*" {
@@ -362,6 +396,7 @@ mod tests {
priority_mode: RoutingSetPriorityMode::GlobalKey,
scheduling_mode: RoutingSchedulingMode::LoadBalance,
keep_priority_on_conversion: true,
sticky_key_attempts: 3,
},
model_policies: vec![RoutingModelPolicy {
model: "special-model".to_string(),
@@ -393,6 +428,7 @@ mod tests {
assert_eq!(special.priority_mode, RoutingSetPriorityMode::GlobalKey);
assert_eq!(special.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(special.keep_priority_on_conversion);
assert_eq!(special.sticky_key_attempts, 3);
assert_eq!(
special.ranking_overlay.allowed_providers,
vec!["provider-special"]
@@ -426,6 +462,7 @@ mod tests {
assert_eq!(ordinary.priority_mode, RoutingSetPriorityMode::GlobalKey);
assert_eq!(ordinary.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(ordinary.keep_priority_on_conversion);
assert_eq!(ordinary.sticky_key_attempts, 3);
assert!(ordinary.ranking_overlay.allowed_providers.is_empty());
assert!(ordinary.ranking_overlay.allowed_keys.is_empty());
assert!(ordinary
@@ -434,6 +471,76 @@ mod tests {
.is_empty());
}
#[test]
fn sticky_key_attempts_defaults_to_two_and_can_be_overridden_by_rule() {
let default_config = RoutingGroupConfig::default();
let default_policy = resolve_routing_policy(
&default_config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("default config should resolve");
assert_eq!(
default_policy.sticky_key_attempts,
crate::DEFAULT_STICKY_KEY_ATTEMPTS
);
let parsed: RoutingGroupConfig =
serde_json::from_value(json!({ "default_policy": { "priority_mode": "provider" } }))
.expect("legacy config without sticky_key_attempts should deserialize");
assert_eq!(
parsed.default_policy.sticky_key_attempts,
crate::DEFAULT_STICKY_KEY_ATTEMPTS
);
let config = RoutingGroupConfig {
rules: vec![RoutingRule {
id: "no-sticky-retry".to_string(),
priority: 1,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::default(),
actions: vec![RoutingAction::SetScheduling {
priority_mode: None,
scheduling_mode: None,
keep_priority_on_conversion: None,
sticky_key_attempts: Some(1),
}],
stop_processing: false,
}],
..RoutingGroupConfig::default()
};
let policy = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("rule config should resolve");
assert_eq!(policy.sticky_key_attempts, 1);
}
#[test]
fn rejects_disallowed_model() {
let config = RoutingGroupConfig {
crates/aether-routing-core/src/ranking.rs
+96 -1
View File
@@ -21,6 +21,9 @@ pub struct RankingOverlay {
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
/// `api_format -> key_id -> priority`; see `RoutingModelPolicy`.
#[serde(default)]
pub key_priority_overrides_by_format: BTreeMap<String, BTreeMap<String, i32>>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
}
@@ -40,6 +43,46 @@ impl RankingOverlay {
.unwrap_or(fallback)
}
/// Format-scoped key priority: a per-format override wins, then the
/// format-agnostic key override, then `fallback`.
pub fn key_priority_for_format(&self, key_id: &str, api_format: &str, fallback: i32) -> i32 {
self.key_priority_override_for_format(key_id, api_format)
.unwrap_or_else(|| self.key_priority(key_id, fallback))
}
/// Format-scoped key override using exact (case-insensitive) format match.
pub fn key_priority_override_for_format(&self, key_id: &str, api_format: &str) -> Option<i32> {
let api_format = api_format.trim();
self.key_priority_override_matching_format(key_id, |format| {
format.trim().eq_ignore_ascii_case(api_format)
})
}
/// Format-scoped key override where the caller decides how configured
/// format names match the candidate format (for alias-aware matching).
pub fn key_priority_override_matching_format(
&self,
key_id: &str,
mut format_matches: impl FnMut(&str) -> bool,
) -> Option<i32> {
self.key_priority_overrides_by_format
.iter()
.find(|(format, _)| format_matches(format))
.and_then(|(_, overrides)| overrides.get(key_id).copied())
}
pub fn insert_key_priority_override_for_format(
&mut self,
api_format: &str,
key_id: String,
priority: i32,
) {
self.key_priority_overrides_by_format
.entry(api_format.trim().to_ascii_lowercase())
.or_default()
.insert(key_id, priority);
}
pub fn pool_priority(&self, provider_id: &str, fallback: i32) -> i32 {
self.pool_priority_overrides
.get(provider_id)
@@ -89,6 +132,9 @@ pub struct RoutingCandidateFacts {
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
/// Candidate API format used to resolve format-scoped key overrides.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_format: Option<String>,
pub provider_priority: i32,
pub key_priority: i32,
}
@@ -114,7 +160,12 @@ pub fn rank_vector_for_candidate(
CandidateKind::Provider => facts
.key_id
.as_deref()
.map(|key_id| overlay.key_priority(key_id, facts.key_priority))
.map(|key_id| match facts.api_format.as_deref() {
Some(api_format) => {
overlay.key_priority_for_format(key_id, api_format, facts.key_priority)
}
None => overlay.key_priority(key_id, facts.key_priority),
})
.unwrap_or(facts.key_priority),
CandidateKind::PoolGroup => {
overlay.pool_priority(&facts.provider_id, facts.key_priority)
@@ -142,6 +193,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: None,
provider_priority: 10,
key_priority: 20,
};
@@ -151,6 +203,47 @@ mod tests {
assert_eq!(vector.key_priority_after, 5);
}
#[test]
fn format_scoped_key_override_wins_over_key_override_for_matching_format() {
let mut overlay = RankingOverlay {
key_priority_overrides: BTreeMap::from([("key-a".to_string(), 5)]),
..RankingOverlay::default()
};
overlay.insert_key_priority_override_for_format("openai:chat", "key-a".to_string(), 1);
assert_eq!(
overlay.key_priority_for_format("key-a", "openai:chat", 20),
1
);
assert_eq!(
overlay.key_priority_for_format("key-a", "OpenAI:Chat", 20),
1
);
assert_eq!(
overlay.key_priority_for_format("key-a", "claude:messages", 20),
5
);
assert_eq!(
overlay.key_priority_for_format("key-b", "openai:chat", 20),
20
);
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: Some("openai:chat".to_string()),
provider_priority: 10,
key_priority: 20,
};
assert_eq!(
rank_vector_for_candidate(&overlay, &facts).key_priority_after,
1
);
}
#[test]
fn rank_vector_falls_back_to_existing_priorities() {
let facts = RoutingCandidateFacts {
@@ -159,6 +252,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: None,
provider_priority: 10,
key_priority: 20,
};
@@ -180,6 +274,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: None,
api_format: None,
provider_priority: 10,
key_priority: 20,
};
crates/aether-routing-core/src/validation.rs
+2
View File
@@ -271,6 +271,7 @@ mod tests {
priority_mode: None,
scheduling_mode: None,
keep_priority_on_conversion: Some(true),
sticky_key_attempts: None,
},
"set_scheduling",
),
@@ -285,6 +286,7 @@ mod tests {
RoutingAction::SetKeyPriority {
key_id: "key-1".to_string(),
priority: 1,
api_format: None,
},
"set_key_priority",
),
frontend/src/features/providers/components/ProviderFormDialog.vue
-11
View File
@@ -155,17 +155,6 @@
</SelectContent>
</Select>
</div>
<div class="space-y-1.5">
<Label>{{ legacyT('最大重试次数') }}</Label>
<Input
:model-value="form.max_retries ?? ''"
type="number"
min="0"
max="999"
:placeholder="legacyT('默认 2')"
@update:model-value="(v) => form.max_retries = parseNumberInput(v)"
/>
</div>
</div>
<!-- 超时配置 -->
frontend/src/features/routing/__tests__/routingPolicy.spec.ts
+62
View File
@@ -12,12 +12,15 @@ import {
getModelScheduling,
modelSchedulingRuleId,
normalizeRoutingGroupConfig,
normalizeStickyKeyAttempts,
parseAllowedModelsInput,
removePerModelRoutingConfig,
resolveModelKeyPriorityOverride,
routingModelScopeLabel,
savePerModelRoutingConfig,
setDefaultPoolPriorityOverrides,
setDefaultProviderPriorityOverrides,
setModelKeyPriorityOverridesForFormat,
setRoutingSortingScope,
updateAllowedModelsFromInput,
upsertModelSchedulingRule,
@@ -75,6 +78,65 @@ describe('routingPolicy', () => {
expect(policy.key_priority_overrides).toEqual({})
})
it('defaults sticky key attempts to 2 and normalizes invalid values', () => {
expect(createEmptyRoutingGroupConfig().default_policy.sticky_key_attempts).toBe(2)
expect(normalizeRoutingGroupConfig({}).default_policy.sticky_key_attempts).toBe(2)
expect(normalizeRoutingGroupConfig({
default_policy: { priority_mode: 'provider', scheduling_mode: 'cache_affinity', keep_priority_on_conversion: false, sticky_key_attempts: 3 },
}).default_policy.sticky_key_attempts).toBe(3)
expect(normalizeStickyKeyAttempts('5')).toBe(5)
expect(normalizeStickyKeyAttempts(-1)).toBe(2)
expect(normalizeStickyKeyAttempts('abc')).toBe(2)
expect(normalizeStickyKeyAttempts(500)).toBe(99)
expect(getModelScheduling(createEmptyRoutingGroupConfig(), 'gpt-5').sticky_key_attempts).toBe(2)
})
it('keeps key priority overrides independent per api format', () => {
let config = setModelKeyPriorityOverridesForFormat(
createEmptyRoutingGroupConfig(),
DEFAULT_ROUTING_POLICY_MODEL,
'OpenAI:Chat',
{ 'key-a': 0, 'key-b': 1 },
)
config = setModelKeyPriorityOverridesForFormat(
config,
DEFAULT_ROUTING_POLICY_MODEL,
'claude:messages',
{ 'key-a': 3 },
)
const policy = getDefaultModelPolicy(config)
expect(policy.key_priority_overrides).toEqual({})
expect(policy.key_priority_overrides_by_format).toEqual({
'openai:chat': { 'key-a': 0, 'key-b': 1 },
'claude:messages': { 'key-a': 3 },
})
expect(resolveModelKeyPriorityOverride(config, DEFAULT_ROUTING_POLICY_MODEL, 'openai:chat', 'key-a')).toBe(0)
expect(resolveModelKeyPriorityOverride(config, DEFAULT_ROUTING_POLICY_MODEL, 'claude:messages', 'key-a')).toBe(3)
expect(resolveModelKeyPriorityOverride(config, DEFAULT_ROUTING_POLICY_MODEL, 'claude:messages', 'key-b')).toBeUndefined()
const cleared = setModelKeyPriorityOverridesForFormat(config, DEFAULT_ROUTING_POLICY_MODEL, 'claude:messages', {})
expect(getDefaultModelPolicy(cleared).key_priority_overrides_by_format).toEqual({
'openai:chat': { 'key-a': 0, 'key-b': 1 },
})
})
it('falls back to format-agnostic key overrides and normalizes legacy configs', () => {
const config = normalizeRoutingGroupConfig({
model_policies: [{
...createEmptyModelPolicy('gpt-5'),
key_priority_overrides: { 'key-a': 7 },
key_priority_overrides_by_format: { ' OpenAI:Chat ': { 'key-a': 1 } },
}],
})
expect(config.model_policies[0].key_priority_overrides_by_format).toEqual({
'openai:chat': { 'key-a': 1 },
})
expect(resolveModelKeyPriorityOverride(config, 'gpt-5', 'openai:chat', 'key-a')).toBe(1)
expect(resolveModelKeyPriorityOverride(config, 'gpt-5', 'gemini:generate_content', 'key-a')).toBe(7)
})
it('stores per-model scheduling as generated routing rules', () => {
const next = upsertModelSchedulingRule(createEmptyRoutingGroupConfig(), 'gpt-5', {
priority_mode: 'global_key',
frontend/src/features/routing/components/RoutingPriorityPolicyEditor.vue
+27 -6
View File
@@ -315,7 +315,8 @@ import {
getDefaultModelPolicy,
getModelPolicy,
normalizeRoutingGroupConfig,
setModelKeyPriorityOverrides,
normalizeRoutingApiFormatKey,
setModelKeyPriorityOverridesForFormat,
setModelPoolPriorityOverrides,
setModelProviderPriorityOverrides,
type RoutingDefaultPolicy,
@@ -452,9 +453,18 @@ const providerRows = computed<ProviderPriorityRow[]>(() => {
.sort(comparePriorityRows)
})
const selectedFormatKey = computed(() => normalizeRoutingApiFormatKey(selectedApiFormat.value))
const selectedFormatKeyOverrides = computed<Record<string, number>>(() => (
targetModelPolicy.value.key_priority_overrides_by_format[selectedFormatKey.value] ?? {}
))
const keyRows = computed<KeyPriorityRow[]>(() => {
const format = selectedApiFormat.value
const keyOverrides = targetModelPolicy.value.key_priority_overrides
// 按格式覆盖优先；旧的不分格式覆盖仅作为兜底展示
const keyOverrides: Record<string, number> = {
...targetModelPolicy.value.key_priority_overrides,
...selectedFormatKeyOverrides.value,
}
const poolOverrides = targetModelPolicy.value.pool_priority_overrides
const normalRows: KeyPriorityRow[] = []
const poolGroups = new Map<string, GlobalKeySource[]>()
@@ -686,7 +696,7 @@ function setKeyPriority(keyId: string, event: Event): void {
})
} else {
updateKeyOverrides({
...targetModelPolicy.value.key_priority_overrides,
...selectedFormatKeyOverrides.value,
[row.target_id]: priority,
})
}
@@ -697,8 +707,14 @@ function moveKey(keyId: string, direction: -1 | 1): void {
updateVisibleKeyAndPoolOverrides(rows)
}
// Key 覆盖始终写入当前选中的 API 格式，不同格式互不影响
function updateKeyOverrides(overrides: Record<string, number>): void {
updateConfig(setModelKeyPriorityOverrides(config.value, targetModel.value, overrides))
updateConfig(setModelKeyPriorityOverridesForFormat(
config.value,
targetModel.value,
selectedApiFormat.value,
overrides,
))
}
function updatePoolOverrides(overrides: Record<string, number>): void {
@@ -710,7 +726,12 @@ function updateKeyAndPoolOverrides(
poolOverrides: Record<string, number>,
): void {
const next = setModelPoolPriorityOverrides(
setModelKeyPriorityOverrides(config.value, targetModel.value, keyOverrides),
setModelKeyPriorityOverridesForFormat(
config.value,
targetModel.value,
selectedApiFormat.value,
keyOverrides,
),
targetModel.value,
poolOverrides,
)
@@ -718,7 +739,7 @@ function updateKeyAndPoolOverrides(
}
function updateVisibleKeyAndPoolOverrides(rows: KeyPriorityRow[]): void {
const keyOverrides = { ...targetModelPolicy.value.key_priority_overrides }
const keyOverrides = { ...selectedFormatKeyOverrides.value }
const poolOverrides = { ...targetModelPolicy.value.pool_priority_overrides }
for (const row of keyRows.value) {
frontend/src/features/routing/utils/routingPolicy.ts
+100
View File
@@ -3,10 +3,15 @@ export type RoutingSchedulingMode = 'fixed_order' | 'cache_affinity' | 'load_bal
export type RoutingRulePhase = 'client_request' | 'provider_request'
export type RoutingSortingScope = 'unified' | 'per_model'
/** 首个候选（粘性 Key）的总尝试次数默认值：失败后同 Key 重试 1 次 */
export const DEFAULT_STICKY_KEY_ATTEMPTS = 2
export interface RoutingDefaultPolicy {
priority_mode: RoutingPriorityMode
scheduling_mode: RoutingSchedulingMode
keep_priority_on_conversion: boolean
/** 首个候选的总尝试次数；后续候选始终只尝试 1 次。0 或 1 表示不重试 */
sticky_key_attempts: number
}
export interface RoutingPoolSchedulingPreset {
@@ -25,6 +30,8 @@ export interface RoutingModelPolicy {
allowed_keys: string[]
provider_priority_overrides: Record<string, number>
key_priority_overrides: Record<string, number>
/** api_format -> key_id -> priority；同一 Key 在不同 API 格式下可独立排序 */
key_priority_overrides_by_format: Record<string, Record<string, number>>
pool_priority_overrides: Record<string, number>
pool_policy_overrides: Record<string, RoutingPoolPolicyOverride>
}
@@ -49,6 +56,7 @@ export interface RoutingSetSchedulingAction {
type: 'set_scheduling'
priority_mode: RoutingPriorityMode
scheduling_mode: RoutingSchedulingMode
sticky_key_attempts?: number
}
export interface RoutingGroupConfig {
@@ -68,12 +76,19 @@ export function createEmptyRoutingGroupConfig(): RoutingGroupConfig {
priority_mode: 'provider',
scheduling_mode: 'cache_affinity',
keep_priority_on_conversion: false,
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
},
model_policies: [],
rules: [],
}
}
export function normalizeStickyKeyAttempts(value: unknown): number {
const parsed = Math.trunc(Number(value))
if (!Number.isFinite(parsed) || parsed < 0) return DEFAULT_STICKY_KEY_ATTEMPTS
return Math.min(parsed, 99)
}
export function createEmptyModelPolicy(model = ''): RoutingModelPolicy {
return {
model,
@@ -81,6 +96,7 @@ export function createEmptyModelPolicy(model = ''): RoutingModelPolicy {
allowed_keys: [],
provider_priority_overrides: {},
key_priority_overrides: {},
key_priority_overrides_by_format: {},
pool_priority_overrides: {},
pool_policy_overrides: {},
}
@@ -94,6 +110,9 @@ export function normalizeRoutingGroupConfig(value: Partial<RoutingGroupConfig> |
default_policy: {
...base.default_policy,
...(value?.default_policy ?? {}),
sticky_key_attempts: normalizeStickyKeyAttempts(
value?.default_policy?.sticky_key_attempts ?? DEFAULT_STICKY_KEY_ATTEMPTS,
),
},
model_policies: Array.isArray(value?.model_policies)
? value.model_policies.map(policy => ({
@@ -103,6 +122,9 @@ export function normalizeRoutingGroupConfig(value: Partial<RoutingGroupConfig> |
allowed_keys: Array.isArray(policy.allowed_keys) ? [...policy.allowed_keys] : [],
provider_priority_overrides: { ...(policy.provider_priority_overrides ?? {}) },
key_priority_overrides: { ...(policy.key_priority_overrides ?? {}) },
key_priority_overrides_by_format: normalizeKeyPriorityOverridesByFormat(
policy.key_priority_overrides_by_format,
),
pool_priority_overrides: { ...(policy.pool_priority_overrides ?? {}) },
pool_policy_overrides: { ...(policy.pool_policy_overrides ?? {}) },
}))
@@ -296,6 +318,64 @@ export function setModelKeyPriorityOverrides(
})
}
export function normalizeRoutingApiFormatKey(apiFormat: string): string {
return apiFormat.trim().toLowerCase()
}
export function getModelKeyPriorityOverridesForFormat(
config: RoutingGroupConfig,
model: string,
apiFormat: string,
): Record<string, number> {
const policy = getModelPolicy(config, model)
const format = normalizeRoutingApiFormatKey(apiFormat)
return { ...(policy.key_priority_overrides_by_format[format] ?? {}) }
}
/**
* 某个 Key 在指定 API 格式下的生效覆盖值：按格式覆盖优先，其次是不分格式的 Key 覆盖。
*/
export function resolveModelKeyPriorityOverride(
config: RoutingGroupConfig,
model: string,
apiFormat: string,
keyId: string,
): number | undefined {
const policy = getModelPolicy(config, model)
const format = normalizeRoutingApiFormatKey(apiFormat)
return policy.key_priority_overrides_by_format[format]?.[keyId]
?? policy.key_priority_overrides[keyId]
}
export function setModelKeyPriorityOverridesForFormat(
config: RoutingGroupConfig,
model: string,
apiFormat: string,
overrides: Record<string, number>,
): RoutingGroupConfig {
const normalizedModel = model.trim() || DEFAULT_ROUTING_POLICY_MODEL
const format = normalizeRoutingApiFormatKey(apiFormat)
if (!format) return normalizeRoutingGroupConfig(config)
const current = getModelPolicy(config, normalizedModel)
const byFormat = { ...current.key_priority_overrides_by_format }
const normalized = normalizePriorityOverrides(overrides)
if (Object.keys(normalized).length > 0) {
byFormat[format] = normalized
} else {
delete byFormat[format]
}
if (normalizedModel === DEFAULT_ROUTING_POLICY_MODEL) {
return upsertDefaultModelPolicy(config, { key_priority_overrides_by_format: byFormat })
}
return upsertModelPolicy(config, {
...current,
model: normalizedModel,
key_priority_overrides_by_format: byFormat,
})
}
export function setModelPoolPriorityOverrides(
config: RoutingGroupConfig,
model: string,
@@ -347,6 +427,7 @@ export function getModelScheduling(
priority_mode: action?.priority_mode ?? normalized.default_policy.priority_mode,
scheduling_mode: action?.scheduling_mode ?? normalized.default_policy.scheduling_mode,
keep_priority_on_conversion: normalized.default_policy.keep_priority_on_conversion,
sticky_key_attempts: action?.sticky_key_attempts ?? normalized.default_policy.sticky_key_attempts,
}
}
@@ -464,6 +545,25 @@ export function normalizePriorityOverrides(overrides: Record<string, number>): R
return normalized
}
function normalizeKeyPriorityOverridesByFormat(
value: Record<string, Record<string, number>> | null | undefined,
): Record<string, Record<string, number>> {
const normalized: Record<string, Record<string, number>> = {}
if (!value || typeof value !== 'object') return normalized
for (const [rawFormat, overrides] of Object.entries(value)) {
const format = normalizeRoutingApiFormatKey(rawFormat)
if (!format || !overrides || typeof overrides !== 'object') continue
const merged = normalizePriorityOverrides({
...(normalized[format] ?? {}),
...overrides,
})
if (Object.keys(merged).length > 0) {
normalized[format] = merged
}
}
return normalized
}
function isSetSchedulingAction(action: unknown): action is RoutingSetSchedulingAction {
if (!action || typeof action !== 'object') return false
const candidate = action as Partial<RoutingSetSchedulingAction>
frontend/src/mocks/handler.ts
+3
View File
@@ -977,6 +977,7 @@ const MOCK_ROUTING_GROUPS: MockRoutingGroup[] = [
priority_mode: 'provider',
scheduling_mode: 'cache_affinity',
keep_priority_on_conversion: false,
sticky_key_attempts: 2,
},
model_policies: [
{
@@ -985,6 +986,8 @@ const MOCK_ROUTING_GROUPS: MockRoutingGroup[] = [
allowed_keys: [],
provider_priority_overrides: { 'provider-002': 0 },
key_priority_overrides: {},
key_priority_overrides_by_format: {},
pool_priority_overrides: {},
pool_policy_overrides: {},
},
],
frontend/src/views/admin/ProviderManagement.vue
+13 -1
View File
@@ -272,6 +272,8 @@ import {
type ProviderWithEndpointsSummary,
} from '@/api/endpoints'
import { adminApi } from '@/api/admin'
import { listRoutingGroups } from '@/api/routing-profiles'
import { normalizeRoutingGroupConfig } from '@/features/routing/utils/routingPolicy'
import { parseApiError } from '@/utils/errorParser'
import { useI18n } from '@/i18n'
@@ -533,8 +535,18 @@ const maxProviderPriority = computed(() => {
return priorities.length > 0 ? Math.max(...priorities) : undefined
})
// 加载优先级模式
// 加载优先级模式：优先使用启用中的系统默认调度策略，旧的系统配置键仅作兜底
async function loadPriorityMode(options: { cacheTtlMs?: number } = {}) {
try {
const groups = await listRoutingGroups()
const systemDefault = groups.items.find(group => group.is_system_default && group.enabled)
if (systemDefault) {
priorityMode.value = normalizeRoutingGroupConfig(systemDefault.config_json).default_policy.priority_mode
return
}
} catch {
// 路由策略不可用时继续尝试旧配置
}
try {
const response = await adminApi.getSystemConfig('provider_priority_mode', {
cacheTtlMs: options.cacheTtlMs ?? 0,
frontend/src/views/admin/RoutingProfiles.vue
+69
View File
@@ -454,6 +454,44 @@
</div>
</div>
</div>
<label
class="flex items-start gap-3 rounded-lg border border-border/60 px-3 py-2 text-sm"
data-testid="keep-priority-on-conversion"
>
<Switch
:model-value="keepPriorityOnConversion"
:disabled="saving"
aria-label="格式转换时保持优先级"
@update:model-value="updateKeepPriorityOnConversion"
/>
<span class="min-w-0">
<span class="block font-medium">格式转换时保持优先级</span>
<span class="mt-0.5 block text-xs text-muted-foreground">
开启后，需要跨 API 格式转换的候选不会被降级到同格式候选之后。作用于本策略范围内的全部模型；Provider 自身的同名开关仍单独生效。
</span>
</span>
</label>
<label
class="flex items-start gap-3 rounded-lg border border-border/60 px-3 py-2 text-sm"
data-testid="sticky-key-attempts"
>
<Input
:model-value="stickyKeyAttempts"
type="number"
min="0"
max="99"
class="w-20 shrink-0"
:disabled="saving"
aria-label="粘性 Key 尝试次数"
@update:model-value="updateStickyKeyAttempts"
/>
<span class="min-w-0">
<span class="block font-medium">粘性 Key 尝试次数</span>
<span class="mt-0.5 block text-xs text-muted-foreground">
首个候选（缓存亲和命中的 Key）的总尝试次数。2 表示失败后同 Key 重试 1 次再转移，避免偶发错误破坏缓存；0 或 1 表示不重试。转移后的候选始终只尝试 1 次。
</span>
</span>
</label>
</div>
<RoutingPriorityPolicyEditor
@@ -749,6 +787,7 @@ import {
Button,
Card,
Input,
Switch,
Table,
TableBody,
TableCard,
@@ -762,6 +801,7 @@ import { DropdownMenu, DropdownMenuTrigger, DropdownMenuContent, DropdownMenuIte
import { AlertDialog } from '@/components/common'
import {
DEFAULT_ROUTING_POLICY_MODEL,
DEFAULT_STICKY_KEY_ATTEMPTS,
allowedModelsMirrorPerModelPolicies,
clearAllowedModels,
copyPerModelRoutingConfig,
@@ -772,6 +812,7 @@ import {
isGeneratedModelSchedulingRule,
modelSchedulingRuleId,
normalizeRoutingGroupConfig,
normalizeStickyKeyAttempts,
removePerModelRoutingConfig,
routingModelScopeLabel,
savePerModelRoutingConfig,
@@ -874,6 +915,12 @@ const firstStepSchedulingMode = computed<RoutingSchedulingMode>(() => {
}
return draft.value?.config_json.default_policy.scheduling_mode ?? 'cache_affinity'
})
const keepPriorityOnConversion = computed<boolean>(() => (
draft.value?.config_json.default_policy.keep_priority_on_conversion ?? false
))
const stickyKeyAttempts = computed<number>(() => (
draft.value?.config_json.default_policy.sticky_key_attempts ?? DEFAULT_STICKY_KEY_ATTEMPTS
))
const allowedModelsLookLikeLegacyMirror = computed(() => {
return draft.value
? allowedModelsMirrorPerModelPolicies(draft.value.config_json)
@@ -1190,6 +1237,28 @@ function updateFirstStepSchedulingMode(mode: RoutingSchedulingMode): void {
})
}
function updateStickyKeyAttempts(value: string | number): void {
if (!draft.value) return
updateDraftConfig({
...draft.value.config_json,
default_policy: {
...draft.value.config_json.default_policy,
sticky_key_attempts: normalizeStickyKeyAttempts(value),
},
})
}
function updateKeepPriorityOnConversion(value: boolean): void {
if (!draft.value) return
updateDraftConfig({
...draft.value.config_json,
default_policy: {
...draft.value.config_json.default_policy,
keep_priority_on_conversion: value,
},
})
}
function removePerModelPolicy(model: string): void {
if (!draft.value) return
if (perModelEditingActive.value && editingDirty.value) {
frontend/src/views/admin/__tests__/RoutingProfiles.allowed-models.spec.ts
+21
View File
@@ -131,11 +131,31 @@ vi.mock('@/components/ui', async () => {
},
})
const Switch = defineComponent({
inheritAttrs: false,
props: {
modelValue: { type: Boolean, default: false },
disabled: Boolean,
},
emits: ['update:modelValue'],
setup(props, { attrs, emit }) {
return () => h('button', {
...attrs,
type: 'button',
role: 'switch',
'aria-checked': props.modelValue,
disabled: props.disabled,
onClick: () => emit('update:modelValue', !props.modelValue),
})
},
})
return {
Badge: wrapper(),
Button,
Card: wrapper('section'),
Input,
Switch,
Table: wrapper('table'),
TableBody: wrapper('tbody'),
TableCard: wrapper(),
@@ -195,6 +215,7 @@ function routingGroup(
priority_mode: 'provider',
scheduling_mode: 'cache_affinity',
keep_priority_on_conversion: false,
sticky_key_attempts: 2,
},
model_policies: [],
rules: [],