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

crates/aether-provider/pool/src/lib.rs

@@ -38,11 +38,11 @@ pub use providers::{
     XAI_BILLING_PATH, XAI_USER_PATH,
 };
 pub use quota::{
-    provider_pool_key_account_quota_exhausted, provider_pool_key_model_quota_exhausted,
-    provider_pool_key_model_quota_hard_blocked, provider_pool_key_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,
+    provider_pool_key_account_quota_exhausted, provider_pool_key_minimum_quota_reached,
+    provider_pool_key_model_quota_exhausted, provider_pool_key_model_quota_hard_blocked,
+    provider_pool_key_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,
 };
 pub use quota_refresh::ProviderPoolQuotaRequestSpec;
 pub use service::ProviderPoolService;
@@ -727,6 +727,275 @@ mod tests {
         assert!(unsupported.is_empty());
     }
 
+    #[test]
+    fn codex_minimum_quota_respects_boundary_and_provider() {
+        for (used_percent, expected) in [
+            (98.9, false),
+            (98.99999, false),
+            (99.0, true),
+            (99.5, true),
+            (100.0, true),
+        ] {
+            let key = sample_key(Some(json!({
+                "codex": { "primary_used_percent": used_percent }
+            })));
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", None),
+                expected,
+                "used_percent={used_percent}"
+            );
+            assert!(!provider_pool_key_minimum_quota_reached(&key, "kiro", None));
+        }
+        assert!(!provider_pool_key_minimum_quota_reached(
+            &sample_key(None),
+            "codex",
+            None
+        ));
+    }
+
+    #[test]
+    fn codex_minimum_quota_short_model_names_still_use_account_windows() {
+        for (used_percent, expected) in [(98.0, false), (99.0, true)] {
+            let key = sample_key(Some(json!({
+                "codex": { "primary_used_percent": used_percent }
+            })));
+            for model in ["o1", "o3", "", "   "] {
+                assert_eq!(
+                    provider_pool_key_minimum_quota_reached(&key, "codex", Some(model)),
+                    expected,
+                    "model={model:?}, used_percent={used_percent}"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_respects_windows_and_reset() {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .expect("system time should be after unix epoch")
+            .as_secs();
+        for (reset_at, expected) in [(now + 3600, true), (now - 60, false)] {
+            let mut key = sample_key(None);
+            key.status_snapshot = Some(json!({
+                "quota": {
+                    "provider_type": "codex",
+                    "updated_at": now - 600,
+                    "windows": [
+                        { "code": "5h", "used_ratio": 0.2 },
+                        { "code": "weekly", "used_ratio": 0.99, "reset_at": reset_at }
+                    ]
+                }
+            }));
+            for model in [None, Some("gpt-5.4")] {
+                assert_eq!(
+                    provider_pool_key_minimum_quota_reached(&key, "codex", model),
+                    expected
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_checks_either_window_and_relative_resets() {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .expect("system time should be after unix epoch")
+            .as_secs();
+        for prefix in ["primary", "secondary"] {
+            for (reset_seconds, expected) in [(3600, true), (60, false)] {
+                let key = sample_key(Some(json!({
+                    "codex": {
+                        "updated_at": now - 600,
+                        "allowed": true,
+                        "limit_reached": false,
+                        format!("{prefix}_used_percent"): 99.0,
+                        format!("{prefix}_reset_after_seconds"): reset_seconds
+                    }
+                })));
+                for model in [None, Some("gpt-5.4")] {
+                    assert_eq!(
+                        provider_pool_key_minimum_quota_reached(&key, "codex", model),
+                        expected,
+                        "prefix={prefix}, reset_seconds={reset_seconds}, model={model:?}"
+                    );
+                }
+                assert!(!provider_pool_key_account_quota_exhausted(&key, "codex"));
+            }
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_uses_newest_source_with_or_without_model() {
+        for (snapshot_at, metadata_at, metadata_wins) in [
+            (Some(200), Some(100), false),
+            (Some(100), Some(200), true),
+            (None, None, false),
+            (Some(100), None, false),
+            (None, Some(100), true),
+        ] {
+            for snapshot_reached in [false, true] {
+                let mut key = sample_key(Some(json!({
+                    "codex": {
+                        "updated_at": metadata_at,
+                        "primary_used_percent": if snapshot_reached { 20.0 } else { 99.0 }
+                    }
+                })));
+                key.status_snapshot = Some(json!({
+                    "quota": {
+                        "provider_type": "codex",
+                        "observed_at": snapshot_at,
+                        "windows": [{
+                            "code": "5h",
+                            "used_ratio": if snapshot_reached { 0.99 } else { 0.2 },
+                            "is_exhausted": false
+                        }]
+                    }
+                }));
+                for model in [None, Some("gpt-5.4")] {
+                    assert_eq!(
+                        provider_pool_key_minimum_quota_reached(&key, "codex", model),
+                        snapshot_reached != metadata_wins,
+                        "snapshot_at={snapshot_at:?}, metadata_at={metadata_at:?}, model={model:?}"
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_isolates_model_buckets_and_checks_each_model_window() {
+        for explicit_model in [false, true] {
+            let mut key = sample_key(None);
+            key.status_snapshot = Some(json!({
+                "quota": {
+                    "provider_type": "codex",
+                    "windows": [
+                        { "code": "5h", "used_ratio": 0.2 },
+                        {
+                            "code": "spark_5h", "used_ratio": 0.99,
+                            "model": if explicit_model { Some("spark") } else { None },
+                            "is_exhausted": false
+                        },
+                        {
+                            "code": "spark_weekly", "used_ratio": 0.2,
+                            "model": if explicit_model { Some("spark") } else { None },
+                            "is_exhausted": false
+                        }
+                    ]
+                }
+            }));
+            assert!(provider_pool_key_minimum_quota_reached(
+                &key,
+                "codex",
+                Some("gpt-5.3-codex-spark")
+            ));
+            for model in [None, Some("gpt-5.4")] {
+                assert!(!provider_pool_key_minimum_quota_reached(
+                    &key, "codex", model
+                ));
+            }
+            assert_eq!(
+                provider_pool_key_model_quota_exhausted(&key, "codex", "gpt-5.3-codex-spark"),
+                Some(false)
+            );
+        }
+        for (account_percent, spark_percent) in [(99.0, 20.0), (20.0, 99.0)] {
+            let key = sample_key(Some(json!({
+                "codex": {
+                    "primary_used_percent": account_percent,
+                    "spark_primary_used_percent": spark_percent,
+                    "spark_secondary_used_percent": 20.0
+                }
+            })));
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
+                spark_percent == 99.0
+            );
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.4")),
+                account_percent == 99.0
+            );
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_keeps_model_bucket_when_account_metadata_is_newer() {
+        for (snapshot_ratio, account_percent) in [(0.2, 99.0), (0.99, 20.0)] {
+            let mut key = sample_key(Some(json!({
+                "codex": { "updated_at": 200, "primary_used_percent": account_percent }
+            })));
+            key.status_snapshot = Some(json!({
+                "quota": {
+                    "provider_type": "codex",
+                    "observed_at": 100,
+                    "windows": [{ "code": "spark_5h", "used_ratio": snapshot_ratio }]
+                }
+            }));
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
+                snapshot_ratio == 0.99
+            );
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.4")),
+                account_percent == 99.0
+            );
+
+            key.upstream_metadata.as_mut().unwrap()["codex"]["spark_primary_used_percent"] =
+                json!(account_percent);
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
+                account_percent == 99.0,
+                "the newer observation for the same model bucket must win"
+            );
+        }
+    }
+
+    #[test]
+    fn codex_minimum_quota_supports_remaining_values_and_ignores_unknown_data() {
+        for (metrics, expected) in [
+            (json!({ "remaining_ratio": 0.01 }), true),
+            (json!({ "remaining_percent": "1" }), true),
+            (json!({ "remaining": 1, "limit": 100 }), true),
+            (json!({ "used_percent": "99" }), true),
+            (json!({ "used_ratio": null }), false),
+            (json!({ "used_ratio": "NaN" }), false),
+            (json!({ "remaining": 0, "limit": 0 }), false),
+            (json!({ "remaining_ratio": 0.01001 }), false),
+        ] {
+            let key = sample_key(Some(json!({
+                "codex": { "quota_by_model": { "spark": metrics } }
+            })));
+            assert_eq!(
+                provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
+                expected,
+                "metrics={metrics}"
+            );
+            assert!(!provider_pool_key_minimum_quota_reached(
+                &key, "codex", None
+            ));
+        }
+        let disabled = sample_key(Some(json!({
+            "codex": { "primary_used_percent": 99.0, "primary_window_minutes": 0 }
+        })));
+        assert!(!provider_pool_key_minimum_quota_reached(
+            &disabled, "codex", None
+        ));
+
+        let mut mismatched = sample_key(None);
+        mismatched.status_snapshot = Some(json!({
+            "quota": {
+                "provider_type": "kiro",
+                "windows": [{ "code": "weekly", "used_ratio": 0.99 }]
+            }
+        }));
+        assert!(!provider_pool_key_minimum_quota_reached(
+            &mismatched,
+            "codex",
+            None
+        ));
+    }
+
     #[test]
     fn provider_quota_exhaustion_is_adapter_owned() {
         assert!(provider_pool_key_account_quota_exhausted(

crates/aether-provider/pool/src/quota.rs

@@ -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