mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
fix(routing): harden routed pool scheduling
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user