i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 08:27:46 +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/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()