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 match explicit_kv.get("github:5xx") {
291 Some(None) => {
292 let mut open = github_5xx_circuit_open_lock();
293 if open.take().is_some() {
294 tracing::info!("GitHub 5xx circuit breaker closed by KV clear");
295 github_5xx_timestamps_lock().clear();
296 }
297 }
298 Some(Some((ts, _))) => {
299 *github_5xx_circuit_open_lock() = Some(*ts);
300 }
301 None => {}
302 }
303
304 let mut map = cooldown_lock();
305 let mut dropped = 0usize;
306 let mut updated = 0usize;
307
308 map.retain(|key, _| {
310 if matches!(explicit_kv.get(key), Some(None)) {
311 dropped += 1;
312 false
313 } else {
314 true
315 }
316 });
317
318 for (key, normalized) in explicit_kv {
320 if let Some((ts, reason)) = normalized {
321 match map.get_mut(&key) {
322 Some(entry) => {
323 if entry.cooldown_until != ts {
324 entry.cooldown_until = ts;
325 updated += 1;
326 }
327 }
328 None => {
329 map.insert(
330 key,
331 CooldownEntry {
332 cooldown_until: ts,
333 reason,
334 },
335 );
336 updated += 1;
337 }
338 }
339 }
340 }
341
342 if dropped > 0 || updated > 0 {
343 tracing::info!(
344 dropped,
345 updated,
346 "cooldown sync_from_kv: reconciled in-memory map with KV"
347 );
348 }
349}
350
351const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
355
356const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
358
359const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
363
364const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
366
367pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
386 if count <= 1 {
387 return base;
388 }
389 let exp = count.saturating_sub(1).min(u32::MAX as u64) as u32;
390 let factor = 3_i64.saturating_pow(exp);
391 let standard = base.saturating_mul(factor).min(max);
392
393 let first_max = if base <= 0 {
402 0
403 } else {
404 let mut accumulated = base;
406 let mut first_max_count = 1u64;
407 while accumulated < max && first_max_count < count {
408 accumulated = accumulated.saturating_mul(3);
409 first_max_count += 1;
410 }
411 first_max_count
412 };
413
414 let times_at_max = count.saturating_sub(first_max);
416
417 if !allow_extended {
418 return standard;
419 }
420
421 if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
422 max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
423 } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
424 max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
425 } else {
426 standard
427 }
428}
429
430async fn read_and_increment_failure_count(
439 store_opt: &Option<Arc<crate::store::TaskStore>>,
440 key: &str,
441) -> u64 {
442 read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
443}
444
445async fn read_and_increment_failure_count_with_prefix(
450 store_opt: &Option<Arc<crate::store::TaskStore>>,
451 prefix: &str,
452 key: &str,
453) -> u64 {
454 let kv_key = format!("{prefix}{key}");
455 if let Some(store) = store_opt {
456 match store.kv_increment(&kv_key).await {
457 Ok(n) => n,
458 Err(e) => {
459 tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
460 u64::MAX
461 }
462 }
463 } else {
464 1
465 }
466}
467
468pub async fn record_agent_success(agent_name: &str, model: &str) {
474 let store_opt = cooldown_store().lock().await.clone();
475 if let Some(store) = store_opt {
476 let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
477 let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
478 if let Err(e) = store.kv_set(&agent_key, "0").await {
479 tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
480 }
481 if !model.is_empty() {
482 let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
483 if let Err(e) = store.kv_set(&model_key, "0").await {
484 tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
485 }
486 }
487 if let Err(e) = store.kv_set(&credit_key, "0").await {
488 tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
489 }
490 }
491}
492
493pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
499 if let Some(cooldown_until) = parse_retry_at(error_message) {
504 let now = chrono::Utc::now().timestamp();
505 if cooldown_until > now {
506 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
507 return;
508 }
509 tracing::debug!(
510 agent = agent_name,
511 cooldown_until,
512 "vendor retry_at is in the past — using backoff instead"
513 );
514 }
515
516 let store_opt = cooldown_store().lock().await.clone();
517 let count = read_and_increment_failure_count(&store_opt, agent_name).await;
518 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
519 let cooldown_until =
520 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
521 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
522}
523
524pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
533 detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
534}
535
536pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
541 let lower = error_message.to_lowercase();
542
543 let overage_patterns = [
544 "overagedisabledreason",
545 "overage_disabled_reason",
546 "overage disabled reason",
547 ];
548
549 if overage_patterns.iter().any(|p| lower.contains(p)) {
550 if lower.contains("out_of_credits") || lower.contains("outofcredits") {
551 return Some(CreditExhaustionReason::OutOfCredits);
552 }
553 if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
554 return Some(CreditExhaustionReason::OrgLevelDisabled);
555 }
556 }
557
558 let billing_patterns = [
559 "credit balance too low",
560 "credit balance insufficient",
561 "insufficient credit balance",
562 "billing quota exceeded",
563 "billing limit exceeded",
564 "organization has been disabled",
565 "org disabled",
566 ];
567
568 if billing_patterns.iter().any(|p| lower.contains(p)) {
569 return Some(CreditExhaustionReason::OrgLevelDisabled);
570 }
571
572 let billing_cycle_patterns = [
575 "billing cycle",
576 "quota refreshed next cycle",
577 "monthly quota",
578 "quota will be refreshed",
579 "next billing cycle",
580 "refreshed in the next cycle",
581 "weekly/monthly limit exhausted",
582 "limit will reset at",
583 "upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
588
589 if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
590 return Some(CreditExhaustionReason::BillingCycleExhausted);
591 }
592
593 if lower.contains("out of credits")
594 || lower.contains("outofcredits")
595 || lower.contains("insufficient funds")
596 || lower.contains("no credits remaining")
597 || lower.contains("insufficient_quota")
598 || lower.contains("quota exceeded")
599 || lower.contains("insufficient balance")
600 || lower.contains("please recharge")
601 || lower.contains("no resource package")
602 || lower.contains("balance exhausted")
603 || lower.contains("402 payment required")
604 {
605 return Some(CreditExhaustionReason::OutOfCredits);
606 }
607
608 None
609}
610
611pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
624 let reason_str = match reason {
625 CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
626 CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
627 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
628 };
629
630 let store_opt = cooldown_store().lock().await.clone();
631 let count = read_and_increment_failure_count_with_prefix(
636 &store_opt,
637 CREDIT_FAILURE_COUNT_PREFIX,
638 agent_name,
639 )
640 .await;
641
642 let (base, max) = match reason {
643 CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
644 CreditExhaustionReason::OrgLevelDisabled => {
645 (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
646 }
647 CreditExhaustionReason::BillingCycleExhausted => {
648 (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
649 }
650 };
651 let cooldown_secs = compute_backoff(
652 count,
653 base,
654 max,
655 !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
656 );
657 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
658 set_cooldown_async(agent_name, cooldown_until, reason_str).await;
659 tracing::warn!(
660 agent = agent_name,
661 reason = reason_str,
662 attempt = count,
663 cooldown_secs,
664 "credit exhaustion detected: applying exponential agent-wide cooldown"
665 );
666}
667
668pub async fn record_model_failure(agent_name: &str, model: &str) {
672 let key = format!("{agent_name}:{model}");
673 let store_opt = cooldown_store().lock().await.clone();
674 let count = read_and_increment_failure_count(&store_opt, &key).await;
675 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
676 let cooldown_until =
677 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
678 set_cooldown_async(&key, cooldown_until, "model_error").await;
679}
680
681pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
689 if let Some(window_secs) = parse_relative_usage_window(error_message) {
690 if let Some(m) = model {
691 let cooldown_until = chrono::Utc::now().timestamp() + window_secs as i64;
692 set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
693 tracing::info!(
694 agent = agent_name,
695 model = m,
696 window_secs,
697 cooldown_until,
698 "relative usage window detected: applying model-specific cooldown"
699 );
700 return;
701 }
702 }
703
704 record_agent_failure_with_message(agent_name, error_message).await;
705 if let Some(m) = model {
706 record_model_failure(agent_name, m).await;
707 }
708}
709
710pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
722 let key = format!("{agent_name}:{model}");
723 let store_opt = cooldown_store().lock().await.clone();
724 let count = read_and_increment_failure_count_with_prefix(
725 &store_opt,
726 PERSISTENT_FAILURE_COUNT_PREFIX,
727 &key,
728 )
729 .await;
730 let cooldown_secs = compute_backoff(
731 count,
732 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
733 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
734 false,
735 );
736 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
737 set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
738}
739
740pub async fn record_transient_model_failure(agent_name: &str, model: &str) {
746 let key = format!("{agent_name}:{model}");
747 let store_opt = cooldown_store().lock().await.clone();
748 let count = read_and_increment_failure_count(&store_opt, &key).await;
749 let cooldown_secs = compute_backoff(
750 count,
751 TRANSIENT_MODEL_BACKOFF_BASE_SECS,
752 TRANSIENT_MODEL_BACKOFF_MAX_SECS,
753 false,
754 );
755 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
756 set_cooldown_async(&key, cooldown_until, "transient_model_error").await;
757}
758
759pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
766 let key = format!("{agent_name}:{model}");
767 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
768 set_cooldown_async(&key, cooldown_until, reason).await;
769}
770
771pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
773 let key = format!("{agent_name}:{model}");
774 set_cooldown_async(&key, until, "vendor_reset_at").await;
775}
776
777pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
784 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
785 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
786}
787
788pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
791 let store_opt = cooldown_store().lock().await.clone();
792 let store = match store_opt {
793 Some(store) => store,
794 None => {
795 tracing::debug!(
796 agent = agent_name,
797 model,
798 "skipping silence count record (no KV store)"
799 );
800 return None;
801 }
802 };
803
804 let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
805 let now = chrono::Utc::now().timestamp();
806 let window_start = now - SILENCE_COUNT_WINDOW_SECS;
807
808 let mut timestamps = match store.kv_get(&key).await {
809 Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
810 Ok(ts) => ts,
811 Err(e) => {
812 tracing::warn!(
813 kv_key = %key,
814 err = %e,
815 "failed to parse silence timestamps from KV — resetting window"
816 );
817 Vec::new()
818 }
819 },
820 Ok(None) => Vec::new(),
821 Err(err) => {
822 tracing::warn!(
823 kv_key = %key,
824 err = %err,
825 "failed to load silence count from KV"
826 );
827 Vec::new()
828 }
829 };
830
831 timestamps.retain(|ts| *ts >= window_start);
832 timestamps.push(now);
833
834 let count = timestamps.len();
835 let mut extended_cooldown_applied = false;
836 if count >= SILENCE_COUNT_THRESHOLD {
837 set_model_cooldown(
838 agent_name,
839 model,
840 SILENCE_EXTENDED_COOLDOWN_SECS,
841 "silence_detected",
842 )
843 .await;
844 extended_cooldown_applied = true;
845 timestamps.clear();
846 }
847
848 let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
849 if let Err(err) = store.kv_set(&key, &value).await {
850 tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
851 }
852
853 Some(SilenceCountResult {
854 count,
855 extended_cooldown_applied,
856 })
857}
858
859pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
861 let key = format!("{agent_name}:{model}");
862 is_in_cooldown(&key)
863}
864
865pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
867 is_in_cooldown(agent_name)
868}
869
870pub fn cooldown_until(key: &str) -> Option<i64> {
872 let map = cooldown_lock();
873 map.get(key).and_then(|entry| {
874 let now = chrono::Utc::now().timestamp();
875 if now < entry.cooldown_until {
876 Some(entry.cooldown_until)
877 } else {
878 None
879 }
880 })
881}
882
883pub fn cooldown_reason(key: &str) -> Option<String> {
885 let map = cooldown_lock();
886 map.get(key).and_then(|entry| {
887 let now = chrono::Utc::now().timestamp();
888 if now < entry.cooldown_until {
889 Some(entry.reason.clone())
890 } else {
891 None
892 }
893 })
894}
895
896pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
901 let now = chrono::Utc::now().timestamp();
902 let map = cooldown_lock();
903 let mut result: Vec<(String, i64, String)> = map
904 .iter()
905 .filter(|(_, entry)| now < entry.cooldown_until)
906 .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
907 .collect();
908 result.sort_by(|a, b| a.0.cmp(&b.0));
909 result
910}
911
912pub async fn mark_rate_limits_cleared(store: &Arc<crate::store::TaskStore>, agents: &[String]) {
915 let now = chrono::Utc::now().timestamp().to_string();
916 let mut seen = std::collections::HashSet::new();
917 for key in agents {
918 let agent = agent_from_key(key);
919 if agent.is_empty() || !seen.insert(agent.to_string()) {
920 continue;
921 }
922 let kv_key = format!("{RATE_LIMITS_CLEARED_AT_PREFIX}{agent}");
923 if let Err(e) = store.kv_set(&kv_key, &now).await {
924 tracing::warn!(agent, err = %e, "failed to mark rate limits cleared");
925 }
926 }
927}
928
929pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
936 let mut agents_to_mark: Vec<String> = Vec::new();
937 if key == "*" {
938 let keys: Vec<String> = {
939 let mut map = cooldown_lock();
940 let keys: Vec<String> = map.keys().cloned().collect();
941 map.clear();
942 keys
943 };
944 for k in &keys {
945 let kv_key = format!("{KV_PREFIX}{k}");
946 if let Err(e) = store.kv_set(&kv_key, "0").await {
947 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
948 }
949 let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
951 if let Err(e) = store.kv_set(&fc_key, "0").await {
952 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
953 }
954 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
956 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
957 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
958 }
959 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
961 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
962 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
963 }
964 }
965 if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
969 for (fc_key, _) in fc_entries {
970 if let Err(e) = store.kv_set(&fc_key, "0").await {
971 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
972 }
973 }
974 }
975 if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
976 for (cfc_key, _) in cfc_entries {
977 if let Err(e) = store.kv_set(&cfc_key, "0").await {
978 tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
979 }
980 }
981 }
982 if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
983 for (pfc_key, _) in pfc_entries {
984 if let Err(e) = store.kv_set(&pfc_key, "0").await {
985 tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
986 }
987 }
988 }
989 tracing::info!(
990 count = keys.len(),
991 "cleared all cooldowns and failure counts"
992 );
993 agents_to_mark.extend(keys.iter().map(|k| agent_from_key(k).to_string()));
994 match store.rate_limit_agents().await {
995 Ok(rl_agents) => agents_to_mark.extend(rl_agents),
996 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits agents while clearing"),
997 }
998 match store.kv_list_prefix(RATE_LIMITS_CLEARED_AT_PREFIX).await {
999 Ok(marks) => {
1000 for (k, _) in marks {
1001 let agent = k.trim_start_matches(RATE_LIMITS_CLEARED_AT_PREFIX);
1002 if !agent.is_empty() {
1003 agents_to_mark.push(agent.to_string());
1004 }
1005 }
1006 }
1007 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits_cleared_at keys"),
1008 }
1009 } else {
1010 {
1011 let mut map = cooldown_lock();
1012 map.remove(key);
1013 }
1014 let kv_key = format!("{KV_PREFIX}{key}");
1015 if let Err(e) = store.kv_set(&kv_key, "0").await {
1016 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
1017 }
1018 let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
1020 if let Err(e) = store.kv_set(&fc_key, "0").await {
1021 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
1022 }
1023 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
1025 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
1026 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
1027 }
1028 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
1030 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
1031 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
1032 }
1033 tracing::info!(key, "cleared cooldown and failure count");
1034 agents_to_mark.push(agent_from_key(key).to_string());
1035 }
1036 mark_rate_limits_cleared(store, &agents_to_mark).await;
1037}
1038
1039fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
1048 let mut map = cooldown_lock();
1049 if let Some(existing) = map.get(key) {
1050 if existing.cooldown_until >= cooldown_until {
1054 return false;
1055 }
1056 }
1057 map.insert(
1058 key.to_string(),
1059 CooldownEntry {
1060 cooldown_until,
1061 reason: reason.to_string(),
1062 },
1063 );
1064 true
1065}
1066
1067async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
1077 if !set_cooldown_in_memory(key, cooldown_until, reason) {
1078 return false;
1079 }
1080 let store_opt = cooldown_store().lock().await.clone();
1081 if let Some(store) = store_opt {
1082 let kv_key = format!("{KV_PREFIX}{key}");
1083 let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
1084 if let Err(e) = store.kv_set(&kv_key, &value).await {
1085 tracing::warn!(
1086 kv_key,
1087 err = %e,
1088 "failed to persist cooldown to KV store — rolling back in-memory entry"
1089 );
1090 let mut map = cooldown_lock();
1091 if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1095 map.remove(key);
1096 }
1097 return false;
1098 }
1099 }
1100 true
1101}
1102
1103fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
1105 cooldowns().lock().unwrap_or_else(|e| {
1106 tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1107 e.into_inner()
1108 })
1109}
1110
1111fn degraded_agents_lock() -> std::sync::MutexGuard<'static, HashMap<String, String>> {
1113 degraded_agents().lock().unwrap_or_else(|e| {
1114 tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1115 e.into_inner()
1116 })
1117}
1118
1119fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
1121 github_5xx_timestamps().lock().unwrap_or_else(|e| {
1122 tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1123 e.into_inner()
1124 })
1125}
1126
1127fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
1129 github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1130 tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1131 e.into_inner()
1132 })
1133}
1134
1135fn is_in_cooldown(key: &str) -> bool {
1137 let map = cooldown_lock();
1138 if let Some(entry) = map.get(key) {
1139 let now = chrono::Utc::now().timestamp();
1140 return now < entry.cooldown_until;
1141 }
1142 false
1143}
1144
1145pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1151
1152pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1156
1157fn degraded_agents() -> &'static Mutex<HashMap<String, String>> {
1160 static DEGRADED: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
1161 DEGRADED.get_or_init(|| Mutex::new(HashMap::new()))
1162}
1163
1164pub fn is_agent_degraded(agent: &str) -> bool {
1169 let set = degraded_agents_lock();
1170 set.contains_key(agent)
1171}
1172
1173pub fn agent_degraded_reason(agent: &str) -> Option<String> {
1175 degraded_agents_lock().get(agent).cloned()
1176}
1177
1178pub fn list_degraded_agents() -> Vec<(String, String)> {
1180 let mut entries: Vec<(String, String)> = degraded_agents_lock()
1181 .iter()
1182 .map(|(agent, reason)| (agent.clone(), reason.clone()))
1183 .collect();
1184 entries.sort_by(|a, b| a.0.cmp(&b.0));
1185 entries
1186}
1187
1188pub fn mark_agent_degraded(agent: &str) {
1190 mark_agent_degraded_with_reason(agent, "degraded");
1191}
1192
1193fn mark_agent_degraded_with_reason(agent: &str, reason: &str) {
1194 let mut set = degraded_agents_lock();
1195 set.insert(agent.to_string(), reason.to_string());
1196}
1197
1198pub fn clear_agent_degraded(agent: &str) {
1200 let mut set = degraded_agents_lock();
1201 set.remove(agent);
1202}
1203
1204pub async fn refresh_degraded_agents(
1218 store: &Arc<crate::store::TaskStore>,
1219 available_agents: &[String],
1220 model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1221 window_hours: u32,
1222 threshold: i64,
1223) {
1224 clear_expired_cooldowns();
1228
1229 let start = chrono::Utc::now();
1232 let counts = match store.recent_rate_limit_counts(window_hours).await {
1233 Ok(c) => {
1234 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1235 tracing::info!(
1236 duration_ms = dur_ms,
1237 window_hours,
1238 "recent_rate_limit_counts query completed"
1239 );
1240 c
1241 }
1242 Err(e) => {
1243 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1244 tracing::warn!(
1245 err = %e,
1246 duration_ms = dur_ms,
1247 "failed to query recent rate limit counts for health check"
1248 );
1249 return;
1250 }
1251 };
1252
1253 for agent in available_agents {
1254 let in_cooldown = is_agent_in_cooldown(agent);
1255 let no_models = !model_checker(agent);
1256 let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1257 let over_threshold = rate_limit_count >= threshold;
1258
1259 let degraded = in_cooldown || no_models || over_threshold;
1260
1261 if degraded {
1262 let reason = if in_cooldown {
1263 "agent in cooldown"
1264 } else if no_models {
1265 "all models cooled"
1266 } else {
1267 "rate limit threshold exceeded"
1268 };
1269 let previous = agent_degraded_reason(agent);
1270 match previous.as_deref() {
1271 None => tracing::warn!(
1272 agent,
1273 reason,
1274 rate_limit_count,
1275 window_hours,
1276 threshold,
1277 "pre-emptive health check: marking agent as degraded"
1278 ),
1279 Some(prev) if prev != reason => tracing::warn!(
1280 agent,
1281 reason,
1282 previous = prev,
1283 rate_limit_count,
1284 window_hours,
1285 threshold,
1286 "pre-emptive health check: degraded reason changed"
1287 ),
1288 Some(_) => {}
1289 }
1290 mark_agent_degraded_with_reason(agent, reason);
1291 } else {
1292 if is_agent_degraded(agent) {
1293 tracing::info!(
1294 agent,
1295 rate_limit_count,
1296 "pre-emptive health check: agent recovered, clearing degraded flag"
1297 );
1298 }
1299 clear_agent_degraded(agent);
1300 }
1301 }
1302}
1303
1304pub fn clear_expired_cooldowns() {
1306 let mut map = cooldown_lock();
1307 let now = chrono::Utc::now().timestamp();
1308 map.retain(|_key, entry| now < entry.cooldown_until);
1309}
1310
1311const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1317
1318const GITHUB_5XX_THRESHOLD: usize = 5;
1320
1321const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1323
1324fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1326 static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1327 TS.get_or_init(|| Mutex::new(Vec::new()))
1328}
1329
1330fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1332 static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1333 OPEN.get_or_init(|| Mutex::new(None))
1334}
1335
1336pub async fn record_github_5xx() {
1341 let now = chrono::Utc::now().timestamp();
1342 let window_start = now - GITHUB_5XX_WINDOW_SECS;
1343
1344 let should_trip = {
1345 let mut ts = github_5xx_timestamps_lock();
1346 ts.retain(|t| *t >= window_start);
1347 ts.push(now);
1348 ts.len() >= GITHUB_5XX_THRESHOLD
1349 };
1350
1351 if should_trip {
1352 let cooldown_until = {
1353 let mut open = github_5xx_circuit_open_lock();
1354 if open.is_none() {
1355 let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1356 *open = Some(cd);
1357 tracing::warn!(
1358 errors_in_window = GITHUB_5XX_THRESHOLD,
1359 window_secs = GITHUB_5XX_WINDOW_SECS,
1360 cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1361 "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1362 );
1363 Some(cd)
1364 } else {
1365 None
1366 }
1367 };
1368 if let Some(cd) = cooldown_until {
1370 let store_opt = cooldown_store().lock().await.clone();
1371 if let Some(store) = store_opt {
1372 let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1373 }
1374 set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1378 }
1379 let mut ts = github_5xx_timestamps_lock();
1381 ts.clear();
1382 }
1383}
1384
1385pub fn is_github_circuit_open() -> bool {
1393 {
1396 let mut open = github_5xx_circuit_open_lock();
1397 if let Some(until) = *open {
1398 let now = chrono::Utc::now().timestamp();
1399 if now >= until {
1400 *open = None;
1401 tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1402 return false;
1403 }
1404 return true;
1405 }
1406 }
1407
1408 if let Some(until) = cooldown_until("github:5xx") {
1412 let mut open = github_5xx_circuit_open_lock();
1415 if open.is_none() {
1416 *open = Some(until);
1417 let now = chrono::Utc::now().timestamp();
1418 if now < until {
1419 let remaining = until - now;
1420 tracing::warn!(
1421 remaining_secs = remaining,
1422 "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1423 );
1424 }
1425 }
1426 return true;
1427 }
1428
1429 false
1430}
1431
1432pub fn github_circuit_remaining_secs() -> u64 {
1435 {
1437 let open = github_5xx_circuit_open_lock();
1438 if let Some(until) = *open {
1439 let now = chrono::Utc::now().timestamp();
1440 if now < until {
1441 return (until - now) as u64;
1442 }
1443 return 0;
1444 }
1445 }
1446
1447 if let Some(until) = cooldown_until("github:5xx") {
1449 let now = chrono::Utc::now().timestamp();
1450 if now < until {
1451 return (until - now) as u64;
1452 }
1453 }
1454 0
1455}
1456
1457pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1461 if error_message.is_empty() {
1462 return None;
1463 }
1464
1465 let lower = error_message.to_lowercase();
1466 for retry_marker in [
1470 "try again at ",
1471 "reset at ",
1472 "resets at ",
1473 "resets ",
1474 "reset ",
1475 ] {
1476 if let Some(idx) = lower.find(retry_marker) {
1477 let date_str = &lower[idx + retry_marker.len()..];
1478 let date_str = date_str
1480 .split(['.', '\n'])
1481 .next()
1482 .unwrap_or(date_str)
1483 .trim();
1484
1485 let cleaned = date_str
1488 .replace("st,", ",")
1489 .replace("nd,", ",")
1490 .replace("rd,", ",")
1491 .replace("th,", ",");
1492
1493 for fmt in &[
1494 "%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", ] {
1501 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1502 use chrono::TimeZone;
1505 let utc = match chrono::Local.from_local_datetime(&dt) {
1506 chrono::LocalResult::Single(local_dt) => {
1507 local_dt.with_timezone(&chrono::Utc)
1508 }
1509 chrono::LocalResult::Ambiguous(earliest, _) => {
1510 earliest.with_timezone(&chrono::Utc)
1512 }
1513 chrono::LocalResult::None => {
1514 dt.and_utc()
1516 }
1517 };
1518 tracing::info!(
1519 retry_at = %utc,
1520 raw = date_str,
1521 marker = retry_marker,
1522 "parsed rate limit retry-at date"
1523 );
1524 return Some(utc.timestamp());
1525 }
1526 }
1527 tracing::debug!(
1528 raw = date_str,
1529 marker = retry_marker,
1530 "could not parse retry-at date"
1531 );
1532
1533 let time_only = date_str
1538 .split_once('(')
1539 .map(|(t, _)| t)
1540 .unwrap_or(date_str)
1541 .trim();
1542 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1544 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1545 let last = &compact[split..];
1546 if !last.contains(':')
1547 && (last.ends_with("am") || last.ends_with("pm"))
1548 && last.len() > 2
1549 {
1550 let at = split + last.len() - 2;
1551 compact.insert_str(at, ":00");
1552 }
1553 let md = compact.replace(" at ", " ");
1555 let now_local = chrono::Local::now();
1556 let this_year = chrono::Datelike::year(&now_local);
1557 for year in [this_year, this_year + 1] {
1558 let with_year = format!("{year} {md}");
1559 if let Ok(dt) =
1560 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1561 {
1562 use chrono::TimeZone;
1563 if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
1564 if local_dt > now_local {
1565 return Some(local_dt.timestamp());
1566 }
1567 }
1568 }
1569 }
1570 let time_only = compact.as_str();
1571 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1572 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1573 let local_now = chrono::Local::now();
1576 let today = local_now.date_naive();
1577 let candidate = chrono::NaiveDateTime::new(today, t);
1578 let candidate_local: chrono::DateTime<chrono::Local> =
1581 match chrono::Local.from_local_datetime(&candidate) {
1582 chrono::LocalResult::Single(dt) => dt,
1583 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1584 chrono::LocalResult::None => {
1585 let utc_dt = candidate.and_utc();
1587 return Some(utc_dt.timestamp());
1588 }
1589 };
1590 let chosen = if candidate_local <= local_now {
1591 let tomorrow = today.succ_opt().unwrap_or(today);
1593 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1594 match chrono::Local.from_local_datetime(&ndt) {
1595 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1596 chrono::LocalResult::Ambiguous(earliest, _) => {
1597 earliest.with_timezone(&chrono::Utc)
1598 }
1599 chrono::LocalResult::None => ndt.and_utc(),
1600 }
1601 } else {
1602 candidate_local.with_timezone(&chrono::Utc)
1603 };
1604 tracing::info!(
1605 retry_at = %chosen,
1606 raw = date_str,
1607 marker = retry_marker,
1608 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1609 );
1610 return Some(chosen.timestamp());
1611 }
1612 }
1613 }
1614 }
1615
1616 None
1620}
1621
1622pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1633 if error_message.is_empty() {
1634 return None;
1635 }
1636
1637 let lower = error_message.to_lowercase();
1638 let has_context = lower.contains("usage limit")
1639 || lower.contains("quota")
1640 || lower.contains("window")
1641 || lower.contains("limit");
1642 if !has_context {
1643 return None;
1644 }
1645
1646 static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
1647 Regex::new(
1648 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1649 )
1650 .expect("usage window regex is valid")
1651 });
1652
1653 let mut max_secs: Option<u64> = None;
1654 for caps in WINDOW_RE.captures_iter(&lower) {
1655 let secs = if let Some(num) = caps.get(1) {
1656 let n: u64 = num.as_str().parse().ok()?;
1657 let unit = caps.get(2)?.as_str().to_lowercase();
1658 let multiplier = match unit.as_str() {
1659 "hour" | "hours" | "hr" | "hrs" => 3600,
1660 "minute" | "minutes" | "min" | "mins" => 60,
1661 "day" | "days" => 86400,
1662 _ => return None,
1663 };
1664 n.checked_mul(multiplier)?
1665 } else if let Some(word) = caps.get(3) {
1666 match word.as_str().to_lowercase().as_str() {
1667 "hourly" => 3600,
1668 "daily" => 86400,
1669 "weekly" => 604800,
1670 "monthly" => 2592000,
1671 _ => return None,
1672 }
1673 } else {
1674 continue;
1675 };
1676 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1677 }
1678
1679 max_secs
1680}
1681
1682#[cfg(test)]
1683pub(crate) async fn reset_global_state() {
1684 cooldowns().lock().unwrap().clear();
1685 *cooldown_store().lock().await = None;
1686 github_5xx_timestamps().lock().unwrap().clear();
1687 *github_5xx_circuit_open().lock().unwrap() = None;
1688 degraded_agents().lock().unwrap().clear();
1689 crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1690}
1691
1692#[cfg(test)]
1693mod tests {
1694 use super::*;
1695 use serial_test::serial;
1696
1697 async fn test_store() -> Arc<crate::store::TaskStore> {
1698 Arc::new(
1699 crate::store::TaskStore::open_memory()
1700 .await
1701 .expect("in-memory store"),
1702 )
1703 }
1704
1705 fn sqlite_datetime_offset(secs: i64) -> String {
1706 (chrono::Utc::now() + chrono::Duration::seconds(secs))
1707 .format("%Y-%m-%d %H:%M:%S")
1708 .to_string()
1709 }
1710
1711 #[serial(cooldown_state)]
1712 #[tokio::test]
1713 async fn init_loads_unexpired_cooldowns_from_kv() {
1714 super::reset_global_state().await;
1715 let store = test_store().await;
1716 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1719 .kv_set(
1720 &format!("{KV_PREFIX}kimi:k2p5"),
1721 &cooldown_until.to_string(),
1722 )
1723 .await
1724 .unwrap();
1725
1726 let expired = chrono::Utc::now().timestamp() - 1;
1728 store
1729 .kv_set(
1730 &format!("{KV_PREFIX}opencode:old-free"),
1731 &expired.to_string(),
1732 )
1733 .await
1734 .unwrap();
1735
1736 {
1737 let mut map = cooldown_lock();
1738 map.clear();
1739 }
1740
1741 init_cooldown_store(store).await;
1742
1743 assert!(
1744 is_model_in_cooldown("kimi", "k2p5"),
1745 "unexpired cooldown should be loaded"
1746 );
1747 assert!(
1748 !is_model_in_cooldown("opencode", "old-free"),
1749 "expired cooldown should not be loaded"
1750 );
1751 }
1752
1753 #[serial(cooldown_state)]
1754 #[tokio::test]
1755 async fn record_model_failure_writes_to_kv() {
1756 super::reset_global_state().await;
1757 let store = test_store().await;
1758
1759 *cooldown_store().lock().await = Some(store.clone());
1760
1761 record_model_failure("testagent_persist", "testmodel_persist").await;
1762
1763 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1765
1766 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1767 let stored = store.kv_get(&kv_key).await.unwrap();
1768 assert!(stored.is_some(), "model failure should be persisted to KV");
1769 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1770 let now = chrono::Utc::now().timestamp();
1771 assert!(ts > now, "persisted timestamp should be in the future");
1773 assert_eq!(reason, "model_error");
1774 }
1775
1776 #[serial(cooldown_state)]
1777 #[tokio::test]
1778 async fn cooldown_reason_survives_restart_via_kv() {
1779 super::reset_global_state().await;
1780 let store = test_store().await;
1781 let agent = "test_reason_restart";
1782 let future = chrono::Utc::now().timestamp() + 600;
1783
1784 init_cooldown_store(store.clone()).await;
1785 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1786
1787 {
1789 let mut map = cooldown_lock();
1790 map.clear();
1791 }
1792 init_cooldown_store(store).await;
1793
1794 let entries = list_all_cooldowns();
1795 let entry = entries.iter().find(|(key, ..)| key == agent);
1796 assert_eq!(
1797 entry.map(|(_, _, reason)| reason.as_str()),
1798 Some("billing_cycle_exhausted"),
1799 "cooldown reason must survive a restart, not collapse to 'persisted'"
1800 );
1801 }
1802
1803 #[serial(cooldown_state)]
1804 #[tokio::test]
1805 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1806 super::reset_global_state().await;
1807 let store = test_store().await;
1808 *cooldown_store().lock().await = Some(store.clone());
1809
1810 let agent = "testagent_persistent";
1811 let model = "testmodel_persistent";
1812 let key = format!("{agent}:{model}");
1813
1814 record_persistent_model_failure(agent, model).await;
1816 let now = chrono::Utc::now().timestamp();
1817 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1818 let first_remaining = first_until.saturating_sub(now);
1819 assert!(
1820 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1821 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1822 );
1823
1824 record_persistent_model_failure(agent, model).await;
1826 let now2 = chrono::Utc::now().timestamp();
1827 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1828 let second_remaining = second_until.saturating_sub(now2);
1829 assert!(
1830 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1831 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1832 );
1833
1834 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1835 let stored_count = store
1836 .kv_get(&count_key)
1837 .await
1838 .expect("kv read should succeed")
1839 .expect("failure count should be written");
1840 assert_eq!(stored_count, "2");
1841 }
1842
1843 #[serial(cooldown_state)]
1844 #[tokio::test]
1845 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1846 super::reset_global_state().await;
1847 let store = test_store().await;
1848 *cooldown_store().lock().await = Some(store.clone());
1849
1850 let agent = "testagent_persistent_success";
1851 let model = "testmodel_persistent_success";
1852 let key = format!("{agent}:{model}");
1853
1854 record_persistent_model_failure(agent, model).await;
1856 record_persistent_model_failure(agent, model).await;
1857
1858 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1859 let stored_count = store
1860 .kv_get(&count_key)
1861 .await
1862 .expect("kv read should succeed")
1863 .expect("failure count should be written");
1864 assert_eq!(stored_count, "2");
1865
1866 record_agent_success(agent, model).await;
1870
1871 let stored_count_after_success = store
1872 .kv_get(&count_key)
1873 .await
1874 .expect("kv read should succeed")
1875 .expect("persistent failure count should still exist");
1876 assert_eq!(
1877 stored_count_after_success, "2",
1878 "persistent failure count must survive an intervening success"
1879 );
1880
1881 record_persistent_model_failure(agent, model).await;
1883 let now = chrono::Utc::now().timestamp();
1884 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1885 let third_remaining = third_until.saturating_sub(now);
1886 assert!(
1887 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1888 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1889 );
1890 }
1891
1892 #[serial(cooldown_state)]
1893 #[tokio::test]
1894 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1895 super::reset_global_state().await;
1896 let store = test_store().await;
1897 *cooldown_store().lock().await = Some(store.clone());
1898
1899 let agent = "testagent_persistent_capped";
1900 let model = "testmodel_persistent_capped";
1901 let key = format!("{agent}:{model}");
1902
1903 for _ in 0..9 {
1904 record_persistent_model_failure(agent, model).await;
1905 }
1906
1907 let remaining = cooldown_until(&key)
1908 .expect("cooldown should exist")
1909 .saturating_sub(chrono::Utc::now().timestamp());
1910 assert!(
1911 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1912 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1913 );
1914 assert!(
1915 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1916 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1917 );
1918 }
1919
1920 #[serial(cooldown_state)]
1921 #[test]
1922 fn parse_retry_at_codex_format() {
1923 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.";
1924 let ts = parse_retry_at(msg);
1925 assert!(ts.is_some(), "should parse codex retry-at date");
1926
1927 use chrono::TimeZone;
1930 let naive = chrono::NaiveDateTime::new(
1931 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1932 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1933 );
1934 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1935 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1936 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1937 chrono::LocalResult::None => naive.and_utc(),
1938 };
1939 assert_eq!(
1940 ts.unwrap(),
1941 expected_utc.timestamp(),
1942 "timestamp should reflect local-timezone interpretation of the vendor date"
1943 );
1944
1945 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1947 use chrono::Datelike;
1948 assert_eq!(dt.month(), 3);
1949 assert_eq!(dt.day(), 26);
1950 }
1951
1952 #[serial(cooldown_state)]
1953 #[test]
1954 fn parse_retry_at_hour_only_and_month_day() {
1955 let now = chrono::Utc::now().timestamp();
1956 for msg in [
1957 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1958 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1959 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1960 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1961 ] {
1962 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1963 assert!(ts > now, "{msg}");
1964 }
1965 }
1966
1967 #[serial(cooldown_state)]
1968 #[test]
1969 fn parse_retry_at_weekly_resets_without_year() {
1970 use chrono::Datelike;
1971 let t = chrono::Local::now() + chrono::Duration::days(2);
1972 let msg = format!(
1973 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1974 t.format("%b"),
1975 t.day()
1976 );
1977 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1978 let remaining = ts - chrono::Utc::now().timestamp();
1979 assert!(
1980 remaining > 3600 && remaining < 3 * 86400,
1981 "expected ~2d retry-at, got {remaining}s"
1982 );
1983 }
1984
1985 #[serial(cooldown_state)]
1986 #[test]
1987 fn parse_retry_at_no_date() {
1988 assert!(parse_retry_at("generic rate limit error").is_none());
1989 assert!(parse_retry_at("").is_none());
1990 }
1991
1992 #[serial(cooldown_state)]
1993 #[test]
1994 fn parse_retry_at_glm_reset_format_with_seconds() {
1995 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
1996 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
1997
1998 use chrono::TimeZone;
1999 let naive = chrono::NaiveDateTime::new(
2000 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
2001 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
2002 );
2003 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
2004 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
2005 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
2006 chrono::LocalResult::None => naive.and_utc(),
2007 };
2008
2009 assert_eq!(
2010 ts,
2011 expected_utc.timestamp(),
2012 "timestamp should reflect local-timezone interpretation of vendor reset date"
2013 );
2014 }
2015
2016 #[serial(cooldown_state)]
2017 #[test]
2018 fn parse_retry_at_billing_cycle_without_date_returns_none() {
2019 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.";
2022 assert!(
2023 parse_retry_at(msg).is_none(),
2024 "billing cycle without date should return None"
2025 );
2026 }
2027
2028 #[serial(cooldown_state)]
2029 #[test]
2030 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
2031 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.";
2032 assert_eq!(
2033 parse_relative_usage_window(msg),
2034 Some(5 * 3600),
2035 "should extract the 5-hour window as seconds"
2036 );
2037 }
2038
2039 #[serial(cooldown_state)]
2040 #[test]
2041 fn parse_relative_usage_window_handles_common_durations() {
2042 assert_eq!(
2043 parse_relative_usage_window("You've reached your hourly usage limit."),
2044 Some(3600)
2045 );
2046 assert_eq!(
2047 parse_relative_usage_window("You've reached your daily usage limit."),
2048 Some(86400)
2049 );
2050 assert_eq!(
2051 parse_relative_usage_window("You've reached your weekly limit."),
2052 Some(604800)
2053 );
2054 assert_eq!(
2055 parse_relative_usage_window("You've reached your 30-minute usage limit."),
2056 Some(1800)
2057 );
2058 assert_eq!(
2059 parse_relative_usage_window("You've reached your 2-day quota."),
2060 Some(172800)
2061 );
2062 }
2063
2064 #[serial(cooldown_state)]
2065 #[test]
2066 fn parse_relative_usage_window_returns_none_without_usage_context() {
2067 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
2070 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
2071 assert!(parse_relative_usage_window("").is_none());
2072 }
2073
2074 #[serial(cooldown_state)]
2075 #[tokio::test]
2076 async fn record_rate_limit_uses_model_specific_window_cooldown() {
2077 super::reset_global_state().await;
2078 let store = test_store().await;
2079 *cooldown_store().lock().await = Some(store.clone());
2080
2081 let agent = "test_record_rl_window_agent";
2082 let model = "test_record_rl_window_model";
2083 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
2084
2085 record_rate_limit(agent, Some(model), msg).await;
2086
2087 assert!(
2088 is_model_in_cooldown(agent, model),
2089 "model should be cooled for the relative usage window"
2090 );
2091 assert!(
2092 !is_agent_in_cooldown(agent),
2093 "agent should not be cooled for a model-scoped quota"
2094 );
2095
2096 let key = format!("{agent}:{model}");
2097 let remaining = cooldown_until(&key)
2098 .expect("model cooldown should exist")
2099 .saturating_sub(chrono::Utc::now().timestamp());
2100 assert!(
2101 remaining >= 5 * 3600 - 5,
2102 "model cooldown should be at least ~5 hours, got {remaining}s"
2103 );
2104 assert_eq!(
2105 cooldown_reason(&key).as_deref(),
2106 Some("rate_limit_usage_window"),
2107 "relative usage-window cooldown must be labeled as a rate limit, not silence"
2108 );
2109 }
2110
2111 #[serial(cooldown_state)]
2112 #[tokio::test]
2113 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
2114 super::reset_global_state().await;
2115 let store = test_store().await;
2116 *cooldown_store().lock().await = Some(store.clone());
2117
2118 let agent = "test_record_rl_fallback_agent";
2119 let model = "test_record_rl_fallback_model";
2120 let msg = "429 Too Many Requests";
2121
2122 record_rate_limit(agent, Some(model), msg).await;
2123
2124 assert!(
2125 is_agent_in_cooldown(agent),
2126 "agent should be cooled when no relative window is detected"
2127 );
2128 assert!(
2129 is_model_in_cooldown(agent, model),
2130 "model should be cooled when no relative window is detected"
2131 );
2132 }
2133
2134 #[serial(cooldown_state)]
2135 #[tokio::test]
2136 async fn agent_cooldown_persists_to_kv() {
2137 super::reset_global_state().await;
2138 let agent = "test_agent_persist_check";
2140 record_agent_failure_with_message(agent, "").await;
2141 assert!(is_agent_in_cooldown(agent));
2142 }
2143
2144 #[serial(cooldown_state)]
2145 #[tokio::test]
2146 async fn cooldown_never_shortens() {
2147 super::reset_global_state().await;
2148 let agent = "test_agent_no_shorten";
2149 set_agent_cooldown(agent, 24 * 60 * 60).await;
2150 let initial = {
2151 let map = cooldowns().lock().unwrap();
2152 map.get(agent)
2153 .expect("cooldown entry should exist")
2154 .cooldown_until
2155 };
2156
2157 record_agent_failure_with_message(agent, "rate limit").await;
2159 let after = {
2160 let map = cooldowns().lock().unwrap();
2161 map.get(agent)
2162 .expect("cooldown entry should exist")
2163 .cooldown_until
2164 };
2165 assert_eq!(
2166 initial, after,
2167 "cooldown should not be shortened by later failures"
2168 );
2169 }
2170
2171 #[serial(cooldown_state)]
2172 #[tokio::test]
2173 async fn silence_agent_cooldown_is_short_lived() {
2174 super::reset_global_state().await;
2175 let agent = "test_silence_agent_cd";
2176 assert!(!is_agent_in_cooldown(agent));
2177
2178 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2179 assert!(is_agent_in_cooldown(agent));
2180
2181 let map = cooldowns().lock().unwrap();
2183 let entry = map.get(agent).expect("should have cooldown entry");
2184 let now = chrono::Utc::now().timestamp();
2185 let remaining = entry.cooldown_until - now;
2186 assert!(
2187 remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
2188 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2189 );
2190 assert!(remaining > 0, "cooldown should still be active");
2191 }
2192
2193 #[serial(cooldown_state)]
2194 #[tokio::test]
2195 async fn record_silence_detection_applies_extended_cooldown() {
2196 super::reset_global_state().await;
2197 let store = test_store().await;
2198 *cooldown_store().lock().await = Some(store.clone());
2199
2200 let agent = "test_silence_count_agent";
2201 let model = "test_silence_count_model";
2202 assert!(!is_model_in_cooldown(agent, model));
2203
2204 for _ in 0..SILENCE_COUNT_THRESHOLD {
2205 let result = record_silence_detection(agent, model).await;
2206 assert!(result.is_some());
2207 }
2208
2209 assert!(is_model_in_cooldown(agent, model));
2210 let key = format!("{agent}:{model}");
2211 let remaining = {
2212 let map = cooldowns().lock().unwrap();
2213 let entry = map.get(&key).expect("cooldown entry should exist");
2214 entry.cooldown_until - chrono::Utc::now().timestamp()
2215 };
2216 assert!(
2217 remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
2218 "extended cooldown should be applied, got {remaining}s"
2219 );
2220
2221 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2222 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2223 assert_eq!(stored, "[]");
2224 }
2225
2226 #[serial(cooldown_state)]
2227 #[test]
2228 fn detect_credit_exhaustion_out_of_credits() {
2229 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2230 assert_eq!(
2231 detect_credit_exhaustion(msg),
2232 Some(CreditExhaustionReason::OutOfCredits)
2233 );
2234 }
2235
2236 #[serial(cooldown_state)]
2237 #[test]
2238 fn detect_credit_exhaustion_out_of_credits_variations() {
2239 assert_eq!(
2240 detect_credit_exhaustion("out of credits"),
2241 Some(CreditExhaustionReason::OutOfCredits)
2242 );
2243 assert_eq!(
2244 detect_credit_exhaustion("outofcredits"),
2245 Some(CreditExhaustionReason::OutOfCredits)
2246 );
2247 assert_eq!(
2248 detect_credit_exhaustion("no credits remaining"),
2249 Some(CreditExhaustionReason::OutOfCredits)
2250 );
2251 assert_eq!(
2252 detect_credit_exhaustion("insufficient funds"),
2253 Some(CreditExhaustionReason::OutOfCredits)
2254 );
2255 }
2256
2257 #[serial(cooldown_state)]
2258 #[test]
2259 fn detect_credit_exhaustion_org_level_disabled() {
2260 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2261 assert_eq!(
2262 detect_credit_exhaustion(msg),
2263 Some(CreditExhaustionReason::OrgLevelDisabled)
2264 );
2265 }
2266
2267 #[serial(cooldown_state)]
2268 #[test]
2269 fn detect_credit_exhaustion_org_level_disabled_variations() {
2270 assert_eq!(
2271 detect_credit_exhaustion("organization has been disabled"),
2272 Some(CreditExhaustionReason::OrgLevelDisabled)
2273 );
2274 assert_eq!(
2275 detect_credit_exhaustion("org disabled"),
2276 Some(CreditExhaustionReason::OrgLevelDisabled)
2277 );
2278 assert_eq!(
2279 detect_credit_exhaustion("billing quota exceeded"),
2280 Some(CreditExhaustionReason::OrgLevelDisabled)
2281 );
2282 assert_eq!(
2283 detect_credit_exhaustion("credit balance too low"),
2284 Some(CreditExhaustionReason::OrgLevelDisabled)
2285 );
2286 }
2287
2288 #[serial(cooldown_state)]
2289 #[test]
2290 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2291 let msg = "overageDisabledReason: out_of_credits";
2292 assert_eq!(
2293 detect_credit_exhaustion(msg),
2294 Some(CreditExhaustionReason::OutOfCredits)
2295 );
2296
2297 let msg = "overageDisabledReason: org_level_disabled";
2298 assert_eq!(
2299 detect_credit_exhaustion(msg),
2300 Some(CreditExhaustionReason::OrgLevelDisabled)
2301 );
2302 }
2303
2304 #[serial(cooldown_state)]
2305 #[test]
2306 fn detect_credit_exhaustion_glm_insufficient_balance() {
2307 assert_eq!(
2309 detect_credit_exhaustion(
2310 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2311 ),
2312 Some(CreditExhaustionReason::OutOfCredits)
2313 );
2314 assert_eq!(
2315 detect_credit_exhaustion("Insufficient balance"),
2316 Some(CreditExhaustionReason::OutOfCredits)
2317 );
2318 assert_eq!(
2319 detect_credit_exhaustion("please recharge your account"),
2320 Some(CreditExhaustionReason::OutOfCredits)
2321 );
2322 assert_eq!(
2323 detect_credit_exhaustion("no resource package available"),
2324 Some(CreditExhaustionReason::OutOfCredits)
2325 );
2326 }
2327
2328 #[serial(cooldown_state)]
2329 #[test]
2330 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2331 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2332 assert!(detect_credit_exhaustion("too many requests").is_none());
2333 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2334 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2335 assert!(detect_credit_exhaustion(
2336 "You've hit your usage limit, try again in a few minutes."
2337 )
2338 .is_none());
2339 }
2340
2341 #[serial(cooldown_state)]
2342 #[tokio::test]
2343 async fn record_credit_exhaustion_applies_agent_cooldown() {
2344 super::reset_global_state().await;
2345 let agent = "test_credit_exhaust_agent";
2346 assert!(!is_agent_in_cooldown(agent));
2347
2348 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2349 assert!(is_agent_in_cooldown(agent));
2350
2351 let remaining = {
2352 let map = cooldowns().lock().unwrap();
2353 let entry = map.get(agent).expect("cooldown entry should exist");
2354 entry.cooldown_until - chrono::Utc::now().timestamp()
2355 };
2356 assert!(
2358 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2359 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2360 );
2361 assert!(
2362 remaining <= CREDIT_BACKOFF_MAX_SECS,
2363 "cooldown should not exceed cap of 8h, got {remaining}s"
2364 );
2365 }
2366
2367 #[serial(cooldown_state)]
2368 #[tokio::test]
2369 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2370 super::reset_global_state().await;
2371 let agent = "test_org_disabled_agent";
2372 assert!(!is_agent_in_cooldown(agent));
2373
2374 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2375 assert!(is_agent_in_cooldown(agent));
2376
2377 let remaining = {
2378 let map = cooldowns().lock().unwrap();
2379 let entry = map.get(agent).expect("cooldown entry should exist");
2380 entry.cooldown_until - chrono::Utc::now().timestamp()
2381 };
2382 assert!(
2384 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2385 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2386 );
2387 assert!(
2388 remaining <= CREDIT_BACKOFF_MAX_SECS,
2389 "cooldown should not exceed cap of 8h, got {remaining}s"
2390 );
2391 }
2392
2393 #[serial(cooldown_state)]
2394 #[tokio::test]
2395 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2396 super::reset_global_state().await;
2397 let store = test_store().await;
2398 *cooldown_store().lock().await = Some(store.clone());
2399
2400 let agent = "test_billing_cycle_agent";
2401 assert!(!is_agent_in_cooldown(agent));
2402
2403 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2404 assert!(is_agent_in_cooldown(agent));
2405
2406 let remaining = {
2407 let map = cooldowns().lock().unwrap();
2408 let entry = map.get(agent).expect("cooldown entry should exist");
2409 entry.cooldown_until - chrono::Utc::now().timestamp()
2410 };
2411 assert!(
2413 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2414 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2415 );
2416 assert!(
2417 remaining <= BILLING_CYCLE_MAX_SECS,
2418 "cooldown should not exceed cap of 7 days, got {remaining}s"
2419 );
2420 }
2421
2422 #[serial(cooldown_state)]
2423 #[tokio::test]
2424 async fn record_credit_exhaustion_billing_cycle_escalates() {
2425 super::reset_global_state().await;
2426 let store = test_store().await;
2427 *cooldown_store().lock().await = Some(store.clone());
2428
2429 let agent = "test_billing_cycle_escalation";
2430
2431 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2433 let remaining_1 = {
2434 let map = cooldowns().lock().unwrap();
2435 let entry = map.get(agent).expect("cooldown entry should exist");
2436 entry.cooldown_until - chrono::Utc::now().timestamp()
2437 };
2438 assert!(
2439 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2440 .contains(&remaining_1),
2441 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2442 );
2443
2444 {
2446 let mut map = cooldowns().lock().unwrap();
2447 map.remove(agent);
2448 }
2449
2450 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2452 let remaining_2 = {
2453 let map = cooldowns().lock().unwrap();
2454 let entry = map.get(agent).expect("cooldown entry should exist");
2455 entry.cooldown_until - chrono::Utc::now().timestamp()
2456 };
2457 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2459 remaining_2 >= expected_2 - 5,
2460 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2461 );
2462
2463 {
2465 let mut map = cooldowns().lock().unwrap();
2466 map.remove(agent);
2467 }
2468
2469 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2471 let remaining_3 = {
2472 let map = cooldowns().lock().unwrap();
2473 let entry = map.get(agent).expect("cooldown entry should exist");
2474 entry.cooldown_until - chrono::Utc::now().timestamp()
2475 };
2476 assert!(
2477 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2478 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2479 );
2480 assert!(
2481 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2482 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2483 );
2484 }
2485
2486 #[serial(cooldown_state)]
2487 #[tokio::test]
2488 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2489 super::reset_global_state().await;
2490 let store = test_store().await;
2491 *cooldown_store().lock().await = Some(store.clone());
2492
2493 let agent = "test_billing_cycle_agent_capped";
2494
2495 for _ in 0..5 {
2496 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2497 }
2498
2499 let remaining = {
2500 let map = cooldowns().lock().unwrap();
2501 let entry = map.get(agent).expect("cooldown entry should exist");
2502 entry.cooldown_until - chrono::Utc::now().timestamp()
2503 };
2504 assert!(
2505 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2506 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2507 );
2508 assert!(
2509 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2510 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2511 );
2512 }
2513
2514 #[serial(cooldown_state)]
2515 #[test]
2516 fn credit_exhaustion_backoff_constants() {
2517 assert_eq!(
2519 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2520 CREDIT_BACKOFF_BASE_SECS
2521 );
2522 assert_eq!(
2524 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2525 ORG_BACKOFF_BASE_SECS
2526 );
2527 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2529 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2530 }
2531
2532 #[serial(cooldown_state)]
2533 #[test]
2534 fn compute_backoff_long_cap_without_extended_stays_capped() {
2535 assert_eq!(
2536 compute_backoff(
2537 3,
2538 BILLING_CYCLE_COOLDOWN_SECS,
2539 BILLING_CYCLE_MAX_SECS,
2540 false
2541 ),
2542 BILLING_CYCLE_MAX_SECS,
2543 "7-day capped paths should stop at the documented max"
2544 );
2545 assert_eq!(
2546 compute_backoff(
2547 9,
2548 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2549 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2550 false,
2551 ),
2552 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2553 "persistent model failures should not extend past 7 days"
2554 );
2555 }
2556
2557 #[serial(cooldown_state)]
2560 #[test]
2561 fn mark_and_check_agent_degraded() {
2562 let agent = "test_degraded_mark";
2563 assert!(!is_agent_degraded(agent));
2564 mark_agent_degraded(agent);
2565 assert!(is_agent_degraded(agent));
2566 clear_agent_degraded(agent);
2567 assert!(!is_agent_degraded(agent));
2568 }
2569
2570 #[serial(cooldown_state)]
2571 #[tokio::test]
2572 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2573 super::reset_global_state().await;
2574 let store = test_store().await;
2575 let agent = "test_sync_external_clear";
2576 let future = chrono::Utc::now().timestamp() + 600;
2577
2578 init_cooldown_store(store.clone()).await;
2580 set_cooldown_async(agent, future, "test").await;
2581 assert!(is_agent_in_cooldown(agent));
2582
2583 store
2587 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2588 .await
2589 .expect("kv_set tombstone");
2590 assert!(
2591 is_agent_in_cooldown(agent),
2592 "in-memory cooldown should still be present before sync"
2593 );
2594
2595 sync_from_kv(&store).await;
2596
2597 assert!(
2598 !is_agent_in_cooldown(agent),
2599 "sync_from_kv must drop in-memory entries the KV tombstoned"
2600 );
2601 }
2602
2603 #[serial(cooldown_state)]
2604 #[tokio::test]
2605 async fn sync_from_kv_closes_github_circuit_when_row_cleared() {
2606 super::reset_global_state().await;
2607 let store = test_store().await;
2608 {
2609 *github_5xx_circuit_open_lock() = Some(chrono::Utc::now().timestamp() + 600);
2610 github_5xx_timestamps_lock().push(chrono::Utc::now().timestamp());
2611 }
2612 assert!(is_github_circuit_open());
2613
2614 store
2615 .kv_set("cooldown:github:5xx", "0")
2616 .await
2617 .expect("kv_set tombstone");
2618 sync_from_kv(&store).await;
2619
2620 assert!(!is_github_circuit_open());
2621 assert!(github_5xx_timestamps_lock().is_empty());
2622 }
2623
2624 #[serial(cooldown_state)]
2625 #[tokio::test]
2626 async fn sync_from_kv_keeps_5xx_window_when_expired_row_and_breaker_closed() {
2627 super::reset_global_state().await;
2628 let store = test_store().await;
2629 let now = chrono::Utc::now().timestamp();
2630 {
2631 github_5xx_timestamps_lock().push(now);
2632 github_5xx_timestamps_lock().push(now);
2633 }
2634
2635 store
2636 .kv_set("cooldown:github:5xx", &(now - 600).to_string())
2637 .await
2638 .expect("kv_set past");
2639 sync_from_kv(&store).await;
2640
2641 assert!(!is_github_circuit_open());
2642 assert_eq!(github_5xx_timestamps_lock().len(), 2);
2643
2644 github_5xx_timestamps_lock().clear();
2645 }
2646
2647 #[serial(cooldown_state)]
2648 #[tokio::test]
2649 async fn sync_from_kv_keeps_github_circuit_open_for_future_row() {
2650 super::reset_global_state().await;
2651 let store = test_store().await;
2652 let future = chrono::Utc::now().timestamp() + 600;
2653 *github_5xx_circuit_open_lock() = Some(future);
2654
2655 store
2656 .kv_set("cooldown:github:5xx", &future.to_string())
2657 .await
2658 .expect("kv_set future");
2659 sync_from_kv(&store).await;
2660
2661 assert!(is_github_circuit_open());
2662
2663 *github_5xx_circuit_open_lock() = None;
2664 }
2665
2666 #[serial(cooldown_state)]
2667 #[tokio::test]
2668 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2669 super::reset_global_state().await;
2673 let store = test_store().await;
2674 let agent = "test_sync_race_window";
2675 let future = chrono::Utc::now().timestamp() + 600;
2676
2677 set_cooldown_in_memory(agent, future, "raced");
2679 assert!(is_agent_in_cooldown(agent));
2680
2681 sync_from_kv(&store).await;
2682
2683 assert!(
2684 is_agent_in_cooldown(agent),
2685 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2686 );
2687 }
2688
2689 #[serial(cooldown_state)]
2690 #[tokio::test]
2691 async fn sync_from_kv_adopts_kv_timestamp() {
2692 super::reset_global_state().await;
2695 let store = test_store().await;
2696 let agent = "test_sync_adopt_kv";
2697 let original = chrono::Utc::now().timestamp() + 600;
2698 let updated = chrono::Utc::now().timestamp() + 1200;
2699
2700 init_cooldown_store(store.clone()).await;
2701 set_cooldown_async(agent, original, "initial").await;
2702
2703 store
2704 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2705 .await
2706 .expect("kv_set updated");
2707
2708 sync_from_kv(&store).await;
2709
2710 let until = cooldown_until(agent).expect("cooldown still active");
2711 assert_eq!(
2712 until, updated,
2713 "sync_from_kv must adopt KV's authoritative timestamp"
2714 );
2715 }
2716
2717 #[serial(cooldown_state)]
2718 #[tokio::test]
2719 async fn refresh_degraded_agents_marks_cooled_agent() {
2720 super::reset_global_state().await;
2721 let store = test_store().await;
2722 let agent = "test_refresh_cooled";
2723
2724 record_agent_failure_with_message(agent, "").await;
2726 assert!(is_agent_in_cooldown(agent));
2727
2728 let agents = vec![agent.to_string()];
2730 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2731
2732 assert!(
2733 is_agent_degraded(agent),
2734 "agent in cooldown should be marked degraded"
2735 );
2736
2737 clear_agent_degraded(agent);
2739 }
2740
2741 #[serial(cooldown_state)]
2742 #[tokio::test]
2743 async fn refresh_degraded_agents_marks_no_models_agent() {
2744 super::reset_global_state().await;
2745 let store = test_store().await;
2746 let agent = "test_refresh_no_models";
2747
2748 let agents = vec![agent.to_string()];
2750 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2751
2752 assert!(
2753 is_agent_degraded(agent),
2754 "agent with no available models should be marked degraded"
2755 );
2756
2757 clear_agent_degraded(agent);
2759 }
2760
2761 #[serial(cooldown_state)]
2762 #[tokio::test]
2763 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2764 super::reset_global_state().await;
2765 let store = test_store().await;
2766 let agent = "test_refresh_rate_limited";
2767
2768 for _ in 0..4 {
2770 store.record_rate_limit(agent, "rate", None).await.unwrap();
2771 }
2772
2773 let agents = vec![agent.to_string()];
2774 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2776
2777 assert!(
2778 is_agent_degraded(agent),
2779 "agent with rate limit count >= threshold should be marked degraded"
2780 );
2781
2782 clear_agent_degraded(agent);
2784 }
2785
2786 #[serial(cooldown_state)]
2787 #[tokio::test]
2788 async fn refresh_degraded_agents_clears_healthy_agent() {
2789 super::reset_global_state().await;
2790 let store = test_store().await;
2791 let agent = "test_refresh_healthy";
2792
2793 mark_agent_degraded(agent);
2795 assert!(is_agent_degraded(agent));
2796
2797 let agents = vec![agent.to_string()];
2799 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2800
2801 assert!(
2802 !is_agent_degraded(agent),
2803 "healthy agent should have degraded flag cleared"
2804 );
2805 }
2806
2807 #[serial(cooldown_state)]
2808 #[tokio::test]
2809 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2810 super::reset_global_state().await;
2811 let store = test_store().await;
2812 let agent = "test_refresh_expired_cooldown";
2813
2814 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2818 let mut map = cooldowns().lock().unwrap();
2819 map.insert(
2820 agent.to_string(),
2821 CooldownEntry {
2822 cooldown_until: expired_ts,
2823 reason: "expired_test".to_string(),
2824 },
2825 );
2826 }
2827
2828 assert!(
2830 !is_agent_in_cooldown(agent),
2831 "expired cooldown should not be active"
2832 );
2833
2834 mark_agent_degraded(agent);
2836 assert!(is_agent_degraded(agent));
2837
2838 let agents = vec![agent.to_string()];
2840 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2841
2842 assert!(
2843 !is_agent_degraded(agent),
2844 "agent with only expired cooldown should NOT be marked degraded"
2845 );
2846
2847 clear_agent_degraded(agent);
2849 {
2850 let mut map = cooldowns().lock().unwrap();
2851 map.remove(agent);
2852 }
2853 }
2854
2855 #[serial(cooldown_state)]
2856 #[test]
2857 fn compute_backoff_grows_exponentially() {
2858 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2860 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2861 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2863 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2865 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2867 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2869 }
2870
2871 #[serial(cooldown_state)]
2872 #[test]
2873 fn compute_backoff_extended_tier_1_applies_24h() {
2874 assert_eq!(
2880 compute_backoff(5, 300, 14400, true),
2881 14400,
2882 "first at max, not extended"
2883 );
2884 assert_eq!(
2885 compute_backoff(6, 300, 14400, true),
2886 14400,
2887 "times_at_max=1, not extended"
2888 );
2889 assert_eq!(
2890 compute_backoff(7, 300, 14400, true),
2891 86400,
2892 "times_at_max=2, extended 6x"
2893 );
2894 assert_eq!(
2895 compute_backoff(8, 300, 14400, true),
2896 86400,
2897 "times_at_max=3, still extended 6x"
2898 );
2899 }
2900
2901 #[serial(cooldown_state)]
2902 #[test]
2903 fn compute_backoff_extended_tier_2_applies_48h() {
2904 assert_eq!(
2907 compute_backoff(9, 300, 14400, true),
2908 86400,
2909 "times_at_max=4, still tier 1"
2910 );
2911 assert_eq!(
2912 compute_backoff(10, 300, 14400, true),
2913 172800,
2914 "times_at_max=5, extended 12x"
2915 );
2916 assert_eq!(
2917 compute_backoff(15, 300, 14400, true),
2918 172800,
2919 "times_at_max=10, still tier 2"
2920 );
2921 assert_eq!(
2922 compute_backoff(20, 300, 14400, true),
2923 172800,
2924 "times_at_max=15, still tier 2"
2925 );
2926 }
2927
2928 #[serial(cooldown_state)]
2929 #[test]
2930 fn compute_backoff_extended_tier_with_different_base_max() {
2931 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2938 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2939 assert_eq!(
2940 compute_backoff(7, 60, 3600, true),
2941 21600,
2942 "times_at_max=2, extended 6x"
2943 );
2944 }
2945
2946 #[serial(cooldown_state)]
2947 #[test]
2948 fn compute_backoff_extended_tier_persistent_fails_example() {
2949 assert_eq!(
2953 compute_backoff(21, 300, 14400, true),
2954 172800,
2955 "21 failures should get 48h extended cooldown"
2956 );
2957
2958 assert_eq!(
2960 compute_backoff(36, 300, 14400, true),
2961 172800,
2962 "36 failures should also get 48h extended cooldown"
2963 );
2964 }
2965
2966 #[serial(cooldown_state)]
2967 #[test]
2968 fn compute_backoff_extended_tier_zero_base() {
2969 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2971 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2972 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2973 }
2974
2975 #[serial(cooldown_state)]
2976 #[test]
2977 fn detect_credit_exhaustion_billing_cycle() {
2978 assert_eq!(
2979 detect_credit_exhaustion("billing cycle"),
2980 Some(CreditExhaustionReason::BillingCycleExhausted)
2981 );
2982 assert_eq!(
2983 detect_credit_exhaustion("quota refreshed next cycle"),
2984 Some(CreditExhaustionReason::BillingCycleExhausted)
2985 );
2986 assert_eq!(
2987 detect_credit_exhaustion(
2988 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
2989 ),
2990 Some(CreditExhaustionReason::BillingCycleExhausted)
2991 );
2992 assert_eq!(
2993 detect_credit_exhaustion("monthly quota exceeded"),
2994 Some(CreditExhaustionReason::BillingCycleExhausted)
2995 );
2996 assert_eq!(
2997 detect_credit_exhaustion(
2998 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
2999 ),
3000 Some(CreditExhaustionReason::BillingCycleExhausted)
3001 );
3002 }
3003
3004 #[serial(cooldown_state)]
3005 #[test]
3006 fn detect_credit_exhaustion_billing_cycle_minimax() {
3007 assert_eq!(
3008 detect_credit_exhaustion(
3009 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
3010 ),
3011 Some(CreditExhaustionReason::BillingCycleExhausted)
3012 );
3013 assert_eq!(
3014 detect_credit_exhaustion("Upgrade your Token Plan"),
3015 Some(CreditExhaustionReason::BillingCycleExhausted)
3016 );
3017 assert_eq!(
3018 detect_credit_exhaustion("Token Plan usage limit reached"),
3019 Some(CreditExhaustionReason::BillingCycleExhausted)
3020 );
3021 assert_eq!(
3022 detect_credit_exhaustion("purchase Credits for more usage"),
3023 Some(CreditExhaustionReason::BillingCycleExhausted)
3024 );
3025 assert_eq!(
3026 detect_credit_exhaustion(
3027 "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"
3028 ),
3029 Some(CreditExhaustionReason::BillingCycleExhausted)
3030 );
3031 }
3032
3033 #[serial(cooldown_state)]
3034 #[test]
3035 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
3036 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
3041 this billing cycle. Your quota will be refreshed soon.";
3042 assert_eq!(
3043 detect_credit_exhaustion_outcome(msg),
3044 Some("billing_cycle_exhausted")
3045 );
3046 }
3047
3048 #[serial(cooldown_state)]
3049 #[test]
3050 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
3051 assert_eq!(
3052 detect_credit_exhaustion_outcome("out of credits"),
3053 Some("credit_exhausted")
3054 );
3055 assert_eq!(
3056 detect_credit_exhaustion_outcome("organization has been disabled"),
3057 Some("org_level_disabled")
3058 );
3059 assert_eq!(
3060 detect_credit_exhaustion_outcome("monthly quota exceeded"),
3061 Some("billing_cycle_exhausted")
3062 );
3063 }
3064
3065 #[test]
3066 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
3067 assert_eq!(
3068 detect_credit_exhaustion_outcome("429 Too Many Requests"),
3069 None
3070 );
3071 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
3072 }
3073
3074 #[serial(cooldown_state)]
3075 #[test]
3076 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
3077 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
3080 assert_eq!(
3081 detect_credit_exhaustion(msg),
3082 Some(CreditExhaustionReason::OrgLevelDisabled)
3083 );
3084 }
3085
3086 #[serial(cooldown_state)]
3087 #[tokio::test]
3088 async fn refresh_degraded_agents_below_threshold_not_degraded() {
3089 super::reset_global_state().await;
3090 let store = test_store().await;
3091 let agent = "test_refresh_below_threshold";
3092
3093 for _ in 0..2 {
3095 store.record_rate_limit(agent, "rate", None).await.unwrap();
3096 }
3097
3098 let agents = vec![agent.to_string()];
3099 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3101
3102 assert!(
3103 !is_agent_degraded(agent),
3104 "agent with rate limit count < threshold should not be degraded"
3105 );
3106 }
3107
3108 #[serial(cooldown_state)]
3109 #[tokio::test]
3110 async fn refresh_degraded_agents_ignores_error_limit_type() {
3111 super::reset_global_state().await;
3112 let store = test_store().await;
3113 let agent = "test_refresh_error_rows";
3114
3115 for _ in 0..3 {
3116 store.record_rate_limit(agent, "error", None).await.unwrap();
3117 }
3118
3119 let agents = vec![agent.to_string()];
3120 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3121
3122 assert!(
3123 !is_agent_degraded(agent),
3124 "three error rows in 6 hours must not mark an agent degraded"
3125 );
3126 }
3127
3128 #[serial(cooldown_state)]
3129 #[tokio::test]
3130 async fn refresh_degraded_agents_marks_three_rate_rows() {
3131 super::reset_global_state().await;
3132 let store = test_store().await;
3133 let agent = "test_refresh_three_rate_rows";
3134
3135 for _ in 0..3 {
3136 store.record_rate_limit(agent, "rate", None).await.unwrap();
3137 }
3138
3139 let agents = vec![agent.to_string()];
3140 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3141
3142 assert!(
3143 is_agent_degraded(agent),
3144 "three rate rows in 6 hours must mark an agent degraded"
3145 );
3146 assert_eq!(
3147 agent_degraded_reason(agent).as_deref(),
3148 Some("rate limit threshold exceeded")
3149 );
3150 let listed = list_degraded_agents();
3151 assert!(
3152 listed
3153 .iter()
3154 .any(|(a, r)| a == agent && r == "rate limit threshold exceeded"),
3155 "orch cooldown list must show the degraded reason, got {listed:?}"
3156 );
3157
3158 clear_agent_degraded(agent);
3159 }
3160
3161 #[serial(cooldown_state)]
3162 #[tokio::test]
3163 async fn refresh_degraded_agents_logs_reason_change() {
3164 super::reset_global_state().await;
3165 let store = test_store().await;
3166 let agent = "test_refresh_reason_change";
3167 *cooldown_store().lock().await = Some(store.clone());
3168
3169 set_agent_cooldown(agent, 600).await;
3170 let agents = vec![agent.to_string()];
3171 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3172 assert_eq!(
3173 agent_degraded_reason(agent).as_deref(),
3174 Some("agent in cooldown")
3175 );
3176
3177 clear_cooldown(agent, &store).await;
3178 let after_clear = sqlite_datetime_offset(60);
3179 for _ in 0..3 {
3180 store
3181 .record_rate_limit_at(agent, "rate", None, &after_clear)
3182 .await
3183 .unwrap();
3184 }
3185 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3186 assert_eq!(
3187 agent_degraded_reason(agent).as_deref(),
3188 Some("rate limit threshold exceeded"),
3189 "degraded reason must update when the hold changes"
3190 );
3191
3192 clear_agent_degraded(agent);
3193 }
3194
3195 #[serial(cooldown_state)]
3196 #[tokio::test]
3197 async fn clear_un_degrades_on_next_health_check() {
3198 super::reset_global_state().await;
3199 let store = test_store().await;
3200 let agent = "test_clear_un_degrades";
3201 let before = sqlite_datetime_offset(-30);
3202 let after = sqlite_datetime_offset(60);
3203
3204 for _ in 0..3 {
3205 store
3206 .record_rate_limit_at(agent, "rate", None, &before)
3207 .await
3208 .unwrap();
3209 }
3210
3211 let agents = vec![agent.to_string()];
3212 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3213 assert!(
3214 is_agent_degraded(agent),
3215 "three rate rows must mark an agent degraded"
3216 );
3217
3218 clear_cooldown(agent, &store).await;
3219 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3220 assert!(
3221 !is_agent_degraded(agent),
3222 "clear must un-degrade on the next health check"
3223 );
3224
3225 store
3226 .record_rate_limit_at(agent, "rate", None, &after)
3227 .await
3228 .unwrap();
3229 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3230 assert!(
3231 !is_agent_degraded(agent),
3232 "one rate row after the mark must not degrade again"
3233 );
3234
3235 for _ in 0..2 {
3236 store
3237 .record_rate_limit_at(agent, "rate", None, &after)
3238 .await
3239 .unwrap();
3240 }
3241 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3242 assert!(
3243 is_agent_degraded(agent),
3244 "three rate rows after the mark must degrade again"
3245 );
3246
3247 clear_agent_degraded(agent);
3248 }
3249
3250 #[serial(cooldown_state)]
3251 #[tokio::test]
3252 async fn clear_all_un_degrades_rate_limited_agent() {
3253 super::reset_global_state().await;
3254 let store = test_store().await;
3255 let agent = "test_clear_all_un_degrades";
3256 let before = sqlite_datetime_offset(-30);
3257
3258 for _ in 0..3 {
3259 store
3260 .record_rate_limit_at(agent, "rate", None, &before)
3261 .await
3262 .unwrap();
3263 }
3264
3265 let agents = vec![agent.to_string()];
3266 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3267 assert!(is_agent_degraded(agent));
3268
3269 clear_cooldown("*", &store).await;
3270 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3271 assert!(
3272 !is_agent_degraded(agent),
3273 "clear --all must un-degrade an agent held only by rate rows"
3274 );
3275 }
3276
3277 #[serial(cooldown_state)]
3280 #[tokio::test]
3281 async fn repeated_agent_failures_escalate_backoff() {
3282 super::reset_global_state().await;
3283 let store = test_store().await;
3284 *cooldown_store().lock().await = Some(store.clone());
3285
3286 let agent = "test_escalation_agent";
3287
3288 record_agent_failure_with_message(agent, "").await;
3290 let remaining_1 = {
3291 let map = cooldowns().lock().unwrap();
3292 let entry = map.get(agent).expect("should have cooldown");
3293 entry.cooldown_until - chrono::Utc::now().timestamp()
3294 };
3295 assert!(
3296 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
3297 "first failure should be ~5 min, got {remaining_1}s"
3298 );
3299
3300 {
3302 let mut map = cooldowns().lock().unwrap();
3303 map.remove(agent);
3304 }
3305
3306 record_agent_failure_with_message(agent, "").await;
3308 let remaining_2 = {
3309 let map = cooldowns().lock().unwrap();
3310 let entry = map.get(agent).expect("should have cooldown");
3311 entry.cooldown_until - chrono::Utc::now().timestamp()
3312 };
3313 assert!(
3314 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
3315 "second failure should be ~15 min, got {remaining_2}s"
3316 );
3317 }
3318
3319 #[serial(cooldown_state)]
3320 #[tokio::test]
3321 async fn vendor_past_retry_falls_back_to_backoff() {
3322 super::reset_global_state().await;
3323 let store = test_store().await;
3327 *cooldown_store().lock().await = Some(store.clone());
3328
3329 let agent = "test_vendor_past_retry";
3330
3331 {
3333 let mut map = cooldowns().lock().unwrap();
3334 map.remove(agent);
3335 }
3336
3337 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
3339
3340 record_agent_failure_with_message(agent, past_msg).await;
3341
3342 assert!(
3346 is_agent_in_cooldown(agent),
3347 "agent should be in cooldown after failure"
3348 );
3349
3350 let remaining = {
3351 let map = cooldowns().lock().unwrap();
3352 let entry = map.get(agent).expect("cooldown entry should exist");
3353 entry.cooldown_until - chrono::Utc::now().timestamp()
3354 };
3355
3356 assert!(
3357 remaining >= BACKOFF_BASE_SECS - 5,
3358 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3359 );
3360
3361 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3363 let stored = store
3364 .kv_get(&kv_key)
3365 .await
3366 .unwrap()
3367 .expect("failure count should be stored");
3368 assert_eq!(stored, "1");
3369 }
3370
3371 #[serial(cooldown_state)]
3372 #[tokio::test]
3373 async fn record_agent_success_resets_backoff() {
3374 super::reset_global_state().await;
3375 let store = test_store().await;
3376 *cooldown_store().lock().await = Some(store.clone());
3377
3378 let agent = "test_success_reset_agent";
3379 let model = "test_model";
3380
3381 for _ in 0..3 {
3383 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3384 }
3385 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3387 let count: u32 = store
3388 .kv_get(&kv_key)
3389 .await
3390 .unwrap()
3391 .unwrap()
3392 .parse()
3393 .unwrap();
3394 assert_eq!(count, 3);
3395
3396 record_agent_success(agent, model).await;
3398
3399 let count_after: u32 = store
3400 .kv_get(&kv_key)
3401 .await
3402 .unwrap()
3403 .unwrap()
3404 .parse()
3405 .unwrap();
3406 assert_eq!(count_after, 0, "success should reset failure count to 0");
3407
3408 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3410 let model_count: u32 = store
3411 .kv_get(&model_kv_key)
3412 .await
3413 .unwrap()
3414 .unwrap()
3415 .parse()
3416 .unwrap();
3417 assert_eq!(
3418 model_count, 0,
3419 "success should reset model failure count to 0"
3420 );
3421 }
3422
3423 #[serial(cooldown_state)]
3424 #[tokio::test]
3425 async fn record_agent_success_without_model_skips_empty_model_key() {
3426 super::reset_global_state().await;
3427 let store = test_store().await;
3428 *cooldown_store().lock().await = Some(store.clone());
3429
3430 let agent = "test_success_without_model";
3431
3432 record_agent_success(agent, "").await;
3433
3434 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3435 let agent_count = store
3436 .kv_get(&agent_kv_key)
3437 .await
3438 .unwrap()
3439 .expect("agent failure count should be stored");
3440 assert_eq!(agent_count, "0");
3441
3442 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3443 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3444 assert!(
3445 empty_model_count.is_none(),
3446 "success without model should not create an empty model failure count key"
3447 );
3448 }
3449
3450 #[serial(cooldown_state)]
3451 #[tokio::test]
3452 async fn clear_cooldown_resets_failure_counts() {
3453 super::reset_global_state().await;
3454 let store = test_store().await;
3455 *cooldown_store().lock().await = Some(store.clone());
3456
3457 let agent = "test_clear_fc_agent";
3458
3459 for _ in 0..5 {
3461 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3462 }
3463 set_agent_cooldown(agent, 3600).await;
3465
3466 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3468 let count: u32 = store
3469 .kv_get(&kv_key)
3470 .await
3471 .unwrap()
3472 .unwrap()
3473 .parse()
3474 .unwrap();
3475 assert_eq!(count, 5);
3476
3477 clear_cooldown(agent, &store).await;
3479
3480 let count_after: u32 = store
3482 .kv_get(&kv_key)
3483 .await
3484 .unwrap()
3485 .unwrap()
3486 .parse()
3487 .unwrap();
3488 assert_eq!(
3489 count_after, 0,
3490 "clear_cooldown should reset failure count to 0"
3491 );
3492
3493 assert!(
3495 !is_agent_in_cooldown(agent),
3496 "agent should not be in cooldown after clear"
3497 );
3498 }
3499
3500 #[serial(cooldown_state)]
3501 #[tokio::test]
3502 async fn clear_all_cooldowns_resets_all_failure_counts() {
3503 super::reset_global_state().await;
3504 let store = test_store().await;
3505 *cooldown_store().lock().await = Some(store.clone());
3506
3507 let agents = ["test_clear_all_a", "test_clear_all_b"];
3508 for agent in &agents {
3509 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3510 set_agent_cooldown(agent, 3600).await;
3511 }
3512
3513 clear_cooldown("*", &store).await;
3514
3515 for agent in &agents {
3516 assert!(!is_agent_in_cooldown(agent));
3517 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3518 let count: u32 = store
3519 .kv_get(&kv_key)
3520 .await
3521 .unwrap()
3522 .unwrap()
3523 .parse()
3524 .unwrap();
3525 assert_eq!(
3526 count, 0,
3527 "clear --all should reset failure count for {agent}"
3528 );
3529 }
3530 }
3531
3532 #[serial(cooldown_state)]
3533 #[test]
3534 fn compute_backoff_saturates_on_large_counts() {
3535 assert_eq!(
3537 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3538 BACKOFF_MAX_SECS * 12
3539 );
3540 assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
3541 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
3542 }
3543
3544 #[serial(cooldown_state)]
3545 #[test]
3546 fn compute_backoff_zero_base_returns_zero() {
3547 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3548 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3549 }
3550
3551 #[serial(cooldown_state)]
3554 #[tokio::test]
3555 async fn github_circuit_stays_closed_below_threshold() {
3556 super::reset_global_state().await;
3557 {
3559 let mut open = github_5xx_circuit_open().lock().unwrap();
3560 *open = None;
3561 let mut ts = github_5xx_timestamps().lock().unwrap();
3562 ts.clear();
3563 let mut map = cooldown_lock();
3567 map.remove("github:5xx");
3568 }
3569
3570 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3572 record_github_5xx().await;
3573 }
3574
3575 assert!(
3576 !is_github_circuit_open(),
3577 "circuit should stay closed below threshold"
3578 );
3579 assert_eq!(github_circuit_remaining_secs(), 0);
3580 }
3581
3582 #[serial(cooldown_state)]
3583 #[tokio::test]
3584 async fn github_circuit_trips_at_threshold() {
3585 super::reset_global_state().await;
3586 {
3588 let mut open = github_5xx_circuit_open_lock();
3589 *open = None;
3590 let mut ts = github_5xx_timestamps_lock();
3591 ts.clear();
3592 let mut map = cooldown_lock();
3593 map.remove("github:5xx");
3594 }
3595
3596 for _ in 0..GITHUB_5XX_THRESHOLD {
3598 record_github_5xx().await;
3599 }
3600
3601 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3602 assert!(
3603 github_circuit_remaining_secs() > 0,
3604 "remaining should be > 0 when open"
3605 );
3606 assert!(
3607 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3608 "remaining should not exceed cooldown duration"
3609 );
3610
3611 {
3613 let mut open = github_5xx_circuit_open().lock().unwrap();
3614 *open = None;
3615 }
3616 }
3617
3618 #[serial(cooldown_state)]
3619 #[tokio::test]
3620 async fn github_circuit_cannot_double_trip() {
3621 super::reset_global_state().await;
3622 {
3624 let mut open = github_5xx_circuit_open_lock();
3625 *open = None;
3626 let mut ts = github_5xx_timestamps_lock();
3627 ts.clear();
3628 let mut map = cooldown_lock();
3629 map.remove("github:5xx");
3630 }
3631
3632 for _ in 0..GITHUB_5XX_THRESHOLD {
3634 record_github_5xx().await;
3635 }
3636 assert!(is_github_circuit_open());
3637
3638 let first_remaining = github_circuit_remaining_secs();
3639
3640 for _ in 0..GITHUB_5XX_THRESHOLD {
3642 record_github_5xx().await;
3643 }
3644
3645 let second_remaining = github_circuit_remaining_secs();
3646 assert!(
3647 second_remaining <= first_remaining,
3648 "additional errors should not extend cooldown"
3649 );
3650
3651 {
3653 let mut open = github_5xx_circuit_open().lock().unwrap();
3654 *open = None;
3655 }
3656 }
3657
3658 #[serial(cooldown_state)]
3659 #[tokio::test]
3660 async fn github_circuit_recovers_after_cooldown() {
3661 super::reset_global_state().await;
3662 {
3664 let mut open = github_5xx_circuit_open_lock();
3665 *open = None;
3666 let mut ts = github_5xx_timestamps_lock();
3667 ts.clear();
3668 let mut map = cooldown_lock();
3669 map.remove("github:5xx");
3670 }
3671
3672 {
3674 let mut open = github_5xx_circuit_open().lock().unwrap();
3675 *open = Some(chrono::Utc::now().timestamp() - 1);
3676 }
3677
3678 assert!(
3680 !is_github_circuit_open(),
3681 "circuit should auto-recover after cooldown"
3682 );
3683 }
3684
3685 #[serial(cooldown_state)]
3686 #[tokio::test]
3687 async fn github_circuit_sliding_window_ages_out() {
3688 super::reset_global_state().await;
3689 {
3691 let mut open = github_5xx_circuit_open_lock();
3692 *open = None;
3693 let mut ts = github_5xx_timestamps_lock();
3694 ts.clear();
3695 let mut map = cooldown_lock();
3696 map.remove("github:5xx");
3697 }
3698
3699 {
3701 let mut ts = github_5xx_timestamps().lock().unwrap();
3702 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3703 for _ in 0..GITHUB_5XX_THRESHOLD {
3704 ts.push(old);
3705 }
3706 }
3707
3708 record_github_5xx().await;
3710
3711 assert!(
3712 !is_github_circuit_open(),
3713 "old errors outside window should not trip circuit"
3714 );
3715 }
3716}