feat: scope planner affinity cache by session

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

@@ -1,6 +1,6 @@
 use aether_scheduler_core::{
-    build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
+    build_scheduler_affinity_cache_key_for_api_key_id_with_client_session, ClientSessionAffinity,
-    SchedulerMinimalCandidateSelectionCandidate,
+    SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
 };
 
 use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
@@ -11,6 +11,7 @@ const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
 pub(crate) fn read_cached_scheduler_affinity_target(
     state: PlannerAppState<'_>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     client_api_format: &str,
     requested_model: Option<&str>,
 ) -> Option<SchedulerAffinityTarget> {
@@ -20,10 +21,11 @@ pub(crate) fn read_cached_scheduler_affinity_target(
     let api_key_id = auth_snapshot
         .map(|snapshot| snapshot.api_key_id.trim())
         .filter(|value| !value.is_empty())?;
-    let cache_key = build_scheduler_affinity_cache_key_for_api_key_id(
+    let cache_key = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
         api_key_id,
         client_api_format,
         requested_model,
+        client_session_affinity,
     )?;
 
     state
@@ -34,6 +36,7 @@ pub(crate) fn read_cached_scheduler_affinity_target(
 pub(crate) fn remember_scheduler_affinity_for_candidate(
     state: PlannerAppState<'_>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     client_api_format: &str,
     requested_model: &str,
     candidate: &SchedulerMinimalCandidateSelectionCandidate,
@@ -44,10 +47,11 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
     else {
         return;
     };
-    let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id(
+    let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
         api_key_id,
         client_api_format,
         requested_model,
+        client_session_affinity,
     ) else {
         return;
     };

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

@@ -14,8 +14,9 @@ use crate::scheduler::config::{
     read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
 };
 use aether_scheduler_core::{
-    matches_affinity_target, SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
+    matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
-    SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingOutcome,
+    SchedulerMinimalCandidateSelectionCandidate, SchedulerRankableCandidate,
+    SchedulerRankingContext, SchedulerRankingOutcome,
 };
 
 use super::candidate_affinity_cache::read_cached_scheduler_affinity_target;
@@ -28,6 +29,7 @@ struct GatewayLocalCandidateRankingPort<'a> {
     state: PlannerAppState<'a>,
     requested_model: Option<&'a str>,
     auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&'a ClientSessionAffinity>,
     required_capabilities: Option<&'a serde_json::Value>,
     ordering_config: SchedulerOrderingConfig,
 }
@@ -58,6 +60,7 @@ impl AiCandidateRankingPort for GatewayLocalCandidateRankingPort<'_> {
         Ok(read_cached_scheduler_affinity_target(
             self.state,
             self.auth_snapshot,
+            self.client_session_affinity,
             normalized_client_api_format,
             affinity_requested_model,
         ))
@@ -116,6 +119,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
     normalized_client_api_format: &str,
     requested_model: Option<&str>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     required_capabilities: Option<&serde_json::Value>,
 ) -> Vec<EligibleLocalExecutionCandidate> {
     let ordering_config = read_scheduler_ordering_config_or_default(state).await;
@@ -123,6 +127,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
         state,
         requested_model,
         auth_snapshot,
+        client_session_affinity,
         required_capabilities,
         ordering_config,
     };
@@ -212,7 +217,9 @@ mod tests {
         StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
     };
     use aether_scheduler_core::{
-        apply_scheduler_candidate_ranking, RANKING_REASON_CACHED_AFFINITY,
+        apply_scheduler_candidate_ranking,
+        build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
+        ClientSessionAffinity, RANKING_REASON_CACHED_AFFINITY,
     };
     use serde_json::json;
 
@@ -1052,6 +1059,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1130,6 +1138,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1204,6 +1213,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1269,6 +1279,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1350,6 +1361,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1398,6 +1410,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "openai:chat",
             "gpt-4.1",
             &cached_candidate,
@@ -1422,6 +1435,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1498,6 +1512,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1567,6 +1582,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "claude:messages",
             "gpt-4.1",
             &cached_cross_format,
@@ -1591,6 +1607,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1660,6 +1677,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "openai:chat",
             "gpt-4.1",
             &cached_candidate,
@@ -1692,6 +1710,7 @@ mod tests {
             None,
             None,
             None,
+            None,
         )
         .await;
 
@@ -1762,6 +1781,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "openai:chat",
             "gpt-4.1",
             &cached_candidate,
@@ -1786,6 +1806,7 @@ mod tests {
             None,
             None,
             None,
+            None,
             aether_ai_serving::AiCandidateResolutionMode::Standard,
         )
         .await;
@@ -1852,6 +1873,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "openai:chat",
             "gpt-4.1",
             &cached_candidate,
@@ -1876,6 +1898,7 @@ mod tests {
             None,
             None,
             None,
+            None,
             aether_ai_serving::AiCandidateResolutionMode::Standard,
         )
         .await;
@@ -1902,6 +1925,7 @@ mod tests {
         remember_scheduler_affinity_for_candidate(
             PlannerAppState::new(&state),
             Some(&auth_snapshot),
+            None,
             "openai:chat",
             "gpt-5",
             &candidate,
@@ -1917,4 +1941,44 @@ mod tests {
         assert_eq!(remembered.endpoint_id, "endpoint-1");
         assert_eq!(remembered.key_id, "key-1");
     }
+
+    #[tokio::test]
+    async fn remembers_scheduler_affinity_for_client_session_scope() {
+        let state = AppState::new().expect("state should build");
+        let auth_snapshot = sample_auth_snapshot();
+        let client_session_affinity = ClientSessionAffinity::new(
+            Some("generic".to_string()),
+            Some("session=conversation-1;agent=coder".to_string()),
+        );
+        let candidate = sample_candidate("endpoint-session", "key-session");
+
+        remember_scheduler_affinity_for_candidate(
+            PlannerAppState::new(&state),
+            Some(&auth_snapshot),
+            Some(&client_session_affinity),
+            "openai:chat",
+            "gpt-5",
+            &candidate,
+        );
+
+        let session_key = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
+            "api-key-1",
+            "openai:chat",
+            "gpt-5",
+            Some(&client_session_affinity),
+        )
+        .expect("session key should build");
+        let remembered = state
+            .read_scheduler_affinity_target(&session_key, SCHEDULER_AFFINITY_TTL)
+            .expect("session affinity target should be cached");
+        assert_eq!(remembered.provider_id, "provider-1");
+        assert_eq!(remembered.endpoint_id, "endpoint-session");
+        assert_eq!(remembered.key_id, "key-session");
+        assert!(state
+            .read_scheduler_affinity_target(
+                "scheduler_affinity:api-key-1:openai:chat:gpt-5",
+                SCHEDULER_AFFINITY_TTL,
+            )
+            .is_none());
+    }
 }