fix(routing): harden routed pool scheduling

This commit is contained in:
elky
2026-07-27 22:06:28 +08:00
parent 550cc36760
commit 4148ab1931
62 changed files with 5510 additions and 850 deletions
@@ -10,6 +10,7 @@ const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
"load_balance",
"single_account",
"priority_first",
"free_team_first",
"free_first",
"team_first",
"plus_first",
@@ -166,8 +167,9 @@ fn parse_pool_score_rules(pool_advanced: &Map<String, Value>) -> PoolMemberScore
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
match preset {
"free_first" | "team_first" | "plus_first" | "pro_first" => {
"free_team_first" | "free_first" | "team_first" | "plus_first" | "pro_first" => {
let default_mode = match preset {
"free_team_first" => "both",
"free_first" => "free_only",
"team_first" => "team_only",
"plus_first" => "plus_only",
@@ -180,6 +182,9 @@ fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase())
.filter(|value| match preset {
"free_team_first" => {
matches!(value.as_str(), "free_only" | "team_only" | "both")
}
"free_first" => value == "free_only",
"team_first" => value == "team_only",
"plus_first" => value == "plus_only",
@@ -786,7 +791,7 @@ mod tests {
}
#[test]
fn retired_free_team_first_preset_is_rejected() {
fn legacy_free_team_first_preset_is_preserved() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
@@ -797,8 +802,11 @@ mod tests {
.expect("pool config should parse");
assert_eq!(config.scheduling_presets.len(), 1);
assert_eq!(config.scheduling_presets[0].preset, "lru");
assert_eq!(config.scheduling_presets[0].mode, None);
assert_eq!(config.scheduling_presets[0].preset, "free_team_first");
assert_eq!(
config.scheduling_presets[0].mode.as_deref(),
Some("team_only")
);
}
#[test]
@@ -65,6 +65,8 @@ use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
};
use aether_pool_core::PoolSchedulingPreset;
use aether_provider_pool::ProviderPoolService;
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use axum::{
body::{to_bytes, Body},
@@ -1106,15 +1108,26 @@ fn provider_query_pool_sort_seed() -> String {
fn provider_query_ai_pool_scheduling_config(
config: &AdminProviderPoolConfig,
provider_type: &str,
) -> AiPoolSchedulingConfig {
let presets = config
.scheduling_presets
.iter()
.map(|preset| PoolSchedulingPreset {
preset: preset.preset.clone(),
enabled: preset.enabled,
mode: preset.mode.clone(),
})
.collect::<Vec<_>>();
let normalized_presets = ProviderPoolService::with_builtin_adapters()
.normalize_scheduling_presets(provider_type, &presets);
AiPoolSchedulingConfig {
scheduling_presets: config
.scheduling_presets
.iter()
scheduling_presets: normalized_presets
.into_iter()
.map(|preset| AiPoolSchedulingPreset {
preset: preset.preset.clone(),
preset: preset.preset,
enabled: preset.enabled,
mode: preset.mode.clone(),
mode: preset.mode,
})
.collect(),
lru_enabled: config.lru_enabled,
@@ -1324,7 +1337,8 @@ async fn provider_query_apply_pool_scheduler_to_test_candidates(
provider.id.clone(),
provider_query_ai_pool_runtime_state(&runtime),
);
let pool_config = provider_query_ai_pool_scheduling_config(pool_config);
let pool_config =
provider_query_ai_pool_scheduling_config(pool_config, provider.provider_type.as_str());
let inputs = keys
.into_iter()
.map(|key| {
@@ -320,6 +320,31 @@ fn provider_query_model_test_empty_selected_key_ids_keep_default_selection() {
assert!(provider_query_extract_api_key_ids(&json!({ "api_key_ids": [] })).is_none());
}
#[test]
fn provider_query_codex_pool_config_injects_recent_refresh() {
let raw_config = json!({
"pool_advanced": {
"scheduling_presets": [{
"preset": "cache_affinity",
"enabled": true
}]
}
});
let config = admin_provider_pool_config_from_config_value(Some(&raw_config))
.expect("pool config should parse");
let normalized = provider_query_ai_pool_scheduling_config(&config, "codex");
assert_eq!(
normalized
.scheduling_presets
.iter()
.map(|preset| preset.preset.as_str())
.collect::<Vec<_>>(),
["cache_affinity", "recent_refresh"]
);
}
#[test]
fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() {
assert!(provider_query_resolve_standard_test_upstream_is_stream(