Unify candidate ranking pipeline

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

@@ -0,0 +1,28 @@
+use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
+
+use crate::data::auth::GatewayAuthApiKeySnapshot;
+use crate::data::candidate_selection::{
+    enumerate_minimal_candidate_selection_with_required_capabilities,
+    MinimalCandidateSelectionRowSource,
+};
+use crate::GatewayError;
+
+pub(super) async fn enumerate_scheduler_candidates(
+    selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
+    api_format: &str,
+    global_model_name: &str,
+    require_streaming: bool,
+    required_capabilities: Option<&serde_json::Value>,
+    auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
+    enumerate_minimal_candidate_selection_with_required_capabilities(
+        selection_row_source,
+        api_format,
+        global_model_name,
+        require_streaming,
+        auth_snapshot,
+        required_capabilities,
+    )
+    .await
+    .map_err(|err| GatewayError::Internal(err.to_string()))
+}

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

@@ -1,10 +1,12 @@
-use self::affinity::candidate_affinity_hash;
 use self::selection::{
     collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
 };
 use super::state::SchedulerRuntimeState;
 
 mod affinity;
+mod enumeration;
+mod ranking;
+mod resolution;
 mod runtime;
 mod selection;
 

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

@@ -0,0 +1,76 @@
+use aether_scheduler_core::{
+    apply_scheduler_candidate_ranking, effective_provider_key_health_score,
+    matches_affinity_target, provider_key_health_bucket,
+    requested_capability_priority_for_candidate, SchedulerAffinityTarget,
+    SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
+};
+
+use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
+
+use super::affinity::candidate_affinity_hash;
+use super::runtime::CandidateRuntimeSelectionSnapshot;
+use super::SchedulerMinimalCandidateSelectionCandidate;
+
+pub(super) fn rank_scheduler_candidates(
+    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
+    runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
+    ordering_config: SchedulerOrderingConfig,
+    required_capabilities: Option<&serde_json::Value>,
+    priority_affinity_key: Option<&str>,
+    cached_affinity_target: Option<&SchedulerAffinityTarget>,
+    now_unix_secs: u64,
+) {
+    let rankables = candidates
+        .iter()
+        .enumerate()
+        .map(|(index, candidate)| {
+            let provider_key = runtime_snapshot
+                .provider_key_rpm_states
+                .get(&candidate.key_id);
+            SchedulerRankableCandidate::from_candidate(candidate, index)
+                .with_capability_priority(requested_capability_priority_for_candidate(
+                    required_capabilities,
+                    candidate,
+                ))
+                .with_cached_affinity_match(
+                    cached_affinity_target
+                        .is_some_and(|target| matches_affinity_target(candidate, target)),
+                )
+                .with_affinity_hash(
+                    priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
+                )
+                .with_health(
+                    provider_key.and_then(|key| {
+                        provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
+                    }),
+                    provider_key
+                        .and_then(|key| {
+                            effective_provider_key_health_score(
+                                key,
+                                candidate.endpoint_api_format.as_str(),
+                            )
+                        })
+                        .unwrap_or(1.0),
+                )
+        })
+        .collect::<Vec<_>>();
+
+    apply_scheduler_candidate_ranking(
+        candidates,
+        &rankables,
+        SchedulerRankingContext {
+            priority_mode: ordering_config.priority_mode,
+            ranking_mode: scheduler_ranking_mode(ordering_config.scheduling_mode),
+            include_health: true,
+            load_balance_seed: now_unix_secs,
+        },
+    );
+}
+
+fn scheduler_ranking_mode(mode: SchedulerSchedulingMode) -> SchedulerRankingMode {
+    match mode {
+        SchedulerSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
+        SchedulerSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
+        SchedulerSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
+    }
+}

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

@@ -0,0 +1,45 @@
+use std::collections::BTreeSet;
+
+use aether_scheduler_core::SchedulerAffinityTarget;
+
+use super::affinity::candidate_key;
+use super::runtime::{current_candidate_runtime_skip_reason, CandidateRuntimeSelectionSnapshot};
+use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerSkippedCandidate};
+
+pub(super) fn resolve_scheduler_candidate_selectability(
+    candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
+    runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
+    now_unix_secs: u64,
+    cached_affinity_target: Option<&SchedulerAffinityTarget>,
+) -> (
+    Vec<SchedulerMinimalCandidateSelectionCandidate>,
+    Vec<SchedulerSkippedCandidate>,
+) {
+    let mut selected = Vec::with_capacity(candidates.len());
+    let mut skipped = Vec::new();
+    let mut emitted_selected_keys = BTreeSet::new();
+    let mut emitted_skipped_keys = BTreeSet::new();
+
+    for candidate in candidates {
+        let key = candidate_key(&candidate);
+        if let Some(skip_reason) = current_candidate_runtime_skip_reason(
+            &candidate,
+            runtime_snapshot,
+            now_unix_secs,
+            cached_affinity_target,
+        ) {
+            if emitted_skipped_keys.insert(key) {
+                skipped.push(SchedulerSkippedCandidate {
+                    candidate,
+                    skip_reason,
+                });
+            }
+            continue;
+        }
+        if emitted_selected_keys.insert(key) {
+            selected.push(candidate);
+        }
+    }
+
+    (selected, skipped)
+}

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

@@ -1,27 +1,15 @@
-use std::collections::{BTreeMap, BTreeSet};
-
-use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
-use aether_scheduler_core::{
-    collect_selectable_candidates_from_keys,
-    reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
-    SchedulerPriorityMode,
-};
-
 use crate::data::auth::GatewayAuthApiKeySnapshot;
-use crate::data::candidate_selection::{
-    read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
-    MinimalCandidateSelectionRowSource,
-};
+use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
 use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
 use crate::scheduler::config::SchedulerSchedulingMode;
 use crate::GatewayError;
 
-use super::affinity::{
-    build_scheduler_affinity_cache_key, candidate_key, remember_scheduler_affinity,
-};
+use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
+use super::enumeration::enumerate_scheduler_candidates;
+use super::ranking::rank_scheduler_candidates;
+use super::resolution::resolve_scheduler_candidate_selectability;
 use super::runtime::{
-    auth_snapshot_concurrency_limit_reached, current_candidate_runtime_skip_reason,
-    read_candidate_runtime_selection_snapshot,
+    auth_snapshot_concurrency_limit_reached, read_candidate_runtime_selection_snapshot,
 };
 use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
 
@@ -44,69 +32,6 @@ pub(super) fn is_exact_all_skipped_by_auth_limit(
             .all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
 }
 
-pub(super) fn reorder_candidates_by_scheduler_health(
-    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
-    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
-    required_capabilities: Option<&serde_json::Value>,
-    affinity_key: Option<&str>,
-    priority_mode: SchedulerPriorityMode,
-) {
-    reorder_candidates_by_scheduler_health_in_core(
-        candidates,
-        provider_key_rpm_states,
-        required_capabilities,
-        affinity_key,
-        priority_mode,
-    );
-}
-
-fn apply_load_balance_rotation(
-    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
-    priority_mode: SchedulerPriorityMode,
-    now_unix_secs: u64,
-) {
-    if candidates.len() < 2 {
-        return;
-    }
-
-    let mut start = 0usize;
-    while start < candidates.len() {
-        let mut end = start + 1;
-        while end < candidates.len()
-            && candidates_share_load_balance_group(
-                &candidates[start],
-                &candidates[end],
-                priority_mode,
-            )
-        {
-            end += 1;
-        }
-
-        let group_len = end - start;
-        if group_len > 1 {
-            let offset = usize::try_from(now_unix_secs).unwrap_or(0) % group_len;
-            candidates[start..end].rotate_left(offset);
-        }
-        start = end;
-    }
-}
-
-fn candidates_share_load_balance_group(
-    left: &SchedulerMinimalCandidateSelectionCandidate,
-    right: &SchedulerMinimalCandidateSelectionCandidate,
-    priority_mode: SchedulerPriorityMode,
-) -> bool {
-    match priority_mode {
-        SchedulerPriorityMode::Provider => {
-            left.provider_priority == right.provider_priority
-                && left.key_internal_priority == right.key_internal_priority
-        }
-        SchedulerPriorityMode::GlobalKey => {
-            left.key_global_priority_for_format == right.key_global_priority_for_format
-        }
-    }
-}
-
 #[cfg_attr(not(test), allow(dead_code))]
 pub(super) async fn select_minimal_candidate(
     selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -182,36 +107,18 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
     let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
     let priority_affinity_key =
         scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
-    let mut candidates =
-        read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
+    let mut candidates = enumerate_scheduler_candidates(
         selection_row_source,
         api_format,
         global_model_name,
         require_streaming,
-        auth_snapshot,
-        ordering_config.priority_mode,
-        priority_affinity_key,
         required_capabilities,
+        auth_snapshot,
     )
-        .await
-        .map_err(|err| GatewayError::Internal(err.to_string()))?;
+    .await?;
     let runtime_snapshot =
         read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
             .await?;
-    reorder_candidates_by_scheduler_health(
-        &mut candidates,
-        &runtime_snapshot.provider_key_rpm_states,
-        required_capabilities,
-        priority_affinity_key,
-        ordering_config.priority_mode,
-    );
-    if ordering_config.scheduling_mode == SchedulerSchedulingMode::LoadBalance {
-        apply_load_balance_rotation(
-            &mut candidates,
-            ordering_config.priority_mode,
-            now_unix_secs,
-        );
-    }
     let affinity_cache_key =
         build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
     let cached_affinity_target = if ordering_config.scheduling_mode
@@ -225,6 +132,15 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
     };
 
     if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
+        rank_scheduler_candidates(
+            &mut candidates,
+            &runtime_snapshot,
+            ordering_config,
+            required_capabilities,
+            priority_affinity_key,
+            cached_affinity_target.as_ref(),
+            now_unix_secs,
+        );
         return Ok((
             Vec::new(),
             candidates
@@ -237,37 +153,23 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
         ));
     }
 
-    let mut selected_keys = BTreeSet::new();
-    let mut skipped = Vec::new();
-    let mut emitted_skipped_keys = BTreeSet::new();
+    let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
+        candidates,
+        &runtime_snapshot,
+        now_unix_secs,
+        cached_affinity_target.as_ref(),
+    );
+    rank_scheduler_candidates(
+        &mut selected,
+        &runtime_snapshot,
+        ordering_config,
+        required_capabilities,
+        priority_affinity_key,
+        cached_affinity_target.as_ref(),
+        now_unix_secs,
+    );
 
-    for candidate in &candidates {
-        let key = candidate_key(candidate);
-        if let Some(skip_reason) = current_candidate_runtime_skip_reason(
-            candidate,
-            &runtime_snapshot,
-            now_unix_secs,
-            cached_affinity_target.as_ref(),
-        ) {
-            if emitted_skipped_keys.insert(key) {
-                skipped.push(SchedulerSkippedCandidate {
-                    candidate: candidate.clone(),
-                    skip_reason,
-                });
-            }
-            continue;
-        }
-        selected_keys.insert(key);
-    }
-
-    Ok((
-        collect_selectable_candidates_from_keys(
-            candidates,
-            &selected_keys,
-            cached_affinity_target.as_ref(),
-        ),
-        skipped,
-    ))
+    Ok((selected, skipped))
 }
 
 fn scheduling_priority_affinity_key<'a>(

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

@@ -144,6 +144,54 @@ async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
     assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
 }
 
+#[tokio::test]
+async fn read_minimal_candidate_selection_keeps_all_rows_supporting_requested_model() {
+    let mut exact = sample_row();
+    exact.provider_id = "provider-exact".to_string();
+    exact.endpoint_id = "endpoint-exact".to_string();
+    exact.key_id = "key-exact".to_string();
+    exact.model_id = "model-exact".to_string();
+    exact.global_model_id = "global-exact".to_string();
+    exact.global_model_name = "gpt-5".to_string();
+    exact.model_provider_model_name = "gpt-5".to_string();
+    exact.model_provider_model_mappings = None;
+
+    let mut mapped = sample_row();
+    mapped.provider_id = "provider-mapped".to_string();
+    mapped.endpoint_id = "endpoint-mapped".to_string();
+    mapped.key_id = "key-mapped".to_string();
+    mapped.model_id = "model-mapped".to_string();
+    mapped.global_model_id = "global-mapped".to_string();
+    mapped.global_model_name = "claude-sonnet".to_string();
+    mapped.global_model_mappings = Some(vec!["gpt-5".to_string()]);
+    mapped.model_provider_model_name = "claude-sonnet-upstream".to_string();
+    mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
+        name: "claude-sonnet-upstream".to_string(),
+        priority: 1,
+        api_formats: Some(vec!["openai:chat".to_string()]),
+    }]);
+
+    let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+        exact, mapped,
+    ]));
+    let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
+    let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
+
+    let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5", false, None)
+        .await
+        .expect("selection should succeed");
+
+    let provider_ids = selection
+        .iter()
+        .map(|candidate| candidate.provider_id.as_str())
+        .collect::<Vec<_>>();
+    assert_eq!(provider_ids, vec!["provider-exact", "provider-mapped"]);
+    assert_eq!(
+        selection[1].selected_provider_model_name,
+        "claude-sonnet-upstream"
+    );
+}
+
 #[tokio::test]
 async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
     let mut row = sample_row();

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

@@ -384,9 +384,9 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
     second.endpoint_id = "endpoint-b".to_string();
     second.key_id = "key-b".to_string();
     second.key_name = "beta".to_string();
-    second.provider_priority = 0;
+    second.provider_priority = 1;
     second.key_internal_priority = 0;
-    second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
+    second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
 
     let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
         first, second,
@@ -438,9 +438,9 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
     second.endpoint_id = "endpoint-b".to_string();
     second.key_id = "key-b".to_string();
     second.key_name = "beta".to_string();
-    second.provider_priority = 1;
+    second.provider_priority = 0;
     second.key_internal_priority = 0;
-    second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
+    second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
 
     let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
         first, second,