feat: thread session affinity through scheduler selection

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

@@ -24,7 +24,7 @@ use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
 use aether_scheduler_core::{
     candidate_model_names, candidate_supports_required_capability, matches_model_mapping,
     normalize_api_format, resolve_provider_model_name, select_provider_model_name,
-    SchedulerMinimalCandidateSelectionCandidate,
+    ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate,
 };
 use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
 use regex::Regex;
@@ -57,6 +57,7 @@ pub(crate) async fn list_selectable_candidates(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -68,6 +69,7 @@ pub(crate) async fn list_selectable_candidates(
         require_streaming,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         enable_model_directives,
     )
@@ -89,6 +91,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<
@@ -106,6 +109,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
         require_streaming,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         enable_model_directives,
     )
@@ -120,6 +124,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
     candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
 ) -> Result<
     (
@@ -138,6 +143,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
         candidates,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         ordering_config,
         priority_affinity_key,
@@ -218,6 +224,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
             require_streaming,
             required_capabilities.as_ref(),
             auth_snapshot,
+            None,
             now_unix_secs,
             false,
         )

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

@@ -3,6 +3,7 @@ use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
 use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
 use crate::scheduler::config::SchedulerSchedulingMode;
 use crate::GatewayError;
+use aether_scheduler_core::ClientSessionAffinity;
 
 use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
 use super::enumeration::enumerate_scheduler_candidates;
@@ -41,11 +42,16 @@ pub(super) async fn select_minimal_candidate(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
-    let affinity_cache_key =
+    let affinity_cache_key = build_scheduler_affinity_cache_key(
-        build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
+        auth_snapshot,
+        api_format,
+        global_model_name,
+        client_session_affinity,
+    );
     let selected = collect_selectable_candidates(
         selection_row_source,
         runtime_state,
@@ -54,6 +60,7 @@ pub(super) async fn select_minimal_candidate(
         require_streaming,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         enable_model_directives,
     )
@@ -74,6 +81,7 @@ pub(super) async fn collect_selectable_candidates(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -85,6 +93,7 @@ pub(super) async fn collect_selectable_candidates(
         require_streaming,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         enable_model_directives,
     )
@@ -100,6 +109,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<
@@ -129,6 +139,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
         candidates,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         now_unix_secs,
         ordering_config,
         priority_affinity_key,
@@ -144,6 +155,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
     mut candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
     ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
     priority_affinity_key: Option<&str>,
@@ -157,8 +169,12 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
     let runtime_snapshot =
         read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
             .await?;
-    let affinity_cache_key =
+    let affinity_cache_key = build_scheduler_affinity_cache_key(
-        build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
+        auth_snapshot,
+        api_format,
+        global_model_name,
+        client_session_affinity,
+    );
     let cached_affinity_target = if ordering_config.scheduling_mode
         == SchedulerSchedulingMode::CacheAffinity
     {
@@ -191,12 +207,8 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
         ));
     }
 
-    let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
+    let (mut selected, skipped) =
-        candidates,
+        resolve_scheduler_candidate_selectability(candidates, &runtime_snapshot, now_unix_secs);
-        &runtime_snapshot,
-        now_unix_secs,
-        cached_affinity_target.as_ref(),
-    );
     rank_scheduler_candidates(
         &mut selected,
         &runtime_snapshot,

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

@@ -47,6 +47,7 @@ async fn select_candidate(
         require_streaming,
         None,
         auth_snapshot,
+        None,
         now_unix_secs,
         false,
     )
@@ -70,6 +71,7 @@ async fn collect_selectable_candidates(
         require_streaming,
         None,
         auth_snapshot,
+        None,
         now_unix_secs,
         false,
     )
@@ -99,6 +101,7 @@ async fn collect_selectable_candidates_with_skip_reasons(
         require_streaming,
         None,
         auth_snapshot,
+        None,
         now_unix_secs,
         false,
     )
@@ -406,6 +409,7 @@ async fn scheduler_selection_prefers_required_capability_matches_before_priority
         false,
         Some(&required_capabilities),
         None,
+        None,
         100,
         false,
     )