1use chrono::TimeZone;
19use once_cell::sync::Lazy;
20use regex::Regex;
21use std::collections::HashMap;
22use std::sync::{Arc, Mutex, OnceLock};
23
24pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
30
31pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
33
34pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
36
37pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
39
40pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
42
43pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
45
46pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
48
49pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
52
53pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
61
62pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum CreditExhaustionReason {
72 OutOfCredits,
74 OrgLevelDisabled,
76 BillingCycleExhausted,
78}
79
80impl CreditExhaustionReason {
81 pub fn outcome_str(self) -> &'static str {
85 match self {
86 CreditExhaustionReason::OutOfCredits => "credit_exhausted",
87 CreditExhaustionReason::OrgLevelDisabled => "org_level_disabled",
88 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
89 }
90 }
91}
92
93pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
100
101pub const SILENCE_COUNT_THRESHOLD: usize = 5;
103
104pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
106
107pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
109
110const KV_PREFIX: &str = "cooldown:";
112
113const SILENCE_COUNT_PREFIX: &str = "silence_count:";
115
116const FAILURE_COUNT_PREFIX: &str = "failure_count:";
118
119const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
122
123const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
127
128const RATE_LIMITS_CLEARED_AT_PREFIX: &str = "rate_limits_cleared_at:";
130
131fn agent_from_key(key: &str) -> &str {
133 key.split_once(':').map(|(agent, _)| agent).unwrap_or(key)
134}
135
136struct CooldownEntry {
137 cooldown_until: i64,
139 reason: String,
140}
141
142pub struct SilenceCountResult {
143 pub count: usize,
144 pub extended_cooldown_applied: bool,
145}
146
147fn cooldowns() -> &'static Mutex<HashMap<String, CooldownEntry>> {
156 static COOLDOWNS: OnceLock<Mutex<HashMap<String, CooldownEntry>>> = OnceLock::new();
157 COOLDOWNS.get_or_init(|| Mutex::new(HashMap::new()))
158}
159
160fn cooldown_store() -> &'static tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>> {
166 static STORE: OnceLock<tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>>> =
167 OnceLock::new();
168 STORE.get_or_init(|| tokio::sync::Mutex::new(None))
169}
170
171fn parse_cooldown_value(value: &str) -> Option<(i64, String)> {
179 if let Ok(serde_json::Value::Object(map)) = serde_json::from_str(value) {
180 let until = map.get("until")?.as_i64()?;
181 let reason = map
182 .get("reason")
183 .and_then(|r| r.as_str())
184 .unwrap_or("persisted")
185 .to_string();
186 return Some((until, reason));
187 }
188 value
189 .parse::<i64>()
190 .ok()
191 .map(|until| (until, "persisted".to_string()))
192}
193
194pub async fn init_cooldown_store(store: Arc<crate::store::TaskStore>) {
195 match store.kv_list_prefix(KV_PREFIX).await {
196 Ok(rows) => {
197 let now = chrono::Utc::now().timestamp();
198 let mut map = cooldown_lock();
199 let mut loaded = 0;
200 for (key, value) in rows {
201 let cooldown_key = key.trim_start_matches(KV_PREFIX);
202 if let Some((cooldown_until, reason)) = parse_cooldown_value(&value) {
203 if now < cooldown_until {
204 map.insert(
205 cooldown_key.to_string(),
206 CooldownEntry {
207 cooldown_until,
208 reason,
209 },
210 );
211 loaded += 1;
212 }
213 }
216 }
217 if loaded > 0 {
218 tracing::info!(loaded, "loaded persisted cooldowns from KV store");
219 }
220 }
221 Err(e) => {
222 tracing::warn!(err = %e, "failed to load cooldowns from KV store");
223 }
224 }
225
226 if let Ok(Some(raw)) = store.kv_get("cooldown:github:5xx").await {
228 if let Ok(cooldown_until) = raw.parse::<i64>() {
229 let now = chrono::Utc::now().timestamp();
230 if now < cooldown_until {
231 let mut open = github_5xx_circuit_open_lock();
232 *open = Some(cooldown_until);
233 let remaining = cooldown_until - now;
234 tracing::warn!(
235 remaining_secs = remaining,
236 "GitHub 5xx circuit breaker restored from KV (still open)"
237 );
238 }
239 }
240 }
241
242 *cooldown_store().lock().await = Some(store);
243}
244
245pub async fn sync_from_kv(store: &Arc<crate::store::TaskStore>) {
262 let rows = match store.kv_list_prefix(KV_PREFIX).await {
263 Ok(rows) => rows,
264 Err(e) => {
265 tracing::warn!(err = %e, "cooldown sync_from_kv: KV list failed");
266 return;
267 }
268 };
269 let now = chrono::Utc::now().timestamp();
270
271 let mut explicit_kv: HashMap<String, Option<(i64, String)>> =
274 HashMap::with_capacity(rows.len());
275 for (kv_key, value) in &rows {
276 let cooldown_key = kv_key.trim_start_matches(KV_PREFIX);
277 if let Some((ts, reason)) = parse_cooldown_value(value) {
278 let normalized = if ts > now { Some((ts, reason)) } else { None };
279 explicit_kv.insert(cooldown_key.to_string(), normalized);
280 }
281 }
282
283 let mut map = cooldown_lock();
284 let mut dropped = 0usize;
285 let mut updated = 0usize;
286
287 map.retain(|key, _| {
289 if matches!(explicit_kv.get(key), Some(None)) {
290 dropped += 1;
291 false
292 } else {
293 true
294 }
295 });
296
297 for (key, normalized) in explicit_kv {
299 if let Some((ts, reason)) = normalized {
300 match map.get_mut(&key) {
301 Some(entry) => {
302 if entry.cooldown_until != ts {
303 entry.cooldown_until = ts;
304 updated += 1;
305 }
306 }
307 None => {
308 map.insert(
309 key,
310 CooldownEntry {
311 cooldown_until: ts,
312 reason,
313 },
314 );
315 updated += 1;
316 }
317 }
318 }
319 }
320
321 if dropped > 0 || updated > 0 {
322 tracing::info!(
323 dropped,
324 updated,
325 "cooldown sync_from_kv: reconciled in-memory map with KV"
326 );
327 }
328}
329
330const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
334
335const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
337
338const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
342
343const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
345
346pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
365 if count <= 1 {
366 return base;
367 }
368 let exp = count.saturating_sub(1).min(u32::MAX as u64) as u32;
369 let factor = 3_i64.saturating_pow(exp);
370 let standard = base.saturating_mul(factor).min(max);
371
372 let first_max = if base <= 0 {
381 0
382 } else {
383 let mut accumulated = base;
385 let mut first_max_count = 1u64;
386 while accumulated < max && first_max_count < count {
387 accumulated = accumulated.saturating_mul(3);
388 first_max_count += 1;
389 }
390 first_max_count
391 };
392
393 let times_at_max = count.saturating_sub(first_max);
395
396 if !allow_extended {
397 return standard;
398 }
399
400 if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
401 max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
402 } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
403 max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
404 } else {
405 standard
406 }
407}
408
409async fn read_and_increment_failure_count(
418 store_opt: &Option<Arc<crate::store::TaskStore>>,
419 key: &str,
420) -> u64 {
421 read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
422}
423
424async fn read_and_increment_failure_count_with_prefix(
429 store_opt: &Option<Arc<crate::store::TaskStore>>,
430 prefix: &str,
431 key: &str,
432) -> u64 {
433 let kv_key = format!("{prefix}{key}");
434 if let Some(store) = store_opt {
435 match store.kv_increment(&kv_key).await {
436 Ok(n) => n,
437 Err(e) => {
438 tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
439 u64::MAX
440 }
441 }
442 } else {
443 1
444 }
445}
446
447pub async fn record_agent_success(agent_name: &str, model: &str) {
453 let store_opt = cooldown_store().lock().await.clone();
454 if let Some(store) = store_opt {
455 let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
456 let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
457 if let Err(e) = store.kv_set(&agent_key, "0").await {
458 tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
459 }
460 if !model.is_empty() {
461 let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
462 if let Err(e) = store.kv_set(&model_key, "0").await {
463 tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
464 }
465 }
466 if let Err(e) = store.kv_set(&credit_key, "0").await {
467 tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
468 }
469 }
470}
471
472pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
478 if let Some(cooldown_until) = parse_retry_at(error_message) {
483 let now = chrono::Utc::now().timestamp();
484 if cooldown_until > now {
485 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
486 return;
487 }
488 tracing::debug!(
489 agent = agent_name,
490 cooldown_until,
491 "vendor retry_at is in the past — using backoff instead"
492 );
493 }
494
495 let store_opt = cooldown_store().lock().await.clone();
496 let count = read_and_increment_failure_count(&store_opt, agent_name).await;
497 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
498 let cooldown_until =
499 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
500 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
501}
502
503pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
512 detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
513}
514
515pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
520 let lower = error_message.to_lowercase();
521
522 let overage_patterns = [
523 "overagedisabledreason",
524 "overage_disabled_reason",
525 "overage disabled reason",
526 ];
527
528 if overage_patterns.iter().any(|p| lower.contains(p)) {
529 if lower.contains("out_of_credits") || lower.contains("outofcredits") {
530 return Some(CreditExhaustionReason::OutOfCredits);
531 }
532 if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
533 return Some(CreditExhaustionReason::OrgLevelDisabled);
534 }
535 }
536
537 let billing_patterns = [
538 "credit balance too low",
539 "credit balance insufficient",
540 "insufficient credit balance",
541 "billing quota exceeded",
542 "billing limit exceeded",
543 "organization has been disabled",
544 "org disabled",
545 ];
546
547 if billing_patterns.iter().any(|p| lower.contains(p)) {
548 return Some(CreditExhaustionReason::OrgLevelDisabled);
549 }
550
551 let billing_cycle_patterns = [
554 "billing cycle",
555 "quota refreshed next cycle",
556 "monthly quota",
557 "quota will be refreshed",
558 "next billing cycle",
559 "refreshed in the next cycle",
560 "weekly/monthly limit exhausted",
561 "limit will reset at",
562 "upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
567
568 if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
569 return Some(CreditExhaustionReason::BillingCycleExhausted);
570 }
571
572 if lower.contains("out of credits")
573 || lower.contains("outofcredits")
574 || lower.contains("insufficient funds")
575 || lower.contains("no credits remaining")
576 || lower.contains("insufficient_quota")
577 || lower.contains("quota exceeded")
578 || lower.contains("insufficient balance")
579 || lower.contains("please recharge")
580 || lower.contains("no resource package")
581 || lower.contains("balance exhausted")
582 || lower.contains("402 payment required")
583 {
584 return Some(CreditExhaustionReason::OutOfCredits);
585 }
586
587 None
588}
589
590pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
603 let reason_str = match reason {
604 CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
605 CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
606 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
607 };
608
609 let store_opt = cooldown_store().lock().await.clone();
610 let count = read_and_increment_failure_count_with_prefix(
615 &store_opt,
616 CREDIT_FAILURE_COUNT_PREFIX,
617 agent_name,
618 )
619 .await;
620
621 let (base, max) = match reason {
622 CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
623 CreditExhaustionReason::OrgLevelDisabled => {
624 (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
625 }
626 CreditExhaustionReason::BillingCycleExhausted => {
627 (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
628 }
629 };
630 let cooldown_secs = compute_backoff(
631 count,
632 base,
633 max,
634 !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
635 );
636 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
637 set_cooldown_async(agent_name, cooldown_until, reason_str).await;
638 tracing::warn!(
639 agent = agent_name,
640 reason = reason_str,
641 attempt = count,
642 cooldown_secs,
643 "credit exhaustion detected: applying exponential agent-wide cooldown"
644 );
645}
646
647pub async fn record_model_failure(agent_name: &str, model: &str) {
651 let key = format!("{agent_name}:{model}");
652 let store_opt = cooldown_store().lock().await.clone();
653 let count = read_and_increment_failure_count(&store_opt, &key).await;
654 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
655 let cooldown_until =
656 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
657 set_cooldown_async(&key, cooldown_until, "model_error").await;
658}
659
660pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
668 if let Some(window_secs) = parse_relative_usage_window(error_message) {
669 if let Some(m) = model {
670 let cooldown_until = chrono::Utc::now().timestamp() + window_secs as i64;
671 set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
672 tracing::info!(
673 agent = agent_name,
674 model = m,
675 window_secs,
676 cooldown_until,
677 "relative usage window detected: applying model-specific cooldown"
678 );
679 return;
680 }
681 }
682
683 record_agent_failure_with_message(agent_name, error_message).await;
684 if let Some(m) = model {
685 record_model_failure(agent_name, m).await;
686 }
687}
688
689pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
701 let key = format!("{agent_name}:{model}");
702 let store_opt = cooldown_store().lock().await.clone();
703 let count = read_and_increment_failure_count_with_prefix(
704 &store_opt,
705 PERSISTENT_FAILURE_COUNT_PREFIX,
706 &key,
707 )
708 .await;
709 let cooldown_secs = compute_backoff(
710 count,
711 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
712 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
713 false,
714 );
715 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
716 set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
717}
718
719pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
726 let key = format!("{agent_name}:{model}");
727 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
728 set_cooldown_async(&key, cooldown_until, reason).await;
729}
730
731pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
733 let key = format!("{agent_name}:{model}");
734 set_cooldown_async(&key, until, "vendor_reset_at").await;
735}
736
737pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
744 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
745 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
746}
747
748pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
751 let store_opt = cooldown_store().lock().await.clone();
752 let store = match store_opt {
753 Some(store) => store,
754 None => {
755 tracing::debug!(
756 agent = agent_name,
757 model,
758 "skipping silence count record (no KV store)"
759 );
760 return None;
761 }
762 };
763
764 let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
765 let now = chrono::Utc::now().timestamp();
766 let window_start = now - SILENCE_COUNT_WINDOW_SECS;
767
768 let mut timestamps = match store.kv_get(&key).await {
769 Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
770 Ok(ts) => ts,
771 Err(e) => {
772 tracing::warn!(
773 kv_key = %key,
774 err = %e,
775 "failed to parse silence timestamps from KV — resetting window"
776 );
777 Vec::new()
778 }
779 },
780 Ok(None) => Vec::new(),
781 Err(err) => {
782 tracing::warn!(
783 kv_key = %key,
784 err = %err,
785 "failed to load silence count from KV"
786 );
787 Vec::new()
788 }
789 };
790
791 timestamps.retain(|ts| *ts >= window_start);
792 timestamps.push(now);
793
794 let count = timestamps.len();
795 let mut extended_cooldown_applied = false;
796 if count >= SILENCE_COUNT_THRESHOLD {
797 set_model_cooldown(
798 agent_name,
799 model,
800 SILENCE_EXTENDED_COOLDOWN_SECS,
801 "silence_detected",
802 )
803 .await;
804 extended_cooldown_applied = true;
805 timestamps.clear();
806 }
807
808 let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
809 if let Err(err) = store.kv_set(&key, &value).await {
810 tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
811 }
812
813 Some(SilenceCountResult {
814 count,
815 extended_cooldown_applied,
816 })
817}
818
819pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
821 let key = format!("{agent_name}:{model}");
822 is_in_cooldown(&key)
823}
824
825pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
827 is_in_cooldown(agent_name)
828}
829
830pub fn cooldown_until(key: &str) -> Option<i64> {
832 let map = cooldown_lock();
833 map.get(key).and_then(|entry| {
834 let now = chrono::Utc::now().timestamp();
835 if now < entry.cooldown_until {
836 Some(entry.cooldown_until)
837 } else {
838 None
839 }
840 })
841}
842
843pub fn cooldown_reason(key: &str) -> Option<String> {
845 let map = cooldown_lock();
846 map.get(key).and_then(|entry| {
847 let now = chrono::Utc::now().timestamp();
848 if now < entry.cooldown_until {
849 Some(entry.reason.clone())
850 } else {
851 None
852 }
853 })
854}
855
856pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
861 let now = chrono::Utc::now().timestamp();
862 let map = cooldown_lock();
863 let mut result: Vec<(String, i64, String)> = map
864 .iter()
865 .filter(|(_, entry)| now < entry.cooldown_until)
866 .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
867 .collect();
868 result.sort_by(|a, b| a.0.cmp(&b.0));
869 result
870}
871
872pub async fn mark_rate_limits_cleared(store: &Arc<crate::store::TaskStore>, agents: &[String]) {
875 let now = chrono::Utc::now().timestamp().to_string();
876 let mut seen = std::collections::HashSet::new();
877 for key in agents {
878 let agent = agent_from_key(key);
879 if agent.is_empty() || !seen.insert(agent.to_string()) {
880 continue;
881 }
882 let kv_key = format!("{RATE_LIMITS_CLEARED_AT_PREFIX}{agent}");
883 if let Err(e) = store.kv_set(&kv_key, &now).await {
884 tracing::warn!(agent, err = %e, "failed to mark rate limits cleared");
885 }
886 }
887}
888
889pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
896 let mut agents_to_mark: Vec<String> = Vec::new();
897 if key == "*" {
898 let keys: Vec<String> = {
899 let mut map = cooldown_lock();
900 let keys: Vec<String> = map.keys().cloned().collect();
901 map.clear();
902 keys
903 };
904 for k in &keys {
905 let kv_key = format!("{KV_PREFIX}{k}");
906 if let Err(e) = store.kv_set(&kv_key, "0").await {
907 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
908 }
909 let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
911 if let Err(e) = store.kv_set(&fc_key, "0").await {
912 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
913 }
914 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
916 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
917 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
918 }
919 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
921 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
922 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
923 }
924 }
925 if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
929 for (fc_key, _) in fc_entries {
930 if let Err(e) = store.kv_set(&fc_key, "0").await {
931 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
932 }
933 }
934 }
935 if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
936 for (cfc_key, _) in cfc_entries {
937 if let Err(e) = store.kv_set(&cfc_key, "0").await {
938 tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
939 }
940 }
941 }
942 if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
943 for (pfc_key, _) in pfc_entries {
944 if let Err(e) = store.kv_set(&pfc_key, "0").await {
945 tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
946 }
947 }
948 }
949 tracing::info!(
950 count = keys.len(),
951 "cleared all cooldowns and failure counts"
952 );
953 agents_to_mark.extend(keys.iter().map(|k| agent_from_key(k).to_string()));
954 match store.rate_limit_agents().await {
955 Ok(rl_agents) => agents_to_mark.extend(rl_agents),
956 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits agents while clearing"),
957 }
958 match store.kv_list_prefix(RATE_LIMITS_CLEARED_AT_PREFIX).await {
959 Ok(marks) => {
960 for (k, _) in marks {
961 let agent = k.trim_start_matches(RATE_LIMITS_CLEARED_AT_PREFIX);
962 if !agent.is_empty() {
963 agents_to_mark.push(agent.to_string());
964 }
965 }
966 }
967 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits_cleared_at keys"),
968 }
969 } else {
970 {
971 let mut map = cooldown_lock();
972 map.remove(key);
973 }
974 let kv_key = format!("{KV_PREFIX}{key}");
975 if let Err(e) = store.kv_set(&kv_key, "0").await {
976 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
977 }
978 let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
980 if let Err(e) = store.kv_set(&fc_key, "0").await {
981 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
982 }
983 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
985 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
986 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
987 }
988 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
990 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
991 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
992 }
993 tracing::info!(key, "cleared cooldown and failure count");
994 agents_to_mark.push(agent_from_key(key).to_string());
995 }
996 mark_rate_limits_cleared(store, &agents_to_mark).await;
997}
998
999fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
1008 let mut map = cooldown_lock();
1009 if let Some(existing) = map.get(key) {
1010 if existing.cooldown_until >= cooldown_until {
1014 return false;
1015 }
1016 }
1017 map.insert(
1018 key.to_string(),
1019 CooldownEntry {
1020 cooldown_until,
1021 reason: reason.to_string(),
1022 },
1023 );
1024 true
1025}
1026
1027async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
1037 if !set_cooldown_in_memory(key, cooldown_until, reason) {
1038 return false;
1039 }
1040 let store_opt = cooldown_store().lock().await.clone();
1041 if let Some(store) = store_opt {
1042 let kv_key = format!("{KV_PREFIX}{key}");
1043 let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
1044 if let Err(e) = store.kv_set(&kv_key, &value).await {
1045 tracing::warn!(
1046 kv_key,
1047 err = %e,
1048 "failed to persist cooldown to KV store — rolling back in-memory entry"
1049 );
1050 let mut map = cooldown_lock();
1051 if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1055 map.remove(key);
1056 }
1057 return false;
1058 }
1059 }
1060 true
1061}
1062
1063fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
1065 cooldowns().lock().unwrap_or_else(|e| {
1066 tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1067 e.into_inner()
1068 })
1069}
1070
1071fn degraded_agents_lock() -> std::sync::MutexGuard<'static, HashMap<String, String>> {
1073 degraded_agents().lock().unwrap_or_else(|e| {
1074 tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1075 e.into_inner()
1076 })
1077}
1078
1079fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
1081 github_5xx_timestamps().lock().unwrap_or_else(|e| {
1082 tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1083 e.into_inner()
1084 })
1085}
1086
1087fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
1089 github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1090 tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1091 e.into_inner()
1092 })
1093}
1094
1095fn is_in_cooldown(key: &str) -> bool {
1097 let map = cooldown_lock();
1098 if let Some(entry) = map.get(key) {
1099 let now = chrono::Utc::now().timestamp();
1100 return now < entry.cooldown_until;
1101 }
1102 false
1103}
1104
1105pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1111
1112pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1116
1117fn degraded_agents() -> &'static Mutex<HashMap<String, String>> {
1120 static DEGRADED: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
1121 DEGRADED.get_or_init(|| Mutex::new(HashMap::new()))
1122}
1123
1124pub fn is_agent_degraded(agent: &str) -> bool {
1129 let set = degraded_agents_lock();
1130 set.contains_key(agent)
1131}
1132
1133pub fn agent_degraded_reason(agent: &str) -> Option<String> {
1135 degraded_agents_lock().get(agent).cloned()
1136}
1137
1138pub fn list_degraded_agents() -> Vec<(String, String)> {
1140 let mut entries: Vec<(String, String)> = degraded_agents_lock()
1141 .iter()
1142 .map(|(agent, reason)| (agent.clone(), reason.clone()))
1143 .collect();
1144 entries.sort_by(|a, b| a.0.cmp(&b.0));
1145 entries
1146}
1147
1148pub fn mark_agent_degraded(agent: &str) {
1150 mark_agent_degraded_with_reason(agent, "degraded");
1151}
1152
1153fn mark_agent_degraded_with_reason(agent: &str, reason: &str) {
1154 let mut set = degraded_agents_lock();
1155 set.insert(agent.to_string(), reason.to_string());
1156}
1157
1158pub fn clear_agent_degraded(agent: &str) {
1160 let mut set = degraded_agents_lock();
1161 set.remove(agent);
1162}
1163
1164pub async fn refresh_degraded_agents(
1178 store: &Arc<crate::store::TaskStore>,
1179 available_agents: &[String],
1180 model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1181 window_hours: u32,
1182 threshold: i64,
1183) {
1184 clear_expired_cooldowns();
1188
1189 let start = chrono::Utc::now();
1192 let counts = match store.recent_rate_limit_counts(window_hours).await {
1193 Ok(c) => {
1194 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1195 tracing::info!(
1196 duration_ms = dur_ms,
1197 window_hours,
1198 "recent_rate_limit_counts query completed"
1199 );
1200 c
1201 }
1202 Err(e) => {
1203 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1204 tracing::warn!(
1205 err = %e,
1206 duration_ms = dur_ms,
1207 "failed to query recent rate limit counts for health check"
1208 );
1209 return;
1210 }
1211 };
1212
1213 for agent in available_agents {
1214 let in_cooldown = is_agent_in_cooldown(agent);
1215 let no_models = !model_checker(agent);
1216 let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1217 let over_threshold = rate_limit_count >= threshold;
1218
1219 let degraded = in_cooldown || no_models || over_threshold;
1220
1221 if degraded {
1222 let reason = if in_cooldown {
1223 "agent in cooldown"
1224 } else if no_models {
1225 "all models cooled"
1226 } else {
1227 "rate limit threshold exceeded"
1228 };
1229 let previous = agent_degraded_reason(agent);
1230 match previous.as_deref() {
1231 None => tracing::warn!(
1232 agent,
1233 reason,
1234 rate_limit_count,
1235 window_hours,
1236 threshold,
1237 "pre-emptive health check: marking agent as degraded"
1238 ),
1239 Some(prev) if prev != reason => tracing::warn!(
1240 agent,
1241 reason,
1242 previous = prev,
1243 rate_limit_count,
1244 window_hours,
1245 threshold,
1246 "pre-emptive health check: degraded reason changed"
1247 ),
1248 Some(_) => {}
1249 }
1250 mark_agent_degraded_with_reason(agent, reason);
1251 } else {
1252 if is_agent_degraded(agent) {
1253 tracing::info!(
1254 agent,
1255 rate_limit_count,
1256 "pre-emptive health check: agent recovered, clearing degraded flag"
1257 );
1258 }
1259 clear_agent_degraded(agent);
1260 }
1261 }
1262}
1263
1264pub fn clear_expired_cooldowns() {
1266 let mut map = cooldown_lock();
1267 let now = chrono::Utc::now().timestamp();
1268 map.retain(|_key, entry| now < entry.cooldown_until);
1269}
1270
1271const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1277
1278const GITHUB_5XX_THRESHOLD: usize = 5;
1280
1281const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1283
1284fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1286 static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1287 TS.get_or_init(|| Mutex::new(Vec::new()))
1288}
1289
1290fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1292 static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1293 OPEN.get_or_init(|| Mutex::new(None))
1294}
1295
1296pub async fn record_github_5xx() {
1301 let now = chrono::Utc::now().timestamp();
1302 let window_start = now - GITHUB_5XX_WINDOW_SECS;
1303
1304 let should_trip = {
1305 let mut ts = github_5xx_timestamps_lock();
1306 ts.retain(|t| *t >= window_start);
1307 ts.push(now);
1308 ts.len() >= GITHUB_5XX_THRESHOLD
1309 };
1310
1311 if should_trip {
1312 let cooldown_until = {
1313 let mut open = github_5xx_circuit_open_lock();
1314 if open.is_none() {
1315 let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1316 *open = Some(cd);
1317 tracing::warn!(
1318 errors_in_window = GITHUB_5XX_THRESHOLD,
1319 window_secs = GITHUB_5XX_WINDOW_SECS,
1320 cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1321 "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1322 );
1323 Some(cd)
1324 } else {
1325 None
1326 }
1327 };
1328 if let Some(cd) = cooldown_until {
1330 let store_opt = cooldown_store().lock().await.clone();
1331 if let Some(store) = store_opt {
1332 let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1333 }
1334 set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1338 }
1339 let mut ts = github_5xx_timestamps_lock();
1341 ts.clear();
1342 }
1343}
1344
1345pub fn is_github_circuit_open() -> bool {
1353 {
1356 let mut open = github_5xx_circuit_open_lock();
1357 if let Some(until) = *open {
1358 let now = chrono::Utc::now().timestamp();
1359 if now >= until {
1360 *open = None;
1361 tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1362 return false;
1363 }
1364 return true;
1365 }
1366 }
1367
1368 if let Some(until) = cooldown_until("github:5xx") {
1372 let mut open = github_5xx_circuit_open_lock();
1375 if open.is_none() {
1376 *open = Some(until);
1377 let now = chrono::Utc::now().timestamp();
1378 if now < until {
1379 let remaining = until - now;
1380 tracing::warn!(
1381 remaining_secs = remaining,
1382 "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1383 );
1384 }
1385 }
1386 return true;
1387 }
1388
1389 false
1390}
1391
1392pub fn github_circuit_remaining_secs() -> u64 {
1395 {
1397 let open = github_5xx_circuit_open_lock();
1398 if let Some(until) = *open {
1399 let now = chrono::Utc::now().timestamp();
1400 if now < until {
1401 return (until - now) as u64;
1402 }
1403 return 0;
1404 }
1405 }
1406
1407 if let Some(until) = cooldown_until("github:5xx") {
1409 let now = chrono::Utc::now().timestamp();
1410 if now < until {
1411 return (until - now) as u64;
1412 }
1413 }
1414 0
1415}
1416
1417pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1421 if error_message.is_empty() {
1422 return None;
1423 }
1424
1425 let lower = error_message.to_lowercase();
1426 for retry_marker in [
1430 "try again at ",
1431 "reset at ",
1432 "resets at ",
1433 "resets ",
1434 "reset ",
1435 ] {
1436 if let Some(idx) = lower.find(retry_marker) {
1437 let date_str = &lower[idx + retry_marker.len()..];
1438 let date_str = date_str
1440 .split(['.', '\n'])
1441 .next()
1442 .unwrap_or(date_str)
1443 .trim();
1444
1445 let cleaned = date_str
1448 .replace("st,", ",")
1449 .replace("nd,", ",")
1450 .replace("rd,", ",")
1451 .replace("th,", ",");
1452
1453 for fmt in &[
1454 "%b %d, %Y %I:%M %p", "%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", ] {
1461 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1462 use chrono::TimeZone;
1465 let utc = match chrono::Local.from_local_datetime(&dt) {
1466 chrono::LocalResult::Single(local_dt) => {
1467 local_dt.with_timezone(&chrono::Utc)
1468 }
1469 chrono::LocalResult::Ambiguous(earliest, _) => {
1470 earliest.with_timezone(&chrono::Utc)
1472 }
1473 chrono::LocalResult::None => {
1474 dt.and_utc()
1476 }
1477 };
1478 tracing::info!(
1479 retry_at = %utc,
1480 raw = date_str,
1481 marker = retry_marker,
1482 "parsed rate limit retry-at date"
1483 );
1484 return Some(utc.timestamp());
1485 }
1486 }
1487 tracing::debug!(
1488 raw = date_str,
1489 marker = retry_marker,
1490 "could not parse retry-at date"
1491 );
1492
1493 let time_only = date_str
1498 .split_once('(')
1499 .map(|(t, _)| t)
1500 .unwrap_or(date_str)
1501 .trim();
1502 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1504 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1505 let last = &compact[split..];
1506 if !last.contains(':')
1507 && (last.ends_with("am") || last.ends_with("pm"))
1508 && last.len() > 2
1509 {
1510 let at = split + last.len() - 2;
1511 compact.insert_str(at, ":00");
1512 }
1513 let md = compact.replace(" at ", " ");
1515 let now_local = chrono::Local::now();
1516 let this_year = chrono::Datelike::year(&now_local);
1517 for year in [this_year, this_year + 1] {
1518 let with_year = format!("{year} {md}");
1519 if let Ok(dt) =
1520 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1521 {
1522 use chrono::TimeZone;
1523 if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
1524 if local_dt > now_local {
1525 return Some(local_dt.timestamp());
1526 }
1527 }
1528 }
1529 }
1530 let time_only = compact.as_str();
1531 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1532 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1533 let local_now = chrono::Local::now();
1536 let today = local_now.date_naive();
1537 let candidate = chrono::NaiveDateTime::new(today, t);
1538 let candidate_local: chrono::DateTime<chrono::Local> =
1541 match chrono::Local.from_local_datetime(&candidate) {
1542 chrono::LocalResult::Single(dt) => dt,
1543 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1544 chrono::LocalResult::None => {
1545 let utc_dt = candidate.and_utc();
1547 return Some(utc_dt.timestamp());
1548 }
1549 };
1550 let chosen = if candidate_local <= local_now {
1551 let tomorrow = today.succ_opt().unwrap_or(today);
1553 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1554 match chrono::Local.from_local_datetime(&ndt) {
1555 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1556 chrono::LocalResult::Ambiguous(earliest, _) => {
1557 earliest.with_timezone(&chrono::Utc)
1558 }
1559 chrono::LocalResult::None => ndt.and_utc(),
1560 }
1561 } else {
1562 candidate_local.with_timezone(&chrono::Utc)
1563 };
1564 tracing::info!(
1565 retry_at = %chosen,
1566 raw = date_str,
1567 marker = retry_marker,
1568 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1569 );
1570 return Some(chosen.timestamp());
1571 }
1572 }
1573 }
1574 }
1575
1576 None
1580}
1581
1582pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1593 if error_message.is_empty() {
1594 return None;
1595 }
1596
1597 let lower = error_message.to_lowercase();
1598 let has_context = lower.contains("usage limit")
1599 || lower.contains("quota")
1600 || lower.contains("window")
1601 || lower.contains("limit");
1602 if !has_context {
1603 return None;
1604 }
1605
1606 static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
1607 Regex::new(
1608 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1609 )
1610 .expect("usage window regex is valid")
1611 });
1612
1613 let mut max_secs: Option<u64> = None;
1614 for caps in WINDOW_RE.captures_iter(&lower) {
1615 let secs = if let Some(num) = caps.get(1) {
1616 let n: u64 = num.as_str().parse().ok()?;
1617 let unit = caps.get(2)?.as_str().to_lowercase();
1618 let multiplier = match unit.as_str() {
1619 "hour" | "hours" | "hr" | "hrs" => 3600,
1620 "minute" | "minutes" | "min" | "mins" => 60,
1621 "day" | "days" => 86400,
1622 _ => return None,
1623 };
1624 n.checked_mul(multiplier)?
1625 } else if let Some(word) = caps.get(3) {
1626 match word.as_str().to_lowercase().as_str() {
1627 "hourly" => 3600,
1628 "daily" => 86400,
1629 "weekly" => 604800,
1630 "monthly" => 2592000,
1631 _ => return None,
1632 }
1633 } else {
1634 continue;
1635 };
1636 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1637 }
1638
1639 max_secs
1640}
1641
1642#[cfg(test)]
1643pub(crate) async fn reset_global_state() {
1644 cooldowns().lock().unwrap().clear();
1645 *cooldown_store().lock().await = None;
1646 github_5xx_timestamps().lock().unwrap().clear();
1647 *github_5xx_circuit_open().lock().unwrap() = None;
1648 degraded_agents().lock().unwrap().clear();
1649 crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1650}
1651
1652#[cfg(test)]
1653mod tests {
1654 use super::*;
1655 use serial_test::serial;
1656
1657 async fn test_store() -> Arc<crate::store::TaskStore> {
1658 Arc::new(
1659 crate::store::TaskStore::open_memory()
1660 .await
1661 .expect("in-memory store"),
1662 )
1663 }
1664
1665 fn sqlite_datetime_offset(secs: i64) -> String {
1666 (chrono::Utc::now() + chrono::Duration::seconds(secs))
1667 .format("%Y-%m-%d %H:%M:%S")
1668 .to_string()
1669 }
1670
1671 #[serial(cooldown_state)]
1672 #[tokio::test]
1673 async fn init_loads_unexpired_cooldowns_from_kv() {
1674 super::reset_global_state().await;
1675 let store = test_store().await;
1676 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1679 .kv_set(
1680 &format!("{KV_PREFIX}kimi:k2p5"),
1681 &cooldown_until.to_string(),
1682 )
1683 .await
1684 .unwrap();
1685
1686 let expired = chrono::Utc::now().timestamp() - 1;
1688 store
1689 .kv_set(
1690 &format!("{KV_PREFIX}opencode:old-free"),
1691 &expired.to_string(),
1692 )
1693 .await
1694 .unwrap();
1695
1696 {
1697 let mut map = cooldown_lock();
1698 map.clear();
1699 }
1700
1701 init_cooldown_store(store).await;
1702
1703 assert!(
1704 is_model_in_cooldown("kimi", "k2p5"),
1705 "unexpired cooldown should be loaded"
1706 );
1707 assert!(
1708 !is_model_in_cooldown("opencode", "old-free"),
1709 "expired cooldown should not be loaded"
1710 );
1711 }
1712
1713 #[serial(cooldown_state)]
1714 #[tokio::test]
1715 async fn record_model_failure_writes_to_kv() {
1716 super::reset_global_state().await;
1717 let store = test_store().await;
1718
1719 *cooldown_store().lock().await = Some(store.clone());
1720
1721 record_model_failure("testagent_persist", "testmodel_persist").await;
1722
1723 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1725
1726 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1727 let stored = store.kv_get(&kv_key).await.unwrap();
1728 assert!(stored.is_some(), "model failure should be persisted to KV");
1729 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1730 let now = chrono::Utc::now().timestamp();
1731 assert!(ts > now, "persisted timestamp should be in the future");
1733 assert_eq!(reason, "model_error");
1734 }
1735
1736 #[serial(cooldown_state)]
1737 #[tokio::test]
1738 async fn cooldown_reason_survives_restart_via_kv() {
1739 super::reset_global_state().await;
1740 let store = test_store().await;
1741 let agent = "test_reason_restart";
1742 let future = chrono::Utc::now().timestamp() + 600;
1743
1744 init_cooldown_store(store.clone()).await;
1745 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1746
1747 {
1749 let mut map = cooldown_lock();
1750 map.clear();
1751 }
1752 init_cooldown_store(store).await;
1753
1754 let entries = list_all_cooldowns();
1755 let entry = entries.iter().find(|(key, ..)| key == agent);
1756 assert_eq!(
1757 entry.map(|(_, _, reason)| reason.as_str()),
1758 Some("billing_cycle_exhausted"),
1759 "cooldown reason must survive a restart, not collapse to 'persisted'"
1760 );
1761 }
1762
1763 #[serial(cooldown_state)]
1764 #[tokio::test]
1765 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1766 super::reset_global_state().await;
1767 let store = test_store().await;
1768 *cooldown_store().lock().await = Some(store.clone());
1769
1770 let agent = "testagent_persistent";
1771 let model = "testmodel_persistent";
1772 let key = format!("{agent}:{model}");
1773
1774 record_persistent_model_failure(agent, model).await;
1776 let now = chrono::Utc::now().timestamp();
1777 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1778 let first_remaining = first_until.saturating_sub(now);
1779 assert!(
1780 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1781 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1782 );
1783
1784 record_persistent_model_failure(agent, model).await;
1786 let now2 = chrono::Utc::now().timestamp();
1787 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1788 let second_remaining = second_until.saturating_sub(now2);
1789 assert!(
1790 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1791 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1792 );
1793
1794 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1795 let stored_count = store
1796 .kv_get(&count_key)
1797 .await
1798 .expect("kv read should succeed")
1799 .expect("failure count should be written");
1800 assert_eq!(stored_count, "2");
1801 }
1802
1803 #[serial(cooldown_state)]
1804 #[tokio::test]
1805 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1806 super::reset_global_state().await;
1807 let store = test_store().await;
1808 *cooldown_store().lock().await = Some(store.clone());
1809
1810 let agent = "testagent_persistent_success";
1811 let model = "testmodel_persistent_success";
1812 let key = format!("{agent}:{model}");
1813
1814 record_persistent_model_failure(agent, model).await;
1816 record_persistent_model_failure(agent, model).await;
1817
1818 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1819 let stored_count = store
1820 .kv_get(&count_key)
1821 .await
1822 .expect("kv read should succeed")
1823 .expect("failure count should be written");
1824 assert_eq!(stored_count, "2");
1825
1826 record_agent_success(agent, model).await;
1830
1831 let stored_count_after_success = store
1832 .kv_get(&count_key)
1833 .await
1834 .expect("kv read should succeed")
1835 .expect("persistent failure count should still exist");
1836 assert_eq!(
1837 stored_count_after_success, "2",
1838 "persistent failure count must survive an intervening success"
1839 );
1840
1841 record_persistent_model_failure(agent, model).await;
1843 let now = chrono::Utc::now().timestamp();
1844 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1845 let third_remaining = third_until.saturating_sub(now);
1846 assert!(
1847 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1848 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1849 );
1850 }
1851
1852 #[serial(cooldown_state)]
1853 #[tokio::test]
1854 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1855 super::reset_global_state().await;
1856 let store = test_store().await;
1857 *cooldown_store().lock().await = Some(store.clone());
1858
1859 let agent = "testagent_persistent_capped";
1860 let model = "testmodel_persistent_capped";
1861 let key = format!("{agent}:{model}");
1862
1863 for _ in 0..9 {
1864 record_persistent_model_failure(agent, model).await;
1865 }
1866
1867 let remaining = cooldown_until(&key)
1868 .expect("cooldown should exist")
1869 .saturating_sub(chrono::Utc::now().timestamp());
1870 assert!(
1871 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1872 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1873 );
1874 assert!(
1875 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1876 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1877 );
1878 }
1879
1880 #[serial(cooldown_state)]
1881 #[test]
1882 fn parse_retry_at_codex_format() {
1883 let msg = "You've hit your usage limit. Upgrade to Pro, visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at Mar 26th, 2026 5:55 AM.";
1884 let ts = parse_retry_at(msg);
1885 assert!(ts.is_some(), "should parse codex retry-at date");
1886
1887 use chrono::TimeZone;
1890 let naive = chrono::NaiveDateTime::new(
1891 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1892 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1893 );
1894 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1895 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1896 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1897 chrono::LocalResult::None => naive.and_utc(),
1898 };
1899 assert_eq!(
1900 ts.unwrap(),
1901 expected_utc.timestamp(),
1902 "timestamp should reflect local-timezone interpretation of the vendor date"
1903 );
1904
1905 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1907 use chrono::Datelike;
1908 assert_eq!(dt.month(), 3);
1909 assert_eq!(dt.day(), 26);
1910 }
1911
1912 #[serial(cooldown_state)]
1913 #[test]
1914 fn parse_retry_at_hour_only_and_month_day() {
1915 let now = chrono::Utc::now().timestamp();
1916 for msg in [
1917 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1918 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1919 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1920 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1921 ] {
1922 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1923 assert!(ts > now, "{msg}");
1924 }
1925 }
1926
1927 #[serial(cooldown_state)]
1928 #[test]
1929 fn parse_retry_at_weekly_resets_without_year() {
1930 use chrono::Datelike;
1931 let t = chrono::Local::now() + chrono::Duration::days(2);
1932 let msg = format!(
1933 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1934 t.format("%b"),
1935 t.day()
1936 );
1937 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1938 let remaining = ts - chrono::Utc::now().timestamp();
1939 assert!(
1940 remaining > 3600 && remaining < 3 * 86400,
1941 "expected ~2d retry-at, got {remaining}s"
1942 );
1943 }
1944
1945 #[serial(cooldown_state)]
1946 #[test]
1947 fn parse_retry_at_no_date() {
1948 assert!(parse_retry_at("generic rate limit error").is_none());
1949 assert!(parse_retry_at("").is_none());
1950 }
1951
1952 #[serial(cooldown_state)]
1953 #[test]
1954 fn parse_retry_at_glm_reset_format_with_seconds() {
1955 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
1956 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
1957
1958 use chrono::TimeZone;
1959 let naive = chrono::NaiveDateTime::new(
1960 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
1961 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
1962 );
1963 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1964 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1965 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1966 chrono::LocalResult::None => naive.and_utc(),
1967 };
1968
1969 assert_eq!(
1970 ts,
1971 expected_utc.timestamp(),
1972 "timestamp should reflect local-timezone interpretation of vendor reset date"
1973 );
1974 }
1975
1976 #[serial(cooldown_state)]
1977 #[test]
1978 fn parse_retry_at_billing_cycle_without_date_returns_none() {
1979 let msg = "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle. Upgrade to get more.";
1982 assert!(
1983 parse_retry_at(msg).is_none(),
1984 "billing cycle without date should return None"
1985 );
1986 }
1987
1988 #[serial(cooldown_state)]
1989 #[test]
1990 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
1991 let msg = "Failed to authenticate. API Error: 403 You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
1992 assert_eq!(
1993 parse_relative_usage_window(msg),
1994 Some(5 * 3600),
1995 "should extract the 5-hour window as seconds"
1996 );
1997 }
1998
1999 #[serial(cooldown_state)]
2000 #[test]
2001 fn parse_relative_usage_window_handles_common_durations() {
2002 assert_eq!(
2003 parse_relative_usage_window("You've reached your hourly usage limit."),
2004 Some(3600)
2005 );
2006 assert_eq!(
2007 parse_relative_usage_window("You've reached your daily usage limit."),
2008 Some(86400)
2009 );
2010 assert_eq!(
2011 parse_relative_usage_window("You've reached your weekly limit."),
2012 Some(604800)
2013 );
2014 assert_eq!(
2015 parse_relative_usage_window("You've reached your 30-minute usage limit."),
2016 Some(1800)
2017 );
2018 assert_eq!(
2019 parse_relative_usage_window("You've reached your 2-day quota."),
2020 Some(172800)
2021 );
2022 }
2023
2024 #[serial(cooldown_state)]
2025 #[test]
2026 fn parse_relative_usage_window_returns_none_without_usage_context() {
2027 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
2030 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
2031 assert!(parse_relative_usage_window("").is_none());
2032 }
2033
2034 #[serial(cooldown_state)]
2035 #[tokio::test]
2036 async fn record_rate_limit_uses_model_specific_window_cooldown() {
2037 super::reset_global_state().await;
2038 let store = test_store().await;
2039 *cooldown_store().lock().await = Some(store.clone());
2040
2041 let agent = "test_record_rl_window_agent";
2042 let model = "test_record_rl_window_model";
2043 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
2044
2045 record_rate_limit(agent, Some(model), msg).await;
2046
2047 assert!(
2048 is_model_in_cooldown(agent, model),
2049 "model should be cooled for the relative usage window"
2050 );
2051 assert!(
2052 !is_agent_in_cooldown(agent),
2053 "agent should not be cooled for a model-scoped quota"
2054 );
2055
2056 let key = format!("{agent}:{model}");
2057 let remaining = cooldown_until(&key)
2058 .expect("model cooldown should exist")
2059 .saturating_sub(chrono::Utc::now().timestamp());
2060 assert!(
2061 remaining >= 5 * 3600 - 5,
2062 "model cooldown should be at least ~5 hours, got {remaining}s"
2063 );
2064 assert_eq!(
2065 cooldown_reason(&key).as_deref(),
2066 Some("rate_limit_usage_window"),
2067 "relative usage-window cooldown must be labeled as a rate limit, not silence"
2068 );
2069 }
2070
2071 #[serial(cooldown_state)]
2072 #[tokio::test]
2073 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
2074 super::reset_global_state().await;
2075 let store = test_store().await;
2076 *cooldown_store().lock().await = Some(store.clone());
2077
2078 let agent = "test_record_rl_fallback_agent";
2079 let model = "test_record_rl_fallback_model";
2080 let msg = "429 Too Many Requests";
2081
2082 record_rate_limit(agent, Some(model), msg).await;
2083
2084 assert!(
2085 is_agent_in_cooldown(agent),
2086 "agent should be cooled when no relative window is detected"
2087 );
2088 assert!(
2089 is_model_in_cooldown(agent, model),
2090 "model should be cooled when no relative window is detected"
2091 );
2092 }
2093
2094 #[serial(cooldown_state)]
2095 #[tokio::test]
2096 async fn agent_cooldown_persists_to_kv() {
2097 super::reset_global_state().await;
2098 let agent = "test_agent_persist_check";
2100 record_agent_failure_with_message(agent, "").await;
2101 assert!(is_agent_in_cooldown(agent));
2102 }
2103
2104 #[serial(cooldown_state)]
2105 #[tokio::test]
2106 async fn cooldown_never_shortens() {
2107 super::reset_global_state().await;
2108 let agent = "test_agent_no_shorten";
2109 set_agent_cooldown(agent, 24 * 60 * 60).await;
2110 let initial = {
2111 let map = cooldowns().lock().unwrap();
2112 map.get(agent)
2113 .expect("cooldown entry should exist")
2114 .cooldown_until
2115 };
2116
2117 record_agent_failure_with_message(agent, "rate limit").await;
2119 let after = {
2120 let map = cooldowns().lock().unwrap();
2121 map.get(agent)
2122 .expect("cooldown entry should exist")
2123 .cooldown_until
2124 };
2125 assert_eq!(
2126 initial, after,
2127 "cooldown should not be shortened by later failures"
2128 );
2129 }
2130
2131 #[serial(cooldown_state)]
2132 #[tokio::test]
2133 async fn silence_agent_cooldown_is_short_lived() {
2134 super::reset_global_state().await;
2135 let agent = "test_silence_agent_cd";
2136 assert!(!is_agent_in_cooldown(agent));
2137
2138 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2139 assert!(is_agent_in_cooldown(agent));
2140
2141 let map = cooldowns().lock().unwrap();
2143 let entry = map.get(agent).expect("should have cooldown entry");
2144 let now = chrono::Utc::now().timestamp();
2145 let remaining = entry.cooldown_until - now;
2146 assert!(
2147 remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
2148 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2149 );
2150 assert!(remaining > 0, "cooldown should still be active");
2151 }
2152
2153 #[serial(cooldown_state)]
2154 #[tokio::test]
2155 async fn record_silence_detection_applies_extended_cooldown() {
2156 super::reset_global_state().await;
2157 let store = test_store().await;
2158 *cooldown_store().lock().await = Some(store.clone());
2159
2160 let agent = "test_silence_count_agent";
2161 let model = "test_silence_count_model";
2162 assert!(!is_model_in_cooldown(agent, model));
2163
2164 for _ in 0..SILENCE_COUNT_THRESHOLD {
2165 let result = record_silence_detection(agent, model).await;
2166 assert!(result.is_some());
2167 }
2168
2169 assert!(is_model_in_cooldown(agent, model));
2170 let key = format!("{agent}:{model}");
2171 let remaining = {
2172 let map = cooldowns().lock().unwrap();
2173 let entry = map.get(&key).expect("cooldown entry should exist");
2174 entry.cooldown_until - chrono::Utc::now().timestamp()
2175 };
2176 assert!(
2177 remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
2178 "extended cooldown should be applied, got {remaining}s"
2179 );
2180
2181 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2182 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2183 assert_eq!(stored, "[]");
2184 }
2185
2186 #[serial(cooldown_state)]
2187 #[test]
2188 fn detect_credit_exhaustion_out_of_credits() {
2189 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2190 assert_eq!(
2191 detect_credit_exhaustion(msg),
2192 Some(CreditExhaustionReason::OutOfCredits)
2193 );
2194 }
2195
2196 #[serial(cooldown_state)]
2197 #[test]
2198 fn detect_credit_exhaustion_out_of_credits_variations() {
2199 assert_eq!(
2200 detect_credit_exhaustion("out of credits"),
2201 Some(CreditExhaustionReason::OutOfCredits)
2202 );
2203 assert_eq!(
2204 detect_credit_exhaustion("outofcredits"),
2205 Some(CreditExhaustionReason::OutOfCredits)
2206 );
2207 assert_eq!(
2208 detect_credit_exhaustion("no credits remaining"),
2209 Some(CreditExhaustionReason::OutOfCredits)
2210 );
2211 assert_eq!(
2212 detect_credit_exhaustion("insufficient funds"),
2213 Some(CreditExhaustionReason::OutOfCredits)
2214 );
2215 }
2216
2217 #[serial(cooldown_state)]
2218 #[test]
2219 fn detect_credit_exhaustion_org_level_disabled() {
2220 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2221 assert_eq!(
2222 detect_credit_exhaustion(msg),
2223 Some(CreditExhaustionReason::OrgLevelDisabled)
2224 );
2225 }
2226
2227 #[serial(cooldown_state)]
2228 #[test]
2229 fn detect_credit_exhaustion_org_level_disabled_variations() {
2230 assert_eq!(
2231 detect_credit_exhaustion("organization has been disabled"),
2232 Some(CreditExhaustionReason::OrgLevelDisabled)
2233 );
2234 assert_eq!(
2235 detect_credit_exhaustion("org disabled"),
2236 Some(CreditExhaustionReason::OrgLevelDisabled)
2237 );
2238 assert_eq!(
2239 detect_credit_exhaustion("billing quota exceeded"),
2240 Some(CreditExhaustionReason::OrgLevelDisabled)
2241 );
2242 assert_eq!(
2243 detect_credit_exhaustion("credit balance too low"),
2244 Some(CreditExhaustionReason::OrgLevelDisabled)
2245 );
2246 }
2247
2248 #[serial(cooldown_state)]
2249 #[test]
2250 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2251 let msg = "overageDisabledReason: out_of_credits";
2252 assert_eq!(
2253 detect_credit_exhaustion(msg),
2254 Some(CreditExhaustionReason::OutOfCredits)
2255 );
2256
2257 let msg = "overageDisabledReason: org_level_disabled";
2258 assert_eq!(
2259 detect_credit_exhaustion(msg),
2260 Some(CreditExhaustionReason::OrgLevelDisabled)
2261 );
2262 }
2263
2264 #[serial(cooldown_state)]
2265 #[test]
2266 fn detect_credit_exhaustion_glm_insufficient_balance() {
2267 assert_eq!(
2269 detect_credit_exhaustion(
2270 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2271 ),
2272 Some(CreditExhaustionReason::OutOfCredits)
2273 );
2274 assert_eq!(
2275 detect_credit_exhaustion("Insufficient balance"),
2276 Some(CreditExhaustionReason::OutOfCredits)
2277 );
2278 assert_eq!(
2279 detect_credit_exhaustion("please recharge your account"),
2280 Some(CreditExhaustionReason::OutOfCredits)
2281 );
2282 assert_eq!(
2283 detect_credit_exhaustion("no resource package available"),
2284 Some(CreditExhaustionReason::OutOfCredits)
2285 );
2286 }
2287
2288 #[serial(cooldown_state)]
2289 #[test]
2290 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2291 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2292 assert!(detect_credit_exhaustion("too many requests").is_none());
2293 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2294 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2295 assert!(detect_credit_exhaustion(
2296 "You've hit your usage limit, try again in a few minutes."
2297 )
2298 .is_none());
2299 }
2300
2301 #[serial(cooldown_state)]
2302 #[tokio::test]
2303 async fn record_credit_exhaustion_applies_agent_cooldown() {
2304 super::reset_global_state().await;
2305 let agent = "test_credit_exhaust_agent";
2306 assert!(!is_agent_in_cooldown(agent));
2307
2308 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2309 assert!(is_agent_in_cooldown(agent));
2310
2311 let remaining = {
2312 let map = cooldowns().lock().unwrap();
2313 let entry = map.get(agent).expect("cooldown entry should exist");
2314 entry.cooldown_until - chrono::Utc::now().timestamp()
2315 };
2316 assert!(
2318 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2319 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2320 );
2321 assert!(
2322 remaining <= CREDIT_BACKOFF_MAX_SECS,
2323 "cooldown should not exceed cap of 8h, got {remaining}s"
2324 );
2325 }
2326
2327 #[serial(cooldown_state)]
2328 #[tokio::test]
2329 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2330 super::reset_global_state().await;
2331 let agent = "test_org_disabled_agent";
2332 assert!(!is_agent_in_cooldown(agent));
2333
2334 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2335 assert!(is_agent_in_cooldown(agent));
2336
2337 let remaining = {
2338 let map = cooldowns().lock().unwrap();
2339 let entry = map.get(agent).expect("cooldown entry should exist");
2340 entry.cooldown_until - chrono::Utc::now().timestamp()
2341 };
2342 assert!(
2344 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2345 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2346 );
2347 assert!(
2348 remaining <= CREDIT_BACKOFF_MAX_SECS,
2349 "cooldown should not exceed cap of 8h, got {remaining}s"
2350 );
2351 }
2352
2353 #[serial(cooldown_state)]
2354 #[tokio::test]
2355 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2356 super::reset_global_state().await;
2357 let store = test_store().await;
2358 *cooldown_store().lock().await = Some(store.clone());
2359
2360 let agent = "test_billing_cycle_agent";
2361 assert!(!is_agent_in_cooldown(agent));
2362
2363 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2364 assert!(is_agent_in_cooldown(agent));
2365
2366 let remaining = {
2367 let map = cooldowns().lock().unwrap();
2368 let entry = map.get(agent).expect("cooldown entry should exist");
2369 entry.cooldown_until - chrono::Utc::now().timestamp()
2370 };
2371 assert!(
2373 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2374 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2375 );
2376 assert!(
2377 remaining <= BILLING_CYCLE_MAX_SECS,
2378 "cooldown should not exceed cap of 7 days, got {remaining}s"
2379 );
2380 }
2381
2382 #[serial(cooldown_state)]
2383 #[tokio::test]
2384 async fn record_credit_exhaustion_billing_cycle_escalates() {
2385 super::reset_global_state().await;
2386 let store = test_store().await;
2387 *cooldown_store().lock().await = Some(store.clone());
2388
2389 let agent = "test_billing_cycle_escalation";
2390
2391 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2393 let remaining_1 = {
2394 let map = cooldowns().lock().unwrap();
2395 let entry = map.get(agent).expect("cooldown entry should exist");
2396 entry.cooldown_until - chrono::Utc::now().timestamp()
2397 };
2398 assert!(
2399 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2400 .contains(&remaining_1),
2401 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2402 );
2403
2404 {
2406 let mut map = cooldowns().lock().unwrap();
2407 map.remove(agent);
2408 }
2409
2410 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2412 let remaining_2 = {
2413 let map = cooldowns().lock().unwrap();
2414 let entry = map.get(agent).expect("cooldown entry should exist");
2415 entry.cooldown_until - chrono::Utc::now().timestamp()
2416 };
2417 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2419 remaining_2 >= expected_2 - 5,
2420 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2421 );
2422
2423 {
2425 let mut map = cooldowns().lock().unwrap();
2426 map.remove(agent);
2427 }
2428
2429 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2431 let remaining_3 = {
2432 let map = cooldowns().lock().unwrap();
2433 let entry = map.get(agent).expect("cooldown entry should exist");
2434 entry.cooldown_until - chrono::Utc::now().timestamp()
2435 };
2436 assert!(
2437 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2438 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2439 );
2440 assert!(
2441 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2442 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2443 );
2444 }
2445
2446 #[serial(cooldown_state)]
2447 #[tokio::test]
2448 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2449 super::reset_global_state().await;
2450 let store = test_store().await;
2451 *cooldown_store().lock().await = Some(store.clone());
2452
2453 let agent = "test_billing_cycle_agent_capped";
2454
2455 for _ in 0..5 {
2456 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2457 }
2458
2459 let remaining = {
2460 let map = cooldowns().lock().unwrap();
2461 let entry = map.get(agent).expect("cooldown entry should exist");
2462 entry.cooldown_until - chrono::Utc::now().timestamp()
2463 };
2464 assert!(
2465 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2466 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2467 );
2468 assert!(
2469 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2470 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2471 );
2472 }
2473
2474 #[serial(cooldown_state)]
2475 #[test]
2476 fn credit_exhaustion_backoff_constants() {
2477 assert_eq!(
2479 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2480 CREDIT_BACKOFF_BASE_SECS
2481 );
2482 assert_eq!(
2484 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2485 ORG_BACKOFF_BASE_SECS
2486 );
2487 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2489 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2490 }
2491
2492 #[serial(cooldown_state)]
2493 #[test]
2494 fn compute_backoff_long_cap_without_extended_stays_capped() {
2495 assert_eq!(
2496 compute_backoff(
2497 3,
2498 BILLING_CYCLE_COOLDOWN_SECS,
2499 BILLING_CYCLE_MAX_SECS,
2500 false
2501 ),
2502 BILLING_CYCLE_MAX_SECS,
2503 "7-day capped paths should stop at the documented max"
2504 );
2505 assert_eq!(
2506 compute_backoff(
2507 9,
2508 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2509 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2510 false,
2511 ),
2512 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2513 "persistent model failures should not extend past 7 days"
2514 );
2515 }
2516
2517 #[serial(cooldown_state)]
2520 #[test]
2521 fn mark_and_check_agent_degraded() {
2522 let agent = "test_degraded_mark";
2523 assert!(!is_agent_degraded(agent));
2524 mark_agent_degraded(agent);
2525 assert!(is_agent_degraded(agent));
2526 clear_agent_degraded(agent);
2527 assert!(!is_agent_degraded(agent));
2528 }
2529
2530 #[serial(cooldown_state)]
2531 #[tokio::test]
2532 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2533 super::reset_global_state().await;
2534 let store = test_store().await;
2535 let agent = "test_sync_external_clear";
2536 let future = chrono::Utc::now().timestamp() + 600;
2537
2538 init_cooldown_store(store.clone()).await;
2540 set_cooldown_async(agent, future, "test").await;
2541 assert!(is_agent_in_cooldown(agent));
2542
2543 store
2547 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2548 .await
2549 .expect("kv_set tombstone");
2550 assert!(
2551 is_agent_in_cooldown(agent),
2552 "in-memory cooldown should still be present before sync"
2553 );
2554
2555 sync_from_kv(&store).await;
2556
2557 assert!(
2558 !is_agent_in_cooldown(agent),
2559 "sync_from_kv must drop in-memory entries the KV tombstoned"
2560 );
2561 }
2562
2563 #[serial(cooldown_state)]
2564 #[tokio::test]
2565 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2566 super::reset_global_state().await;
2570 let store = test_store().await;
2571 let agent = "test_sync_race_window";
2572 let future = chrono::Utc::now().timestamp() + 600;
2573
2574 set_cooldown_in_memory(agent, future, "raced");
2576 assert!(is_agent_in_cooldown(agent));
2577
2578 sync_from_kv(&store).await;
2579
2580 assert!(
2581 is_agent_in_cooldown(agent),
2582 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2583 );
2584 }
2585
2586 #[serial(cooldown_state)]
2587 #[tokio::test]
2588 async fn sync_from_kv_adopts_kv_timestamp() {
2589 super::reset_global_state().await;
2592 let store = test_store().await;
2593 let agent = "test_sync_adopt_kv";
2594 let original = chrono::Utc::now().timestamp() + 600;
2595 let updated = chrono::Utc::now().timestamp() + 1200;
2596
2597 init_cooldown_store(store.clone()).await;
2598 set_cooldown_async(agent, original, "initial").await;
2599
2600 store
2601 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2602 .await
2603 .expect("kv_set updated");
2604
2605 sync_from_kv(&store).await;
2606
2607 let until = cooldown_until(agent).expect("cooldown still active");
2608 assert_eq!(
2609 until, updated,
2610 "sync_from_kv must adopt KV's authoritative timestamp"
2611 );
2612 }
2613
2614 #[serial(cooldown_state)]
2615 #[tokio::test]
2616 async fn refresh_degraded_agents_marks_cooled_agent() {
2617 super::reset_global_state().await;
2618 let store = test_store().await;
2619 let agent = "test_refresh_cooled";
2620
2621 record_agent_failure_with_message(agent, "").await;
2623 assert!(is_agent_in_cooldown(agent));
2624
2625 let agents = vec![agent.to_string()];
2627 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2628
2629 assert!(
2630 is_agent_degraded(agent),
2631 "agent in cooldown should be marked degraded"
2632 );
2633
2634 clear_agent_degraded(agent);
2636 }
2637
2638 #[serial(cooldown_state)]
2639 #[tokio::test]
2640 async fn refresh_degraded_agents_marks_no_models_agent() {
2641 super::reset_global_state().await;
2642 let store = test_store().await;
2643 let agent = "test_refresh_no_models";
2644
2645 let agents = vec![agent.to_string()];
2647 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2648
2649 assert!(
2650 is_agent_degraded(agent),
2651 "agent with no available models should be marked degraded"
2652 );
2653
2654 clear_agent_degraded(agent);
2656 }
2657
2658 #[serial(cooldown_state)]
2659 #[tokio::test]
2660 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2661 super::reset_global_state().await;
2662 let store = test_store().await;
2663 let agent = "test_refresh_rate_limited";
2664
2665 for _ in 0..4 {
2667 store.record_rate_limit(agent, "rate", None).await.unwrap();
2668 }
2669
2670 let agents = vec![agent.to_string()];
2671 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2673
2674 assert!(
2675 is_agent_degraded(agent),
2676 "agent with rate limit count >= threshold should be marked degraded"
2677 );
2678
2679 clear_agent_degraded(agent);
2681 }
2682
2683 #[serial(cooldown_state)]
2684 #[tokio::test]
2685 async fn refresh_degraded_agents_clears_healthy_agent() {
2686 super::reset_global_state().await;
2687 let store = test_store().await;
2688 let agent = "test_refresh_healthy";
2689
2690 mark_agent_degraded(agent);
2692 assert!(is_agent_degraded(agent));
2693
2694 let agents = vec![agent.to_string()];
2696 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2697
2698 assert!(
2699 !is_agent_degraded(agent),
2700 "healthy agent should have degraded flag cleared"
2701 );
2702 }
2703
2704 #[serial(cooldown_state)]
2705 #[tokio::test]
2706 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2707 super::reset_global_state().await;
2708 let store = test_store().await;
2709 let agent = "test_refresh_expired_cooldown";
2710
2711 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2715 let mut map = cooldowns().lock().unwrap();
2716 map.insert(
2717 agent.to_string(),
2718 CooldownEntry {
2719 cooldown_until: expired_ts,
2720 reason: "expired_test".to_string(),
2721 },
2722 );
2723 }
2724
2725 assert!(
2727 !is_agent_in_cooldown(agent),
2728 "expired cooldown should not be active"
2729 );
2730
2731 mark_agent_degraded(agent);
2733 assert!(is_agent_degraded(agent));
2734
2735 let agents = vec![agent.to_string()];
2737 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2738
2739 assert!(
2740 !is_agent_degraded(agent),
2741 "agent with only expired cooldown should NOT be marked degraded"
2742 );
2743
2744 clear_agent_degraded(agent);
2746 {
2747 let mut map = cooldowns().lock().unwrap();
2748 map.remove(agent);
2749 }
2750 }
2751
2752 #[serial(cooldown_state)]
2753 #[test]
2754 fn compute_backoff_grows_exponentially() {
2755 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2757 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2758 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2760 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2762 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2764 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2766 }
2767
2768 #[serial(cooldown_state)]
2769 #[test]
2770 fn compute_backoff_extended_tier_1_applies_24h() {
2771 assert_eq!(
2777 compute_backoff(5, 300, 14400, true),
2778 14400,
2779 "first at max, not extended"
2780 );
2781 assert_eq!(
2782 compute_backoff(6, 300, 14400, true),
2783 14400,
2784 "times_at_max=1, not extended"
2785 );
2786 assert_eq!(
2787 compute_backoff(7, 300, 14400, true),
2788 86400,
2789 "times_at_max=2, extended 6x"
2790 );
2791 assert_eq!(
2792 compute_backoff(8, 300, 14400, true),
2793 86400,
2794 "times_at_max=3, still extended 6x"
2795 );
2796 }
2797
2798 #[serial(cooldown_state)]
2799 #[test]
2800 fn compute_backoff_extended_tier_2_applies_48h() {
2801 assert_eq!(
2804 compute_backoff(9, 300, 14400, true),
2805 86400,
2806 "times_at_max=4, still tier 1"
2807 );
2808 assert_eq!(
2809 compute_backoff(10, 300, 14400, true),
2810 172800,
2811 "times_at_max=5, extended 12x"
2812 );
2813 assert_eq!(
2814 compute_backoff(15, 300, 14400, true),
2815 172800,
2816 "times_at_max=10, still tier 2"
2817 );
2818 assert_eq!(
2819 compute_backoff(20, 300, 14400, true),
2820 172800,
2821 "times_at_max=15, still tier 2"
2822 );
2823 }
2824
2825 #[serial(cooldown_state)]
2826 #[test]
2827 fn compute_backoff_extended_tier_with_different_base_max() {
2828 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2835 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2836 assert_eq!(
2837 compute_backoff(7, 60, 3600, true),
2838 21600,
2839 "times_at_max=2, extended 6x"
2840 );
2841 }
2842
2843 #[serial(cooldown_state)]
2844 #[test]
2845 fn compute_backoff_extended_tier_persistent_fails_example() {
2846 assert_eq!(
2850 compute_backoff(21, 300, 14400, true),
2851 172800,
2852 "21 failures should get 48h extended cooldown"
2853 );
2854
2855 assert_eq!(
2857 compute_backoff(36, 300, 14400, true),
2858 172800,
2859 "36 failures should also get 48h extended cooldown"
2860 );
2861 }
2862
2863 #[serial(cooldown_state)]
2864 #[test]
2865 fn compute_backoff_extended_tier_zero_base() {
2866 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2868 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2869 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2870 }
2871
2872 #[serial(cooldown_state)]
2873 #[test]
2874 fn detect_credit_exhaustion_billing_cycle() {
2875 assert_eq!(
2876 detect_credit_exhaustion("billing cycle"),
2877 Some(CreditExhaustionReason::BillingCycleExhausted)
2878 );
2879 assert_eq!(
2880 detect_credit_exhaustion("quota refreshed next cycle"),
2881 Some(CreditExhaustionReason::BillingCycleExhausted)
2882 );
2883 assert_eq!(
2884 detect_credit_exhaustion(
2885 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
2886 ),
2887 Some(CreditExhaustionReason::BillingCycleExhausted)
2888 );
2889 assert_eq!(
2890 detect_credit_exhaustion("monthly quota exceeded"),
2891 Some(CreditExhaustionReason::BillingCycleExhausted)
2892 );
2893 assert_eq!(
2894 detect_credit_exhaustion(
2895 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
2896 ),
2897 Some(CreditExhaustionReason::BillingCycleExhausted)
2898 );
2899 }
2900
2901 #[serial(cooldown_state)]
2902 #[test]
2903 fn detect_credit_exhaustion_billing_cycle_minimax() {
2904 assert_eq!(
2905 detect_credit_exhaustion(
2906 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
2907 ),
2908 Some(CreditExhaustionReason::BillingCycleExhausted)
2909 );
2910 assert_eq!(
2911 detect_credit_exhaustion("Upgrade your Token Plan"),
2912 Some(CreditExhaustionReason::BillingCycleExhausted)
2913 );
2914 assert_eq!(
2915 detect_credit_exhaustion("Token Plan usage limit reached"),
2916 Some(CreditExhaustionReason::BillingCycleExhausted)
2917 );
2918 assert_eq!(
2919 detect_credit_exhaustion("purchase Credits for more usage"),
2920 Some(CreditExhaustionReason::BillingCycleExhausted)
2921 );
2922 assert_eq!(
2923 detect_credit_exhaustion(
2924 "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits"
2925 ),
2926 Some(CreditExhaustionReason::BillingCycleExhausted)
2927 );
2928 }
2929
2930 #[serial(cooldown_state)]
2931 #[test]
2932 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
2933 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
2938 this billing cycle. Your quota will be refreshed soon.";
2939 assert_eq!(
2940 detect_credit_exhaustion_outcome(msg),
2941 Some("billing_cycle_exhausted")
2942 );
2943 }
2944
2945 #[serial(cooldown_state)]
2946 #[test]
2947 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
2948 assert_eq!(
2949 detect_credit_exhaustion_outcome("out of credits"),
2950 Some("credit_exhausted")
2951 );
2952 assert_eq!(
2953 detect_credit_exhaustion_outcome("organization has been disabled"),
2954 Some("org_level_disabled")
2955 );
2956 assert_eq!(
2957 detect_credit_exhaustion_outcome("monthly quota exceeded"),
2958 Some("billing_cycle_exhausted")
2959 );
2960 }
2961
2962 #[test]
2963 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
2964 assert_eq!(
2965 detect_credit_exhaustion_outcome("429 Too Many Requests"),
2966 None
2967 );
2968 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
2969 }
2970
2971 #[serial(cooldown_state)]
2972 #[test]
2973 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
2974 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
2977 assert_eq!(
2978 detect_credit_exhaustion(msg),
2979 Some(CreditExhaustionReason::OrgLevelDisabled)
2980 );
2981 }
2982
2983 #[serial(cooldown_state)]
2984 #[tokio::test]
2985 async fn refresh_degraded_agents_below_threshold_not_degraded() {
2986 super::reset_global_state().await;
2987 let store = test_store().await;
2988 let agent = "test_refresh_below_threshold";
2989
2990 for _ in 0..2 {
2992 store.record_rate_limit(agent, "rate", None).await.unwrap();
2993 }
2994
2995 let agents = vec![agent.to_string()];
2996 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2998
2999 assert!(
3000 !is_agent_degraded(agent),
3001 "agent with rate limit count < threshold should not be degraded"
3002 );
3003 }
3004
3005 #[serial(cooldown_state)]
3006 #[tokio::test]
3007 async fn refresh_degraded_agents_ignores_error_limit_type() {
3008 super::reset_global_state().await;
3009 let store = test_store().await;
3010 let agent = "test_refresh_error_rows";
3011
3012 for _ in 0..3 {
3013 store.record_rate_limit(agent, "error", None).await.unwrap();
3014 }
3015
3016 let agents = vec![agent.to_string()];
3017 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3018
3019 assert!(
3020 !is_agent_degraded(agent),
3021 "three error rows in 6 hours must not mark an agent degraded"
3022 );
3023 }
3024
3025 #[serial(cooldown_state)]
3026 #[tokio::test]
3027 async fn refresh_degraded_agents_marks_three_rate_rows() {
3028 super::reset_global_state().await;
3029 let store = test_store().await;
3030 let agent = "test_refresh_three_rate_rows";
3031
3032 for _ in 0..3 {
3033 store.record_rate_limit(agent, "rate", None).await.unwrap();
3034 }
3035
3036 let agents = vec![agent.to_string()];
3037 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3038
3039 assert!(
3040 is_agent_degraded(agent),
3041 "three rate rows in 6 hours must mark an agent degraded"
3042 );
3043 assert_eq!(
3044 agent_degraded_reason(agent).as_deref(),
3045 Some("rate limit threshold exceeded")
3046 );
3047 let listed = list_degraded_agents();
3048 assert!(
3049 listed
3050 .iter()
3051 .any(|(a, r)| a == agent && r == "rate limit threshold exceeded"),
3052 "orch cooldown list must show the degraded reason, got {listed:?}"
3053 );
3054
3055 clear_agent_degraded(agent);
3056 }
3057
3058 #[serial(cooldown_state)]
3059 #[tokio::test]
3060 async fn refresh_degraded_agents_logs_reason_change() {
3061 super::reset_global_state().await;
3062 let store = test_store().await;
3063 let agent = "test_refresh_reason_change";
3064 *cooldown_store().lock().await = Some(store.clone());
3065
3066 set_agent_cooldown(agent, 600).await;
3067 let agents = vec![agent.to_string()];
3068 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3069 assert_eq!(
3070 agent_degraded_reason(agent).as_deref(),
3071 Some("agent in cooldown")
3072 );
3073
3074 clear_cooldown(agent, &store).await;
3075 let after_clear = sqlite_datetime_offset(60);
3076 for _ in 0..3 {
3077 store
3078 .record_rate_limit_at(agent, "rate", None, &after_clear)
3079 .await
3080 .unwrap();
3081 }
3082 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3083 assert_eq!(
3084 agent_degraded_reason(agent).as_deref(),
3085 Some("rate limit threshold exceeded"),
3086 "degraded reason must update when the hold changes"
3087 );
3088
3089 clear_agent_degraded(agent);
3090 }
3091
3092 #[serial(cooldown_state)]
3093 #[tokio::test]
3094 async fn clear_un_degrades_on_next_health_check() {
3095 super::reset_global_state().await;
3096 let store = test_store().await;
3097 let agent = "test_clear_un_degrades";
3098 let before = sqlite_datetime_offset(-30);
3099 let after = sqlite_datetime_offset(60);
3100
3101 for _ in 0..3 {
3102 store
3103 .record_rate_limit_at(agent, "rate", None, &before)
3104 .await
3105 .unwrap();
3106 }
3107
3108 let agents = vec![agent.to_string()];
3109 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3110 assert!(
3111 is_agent_degraded(agent),
3112 "three rate rows must mark an agent degraded"
3113 );
3114
3115 clear_cooldown(agent, &store).await;
3116 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3117 assert!(
3118 !is_agent_degraded(agent),
3119 "clear must un-degrade on the next health check"
3120 );
3121
3122 store
3123 .record_rate_limit_at(agent, "rate", None, &after)
3124 .await
3125 .unwrap();
3126 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3127 assert!(
3128 !is_agent_degraded(agent),
3129 "one rate row after the mark must not degrade again"
3130 );
3131
3132 for _ in 0..2 {
3133 store
3134 .record_rate_limit_at(agent, "rate", None, &after)
3135 .await
3136 .unwrap();
3137 }
3138 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3139 assert!(
3140 is_agent_degraded(agent),
3141 "three rate rows after the mark must degrade again"
3142 );
3143
3144 clear_agent_degraded(agent);
3145 }
3146
3147 #[serial(cooldown_state)]
3148 #[tokio::test]
3149 async fn clear_all_un_degrades_rate_limited_agent() {
3150 super::reset_global_state().await;
3151 let store = test_store().await;
3152 let agent = "test_clear_all_un_degrades";
3153 let before = sqlite_datetime_offset(-30);
3154
3155 for _ in 0..3 {
3156 store
3157 .record_rate_limit_at(agent, "rate", None, &before)
3158 .await
3159 .unwrap();
3160 }
3161
3162 let agents = vec![agent.to_string()];
3163 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3164 assert!(is_agent_degraded(agent));
3165
3166 clear_cooldown("*", &store).await;
3167 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3168 assert!(
3169 !is_agent_degraded(agent),
3170 "clear --all must un-degrade an agent held only by rate rows"
3171 );
3172 }
3173
3174 #[serial(cooldown_state)]
3177 #[tokio::test]
3178 async fn repeated_agent_failures_escalate_backoff() {
3179 super::reset_global_state().await;
3180 let store = test_store().await;
3181 *cooldown_store().lock().await = Some(store.clone());
3182
3183 let agent = "test_escalation_agent";
3184
3185 record_agent_failure_with_message(agent, "").await;
3187 let remaining_1 = {
3188 let map = cooldowns().lock().unwrap();
3189 let entry = map.get(agent).expect("should have cooldown");
3190 entry.cooldown_until - chrono::Utc::now().timestamp()
3191 };
3192 assert!(
3193 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
3194 "first failure should be ~5 min, got {remaining_1}s"
3195 );
3196
3197 {
3199 let mut map = cooldowns().lock().unwrap();
3200 map.remove(agent);
3201 }
3202
3203 record_agent_failure_with_message(agent, "").await;
3205 let remaining_2 = {
3206 let map = cooldowns().lock().unwrap();
3207 let entry = map.get(agent).expect("should have cooldown");
3208 entry.cooldown_until - chrono::Utc::now().timestamp()
3209 };
3210 assert!(
3211 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
3212 "second failure should be ~15 min, got {remaining_2}s"
3213 );
3214 }
3215
3216 #[serial(cooldown_state)]
3217 #[tokio::test]
3218 async fn vendor_past_retry_falls_back_to_backoff() {
3219 super::reset_global_state().await;
3220 let store = test_store().await;
3224 *cooldown_store().lock().await = Some(store.clone());
3225
3226 let agent = "test_vendor_past_retry";
3227
3228 {
3230 let mut map = cooldowns().lock().unwrap();
3231 map.remove(agent);
3232 }
3233
3234 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
3236
3237 record_agent_failure_with_message(agent, past_msg).await;
3238
3239 assert!(
3243 is_agent_in_cooldown(agent),
3244 "agent should be in cooldown after failure"
3245 );
3246
3247 let remaining = {
3248 let map = cooldowns().lock().unwrap();
3249 let entry = map.get(agent).expect("cooldown entry should exist");
3250 entry.cooldown_until - chrono::Utc::now().timestamp()
3251 };
3252
3253 assert!(
3254 remaining >= BACKOFF_BASE_SECS - 5,
3255 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3256 );
3257
3258 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3260 let stored = store
3261 .kv_get(&kv_key)
3262 .await
3263 .unwrap()
3264 .expect("failure count should be stored");
3265 assert_eq!(stored, "1");
3266 }
3267
3268 #[serial(cooldown_state)]
3269 #[tokio::test]
3270 async fn record_agent_success_resets_backoff() {
3271 super::reset_global_state().await;
3272 let store = test_store().await;
3273 *cooldown_store().lock().await = Some(store.clone());
3274
3275 let agent = "test_success_reset_agent";
3276 let model = "test_model";
3277
3278 for _ in 0..3 {
3280 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3281 }
3282 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3284 let count: u32 = store
3285 .kv_get(&kv_key)
3286 .await
3287 .unwrap()
3288 .unwrap()
3289 .parse()
3290 .unwrap();
3291 assert_eq!(count, 3);
3292
3293 record_agent_success(agent, model).await;
3295
3296 let count_after: u32 = store
3297 .kv_get(&kv_key)
3298 .await
3299 .unwrap()
3300 .unwrap()
3301 .parse()
3302 .unwrap();
3303 assert_eq!(count_after, 0, "success should reset failure count to 0");
3304
3305 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3307 let model_count: u32 = store
3308 .kv_get(&model_kv_key)
3309 .await
3310 .unwrap()
3311 .unwrap()
3312 .parse()
3313 .unwrap();
3314 assert_eq!(
3315 model_count, 0,
3316 "success should reset model failure count to 0"
3317 );
3318 }
3319
3320 #[serial(cooldown_state)]
3321 #[tokio::test]
3322 async fn record_agent_success_without_model_skips_empty_model_key() {
3323 super::reset_global_state().await;
3324 let store = test_store().await;
3325 *cooldown_store().lock().await = Some(store.clone());
3326
3327 let agent = "test_success_without_model";
3328
3329 record_agent_success(agent, "").await;
3330
3331 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3332 let agent_count = store
3333 .kv_get(&agent_kv_key)
3334 .await
3335 .unwrap()
3336 .expect("agent failure count should be stored");
3337 assert_eq!(agent_count, "0");
3338
3339 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3340 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3341 assert!(
3342 empty_model_count.is_none(),
3343 "success without model should not create an empty model failure count key"
3344 );
3345 }
3346
3347 #[serial(cooldown_state)]
3348 #[tokio::test]
3349 async fn clear_cooldown_resets_failure_counts() {
3350 super::reset_global_state().await;
3351 let store = test_store().await;
3352 *cooldown_store().lock().await = Some(store.clone());
3353
3354 let agent = "test_clear_fc_agent";
3355
3356 for _ in 0..5 {
3358 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3359 }
3360 set_agent_cooldown(agent, 3600).await;
3362
3363 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3365 let count: u32 = store
3366 .kv_get(&kv_key)
3367 .await
3368 .unwrap()
3369 .unwrap()
3370 .parse()
3371 .unwrap();
3372 assert_eq!(count, 5);
3373
3374 clear_cooldown(agent, &store).await;
3376
3377 let count_after: u32 = store
3379 .kv_get(&kv_key)
3380 .await
3381 .unwrap()
3382 .unwrap()
3383 .parse()
3384 .unwrap();
3385 assert_eq!(
3386 count_after, 0,
3387 "clear_cooldown should reset failure count to 0"
3388 );
3389
3390 assert!(
3392 !is_agent_in_cooldown(agent),
3393 "agent should not be in cooldown after clear"
3394 );
3395 }
3396
3397 #[serial(cooldown_state)]
3398 #[tokio::test]
3399 async fn clear_all_cooldowns_resets_all_failure_counts() {
3400 super::reset_global_state().await;
3401 let store = test_store().await;
3402 *cooldown_store().lock().await = Some(store.clone());
3403
3404 let agents = ["test_clear_all_a", "test_clear_all_b"];
3405 for agent in &agents {
3406 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3407 set_agent_cooldown(agent, 3600).await;
3408 }
3409
3410 clear_cooldown("*", &store).await;
3411
3412 for agent in &agents {
3413 assert!(!is_agent_in_cooldown(agent));
3414 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3415 let count: u32 = store
3416 .kv_get(&kv_key)
3417 .await
3418 .unwrap()
3419 .unwrap()
3420 .parse()
3421 .unwrap();
3422 assert_eq!(
3423 count, 0,
3424 "clear --all should reset failure count for {agent}"
3425 );
3426 }
3427 }
3428
3429 #[serial(cooldown_state)]
3430 #[test]
3431 fn compute_backoff_saturates_on_large_counts() {
3432 assert_eq!(
3434 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3435 BACKOFF_MAX_SECS * 12
3436 );
3437 assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
3438 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
3439 }
3440
3441 #[serial(cooldown_state)]
3442 #[test]
3443 fn compute_backoff_zero_base_returns_zero() {
3444 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3445 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3446 }
3447
3448 #[serial(cooldown_state)]
3451 #[tokio::test]
3452 async fn github_circuit_stays_closed_below_threshold() {
3453 super::reset_global_state().await;
3454 {
3456 let mut open = github_5xx_circuit_open().lock().unwrap();
3457 *open = None;
3458 let mut ts = github_5xx_timestamps().lock().unwrap();
3459 ts.clear();
3460 let mut map = cooldown_lock();
3464 map.remove("github:5xx");
3465 }
3466
3467 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3469 record_github_5xx().await;
3470 }
3471
3472 assert!(
3473 !is_github_circuit_open(),
3474 "circuit should stay closed below threshold"
3475 );
3476 assert_eq!(github_circuit_remaining_secs(), 0);
3477 }
3478
3479 #[serial(cooldown_state)]
3480 #[tokio::test]
3481 async fn github_circuit_trips_at_threshold() {
3482 super::reset_global_state().await;
3483 {
3485 let mut open = github_5xx_circuit_open_lock();
3486 *open = None;
3487 let mut ts = github_5xx_timestamps_lock();
3488 ts.clear();
3489 let mut map = cooldown_lock();
3490 map.remove("github:5xx");
3491 }
3492
3493 for _ in 0..GITHUB_5XX_THRESHOLD {
3495 record_github_5xx().await;
3496 }
3497
3498 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3499 assert!(
3500 github_circuit_remaining_secs() > 0,
3501 "remaining should be > 0 when open"
3502 );
3503 assert!(
3504 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3505 "remaining should not exceed cooldown duration"
3506 );
3507
3508 {
3510 let mut open = github_5xx_circuit_open().lock().unwrap();
3511 *open = None;
3512 }
3513 }
3514
3515 #[serial(cooldown_state)]
3516 #[tokio::test]
3517 async fn github_circuit_cannot_double_trip() {
3518 super::reset_global_state().await;
3519 {
3521 let mut open = github_5xx_circuit_open_lock();
3522 *open = None;
3523 let mut ts = github_5xx_timestamps_lock();
3524 ts.clear();
3525 let mut map = cooldown_lock();
3526 map.remove("github:5xx");
3527 }
3528
3529 for _ in 0..GITHUB_5XX_THRESHOLD {
3531 record_github_5xx().await;
3532 }
3533 assert!(is_github_circuit_open());
3534
3535 let first_remaining = github_circuit_remaining_secs();
3536
3537 for _ in 0..GITHUB_5XX_THRESHOLD {
3539 record_github_5xx().await;
3540 }
3541
3542 let second_remaining = github_circuit_remaining_secs();
3543 assert!(
3544 second_remaining <= first_remaining,
3545 "additional errors should not extend cooldown"
3546 );
3547
3548 {
3550 let mut open = github_5xx_circuit_open().lock().unwrap();
3551 *open = None;
3552 }
3553 }
3554
3555 #[serial(cooldown_state)]
3556 #[tokio::test]
3557 async fn github_circuit_recovers_after_cooldown() {
3558 super::reset_global_state().await;
3559 {
3561 let mut open = github_5xx_circuit_open_lock();
3562 *open = None;
3563 let mut ts = github_5xx_timestamps_lock();
3564 ts.clear();
3565 let mut map = cooldown_lock();
3566 map.remove("github:5xx");
3567 }
3568
3569 {
3571 let mut open = github_5xx_circuit_open().lock().unwrap();
3572 *open = Some(chrono::Utc::now().timestamp() - 1);
3573 }
3574
3575 assert!(
3577 !is_github_circuit_open(),
3578 "circuit should auto-recover after cooldown"
3579 );
3580 }
3581
3582 #[serial(cooldown_state)]
3583 #[tokio::test]
3584 async fn github_circuit_sliding_window_ages_out() {
3585 super::reset_global_state().await;
3586 {
3588 let mut open = github_5xx_circuit_open_lock();
3589 *open = None;
3590 let mut ts = github_5xx_timestamps_lock();
3591 ts.clear();
3592 let mut map = cooldown_lock();
3593 map.remove("github:5xx");
3594 }
3595
3596 {
3598 let mut ts = github_5xx_timestamps().lock().unwrap();
3599 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3600 for _ in 0..GITHUB_5XX_THRESHOLD {
3601 ts.push(old);
3602 }
3603 }
3604
3605 record_github_5xx().await;
3607
3608 assert!(
3609 !is_github_circuit_open(),
3610 "old errors outside window should not trip circuit"
3611 );
3612 }
3613}