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

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": [{
crates/aether-admin/src/provider/pool.rs
+24
View File
@@ -118,6 +118,30 @@ pub fn admin_pool_key_quota_hard_blocked(
aether_provider_pool::provider_pool_key_quota_hard_blocked(key, provider_type)
}
pub fn admin_pool_key_model_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
aether_provider_pool::provider_pool_key_model_quota_exhausted(
key,
provider_type,
provider_model_name,
)
}
pub fn admin_pool_key_model_quota_hard_blocked(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> bool {
aether_provider_pool::provider_pool_key_model_quota_hard_blocked(
key,
provider_type,
provider_model_name,
)
}
fn admin_pool_has_proxy(key: &StoredProviderCatalogKey) -> bool {
match key.proxy.as_ref() {
Some(Value::Object(values)) => !values.is_empty(),
crates/aether-admin/src/provider/quota.rs
+153
View File
@@ -10,6 +10,7 @@ const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
const CODEX_SPARK_LIMIT_NAME: &str = "GPT-5.3-Codex-Spark";
const CODEX_ADDITIONAL_QUOTA_WINDOWS_KEY: &str = "additional_quota_windows";
pub fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
config
@@ -1893,6 +1894,112 @@ fn codex_find_spark_rate_limit(
.and_then(serde_json::Value::as_object)
}
/// Preserve every additional Codex rate-limit bucket in a model-addressable
/// representation. The upstream can add independent limits over time; do
/// not discard them merely because they are not one of the legacy flat
/// primary/secondary slots.
fn codex_additional_quota_windows(
root: &serde_json::Map<String, serde_json::Value>,
) -> Vec<serde_json::Value> {
let mut windows = Vec::new();
let Some(items) = root
.get("additional_rate_limits")
.and_then(serde_json::Value::as_array)
else {
return windows;
};
for (index, item) in items.iter().enumerate() {
let Some(item_object) = item.as_object() else {
continue;
};
let Some(limit_name) = ["limit_name", "metered_feature", "name", "id"]
.iter()
.find_map(|key| item_object.get(*key).and_then(serde_json::Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let model_identity = item_object
.get("model")
.or_else(|| item_object.get("model_name"))
.or_else(|| item_object.get("model_id"))
.cloned()
.unwrap_or_else(|| json!(limit_name));
let model_identities = item_object
.get("models")
.or_else(|| item_object.get("model_ids"))
.cloned();
let Some(rate_limit) = item_object
.get("rate_limit")
.and_then(serde_json::Value::as_object)
else {
continue;
};
for (slot, slot_name) in [
("primary_window", "primary"),
("secondary_window", "secondary"),
] {
let Some(source) = rate_limit.get(slot).and_then(serde_json::Value::as_object) else {
continue;
};
let used_percent = source.get("used_percent").and_then(coerce_json_f64);
let reset_after_seconds = source
.get("reset_after_seconds")
.or_else(|| source.get("reset_seconds"))
.and_then(coerce_json_u64);
let reset_at = source.get("reset_at").and_then(coerce_json_u64);
let window_minutes = source
.get("limit_window_seconds")
.or_else(|| source.get("window_seconds"))
.and_then(coerce_json_u64)
.map(|seconds| seconds.saturating_add(59) / 60)
.filter(|minutes| *minutes > 0);
if used_percent.is_none()
&& reset_after_seconds.is_none()
&& reset_at.is_none()
&& window_minutes.is_none()
{
continue;
}
let used_ratio = used_percent.map(|value| (value / 100.0).clamp(0.0, 1.0));
let mut window = serde_json::Map::new();
window.insert(
"code".to_string(),
json!(format!("additional_{index}_{slot_name}")),
);
window.insert("label".to_string(), json!(limit_name));
window.insert("scope".to_string(), json!("model"));
// The upstream limit name is the only stable identity available
// for an additional bucket. Keeping it in both fields allows
// exact model matching and generic family matching downstream.
window.insert("model".to_string(), model_identity.clone());
if let Some(model_identities) = model_identities.clone() {
window.insert("models".to_string(), model_identities);
}
window.insert("quota_group".to_string(), json!(limit_name));
window.insert("limit_name".to_string(), json!(limit_name));
window.insert("used_ratio".to_string(), json!(used_ratio));
window.insert(
"remaining_ratio".to_string(),
json!(used_ratio.map(|value| (1.0 - value).max(0.0))),
);
window.insert(
"is_exhausted".to_string(),
json!(used_ratio.is_some_and(|value| value >= 1.0 - 1e-6)),
);
window.insert("reset_at".to_string(), json!(reset_at));
window.insert("reset_seconds".to_string(), json!(reset_after_seconds));
window.insert("window_minutes".to_string(), json!(window_minutes));
windows.push(serde_json::Value::Object(window));
}
}
windows
}
fn codex_reset_credits_container(
root: &serde_json::Map<String, serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
@@ -1993,6 +2100,14 @@ pub fn parse_codex_wham_usage_response(
}
}
// Keep all additional limits, including future ones with names unknown to
// this version of the gateway, so scheduling can select their bucket by
// metadata rather than by a hard-coded model name.
result.insert(
CODEX_ADDITIONAL_QUOTA_WINDOWS_KEY.to_string(),
serde_json::Value::Array(codex_additional_quota_windows(root)),
);
if let Some(credits) = root.get("credits").and_then(serde_json::Value::as_object) {
if let Some(value) = credits.get("has_credits").and_then(coerce_json_bool) {
result.insert("has_credits".to_string(), json!(value));
@@ -5413,6 +5528,44 @@ mod tests {
parsed.get("spark_secondary_window_minutes"),
Some(&json!(10_080u64))
);
let additional = parsed
.get("additional_quota_windows")
.and_then(serde_json::Value::as_array)
.expect("all additional quota windows should be preserved");
assert_eq!(additional.len(), 2);
assert!(additional.iter().all(|window| {
window.get("model") == Some(&json!("GPT-5.3-Codex-Spark"))
&& window.get("scope") == Some(&json!("model"))
}));
}
#[test]
fn parses_unknown_additional_quota_buckets_without_model_specific_code() {
let parsed = parse_codex_wham_usage_response(
&json!({
"rate_limit": {
"primary_window": {"used_percent": 100.0},
"secondary_window": {"used_percent": 0.0}
},
"additional_rate_limits": [{
"limit_name": "Future-Model-Limit",
"rate_limit": {
"primary_window": {
"used_percent": 100.0,
"limit_window_seconds": 3600
}
}
}]
}),
1_700_000_000,
)
.expect("quota response should parse");
let additional = parsed["additional_quota_windows"]
.as_array()
.expect("additional windows should be an array");
assert_eq!(additional.len(), 1);
assert_eq!(additional[0]["model"], json!("Future-Model-Limit"));
assert_eq!(additional[0]["is_exhausted"], json!(true));
}
#[test]
crates/aether-ai/formats/src/api.rs
+1
View File
@@ -78,6 +78,7 @@ pub use crate::formats::openai::{
OpenAiProviderRequestFinalization,
},
responses::{
normalize_openai_responses_message_item_ids, openai_responses_message_item_id,
openai_responses_synthetic_reasoning_item_id,
strip_incompatible_openai_responses_reasoning_items,
strip_incompatible_openai_responses_reasoning_items_with_policy,
crates/aether-ai/formats/src/formats/conversion/response.rs
+2 -2
View File
@@ -9,7 +9,7 @@ use serde_json::{json, Value};
use crate::formats::{
context::FormatContext,
openai::responses::{
openai_responses_synthetic_reasoning_item_id,
openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id,
response::ensure_modern_openai_responses_response_fields,
},
registry,
@@ -218,7 +218,7 @@ pub fn build_openai_responses_response_with_content(
if !content.is_empty() {
output.push(json!({
"type": "message",
"id": format!("{response_id}_msg"),
"id": openai_responses_message_item_id(response_id, 0),
"role": "assistant",
"status": "completed",
"content": content
crates/aether-ai/formats/src/formats/openai/chat/stream.rs
+6 -6
View File
@@ -3,7 +3,7 @@ use std::collections::{BTreeMap, BTreeSet};
use serde_json::{json, Map, Value};
use crate::formats::openai::responses::{
openai_responses_synthetic_reasoning_item_id,
openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id,
response::{
ensure_modern_openai_responses_response_fields, openai_responses_current_timestamp,
},
@@ -2240,7 +2240,7 @@ impl OpenAIResponsesClientEmitter {
fn message_item_id(&self) -> String {
self.message_item_id
.clone()
.unwrap_or_else(|| format!("{}_msg", self.response_id()))
.unwrap_or_else(|| openai_responses_message_item_id(self.response_id(), 0))
}
fn reasoning_item_id(&self) -> String {
@@ -2259,7 +2259,7 @@ impl OpenAIResponsesClientEmitter {
fn ensure_message_item_id(&mut self) -> String {
if self.message_item_id.is_none() {
self.message_item_id = Some(format!("{}_msg", self.response_id()));
self.message_item_id = Some(openai_responses_message_item_id(self.response_id(), 0));
}
self.message_item_id()
}
@@ -4167,7 +4167,7 @@ mod tests {
assert!(sse.contains("event: response.output_item.done\n"));
assert!(sse.contains("event: response.completed\n"));
assert!(sse.contains("\"response_id\":\"resp_stream_123\""));
assert!(sse.contains("\"item_id\":\"resp_stream_123_msg\""));
assert!(sse.contains("\"item_id\":\"msg_aether_"));
assert!(sse.contains("\"text\":\"Hello\""));
assert!(sse.contains("\"output_text\":\"Hello\""));
assert!(sse.contains("\"created_at\":"));
@@ -4231,8 +4231,8 @@ mod tests {
);
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert!(sse.contains("\"item_id\":\"msg_first_msg\""));
assert!(!sse.contains("\"item_id\":\"msg_second_msg\""));
assert!(sse.contains("\"item_id\":\"msg_aether_"));
assert!(!sse.contains("msg_second_msg"));
}
#[test]
crates/aether-ai/formats/src/formats/openai/request_contract.rs
+40
View File
@@ -147,6 +147,13 @@ fn finalize_openai_provider_request_with_codex_model_capabilities_and_reasoning_
finalization.provider_api_format,
reasoning_replay_policy,
);
if finalization
.provider_api_format
.trim()
.eq_ignore_ascii_case("openai:responses")
{
super::responses::normalize_openai_responses_message_item_ids(body);
}
crate::enforce_request_body_stream_field(
body,
finalization.provider_api_format,
@@ -408,6 +415,39 @@ mod tests {
assert_eq!(input[1]["type"], "message");
}
#[test]
fn finalization_repairs_legacy_responses_message_ids_for_same_format_upstream() {
let mut body = json!({
"model": "gpt-5.4",
"input": [{
"type": "message",
"id": "1c938e58-32a8-4d28-9c34-538d78076895_msg",
"role": "assistant",
"content": [{"type": "input_text", "text": "previous answer"}]
}]
});
finalize_openai_provider_request(
&mut body,
OpenAiProviderRequestFinalization {
source_api_format: "openai:responses",
provider_api_format: "openai:responses",
provider_type: "codex",
provider_model: "gpt-5.4",
source_model: "gpt-5.4",
body_rules: None,
upstream_is_stream: false,
require_body_stream_field: false,
},
)
.expect("legacy message IDs should be repaired before provider validation");
let repaired_id = body["input"][0]["id"]
.as_str()
.expect("message ID should be a string");
assert!(repaired_id.starts_with("msg_"));
}
#[test]
fn non_responses_sources_receive_codex_responses_reasoning_defaults() {
for source_api_format in ["openai:chat", "claude:messages", "gemini:generate_content"] {
crates/aether-ai/formats/src/formats/openai/responses/mod.rs
+101 -1
View File
@@ -9,6 +9,7 @@ pub mod stream;
const TOOL_ERROR_PREFIX: &str = "[tool error]";
const AETHER_REASONING_ITEM_ID_PREFIX: &str = "rs_aether_";
const AETHER_MESSAGE_ITEM_ID_PREFIX: &str = "msg_aether_";
/// Controls which provider-owned reasoning items may be replayed on a Responses request.
///
@@ -38,6 +39,67 @@ pub fn openai_responses_synthetic_reasoning_item_id(
)
}
/// Builds a stable, wire-compatible ID for a message item synthesized by Aether.
///
/// Responses clients replay assistant message items verbatim on the next turn and
/// OpenAI requires those IDs to begin with `msg`. Upstream response IDs are not
/// guaranteed to have that prefix (Chat completion IDs and UUIDs are common), so
/// appending a suffix to the response ID is not sufficient. A deterministic UUID
/// keeps the ID stable across sync/stream projections while avoiding assumptions
/// about the upstream ID's shape or length.
pub fn openai_responses_message_item_id(response_id: &str, output_index: usize) -> String {
let seed = format!("{response_id}:{output_index}");
format!(
"{AETHER_MESSAGE_ITEM_ID_PREFIX}{}",
uuid::Uuid::new_v5(&uuid::Uuid::NAMESPACE_OID, seed.as_bytes()).simple()
)
}
/// Repairs legacy/non-OpenAI message IDs in a Responses request in place.
///
/// Aether versions before the `msg_` contract emitted IDs such as
/// `<response-id>_msg`. Clients legitimately replay those assistant items on
/// the next turn, so merely fixing newly generated responses leaves existing
/// conversations broken. Preserve already-valid provider IDs and deterministically
/// remap only message items that do not begin with `msg`.
pub fn normalize_openai_responses_message_item_ids(body: &mut Value) -> usize {
let Some(items) = body.get_mut("input").and_then(Value::as_array_mut) else {
return 0;
};
let mut repaired = 0usize;
for (index, item) in items.iter_mut().enumerate() {
let Some(object) = item.as_object_mut() else {
continue;
};
if object.get("type").and_then(Value::as_str) != Some("message") {
continue;
}
let Some(raw_id) = object.get("id") else {
// IDs are optional for newly-authored input messages. Only repair
// an ID that a previous response actually supplied.
continue;
};
let valid = object
.get("id")
.and_then(Value::as_str)
.is_some_and(|id| id.starts_with("msg"));
if valid {
continue;
}
let source_id = raw_id
.as_str()
.filter(|id| !id.trim().is_empty())
.unwrap_or("missing")
.to_string();
object.insert(
"id".to_string(),
Value::String(openai_responses_message_item_id(source_id.as_str(), index)),
);
repaired += 1;
}
repaired
}
/// Removes reasoning history items that cannot be replayed against an OpenAI Responses backend.
///
/// Reasoning IDs are opaque provider references and must never be repaired by changing their
@@ -190,7 +252,9 @@ mod tests {
use serde_json::json;
use super::{
openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id,
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,
OpenAiResponsesReasoningReplayPolicy, OPENAI_RESPONSES_OPERATION_COMPACT,
@@ -238,6 +302,42 @@ mod tests {
assert_ne!(first, other);
}
#[test]
fn synthetic_message_item_ids_are_stable_and_start_with_msg() {
let first = openai_responses_message_item_id(
"1c938e58-32a8-4d28-9c34-538d78076895",
0,
);
let second = openai_responses_message_item_id(
"1c938e58-32a8-4d28-9c34-538d78076895",
0,
);
let other = openai_responses_message_item_id("chatcmpl-123", 1);
assert!(first.starts_with("msg_"));
assert_eq!(first, second);
assert_ne!(first, other);
}
#[test]
fn normalizes_legacy_message_ids_but_preserves_valid_ids() {
let mut body = json!({
"input": [
{"type": "message", "id": "1c938e58-32a8-4d28-9c34-538d78076895_msg", "role": "assistant"},
{"type": "message", "id": "msg_provider_123", "role": "assistant"},
{"type": "function_call", "id": "legacy_call"},
{"type": "message", "role": "user"}
]
});
assert_eq!(normalize_openai_responses_message_item_ids(&mut body), 1);
let input = body["input"].as_array().expect("input array");
assert!(input[0]["id"].as_str().is_some_and(|id| id.starts_with("msg_")));
assert_eq!(input[1]["id"], "msg_provider_123");
assert_eq!(input[2].get("id"), Some(&json!("legacy_call")));
assert!(input[3].get("id").is_none());
}
#[test]
fn strips_foreign_and_non_replayable_synthetic_reasoning_items() {
let portable_synthetic = openai_responses_synthetic_reasoning_item_id("resp_123", 1);
crates/aether-ai/formats/src/formats/shared/sync_products.rs
+24 -3
View File
@@ -7,7 +7,9 @@ use aether_ai_formats::formats::conversion::response::{
convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat,
};
use aether_ai_formats::formats::openai::responses::openai_responses_synthetic_reasoning_item_id;
use aether_ai_formats::formats::openai::responses::{
openai_responses_message_item_id, openai_responses_synthetic_reasoning_item_id,
};
use aether_ai_formats::formats::openai::responses::response::ensure_modern_openai_responses_response_fields;
use aether_ai_formats::formats::registry::{convert_response, FormatContext, FormatError};
use aether_ai_formats::{
@@ -2696,6 +2698,7 @@ fn aggregate_openai_responses_stream_sync_response_from_validated_terminal(
if let Some(state) = message_states.remove(&output_index) {
output.push(materialize_openai_responses_message_item(
&response_id,
output_index,
state,
));
}
@@ -3152,13 +3155,31 @@ fn resolve_openai_responses_tool_output_index(
fn materialize_openai_responses_message_item(
response_id: &str,
output_index: usize,
state: OpenAIResponsesSyncMessageState,
) -> Value {
let mut item = state.item;
item.entry("type".to_string())
.or_insert_with(|| Value::String("message".to_string()));
item.entry("id".to_string())
.or_insert_with(|| Value::String(format!("{response_id}_msg")));
let message_id_is_valid = item
.get("id")
.and_then(Value::as_str)
.is_some_and(|id| id.starts_with("msg"));
if !message_id_is_valid {
let source_id = item
.get("id")
.and_then(Value::as_str)
.filter(|id| !id.trim().is_empty())
.unwrap_or(response_id)
.to_string();
item.insert(
"id".to_string(),
Value::String(openai_responses_message_item_id(
source_id.as_str(),
output_index,
)),
);
}
item.entry("role".to_string())
.or_insert_with(|| Value::String("assistant".to_string()));
item.entry("status".to_string())
crates/aether-ai/formats/src/lib.rs
+3 -1
View File
@@ -56,7 +56,9 @@ pub use formats::openai::responses::request::{
validate_openai_responses_request_contract, OpenAiResponsesRequestContractViolation,
};
pub use formats::openai::responses::{
openai_responses_request_operation, openai_responses_synthetic_reasoning_item_id,
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,
OpenAiResponsesReasoningReplayPolicy, OPENAI_RESPONSES_OPERATION_COMPACT,
crates/aether-ai/formats/src/protocol/canonical.rs
+2 -5
View File
@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use crate::formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort;
use crate::formats::openai::responses::openai_responses_message_item_id;
use crate::formats::shared::model_directives::ReasoningEffort;
use crate::formats::shared::response::remove_empty_pages_from_tool_input_value;
@@ -3948,11 +3949,7 @@ pub(crate) fn flush_openai_responses_message_item(
if message_content.is_empty() {
return;
}
let id = if *message_index == 0 {
format!("{response_id}_msg")
} else {
format!("{response_id}_msg_{message_index}")
};
let id = openai_responses_message_item_id(response_id, *message_index);
output.push(json!({
"type": "message",
"id": id,
crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs
+2
View File
@@ -8,6 +8,8 @@ use tracing::{error, info, warn};
use super::types::PendingBackfillInfo;
// Historical backfill manifests must remain available after they are applied;
// deployed databases retain their versions for startup compatibility checks.
static BACKFILL_MIGRATOR: Migrator = sqlx::migrate!("./backfills/postgres");
const SCHEMA_BACKFILLS_TABLE_EXISTS_SQL: &str =
crates/aether-gateway/control/src/public/request_context.rs
+2
View File
@@ -8,6 +8,7 @@ pub struct PublicRequestContext<Decision> {
pub request_query_string: Option<String>,
pub request_content_type: Option<String>,
pub host_header: Option<String>,
pub client_ip: Option<String>,
pub control_decision: Option<Decision>,
}
@@ -32,6 +33,7 @@ impl<Decision> PublicRequestContext<Decision> {
request_query_string: uri.query().map(ToOwned::to_owned),
request_content_type: header_value(headers, http::header::CONTENT_TYPE),
host_header: header_value(headers, http::header::HOST),
client_ip: None,
control_decision,
}
}
crates/aether-pool-core/src/scheduler.rs
+33 -3
View File
@@ -18,6 +18,8 @@ pub struct PoolSchedulingPreset {
pub struct PoolSchedulingConfig {
pub scheduling_presets: Vec<PoolSchedulingPreset>,
pub lru_enabled: bool,
/// Retained for configuration/API compatibility. Active quota exhaustion is
/// always an admission block; reset-aware adapters decide when it clears.
pub skip_exhausted_accounts: bool,
pub cost_limit_per_key_tokens: Option<u64>,
}
@@ -225,9 +227,14 @@ fn schedule_pool_group<Candidate>(
continue;
}
if item.key_context.quota_hard_blocked
|| (pool_config.skip_exhausted_accounts && item.key_context.quota_exhausted)
{
// A quota snapshot is an account-level admission signal, not merely a
// ranking hint. Continuing to schedule a member whose quota is known to
// be exhausted causes a request-wide retry storm (the upstream returns
// 429 for every attempt). `quota_hard_blocked` remains available for
// providers that can distinguish an explicit permanent block, but every
// active exhaustion must be removed from the request's candidate set;
// reset-aware provider adapters clear the signal once capacity returns.
if item.key_context.quota_hard_blocked || item.key_context.quota_exhausted {
skipped.push(PoolSkippedCandidate {
candidate: item.candidate,
skip_reason: POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
@@ -928,6 +935,29 @@ mod tests {
);
}
#[test]
fn pool_scheduler_skips_exhausted_accounts_even_when_legacy_flag_is_false() {
let ready = sample_candidate("provider-pool", "endpoint-1", "key-ready", 10, true);
let mut exhausted =
sample_candidate("provider-pool", "endpoint-1", "key-exhausted", 10, true);
exhausted.key_context.quota_exhausted = true;
let outcome = run_pool_scheduler(vec![ready, exhausted], &BTreeMap::new(), "seed");
assert_eq!(
outcome
.candidates
.iter()
.map(|item| item.candidate.as_str())
.collect::<Vec<_>>(),
vec!["key-ready"]
);
assert_eq!(
outcome.skipped_candidates[0].skip_reason,
POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
);
}
#[test]
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true)
crates/aether-provider/pool/src/lib.rs
+282 -2
View File
@@ -35,6 +35,7 @@ pub use providers::{
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_quota_hard_blocked,
provider_pool_key_model_quota_exhausted, provider_pool_key_model_quota_hard_blocked,
provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot,
provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at,
provider_pool_quota_snapshot_updated_at,
@@ -561,7 +562,7 @@ mod tests {
}
})));
let signals = service.member_signals("windsurf", &key, None);
let signals = service.member_signals("windsurf", &key, None, None);
assert!(!signals.quota_exhausted);
}
@@ -588,7 +589,7 @@ mod tests {
}
}));
let signals = service.member_signals("windsurf", &key, None);
let signals = service.member_signals("windsurf", &key, None, None);
assert!(signals.quota_exhausted);
}
@@ -876,6 +877,285 @@ mod tests {
));
}
#[test]
fn model_quota_windows_are_isolated_without_provider_specific_names() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{
"code": "alpha_short",
"quota_group": "alpha",
"model": "vendor-alpha-model",
"used_ratio": 1.0,
"is_exhausted": true
},
{
"code": "alpha_long",
"quota_group": "alpha",
"model": "vendor-alpha-model",
"used_ratio": 0.2,
"is_exhausted": false
},
{
"code": "beta_short",
"quota_group": "beta",
"model": "vendor-beta-model",
"used_ratio": 1.0,
"is_exhausted": true
}
]
}
}));
let alpha = service.member_signals(
"codex",
&key,
None,
Some("vendor-alpha-model"),
);
let beta = service.member_signals("codex", &key, None, Some("vendor-beta-model"));
assert!(!alpha.quota_exhausted, "one available alpha window must keep it usable");
assert!(beta.quota_exhausted);
assert!(!beta.quota_hard_blocked);
}
#[test]
fn legacy_family_prefix_is_matched_by_model_tokens() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{ "code": "alpha_weekly", "used_ratio": 1.0, "is_exhausted": true },
{ "code": "default_weekly", "used_ratio": 0.1, "is_exhausted": false }
]
}
}));
let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-v2"));
assert!(alpha.quota_exhausted);
// An unrelated model has no identifiable bucket and therefore falls
// back to the account-level snapshot rather than guessing a family.
let unknown = service.member_signals("codex", &key, None, Some("vendor-gamma-v2"));
assert!(unknown.quota_exhausted);
}
#[test]
fn compact_model_bucket_identity_matches_request_tokens() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{
"code": "additional_0_primary_window",
"scope": "model",
"model": "spark",
"used_ratio": 0.0,
"is_exhausted": false
}
]
}
}));
let signals = service.member_signals(
"codex",
&key,
None,
Some("gpt-5.3-codex-spark"),
);
assert!(
!signals.quota_exhausted,
"a compact bucket name should match a token in the selected model"
);
}
#[test]
fn model_family_token_matches_versioned_alias_without_hardcoding_name() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "additional_0_primary_window",
"scope": "model",
"model": "gpt-5.3-codex-spark",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals(
"codex",
&key,
None,
Some("gpt-5.4-codex-spark"),
);
assert!(signals.quota_exhausted);
}
#[test]
fn model_only_snapshot_does_not_poison_account_fallback() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "model:vendor-alpha",
"scope": "model",
"model": "vendor-alpha",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, None);
assert!(
!signals.quota_exhausted,
"account-level inspection must ignore model-only buckets"
);
}
#[test]
fn model_map_quota_is_resolved_without_materialized_windows() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "grok",
"exhausted": false,
"quota_by_model": {
"vendor-alpha": {
"remaining": 0.0,
"total": 10.0
},
"vendor-beta": {
"remaining": 5.0,
"total": 10.0
}
}
}
}));
let alpha = service.member_signals("grok", &key, None, Some("vendor-alpha"));
let beta = service.member_signals("grok", &key, None, Some("vendor-beta"));
assert!(alpha.quota_exhausted);
assert!(!beta.quota_exhausted);
}
#[test]
fn raw_provider_metadata_model_bucket_overrides_account_snapshot() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(Some(json!({
"codex": {
"updated_at": 1_700_000_001u64,
"additional_quota_windows": [{
"scope": "model",
"model": "future-spark",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
})));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "primary",
"scope": "account",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-spark"));
assert!(
!signals.quota_exhausted,
"a newer model bucket in provider metadata must not inherit an exhausted account bucket"
);
}
#[test]
fn model_quota_availability_suppresses_account_hard_block() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"allowed": false,
"limit_reached": true,
"exhausted": true,
"windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-model"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn newer_model_quota_observation_wins_over_stale_snapshot() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(Some(json!({
"codex": {
"updated_at": 1_700_000_200u64,
"additional_quota_windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
})));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"observed_at": 1_700_000_100u64,
"exhausted": true,
"windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-model"));
assert!(!signals.quota_exhausted);
}
#[test]
fn codex_explicit_quota_block_is_hard_until_reset() {
let now = std::time::SystemTime::now()
crates/aether-provider/pool/src/provider.rs
+12 -1
View File
@@ -8,7 +8,8 @@ use crate::capability::{ProviderPoolCapabilities, ProviderPoolCapability};
use crate::plan::{derive_plan_tier, normalize_provider_plan_tier};
use crate::quota::{
provider_pool_account_blocked, provider_pool_quota_reset_seconds,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_quota_usage_ratio,
};
#[derive(Debug, Clone)]
@@ -16,6 +17,11 @@ pub struct ProviderPoolMemberInput<'a> {
pub provider_type: &'a str,
pub key: &'a StoredProviderCatalogKey,
pub auth_config: Option<&'a Map<String, Value>>,
/// The provider-side model selected for the current request. Quota
/// snapshots may contain several independent windows, so adapters use
/// this value to select the applicable bucket instead of treating the
/// whole account as one quota.
pub provider_model_name: Option<&'a str>,
}
pub trait ProviderPoolAdapter: Send + Sync {
@@ -70,6 +76,11 @@ pub trait ProviderPoolAdapter: Send + Sync {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
.unwrap_or(false)
}
crates/aether-provider/pool/src/providers/antigravity.rs
+17
View File
@@ -6,7 +6,9 @@ use serde_json::json;
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
ProviderPoolMemberInput,
};
use crate::quota::provider_pool_model_quota_exhausted;
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
@@ -26,6 +28,21 @@ impl ProviderPoolAdapter for AntigravityProviderPoolAdapter {
}
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
input
.provider_model_name
.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
})
.unwrap_or_else(|| {
crate::quota::provider_pool_quota_snapshot_exhausted_decision(
input.key,
input.provider_type,
)
.unwrap_or(false)
})
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
crates/aether-provider/pool/src/providers/chatgpt_web.rs
+7 -1
View File
@@ -12,7 +12,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -40,6 +41,11 @@ impl ProviderPoolAdapter for ChatGptWebProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
crates/aether-provider/pool/src/providers/codex.rs
+12 -2
View File
@@ -12,8 +12,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_member_quota_snapshot, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -49,6 +49,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(quota_snapshot) =
provider_pool_member_quota_snapshot(input.key, input.provider_type)
{
@@ -76,6 +81,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
fn quota_hard_blocked(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if input.provider_model_name.is_some_and(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model).is_some()
}) {
return false;
}
codex_explicit_quota_block_active(input.key, input.provider_type)
}
crates/aether-provider/pool/src/providers/grok.rs
+7 -1
View File
@@ -9,7 +9,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
@@ -44,6 +45,11 @@ impl ProviderPoolAdapter for GrokProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
crates/aether-provider/pool/src/providers/kiro.rs
+7 -2
View File
@@ -12,8 +12,8 @@ use crate::provider::{
};
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -46,6 +46,11 @@ impl ProviderPoolAdapter for KiroProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
crates/aether-provider/pool/src/providers/windsurf.rs
+7 -1
View File
@@ -13,7 +13,8 @@ use crate::provider::{
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_member_quota_snapshot,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_metadata_bucket, provider_pool_model_quota_exhausted,
provider_pool_quota_snapshot_exhausted_decision,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -50,6 +51,11 @@ impl ProviderPoolAdapter for WindsurfProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if windsurf_quota_snapshot_hard_exhausted(input.key, input.provider_type) {
return true;
}
crates/aether-provider/pool/src/quota.rs
+415 -7
View File
@@ -15,6 +15,7 @@ pub fn provider_pool_key_account_quota_exhausted(
provider_type,
key,
auth_config: None,
provider_model_name: None,
})
}
@@ -27,6 +28,25 @@ pub fn provider_pool_key_quota_hard_blocked(
provider_type,
key,
auth_config: None,
provider_model_name: None,
})
}
/// Model-aware hard-block lookup for pre-scheduler candidate filtering. Some
/// providers expose permanent account flags alongside independent model
/// buckets; adapters can suppress the account flag when the selected model
/// has its own usable quota.
pub fn provider_pool_key_model_quota_hard_blocked(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> bool {
let adapter = ProviderPoolService::with_builtin_adapters().adapter(provider_type);
adapter.quota_hard_blocked(&ProviderPoolMemberInput {
provider_type,
key,
auth_config: None,
provider_model_name: Some(provider_model_name),
})
}
@@ -44,6 +64,331 @@ pub fn provider_pool_member_quota_snapshot<'a>(
.then_some(quota_snapshot)
}
/// Resolve exhaustion for the quota bucket applicable to one provider model.
///
/// Providers are free to expose quota windows in different shapes. Newer
/// snapshots should put an explicit `model`/`models` (or `quota_group`) on a
/// window; legacy Codex snapshots use a family prefix in `code` (for example
/// `spark_5h`). We deliberately do not name any product or model here: the
/// resolver compares the metadata supplied by the provider with the selected
/// model and only falls back to account-level exhaustion when no model bucket
/// can be identified.
pub(crate) fn provider_pool_model_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
let requested = provider_pool_identifier_tokens(provider_model_name);
if requested.is_empty() {
return None;
}
// Prefer the materialized status snapshot, but also inspect the raw
// provider metadata. A quota refresh and a request can race, leaving the
// latter newer than the snapshot; resolving both avoids falling back to an
// account-wide signal and incorrectly blocking an unrelated model bucket.
let sources = [
provider_pool_member_quota_snapshot(key, provider_type),
provider_pool_metadata_bucket(key.upstream_metadata.as_ref(), provider_type),
];
let mut resolved = None::<(Option<u64>, bool)>;
for source in sources.into_iter().flatten() {
let windows = provider_pool_collect_quota_windows(source);
if windows.is_empty() {
continue;
}
let observed_at = provider_pool_timestamp_unix_secs(source.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(source.get("updated_at")));
let model_matches = windows
.iter()
.filter(|window| {
provider_pool_window_explicitly_matches_model(window, provider_model_name)
})
.collect::<Vec<_>>();
if !model_matches.is_empty() {
let exhausted = provider_pool_matching_windows_exhausted(
model_matches,
observed_at,
);
if resolved.is_none()
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
}
continue;
}
// Legacy snapshots may not carry a model field. Match an opaque
// family token (the prefix before `_`/`:` in `code`, or an explicit
// family key) against the model's tokens. This keeps independent
// windows isolated without baking in names such as "spark".
let family_matches = windows
.iter()
.filter(|window| provider_pool_window_family_matches_model(window, &requested))
.collect::<Vec<_>>();
if !family_matches.is_empty() {
let exhausted = provider_pool_matching_windows_exhausted(
family_matches,
observed_at,
);
if resolved.is_none()
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
}
}
}
resolved.map(|(_, exhausted)| exhausted)
}
fn provider_pool_should_replace_model_quota_resolution(
previous_observed_at: Option<u64>,
next_observed_at: Option<u64>,
) -> bool {
match (previous_observed_at, next_observed_at) {
(Some(previous), Some(next)) => next >= previous,
(None, Some(_)) => true,
(Some(_), None) => false,
// Preserve source order when neither side carries freshness metadata;
// the materialized status snapshot is preferred over raw metadata.
(None, None) => false,
}
}
/// Public adapter-independent model quota lookup used by schedulers that need
/// to prefilter candidates before constructing provider-pool signals.
pub fn provider_pool_key_model_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
provider_pool_model_quota_exhausted(key, provider_type, provider_model_name)
}
fn provider_pool_matching_windows_exhausted(
windows: Vec<&Map<String, Value>>,
snapshot_observed_at: Option<u64>,
) -> bool {
let now_unix_secs = provider_pool_current_unix_secs();
// A request can use a bucket as long as at least one of its independent
// windows still has capacity (e.g. a short and a long rolling window).
windows.iter().all(|window| {
provider_pool_quota_window_is_exhausted(window)
&& !now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
})
})
}
fn provider_pool_window_explicitly_matches_model(
window: &Map<String, Value>,
requested_model: &str,
) -> bool {
let requested = provider_pool_normalize_identifier(requested_model);
let requested_tokens = provider_pool_identifier_tokens(requested_model);
if requested.is_empty() {
return false;
}
[
"model",
"model_name",
"model_id",
"quota_model",
"quota_model_name",
"target_model",
"limit_name",
]
.iter()
.filter_map(|key| window.get(*key))
.any(|value| match value {
Value::String(value) => {
provider_pool_identifiers_match(&requested, value, &requested_tokens)
}
Value::Array(values) => values
.iter()
.filter_map(Value::as_str)
.any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens)),
_ => false,
}) || ["models", "model_ids"]
.iter()
.filter_map(|key| window.get(*key).and_then(Value::as_array))
.flatten()
.filter_map(Value::as_str)
.any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens))
}
fn provider_pool_window_family_matches_model(
window: &Map<String, Value>,
requested_tokens: &std::collections::BTreeSet<String>,
) -> bool {
let explicit_scope = window
.get("scope")
.and_then(Value::as_str)
.map(|scope| scope.trim().to_ascii_lowercase());
// A model-scoped window without an explicit model must not accidentally
// match a token from its opaque code.
if explicit_scope.as_deref() == Some("model")
|| provider_pool_window_has_explicit_model(window)
{
return false;
}
let mut families = Vec::new();
for key in ["quota_group", "quota_family", "family", "bucket"] {
if let Some(value) = window.get(key).and_then(Value::as_str) {
families.push(value.to_string());
}
}
if let Some(code) = window.get("code").and_then(Value::as_str) {
let code = code.trim();
if let Some((prefix, _)) = code.split_once(['_', ':', '/']) {
families.push(prefix.to_string());
}
}
families.into_iter().any(|family| {
let normalized = provider_pool_normalize_identifier(&family);
if normalized.is_empty()
|| [
"account",
"quota",
"window",
"primary",
"secondary",
"rate",
"reset",
]
.iter()
.any(|generic| normalized == *generic)
{
return false;
}
requested_tokens.iter().any(|token| {
token.len() >= 3
&& (token == &normalized
|| token.contains(&normalized)
|| normalized.contains(token))
})
})
}
fn provider_pool_identifiers_match(
requested: &str,
candidate: &str,
requested_tokens: &std::collections::BTreeSet<String>,
) -> bool {
let candidate_tokens = provider_pool_identifier_tokens(candidate);
let candidate = provider_pool_normalize_identifier(candidate);
if candidate.is_empty() {
return false;
}
requested == candidate
|| candidate_tokens
.iter()
.any(|token| {
requested_tokens.contains(token) && provider_pool_is_specific_model_token(token)
})
// Handle compact upstream identifiers such as `spark` embedded in a
// provider model name (`vendor-codex-spark`) while avoiding accidental
// one/two-character matches.
|| (candidate.len() >= 4
&& requested.len() >= 6
&& (requested.contains(&candidate) || candidate.contains(requested)))
}
fn provider_pool_is_specific_model_token(token: &str) -> bool {
token.len() >= 4
&& !token.chars().all(|character| character.is_ascii_digit())
&& ![
"auto", "base", "claude", "codex", "default", "fast", "flash", "free",
"gemini", "gpt", "latest", "mini", "model", "plus", "pro", "reasoning",
"team", "think", "thinking", "vendor",
]
.contains(&token)
}
fn provider_pool_window_has_explicit_model(window: &Map<String, Value>) -> bool {
[
"model",
"model_name",
"model_id",
"quota_model",
"quota_model_name",
"target_model",
"models",
"model_ids",
]
.iter()
.any(|key| match window.get(*key) {
Some(Value::String(value)) => !value.trim().is_empty(),
Some(Value::Array(values)) => values.iter().any(|value| {
value.as_str().is_some_and(|value| !value.trim().is_empty())
}),
_ => false,
})
}
fn provider_pool_normalize_identifier(value: &str) -> String {
value
.trim()
.to_ascii_lowercase()
.chars()
.filter(|character| character.is_ascii_alphanumeric())
.collect()
}
fn provider_pool_identifier_tokens(value: &str) -> std::collections::BTreeSet<String> {
value
.split(|character: char| !character.is_ascii_alphanumeric())
.map(|token| token.trim().to_ascii_lowercase())
.filter(|token| token.len() >= 3)
.collect()
}
/// Materialize quota windows from the small set of shapes emitted by current
/// and legacy adapters. Model maps (`quota_by_model`/`models`) are converted
/// to the same window representation used by status snapshots, with the map
/// key retained as the model identity. Keeping this normalization here means
/// provider adapters do not need to grow model-specific quota code whenever an
/// upstream introduces another independent bucket.
fn provider_pool_collect_quota_windows(source: &Map<String, Value>) -> Vec<Map<String, Value>> {
let mut windows = Vec::new();
for key in ["windows", "additional_quota_windows"] {
if let Some(values) = source.get(key).and_then(Value::as_array) {
windows.extend(values.iter().filter_map(Value::as_object).cloned());
}
}
for key in ["quota_by_model", "models", "model_quotas"] {
let Some(models) = source.get(key).and_then(Value::as_object) else {
continue;
};
for (model_name, item) in models {
let Some(item) = item.as_object() else {
continue;
};
let mut window = item.clone();
window
.entry("model".to_string())
.or_insert_with(|| json!(model_name));
window
.entry("scope".to_string())
.or_insert_with(|| json!("model"));
window
.entry("code".to_string())
.or_insert_with(|| json!(format!("model:{model_name}")));
windows.push(window);
}
}
windows
}
pub fn provider_pool_quota_snapshot_updated_at(
key: &StoredProviderCatalogKey,
provider_type: &str,
@@ -197,12 +542,51 @@ pub(crate) fn provider_pool_reset_deadline_elapsed(
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> bool {
provider_pool_json_bool(window.get("is_exhausted"))
.or_else(|| provider_pool_json_bool(window.get("exhausted")))
.or_else(|| {
provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
provider_pool_json_f64(window.get("used_ratio").or_else(|| window.get("usage_ratio")))
.map(|value| value >= 1.0 - 1e-6)
})
.or_else(|| {
provider_pool_json_f64(window.get("used_percent")).map(|value| value >= 100.0 - 1e-6)
})
.or_else(|| {
provider_pool_json_f64(
window
.get("remaining_ratio")
.or_else(|| window.get("remaining_fraction")),
)
.map(|value| value <= 1e-6)
})
.or_else(|| {
provider_pool_json_f64(window.get("remaining_percent"))
.map(|value| value <= 1e-6)
})
.or_else(|| {
let remaining = provider_pool_json_f64(
window
.get("remaining")
.or_else(|| window.get("remaining_value")),
)?;
let limit = provider_pool_json_f64(
window
.get("limit")
.or_else(|| window.get("limit_value"))
.or_else(|| window.get("total")),
)?;
(limit > 0.0).then_some(remaining <= 0.0)
})
.unwrap_or(false)
}
fn provider_pool_window_is_model_scoped(window: &Map<String, Value>) -> bool {
window
.get("scope")
.and_then(Value::as_str)
.is_some_and(|scope| scope.trim().eq_ignore_ascii_case("model"))
|| provider_pool_window_has_explicit_model(window)
}
fn provider_pool_quota_snapshot_matches_provider(
quota_snapshot: &Map<String, Value>,
provider_type: &str,
@@ -237,6 +621,18 @@ fn provider_pool_quota_snapshot_matches_provider(
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("additional_quota_windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| ["quota_by_model", "models", "model_quotas"]
.iter()
.any(|key| {
quota_snapshot
.get(*key)
.and_then(Value::as_object)
.is_some_and(|models| !models.is_empty())
})
|| quota_snapshot
.get("credits")
.and_then(Value::as_object)
@@ -265,16 +661,28 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
if let Some(windows) = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.filter(|windows| !windows.is_empty())
{
let materialized_windows = provider_pool_collect_quota_windows(quota_snapshot);
if !materialized_windows.is_empty() {
// Model-scoped windows are evaluated only when the request model
// is known. They must not turn the account-level fallback into
// an exhausted state for unrelated models.
let account_scoped_windows = materialized_windows
.iter()
.filter(|window| !provider_pool_window_is_model_scoped(window))
.collect::<Vec<_>>();
// A snapshot containing only model-scoped buckets has no
// account-wide signal to apply when the caller did not provide a
// model name (for example, an admin status listing). Do not let a
// single exhausted model poison every sibling bucket.
if account_scoped_windows.is_empty() {
return Some(false);
}
let windows = account_scoped_windows;
let mut saw_exhausted_window = false;
let mut saw_active_exhausted_window = false;
let mut windows_max_ratio = None::<f64>;
for window in windows.iter().filter_map(Value::as_object) {
for window in windows.iter() {
if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
windows_max_ratio =
Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio)));
crates/aether-provider/pool/src/service.rs
+2
View File
@@ -123,12 +123,14 @@ impl ProviderPoolService {
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&Map<String, Value>>,
provider_model_name: Option<&str>,
) -> aether_pool_core::PoolMemberSignals {
let adapter = self.adapter(provider_type);
let input = ProviderPoolMemberInput {
provider_type,
key,
auth_config,
provider_model_name,
};
adapter.member_signals(&input)
}