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

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
+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(),
+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]