feat: pass session affinity through planner inputs

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

@@ -3,7 +3,7 @@ use aether_ai_serving::{
 };
 use aether_scheduler_core::{
     enumerate_minimal_candidate_selection_with_model_directives, normalize_api_format,
-    resolve_requested_global_model_name_with_model_directives,
+    resolve_requested_global_model_name_with_model_directives, ClientSessionAffinity,
     EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
 };
 use async_trait::async_trait;
@@ -32,6 +32,7 @@ struct GatewayLocalCandidatePreselectionPort<'a> {
     require_streaming: bool,
     required_capabilities: Option<&'a serde_json::Value>,
     auth_snapshot: &'a GatewayAuthApiKeySnapshot,
+    client_session_affinity: Option<&'a ClientSessionAffinity>,
     use_api_format_alias_match: bool,
     key_mode: LocalCandidatePreselectionKeyMode,
     candidate_api_formats: Vec<String>,
@@ -73,6 +74,7 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
                 self.require_streaming,
                 self.required_capabilities,
                 auth_snapshot,
+                self.client_session_affinity,
                 current_unix_secs(),
             )
             .await?;
@@ -139,6 +141,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
     require_streaming: bool,
     required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: &GatewayAuthApiKeySnapshot,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     use_api_format_alias_match: bool,
     key_mode: LocalCandidatePreselectionKeyMode,
 ) -> Result<
@@ -173,6 +176,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
         require_streaming,
         required_capabilities,
         auth_snapshot,
+        client_session_affinity,
         use_api_format_alias_match,
         key_mode,
         candidate_api_formats,
@@ -189,6 +193,7 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
     require_streaming: bool,
     required_capabilities: Option<serde_json::Value>,
     auth_snapshot: GatewayAuthApiKeySnapshot,
+    client_session_affinity: Option<ClientSessionAffinity>,
     use_api_format_alias_match: bool,
     key_mode: LocalCandidatePreselectionKeyMode,
     candidate_api_formats: Vec<String>,
@@ -210,6 +215,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
         require_streaming: bool,
         required_capabilities: Option<&serde_json::Value>,
         auth_snapshot: &GatewayAuthApiKeySnapshot,
+        client_session_affinity: Option<&ClientSessionAffinity>,
         use_api_format_alias_match: bool,
         key_mode: LocalCandidatePreselectionKeyMode,
     ) -> Self {
@@ -239,6 +245,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
             require_streaming,
             required_capabilities: required_capabilities.cloned(),
             auth_snapshot: auth_snapshot.clone(),
+            client_session_affinity: client_session_affinity.cloned(),
             use_api_format_alias_match,
             key_mode,
             candidate_api_formats,
@@ -449,6 +456,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
                     candidates,
                     self.required_capabilities.as_ref(),
                     auth_snapshot,
+                    self.client_session_affinity.as_ref(),
                     current_unix_secs(),
                 )
                 .await?;

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

@@ -1,4 +1,5 @@
 use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
+use aether_scheduler_core::ClientSessionAffinity;
 use async_trait::async_trait;
 
 use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState};
@@ -19,6 +20,7 @@ pub(crate) struct LocalRequestedModelDecisionInput {
     pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
     pub(crate) required_capabilities: Option<serde_json::Value>,
     pub(crate) request_auth_channel: Option<String>,
+    pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
 }
 
 #[derive(Debug, Clone)]
@@ -95,6 +97,7 @@ pub(crate) fn build_local_requested_model_decision_input(
         auth_snapshot: resolved_input.auth_snapshot,
         required_capabilities: resolved_input.required_capabilities,
         request_auth_channel: None,
+        client_session_affinity: None,
     }
 }
 

apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs

@@ -1,4 +1,4 @@
-use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
+use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
 use std::time::Duration;
 use tokio::time::Instant;
 
@@ -18,6 +18,7 @@ impl<'a> PlannerAppState<'a> {
         require_streaming: bool,
         required_capabilities: Option<&serde_json::Value>,
         auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        client_session_affinity: Option<&ClientSessionAffinity>,
         now_unix_secs: u64,
     ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
         let enable_model_directives =
@@ -35,6 +36,7 @@ impl<'a> PlannerAppState<'a> {
             require_streaming,
             required_capabilities,
             auth_snapshot,
+            client_session_affinity,
             now_unix_secs,
             enable_model_directives,
         )
@@ -48,6 +50,7 @@ impl<'a> PlannerAppState<'a> {
         require_streaming: bool,
         required_capabilities: Option<&serde_json::Value>,
         auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        client_session_affinity: Option<&ClientSessionAffinity>,
         now_unix_secs: u64,
     ) -> Result<
         (
@@ -77,6 +80,7 @@ impl<'a> PlannerAppState<'a> {
                 require_streaming,
                 required_capabilities,
                 auth_snapshot,
+                client_session_affinity,
                 attempt_now_unix_secs,
                 enable_model_directives,
             )
@@ -106,6 +110,7 @@ impl<'a> PlannerAppState<'a> {
         candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
         required_capabilities: Option<&serde_json::Value>,
         auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        client_session_affinity: Option<&ClientSessionAffinity>,
         now_unix_secs: u64,
     ) -> Result<
         (
@@ -121,6 +126,7 @@ impl<'a> PlannerAppState<'a> {
             candidates,
             required_capabilities,
             auth_snapshot,
+            client_session_affinity,
             now_unix_secs,
         )
         .await