1use chrono::TimeZone;
19use once_cell::sync::Lazy;
20use regex::Regex;
21use std::collections::HashMap;
22use std::sync::{Arc, Mutex, OnceLock};
23
24pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
30
31pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
33
34pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
36
37pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
39
40pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
42
43pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
45
46pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
48
49pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
52
53pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
61
62pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum CreditExhaustionReason {
72 OutOfCredits,
74 OrgLevelDisabled,
76 BillingCycleExhausted,
78}
79
80impl CreditExhaustionReason {
81 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
93pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
100
101pub const SILENCE_COUNT_THRESHOLD: usize = 5;
103
104pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
106
107pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
109
110const KV_PREFIX: &str = "cooldown:";
112
113const SILENCE_COUNT_PREFIX: &str = "silence_count:";
115
116const FAILURE_COUNT_PREFIX: &str = "failure_count:";
118
119const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
122
123const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
127
128struct CooldownEntry {
129 cooldown_until: i64,
131 reason: String,
132}
133
134pub struct SilenceCountResult {
135 pub count: usize,
136 pub extended_cooldown_applied: bool,
137}
138
139fn cooldowns() -> &'static Mutex<HashMap<String, CooldownEntry>> {
148 static COOLDOWNS: OnceLock<Mutex<HashMap<String, CooldownEntry>>> = OnceLock::new();
149 COOLDOWNS.get_or_init(|| Mutex::new(HashMap::new()))
150}
151
152fn cooldown_store() -> &'static tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>> {
158 static STORE: OnceLock<tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>>> =
159 OnceLock::new();
160 STORE.get_or_init(|| tokio::sync::Mutex::new(None))
161}
162
163fn parse_cooldown_value(value: &str) -> Option<(i64, String)> {
171 if let Ok(serde_json::Value::Object(map)) = serde_json::from_str(value) {
172 let until = map.get("until")?.as_i64()?;
173 let reason = map
174 .get("reason")
175 .and_then(|r| r.as_str())
176 .unwrap_or("persisted")
177 .to_string();
178 return Some((until, reason));
179 }
180 value
181 .parse::<i64>()
182 .ok()
183 .map(|until| (until, "persisted".to_string()))
184}
185
186pub async fn init_cooldown_store(store: Arc<crate::store::TaskStore>) {
187 match store.kv_list_prefix(KV_PREFIX).await {
188 Ok(rows) => {
189 let now = chrono::Utc::now().timestamp();
190 let mut map = cooldown_lock();
191 let mut loaded = 0;
192 for (key, value) in rows {
193 let cooldown_key = key.trim_start_matches(KV_PREFIX);
194 if let Some((cooldown_until, reason)) = parse_cooldown_value(&value) {
195 if now < cooldown_until {
196 map.insert(
197 cooldown_key.to_string(),
198 CooldownEntry {
199 cooldown_until,
200 reason,
201 },
202 );
203 loaded += 1;
204 }
205 }
208 }
209 if loaded > 0 {
210 tracing::info!(loaded, "loaded persisted cooldowns from KV store");
211 }
212 }
213 Err(e) => {
214 tracing::warn!(err = %e, "failed to load cooldowns from KV store");
215 }
216 }
217
218 if let Ok(Some(raw)) = store.kv_get("cooldown:github:5xx").await {
220 if let Ok(cooldown_until) = raw.parse::<i64>() {
221 let now = chrono::Utc::now().timestamp();
222 if now < cooldown_until {
223 let mut open = github_5xx_circuit_open_lock();
224 *open = Some(cooldown_until);
225 let remaining = cooldown_until - now;
226 tracing::warn!(
227 remaining_secs = remaining,
228 "GitHub 5xx circuit breaker restored from KV (still open)"
229 );
230 }
231 }
232 }
233
234 *cooldown_store().lock().await = Some(store);
235}
236
237pub async fn sync_from_kv(store: &Arc<crate::store::TaskStore>) {
254 let rows = match store.kv_list_prefix(KV_PREFIX).await {
255 Ok(rows) => rows,
256 Err(e) => {
257 tracing::warn!(err = %e, "cooldown sync_from_kv: KV list failed");
258 return;
259 }
260 };
261 let now = chrono::Utc::now().timestamp();
262
263 let mut explicit_kv: HashMap<String, Option<(i64, String)>> =
266 HashMap::with_capacity(rows.len());
267 for (kv_key, value) in &rows {
268 let cooldown_key = kv_key.trim_start_matches(KV_PREFIX);
269 if let Some((ts, reason)) = parse_cooldown_value(value) {
270 let normalized = if ts > now { Some((ts, reason)) } else { None };
271 explicit_kv.insert(cooldown_key.to_string(), normalized);
272 }
273 }
274
275 let mut map = cooldown_lock();
276 let mut dropped = 0usize;
277 let mut updated = 0usize;
278
279 map.retain(|key, _| {
281 if matches!(explicit_kv.get(key), Some(None)) {
282 dropped += 1;
283 false
284 } else {
285 true
286 }
287 });
288
289 for (key, normalized) in explicit_kv {
291 if let Some((ts, reason)) = normalized {
292 match map.get_mut(&key) {
293 Some(entry) => {
294 if entry.cooldown_until != ts {
295 entry.cooldown_until = ts;
296 updated += 1;
297 }
298 }
299 None => {
300 map.insert(
301 key,
302 CooldownEntry {
303 cooldown_until: ts,
304 reason,
305 },
306 );
307 updated += 1;
308 }
309 }
310 }
311 }
312
313 if dropped > 0 || updated > 0 {
314 tracing::info!(
315 dropped,
316 updated,
317 "cooldown sync_from_kv: reconciled in-memory map with KV"
318 );
319 }
320}
321
322const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
326
327const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
329
330const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
334
335const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
337
338pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
357 if count <= 1 {
358 return base;
359 }
360 let exp = count.saturating_sub(1).min(u32::MAX as u64) as u32;
361 let factor = 3_i64.saturating_pow(exp);
362 let standard = base.saturating_mul(factor).min(max);
363
364 let first_max = if base <= 0 {
373 0
374 } else {
375 let mut accumulated = base;
377 let mut first_max_count = 1u64;
378 while accumulated < max && first_max_count < count {
379 accumulated = accumulated.saturating_mul(3);
380 first_max_count += 1;
381 }
382 first_max_count
383 };
384
385 let times_at_max = count.saturating_sub(first_max);
387
388 if !allow_extended {
389 return standard;
390 }
391
392 if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
393 max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
394 } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
395 max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
396 } else {
397 standard
398 }
399}
400
401async fn read_and_increment_failure_count(
410 store_opt: &Option<Arc<crate::store::TaskStore>>,
411 key: &str,
412) -> u64 {
413 read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
414}
415
416async fn read_and_increment_failure_count_with_prefix(
421 store_opt: &Option<Arc<crate::store::TaskStore>>,
422 prefix: &str,
423 key: &str,
424) -> u64 {
425 let kv_key = format!("{prefix}{key}");
426 if let Some(store) = store_opt {
427 match store.kv_increment(&kv_key).await {
428 Ok(n) => n,
429 Err(e) => {
430 tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
431 u64::MAX
432 }
433 }
434 } else {
435 1
436 }
437}
438
439pub async fn record_agent_success(agent_name: &str, model: &str) {
445 let store_opt = cooldown_store().lock().await.clone();
446 if let Some(store) = store_opt {
447 let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
448 let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
449 if let Err(e) = store.kv_set(&agent_key, "0").await {
450 tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
451 }
452 if !model.is_empty() {
453 let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
454 if let Err(e) = store.kv_set(&model_key, "0").await {
455 tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
456 }
457 }
458 if let Err(e) = store.kv_set(&credit_key, "0").await {
459 tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
460 }
461 }
462}
463
464pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
470 if let Some(cooldown_until) = parse_retry_at(error_message) {
475 let now = chrono::Utc::now().timestamp();
476 if cooldown_until > now {
477 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
478 return;
479 }
480 tracing::debug!(
481 agent = agent_name,
482 cooldown_until,
483 "vendor retry_at is in the past — using backoff instead"
484 );
485 }
486
487 let store_opt = cooldown_store().lock().await.clone();
488 let count = read_and_increment_failure_count(&store_opt, agent_name).await;
489 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
490 let cooldown_until =
491 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
492 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
493}
494
495pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
504 detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
505}
506
507pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
512 let lower = error_message.to_lowercase();
513
514 let overage_patterns = [
515 "overagedisabledreason",
516 "overage_disabled_reason",
517 "overage disabled reason",
518 ];
519
520 if overage_patterns.iter().any(|p| lower.contains(p)) {
521 if lower.contains("out_of_credits") || lower.contains("outofcredits") {
522 return Some(CreditExhaustionReason::OutOfCredits);
523 }
524 if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
525 return Some(CreditExhaustionReason::OrgLevelDisabled);
526 }
527 }
528
529 let billing_patterns = [
530 "credit balance too low",
531 "credit balance insufficient",
532 "insufficient credit balance",
533 "billing quota exceeded",
534 "billing limit exceeded",
535 "organization has been disabled",
536 "org disabled",
537 ];
538
539 if billing_patterns.iter().any(|p| lower.contains(p)) {
540 return Some(CreditExhaustionReason::OrgLevelDisabled);
541 }
542
543 let billing_cycle_patterns = [
546 "billing cycle",
547 "quota refreshed next cycle",
548 "monthly quota",
549 "quota will be refreshed",
550 "next billing cycle",
551 "refreshed in the next cycle",
552 "weekly/monthly limit exhausted",
553 "limit will reset at",
554 "upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
559
560 if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
561 return Some(CreditExhaustionReason::BillingCycleExhausted);
562 }
563
564 if lower.contains("out of credits")
565 || lower.contains("outofcredits")
566 || lower.contains("insufficient funds")
567 || lower.contains("no credits remaining")
568 || lower.contains("insufficient_quota")
569 || lower.contains("quota exceeded")
570 || lower.contains("insufficient balance")
571 || lower.contains("please recharge")
572 || lower.contains("no resource package")
573 || lower.contains("balance exhausted")
574 || lower.contains("402 payment required")
575 {
576 return Some(CreditExhaustionReason::OutOfCredits);
577 }
578
579 None
580}
581
582pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
595 let reason_str = match reason {
596 CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
597 CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
598 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
599 };
600
601 let store_opt = cooldown_store().lock().await.clone();
602 let count = read_and_increment_failure_count_with_prefix(
607 &store_opt,
608 CREDIT_FAILURE_COUNT_PREFIX,
609 agent_name,
610 )
611 .await;
612
613 let (base, max) = match reason {
614 CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
615 CreditExhaustionReason::OrgLevelDisabled => {
616 (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
617 }
618 CreditExhaustionReason::BillingCycleExhausted => {
619 (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
620 }
621 };
622 let cooldown_secs = compute_backoff(
623 count,
624 base,
625 max,
626 !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
627 );
628 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
629 set_cooldown_async(agent_name, cooldown_until, reason_str).await;
630 tracing::warn!(
631 agent = agent_name,
632 reason = reason_str,
633 attempt = count,
634 cooldown_secs,
635 "credit exhaustion detected: applying exponential agent-wide cooldown"
636 );
637}
638
639pub async fn record_model_failure(agent_name: &str, model: &str) {
643 let key = format!("{agent_name}:{model}");
644 let store_opt = cooldown_store().lock().await.clone();
645 let count = read_and_increment_failure_count(&store_opt, &key).await;
646 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
647 let cooldown_until =
648 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
649 set_cooldown_async(&key, cooldown_until, "model_error").await;
650}
651
652pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
660 if let Some(window_secs) = parse_relative_usage_window(error_message) {
661 if let Some(m) = model {
662 let cooldown_until = chrono::Utc::now().timestamp() + window_secs as i64;
663 set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
664 tracing::info!(
665 agent = agent_name,
666 model = m,
667 window_secs,
668 cooldown_until,
669 "relative usage window detected: applying model-specific cooldown"
670 );
671 return;
672 }
673 }
674
675 record_agent_failure_with_message(agent_name, error_message).await;
676 if let Some(m) = model {
677 record_model_failure(agent_name, m).await;
678 }
679}
680
681pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
693 let key = format!("{agent_name}:{model}");
694 let store_opt = cooldown_store().lock().await.clone();
695 let count = read_and_increment_failure_count_with_prefix(
696 &store_opt,
697 PERSISTENT_FAILURE_COUNT_PREFIX,
698 &key,
699 )
700 .await;
701 let cooldown_secs = compute_backoff(
702 count,
703 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
704 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
705 false,
706 );
707 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
708 set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
709}
710
711pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
718 let key = format!("{agent_name}:{model}");
719 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
720 set_cooldown_async(&key, cooldown_until, reason).await;
721}
722
723pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
725 let key = format!("{agent_name}:{model}");
726 set_cooldown_async(&key, until, "vendor_reset_at").await;
727}
728
729pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
736 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
737 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
738}
739
740pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
743 let store_opt = cooldown_store().lock().await.clone();
744 let store = match store_opt {
745 Some(store) => store,
746 None => {
747 tracing::debug!(
748 agent = agent_name,
749 model,
750 "skipping silence count record (no KV store)"
751 );
752 return None;
753 }
754 };
755
756 let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
757 let now = chrono::Utc::now().timestamp();
758 let window_start = now - SILENCE_COUNT_WINDOW_SECS;
759
760 let mut timestamps = match store.kv_get(&key).await {
761 Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
762 Ok(ts) => ts,
763 Err(e) => {
764 tracing::warn!(
765 kv_key = %key,
766 err = %e,
767 "failed to parse silence timestamps from KV — resetting window"
768 );
769 Vec::new()
770 }
771 },
772 Ok(None) => Vec::new(),
773 Err(err) => {
774 tracing::warn!(
775 kv_key = %key,
776 err = %err,
777 "failed to load silence count from KV"
778 );
779 Vec::new()
780 }
781 };
782
783 timestamps.retain(|ts| *ts >= window_start);
784 timestamps.push(now);
785
786 let count = timestamps.len();
787 let mut extended_cooldown_applied = false;
788 if count >= SILENCE_COUNT_THRESHOLD {
789 set_model_cooldown(
790 agent_name,
791 model,
792 SILENCE_EXTENDED_COOLDOWN_SECS,
793 "silence_detected",
794 )
795 .await;
796 extended_cooldown_applied = true;
797 timestamps.clear();
798 }
799
800 let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
801 if let Err(err) = store.kv_set(&key, &value).await {
802 tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
803 }
804
805 Some(SilenceCountResult {
806 count,
807 extended_cooldown_applied,
808 })
809}
810
811pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
813 let key = format!("{agent_name}:{model}");
814 is_in_cooldown(&key)
815}
816
817pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
819 is_in_cooldown(agent_name)
820}
821
822pub fn cooldown_until(key: &str) -> Option<i64> {
824 let map = cooldown_lock();
825 map.get(key).and_then(|entry| {
826 let now = chrono::Utc::now().timestamp();
827 if now < entry.cooldown_until {
828 Some(entry.cooldown_until)
829 } else {
830 None
831 }
832 })
833}
834
835pub fn cooldown_reason(key: &str) -> Option<String> {
837 let map = cooldown_lock();
838 map.get(key).and_then(|entry| {
839 let now = chrono::Utc::now().timestamp();
840 if now < entry.cooldown_until {
841 Some(entry.reason.clone())
842 } else {
843 None
844 }
845 })
846}
847
848pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
853 let now = chrono::Utc::now().timestamp();
854 let map = cooldown_lock();
855 let mut result: Vec<(String, i64, String)> = map
856 .iter()
857 .filter(|(_, entry)| now < entry.cooldown_until)
858 .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
859 .collect();
860 result.sort_by(|a, b| a.0.cmp(&b.0));
861 result
862}
863
864pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
870 if key == "*" {
871 let keys: Vec<String> = {
872 let mut map = cooldown_lock();
873 let keys: Vec<String> = map.keys().cloned().collect();
874 map.clear();
875 keys
876 };
877 for k in &keys {
878 let kv_key = format!("{KV_PREFIX}{k}");
879 if let Err(e) = store.kv_set(&kv_key, "0").await {
880 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
881 }
882 let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
884 if let Err(e) = store.kv_set(&fc_key, "0").await {
885 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
886 }
887 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
889 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
890 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
891 }
892 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
894 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
895 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
896 }
897 }
898 if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
902 for (fc_key, _) in fc_entries {
903 if let Err(e) = store.kv_set(&fc_key, "0").await {
904 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
905 }
906 }
907 }
908 if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
909 for (cfc_key, _) in cfc_entries {
910 if let Err(e) = store.kv_set(&cfc_key, "0").await {
911 tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
912 }
913 }
914 }
915 if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
916 for (pfc_key, _) in pfc_entries {
917 if let Err(e) = store.kv_set(&pfc_key, "0").await {
918 tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
919 }
920 }
921 }
922 tracing::info!(
923 count = keys.len(),
924 "cleared all cooldowns and failure counts"
925 );
926 } else {
927 {
928 let mut map = cooldown_lock();
929 map.remove(key);
930 }
931 let kv_key = format!("{KV_PREFIX}{key}");
932 if let Err(e) = store.kv_set(&kv_key, "0").await {
933 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
934 }
935 let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
937 if let Err(e) = store.kv_set(&fc_key, "0").await {
938 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
939 }
940 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
942 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
943 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
944 }
945 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
947 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
948 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
949 }
950 tracing::info!(key, "cleared cooldown and failure count");
951 }
952}
953
954fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
963 let mut map = cooldown_lock();
964 if let Some(existing) = map.get(key) {
965 if existing.cooldown_until >= cooldown_until {
969 return false;
970 }
971 }
972 map.insert(
973 key.to_string(),
974 CooldownEntry {
975 cooldown_until,
976 reason: reason.to_string(),
977 },
978 );
979 true
980}
981
982async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
992 if !set_cooldown_in_memory(key, cooldown_until, reason) {
993 return false;
994 }
995 let store_opt = cooldown_store().lock().await.clone();
996 if let Some(store) = store_opt {
997 let kv_key = format!("{KV_PREFIX}{key}");
998 let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
999 if let Err(e) = store.kv_set(&kv_key, &value).await {
1000 tracing::warn!(
1001 kv_key,
1002 err = %e,
1003 "failed to persist cooldown to KV store — rolling back in-memory entry"
1004 );
1005 let mut map = cooldown_lock();
1006 if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1010 map.remove(key);
1011 }
1012 return false;
1013 }
1014 }
1015 true
1016}
1017
1018fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
1020 cooldowns().lock().unwrap_or_else(|e| {
1021 tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1022 e.into_inner()
1023 })
1024}
1025
1026fn degraded_agents_lock() -> std::sync::MutexGuard<'static, std::collections::HashSet<String>> {
1028 degraded_agents().lock().unwrap_or_else(|e| {
1029 tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1030 e.into_inner()
1031 })
1032}
1033
1034fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
1036 github_5xx_timestamps().lock().unwrap_or_else(|e| {
1037 tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1038 e.into_inner()
1039 })
1040}
1041
1042fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
1044 github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1045 tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1046 e.into_inner()
1047 })
1048}
1049
1050fn is_in_cooldown(key: &str) -> bool {
1052 let map = cooldown_lock();
1053 if let Some(entry) = map.get(key) {
1054 let now = chrono::Utc::now().timestamp();
1055 return now < entry.cooldown_until;
1056 }
1057 false
1058}
1059
1060pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1066
1067pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1071
1072fn degraded_agents() -> &'static Mutex<std::collections::HashSet<String>> {
1074 static DEGRADED: OnceLock<Mutex<std::collections::HashSet<String>>> = OnceLock::new();
1075 DEGRADED.get_or_init(|| Mutex::new(std::collections::HashSet::new()))
1076}
1077
1078pub fn is_agent_degraded(agent: &str) -> bool {
1084 let set = degraded_agents_lock();
1085 set.contains(agent)
1086}
1087
1088pub fn mark_agent_degraded(agent: &str) {
1090 let mut set = degraded_agents_lock();
1091 set.insert(agent.to_string());
1092}
1093
1094pub fn clear_agent_degraded(agent: &str) {
1096 let mut set = degraded_agents_lock();
1097 set.remove(agent);
1098}
1099
1100pub async fn refresh_degraded_agents(
1113 store: &Arc<crate::store::TaskStore>,
1114 available_agents: &[String],
1115 model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1116 window_hours: u32,
1117 threshold: i64,
1118) {
1119 clear_expired_cooldowns();
1123
1124 let start = chrono::Utc::now();
1127 let counts = match store.recent_rate_limit_counts(window_hours).await {
1128 Ok(c) => {
1129 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1130 tracing::info!(
1131 duration_ms = dur_ms,
1132 window_hours,
1133 "recent_rate_limit_counts query completed"
1134 );
1135 c
1136 }
1137 Err(e) => {
1138 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1139 tracing::warn!(
1140 err = %e,
1141 duration_ms = dur_ms,
1142 "failed to query recent rate limit counts for health check"
1143 );
1144 return;
1145 }
1146 };
1147
1148 for agent in available_agents {
1149 let in_cooldown = is_agent_in_cooldown(agent);
1150 let no_models = !model_checker(agent);
1151 let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1152 let over_threshold = rate_limit_count >= threshold;
1153
1154 let degraded = in_cooldown || no_models || over_threshold;
1155
1156 if degraded {
1157 if !is_agent_degraded(agent) {
1158 let reason = if in_cooldown {
1159 "agent in cooldown"
1160 } else if no_models {
1161 "all models cooled"
1162 } else {
1163 "rate limit threshold exceeded"
1164 };
1165 tracing::warn!(
1166 agent,
1167 reason,
1168 rate_limit_count,
1169 window_hours,
1170 threshold,
1171 "pre-emptive health check: marking agent as degraded"
1172 );
1173 }
1174 mark_agent_degraded(agent);
1175 } else {
1176 if is_agent_degraded(agent) {
1177 tracing::info!(
1178 agent,
1179 rate_limit_count,
1180 "pre-emptive health check: agent recovered, clearing degraded flag"
1181 );
1182 }
1183 clear_agent_degraded(agent);
1184 }
1185 }
1186}
1187
1188pub fn clear_expired_cooldowns() {
1190 let mut map = cooldown_lock();
1191 let now = chrono::Utc::now().timestamp();
1192 map.retain(|_key, entry| now < entry.cooldown_until);
1193}
1194
1195const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1201
1202const GITHUB_5XX_THRESHOLD: usize = 5;
1204
1205const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1207
1208fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1210 static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1211 TS.get_or_init(|| Mutex::new(Vec::new()))
1212}
1213
1214fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1216 static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1217 OPEN.get_or_init(|| Mutex::new(None))
1218}
1219
1220pub async fn record_github_5xx() {
1225 let now = chrono::Utc::now().timestamp();
1226 let window_start = now - GITHUB_5XX_WINDOW_SECS;
1227
1228 let should_trip = {
1229 let mut ts = github_5xx_timestamps_lock();
1230 ts.retain(|t| *t >= window_start);
1231 ts.push(now);
1232 ts.len() >= GITHUB_5XX_THRESHOLD
1233 };
1234
1235 if should_trip {
1236 let cooldown_until = {
1237 let mut open = github_5xx_circuit_open_lock();
1238 if open.is_none() {
1239 let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1240 *open = Some(cd);
1241 tracing::warn!(
1242 errors_in_window = GITHUB_5XX_THRESHOLD,
1243 window_secs = GITHUB_5XX_WINDOW_SECS,
1244 cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1245 "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1246 );
1247 Some(cd)
1248 } else {
1249 None
1250 }
1251 };
1252 if let Some(cd) = cooldown_until {
1254 let store_opt = cooldown_store().lock().await.clone();
1255 if let Some(store) = store_opt {
1256 let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1257 }
1258 set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1262 }
1263 let mut ts = github_5xx_timestamps_lock();
1265 ts.clear();
1266 }
1267}
1268
1269pub fn is_github_circuit_open() -> bool {
1277 {
1280 let mut open = github_5xx_circuit_open_lock();
1281 if let Some(until) = *open {
1282 let now = chrono::Utc::now().timestamp();
1283 if now >= until {
1284 *open = None;
1285 tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1286 return false;
1287 }
1288 return true;
1289 }
1290 }
1291
1292 if let Some(until) = cooldown_until("github:5xx") {
1296 let mut open = github_5xx_circuit_open_lock();
1299 if open.is_none() {
1300 *open = Some(until);
1301 let now = chrono::Utc::now().timestamp();
1302 if now < until {
1303 let remaining = until - now;
1304 tracing::warn!(
1305 remaining_secs = remaining,
1306 "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1307 );
1308 }
1309 }
1310 return true;
1311 }
1312
1313 false
1314}
1315
1316pub fn github_circuit_remaining_secs() -> u64 {
1319 {
1321 let open = github_5xx_circuit_open_lock();
1322 if let Some(until) = *open {
1323 let now = chrono::Utc::now().timestamp();
1324 if now < until {
1325 return (until - now) as u64;
1326 }
1327 return 0;
1328 }
1329 }
1330
1331 if let Some(until) = cooldown_until("github:5xx") {
1333 let now = chrono::Utc::now().timestamp();
1334 if now < until {
1335 return (until - now) as u64;
1336 }
1337 }
1338 0
1339}
1340
1341pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1345 if error_message.is_empty() {
1346 return None;
1347 }
1348
1349 let lower = error_message.to_lowercase();
1350 for retry_marker in [
1354 "try again at ",
1355 "reset at ",
1356 "resets at ",
1357 "resets ",
1358 "reset ",
1359 ] {
1360 if let Some(idx) = lower.find(retry_marker) {
1361 let date_str = &lower[idx + retry_marker.len()..];
1362 let date_str = date_str
1364 .split(['.', '\n'])
1365 .next()
1366 .unwrap_or(date_str)
1367 .trim();
1368
1369 let cleaned = date_str
1372 .replace("st,", ",")
1373 .replace("nd,", ",")
1374 .replace("rd,", ",")
1375 .replace("th,", ",");
1376
1377 for fmt in &[
1378 "%b %d, %Y %I:%M %p", "%b %d, %Y %I:%M%p", "%B %d, %Y %I:%M %p", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", ] {
1385 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1386 use chrono::TimeZone;
1389 let utc = match chrono::Local.from_local_datetime(&dt) {
1390 chrono::LocalResult::Single(local_dt) => {
1391 local_dt.with_timezone(&chrono::Utc)
1392 }
1393 chrono::LocalResult::Ambiguous(earliest, _) => {
1394 earliest.with_timezone(&chrono::Utc)
1396 }
1397 chrono::LocalResult::None => {
1398 dt.and_utc()
1400 }
1401 };
1402 tracing::info!(
1403 retry_at = %utc,
1404 raw = date_str,
1405 marker = retry_marker,
1406 "parsed rate limit retry-at date"
1407 );
1408 return Some(utc.timestamp());
1409 }
1410 }
1411 tracing::debug!(
1412 raw = date_str,
1413 marker = retry_marker,
1414 "could not parse retry-at date"
1415 );
1416
1417 let time_only = date_str
1422 .split_once('(')
1423 .map(|(t, _)| t)
1424 .unwrap_or(date_str)
1425 .trim();
1426 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1428 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1429 let last = &compact[split..];
1430 if !last.contains(':')
1431 && (last.ends_with("am") || last.ends_with("pm"))
1432 && last.len() > 2
1433 {
1434 let at = split + last.len() - 2;
1435 compact.insert_str(at, ":00");
1436 }
1437 let md = compact.replace(" at ", " ");
1439 let now_local = chrono::Local::now();
1440 let this_year = chrono::Datelike::year(&now_local);
1441 for year in [this_year, this_year + 1] {
1442 let with_year = format!("{year} {md}");
1443 if let Ok(dt) =
1444 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1445 {
1446 use chrono::TimeZone;
1447 if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
1448 if local_dt > now_local {
1449 return Some(local_dt.timestamp());
1450 }
1451 }
1452 }
1453 }
1454 let time_only = compact.as_str();
1455 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1456 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1457 let local_now = chrono::Local::now();
1460 let today = local_now.date_naive();
1461 let candidate = chrono::NaiveDateTime::new(today, t);
1462 let candidate_local: chrono::DateTime<chrono::Local> =
1465 match chrono::Local.from_local_datetime(&candidate) {
1466 chrono::LocalResult::Single(dt) => dt,
1467 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1468 chrono::LocalResult::None => {
1469 let utc_dt = candidate.and_utc();
1471 return Some(utc_dt.timestamp());
1472 }
1473 };
1474 let chosen = if candidate_local <= local_now {
1475 let tomorrow = today.succ_opt().unwrap_or(today);
1477 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1478 match chrono::Local.from_local_datetime(&ndt) {
1479 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1480 chrono::LocalResult::Ambiguous(earliest, _) => {
1481 earliest.with_timezone(&chrono::Utc)
1482 }
1483 chrono::LocalResult::None => ndt.and_utc(),
1484 }
1485 } else {
1486 candidate_local.with_timezone(&chrono::Utc)
1487 };
1488 tracing::info!(
1489 retry_at = %chosen,
1490 raw = date_str,
1491 marker = retry_marker,
1492 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1493 );
1494 return Some(chosen.timestamp());
1495 }
1496 }
1497 }
1498 }
1499
1500 None
1504}
1505
1506pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1517 if error_message.is_empty() {
1518 return None;
1519 }
1520
1521 let lower = error_message.to_lowercase();
1522 let has_context = lower.contains("usage limit")
1523 || lower.contains("quota")
1524 || lower.contains("window")
1525 || lower.contains("limit");
1526 if !has_context {
1527 return None;
1528 }
1529
1530 static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
1531 Regex::new(
1532 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1533 )
1534 .expect("usage window regex is valid")
1535 });
1536
1537 let mut max_secs: Option<u64> = None;
1538 for caps in WINDOW_RE.captures_iter(&lower) {
1539 let secs = if let Some(num) = caps.get(1) {
1540 let n: u64 = num.as_str().parse().ok()?;
1541 let unit = caps.get(2)?.as_str().to_lowercase();
1542 let multiplier = match unit.as_str() {
1543 "hour" | "hours" | "hr" | "hrs" => 3600,
1544 "minute" | "minutes" | "min" | "mins" => 60,
1545 "day" | "days" => 86400,
1546 _ => return None,
1547 };
1548 n.checked_mul(multiplier)?
1549 } else if let Some(word) = caps.get(3) {
1550 match word.as_str().to_lowercase().as_str() {
1551 "hourly" => 3600,
1552 "daily" => 86400,
1553 "weekly" => 604800,
1554 "monthly" => 2592000,
1555 _ => return None,
1556 }
1557 } else {
1558 continue;
1559 };
1560 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1561 }
1562
1563 max_secs
1564}
1565
1566#[cfg(test)]
1567pub(crate) async fn reset_global_state() {
1568 cooldowns().lock().unwrap().clear();
1569 *cooldown_store().lock().await = None;
1570 github_5xx_timestamps().lock().unwrap().clear();
1571 *github_5xx_circuit_open().lock().unwrap() = None;
1572 degraded_agents().lock().unwrap().clear();
1573 crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1574}
1575
1576#[cfg(test)]
1577mod tests {
1578 use super::*;
1579 use serial_test::serial;
1580
1581 async fn test_store() -> Arc<crate::store::TaskStore> {
1582 Arc::new(
1583 crate::store::TaskStore::open_memory()
1584 .await
1585 .expect("in-memory store"),
1586 )
1587 }
1588
1589 #[serial(cooldown_state)]
1590 #[tokio::test]
1591 async fn init_loads_unexpired_cooldowns_from_kv() {
1592 super::reset_global_state().await;
1593 let store = test_store().await;
1594 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1597 .kv_set(
1598 &format!("{KV_PREFIX}kimi:k2p5"),
1599 &cooldown_until.to_string(),
1600 )
1601 .await
1602 .unwrap();
1603
1604 let expired = chrono::Utc::now().timestamp() - 1;
1606 store
1607 .kv_set(
1608 &format!("{KV_PREFIX}opencode:old-free"),
1609 &expired.to_string(),
1610 )
1611 .await
1612 .unwrap();
1613
1614 {
1615 let mut map = cooldown_lock();
1616 map.clear();
1617 }
1618
1619 init_cooldown_store(store).await;
1620
1621 assert!(
1622 is_model_in_cooldown("kimi", "k2p5"),
1623 "unexpired cooldown should be loaded"
1624 );
1625 assert!(
1626 !is_model_in_cooldown("opencode", "old-free"),
1627 "expired cooldown should not be loaded"
1628 );
1629 }
1630
1631 #[serial(cooldown_state)]
1632 #[tokio::test]
1633 async fn record_model_failure_writes_to_kv() {
1634 super::reset_global_state().await;
1635 let store = test_store().await;
1636
1637 *cooldown_store().lock().await = Some(store.clone());
1638
1639 record_model_failure("testagent_persist", "testmodel_persist").await;
1640
1641 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1643
1644 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1645 let stored = store.kv_get(&kv_key).await.unwrap();
1646 assert!(stored.is_some(), "model failure should be persisted to KV");
1647 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1648 let now = chrono::Utc::now().timestamp();
1649 assert!(ts > now, "persisted timestamp should be in the future");
1651 assert_eq!(reason, "model_error");
1652 }
1653
1654 #[serial(cooldown_state)]
1655 #[tokio::test]
1656 async fn cooldown_reason_survives_restart_via_kv() {
1657 super::reset_global_state().await;
1658 let store = test_store().await;
1659 let agent = "test_reason_restart";
1660 let future = chrono::Utc::now().timestamp() + 600;
1661
1662 init_cooldown_store(store.clone()).await;
1663 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1664
1665 {
1667 let mut map = cooldown_lock();
1668 map.clear();
1669 }
1670 init_cooldown_store(store).await;
1671
1672 let entries = list_all_cooldowns();
1673 let entry = entries.iter().find(|(key, ..)| key == agent);
1674 assert_eq!(
1675 entry.map(|(_, _, reason)| reason.as_str()),
1676 Some("billing_cycle_exhausted"),
1677 "cooldown reason must survive a restart, not collapse to 'persisted'"
1678 );
1679 }
1680
1681 #[serial(cooldown_state)]
1682 #[tokio::test]
1683 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1684 super::reset_global_state().await;
1685 let store = test_store().await;
1686 *cooldown_store().lock().await = Some(store.clone());
1687
1688 let agent = "testagent_persistent";
1689 let model = "testmodel_persistent";
1690 let key = format!("{agent}:{model}");
1691
1692 record_persistent_model_failure(agent, model).await;
1694 let now = chrono::Utc::now().timestamp();
1695 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1696 let first_remaining = first_until.saturating_sub(now);
1697 assert!(
1698 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1699 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1700 );
1701
1702 record_persistent_model_failure(agent, model).await;
1704 let now2 = chrono::Utc::now().timestamp();
1705 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1706 let second_remaining = second_until.saturating_sub(now2);
1707 assert!(
1708 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1709 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1710 );
1711
1712 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1713 let stored_count = store
1714 .kv_get(&count_key)
1715 .await
1716 .expect("kv read should succeed")
1717 .expect("failure count should be written");
1718 assert_eq!(stored_count, "2");
1719 }
1720
1721 #[serial(cooldown_state)]
1722 #[tokio::test]
1723 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1724 super::reset_global_state().await;
1725 let store = test_store().await;
1726 *cooldown_store().lock().await = Some(store.clone());
1727
1728 let agent = "testagent_persistent_success";
1729 let model = "testmodel_persistent_success";
1730 let key = format!("{agent}:{model}");
1731
1732 record_persistent_model_failure(agent, model).await;
1734 record_persistent_model_failure(agent, model).await;
1735
1736 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1737 let stored_count = store
1738 .kv_get(&count_key)
1739 .await
1740 .expect("kv read should succeed")
1741 .expect("failure count should be written");
1742 assert_eq!(stored_count, "2");
1743
1744 record_agent_success(agent, model).await;
1748
1749 let stored_count_after_success = store
1750 .kv_get(&count_key)
1751 .await
1752 .expect("kv read should succeed")
1753 .expect("persistent failure count should still exist");
1754 assert_eq!(
1755 stored_count_after_success, "2",
1756 "persistent failure count must survive an intervening success"
1757 );
1758
1759 record_persistent_model_failure(agent, model).await;
1761 let now = chrono::Utc::now().timestamp();
1762 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1763 let third_remaining = third_until.saturating_sub(now);
1764 assert!(
1765 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1766 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1767 );
1768 }
1769
1770 #[serial(cooldown_state)]
1771 #[tokio::test]
1772 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1773 super::reset_global_state().await;
1774 let store = test_store().await;
1775 *cooldown_store().lock().await = Some(store.clone());
1776
1777 let agent = "testagent_persistent_capped";
1778 let model = "testmodel_persistent_capped";
1779 let key = format!("{agent}:{model}");
1780
1781 for _ in 0..9 {
1782 record_persistent_model_failure(agent, model).await;
1783 }
1784
1785 let remaining = cooldown_until(&key)
1786 .expect("cooldown should exist")
1787 .saturating_sub(chrono::Utc::now().timestamp());
1788 assert!(
1789 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1790 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1791 );
1792 assert!(
1793 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1794 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1795 );
1796 }
1797
1798 #[serial(cooldown_state)]
1799 #[test]
1800 fn parse_retry_at_codex_format() {
1801 let msg = "You've hit your usage limit. Upgrade to Pro, visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at Mar 26th, 2026 5:55 AM.";
1802 let ts = parse_retry_at(msg);
1803 assert!(ts.is_some(), "should parse codex retry-at date");
1804
1805 use chrono::TimeZone;
1808 let naive = chrono::NaiveDateTime::new(
1809 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1810 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1811 );
1812 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1813 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1814 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1815 chrono::LocalResult::None => naive.and_utc(),
1816 };
1817 assert_eq!(
1818 ts.unwrap(),
1819 expected_utc.timestamp(),
1820 "timestamp should reflect local-timezone interpretation of the vendor date"
1821 );
1822
1823 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1825 use chrono::Datelike;
1826 assert_eq!(dt.month(), 3);
1827 assert_eq!(dt.day(), 26);
1828 }
1829
1830 #[serial(cooldown_state)]
1831 #[test]
1832 fn parse_retry_at_hour_only_and_month_day() {
1833 let now = chrono::Utc::now().timestamp();
1834 for msg in [
1835 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1836 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1837 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1838 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1839 ] {
1840 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1841 assert!(ts > now, "{msg}");
1842 }
1843 }
1844
1845 #[serial(cooldown_state)]
1846 #[test]
1847 fn parse_retry_at_weekly_resets_without_year() {
1848 use chrono::Datelike;
1849 let t = chrono::Local::now() + chrono::Duration::days(2);
1850 let msg = format!(
1851 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1852 t.format("%b"),
1853 t.day()
1854 );
1855 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1856 let remaining = ts - chrono::Utc::now().timestamp();
1857 assert!(
1858 remaining > 3600 && remaining < 3 * 86400,
1859 "expected ~2d retry-at, got {remaining}s"
1860 );
1861 }
1862
1863 #[serial(cooldown_state)]
1864 #[test]
1865 fn parse_retry_at_no_date() {
1866 assert!(parse_retry_at("generic rate limit error").is_none());
1867 assert!(parse_retry_at("").is_none());
1868 }
1869
1870 #[serial(cooldown_state)]
1871 #[test]
1872 fn parse_retry_at_glm_reset_format_with_seconds() {
1873 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
1874 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
1875
1876 use chrono::TimeZone;
1877 let naive = chrono::NaiveDateTime::new(
1878 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
1879 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
1880 );
1881 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1882 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1883 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1884 chrono::LocalResult::None => naive.and_utc(),
1885 };
1886
1887 assert_eq!(
1888 ts,
1889 expected_utc.timestamp(),
1890 "timestamp should reflect local-timezone interpretation of vendor reset date"
1891 );
1892 }
1893
1894 #[serial(cooldown_state)]
1895 #[test]
1896 fn parse_retry_at_billing_cycle_without_date_returns_none() {
1897 let msg = "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle. Upgrade to get more.";
1900 assert!(
1901 parse_retry_at(msg).is_none(),
1902 "billing cycle without date should return None"
1903 );
1904 }
1905
1906 #[serial(cooldown_state)]
1907 #[test]
1908 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
1909 let msg = "Failed to authenticate. API Error: 403 You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
1910 assert_eq!(
1911 parse_relative_usage_window(msg),
1912 Some(5 * 3600),
1913 "should extract the 5-hour window as seconds"
1914 );
1915 }
1916
1917 #[serial(cooldown_state)]
1918 #[test]
1919 fn parse_relative_usage_window_handles_common_durations() {
1920 assert_eq!(
1921 parse_relative_usage_window("You've reached your hourly usage limit."),
1922 Some(3600)
1923 );
1924 assert_eq!(
1925 parse_relative_usage_window("You've reached your daily usage limit."),
1926 Some(86400)
1927 );
1928 assert_eq!(
1929 parse_relative_usage_window("You've reached your weekly limit."),
1930 Some(604800)
1931 );
1932 assert_eq!(
1933 parse_relative_usage_window("You've reached your 30-minute usage limit."),
1934 Some(1800)
1935 );
1936 assert_eq!(
1937 parse_relative_usage_window("You've reached your 2-day quota."),
1938 Some(172800)
1939 );
1940 }
1941
1942 #[serial(cooldown_state)]
1943 #[test]
1944 fn parse_relative_usage_window_returns_none_without_usage_context() {
1945 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
1948 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
1949 assert!(parse_relative_usage_window("").is_none());
1950 }
1951
1952 #[serial(cooldown_state)]
1953 #[tokio::test]
1954 async fn record_rate_limit_uses_model_specific_window_cooldown() {
1955 super::reset_global_state().await;
1956 let store = test_store().await;
1957 *cooldown_store().lock().await = Some(store.clone());
1958
1959 let agent = "test_record_rl_window_agent";
1960 let model = "test_record_rl_window_model";
1961 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
1962
1963 record_rate_limit(agent, Some(model), msg).await;
1964
1965 assert!(
1966 is_model_in_cooldown(agent, model),
1967 "model should be cooled for the relative usage window"
1968 );
1969 assert!(
1970 !is_agent_in_cooldown(agent),
1971 "agent should not be cooled for a model-scoped quota"
1972 );
1973
1974 let key = format!("{agent}:{model}");
1975 let remaining = cooldown_until(&key)
1976 .expect("model cooldown should exist")
1977 .saturating_sub(chrono::Utc::now().timestamp());
1978 assert!(
1979 remaining >= 5 * 3600 - 5,
1980 "model cooldown should be at least ~5 hours, got {remaining}s"
1981 );
1982 assert_eq!(
1983 cooldown_reason(&key).as_deref(),
1984 Some("rate_limit_usage_window"),
1985 "relative usage-window cooldown must be labeled as a rate limit, not silence"
1986 );
1987 }
1988
1989 #[serial(cooldown_state)]
1990 #[tokio::test]
1991 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
1992 super::reset_global_state().await;
1993 let store = test_store().await;
1994 *cooldown_store().lock().await = Some(store.clone());
1995
1996 let agent = "test_record_rl_fallback_agent";
1997 let model = "test_record_rl_fallback_model";
1998 let msg = "429 Too Many Requests";
1999
2000 record_rate_limit(agent, Some(model), msg).await;
2001
2002 assert!(
2003 is_agent_in_cooldown(agent),
2004 "agent should be cooled when no relative window is detected"
2005 );
2006 assert!(
2007 is_model_in_cooldown(agent, model),
2008 "model should be cooled when no relative window is detected"
2009 );
2010 }
2011
2012 #[serial(cooldown_state)]
2013 #[tokio::test]
2014 async fn agent_cooldown_persists_to_kv() {
2015 super::reset_global_state().await;
2016 let agent = "test_agent_persist_check";
2018 record_agent_failure_with_message(agent, "").await;
2019 assert!(is_agent_in_cooldown(agent));
2020 }
2021
2022 #[serial(cooldown_state)]
2023 #[tokio::test]
2024 async fn cooldown_never_shortens() {
2025 super::reset_global_state().await;
2026 let agent = "test_agent_no_shorten";
2027 set_agent_cooldown(agent, 24 * 60 * 60).await;
2028 let initial = {
2029 let map = cooldowns().lock().unwrap();
2030 map.get(agent)
2031 .expect("cooldown entry should exist")
2032 .cooldown_until
2033 };
2034
2035 record_agent_failure_with_message(agent, "rate limit").await;
2037 let after = {
2038 let map = cooldowns().lock().unwrap();
2039 map.get(agent)
2040 .expect("cooldown entry should exist")
2041 .cooldown_until
2042 };
2043 assert_eq!(
2044 initial, after,
2045 "cooldown should not be shortened by later failures"
2046 );
2047 }
2048
2049 #[serial(cooldown_state)]
2050 #[tokio::test]
2051 async fn silence_agent_cooldown_is_short_lived() {
2052 super::reset_global_state().await;
2053 let agent = "test_silence_agent_cd";
2054 assert!(!is_agent_in_cooldown(agent));
2055
2056 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2057 assert!(is_agent_in_cooldown(agent));
2058
2059 let map = cooldowns().lock().unwrap();
2061 let entry = map.get(agent).expect("should have cooldown entry");
2062 let now = chrono::Utc::now().timestamp();
2063 let remaining = entry.cooldown_until - now;
2064 assert!(
2065 remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
2066 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2067 );
2068 assert!(remaining > 0, "cooldown should still be active");
2069 }
2070
2071 #[serial(cooldown_state)]
2072 #[tokio::test]
2073 async fn record_silence_detection_applies_extended_cooldown() {
2074 super::reset_global_state().await;
2075 let store = test_store().await;
2076 *cooldown_store().lock().await = Some(store.clone());
2077
2078 let agent = "test_silence_count_agent";
2079 let model = "test_silence_count_model";
2080 assert!(!is_model_in_cooldown(agent, model));
2081
2082 for _ in 0..SILENCE_COUNT_THRESHOLD {
2083 let result = record_silence_detection(agent, model).await;
2084 assert!(result.is_some());
2085 }
2086
2087 assert!(is_model_in_cooldown(agent, model));
2088 let key = format!("{agent}:{model}");
2089 let remaining = {
2090 let map = cooldowns().lock().unwrap();
2091 let entry = map.get(&key).expect("cooldown entry should exist");
2092 entry.cooldown_until - chrono::Utc::now().timestamp()
2093 };
2094 assert!(
2095 remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
2096 "extended cooldown should be applied, got {remaining}s"
2097 );
2098
2099 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2100 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2101 assert_eq!(stored, "[]");
2102 }
2103
2104 #[serial(cooldown_state)]
2105 #[test]
2106 fn detect_credit_exhaustion_out_of_credits() {
2107 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2108 assert_eq!(
2109 detect_credit_exhaustion(msg),
2110 Some(CreditExhaustionReason::OutOfCredits)
2111 );
2112 }
2113
2114 #[serial(cooldown_state)]
2115 #[test]
2116 fn detect_credit_exhaustion_out_of_credits_variations() {
2117 assert_eq!(
2118 detect_credit_exhaustion("out of credits"),
2119 Some(CreditExhaustionReason::OutOfCredits)
2120 );
2121 assert_eq!(
2122 detect_credit_exhaustion("outofcredits"),
2123 Some(CreditExhaustionReason::OutOfCredits)
2124 );
2125 assert_eq!(
2126 detect_credit_exhaustion("no credits remaining"),
2127 Some(CreditExhaustionReason::OutOfCredits)
2128 );
2129 assert_eq!(
2130 detect_credit_exhaustion("insufficient funds"),
2131 Some(CreditExhaustionReason::OutOfCredits)
2132 );
2133 }
2134
2135 #[serial(cooldown_state)]
2136 #[test]
2137 fn detect_credit_exhaustion_org_level_disabled() {
2138 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2139 assert_eq!(
2140 detect_credit_exhaustion(msg),
2141 Some(CreditExhaustionReason::OrgLevelDisabled)
2142 );
2143 }
2144
2145 #[serial(cooldown_state)]
2146 #[test]
2147 fn detect_credit_exhaustion_org_level_disabled_variations() {
2148 assert_eq!(
2149 detect_credit_exhaustion("organization has been disabled"),
2150 Some(CreditExhaustionReason::OrgLevelDisabled)
2151 );
2152 assert_eq!(
2153 detect_credit_exhaustion("org disabled"),
2154 Some(CreditExhaustionReason::OrgLevelDisabled)
2155 );
2156 assert_eq!(
2157 detect_credit_exhaustion("billing quota exceeded"),
2158 Some(CreditExhaustionReason::OrgLevelDisabled)
2159 );
2160 assert_eq!(
2161 detect_credit_exhaustion("credit balance too low"),
2162 Some(CreditExhaustionReason::OrgLevelDisabled)
2163 );
2164 }
2165
2166 #[serial(cooldown_state)]
2167 #[test]
2168 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2169 let msg = "overageDisabledReason: out_of_credits";
2170 assert_eq!(
2171 detect_credit_exhaustion(msg),
2172 Some(CreditExhaustionReason::OutOfCredits)
2173 );
2174
2175 let msg = "overageDisabledReason: org_level_disabled";
2176 assert_eq!(
2177 detect_credit_exhaustion(msg),
2178 Some(CreditExhaustionReason::OrgLevelDisabled)
2179 );
2180 }
2181
2182 #[serial(cooldown_state)]
2183 #[test]
2184 fn detect_credit_exhaustion_glm_insufficient_balance() {
2185 assert_eq!(
2187 detect_credit_exhaustion(
2188 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2189 ),
2190 Some(CreditExhaustionReason::OutOfCredits)
2191 );
2192 assert_eq!(
2193 detect_credit_exhaustion("Insufficient balance"),
2194 Some(CreditExhaustionReason::OutOfCredits)
2195 );
2196 assert_eq!(
2197 detect_credit_exhaustion("please recharge your account"),
2198 Some(CreditExhaustionReason::OutOfCredits)
2199 );
2200 assert_eq!(
2201 detect_credit_exhaustion("no resource package available"),
2202 Some(CreditExhaustionReason::OutOfCredits)
2203 );
2204 }
2205
2206 #[serial(cooldown_state)]
2207 #[test]
2208 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2209 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2210 assert!(detect_credit_exhaustion("too many requests").is_none());
2211 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2212 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2213 assert!(detect_credit_exhaustion(
2214 "You've hit your usage limit, try again in a few minutes."
2215 )
2216 .is_none());
2217 }
2218
2219 #[serial(cooldown_state)]
2220 #[tokio::test]
2221 async fn record_credit_exhaustion_applies_agent_cooldown() {
2222 super::reset_global_state().await;
2223 let agent = "test_credit_exhaust_agent";
2224 assert!(!is_agent_in_cooldown(agent));
2225
2226 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2227 assert!(is_agent_in_cooldown(agent));
2228
2229 let remaining = {
2230 let map = cooldowns().lock().unwrap();
2231 let entry = map.get(agent).expect("cooldown entry should exist");
2232 entry.cooldown_until - chrono::Utc::now().timestamp()
2233 };
2234 assert!(
2236 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2237 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2238 );
2239 assert!(
2240 remaining <= CREDIT_BACKOFF_MAX_SECS,
2241 "cooldown should not exceed cap of 8h, got {remaining}s"
2242 );
2243 }
2244
2245 #[serial(cooldown_state)]
2246 #[tokio::test]
2247 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2248 super::reset_global_state().await;
2249 let agent = "test_org_disabled_agent";
2250 assert!(!is_agent_in_cooldown(agent));
2251
2252 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2253 assert!(is_agent_in_cooldown(agent));
2254
2255 let remaining = {
2256 let map = cooldowns().lock().unwrap();
2257 let entry = map.get(agent).expect("cooldown entry should exist");
2258 entry.cooldown_until - chrono::Utc::now().timestamp()
2259 };
2260 assert!(
2262 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2263 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2264 );
2265 assert!(
2266 remaining <= CREDIT_BACKOFF_MAX_SECS,
2267 "cooldown should not exceed cap of 8h, got {remaining}s"
2268 );
2269 }
2270
2271 #[serial(cooldown_state)]
2272 #[tokio::test]
2273 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2274 super::reset_global_state().await;
2275 let store = test_store().await;
2276 *cooldown_store().lock().await = Some(store.clone());
2277
2278 let agent = "test_billing_cycle_agent";
2279 assert!(!is_agent_in_cooldown(agent));
2280
2281 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2282 assert!(is_agent_in_cooldown(agent));
2283
2284 let remaining = {
2285 let map = cooldowns().lock().unwrap();
2286 let entry = map.get(agent).expect("cooldown entry should exist");
2287 entry.cooldown_until - chrono::Utc::now().timestamp()
2288 };
2289 assert!(
2291 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2292 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2293 );
2294 assert!(
2295 remaining <= BILLING_CYCLE_MAX_SECS,
2296 "cooldown should not exceed cap of 7 days, got {remaining}s"
2297 );
2298 }
2299
2300 #[serial(cooldown_state)]
2301 #[tokio::test]
2302 async fn record_credit_exhaustion_billing_cycle_escalates() {
2303 super::reset_global_state().await;
2304 let store = test_store().await;
2305 *cooldown_store().lock().await = Some(store.clone());
2306
2307 let agent = "test_billing_cycle_escalation";
2308
2309 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2311 let remaining_1 = {
2312 let map = cooldowns().lock().unwrap();
2313 let entry = map.get(agent).expect("cooldown entry should exist");
2314 entry.cooldown_until - chrono::Utc::now().timestamp()
2315 };
2316 assert!(
2317 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2318 .contains(&remaining_1),
2319 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2320 );
2321
2322 {
2324 let mut map = cooldowns().lock().unwrap();
2325 map.remove(agent);
2326 }
2327
2328 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2330 let remaining_2 = {
2331 let map = cooldowns().lock().unwrap();
2332 let entry = map.get(agent).expect("cooldown entry should exist");
2333 entry.cooldown_until - chrono::Utc::now().timestamp()
2334 };
2335 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2337 remaining_2 >= expected_2 - 5,
2338 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2339 );
2340
2341 {
2343 let mut map = cooldowns().lock().unwrap();
2344 map.remove(agent);
2345 }
2346
2347 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2349 let remaining_3 = {
2350 let map = cooldowns().lock().unwrap();
2351 let entry = map.get(agent).expect("cooldown entry should exist");
2352 entry.cooldown_until - chrono::Utc::now().timestamp()
2353 };
2354 assert!(
2355 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2356 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2357 );
2358 assert!(
2359 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2360 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2361 );
2362 }
2363
2364 #[serial(cooldown_state)]
2365 #[tokio::test]
2366 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2367 super::reset_global_state().await;
2368 let store = test_store().await;
2369 *cooldown_store().lock().await = Some(store.clone());
2370
2371 let agent = "test_billing_cycle_agent_capped";
2372
2373 for _ in 0..5 {
2374 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2375 }
2376
2377 let remaining = {
2378 let map = cooldowns().lock().unwrap();
2379 let entry = map.get(agent).expect("cooldown entry should exist");
2380 entry.cooldown_until - chrono::Utc::now().timestamp()
2381 };
2382 assert!(
2383 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2384 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2385 );
2386 assert!(
2387 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2388 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2389 );
2390 }
2391
2392 #[serial(cooldown_state)]
2393 #[test]
2394 fn credit_exhaustion_backoff_constants() {
2395 assert_eq!(
2397 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2398 CREDIT_BACKOFF_BASE_SECS
2399 );
2400 assert_eq!(
2402 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2403 ORG_BACKOFF_BASE_SECS
2404 );
2405 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2407 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2408 }
2409
2410 #[serial(cooldown_state)]
2411 #[test]
2412 fn compute_backoff_long_cap_without_extended_stays_capped() {
2413 assert_eq!(
2414 compute_backoff(
2415 3,
2416 BILLING_CYCLE_COOLDOWN_SECS,
2417 BILLING_CYCLE_MAX_SECS,
2418 false
2419 ),
2420 BILLING_CYCLE_MAX_SECS,
2421 "7-day capped paths should stop at the documented max"
2422 );
2423 assert_eq!(
2424 compute_backoff(
2425 9,
2426 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2427 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2428 false,
2429 ),
2430 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2431 "persistent model failures should not extend past 7 days"
2432 );
2433 }
2434
2435 #[serial(cooldown_state)]
2438 #[test]
2439 fn mark_and_check_agent_degraded() {
2440 let agent = "test_degraded_mark";
2441 assert!(!is_agent_degraded(agent));
2442 mark_agent_degraded(agent);
2443 assert!(is_agent_degraded(agent));
2444 clear_agent_degraded(agent);
2445 assert!(!is_agent_degraded(agent));
2446 }
2447
2448 #[serial(cooldown_state)]
2449 #[tokio::test]
2450 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2451 super::reset_global_state().await;
2452 let store = test_store().await;
2453 let agent = "test_sync_external_clear";
2454 let future = chrono::Utc::now().timestamp() + 600;
2455
2456 init_cooldown_store(store.clone()).await;
2458 set_cooldown_async(agent, future, "test").await;
2459 assert!(is_agent_in_cooldown(agent));
2460
2461 store
2465 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2466 .await
2467 .expect("kv_set tombstone");
2468 assert!(
2469 is_agent_in_cooldown(agent),
2470 "in-memory cooldown should still be present before sync"
2471 );
2472
2473 sync_from_kv(&store).await;
2474
2475 assert!(
2476 !is_agent_in_cooldown(agent),
2477 "sync_from_kv must drop in-memory entries the KV tombstoned"
2478 );
2479 }
2480
2481 #[serial(cooldown_state)]
2482 #[tokio::test]
2483 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2484 super::reset_global_state().await;
2488 let store = test_store().await;
2489 let agent = "test_sync_race_window";
2490 let future = chrono::Utc::now().timestamp() + 600;
2491
2492 set_cooldown_in_memory(agent, future, "raced");
2494 assert!(is_agent_in_cooldown(agent));
2495
2496 sync_from_kv(&store).await;
2497
2498 assert!(
2499 is_agent_in_cooldown(agent),
2500 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2501 );
2502 }
2503
2504 #[serial(cooldown_state)]
2505 #[tokio::test]
2506 async fn sync_from_kv_adopts_kv_timestamp() {
2507 super::reset_global_state().await;
2510 let store = test_store().await;
2511 let agent = "test_sync_adopt_kv";
2512 let original = chrono::Utc::now().timestamp() + 600;
2513 let updated = chrono::Utc::now().timestamp() + 1200;
2514
2515 init_cooldown_store(store.clone()).await;
2516 set_cooldown_async(agent, original, "initial").await;
2517
2518 store
2519 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2520 .await
2521 .expect("kv_set updated");
2522
2523 sync_from_kv(&store).await;
2524
2525 let until = cooldown_until(agent).expect("cooldown still active");
2526 assert_eq!(
2527 until, updated,
2528 "sync_from_kv must adopt KV's authoritative timestamp"
2529 );
2530 }
2531
2532 #[serial(cooldown_state)]
2533 #[tokio::test]
2534 async fn refresh_degraded_agents_marks_cooled_agent() {
2535 super::reset_global_state().await;
2536 let store = test_store().await;
2537 let agent = "test_refresh_cooled";
2538
2539 record_agent_failure_with_message(agent, "").await;
2541 assert!(is_agent_in_cooldown(agent));
2542
2543 let agents = vec![agent.to_string()];
2545 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2546
2547 assert!(
2548 is_agent_degraded(agent),
2549 "agent in cooldown should be marked degraded"
2550 );
2551
2552 clear_agent_degraded(agent);
2554 }
2555
2556 #[serial(cooldown_state)]
2557 #[tokio::test]
2558 async fn refresh_degraded_agents_marks_no_models_agent() {
2559 super::reset_global_state().await;
2560 let store = test_store().await;
2561 let agent = "test_refresh_no_models";
2562
2563 let agents = vec![agent.to_string()];
2565 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2566
2567 assert!(
2568 is_agent_degraded(agent),
2569 "agent with no available models should be marked degraded"
2570 );
2571
2572 clear_agent_degraded(agent);
2574 }
2575
2576 #[serial(cooldown_state)]
2577 #[tokio::test]
2578 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2579 super::reset_global_state().await;
2580 let store = test_store().await;
2581 let agent = "test_refresh_rate_limited";
2582
2583 for _ in 0..4 {
2585 store
2586 .record_rate_limit(agent, "rate_limit", None)
2587 .await
2588 .unwrap();
2589 }
2590
2591 let agents = vec![agent.to_string()];
2592 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2594
2595 assert!(
2596 is_agent_degraded(agent),
2597 "agent with rate limit count >= threshold should be marked degraded"
2598 );
2599
2600 clear_agent_degraded(agent);
2602 }
2603
2604 #[serial(cooldown_state)]
2605 #[tokio::test]
2606 async fn refresh_degraded_agents_clears_healthy_agent() {
2607 super::reset_global_state().await;
2608 let store = test_store().await;
2609 let agent = "test_refresh_healthy";
2610
2611 mark_agent_degraded(agent);
2613 assert!(is_agent_degraded(agent));
2614
2615 let agents = vec![agent.to_string()];
2617 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2618
2619 assert!(
2620 !is_agent_degraded(agent),
2621 "healthy agent should have degraded flag cleared"
2622 );
2623 }
2624
2625 #[serial(cooldown_state)]
2626 #[tokio::test]
2627 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2628 super::reset_global_state().await;
2629 let store = test_store().await;
2630 let agent = "test_refresh_expired_cooldown";
2631
2632 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2636 let mut map = cooldowns().lock().unwrap();
2637 map.insert(
2638 agent.to_string(),
2639 CooldownEntry {
2640 cooldown_until: expired_ts,
2641 reason: "expired_test".to_string(),
2642 },
2643 );
2644 }
2645
2646 assert!(
2648 !is_agent_in_cooldown(agent),
2649 "expired cooldown should not be active"
2650 );
2651
2652 mark_agent_degraded(agent);
2654 assert!(is_agent_degraded(agent));
2655
2656 let agents = vec![agent.to_string()];
2658 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2659
2660 assert!(
2661 !is_agent_degraded(agent),
2662 "agent with only expired cooldown should NOT be marked degraded"
2663 );
2664
2665 clear_agent_degraded(agent);
2667 {
2668 let mut map = cooldowns().lock().unwrap();
2669 map.remove(agent);
2670 }
2671 }
2672
2673 #[serial(cooldown_state)]
2674 #[test]
2675 fn compute_backoff_grows_exponentially() {
2676 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2678 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2679 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2681 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2683 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2685 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2687 }
2688
2689 #[serial(cooldown_state)]
2690 #[test]
2691 fn compute_backoff_extended_tier_1_applies_24h() {
2692 assert_eq!(
2698 compute_backoff(5, 300, 14400, true),
2699 14400,
2700 "first at max, not extended"
2701 );
2702 assert_eq!(
2703 compute_backoff(6, 300, 14400, true),
2704 14400,
2705 "times_at_max=1, not extended"
2706 );
2707 assert_eq!(
2708 compute_backoff(7, 300, 14400, true),
2709 86400,
2710 "times_at_max=2, extended 6x"
2711 );
2712 assert_eq!(
2713 compute_backoff(8, 300, 14400, true),
2714 86400,
2715 "times_at_max=3, still extended 6x"
2716 );
2717 }
2718
2719 #[serial(cooldown_state)]
2720 #[test]
2721 fn compute_backoff_extended_tier_2_applies_48h() {
2722 assert_eq!(
2725 compute_backoff(9, 300, 14400, true),
2726 86400,
2727 "times_at_max=4, still tier 1"
2728 );
2729 assert_eq!(
2730 compute_backoff(10, 300, 14400, true),
2731 172800,
2732 "times_at_max=5, extended 12x"
2733 );
2734 assert_eq!(
2735 compute_backoff(15, 300, 14400, true),
2736 172800,
2737 "times_at_max=10, still tier 2"
2738 );
2739 assert_eq!(
2740 compute_backoff(20, 300, 14400, true),
2741 172800,
2742 "times_at_max=15, still tier 2"
2743 );
2744 }
2745
2746 #[serial(cooldown_state)]
2747 #[test]
2748 fn compute_backoff_extended_tier_with_different_base_max() {
2749 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2756 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2757 assert_eq!(
2758 compute_backoff(7, 60, 3600, true),
2759 21600,
2760 "times_at_max=2, extended 6x"
2761 );
2762 }
2763
2764 #[serial(cooldown_state)]
2765 #[test]
2766 fn compute_backoff_extended_tier_persistent_fails_example() {
2767 assert_eq!(
2771 compute_backoff(21, 300, 14400, true),
2772 172800,
2773 "21 failures should get 48h extended cooldown"
2774 );
2775
2776 assert_eq!(
2778 compute_backoff(36, 300, 14400, true),
2779 172800,
2780 "36 failures should also get 48h extended cooldown"
2781 );
2782 }
2783
2784 #[serial(cooldown_state)]
2785 #[test]
2786 fn compute_backoff_extended_tier_zero_base() {
2787 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2789 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2790 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2791 }
2792
2793 #[serial(cooldown_state)]
2794 #[test]
2795 fn detect_credit_exhaustion_billing_cycle() {
2796 assert_eq!(
2797 detect_credit_exhaustion("billing cycle"),
2798 Some(CreditExhaustionReason::BillingCycleExhausted)
2799 );
2800 assert_eq!(
2801 detect_credit_exhaustion("quota refreshed next cycle"),
2802 Some(CreditExhaustionReason::BillingCycleExhausted)
2803 );
2804 assert_eq!(
2805 detect_credit_exhaustion(
2806 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
2807 ),
2808 Some(CreditExhaustionReason::BillingCycleExhausted)
2809 );
2810 assert_eq!(
2811 detect_credit_exhaustion("monthly quota exceeded"),
2812 Some(CreditExhaustionReason::BillingCycleExhausted)
2813 );
2814 assert_eq!(
2815 detect_credit_exhaustion(
2816 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
2817 ),
2818 Some(CreditExhaustionReason::BillingCycleExhausted)
2819 );
2820 }
2821
2822 #[serial(cooldown_state)]
2823 #[test]
2824 fn detect_credit_exhaustion_billing_cycle_minimax() {
2825 assert_eq!(
2826 detect_credit_exhaustion(
2827 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
2828 ),
2829 Some(CreditExhaustionReason::BillingCycleExhausted)
2830 );
2831 assert_eq!(
2832 detect_credit_exhaustion("Upgrade your Token Plan"),
2833 Some(CreditExhaustionReason::BillingCycleExhausted)
2834 );
2835 assert_eq!(
2836 detect_credit_exhaustion("Token Plan usage limit reached"),
2837 Some(CreditExhaustionReason::BillingCycleExhausted)
2838 );
2839 assert_eq!(
2840 detect_credit_exhaustion("purchase Credits for more usage"),
2841 Some(CreditExhaustionReason::BillingCycleExhausted)
2842 );
2843 assert_eq!(
2844 detect_credit_exhaustion(
2845 "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits"
2846 ),
2847 Some(CreditExhaustionReason::BillingCycleExhausted)
2848 );
2849 }
2850
2851 #[serial(cooldown_state)]
2852 #[test]
2853 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
2854 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
2859 this billing cycle. Your quota will be refreshed soon.";
2860 assert_eq!(
2861 detect_credit_exhaustion_outcome(msg),
2862 Some("billing_cycle_exhausted")
2863 );
2864 }
2865
2866 #[serial(cooldown_state)]
2867 #[test]
2868 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
2869 assert_eq!(
2870 detect_credit_exhaustion_outcome("out of credits"),
2871 Some("credit_exhausted")
2872 );
2873 assert_eq!(
2874 detect_credit_exhaustion_outcome("organization has been disabled"),
2875 Some("org_level_disabled")
2876 );
2877 assert_eq!(
2878 detect_credit_exhaustion_outcome("monthly quota exceeded"),
2879 Some("billing_cycle_exhausted")
2880 );
2881 }
2882
2883 #[test]
2884 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
2885 assert_eq!(
2886 detect_credit_exhaustion_outcome("429 Too Many Requests"),
2887 None
2888 );
2889 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
2890 }
2891
2892 #[serial(cooldown_state)]
2893 #[test]
2894 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
2895 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
2898 assert_eq!(
2899 detect_credit_exhaustion(msg),
2900 Some(CreditExhaustionReason::OrgLevelDisabled)
2901 );
2902 }
2903
2904 #[serial(cooldown_state)]
2905 #[tokio::test]
2906 async fn refresh_degraded_agents_below_threshold_not_degraded() {
2907 super::reset_global_state().await;
2908 let store = test_store().await;
2909 let agent = "test_refresh_below_threshold";
2910
2911 for _ in 0..2 {
2913 store
2914 .record_rate_limit(agent, "rate_limit", None)
2915 .await
2916 .unwrap();
2917 }
2918
2919 let agents = vec![agent.to_string()];
2920 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2922
2923 assert!(
2924 !is_agent_degraded(agent),
2925 "agent with rate limit count < threshold should not be degraded"
2926 );
2927 }
2928
2929 #[serial(cooldown_state)]
2932 #[tokio::test]
2933 async fn repeated_agent_failures_escalate_backoff() {
2934 super::reset_global_state().await;
2935 let store = test_store().await;
2936 *cooldown_store().lock().await = Some(store.clone());
2937
2938 let agent = "test_escalation_agent";
2939
2940 record_agent_failure_with_message(agent, "").await;
2942 let remaining_1 = {
2943 let map = cooldowns().lock().unwrap();
2944 let entry = map.get(agent).expect("should have cooldown");
2945 entry.cooldown_until - chrono::Utc::now().timestamp()
2946 };
2947 assert!(
2948 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
2949 "first failure should be ~5 min, got {remaining_1}s"
2950 );
2951
2952 {
2954 let mut map = cooldowns().lock().unwrap();
2955 map.remove(agent);
2956 }
2957
2958 record_agent_failure_with_message(agent, "").await;
2960 let remaining_2 = {
2961 let map = cooldowns().lock().unwrap();
2962 let entry = map.get(agent).expect("should have cooldown");
2963 entry.cooldown_until - chrono::Utc::now().timestamp()
2964 };
2965 assert!(
2966 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
2967 "second failure should be ~15 min, got {remaining_2}s"
2968 );
2969 }
2970
2971 #[serial(cooldown_state)]
2972 #[tokio::test]
2973 async fn vendor_past_retry_falls_back_to_backoff() {
2974 super::reset_global_state().await;
2975 let store = test_store().await;
2979 *cooldown_store().lock().await = Some(store.clone());
2980
2981 let agent = "test_vendor_past_retry";
2982
2983 {
2985 let mut map = cooldowns().lock().unwrap();
2986 map.remove(agent);
2987 }
2988
2989 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
2991
2992 record_agent_failure_with_message(agent, past_msg).await;
2993
2994 assert!(
2998 is_agent_in_cooldown(agent),
2999 "agent should be in cooldown after failure"
3000 );
3001
3002 let remaining = {
3003 let map = cooldowns().lock().unwrap();
3004 let entry = map.get(agent).expect("cooldown entry should exist");
3005 entry.cooldown_until - chrono::Utc::now().timestamp()
3006 };
3007
3008 assert!(
3009 remaining >= BACKOFF_BASE_SECS - 5,
3010 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3011 );
3012
3013 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3015 let stored = store
3016 .kv_get(&kv_key)
3017 .await
3018 .unwrap()
3019 .expect("failure count should be stored");
3020 assert_eq!(stored, "1");
3021 }
3022
3023 #[serial(cooldown_state)]
3024 #[tokio::test]
3025 async fn record_agent_success_resets_backoff() {
3026 super::reset_global_state().await;
3027 let store = test_store().await;
3028 *cooldown_store().lock().await = Some(store.clone());
3029
3030 let agent = "test_success_reset_agent";
3031 let model = "test_model";
3032
3033 for _ in 0..3 {
3035 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3036 }
3037 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3039 let count: u32 = store
3040 .kv_get(&kv_key)
3041 .await
3042 .unwrap()
3043 .unwrap()
3044 .parse()
3045 .unwrap();
3046 assert_eq!(count, 3);
3047
3048 record_agent_success(agent, model).await;
3050
3051 let count_after: u32 = store
3052 .kv_get(&kv_key)
3053 .await
3054 .unwrap()
3055 .unwrap()
3056 .parse()
3057 .unwrap();
3058 assert_eq!(count_after, 0, "success should reset failure count to 0");
3059
3060 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3062 let model_count: u32 = store
3063 .kv_get(&model_kv_key)
3064 .await
3065 .unwrap()
3066 .unwrap()
3067 .parse()
3068 .unwrap();
3069 assert_eq!(
3070 model_count, 0,
3071 "success should reset model failure count to 0"
3072 );
3073 }
3074
3075 #[serial(cooldown_state)]
3076 #[tokio::test]
3077 async fn record_agent_success_without_model_skips_empty_model_key() {
3078 super::reset_global_state().await;
3079 let store = test_store().await;
3080 *cooldown_store().lock().await = Some(store.clone());
3081
3082 let agent = "test_success_without_model";
3083
3084 record_agent_success(agent, "").await;
3085
3086 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3087 let agent_count = store
3088 .kv_get(&agent_kv_key)
3089 .await
3090 .unwrap()
3091 .expect("agent failure count should be stored");
3092 assert_eq!(agent_count, "0");
3093
3094 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3095 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3096 assert!(
3097 empty_model_count.is_none(),
3098 "success without model should not create an empty model failure count key"
3099 );
3100 }
3101
3102 #[serial(cooldown_state)]
3103 #[tokio::test]
3104 async fn clear_cooldown_resets_failure_counts() {
3105 super::reset_global_state().await;
3106 let store = test_store().await;
3107 *cooldown_store().lock().await = Some(store.clone());
3108
3109 let agent = "test_clear_fc_agent";
3110
3111 for _ in 0..5 {
3113 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3114 }
3115 set_agent_cooldown(agent, 3600).await;
3117
3118 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3120 let count: u32 = store
3121 .kv_get(&kv_key)
3122 .await
3123 .unwrap()
3124 .unwrap()
3125 .parse()
3126 .unwrap();
3127 assert_eq!(count, 5);
3128
3129 clear_cooldown(agent, &store).await;
3131
3132 let count_after: u32 = store
3134 .kv_get(&kv_key)
3135 .await
3136 .unwrap()
3137 .unwrap()
3138 .parse()
3139 .unwrap();
3140 assert_eq!(
3141 count_after, 0,
3142 "clear_cooldown should reset failure count to 0"
3143 );
3144
3145 assert!(
3147 !is_agent_in_cooldown(agent),
3148 "agent should not be in cooldown after clear"
3149 );
3150 }
3151
3152 #[serial(cooldown_state)]
3153 #[tokio::test]
3154 async fn clear_all_cooldowns_resets_all_failure_counts() {
3155 super::reset_global_state().await;
3156 let store = test_store().await;
3157 *cooldown_store().lock().await = Some(store.clone());
3158
3159 let agents = ["test_clear_all_a", "test_clear_all_b"];
3160 for agent in &agents {
3161 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3162 set_agent_cooldown(agent, 3600).await;
3163 }
3164
3165 clear_cooldown("*", &store).await;
3166
3167 for agent in &agents {
3168 assert!(!is_agent_in_cooldown(agent));
3169 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3170 let count: u32 = store
3171 .kv_get(&kv_key)
3172 .await
3173 .unwrap()
3174 .unwrap()
3175 .parse()
3176 .unwrap();
3177 assert_eq!(
3178 count, 0,
3179 "clear --all should reset failure count for {agent}"
3180 );
3181 }
3182 }
3183
3184 #[serial(cooldown_state)]
3185 #[test]
3186 fn compute_backoff_saturates_on_large_counts() {
3187 assert_eq!(
3189 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3190 BACKOFF_MAX_SECS * 12
3191 );
3192 assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
3193 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
3194 }
3195
3196 #[serial(cooldown_state)]
3197 #[test]
3198 fn compute_backoff_zero_base_returns_zero() {
3199 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3200 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3201 }
3202
3203 #[serial(cooldown_state)]
3206 #[tokio::test]
3207 async fn github_circuit_stays_closed_below_threshold() {
3208 super::reset_global_state().await;
3209 {
3211 let mut open = github_5xx_circuit_open().lock().unwrap();
3212 *open = None;
3213 let mut ts = github_5xx_timestamps().lock().unwrap();
3214 ts.clear();
3215 let mut map = cooldown_lock();
3219 map.remove("github:5xx");
3220 }
3221
3222 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3224 record_github_5xx().await;
3225 }
3226
3227 assert!(
3228 !is_github_circuit_open(),
3229 "circuit should stay closed below threshold"
3230 );
3231 assert_eq!(github_circuit_remaining_secs(), 0);
3232 }
3233
3234 #[serial(cooldown_state)]
3235 #[tokio::test]
3236 async fn github_circuit_trips_at_threshold() {
3237 super::reset_global_state().await;
3238 {
3240 let mut open = github_5xx_circuit_open_lock();
3241 *open = None;
3242 let mut ts = github_5xx_timestamps_lock();
3243 ts.clear();
3244 let mut map = cooldown_lock();
3245 map.remove("github:5xx");
3246 }
3247
3248 for _ in 0..GITHUB_5XX_THRESHOLD {
3250 record_github_5xx().await;
3251 }
3252
3253 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3254 assert!(
3255 github_circuit_remaining_secs() > 0,
3256 "remaining should be > 0 when open"
3257 );
3258 assert!(
3259 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3260 "remaining should not exceed cooldown duration"
3261 );
3262
3263 {
3265 let mut open = github_5xx_circuit_open().lock().unwrap();
3266 *open = None;
3267 }
3268 }
3269
3270 #[serial(cooldown_state)]
3271 #[tokio::test]
3272 async fn github_circuit_cannot_double_trip() {
3273 super::reset_global_state().await;
3274 {
3276 let mut open = github_5xx_circuit_open_lock();
3277 *open = None;
3278 let mut ts = github_5xx_timestamps_lock();
3279 ts.clear();
3280 let mut map = cooldown_lock();
3281 map.remove("github:5xx");
3282 }
3283
3284 for _ in 0..GITHUB_5XX_THRESHOLD {
3286 record_github_5xx().await;
3287 }
3288 assert!(is_github_circuit_open());
3289
3290 let first_remaining = github_circuit_remaining_secs();
3291
3292 for _ in 0..GITHUB_5XX_THRESHOLD {
3294 record_github_5xx().await;
3295 }
3296
3297 let second_remaining = github_circuit_remaining_secs();
3298 assert!(
3299 second_remaining <= first_remaining,
3300 "additional errors should not extend cooldown"
3301 );
3302
3303 {
3305 let mut open = github_5xx_circuit_open().lock().unwrap();
3306 *open = None;
3307 }
3308 }
3309
3310 #[serial(cooldown_state)]
3311 #[tokio::test]
3312 async fn github_circuit_recovers_after_cooldown() {
3313 super::reset_global_state().await;
3314 {
3316 let mut open = github_5xx_circuit_open_lock();
3317 *open = None;
3318 let mut ts = github_5xx_timestamps_lock();
3319 ts.clear();
3320 let mut map = cooldown_lock();
3321 map.remove("github:5xx");
3322 }
3323
3324 {
3326 let mut open = github_5xx_circuit_open().lock().unwrap();
3327 *open = Some(chrono::Utc::now().timestamp() - 1);
3328 }
3329
3330 assert!(
3332 !is_github_circuit_open(),
3333 "circuit should auto-recover after cooldown"
3334 );
3335 }
3336
3337 #[serial(cooldown_state)]
3338 #[tokio::test]
3339 async fn github_circuit_sliding_window_ages_out() {
3340 super::reset_global_state().await;
3341 {
3343 let mut open = github_5xx_circuit_open_lock();
3344 *open = None;
3345 let mut ts = github_5xx_timestamps_lock();
3346 ts.clear();
3347 let mut map = cooldown_lock();
3348 map.remove("github:5xx");
3349 }
3350
3351 {
3353 let mut ts = github_5xx_timestamps().lock().unwrap();
3354 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3355 for _ in 0..GITHUB_5XX_THRESHOLD {
3356 ts.push(old);
3357 }
3358 }
3359
3360 record_github_5xx().await;
3362
3363 assert!(
3364 !is_github_circuit_open(),
3365 "old errors outside window should not trip circuit"
3366 );
3367 }
3368}