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