fix: carry session affinity through Gemini files

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

@@ -28,6 +28,7 @@ pub(crate) struct LocalAuthenticatedDecisionInput {
     pub(crate) auth_context: ExecutionRuntimeAuthContext,
     pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
     pub(crate) required_capabilities: Option<serde_json::Value>,
+    pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
 }
 
 struct GatewayAuthenticatedDecisionInputPort<'a> {
@@ -108,6 +109,7 @@ pub(crate) fn build_local_authenticated_decision_input(
         auth_context: resolved_input.auth_context,
         auth_snapshot: resolved_input.auth_snapshot,
         required_capabilities: resolved_input.required_capabilities,
+        client_session_affinity: None,
     }
 }
 

apps/aether-gateway/src/ai_serving/planner/specialized/files.rs

@@ -103,7 +103,14 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
         return Ok(None);
     };
 
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) = resolve_local_gemini_files_decision_input(
+        state,
+        parts,
+        Some(body_json),
+        trace_id,
+        decision,
+    )
+    .await
     else {
         return Ok(None);
     };
@@ -140,7 +147,8 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
         return Ok(None);
     };
 
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) =
+        resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
     else {
         return Ok(None);
     };
@@ -298,7 +306,14 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
         return Ok(None);
     };
 
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) = resolve_local_gemini_files_decision_input(
+        state,
+        parts,
+        Some(body_json),
+        trace_id,
+        decision,
+    )
+    .await
     else {
         return Ok(None);
     };
@@ -338,7 +353,8 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
         return Ok(None);
     };
 
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) =
+        resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
     else {
         return Ok(None);
     };
@@ -379,7 +395,14 @@ async fn build_local_sync_plan_and_reports(
     spec: LocalGeminiFilesSpec,
 ) -> Result<Vec<AiSyncAttempt>, GatewayError> {
     let spec_metadata = local_gemini_files_spec_metadata(spec);
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) = resolve_local_gemini_files_decision_input(
+        state,
+        parts,
+        Some(body_json),
+        trace_id,
+        decision,
+    )
+    .await
     else {
         return Ok(Vec::new());
     };
@@ -430,7 +453,8 @@ async fn build_local_stream_plan_and_reports(
     spec: LocalGeminiFilesSpec,
 ) -> Result<Vec<AiStreamAttempt>, GatewayError> {
     let spec_metadata = local_gemini_files_spec_metadata(spec);
-    let Some(input) = resolve_local_gemini_files_decision_input(state, trace_id, decision).await
+    let Some(input) =
+        resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
     else {
         return Ok(Vec::new());
     };

apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs

@@ -23,6 +23,7 @@ use crate::ai_serving::{
     resolve_local_decision_execution_runtime_auth_context, CandidateFailureDiagnostic,
     ExecutionRuntimeAuthContext, GatewayControlDecision, PlannerAppState,
 };
+use crate::client_session_affinity::client_session_affinity_from_parts;
 use crate::clock::current_unix_secs;
 use crate::{AppState, GatewayError};
 
@@ -36,6 +37,8 @@ pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
 
 pub(super) async fn resolve_local_gemini_files_decision_input(
     state: &AppState,
+    parts: &http::request::Parts,
+    body_json: Option<&serde_json::Value>,
     trace_id: &str,
     decision: &GatewayControlDecision,
 ) -> Option<LocalGeminiFilesDecisionInput> {
@@ -64,7 +67,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
         }
     };
 
-    Some(build_local_authenticated_decision_input(resolved_input))
+    let mut input = build_local_authenticated_decision_input(resolved_input);
+    input.client_session_affinity = client_session_affinity_from_parts(parts, body_json);
+    Some(input)
 }
 
 pub(super) async fn materialize_local_gemini_files_candidate_attempts(
@@ -84,6 +89,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
             GEMINI_FILES_REQUIRED_CAPABILITY,
             false,
             Some(&input.auth_snapshot),
+            input.client_session_affinity.as_ref(),
             current_unix_secs(),
         )
         .await?;
@@ -93,7 +99,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
         GEMINI_FILES_CLIENT_API_FORMAT,
         None,
         Some(&input.auth_snapshot),
-        None,
+        input.client_session_affinity.as_ref(),
         input.required_capabilities.as_ref(),
         None,
         None,
@@ -155,6 +161,7 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
             GEMINI_FILES_REQUIRED_CAPABILITY,
             false,
             Some(&input.auth_snapshot),
+            input.client_session_affinity.as_ref(),
             current_unix_secs(),
         )
         .await?;
@@ -164,7 +171,7 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
         GEMINI_FILES_CLIENT_API_FORMAT,
         None,
         Some(&input.auth_snapshot),
-        None,
+        input.client_session_affinity.as_ref(),
         input.required_capabilities.as_ref(),
         None,
         None,

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

@@ -138,6 +138,7 @@ impl<'a> PlannerAppState<'a> {
         required_capability: &str,
         require_streaming: bool,
         auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        client_session_affinity: Option<&ClientSessionAffinity>,
         now_unix_secs: u64,
     ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
         let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
@@ -153,6 +154,7 @@ impl<'a> PlannerAppState<'a> {
                 required_capability,
                 require_streaming,
                 auth_snapshot,
+                client_session_affinity,
                 attempt_now_unix_secs,
             )
             .await?;