Implement independent provider pool scheduling runtime

apps/aether-gateway/src/ai_pipeline/mod.rs

@@ -33,10 +33,11 @@ pub(crate) use self::planner::{
     build_openai_cli_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
     build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
     build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
-    maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
-    maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
-    set_local_openai_chat_execution_exhausted_diagnostic, GatewayAuthApiKeySnapshot,
-    GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
+    extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
+    maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
+    maybe_build_sync_plan_payload, set_local_openai_chat_execution_exhausted_diagnostic,
+    GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
+    PlannerAppState,
 };
 pub(crate) use self::pure::*;
 pub(crate) use crate::control::GatewayControlDecision;

apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs

@@ -1040,6 +1040,7 @@ mod tests {
             "openai:chat",
             "gpt-4.1",
             None,
+            None,
         )
         .await;
 
@@ -1115,6 +1116,7 @@ mod tests {
             "openai:chat",
             "gpt-4.1",
             None,
+            None,
         )
         .await;
 
@@ -1186,6 +1188,7 @@ mod tests {
             "openai:chat",
             "gpt-4.1",
             None,
+            None,
         )
         .await;
 
@@ -1244,6 +1247,7 @@ mod tests {
             "openai:chat",
             "gpt-4.1",
             None,
+            None,
         )
         .await;
 
@@ -1320,6 +1324,7 @@ mod tests {
             "openai:chat",
             "gpt-4.1",
             None,
+            None,
         )
         .await;
 

apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs

@@ -6,6 +6,7 @@ use std::collections::BTreeSet;
 use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
 
 use super::candidate_affinity::rank_eligible_local_execution_candidates;
+use super::pool_scheduler::apply_local_execution_pool_scheduler;
 
 #[derive(Debug, Clone, PartialEq)]
 pub(crate) struct EligibleLocalExecutionCandidate {
@@ -28,6 +29,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
     client_api_format: &str,
     requested_model: &str,
     required_capabilities: Option<&serde_json::Value>,
+    sticky_session_token: Option<&str>,
 ) -> (
     Vec<EligibleLocalExecutionCandidate>,
     Vec<SkippedLocalExecutionCandidate>,
@@ -37,6 +39,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates(
         candidates,
         client_api_format,
         required_capabilities,
+        sticky_session_token,
         |candidate, transport| {
             current_local_execution_candidate_skip_reason_with_transport(
                 candidate,
@@ -55,6 +58,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
     client_api_format: &str,
     requested_model: Option<&str>,
     required_capabilities: Option<&serde_json::Value>,
+    sticky_session_token: Option<&str>,
 ) -> (
     Vec<EligibleLocalExecutionCandidate>,
     Vec<SkippedLocalExecutionCandidate>,
@@ -64,6 +68,7 @@ pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport
         candidates,
         client_api_format,
         required_capabilities,
+        sticky_session_token,
         |candidate, transport| {
             current_local_execution_candidate_common_skip_reason_with_transport(
                 candidate,
@@ -80,6 +85,7 @@ async fn filter_and_rank_local_execution_candidates_with_gate<F>(
     candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
     client_api_format: &str,
     required_capabilities: Option<&serde_json::Value>,
+    sticky_session_token: Option<&str>,
     runtime_skip_reason: F,
 ) -> (
     Vec<EligibleLocalExecutionCandidate>,
@@ -149,10 +155,49 @@ where
         required_capabilities,
     )
     .await;
+    let (ranked, pool_skipped) =
+        apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
+    skipped.extend(pool_skipped);
 
     (ranked, skipped)
 }
 
+pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
+    fn non_empty_string(value: Option<&serde_json::Value>) -> Option<String> {
+        value
+            .and_then(serde_json::Value::as_str)
+            .map(str::trim)
+            .filter(|value| !value.is_empty())
+            .map(ToOwned::to_owned)
+    }
+
+    let object = body_json.as_object()?;
+
+    non_empty_string(object.get("prompt_cache_key"))
+        .or_else(|| non_empty_string(object.get("conversation_id")))
+        .or_else(|| non_empty_string(object.get("conversationId")))
+        .or_else(|| non_empty_string(object.get("session_id")))
+        .or_else(|| non_empty_string(object.get("sessionId")))
+        .or_else(|| {
+            object
+                .get("metadata")
+                .and_then(serde_json::Value::as_object)
+                .and_then(|metadata| {
+                    non_empty_string(metadata.get("session_id"))
+                        .or_else(|| non_empty_string(metadata.get("conversation_id")))
+                })
+        })
+        .or_else(|| {
+            object
+                .get("conversationState")
+                .and_then(serde_json::Value::as_object)
+                .and_then(|state| {
+                    non_empty_string(state.get("conversationId"))
+                        .or_else(|| non_empty_string(state.get("sessionId")))
+                })
+        })
+}
+
 fn current_local_execution_candidate_common_skip_reason_with_transport(
     candidate: &SchedulerMinimalCandidateSelectionCandidate,
     transport: &GatewayProviderTransportSnapshot,

apps/aether-gateway/src/ai_pipeline/planner/mod.rs

@@ -17,6 +17,7 @@ mod materialization_policy;
 mod passthrough;
 mod payload_metadata;
 mod plan_builders;
+mod pool_scheduler;
 mod report_context;
 mod route;
 mod runtime_miss;
@@ -25,6 +26,7 @@ mod specialized;
 mod standard;
 mod state;
 
+pub(crate) use self::candidate_eligibility::extract_pool_sticky_session_token;
 pub(crate) use self::passthrough::{
     build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
 };

apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/build.rs

@@ -55,9 +55,10 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
         Some(input.requested_model.as_str()),
         "candidate_evaluation_incomplete",
     );
-    let (attempts, candidate_count) =
-        materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
-            .await?;
+    let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
+        state, trace_id, &input, body_json, spec,
+    )
+    .await?;
     apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
         state,
         trace_id,
@@ -121,9 +122,10 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
         Some(input.requested_model.as_str()),
         "candidate_evaluation_incomplete",
     );
-    let (attempts, candidate_count) =
-        materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
-            .await?;
+    let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
+        state, trace_id, &input, body_json, spec,
+    )
+    .await?;
     apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
         state,
         trace_id,

apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/candidates.rs

@@ -1,7 +1,8 @@
 use tracing::warn;
 
 use crate::ai_pipeline::planner::candidate_eligibility::{
-    filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
+    extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
+    SkippedLocalExecutionCandidate,
 };
 use crate::ai_pipeline::planner::candidate_materialization::{
     persist_available_local_execution_candidates_with_context,
@@ -85,10 +86,12 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
     state: &AppState,
     trace_id: &str,
     input: &LocalSameFormatProviderDecisionInput,
+    body_json: &serde_json::Value,
     spec: LocalSameFormatProviderSpec,
 ) -> Result<(Vec<LocalSameFormatProviderCandidateAttempt>, usize), GatewayError> {
     let spec_metadata = local_same_format_provider_spec_metadata(spec);
     let planner_state = PlannerAppState::new(state);
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
     let persistence_policy = build_local_candidate_persistence_policy(
         &input.auth_context,
         input.required_capabilities.as_ref(),
@@ -110,6 +113,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
         spec_metadata.api_format,
         &input.requested_model,
         input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
     )
     .await;
     let skipped_candidates = preselection_skipped

apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/plans.rs

@@ -57,9 +57,10 @@ pub(crate) async fn build_local_sync_plan_and_reports(
         Some(input.requested_model.as_str()),
         "candidate_evaluation_incomplete",
     );
-    let (attempts, candidate_count) =
-        materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
-            .await?;
+    let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
+        state, trace_id, &input, body_json, spec,
+    )
+    .await?;
     apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
         state,
         trace_id,
@@ -141,9 +142,10 @@ pub(crate) async fn build_local_stream_plan_and_reports(
         Some(input.requested_model.as_str()),
         "candidate_evaluation_incomplete",
     );
-    let (attempts, candidate_count) =
-        materialize_local_same_format_provider_candidate_attempts(state, trace_id, &input, spec)
-            .await?;
+    let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
+        state, trace_id, &input, body_json, spec,
+    )
+    .await?;
     apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
         state,
         trace_id,

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

@@ -94,6 +94,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
             GEMINI_FILES_CLIENT_API_FORMAT,
             None,
             input.required_capabilities.as_ref(),
+            None,
         )
         .await;
 

apps/aether-gateway/src/ai_pipeline/planner/specialized/video.rs

@@ -60,6 +60,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
         state,
         trace_id,
         &input,
+        body_json,
         spec_metadata.api_format,
         spec_metadata.decision_kind,
     )
@@ -102,6 +103,7 @@ async fn build_local_sync_plan_and_reports(
         state,
         trace_id,
         &input,
+        body_json,
         spec_metadata.api_format,
         spec_metadata.decision_kind,
     )

apps/aether-gateway/src/ai_pipeline/planner/specialized/video/support.rs

@@ -4,7 +4,8 @@ use tracing::warn;
 use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
 use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
 use crate::ai_pipeline::planner::candidate_eligibility::{
-    filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
+    extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
+    SkippedLocalExecutionCandidate,
 };
 use crate::ai_pipeline::planner::candidate_materialization::{
     mark_skipped_local_execution_candidate,
@@ -96,6 +97,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
     state: &AppState,
     trace_id: &str,
     input: &LocalVideoCreateDecisionInput,
+    body_json: &serde_json::Value,
     api_format: &str,
     decision_kind: &str,
 ) -> Option<Vec<LocalVideoCreateCandidateAttempt>> {
@@ -128,6 +130,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
             planner_state,
             trace_id,
             input,
+            body_json,
             candidates,
             preselection_skipped
                 .into_iter()
@@ -148,10 +151,12 @@ async fn materialize_local_video_create_candidate_attempts(
     state: PlannerAppState<'_>,
     trace_id: &str,
     input: &LocalVideoCreateDecisionInput,
+    body_json: &serde_json::Value,
     candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
     preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
     api_format: &str,
 ) -> Vec<LocalVideoCreateCandidateAttempt> {
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
     let persistence_policy = build_local_candidate_persistence_policy(
         &input.auth_context,
         input.required_capabilities.as_ref(),
@@ -163,6 +168,7 @@ async fn materialize_local_video_create_candidate_attempts(
         api_format,
         &input.requested_model,
         input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
     )
     .await;
     let skipped_candidates = preselection_skipped

apps/aether-gateway/src/ai_pipeline/planner/standard/family/build.rs

@@ -51,7 +51,8 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
         "candidate_evaluation_incomplete",
     );
     let (attempts, candidate_count) =
-        materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
 
     for attempt in attempts {
@@ -99,7 +100,8 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
         "candidate_evaluation_incomplete",
     );
     let (attempts, candidate_count) =
-        materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
 
     for attempt in attempts {
@@ -153,7 +155,8 @@ pub(crate) async fn build_local_sync_plan_and_reports(
         "candidate_evaluation_incomplete",
     );
     let (attempts, candidate_count) =
-        materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
     if candidate_count == 0 {
         return Ok(Vec::new());
@@ -226,7 +229,8 @@ pub(crate) async fn build_local_stream_plan_and_reports(
         "candidate_evaluation_incomplete",
     );
     let (attempts, candidate_count) =
-        materialize_local_standard_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
     apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
     if candidate_count == 0 {
         return Ok(Vec::new());

apps/aether-gateway/src/ai_pipeline/planner/standard/family/candidates.rs

@@ -4,7 +4,8 @@ use tracing::warn;
 
 use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
 use crate::ai_pipeline::planner::candidate_eligibility::{
-    filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
+    extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
+    SkippedLocalExecutionCandidate,
 };
 use crate::ai_pipeline::planner::candidate_materialization::{
     persist_available_local_execution_candidates_with_context,
@@ -87,10 +88,12 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
     state: &AppState,
     trace_id: &str,
     input: &LocalStandardDecisionInput,
+    body_json: &serde_json::Value,
     spec: LocalStandardSpec,
 ) -> Result<(Vec<LocalStandardCandidateAttempt>, usize), GatewayError> {
     let spec_metadata = local_standard_spec_metadata(spec);
     let planner_state = PlannerAppState::new(state);
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
     let persistence_policy = build_local_candidate_persistence_policy(
         &input.auth_context,
         input.required_capabilities.as_ref(),
@@ -176,6 +179,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
         spec_metadata.api_format,
         &input.requested_model,
         input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
     )
     .await;
     let skipped_candidates = preselection_skipped

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/support.rs

@@ -2,7 +2,8 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
 
 use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
 use crate::ai_pipeline::planner::candidate_eligibility::{
-    filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
+    extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
+    SkippedLocalExecutionCandidate,
 };
 use crate::ai_pipeline::planner::candidate_materialization::{
     mark_skipped_local_execution_candidate,
@@ -55,10 +56,12 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
     state: &AppState,
     trace_id: &str,
     input: &LocalOpenAiChatDecisionInput,
+    body_json: &serde_json::Value,
     candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
     preselection_skipped: Vec<SkippedLocalExecutionCandidate>,
 ) -> Vec<LocalOpenAiChatCandidateAttempt> {
     let planner_state = PlannerAppState::new(state);
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
     let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
     let persistence_policy = build_local_candidate_persistence_policy(
         auth_context,
@@ -71,6 +74,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
         "openai:chat",
         &input.requested_model,
         input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
     )
     .await;
     let skipped_candidates = preselection_skipped

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/mod.rs

@@ -118,6 +118,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
         state,
         trace_id,
         &input,
+        body_json,
         candidates,
         skipped_candidates,
     )
@@ -183,6 +184,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
         state,
         trace_id,
         &input,
+        body_json,
         candidates,
         skipped_candidates,
     )

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/stream.rs

@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
         state,
         trace_id,
         &input,
+        body_json,
         candidates,
         skipped_candidates,
     )

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/sync.rs

@@ -80,6 +80,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
         state,
         trace_id,
         &input,
+        body_json,
         candidates,
         skipped_candidates,
     )

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/support.rs

@@ -6,7 +6,8 @@ use tracing::warn;
 use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
 use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
 use crate::ai_pipeline::planner::candidate_eligibility::{
-    filter_and_rank_local_execution_candidates, SkippedLocalExecutionCandidate,
+    extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
+    SkippedLocalExecutionCandidate,
 };
 use crate::ai_pipeline::planner::candidate_materialization::{
     mark_skipped_local_execution_candidate,
@@ -82,11 +83,13 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
     state: &AppState,
     trace_id: &str,
     input: &LocalOpenAiCliDecisionInput,
+    body_json: &serde_json::Value,
     spec: LocalOpenAiCliSpec,
 ) -> Result<Vec<LocalOpenAiCliCandidateAttempt>, GatewayError> {
     let spec_metadata = local_openai_cli_spec_metadata(spec);
     let client_api_format = spec_metadata.api_format.trim().to_ascii_lowercase();
     let planner_state = PlannerAppState::new(state);
+    let sticky_session_token = extract_pool_sticky_session_token(body_json);
     let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
     let persistence_policy = build_local_candidate_persistence_policy(
         auth_context,
@@ -173,6 +176,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
         spec_metadata.api_format,
         &input.requested_model,
         input.required_capabilities.as_ref(),
+        sticky_session_token.as_deref(),
     )
     .await;
     let skipped_candidates = preselection_skipped

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/mod.rs

@@ -66,7 +66,8 @@ pub(crate) async fn maybe_build_sync_local_openai_cli_decision_payload(
     };
 
     let attempts =
-        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
 
     for attempt in attempts {
         if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(
@@ -100,7 +101,8 @@ pub(crate) async fn maybe_build_stream_local_openai_cli_decision_payload(
     };
 
     let attempts =
-        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
 
     for attempt in attempts {
         if let Some(payload) = maybe_build_local_openai_cli_decision_payload_for_candidate(

apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/plans.rs

@@ -33,7 +33,8 @@ pub(super) async fn build_local_sync_plan_and_reports(
     };
 
     let attempts =
-        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
 
     let mut plans = Vec::new();
     for attempt in attempts {
@@ -78,7 +79,8 @@ pub(super) async fn build_local_stream_plan_and_reports(
     };
 
     let attempts =
-        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, spec).await?;
+        materialize_local_openai_cli_candidate_attempts(state, trace_id, &input, body_json, spec)
+            .await?;
 
     let mut plans = Vec::new();
     for attempt in attempts {