diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs index a416eff57..8665342e9 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs @@ -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 { 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( where F: Fn(&EligibleLocalExecutionCandidate) -> Option + 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( @@ -1924,32 +1901,15 @@ fn build_unpersisted_local_execution_candidate_attempts( candidate: EligibleLocalExecutionCandidate, candidate_index: u32, ) -> VecDeque { - 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( @@ -2277,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, } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs index 57ee9f2b7..b1e12057e 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs @@ -371,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(), @@ -406,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(), @@ -443,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)]), diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs index fdf7b00cf..91d165041 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs @@ -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) } diff --git a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs index f09dcdf72..74e1a1b7f 100644 --- a/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs +++ b/apps/aether-gateway/src/ai_serving/planner/candidate_source.rs @@ -1889,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(), @@ -1952,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(), @@ -2689,6 +2691,7 @@ mod tests { priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider, scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder, 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(), diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs index ed3797c25..e65590f58 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs @@ -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) diff --git a/apps/aether-gateway/src/ai_serving/planner/report_context.rs b/apps/aether-gateway/src/ai_serving/planner/report_context.rs index 7e3a29343..3c5d6e53b 100644 --- a/apps/aether-gateway/src/ai_serving/planner/report_context.rs +++ b/apps/aether-gateway/src/ai_serving/planner/report_context.rs @@ -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, + /// Routing policy sticky-key attempt budget; read back by the attempt + /// loop to derive same-key retries lazily. + pub(crate) sticky_key_attempts: Option, 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, diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs index 8f329cf8b..9dffd02df 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs @@ -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, diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs index ea13663db..c45cf08c8 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs @@ -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, diff --git a/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs b/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs index f252354db..c1207443e 100644 --- a/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs +++ b/apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs @@ -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, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs index a025c1349..85be2f3fe 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs @@ -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) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs index 4b6644845..496e57153 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs @@ -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) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs index dfc950f2d..3769ad52d 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs @@ -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) diff --git a/apps/aether-gateway/src/dispatch/pool_scheduler.rs b/apps/aether-gateway/src/dispatch/pool_scheduler.rs index 5dc33514d..f1aa9b5af 100644 --- a/apps/aether-gateway/src/dispatch/pool_scheduler.rs +++ b/apps/aether-gateway/src/dispatch/pool_scheduler.rs @@ -1951,11 +1951,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 } @@ -2159,6 +2161,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)); @@ -2176,6 +2179,7 @@ mod tests { pool_key_index: None, pool_key_lease: None, scheduler_affinity_epoch: None, + sticky_key_attempts: None, } ); } @@ -5072,6 +5076,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() diff --git a/apps/aether-gateway/src/dispatch/refs.rs b/apps/aether-gateway/src/dispatch/refs.rs index e8acb0a1d..4c689cfb3 100644 --- a/apps/aether-gateway/src/dispatch/refs.rs +++ b/apps/aether-gateway/src/dispatch/refs.rs @@ -208,6 +208,7 @@ mod tests { pool_key_index: None, pool_key_lease: None, scheduler_affinity_epoch: None, + sticky_key_attempts: None, }, ranking: None, } diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 5b0e5409f..9a689089d 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -252,6 +252,10 @@ where Ok(()) } + async fn next_same_key_retry(&self, attempt: &T) -> Result, 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 = 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, 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, diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index 366690283..e054ea9ef 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -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)>::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::>(), + [ + "endpoint-retry", + "endpoint-retry", + "endpoint-retry", + "endpoint-success" + ] + ); + let sticky_candidate_ids = seen_plans[..3] + .iter() + .map(|(_, candidate_id)| candidate_id.clone()) + .collect::>(); + 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 diff --git a/apps/aether-gateway/src/orchestration/attempt.rs b/apps/aether-gateway/src/orchestration/attempt.rs index 1ae62662f..aa4c1d0be 100644 --- a/apps/aether-gateway/src/orchestration/attempt.rs +++ b/apps/aether-gateway/src/orchestration/attempt.rs @@ -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, pub(crate) pool_key_lease: Option, pub(crate) scheduler_affinity_epoch: Option, + /// Routing-policy `sticky_key_attempts` in effect for this request. `None` + /// means the policy default applies. + pub(crate) sticky_key_attempts: Option, } 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 { @@ -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 { - 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, +) -> Option { + 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 { - 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(attempt: &A) -> Option { + 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, - endpoint_max_retries: Option, - provider_config: Option, - ) -> 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), } ); } diff --git a/apps/aether-gateway/src/orchestration/mod.rs b/apps/aether-gateway/src/orchestration/mod.rs index 8b86eec98..ef339f46f 100644 --- a/apps/aether-gateway/src/orchestration/mod.rs +++ b/apps/aether-gateway/src/orchestration/mod.rs @@ -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, diff --git a/apps/aether-gateway/src/routing/resolver.rs b/apps/aether-gateway/src/routing/resolver.rs index 451bbedd7..fc73a547e 100644 --- a/apps/aether-gateway/src/routing/resolver.rs +++ b/apps/aether-gateway/src/routing/resolver.rs @@ -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, 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, GatewayError> { +) -> Result, 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( diff --git a/apps/aether-gateway/src/scheduler/config.rs b/apps/aether-gateway/src/scheduler/config.rs index d566827df..fe743f3e4 100644 --- a/apps/aether-gateway/src/scheduler/config.rs +++ b/apps/aether-gateway/src/scheduler/config.rs @@ -1,6 +1,7 @@ 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; @@ -37,6 +38,8 @@ 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 { @@ -45,6 +48,7 @@ impl Default for SchedulerOrderingConfig { priority_mode: SchedulerPriorityMode::Provider, scheduling_mode: SchedulerSchedulingMode::CacheAffinity, keep_priority_on_conversion: false, + sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS, } } } @@ -57,6 +61,7 @@ impl SchedulerOrderingConfig { 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, } } @@ -65,6 +70,7 @@ impl SchedulerOrderingConfig { 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, } } @@ -80,6 +86,7 @@ impl SchedulerOrderingConfig { SchedulerSchedulingMode::LoadBalance => RoutingSchedulingMode::LoadBalance, }, keep_priority_on_conversion: self.keep_priority_on_conversion, + sticky_key_attempts: self.sticky_key_attempts, } } @@ -221,6 +228,8 @@ pub(crate) async fn read_legacy_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, }) } @@ -358,7 +367,8 @@ mod tests { json!({ "priority_mode": "global_key", "scheduling_mode": "load_balance", - "keep_priority_on_conversion": true + "keep_priority_on_conversion": true, + "sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS }) ); diff --git a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs index 4759ab296..35b0623ad 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream/decision.rs @@ -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); diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs b/apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs index cc8de1774..345a0f804 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/chat/failover.rs @@ -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); diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/search.rs b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs index 3e3cfd577..5b2ec2067 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/search.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/search.rs @@ -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::>(); + // 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(); diff --git a/apps/aether-gateway/src/tests/architecture/ai_serving.rs b/apps/aether-gateway/src/tests/architecture/ai_serving.rs index eed95eceb..3859d68e8 100644 --- a/apps/aether-gateway/src/tests/architecture/ai_serving.rs +++ b/apps/aether-gateway/src/tests/architecture/ai_serving.rs @@ -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", diff --git a/apps/aether-gateway/src/tests/usage.rs b/apps/aether-gateway/src/tests/usage.rs index 77e865e61..43cd56bc0 100644 --- a/apps/aether-gateway/src/tests/usage.rs +++ b/apps/aether-gateway/src/tests/usage.rs @@ -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, diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index d59d43a95..af9937cb0 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -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, diff --git a/crates/aether-ai/serving/src/attempt_loop.rs b/crates/aether-ai/serving/src/attempt_loop.rs index c2df49cfc..ed3c21a45 100644 --- a/crates/aether-ai/serving/src/attempt_loop.rs +++ b/crates/aether-ai/serving/src/attempt_loop.rs @@ -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 + 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 { 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, Self::Error> { + Ok(None) + } + async fn mark_unused_attempts(&self, attempts: Vec) -> 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 = None; let mut last_attempted = None; let mut retry_filters: Vec = 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) { + 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 { + 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 { + 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)] diff --git a/crates/aether-ai/serving/src/candidate_persistence.rs b/crates/aether-ai/serving/src/candidate_persistence.rs index 287b67d2f..7fe326f6c 100644 --- a/crates/aether-ai/serving/src/candidate_persistence.rs +++ b/crates/aether-ai/serving/src/candidate_persistence.rs @@ -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; fn generate_candidate_id(&self) -> String; @@ -42,50 +40,28 @@ pub async fn run_ai_available_candidate_persistence( 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 { 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"] ); } diff --git a/crates/aether-ai/serving/src/lib.rs b/crates/aether-ai/serving/src/lib.rs index 502bfef9a..99dbac00d 100644 --- a/crates/aether-ai/serving/src/lib.rs +++ b/crates/aether-ai/serving/src/lib.rs @@ -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, diff --git a/crates/aether-routing-core/src/actions.rs b/crates/aether-routing-core/src/actions.rs index 370c9e618..71bb09f80 100644 --- a/crates/aether-routing-core/src/actions.rs +++ b/crates/aether-routing-core/src/actions.rs @@ -73,6 +73,8 @@ pub enum RoutingAction { priority_mode: Option, scheduling_mode: Option, keep_priority_on_conversion: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + sticky_key_attempts: Option, }, SetProviderPriority { provider_id: String, diff --git a/crates/aether-routing-core/src/lib.rs b/crates/aether-routing-core/src/lib.rs index d92d7010e..a8dceb721 100644 --- a/crates/aether-routing-core/src/lib.rs +++ b/crates/aether-routing-core/src/lib.rs @@ -15,7 +15,7 @@ pub use conditions::{RoutingCondition, RoutingConditionContext, RoutingCondition pub use model::{ RoutingDefaultPolicy, RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord, RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride, - RoutingRule, RoutingSchedulingPreset, + RoutingRule, RoutingSchedulingPreset, DEFAULT_STICKY_KEY_ATTEMPTS, }; pub use mutations::{ apply_json_patch_operations, validate_header_patch, validate_json_patch_operations, diff --git a/crates/aether-routing-core/src/model.rs b/crates/aether-routing-core/src/model.rs index 683d527aa..4526ece7f 100644 --- a/crates/aether-routing-core/src/model.rs +++ b/crates/aether-routing-core/src/model.rs @@ -23,7 +23,11 @@ pub struct RoutingPoolPolicyOverride { pub scheduling_presets: Vec, } -#[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)] diff --git a/crates/aether-routing-core/src/policy.rs b/crates/aether-routing-core/src/policy.rs index 70b66157e..4ee70067b 100644 --- a/crates/aether-routing-core/src/policy.rs +++ b/crates/aether-routing-core/src/policy.rs @@ -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(), @@ -205,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; @@ -215,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, @@ -282,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 == "*" { @@ -384,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(), @@ -415,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"] @@ -448,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 @@ -456,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 { diff --git a/crates/aether-routing-core/src/validation.rs b/crates/aether-routing-core/src/validation.rs index 54274907a..8314cc8c8 100644 --- a/crates/aether-routing-core/src/validation.rs +++ b/crates/aether-routing-core/src/validation.rs @@ -271,6 +271,7 @@ mod tests { priority_mode: None, scheduling_mode: None, keep_priority_on_conversion: Some(true), + sticky_key_attempts: None, }, "set_scheduling", ), diff --git a/frontend/src/features/providers/components/ProviderFormDialog.vue b/frontend/src/features/providers/components/ProviderFormDialog.vue index ea877c215..b1f90f2a5 100644 --- a/frontend/src/features/providers/components/ProviderFormDialog.vue +++ b/frontend/src/features/providers/components/ProviderFormDialog.vue @@ -155,17 +155,6 @@ -
- - -
diff --git a/frontend/src/features/routing/__tests__/routingPolicy.spec.ts b/frontend/src/features/routing/__tests__/routingPolicy.spec.ts index 0618bce86..88b7d7913 100644 --- a/frontend/src/features/routing/__tests__/routingPolicy.spec.ts +++ b/frontend/src/features/routing/__tests__/routingPolicy.spec.ts @@ -12,6 +12,7 @@ import { getModelScheduling, modelSchedulingRuleId, normalizeRoutingGroupConfig, + normalizeStickyKeyAttempts, parseAllowedModelsInput, removePerModelRoutingConfig, resolveModelKeyPriorityOverride, @@ -77,6 +78,19 @@ 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(), diff --git a/frontend/src/features/routing/utils/routingPolicy.ts b/frontend/src/features/routing/utils/routingPolicy.ts index 72d6afd43..dc3e5749f 100644 --- a/frontend/src/features/routing/utils/routingPolicy.ts +++ b/frontend/src/features/routing/utils/routingPolicy.ts @@ -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 { @@ -51,6 +56,7 @@ export interface RoutingSetSchedulingAction { type: 'set_scheduling' priority_mode: RoutingPriorityMode scheduling_mode: RoutingSchedulingMode + sticky_key_attempts?: number } export interface RoutingGroupConfig { @@ -70,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, @@ -97,6 +110,9 @@ export function normalizeRoutingGroupConfig(value: Partial | 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 => ({ @@ -411,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, } } diff --git a/frontend/src/mocks/handler.ts b/frontend/src/mocks/handler.ts index 999c87c03..af5af0ed9 100644 --- a/frontend/src/mocks/handler.ts +++ b/frontend/src/mocks/handler.ts @@ -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: [ { diff --git a/frontend/src/views/admin/RoutingProfiles.vue b/frontend/src/views/admin/RoutingProfiles.vue index 07074104f..308d2a77e 100644 --- a/frontend/src/views/admin/RoutingProfiles.vue +++ b/frontend/src/views/admin/RoutingProfiles.vue @@ -471,6 +471,27 @@ + (() => { const keepPriorityOnConversion = computed(() => ( draft.value?.config_json.default_policy.keep_priority_on_conversion ?? false )) +const stickyKeyAttempts = computed(() => ( + draft.value?.config_json.default_policy.sticky_key_attempts ?? DEFAULT_STICKY_KEY_ATTEMPTS +)) const allowedModelsLookLikeLegacyMirror = computed(() => { return draft.value ? allowedModelsMirrorPerModelPolicies(draft.value.config_json) @@ -1211,6 +1237,17 @@ 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({ diff --git a/frontend/src/views/admin/__tests__/RoutingProfiles.allowed-models.spec.ts b/frontend/src/views/admin/__tests__/RoutingProfiles.allowed-models.spec.ts index e2f94adb6..c7daf5c48 100644 --- a/frontend/src/views/admin/__tests__/RoutingProfiles.allowed-models.spec.ts +++ b/frontend/src/views/admin/__tests__/RoutingProfiles.allowed-models.spec.ts @@ -215,6 +215,7 @@ function routingGroup( priority_mode: 'provider', scheduling_mode: 'cache_affinity', keep_priority_on_conversion: false, + sticky_key_attempts: 2, }, model_policies: [], rules: [],