Restrict scheduler affinity to cache affinity mode

apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs

@@ -31,6 +31,7 @@ use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
 use crate::clock::current_unix_ms;
 use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
 use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
+use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
 use crate::{AppState, GatewayError};
 
 const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
@@ -182,6 +183,7 @@ struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
     request_auth_channel: Option<&'a str>,
     persistence_policy: LocalCandidatePersistencePolicy<'a>,
     resolution_mode: LocalCandidateResolutionMode,
+    scheduler_cache_affinity_enabled: bool,
     build_available_extra_data: F,
     decorate_skipped_candidate: G,
 }
@@ -244,6 +246,9 @@ where
     }
 
     fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
+        if !self.scheduler_cache_affinity_enabled {
+            return;
+        }
         remember_first_local_candidate_affinity(
             self.state,
             self.auth_snapshot,
@@ -430,6 +435,7 @@ where
     F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
     G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
 {
+    let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
     let port = GatewayLocalCandidateMaterializationPort {
         state,
         trace_id,
@@ -442,6 +448,7 @@ where
         request_auth_channel,
         persistence_policy,
         resolution_mode,
+        scheduler_cache_affinity_enabled,
         build_available_extra_data,
         decorate_skipped_candidate,
     };
@@ -474,6 +481,7 @@ where
     F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
     G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
 {
+    let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
     let _ = build_available_extra_data;
     let (candidates, resolved_skipped) = resolve_and_rank_logical_local_execution_candidates(
         state,
@@ -505,14 +513,16 @@ where
         }
     }
 
-    remember_first_local_candidate_affinity(
+    if scheduler_cache_affinity_enabled {
-        state,
+        remember_first_local_candidate_affinity(
-        auth_snapshot,
+            state,
-        client_session_affinity,
+            auth_snapshot,
-        client_api_format,
+            client_session_affinity,
-        requested_model,
+            client_api_format,
-        &candidates,
+            requested_model,
-    );
+            &candidates,
+        );
+    }
 
     let (items, _) = build_logical_candidate_items(
         state,
@@ -597,6 +607,7 @@ where
     F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync + 'a,
     G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
 {
+    let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
     let _ = build_available_extra_data;
     let decorate_skipped_candidate = Arc::new(decorate_skipped_candidate);
     let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
@@ -630,6 +641,7 @@ where
         candidate_count: 0,
         next_candidate_index: 0,
         remembered_affinity: false,
+        scheduler_cache_affinity_enabled,
     };
     cursor.load_next_page().await;
     let candidate_count = cursor.candidate_count;
@@ -665,6 +677,7 @@ struct RequestedModelAttemptPageCursor<'a> {
     candidate_count: usize,
     next_candidate_index: u32,
     remembered_affinity: bool,
+    scheduler_cache_affinity_enabled: bool,
 }
 
 impl<'a> RequestedModelAttemptPageCursor<'a> {
@@ -726,7 +739,10 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
                     );
                 }
             }
-            if !self.remembered_affinity && !candidates.is_empty() {
+            if self.scheduler_cache_affinity_enabled
+                && !self.remembered_affinity
+                && !candidates.is_empty()
+            {
                 remember_first_local_candidate_affinity(
                     self.state,
                     Some(&self.auth_snapshot),
@@ -795,6 +811,21 @@ async fn pop_attempt_from_items(
     }
 }
 
+async fn scheduler_cache_affinity_enabled(state: PlannerAppState<'_>) -> bool {
+    match read_scheduler_ordering_config(state.app()).await {
+        Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
+        Err(error) => {
+            warn!(
+                event_name = "planner_scheduler_affinity_config_load_failed",
+                log_type = "event",
+                error = ?error,
+                "failed to load scheduler config while checking cache affinity mode"
+            );
+            SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
+        }
+    }
+}
+
 pub(crate) fn remember_first_local_candidate_affinity(
     state: PlannerAppState<'_>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
@@ -1211,6 +1242,7 @@ mod tests {
     use std::sync::Arc;
 
     use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
+    use aether_data::repository::auth::StoredAuthApiKeySnapshot;
     use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
     use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
     use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -1219,6 +1251,7 @@ mod tests {
         GatewayProviderTransportProvider,
     };
     use aether_scheduler_core::{
+        build_scheduler_affinity_cache_key_for_api_key_id,
         SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankingMode,
         SchedulerRankingOutcome,
     };
@@ -1227,6 +1260,7 @@ mod tests {
     use super::*;
     use crate::data::GatewayDataState;
     use crate::orchestration::LocalExecutionCandidateMetadata;
+    use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
 
     fn sample_candidate(key_id: &str) -> SchedulerMinimalCandidateSelectionCandidate {
         SchedulerMinimalCandidateSelectionCandidate {
@@ -1331,6 +1365,46 @@ mod tests {
         }
     }
 
+    fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
+        GatewayAuthApiKeySnapshot::from_stored(
+            StoredAuthApiKeySnapshot::new(
+                "user-1".to_string(),
+                "alice".to_string(),
+                Some("alice@example.com".to_string()),
+                "user".to_string(),
+                "local".to_string(),
+                true,
+                false,
+                None,
+                None,
+                None,
+                "api-key-1".to_string(),
+                Some("default".to_string()),
+                true,
+                false,
+                false,
+                Some(60),
+                Some(5),
+                Some(4_102_444_800),
+                None,
+                None,
+                None,
+            )
+            .expect("stored auth snapshot should build"),
+            0,
+        )
+    }
+
+    fn no_extra_data(_: &EligibleLocalExecutionCandidate) -> Option<Value> {
+        None
+    }
+
+    fn identity_skipped_candidate(
+        candidate: SkippedLocalExecutionCandidate,
+    ) -> SkippedLocalExecutionCandidate {
+        candidate
+    }
+
     #[tokio::test]
     async fn pool_group_keys_are_not_persisted_as_available_before_attempt() {
         let repository = Arc::new(InMemoryRequestCandidateRepository::default());
@@ -1368,6 +1442,55 @@ mod tests {
         assert_eq!(stored[0].candidate_index, 2);
     }
 
+    #[test]
+    fn materialization_port_ignores_scheduler_affinity_when_cache_affinity_disabled() {
+        let app = AppState::new().expect("state should build");
+        let auth_snapshot = sample_auth_snapshot();
+        let port = GatewayLocalCandidateMaterializationPort {
+            state: PlannerAppState::new(&app),
+            trace_id: "trace-affinity-disabled",
+            client_api_format: "openai:chat",
+            requested_model: Some("gpt-5"),
+            auth_snapshot: Some(&auth_snapshot),
+            client_session_affinity: None,
+            required_capabilities: None,
+            sticky_session_token: None,
+            request_auth_channel: None,
+            persistence_policy: LocalCandidatePersistencePolicy {
+                available: LocalAvailableCandidatePersistenceContext {
+                    user_id: "user-1",
+                    api_key_id: "api-key-1",
+                    required_capabilities: None,
+                    error_context: "test available",
+                },
+                skipped: LocalSkippedCandidatePersistenceContext {
+                    user_id: "user-1",
+                    api_key_id: "api-key-1",
+                    required_capabilities: None,
+                    error_context: "test skipped",
+                    record_runtime_miss_diagnostic: false,
+                },
+            },
+            resolution_mode: LocalCandidateResolutionMode::Standard,
+            scheduler_cache_affinity_enabled: false,
+            build_available_extra_data: no_extra_data,
+            decorate_skipped_candidate: identity_skipped_candidate,
+        };
+        let candidate = sample_eligible("key-a", None);
+        let cache_key =
+            build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
+                .expect("scheduler affinity cache key should build");
+
+        aether_ai_serving::AiCandidateMaterializationPort::remember_first_candidate_affinity(
+            &port,
+            &[candidate],
+        );
+
+        assert!(app
+            .read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
+            .is_none());
+    }
+
     #[tokio::test]
     async fn logical_materialization_does_not_persist_pool_group_representative() {
         let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());

apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs

@@ -23,8 +23,10 @@ use crate::ai_serving::{
     candidate_common_transport_skip_reason, CandidateTransportPolicyFacts, PlannerAppState,
 };
 use crate::clock::current_unix_ms;
-use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
 use crate::handlers::shared::provider_pool::read_admin_provider_pool_runtime_state;
+use crate::handlers::shared::provider_pool::{
+    admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
+};
 use crate::handlers::shared::provider_pool::{
     try_claim_admin_provider_pool_key, AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
 };
@@ -306,6 +308,9 @@ impl<'a> PoolKeyCursor<'a> {
 
     async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
         let pool_config = pool_config_for_candidate(&self.group)?;
+        if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
+            return None;
+        }
         let runtime = read_admin_provider_pool_runtime_state(
             self.state.app().runtime_state.as_ref(),
             self.group.candidate.provider_id.as_str(),
@@ -1206,7 +1211,7 @@ mod tests {
     }
 
     #[test]
-    fn pool_scheduler_promotes_sticky_hit_regardless_distribution_mode() {
+    fn pool_scheduler_ignores_sticky_hit_without_cache_affinity() {
         let key_a = sample_eligible_candidate(
             "provider-pool",
             "endpoint-1",
@@ -1250,7 +1255,7 @@ mod tests {
                 .iter()
                 .map(|item| item.candidate.key_id.as_str())
                 .collect::<Vec<_>>(),
-            vec!["key-a", "key-b"]
+            vec!["key-b", "key-a"]
         );
     }
 

apps/aether-gateway/src/handlers/admin/provider/pool/config.rs

@@ -215,6 +215,28 @@ fn admin_provider_pool_lru_enabled(
         .any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
 }
 
+pub(crate) fn admin_provider_pool_cache_affinity_enabled(
+    pool_config: &AdminProviderPoolConfig,
+) -> bool {
+    let mut seen = std::collections::BTreeSet::new();
+    for item in &pool_config.scheduling_presets {
+        let preset = item.preset.trim().to_ascii_lowercase();
+        if preset.is_empty() || !seen.insert(preset.clone()) {
+            continue;
+        }
+        if !item.enabled {
+            continue;
+        }
+        if matches!(
+            preset.as_str(),
+            "lru" | "cache_affinity" | "load_balance" | "single_account"
+        ) {
+            return preset == "cache_affinity";
+        }
+    }
+    false
+}
+
 pub(crate) fn admin_provider_pool_config(
     provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
 ) -> Option<AdminProviderPoolConfig> {
@@ -331,7 +353,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
 
 #[cfg(test)]
 mod tests {
-    use super::{admin_provider_pool_config, admin_provider_pool_config_from_config_value};
+    use super::{
+        admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config,
+        admin_provider_pool_config_from_config_value,
+    };
     use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
     use serde_json::json;
 
@@ -551,4 +576,29 @@ mod tests {
         assert_eq!(config.unschedulable_rules[1].keyword, "review_required");
         assert_eq!(config.unschedulable_rules[1].duration_minutes, 5);
     }
+
+    #[test]
+    fn cache_affinity_enabled_only_when_it_is_distribution_mode() {
+        let cache_affinity = admin_provider_pool_config_from_config_value(Some(&json!({
+            "pool_advanced": {
+                "scheduling_presets": [
+                    {"preset": "cache_affinity", "enabled": true},
+                    {"preset": "priority_first", "enabled": true}
+                ]
+            }
+        })))
+        .expect("pool config should parse");
+        assert!(admin_provider_pool_cache_affinity_enabled(&cache_affinity));
+
+        let load_balance = admin_provider_pool_config_from_config_value(Some(&json!({
+            "pool_advanced": {
+                "scheduling_presets": [
+                    {"preset": "load_balance", "enabled": true},
+                    {"preset": "cache_affinity", "enabled": true}
+                ]
+            }
+        })))
+        .expect("pool config should parse");
+        assert!(!admin_provider_pool_cache_affinity_enabled(&load_balance));
+    }
 }

apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs

@@ -3,6 +3,7 @@ use super::keys::{
     pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
     pool_sticky_pattern,
 };
+use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
 use crate::handlers::admin::provider::shared::support::{
     AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
 };
@@ -44,11 +45,13 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
     let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
     let cost_keys = pool_cost_keys(provider_id, key_ids);
     let latency_keys = pool_latency_keys(provider_id, key_ids);
+    let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
+        && admin_provider_pool_cache_affinity_enabled(pool_config);
 
     if let Some(sticky_session_token) = sticky_session_token
         .map(str::trim)
         .filter(|value| !value.is_empty())
-        .filter(|_| pool_config.sticky_session_ttl_seconds > 0)
+        .filter(|_| sticky_sessions_enabled)
     {
         let sticky_key = pool_sticky_key(provider_id, sticky_session_token);
         if let Ok(Some(bound_key_id)) = runtime.kv_get(&sticky_key).await {
@@ -77,22 +80,24 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
         }
     }
 
-    let sticky_keys = runtime
+    if sticky_sessions_enabled {
-        .scan_keys(&pool_sticky_pattern(provider_id), 200)
+        let sticky_keys = runtime
-        .await
+            .scan_keys(&pool_sticky_pattern(provider_id), 200)
-        .unwrap_or_default();
+            .await
-    state.total_sticky_sessions = sticky_keys.len();
+            .unwrap_or_default();
-    if !sticky_keys.is_empty() {
+        state.total_sticky_sessions = sticky_keys.len();
-        let raw_keys = sticky_keys
+        if !sticky_keys.is_empty() {
-            .iter()
+            let raw_keys = sticky_keys
-            .map(|key| runtime.strip_namespace(key).to_string())
+                .iter()
-            .collect::<Vec<_>>();
+                .map(|key| runtime.strip_namespace(key).to_string())
-        if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
+                .collect::<Vec<_>>();
-            for bound_key_id in values.into_iter().flatten() {
+            if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
-                *state
+                for bound_key_id in values.into_iter().flatten() {
-                    .sticky_sessions_by_key
+                    *state
-                    .entry(bound_key_id)
+                        .sticky_sessions_by_key
-                    .or_insert(0) += 1;
+                        .entry(bound_key_id)
+                        .or_insert(0) += 1;
+                }
             }
         }
     }

apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs

@@ -2,6 +2,7 @@ use super::keys::{
     pool_cooldown_index_key, pool_cooldown_key, pool_cost_key, pool_latency_key, pool_lru_key,
     pool_sticky_key, pool_stream_timeout_key,
 };
+use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
 use crate::handlers::admin::provider::shared::support::{
     AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
 };
@@ -346,6 +347,7 @@ pub(crate) async fn record_admin_provider_pool_success(
         .map(str::trim)
         .filter(|value| !value.is_empty())
         .filter(|_| pool_config.sticky_session_ttl_seconds > 0)
+        .filter(|_| admin_provider_pool_cache_affinity_enabled(pool_config))
     {
         let _ = runtime
             .kv_set(
@@ -810,6 +812,50 @@ mod tests {
         assert!(runtime.lru_score_by_key.contains_key("key-1"));
     }
 
+    #[tokio::test]
+    async fn success_feedback_does_not_write_sticky_without_cache_affinity() {
+        let Some(redis) = start_managed_redis_or_skip().await else {
+            return;
+        };
+        let app = build_runner_app(redis.redis_url(), "pool_runtime_no_sticky_load_balance").await;
+        let runtime = app.runtime_state.as_ref();
+        let mut pool_config = sample_pool_config();
+        pool_config.scheduling_presets = vec![AdminProviderPoolSchedulingPreset {
+            preset: "load_balance".to_string(),
+            enabled: true,
+            mode: None,
+        }];
+        pool_config.lru_enabled = false;
+        let key_ids = vec!["key-1".to_string()];
+
+        record_admin_provider_pool_success(
+            runtime,
+            "provider-1",
+            "key-1",
+            &pool_config,
+            Some("session-1"),
+            120,
+            Some(80),
+        )
+        .await;
+
+        let runtime = read_admin_provider_pool_runtime_state(
+            runtime,
+            "provider-1",
+            &key_ids,
+            &pool_config,
+            Some("session-1"),
+        )
+        .await;
+
+        assert_eq!(runtime.total_sticky_sessions, 0);
+        assert_eq!(runtime.sticky_bound_key_id, None);
+        assert_eq!(runtime.sticky_sessions_by_key.get("key-1"), None);
+        assert_eq!(runtime.cost_window_usage_by_key.get("key-1"), Some(&120));
+        assert_eq!(runtime.latency_avg_ms_by_key.get("key-1"), Some(&80.0));
+        assert!(runtime.lru_score_by_key.contains_key("key-1"));
+    }
+
     #[tokio::test]
     async fn error_feedback_respects_retry_after_for_rate_limits() {
         let Some(redis) = start_managed_redis_or_skip().await else {

apps/aether-gateway/src/handlers/proxy/mod.rs

@@ -54,6 +54,7 @@ use crate::headers::{
     should_skip_request_header,
 };
 use crate::router::RequestAdmissionError;
+use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
 use crate::{
     AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
     GatewayFallbackReason, LocalExecutionRuntimeMissDiagnostic,
@@ -297,6 +298,20 @@ async fn maybe_forward_public_request_to_tunnel_owner(
     }) else {
         return Ok(None);
     };
+    let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
+        Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
+        Err(err) => {
+            warn!(
+                trace_id = %request_context.trace_id,
+                error = ?err,
+                "gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
+            );
+            SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
+        }
+    };
+    if !cache_affinity_enabled {
+        return Ok(None);
+    }
     let Some(api_format) = decision
         .auth_endpoint_signature
         .as_deref()

apps/aether-gateway/src/handlers/shared/provider_pool.rs

@@ -1,4 +1,6 @@
-pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
+pub(crate) use super::super::admin::provider::pool::config::{
+    admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
+};
 pub(crate) use super::super::admin::provider::pool::runtime::{
     admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
     record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,

apps/aether-gateway/src/orchestration/effects.rs

@@ -31,6 +31,7 @@ use crate::handlers::shared::provider_pool::{
 };
 use crate::orchestration::local_execution_candidate_metadata_from_report_context;
 use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
+use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
 use crate::AppState;
 
 #[derive(Debug, Clone, Copy)]
@@ -238,10 +239,28 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
     })
 }
 
-fn remember_successful_local_scheduler_affinity(
+async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
+    match read_scheduler_ordering_config(state).await {
+        Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
+        Err(error) => {
+            warn!(
+                event_name = "orchestration_scheduler_affinity_config_load_failed",
+                log_type = "event",
+                error = ?error,
+                "failed to load scheduler config while checking cache affinity mode"
+            );
+            SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
+        }
+    }
+}
+
+async fn remember_successful_local_scheduler_affinity(
     state: &AppState,
     context: LocalExecutionEffectContext<'_>,
 ) {
+    if !scheduler_cache_affinity_enabled(state).await {
+        return;
+    }
     let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) else {
         return;
     };
@@ -516,7 +535,7 @@ async fn record_health_success_effect(
     context: LocalExecutionEffectContext<'_>,
     _effect: LocalHealthSuccessEffect,
 ) {
-    remember_successful_local_scheduler_affinity(state, context);
+    remember_successful_local_scheduler_affinity(state, context).await;
 
     let api_format = context.plan.provider_api_format.trim();
     if api_format.is_empty() {
@@ -1362,6 +1381,41 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn load_balance_success_does_not_remember_scheduler_affinity_cache() {
+        let state = AppState::new()
+            .expect("gateway state should build")
+            .with_data_state_for_tests(
+                GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
+                    "scheduling_mode".to_string(),
+                    json!("load_balance"),
+                )]),
+            );
+        let plan = sample_plan();
+        let report_context = json!({
+            "api_key_id": "api-key-1",
+            "client_api_format": "openai:chat",
+            "model": "gpt-5",
+        });
+        let cache_key =
+            build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
+                .expect("scheduler affinity cache key should build");
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
+        )
+        .await;
+
+        assert!(state
+            .read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
+            .is_none());
+    }
+
     #[tokio::test]
     async fn success_remembers_session_scoped_scheduler_affinity_cache() {
         let state = AppState::new().expect("gateway state should build");

apps/aether-gateway/src/scheduler/candidate/selection.rs

@@ -47,34 +47,50 @@ pub(super) async fn select_minimal_candidate(
     enable_model_directives: bool,
 ) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
     let affinity_epoch = runtime_state.scheduler_affinity_epoch();
+    let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
     let affinity_cache_key = build_scheduler_affinity_cache_key(
         auth_snapshot,
         api_format,
         global_model_name,
         client_session_affinity,
     );
-    let selected = collect_selectable_candidates(
+    let priority_affinity_key =
+        scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
+    let candidates = enumerate_scheduler_candidates(
         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,
     )
+    .await?;
+    let selected = collect_selectable_enumerated_candidates_with_skip_reasons(
+        runtime_state,
+        api_format,
+        global_model_name,
+        candidates,
+        required_capabilities,
+        auth_snapshot,
+        client_session_affinity,
+        now_unix_secs,
+        ordering_config,
+        priority_affinity_key,
+    )
     .await?
+    .0
     .into_iter()
     .next();
-    if let Some(candidate) = selected.as_ref() {
+    if ordering_config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity {
-        remember_scheduler_affinity(
+        if let Some(candidate) = selected.as_ref() {
-            affinity_cache_key.as_deref(),
+            remember_scheduler_affinity(
-            runtime_state,
+                affinity_cache_key.as_deref(),
-            candidate,
+                runtime_state,
-            Some(affinity_epoch),
+                candidate,
-        );
+                Some(affinity_epoch),
+            );
+        }
     }
     Ok(selected)
 }

apps/aether-gateway/src/scheduler/candidate/tests/selection.rs

@@ -22,6 +22,7 @@ use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
 use crate::data::GatewayDataState;
 use crate::{AppState, GatewayError};
 
+use super::super::affinity::build_scheduler_affinity_cache_key;
 use super::super::runtime::should_skip_provider_quota;
 use super::super::selection::{
     collect_selectable_candidates as collect_selectable_candidates_impl,
@@ -605,6 +606,46 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
     assert_eq!(selected.key_id, "key-b");
 }
 
+#[tokio::test]
+async fn load_balance_selection_does_not_remember_scheduler_affinity() {
+    let row = sample_row();
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        row,
+    ]));
+    let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
+    let state = AppState::new()
+        .expect("state should build")
+        .with_data_state_for_tests(
+            GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
+                .with_system_config_values_for_tests(vec![(
+                    "scheduling_mode".to_string(),
+                    json!("load_balance"),
+                )]),
+        );
+    let auth_snapshot = sample_auth_snapshot("affinity-key-1");
+    let cache_key =
+        build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
+            .expect("scheduler affinity cache key should build");
+
+    let selected = select_candidate(
+        state.data.as_ref(),
+        &state,
+        "openai:chat",
+        "gpt-4.1",
+        false,
+        Some(&auth_snapshot),
+        100,
+    )
+    .await
+    .expect("selection should succeed")
+    .expect("candidate should exist");
+
+    assert_eq!(selected.key_id, "key-1");
+    assert!(state
+        .read_scheduler_affinity_target(cache_key.as_str(), Duration::from_secs(300))
+        .is_none());
+}
+
 #[tokio::test]
 async fn load_balance_ignores_provider_priority_and_cached_affinity() {
     let mut first = sample_row();

crates/aether-ai-serving/src/candidate_ranking.rs

@@ -120,13 +120,18 @@ pub async fn run_ai_candidate_ranking<Port>(
 where
     Port: AiCandidateRankingPort,
 {
-    let affinity_requested_model = port.affinity_requested_model(&candidates);
+    let ranking_context = port.ranking_context();
-    let cached_affinity_target = port
+    let cached_affinity_target =
-        .read_cached_affinity_target(
+        if ranking_context.ranking_mode == SchedulerRankingMode::CacheAffinity {
-            normalized_client_api_format,
+            let affinity_requested_model = port.affinity_requested_model(&candidates);
-            affinity_requested_model.as_deref(),
+            port.read_cached_affinity_target(
-        )
+                normalized_client_api_format,
-        .await?;
+                affinity_requested_model.as_deref(),
+            )
+            .await?
+        } else {
+            None
+        };
 
     let mut rankables = Vec::with_capacity(candidates.len());
     for (original_index, candidate) in candidates.iter().enumerate() {
@@ -144,8 +149,7 @@ where
         );
     }
 
-    let outcomes =
+    let outcomes = apply_scheduler_candidate_ranking(&mut candidates, &rankables, ranking_context);
-        apply_scheduler_candidate_ranking(&mut candidates, &rankables, port.ranking_context());
     for outcome in outcomes {
         let ranking_index = outcome.ranking_index;
         if let Some(candidate) = candidates.get_mut(ranking_index) {
@@ -172,6 +176,7 @@ mod tests {
 
     #[derive(Default)]
     struct TestPort {
+        ranking_mode: SchedulerRankingMode,
         calls: Mutex<Vec<String>>,
     }
 
@@ -239,7 +244,7 @@ mod tests {
         fn ranking_context(&self) -> SchedulerRankingContext {
             SchedulerRankingContext {
                 priority_mode: SchedulerPriorityMode::Provider,
-                ranking_mode: SchedulerRankingMode::CacheAffinity,
+                ranking_mode: self.ranking_mode,
                 include_health: false,
                 load_balance_seed: 0,
             }
@@ -291,4 +296,29 @@ mod tests {
             ]
         );
     }
+
+    #[tokio::test]
+    async fn non_cache_affinity_ranking_does_not_read_affinity_target() {
+        let port = TestPort {
+            ranking_mode: SchedulerRankingMode::LoadBalance,
+            calls: Mutex::new(Vec::new()),
+        };
+        let candidates = vec![TestCandidate {
+            id: "candidate-a",
+            priority: 10,
+            ranking_index: None,
+            cached_affinity: false,
+        }];
+
+        let ranked = run_ai_candidate_ranking(&port, candidates, "openai:chat")
+            .await
+            .unwrap();
+
+        assert_eq!(ranked[0].id, "candidate-a");
+        assert!(!ranked[0].cached_affinity);
+        assert_eq!(
+            port.calls.lock().unwrap().as_slice(),
+            ["rankable:candidate-a:false"]
+        );
+    }
 }

crates/aether-ai-serving/src/pool_scheduler.rs

@@ -275,15 +275,19 @@ fn schedule_pool_group<Candidate>(
         };
     }
 
-    let sticky_candidate = runtime
+    let sticky_candidate = if pool_sticky_enabled(&active_presets) {
-        .sticky_bound_key_id
+        runtime
-        .as_ref()
+            .sticky_bound_key_id
-        .and_then(|sticky_key_id| {
+            .as_ref()
-            available
+            .and_then(|sticky_key_id| {
-                .iter()
+                available
-                .position(|item| item.item.facts.key_id == *sticky_key_id)
+                    .iter()
-        })
+                    .position(|item| item.item.facts.key_id == *sticky_key_id)
-        .map(|index| available.remove(index));
+            })
+            .map(|index| available.remove(index))
+    } else {
+        None
+    };
 
     if !active_presets.is_empty() {
         let sort_vectors = build_pool_sort_vectors(
@@ -402,6 +406,12 @@ fn build_pool_sort_vectors<Candidate>(
     vectors
 }
 
+fn pool_sticky_enabled(presets: &[NormalizedPoolPreset]) -> bool {
+    presets
+        .iter()
+        .any(|preset| preset.preset == "cache_affinity")
+}
+
 fn lru_rank_indices<Candidate>(
     items: &[PoolGroupCandidateOrdering<Candidate>],
     descending: bool,
@@ -916,8 +926,18 @@ mod tests {
 
     #[test]
     fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
-        let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true);
+        let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true)
-        let key_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true);
+            .with_presets(vec![AiPoolSchedulingPreset {
+                preset: "cache_affinity".to_string(),
+                enabled: true,
+                mode: None,
+            }]);
+        let key_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true)
+            .with_presets(vec![AiPoolSchedulingPreset {
+                preset: "cache_affinity".to_string(),
+                enabled: true,
+                mode: None,
+            }]);
 
         let runtime_by_provider = BTreeMap::from([(
             "provider-pool".to_string(),
@@ -944,6 +964,49 @@ mod tests {
         );
     }
 
+    #[test]
+    fn load_balance_distribution_ignores_sticky_hit() {
+        let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true)
+            .with_presets(vec![AiPoolSchedulingPreset {
+                preset: "load_balance".to_string(),
+                enabled: true,
+                mode: None,
+            }]);
+        let key_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true)
+            .with_presets(vec![AiPoolSchedulingPreset {
+                preset: "load_balance".to_string(),
+                enabled: true,
+                mode: None,
+            }]);
+        let nonce = (0..1000)
+            .map(|index| format!("seed-{index}"))
+            .find(|nonce| {
+                let group_seed = format!("codex:provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
+                stable_hash_score(format!("{group_seed}:key-b").as_str())
+                    < stable_hash_score(format!("{group_seed}:key-a").as_str())
+            })
+            .expect("test seed should exist");
+        let runtime_by_provider = BTreeMap::from([(
+            "provider-pool".to_string(),
+            AiPoolRuntimeState {
+                sticky_bound_key_id: Some("key-a".to_string()),
+                ..AiPoolRuntimeState::default()
+            },
+        )]);
+
+        let outcome = run_ai_pool_scheduler(vec![key_a, key_b], &runtime_by_provider, &nonce);
+
+        assert!(outcome.skipped_candidates.is_empty());
+        assert_eq!(
+            outcome
+                .candidates
+                .iter()
+                .map(|item| item.candidate.as_str())
+                .collect::<Vec<_>>(),
+            vec!["key-b", "key-a"]
+        );
+    }
+
     #[test]
     fn pool_scheduler_uses_plan_preset_with_catalog_context() {
         let key_free = sample_candidate("provider-pool", "endpoint-1", "key-free", 10, true)