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

Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota

Browse Source
This commit is contained in:
ZheFox
2026-09-03 10:46:49 +08:00
parent 45c840b8d3 cae9aa4134
commit f822df6cce
57 changed files with 1850 additions and 696 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-ai/serving/src/attempt_loop.rs
+88 -2
View File
@@ -13,8 +13,23 @@ pub trait AiExecutionAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
None
}
/// Re-issue this attempt against the same key as a fresh attempt with the
/// given retry index and candidate id. Attempt types that cannot be
/// re-issued return `None`, which disables same-key retries for them.
fn with_same_key_retry(&self, _retry_index: u32, _candidate_id: String) -> Option<Self>
where
Self: Sized,
{
None
}
}
/// Report-context field carrying the routing policy's sticky-key attempt
/// budget for the request, so the attempt loop can derive same-key retries
/// lazily instead of pre-materializing them.
pub const STICKY_KEY_ATTEMPTS_REPORT_FIELD: &str = "sticky_key_attempts";
#[derive(Debug)]
pub enum AiAttemptLoopOutcome<Response, Exhaustion> {
Responded(Response),
@@ -83,6 +98,17 @@ where
Ok(())
}
/// After `attempt` failed with candidate scope, return the next attempt on
/// the same key, or `None` once the sticky-key budget is used up. Retries
/// are derived here on demand so no attempt is materialized before it is
/// actually needed.
async fn next_same_key_retry(
&self,
_attempt: &Attempt,
) -> Result<Option<Attempt>, Self::Error> {
Ok(None)
}
async fn mark_unused_attempts(&self, attempts: Vec<Attempt>) -> Result<(), Self::Error>;
async fn build_exhaustion(
@@ -101,11 +127,15 @@ where
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
let mut remaining = attempts.into_iter();
let mut pending_same_key_retry: Option<Attempt> = None;
let mut last_attempted = None;
let mut retry_filters: Vec<AiAttemptRetryFilter> = Vec::new();
let mut fallback_response = None;
while let Some(attempt) = remaining.next() {
loop {
let Some(attempt) = pending_same_key_retry.take().or_else(|| remaining.next()) else {
break;
};
if retry_filters.iter().any(|filter| filter.matches(&attempt))
|| port.should_skip_attempt(&attempt).await?
{
@@ -133,7 +163,9 @@ where
if attempt_fallback_response.is_some() {
fallback_response = attempt_fallback_response;
}
if scope != AiAttemptRetryScope::Candidate {
if scope == AiAttemptRetryScope::Candidate {
pending_same_key_retry = port.next_same_key_retry(&attempt).await?;
} else {
retry_filters.push(AiAttemptRetryFilter::new(&attempt, scope));
}
}
@@ -188,6 +220,32 @@ impl AiAttemptRetryFilter {
}
}
/// Clone `plan`/`report_context` for a same-key retry: only the candidate id
/// and retry index change, everything else (url, headers, body) is reused.
fn same_key_retry_parts(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
retry_index: u32,
candidate_id: String,
) -> (aether_contracts::ExecutionPlan, Option<serde_json::Value>) {
let mut plan = plan.clone();
plan.candidate_id = Some(candidate_id.clone());
let report_context = report_context.cloned().map(|mut value| {
if let Some(object) = value.as_object_mut() {
object.insert(
"candidate_id".to_string(),
serde_json::Value::String(candidate_id),
);
object.insert(
"retry_index".to_string(),
serde_json::Value::Number(retry_index.into()),
);
}
value
});
(plan, report_context)
}
impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
fn execution_plan(&self) -> &aether_contracts::ExecutionPlan {
&self.plan
@@ -204,6 +262,20 @@ impl AiExecutionAttempt for crate::dto::AiSyncAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
self.report_context.as_ref()
}
fn with_same_key_retry(&self, retry_index: u32, candidate_id: String) -> Option<Self> {
let (plan, report_context) = same_key_retry_parts(
&self.plan,
self.report_context.as_ref(),
retry_index,
candidate_id,
);
Some(Self {
plan,
report_kind: self.report_kind.clone(),
report_context,
})
}
}
impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
@@ -222,6 +294,20 @@ impl AiExecutionAttempt for crate::dto::AiStreamAttempt {
fn report_context_ref(&self) -> Option<&serde_json::Value> {
self.report_context.as_ref()
}
fn with_same_key_retry(&self, retry_index: u32, candidate_id: String) -> Option<Self> {
let (plan, report_context) = same_key_retry_parts(
&self.plan,
self.report_context.as_ref(),
retry_index,
candidate_id,
);
Some(Self {
plan,
report_kind: self.report_kind.clone(),
report_context,
})
}
}
#[cfg(test)]
crates/aether-ai/serving/src/candidate_persistence.rs
+22 -60
View File
@@ -9,8 +9,6 @@ pub trait AiAvailableCandidatePersistencePort: Send + Sync {
type ExtraData: Clone + Send + Sync;
type Error: Send;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32;
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData>;
fn generate_candidate_id(&self) -> String;
@@ -42,50 +40,28 @@ pub async fn run_ai_available_candidate_persistence<Port>(
where
Port: AiAvailableCandidatePersistencePort,
{
let total_attempts = candidates
.iter()
.map(|candidate| port.attempt_slot_count(candidate) as usize)
.sum();
let mut materialized = Vec::with_capacity(total_attempts);
// One attempt per candidate. Same-key retries are derived lazily by the
// attempt loop (`AiAttemptLoopPort::next_same_key_retry`) after a failure,
// so the sticky-key budget never inflates up-front materialization.
let mut materialized = Vec::with_capacity(candidates.len());
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_index = candidate_index as u32;
let attempt_slots = port.attempt_slot_count(&candidate).max(1);
let extra_data = port.build_extra_data(&candidate);
let mut owned_candidate = Some(candidate);
for retry_index in 0..attempt_slots {
let candidate = owned_candidate
.as_ref()
.expect("candidate should remain available until final retry");
let generated_candidate_id = port.generate_candidate_id();
let candidate_id = if port.should_persist_available_candidate(candidate) {
port.persist_available_candidate(
candidate,
candidate_index,
retry_index,
generated_candidate_id.as_str(),
extra_data.clone(),
)
.await?
} else {
generated_candidate_id
};
let candidate = if retry_index + 1 == attempt_slots {
owned_candidate
.take()
.expect("final retry should consume owned candidate")
} else {
candidate.clone()
};
materialized.push(port.build_attempt(
candidate,
let generated_candidate_id = port.generate_candidate_id();
let candidate_id = if port.should_persist_available_candidate(&candidate) {
port.persist_available_candidate(
&candidate,
candidate_index,
retry_index,
candidate_id,
));
}
0,
generated_candidate_id.as_str(),
extra_data,
)
.await?
} else {
generated_candidate_id
};
materialized.push(port.build_attempt(candidate, candidate_index, 0, candidate_id));
}
Ok(materialized)
@@ -177,7 +153,6 @@ mod tests {
#[derive(Debug, Clone, PartialEq, Eq)]
struct TestCandidate {
id: &'static str,
attempt_slots: u32,
persist: bool,
}
@@ -216,10 +191,6 @@ mod tests {
type ExtraData = String;
type Error = std::convert::Infallible;
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
candidate.attempt_slots
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
Some(format!("extra:{}", candidate.id))
}
@@ -299,7 +270,7 @@ mod tests {
}
#[tokio::test]
async fn available_persistence_expands_candidates_into_retry_attempts() {
async fn available_persistence_materializes_one_attempt_per_candidate() {
let port = TestPort::default();
let attempts = run_ai_available_candidate_persistence(
@@ -307,12 +278,10 @@ mod tests {
vec![
TestCandidate {
id: "a",
attempt_slots: 2,
persist: true,
},
TestCandidate {
id: "b",
attempt_slots: 1,
persist: false,
},
],
@@ -320,6 +289,8 @@ mod tests {
.await
.unwrap();
// Same-key retries are never pre-materialized; the attempt loop
// derives them on demand after a failure.
assert_eq!(
attempts,
[
@@ -329,26 +300,17 @@ mod tests {
retry_index: 0,
candidate_id: "stored-candidate-1".to_string(),
},
TestAttempt {
id: "a",
candidate_index: 0,
retry_index: 1,
candidate_id: "stored-candidate-2".to_string(),
},
TestAttempt {
id: "b",
candidate_index: 1,
retry_index: 0,
candidate_id: "candidate-3".to_string(),
candidate_id: "candidate-2".to_string(),
},
]
);
assert_eq!(
port.calls.lock().unwrap().as_slice(),
[
"available:a:0:0:candidate-1:extra:a",
"available:a:0:1:candidate-2:extra:a",
]
["available:a:0:0:candidate-1:extra:a"]
);
}
crates/aether-ai/serving/src/lib.rs
+1 -1
View File
@@ -54,7 +54,7 @@ pub use aether_pool_core::{
};
pub use attempt_loop::{
run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopOutcome, AiAttemptLoopPort,
AiAttemptRetryScope, AiExecutionAttempt,
AiAttemptRetryScope, AiExecutionAttempt, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
};
pub use attempt_plan::{
build_ai_execution_decision_from_plan, build_ai_execution_plan_from_decision,
crates/aether-routing-core/src/actions.rs
+6
View File
@@ -73,6 +73,8 @@ pub enum RoutingAction {
priority_mode: Option<RoutingSetPriorityMode>,
scheduling_mode: Option<RoutingSchedulingMode>,
keep_priority_on_conversion: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
sticky_key_attempts: Option<u32>,
},
SetProviderPriority {
provider_id: String,
@@ -81,6 +83,10 @@ pub enum RoutingAction {
SetKeyPriority {
key_id: String,
priority: i32,
/// When set, the override only applies to candidates served through
/// this API format; otherwise it applies to the key on every format.
#[serde(default, skip_serializing_if = "Option::is_none")]
api_format: Option<String>,
},
JsonPatchBody {
patch: Vec<RoutingJsonPatchOperation>,
crates/aether-routing-core/src/lib.rs
+3 -3
View File
@@ -13,9 +13,9 @@ pub use actions::{
};
pub use conditions::{RoutingCondition, RoutingConditionContext, RoutingConditionOp};
pub use model::{
RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord,
RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride, RoutingRule,
RoutingSchedulingPreset,
RoutingDefaultPolicy, RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig,
RoutingGroupRecord, RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride,
RoutingRule, RoutingSchedulingPreset, DEFAULT_STICKY_KEY_ATTEMPTS,
};
pub use mutations::{
apply_json_patch_operations, validate_header_patch, validate_json_patch_operations,
crates/aether-routing-core/src/model.rs
+32 -1
View File
@@ -23,7 +23,11 @@ pub struct RoutingPoolPolicyOverride {
pub scheduling_presets: Vec<RoutingSchedulingPreset>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
/// Default number of attempts on the first-ranked (sticky) candidate before
/// failing over: one retry on the same key.
pub const DEFAULT_STICKY_KEY_ATTEMPTS: u32 = 2;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingDefaultPolicy {
#[serde(default)]
pub priority_mode: RoutingSetPriorityMode,
@@ -31,6 +35,26 @@ pub struct RoutingDefaultPolicy {
pub scheduling_mode: RoutingSchedulingMode,
#[serde(default)]
pub keep_priority_on_conversion: bool,
/// Total attempts on the first-ranked candidate before moving on. Later
/// candidates always get a single attempt so failover keeps advancing.
/// `0` and `1` both mean no same-key retry.
#[serde(default = "default_sticky_key_attempts")]
pub sticky_key_attempts: u32,
}
impl Default for RoutingDefaultPolicy {
fn default() -> Self {
Self {
priority_mode: RoutingSetPriorityMode::default(),
scheduling_mode: RoutingSchedulingMode::default(),
keep_priority_on_conversion: false,
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
}
}
}
fn default_sticky_key_attempts() -> u32 {
DEFAULT_STICKY_KEY_ATTEMPTS
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
@@ -44,6 +68,13 @@ pub struct RoutingModelPolicy {
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
/// Key priority overrides scoped to one API format: `api_format -> key_id -> priority`.
///
/// A key can serve several API formats and legacy `global_priority_by_format`
/// ranks it independently per format. Entries here take precedence over
/// `key_priority_overrides` when the candidate format matches.
#[serde(default)]
pub key_priority_overrides_by_format: BTreeMap<String, BTreeMap<String, i32>>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
crates/aether-routing-core/src/policy.rs
+113 -6
View File
@@ -57,6 +57,9 @@ pub struct ResolvedRoutingPolicy {
pub priority_mode: RoutingSetPriorityMode,
pub scheduling_mode: RoutingSchedulingMode,
pub keep_priority_on_conversion: bool,
/// See `RoutingDefaultPolicy::sticky_key_attempts`.
#[serde(default = "default_sticky_key_attempts")]
pub sticky_key_attempts: u32,
pub ranking_overlay: RankingOverlay,
pub mutation_plan: MutationPlan,
#[serde(default)]
@@ -89,6 +92,7 @@ pub fn resolve_routing_policy(
priority_mode: config.default_policy.priority_mode,
scheduling_mode: config.default_policy.scheduling_mode,
keep_priority_on_conversion: config.default_policy.keep_priority_on_conversion,
sticky_key_attempts: config.default_policy.sticky_key_attempts,
ranking_overlay: RankingOverlay::default(),
mutation_plan: MutationPlan::default(),
pool_policy_overrides: BTreeMap::new(),
@@ -163,6 +167,13 @@ fn apply_model_policy(policy: &mut ResolvedRoutingPolicy, model_policy: &Routing
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
for (api_format, overrides) in &model_policy.key_priority_overrides_by_format {
for (key_id, priority) in overrides {
policy
.ranking_overlay
.insert_key_priority_override_for_format(api_format, key_id.clone(), *priority);
}
}
policy.ranking_overlay.pool_priority_overrides.extend(
model_policy
.pool_priority_overrides
@@ -198,6 +209,7 @@ fn apply_action(
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
sticky_key_attempts,
} => {
if let Some(priority_mode) = priority_mode {
policy.priority_mode = *priority_mode;
@@ -208,6 +220,9 @@ fn apply_action(
if let Some(keep_priority_on_conversion) = keep_priority_on_conversion {
policy.keep_priority_on_conversion = *keep_priority_on_conversion;
}
if let Some(sticky_key_attempts) = sticky_key_attempts {
policy.sticky_key_attempts = *sticky_key_attempts;
}
}
RoutingAction::SetProviderPriority {
provider_id,
@@ -218,12 +233,27 @@ fn apply_action(
.provider_priority_overrides
.insert(provider_id.clone(), *priority);
}
RoutingAction::SetKeyPriority { key_id, priority } => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
RoutingAction::SetKeyPriority {
key_id,
priority,
api_format,
} => match api_format
.as_deref()
.map(str::trim)
.filter(|f| !f.is_empty())
{
Some(api_format) => {
policy
.ranking_overlay
.insert_key_priority_override_for_format(api_format, key_id.clone(), *priority);
}
None => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
},
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
@@ -260,6 +290,10 @@ fn model_allowed(patterns: &[String], requested_model: &str) -> bool {
.any(|pattern| model_pattern_matches(pattern, requested_model))
}
fn default_sticky_key_attempts() -> u32 {
crate::model::DEFAULT_STICKY_KEY_ATTEMPTS
}
fn model_pattern_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.trim();
if pattern == "*" {
@@ -362,6 +396,7 @@ mod tests {
priority_mode: RoutingSetPriorityMode::GlobalKey,
scheduling_mode: RoutingSchedulingMode::LoadBalance,
keep_priority_on_conversion: true,
sticky_key_attempts: 3,
},
model_policies: vec![RoutingModelPolicy {
model: "special-model".to_string(),
@@ -393,6 +428,7 @@ mod tests {
assert_eq!(special.priority_mode, RoutingSetPriorityMode::GlobalKey);
assert_eq!(special.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(special.keep_priority_on_conversion);
assert_eq!(special.sticky_key_attempts, 3);
assert_eq!(
special.ranking_overlay.allowed_providers,
vec!["provider-special"]
@@ -426,6 +462,7 @@ mod tests {
assert_eq!(ordinary.priority_mode, RoutingSetPriorityMode::GlobalKey);
assert_eq!(ordinary.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(ordinary.keep_priority_on_conversion);
assert_eq!(ordinary.sticky_key_attempts, 3);
assert!(ordinary.ranking_overlay.allowed_providers.is_empty());
assert!(ordinary.ranking_overlay.allowed_keys.is_empty());
assert!(ordinary
@@ -434,6 +471,76 @@ mod tests {
.is_empty());
}
#[test]
fn sticky_key_attempts_defaults_to_two_and_can_be_overridden_by_rule() {
let default_config = RoutingGroupConfig::default();
let default_policy = resolve_routing_policy(
&default_config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("default config should resolve");
assert_eq!(
default_policy.sticky_key_attempts,
crate::DEFAULT_STICKY_KEY_ATTEMPTS
);
let parsed: RoutingGroupConfig =
serde_json::from_value(json!({ "default_policy": { "priority_mode": "provider" } }))
.expect("legacy config without sticky_key_attempts should deserialize");
assert_eq!(
parsed.default_policy.sticky_key_attempts,
crate::DEFAULT_STICKY_KEY_ATTEMPTS
);
let config = RoutingGroupConfig {
rules: vec![RoutingRule {
id: "no-sticky-retry".to_string(),
priority: 1,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::default(),
actions: vec![RoutingAction::SetScheduling {
priority_mode: None,
scheduling_mode: None,
keep_priority_on_conversion: None,
sticky_key_attempts: Some(1),
}],
stop_processing: false,
}],
..RoutingGroupConfig::default()
};
let policy = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("rule config should resolve");
assert_eq!(policy.sticky_key_attempts, 1);
}
#[test]
fn rejects_disallowed_model() {
let config = RoutingGroupConfig {
crates/aether-routing-core/src/ranking.rs
+96 -1
View File
@@ -21,6 +21,9 @@ pub struct RankingOverlay {
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
/// `api_format -> key_id -> priority`; see `RoutingModelPolicy`.
#[serde(default)]
pub key_priority_overrides_by_format: BTreeMap<String, BTreeMap<String, i32>>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
}
@@ -40,6 +43,46 @@ impl RankingOverlay {
.unwrap_or(fallback)
}
/// Format-scoped key priority: a per-format override wins, then the
/// format-agnostic key override, then `fallback`.
pub fn key_priority_for_format(&self, key_id: &str, api_format: &str, fallback: i32) -> i32 {
self.key_priority_override_for_format(key_id, api_format)
.unwrap_or_else(|| self.key_priority(key_id, fallback))
}
/// Format-scoped key override using exact (case-insensitive) format match.
pub fn key_priority_override_for_format(&self, key_id: &str, api_format: &str) -> Option<i32> {
let api_format = api_format.trim();
self.key_priority_override_matching_format(key_id, |format| {
format.trim().eq_ignore_ascii_case(api_format)
})
}
/// Format-scoped key override where the caller decides how configured
/// format names match the candidate format (for alias-aware matching).
pub fn key_priority_override_matching_format(
&self,
key_id: &str,
mut format_matches: impl FnMut(&str) -> bool,
) -> Option<i32> {
self.key_priority_overrides_by_format
.iter()
.find(|(format, _)| format_matches(format))
.and_then(|(_, overrides)| overrides.get(key_id).copied())
}
pub fn insert_key_priority_override_for_format(
&mut self,
api_format: &str,
key_id: String,
priority: i32,
) {
self.key_priority_overrides_by_format
.entry(api_format.trim().to_ascii_lowercase())
.or_default()
.insert(key_id, priority);
}
pub fn pool_priority(&self, provider_id: &str, fallback: i32) -> i32 {
self.pool_priority_overrides
.get(provider_id)
@@ -89,6 +132,9 @@ pub struct RoutingCandidateFacts {
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
/// Candidate API format used to resolve format-scoped key overrides.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_format: Option<String>,
pub provider_priority: i32,
pub key_priority: i32,
}
@@ -114,7 +160,12 @@ pub fn rank_vector_for_candidate(
CandidateKind::Provider => facts
.key_id
.as_deref()
.map(|key_id| overlay.key_priority(key_id, facts.key_priority))
.map(|key_id| match facts.api_format.as_deref() {
Some(api_format) => {
overlay.key_priority_for_format(key_id, api_format, facts.key_priority)
}
None => 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,6 +193,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: None,
provider_priority: 10,
key_priority: 20,
};
@@ -151,6 +203,47 @@ mod tests {
assert_eq!(vector.key_priority_after, 5);
}
#[test]
fn format_scoped_key_override_wins_over_key_override_for_matching_format() {
let mut overlay = RankingOverlay {
key_priority_overrides: BTreeMap::from([("key-a".to_string(), 5)]),
..RankingOverlay::default()
};
overlay.insert_key_priority_override_for_format("openai:chat", "key-a".to_string(), 1);
assert_eq!(
overlay.key_priority_for_format("key-a", "openai:chat", 20),
1
);
assert_eq!(
overlay.key_priority_for_format("key-a", "OpenAI:Chat", 20),
1
);
assert_eq!(
overlay.key_priority_for_format("key-a", "claude:messages", 20),
5
);
assert_eq!(
overlay.key_priority_for_format("key-b", "openai:chat", 20),
20
);
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: Some("openai:chat".to_string()),
provider_priority: 10,
key_priority: 20,
};
assert_eq!(
rank_vector_for_candidate(&overlay, &facts).key_priority_after,
1
);
}
#[test]
fn rank_vector_falls_back_to_existing_priorities() {
let facts = RoutingCandidateFacts {
@@ -159,6 +252,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
api_format: None,
provider_priority: 10,
key_priority: 20,
};
@@ -180,6 +274,7 @@ mod tests {
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: None,
api_format: None,
provider_priority: 10,
key_priority: 20,
};
crates/aether-routing-core/src/validation.rs
+2
View File
@@ -271,6 +271,7 @@ mod tests {
priority_mode: None,
scheduling_mode: None,
keep_priority_on_conversion: Some(true),
sticky_key_attempts: None,
},
"set_scheduling",
),
@@ -285,6 +286,7 @@ mod tests {
RoutingAction::SetKeyPriority {
key_id: "key-1".to_string(),
priority: 1,
api_format: None,
},
"set_key_priority",
),