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

feat(codex): add optional minimum quota reserve for pool scheduling

Browse Source
This commit is contained in:
ZheFox
2026-09-23 14:44:29 +08:00
parent 2a9d8d3b25
commit 1a4eba1005
16 changed files with 951 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-provider/pool/src/quota.rs
+125 -10
View File
@@ -78,8 +78,17 @@ pub(crate) fn provider_pool_model_quota_exhausted(
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
let requested = provider_pool_identifier_tokens(provider_model_name);
if requested.is_empty() {
provider_pool_quota_reaches_reserve(key, provider_type, Some(provider_model_name), 0.0)
}
fn provider_pool_quota_reaches_reserve(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: Option<&str>,
reserve_ratio: f64,
) -> Option<bool> {
let requested = provider_model_name.map(provider_pool_identifier_tokens);
if reserve_ratio <= 0.0 && requested.as_ref().is_some_and(|tokens| tokens.is_empty()) {
return None;
}
@@ -92,8 +101,38 @@ pub(crate) fn provider_pool_model_quota_exhausted(
provider_pool_metadata_bucket(key.upstream_metadata.as_ref(), provider_type),
];
let mut resolved = None::<(Option<u64>, bool)>;
let mut resolved_specific_bucket = false;
for source in sources.into_iter().flatten() {
let windows = provider_pool_collect_quota_windows(source);
let mut windows = provider_pool_collect_quota_windows(source);
if reserve_ratio > 0.0 && provider_type.trim().eq_ignore_ascii_case("codex") {
// Raw refresh metadata can be newer than the materialized windows.
// Include all four legacy slots before selecting the model bucket.
for prefix in ["primary", "secondary", "spark_primary", "spark_secondary"] {
let Some(used_percent) = source.get(&format!("{prefix}_used_percent")) else {
continue;
};
if provider_pool_json_f64(source.get(&format!("{prefix}_window_minutes")))
== Some(0.0)
{
continue;
}
let mut window = Map::from_iter([
("code".to_string(), json!(prefix)),
("used_percent".to_string(), used_percent.clone()),
]);
for field in [
"reset_at",
"next_reset_at",
"reset_seconds",
"reset_after_seconds",
] {
if let Some(value) = source.get(&format!("{prefix}_{field}")) {
window.insert(field.to_string(), value.clone());
}
}
windows.push(window);
}
}
if windows.is_empty() {
continue;
}
@@ -102,19 +141,28 @@ pub(crate) fn provider_pool_model_quota_exhausted(
let model_matches = windows
.iter()
.filter(|window| {
provider_pool_window_explicitly_matches_model(window, provider_model_name)
provider_model_name.is_some_and(|model| {
provider_pool_window_explicitly_matches_model(window, model)
})
})
.collect::<Vec<_>>();
if !model_matches.is_empty() {
let exhausted =
provider_pool_explicit_model_windows_exhausted(model_matches, observed_at);
let exhausted = if reserve_ratio > 0.0 {
// Reserving quota protects every applicable rate-limit window,
// including short and weekly limits for an explicit model.
provider_pool_any_window_exhausted(model_matches, observed_at, reserve_ratio)
} else {
provider_pool_explicit_model_windows_exhausted(model_matches, observed_at)
};
if resolved.is_none()
|| (reserve_ratio > 0.0 && !resolved_specific_bucket)
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
resolved_specific_bucket = true;
}
continue;
}
@@ -125,21 +173,34 @@ pub(crate) fn provider_pool_model_quota_exhausted(
// windows isolated without baking in names such as "spark".
let family_matches = windows
.iter()
.filter(|window| provider_pool_window_family_matches_model(window, &requested))
.filter(|window| {
requested
.as_ref()
.is_some_and(|tokens| provider_pool_window_family_matches_model(window, tokens))
})
.collect::<Vec<_>>();
if !family_matches.is_empty() {
let exhausted = provider_pool_any_window_exhausted(family_matches, observed_at);
let exhausted =
provider_pool_any_window_exhausted(family_matches, observed_at, reserve_ratio);
if resolved.is_none()
|| (reserve_ratio > 0.0 && !resolved_specific_bucket)
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
resolved_specific_bucket = true;
}
continue;
}
// A newer account observation does not update an independent model
// bucket. Only compare freshness between applicable model sources.
if reserve_ratio > 0.0 && resolved_specific_bucket {
continue;
}
// Account-scoped windows (for example the ordinary weekly and short
// windows emitted by Codex) apply to every model that has no more
// specific family. Restrict this fallback to well-known structural
@@ -150,7 +211,8 @@ pub(crate) fn provider_pool_model_quota_exhausted(
.filter(|window| provider_pool_window_is_generic(window))
.collect::<Vec<_>>();
if !generic_matches.is_empty() {
let exhausted = provider_pool_any_window_exhausted(generic_matches, observed_at);
let exhausted =
provider_pool_any_window_exhausted(generic_matches, observed_at, reserve_ratio);
if resolved.is_none()
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
@@ -209,16 +271,69 @@ fn provider_pool_explicit_model_windows_exhausted(
fn provider_pool_any_window_exhausted(
windows: Vec<&Map<String, Value>>,
snapshot_observed_at: Option<u64>,
reserve_ratio: f64,
) -> bool {
let now_unix_secs = provider_pool_current_unix_secs();
windows.iter().any(|window| {
provider_pool_quota_window_is_exhausted(window)
provider_pool_quota_window_reaches_reserve(window, reserve_ratio)
&& !now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
})
})
}
/// Whether an applicable Codex window has at most 1% remaining. Missing quota
/// data and windows whose reset has elapsed do not trigger this opt-in guard.
pub fn provider_pool_key_minimum_quota_reached(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: Option<&str>,
) -> bool {
if !provider_type.trim().eq_ignore_ascii_case("codex") {
return false;
}
provider_pool_quota_reaches_reserve(key, provider_type, provider_model_name, 0.01)
.unwrap_or(false)
}
fn provider_pool_quota_window_reaches_reserve(
window: &Map<String, Value>,
reserve_ratio: f64,
) -> bool {
if provider_pool_quota_window_is_exhausted(window) {
return true;
}
if reserve_ratio <= 0.0 {
return false;
}
let used_ratio = provider_pool_json_f64(window.get("used_ratio"))
.or_else(|| provider_pool_json_f64(window.get("usage_ratio")))
.or_else(|| provider_pool_json_f64(window.get("used_percent")).map(|value| value / 100.0))
.or_else(|| {
provider_pool_json_f64(window.get("remaining_ratio"))
.or_else(|| provider_pool_json_f64(window.get("remaining_fraction")))
.map(|value| 1.0 - value)
})
.or_else(|| {
provider_pool_json_f64(window.get("remaining_percent")).map(|value| 1.0 - value / 100.0)
})
.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(1.0 - remaining / limit)
});
used_ratio.is_some_and(|used_ratio| used_ratio >= 1.0 - reserve_ratio)
}
fn provider_pool_window_is_generic(window: &Map<String, Value>) -> bool {
if provider_pool_window_has_explicit_model(window)
|| window