Skip to main content

orch/engine/
cooldown.rs

1//! Per-agent and per-model cooldown tracking.
2//!
3//! Shared by `engine::runner::response` (recording failures) and
4//! `engine::router::config` (avoiding cooled models during pool selection).
5//!
6//! # Persistence
7//!
8//! All cooldowns are persisted to the KV store in SQLite so they survive
9//! service restarts. Call [`init_cooldown_store`] once at startup (after
10//! the `TaskStore` is open) to register the store and pre-load unexpired
11//! cooldowns.
12//!
13//! When a rate limit includes a "try again at {date}" message, the cooldown
14//! is set to that specific timestamp instead of the default duration.
15//! Relative usage windows (e.g. "5-hour usage limit") are also detected and
16//! used to set a model-specific cooldown matching the advertised window.
17
18use chrono::TimeZone;
19use once_cell::sync::Lazy;
20use regex::Regex;
21use std::collections::HashMap;
22use std::sync::{Arc, Mutex, OnceLock};
23
24/// Default fallback cooldown for agents when no store is available: 30 minutes.
25///
26/// Used in contexts that need a concrete duration and cannot compute exponential
27/// backoff (e.g. computing how long to wait before retrying when a cooldown
28/// expiry timestamp is unknown). New code should use [`BACKOFF_BASE_SECS`] instead.
29pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
30
31/// Base backoff for generic agent/model failures: 5 minutes.
32pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
33
34/// Maximum backoff for generic agent/model failures: 4 hours.
35pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
36
37/// Base backoff for persistent model failures (silent exit/provider errors): 4 hours.
38pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
39
40/// Maximum backoff for persistent model failures: 7 days.
41pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
42
43/// Base backoff for credit exhaustion (out_of_credits): 1 hour.
44pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
45
46/// Base backoff for org-level disabling: 2 hours.
47pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
48
49/// Standard cap for credit exhaustion and org-level disabling: 8 hours.
50/// Repeated failures at the cap move to the extended tiers (48h, then 96h)
51pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
52
53/// Base cooldown for billing cycle exhaustion: 24 hours.
54///
55/// Unlike per-request rate limits, billing cycle exhaustion is a calendar event
56/// that won't self-resolve for days or weeks. The first failure starts at 24h;
57/// subsequent failures escalate via `compute_backoff()` up to
58/// `BILLING_CYCLE_MAX_SECS` (7 days). This prevents daily retry-and-fail cycles
59/// when a monthly billing quota is exhausted.
60pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
61
62/// Maximum cooldown for billing cycle exhaustion: 7 days.
63///
64/// Monthly billing cycles typically reset on a fixed date. 7 days is a
65/// reasonable cap — the cooldown will expire before the next monthly reset
66/// in most cases, allowing a single probe to re-check availability.
67pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
68
69/// Credit exhaustion reason detected from error message.
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum CreditExhaustionReason {
72    /// Per-model credit exhaustion (out_of_credits).
73    OutOfCredits,
74    /// Organization-level disabling (org_level_disabled).
75    OrgLevelDisabled,
76    /// Monthly billing cycle exhaustion — escalating cooldown (24h base, 7d cap).
77    BillingCycleExhausted,
78}
79
80impl CreditExhaustionReason {
81    /// Stable string for `task_runs.outcome`, distinct from the cooldown reason
82    /// string used by `record_credit_exhaustion`. Lets callers persist the actual
83    /// failure mode instead of collapsing it into a generic "rate_limit"/"failed".
84    pub fn outcome_str(self) -> &'static str {
85        match self {
86            CreditExhaustionReason::OutOfCredits => "credit_exhausted",
87            CreditExhaustionReason::OrgLevelDisabled => "org_level_disabled",
88            CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
89        }
90    }
91}
92
93/// Short agent cooldown applied on silence detection (120 seconds).
94///
95/// Forces the router to pick a different agent immediately on re-route,
96/// without long-term blocking. The model cooldown (exponential, starting at
97/// 5 min) keeps the dead model out while this short cooldown just breaks
98/// the same-agent loop.
99pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
100
101/// Silence detections before applying an extended cooldown (rolling window).
102pub const SILENCE_COUNT_THRESHOLD: usize = 5;
103
104/// Rolling window for silence detections (24 hours).
105pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
106
107/// Extended cooldown applied after repeated silence detections (4 hours).
108pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
109
110/// KV key prefix for persisted cooldowns (both agent and model).
111const KV_PREFIX: &str = "cooldown:";
112
113/// KV key prefix for rolling silence counters.
114const SILENCE_COUNT_PREFIX: &str = "silence_count:";
115
116/// KV key prefix for per-agent and per-agent:model failure counts (drives exponential backoff).
117const FAILURE_COUNT_PREFIX: &str = "failure_count:";
118
119/// KV key prefix for credit-exhaustion failure counts (separate from generic failures to prevent
120/// cross-contamination of backoff escalation — credit exhaustion uses different base/max durations).
121const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
122
123/// KV key prefix for persistent model failure counts (separate from [`FAILURE_COUNT_PREFIX`] so
124/// `record_agent_success()` resetting the generic counter on an intervening success doesn't erase
125/// escalation progress for structurally unreliable models — see `record_persistent_model_failure`).
126const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
127
128struct CooldownEntry {
129    /// Unix timestamp when the cooldown expires.
130    cooldown_until: i64,
131    reason: String,
132}
133
134pub struct SilenceCountResult {
135    pub count: usize,
136    pub extended_cooldown_applied: bool,
137}
138
139/// Global in-memory cooldown map, protected by a Mutex.
140///
141/// Uses `std::sync::Mutex` because this map is also read from sync helper
142/// functions (`is_in_cooldown`, `is_model_in_cooldown`, etc.).  All lock
143/// holders perform only in-memory HashMap operations — **no `.await` is ever
144/// called while a guard is held**.  This invariant must be preserved: if a
145/// future caller needs to await while holding this lock, switch to
146/// `tokio::sync::Mutex` and make the helpers async.
147fn cooldowns() -> &'static Mutex<HashMap<String, CooldownEntry>> {
148    static COOLDOWNS: OnceLock<Mutex<HashMap<String, CooldownEntry>>> = OnceLock::new();
149    COOLDOWNS.get_or_init(|| Mutex::new(HashMap::new()))
150}
151
152/// Global store reference used for background KV writes.
153///
154/// Uses `tokio::sync::Mutex` because this value is only ever accessed from
155/// async functions. A `tokio::sync::Mutex` yields the Tokio task instead of
156/// blocking the worker thread if the lock is contended.
157fn cooldown_store() -> &'static tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>> {
158    static STORE: OnceLock<tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>>> =
159        OnceLock::new();
160    STORE.get_or_init(|| tokio::sync::Mutex::new(None))
161}
162
163/// Initialise persistent cooldowns.
164///
165/// Must be called once at engine startup, after the `TaskStore` is opened.
166/// Loads unexpired cooldowns from KV and registers the store for future writes.
167/* Parses the `{"until": N, "reason": "..."}` cooldown KV format, falling back
168to the legacy bare-integer format (reason unknown) so existing rows keep
169working across the upgrade without a migration. */
170fn parse_cooldown_value(value: &str) -> Option<(i64, String)> {
171    if let Ok(serde_json::Value::Object(map)) = serde_json::from_str(value) {
172        let until = map.get("until")?.as_i64()?;
173        let reason = map
174            .get("reason")
175            .and_then(|r| r.as_str())
176            .unwrap_or("persisted")
177            .to_string();
178        return Some((until, reason));
179    }
180    value
181        .parse::<i64>()
182        .ok()
183        .map(|until| (until, "persisted".to_string()))
184}
185
186pub async fn init_cooldown_store(store: Arc<crate::store::TaskStore>) {
187    match store.kv_list_prefix(KV_PREFIX).await {
188        Ok(rows) => {
189            let now = chrono::Utc::now().timestamp();
190            let mut map = cooldown_lock();
191            let mut loaded = 0;
192            for (key, value) in rows {
193                let cooldown_key = key.trim_start_matches(KV_PREFIX);
194                if let Some((cooldown_until, reason)) = parse_cooldown_value(&value) {
195                    if now < cooldown_until {
196                        map.insert(
197                            cooldown_key.to_string(),
198                            CooldownEntry {
199                                cooldown_until,
200                                reason,
201                            },
202                        );
203                        loaded += 1;
204                    }
205                    // Expired entries are simply not loaded — they'll be overwritten
206                    // on the next failure or cleaned up by future writes.
207                }
208            }
209            if loaded > 0 {
210                tracing::info!(loaded, "loaded persisted cooldowns from KV store");
211            }
212        }
213        Err(e) => {
214            tracing::warn!(err = %e, "failed to load cooldowns from KV store");
215        }
216    }
217
218    // Load GitHub 5xx circuit breaker state if it was persisted before restart.
219    if let Ok(Some(raw)) = store.kv_get("cooldown:github:5xx").await {
220        if let Ok(cooldown_until) = raw.parse::<i64>() {
221            let now = chrono::Utc::now().timestamp();
222            if now < cooldown_until {
223                let mut open = github_5xx_circuit_open_lock();
224                *open = Some(cooldown_until);
225                let remaining = cooldown_until - now;
226                tracing::warn!(
227                    remaining_secs = remaining,
228                    "GitHub 5xx circuit breaker restored from KV (still open)"
229                );
230            }
231        }
232    }
233
234    *cooldown_store().lock().await = Some(store);
235}
236
237/// Reconcile the in-memory cooldown map with the persisted KV state.
238///
239/// KV is the source of truth. This is called once per tick by
240/// `router.refresh_health` so the running service notices state mutated by a
241/// separate process (e.g. `orch cooldown clear`). Without this sync, the CLI
242/// clears KV but the engine's in-memory map keeps the stale entries until they
243/// expire on their own, producing the "all agents cooled" + "No active
244/// cooldowns" inconsistency.
245///
246/// Reconciliation rules (KV authoritative for every key it has):
247/// - KV row missing for an in-memory key: leave memory alone — likely a
248///   `set_cooldown_async` race where memory was written but the KV row isn't
249///   yet visible. The next tick will catch up.
250/// - KV row present with `cooldown_until <= now` (including the "0" tombstone
251///   written by `clear_cooldown`): drop from memory.
252/// - KV row present with `cooldown_until > now`: insert/update memory to match.
253pub async fn sync_from_kv(store: &Arc<crate::store::TaskStore>) {
254    let rows = match store.kv_list_prefix(KV_PREFIX).await {
255        Ok(rows) => rows,
256        Err(e) => {
257            tracing::warn!(err = %e, "cooldown sync_from_kv: KV list failed");
258            return;
259        }
260    };
261    let now = chrono::Utc::now().timestamp();
262
263    // Some((ts, reason)) for active cooldowns, None for cleared/expired rows.
264    // Missing key means no KV row at all, left alone below.
265    let mut explicit_kv: HashMap<String, Option<(i64, String)>> =
266        HashMap::with_capacity(rows.len());
267    for (kv_key, value) in &rows {
268        let cooldown_key = kv_key.trim_start_matches(KV_PREFIX);
269        if let Some((ts, reason)) = parse_cooldown_value(value) {
270            let normalized = if ts > now { Some((ts, reason)) } else { None };
271            explicit_kv.insert(cooldown_key.to_string(), normalized);
272        }
273    }
274
275    let mut map = cooldown_lock();
276    let mut dropped = 0usize;
277    let mut updated = 0usize;
278
279    // Drop in-memory entries that KV explicitly marks cleared/expired.
280    map.retain(|key, _| {
281        if matches!(explicit_kv.get(key), Some(None)) {
282            dropped += 1;
283            false
284        } else {
285            true
286        }
287    });
288
289    // Adopt KV's authoritative timestamp for active cooldowns.
290    for (key, normalized) in explicit_kv {
291        if let Some((ts, reason)) = normalized {
292            match map.get_mut(&key) {
293                Some(entry) => {
294                    if entry.cooldown_until != ts {
295                        entry.cooldown_until = ts;
296                        updated += 1;
297                    }
298                }
299                None => {
300                    map.insert(
301                        key,
302                        CooldownEntry {
303                            cooldown_until: ts,
304                            reason,
305                        },
306                    );
307                    updated += 1;
308                }
309            }
310        }
311    }
312
313    if dropped > 0 || updated > 0 {
314        tracing::info!(
315            dropped,
316            updated,
317            "cooldown sync_from_kv: reconciled in-memory map with KV"
318        );
319    }
320}
321
322/// Failure count threshold (times at max backoff) for first extended tier (24h).
323///
324/// After a model has been at the max cap this many times, apply a 6x multiplier.
325const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
326
327/// Multiplier for first extended tier (6x = 24h when max=4h).
328const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
329
330/// Failure count threshold (times at max backoff) for second extended tier (48h).
331///
332/// After a model has been at the max cap this many times, apply a 12x multiplier.
333const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
334
335/// Multiplier for second extended tier (12x = 48h when max=4h).
336const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
337
338/// Compute exponential backoff duration using base-3 growth.
339///
340/// Returns `min(base * 3^(count-1), max)` for count >= 1.
341/// For count == 0 or 1 (first failure), returns `base`.
342///
343/// | count | base=300s | result |
344/// |-------|-----------|--------|
345/// | 0     | 300       | 300    |
346/// | 1     | 300       | 300    |
347/// | 2     | 300       | 900    |
348/// | 3     | 300       | 2700   |
349/// | 4     | 300       | 8100   |
350/// | 5-6   | 14400     | 14400 (max, not yet extended) |
351/// | 7-9   | 86400     | 86400 (24h extended tier) |
352/// | 10+   | 172800    | 172800 (48h extended tier) |
353///
354/// The extended tier applies when the model has hit the max cap multiple times,
355/// indicating a persistently broken model that won't self-resolve quickly.
356pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
357    if count <= 1 {
358        return base;
359    }
360    let exp = count.saturating_sub(1).min(u32::MAX as u64) as u32;
361    let factor = 3_i64.saturating_pow(exp);
362    let standard = base.saturating_mul(factor).min(max);
363
364    // Extended tier: after N failures at max backoff, apply longer window
365    // to reduce noise from persistently broken models.
366    //
367    // First find the count at which exponential growth first hits max:
368    // solve: base * 3^(first_max-1) >= max
369    // => 3^(first_max-1) >= max / base
370    // => first_max-1 >= log_3(max / base)
371    // => first_max >= 1 + log_3(max / base)
372    let first_max = if base <= 0 {
373        0
374    } else {
375        // Use integer approximation: count where base * 3^(count-1) would first exceed max
376        let mut accumulated = base;
377        let mut first_max_count = 1u64;
378        while accumulated < max && first_max_count < count {
379            accumulated = accumulated.saturating_mul(3);
380            first_max_count += 1;
381        }
382        first_max_count
383    };
384
385    // Times we've been at the max cap (excluding the first time we hit it)
386    let times_at_max = count.saturating_sub(first_max);
387
388    if !allow_extended {
389        return standard;
390    }
391
392    if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
393        max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
394    } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
395        max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
396    } else {
397        standard
398    }
399}
400
401/// Read the failure count for a key from KV, increment it, write it back, and return the new count.
402///
403/// Always increments — the backoff formula's `max` parameter prevents the
404/// cooldown duration from growing beyond the cap, so unbounded counts are safe.
405///
406/// Returns 1 when no store is configured (unit-test contexts without a store).
407/// Returns `u64::MAX` on store errors (e.g. lock contention) so backoff applies the cap duration
408/// instead of resetting to the base, preventing rapid re-dispatch during store outages.
409async fn read_and_increment_failure_count(
410    store_opt: &Option<Arc<crate::store::TaskStore>>,
411    key: &str,
412) -> u64 {
413    read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
414}
415
416/// Like [`read_and_increment_failure_count`] but with an explicit KV prefix.
417///
418/// Credit exhaustion uses [`CREDIT_FAILURE_COUNT_PREFIX`] so its backoff escalation
419/// is independent of generic agent failures (which use [`FAILURE_COUNT_PREFIX`]).
420async fn read_and_increment_failure_count_with_prefix(
421    store_opt: &Option<Arc<crate::store::TaskStore>>,
422    prefix: &str,
423    key: &str,
424) -> u64 {
425    let kv_key = format!("{prefix}{key}");
426    if let Some(store) = store_opt {
427        match store.kv_increment(&kv_key).await {
428            Ok(n) => n,
429            Err(e) => {
430                tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
431                u64::MAX
432            }
433        }
434    } else {
435        1
436    }
437}
438
439/// Reset failure counts for an agent and agent:model combo after a successful run.
440///
441/// Call this from the runner's success path so that the next failure starts
442/// backoff from the base duration again, not from wherever it left off.
443/// Resets both generic and credit-exhaustion failure counters.
444pub async fn record_agent_success(agent_name: &str, model: &str) {
445    let store_opt = cooldown_store().lock().await.clone();
446    if let Some(store) = store_opt {
447        let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
448        let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
449        if let Err(e) = store.kv_set(&agent_key, "0").await {
450            tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
451        }
452        if !model.is_empty() {
453            let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
454            if let Err(e) = store.kv_set(&model_key, "0").await {
455                tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
456            }
457        }
458        if let Err(e) = store.kv_set(&credit_key, "0").await {
459            tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
460        }
461    }
462}
463
464/// Record that an agent has failed and should be temporarily avoided.
465///
466/// Applies exponential backoff based on the agent's failure count in KV.
467/// If `error_message` contains "try again at {date}", that vendor date is used
468/// instead of the computed backoff (vendor dates are always authoritative).
469pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
470    // Vendor-specified retry date takes priority over backoff — but only
471    // if it is in the present or future. If the parsed timestamp is in the
472    // past (stale vendor message), fall through to the exponential backoff
473    // path instead of returning early and bypassing backoff entirely.
474    if let Some(cooldown_until) = parse_retry_at(error_message) {
475        let now = chrono::Utc::now().timestamp();
476        if cooldown_until > now {
477            set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
478            return;
479        }
480        tracing::debug!(
481            agent = agent_name,
482            cooldown_until,
483            "vendor retry_at is in the past — using backoff instead"
484        );
485    }
486
487    let store_opt = cooldown_store().lock().await.clone();
488    let count = read_and_increment_failure_count(&store_opt, agent_name).await;
489    let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
490    let cooldown_until =
491        chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
492    set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
493}
494
495/// Detect credit exhaustion from an error message and return the stable
496/// `task_runs.outcome` string for it (see `CreditExhaustionReason::outcome_str`).
497///
498/// Callers that classify `AgentError::RateLimit`/`AgentError::Auth` into an
499/// outcome string should consult this first — generic rate-limit patterns like
500/// "usage limit" also match billing-cycle exhaustion messages (e.g. Kimi's
501/// "You've reached your usage limit for this billing cycle"), which would
502/// otherwise be stored as a plain "rate_limit" outcome.
503pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
504    detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
505}
506
507/// Detect credit exhaustion reasons from error message.
508///
509/// Returns `Some(CreditExhaustionReason)` if the error indicates
510/// out_of_credits or org_level_disabled, which require longer agent-wide cooldowns.
511pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
512    let lower = error_message.to_lowercase();
513
514    let overage_patterns = [
515        "overagedisabledreason",
516        "overage_disabled_reason",
517        "overage disabled reason",
518    ];
519
520    if overage_patterns.iter().any(|p| lower.contains(p)) {
521        if lower.contains("out_of_credits") || lower.contains("outofcredits") {
522            return Some(CreditExhaustionReason::OutOfCredits);
523        }
524        if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
525            return Some(CreditExhaustionReason::OrgLevelDisabled);
526        }
527    }
528
529    let billing_patterns = [
530        "credit balance too low",
531        "credit balance insufficient",
532        "insufficient credit balance",
533        "billing quota exceeded",
534        "billing limit exceeded",
535        "organization has been disabled",
536        "org disabled",
537    ];
538
539    if billing_patterns.iter().any(|p| lower.contains(p)) {
540        return Some(CreditExhaustionReason::OrgLevelDisabled);
541    }
542
543    // Billing cycle exhaustion: monthly quota reset — check before generic quota
544    // patterns because "monthly quota exceeded" would otherwise match "quota exceeded".
545    let billing_cycle_patterns = [
546        "billing cycle",
547        "quota refreshed next cycle",
548        "monthly quota",
549        "quota will be refreshed",
550        "next billing cycle",
551        "refreshed in the next cycle",
552        "weekly/monthly limit exhausted",
553        "limit will reset at",
554        "upgrade to pro", // codex/chatgpt plan usage cap: requires billing upgrade
555        "upgrade your token plan", // minimax: Token Plan usage limit reached
556        "token plan usage limit", // minimax: Token Plan usage limit reached
557        "purchase credits for more", // minimax: ...purchase Credits for more usage
558    ];
559
560    if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
561        return Some(CreditExhaustionReason::BillingCycleExhausted);
562    }
563
564    if lower.contains("out of credits")
565        || lower.contains("outofcredits")
566        || lower.contains("insufficient funds")
567        || lower.contains("no credits remaining")
568        || lower.contains("insufficient_quota")
569        || lower.contains("quota exceeded")
570        || lower.contains("insufficient balance")
571        || lower.contains("please recharge")
572        || lower.contains("no resource package")
573        || lower.contains("balance exhausted")
574        || lower.contains("402 payment required")
575    {
576        return Some(CreditExhaustionReason::OutOfCredits);
577    }
578
579    None
580}
581
582/// Record an agent-level cooldown for credit exhaustion.
583///
584/// Applies exponential backoff based on the agent's failure count.
585/// All variants use escalating backoff:
586/// - `OutOfCredits`: 1h, 3h, then 8h (credits can be refilled any time), 48h from the 5th failure, 96h from the 8th
587/// - `OrgLevelDisabled`: 2h, 6h, then 8h, 48h from the 5th failure, 96h from the 8th
588/// - `BillingCycleExhausted`: starts at 24h, caps at 7 days (monthly event)
589///
590/// Billing cycle exhaustion was previously a flat 24h, but this caused daily
591/// retry-and-fail cycles when a monthly quota was exhausted (the 24h cooldown
592/// would expire, kimi would be retried, fail immediately, and get another 24h).
593/// Now each recurrence escalates: 24h → 72h → 7d (capped).
594pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
595    let reason_str = match reason {
596        CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
597        CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
598        CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
599    };
600
601    let store_opt = cooldown_store().lock().await.clone();
602    // Use a distinct failure counter (CREDIT_FAILURE_COUNT_PREFIX) so credit-
603    // exhaustion backoff escalation is independent of generic agent failures.
604    // Previously both paths shared FAILURE_COUNT_PREFIX, causing generic failures
605    // to inflate credit-exhaustion backoff (and vice versa).
606    let count = read_and_increment_failure_count_with_prefix(
607        &store_opt,
608        CREDIT_FAILURE_COUNT_PREFIX,
609        agent_name,
610    )
611    .await;
612
613    let (base, max) = match reason {
614        CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
615        CreditExhaustionReason::OrgLevelDisabled => {
616            (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
617        }
618        CreditExhaustionReason::BillingCycleExhausted => {
619            (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
620        }
621    };
622    let cooldown_secs = compute_backoff(
623        count,
624        base,
625        max,
626        !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
627    );
628    let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
629    set_cooldown_async(agent_name, cooldown_until, reason_str).await;
630    tracing::warn!(
631        agent = agent_name,
632        reason = reason_str,
633        attempt = count,
634        cooldown_secs,
635        "credit exhaustion detected: applying exponential agent-wide cooldown"
636    );
637}
638
639/// Record that a specific agent+model combo has failed.
640///
641/// Applies exponential backoff based on the model's failure count in KV.
642pub async fn record_model_failure(agent_name: &str, model: &str) {
643    let key = format!("{agent_name}:{model}");
644    let store_opt = cooldown_store().lock().await.clone();
645    let count = read_and_increment_failure_count(&store_opt, &key).await;
646    let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
647    let cooldown_until =
648        chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
649    set_cooldown_async(&key, cooldown_until, "model_error").await;
650}
651
652/// Record a rate-limit event for an agent+model.
653///
654/// If the error message advertises a relative usage window (e.g.
655/// "5-hour usage limit"), the cooldown is applied to the specific model for
656/// that window length. The quota is model-scoped, so the whole agent is not
657/// cooled. If no window is detected, fall back to the generic agent+model
658/// exponential backoff path.
659pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
660    if let Some(window_secs) = parse_relative_usage_window(error_message) {
661        if let Some(m) = model {
662            let cooldown_until = chrono::Utc::now().timestamp() + window_secs as i64;
663            set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
664            tracing::info!(
665                agent = agent_name,
666                model = m,
667                window_secs,
668                cooldown_until,
669                "relative usage window detected: applying model-specific cooldown"
670            );
671            return;
672        }
673    }
674
675    record_agent_failure_with_message(agent_name, error_message).await;
676    if let Some(m) = model {
677        record_model_failure(agent_name, m).await;
678    }
679}
680
681/// Record a persistent model-specific failure and apply longer exponential backoff.
682///
683/// Used for failure modes that are commonly broken for long periods (for example
684/// silent exit-0 or opaque provider errors). Starts at 4h and escalates with
685/// the generic backoff formula up to 7 days.
686///
687/// Uses [`PERSISTENT_FAILURE_COUNT_PREFIX`], a counter independent of the generic
688/// `failure_count:` key that `record_agent_success()` resets on every successful
689/// run. Otherwise a single intervening success between persistent failures would
690/// reset escalation back to the 4h base, letting structurally unreliable models
691/// (intermittent parse errors, truncated runs) stay in rotation indefinitely.
692pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
693    let key = format!("{agent_name}:{model}");
694    let store_opt = cooldown_store().lock().await.clone();
695    let count = read_and_increment_failure_count_with_prefix(
696        &store_opt,
697        PERSISTENT_FAILURE_COUNT_PREFIX,
698        &key,
699    )
700    .await;
701    let cooldown_secs = compute_backoff(
702        count,
703        PERSISTENT_MODEL_BACKOFF_BASE_SECS,
704        PERSISTENT_MODEL_BACKOFF_MAX_SECS,
705        false,
706    );
707    let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
708    set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
709}
710
711/// Set a model cooldown with a custom duration (in seconds) and reason.
712///
713/// Generic model-cooldown setter shared by silence detection and the
714/// rate-limit relative-usage-window path. The reason is persisted with
715/// the cooldown (KV + in-memory entry) so operators can tell the two
716/// apart in `orch cooldown list`.
717pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
718    let key = format!("{agent_name}:{model}");
719    let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
720    set_cooldown_async(&key, cooldown_until, reason).await;
721}
722
723/// Set a model cooldown that ends at a vendor-provided Unix timestamp.
724pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
725    let key = format!("{agent_name}:{model}");
726    set_cooldown_async(&key, until, "vendor_reset_at").await;
727}
728
729/// Set a short agent-level cooldown (in seconds).
730///
731/// Used by silence detection to temporarily block the whole agent so the
732/// router picks a different one on re-route. Unlike `record_agent_failure`
733/// (exponential backoff), this uses a short duration (typically 120s) —
734/// just enough to force one re-route cycle to a different agent.
735pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
736    let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
737    set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
738}
739
740/// Record a silence detection for an agent+model and apply extended cooldowns
741/// when repeated silences exceed the threshold within the rolling window.
742pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
743    let store_opt = cooldown_store().lock().await.clone();
744    let store = match store_opt {
745        Some(store) => store,
746        None => {
747            tracing::debug!(
748                agent = agent_name,
749                model,
750                "skipping silence count record (no KV store)"
751            );
752            return None;
753        }
754    };
755
756    let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
757    let now = chrono::Utc::now().timestamp();
758    let window_start = now - SILENCE_COUNT_WINDOW_SECS;
759
760    let mut timestamps = match store.kv_get(&key).await {
761        Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
762            Ok(ts) => ts,
763            Err(e) => {
764                tracing::warn!(
765                    kv_key = %key,
766                    err = %e,
767                    "failed to parse silence timestamps from KV — resetting window"
768                );
769                Vec::new()
770            }
771        },
772        Ok(None) => Vec::new(),
773        Err(err) => {
774            tracing::warn!(
775                kv_key = %key,
776                err = %err,
777                "failed to load silence count from KV"
778            );
779            Vec::new()
780        }
781    };
782
783    timestamps.retain(|ts| *ts >= window_start);
784    timestamps.push(now);
785
786    let count = timestamps.len();
787    let mut extended_cooldown_applied = false;
788    if count >= SILENCE_COUNT_THRESHOLD {
789        set_model_cooldown(
790            agent_name,
791            model,
792            SILENCE_EXTENDED_COOLDOWN_SECS,
793            "silence_detected",
794        )
795        .await;
796        extended_cooldown_applied = true;
797        timestamps.clear();
798    }
799
800    let value = serde_json::to_string(&timestamps).unwrap_or_else(|_| "[]".to_string());
801    if let Err(err) = store.kv_set(&key, &value).await {
802        tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
803    }
804
805    Some(SilenceCountResult {
806        count,
807        extended_cooldown_applied,
808    })
809}
810
811/// Check if a specific agent+model combo is in cooldown.
812pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
813    let key = format!("{agent_name}:{model}");
814    is_in_cooldown(&key)
815}
816
817/// Check if an agent is currently in cooldown period.
818pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
819    is_in_cooldown(agent_name)
820}
821
822/// Return the cooldown expiry timestamp for a key (agent or agent:model), if active.
823pub fn cooldown_until(key: &str) -> Option<i64> {
824    let map = cooldown_lock();
825    map.get(key).and_then(|entry| {
826        let now = chrono::Utc::now().timestamp();
827        if now < entry.cooldown_until {
828            Some(entry.cooldown_until)
829        } else {
830            None
831        }
832    })
833}
834
835/// Return the reason string recorded when a cooldown was set, if the cooldown is still active.
836pub fn cooldown_reason(key: &str) -> Option<String> {
837    let map = cooldown_lock();
838    map.get(key).and_then(|entry| {
839        let now = chrono::Utc::now().timestamp();
840        if now < entry.cooldown_until {
841            Some(entry.reason.clone())
842        } else {
843            None
844        }
845    })
846}
847
848/// Return all currently active cooldowns as `(key, cooldown_until_unix, reason)`.
849///
850/// Reads from the in-memory map only — always fast, no async needed.
851/// Expired entries are filtered out. Used by `orch cooldown list`.
852pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
853    let now = chrono::Utc::now().timestamp();
854    let map = cooldown_lock();
855    let mut result: Vec<(String, i64, String)> = map
856        .iter()
857        .filter(|(_, entry)| now < entry.cooldown_until)
858        .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
859        .collect();
860    result.sort_by(|a, b| a.0.cmp(&b.0));
861    result
862}
863
864/// Clear a specific cooldown by key, or all cooldowns when `key == "*"`.
865///
866/// Removes from the in-memory map, writes a past timestamp to KV, and resets
867/// the corresponding `failure_count:` entries so the next failure starts
868/// backoff from the base duration.
869pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
870    if key == "*" {
871        let keys: Vec<String> = {
872            let mut map = cooldown_lock();
873            let keys: Vec<String> = map.keys().cloned().collect();
874            map.clear();
875            keys
876        };
877        for k in &keys {
878            let kv_key = format!("{KV_PREFIX}{k}");
879            if let Err(e) = store.kv_set(&kv_key, "0").await {
880                tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
881            }
882            // Reset failure count so backoff restarts from base.
883            let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
884            if let Err(e) = store.kv_set(&fc_key, "0").await {
885                tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
886            }
887            // Also reset credit-specific failure count.
888            let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
889            if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
890                tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
891            }
892            // Also reset the persistent-model-failure counter.
893            let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
894            if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
895                tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
896            }
897        }
898        // Also reset any persisted failure_count keys that are not in the
899        // in-memory map (e.g. survived a restart or were set without a
900        // corresponding in-memory cooldown entry).
901        if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
902            for (fc_key, _) in fc_entries {
903                if let Err(e) = store.kv_set(&fc_key, "0").await {
904                    tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
905                }
906            }
907        }
908        if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
909            for (cfc_key, _) in cfc_entries {
910                if let Err(e) = store.kv_set(&cfc_key, "0").await {
911                    tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
912                }
913            }
914        }
915        if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
916            for (pfc_key, _) in pfc_entries {
917                if let Err(e) = store.kv_set(&pfc_key, "0").await {
918                    tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
919                }
920            }
921        }
922        tracing::info!(
923            count = keys.len(),
924            "cleared all cooldowns and failure counts"
925        );
926    } else {
927        {
928            let mut map = cooldown_lock();
929            map.remove(key);
930        }
931        let kv_key = format!("{KV_PREFIX}{key}");
932        if let Err(e) = store.kv_set(&kv_key, "0").await {
933            tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
934        }
935        // Reset failure count so backoff restarts from base.
936        let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
937        if let Err(e) = store.kv_set(&fc_key, "0").await {
938            tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
939        }
940        // Also reset credit-specific failure count.
941        let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
942        if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
943            tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
944        }
945        // Also reset the persistent-model-failure counter.
946        let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
947        if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
948            tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
949        }
950        tracing::info!(key, "cleared cooldown and failure count");
951    }
952}
953
954/// Set a cooldown with a specific expiry timestamp. Updates in-memory state only.
955///
956/// Returns `true` if the cooldown was applied, `false` if an existing longer
957/// cooldown prevented the update (never-shorten rule).
958///
959/// **Callers are responsible for persisting to KV.** Use [`set_cooldown_async`]
960/// for critical cooldowns that must survive a crash, or this function combined
961/// with a fire-and-forget `tokio::spawn` for non-critical short cooldowns.
962fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
963    let mut map = cooldown_lock();
964    if let Some(existing) = map.get(key) {
965        // Never shorten an existing cooldown. This prevents a short
966        // retry window (e.g., generic rate limit) from overriding a longer
967        // billing-cycle cooldown.
968        if existing.cooldown_until >= cooldown_until {
969            return false;
970        }
971    }
972    map.insert(
973        key.to_string(),
974        CooldownEntry {
975            cooldown_until,
976            reason: reason.to_string(),
977        },
978    );
979    true
980}
981
982/// Set a cooldown with a specific expiry timestamp and persist to KV inline.
983///
984/// This is the preferred function for critical cooldowns (credit exhaustion,
985/// billing cycle, agent/model failures) that must survive a service crash.
986/// The KV write is awaited before returning, so the caller can be confident
987/// the cooldown is durable once this function returns `true`.
988///
989/// Returns `true` if the cooldown was applied, `false` if an existing longer
990/// cooldown prevented the update (never-shorten rule).
991async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
992    if !set_cooldown_in_memory(key, cooldown_until, reason) {
993        return false;
994    }
995    let store_opt = cooldown_store().lock().await.clone();
996    if let Some(store) = store_opt {
997        let kv_key = format!("{KV_PREFIX}{key}");
998        let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
999        if let Err(e) = store.kv_set(&kv_key, &value).await {
1000            tracing::warn!(
1001                kv_key,
1002                err = %e,
1003                "failed to persist cooldown to KV store — rolling back in-memory entry"
1004            );
1005            let mut map = cooldown_lock();
1006            // Only roll back if nothing has superseded our write.
1007            // A racing thread may have set a longer cooldown in the meantime;
1008            // removing unconditionally would erase that valid entry.
1009            if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1010                map.remove(key);
1011            }
1012            return false;
1013        }
1014    }
1015    true
1016}
1017
1018/// Helper to acquire the cooldown map lock, logging and recovering on poison.
1019fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
1020    cooldowns().lock().unwrap_or_else(|e| {
1021        tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1022        e.into_inner()
1023    })
1024}
1025
1026/// Helper to acquire the degraded agents lock, logging and recovering on poison.
1027fn degraded_agents_lock() -> std::sync::MutexGuard<'static, std::collections::HashSet<String>> {
1028    degraded_agents().lock().unwrap_or_else(|e| {
1029        tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1030        e.into_inner()
1031    })
1032}
1033
1034/// Helper to acquire the GitHub 5xx timestamps lock, logging and recovering on poison.
1035fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
1036    github_5xx_timestamps().lock().unwrap_or_else(|e| {
1037        tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1038        e.into_inner()
1039    })
1040}
1041
1042/// Helper to acquire the GitHub 5xx circuit open lock, logging and recovering on poison.
1043fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
1044    github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1045        tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1046        e.into_inner()
1047    })
1048}
1049
1050/// Check if a key is currently in cooldown.
1051fn is_in_cooldown(key: &str) -> bool {
1052    let map = cooldown_lock();
1053    if let Some(entry) = map.get(key) {
1054        let now = chrono::Utc::now().timestamp();
1055        return now < entry.cooldown_until;
1056    }
1057    false
1058}
1059
1060// ---------------------------------------------------------------------------
1061// Degraded-agent tracking (pre-emptive health check)
1062// ---------------------------------------------------------------------------
1063
1064/// Default health check window for recent rate-limit events (hours).
1065pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1066
1067/// Default threshold: minimum rate-limit events within the window to consider
1068/// an agent degraded.  A single transient 429 shouldn't mark an agent; we need
1069/// a pattern of repeated failures.
1070pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1071
1072/// In-memory set of agents currently flagged as degraded.
1073fn degraded_agents() -> &'static Mutex<std::collections::HashSet<String>> {
1074    static DEGRADED: OnceLock<Mutex<std::collections::HashSet<String>>> = OnceLock::new();
1075    DEGRADED.get_or_init(|| Mutex::new(std::collections::HashSet::new()))
1076}
1077
1078/// Check if an agent is currently flagged as degraded by the health check.
1079///
1080/// Degraded means: all configured models are cooled **or** the agent has
1081/// exceeded the rate-limit/out_of_credits threshold within the configured
1082/// lookback window.
1083pub fn is_agent_degraded(agent: &str) -> bool {
1084    let set = degraded_agents_lock();
1085    set.contains(agent)
1086}
1087
1088/// Mark an agent as degraded (called by [`refresh_degraded_agents`]).
1089pub fn mark_agent_degraded(agent: &str) {
1090    let mut set = degraded_agents_lock();
1091    set.insert(agent.to_string());
1092}
1093
1094/// Clear the degraded flag for an agent.
1095pub fn clear_agent_degraded(agent: &str) {
1096    let mut set = degraded_agents_lock();
1097    set.remove(agent);
1098}
1099
1100/// Refresh the degraded-agent set from the rate_limits table and cooldown state.
1101///
1102/// An agent is marked degraded when:
1103/// 1. It is in agent-level cooldown (explicit cooldown KV key exists and hasn't expired), **or**
1104/// 2. All its configured models are in cooldown, **or**
1105/// 3. It has >= `threshold` rate_limit/out_of_credits events within `window_hours`.
1106///
1107/// Agents that no longer meet these criteria are cleared.
1108///
1109/// Note: This function calls `clear_expired_cooldowns()` first to ensure the
1110/// in-memory cooldown map reflects current state. An agent should NOT be marked
1111/// degraded simply because an expired cooldown entry hasn't been cleaned up yet.
1112pub async fn refresh_degraded_agents(
1113    store: &Arc<crate::store::TaskStore>,
1114    available_agents: &[String],
1115    model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1116    window_hours: u32,
1117    threshold: i64,
1118) {
1119    // Clear expired cooldowns from in-memory map before evaluating agent state.
1120    // This ensures we don't incorrectly mark agents as degraded due to stale
1121    // in-memory entries that should have expired but weren't cleaned up.
1122    clear_expired_cooldowns();
1123
1124    // Time the DB query to surface slow health-checks that can dominate the
1125    // engine tick latency. If the query fails, log and return early.
1126    let start = chrono::Utc::now();
1127    let counts = match store.recent_rate_limit_counts(window_hours).await {
1128        Ok(c) => {
1129            let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1130            tracing::info!(
1131                duration_ms = dur_ms,
1132                window_hours,
1133                "recent_rate_limit_counts query completed"
1134            );
1135            c
1136        }
1137        Err(e) => {
1138            let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1139            tracing::warn!(
1140                err = %e,
1141                duration_ms = dur_ms,
1142                "failed to query recent rate limit counts for health check"
1143            );
1144            return;
1145        }
1146    };
1147
1148    for agent in available_agents {
1149        let in_cooldown = is_agent_in_cooldown(agent);
1150        let no_models = !model_checker(agent);
1151        let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1152        let over_threshold = rate_limit_count >= threshold;
1153
1154        let degraded = in_cooldown || no_models || over_threshold;
1155
1156        if degraded {
1157            if !is_agent_degraded(agent) {
1158                let reason = if in_cooldown {
1159                    "agent in cooldown"
1160                } else if no_models {
1161                    "all models cooled"
1162                } else {
1163                    "rate limit threshold exceeded"
1164                };
1165                tracing::warn!(
1166                    agent,
1167                    reason,
1168                    rate_limit_count,
1169                    window_hours,
1170                    threshold,
1171                    "pre-emptive health check: marking agent as degraded"
1172                );
1173            }
1174            mark_agent_degraded(agent);
1175        } else {
1176            if is_agent_degraded(agent) {
1177                tracing::info!(
1178                    agent,
1179                    rate_limit_count,
1180                    "pre-emptive health check: agent recovered, clearing degraded flag"
1181                );
1182            }
1183            clear_agent_degraded(agent);
1184        }
1185    }
1186}
1187
1188/// Clear expired cooldowns from the in-memory map.
1189pub fn clear_expired_cooldowns() {
1190    let mut map = cooldown_lock();
1191    let now = chrono::Utc::now().timestamp();
1192    map.retain(|_key, entry| now < entry.cooldown_until);
1193}
1194
1195// ---------------------------------------------------------------------------
1196// GitHub 5xx circuit breaker
1197// ---------------------------------------------------------------------------
1198
1199/// Sliding window for counting GitHub 5xx errors (seconds).
1200const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1201
1202/// Number of 5xx errors within the window to trip the circuit breaker.
1203const GITHUB_5XX_THRESHOLD: usize = 5;
1204
1205/// Duration the circuit breaker stays open after tripping (seconds).
1206const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1207
1208/// Global sliding window of 5xx error timestamps.
1209fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1210    static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1211    TS.get_or_init(|| Mutex::new(Vec::new()))
1212}
1213
1214/// Whether the GitHub 5xx circuit breaker is currently open (tripped).
1215fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1216    static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1217    OPEN.get_or_init(|| Mutex::new(None))
1218}
1219
1220/// Record a GitHub 5xx server error and trip the circuit breaker if the
1221/// sliding window threshold is exceeded.
1222///
1223/// Emits a single high-level log entry when the circuit transitions to open.
1224pub async fn record_github_5xx() {
1225    let now = chrono::Utc::now().timestamp();
1226    let window_start = now - GITHUB_5XX_WINDOW_SECS;
1227
1228    let should_trip = {
1229        let mut ts = github_5xx_timestamps_lock();
1230        ts.retain(|t| *t >= window_start);
1231        ts.push(now);
1232        ts.len() >= GITHUB_5XX_THRESHOLD
1233    };
1234
1235    if should_trip {
1236        let cooldown_until = {
1237            let mut open = github_5xx_circuit_open_lock();
1238            if open.is_none() {
1239                let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1240                *open = Some(cd);
1241                tracing::warn!(
1242                    errors_in_window = GITHUB_5XX_THRESHOLD,
1243                    window_secs = GITHUB_5XX_WINDOW_SECS,
1244                    cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1245                    "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1246                );
1247                Some(cd)
1248            } else {
1249                None
1250            }
1251        };
1252        // Persist to KV (outside the mutex lock so the future is Send).
1253        if let Some(cd) = cooldown_until {
1254            let store_opt = cooldown_store().lock().await.clone();
1255            if let Some(store) = store_opt {
1256                let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1257            }
1258            // Also set the generic cooldown so is_agent_in_cooldown("github:5xx")
1259            // returns true immediately, avoiding a race window where concurrent
1260            // requests bypass the circuit breaker.
1261            set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1262        }
1263        // Clear the sliding window after tripping so we don't re-trip immediately
1264        let mut ts = github_5xx_timestamps_lock();
1265        ts.clear();
1266    }
1267}
1268
1269/// Check if the GitHub 5xx circuit breaker is currently open.
1270///
1271/// When open, non-critical background work (cleanup, polling, mention scans)
1272/// should be skipped. Critical operations (PR creation) should still attempt
1273/// with their existing exponential backoff.
1274///
1275/// Automatically recovers (logs a recovery message) when the cooldown expires.
1276pub fn is_github_circuit_open() -> bool {
1277    // Check the dedicated in-memory circuit flag first. If present, honour it
1278    // and perform auto-recovery when it expires (logging once on close).
1279    {
1280        let mut open = github_5xx_circuit_open_lock();
1281        if let Some(until) = *open {
1282            let now = chrono::Utc::now().timestamp();
1283            if now >= until {
1284                *open = None;
1285                tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1286                return false;
1287            }
1288            return true;
1289        }
1290    }
1291
1292    // If the dedicated flag is not set, consult the generic cooldown map so
1293    // callers see a consistent view regardless of which code path set the
1294    // persisted cooldown (e.g. send_with_retries -> set_agent_cooldown).
1295    if let Some(until) = cooldown_until("github:5xx") {
1296        // Populate the dedicated in-memory flag so future checks and the
1297        // auto-recovery path go through the same code and logging.
1298        let mut open = github_5xx_circuit_open_lock();
1299        if open.is_none() {
1300            *open = Some(until);
1301            let now = chrono::Utc::now().timestamp();
1302            if now < until {
1303                let remaining = until - now;
1304                tracing::warn!(
1305                    remaining_secs = remaining,
1306                    "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1307                );
1308            }
1309        }
1310        return true;
1311    }
1312
1313    false
1314}
1315
1316/// Returns the remaining seconds of the GitHub 5xx circuit breaker cooldown,
1317/// or 0 if the circuit is closed.
1318pub fn github_circuit_remaining_secs() -> u64 {
1319    // Prefer the dedicated in-memory value when present.
1320    {
1321        let open = github_5xx_circuit_open_lock();
1322        if let Some(until) = *open {
1323            let now = chrono::Utc::now().timestamp();
1324            if now < until {
1325                return (until - now) as u64;
1326            }
1327            return 0;
1328        }
1329    }
1330
1331    // Fall back to the generic cooldown map if the dedicated flag isn't set.
1332    if let Some(until) = cooldown_until("github:5xx") {
1333        let now = chrono::Utc::now().timestamp();
1334        if now < until {
1335            return (until - now) as u64;
1336        }
1337    }
1338    0
1339}
1340
1341/// Parse a "try again at {date}" or "reset at {date}" from an error message.
1342///
1343/// Returns a Unix timestamp if a retry-at date is found.
1344pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1345    if error_message.is_empty() {
1346        return None;
1347    }
1348
1349    let lower = error_message.to_lowercase();
1350    // Accept a wider set of vendor-provided markers. Claude commonly uses
1351    // variations like "resets 10:30am" (no explicit "at"), so include
1352    // both "resets" and the older "reset at"/"try again at" markers.
1353    for retry_marker in [
1354        "try again at ",
1355        "reset at ",
1356        "resets at ",
1357        "resets ",
1358        "reset ",
1359    ] {
1360        if let Some(idx) = lower.find(retry_marker) {
1361            let date_str = &lower[idx + retry_marker.len()..];
1362            // Take until period, newline, or end
1363            let date_str = date_str
1364                .split(['.', '\n'])
1365                .next()
1366                .unwrap_or(date_str)
1367                .trim();
1368
1369            // Try common date formats
1370            // "Mar 26th, 2026 5:55 AM" → strip ordinal suffixes
1371            let cleaned = date_str
1372                .replace("st,", ",")
1373                .replace("nd,", ",")
1374                .replace("rd,", ",")
1375                .replace("th,", ",");
1376
1377            for fmt in &[
1378                "%b %d, %Y %I:%M %p", // Mar 26, 2026 5:55 AM
1379                "%b %d, %Y %I:%M%p",  // Mar 26, 2026 5:55AM
1380                "%B %d, %Y %I:%M %p", // March 26, 2026 5:55 AM
1381                "%Y-%m-%dT%H:%M:%S",  // 2026-03-26T05:55:00
1382                "%Y-%m-%d %H:%M:%S",  // 2026-03-26 05:55:00
1383                "%Y-%m-%d %H:%M",     // 2026-03-26 05:55
1384            ] {
1385                if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1386                    // Vendor messages (e.g. Codex/GLM) emit times in the machine's
1387                    // local timezone, not UTC. Interpret naive datetime as local time.
1388                    use chrono::TimeZone;
1389                    let utc = match chrono::Local.from_local_datetime(&dt) {
1390                        chrono::LocalResult::Single(local_dt) => {
1391                            local_dt.with_timezone(&chrono::Utc)
1392                        }
1393                        chrono::LocalResult::Ambiguous(earliest, _) => {
1394                            // DST fold: pick the earlier (conservative) interpretation
1395                            earliest.with_timezone(&chrono::Utc)
1396                        }
1397                        chrono::LocalResult::None => {
1398                            // DST gap: fall back to treating as UTC to avoid panic
1399                            dt.and_utc()
1400                        }
1401                    };
1402                    tracing::info!(
1403                        retry_at = %utc,
1404                        raw = date_str,
1405                        marker = retry_marker,
1406                        "parsed rate limit retry-at date"
1407                    );
1408                    return Some(utc.timestamp());
1409                }
1410            }
1411            tracing::debug!(
1412                raw = date_str,
1413                marker = retry_marker,
1414                "could not parse retry-at date"
1415            );
1416
1417            // If the vendor provided only a time (e.g. "resets 10:30am"),
1418            // attempt to parse a time-only and interpret it as the next
1419            // occurrence of that time in the local timezone.
1420            // Strip parenthesised timezone hints like "(America/Sao_Paulo)".
1421            let time_only = date_str
1422                .split_once('(')
1423                .map(|(t, _)| t)
1424                .unwrap_or(date_str)
1425                .trim();
1426            // Hour-only values like "8am" have no minute field, which chrono needs, so insert ":00"
1427            let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1428            let split = compact.rfind(' ').map_or(0, |i| i + 1);
1429            let last = &compact[split..];
1430            if !last.contains(':')
1431                && (last.ends_with("am") || last.ends_with("pm"))
1432                && last.len() > 2
1433            {
1434                let at = split + last.len() - 2;
1435                compact.insert_str(at, ":00");
1436            }
1437            // "Jun 9 at 1am" has no year, so take the next future occurrence
1438            let md = compact.replace(" at ", " ");
1439            let now_local = chrono::Local::now();
1440            let this_year = chrono::Datelike::year(&now_local);
1441            for year in [this_year, this_year + 1] {
1442                let with_year = format!("{year} {md}");
1443                if let Ok(dt) =
1444                    chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1445                {
1446                    use chrono::TimeZone;
1447                    if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
1448                        if local_dt > now_local {
1449                            return Some(local_dt.timestamp());
1450                        }
1451                    }
1452                }
1453            }
1454            let time_only = compact.as_str();
1455            for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1456                if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1457                    // Build a NaiveDateTime for today at that time, then choose
1458                    // today or tomorrow depending on whether the time has passed.
1459                    let local_now = chrono::Local::now();
1460                    let today = local_now.date_naive();
1461                    let candidate = chrono::NaiveDateTime::new(today, t);
1462                    // Use the TimeZone trait methods; annotate types so the
1463                    // compiler can infer the concrete DateTime types.
1464                    let candidate_local: chrono::DateTime<chrono::Local> =
1465                        match chrono::Local.from_local_datetime(&candidate) {
1466                            chrono::LocalResult::Single(dt) => dt,
1467                            chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1468                            chrono::LocalResult::None => {
1469                                // DST gap — fall back to UTC interpretation
1470                                let utc_dt = candidate.and_utc();
1471                                return Some(utc_dt.timestamp());
1472                            }
1473                        };
1474                    let chosen = if candidate_local <= local_now {
1475                        // Use tomorrow
1476                        let tomorrow = today.succ_opt().unwrap_or(today);
1477                        let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1478                        match chrono::Local.from_local_datetime(&ndt) {
1479                            chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1480                            chrono::LocalResult::Ambiguous(earliest, _) => {
1481                                earliest.with_timezone(&chrono::Utc)
1482                            }
1483                            chrono::LocalResult::None => ndt.and_utc(),
1484                        }
1485                    } else {
1486                        candidate_local.with_timezone(&chrono::Utc)
1487                    };
1488                    tracing::info!(
1489                        retry_at = %chosen,
1490                        raw = date_str,
1491                        marker = retry_marker,
1492                        "parsed rate limit retry-at time-only; interpreted as next occurrence"
1493                    );
1494                    return Some(chosen.timestamp());
1495                }
1496            }
1497        }
1498    }
1499
1500    // Without a specific retry/reset date, return None and let the caller use
1501    // its default cooldown (exponential backoff). A blanket long fallback here
1502    // would cause false billing-cycle cooldowns on agents without those limits.
1503    None
1504}
1505
1506/// Parse a relative usage window (e.g. "5-hour usage limit") from an error
1507/// message.
1508///
1509/// Returns the window length in seconds. Only matches in usage/quota/limit
1510/// contexts ("usage limit", "quota", "window", "limit") to avoid interpreting
1511/// unrelated numbers as cooldown durations.
1512///
1513/// `pub(crate)` so `engine::runner::fallback::handle_error()` can apply the
1514/// same window-detection the `needs_review` call site
1515/// (`record_rate_limit()`, below) already uses — see issue #3580.
1516pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1517    if error_message.is_empty() {
1518        return None;
1519    }
1520
1521    let lower = error_message.to_lowercase();
1522    let has_context = lower.contains("usage limit")
1523        || lower.contains("quota")
1524        || lower.contains("window")
1525        || lower.contains("limit");
1526    if !has_context {
1527        return None;
1528    }
1529
1530    static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
1531        Regex::new(
1532            r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1533        )
1534        .expect("usage window regex is valid")
1535    });
1536
1537    let mut max_secs: Option<u64> = None;
1538    for caps in WINDOW_RE.captures_iter(&lower) {
1539        let secs = if let Some(num) = caps.get(1) {
1540            let n: u64 = num.as_str().parse().ok()?;
1541            let unit = caps.get(2)?.as_str().to_lowercase();
1542            let multiplier = match unit.as_str() {
1543                "hour" | "hours" | "hr" | "hrs" => 3600,
1544                "minute" | "minutes" | "min" | "mins" => 60,
1545                "day" | "days" => 86400,
1546                _ => return None,
1547            };
1548            n.checked_mul(multiplier)?
1549        } else if let Some(word) = caps.get(3) {
1550            match word.as_str().to_lowercase().as_str() {
1551                "hourly" => 3600,
1552                "daily" => 86400,
1553                "weekly" => 604800,
1554                "monthly" => 2592000,
1555                _ => return None,
1556            }
1557        } else {
1558            continue;
1559        };
1560        max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1561    }
1562
1563    max_secs
1564}
1565
1566#[cfg(test)]
1567pub(crate) async fn reset_global_state() {
1568    cooldowns().lock().unwrap().clear();
1569    *cooldown_store().lock().await = None;
1570    github_5xx_timestamps().lock().unwrap().clear();
1571    *github_5xx_circuit_open().lock().unwrap() = None;
1572    degraded_agents().lock().unwrap().clear();
1573    crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1574}
1575
1576#[cfg(test)]
1577mod tests {
1578    use super::*;
1579    use serial_test::serial;
1580
1581    async fn test_store() -> Arc<crate::store::TaskStore> {
1582        Arc::new(
1583            crate::store::TaskStore::open_memory()
1584                .await
1585                .expect("in-memory store"),
1586        )
1587    }
1588
1589    #[serial(cooldown_state)]
1590    #[tokio::test]
1591    async fn init_loads_unexpired_cooldowns_from_kv() {
1592        super::reset_global_state().await;
1593        let store = test_store().await;
1594        let cooldown_until = chrono::Utc::now().timestamp() + 3600; // 1 hour from now
1595
1596        store
1597            .kv_set(
1598                &format!("{KV_PREFIX}kimi:k2p5"),
1599                &cooldown_until.to_string(),
1600            )
1601            .await
1602            .unwrap();
1603
1604        // Also add an already-expired entry
1605        let expired = chrono::Utc::now().timestamp() - 1;
1606        store
1607            .kv_set(
1608                &format!("{KV_PREFIX}opencode:old-free"),
1609                &expired.to_string(),
1610            )
1611            .await
1612            .unwrap();
1613
1614        {
1615            let mut map = cooldown_lock();
1616            map.clear();
1617        }
1618
1619        init_cooldown_store(store).await;
1620
1621        assert!(
1622            is_model_in_cooldown("kimi", "k2p5"),
1623            "unexpired cooldown should be loaded"
1624        );
1625        assert!(
1626            !is_model_in_cooldown("opencode", "old-free"),
1627            "expired cooldown should not be loaded"
1628        );
1629    }
1630
1631    #[serial(cooldown_state)]
1632    #[tokio::test]
1633    async fn record_model_failure_writes_to_kv() {
1634        super::reset_global_state().await;
1635        let store = test_store().await;
1636
1637        *cooldown_store().lock().await = Some(store.clone());
1638
1639        record_model_failure("testagent_persist", "testmodel_persist").await;
1640
1641        // Give the async KV write a moment to settle before querying.
1642        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1643
1644        let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1645        let stored = store.kv_get(&kv_key).await.unwrap();
1646        assert!(stored.is_some(), "model failure should be persisted to KV");
1647        let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1648        let now = chrono::Utc::now().timestamp();
1649        // Should be in the future (cooldown_until, not failed_at)
1650        assert!(ts > now, "persisted timestamp should be in the future");
1651        assert_eq!(reason, "model_error");
1652    }
1653
1654    #[serial(cooldown_state)]
1655    #[tokio::test]
1656    async fn cooldown_reason_survives_restart_via_kv() {
1657        super::reset_global_state().await;
1658        let store = test_store().await;
1659        let agent = "test_reason_restart";
1660        let future = chrono::Utc::now().timestamp() + 600;
1661
1662        init_cooldown_store(store.clone()).await;
1663        set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1664
1665        // Simulate a service restart: drop the in-memory map, reload from KV.
1666        {
1667            let mut map = cooldown_lock();
1668            map.clear();
1669        }
1670        init_cooldown_store(store).await;
1671
1672        let entries = list_all_cooldowns();
1673        let entry = entries.iter().find(|(key, ..)| key == agent);
1674        assert_eq!(
1675            entry.map(|(_, _, reason)| reason.as_str()),
1676            Some("billing_cycle_exhausted"),
1677            "cooldown reason must survive a restart, not collapse to 'persisted'"
1678        );
1679    }
1680
1681    #[serial(cooldown_state)]
1682    #[tokio::test]
1683    async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1684        super::reset_global_state().await;
1685        let store = test_store().await;
1686        *cooldown_store().lock().await = Some(store.clone());
1687
1688        let agent = "testagent_persistent";
1689        let model = "testmodel_persistent";
1690        let key = format!("{agent}:{model}");
1691
1692        // First failure -> base 4h.
1693        record_persistent_model_failure(agent, model).await;
1694        let now = chrono::Utc::now().timestamp();
1695        let first_until = cooldown_until(&key).expect("first cooldown should exist");
1696        let first_remaining = first_until.saturating_sub(now);
1697        assert!(
1698            first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1699            "first persistent model cooldown should be ~4h, got {first_remaining}s"
1700        );
1701
1702        // Second failure -> 12h (3x growth).
1703        record_persistent_model_failure(agent, model).await;
1704        let now2 = chrono::Utc::now().timestamp();
1705        let second_until = cooldown_until(&key).expect("second cooldown should exist");
1706        let second_remaining = second_until.saturating_sub(now2);
1707        assert!(
1708            second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1709            "second persistent model cooldown should be ~12h, got {second_remaining}s"
1710        );
1711
1712        let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1713        let stored_count = store
1714            .kv_get(&count_key)
1715            .await
1716            .expect("kv read should succeed")
1717            .expect("failure count should be written");
1718        assert_eq!(stored_count, "2");
1719    }
1720
1721    #[serial(cooldown_state)]
1722    #[tokio::test]
1723    async fn record_agent_success_does_not_reset_persistent_failure_count() {
1724        super::reset_global_state().await;
1725        let store = test_store().await;
1726        *cooldown_store().lock().await = Some(store.clone());
1727
1728        let agent = "testagent_persistent_success";
1729        let model = "testmodel_persistent_success";
1730        let key = format!("{agent}:{model}");
1731
1732        // Two persistent failures escalate the counter to 2 (~12h cooldown).
1733        record_persistent_model_failure(agent, model).await;
1734        record_persistent_model_failure(agent, model).await;
1735
1736        let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1737        let stored_count = store
1738            .kv_get(&count_key)
1739            .await
1740            .expect("kv read should succeed")
1741            .expect("failure count should be written");
1742        assert_eq!(stored_count, "2");
1743
1744        // A successful run resets the generic failure_count but must NOT touch
1745        // the persistent counter — otherwise an intervening success would erase
1746        // escalation progress for a structurally unreliable model.
1747        record_agent_success(agent, model).await;
1748
1749        let stored_count_after_success = store
1750            .kv_get(&count_key)
1751            .await
1752            .expect("kv read should succeed")
1753            .expect("persistent failure count should still exist");
1754        assert_eq!(
1755            stored_count_after_success, "2",
1756            "persistent failure count must survive an intervening success"
1757        );
1758
1759        // The next persistent failure should escalate from 3 (~36h), not reset to base (4h).
1760        record_persistent_model_failure(agent, model).await;
1761        let now = chrono::Utc::now().timestamp();
1762        let third_until = cooldown_until(&key).expect("third cooldown should exist");
1763        let third_remaining = third_until.saturating_sub(now);
1764        assert!(
1765            third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1766            "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1767        );
1768    }
1769
1770    #[serial(cooldown_state)]
1771    #[tokio::test]
1772    async fn record_persistent_model_failure_stays_capped_at_7_days() {
1773        super::reset_global_state().await;
1774        let store = test_store().await;
1775        *cooldown_store().lock().await = Some(store.clone());
1776
1777        let agent = "testagent_persistent_capped";
1778        let model = "testmodel_persistent_capped";
1779        let key = format!("{agent}:{model}");
1780
1781        for _ in 0..9 {
1782            record_persistent_model_failure(agent, model).await;
1783        }
1784
1785        let remaining = cooldown_until(&key)
1786            .expect("cooldown should exist")
1787            .saturating_sub(chrono::Utc::now().timestamp());
1788        assert!(
1789            remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1790            "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1791        );
1792        assert!(
1793            remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1794            "persistent model cooldown should not exceed 7 days, got {remaining}s"
1795        );
1796    }
1797
1798    #[serial(cooldown_state)]
1799    #[test]
1800    fn parse_retry_at_codex_format() {
1801        let msg = "You've hit your usage limit. Upgrade to Pro, visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at Mar 26th, 2026 5:55 AM.";
1802        let ts = parse_retry_at(msg);
1803        assert!(ts.is_some(), "should parse codex retry-at date");
1804
1805        // The parsed timestamp must equal "Mar 26 2026 05:55 local" in UTC.
1806        // We reconstruct what the function should produce and compare directly.
1807        use chrono::TimeZone;
1808        let naive = chrono::NaiveDateTime::new(
1809            chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1810            chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1811        );
1812        let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1813            chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1814            chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1815            chrono::LocalResult::None => naive.and_utc(),
1816        };
1817        assert_eq!(
1818            ts.unwrap(),
1819            expected_utc.timestamp(),
1820            "timestamp should reflect local-timezone interpretation of the vendor date"
1821        );
1822
1823        // Sanity: the UTC date must still be March 26 (true for all timezones within ±14h)
1824        let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1825        use chrono::Datelike;
1826        assert_eq!(dt.month(), 3);
1827        assert_eq!(dt.day(), 26);
1828    }
1829
1830    #[serial(cooldown_state)]
1831    #[test]
1832    fn parse_retry_at_hour_only_and_month_day() {
1833        let now = chrono::Utc::now().timestamp();
1834        for msg in [
1835            "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1836            "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1837            "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1838            "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1839        ] {
1840            let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1841            assert!(ts > now, "{msg}");
1842        }
1843    }
1844
1845    #[serial(cooldown_state)]
1846    #[test]
1847    fn parse_retry_at_weekly_resets_without_year() {
1848        use chrono::Datelike;
1849        let t = chrono::Local::now() + chrono::Duration::days(2);
1850        let msg = format!(
1851            "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1852            t.format("%b"),
1853            t.day()
1854        );
1855        let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1856        let remaining = ts - chrono::Utc::now().timestamp();
1857        assert!(
1858            remaining > 3600 && remaining < 3 * 86400,
1859            "expected ~2d retry-at, got {remaining}s"
1860        );
1861    }
1862
1863    #[serial(cooldown_state)]
1864    #[test]
1865    fn parse_retry_at_no_date() {
1866        assert!(parse_retry_at("generic rate limit error").is_none());
1867        assert!(parse_retry_at("").is_none());
1868    }
1869
1870    #[serial(cooldown_state)]
1871    #[test]
1872    fn parse_retry_at_glm_reset_format_with_seconds() {
1873        let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
1874        let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
1875
1876        use chrono::TimeZone;
1877        let naive = chrono::NaiveDateTime::new(
1878            chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
1879            chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
1880        );
1881        let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1882            chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1883            chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1884            chrono::LocalResult::None => naive.and_utc(),
1885        };
1886
1887        assert_eq!(
1888            ts,
1889            expected_utc.timestamp(),
1890            "timestamp should reflect local-timezone interpretation of vendor reset date"
1891        );
1892    }
1893
1894    #[serial(cooldown_state)]
1895    #[test]
1896    fn parse_retry_at_billing_cycle_without_date_returns_none() {
1897        // Without a specific "try again at" date, parse_retry_at should return
1898        // None so callers use their default cooldown (30m agent / 1h model).
1899        let msg = "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle. Upgrade to get more.";
1900        assert!(
1901            parse_retry_at(msg).is_none(),
1902            "billing cycle without date should return None"
1903        );
1904    }
1905
1906    #[serial(cooldown_state)]
1907    #[test]
1908    fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
1909        let msg = "Failed to authenticate. API Error: 403 You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
1910        assert_eq!(
1911            parse_relative_usage_window(msg),
1912            Some(5 * 3600),
1913            "should extract the 5-hour window as seconds"
1914        );
1915    }
1916
1917    #[serial(cooldown_state)]
1918    #[test]
1919    fn parse_relative_usage_window_handles_common_durations() {
1920        assert_eq!(
1921            parse_relative_usage_window("You've reached your hourly usage limit."),
1922            Some(3600)
1923        );
1924        assert_eq!(
1925            parse_relative_usage_window("You've reached your daily usage limit."),
1926            Some(86400)
1927        );
1928        assert_eq!(
1929            parse_relative_usage_window("You've reached your weekly limit."),
1930            Some(604800)
1931        );
1932        assert_eq!(
1933            parse_relative_usage_window("You've reached your 30-minute usage limit."),
1934            Some(1800)
1935        );
1936        assert_eq!(
1937            parse_relative_usage_window("You've reached your 2-day quota."),
1938            Some(172800)
1939        );
1940    }
1941
1942    #[serial(cooldown_state)]
1943    #[test]
1944    fn parse_relative_usage_window_returns_none_without_usage_context() {
1945        // A bare rate-limit message with a number should not be interpreted as a
1946        // relative usage window unless usage/quota/window context is present.
1947        assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
1948        assert!(parse_relative_usage_window("rate limit exceeded").is_none());
1949        assert!(parse_relative_usage_window("").is_none());
1950    }
1951
1952    #[serial(cooldown_state)]
1953    #[tokio::test]
1954    async fn record_rate_limit_uses_model_specific_window_cooldown() {
1955        super::reset_global_state().await;
1956        let store = test_store().await;
1957        *cooldown_store().lock().await = Some(store.clone());
1958
1959        let agent = "test_record_rl_window_agent";
1960        let model = "test_record_rl_window_model";
1961        let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
1962
1963        record_rate_limit(agent, Some(model), msg).await;
1964
1965        assert!(
1966            is_model_in_cooldown(agent, model),
1967            "model should be cooled for the relative usage window"
1968        );
1969        assert!(
1970            !is_agent_in_cooldown(agent),
1971            "agent should not be cooled for a model-scoped quota"
1972        );
1973
1974        let key = format!("{agent}:{model}");
1975        let remaining = cooldown_until(&key)
1976            .expect("model cooldown should exist")
1977            .saturating_sub(chrono::Utc::now().timestamp());
1978        assert!(
1979            remaining >= 5 * 3600 - 5,
1980            "model cooldown should be at least ~5 hours, got {remaining}s"
1981        );
1982        assert_eq!(
1983            cooldown_reason(&key).as_deref(),
1984            Some("rate_limit_usage_window"),
1985            "relative usage-window cooldown must be labeled as a rate limit, not silence"
1986        );
1987    }
1988
1989    #[serial(cooldown_state)]
1990    #[tokio::test]
1991    async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
1992        super::reset_global_state().await;
1993        let store = test_store().await;
1994        *cooldown_store().lock().await = Some(store.clone());
1995
1996        let agent = "test_record_rl_fallback_agent";
1997        let model = "test_record_rl_fallback_model";
1998        let msg = "429 Too Many Requests";
1999
2000        record_rate_limit(agent, Some(model), msg).await;
2001
2002        assert!(
2003            is_agent_in_cooldown(agent),
2004            "agent should be cooled when no relative window is detected"
2005        );
2006        assert!(
2007            is_model_in_cooldown(agent, model),
2008            "model should be cooled when no relative window is detected"
2009        );
2010    }
2011
2012    #[serial(cooldown_state)]
2013    #[tokio::test]
2014    async fn agent_cooldown_persists_to_kv() {
2015        super::reset_global_state().await;
2016        // Agent-level cooldowns persist to KV via set_cooldown_async
2017        let agent = "test_agent_persist_check";
2018        record_agent_failure_with_message(agent, "").await;
2019        assert!(is_agent_in_cooldown(agent));
2020    }
2021
2022    #[serial(cooldown_state)]
2023    #[tokio::test]
2024    async fn cooldown_never_shortens() {
2025        super::reset_global_state().await;
2026        let agent = "test_agent_no_shorten";
2027        set_agent_cooldown(agent, 24 * 60 * 60).await;
2028        let initial = {
2029            let map = cooldowns().lock().unwrap();
2030            map.get(agent)
2031                .expect("cooldown entry should exist")
2032                .cooldown_until
2033        };
2034
2035        // Shorter cooldown attempt should not override the existing one.
2036        record_agent_failure_with_message(agent, "rate limit").await;
2037        let after = {
2038            let map = cooldowns().lock().unwrap();
2039            map.get(agent)
2040                .expect("cooldown entry should exist")
2041                .cooldown_until
2042        };
2043        assert_eq!(
2044            initial, after,
2045            "cooldown should not be shortened by later failures"
2046        );
2047    }
2048
2049    #[serial(cooldown_state)]
2050    #[tokio::test]
2051    async fn silence_agent_cooldown_is_short_lived() {
2052        super::reset_global_state().await;
2053        let agent = "test_silence_agent_cd";
2054        assert!(!is_agent_in_cooldown(agent));
2055
2056        set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2057        assert!(is_agent_in_cooldown(agent));
2058
2059        // Verify it's a short cooldown (120s), not the exponential backoff one
2060        let map = cooldowns().lock().unwrap();
2061        let entry = map.get(agent).expect("should have cooldown entry");
2062        let now = chrono::Utc::now().timestamp();
2063        let remaining = entry.cooldown_until - now;
2064        assert!(
2065            remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
2066            "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2067        );
2068        assert!(remaining > 0, "cooldown should still be active");
2069    }
2070
2071    #[serial(cooldown_state)]
2072    #[tokio::test]
2073    async fn record_silence_detection_applies_extended_cooldown() {
2074        super::reset_global_state().await;
2075        let store = test_store().await;
2076        *cooldown_store().lock().await = Some(store.clone());
2077
2078        let agent = "test_silence_count_agent";
2079        let model = "test_silence_count_model";
2080        assert!(!is_model_in_cooldown(agent, model));
2081
2082        for _ in 0..SILENCE_COUNT_THRESHOLD {
2083            let result = record_silence_detection(agent, model).await;
2084            assert!(result.is_some());
2085        }
2086
2087        assert!(is_model_in_cooldown(agent, model));
2088        let key = format!("{agent}:{model}");
2089        let remaining = {
2090            let map = cooldowns().lock().unwrap();
2091            let entry = map.get(&key).expect("cooldown entry should exist");
2092            entry.cooldown_until - chrono::Utc::now().timestamp()
2093        };
2094        assert!(
2095            remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
2096            "extended cooldown should be applied, got {remaining}s"
2097        );
2098
2099        let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2100        let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2101        assert_eq!(stored, "[]");
2102    }
2103
2104    #[serial(cooldown_state)]
2105    #[test]
2106    fn detect_credit_exhaustion_out_of_credits() {
2107        let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2108        assert_eq!(
2109            detect_credit_exhaustion(msg),
2110            Some(CreditExhaustionReason::OutOfCredits)
2111        );
2112    }
2113
2114    #[serial(cooldown_state)]
2115    #[test]
2116    fn detect_credit_exhaustion_out_of_credits_variations() {
2117        assert_eq!(
2118            detect_credit_exhaustion("out of credits"),
2119            Some(CreditExhaustionReason::OutOfCredits)
2120        );
2121        assert_eq!(
2122            detect_credit_exhaustion("outofcredits"),
2123            Some(CreditExhaustionReason::OutOfCredits)
2124        );
2125        assert_eq!(
2126            detect_credit_exhaustion("no credits remaining"),
2127            Some(CreditExhaustionReason::OutOfCredits)
2128        );
2129        assert_eq!(
2130            detect_credit_exhaustion("insufficient funds"),
2131            Some(CreditExhaustionReason::OutOfCredits)
2132        );
2133    }
2134
2135    #[serial(cooldown_state)]
2136    #[test]
2137    fn detect_credit_exhaustion_org_level_disabled() {
2138        let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2139        assert_eq!(
2140            detect_credit_exhaustion(msg),
2141            Some(CreditExhaustionReason::OrgLevelDisabled)
2142        );
2143    }
2144
2145    #[serial(cooldown_state)]
2146    #[test]
2147    fn detect_credit_exhaustion_org_level_disabled_variations() {
2148        assert_eq!(
2149            detect_credit_exhaustion("organization has been disabled"),
2150            Some(CreditExhaustionReason::OrgLevelDisabled)
2151        );
2152        assert_eq!(
2153            detect_credit_exhaustion("org disabled"),
2154            Some(CreditExhaustionReason::OrgLevelDisabled)
2155        );
2156        assert_eq!(
2157            detect_credit_exhaustion("billing quota exceeded"),
2158            Some(CreditExhaustionReason::OrgLevelDisabled)
2159        );
2160        assert_eq!(
2161            detect_credit_exhaustion("credit balance too low"),
2162            Some(CreditExhaustionReason::OrgLevelDisabled)
2163        );
2164    }
2165
2166    #[serial(cooldown_state)]
2167    #[test]
2168    fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2169        let msg = "overageDisabledReason: out_of_credits";
2170        assert_eq!(
2171            detect_credit_exhaustion(msg),
2172            Some(CreditExhaustionReason::OutOfCredits)
2173        );
2174
2175        let msg = "overageDisabledReason: org_level_disabled";
2176        assert_eq!(
2177            detect_credit_exhaustion(msg),
2178            Some(CreditExhaustionReason::OrgLevelDisabled)
2179        );
2180    }
2181
2182    #[serial(cooldown_state)]
2183    #[test]
2184    fn detect_credit_exhaustion_glm_insufficient_balance() {
2185        // GLM/ZhipuAI 429: "Insufficient balance or no resource package. Please recharge."
2186        assert_eq!(
2187            detect_credit_exhaustion(
2188                "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2189            ),
2190            Some(CreditExhaustionReason::OutOfCredits)
2191        );
2192        assert_eq!(
2193            detect_credit_exhaustion("Insufficient balance"),
2194            Some(CreditExhaustionReason::OutOfCredits)
2195        );
2196        assert_eq!(
2197            detect_credit_exhaustion("please recharge your account"),
2198            Some(CreditExhaustionReason::OutOfCredits)
2199        );
2200        assert_eq!(
2201            detect_credit_exhaustion("no resource package available"),
2202            Some(CreditExhaustionReason::OutOfCredits)
2203        );
2204    }
2205
2206    #[serial(cooldown_state)]
2207    #[test]
2208    fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2209        assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2210        assert!(detect_credit_exhaustion("too many requests").is_none());
2211        assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2212        assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2213        assert!(detect_credit_exhaustion(
2214            "You've hit your usage limit, try again in a few minutes."
2215        )
2216        .is_none());
2217    }
2218
2219    #[serial(cooldown_state)]
2220    #[tokio::test]
2221    async fn record_credit_exhaustion_applies_agent_cooldown() {
2222        super::reset_global_state().await;
2223        let agent = "test_credit_exhaust_agent";
2224        assert!(!is_agent_in_cooldown(agent));
2225
2226        record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2227        assert!(is_agent_in_cooldown(agent));
2228
2229        let remaining = {
2230            let map = cooldowns().lock().unwrap();
2231            let entry = map.get(agent).expect("cooldown entry should exist");
2232            entry.cooldown_until - chrono::Utc::now().timestamp()
2233        };
2234        // First failure: base = 1h
2235        assert!(
2236            remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2237            "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2238        );
2239        assert!(
2240            remaining <= CREDIT_BACKOFF_MAX_SECS,
2241            "cooldown should not exceed cap of 8h, got {remaining}s"
2242        );
2243    }
2244
2245    #[serial(cooldown_state)]
2246    #[tokio::test]
2247    async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2248        super::reset_global_state().await;
2249        let agent = "test_org_disabled_agent";
2250        assert!(!is_agent_in_cooldown(agent));
2251
2252        record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2253        assert!(is_agent_in_cooldown(agent));
2254
2255        let remaining = {
2256            let map = cooldowns().lock().unwrap();
2257            let entry = map.get(agent).expect("cooldown entry should exist");
2258            entry.cooldown_until - chrono::Utc::now().timestamp()
2259        };
2260        // First failure: base = 2h (ORG_BACKOFF_BASE_SECS)
2261        assert!(
2262            remaining >= ORG_BACKOFF_BASE_SECS - 5,
2263            "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2264        );
2265        assert!(
2266            remaining <= CREDIT_BACKOFF_MAX_SECS,
2267            "cooldown should not exceed cap of 8h, got {remaining}s"
2268        );
2269    }
2270
2271    #[serial(cooldown_state)]
2272    #[tokio::test]
2273    async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2274        super::reset_global_state().await;
2275        let store = test_store().await;
2276        *cooldown_store().lock().await = Some(store.clone());
2277
2278        let agent = "test_billing_cycle_agent";
2279        assert!(!is_agent_in_cooldown(agent));
2280
2281        record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2282        assert!(is_agent_in_cooldown(agent));
2283
2284        let remaining = {
2285            let map = cooldowns().lock().unwrap();
2286            let entry = map.get(agent).expect("cooldown entry should exist");
2287            entry.cooldown_until - chrono::Utc::now().timestamp()
2288        };
2289        // First failure: base = 24h
2290        assert!(
2291            remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2292            "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2293        );
2294        assert!(
2295            remaining <= BILLING_CYCLE_MAX_SECS,
2296            "cooldown should not exceed cap of 7 days, got {remaining}s"
2297        );
2298    }
2299
2300    #[serial(cooldown_state)]
2301    #[tokio::test]
2302    async fn record_credit_exhaustion_billing_cycle_escalates() {
2303        super::reset_global_state().await;
2304        let store = test_store().await;
2305        *cooldown_store().lock().await = Some(store.clone());
2306
2307        let agent = "test_billing_cycle_escalation";
2308
2309        // First failure: 24h
2310        record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2311        let remaining_1 = {
2312            let map = cooldowns().lock().unwrap();
2313            let entry = map.get(agent).expect("cooldown entry should exist");
2314            entry.cooldown_until - chrono::Utc::now().timestamp()
2315        };
2316        assert!(
2317            (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2318                .contains(&remaining_1),
2319            "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2320        );
2321
2322        // Clear in-memory cooldown (but NOT failure count) to allow next set_cooldown_async
2323        {
2324            let mut map = cooldowns().lock().unwrap();
2325            map.remove(agent);
2326        }
2327
2328        // Second failure: 24h * 3 = 72h (3 days)
2329        record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2330        let remaining_2 = {
2331            let map = cooldowns().lock().unwrap();
2332            let entry = map.get(agent).expect("cooldown entry should exist");
2333            entry.cooldown_until - chrono::Utc::now().timestamp()
2334        };
2335        let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; // 72h
2336        assert!(
2337            remaining_2 >= expected_2 - 5,
2338            "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2339        );
2340
2341        // Clear in-memory cooldown again
2342        {
2343            let mut map = cooldowns().lock().unwrap();
2344            map.remove(agent);
2345        }
2346
2347        // Third failure: 24h * 9 = 216h → capped at 7 days (168h)
2348        record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2349        let remaining_3 = {
2350            let map = cooldowns().lock().unwrap();
2351            let entry = map.get(agent).expect("cooldown entry should exist");
2352            entry.cooldown_until - chrono::Utc::now().timestamp()
2353        };
2354        assert!(
2355            remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2356            "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2357        );
2358        assert!(
2359            remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2360            "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2361        );
2362    }
2363
2364    #[serial(cooldown_state)]
2365    #[tokio::test]
2366    async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2367        super::reset_global_state().await;
2368        let store = test_store().await;
2369        *cooldown_store().lock().await = Some(store.clone());
2370
2371        let agent = "test_billing_cycle_agent_capped";
2372
2373        for _ in 0..5 {
2374            record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2375        }
2376
2377        let remaining = {
2378            let map = cooldowns().lock().unwrap();
2379            let entry = map.get(agent).expect("cooldown entry should exist");
2380            entry.cooldown_until - chrono::Utc::now().timestamp()
2381        };
2382        assert!(
2383            remaining >= BILLING_CYCLE_MAX_SECS - 5,
2384            "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2385        );
2386        assert!(
2387            remaining <= BILLING_CYCLE_MAX_SECS + 5,
2388            "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2389        );
2390    }
2391
2392    #[serial(cooldown_state)]
2393    #[test]
2394    fn credit_exhaustion_backoff_constants() {
2395        // OutOfCredits first failure: 1h base
2396        assert_eq!(
2397            compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2398            CREDIT_BACKOFF_BASE_SECS
2399        );
2400        // OrgLevelDisabled first failure: 2h base
2401        assert_eq!(
2402            compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2403            ORG_BACKOFF_BASE_SECS
2404        );
2405        // BillingCycleExhausted: 24h base, 7d cap
2406        assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2407        assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2408    }
2409
2410    #[serial(cooldown_state)]
2411    #[test]
2412    fn compute_backoff_long_cap_without_extended_stays_capped() {
2413        assert_eq!(
2414            compute_backoff(
2415                3,
2416                BILLING_CYCLE_COOLDOWN_SECS,
2417                BILLING_CYCLE_MAX_SECS,
2418                false
2419            ),
2420            BILLING_CYCLE_MAX_SECS,
2421            "7-day capped paths should stop at the documented max"
2422        );
2423        assert_eq!(
2424            compute_backoff(
2425                9,
2426                PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2427                PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2428                false,
2429            ),
2430            PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2431            "persistent model failures should not extend past 7 days"
2432        );
2433    }
2434
2435    // ---- Degraded-agent tracking tests ----
2436
2437    #[serial(cooldown_state)]
2438    #[test]
2439    fn mark_and_check_agent_degraded() {
2440        let agent = "test_degraded_mark";
2441        assert!(!is_agent_degraded(agent));
2442        mark_agent_degraded(agent);
2443        assert!(is_agent_degraded(agent));
2444        clear_agent_degraded(agent);
2445        assert!(!is_agent_degraded(agent));
2446    }
2447
2448    #[serial(cooldown_state)]
2449    #[tokio::test]
2450    async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2451        super::reset_global_state().await;
2452        let store = test_store().await;
2453        let agent = "test_sync_external_clear";
2454        let future = chrono::Utc::now().timestamp() + 600;
2455
2456        // Register the store and put a cooldown in both memory and KV.
2457        init_cooldown_store(store.clone()).await;
2458        set_cooldown_async(agent, future, "test").await;
2459        assert!(is_agent_in_cooldown(agent));
2460
2461        // Simulate `orch cooldown clear --all` from a separate CLI process: it
2462        // can only write KV (to "0"); the running engine's in-memory map is
2463        // untouched.
2464        store
2465            .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2466            .await
2467            .expect("kv_set tombstone");
2468        assert!(
2469            is_agent_in_cooldown(agent),
2470            "in-memory cooldown should still be present before sync"
2471        );
2472
2473        sync_from_kv(&store).await;
2474
2475        assert!(
2476            !is_agent_in_cooldown(agent),
2477            "sync_from_kv must drop in-memory entries the KV tombstoned"
2478        );
2479    }
2480
2481    #[serial(cooldown_state)]
2482    #[tokio::test]
2483    async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2484        // Race window: set_cooldown_async writes memory first, then KV. If a
2485        // sync runs between those two writes, the KV row is missing — we must
2486        // leave the in-memory entry alone, not drop it.
2487        super::reset_global_state().await;
2488        let store = test_store().await;
2489        let agent = "test_sync_race_window";
2490        let future = chrono::Utc::now().timestamp() + 600;
2491
2492        // Simulate the race: memory has the entry, KV does not.
2493        set_cooldown_in_memory(agent, future, "raced");
2494        assert!(is_agent_in_cooldown(agent));
2495
2496        sync_from_kv(&store).await;
2497
2498        assert!(
2499            is_agent_in_cooldown(agent),
2500            "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2501        );
2502    }
2503
2504    #[serial(cooldown_state)]
2505    #[tokio::test]
2506    async fn sync_from_kv_adopts_kv_timestamp() {
2507        // KV is the source of truth: if KV has a different active timestamp,
2508        // memory adopts it (covers external manual KV edits).
2509        super::reset_global_state().await;
2510        let store = test_store().await;
2511        let agent = "test_sync_adopt_kv";
2512        let original = chrono::Utc::now().timestamp() + 600;
2513        let updated = chrono::Utc::now().timestamp() + 1200;
2514
2515        init_cooldown_store(store.clone()).await;
2516        set_cooldown_async(agent, original, "initial").await;
2517
2518        store
2519            .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2520            .await
2521            .expect("kv_set updated");
2522
2523        sync_from_kv(&store).await;
2524
2525        let until = cooldown_until(agent).expect("cooldown still active");
2526        assert_eq!(
2527            until, updated,
2528            "sync_from_kv must adopt KV's authoritative timestamp"
2529        );
2530    }
2531
2532    #[serial(cooldown_state)]
2533    #[tokio::test]
2534    async fn refresh_degraded_agents_marks_cooled_agent() {
2535        super::reset_global_state().await;
2536        let store = test_store().await;
2537        let agent = "test_refresh_cooled";
2538
2539        // Put agent in cooldown
2540        record_agent_failure_with_message(agent, "").await;
2541        assert!(is_agent_in_cooldown(agent));
2542
2543        // model_checker always returns true (agent has models)
2544        let agents = vec![agent.to_string()];
2545        refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2546
2547        assert!(
2548            is_agent_degraded(agent),
2549            "agent in cooldown should be marked degraded"
2550        );
2551
2552        // Cleanup
2553        clear_agent_degraded(agent);
2554    }
2555
2556    #[serial(cooldown_state)]
2557    #[tokio::test]
2558    async fn refresh_degraded_agents_marks_no_models_agent() {
2559        super::reset_global_state().await;
2560        let store = test_store().await;
2561        let agent = "test_refresh_no_models";
2562
2563        // model_checker returns false (no available models)
2564        let agents = vec![agent.to_string()];
2565        refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2566
2567        assert!(
2568            is_agent_degraded(agent),
2569            "agent with no available models should be marked degraded"
2570        );
2571
2572        // Cleanup
2573        clear_agent_degraded(agent);
2574    }
2575
2576    #[serial(cooldown_state)]
2577    #[tokio::test]
2578    async fn refresh_degraded_agents_marks_rate_limited_agent() {
2579        super::reset_global_state().await;
2580        let store = test_store().await;
2581        let agent = "test_refresh_rate_limited";
2582
2583        // Insert rate limit events exceeding threshold
2584        for _ in 0..4 {
2585            store
2586                .record_rate_limit(agent, "rate_limit", None)
2587                .await
2588                .unwrap();
2589        }
2590
2591        let agents = vec![agent.to_string()];
2592        // threshold=3, window=6h — 4 events should trigger degraded
2593        refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2594
2595        assert!(
2596            is_agent_degraded(agent),
2597            "agent with rate limit count >= threshold should be marked degraded"
2598        );
2599
2600        // Cleanup
2601        clear_agent_degraded(agent);
2602    }
2603
2604    #[serial(cooldown_state)]
2605    #[tokio::test]
2606    async fn refresh_degraded_agents_clears_healthy_agent() {
2607        super::reset_global_state().await;
2608        let store = test_store().await;
2609        let agent = "test_refresh_healthy";
2610
2611        // Pre-mark as degraded
2612        mark_agent_degraded(agent);
2613        assert!(is_agent_degraded(agent));
2614
2615        // No cooldown, has models, no rate limit events
2616        let agents = vec![agent.to_string()];
2617        refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2618
2619        assert!(
2620            !is_agent_degraded(agent),
2621            "healthy agent should have degraded flag cleared"
2622        );
2623    }
2624
2625    #[serial(cooldown_state)]
2626    #[tokio::test]
2627    async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2628        super::reset_global_state().await;
2629        let store = test_store().await;
2630        let agent = "test_refresh_expired_cooldown";
2631
2632        // Manually insert an already-expired cooldown entry directly into the map
2633        // (simulating a stale entry that wasn't cleaned up)
2634        let expired_ts = chrono::Utc::now().timestamp() - 100; // 100 seconds ago
2635        {
2636            let mut map = cooldowns().lock().unwrap();
2637            map.insert(
2638                agent.to_string(),
2639                CooldownEntry {
2640                    cooldown_until: expired_ts,
2641                    reason: "expired_test".to_string(),
2642                },
2643            );
2644        }
2645
2646        // Verify the expired entry is in the map but should not cause degraded status
2647        assert!(
2648            !is_agent_in_cooldown(agent),
2649            "expired cooldown should not be active"
2650        );
2651
2652        // Pre-mark as degraded to verify it gets cleared
2653        mark_agent_degraded(agent);
2654        assert!(is_agent_degraded(agent));
2655
2656        // Run refresh — should clear degraded because cooldown is expired
2657        let agents = vec![agent.to_string()];
2658        refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2659
2660        assert!(
2661            !is_agent_degraded(agent),
2662            "agent with only expired cooldown should NOT be marked degraded"
2663        );
2664
2665        // Cleanup
2666        clear_agent_degraded(agent);
2667        {
2668            let mut map = cooldowns().lock().unwrap();
2669            map.remove(agent);
2670        }
2671    }
2672
2673    #[serial(cooldown_state)]
2674    #[test]
2675    fn compute_backoff_grows_exponentially() {
2676        // count=0 or 1: base
2677        assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2678        assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2679        // count=2: base * 3
2680        assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2681        // count=3: base * 9
2682        assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2683        // count=4: base * 27 = 8100 < 14400
2684        assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2685        // count=5: base * 81 = 24300 → capped at 14400
2686        assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2687    }
2688
2689    #[serial(cooldown_state)]
2690    #[test]
2691    fn compute_backoff_extended_tier_1_applies_24h() {
2692        // With base=300, max=14400:
2693        // count=4: 8100s (not yet at max)
2694        // count=5: 14400s (first time at max, times_at_max=0, no extension)
2695        // count=6: 14400s (times_at_max=1, no extension yet, need >=2)
2696        // count=7: 86400s (times_at_max=2, extended tier 1: 6x = 24h)
2697        assert_eq!(
2698            compute_backoff(5, 300, 14400, true),
2699            14400,
2700            "first at max, not extended"
2701        );
2702        assert_eq!(
2703            compute_backoff(6, 300, 14400, true),
2704            14400,
2705            "times_at_max=1, not extended"
2706        );
2707        assert_eq!(
2708            compute_backoff(7, 300, 14400, true),
2709            86400,
2710            "times_at_max=2, extended 6x"
2711        );
2712        assert_eq!(
2713            compute_backoff(8, 300, 14400, true),
2714            86400,
2715            "times_at_max=3, still extended 6x"
2716        );
2717    }
2718
2719    #[serial(cooldown_state)]
2720    #[test]
2721    fn compute_backoff_extended_tier_2_applies_48h() {
2722        // With base=300, max=14400:
2723        // count=10: times_at_max=5 (first_max=5, count=10), extended tier 2: 12x = 48h
2724        assert_eq!(
2725            compute_backoff(9, 300, 14400, true),
2726            86400,
2727            "times_at_max=4, still tier 1"
2728        );
2729        assert_eq!(
2730            compute_backoff(10, 300, 14400, true),
2731            172800,
2732            "times_at_max=5, extended 12x"
2733        );
2734        assert_eq!(
2735            compute_backoff(15, 300, 14400, true),
2736            172800,
2737            "times_at_max=10, still tier 2"
2738        );
2739        assert_eq!(
2740            compute_backoff(20, 300, 14400, true),
2741            172800,
2742            "times_at_max=15, still tier 2"
2743        );
2744    }
2745
2746    #[serial(cooldown_state)]
2747    #[test]
2748    fn compute_backoff_extended_tier_with_different_base_max() {
2749        // With base=60, max=3600:
2750        // count=3: 60*9=540 < 3600
2751        // count=4: 60*27=1620 < 3600
2752        // count=5: 60*81=4860 → capped at 3600
2753        // count=6: times_at_max=1, still 3600
2754        // count=7: times_at_max=2, extended 6x = 21600 (6h)
2755        assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2756        assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2757        assert_eq!(
2758            compute_backoff(7, 60, 3600, true),
2759            21600,
2760            "times_at_max=2, extended 6x"
2761        );
2762    }
2763
2764    #[serial(cooldown_state)]
2765    #[test]
2766    fn compute_backoff_extended_tier_persistent_fails_example() {
2767        // Real-world example from issue: glm:haiku with 21 failures
2768        // With base=300, max=14400:
2769        // count=21 → times_at_max=16 (first_max=5), extended tier 2: 12x = 48h
2770        assert_eq!(
2771            compute_backoff(21, 300, 14400, true),
2772            172800,
2773            "21 failures should get 48h extended cooldown"
2774        );
2775
2776        // kimi:haiku with 36 failures
2777        assert_eq!(
2778            compute_backoff(36, 300, 14400, true),
2779            172800,
2780            "36 failures should also get 48h extended cooldown"
2781        );
2782    }
2783
2784    #[serial(cooldown_state)]
2785    #[test]
2786    fn compute_backoff_extended_tier_zero_base() {
2787        // Zero base should still work without panicking
2788        assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2789        assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2790        assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2791    }
2792
2793    #[serial(cooldown_state)]
2794    #[test]
2795    fn detect_credit_exhaustion_billing_cycle() {
2796        assert_eq!(
2797            detect_credit_exhaustion("billing cycle"),
2798            Some(CreditExhaustionReason::BillingCycleExhausted)
2799        );
2800        assert_eq!(
2801            detect_credit_exhaustion("quota refreshed next cycle"),
2802            Some(CreditExhaustionReason::BillingCycleExhausted)
2803        );
2804        assert_eq!(
2805            detect_credit_exhaustion(
2806                "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
2807            ),
2808            Some(CreditExhaustionReason::BillingCycleExhausted)
2809        );
2810        assert_eq!(
2811            detect_credit_exhaustion("monthly quota exceeded"),
2812            Some(CreditExhaustionReason::BillingCycleExhausted)
2813        );
2814        assert_eq!(
2815            detect_credit_exhaustion(
2816                "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
2817            ),
2818            Some(CreditExhaustionReason::BillingCycleExhausted)
2819        );
2820    }
2821
2822    #[serial(cooldown_state)]
2823    #[test]
2824    fn detect_credit_exhaustion_billing_cycle_minimax() {
2825        assert_eq!(
2826            detect_credit_exhaustion(
2827                "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
2828            ),
2829            Some(CreditExhaustionReason::BillingCycleExhausted)
2830        );
2831        assert_eq!(
2832            detect_credit_exhaustion("Upgrade your Token Plan"),
2833            Some(CreditExhaustionReason::BillingCycleExhausted)
2834        );
2835        assert_eq!(
2836            detect_credit_exhaustion("Token Plan usage limit reached"),
2837            Some(CreditExhaustionReason::BillingCycleExhausted)
2838        );
2839        assert_eq!(
2840            detect_credit_exhaustion("purchase Credits for more usage"),
2841            Some(CreditExhaustionReason::BillingCycleExhausted)
2842        );
2843        assert_eq!(
2844            detect_credit_exhaustion(
2845                "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits"
2846            ),
2847            Some(CreditExhaustionReason::BillingCycleExhausted)
2848        );
2849    }
2850
2851    #[serial(cooldown_state)]
2852    #[test]
2853    fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
2854        // Regression test for issue #3529: the exact Kimi billing-cycle exhaustion
2855        // message from task_runs (tasks 156734/156678) must resolve to the
2856        // "billing_cycle_exhausted" outcome string, not fall through to a generic
2857        // rate-limit/auth classification.
2858        let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
2859                    this billing cycle. Your quota will be refreshed soon.";
2860        assert_eq!(
2861            detect_credit_exhaustion_outcome(msg),
2862            Some("billing_cycle_exhausted")
2863        );
2864    }
2865
2866    #[serial(cooldown_state)]
2867    #[test]
2868    fn detect_credit_exhaustion_outcome_maps_all_reasons() {
2869        assert_eq!(
2870            detect_credit_exhaustion_outcome("out of credits"),
2871            Some("credit_exhausted")
2872        );
2873        assert_eq!(
2874            detect_credit_exhaustion_outcome("organization has been disabled"),
2875            Some("org_level_disabled")
2876        );
2877        assert_eq!(
2878            detect_credit_exhaustion_outcome("monthly quota exceeded"),
2879            Some("billing_cycle_exhausted")
2880        );
2881    }
2882
2883    #[test]
2884    fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
2885        assert_eq!(
2886            detect_credit_exhaustion_outcome("429 Too Many Requests"),
2887            None
2888        );
2889        assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
2890    }
2891
2892    #[serial(cooldown_state)]
2893    #[test]
2894    fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
2895        // A message containing both overage and billing cycle language should resolve
2896        // to OrgLevelDisabled/OutOfCredits (detected first) not BillingCycleExhausted.
2897        let msg = "overageDisabledReason: org_level_disabled for billing cycle";
2898        assert_eq!(
2899            detect_credit_exhaustion(msg),
2900            Some(CreditExhaustionReason::OrgLevelDisabled)
2901        );
2902    }
2903
2904    #[serial(cooldown_state)]
2905    #[tokio::test]
2906    async fn refresh_degraded_agents_below_threshold_not_degraded() {
2907        super::reset_global_state().await;
2908        let store = test_store().await;
2909        let agent = "test_refresh_below_threshold";
2910
2911        // Insert rate limit events below threshold
2912        for _ in 0..2 {
2913            store
2914                .record_rate_limit(agent, "rate_limit", None)
2915                .await
2916                .unwrap();
2917        }
2918
2919        let agents = vec![agent.to_string()];
2920        // threshold=3, only 2 events — should NOT be degraded
2921        refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2922
2923        assert!(
2924            !is_agent_degraded(agent),
2925            "agent with rate limit count < threshold should not be degraded"
2926        );
2927    }
2928
2929    // ---- Exponential backoff integration tests ----
2930
2931    #[serial(cooldown_state)]
2932    #[tokio::test]
2933    async fn repeated_agent_failures_escalate_backoff() {
2934        super::reset_global_state().await;
2935        let store = test_store().await;
2936        *cooldown_store().lock().await = Some(store.clone());
2937
2938        let agent = "test_escalation_agent";
2939
2940        // First failure: should be ~5 min (BACKOFF_BASE_SECS)
2941        record_agent_failure_with_message(agent, "").await;
2942        let remaining_1 = {
2943            let map = cooldowns().lock().unwrap();
2944            let entry = map.get(agent).expect("should have cooldown");
2945            entry.cooldown_until - chrono::Utc::now().timestamp()
2946        };
2947        assert!(
2948            (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
2949            "first failure should be ~5 min, got {remaining_1}s"
2950        );
2951
2952        // Clear the cooldown (but NOT failure count) to allow the next set_cooldown_async to apply
2953        {
2954            let mut map = cooldowns().lock().unwrap();
2955            map.remove(agent);
2956        }
2957
2958        // Second failure: should be ~15 min (5 * 3)
2959        record_agent_failure_with_message(agent, "").await;
2960        let remaining_2 = {
2961            let map = cooldowns().lock().unwrap();
2962            let entry = map.get(agent).expect("should have cooldown");
2963            entry.cooldown_until - chrono::Utc::now().timestamp()
2964        };
2965        assert!(
2966            remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
2967            "second failure should be ~15 min, got {remaining_2}s"
2968        );
2969    }
2970
2971    #[serial(cooldown_state)]
2972    #[tokio::test]
2973    async fn vendor_past_retry_falls_back_to_backoff() {
2974        super::reset_global_state().await;
2975        // If a vendor-provided "try again at" date is in the past, the
2976        // code must NOT early-return and must instead apply the normal
2977        // exponential backoff (increment failure count and set a future cooldown).
2978        let store = test_store().await;
2979        *cooldown_store().lock().await = Some(store.clone());
2980
2981        let agent = "test_vendor_past_retry";
2982
2983        // Ensure no pre-existing cooldown
2984        {
2985            let mut map = cooldowns().lock().unwrap();
2986            map.remove(agent);
2987        }
2988
2989        // A clearly past date that parse_retry_at can parse
2990        let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
2991
2992        record_agent_failure_with_message(agent, past_msg).await;
2993
2994        // After handling a past vendor timestamp, we should have applied
2995        // exponential backoff (base = BACKOFF_BASE_SECS) and a persisted
2996        // failure count of 1.
2997        assert!(
2998            is_agent_in_cooldown(agent),
2999            "agent should be in cooldown after failure"
3000        );
3001
3002        let remaining = {
3003            let map = cooldowns().lock().unwrap();
3004            let entry = map.get(agent).expect("cooldown entry should exist");
3005            entry.cooldown_until - chrono::Utc::now().timestamp()
3006        };
3007
3008        assert!(
3009            remaining >= BACKOFF_BASE_SECS - 5,
3010            "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3011        );
3012
3013        // Verify the failure count was incremented and persisted to KV
3014        let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3015        let stored = store
3016            .kv_get(&kv_key)
3017            .await
3018            .unwrap()
3019            .expect("failure count should be stored");
3020        assert_eq!(stored, "1");
3021    }
3022
3023    #[serial(cooldown_state)]
3024    #[tokio::test]
3025    async fn record_agent_success_resets_backoff() {
3026        super::reset_global_state().await;
3027        let store = test_store().await;
3028        *cooldown_store().lock().await = Some(store.clone());
3029
3030        let agent = "test_success_reset_agent";
3031        let model = "test_model";
3032
3033        // Accumulate 3 failures to escalate the backoff
3034        for _ in 0..3 {
3035            let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3036        }
3037        // Verify count is 3
3038        let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3039        let count: u32 = store
3040            .kv_get(&kv_key)
3041            .await
3042            .unwrap()
3043            .unwrap()
3044            .parse()
3045            .unwrap();
3046        assert_eq!(count, 3);
3047
3048        // Success should reset to 0
3049        record_agent_success(agent, model).await;
3050
3051        let count_after: u32 = store
3052            .kv_get(&kv_key)
3053            .await
3054            .unwrap()
3055            .unwrap()
3056            .parse()
3057            .unwrap();
3058        assert_eq!(count_after, 0, "success should reset failure count to 0");
3059
3060        // Model-specific count should also be reset
3061        let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3062        let model_count: u32 = store
3063            .kv_get(&model_kv_key)
3064            .await
3065            .unwrap()
3066            .unwrap()
3067            .parse()
3068            .unwrap();
3069        assert_eq!(
3070            model_count, 0,
3071            "success should reset model failure count to 0"
3072        );
3073    }
3074
3075    #[serial(cooldown_state)]
3076    #[tokio::test]
3077    async fn record_agent_success_without_model_skips_empty_model_key() {
3078        super::reset_global_state().await;
3079        let store = test_store().await;
3080        *cooldown_store().lock().await = Some(store.clone());
3081
3082        let agent = "test_success_without_model";
3083
3084        record_agent_success(agent, "").await;
3085
3086        let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3087        let agent_count = store
3088            .kv_get(&agent_kv_key)
3089            .await
3090            .unwrap()
3091            .expect("agent failure count should be stored");
3092        assert_eq!(agent_count, "0");
3093
3094        let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3095        let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3096        assert!(
3097            empty_model_count.is_none(),
3098            "success without model should not create an empty model failure count key"
3099        );
3100    }
3101
3102    #[serial(cooldown_state)]
3103    #[tokio::test]
3104    async fn clear_cooldown_resets_failure_counts() {
3105        super::reset_global_state().await;
3106        let store = test_store().await;
3107        *cooldown_store().lock().await = Some(store.clone());
3108
3109        let agent = "test_clear_fc_agent";
3110
3111        // Accumulate failures
3112        for _ in 0..5 {
3113            let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3114        }
3115        // Set a cooldown so clear_cooldown has something to clear
3116        set_agent_cooldown(agent, 3600).await;
3117
3118        // Verify failure count is 5
3119        let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3120        let count: u32 = store
3121            .kv_get(&kv_key)
3122            .await
3123            .unwrap()
3124            .unwrap()
3125            .parse()
3126            .unwrap();
3127        assert_eq!(count, 5);
3128
3129        // Clear the cooldown
3130        clear_cooldown(agent, &store).await;
3131
3132        // Failure count should be reset to 0
3133        let count_after: u32 = store
3134            .kv_get(&kv_key)
3135            .await
3136            .unwrap()
3137            .unwrap()
3138            .parse()
3139            .unwrap();
3140        assert_eq!(
3141            count_after, 0,
3142            "clear_cooldown should reset failure count to 0"
3143        );
3144
3145        // Cooldown should be removed
3146        assert!(
3147            !is_agent_in_cooldown(agent),
3148            "agent should not be in cooldown after clear"
3149        );
3150    }
3151
3152    #[serial(cooldown_state)]
3153    #[tokio::test]
3154    async fn clear_all_cooldowns_resets_all_failure_counts() {
3155        super::reset_global_state().await;
3156        let store = test_store().await;
3157        *cooldown_store().lock().await = Some(store.clone());
3158
3159        let agents = ["test_clear_all_a", "test_clear_all_b"];
3160        for agent in &agents {
3161            let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3162            set_agent_cooldown(agent, 3600).await;
3163        }
3164
3165        clear_cooldown("*", &store).await;
3166
3167        for agent in &agents {
3168            assert!(!is_agent_in_cooldown(agent));
3169            let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3170            let count: u32 = store
3171                .kv_get(&kv_key)
3172                .await
3173                .unwrap()
3174                .unwrap()
3175                .parse()
3176                .unwrap();
3177            assert_eq!(
3178                count, 0,
3179                "clear --all should reset failure count for {agent}"
3180            );
3181        }
3182    }
3183
3184    #[serial(cooldown_state)]
3185    #[test]
3186    fn compute_backoff_saturates_on_large_counts() {
3187        // Very large count should not overflow — extended tier applies 48h
3188        assert_eq!(
3189            compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3190            BACKOFF_MAX_SECS * 12
3191        );
3192        assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
3193        assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
3194    }
3195
3196    #[serial(cooldown_state)]
3197    #[test]
3198    fn compute_backoff_zero_base_returns_zero() {
3199        assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3200        assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3201    }
3202
3203    // ---- GitHub 5xx circuit breaker tests ----
3204
3205    #[serial(cooldown_state)]
3206    #[tokio::test]
3207    async fn github_circuit_stays_closed_below_threshold() {
3208        super::reset_global_state().await;
3209        // Reset circuit breaker state
3210        {
3211            let mut open = github_5xx_circuit_open().lock().unwrap();
3212            *open = None;
3213            let mut ts = github_5xx_timestamps().lock().unwrap();
3214            ts.clear();
3215            // Also clear any generic cooldown map entry for the github 5xx
3216            // circuit so tests are isolated from other tests that may have
3217            // set the generic cooldown via set_agent_cooldown("github:5xx", ...).
3218            let mut map = cooldown_lock();
3219            map.remove("github:5xx");
3220        }
3221
3222        // Record fewer errors than threshold
3223        for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3224            record_github_5xx().await;
3225        }
3226
3227        assert!(
3228            !is_github_circuit_open(),
3229            "circuit should stay closed below threshold"
3230        );
3231        assert_eq!(github_circuit_remaining_secs(), 0);
3232    }
3233
3234    #[serial(cooldown_state)]
3235    #[tokio::test]
3236    async fn github_circuit_trips_at_threshold() {
3237        super::reset_global_state().await;
3238        // Reset circuit breaker state
3239        {
3240            let mut open = github_5xx_circuit_open_lock();
3241            *open = None;
3242            let mut ts = github_5xx_timestamps_lock();
3243            ts.clear();
3244            let mut map = cooldown_lock();
3245            map.remove("github:5xx");
3246        }
3247
3248        // Record exactly threshold errors
3249        for _ in 0..GITHUB_5XX_THRESHOLD {
3250            record_github_5xx().await;
3251        }
3252
3253        assert!(is_github_circuit_open(), "circuit should trip at threshold");
3254        assert!(
3255            github_circuit_remaining_secs() > 0,
3256            "remaining should be > 0 when open"
3257        );
3258        assert!(
3259            github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3260            "remaining should not exceed cooldown duration"
3261        );
3262
3263        // Clean up
3264        {
3265            let mut open = github_5xx_circuit_open().lock().unwrap();
3266            *open = None;
3267        }
3268    }
3269
3270    #[serial(cooldown_state)]
3271    #[tokio::test]
3272    async fn github_circuit_cannot_double_trip() {
3273        super::reset_global_state().await;
3274        // Reset circuit breaker state
3275        {
3276            let mut open = github_5xx_circuit_open_lock();
3277            *open = None;
3278            let mut ts = github_5xx_timestamps_lock();
3279            ts.clear();
3280            let mut map = cooldown_lock();
3281            map.remove("github:5xx");
3282        }
3283
3284        // Trip the circuit
3285        for _ in 0..GITHUB_5XX_THRESHOLD {
3286            record_github_5xx().await;
3287        }
3288        assert!(is_github_circuit_open());
3289
3290        let first_remaining = github_circuit_remaining_secs();
3291
3292        // More 5xx errors should NOT reset the cooldown (never-shorten)
3293        for _ in 0..GITHUB_5XX_THRESHOLD {
3294            record_github_5xx().await;
3295        }
3296
3297        let second_remaining = github_circuit_remaining_secs();
3298        assert!(
3299            second_remaining <= first_remaining,
3300            "additional errors should not extend cooldown"
3301        );
3302
3303        // Clean up
3304        {
3305            let mut open = github_5xx_circuit_open().lock().unwrap();
3306            *open = None;
3307        }
3308    }
3309
3310    #[serial(cooldown_state)]
3311    #[tokio::test]
3312    async fn github_circuit_recovers_after_cooldown() {
3313        super::reset_global_state().await;
3314        // Reset circuit breaker state
3315        {
3316            let mut open = github_5xx_circuit_open_lock();
3317            *open = None;
3318            let mut ts = github_5xx_timestamps_lock();
3319            ts.clear();
3320            let mut map = cooldown_lock();
3321            map.remove("github:5xx");
3322        }
3323
3324        // Set the circuit to an expired cooldown
3325        {
3326            let mut open = github_5xx_circuit_open().lock().unwrap();
3327            *open = Some(chrono::Utc::now().timestamp() - 1);
3328        }
3329
3330        // Should be closed now (auto-recovery)
3331        assert!(
3332            !is_github_circuit_open(),
3333            "circuit should auto-recover after cooldown"
3334        );
3335    }
3336
3337    #[serial(cooldown_state)]
3338    #[tokio::test]
3339    async fn github_circuit_sliding_window_ages_out() {
3340        super::reset_global_state().await;
3341        // Reset circuit breaker state
3342        {
3343            let mut open = github_5xx_circuit_open_lock();
3344            *open = None;
3345            let mut ts = github_5xx_timestamps_lock();
3346            ts.clear();
3347            let mut map = cooldown_lock();
3348            map.remove("github:5xx");
3349        }
3350
3351        // Record errors that are outside the window
3352        {
3353            let mut ts = github_5xx_timestamps().lock().unwrap();
3354            let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3355            for _ in 0..GITHUB_5XX_THRESHOLD {
3356                ts.push(old);
3357            }
3358        }
3359
3360        // A new error should clear old ones and not trip the circuit
3361        record_github_5xx().await;
3362
3363        assert!(
3364            !is_github_circuit_open(),
3365            "old errors outside window should not trip circuit"
3366        );
3367    }
3368}