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

Cap Codex pool cooldowns and add key circuit breaker

Browse Source
This commit is contained in:
fawney19
2026-05-04 02:15:32 +08:00
parent 1d62722d47
commit 099653f732
10 changed files with 536 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
View File

@@ -13,6 +13,6 @@ pub(crate) use self::reads::{
};
pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
};

170
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -12,6 +12,8 @@ use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use uuid::Uuid;
const MAX_POOL_COOLDOWN_SECONDS: u64 = 32 * 60;
const ACCOUNT_DISABLE_PATTERNS: &[&str] = &[
"organization has been disabled",
"organization_disabled",
@@ -68,7 +70,7 @@ fn parse_retry_after_seconds(headers: Option<&BTreeMap<String, String>>) -> Opti
.filter(|value| !value.is_empty())
})?;
let seconds = raw.parse::<u64>().ok()?;
Some(seconds.clamp(1, 3600))
Some(seconds.clamp(1, MAX_POOL_COOLDOWN_SECONDS))
}
fn parse_google_quota_duration_seconds(raw: &serde_json::Value) -> Option<u64> {
@@ -187,6 +189,33 @@ fn extract_error_message(error_body: Option<&str>) -> String {
.unwrap_or_else(|| error_body.chars().take(500).collect())
}
pub(crate) fn admin_provider_pool_key_circuit_breaker_reason(
status_code: u16,
error_body: Option<&str>,
) -> Option<String> {
let error_message = extract_error_message(error_body).to_ascii_lowercase();
match status_code {
401 if ACCOUNT_DISABLE_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern)) =>
{
Some("account_deactivated_401".to_string())
}
402 => Some("payment_required_402".to_string()),
403 if FORBIDDEN_ACCOUNT_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern)) =>
{
Some("forbidden_403".to_string())
}
400 => ACCOUNT_DISABLE_PATTERNS
.iter()
.find(|pattern| error_message.contains(**pattern))
.map(|pattern| format!("account_disabled_400:{pattern}")),
_ => None,
}
}
fn resolve_transient_cooldown_ttl(
status_code: u16,
retry_after_seconds: Option<u64>,
@@ -213,6 +242,7 @@ async fn set_pool_cooldown(
if ttl_seconds == 0 {
return;
}
let ttl_seconds = ttl_seconds.min(MAX_POOL_COOLDOWN_SECONDS);
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else {
warn!(
@@ -393,46 +423,33 @@ pub(crate) async fn record_admin_provider_pool_error(
if status_code == 401 {
invalidate_pool_oauth_cache(runner, key_id).await;
if ACCOUNT_DISABLE_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern))
{
set_pool_cooldown(runner, provider_id, key_id, "account_deactivated_401", 3600).await;
}
return;
}
if status_code == 402 {
set_pool_cooldown(runner, provider_id, key_id, "payment_required_402", 3600).await;
return;
}
if status_code == 403 {
let severe = FORBIDDEN_ACCOUNT_PATTERNS
if FORBIDDEN_ACCOUNT_PATTERNS
.iter()
.any(|pattern| error_message.contains(pattern));
let ttl_seconds = if severe {
3600
} else {
pool_config.rate_limit_cooldown_seconds.max(300)
};
set_pool_cooldown(runner, provider_id, key_id, "forbidden_403", ttl_seconds).await;
.any(|pattern| error_message.contains(pattern))
{
return;
}
set_pool_cooldown(
runner,
provider_id,
key_id,
"forbidden_403",
pool_config.rate_limit_cooldown_seconds.max(300),
)
.await;
return;
}
if status_code == 400 {
if let Some(pattern) = ACCOUNT_DISABLE_PATTERNS
.iter()
.find(|pattern| error_message.contains(**pattern))
{
set_pool_cooldown(
runner,
provider_id,
key_id,
&format!("account_disabled_400:{pattern}"),
3600,
)
.await;
if admin_provider_pool_key_circuit_breaker_reason(status_code, error_body).is_some() {
return;
}
}
@@ -563,8 +580,9 @@ pub(crate) async fn record_admin_provider_pool_stream_timeout(
#[cfg(test)]
mod tests {
use super::{
parse_google_quota_cooldown_seconds_at, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
admin_provider_pool_key_circuit_breaker_reason, parse_google_quota_cooldown_seconds_at,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
};
use crate::data::{GatewayDataConfig, GatewayDataState};
use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state;
@@ -868,6 +886,98 @@ mod tests {
.is_some_and(|ttl| *ttl <= 45 && *ttl >= 30));
}
#[tokio::test]
async fn error_feedback_caps_long_retry_after_cooldowns_at_32_minutes() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_capped_cooldown");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-long-cooldown".to_string()];
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-long-cooldown",
&pool_config,
429,
Some(r#"{"error":{"message":"rate limited"}}"#),
Some(&BTreeMap::from([(
"Retry-After".to_string(),
"3600".to_string(),
)])),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert_eq!(
runtime
.cooldown_reason_by_key
.get("key-long-cooldown")
.map(String::as_str),
Some("rate_limited_429")
);
assert!(runtime
.cooldown_ttl_by_key
.get("key-long-cooldown")
.is_some_and(|ttl| *ttl <= 32 * 60 && *ttl >= 31 * 60));
}
#[tokio::test]
async fn severe_account_errors_use_circuit_breaker_instead_of_pool_cooldown() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_circuit_no_cooldown");
let runner = app.redis_kv_runner().expect("redis runner should exist");
let pool_config = sample_pool_config();
let key_ids = vec!["key-account-disabled".to_string()];
assert_eq!(
admin_provider_pool_key_circuit_breaker_reason(
401,
Some(r#"{"error":{"message":"account has been deactivated"}}"#),
)
.as_deref(),
Some("account_deactivated_401")
);
record_admin_provider_pool_error(
&runner,
"provider-1",
"key-account-disabled",
&pool_config,
401,
Some(r#"{"error":{"message":"account has been deactivated"}}"#),
None,
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
&runner,
"provider-1",
&key_ids,
&pool_config,
None,
)
.await;
assert!(!runtime
.cooldown_reason_by_key
.contains_key("key-account-disabled"));
assert!(!runtime
.cooldown_ttl_by_key
.contains_key("key-account-disabled"));
}
#[tokio::test]
async fn error_feedback_applies_unschedulable_rule_cooldown() {
let Some(redis) = start_managed_redis_or_skip().await else {

5
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -1,7 +1,8 @@
pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
pub(crate) use super::super::admin::provider::pool::runtime::{
read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,

209
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -15,15 +15,16 @@ use tracing::warn;
use super::{
local_failover_error_message, project_local_adaptive_rate_limit,
project_local_adaptive_success, project_local_failure_health, project_local_success_health,
LocalFailoverClassification,
project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification,
};
use crate::ai_serving::extract_pool_sticky_session_token;
use crate::clock::current_unix_secs;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, AdminProviderPoolConfig,
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
AdminProviderPoolConfig,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
@@ -472,12 +473,20 @@ async fn record_health_success_effect(
else {
return;
};
let circuit_breaker_by_format = current_key
.circuit_breaker_by_format
.as_ref()
.and_then(|current| project_local_key_circuit_closed(Some(current), api_format));
let circuit_breaker_update = circuit_breaker_by_format
.as_ref()
.or(current_key.circuit_breaker_by_format.as_ref());
if let Err(err) = state
.update_provider_catalog_key_format_health(
.update_provider_catalog_key_health_state(
&context.plan.key_id,
api_format,
&health_by_format,
current_key.is_active,
Some(&health_by_format),
circuit_breaker_update,
)
.await
{
@@ -516,7 +525,13 @@ async fn record_pool_error_effect(
context: LocalExecutionEffectContext<'_>,
effect: LocalPoolErrorEffect<'_>,
) {
if !local_candidate_failure_should_record_pool_error(effect.classification, effect.status_code)
let circuit_reason =
admin_provider_pool_key_circuit_breaker_reason(effect.status_code, effect.error_body);
if circuit_reason.is_none()
&& !local_candidate_failure_should_record_pool_error(
effect.classification,
effect.status_code,
)
{
return;
}
@@ -525,6 +540,10 @@ async fn record_pool_error_effect(
return;
};
if let Some(reason) = circuit_reason {
open_pool_key_circuit_breaker(state, context, &reason).await;
}
record_admin_provider_pool_error(
&pool_context.runner,
&context.plan.provider_id,
@@ -537,6 +556,49 @@ async fn record_pool_error_effect(
.await;
}
async fn open_pool_key_circuit_breaker(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
reason: &str,
) {
let api_format = context.plan.provider_api_format.trim();
if api_format.is_empty() {
return;
}
let Some(current_key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&context.plan.key_id))
.await
.ok()
.and_then(|mut keys| keys.drain(..).next())
else {
return;
};
let Some(circuit_breaker_by_format) = project_local_key_circuit_open(
current_key.circuit_breaker_by_format.as_ref(),
api_format,
reason,
current_unix_secs(),
) else {
return;
};
if let Err(err) = state
.update_provider_catalog_key_health_state(
&context.plan.key_id,
current_key.is_active,
current_key.health_by_format.as_ref(),
Some(&circuit_breaker_by_format),
)
.await
{
warn!(
"gateway orchestration effects: failed to open pool key circuit for provider {} endpoint {} key {}: {:?}",
context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
);
}
}
async fn record_oauth_invalidation_effect(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
@@ -661,15 +723,16 @@ mod tests {
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_testkit::ManagedRedisServer;
use serde_json::{json, Value};
use super::{
apply_local_execution_effect, local_candidate_failure_should_record_pool_error,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::data::GatewayDataState;
use crate::data::{GatewayDataConfig, GatewayDataState};
use crate::orchestration::LocalFailoverClassification;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
@@ -677,6 +740,17 @@ mod tests {
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
};
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
match ManagedRedisServer::start().await {
Ok(server) => Some(server),
Err(err) if err.to_string().contains("No such file or directory") => {
eprintln!("skipping redis-backed orchestration effect test: {err}");
None
}
Err(err) => panic!("redis server should start: {err}"),
}
}
fn sample_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req-1".to_string(),
@@ -742,7 +816,7 @@ mod tests {
None,
Some(20.0),
None,
None,
Some(json!({"pool_advanced": {}})),
)
}
@@ -813,6 +887,24 @@ mod tests {
)
}
fn codex_state_with_redis(redis_url: &str, redis_key_prefix: &str) -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_codex_provider()],
vec![sample_codex_endpoint()],
vec![sample_codex_key()],
));
let data_state = GatewayDataState::from_config(
GatewayDataConfig::disabled()
.with_redis_url(redis_url, Some(redis_key_prefix))
.with_encryption_key(DEVELOPMENT_ENCRYPTION_KEY),
)
.expect("data state should build")
.attach_provider_catalog_repository_for_tests(repository);
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state)
}
fn sample_health_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"prov-1".to_string(),
@@ -885,6 +977,20 @@ mod tests {
)
}
fn health_state_with_key(key: StoredProviderCatalogKey) -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_health_provider()],
vec![sample_health_endpoint()],
vec![key],
));
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
}
fn sample_adaptive_key() -> StoredProviderCatalogKey {
let mut key = sample_health_key();
key.name = "adaptive".to_string();
@@ -1182,6 +1288,47 @@ mod tests {
));
}
#[tokio::test]
async fn pool_account_error_opens_key_circuit() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let state = codex_state_with_redis(redis.redis_url(), "orchestration_pool_circuit");
let plan = sample_codex_plan();
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: None,
},
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code: 401,
classification: LocalFailoverClassification::StopErrorPattern,
headers: &BTreeMap::new(),
error_body: Some(r#"{"error":{"message":"account has been deactivated"}}"#),
}),
)
.await;
let stored_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load")
.into_iter()
.next()
.expect("stored key should exist");
let circuit = stored_key
.circuit_breaker_by_format
.as_ref()
.and_then(|value| value.get("openai:responses"))
.expect("format circuit should be stored");
assert_eq!(circuit["open"], json!(true));
assert_eq!(circuit["reason"], json!("account_deactivated_401"));
assert!(circuit["next_probe_at"].is_string());
assert!(circuit["next_probe_at_unix_secs"].as_u64().is_some());
}
#[tokio::test]
async fn oauth_invalidation_marks_codex_key_invalid() {
let state = codex_state();
@@ -1343,6 +1490,46 @@ mod tests {
);
}
#[tokio::test]
async fn health_success_projection_closes_key_circuit_for_format() {
let mut key = sample_health_key();
key.circuit_breaker_by_format = Some(json!({
"openai:chat": {
"open": true,
"reason": "account_deactivated_401",
"next_probe_at_unix_secs": 1_760_001_920u64
}
}));
let state = health_state_with_key(key);
let plan = sample_plan();
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: None,
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
let stored_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load")
.into_iter()
.next()
.expect("stored key should exist");
let circuit = stored_key
.circuit_breaker_by_format
.as_ref()
.and_then(|value| value.get("openai:chat"))
.expect("format circuit should be stored");
assert_eq!(circuit["open"], json!(false));
assert_eq!(circuit["reason"], Value::Null);
assert_eq!(circuit["next_probe_at_unix_secs"], Value::Null);
}
#[tokio::test]
async fn adaptive_rate_limit_effect_updates_adaptive_key_observation() {
let state = adaptive_state();

112
apps/aether-gateway/src/orchestration/health.rs
View File

@@ -4,6 +4,7 @@ use super::LocalFailoverClassification;
use crate::handlers::shared::unix_secs_to_rfc3339;
const LOCAL_HEALTH_SCORE_FLOOR: f64 = 0.2;
const LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS: u64 = 32 * 60;
pub(crate) fn project_local_failure_health(
current_health_by_format: Option<&Value>,
@@ -73,6 +74,69 @@ pub(crate) fn project_local_success_health(
Some(Value::Object(health_by_format))
}
pub(crate) fn project_local_key_circuit_open(
current_circuit_by_format: Option<&Value>,
api_format: &str,
reason: &str,
observed_at_unix_secs: u64,
) -> Option<Value> {
let api_format = api_format.trim();
let reason = reason.trim();
if api_format.is_empty() || reason.is_empty() {
return None;
}
let next_probe_at_unix_secs =
observed_at_unix_secs.saturating_add(LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS);
let mut circuit_by_format = current_circuit_by_format
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
circuit_by_format.insert(
api_format.to_string(),
json!({
"open": true,
"open_at": unix_secs_to_rfc3339(observed_at_unix_secs),
"reason": reason,
"next_probe_at": unix_secs_to_rfc3339(next_probe_at_unix_secs),
"next_probe_at_unix_secs": next_probe_at_unix_secs,
"half_open_until": Value::Null,
"half_open_successes": 0,
"half_open_failures": 0,
}),
);
Some(Value::Object(circuit_by_format))
}
pub(crate) fn project_local_key_circuit_closed(
current_circuit_by_format: Option<&Value>,
api_format: &str,
) -> Option<Value> {
let api_format = api_format.trim();
if api_format.is_empty() {
return None;
}
let mut circuit_by_format = current_circuit_by_format
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
circuit_by_format.insert(
api_format.to_string(),
json!({
"open": false,
"open_at": Value::Null,
"reason": Value::Null,
"next_probe_at": Value::Null,
"next_probe_at_unix_secs": Value::Null,
"half_open_until": Value::Null,
"half_open_successes": 0,
"half_open_failures": 0,
}),
);
Some(Value::Object(circuit_by_format))
}
fn local_candidate_failure_should_project_health(
classification: LocalFailoverClassification,
status_code: u16,
@@ -112,7 +176,10 @@ fn projected_failure_health_score(
mod tests {
use serde_json::{json, Value};
use super::{project_local_failure_health, project_local_success_health};
use super::{
project_local_failure_health, project_local_key_circuit_closed,
project_local_key_circuit_open, project_local_success_health,
};
use crate::orchestration::LocalFailoverClassification;
#[test]
@@ -178,4 +245,47 @@ mod tests {
);
assert_eq!(projected["openai:responses"]["health_score"], json!(0.8));
}
#[test]
fn circuit_open_projection_sets_probe_deadline() {
let projected = project_local_key_circuit_open(
None,
"openai:chat",
"account_deactivated_401",
1_760_000_000,
)
.expect("projection should exist");
assert_eq!(projected["openai:chat"]["open"], json!(true));
assert_eq!(
projected["openai:chat"]["reason"],
json!("account_deactivated_401")
);
assert_eq!(
projected["openai:chat"]["next_probe_at_unix_secs"],
json!(1_760_001_920u64)
);
}
#[test]
fn circuit_closed_projection_resets_format_circuit() {
let projected = project_local_key_circuit_closed(
Some(&json!({
"openai:chat": {
"open": true,
"reason": "account_deactivated_401",
"next_probe_at_unix_secs": 1_760_001_920u64
}
})),
"openai:chat",
)
.expect("projection should exist");
assert_eq!(projected["openai:chat"]["open"], json!(false));
assert_eq!(projected["openai:chat"]["reason"], Value::Null);
assert_eq!(
projected["openai:chat"]["next_probe_at_unix_secs"],
Value::Null
);
}
}

5
apps/aether-gateway/src/orchestration/mod.rs
View File

@@ -31,7 +31,10 @@ pub(crate) use self::effects::{
LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect,
LocalPoolErrorEffect,
};
pub(crate) use self::health::{project_local_failure_health, project_local_success_health};
pub(crate) use self::health::{
project_local_failure_health, project_local_key_circuit_closed, project_local_key_circuit_open,
project_local_success_health,
};
pub(crate) use self::policy::{
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
local_failover_policy_from_transport, resolve_local_failover_policy, LocalFailoverPolicy,

43
crates/aether-scheduler-core/src/candidate/mod.rs
View File

@@ -506,6 +506,49 @@ mod tests {
);
}
#[test]
fn key_circuit_allows_probe_after_next_probe_time() {
let mut circuit_open_key = sample_key_with_concurrent_limit("1", Some(2));
circuit_open_key.circuit_breaker_by_format = Some(serde_json::json!({
"openai:chat": {
"open": true,
"next_probe_at_unix_secs": 100
}
}));
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), circuit_open_key)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 99,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("key_circuit_open")
);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
None
);
}
#[test]
fn candidate_selectability_rejects_quota_or_zero_health() {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);

7
crates/aether-scheduler-core/src/candidate/selectability.rs
View File

@@ -107,8 +107,11 @@ pub fn candidate_runtime_skip_reason_with_state(
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
if crate::is_provider_key_circuit_open_at(
provider_key,
candidate.endpoint_api_format.as_str(),
now_unix_secs,
) {
return Some("key_circuit_open");
}
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())

27
crates/aether-scheduler-core/src/health.rs
View File

@@ -284,6 +284,33 @@ pub fn is_provider_key_circuit_open(key: &StoredProviderCatalogKey, api_format:
.unwrap_or(false)
}
pub fn is_provider_key_circuit_open_at(
key: &StoredProviderCatalogKey,
api_format: &str,
now_unix_secs: u64,
) -> bool {
let Some(payload) = key
.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
else {
return false;
};
if !payload
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return false;
}
payload
.get("next_probe_at_unix_secs")
.and_then(serde_json::Value::as_u64)
.is_none_or(|next_probe_at| now_unix_secs < next_probe_at)
}
fn available_provider_key_rpm_slots_for_new_user(
key: &StoredProviderCatalogKey,
current_usage: usize,

5
crates/aether-scheduler-core/src/lib.rs
View File

@@ -31,8 +31,9 @@ pub use health::{
count_recent_rpm_requests_for_provider_key, count_recent_rpm_requests_for_provider_key_since,
effective_provider_key_health_score, effective_provider_key_rpm_limit,
is_candidate_in_recent_failure_cooldown, is_provider_key_circuit_open,
provider_key_health_bucket, provider_key_health_score, provider_key_rpm_allows_request,
provider_key_rpm_allows_request_since, ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
is_provider_key_circuit_open_at, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
};
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,