Fix provider pool exhaustion scheduling

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

@@ -380,6 +380,81 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .expect("system time should be after unix epoch")
+            .as_secs();
+
+        let mut expired = sample_key(None);
+        expired.status_snapshot = Some(json!({
+            "quota": {
+                "version": 2,
+                "provider_type": "codex",
+                "code": "exhausted",
+                "exhausted": true,
+                "updated_at": now.saturating_sub(600),
+                "windows": [{
+                    "code": "5h",
+                    "used_ratio": 1.0,
+                    "reset_at": now.saturating_sub(60),
+                    "is_exhausted": true
+                }]
+            }
+        }));
+        assert!(!provider_pool_key_account_quota_exhausted(
+            &expired, "codex"
+        ));
+
+        let mut active = sample_key(None);
+        active.status_snapshot = Some(json!({
+            "quota": {
+                "version": 2,
+                "provider_type": "codex",
+                "code": "exhausted",
+                "exhausted": true,
+                "updated_at": now,
+                "windows": [{
+                    "code": "5h",
+                    "used_ratio": 1.0,
+                    "reset_at": now.saturating_add(3600),
+                    "is_exhausted": true
+                }]
+            }
+        }));
+        assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
+    }
+
+    #[test]
+    fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .expect("system time should be after unix epoch")
+            .as_secs();
+
+        assert!(!provider_pool_key_account_quota_exhausted(
+            &sample_key(Some(json!({
+                "codex": {
+                    "updated_at": now.saturating_sub(600),
+                    "primary_used_percent": 100.0,
+                    "primary_reset_at": now.saturating_sub(60)
+                }
+            }))),
+            "codex",
+        ));
+        assert!(provider_pool_key_account_quota_exhausted(
+            &sample_key(Some(json!({
+                "codex": {
+                    "updated_at": now,
+                    "primary_used_percent": 100.0,
+                    "primary_reset_at": now.saturating_add(3600)
+                }
+            }))),
+            "codex",
+        ));
+    }
+
     #[test]
     fn grok_quota_tier_boundaries_match_pool_modes() {
         assert_eq!(

crates/aether-provider-pool/src/providers/chatgpt_web.rs

@@ -10,8 +10,9 @@ use crate::provider::{
     ProviderPoolMemberInput,
 };
 use crate::quota::{
-    provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
-    provider_pool_quota_snapshot_exhausted_decision,
+    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,
 };
 use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
 
@@ -255,6 +256,19 @@ fn infer_chatgpt_web_image_quota_limit(
 }
 
 pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
+    if provider_pool_current_unix_secs().is_some_and(|now| {
+        let mut image_quota = Map::new();
+        if let Some(value) = bucket.get("image_quota_reset_at") {
+            image_quota.insert("reset_at".to_string(), value.clone());
+        }
+        provider_pool_reset_deadline_elapsed(
+            &image_quota,
+            provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
+            now,
+        )
+    }) {
+        return false;
+    }
     if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
         return true;
     }

crates/aether-provider-pool/src/providers/codex.rs

@@ -10,8 +10,9 @@ use crate::provider::{
     ProviderPoolMemberInput,
 };
 use crate::quota::{
-    provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
-    provider_pool_quota_snapshot_exhausted_decision,
+    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,
 };
 use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
 
@@ -121,6 +122,37 @@ pub fn build_codex_pool_quota_request(
     })
 }
 
+fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
+    let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
+        return false;
+    };
+    let mut window = Map::new();
+    for (target, source) in [
+        ("reset_at", format!("{prefix}_reset_at")),
+        ("next_reset_at", format!("{prefix}_next_reset_at")),
+        ("reset_seconds", format!("{prefix}_reset_seconds")),
+        (
+            "reset_after_seconds",
+            format!("{prefix}_reset_after_seconds"),
+        ),
+    ] {
+        if let Some(value) = bucket.get(source.as_str()) {
+            window.insert(target.to_string(), value.clone());
+        }
+    }
+    provider_pool_reset_deadline_elapsed(
+        &window,
+        provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
+        now_unix_secs,
+    )
+}
+
+fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
+    let used_percent_key = format!("{prefix}_used_percent");
+    provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
+        .is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
+}
+
 pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
     if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
         return false;
@@ -130,7 +162,6 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
     if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
         return true;
     }
-    provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
-        || provider_pool_json_f64(bucket.get("secondary_used_percent"))
-            .is_some_and(|value| value >= 100.0)
+    codex_window_used_percent_exhausted(bucket, "primary")
+        || codex_window_used_percent_exhausted(bucket, "secondary")
 }

crates/aether-provider-pool/src/providers/grok.rs

@@ -7,8 +7,9 @@ use crate::provider::{
     ProviderPoolMemberInput,
 };
 use crate::quota::{
-    provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
-    provider_pool_quota_snapshot_exhausted_decision,
+    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,
 };
 
 pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
@@ -191,6 +192,8 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
 
     let mut model_count = 0usize;
     let mut exhausted_count = 0usize;
+    let now_unix_secs = provider_pool_current_unix_secs();
+    let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
     for (model_key, item) in models.iter() {
         if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
             continue;
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
             continue;
         }
         model_count += 1;
-        if provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
+        let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
             || provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
             || provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
             || provider_pool_json_f64(item.get("remaining_fraction"))
-                .is_some_and(|value| value <= 0.0)
-        {
+                .is_some_and(|value| value <= 0.0);
+        let reset_elapsed = now_unix_secs
+            .is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
+        if quota_exhausted && !reset_elapsed {
             exhausted_count += 1;
         }
     }

crates/aether-provider-pool/src/providers/kiro.rs

@@ -11,8 +11,9 @@ use crate::provider::{
     ProviderPoolMemberInput,
 };
 use crate::quota::{
-    provider_pool_json_f64, provider_pool_metadata_bucket,
-    provider_pool_quota_snapshot_exhausted_decision,
+    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,
 };
 use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
 
@@ -151,6 +152,15 @@ fn normalize_kiro_version(value: &str) -> &str {
 }
 
 pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
+    if provider_pool_current_unix_secs().is_some_and(|now| {
+        provider_pool_reset_deadline_elapsed(
+            bucket,
+            provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
+            now,
+        )
+    }) {
+        return false;
+    }
     if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
         return true;
     }

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

@@ -1,3 +1,5 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
 use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
 use serde_json::{json, Map, Value};
 
@@ -135,7 +137,7 @@ pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
     .filter(|value| value.is_finite())
 }
 
-fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
+pub(crate) fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
     let mut timestamp = provider_pool_json_f64(value)?;
     if timestamp <= 0.0 {
         return None;
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
     Some(timestamp as u64)
 }
 
+pub(crate) fn provider_pool_current_unix_secs() -> Option<u64> {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .ok()
+        .map(|duration| duration.as_secs())
+}
+
+fn provider_pool_reset_deadline_unix_secs(
+    item: &Map<String, Value>,
+    fallback_observed_at: Option<u64>,
+) -> Option<u64> {
+    provider_pool_timestamp_unix_secs(item.get("reset_at"))
+        .or_else(|| provider_pool_timestamp_unix_secs(item.get("next_reset_at")))
+        .or_else(|| {
+            let reset_seconds = provider_pool_json_f64(item.get("reset_seconds"))
+                .or_else(|| provider_pool_json_f64(item.get("reset_after_seconds")))?;
+            if reset_seconds < 0.0 {
+                return None;
+            }
+            let base = provider_pool_timestamp_unix_secs(item.get("observed_at"))
+                .or_else(|| provider_pool_timestamp_unix_secs(item.get("updated_at")))
+                .or(fallback_observed_at)?;
+            Some(base.saturating_add(reset_seconds.ceil() as u64))
+        })
+}
+
+pub(crate) fn provider_pool_reset_deadline_elapsed(
+    item: &Map<String, Value>,
+    fallback_observed_at: Option<u64>,
+    now_unix_secs: u64,
+) -> bool {
+    provider_pool_reset_deadline_unix_secs(item, fallback_observed_at)
+        .is_some_and(|reset_at| reset_at <= now_unix_secs)
+}
+
+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_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
+        })
+        .unwrap_or(false)
+}
+
 fn provider_pool_quota_snapshot_matches_provider(
     quota_snapshot: &Map<String, Value>,
     provider_type: &str,
@@ -203,19 +248,45 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
     }
     let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
     if exhausted {
-        let windows_max_ratio = quota_snapshot
+        let now_unix_secs = provider_pool_current_unix_secs();
+        let snapshot_observed_at =
+            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(|w| !w.is_empty())
-            .and_then(|windows| {
-                windows
-                    .iter()
-                    .filter_map(Value::as_object)
-                    .filter_map(|w| w.get("used_ratio"))
-                    .filter_map(Value::as_f64)
-                    .max_by(f64::total_cmp)
-            });
-        if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
+            .filter(|windows| !windows.is_empty())
+        {
+            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) {
+                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)));
+                }
+                if provider_pool_quota_window_is_exhausted(window) {
+                    saw_exhausted_window = true;
+                    let reset_elapsed = now_unix_secs.is_some_and(|now| {
+                        provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
+                    });
+                    if !reset_elapsed {
+                        saw_active_exhausted_window = true;
+                    }
+                }
+            }
+
+            if saw_exhausted_window {
+                return Some(saw_active_exhausted_window);
+            }
+            if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
+                return Some(false);
+            }
+        } else if now_unix_secs.is_some_and(|now| {
+            provider_pool_reset_deadline_elapsed(quota_snapshot, snapshot_observed_at, now)
+        }) {
             return Some(false);
         }
     }