i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 抽离 AI pipeline 与调度共享能力逻辑

Browse Source
This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -28,19 +28,27 @@ use std::collections::BTreeMap;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
read_global_model_names_for_required_capability, MinimalCandidateSelectionRowSource,
read_global_model_names_for_api_format, read_global_model_names_for_required_capability,
MinimalCandidateSelectionRowSource,
};
use crate::GatewayError;
#[cfg_attr(not(test), allow(dead_code))]
const SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RequiredCapabilityMatchMode {
Compatible,
Exclusive,
}
pub(crate) async fn list_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -50,6 +58,7 @@ pub(crate) async fn list_selectable_candidates(
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
now_unix_secs,
)
@@ -70,37 +79,87 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
return Ok(Vec::new());
}
let model_names = read_global_model_names_for_required_capability(
selection_row_source,
&normalized_api_format,
required_capability,
require_streaming,
auth_snapshot,
)
.await
let capability_mode = required_capability_match_mode(required_capability);
let model_names = match capability_mode {
RequiredCapabilityMatchMode::Exclusive => {
read_global_model_names_for_required_capability(
selection_row_source,
&normalized_api_format,
required_capability,
require_streaming,
auth_snapshot,
)
.await
}
RequiredCapabilityMatchMode::Compatible => {
read_global_model_names_for_api_format(
selection_row_source,
&normalized_api_format,
require_streaming,
auth_snapshot,
)
.await
}
}
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let required_capabilities = build_required_capabilities_object(required_capability);
for global_model_name in model_names {
let candidates = list_selectable_candidates(
let mut candidates = list_selectable_candidates(
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
now_unix_secs,
)
.await?;
let filtered = candidates
.into_iter()
.filter(|candidate| {
candidate_supports_required_capability(candidate, required_capability)
})
.collect::<Vec<_>>();
if !filtered.is_empty() {
return Ok(filtered);
match capability_mode {
RequiredCapabilityMatchMode::Exclusive => {
let filtered = candidates
.into_iter()
.filter(|candidate| {
candidate_supports_required_capability(candidate, required_capability)
})
.collect::<Vec<_>>();
if !filtered.is_empty() {
return Ok(filtered);
}
}
RequiredCapabilityMatchMode::Compatible => {
if candidates.is_empty() {
continue;
}
candidates.sort_by_key(|candidate| {
!candidate_supports_required_capability(candidate, required_capability)
});
return Ok(candidates);
}
}
}
Ok(Vec::new())
}
fn required_capability_match_mode(required_capability: &str) -> RequiredCapabilityMatchMode {
match required_capability.trim().to_ascii_lowercase().as_str() {
"cache_1h" | "context_1m" => RequiredCapabilityMatchMode::Compatible,
_ => RequiredCapabilityMatchMode::Exclusive,
}
}
fn build_required_capabilities_object(required_capability: &str) -> Option<serde_json::Value> {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return None;
}
let mut capabilities = serde_json::Map::new();
capabilities.insert(
required_capability.to_string(),
serde_json::Value::Bool(true),
);
Some(serde_json::Value::Object(capabilities))
}

113
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -4,13 +4,16 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
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, MinimalCandidateSelectionRowSource,
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::GatewayError;
use super::affinity::{
@@ -25,18 +28,66 @@ use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
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),
@@ -44,6 +95,7 @@ pub(super) async fn select_minimal_candidate(
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -55,6 +107,7 @@ pub(super) async fn select_minimal_candidate(
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
now_unix_secs,
)
@@ -73,31 +126,54 @@ pub(super) async fn collect_selectable_candidates(
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let mut candidates = read_minimal_candidate_selection(
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(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
ordering_config.priority_mode,
priority_affinity_key,
required_capabilities,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
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,
auth_snapshot,
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 = affinity_cache_key.as_deref().and_then(|cache_key| {
runtime_state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
});
let cached_affinity_target = if ordering_config.scheduling_mode
== SchedulerSchedulingMode::CacheAffinity
{
affinity_cache_key.as_deref().and_then(|cache_key| {
runtime_state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
})
} else {
None
};
if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
return Ok(Vec::new());
@@ -123,3 +199,16 @@ pub(super) async fn collect_selectable_candidates(
cached_affinity_target.as_ref(),
))
}
fn scheduling_priority_affinity_key<'a>(
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
scheduling_mode: SchedulerSchedulingMode,
) -> Option<&'a str> {
if scheduling_mode == SchedulerSchedulingMode::FixedOrder {
return None;
}
auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())
}

38
apps/aether-gateway/src/scheduler/candidate/tests/affinity.rs
View File

@@ -9,16 +9,42 @@ use aether_data_contracts::repository::candidate_selection::StoredProviderModelM
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::cache::SchedulerAffinityTarget;
use crate::data::candidate_selection::read_minimal_candidate_selection;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
read_minimal_candidate_selection, MinimalCandidateSelectionRowSource,
};
use crate::data::GatewayDataState;
use crate::AppState;
use crate::{AppState, GatewayError};
use super::super::affinity::{build_scheduler_affinity_cache_key, candidate_affinity_hash};
use super::super::selection::select_minimal_candidate as select_candidate;
use super::super::selection::select_minimal_candidate as select_candidate_impl;
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
async fn select_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
select_candidate_impl(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
None,
auth_snapshot,
now_unix_secs,
)
.await
}
#[tokio::test]
async fn same_priority_candidates_are_distributed_by_affinity_key() {
let mut first = sample_row();
@@ -196,9 +222,9 @@ async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity()
Some(9),
None,
None,
95,
Some(95),
Some(96),
95_000,
Some(95_000),
Some(96_000),
)
.expect("candidate should build"),
]));

1
apps/aether-gateway/src/scheduler/candidate/tests/mod.rs
View File

@@ -1,4 +1,5 @@
mod affinity;
mod model;
mod required_capability;
mod selection;
mod support;

116
apps/aether-gateway/src/scheduler/candidate/tests/required_capability.rs Normal file
View File

@@ -0,0 +1,116 @@
use std::sync::Arc;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use crate::data::GatewayDataState;
use crate::AppState;
use super::super::list_selectable_candidates_for_required_capability_without_requested_model;
use super::support::sample_row;
#[tokio::test]
async fn compatible_required_capability_prefers_matching_keys_without_hard_filtering() {
let mut higher_priority = sample_row();
higher_priority.provider_id = "provider-a".to_string();
higher_priority.provider_name = "provider-a".to_string();
higher_priority.endpoint_id = "endpoint-a".to_string();
higher_priority.key_id = "key-a".to_string();
higher_priority.key_name = "alpha".to_string();
higher_priority.provider_priority = 0;
higher_priority.key_internal_priority = 0;
higher_priority.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 0}));
higher_priority.key_capabilities = Some(serde_json::json!({}));
let mut capability_match = sample_row();
capability_match.provider_id = "provider-b".to_string();
capability_match.provider_name = "provider-b".to_string();
capability_match.endpoint_id = "endpoint-b".to_string();
capability_match.key_id = "key-b".to_string();
capability_match.key_name = "beta".to_string();
capability_match.provider_priority = 10;
capability_match.key_internal_priority = 10;
capability_match.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": 10}));
capability_match.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
higher_priority,
capability_match,
]));
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 selection = list_selectable_candidates_for_required_capability_without_requested_model(
state.data.as_ref(),
&state,
"openai:chat",
"cache_1h",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 2);
assert_eq!(selection[0].provider_id, "provider-b");
assert_eq!(selection[1].provider_id, "provider-a");
}
#[tokio::test]
async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys() {
let mut incompatible = sample_row();
incompatible.provider_id = "provider-a".to_string();
incompatible.provider_name = "provider-a".to_string();
incompatible.endpoint_id = "endpoint-a".to_string();
incompatible.endpoint_api_format = "gemini:chat".to_string();
incompatible.endpoint_api_family = Some("gemini".to_string());
incompatible.key_api_formats = Some(vec!["gemini:chat".to_string()]);
incompatible.key_id = "key-a".to_string();
incompatible.key_name = "alpha".to_string();
incompatible.global_model_name = "gemini-2.5-pro".to_string();
incompatible.key_capabilities = Some(serde_json::json!({}));
let mut compatible = sample_row();
compatible.provider_id = "provider-b".to_string();
compatible.provider_name = "provider-b".to_string();
compatible.endpoint_id = "endpoint-b".to_string();
compatible.endpoint_api_format = "gemini:chat".to_string();
compatible.endpoint_api_family = Some("gemini".to_string());
compatible.key_api_formats = Some(vec!["gemini:chat".to_string()]);
compatible.key_id = "key-b".to_string();
compatible.key_name = "beta".to_string();
compatible.global_model_name = "gemini-2.5-pro".to_string();
compatible.key_capabilities = Some(serde_json::json!({"gemini_files": true}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
incompatible,
compatible,
]));
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 selection = list_selectable_candidates_for_required_capability_without_requested_model(
state.data.as_ref(),
&state,
"gemini:chat",
"gemini_files",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 1);
assert_eq!(selection[0].provider_id, "provider-b");
assert_eq!(selection[0].key_id, "key-b");
}

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

@@ -1,4 +1,5 @@
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
@@ -9,14 +10,66 @@ use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use crate::cache::SchedulerAffinityTarget;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::data::GatewayDataState;
use crate::AppState;
use crate::{AppState, GatewayError};
use super::super::runtime::should_skip_provider_quota;
use super::super::selection::select_minimal_candidate as select_candidate;
use super::super::selection::{
collect_selectable_candidates as collect_selectable_candidates_impl,
select_minimal_candidate as select_candidate_impl,
};
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
async fn select_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
select_candidate_impl(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
None,
auth_snapshot,
now_unix_secs,
)
.await
}
async fn collect_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &AppState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates_impl(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
None,
auth_snapshot,
now_unix_secs,
)
.await
}
#[test]
fn skips_inactive_or_exhausted_monthly_quota_provider() {
let inactive = StoredProviderQuotaSnapshot::new(
@@ -59,6 +112,439 @@ fn skips_inactive_or_exhausted_monthly_quota_provider() {
assert!(!should_skip_provider_quota(&payg, 2_000));
}
#[tokio::test]
async fn selects_by_provider_priority_when_priority_mode_is_provider() {
let mut provider_first = sample_row();
provider_first.provider_id = "provider-a".to_string();
provider_first.provider_name = "provider-a".to_string();
provider_first.endpoint_id = "endpoint-a".to_string();
provider_first.key_id = "key-a".to_string();
provider_first.key_name = "alpha".to_string();
provider_first.provider_priority = 0;
provider_first.key_internal_priority = 20;
provider_first.key_global_priority_by_format = Some(json!({"openai:chat": 10}));
let mut global_key_first = sample_row();
global_key_first.provider_id = "provider-b".to_string();
global_key_first.provider_name = "provider-b".to_string();
global_key_first.endpoint_id = "endpoint-b".to_string();
global_key_first.key_id = "key-b".to_string();
global_key_first.key_name = "beta".to_string();
global_key_first.provider_priority = 10;
global_key_first.key_internal_priority = 0;
global_key_first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
provider_first,
global_key_first,
]));
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)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
json!("provider"),
)]),
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-a");
assert_eq!(selected.key_id, "key-a");
}
#[tokio::test]
async fn selects_by_global_key_priority_when_priority_mode_is_global_key() {
let mut provider_first = sample_row();
provider_first.provider_id = "provider-a".to_string();
provider_first.provider_name = "provider-a".to_string();
provider_first.endpoint_id = "endpoint-a".to_string();
provider_first.key_id = "key-a".to_string();
provider_first.key_name = "alpha".to_string();
provider_first.provider_priority = 0;
provider_first.key_internal_priority = 20;
provider_first.key_global_priority_by_format = Some(json!({"openai:chat": 10}));
let mut global_key_first = sample_row();
global_key_first.provider_id = "provider-b".to_string();
global_key_first.provider_name = "provider-b".to_string();
global_key_first.endpoint_id = "endpoint-b".to_string();
global_key_first.key_id = "key-b".to_string();
global_key_first.key_name = "beta".to_string();
global_key_first.provider_priority = 10;
global_key_first.key_internal_priority = 0;
global_key_first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
provider_first,
global_key_first,
]));
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)
.with_system_config_values_for_tests(vec![(
"provider_priority_mode".to_string(),
json!("global_key"),
)]),
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn scheduler_selection_prefers_required_capability_matches_before_priority_fallback() {
let mut higher_priority_missing_capability = sample_row();
higher_priority_missing_capability.provider_id = "provider-a".to_string();
higher_priority_missing_capability.provider_name = "provider-a".to_string();
higher_priority_missing_capability.endpoint_id = "endpoint-a".to_string();
higher_priority_missing_capability.key_id = "key-a".to_string();
higher_priority_missing_capability.key_name = "alpha".to_string();
higher_priority_missing_capability.provider_priority = 0;
higher_priority_missing_capability.key_internal_priority = 0;
higher_priority_missing_capability.key_capabilities = Some(json!({"cache_1h": false}));
let mut lower_priority_matching_capability = sample_row();
lower_priority_matching_capability.provider_id = "provider-b".to_string();
lower_priority_matching_capability.provider_name = "provider-b".to_string();
lower_priority_matching_capability.endpoint_id = "endpoint-b".to_string();
lower_priority_matching_capability.key_id = "key-b".to_string();
lower_priority_matching_capability.key_name = "beta".to_string();
lower_priority_matching_capability.provider_priority = 10;
lower_priority_matching_capability.key_internal_priority = 10;
lower_priority_matching_capability.key_capabilities = Some(json!({"cache_1h": true}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
higher_priority_missing_capability,
lower_priority_matching_capability,
]));
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 required_capabilities = json!({"cache_1h": true});
let selected = select_candidate_impl(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&required_capabilities),
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn fixed_order_ignores_cached_scheduler_affinity_promotion() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "provider-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.provider_priority = 0;
first.key_internal_priority = 0;
first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "provider-b".to_string();
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.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
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)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("fixed_order"),
)]),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
key_id: "key-b".to_string(),
},
Duration::from_secs(300),
100,
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-a");
assert_eq!(selected.key_id, "key-a");
}
#[tokio::test]
async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "provider-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.provider_priority = 0;
first.key_internal_priority = 0;
first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "provider-b".to_string();
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.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
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)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("fixed_order"),
)]),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
let selection = collect_selectable_candidates(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed");
assert_eq!(selection.len(), 2);
assert_eq!(selection[0].provider_id, "provider-a");
assert_eq!(selection[1].provider_id, "provider-b");
}
#[tokio::test]
async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enabled() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "provider-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.provider_priority = 0;
first.key_internal_priority = 0;
first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "provider-b".to_string();
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.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
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)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("cache_affinity"),
)]),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
key_id: "key-b".to_string(),
},
Duration::from_secs(300),
100,
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "provider-b");
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn load_balance_rotates_same_priority_group_and_ignores_cached_affinity() {
let mut first = sample_row();
first.provider_id = "provider-a".to_string();
first.provider_name = "provider-a".to_string();
first.endpoint_id = "endpoint-a".to_string();
first.key_id = "key-a".to_string();
first.key_name = "alpha".to_string();
first.provider_priority = 0;
first.key_internal_priority = 0;
first.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let mut second = sample_row();
second.provider_id = "provider-b".to_string();
second.provider_name = "provider-b".to_string();
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.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
]));
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)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("load_balance"),
)]),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
key_id: "key-b".to_string(),
},
Duration::from_secs(300),
100,
);
let first_pass = collect_selectable_candidates(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("first pass should succeed");
let second_pass = collect_selectable_candidates(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
101,
)
.await
.expect("second pass should succeed");
assert_eq!(first_pass.len(), 2);
assert_eq!(second_pass.len(), 2);
assert_ne!(first_pass[0].provider_id, second_pass[0].provider_id);
assert!(
first_pass[0].provider_id != "provider-b" || second_pass[0].provider_id != "provider-b"
);
assert_ne!(
first_pass
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>(),
second_pass
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>()
);
}
#[tokio::test]
async fn selects_next_candidate_when_first_provider_quota_is_exhausted() {
let mut first = sample_row();
@@ -185,9 +671,9 @@ async fn cooled_down_when_recent_failures_are_recorded_for_same_key() {
None,
None,
None,
95,
Some(95),
Some(95),
95_000,
Some(95_000),
Some(95_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
@@ -212,9 +698,9 @@ async fn cooled_down_when_recent_failures_are_recorded_for_same_key() {
None,
None,
None,
98,
Some(98),
Some(98),
98_000,
Some(98_000),
Some(98_000),
)
.expect("candidate should build"),
]));
@@ -298,8 +784,8 @@ async fn selects_next_candidate_when_first_provider_concurrent_limit_is_reached(
None,
None,
None,
95,
Some(95),
95_000,
Some(95_000),
None,
)
.expect("candidate should build"),
@@ -366,8 +852,8 @@ async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
None,
None,
None,
95,
Some(95),
95_000,
Some(95_000),
None,
)
.expect("candidate should build"),
@@ -457,9 +943,9 @@ async fn selects_next_candidate_when_first_provider_key_rpm_slots_are_reserved_f
Some(9),
None,
None,
95,
Some(95),
Some(96),
95_000,
Some(95_000),
Some(96_000),
)
.expect("candidate should build"),
]));