Skip to main content

orch/engine/
cooldown.rs

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