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