i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 18:59:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(routing): harden routed pool scheduling

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-routing-core/src/lib.rs
+3 -1
View File
@@ -33,4 +33,6 @@ pub use trace::{
RoutingCandidateTrace, RoutingDecisionTrace, RoutingPatchSummary, RoutingPoolExpansionTrace,
RoutingRuntimeFacts,
};
pub use validation::{validate_routing_group_config, RoutingValidationError};
pub use validation::{
validate_routing_group_config, RoutingValidationError, MAX_ROUTING_ALLOWED_KEYS,
};
crates/aether-routing-core/src/policy.rs
+4
View File
@@ -11,6 +11,7 @@ use crate::conditions::RoutingConditionContext;
use crate::model::{RoutingGroupConfig, RoutingModelPolicy, RoutingPoolPolicyOverride};
use crate::mutations::{validate_header_patch, validate_json_patch_operations, MutationPlan};
use crate::ranking::RankingOverlay;
use crate::validation::validate_routing_group_config;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingPolicyError {
@@ -68,6 +69,9 @@ pub fn resolve_routing_policy(
config: &RoutingGroupConfig,
input: RoutingPolicyInput<'_>,
) -> Result<ResolvedRoutingPolicy, RoutingPolicyError> {
validate_routing_group_config(config)
.map_err(|error| RoutingPolicyError::InvalidConfig(error.to_string()))?;
if !model_allowed(&config.allowed_models, input.requested_model)
&& !model_allowed(&config.allowed_models, input.resolved_model)
{
crates/aether-routing-core/src/ranking.rs
+17 -7
View File
@@ -40,6 +40,13 @@ impl RankingOverlay {
.unwrap_or(fallback)
}
pub fn pool_priority(&self, provider_id: &str, fallback: i32) -> i32 {
self.pool_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(fallback)
}
pub fn provider_priority_or_unspecified(&self, provider_id: &str) -> i32 {
self.provider_priority_overrides
.get(provider_id)
@@ -100,15 +107,18 @@ pub fn rank_vector_for_candidate(
) -> RoutingCandidateRankVector {
RoutingCandidateRankVector {
provider_priority_before: facts.provider_priority,
provider_priority_after: overlay.provider_priority_or_unspecified(&facts.provider_id),
provider_priority_after: overlay
.provider_priority(&facts.provider_id, facts.provider_priority),
key_priority_before: facts.key_priority,
key_priority_after: match facts.candidate_kind {
CandidateKind::Provider => facts
.key_id
.as_deref()
.map(|key_id| overlay.key_priority_or_unspecified(key_id))
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED),
CandidateKind::PoolGroup => overlay.pool_priority_or_unspecified(&facts.provider_id),
.map(|key_id| overlay.key_priority(key_id, facts.key_priority))
.unwrap_or(facts.key_priority),
CandidateKind::PoolGroup => {
overlay.pool_priority(&facts.provider_id, facts.key_priority)
}
},
}
}
@@ -142,7 +152,7 @@ mod tests {
}
#[test]
fn rank_vector_marks_missing_routing_priorities_unspecified() {
fn rank_vector_falls_back_to_existing_priorities() {
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
@@ -154,8 +164,8 @@ mod tests {
};
let vector = rank_vector_for_candidate(&RankingOverlay::default(), &facts);
assert_eq!(vector.provider_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
assert_eq!(vector.key_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
assert_eq!(vector.provider_priority_after, 10);
assert_eq!(vector.key_priority_after, 20);
}
#[test]
crates/aether-routing-core/src/validation.rs
+440 -2
View File
@@ -2,9 +2,29 @@ use std::collections::BTreeSet;
use thiserror::Error;
use crate::model::RoutingGroupConfig;
use crate::model::{RoutingGroupConfig, RoutingPoolPolicyOverride};
use crate::mutations::{validate_header_patch, validate_json_patch_operations};
use crate::RoutingAction;
use crate::{RoutingAction, RoutingRulePhase};
pub const MAX_ROUTING_ALLOWED_KEYS: usize = 512;
const ROUTING_POOL_PRESETS: &[&str] = &[
"lru",
"cache_affinity",
"load_balance",
"single_account",
"priority_first",
"free_team_first",
"free_first",
"team_first",
"plus_first",
"pro_first",
"health_first",
"latency_first",
"cost_first",
"quota_balanced",
"recent_refresh",
];
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingValidationError {
@@ -16,6 +36,34 @@ pub enum RoutingValidationError {
EmptyModelSelector,
#[error("invalid mutation action: {0}")]
InvalidMutation(String),
#[error("routing rule {rule_id} uses unsupported {action} action in provider_request phase")]
ProviderRequestActionNotAllowed {
rule_id: String,
action: &'static str,
},
#[error("routing key selector {selector} contains {count} entries; maximum is {max}")]
TooManyAllowedKeys {
selector: String,
count: usize,
max: usize,
},
#[error("routing pool policy {selector} has an empty provider id")]
EmptyPoolProviderId { selector: String },
#[error("routing pool policy {selector} uses unsupported preset: {preset}")]
UnsupportedPoolPreset { selector: String, preset: String },
#[error("routing pool policy {selector} contains duplicate preset: {preset}")]
DuplicatePoolPreset { selector: String, preset: String },
#[error("routing pool policy {selector} preset {preset} has invalid mode: {mode}")]
InvalidPoolPresetMode {
selector: String,
preset: String,
mode: String,
},
#[error("routing pool policy {selector} enables mutually exclusive distribution presets: {presets:?}")]
ConflictingPoolDistributionPresets {
selector: String,
presets: Vec<String>,
},
}
pub fn validate_routing_group_config(
@@ -26,6 +74,22 @@ pub fn validate_routing_group_config(
if model_policy.model.trim().is_empty() {
return Err(RoutingValidationError::EmptyModelSelector);
}
validate_allowed_key_count(
format!("model:{}", model_policy.model.trim()),
model_policy.allowed_keys.len(),
)?;
for (provider_id, override_policy) in &model_policy.pool_policy_overrides {
let model_selector = format!("model:{}", model_policy.model.trim());
if provider_id.trim().is_empty() {
return Err(RoutingValidationError::EmptyPoolProviderId {
selector: model_selector,
});
}
validate_pool_policy_override(
format!("{model_selector}:provider:{}", provider_id.trim()),
override_policy,
)?;
}
}
for rule in &config.rules {
if rule.id.trim().is_empty() {
@@ -35,6 +99,17 @@ pub fn validate_routing_group_config(
return Err(RoutingValidationError::DuplicateRuleId(rule.id.clone()));
}
for action in &rule.actions {
if rule.phase == RoutingRulePhase::ProviderRequest
&& !matches!(
action,
RoutingAction::JsonPatchBody { .. } | RoutingAction::PatchHeaders { .. }
)
{
return Err(RoutingValidationError::ProviderRequestActionNotAllowed {
rule_id: rule.id.clone(),
action: routing_action_name(action),
});
}
match action {
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch).map_err(|error| {
@@ -46,9 +121,372 @@ pub fn validate_routing_group_config(
RoutingValidationError::InvalidMutation(error.to_string())
})?;
}
RoutingAction::RestrictKeys { key_ids } => {
validate_allowed_key_count(format!("rule:{}", rule.id), key_ids.len())?;
}
_ => {}
}
}
}
Ok(())
}
fn validate_pool_policy_override(
selector: String,
override_policy: &RoutingPoolPolicyOverride,
) -> Result<(), RoutingValidationError> {
let mut seen = BTreeSet::new();
let mut enabled_distribution_presets = Vec::new();
for preset in &override_policy.scheduling_presets {
let normalized = preset.preset.trim().to_ascii_lowercase();
if !ROUTING_POOL_PRESETS.contains(&normalized.as_str()) {
return Err(RoutingValidationError::UnsupportedPoolPreset {
selector,
preset: normalized,
});
}
if !seen.insert(normalized.clone()) {
return Err(RoutingValidationError::DuplicatePoolPreset {
selector,
preset: normalized,
});
}
if let Some(mode) = preset.mode.as_deref() {
let mode = mode.trim().to_ascii_lowercase();
if !routing_pool_preset_mode_valid(&normalized, &mode) {
return Err(RoutingValidationError::InvalidPoolPresetMode {
selector,
preset: normalized,
mode,
});
}
}
if preset.enabled && routing_pool_distribution_preset(&normalized) {
enabled_distribution_presets.push(normalized);
}
}
if enabled_distribution_presets.len() > 1 {
return Err(RoutingValidationError::ConflictingPoolDistributionPresets {
selector,
presets: enabled_distribution_presets,
});
}
Ok(())
}
fn routing_pool_preset_mode_valid(preset: &str, mode: &str) -> bool {
match preset {
"free_team_first" => matches!(mode, "free_only" | "team_only" | "both"),
"free_first" => mode == "free_only",
"team_first" => mode == "team_only",
"plus_first" => mode == "plus_only",
"pro_first" => mode == "pro_only",
_ => false,
}
}
fn routing_pool_distribution_preset(preset: &str) -> bool {
matches!(
preset,
"lru" | "cache_affinity" | "load_balance" | "single_account"
)
}
fn validate_allowed_key_count(
selector: String,
count: usize,
) -> Result<(), RoutingValidationError> {
if count > MAX_ROUTING_ALLOWED_KEYS {
return Err(RoutingValidationError::TooManyAllowedKeys {
selector,
count,
max: MAX_ROUTING_ALLOWED_KEYS,
});
}
Ok(())
}
fn routing_action_name(action: &RoutingAction) -> &'static str {
match action {
RoutingAction::RestrictModels { .. } => "restrict_models",
RoutingAction::RestrictProviders { .. } => "restrict_providers",
RoutingAction::RestrictKeys { .. } => "restrict_keys",
RoutingAction::SetScheduling { .. } => "set_scheduling",
RoutingAction::SetProviderPriority { .. } => "set_provider_priority",
RoutingAction::SetKeyPriority { .. } => "set_key_priority",
RoutingAction::JsonPatchBody { .. } => "json_patch_body",
RoutingAction::PatchHeaders { .. } => "patch_headers",
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::{
RoutingCondition, RoutingGroupConfig, RoutingHeaderPatch, RoutingJsonPatchOperation,
RoutingModelPolicy, RoutingRule,
};
use super::*;
fn provider_request_config(actions: Vec<RoutingAction>) -> RoutingGroupConfig {
RoutingGroupConfig {
rules: vec![RoutingRule {
id: "provider-rule".to_string(),
priority: 0,
enabled: true,
phase: RoutingRulePhase::ProviderRequest,
conditions: RoutingCondition::default(),
actions,
stop_processing: false,
}],
..RoutingGroupConfig::default()
}
}
#[test]
fn rejects_non_mutation_actions_in_provider_request_phase() {
let actions = [
(
RoutingAction::RestrictModels {
models: vec!["gpt-5".to_string()],
},
"restrict_models",
),
(
RoutingAction::RestrictProviders {
provider_ids: vec!["provider-1".to_string()],
},
"restrict_providers",
),
(
RoutingAction::RestrictKeys {
key_ids: vec!["key-1".to_string()],
},
"restrict_keys",
),
(
RoutingAction::SetScheduling {
priority_mode: None,
scheduling_mode: None,
keep_priority_on_conversion: Some(true),
},
"set_scheduling",
),
(
RoutingAction::SetProviderPriority {
provider_id: "provider-1".to_string(),
priority: 1,
},
"set_provider_priority",
),
(
RoutingAction::SetKeyPriority {
key_id: "key-1".to_string(),
priority: 1,
},
"set_key_priority",
),
];
for (action, action_name) in actions {
let error = validate_routing_group_config(&provider_request_config(vec![action]))
.expect_err("provider_request must only accept mutations");
assert_eq!(
error,
RoutingValidationError::ProviderRequestActionNotAllowed {
rule_id: "provider-rule".to_string(),
action: action_name,
}
);
}
}
#[test]
fn accepts_mutation_actions_in_provider_request_phase() {
let config = provider_request_config(vec![
RoutingAction::JsonPatchBody {
patch: vec![RoutingJsonPatchOperation::Add {
path: "/metadata/routed".to_string(),
value: json!(true),
}],
},
RoutingAction::PatchHeaders {
patch: vec![RoutingHeaderPatch::Set {
name: "x-routing-profile".to_string(),
value: "provider-rule".to_string(),
}],
},
]);
validate_routing_group_config(&config)
.expect("provider_request mutation actions should remain valid");
}
#[test]
fn rejects_oversized_model_allowed_key_selector() {
let config = RoutingGroupConfig {
model_policies: vec![RoutingModelPolicy {
model: "gpt-5".to_string(),
allowed_keys: key_ids(MAX_ROUTING_ALLOWED_KEYS + 1),
..RoutingModelPolicy::default()
}],
..RoutingGroupConfig::default()
};
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::TooManyAllowedKeys {
selector: "model:gpt-5".to_string(),
count: MAX_ROUTING_ALLOWED_KEYS + 1,
max: MAX_ROUTING_ALLOWED_KEYS,
})
);
}
#[test]
fn rejects_oversized_rule_allowed_key_selector() {
let config = RoutingGroupConfig {
rules: vec![RoutingRule {
id: "restrict-keys".to_string(),
priority: 0,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::default(),
actions: vec![RoutingAction::RestrictKeys {
key_ids: key_ids(MAX_ROUTING_ALLOWED_KEYS + 1),
}],
stop_processing: false,
}],
..RoutingGroupConfig::default()
};
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::TooManyAllowedKeys {
selector: "rule:restrict-keys".to_string(),
count: MAX_ROUTING_ALLOWED_KEYS + 1,
max: MAX_ROUTING_ALLOWED_KEYS,
})
);
}
#[test]
fn accepts_allowed_key_selectors_at_the_limit() {
let config = RoutingGroupConfig {
model_policies: vec![RoutingModelPolicy {
model: "gpt-5".to_string(),
allowed_keys: key_ids(MAX_ROUTING_ALLOWED_KEYS),
..RoutingModelPolicy::default()
}],
..RoutingGroupConfig::default()
};
validate_routing_group_config(&config)
.expect("allowed key selectors at the scan limit should remain valid");
}
#[test]
fn rejects_unknown_pool_override_preset() {
let config = pool_override_config(vec![pool_preset("typo_priority", true, None)]);
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::UnsupportedPoolPreset {
selector: "model:gpt-5:provider:provider-1".to_string(),
preset: "typo_priority".to_string(),
})
);
}
#[test]
fn rejects_duplicate_pool_override_presets() {
let config = pool_override_config(vec![
pool_preset("health_first", true, None),
pool_preset(" HEALTH_FIRST ", false, None),
]);
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::DuplicatePoolPreset {
selector: "model:gpt-5:provider:provider-1".to_string(),
preset: "health_first".to_string(),
})
);
}
#[test]
fn rejects_invalid_pool_override_mode() {
let config =
pool_override_config(vec![pool_preset("free_team_first", true, Some("pro_only"))]);
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::InvalidPoolPresetMode {
selector: "model:gpt-5:provider:provider-1".to_string(),
preset: "free_team_first".to_string(),
mode: "pro_only".to_string(),
})
);
}
#[test]
fn rejects_conflicting_pool_override_distribution_modes() {
let config = pool_override_config(vec![
pool_preset("lru", true, None),
pool_preset("cache_affinity", true, None),
]);
assert_eq!(
validate_routing_group_config(&config),
Err(RoutingValidationError::ConflictingPoolDistributionPresets {
selector: "model:gpt-5:provider:provider-1".to_string(),
presets: vec!["lru".to_string(), "cache_affinity".to_string()],
})
);
}
#[test]
fn accepts_valid_pool_override_modes_and_disabled_alternatives() {
let config = pool_override_config(vec![
pool_preset("cache_affinity", true, None),
pool_preset("lru", false, None),
pool_preset("free_team_first", true, Some("team_only")),
]);
validate_routing_group_config(&config).expect("valid pool override should pass");
}
fn pool_override_config(
scheduling_presets: Vec<crate::RoutingSchedulingPreset>,
) -> RoutingGroupConfig {
RoutingGroupConfig {
model_policies: vec![RoutingModelPolicy {
model: "gpt-5".to_string(),
pool_policy_overrides: std::collections::BTreeMap::from([(
"provider-1".to_string(),
RoutingPoolPolicyOverride { scheduling_presets },
)]),
..RoutingModelPolicy::default()
}],
..RoutingGroupConfig::default()
}
}
fn pool_preset(
preset: &str,
enabled: bool,
mode: Option<&str>,
) -> crate::RoutingSchedulingPreset {
crate::RoutingSchedulingPreset {
preset: preset.to_string(),
enabled,
mode: mode.map(str::to_string),
}
}
fn key_ids(count: usize) -> Vec<String> {
(0..count).map(|index| format!("key-{index}")).collect()
}
}