1use chrono::TimeZone;
19use regex::Regex;
20use std::collections::HashMap;
21use std::collections::HashSet;
22use std::sync::LazyLock;
23use std::sync::MutexGuard;
24use std::sync::{Arc, Mutex, OnceLock};
25use tokio::sync;
26
27use crate::engine::router::config;
28#[cfg(test)]
29use crate::engine::runner::agents::opencode;
30use crate::store::TaskStore;
31
32pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
38
39pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
41
42pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
44
45pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
47
48pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
50
51pub const TRANSIENT_MODEL_BACKOFF_BASE_SECS: i64 = 5 * 60;
53
54pub const TRANSIENT_MODEL_BACKOFF_MAX_SECS: i64 = 15 * 60;
56
57pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
59
60pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
62
63pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
66
67pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
75
76pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum CreditExhaustionReason {
86 OutOfCredits,
88 OrgLevelDisabled,
90 BillingCycleExhausted,
92}
93
94impl CreditExhaustionReason {
95 pub fn outcome_str(self) -> &'static str {
99 match self {
100 CreditExhaustionReason::OutOfCredits => "credit_exhausted",
101 CreditExhaustionReason::OrgLevelDisabled => "org_level_disabled",
102 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
103 }
104 }
105}
106
107pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
114
115pub const SILENCE_COUNT_THRESHOLD: usize = 5;
117
118pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
120
121pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
123
124const KV_PREFIX: &str = "cooldown:";
126
127const SILENCE_COUNT_PREFIX: &str = "silence_count:";
129
130const FAILURE_COUNT_PREFIX: &str = "failure_count:";
132
133const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
136
137const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
141
142const RATE_LIMITS_CLEARED_AT_PREFIX: &str = "rate_limits_cleared_at:";
144
145fn agent_from_key(key: &str) -> &str {
147 key.split_once(':').map_or(key, |(agent, _)| agent)
148}
149
150struct CooldownEntry {
151 cooldown_until: i64,
153 reason: String,
154}
155
156#[derive(Debug)]
157pub struct SilenceCountResult {
158 pub count: usize,
159 pub extended_cooldown_applied: bool,
160}
161
162fn cooldowns() -> &'static Mutex<HashMap<String, CooldownEntry>> {
171 static COOLDOWNS: OnceLock<Mutex<HashMap<String, CooldownEntry>>> = OnceLock::new();
172 COOLDOWNS.get_or_init(|| Mutex::new(HashMap::new()))
173}
174
175fn cooldown_store() -> &'static sync::Mutex<Option<Arc<TaskStore>>> {
181 static STORE: OnceLock<sync::Mutex<Option<Arc<TaskStore>>>> = OnceLock::new();
182 STORE.get_or_init(|| sync::Mutex::new(None))
183}
184
185fn parse_cooldown_value(value: &str) -> Option<(i64, String)> {
193 if let Ok(serde_json::Value::Object(map)) = serde_json::from_str(value) {
194 let until = map.get("until")?.as_i64()?;
195 let reason = map
196 .get("reason")
197 .and_then(|r| r.as_str())
198 .unwrap_or("persisted")
199 .to_string();
200 return Some((until, reason));
201 }
202 value
203 .parse::<i64>()
204 .ok()
205 .map(|until| (until, "persisted".to_string()))
206}
207
208pub async fn init_cooldown_store(store: Arc<TaskStore>) {
209 match store.kv_list_prefix(KV_PREFIX).await {
210 Ok(rows) => {
211 let now = chrono::Utc::now().timestamp();
212 let mut map = cooldown_lock();
213 let mut loaded = 0;
214 for (key, value) in rows {
215 let cooldown_key = key.trim_start_matches(KV_PREFIX);
216 if let Some((cooldown_until, reason)) = parse_cooldown_value(&value)
217 && now < cooldown_until
218 {
219 map.insert(
220 cooldown_key.to_string(),
221 CooldownEntry {
222 cooldown_until,
223 reason,
224 },
225 );
226 loaded += 1;
227 }
228 }
231 if loaded > 0 {
232 tracing::info!(loaded, "loaded persisted cooldowns from KV store");
233 }
234 }
235 Err(e) => {
236 tracing::warn!(err = %e, "failed to load cooldowns from KV store");
237 }
238 }
239
240 if let Ok(Some(raw)) = store.kv_get("cooldown:github:5xx").await
242 && let Ok(cooldown_until) = raw.parse::<i64>()
243 {
244 let now = chrono::Utc::now().timestamp();
245 if now < cooldown_until {
246 let mut open = github_5xx_circuit_open_lock();
247 *open = Some(cooldown_until);
248 let remaining = cooldown_until - now;
249 tracing::warn!(
250 remaining_secs = remaining,
251 "GitHub 5xx circuit breaker restored from KV (still open)"
252 );
253 }
254 }
255
256 *cooldown_store().lock().await = Some(store);
257}
258
259pub async fn sync_from_kv(store: &Arc<TaskStore>) {
276 let rows = match store.kv_list_prefix(KV_PREFIX).await {
277 Ok(rows) => rows,
278 Err(e) => {
279 tracing::warn!(err = %e, "cooldown sync_from_kv: KV list failed");
280 return;
281 }
282 };
283 let now = chrono::Utc::now().timestamp();
284
285 let mut explicit_kv: HashMap<String, Option<(i64, String)>> =
288 HashMap::with_capacity(rows.len());
289 for (kv_key, value) in &rows {
290 let cooldown_key = kv_key.trim_start_matches(KV_PREFIX);
291 if let Some((ts, reason)) = parse_cooldown_value(value) {
292 let normalized = if ts > now { Some((ts, reason)) } else { None };
293 explicit_kv.insert(cooldown_key.to_string(), normalized);
294 }
295 }
296
297 match explicit_kv.get("github:5xx") {
299 Some(None) => {
300 let mut open = github_5xx_circuit_open_lock();
301 if open.take().is_some() {
302 tracing::info!("GitHub 5xx circuit breaker closed by KV clear");
303 github_5xx_timestamps_lock().clear();
304 }
305 }
306 Some(Some((ts, _))) => {
307 *github_5xx_circuit_open_lock() = Some(*ts);
308 }
309 None => {}
310 }
311
312 let mut map = cooldown_lock();
313 let mut dropped = 0usize;
314 let mut updated = 0usize;
315
316 map.retain(|key, _| {
318 if matches!(explicit_kv.get(key), Some(None)) {
319 dropped += 1;
320 false
321 } else {
322 true
323 }
324 });
325
326 for (key, normalized) in explicit_kv {
328 if let Some((ts, reason)) = normalized {
329 if let Some(entry) = map.get_mut(&key) {
330 if entry.cooldown_until != ts {
331 entry.cooldown_until = ts;
332 updated += 1;
333 }
334 } else {
335 map.insert(
336 key,
337 CooldownEntry {
338 cooldown_until: ts,
339 reason,
340 },
341 );
342 updated += 1;
343 }
344 }
345 }
346
347 if dropped > 0 || updated > 0 {
348 tracing::info!(
349 dropped,
350 updated,
351 "cooldown sync_from_kv: reconciled in-memory map with KV"
352 );
353 }
354}
355
356const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
360
361const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
363
364const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
368
369const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
371
372pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
391 if count <= 1 {
392 return base;
393 }
394 let exp = u32::try_from(count.saturating_sub(1)).unwrap_or(u32::MAX);
395 let factor = 3_i64.saturating_pow(exp);
396 let standard = base.saturating_mul(factor).min(max);
397
398 let first_max = if base <= 0 {
407 0
408 } else {
409 let mut accumulated = base;
411 let mut first_max_count = 1u64;
412 while accumulated < max && first_max_count < count {
413 accumulated = accumulated.saturating_mul(3);
414 first_max_count += 1;
415 }
416 first_max_count
417 };
418
419 let times_at_max = count.saturating_sub(first_max);
421
422 if !allow_extended {
423 return standard;
424 }
425
426 if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
427 max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
428 } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
429 max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
430 } else {
431 standard
432 }
433}
434
435async fn read_and_increment_failure_count(store_opt: Option<&Arc<TaskStore>>, key: &str) -> u64 {
444 read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
445}
446
447async fn read_and_increment_failure_count_with_prefix(
452 store_opt: Option<&Arc<TaskStore>>,
453 prefix: &str,
454 key: &str,
455) -> u64 {
456 let kv_key = format!("{prefix}{key}");
457 if let Some(store) = store_opt {
458 match store.kv_increment(&kv_key).await {
459 Ok(n) => n,
460 Err(e) => {
461 tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
462 u64::MAX
463 }
464 }
465 } else {
466 1
467 }
468}
469
470pub async fn record_agent_success(agent_name: &str, model: &str) {
476 let store_opt = cooldown_store().lock().await.clone();
477 if let Some(store) = store_opt {
478 let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
479 let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
480 if let Err(e) = store.kv_set(&agent_key, "0").await {
481 tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
482 }
483 if !model.is_empty() {
484 let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
485 if let Err(e) = store.kv_set(&model_key, "0").await {
486 tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
487 }
488 }
489 if let Err(e) = store.kv_set(&credit_key, "0").await {
490 tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
491 }
492 }
493}
494
495pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
501 if let Some(cooldown_until) = parse_retry_at(error_message) {
506 let now = chrono::Utc::now().timestamp();
507 if cooldown_until > now {
508 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
509 return;
510 }
511 tracing::debug!(
512 agent = agent_name,
513 cooldown_until,
514 "vendor retry_at is in the past — using backoff instead"
515 );
516 }
517
518 let store_opt = cooldown_store().lock().await.clone();
519 let count = read_and_increment_failure_count(store_opt.as_ref(), agent_name).await;
520 let base = config::get_agent_backoff_base(agent_name);
521 let cooldown_until =
522 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
523 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
524}
525
526pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
535 detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
536}
537
538pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
543 let lower = error_message.to_lowercase();
544
545 let overage_patterns = [
546 "overagedisabledreason",
547 "overage_disabled_reason",
548 "overage disabled reason",
549 ];
550
551 if overage_patterns.iter().any(|p| lower.contains(p)) {
552 if lower.contains("out_of_credits") || lower.contains("outofcredits") {
553 return Some(CreditExhaustionReason::OutOfCredits);
554 }
555 if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
556 return Some(CreditExhaustionReason::OrgLevelDisabled);
557 }
558 }
559
560 let billing_patterns = [
561 "credit balance too low",
562 "credit balance insufficient",
563 "insufficient credit balance",
564 "billing quota exceeded",
565 "billing limit exceeded",
566 "organization has been disabled",
567 "org disabled",
568 ];
569
570 if billing_patterns.iter().any(|p| lower.contains(p)) {
571 return Some(CreditExhaustionReason::OrgLevelDisabled);
572 }
573
574 let billing_cycle_patterns = [
577 "billing cycle",
578 "quota refreshed next cycle",
579 "monthly quota",
580 "quota will be refreshed",
581 "next billing cycle",
582 "refreshed in the next cycle",
583 "weekly/monthly limit exhausted",
584 "limit will reset at",
585 "upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
590
591 if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
592 return Some(CreditExhaustionReason::BillingCycleExhausted);
593 }
594
595 if lower.contains("out of credits")
596 || lower.contains("outofcredits")
597 || lower.contains("insufficient funds")
598 || lower.contains("no credits remaining")
599 || lower.contains("insufficient_quota")
600 || lower.contains("quota exceeded")
601 || lower.contains("insufficient balance")
602 || lower.contains("please recharge")
603 || lower.contains("no resource package")
604 || lower.contains("balance exhausted")
605 || lower.contains("402 payment required")
606 {
607 return Some(CreditExhaustionReason::OutOfCredits);
608 }
609
610 None
611}
612
613pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
626 let reason_str = match reason {
627 CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
628 CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
629 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
630 };
631
632 if is_in_cooldown_with_reason(agent_name, reason_str) {
634 tracing::debug!(
635 agent = agent_name,
636 reason = reason_str,
637 "credit exhaustion already cooled, not escalating"
638 );
639 return;
640 }
641
642 let store_opt = cooldown_store().lock().await.clone();
643 let count = read_and_increment_failure_count_with_prefix(
648 store_opt.as_ref(),
649 CREDIT_FAILURE_COUNT_PREFIX,
650 agent_name,
651 )
652 .await;
653
654 let (base, max) = match reason {
655 CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
656 CreditExhaustionReason::OrgLevelDisabled => {
657 (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
658 }
659 CreditExhaustionReason::BillingCycleExhausted => {
660 (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
661 }
662 };
663 let cooldown_secs = compute_backoff(
664 count,
665 base,
666 max,
667 !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
668 );
669 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
670 set_cooldown_async(agent_name, cooldown_until, reason_str).await;
671 tracing::warn!(
672 agent = agent_name,
673 reason = reason_str,
674 attempt = count,
675 cooldown_secs,
676 "credit exhaustion detected: applying exponential agent-wide cooldown"
677 );
678}
679
680pub async fn record_credit_exhaustion_from_message(
682 agent_name: &str,
683 model: Option<&str>,
684 message: &str,
685) -> bool {
686 let Some(reason) = detect_credit_exhaustion(message) else {
687 return false;
688 };
689 match (reason, model) {
690 (CreditExhaustionReason::BillingCycleExhausted, Some(model)) => {
691 record_persistent_model_failure(agent_name, model).await;
692 }
693 _ => record_credit_exhaustion(agent_name, reason).await,
694 }
695 true
696}
697
698pub async fn record_model_failure(agent_name: &str, model: &str) {
702 let key = format!("{agent_name}:{model}");
703 let store_opt = cooldown_store().lock().await.clone();
704 let count = read_and_increment_failure_count(store_opt.as_ref(), &key).await;
705 let base = config::get_agent_backoff_base(agent_name);
706 let cooldown_until =
707 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
708 set_cooldown_async(&key, cooldown_until, "model_error").await;
709}
710
711pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
719 if let Some(window_secs) = parse_relative_usage_window(error_message)
720 && let Some(m) = model
721 {
722 let cooldown_until = chrono::Utc::now()
723 .timestamp()
724 .saturating_add(i64::try_from(window_secs).unwrap_or(i64::MAX));
725 set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
726 tracing::info!(
727 agent = agent_name,
728 model = m,
729 window_secs,
730 cooldown_until,
731 "relative usage window detected: applying model-specific cooldown"
732 );
733 return;
734 }
735
736 record_agent_failure_with_message(agent_name, error_message).await;
737 if let Some(m) = model {
738 record_model_failure(agent_name, m).await;
739 }
740}
741
742pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
754 let key = format!("{agent_name}:{model}");
755 let store_opt = cooldown_store().lock().await.clone();
756 let count = read_and_increment_failure_count_with_prefix(
757 store_opt.as_ref(),
758 PERSISTENT_FAILURE_COUNT_PREFIX,
759 &key,
760 )
761 .await;
762 let cooldown_secs = compute_backoff(
763 count,
764 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
765 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
766 false,
767 );
768 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
769 set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
770}
771
772pub async fn record_model_unavailable(agent_name: &str, model: &str, message: &str) {
777 if is_model_permanently_gone(message) {
778 record_persistent_model_failure(agent_name, model).await;
779 } else {
780 record_model_failure(agent_name, model).await;
781 }
782}
783
784fn is_model_permanently_gone(message: &str) -> bool {
785 let lower = message.to_lowercase();
786 lower.contains("not found")
787 || lower.contains("not supported")
788 || lower.contains("not available in your country")
789 || lower.contains("not available in your region")
790}
791
792pub async fn record_transient_model_failure(agent_name: &str, model: &str) {
798 let key = format!("{agent_name}:{model}");
799 let store_opt = cooldown_store().lock().await.clone();
800 let count = read_and_increment_failure_count(store_opt.as_ref(), &key).await;
801 let cooldown_secs = compute_backoff(
802 count,
803 TRANSIENT_MODEL_BACKOFF_BASE_SECS,
804 TRANSIENT_MODEL_BACKOFF_MAX_SECS,
805 false,
806 );
807 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
808 set_cooldown_async(&key, cooldown_until, "transient_model_error").await;
809}
810
811pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
818 let key = format!("{agent_name}:{model}");
819 let cooldown_until = chrono::Utc::now()
820 .timestamp()
821 .saturating_add(i64::try_from(duration_secs).unwrap_or(i64::MAX));
822 set_cooldown_async(&key, cooldown_until, reason).await;
823}
824
825pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
827 let key = format!("{agent_name}:{model}");
828 set_cooldown_async(&key, until, "vendor_reset_at").await;
829}
830
831pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
838 let cooldown_until = chrono::Utc::now()
839 .timestamp()
840 .saturating_add(i64::try_from(duration_secs).unwrap_or(i64::MAX));
841 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
842}
843
844pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
847 let store_opt = cooldown_store().lock().await.clone();
848 let Some(store) = store_opt else {
849 tracing::debug!(
850 agent = agent_name,
851 model,
852 "skipping silence count record (no KV store)"
853 );
854 return None;
855 };
856
857 let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
858 let now = chrono::Utc::now().timestamp();
859 let window_start = now - SILENCE_COUNT_WINDOW_SECS;
860
861 let mut timestamps = match store.kv_get(&key).await {
862 Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
863 Ok(ts) => ts,
864 Err(e) => {
865 tracing::warn!(
866 kv_key = %key,
867 err = %e,
868 "failed to parse silence timestamps from KV — resetting window"
869 );
870 Vec::new()
871 }
872 },
873 Ok(None) => Vec::new(),
874 Err(err) => {
875 tracing::warn!(
876 kv_key = %key,
877 err = %err,
878 "failed to load silence count from KV"
879 );
880 Vec::new()
881 }
882 };
883
884 timestamps.retain(|ts| *ts >= window_start);
885 timestamps.push(now);
886
887 let count = timestamps.len();
888 let mut extended_cooldown_applied = false;
889 if count >= SILENCE_COUNT_THRESHOLD {
890 set_model_cooldown(
891 agent_name,
892 model,
893 SILENCE_EXTENDED_COOLDOWN_SECS,
894 "silence_detected",
895 )
896 .await;
897 extended_cooldown_applied = true;
898 timestamps.clear();
899 }
900
901 let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
902 if let Err(err) = store.kv_set(&key, &value).await {
903 tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
904 }
905
906 Some(SilenceCountResult {
907 count,
908 extended_cooldown_applied,
909 })
910}
911
912pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
914 let key = format!("{agent_name}:{model}");
915 is_in_cooldown(&key)
916}
917
918pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
920 is_in_cooldown(agent_name)
921}
922
923pub fn cooldown_until(key: &str) -> Option<i64> {
925 let map = cooldown_lock();
926 map.get(key).and_then(|entry| {
927 let now = chrono::Utc::now().timestamp();
928 if now < entry.cooldown_until {
929 Some(entry.cooldown_until)
930 } else {
931 None
932 }
933 })
934}
935
936pub fn cooldown_reason(key: &str) -> Option<String> {
938 let map = cooldown_lock();
939 map.get(key).and_then(|entry| {
940 let now = chrono::Utc::now().timestamp();
941 if now < entry.cooldown_until {
942 Some(entry.reason.clone())
943 } else {
944 None
945 }
946 })
947}
948
949pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
954 let now = chrono::Utc::now().timestamp();
955 let map = cooldown_lock();
956 let mut result: Vec<(String, i64, String)> = map
957 .iter()
958 .filter(|(_, entry)| now < entry.cooldown_until)
959 .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
960 .collect();
961 result.sort_by(|a, b| a.0.cmp(&b.0));
962 result
963}
964
965pub async fn mark_rate_limits_cleared(store: &Arc<TaskStore>, agents: &[String]) {
968 let now = chrono::Utc::now().timestamp().to_string();
969 let mut seen = HashSet::new();
970 for key in agents {
971 let agent = agent_from_key(key);
972 if agent.is_empty() || !seen.insert(agent.to_string()) {
973 continue;
974 }
975 let kv_key = format!("{RATE_LIMITS_CLEARED_AT_PREFIX}{agent}");
976 if let Err(e) = store.kv_set(&kv_key, &now).await {
977 tracing::warn!(agent, err = %e, "failed to mark rate limits cleared");
978 }
979 }
980}
981
982pub async fn clear_cooldown(key: &str, store: &Arc<TaskStore>) {
989 let mut agents_to_mark: Vec<String> = Vec::new();
990 if key == "*" {
991 let keys: Vec<String> = {
992 let mut map = cooldown_lock();
993 let keys: Vec<String> = map.keys().cloned().collect();
994 map.clear();
995 keys
996 };
997 for k in &keys {
998 let kv_key = format!("{KV_PREFIX}{k}");
999 if let Err(e) = store.kv_set(&kv_key, "0").await {
1000 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
1001 }
1002 let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
1004 if let Err(e) = store.kv_set(&fc_key, "0").await {
1005 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
1006 }
1007 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
1009 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
1010 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
1011 }
1012 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
1014 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
1015 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
1016 }
1017 }
1018 if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
1022 for (fc_key, _) in fc_entries {
1023 if let Err(e) = store.kv_set(&fc_key, "0").await {
1024 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
1025 }
1026 }
1027 }
1028 if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
1029 for (cfc_key, _) in cfc_entries {
1030 if let Err(e) = store.kv_set(&cfc_key, "0").await {
1031 tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
1032 }
1033 }
1034 }
1035 if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
1036 for (pfc_key, _) in pfc_entries {
1037 if let Err(e) = store.kv_set(&pfc_key, "0").await {
1038 tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
1039 }
1040 }
1041 }
1042 tracing::info!(
1043 count = keys.len(),
1044 "cleared all cooldowns and failure counts"
1045 );
1046 agents_to_mark.extend(keys.iter().map(|k| agent_from_key(k).to_string()));
1047 match store.rate_limit_agents().await {
1048 Ok(rl_agents) => agents_to_mark.extend(rl_agents),
1049 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits agents while clearing"),
1050 }
1051 match store.kv_list_prefix(RATE_LIMITS_CLEARED_AT_PREFIX).await {
1052 Ok(marks) => {
1053 for (k, _) in marks {
1054 let agent = k.trim_start_matches(RATE_LIMITS_CLEARED_AT_PREFIX);
1055 if !agent.is_empty() {
1056 agents_to_mark.push(agent.to_string());
1057 }
1058 }
1059 }
1060 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits_cleared_at keys"),
1061 }
1062 } else {
1063 {
1064 let mut map = cooldown_lock();
1065 map.remove(key);
1066 }
1067 let kv_key = format!("{KV_PREFIX}{key}");
1068 if let Err(e) = store.kv_set(&kv_key, "0").await {
1069 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
1070 }
1071 let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
1073 if let Err(e) = store.kv_set(&fc_key, "0").await {
1074 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
1075 }
1076 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
1078 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
1079 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
1080 }
1081 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
1083 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
1084 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
1085 }
1086 tracing::info!(key, "cleared cooldown and failure count");
1087 agents_to_mark.push(agent_from_key(key).to_string());
1088 }
1089 mark_rate_limits_cleared(store, &agents_to_mark).await;
1090}
1091
1092fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
1101 let mut map = cooldown_lock();
1102 if let Some(existing) = map.get(key) {
1103 if existing.cooldown_until >= cooldown_until {
1107 return false;
1108 }
1109 }
1110 map.insert(
1111 key.to_string(),
1112 CooldownEntry {
1113 cooldown_until,
1114 reason: reason.to_string(),
1115 },
1116 );
1117 true
1118}
1119
1120async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
1130 if !set_cooldown_in_memory(key, cooldown_until, reason) {
1131 return false;
1132 }
1133 let store_opt = cooldown_store().lock().await.clone();
1134 if let Some(store) = store_opt {
1135 let kv_key = format!("{KV_PREFIX}{key}");
1136 let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
1137 if let Err(e) = store.kv_set(&kv_key, &value).await {
1138 tracing::warn!(
1139 kv_key,
1140 err = %e,
1141 "failed to persist cooldown to KV store — rolling back in-memory entry"
1142 );
1143 let mut map = cooldown_lock();
1144 if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1148 map.remove(key);
1149 }
1150 return false;
1151 }
1152 }
1153 true
1154}
1155
1156fn cooldown_lock() -> MutexGuard<'static, HashMap<String, CooldownEntry>> {
1158 cooldowns().lock().unwrap_or_else(|e| {
1159 tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1160 e.into_inner()
1161 })
1162}
1163
1164fn degraded_agents_lock() -> MutexGuard<'static, HashMap<String, String>> {
1166 degraded_agents().lock().unwrap_or_else(|e| {
1167 tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1168 e.into_inner()
1169 })
1170}
1171
1172fn github_5xx_timestamps_lock() -> MutexGuard<'static, Vec<i64>> {
1174 github_5xx_timestamps().lock().unwrap_or_else(|e| {
1175 tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1176 e.into_inner()
1177 })
1178}
1179
1180fn github_5xx_circuit_open_lock() -> MutexGuard<'static, Option<i64>> {
1182 github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1183 tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1184 e.into_inner()
1185 })
1186}
1187
1188fn is_in_cooldown_with_reason(key: &str, reason: &str) -> bool {
1190 let map = cooldown_lock();
1191 map.get(key).is_some_and(|entry| {
1192 entry.reason == reason && chrono::Utc::now().timestamp() < entry.cooldown_until
1193 })
1194}
1195
1196fn is_in_cooldown(key: &str) -> bool {
1198 let map = cooldown_lock();
1199 if let Some(entry) = map.get(key) {
1200 let now = chrono::Utc::now().timestamp();
1201 return now < entry.cooldown_until;
1202 }
1203 false
1204}
1205
1206pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1212
1213pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1217
1218fn degraded_agents() -> &'static Mutex<HashMap<String, String>> {
1221 static DEGRADED: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
1222 DEGRADED.get_or_init(|| Mutex::new(HashMap::new()))
1223}
1224
1225pub fn is_agent_degraded(agent: &str) -> bool {
1230 let set = degraded_agents_lock();
1231 set.contains_key(agent)
1232}
1233
1234pub fn agent_degraded_reason(agent: &str) -> Option<String> {
1236 degraded_agents_lock().get(agent).cloned()
1237}
1238
1239pub fn list_degraded_agents() -> Vec<(String, String)> {
1241 let mut entries: Vec<(String, String)> = degraded_agents_lock()
1242 .iter()
1243 .map(|(agent, reason)| (agent.clone(), reason.clone()))
1244 .collect();
1245 entries.sort_by(|a, b| a.0.cmp(&b.0));
1246 entries
1247}
1248
1249pub fn mark_agent_degraded(agent: &str) {
1251 mark_agent_degraded_with_reason(agent, "degraded");
1252}
1253
1254fn mark_agent_degraded_with_reason(agent: &str, reason: &str) {
1255 let mut set = degraded_agents_lock();
1256 set.insert(agent.to_string(), reason.to_string());
1257}
1258
1259pub fn clear_agent_degraded(agent: &str) {
1261 let mut set = degraded_agents_lock();
1262 set.remove(agent);
1263}
1264
1265pub async fn refresh_degraded_agents(
1279 store: &Arc<TaskStore>,
1280 available_agents: &[String],
1281 model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1282 window_hours: u32,
1283 threshold: i64,
1284) {
1285 clear_expired_cooldowns();
1289
1290 let start = chrono::Utc::now();
1293 let counts = match store.recent_rate_limit_counts(window_hours).await {
1294 Ok(c) => {
1295 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1296 tracing::info!(
1297 duration_ms = dur_ms,
1298 window_hours,
1299 "recent_rate_limit_counts query completed"
1300 );
1301 c
1302 }
1303 Err(e) => {
1304 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1305 tracing::warn!(
1306 err = %e,
1307 duration_ms = dur_ms,
1308 "failed to query recent rate limit counts for health check"
1309 );
1310 return;
1311 }
1312 };
1313
1314 for agent in available_agents {
1315 let in_cooldown = is_agent_in_cooldown(agent);
1316 let no_models = !model_checker(agent);
1317 let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1318 let over_threshold = rate_limit_count >= threshold;
1319
1320 let degraded = in_cooldown || no_models || over_threshold;
1321
1322 if degraded {
1323 let reason = if in_cooldown {
1324 "agent in cooldown"
1325 } else if no_models {
1326 "all models cooled"
1327 } else {
1328 "rate limit threshold exceeded"
1329 };
1330 let previous = agent_degraded_reason(agent);
1331 match previous.as_deref() {
1332 None => tracing::warn!(
1333 agent,
1334 reason,
1335 rate_limit_count,
1336 window_hours,
1337 threshold,
1338 "pre-emptive health check: marking agent as degraded"
1339 ),
1340 Some(prev) if prev != reason => tracing::warn!(
1341 agent,
1342 reason,
1343 previous = prev,
1344 rate_limit_count,
1345 window_hours,
1346 threshold,
1347 "pre-emptive health check: degraded reason changed"
1348 ),
1349 Some(_) => {}
1350 }
1351 mark_agent_degraded_with_reason(agent, reason);
1352 } else {
1353 if is_agent_degraded(agent) {
1354 tracing::info!(
1355 agent,
1356 rate_limit_count,
1357 "pre-emptive health check: agent recovered, clearing degraded flag"
1358 );
1359 }
1360 clear_agent_degraded(agent);
1361 }
1362 }
1363}
1364
1365pub fn clear_expired_cooldowns() {
1367 let mut map = cooldown_lock();
1368 let now = chrono::Utc::now().timestamp();
1369 map.retain(|_key, entry| now < entry.cooldown_until);
1370}
1371
1372const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1378
1379const GITHUB_5XX_THRESHOLD: usize = 5;
1381
1382const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1384
1385fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1387 static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1388 TS.get_or_init(|| Mutex::new(Vec::new()))
1389}
1390
1391fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1393 static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1394 OPEN.get_or_init(|| Mutex::new(None))
1395}
1396
1397pub async fn record_github_5xx() {
1402 let now = chrono::Utc::now().timestamp();
1403 let window_start = now - GITHUB_5XX_WINDOW_SECS;
1404
1405 let should_trip = {
1406 let mut ts = github_5xx_timestamps_lock();
1407 ts.retain(|t| *t >= window_start);
1408 ts.push(now);
1409 ts.len() >= GITHUB_5XX_THRESHOLD
1410 };
1411
1412 if should_trip {
1413 let cooldown_until = {
1414 let mut open = github_5xx_circuit_open_lock();
1415 if open.is_none() {
1416 let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1417 *open = Some(cd);
1418 tracing::warn!(
1419 errors_in_window = GITHUB_5XX_THRESHOLD,
1420 window_secs = GITHUB_5XX_WINDOW_SECS,
1421 cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1422 "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1423 );
1424 Some(cd)
1425 } else {
1426 None
1427 }
1428 };
1429 if let Some(cd) = cooldown_until {
1431 let store_opt = cooldown_store().lock().await.clone();
1432 if let Some(store) = store_opt {
1433 let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1434 }
1435 set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1439 }
1440 let mut ts = github_5xx_timestamps_lock();
1442 ts.clear();
1443 }
1444}
1445
1446pub fn is_github_circuit_open() -> bool {
1454 {
1457 let mut open = github_5xx_circuit_open_lock();
1458 if let Some(until) = *open {
1459 let now = chrono::Utc::now().timestamp();
1460 if now >= until {
1461 *open = None;
1462 tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1463 return false;
1464 }
1465 return true;
1466 }
1467 }
1468
1469 if let Some(until) = cooldown_until("github:5xx") {
1473 let mut open = github_5xx_circuit_open_lock();
1476 if open.is_none() {
1477 *open = Some(until);
1478 let now = chrono::Utc::now().timestamp();
1479 if now < until {
1480 let remaining = until - now;
1481 tracing::warn!(
1482 remaining_secs = remaining,
1483 "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1484 );
1485 }
1486 }
1487 return true;
1488 }
1489
1490 false
1491}
1492
1493pub fn github_circuit_remaining_secs() -> u64 {
1496 {
1498 let open = github_5xx_circuit_open_lock();
1499 if let Some(until) = *open {
1500 let now = chrono::Utc::now().timestamp();
1501 if now < until {
1502 return u64::try_from(until - now).unwrap_or(0);
1503 }
1504 return 0;
1505 }
1506 }
1507
1508 if let Some(until) = cooldown_until("github:5xx") {
1510 let now = chrono::Utc::now().timestamp();
1511 if now < until {
1512 return u64::try_from(until - now).unwrap_or(0);
1513 }
1514 }
1515 0
1516}
1517
1518pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1522 if error_message.is_empty() {
1523 return None;
1524 }
1525
1526 let lower = error_message.to_lowercase();
1527 for retry_marker in [
1531 "try again at ",
1532 "reset at ",
1533 "resets at ",
1534 "resets ",
1535 "reset ",
1536 ] {
1537 if let Some(idx) = lower.find(retry_marker) {
1538 let date_str = &lower[idx + retry_marker.len()..];
1539 let date_str = date_str
1541 .split(['.', '\n'])
1542 .next()
1543 .unwrap_or(date_str)
1544 .trim();
1545
1546 let cleaned = date_str
1549 .replace("st,", ",")
1550 .replace("nd,", ",")
1551 .replace("rd,", ",")
1552 .replace("th,", ",");
1553
1554 for fmt in &[
1555 "%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", ] {
1562 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1563 let utc = match chrono::Local.from_local_datetime(&dt) {
1566 chrono::LocalResult::Single(local_dt) => {
1567 local_dt.with_timezone(&chrono::Utc)
1568 }
1569 chrono::LocalResult::Ambiguous(earliest, _) => {
1570 earliest.with_timezone(&chrono::Utc)
1572 }
1573 chrono::LocalResult::None => {
1574 dt.and_utc()
1576 }
1577 };
1578 tracing::info!(
1579 retry_at = %utc,
1580 raw = date_str,
1581 marker = retry_marker,
1582 "parsed rate limit retry-at date"
1583 );
1584 return Some(utc.timestamp());
1585 }
1586 }
1587 tracing::debug!(
1588 raw = date_str,
1589 marker = retry_marker,
1590 "could not parse retry-at date"
1591 );
1592
1593 let time_only = date_str.split_once('(').map_or(date_str, |(t, _)| t).trim();
1598 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1600 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1601 let last = &compact[split..];
1602 if !last.contains(':')
1603 && (last.ends_with("am") || last.ends_with("pm"))
1604 && last.len() > 2
1605 {
1606 let at = split + last.len() - 2;
1607 compact.insert_str(at, ":00");
1608 }
1609 let md = compact.replace(" at ", " ");
1611 let now_local = chrono::Local::now();
1612 let this_year = chrono::Datelike::year(&now_local);
1613 for year in [this_year, this_year + 1] {
1614 let with_year = format!("{year} {md}");
1615 if let Ok(dt) =
1616 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1617 && let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest()
1618 && local_dt > now_local
1619 {
1620 return Some(local_dt.timestamp());
1621 }
1622 }
1623 let time_only = compact.as_str();
1624 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1625 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1626 let local_now = chrono::Local::now();
1629 let today = local_now.date_naive();
1630 let candidate = chrono::NaiveDateTime::new(today, t);
1631 let candidate_local: chrono::DateTime<chrono::Local> =
1634 match chrono::Local.from_local_datetime(&candidate) {
1635 chrono::LocalResult::Single(dt) => dt,
1636 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1637 chrono::LocalResult::None => {
1638 let utc_dt = candidate.and_utc();
1640 return Some(utc_dt.timestamp());
1641 }
1642 };
1643 let chosen = if candidate_local <= local_now {
1644 let tomorrow = today.succ_opt().unwrap_or(today);
1646 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1647 match chrono::Local.from_local_datetime(&ndt) {
1648 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1649 chrono::LocalResult::Ambiguous(earliest, _) => {
1650 earliest.with_timezone(&chrono::Utc)
1651 }
1652 chrono::LocalResult::None => ndt.and_utc(),
1653 }
1654 } else {
1655 candidate_local.with_timezone(&chrono::Utc)
1656 };
1657 tracing::info!(
1658 retry_at = %chosen,
1659 raw = date_str,
1660 marker = retry_marker,
1661 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1662 );
1663 return Some(chosen.timestamp());
1664 }
1665 }
1666 }
1667 }
1668
1669 None
1673}
1674
1675pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1686 static WINDOW_RE: LazyLock<Regex> = LazyLock::new(|| {
1687 Regex::new(
1688 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1689 )
1690 .expect("usage window regex is valid")
1691 });
1692
1693 if error_message.is_empty() {
1694 return None;
1695 }
1696
1697 let lower = error_message.to_lowercase();
1698 let has_context = lower.contains("usage limit")
1699 || lower.contains("quota")
1700 || lower.contains("window")
1701 || lower.contains("limit");
1702 if !has_context {
1703 return None;
1704 }
1705
1706 let mut max_secs: Option<u64> = None;
1707 for caps in WINDOW_RE.captures_iter(&lower) {
1708 let secs = if let Some(num) = caps.get(1) {
1709 let n: u64 = num.as_str().parse().ok()?;
1710 let unit = caps.get(2)?.as_str().to_lowercase();
1711 let multiplier = match unit.as_str() {
1712 "hour" | "hours" | "hr" | "hrs" => 3600,
1713 "minute" | "minutes" | "min" | "mins" => 60,
1714 "day" | "days" => 86400,
1715 _ => return None,
1716 };
1717 n.checked_mul(multiplier)?
1718 } else if let Some(word) = caps.get(3) {
1719 match word.as_str().to_lowercase().as_str() {
1720 "hourly" => 3600,
1721 "daily" => 86400,
1722 "weekly" => 604_800,
1723 "monthly" => 2_592_000,
1724 _ => return None,
1725 }
1726 } else {
1727 continue;
1728 };
1729 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1730 }
1731
1732 max_secs
1733}
1734
1735#[cfg(test)]
1736pub(crate) async fn reset_global_state() {
1737 cooldowns().lock().unwrap().clear();
1738 *cooldown_store().lock().await = None;
1739 github_5xx_timestamps().lock().unwrap().clear();
1740 *github_5xx_circuit_open().lock().unwrap() = None;
1741 degraded_agents().lock().unwrap().clear();
1742 opencode::reset_model_caches_for_test();
1743}
1744
1745#[cfg(test)]
1746mod tests {
1747 use std::time::Duration;
1748
1749 use super::*;
1750 use chrono::Datelike;
1751 use chrono::TimeZone;
1752 use serial_test::serial;
1753 use tokio::time;
1754
1755 async fn test_store() -> Arc<TaskStore> {
1756 Arc::new(TaskStore::open_memory().await.expect("in-memory store"))
1757 }
1758
1759 fn sqlite_datetime_offset(secs: i64) -> String {
1760 (chrono::Utc::now() + chrono::Duration::seconds(secs))
1761 .format("%Y-%m-%d %H:%M:%S")
1762 .to_string()
1763 }
1764
1765 #[serial(cooldown_state)]
1766 #[tokio::test]
1767 async fn init_loads_unexpired_cooldowns_from_kv() {
1768 super::reset_global_state().await;
1769 let store = test_store().await;
1770 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1773 .kv_set(
1774 &format!("{KV_PREFIX}kimi:k2p5"),
1775 &cooldown_until.to_string(),
1776 )
1777 .await
1778 .unwrap();
1779
1780 let expired = chrono::Utc::now().timestamp() - 1;
1782 store
1783 .kv_set(
1784 &format!("{KV_PREFIX}opencode:old-free"),
1785 &expired.to_string(),
1786 )
1787 .await
1788 .unwrap();
1789
1790 {
1791 let mut map = cooldown_lock();
1792 map.clear();
1793 }
1794
1795 init_cooldown_store(store).await;
1796
1797 assert!(
1798 is_model_in_cooldown("kimi", "k2p5"),
1799 "unexpired cooldown should be loaded"
1800 );
1801 assert!(
1802 !is_model_in_cooldown("opencode", "old-free"),
1803 "expired cooldown should not be loaded"
1804 );
1805 }
1806
1807 #[serial(cooldown_state)]
1808 #[tokio::test]
1809 async fn credit_exhaustion_inside_active_cooldown_keeps_counter_at_one() {
1810 super::reset_global_state().await;
1811 let store = test_store().await;
1812 *cooldown_store().lock().await = Some(Arc::clone(&store));
1813
1814 let msg =
1815 "API error (status 402 Payment Required): usage balance exhausted, out of credits";
1816 assert!(record_credit_exhaustion_from_message("credit_dup_agent", Some("m"), msg).await);
1817 assert!(record_credit_exhaustion_from_message("credit_dup_agent", Some("m"), msg).await);
1818
1819 assert!(is_agent_in_cooldown("credit_dup_agent"));
1820 let count = store
1821 .kv_get(&format!("{CREDIT_FAILURE_COUNT_PREFIX}credit_dup_agent"))
1822 .await
1823 .unwrap();
1824 assert_eq!(count.as_deref(), Some("1"));
1825 }
1826
1827 #[serial(cooldown_state)]
1828 #[tokio::test]
1829 async fn non_credit_message_is_not_recorded_as_credit_exhaustion() {
1830 super::reset_global_state().await;
1831 assert!(!record_credit_exhaustion_from_message("x_agent", Some("m"), "boom").await);
1832 assert!(!is_agent_in_cooldown("x_agent"));
1833 }
1834
1835 #[serial(cooldown_state)]
1836 #[tokio::test]
1837 async fn record_model_failure_writes_to_kv() {
1838 super::reset_global_state().await;
1839 let store = test_store().await;
1840
1841 *cooldown_store().lock().await = Some(Arc::clone(&store));
1842
1843 record_model_failure("testagent_persist", "testmodel_persist").await;
1844
1845 time::sleep(Duration::from_millis(50)).await;
1847
1848 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1849 let stored = store.kv_get(&kv_key).await.unwrap();
1850 assert!(stored.is_some(), "model failure should be persisted to KV");
1851 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1852 let now = chrono::Utc::now().timestamp();
1853 assert!(ts > now, "persisted timestamp should be in the future");
1855 assert_eq!(reason, "model_error");
1856 }
1857
1858 #[serial(cooldown_state)]
1859 #[tokio::test]
1860 async fn cooldown_reason_survives_restart_via_kv() {
1861 super::reset_global_state().await;
1862 let store = test_store().await;
1863 let agent = "test_reason_restart";
1864 let future = chrono::Utc::now().timestamp() + 600;
1865
1866 init_cooldown_store(Arc::clone(&store)).await;
1867 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1868
1869 {
1871 let mut map = cooldown_lock();
1872 map.clear();
1873 }
1874 init_cooldown_store(store).await;
1875
1876 let entries = list_all_cooldowns();
1877 let entry = entries.iter().find(|(key, ..)| key == agent);
1878 assert_eq!(
1879 entry.map(|(_, _, reason)| reason.as_str()),
1880 Some("billing_cycle_exhausted"),
1881 "cooldown reason must survive a restart, not collapse to 'persisted'"
1882 );
1883 }
1884
1885 #[serial(cooldown_state)]
1886 #[tokio::test]
1887 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1888 super::reset_global_state().await;
1889 let store = test_store().await;
1890 *cooldown_store().lock().await = Some(Arc::clone(&store));
1891
1892 let agent = "testagent_persistent";
1893 let model = "testmodel_persistent";
1894 let key = format!("{agent}:{model}");
1895
1896 record_persistent_model_failure(agent, model).await;
1898 let now = chrono::Utc::now().timestamp();
1899 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1900 let first_remaining = first_until.saturating_sub(now);
1901 assert!(
1902 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1903 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1904 );
1905
1906 record_persistent_model_failure(agent, model).await;
1908 let now2 = chrono::Utc::now().timestamp();
1909 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1910 let second_remaining = second_until.saturating_sub(now2);
1911 assert!(
1912 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1913 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1914 );
1915
1916 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1917 let stored_count = store
1918 .kv_get(&count_key)
1919 .await
1920 .expect("kv read should succeed")
1921 .expect("failure count should be written");
1922 assert_eq!(stored_count, "2");
1923 }
1924
1925 #[serial(cooldown_state)]
1926 #[tokio::test]
1927 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1928 super::reset_global_state().await;
1929 let store = test_store().await;
1930 *cooldown_store().lock().await = Some(Arc::clone(&store));
1931
1932 let agent = "testagent_persistent_success";
1933 let model = "testmodel_persistent_success";
1934 let key = format!("{agent}:{model}");
1935
1936 record_persistent_model_failure(agent, model).await;
1938 record_persistent_model_failure(agent, model).await;
1939
1940 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1941 let stored_count = store
1942 .kv_get(&count_key)
1943 .await
1944 .expect("kv read should succeed")
1945 .expect("failure count should be written");
1946 assert_eq!(stored_count, "2");
1947
1948 record_agent_success(agent, model).await;
1952
1953 let stored_count_after_success = store
1954 .kv_get(&count_key)
1955 .await
1956 .expect("kv read should succeed")
1957 .expect("persistent failure count should still exist");
1958 assert_eq!(
1959 stored_count_after_success, "2",
1960 "persistent failure count must survive an intervening success"
1961 );
1962
1963 record_persistent_model_failure(agent, model).await;
1965 let now = chrono::Utc::now().timestamp();
1966 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1967 let third_remaining = third_until.saturating_sub(now);
1968 assert!(
1969 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1970 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1971 );
1972 }
1973
1974 #[serial(cooldown_state)]
1975 #[tokio::test]
1976 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1977 super::reset_global_state().await;
1978 let store = test_store().await;
1979 *cooldown_store().lock().await = Some(Arc::clone(&store));
1980
1981 let agent = "testagent_persistent_capped";
1982 let model = "testmodel_persistent_capped";
1983 let key = format!("{agent}:{model}");
1984
1985 for _ in 0..9 {
1986 record_persistent_model_failure(agent, model).await;
1987 }
1988
1989 let remaining = cooldown_until(&key)
1990 .expect("cooldown should exist")
1991 .saturating_sub(chrono::Utc::now().timestamp());
1992 assert!(
1993 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1994 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1995 );
1996 assert!(
1997 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1998 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1999 );
2000 }
2001
2002 #[serial(cooldown_state)]
2003 #[test]
2004 fn parse_retry_at_codex_format() {
2005 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.";
2006 let ts = parse_retry_at(msg);
2007 assert!(ts.is_some(), "should parse codex retry-at date");
2008
2009 let naive = chrono::NaiveDateTime::new(
2012 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
2013 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
2014 );
2015 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
2016 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
2017 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
2018 chrono::LocalResult::None => naive.and_utc(),
2019 };
2020 assert_eq!(
2021 ts.unwrap(),
2022 expected_utc.timestamp(),
2023 "timestamp should reflect local-timezone interpretation of the vendor date"
2024 );
2025
2026 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
2028 assert_eq!(dt.month(), 3);
2029 assert_eq!(dt.day(), 26);
2030 }
2031
2032 #[serial(cooldown_state)]
2033 #[test]
2034 fn parse_retry_at_hour_only_and_month_day() {
2035 let now = chrono::Utc::now().timestamp();
2036 for msg in [
2037 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
2038 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
2039 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
2040 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
2041 ] {
2042 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
2043 assert!(ts > now, "{msg}");
2044 }
2045 }
2046
2047 #[serial(cooldown_state)]
2048 #[test]
2049 fn parse_retry_at_weekly_resets_without_year() {
2050 let t = chrono::Local::now() + chrono::Duration::days(2);
2051 let msg = format!(
2052 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
2053 t.format("%b"),
2054 t.day()
2055 );
2056 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
2057 let remaining = ts - chrono::Utc::now().timestamp();
2058 assert!(
2059 remaining > 3600 && remaining < 3 * 86400,
2060 "expected ~2d retry-at, got {remaining}s"
2061 );
2062 }
2063
2064 #[serial(cooldown_state)]
2065 #[test]
2066 fn parse_retry_at_no_date() {
2067 assert!(parse_retry_at("generic rate limit error").is_none());
2068 assert!(parse_retry_at("").is_none());
2069 }
2070
2071 #[serial(cooldown_state)]
2072 #[test]
2073 fn parse_retry_at_glm_reset_format_with_seconds() {
2074 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
2075 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
2076
2077 let naive = chrono::NaiveDateTime::new(
2078 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
2079 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
2080 );
2081 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
2082 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
2083 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
2084 chrono::LocalResult::None => naive.and_utc(),
2085 };
2086
2087 assert_eq!(
2088 ts,
2089 expected_utc.timestamp(),
2090 "timestamp should reflect local-timezone interpretation of vendor reset date"
2091 );
2092 }
2093
2094 #[serial(cooldown_state)]
2095 #[test]
2096 fn parse_retry_at_billing_cycle_without_date_returns_none() {
2097 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.";
2100 assert!(
2101 parse_retry_at(msg).is_none(),
2102 "billing cycle without date should return None"
2103 );
2104 }
2105
2106 #[serial(cooldown_state)]
2107 #[test]
2108 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
2109 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.";
2110 assert_eq!(
2111 parse_relative_usage_window(msg),
2112 Some(5 * 3600),
2113 "should extract the 5-hour window as seconds"
2114 );
2115 }
2116
2117 #[serial(cooldown_state)]
2118 #[test]
2119 fn parse_relative_usage_window_handles_common_durations() {
2120 assert_eq!(
2121 parse_relative_usage_window("You've reached your hourly usage limit."),
2122 Some(3600)
2123 );
2124 assert_eq!(
2125 parse_relative_usage_window("You've reached your daily usage limit."),
2126 Some(86400)
2127 );
2128 assert_eq!(
2129 parse_relative_usage_window("You've reached your weekly limit."),
2130 Some(604_800)
2131 );
2132 assert_eq!(
2133 parse_relative_usage_window("You've reached your 30-minute usage limit."),
2134 Some(1800)
2135 );
2136 assert_eq!(
2137 parse_relative_usage_window("You've reached your 2-day quota."),
2138 Some(172_800)
2139 );
2140 }
2141
2142 #[serial(cooldown_state)]
2143 #[test]
2144 fn parse_relative_usage_window_returns_none_without_usage_context() {
2145 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
2148 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
2149 assert!(parse_relative_usage_window("").is_none());
2150 }
2151
2152 #[serial(cooldown_state)]
2153 #[tokio::test]
2154 async fn record_rate_limit_uses_model_specific_window_cooldown() {
2155 super::reset_global_state().await;
2156 let store = test_store().await;
2157 *cooldown_store().lock().await = Some(Arc::clone(&store));
2158
2159 let agent = "test_record_rl_window_agent";
2160 let model = "test_record_rl_window_model";
2161 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
2162
2163 record_rate_limit(agent, Some(model), msg).await;
2164
2165 assert!(
2166 is_model_in_cooldown(agent, model),
2167 "model should be cooled for the relative usage window"
2168 );
2169 assert!(
2170 !is_agent_in_cooldown(agent),
2171 "agent should not be cooled for a model-scoped quota"
2172 );
2173
2174 let key = format!("{agent}:{model}");
2175 let remaining = cooldown_until(&key)
2176 .expect("model cooldown should exist")
2177 .saturating_sub(chrono::Utc::now().timestamp());
2178 assert!(
2179 remaining >= 5 * 3600 - 5,
2180 "model cooldown should be at least ~5 hours, got {remaining}s"
2181 );
2182 assert_eq!(
2183 cooldown_reason(&key).as_deref(),
2184 Some("rate_limit_usage_window"),
2185 "relative usage-window cooldown must be labeled as a rate limit, not silence"
2186 );
2187 }
2188
2189 #[serial(cooldown_state)]
2190 #[tokio::test]
2191 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
2192 super::reset_global_state().await;
2193 let store = test_store().await;
2194 *cooldown_store().lock().await = Some(Arc::clone(&store));
2195
2196 let agent = "test_record_rl_fallback_agent";
2197 let model = "test_record_rl_fallback_model";
2198 let msg = "429 Too Many Requests";
2199
2200 record_rate_limit(agent, Some(model), msg).await;
2201
2202 assert!(
2203 is_agent_in_cooldown(agent),
2204 "agent should be cooled when no relative window is detected"
2205 );
2206 assert!(
2207 is_model_in_cooldown(agent, model),
2208 "model should be cooled when no relative window is detected"
2209 );
2210 }
2211
2212 #[serial(cooldown_state)]
2213 #[tokio::test]
2214 async fn agent_cooldown_persists_to_kv() {
2215 super::reset_global_state().await;
2216 let agent = "test_agent_persist_check";
2218 record_agent_failure_with_message(agent, "").await;
2219 assert!(is_agent_in_cooldown(agent));
2220 }
2221
2222 #[serial(cooldown_state)]
2223 #[tokio::test]
2224 async fn cooldown_never_shortens() {
2225 super::reset_global_state().await;
2226 let agent = "test_agent_no_shorten";
2227 set_agent_cooldown(agent, 24 * 60 * 60).await;
2228 let initial = {
2229 let map = cooldowns().lock().unwrap();
2230 map.get(agent)
2231 .expect("cooldown entry should exist")
2232 .cooldown_until
2233 };
2234
2235 record_agent_failure_with_message(agent, "rate limit").await;
2237 let after = {
2238 let map = cooldowns().lock().unwrap();
2239 map.get(agent)
2240 .expect("cooldown entry should exist")
2241 .cooldown_until
2242 };
2243 assert_eq!(
2244 initial, after,
2245 "cooldown should not be shortened by later failures"
2246 );
2247 }
2248
2249 #[serial(cooldown_state)]
2250 #[tokio::test]
2251 async fn silence_agent_cooldown_is_short_lived() {
2252 super::reset_global_state().await;
2253 let agent = "test_silence_agent_cd";
2254 assert!(!is_agent_in_cooldown(agent));
2255
2256 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2257 assert!(is_agent_in_cooldown(agent));
2258
2259 let map = cooldowns().lock().unwrap();
2261 let entry = map.get(agent).expect("should have cooldown entry");
2262 let now = chrono::Utc::now().timestamp();
2263 let remaining = entry.cooldown_until - now;
2264 assert!(
2265 remaining <= i64::try_from(SILENCE_AGENT_COOLDOWN_SECS).unwrap(),
2266 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2267 );
2268 assert!(remaining > 0, "cooldown should still be active");
2269 }
2270
2271 #[serial(cooldown_state)]
2272 #[tokio::test]
2273 async fn record_silence_detection_applies_extended_cooldown() {
2274 super::reset_global_state().await;
2275 let store = test_store().await;
2276 *cooldown_store().lock().await = Some(Arc::clone(&store));
2277
2278 let agent = "test_silence_count_agent";
2279 let model = "test_silence_count_model";
2280 assert!(!is_model_in_cooldown(agent, model));
2281
2282 for _ in 0..SILENCE_COUNT_THRESHOLD {
2283 let result = record_silence_detection(agent, model).await;
2284 assert!(result.is_some());
2285 }
2286
2287 assert!(is_model_in_cooldown(agent, model));
2288 let key = format!("{agent}:{model}");
2289 let remaining = {
2290 let map = cooldowns().lock().unwrap();
2291 let entry = map.get(&key).expect("cooldown entry should exist");
2292 entry.cooldown_until - chrono::Utc::now().timestamp()
2293 };
2294 assert!(
2295 remaining >= i64::try_from(SILENCE_EXTENDED_COOLDOWN_SECS).unwrap() - 5,
2296 "extended cooldown should be applied, got {remaining}s"
2297 );
2298
2299 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2300 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2301 assert_eq!(stored, "[]");
2302 }
2303
2304 #[serial(cooldown_state)]
2305 #[test]
2306 fn detect_credit_exhaustion_out_of_credits() {
2307 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2308 assert_eq!(
2309 detect_credit_exhaustion(msg),
2310 Some(CreditExhaustionReason::OutOfCredits)
2311 );
2312 }
2313
2314 #[serial(cooldown_state)]
2315 #[test]
2316 fn detect_credit_exhaustion_out_of_credits_variations() {
2317 assert_eq!(
2318 detect_credit_exhaustion("out of credits"),
2319 Some(CreditExhaustionReason::OutOfCredits)
2320 );
2321 assert_eq!(
2322 detect_credit_exhaustion("outofcredits"),
2323 Some(CreditExhaustionReason::OutOfCredits)
2324 );
2325 assert_eq!(
2326 detect_credit_exhaustion("no credits remaining"),
2327 Some(CreditExhaustionReason::OutOfCredits)
2328 );
2329 assert_eq!(
2330 detect_credit_exhaustion("insufficient funds"),
2331 Some(CreditExhaustionReason::OutOfCredits)
2332 );
2333 }
2334
2335 #[serial(cooldown_state)]
2336 #[test]
2337 fn detect_credit_exhaustion_org_level_disabled() {
2338 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2339 assert_eq!(
2340 detect_credit_exhaustion(msg),
2341 Some(CreditExhaustionReason::OrgLevelDisabled)
2342 );
2343 }
2344
2345 #[serial(cooldown_state)]
2346 #[test]
2347 fn detect_credit_exhaustion_org_level_disabled_variations() {
2348 assert_eq!(
2349 detect_credit_exhaustion("organization has been disabled"),
2350 Some(CreditExhaustionReason::OrgLevelDisabled)
2351 );
2352 assert_eq!(
2353 detect_credit_exhaustion("org disabled"),
2354 Some(CreditExhaustionReason::OrgLevelDisabled)
2355 );
2356 assert_eq!(
2357 detect_credit_exhaustion("billing quota exceeded"),
2358 Some(CreditExhaustionReason::OrgLevelDisabled)
2359 );
2360 assert_eq!(
2361 detect_credit_exhaustion("credit balance too low"),
2362 Some(CreditExhaustionReason::OrgLevelDisabled)
2363 );
2364 }
2365
2366 #[serial(cooldown_state)]
2367 #[test]
2368 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2369 let msg = "overageDisabledReason: out_of_credits";
2370 assert_eq!(
2371 detect_credit_exhaustion(msg),
2372 Some(CreditExhaustionReason::OutOfCredits)
2373 );
2374
2375 let msg = "overageDisabledReason: org_level_disabled";
2376 assert_eq!(
2377 detect_credit_exhaustion(msg),
2378 Some(CreditExhaustionReason::OrgLevelDisabled)
2379 );
2380 }
2381
2382 #[serial(cooldown_state)]
2383 #[test]
2384 fn detect_credit_exhaustion_glm_insufficient_balance() {
2385 assert_eq!(
2387 detect_credit_exhaustion(
2388 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2389 ),
2390 Some(CreditExhaustionReason::OutOfCredits)
2391 );
2392 assert_eq!(
2393 detect_credit_exhaustion("Insufficient balance"),
2394 Some(CreditExhaustionReason::OutOfCredits)
2395 );
2396 assert_eq!(
2397 detect_credit_exhaustion("please recharge your account"),
2398 Some(CreditExhaustionReason::OutOfCredits)
2399 );
2400 assert_eq!(
2401 detect_credit_exhaustion("no resource package available"),
2402 Some(CreditExhaustionReason::OutOfCredits)
2403 );
2404 }
2405
2406 #[serial(cooldown_state)]
2407 #[test]
2408 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2409 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2410 assert!(detect_credit_exhaustion("too many requests").is_none());
2411 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2412 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2413 assert!(
2414 detect_credit_exhaustion("You've hit your usage limit, try again in a few minutes.")
2415 .is_none()
2416 );
2417 }
2418
2419 #[serial(cooldown_state)]
2420 #[tokio::test]
2421 async fn record_credit_exhaustion_applies_agent_cooldown() {
2422 super::reset_global_state().await;
2423 let agent = "test_credit_exhaust_agent";
2424 assert!(!is_agent_in_cooldown(agent));
2425
2426 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2427 assert!(is_agent_in_cooldown(agent));
2428
2429 let remaining = {
2430 let map = cooldowns().lock().unwrap();
2431 let entry = map.get(agent).expect("cooldown entry should exist");
2432 entry.cooldown_until - chrono::Utc::now().timestamp()
2433 };
2434 assert!(
2436 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2437 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2438 );
2439 assert!(
2440 remaining <= CREDIT_BACKOFF_MAX_SECS,
2441 "cooldown should not exceed cap of 8h, got {remaining}s"
2442 );
2443 }
2444
2445 #[serial(cooldown_state)]
2446 #[tokio::test]
2447 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2448 super::reset_global_state().await;
2449 let agent = "test_org_disabled_agent";
2450 assert!(!is_agent_in_cooldown(agent));
2451
2452 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2453 assert!(is_agent_in_cooldown(agent));
2454
2455 let remaining = {
2456 let map = cooldowns().lock().unwrap();
2457 let entry = map.get(agent).expect("cooldown entry should exist");
2458 entry.cooldown_until - chrono::Utc::now().timestamp()
2459 };
2460 assert!(
2462 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2463 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2464 );
2465 assert!(
2466 remaining <= CREDIT_BACKOFF_MAX_SECS,
2467 "cooldown should not exceed cap of 8h, got {remaining}s"
2468 );
2469 }
2470
2471 #[serial(cooldown_state)]
2472 #[tokio::test]
2473 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2474 super::reset_global_state().await;
2475 let store = test_store().await;
2476 *cooldown_store().lock().await = Some(Arc::clone(&store));
2477
2478 let agent = "test_billing_cycle_agent";
2479 assert!(!is_agent_in_cooldown(agent));
2480
2481 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2482 assert!(is_agent_in_cooldown(agent));
2483
2484 let remaining = {
2485 let map = cooldowns().lock().unwrap();
2486 let entry = map.get(agent).expect("cooldown entry should exist");
2487 entry.cooldown_until - chrono::Utc::now().timestamp()
2488 };
2489 assert!(
2491 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2492 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2493 );
2494 assert!(
2495 remaining <= BILLING_CYCLE_MAX_SECS,
2496 "cooldown should not exceed cap of 7 days, got {remaining}s"
2497 );
2498 }
2499
2500 #[serial(cooldown_state)]
2501 #[tokio::test]
2502 async fn record_credit_exhaustion_billing_cycle_escalates() {
2503 super::reset_global_state().await;
2504 let store = test_store().await;
2505 *cooldown_store().lock().await = Some(Arc::clone(&store));
2506
2507 let agent = "test_billing_cycle_escalation";
2508
2509 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2511 let remaining_1 = {
2512 let map = cooldowns().lock().unwrap();
2513 let entry = map.get(agent).expect("cooldown entry should exist");
2514 entry.cooldown_until - chrono::Utc::now().timestamp()
2515 };
2516 assert!(
2517 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2518 .contains(&remaining_1),
2519 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2520 );
2521
2522 {
2524 let mut map = cooldowns().lock().unwrap();
2525 map.remove(agent);
2526 }
2527
2528 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2530 let remaining_2 = {
2531 let map = cooldowns().lock().unwrap();
2532 let entry = map.get(agent).expect("cooldown entry should exist");
2533 entry.cooldown_until - chrono::Utc::now().timestamp()
2534 };
2535 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2537 remaining_2 >= expected_2 - 5,
2538 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2539 );
2540
2541 {
2543 let mut map = cooldowns().lock().unwrap();
2544 map.remove(agent);
2545 }
2546
2547 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2549 let remaining_3 = {
2550 let map = cooldowns().lock().unwrap();
2551 let entry = map.get(agent).expect("cooldown entry should exist");
2552 entry.cooldown_until - chrono::Utc::now().timestamp()
2553 };
2554 assert!(
2555 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2556 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2557 );
2558 assert!(
2559 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2560 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2561 );
2562 }
2563
2564 #[serial(cooldown_state)]
2565 #[tokio::test]
2566 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2567 super::reset_global_state().await;
2568 let store = test_store().await;
2569 *cooldown_store().lock().await = Some(Arc::clone(&store));
2570
2571 let agent = "test_billing_cycle_agent_capped";
2572
2573 for i in 0..5 {
2574 if i > 0 {
2575 cooldowns().lock().unwrap().remove(agent);
2577 }
2578 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2579 }
2580
2581 let remaining = {
2582 let map = cooldowns().lock().unwrap();
2583 let entry = map.get(agent).expect("cooldown entry should exist");
2584 entry.cooldown_until - chrono::Utc::now().timestamp()
2585 };
2586 assert!(
2587 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2588 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2589 );
2590 assert!(
2591 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2592 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2593 );
2594 }
2595
2596 #[serial(cooldown_state)]
2597 #[test]
2598 fn credit_exhaustion_backoff_constants() {
2599 assert_eq!(
2601 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2602 CREDIT_BACKOFF_BASE_SECS
2603 );
2604 assert_eq!(
2606 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2607 ORG_BACKOFF_BASE_SECS
2608 );
2609 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2611 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2612 }
2613
2614 #[serial(cooldown_state)]
2615 #[test]
2616 fn compute_backoff_long_cap_without_extended_stays_capped() {
2617 assert_eq!(
2618 compute_backoff(
2619 3,
2620 BILLING_CYCLE_COOLDOWN_SECS,
2621 BILLING_CYCLE_MAX_SECS,
2622 false
2623 ),
2624 BILLING_CYCLE_MAX_SECS,
2625 "7-day capped paths should stop at the documented max"
2626 );
2627 assert_eq!(
2628 compute_backoff(
2629 9,
2630 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2631 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2632 false,
2633 ),
2634 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2635 "persistent model failures should not extend past 7 days"
2636 );
2637 }
2638
2639 #[serial(cooldown_state)]
2642 #[test]
2643 fn mark_and_check_agent_degraded() {
2644 let agent = "test_degraded_mark";
2645 assert!(!is_agent_degraded(agent));
2646 mark_agent_degraded(agent);
2647 assert!(is_agent_degraded(agent));
2648 clear_agent_degraded(agent);
2649 assert!(!is_agent_degraded(agent));
2650 }
2651
2652 #[serial(cooldown_state)]
2653 #[tokio::test]
2654 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2655 super::reset_global_state().await;
2656 let store = test_store().await;
2657 let agent = "test_sync_external_clear";
2658 let future = chrono::Utc::now().timestamp() + 600;
2659
2660 init_cooldown_store(Arc::clone(&store)).await;
2662 set_cooldown_async(agent, future, "test").await;
2663 assert!(is_agent_in_cooldown(agent));
2664
2665 store
2669 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2670 .await
2671 .expect("kv_set tombstone");
2672 assert!(
2673 is_agent_in_cooldown(agent),
2674 "in-memory cooldown should still be present before sync"
2675 );
2676
2677 sync_from_kv(&store).await;
2678
2679 assert!(
2680 !is_agent_in_cooldown(agent),
2681 "sync_from_kv must drop in-memory entries the KV tombstoned"
2682 );
2683 }
2684
2685 #[serial(cooldown_state)]
2686 #[tokio::test]
2687 async fn sync_from_kv_closes_github_circuit_when_row_cleared() {
2688 super::reset_global_state().await;
2689 let store = test_store().await;
2690 {
2691 *github_5xx_circuit_open_lock() = Some(chrono::Utc::now().timestamp() + 600);
2692 github_5xx_timestamps_lock().push(chrono::Utc::now().timestamp());
2693 }
2694 assert!(is_github_circuit_open());
2695
2696 store
2697 .kv_set("cooldown:github:5xx", "0")
2698 .await
2699 .expect("kv_set tombstone");
2700 sync_from_kv(&store).await;
2701
2702 assert!(!is_github_circuit_open());
2703 assert!(github_5xx_timestamps_lock().is_empty());
2704 }
2705
2706 #[serial(cooldown_state)]
2707 #[tokio::test]
2708 async fn sync_from_kv_keeps_5xx_window_when_expired_row_and_breaker_closed() {
2709 super::reset_global_state().await;
2710 let store = test_store().await;
2711 let now = chrono::Utc::now().timestamp();
2712 {
2713 github_5xx_timestamps_lock().push(now);
2714 github_5xx_timestamps_lock().push(now);
2715 }
2716
2717 store
2718 .kv_set("cooldown:github:5xx", &(now - 600).to_string())
2719 .await
2720 .expect("kv_set past");
2721 sync_from_kv(&store).await;
2722
2723 assert!(!is_github_circuit_open());
2724 assert_eq!(github_5xx_timestamps_lock().len(), 2);
2725
2726 github_5xx_timestamps_lock().clear();
2727 }
2728
2729 #[serial(cooldown_state)]
2730 #[tokio::test]
2731 async fn sync_from_kv_keeps_github_circuit_open_for_future_row() {
2732 super::reset_global_state().await;
2733 let store = test_store().await;
2734 let future = chrono::Utc::now().timestamp() + 600;
2735 *github_5xx_circuit_open_lock() = Some(future);
2736
2737 store
2738 .kv_set("cooldown:github:5xx", &future.to_string())
2739 .await
2740 .expect("kv_set future");
2741 sync_from_kv(&store).await;
2742
2743 assert!(is_github_circuit_open());
2744
2745 *github_5xx_circuit_open_lock() = None;
2746 }
2747
2748 #[serial(cooldown_state)]
2749 #[tokio::test]
2750 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2751 super::reset_global_state().await;
2755 let store = test_store().await;
2756 let agent = "test_sync_race_window";
2757 let future = chrono::Utc::now().timestamp() + 600;
2758
2759 set_cooldown_in_memory(agent, future, "raced");
2761 assert!(is_agent_in_cooldown(agent));
2762
2763 sync_from_kv(&store).await;
2764
2765 assert!(
2766 is_agent_in_cooldown(agent),
2767 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2768 );
2769 }
2770
2771 #[serial(cooldown_state)]
2772 #[tokio::test]
2773 async fn sync_from_kv_adopts_kv_timestamp() {
2774 super::reset_global_state().await;
2777 let store = test_store().await;
2778 let agent = "test_sync_adopt_kv";
2779 let original = chrono::Utc::now().timestamp() + 600;
2780 let updated = chrono::Utc::now().timestamp() + 1200;
2781
2782 init_cooldown_store(Arc::clone(&store)).await;
2783 set_cooldown_async(agent, original, "initial").await;
2784
2785 store
2786 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2787 .await
2788 .expect("kv_set updated");
2789
2790 sync_from_kv(&store).await;
2791
2792 let until = cooldown_until(agent).expect("cooldown still active");
2793 assert_eq!(
2794 until, updated,
2795 "sync_from_kv must adopt KV's authoritative timestamp"
2796 );
2797 }
2798
2799 #[serial(cooldown_state)]
2800 #[tokio::test]
2801 async fn refresh_degraded_agents_marks_cooled_agent() {
2802 super::reset_global_state().await;
2803 let store = test_store().await;
2804 let agent = "test_refresh_cooled";
2805
2806 record_agent_failure_with_message(agent, "").await;
2808 assert!(is_agent_in_cooldown(agent));
2809
2810 let agents = vec![agent.to_string()];
2812 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2813
2814 assert!(
2815 is_agent_degraded(agent),
2816 "agent in cooldown should be marked degraded"
2817 );
2818
2819 clear_agent_degraded(agent);
2821 }
2822
2823 #[serial(cooldown_state)]
2824 #[tokio::test]
2825 async fn refresh_degraded_agents_marks_no_models_agent() {
2826 super::reset_global_state().await;
2827 let store = test_store().await;
2828 let agent = "test_refresh_no_models";
2829
2830 let agents = vec![agent.to_string()];
2832 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2833
2834 assert!(
2835 is_agent_degraded(agent),
2836 "agent with no available models should be marked degraded"
2837 );
2838
2839 clear_agent_degraded(agent);
2841 }
2842
2843 #[serial(cooldown_state)]
2844 #[tokio::test]
2845 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2846 super::reset_global_state().await;
2847 let store = test_store().await;
2848 let agent = "test_refresh_rate_limited";
2849
2850 for _ in 0..4 {
2852 store.record_rate_limit(agent, "rate", None).await.unwrap();
2853 }
2854
2855 let agents = vec![agent.to_string()];
2856 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2858
2859 assert!(
2860 is_agent_degraded(agent),
2861 "agent with rate limit count >= threshold should be marked degraded"
2862 );
2863
2864 clear_agent_degraded(agent);
2866 }
2867
2868 #[serial(cooldown_state)]
2869 #[tokio::test]
2870 async fn refresh_degraded_agents_clears_healthy_agent() {
2871 super::reset_global_state().await;
2872 let store = test_store().await;
2873 let agent = "test_refresh_healthy";
2874
2875 mark_agent_degraded(agent);
2877 assert!(is_agent_degraded(agent));
2878
2879 let agents = vec![agent.to_string()];
2881 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2882
2883 assert!(
2884 !is_agent_degraded(agent),
2885 "healthy agent should have degraded flag cleared"
2886 );
2887 }
2888
2889 #[serial(cooldown_state)]
2890 #[tokio::test]
2891 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2892 super::reset_global_state().await;
2893 let store = test_store().await;
2894 let agent = "test_refresh_expired_cooldown";
2895
2896 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2900 let mut map = cooldowns().lock().unwrap();
2901 map.insert(
2902 agent.to_string(),
2903 CooldownEntry {
2904 cooldown_until: expired_ts,
2905 reason: "expired_test".to_string(),
2906 },
2907 );
2908 }
2909
2910 assert!(
2912 !is_agent_in_cooldown(agent),
2913 "expired cooldown should not be active"
2914 );
2915
2916 mark_agent_degraded(agent);
2918 assert!(is_agent_degraded(agent));
2919
2920 let agents = vec![agent.to_string()];
2922 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2923
2924 assert!(
2925 !is_agent_degraded(agent),
2926 "agent with only expired cooldown should NOT be marked degraded"
2927 );
2928
2929 clear_agent_degraded(agent);
2931 {
2932 let mut map = cooldowns().lock().unwrap();
2933 map.remove(agent);
2934 }
2935 }
2936
2937 #[serial(cooldown_state)]
2938 #[test]
2939 fn compute_backoff_grows_exponentially() {
2940 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2942 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2943 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2945 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2947 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2949 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2951 }
2952
2953 #[serial(cooldown_state)]
2954 #[test]
2955 fn compute_backoff_extended_tier_1_applies_24h() {
2956 assert_eq!(
2962 compute_backoff(5, 300, 14400, true),
2963 14400,
2964 "first at max, not extended"
2965 );
2966 assert_eq!(
2967 compute_backoff(6, 300, 14400, true),
2968 14400,
2969 "times_at_max=1, not extended"
2970 );
2971 assert_eq!(
2972 compute_backoff(7, 300, 14400, true),
2973 86400,
2974 "times_at_max=2, extended 6x"
2975 );
2976 assert_eq!(
2977 compute_backoff(8, 300, 14400, true),
2978 86400,
2979 "times_at_max=3, still extended 6x"
2980 );
2981 }
2982
2983 #[serial(cooldown_state)]
2984 #[test]
2985 fn compute_backoff_extended_tier_2_applies_48h() {
2986 assert_eq!(
2989 compute_backoff(9, 300, 14400, true),
2990 86400,
2991 "times_at_max=4, still tier 1"
2992 );
2993 assert_eq!(
2994 compute_backoff(10, 300, 14400, true),
2995 172_800,
2996 "times_at_max=5, extended 12x"
2997 );
2998 assert_eq!(
2999 compute_backoff(15, 300, 14400, true),
3000 172_800,
3001 "times_at_max=10, still tier 2"
3002 );
3003 assert_eq!(
3004 compute_backoff(20, 300, 14400, true),
3005 172_800,
3006 "times_at_max=15, still tier 2"
3007 );
3008 }
3009
3010 #[serial(cooldown_state)]
3011 #[test]
3012 fn compute_backoff_extended_tier_with_different_base_max() {
3013 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
3020 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
3021 assert_eq!(
3022 compute_backoff(7, 60, 3600, true),
3023 21600,
3024 "times_at_max=2, extended 6x"
3025 );
3026 }
3027
3028 #[serial(cooldown_state)]
3029 #[test]
3030 fn compute_backoff_extended_tier_persistent_fails_example() {
3031 assert_eq!(
3035 compute_backoff(21, 300, 14400, true),
3036 172_800,
3037 "21 failures should get 48h extended cooldown"
3038 );
3039
3040 assert_eq!(
3042 compute_backoff(36, 300, 14400, true),
3043 172_800,
3044 "36 failures should also get 48h extended cooldown"
3045 );
3046 }
3047
3048 #[serial(cooldown_state)]
3049 #[test]
3050 fn compute_backoff_extended_tier_zero_base() {
3051 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3053 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3054 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
3055 }
3056
3057 #[serial(cooldown_state)]
3058 #[test]
3059 fn detect_credit_exhaustion_billing_cycle() {
3060 assert_eq!(
3061 detect_credit_exhaustion("billing cycle"),
3062 Some(CreditExhaustionReason::BillingCycleExhausted)
3063 );
3064 assert_eq!(
3065 detect_credit_exhaustion("quota refreshed next cycle"),
3066 Some(CreditExhaustionReason::BillingCycleExhausted)
3067 );
3068 assert_eq!(
3069 detect_credit_exhaustion(
3070 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
3071 ),
3072 Some(CreditExhaustionReason::BillingCycleExhausted)
3073 );
3074 assert_eq!(
3075 detect_credit_exhaustion("monthly quota exceeded"),
3076 Some(CreditExhaustionReason::BillingCycleExhausted)
3077 );
3078 assert_eq!(
3079 detect_credit_exhaustion(
3080 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
3081 ),
3082 Some(CreditExhaustionReason::BillingCycleExhausted)
3083 );
3084 }
3085
3086 #[serial(cooldown_state)]
3087 #[test]
3088 fn detect_credit_exhaustion_billing_cycle_minimax() {
3089 assert_eq!(
3090 detect_credit_exhaustion(
3091 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
3092 ),
3093 Some(CreditExhaustionReason::BillingCycleExhausted)
3094 );
3095 assert_eq!(
3096 detect_credit_exhaustion("Upgrade your Token Plan"),
3097 Some(CreditExhaustionReason::BillingCycleExhausted)
3098 );
3099 assert_eq!(
3100 detect_credit_exhaustion("Token Plan usage limit reached"),
3101 Some(CreditExhaustionReason::BillingCycleExhausted)
3102 );
3103 assert_eq!(
3104 detect_credit_exhaustion("purchase Credits for more usage"),
3105 Some(CreditExhaustionReason::BillingCycleExhausted)
3106 );
3107 assert_eq!(
3108 detect_credit_exhaustion(
3109 "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"
3110 ),
3111 Some(CreditExhaustionReason::BillingCycleExhausted)
3112 );
3113 }
3114
3115 #[serial(cooldown_state)]
3116 #[test]
3117 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
3118 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
3123 this billing cycle. Your quota will be refreshed soon.";
3124 assert_eq!(
3125 detect_credit_exhaustion_outcome(msg),
3126 Some("billing_cycle_exhausted")
3127 );
3128 }
3129
3130 #[serial(cooldown_state)]
3131 #[test]
3132 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
3133 assert_eq!(
3134 detect_credit_exhaustion_outcome("out of credits"),
3135 Some("credit_exhausted")
3136 );
3137 assert_eq!(
3138 detect_credit_exhaustion_outcome("organization has been disabled"),
3139 Some("org_level_disabled")
3140 );
3141 assert_eq!(
3142 detect_credit_exhaustion_outcome("monthly quota exceeded"),
3143 Some("billing_cycle_exhausted")
3144 );
3145 }
3146
3147 #[test]
3148 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
3149 assert_eq!(
3150 detect_credit_exhaustion_outcome("429 Too Many Requests"),
3151 None
3152 );
3153 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
3154 }
3155
3156 #[serial(cooldown_state)]
3157 #[test]
3158 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
3159 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
3162 assert_eq!(
3163 detect_credit_exhaustion(msg),
3164 Some(CreditExhaustionReason::OrgLevelDisabled)
3165 );
3166 }
3167
3168 #[serial(cooldown_state)]
3169 #[tokio::test]
3170 async fn refresh_degraded_agents_below_threshold_not_degraded() {
3171 super::reset_global_state().await;
3172 let store = test_store().await;
3173 let agent = "test_refresh_below_threshold";
3174
3175 for _ in 0..2 {
3177 store.record_rate_limit(agent, "rate", None).await.unwrap();
3178 }
3179
3180 let agents = vec![agent.to_string()];
3181 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3183
3184 assert!(
3185 !is_agent_degraded(agent),
3186 "agent with rate limit count < threshold should not be degraded"
3187 );
3188 }
3189
3190 #[serial(cooldown_state)]
3191 #[tokio::test]
3192 async fn refresh_degraded_agents_ignores_error_limit_type() {
3193 super::reset_global_state().await;
3194 let store = test_store().await;
3195 let agent = "test_refresh_error_rows";
3196
3197 for _ in 0..3 {
3198 store.record_rate_limit(agent, "error", None).await.unwrap();
3199 }
3200
3201 let agents = vec![agent.to_string()];
3202 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3203
3204 assert!(
3205 !is_agent_degraded(agent),
3206 "three error rows in 6 hours must not mark an agent degraded"
3207 );
3208 }
3209
3210 #[serial(cooldown_state)]
3211 #[tokio::test]
3212 async fn refresh_degraded_agents_marks_three_rate_rows() {
3213 super::reset_global_state().await;
3214 let store = test_store().await;
3215 let agent = "test_refresh_three_rate_rows";
3216
3217 for _ in 0..3 {
3218 store.record_rate_limit(agent, "rate", None).await.unwrap();
3219 }
3220
3221 let agents = vec![agent.to_string()];
3222 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3223
3224 assert!(
3225 is_agent_degraded(agent),
3226 "three rate rows in 6 hours must mark an agent degraded"
3227 );
3228 assert_eq!(
3229 agent_degraded_reason(agent).as_deref(),
3230 Some("rate limit threshold exceeded")
3231 );
3232 let listed = list_degraded_agents();
3233 assert!(
3234 listed
3235 .iter()
3236 .any(|(a, r)| a == agent && r == "rate limit threshold exceeded"),
3237 "orch cooldown list must show the degraded reason, got {listed:?}"
3238 );
3239
3240 clear_agent_degraded(agent);
3241 }
3242
3243 #[serial(cooldown_state)]
3244 #[tokio::test]
3245 async fn refresh_degraded_agents_logs_reason_change() {
3246 super::reset_global_state().await;
3247 let store = test_store().await;
3248 let agent = "test_refresh_reason_change";
3249 *cooldown_store().lock().await = Some(Arc::clone(&store));
3250
3251 set_agent_cooldown(agent, 600).await;
3252 let agents = vec![agent.to_string()];
3253 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3254 assert_eq!(
3255 agent_degraded_reason(agent).as_deref(),
3256 Some("agent in cooldown")
3257 );
3258
3259 clear_cooldown(agent, &store).await;
3260 let after_clear = sqlite_datetime_offset(60);
3261 for _ in 0..3 {
3262 store
3263 .record_rate_limit_at(agent, "rate", None, &after_clear)
3264 .await
3265 .unwrap();
3266 }
3267 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3268 assert_eq!(
3269 agent_degraded_reason(agent).as_deref(),
3270 Some("rate limit threshold exceeded"),
3271 "degraded reason must update when the hold changes"
3272 );
3273
3274 clear_agent_degraded(agent);
3275 }
3276
3277 #[serial(cooldown_state)]
3278 #[tokio::test]
3279 async fn clear_un_degrades_on_next_health_check() {
3280 super::reset_global_state().await;
3281 let store = test_store().await;
3282 let agent = "test_clear_un_degrades";
3283 let before = sqlite_datetime_offset(-30);
3284 let after = sqlite_datetime_offset(60);
3285
3286 for _ in 0..3 {
3287 store
3288 .record_rate_limit_at(agent, "rate", None, &before)
3289 .await
3290 .unwrap();
3291 }
3292
3293 let agents = vec![agent.to_string()];
3294 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3295 assert!(
3296 is_agent_degraded(agent),
3297 "three rate rows must mark an agent degraded"
3298 );
3299
3300 clear_cooldown(agent, &store).await;
3301 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3302 assert!(
3303 !is_agent_degraded(agent),
3304 "clear must un-degrade on the next health check"
3305 );
3306
3307 store
3308 .record_rate_limit_at(agent, "rate", None, &after)
3309 .await
3310 .unwrap();
3311 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3312 assert!(
3313 !is_agent_degraded(agent),
3314 "one rate row after the mark must not degrade again"
3315 );
3316
3317 for _ in 0..2 {
3318 store
3319 .record_rate_limit_at(agent, "rate", None, &after)
3320 .await
3321 .unwrap();
3322 }
3323 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3324 assert!(
3325 is_agent_degraded(agent),
3326 "three rate rows after the mark must degrade again"
3327 );
3328
3329 clear_agent_degraded(agent);
3330 }
3331
3332 #[serial(cooldown_state)]
3333 #[tokio::test]
3334 async fn clear_all_un_degrades_rate_limited_agent() {
3335 super::reset_global_state().await;
3336 let store = test_store().await;
3337 let agent = "test_clear_all_un_degrades";
3338 let before = sqlite_datetime_offset(-30);
3339
3340 for _ in 0..3 {
3341 store
3342 .record_rate_limit_at(agent, "rate", None, &before)
3343 .await
3344 .unwrap();
3345 }
3346
3347 let agents = vec![agent.to_string()];
3348 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3349 assert!(is_agent_degraded(agent));
3350
3351 clear_cooldown("*", &store).await;
3352 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3353 assert!(
3354 !is_agent_degraded(agent),
3355 "clear --all must un-degrade an agent held only by rate rows"
3356 );
3357 }
3358
3359 #[serial(cooldown_state)]
3362 #[tokio::test]
3363 async fn repeated_agent_failures_escalate_backoff() {
3364 super::reset_global_state().await;
3365 let store = test_store().await;
3366 *cooldown_store().lock().await = Some(Arc::clone(&store));
3367
3368 let agent = "test_escalation_agent";
3369
3370 record_agent_failure_with_message(agent, "").await;
3372 let remaining_1 = {
3373 let map = cooldowns().lock().unwrap();
3374 let entry = map.get(agent).expect("should have cooldown");
3375 entry.cooldown_until - chrono::Utc::now().timestamp()
3376 };
3377 assert!(
3378 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
3379 "first failure should be ~5 min, got {remaining_1}s"
3380 );
3381
3382 {
3384 let mut map = cooldowns().lock().unwrap();
3385 map.remove(agent);
3386 }
3387
3388 record_agent_failure_with_message(agent, "").await;
3390 let remaining_2 = {
3391 let map = cooldowns().lock().unwrap();
3392 let entry = map.get(agent).expect("should have cooldown");
3393 entry.cooldown_until - chrono::Utc::now().timestamp()
3394 };
3395 assert!(
3396 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
3397 "second failure should be ~15 min, got {remaining_2}s"
3398 );
3399 }
3400
3401 #[serial(cooldown_state)]
3402 #[tokio::test]
3403 async fn vendor_past_retry_falls_back_to_backoff() {
3404 super::reset_global_state().await;
3405 let store = test_store().await;
3409 *cooldown_store().lock().await = Some(Arc::clone(&store));
3410
3411 let agent = "test_vendor_past_retry";
3412
3413 {
3415 let mut map = cooldowns().lock().unwrap();
3416 map.remove(agent);
3417 }
3418
3419 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
3421
3422 record_agent_failure_with_message(agent, past_msg).await;
3423
3424 assert!(
3428 is_agent_in_cooldown(agent),
3429 "agent should be in cooldown after failure"
3430 );
3431
3432 let remaining = {
3433 let map = cooldowns().lock().unwrap();
3434 let entry = map.get(agent).expect("cooldown entry should exist");
3435 entry.cooldown_until - chrono::Utc::now().timestamp()
3436 };
3437
3438 assert!(
3439 remaining >= BACKOFF_BASE_SECS - 5,
3440 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3441 );
3442
3443 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3445 let stored = store
3446 .kv_get(&kv_key)
3447 .await
3448 .unwrap()
3449 .expect("failure count should be stored");
3450 assert_eq!(stored, "1");
3451 }
3452
3453 #[serial(cooldown_state)]
3454 #[tokio::test]
3455 async fn record_agent_success_resets_backoff() {
3456 super::reset_global_state().await;
3457 let store = test_store().await;
3458 *cooldown_store().lock().await = Some(Arc::clone(&store));
3459
3460 let agent = "test_success_reset_agent";
3461 let model = "test_model";
3462
3463 for _ in 0..3 {
3465 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3466 }
3467 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3469 let count: u32 = store
3470 .kv_get(&kv_key)
3471 .await
3472 .unwrap()
3473 .unwrap()
3474 .parse()
3475 .unwrap();
3476 assert_eq!(count, 3);
3477
3478 record_agent_success(agent, model).await;
3480
3481 let count_after: u32 = store
3482 .kv_get(&kv_key)
3483 .await
3484 .unwrap()
3485 .unwrap()
3486 .parse()
3487 .unwrap();
3488 assert_eq!(count_after, 0, "success should reset failure count to 0");
3489
3490 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3492 let model_count: u32 = store
3493 .kv_get(&model_kv_key)
3494 .await
3495 .unwrap()
3496 .unwrap()
3497 .parse()
3498 .unwrap();
3499 assert_eq!(
3500 model_count, 0,
3501 "success should reset model failure count to 0"
3502 );
3503 }
3504
3505 #[serial(cooldown_state)]
3506 #[tokio::test]
3507 async fn record_agent_success_without_model_skips_empty_model_key() {
3508 super::reset_global_state().await;
3509 let store = test_store().await;
3510 *cooldown_store().lock().await = Some(Arc::clone(&store));
3511
3512 let agent = "test_success_without_model";
3513
3514 record_agent_success(agent, "").await;
3515
3516 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3517 let agent_count = store
3518 .kv_get(&agent_kv_key)
3519 .await
3520 .unwrap()
3521 .expect("agent failure count should be stored");
3522 assert_eq!(agent_count, "0");
3523
3524 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3525 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3526 assert!(
3527 empty_model_count.is_none(),
3528 "success without model should not create an empty model failure count key"
3529 );
3530 }
3531
3532 #[serial(cooldown_state)]
3533 #[tokio::test]
3534 async fn clear_cooldown_resets_failure_counts() {
3535 super::reset_global_state().await;
3536 let store = test_store().await;
3537 *cooldown_store().lock().await = Some(Arc::clone(&store));
3538
3539 let agent = "test_clear_fc_agent";
3540
3541 for _ in 0..5 {
3543 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3544 }
3545 set_agent_cooldown(agent, 3600).await;
3547
3548 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3550 let count: u32 = store
3551 .kv_get(&kv_key)
3552 .await
3553 .unwrap()
3554 .unwrap()
3555 .parse()
3556 .unwrap();
3557 assert_eq!(count, 5);
3558
3559 clear_cooldown(agent, &store).await;
3561
3562 let count_after: u32 = store
3564 .kv_get(&kv_key)
3565 .await
3566 .unwrap()
3567 .unwrap()
3568 .parse()
3569 .unwrap();
3570 assert_eq!(
3571 count_after, 0,
3572 "clear_cooldown should reset failure count to 0"
3573 );
3574
3575 assert!(
3577 !is_agent_in_cooldown(agent),
3578 "agent should not be in cooldown after clear"
3579 );
3580 }
3581
3582 #[serial(cooldown_state)]
3583 #[tokio::test]
3584 async fn clear_all_cooldowns_resets_all_failure_counts() {
3585 super::reset_global_state().await;
3586 let store = test_store().await;
3587 *cooldown_store().lock().await = Some(Arc::clone(&store));
3588
3589 let agents = ["test_clear_all_a", "test_clear_all_b"];
3590 for agent in &agents {
3591 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3592 set_agent_cooldown(agent, 3600).await;
3593 }
3594
3595 clear_cooldown("*", &store).await;
3596
3597 for agent in &agents {
3598 assert!(!is_agent_in_cooldown(agent));
3599 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3600 let count: u32 = store
3601 .kv_get(&kv_key)
3602 .await
3603 .unwrap()
3604 .unwrap()
3605 .parse()
3606 .unwrap();
3607 assert_eq!(
3608 count, 0,
3609 "clear --all should reset failure count for {agent}"
3610 );
3611 }
3612 }
3613
3614 #[serial(cooldown_state)]
3615 #[test]
3616 fn compute_backoff_saturates_on_large_counts() {
3617 assert_eq!(
3619 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3620 BACKOFF_MAX_SECS * 12
3621 );
3622 assert_eq!(
3623 compute_backoff(u64::from(u32::MAX), 300, 14400, true),
3624 172_800
3625 );
3626 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172_800);
3627 }
3628
3629 #[serial(cooldown_state)]
3630 #[test]
3631 fn compute_backoff_zero_base_returns_zero() {
3632 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3633 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3634 }
3635
3636 #[serial(cooldown_state)]
3639 #[tokio::test]
3640 async fn github_circuit_stays_closed_below_threshold() {
3641 super::reset_global_state().await;
3642 {
3644 let mut open = github_5xx_circuit_open().lock().unwrap();
3645 *open = None;
3646 let mut ts = github_5xx_timestamps().lock().unwrap();
3647 ts.clear();
3648 let mut map = cooldown_lock();
3652 map.remove("github:5xx");
3653 }
3654
3655 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3657 record_github_5xx().await;
3658 }
3659
3660 assert!(
3661 !is_github_circuit_open(),
3662 "circuit should stay closed below threshold"
3663 );
3664 assert_eq!(github_circuit_remaining_secs(), 0);
3665 }
3666
3667 #[serial(cooldown_state)]
3668 #[tokio::test]
3669 async fn github_circuit_trips_at_threshold() {
3670 super::reset_global_state().await;
3671 {
3673 let mut open = github_5xx_circuit_open_lock();
3674 *open = None;
3675 let mut ts = github_5xx_timestamps_lock();
3676 ts.clear();
3677 let mut map = cooldown_lock();
3678 map.remove("github:5xx");
3679 }
3680
3681 for _ in 0..GITHUB_5XX_THRESHOLD {
3683 record_github_5xx().await;
3684 }
3685
3686 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3687 assert!(
3688 github_circuit_remaining_secs() > 0,
3689 "remaining should be > 0 when open"
3690 );
3691 assert!(
3692 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3693 "remaining should not exceed cooldown duration"
3694 );
3695
3696 {
3698 let mut open = github_5xx_circuit_open().lock().unwrap();
3699 *open = None;
3700 }
3701 }
3702
3703 #[serial(cooldown_state)]
3704 #[tokio::test]
3705 async fn github_circuit_cannot_double_trip() {
3706 super::reset_global_state().await;
3707 {
3709 let mut open = github_5xx_circuit_open_lock();
3710 *open = None;
3711 let mut ts = github_5xx_timestamps_lock();
3712 ts.clear();
3713 let mut map = cooldown_lock();
3714 map.remove("github:5xx");
3715 }
3716
3717 for _ in 0..GITHUB_5XX_THRESHOLD {
3719 record_github_5xx().await;
3720 }
3721 assert!(is_github_circuit_open());
3722
3723 let first_remaining = github_circuit_remaining_secs();
3724
3725 for _ in 0..GITHUB_5XX_THRESHOLD {
3727 record_github_5xx().await;
3728 }
3729
3730 let second_remaining = github_circuit_remaining_secs();
3731 assert!(
3732 second_remaining <= first_remaining,
3733 "additional errors should not extend cooldown"
3734 );
3735
3736 {
3738 let mut open = github_5xx_circuit_open().lock().unwrap();
3739 *open = None;
3740 }
3741 }
3742
3743 #[serial(cooldown_state)]
3744 #[tokio::test]
3745 async fn github_circuit_recovers_after_cooldown() {
3746 super::reset_global_state().await;
3747 {
3749 let mut open = github_5xx_circuit_open_lock();
3750 *open = None;
3751 let mut ts = github_5xx_timestamps_lock();
3752 ts.clear();
3753 let mut map = cooldown_lock();
3754 map.remove("github:5xx");
3755 }
3756
3757 {
3759 let mut open = github_5xx_circuit_open().lock().unwrap();
3760 *open = Some(chrono::Utc::now().timestamp() - 1);
3761 }
3762
3763 assert!(
3765 !is_github_circuit_open(),
3766 "circuit should auto-recover after cooldown"
3767 );
3768 }
3769
3770 #[serial(cooldown_state)]
3771 #[tokio::test]
3772 async fn github_circuit_sliding_window_ages_out() {
3773 super::reset_global_state().await;
3774 {
3776 let mut open = github_5xx_circuit_open_lock();
3777 *open = None;
3778 let mut ts = github_5xx_timestamps_lock();
3779 ts.clear();
3780 let mut map = cooldown_lock();
3781 map.remove("github:5xx");
3782 }
3783
3784 {
3786 let mut ts = github_5xx_timestamps().lock().unwrap();
3787 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3788 for _ in 0..GITHUB_5XX_THRESHOLD {
3789 ts.push(old);
3790 }
3791 }
3792
3793 record_github_5xx().await;
3795
3796 assert!(
3797 !is_github_circuit_open(),
3798 "old errors outside window should not trip circuit"
3799 );
3800 }
3801}