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

fix(pool): isolate dynamic model quota buckets and 429 scheduling

Browse Source
This commit is contained in:
zhefox
2026-09-02 15:23:23 +08:00
parent 2cd20da1ec
commit dbbe7b22ab
37 changed files with 1491 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/finalize/tests_sync.rs
+5 -4
View File
@@ -1,5 +1,6 @@
use std::collections::BTreeMap;
use aether_ai_formats::openai_responses_message_item_id;
use axum::body::to_bytes;
use base64::Engine as _;
use serde_json::json;
@@ -192,7 +193,7 @@ fn aggregates_openai_responses_stream_completed_event_to_final_response() {
"output_text": "Hello",
"output": [{
"type": "message",
"id": "resp_123_msg",
"id": openai_responses_message_item_id("resp_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -977,7 +978,7 @@ fn converts_gemini_cli_response_to_openai_responses_response() {
"output_text": "Hello Gemini CLI",
"output": [{
"type": "message",
"id": "resp_cli_123_msg",
"id": openai_responses_message_item_id("resp_cli_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -1046,7 +1047,7 @@ fn converts_gemini_cli_function_call_to_openai_responses_function_call() {
"output": [
{
"type": "message",
"id": "resp_cli_tool_123_msg",
"id": openai_responses_message_item_id("resp_cli_tool_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
@@ -1252,7 +1253,7 @@ fn local_finalize_handles_openai_responses_openai_family_sync_response_even_when
"model": "gpt-5",
"output": [{
"type": "message",
"id": "resp_cli_family_123_msg",
"id": openai_responses_message_item_id("resp_cli_family_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
apps/aether-gateway/src/ai_serving/pure/mod.rs
+1
View File
@@ -177,6 +177,7 @@ pub(crate) use aether_ai_formats::{
api_format_defaults_to_client_error_failover, api_format_defaults_to_non_stream,
api_format_permission_covers, codex_responses_lite_tool_is_client_executed,
intersect_api_format_allowed_lists, is_embedding_api_format, is_rerank_api_format,
normalize_openai_responses_message_item_ids, openai_responses_message_item_id,
openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy, ApiOperation, ClientSurface,
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+124 -10
View File
@@ -105,6 +105,7 @@ async fn schedule_pool_page_candidates(
candidates: Vec<EligibleLocalExecutionCandidate>,
sticky_session_token: Option<&str>,
effective_pool_config: Option<&AdminProviderPoolConfig>,
provider_model_name: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -128,7 +129,8 @@ async fn schedule_pool_page_candidates(
entry.1.insert(candidate.candidate.key_id.clone());
}
let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await;
let key_context_by_id =
read_pool_catalog_key_contexts_by_id(state, &candidates, provider_model_name).await;
let mut runtime_by_provider = BTreeMap::new();
let mut pool_config_by_provider = BTreeMap::new();
@@ -316,7 +318,9 @@ fn active_probe_member_is_unschedulable_for_request(
}) {
return true;
}
key_context.is_some_and(|context| context.account_blocked || context.quota_exhausted)
key_context.is_some_and(|context| {
context.account_blocked || context.quota_exhausted || context.quota_hard_blocked
})
}
async fn expand_pool_group_candidate(
@@ -1034,6 +1038,7 @@ impl<'a> PoolKeyCursor<'a> {
candidates,
self.sticky_session_token.as_deref(),
self.effective_pool_config.as_ref(),
Some(self.group.candidate.selected_provider_model_name.as_str()),
)
.await;
self.record_skipped_candidates(&skipped);
@@ -1406,6 +1411,7 @@ fn pool_candidate_from_catalog_key(
async fn read_pool_catalog_key_contexts_by_id(
state: PlannerAppState<'_>,
candidates: &[EligibleLocalExecutionCandidate],
provider_model_name: Option<&str>,
) -> BTreeMap<String, PoolCatalogKeyContext> {
let mut key_ids = Vec::new();
let mut provider_type_by_key_id = BTreeMap::<String, String>::new();
@@ -1437,13 +1443,30 @@ async fn read_pool_catalog_key_contexts_by_id(
key_count = key_ids.len(),
"gateway pool scheduler: failed to read catalog key metadata"
);
return BTreeMap::new();
// Do not fail open when the quota metadata read is unavailable. A
// missing context must never turn an exhausted account into an
// eligible candidate and produce another upstream 429. The caller
// treats this marker as a pool quota skip and the next request will
// retry the metadata read.
return key_ids
.into_iter()
.map(|key_id| {
(
key_id,
PoolCatalogKeyContext {
quota_hard_blocked: true,
..PoolCatalogKeyContext::default()
},
)
})
.collect();
}
};
let provider_pool_service = ProviderPoolService::with_builtin_adapters();
keys.into_iter()
let mut contexts = keys
.into_iter()
.map(|key| {
let provider_type = provider_type_by_key_id
.get(&key.id)
@@ -1451,10 +1474,28 @@ async fn read_pool_catalog_key_contexts_by_id(
.unwrap_or_default();
(
key.id.clone(),
build_pool_catalog_key_context(state, &provider_pool_service, &key, provider_type),
build_pool_catalog_key_context(
state,
&provider_pool_service,
&key,
provider_type,
provider_model_name,
),
)
})
.collect()
.collect::<BTreeMap<_, _>>();
// A key can disappear between the candidate-row and catalog reads. Keep
// the snapshot non-empty and fail closed for those IDs so the caller does
// not interpret an incomplete read as "all accounts are healthy".
for key_id in key_ids {
contexts
.entry(key_id)
.or_insert_with(|| PoolCatalogKeyContext {
quota_hard_blocked: true,
..PoolCatalogKeyContext::default()
});
}
contexts
}
fn build_pool_catalog_key_context(
@@ -1462,6 +1503,7 @@ fn build_pool_catalog_key_context(
provider_pool_service: &ProviderPoolService,
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: Option<&str>,
) -> PoolCatalogKeyContext {
let (health_score, _, _, _, _) = provider_key_health_summary(key);
let health_score = key
@@ -1480,8 +1522,12 @@ fn build_pool_catalog_key_context(
.filter(|value| value.is_finite() && *value >= 0.0);
let auth_config = parse_catalog_auth_config_json(state.app(), key);
let mut signals =
provider_pool_service.member_signals(provider_type, key, auth_config.as_ref());
let mut signals = provider_pool_service.member_signals(
provider_type,
key,
auth_config.as_ref(),
provider_model_name,
);
signals.account_blocked |= admin_provider_pool_pure::admin_pool_key_is_known_banned(key);
signals.account_blocked |=
pool_key_requires_reauth_for_scheduling(key, current_unix_ms().saturating_div(1000));
@@ -1694,7 +1740,21 @@ fn run_local_execution_pool_scheduler_with_runtime_map(
let key_context = key_context_by_id
.get(&candidate.candidate.key_id)
.cloned()
.unwrap_or_default();
.unwrap_or_else(|| {
// An explicitly non-empty metadata snapshot should contain
// every catalog key in this page. If one disappeared between
// reads, fail closed for that key instead of sending traffic
// with an unknown quota state. Empty maps are retained for
// callers/tests that intentionally provide no runtime context.
if key_context_by_id.is_empty() {
PoolCatalogKeyContext::default()
} else {
PoolCatalogKeyContext {
quota_hard_blocked: true,
..PoolCatalogKeyContext::default()
}
}
});
let admin_pool_config = effective_pool_config_by_provider
.get(&candidate.candidate.provider_id)
.cloned()
@@ -1968,7 +2028,7 @@ mod tests {
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_pool_core::PoolSchedulingPreset;
use aether_pool_core::{PoolSchedulingPreset, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON};
use aether_provider_pool::ProviderPoolService;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -2080,6 +2140,55 @@ mod tests {
);
}
#[test]
fn pool_scheduler_skips_quota_exhausted_key_when_flag_is_false() {
let ready = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-ready",
10,
Some(json!({ "pool_advanced": {} })),
);
let exhausted = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-exhausted",
10,
Some(json!({ "pool_advanced": { "skip_exhausted_accounts": false } })),
);
let key_context_by_id = BTreeMap::from([
("key-ready".to_string(), PoolCatalogKeyContext::default()),
(
"key-exhausted".to_string(),
PoolCatalogKeyContext {
quota_exhausted: true,
..PoolCatalogKeyContext::default()
},
),
]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![ready, exhausted],
&BTreeMap::new(),
&key_context_by_id,
);
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-ready"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-exhausted", POOL_ACCOUNT_EXHAUSTED_SKIP_REASON)]
);
}
#[test]
fn pool_scheduler_attaches_group_and_pool_metadata_to_ranked_candidates() {
let pool_first = sample_eligible_candidate(
@@ -4426,6 +4535,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert_eq!(context.plan_tier.as_deref(), Some("team"));
@@ -4471,6 +4581,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(!context.quota_exhausted);
@@ -4491,6 +4602,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(context.quota_exhausted);
@@ -4521,6 +4633,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"antigravity",
None,
);
assert!(context.quota_exhausted);
@@ -4542,6 +4655,7 @@ mod tests {
&ProviderPoolService::with_builtin_adapters(),
&key,
"codex",
None,
);
assert!(context.account_blocked);
apps/aether-gateway/src/execution_runtime/grok.rs
+7 -2
View File
@@ -26,7 +26,9 @@ use crate::ai_serving::api::{
CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter,
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, StreamingCanonicalUsage,
};
use crate::ai_serving::openai_responses_synthetic_reasoning_item_id;
use crate::ai_serving::{
openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id,
};
use crate::clock::current_unix_secs;
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::{
@@ -2725,7 +2727,7 @@ fn openai_responses_body(
};
if !message_text.trim().is_empty() {
output.push(json!({
"id": format!("{response_id}_msg"),
"id": openai_responses_message_item_id(response_id.as_str(), output.len()),
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": message_text, "annotations": []}],
@@ -3774,6 +3776,9 @@ mod tests {
);
assert_eq!(body["output"][0]["type"], serde_json::json!("reasoning"));
assert_eq!(body["output"][1]["type"], serde_json::json!("message"));
assert!(body["output"][1]["id"]
.as_str()
.is_some_and(|id| id.starts_with("msg_")));
}
#[test]
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
+42 -12
View File
@@ -502,14 +502,12 @@ pub(crate) async fn record_admin_provider_pool_error(
.or_else(|| parse_google_quota_cooldown_seconds(error_body)),
pool_config,
);
set_pool_cooldown(
runtime,
provider_id,
key_id,
"rate_limited_429",
ttl_seconds,
)
.await;
let reason = if error_body_indicates_quota_exhaustion(error_body) {
"quota_exhausted_429"
} else {
"rate_limited_429"
};
set_pool_cooldown(runtime, provider_id, key_id, reason, ttl_seconds).await;
return;
}
@@ -548,6 +546,27 @@ pub(crate) async fn record_admin_provider_pool_error(
}
}
fn error_body_indicates_quota_exhaustion(error_body: Option<&str>) -> bool {
let body = error_body.unwrap_or_default().to_ascii_lowercase();
[
"quota exhausted",
"quota_exhausted",
"quota exceeded",
"quota_exceeded",
"insufficient_quota",
"resource exhausted",
"resource has been exhausted",
"resource_exhausted",
"usage_limit_reached",
"limit_reached",
"quota limit reached",
"credits exhausted",
"insufficient credits",
]
.iter()
.any(|marker| body.contains(marker))
}
pub(crate) async fn record_admin_provider_pool_stream_timeout(
runtime: &RuntimeState,
provider_id: &str,
@@ -589,9 +608,10 @@ pub(crate) async fn record_admin_provider_pool_stream_timeout(
#[cfg(test)]
mod tests {
use super::{
admin_provider_pool_key_terminal_error_reason, parse_google_quota_cooldown_seconds_at,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, resolve_transient_cooldown_ttl,
admin_provider_pool_key_terminal_error_reason, error_body_indicates_quota_exhaustion,
parse_google_quota_cooldown_seconds_at, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
resolve_transient_cooldown_ttl,
};
use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state;
use crate::handlers::admin::provider::shared::support::{
@@ -803,6 +823,16 @@ mod tests {
assert_eq!(resolve_transient_cooldown_ttl(500, None, &pool_config), 0);
}
#[test]
fn detects_quota_exhaustion_markers_in_429_bodies() {
assert!(error_body_indicates_quota_exhaustion(Some(
r#"{"error":{"status":"RESOURCE_EXHAUSTED","message":"quota exceeded"}}"#,
)));
assert!(!error_body_indicates_quota_exhaustion(Some(
r#"{"error":{"message":"temporary rate limit"}}"#,
)));
}
#[tokio::test]
async fn success_feedback_writes_sticky_lru_cost_and_latency() {
let Some(redis) = start_managed_redis_or_skip().await else {
@@ -1110,7 +1140,7 @@ mod tests {
.cooldown_reason_by_key
.get("key-google-429")
.map(String::as_str),
Some("rate_limited_429")
Some("quota_exhausted_429")
);
assert!(runtime
.cooldown_ttl_by_key
apps/aether-gateway/src/handlers/proxy/finalize.rs
+12
View File
@@ -50,6 +50,7 @@ pub(super) fn finalize_gateway_response(
mut response: Response<Body>,
trace_id: &str,
remote_addr: &std::net::SocketAddr,
client_ip: std::net::IpAddr,
method: &http::Method,
path_and_query: &str,
control_decision: Option<&GatewayControlDecision>,
@@ -117,6 +118,7 @@ pub(super) fn finalize_gateway_response(
trace_id = %trace_id,
request_id,
remote_addr = %remote_addr,
client_ip = %client_ip,
method = %method,
path = %sanitized_path_and_query,
user_id,
@@ -137,6 +139,7 @@ pub(super) fn finalize_gateway_response(
trace_id = %trace_id,
request_id,
remote_addr = %remote_addr,
client_ip = %client_ip,
method = %method,
path = %sanitized_path_and_query,
user_id,
@@ -157,6 +160,7 @@ pub(super) fn finalize_gateway_response(
trace_id = %trace_id,
request_id,
remote_addr = %remote_addr,
client_ip = %client_ip,
method = %method,
path = %sanitized_path_and_query,
user_id,
@@ -247,11 +251,17 @@ pub(super) fn finalize_gateway_response_with_context(
started_at: &Instant,
request_permit: Option<AdmissionPermit>,
) -> Response<Body> {
let client_ip = request_context
.client_ip
.as_deref()
.and_then(|value| value.parse().ok())
.unwrap_or_else(|| remote_addr.ip());
finalize_gateway_response(
state,
response,
&request_context.trace_id,
remote_addr,
client_ip,
&request_context.request_method,
&request_context.request_path_and_query(),
request_context.control_decision.as_ref(),
@@ -320,6 +330,7 @@ mod tests {
request_query_string: None,
request_content_type: Some("application/json".to_string()),
host_header: None,
client_ip: None,
control_decision: None,
};
let mut headers = HeaderMap::new();
@@ -373,6 +384,7 @@ mod tests {
response,
"trace-finalize",
&remote_addr,
remote_addr.ip(),
&Method::GET,
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse",
Some(&control_decision),
apps/aether-gateway/src/handlers/proxy/mod.rs
+5
View File
@@ -993,6 +993,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1029,6 +1030,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1069,6 +1071,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
request.method(),
request
.uri()
@@ -1128,6 +1131,7 @@ async fn proxy_request_inner(
response,
&trace_id,
&remote_addr,
client_ip,
&parts.method,
parts
.uri
@@ -1149,6 +1153,7 @@ async fn proxy_request_inner(
&trace_id,
)
.await?;
request_context.client_ip = Some(client_ip.to_string());
maybe_promote_management_token_admin_principal(
&state,
client_ip,
apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs
+5
View File
@@ -1001,6 +1001,7 @@ mod tests {
request_query_string: None,
request_content_type: None,
host_header: None,
client_ip: None,
control_decision: None,
};
let (parts, _) = http::Request::builder()
@@ -1036,6 +1037,7 @@ mod tests {
request_query_string: None,
request_content_type: Some("application/sdp".to_string()),
host_header: None,
client_ip: None,
control_decision: Some(decision),
};
let (parts, _) = http::Request::builder()
@@ -1074,6 +1076,7 @@ mod tests {
request_query_string: Some("intent=quicksilver&architecture=avas".to_string()),
request_content_type: Some("multipart/form-data".to_string()),
host_header: None,
client_ip: None,
control_decision: Some(decision),
};
let (parts, _) = http::Request::builder()
@@ -1128,6 +1131,7 @@ mod tests {
request_query_string: None,
request_content_type: None,
host_header: None,
client_ip: None,
control_decision: Some(decision),
};
let (parts, _) = http::Request::builder()
@@ -1194,6 +1198,7 @@ mod tests {
request_query_string: None,
request_content_type: Some("multipart/form-data".to_string()),
host_header: None,
client_ip: None,
control_decision: Some(decision),
};
let (content_type, body) = build_live_multipart(
apps/aether-gateway/src/handlers/shared/catalog.rs
+12 -1
View File
@@ -975,7 +975,7 @@ fn build_codex_quota_status_snapshot(
.and_then(admin_provider_quota_pure::coerce_json_bool);
let reset_credits = build_codex_reset_credits_status_snapshot(metadata, observed_at_unix_secs);
let windows = [
let mut windows = [
codex_quota_window_snapshot(metadata, "primary", "weekly", "周", observed_at_unix_secs),
codex_quota_window_snapshot(metadata, "secondary", "5h", "5H", observed_at_unix_secs),
codex_quota_window_snapshot(
@@ -996,6 +996,17 @@ fn build_codex_quota_status_snapshot(
.into_iter()
.flatten()
.collect::<Vec<_>>();
if let Some(additional_windows) = metadata
.get("additional_quota_windows")
.and_then(Value::as_array)
{
windows.extend(
additional_windows
.iter()
.filter(|window| window.is_object())
.cloned(),
);
}
if windows.is_empty()
&& plan_type.is_none()
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+52
View File
@@ -1213,6 +1213,18 @@ fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
// Providers frequently encode an exhausted account as HTTP 429 with a
// provider-specific status/message (for example RESOURCE_EXHAUSTED or
// `quota exceeded`) rather than the normalized `quota_exhausted` status.
// Preserve that distinction in the score so the member stays out of the
// scheduler until a successful quota probe observes a reset.
let serialized_item = item.to_string().to_ascii_lowercase();
let status_code = item.get("status_code").and_then(Value::as_u64);
if status == "quota_exhausted"
|| (status_code == Some(429) && contains_quota_exhaustion_marker(&serialized_item))
{
return Some(PoolMemberHardState::QuotaExhausted);
}
match status.as_str() {
"auth_invalid" | "forbidden" => Some(PoolMemberHardState::AuthInvalid),
"workspace_deactivated" => Some(PoolMemberHardState::Banned),
@@ -1226,6 +1238,26 @@ fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
}
}
fn contains_quota_exhaustion_marker(value: &str) -> bool {
[
"quota exhausted",
"quota_exhausted",
"quota exceeded",
"quota_exceeded",
"insufficient_quota",
"resource exhausted",
"resource has been exhausted",
"resource_exhausted",
"usage_limit_reached",
"limit_reached",
"quota limit reached",
"credits exhausted",
"insufficient credits",
]
.iter()
.any(|marker| value.contains(marker))
}
async fn perform_pool_quota_probe_for_provider(
state: &AppState,
admin_state: &AdminAppState<'_>,
@@ -1771,6 +1803,26 @@ mod tests {
assert_eq!(selected, vec!["never".to_string(), "old".to_string()]);
}
#[test]
fn quota_markers_in_429_probe_results_are_hard_exhaustion() {
assert_eq!(
probe_result_hard_state(&json!({
"status": "rate_limited",
"status_code": 429,
"message": "RESOURCE_EXHAUSTED: quota exceeded"
})),
Some(PoolMemberHardState::QuotaExhausted)
);
assert_eq!(
probe_result_hard_state(&json!({
"status": "rate_limited",
"status_code": 429,
"message": "temporary rate limit"
})),
Some(PoolMemberHardState::Cooldown)
);
}
fn score(
member_id: &str,
hard_state: PoolMemberHardState,
apps/aether-gateway/src/orchestration/effects.rs
+37
View File
@@ -2060,6 +2060,15 @@ fn pool_score_hard_state_for_status(
return Some(pool_score_hard_state_for_terminal_error_reason(&reason));
}
// A number of providers report account quota exhaustion as HTTP 429 rather
// than 402. Keep those members out of the score-based pool fallback until
// the provider's quota probe observes a reset; treating every 429 as a
// generic cooldown otherwise lets the member re-enter as soon as the short
// transient cooldown expires.
if status_code == 429 && error_body_indicates_quota_exhaustion(error_body) {
return Some(PoolMemberHardState::QuotaExhausted);
}
match status_code {
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
402 => Some(PoolMemberHardState::QuotaExhausted),
@@ -2082,6 +2091,27 @@ fn pool_score_hard_state_for_status(
}
}
fn error_body_indicates_quota_exhaustion(error_body: Option<&str>) -> bool {
let body = error_body.unwrap_or_default().to_ascii_lowercase();
[
"quota exhausted",
"quota_exhausted",
"quota exceeded",
"quota_exceeded",
"insufficient_quota",
"resource exhausted",
"resource has been exhausted",
"resource_exhausted",
"usage_limit_reached",
"limit_reached",
"quota limit reached",
"credits exhausted",
"insufficient credits",
]
.iter()
.any(|marker| body.contains(marker))
}
fn pool_score_hard_state_for_terminal_error_reason(reason: &str) -> PoolMemberHardState {
if reason.starts_with("payment_required_") {
PoolMemberHardState::QuotaExhausted
@@ -3659,6 +3689,13 @@ mod tests {
),
Some(PoolMemberHardState::QuotaExhausted)
);
assert_eq!(
pool_score_hard_state_for_status(
429,
Some(r#"{"error":{"status":"RESOURCE_EXHAUSTED","message":"quota exhausted"}}"#),
),
Some(PoolMemberHardState::QuotaExhausted)
);
}
#[tokio::test]
apps/aether-gateway/src/scheduler/candidate/runtime.rs
+24 -9
View File
@@ -341,18 +341,33 @@ fn read_key_account_quota_exhaustion_map(
candidates
.iter()
.map(|candidate| {
let exhausted = provider_skip_exhausted_accounts
.get(candidate.provider_id.as_str())
.copied()
.unwrap_or(false)
&& provider_key_rpm_states
.get(candidate.key_id.as_str())
.is_some_and(|key| {
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(
let exhausted = provider_key_rpm_states
.get(candidate.key_id.as_str())
.is_some_and(|key| {
let account_exhausted =
admin_provider_pool_pure::admin_pool_key_model_quota_exhausted(
key,
candidate.provider_type.as_str(),
candidate.selected_provider_model_name.as_str(),
)
});
.unwrap_or_else(|| {
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(
key,
candidate.provider_type.as_str(),
)
});
let hard_blocked =
admin_provider_pool_pure::admin_pool_key_model_quota_hard_blocked(
key,
candidate.provider_type.as_str(),
candidate.selected_provider_model_name.as_str(),
);
let skip_configured = provider_skip_exhausted_accounts
.get(candidate.provider_id.as_str())
.copied()
.unwrap_or(false);
hard_blocked || (skip_configured && account_exhausted)
});
(candidate.key_id.clone(), exhausted)
})
.collect()
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/compact.rs
+2 -1
View File
@@ -6,6 +6,7 @@ use super::{
EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;
use aether_ai_formats::openai_responses_message_item_id;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
@@ -413,7 +414,7 @@ async fn gateway_executes_openai_responses_compact_openai_family_upstream_stream
"output_text": "Hello Compact",
"output": [{
"type": "message",
"id": "resp_compact_openai_family_123_msg",
"id": openai_responses_message_item_id("resp_compact_openai_family_123", 0),
"role": "assistant",
"status": "completed",
"content": [{
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/direct.rs
+2 -1
View File
@@ -6,6 +6,7 @@ use super::{
TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;
use aether_ai_formats::openai_responses_message_item_id;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
@@ -428,7 +429,7 @@ async fn gateway_executes_openai_responses_sync_upstream_stream_via_local_finali
"output_text": "Hello",
"output": [{
"type": "message",
"id": "resp_stream_001_msg",
"id": openai_responses_message_item_id("resp_stream_001", 0),
"role": "assistant",
"status": "completed",
"content": [{