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?;

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

@@ -158,6 +158,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
     required_capability: &str,
     require_streaming: bool,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
 ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
     Ok(
@@ -168,6 +169,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
             required_capability,
             require_streaming,
             auth_snapshot,
+            client_session_affinity,
             now_unix_secs,
         )
         .await?
@@ -182,6 +184,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
     required_capability: &str,
     require_streaming: bool,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    client_session_affinity: Option<&ClientSessionAffinity>,
     now_unix_secs: u64,
 ) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
     let normalized_api_format = normalize_api_format(candidate_api_format);
@@ -224,7 +227,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
             require_streaming,
             required_capabilities.as_ref(),
             auth_snapshot,
-            None,
+            client_session_affinity,
             now_unix_secs,
             false,
         )

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

@@ -7,10 +7,14 @@ use aether_data::repository::quota::InMemoryProviderQuotaRepository;
 use aether_data_contracts::repository::candidates::{
     RequestCandidateStatus, StoredRequestCandidate,
 };
+use aether_scheduler_core::ClientSessionAffinity;
 
+use crate::cache::SchedulerAffinityTarget;
 use crate::data::GatewayDataState;
+use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
 use crate::AppState;
 
+use super::super::affinity::build_scheduler_affinity_cache_key;
 use super::super::{
     list_selectable_candidates_for_required_capability_without_requested_model,
     list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal,
@@ -59,6 +63,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
         "cache_1h",
         false,
         None,
+        None,
         100,
     )
     .await
@@ -113,6 +118,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
         "gemini_files",
         false,
         None,
+        None,
         100,
     )
     .await
@@ -123,6 +129,82 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
     assert_eq!(selection[0].key_id, "key-b");
 }
 
+#[tokio::test]
+async fn required_capability_without_model_uses_session_scoped_affinity() {
+    let mut fallback = sample_row();
+    fallback.provider_id = "provider-a".to_string();
+    fallback.provider_name = "provider-a".to_string();
+    fallback.endpoint_id = "endpoint-a".to_string();
+    fallback.endpoint_api_format = "gemini:generate_content".to_string();
+    fallback.endpoint_api_family = Some("gemini".to_string());
+    fallback.key_api_formats = Some(vec!["gemini:generate_content".to_string()]);
+    fallback.key_id = "key-a".to_string();
+    fallback.key_name = "alpha".to_string();
+    fallback.global_model_name = "gemini-2.5-pro".to_string();
+    fallback.key_capabilities = Some(serde_json::json!({"gemini_files": true}));
+    fallback.key_global_priority_by_format =
+        Some(serde_json::json!({"gemini:generate_content": 0}));
+
+    let mut session_target = fallback.clone();
+    session_target.provider_id = "provider-b".to_string();
+    session_target.provider_name = "provider-b".to_string();
+    session_target.endpoint_id = "endpoint-b".to_string();
+    session_target.key_id = "key-b".to_string();
+    session_target.key_name = "beta".to_string();
+    session_target.key_global_priority_by_format =
+        Some(serde_json::json!({"gemini:generate_content": 10}));
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        fallback,
+        session_target,
+    ]));
+    let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
+    let state = AppState::new()
+        .expect("state should build")
+        .with_data_state_for_tests(
+            GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
+        );
+    let auth_snapshot = sample_auth_snapshot("api-key-1");
+    let client_session_affinity = ClientSessionAffinity::new(
+        Some("generic".to_string()),
+        Some("session=conversation-1".to_string()),
+    );
+    let cache_key = build_scheduler_affinity_cache_key(
+        Some(&auth_snapshot),
+        "gemini:generate_content",
+        "gemini-2.5-pro",
+        Some(&client_session_affinity),
+    )
+    .expect("session affinity cache key should build");
+    state.scheduler_affinity_cache.insert(
+        cache_key,
+        SchedulerAffinityTarget {
+            provider_id: "provider-b".to_string(),
+            endpoint_id: "endpoint-b".to_string(),
+            key_id: "key-b".to_string(),
+        },
+        SCHEDULER_AFFINITY_TTL,
+        16,
+    );
+
+    let selection = list_selectable_candidates_for_required_capability_without_requested_model(
+        state.data.as_ref(),
+        &state,
+        "gemini:generate_content",
+        "gemini_files",
+        false,
+        Some(&auth_snapshot),
+        Some(&client_session_affinity),
+        100,
+    )
+    .await
+    .expect("selection should succeed");
+
+    assert_eq!(selection.len(), 2);
+    assert_eq!(selection[0].provider_id, "provider-b");
+    assert_eq!(selection[0].key_id, "key-b");
+}
+
 #[tokio::test]
 async fn required_capability_reports_auth_limit_signal_when_every_model_is_blocked_by_api_key_concurrency(
 ) {
@@ -189,6 +271,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
             "cache_1h",
             false,
             Some(&auth_snapshot),
+            None,
             100,
         )
         .await