Restrict scheduler affinity to cache affinity mode

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 cached_affinity_target = port
-        .read_cached_affinity_target(
-            normalized_client_api_format,
-            affinity_requested_model.as_deref(),
-        )
-        .await?;
+    let ranking_context = port.ranking_context();
+    let cached_affinity_target =
+        if ranking_context.ranking_mode == SchedulerRankingMode::CacheAffinity {
+            let affinity_requested_model = port.affinity_requested_model(&candidates);
+            port.read_cached_affinity_target(
+                normalized_client_api_format,
+                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 =
-        apply_scheduler_candidate_ranking(&mut candidates, &rankables, port.ranking_context());
+    let outcomes = apply_scheduler_candidate_ranking(&mut candidates, &rankables, 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
-        .sticky_bound_key_id
-        .as_ref()
-        .and_then(|sticky_key_id| {
-            available
-                .iter()
-                .position(|item| item.item.facts.key_id == *sticky_key_id)
-        })
-        .map(|index| available.remove(index));
+    let sticky_candidate = if pool_sticky_enabled(&active_presets) {
+        runtime
+            .sticky_bound_key_id
+            .as_ref()
+            .and_then(|sticky_key_id| {
+                available
+                    .iter()
+                    .position(|item| item.item.facts.key_id == *sticky_key_id)
+            })
+            .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_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true);
+        let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 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)