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_model_unavailable(agent_name: &str, model: &str, message: &str) {
745 if is_model_permanently_gone(message) {
746 record_persistent_model_failure(agent_name, model).await;
747 } else {
748 record_model_failure(agent_name, model).await;
749 }
750}
751
752fn is_model_permanently_gone(message: &str) -> bool {
753 let lower = message.to_lowercase();
754 lower.contains("not found")
755 || lower.contains("not supported")
756 || lower.contains("not available in your country")
757 || lower.contains("not available in your region")
758}
759
760pub async fn record_transient_model_failure(agent_name: &str, model: &str) {
766 let key = format!("{agent_name}:{model}");
767 let store_opt = cooldown_store().lock().await.clone();
768 let count = read_and_increment_failure_count(&store_opt, &key).await;
769 let cooldown_secs = compute_backoff(
770 count,
771 TRANSIENT_MODEL_BACKOFF_BASE_SECS,
772 TRANSIENT_MODEL_BACKOFF_MAX_SECS,
773 false,
774 );
775 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
776 set_cooldown_async(&key, cooldown_until, "transient_model_error").await;
777}
778
779pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
786 let key = format!("{agent_name}:{model}");
787 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
788 set_cooldown_async(&key, cooldown_until, reason).await;
789}
790
791pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
793 let key = format!("{agent_name}:{model}");
794 set_cooldown_async(&key, until, "vendor_reset_at").await;
795}
796
797pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
804 let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
805 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
806}
807
808pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
811 let store_opt = cooldown_store().lock().await.clone();
812 let store = match store_opt {
813 Some(store) => store,
814 None => {
815 tracing::debug!(
816 agent = agent_name,
817 model,
818 "skipping silence count record (no KV store)"
819 );
820 return None;
821 }
822 };
823
824 let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
825 let now = chrono::Utc::now().timestamp();
826 let window_start = now - SILENCE_COUNT_WINDOW_SECS;
827
828 let mut timestamps = match store.kv_get(&key).await {
829 Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
830 Ok(ts) => ts,
831 Err(e) => {
832 tracing::warn!(
833 kv_key = %key,
834 err = %e,
835 "failed to parse silence timestamps from KV — resetting window"
836 );
837 Vec::new()
838 }
839 },
840 Ok(None) => Vec::new(),
841 Err(err) => {
842 tracing::warn!(
843 kv_key = %key,
844 err = %err,
845 "failed to load silence count from KV"
846 );
847 Vec::new()
848 }
849 };
850
851 timestamps.retain(|ts| *ts >= window_start);
852 timestamps.push(now);
853
854 let count = timestamps.len();
855 let mut extended_cooldown_applied = false;
856 if count >= SILENCE_COUNT_THRESHOLD {
857 set_model_cooldown(
858 agent_name,
859 model,
860 SILENCE_EXTENDED_COOLDOWN_SECS,
861 "silence_detected",
862 )
863 .await;
864 extended_cooldown_applied = true;
865 timestamps.clear();
866 }
867
868 let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
869 if let Err(err) = store.kv_set(&key, &value).await {
870 tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
871 }
872
873 Some(SilenceCountResult {
874 count,
875 extended_cooldown_applied,
876 })
877}
878
879pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
881 let key = format!("{agent_name}:{model}");
882 is_in_cooldown(&key)
883}
884
885pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
887 is_in_cooldown(agent_name)
888}
889
890pub fn cooldown_until(key: &str) -> Option<i64> {
892 let map = cooldown_lock();
893 map.get(key).and_then(|entry| {
894 let now = chrono::Utc::now().timestamp();
895 if now < entry.cooldown_until {
896 Some(entry.cooldown_until)
897 } else {
898 None
899 }
900 })
901}
902
903pub fn cooldown_reason(key: &str) -> Option<String> {
905 let map = cooldown_lock();
906 map.get(key).and_then(|entry| {
907 let now = chrono::Utc::now().timestamp();
908 if now < entry.cooldown_until {
909 Some(entry.reason.clone())
910 } else {
911 None
912 }
913 })
914}
915
916pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
921 let now = chrono::Utc::now().timestamp();
922 let map = cooldown_lock();
923 let mut result: Vec<(String, i64, String)> = map
924 .iter()
925 .filter(|(_, entry)| now < entry.cooldown_until)
926 .map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
927 .collect();
928 result.sort_by(|a, b| a.0.cmp(&b.0));
929 result
930}
931
932pub async fn mark_rate_limits_cleared(store: &Arc<crate::store::TaskStore>, agents: &[String]) {
935 let now = chrono::Utc::now().timestamp().to_string();
936 let mut seen = std::collections::HashSet::new();
937 for key in agents {
938 let agent = agent_from_key(key);
939 if agent.is_empty() || !seen.insert(agent.to_string()) {
940 continue;
941 }
942 let kv_key = format!("{RATE_LIMITS_CLEARED_AT_PREFIX}{agent}");
943 if let Err(e) = store.kv_set(&kv_key, &now).await {
944 tracing::warn!(agent, err = %e, "failed to mark rate limits cleared");
945 }
946 }
947}
948
949pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
956 let mut agents_to_mark: Vec<String> = Vec::new();
957 if key == "*" {
958 let keys: Vec<String> = {
959 let mut map = cooldown_lock();
960 let keys: Vec<String> = map.keys().cloned().collect();
961 map.clear();
962 keys
963 };
964 for k in &keys {
965 let kv_key = format!("{KV_PREFIX}{k}");
966 if let Err(e) = store.kv_set(&kv_key, "0").await {
967 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
968 }
969 let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
971 if let Err(e) = store.kv_set(&fc_key, "0").await {
972 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
973 }
974 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
976 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
977 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
978 }
979 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
981 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
982 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
983 }
984 }
985 if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
989 for (fc_key, _) in fc_entries {
990 if let Err(e) = store.kv_set(&fc_key, "0").await {
991 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
992 }
993 }
994 }
995 if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
996 for (cfc_key, _) in cfc_entries {
997 if let Err(e) = store.kv_set(&cfc_key, "0").await {
998 tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
999 }
1000 }
1001 }
1002 if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
1003 for (pfc_key, _) in pfc_entries {
1004 if let Err(e) = store.kv_set(&pfc_key, "0").await {
1005 tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
1006 }
1007 }
1008 }
1009 tracing::info!(
1010 count = keys.len(),
1011 "cleared all cooldowns and failure counts"
1012 );
1013 agents_to_mark.extend(keys.iter().map(|k| agent_from_key(k).to_string()));
1014 match store.rate_limit_agents().await {
1015 Ok(rl_agents) => agents_to_mark.extend(rl_agents),
1016 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits agents while clearing"),
1017 }
1018 match store.kv_list_prefix(RATE_LIMITS_CLEARED_AT_PREFIX).await {
1019 Ok(marks) => {
1020 for (k, _) in marks {
1021 let agent = k.trim_start_matches(RATE_LIMITS_CLEARED_AT_PREFIX);
1022 if !agent.is_empty() {
1023 agents_to_mark.push(agent.to_string());
1024 }
1025 }
1026 }
1027 Err(e) => tracing::warn!(err = %e, "failed to list rate_limits_cleared_at keys"),
1028 }
1029 } else {
1030 {
1031 let mut map = cooldown_lock();
1032 map.remove(key);
1033 }
1034 let kv_key = format!("{KV_PREFIX}{key}");
1035 if let Err(e) = store.kv_set(&kv_key, "0").await {
1036 tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
1037 }
1038 let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
1040 if let Err(e) = store.kv_set(&fc_key, "0").await {
1041 tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
1042 }
1043 let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
1045 if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
1046 tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
1047 }
1048 let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
1050 if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
1051 tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
1052 }
1053 tracing::info!(key, "cleared cooldown and failure count");
1054 agents_to_mark.push(agent_from_key(key).to_string());
1055 }
1056 mark_rate_limits_cleared(store, &agents_to_mark).await;
1057}
1058
1059fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
1068 let mut map = cooldown_lock();
1069 if let Some(existing) = map.get(key) {
1070 if existing.cooldown_until >= cooldown_until {
1074 return false;
1075 }
1076 }
1077 map.insert(
1078 key.to_string(),
1079 CooldownEntry {
1080 cooldown_until,
1081 reason: reason.to_string(),
1082 },
1083 );
1084 true
1085}
1086
1087async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
1097 if !set_cooldown_in_memory(key, cooldown_until, reason) {
1098 return false;
1099 }
1100 let store_opt = cooldown_store().lock().await.clone();
1101 if let Some(store) = store_opt {
1102 let kv_key = format!("{KV_PREFIX}{key}");
1103 let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
1104 if let Err(e) = store.kv_set(&kv_key, &value).await {
1105 tracing::warn!(
1106 kv_key,
1107 err = %e,
1108 "failed to persist cooldown to KV store — rolling back in-memory entry"
1109 );
1110 let mut map = cooldown_lock();
1111 if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
1115 map.remove(key);
1116 }
1117 return false;
1118 }
1119 }
1120 true
1121}
1122
1123fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
1125 cooldowns().lock().unwrap_or_else(|e| {
1126 tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
1127 e.into_inner()
1128 })
1129}
1130
1131fn degraded_agents_lock() -> std::sync::MutexGuard<'static, HashMap<String, String>> {
1133 degraded_agents().lock().unwrap_or_else(|e| {
1134 tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
1135 e.into_inner()
1136 })
1137}
1138
1139fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
1141 github_5xx_timestamps().lock().unwrap_or_else(|e| {
1142 tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1143 e.into_inner()
1144 })
1145}
1146
1147fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
1149 github_5xx_circuit_open().lock().unwrap_or_else(|e| {
1150 tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
1151 e.into_inner()
1152 })
1153}
1154
1155fn is_in_cooldown(key: &str) -> bool {
1157 let map = cooldown_lock();
1158 if let Some(entry) = map.get(key) {
1159 let now = chrono::Utc::now().timestamp();
1160 return now < entry.cooldown_until;
1161 }
1162 false
1163}
1164
1165pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
1171
1172pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
1176
1177fn degraded_agents() -> &'static Mutex<HashMap<String, String>> {
1180 static DEGRADED: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
1181 DEGRADED.get_or_init(|| Mutex::new(HashMap::new()))
1182}
1183
1184pub fn is_agent_degraded(agent: &str) -> bool {
1189 let set = degraded_agents_lock();
1190 set.contains_key(agent)
1191}
1192
1193pub fn agent_degraded_reason(agent: &str) -> Option<String> {
1195 degraded_agents_lock().get(agent).cloned()
1196}
1197
1198pub fn list_degraded_agents() -> Vec<(String, String)> {
1200 let mut entries: Vec<(String, String)> = degraded_agents_lock()
1201 .iter()
1202 .map(|(agent, reason)| (agent.clone(), reason.clone()))
1203 .collect();
1204 entries.sort_by(|a, b| a.0.cmp(&b.0));
1205 entries
1206}
1207
1208pub fn mark_agent_degraded(agent: &str) {
1210 mark_agent_degraded_with_reason(agent, "degraded");
1211}
1212
1213fn mark_agent_degraded_with_reason(agent: &str, reason: &str) {
1214 let mut set = degraded_agents_lock();
1215 set.insert(agent.to_string(), reason.to_string());
1216}
1217
1218pub fn clear_agent_degraded(agent: &str) {
1220 let mut set = degraded_agents_lock();
1221 set.remove(agent);
1222}
1223
1224pub async fn refresh_degraded_agents(
1238 store: &Arc<crate::store::TaskStore>,
1239 available_agents: &[String],
1240 model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
1241 window_hours: u32,
1242 threshold: i64,
1243) {
1244 clear_expired_cooldowns();
1248
1249 let start = chrono::Utc::now();
1252 let counts = match store.recent_rate_limit_counts(window_hours).await {
1253 Ok(c) => {
1254 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1255 tracing::info!(
1256 duration_ms = dur_ms,
1257 window_hours,
1258 "recent_rate_limit_counts query completed"
1259 );
1260 c
1261 }
1262 Err(e) => {
1263 let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
1264 tracing::warn!(
1265 err = %e,
1266 duration_ms = dur_ms,
1267 "failed to query recent rate limit counts for health check"
1268 );
1269 return;
1270 }
1271 };
1272
1273 for agent in available_agents {
1274 let in_cooldown = is_agent_in_cooldown(agent);
1275 let no_models = !model_checker(agent);
1276 let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
1277 let over_threshold = rate_limit_count >= threshold;
1278
1279 let degraded = in_cooldown || no_models || over_threshold;
1280
1281 if degraded {
1282 let reason = if in_cooldown {
1283 "agent in cooldown"
1284 } else if no_models {
1285 "all models cooled"
1286 } else {
1287 "rate limit threshold exceeded"
1288 };
1289 let previous = agent_degraded_reason(agent);
1290 match previous.as_deref() {
1291 None => tracing::warn!(
1292 agent,
1293 reason,
1294 rate_limit_count,
1295 window_hours,
1296 threshold,
1297 "pre-emptive health check: marking agent as degraded"
1298 ),
1299 Some(prev) if prev != reason => tracing::warn!(
1300 agent,
1301 reason,
1302 previous = prev,
1303 rate_limit_count,
1304 window_hours,
1305 threshold,
1306 "pre-emptive health check: degraded reason changed"
1307 ),
1308 Some(_) => {}
1309 }
1310 mark_agent_degraded_with_reason(agent, reason);
1311 } else {
1312 if is_agent_degraded(agent) {
1313 tracing::info!(
1314 agent,
1315 rate_limit_count,
1316 "pre-emptive health check: agent recovered, clearing degraded flag"
1317 );
1318 }
1319 clear_agent_degraded(agent);
1320 }
1321 }
1322}
1323
1324pub fn clear_expired_cooldowns() {
1326 let mut map = cooldown_lock();
1327 let now = chrono::Utc::now().timestamp();
1328 map.retain(|_key, entry| now < entry.cooldown_until);
1329}
1330
1331const GITHUB_5XX_WINDOW_SECS: i64 = 120;
1337
1338const GITHUB_5XX_THRESHOLD: usize = 5;
1340
1341const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
1343
1344fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
1346 static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
1347 TS.get_or_init(|| Mutex::new(Vec::new()))
1348}
1349
1350fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
1352 static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
1353 OPEN.get_or_init(|| Mutex::new(None))
1354}
1355
1356pub async fn record_github_5xx() {
1361 let now = chrono::Utc::now().timestamp();
1362 let window_start = now - GITHUB_5XX_WINDOW_SECS;
1363
1364 let should_trip = {
1365 let mut ts = github_5xx_timestamps_lock();
1366 ts.retain(|t| *t >= window_start);
1367 ts.push(now);
1368 ts.len() >= GITHUB_5XX_THRESHOLD
1369 };
1370
1371 if should_trip {
1372 let cooldown_until = {
1373 let mut open = github_5xx_circuit_open_lock();
1374 if open.is_none() {
1375 let cd = now + GITHUB_5XX_COOLDOWN_SECS;
1376 *open = Some(cd);
1377 tracing::warn!(
1378 errors_in_window = GITHUB_5XX_THRESHOLD,
1379 window_secs = GITHUB_5XX_WINDOW_SECS,
1380 cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
1381 "GitHub 5xx circuit breaker OPEN — throttling non-critical work"
1382 );
1383 Some(cd)
1384 } else {
1385 None
1386 }
1387 };
1388 if let Some(cd) = cooldown_until {
1390 let store_opt = cooldown_store().lock().await.clone();
1391 if let Some(store) = store_opt {
1392 let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
1393 }
1394 set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
1398 }
1399 let mut ts = github_5xx_timestamps_lock();
1401 ts.clear();
1402 }
1403}
1404
1405pub fn is_github_circuit_open() -> bool {
1413 {
1416 let mut open = github_5xx_circuit_open_lock();
1417 if let Some(until) = *open {
1418 let now = chrono::Utc::now().timestamp();
1419 if now >= until {
1420 *open = None;
1421 tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
1422 return false;
1423 }
1424 return true;
1425 }
1426 }
1427
1428 if let Some(until) = cooldown_until("github:5xx") {
1432 let mut open = github_5xx_circuit_open_lock();
1435 if open.is_none() {
1436 *open = Some(until);
1437 let now = chrono::Utc::now().timestamp();
1438 if now < until {
1439 let remaining = until - now;
1440 tracing::warn!(
1441 remaining_secs = remaining,
1442 "GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
1443 );
1444 }
1445 }
1446 return true;
1447 }
1448
1449 false
1450}
1451
1452pub fn github_circuit_remaining_secs() -> u64 {
1455 {
1457 let open = github_5xx_circuit_open_lock();
1458 if let Some(until) = *open {
1459 let now = chrono::Utc::now().timestamp();
1460 if now < until {
1461 return (until - now) as u64;
1462 }
1463 return 0;
1464 }
1465 }
1466
1467 if let Some(until) = cooldown_until("github:5xx") {
1469 let now = chrono::Utc::now().timestamp();
1470 if now < until {
1471 return (until - now) as u64;
1472 }
1473 }
1474 0
1475}
1476
1477pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
1481 if error_message.is_empty() {
1482 return None;
1483 }
1484
1485 let lower = error_message.to_lowercase();
1486 for retry_marker in [
1490 "try again at ",
1491 "reset at ",
1492 "resets at ",
1493 "resets ",
1494 "reset ",
1495 ] {
1496 if let Some(idx) = lower.find(retry_marker) {
1497 let date_str = &lower[idx + retry_marker.len()..];
1498 let date_str = date_str
1500 .split(['.', '\n'])
1501 .next()
1502 .unwrap_or(date_str)
1503 .trim();
1504
1505 let cleaned = date_str
1508 .replace("st,", ",")
1509 .replace("nd,", ",")
1510 .replace("rd,", ",")
1511 .replace("th,", ",");
1512
1513 for fmt in &[
1514 "%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", ] {
1521 if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
1522 use chrono::TimeZone;
1525 let utc = match chrono::Local.from_local_datetime(&dt) {
1526 chrono::LocalResult::Single(local_dt) => {
1527 local_dt.with_timezone(&chrono::Utc)
1528 }
1529 chrono::LocalResult::Ambiguous(earliest, _) => {
1530 earliest.with_timezone(&chrono::Utc)
1532 }
1533 chrono::LocalResult::None => {
1534 dt.and_utc()
1536 }
1537 };
1538 tracing::info!(
1539 retry_at = %utc,
1540 raw = date_str,
1541 marker = retry_marker,
1542 "parsed rate limit retry-at date"
1543 );
1544 return Some(utc.timestamp());
1545 }
1546 }
1547 tracing::debug!(
1548 raw = date_str,
1549 marker = retry_marker,
1550 "could not parse retry-at date"
1551 );
1552
1553 let time_only = date_str
1558 .split_once('(')
1559 .map(|(t, _)| t)
1560 .unwrap_or(date_str)
1561 .trim();
1562 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1564 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1565 let last = &compact[split..];
1566 if !last.contains(':')
1567 && (last.ends_with("am") || last.ends_with("pm"))
1568 && last.len() > 2
1569 {
1570 let at = split + last.len() - 2;
1571 compact.insert_str(at, ":00");
1572 }
1573 let md = compact.replace(" at ", " ");
1575 let now_local = chrono::Local::now();
1576 let this_year = chrono::Datelike::year(&now_local);
1577 for year in [this_year, this_year + 1] {
1578 let with_year = format!("{year} {md}");
1579 if let Ok(dt) =
1580 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1581 {
1582 use chrono::TimeZone;
1583 if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
1584 if local_dt > now_local {
1585 return Some(local_dt.timestamp());
1586 }
1587 }
1588 }
1589 }
1590 let time_only = compact.as_str();
1591 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1592 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1593 let local_now = chrono::Local::now();
1596 let today = local_now.date_naive();
1597 let candidate = chrono::NaiveDateTime::new(today, t);
1598 let candidate_local: chrono::DateTime<chrono::Local> =
1601 match chrono::Local.from_local_datetime(&candidate) {
1602 chrono::LocalResult::Single(dt) => dt,
1603 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1604 chrono::LocalResult::None => {
1605 let utc_dt = candidate.and_utc();
1607 return Some(utc_dt.timestamp());
1608 }
1609 };
1610 let chosen = if candidate_local <= local_now {
1611 let tomorrow = today.succ_opt().unwrap_or(today);
1613 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1614 match chrono::Local.from_local_datetime(&ndt) {
1615 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1616 chrono::LocalResult::Ambiguous(earliest, _) => {
1617 earliest.with_timezone(&chrono::Utc)
1618 }
1619 chrono::LocalResult::None => ndt.and_utc(),
1620 }
1621 } else {
1622 candidate_local.with_timezone(&chrono::Utc)
1623 };
1624 tracing::info!(
1625 retry_at = %chosen,
1626 raw = date_str,
1627 marker = retry_marker,
1628 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1629 );
1630 return Some(chosen.timestamp());
1631 }
1632 }
1633 }
1634 }
1635
1636 None
1640}
1641
1642pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1653 if error_message.is_empty() {
1654 return None;
1655 }
1656
1657 let lower = error_message.to_lowercase();
1658 let has_context = lower.contains("usage limit")
1659 || lower.contains("quota")
1660 || lower.contains("window")
1661 || lower.contains("limit");
1662 if !has_context {
1663 return None;
1664 }
1665
1666 static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
1667 Regex::new(
1668 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1669 )
1670 .expect("usage window regex is valid")
1671 });
1672
1673 let mut max_secs: Option<u64> = None;
1674 for caps in WINDOW_RE.captures_iter(&lower) {
1675 let secs = if let Some(num) = caps.get(1) {
1676 let n: u64 = num.as_str().parse().ok()?;
1677 let unit = caps.get(2)?.as_str().to_lowercase();
1678 let multiplier = match unit.as_str() {
1679 "hour" | "hours" | "hr" | "hrs" => 3600,
1680 "minute" | "minutes" | "min" | "mins" => 60,
1681 "day" | "days" => 86400,
1682 _ => return None,
1683 };
1684 n.checked_mul(multiplier)?
1685 } else if let Some(word) = caps.get(3) {
1686 match word.as_str().to_lowercase().as_str() {
1687 "hourly" => 3600,
1688 "daily" => 86400,
1689 "weekly" => 604800,
1690 "monthly" => 2592000,
1691 _ => return None,
1692 }
1693 } else {
1694 continue;
1695 };
1696 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1697 }
1698
1699 max_secs
1700}
1701
1702#[cfg(test)]
1703pub(crate) async fn reset_global_state() {
1704 cooldowns().lock().unwrap().clear();
1705 *cooldown_store().lock().await = None;
1706 github_5xx_timestamps().lock().unwrap().clear();
1707 *github_5xx_circuit_open().lock().unwrap() = None;
1708 degraded_agents().lock().unwrap().clear();
1709 crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1710}
1711
1712#[cfg(test)]
1713mod tests {
1714 use super::*;
1715 use serial_test::serial;
1716
1717 async fn test_store() -> Arc<crate::store::TaskStore> {
1718 Arc::new(
1719 crate::store::TaskStore::open_memory()
1720 .await
1721 .expect("in-memory store"),
1722 )
1723 }
1724
1725 fn sqlite_datetime_offset(secs: i64) -> String {
1726 (chrono::Utc::now() + chrono::Duration::seconds(secs))
1727 .format("%Y-%m-%d %H:%M:%S")
1728 .to_string()
1729 }
1730
1731 #[serial(cooldown_state)]
1732 #[tokio::test]
1733 async fn init_loads_unexpired_cooldowns_from_kv() {
1734 super::reset_global_state().await;
1735 let store = test_store().await;
1736 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1739 .kv_set(
1740 &format!("{KV_PREFIX}kimi:k2p5"),
1741 &cooldown_until.to_string(),
1742 )
1743 .await
1744 .unwrap();
1745
1746 let expired = chrono::Utc::now().timestamp() - 1;
1748 store
1749 .kv_set(
1750 &format!("{KV_PREFIX}opencode:old-free"),
1751 &expired.to_string(),
1752 )
1753 .await
1754 .unwrap();
1755
1756 {
1757 let mut map = cooldown_lock();
1758 map.clear();
1759 }
1760
1761 init_cooldown_store(store).await;
1762
1763 assert!(
1764 is_model_in_cooldown("kimi", "k2p5"),
1765 "unexpired cooldown should be loaded"
1766 );
1767 assert!(
1768 !is_model_in_cooldown("opencode", "old-free"),
1769 "expired cooldown should not be loaded"
1770 );
1771 }
1772
1773 #[serial(cooldown_state)]
1774 #[tokio::test]
1775 async fn record_model_failure_writes_to_kv() {
1776 super::reset_global_state().await;
1777 let store = test_store().await;
1778
1779 *cooldown_store().lock().await = Some(store.clone());
1780
1781 record_model_failure("testagent_persist", "testmodel_persist").await;
1782
1783 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1785
1786 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1787 let stored = store.kv_get(&kv_key).await.unwrap();
1788 assert!(stored.is_some(), "model failure should be persisted to KV");
1789 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1790 let now = chrono::Utc::now().timestamp();
1791 assert!(ts > now, "persisted timestamp should be in the future");
1793 assert_eq!(reason, "model_error");
1794 }
1795
1796 #[serial(cooldown_state)]
1797 #[tokio::test]
1798 async fn cooldown_reason_survives_restart_via_kv() {
1799 super::reset_global_state().await;
1800 let store = test_store().await;
1801 let agent = "test_reason_restart";
1802 let future = chrono::Utc::now().timestamp() + 600;
1803
1804 init_cooldown_store(store.clone()).await;
1805 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1806
1807 {
1809 let mut map = cooldown_lock();
1810 map.clear();
1811 }
1812 init_cooldown_store(store).await;
1813
1814 let entries = list_all_cooldowns();
1815 let entry = entries.iter().find(|(key, ..)| key == agent);
1816 assert_eq!(
1817 entry.map(|(_, _, reason)| reason.as_str()),
1818 Some("billing_cycle_exhausted"),
1819 "cooldown reason must survive a restart, not collapse to 'persisted'"
1820 );
1821 }
1822
1823 #[serial(cooldown_state)]
1824 #[tokio::test]
1825 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1826 super::reset_global_state().await;
1827 let store = test_store().await;
1828 *cooldown_store().lock().await = Some(store.clone());
1829
1830 let agent = "testagent_persistent";
1831 let model = "testmodel_persistent";
1832 let key = format!("{agent}:{model}");
1833
1834 record_persistent_model_failure(agent, model).await;
1836 let now = chrono::Utc::now().timestamp();
1837 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1838 let first_remaining = first_until.saturating_sub(now);
1839 assert!(
1840 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1841 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1842 );
1843
1844 record_persistent_model_failure(agent, model).await;
1846 let now2 = chrono::Utc::now().timestamp();
1847 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1848 let second_remaining = second_until.saturating_sub(now2);
1849 assert!(
1850 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1851 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1852 );
1853
1854 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1855 let stored_count = store
1856 .kv_get(&count_key)
1857 .await
1858 .expect("kv read should succeed")
1859 .expect("failure count should be written");
1860 assert_eq!(stored_count, "2");
1861 }
1862
1863 #[serial(cooldown_state)]
1864 #[tokio::test]
1865 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1866 super::reset_global_state().await;
1867 let store = test_store().await;
1868 *cooldown_store().lock().await = Some(store.clone());
1869
1870 let agent = "testagent_persistent_success";
1871 let model = "testmodel_persistent_success";
1872 let key = format!("{agent}:{model}");
1873
1874 record_persistent_model_failure(agent, model).await;
1876 record_persistent_model_failure(agent, model).await;
1877
1878 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1879 let stored_count = store
1880 .kv_get(&count_key)
1881 .await
1882 .expect("kv read should succeed")
1883 .expect("failure count should be written");
1884 assert_eq!(stored_count, "2");
1885
1886 record_agent_success(agent, model).await;
1890
1891 let stored_count_after_success = store
1892 .kv_get(&count_key)
1893 .await
1894 .expect("kv read should succeed")
1895 .expect("persistent failure count should still exist");
1896 assert_eq!(
1897 stored_count_after_success, "2",
1898 "persistent failure count must survive an intervening success"
1899 );
1900
1901 record_persistent_model_failure(agent, model).await;
1903 let now = chrono::Utc::now().timestamp();
1904 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1905 let third_remaining = third_until.saturating_sub(now);
1906 assert!(
1907 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1908 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1909 );
1910 }
1911
1912 #[serial(cooldown_state)]
1913 #[tokio::test]
1914 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1915 super::reset_global_state().await;
1916 let store = test_store().await;
1917 *cooldown_store().lock().await = Some(store.clone());
1918
1919 let agent = "testagent_persistent_capped";
1920 let model = "testmodel_persistent_capped";
1921 let key = format!("{agent}:{model}");
1922
1923 for _ in 0..9 {
1924 record_persistent_model_failure(agent, model).await;
1925 }
1926
1927 let remaining = cooldown_until(&key)
1928 .expect("cooldown should exist")
1929 .saturating_sub(chrono::Utc::now().timestamp());
1930 assert!(
1931 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1932 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1933 );
1934 assert!(
1935 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1936 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1937 );
1938 }
1939
1940 #[serial(cooldown_state)]
1941 #[test]
1942 fn parse_retry_at_codex_format() {
1943 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.";
1944 let ts = parse_retry_at(msg);
1945 assert!(ts.is_some(), "should parse codex retry-at date");
1946
1947 use chrono::TimeZone;
1950 let naive = chrono::NaiveDateTime::new(
1951 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1952 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1953 );
1954 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1955 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1956 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1957 chrono::LocalResult::None => naive.and_utc(),
1958 };
1959 assert_eq!(
1960 ts.unwrap(),
1961 expected_utc.timestamp(),
1962 "timestamp should reflect local-timezone interpretation of the vendor date"
1963 );
1964
1965 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1967 use chrono::Datelike;
1968 assert_eq!(dt.month(), 3);
1969 assert_eq!(dt.day(), 26);
1970 }
1971
1972 #[serial(cooldown_state)]
1973 #[test]
1974 fn parse_retry_at_hour_only_and_month_day() {
1975 let now = chrono::Utc::now().timestamp();
1976 for msg in [
1977 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1978 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1979 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1980 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1981 ] {
1982 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1983 assert!(ts > now, "{msg}");
1984 }
1985 }
1986
1987 #[serial(cooldown_state)]
1988 #[test]
1989 fn parse_retry_at_weekly_resets_without_year() {
1990 use chrono::Datelike;
1991 let t = chrono::Local::now() + chrono::Duration::days(2);
1992 let msg = format!(
1993 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1994 t.format("%b"),
1995 t.day()
1996 );
1997 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1998 let remaining = ts - chrono::Utc::now().timestamp();
1999 assert!(
2000 remaining > 3600 && remaining < 3 * 86400,
2001 "expected ~2d retry-at, got {remaining}s"
2002 );
2003 }
2004
2005 #[serial(cooldown_state)]
2006 #[test]
2007 fn parse_retry_at_no_date() {
2008 assert!(parse_retry_at("generic rate limit error").is_none());
2009 assert!(parse_retry_at("").is_none());
2010 }
2011
2012 #[serial(cooldown_state)]
2013 #[test]
2014 fn parse_retry_at_glm_reset_format_with_seconds() {
2015 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
2016 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
2017
2018 use chrono::TimeZone;
2019 let naive = chrono::NaiveDateTime::new(
2020 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
2021 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
2022 );
2023 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
2024 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
2025 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
2026 chrono::LocalResult::None => naive.and_utc(),
2027 };
2028
2029 assert_eq!(
2030 ts,
2031 expected_utc.timestamp(),
2032 "timestamp should reflect local-timezone interpretation of vendor reset date"
2033 );
2034 }
2035
2036 #[serial(cooldown_state)]
2037 #[test]
2038 fn parse_retry_at_billing_cycle_without_date_returns_none() {
2039 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.";
2042 assert!(
2043 parse_retry_at(msg).is_none(),
2044 "billing cycle without date should return None"
2045 );
2046 }
2047
2048 #[serial(cooldown_state)]
2049 #[test]
2050 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
2051 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.";
2052 assert_eq!(
2053 parse_relative_usage_window(msg),
2054 Some(5 * 3600),
2055 "should extract the 5-hour window as seconds"
2056 );
2057 }
2058
2059 #[serial(cooldown_state)]
2060 #[test]
2061 fn parse_relative_usage_window_handles_common_durations() {
2062 assert_eq!(
2063 parse_relative_usage_window("You've reached your hourly usage limit."),
2064 Some(3600)
2065 );
2066 assert_eq!(
2067 parse_relative_usage_window("You've reached your daily usage limit."),
2068 Some(86400)
2069 );
2070 assert_eq!(
2071 parse_relative_usage_window("You've reached your weekly limit."),
2072 Some(604800)
2073 );
2074 assert_eq!(
2075 parse_relative_usage_window("You've reached your 30-minute usage limit."),
2076 Some(1800)
2077 );
2078 assert_eq!(
2079 parse_relative_usage_window("You've reached your 2-day quota."),
2080 Some(172800)
2081 );
2082 }
2083
2084 #[serial(cooldown_state)]
2085 #[test]
2086 fn parse_relative_usage_window_returns_none_without_usage_context() {
2087 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
2090 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
2091 assert!(parse_relative_usage_window("").is_none());
2092 }
2093
2094 #[serial(cooldown_state)]
2095 #[tokio::test]
2096 async fn record_rate_limit_uses_model_specific_window_cooldown() {
2097 super::reset_global_state().await;
2098 let store = test_store().await;
2099 *cooldown_store().lock().await = Some(store.clone());
2100
2101 let agent = "test_record_rl_window_agent";
2102 let model = "test_record_rl_window_model";
2103 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
2104
2105 record_rate_limit(agent, Some(model), msg).await;
2106
2107 assert!(
2108 is_model_in_cooldown(agent, model),
2109 "model should be cooled for the relative usage window"
2110 );
2111 assert!(
2112 !is_agent_in_cooldown(agent),
2113 "agent should not be cooled for a model-scoped quota"
2114 );
2115
2116 let key = format!("{agent}:{model}");
2117 let remaining = cooldown_until(&key)
2118 .expect("model cooldown should exist")
2119 .saturating_sub(chrono::Utc::now().timestamp());
2120 assert!(
2121 remaining >= 5 * 3600 - 5,
2122 "model cooldown should be at least ~5 hours, got {remaining}s"
2123 );
2124 assert_eq!(
2125 cooldown_reason(&key).as_deref(),
2126 Some("rate_limit_usage_window"),
2127 "relative usage-window cooldown must be labeled as a rate limit, not silence"
2128 );
2129 }
2130
2131 #[serial(cooldown_state)]
2132 #[tokio::test]
2133 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
2134 super::reset_global_state().await;
2135 let store = test_store().await;
2136 *cooldown_store().lock().await = Some(store.clone());
2137
2138 let agent = "test_record_rl_fallback_agent";
2139 let model = "test_record_rl_fallback_model";
2140 let msg = "429 Too Many Requests";
2141
2142 record_rate_limit(agent, Some(model), msg).await;
2143
2144 assert!(
2145 is_agent_in_cooldown(agent),
2146 "agent should be cooled when no relative window is detected"
2147 );
2148 assert!(
2149 is_model_in_cooldown(agent, model),
2150 "model should be cooled when no relative window is detected"
2151 );
2152 }
2153
2154 #[serial(cooldown_state)]
2155 #[tokio::test]
2156 async fn agent_cooldown_persists_to_kv() {
2157 super::reset_global_state().await;
2158 let agent = "test_agent_persist_check";
2160 record_agent_failure_with_message(agent, "").await;
2161 assert!(is_agent_in_cooldown(agent));
2162 }
2163
2164 #[serial(cooldown_state)]
2165 #[tokio::test]
2166 async fn cooldown_never_shortens() {
2167 super::reset_global_state().await;
2168 let agent = "test_agent_no_shorten";
2169 set_agent_cooldown(agent, 24 * 60 * 60).await;
2170 let initial = {
2171 let map = cooldowns().lock().unwrap();
2172 map.get(agent)
2173 .expect("cooldown entry should exist")
2174 .cooldown_until
2175 };
2176
2177 record_agent_failure_with_message(agent, "rate limit").await;
2179 let after = {
2180 let map = cooldowns().lock().unwrap();
2181 map.get(agent)
2182 .expect("cooldown entry should exist")
2183 .cooldown_until
2184 };
2185 assert_eq!(
2186 initial, after,
2187 "cooldown should not be shortened by later failures"
2188 );
2189 }
2190
2191 #[serial(cooldown_state)]
2192 #[tokio::test]
2193 async fn silence_agent_cooldown_is_short_lived() {
2194 super::reset_global_state().await;
2195 let agent = "test_silence_agent_cd";
2196 assert!(!is_agent_in_cooldown(agent));
2197
2198 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2199 assert!(is_agent_in_cooldown(agent));
2200
2201 let map = cooldowns().lock().unwrap();
2203 let entry = map.get(agent).expect("should have cooldown entry");
2204 let now = chrono::Utc::now().timestamp();
2205 let remaining = entry.cooldown_until - now;
2206 assert!(
2207 remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
2208 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2209 );
2210 assert!(remaining > 0, "cooldown should still be active");
2211 }
2212
2213 #[serial(cooldown_state)]
2214 #[tokio::test]
2215 async fn record_silence_detection_applies_extended_cooldown() {
2216 super::reset_global_state().await;
2217 let store = test_store().await;
2218 *cooldown_store().lock().await = Some(store.clone());
2219
2220 let agent = "test_silence_count_agent";
2221 let model = "test_silence_count_model";
2222 assert!(!is_model_in_cooldown(agent, model));
2223
2224 for _ in 0..SILENCE_COUNT_THRESHOLD {
2225 let result = record_silence_detection(agent, model).await;
2226 assert!(result.is_some());
2227 }
2228
2229 assert!(is_model_in_cooldown(agent, model));
2230 let key = format!("{agent}:{model}");
2231 let remaining = {
2232 let map = cooldowns().lock().unwrap();
2233 let entry = map.get(&key).expect("cooldown entry should exist");
2234 entry.cooldown_until - chrono::Utc::now().timestamp()
2235 };
2236 assert!(
2237 remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
2238 "extended cooldown should be applied, got {remaining}s"
2239 );
2240
2241 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2242 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2243 assert_eq!(stored, "[]");
2244 }
2245
2246 #[serial(cooldown_state)]
2247 #[test]
2248 fn detect_credit_exhaustion_out_of_credits() {
2249 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2250 assert_eq!(
2251 detect_credit_exhaustion(msg),
2252 Some(CreditExhaustionReason::OutOfCredits)
2253 );
2254 }
2255
2256 #[serial(cooldown_state)]
2257 #[test]
2258 fn detect_credit_exhaustion_out_of_credits_variations() {
2259 assert_eq!(
2260 detect_credit_exhaustion("out of credits"),
2261 Some(CreditExhaustionReason::OutOfCredits)
2262 );
2263 assert_eq!(
2264 detect_credit_exhaustion("outofcredits"),
2265 Some(CreditExhaustionReason::OutOfCredits)
2266 );
2267 assert_eq!(
2268 detect_credit_exhaustion("no credits remaining"),
2269 Some(CreditExhaustionReason::OutOfCredits)
2270 );
2271 assert_eq!(
2272 detect_credit_exhaustion("insufficient funds"),
2273 Some(CreditExhaustionReason::OutOfCredits)
2274 );
2275 }
2276
2277 #[serial(cooldown_state)]
2278 #[test]
2279 fn detect_credit_exhaustion_org_level_disabled() {
2280 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2281 assert_eq!(
2282 detect_credit_exhaustion(msg),
2283 Some(CreditExhaustionReason::OrgLevelDisabled)
2284 );
2285 }
2286
2287 #[serial(cooldown_state)]
2288 #[test]
2289 fn detect_credit_exhaustion_org_level_disabled_variations() {
2290 assert_eq!(
2291 detect_credit_exhaustion("organization has been disabled"),
2292 Some(CreditExhaustionReason::OrgLevelDisabled)
2293 );
2294 assert_eq!(
2295 detect_credit_exhaustion("org disabled"),
2296 Some(CreditExhaustionReason::OrgLevelDisabled)
2297 );
2298 assert_eq!(
2299 detect_credit_exhaustion("billing quota exceeded"),
2300 Some(CreditExhaustionReason::OrgLevelDisabled)
2301 );
2302 assert_eq!(
2303 detect_credit_exhaustion("credit balance too low"),
2304 Some(CreditExhaustionReason::OrgLevelDisabled)
2305 );
2306 }
2307
2308 #[serial(cooldown_state)]
2309 #[test]
2310 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2311 let msg = "overageDisabledReason: out_of_credits";
2312 assert_eq!(
2313 detect_credit_exhaustion(msg),
2314 Some(CreditExhaustionReason::OutOfCredits)
2315 );
2316
2317 let msg = "overageDisabledReason: org_level_disabled";
2318 assert_eq!(
2319 detect_credit_exhaustion(msg),
2320 Some(CreditExhaustionReason::OrgLevelDisabled)
2321 );
2322 }
2323
2324 #[serial(cooldown_state)]
2325 #[test]
2326 fn detect_credit_exhaustion_glm_insufficient_balance() {
2327 assert_eq!(
2329 detect_credit_exhaustion(
2330 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2331 ),
2332 Some(CreditExhaustionReason::OutOfCredits)
2333 );
2334 assert_eq!(
2335 detect_credit_exhaustion("Insufficient balance"),
2336 Some(CreditExhaustionReason::OutOfCredits)
2337 );
2338 assert_eq!(
2339 detect_credit_exhaustion("please recharge your account"),
2340 Some(CreditExhaustionReason::OutOfCredits)
2341 );
2342 assert_eq!(
2343 detect_credit_exhaustion("no resource package available"),
2344 Some(CreditExhaustionReason::OutOfCredits)
2345 );
2346 }
2347
2348 #[serial(cooldown_state)]
2349 #[test]
2350 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2351 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2352 assert!(detect_credit_exhaustion("too many requests").is_none());
2353 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2354 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2355 assert!(detect_credit_exhaustion(
2356 "You've hit your usage limit, try again in a few minutes."
2357 )
2358 .is_none());
2359 }
2360
2361 #[serial(cooldown_state)]
2362 #[tokio::test]
2363 async fn record_credit_exhaustion_applies_agent_cooldown() {
2364 super::reset_global_state().await;
2365 let agent = "test_credit_exhaust_agent";
2366 assert!(!is_agent_in_cooldown(agent));
2367
2368 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2369 assert!(is_agent_in_cooldown(agent));
2370
2371 let remaining = {
2372 let map = cooldowns().lock().unwrap();
2373 let entry = map.get(agent).expect("cooldown entry should exist");
2374 entry.cooldown_until - chrono::Utc::now().timestamp()
2375 };
2376 assert!(
2378 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2379 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2380 );
2381 assert!(
2382 remaining <= CREDIT_BACKOFF_MAX_SECS,
2383 "cooldown should not exceed cap of 8h, got {remaining}s"
2384 );
2385 }
2386
2387 #[serial(cooldown_state)]
2388 #[tokio::test]
2389 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2390 super::reset_global_state().await;
2391 let agent = "test_org_disabled_agent";
2392 assert!(!is_agent_in_cooldown(agent));
2393
2394 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2395 assert!(is_agent_in_cooldown(agent));
2396
2397 let remaining = {
2398 let map = cooldowns().lock().unwrap();
2399 let entry = map.get(agent).expect("cooldown entry should exist");
2400 entry.cooldown_until - chrono::Utc::now().timestamp()
2401 };
2402 assert!(
2404 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2405 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2406 );
2407 assert!(
2408 remaining <= CREDIT_BACKOFF_MAX_SECS,
2409 "cooldown should not exceed cap of 8h, got {remaining}s"
2410 );
2411 }
2412
2413 #[serial(cooldown_state)]
2414 #[tokio::test]
2415 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2416 super::reset_global_state().await;
2417 let store = test_store().await;
2418 *cooldown_store().lock().await = Some(store.clone());
2419
2420 let agent = "test_billing_cycle_agent";
2421 assert!(!is_agent_in_cooldown(agent));
2422
2423 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2424 assert!(is_agent_in_cooldown(agent));
2425
2426 let remaining = {
2427 let map = cooldowns().lock().unwrap();
2428 let entry = map.get(agent).expect("cooldown entry should exist");
2429 entry.cooldown_until - chrono::Utc::now().timestamp()
2430 };
2431 assert!(
2433 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2434 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2435 );
2436 assert!(
2437 remaining <= BILLING_CYCLE_MAX_SECS,
2438 "cooldown should not exceed cap of 7 days, got {remaining}s"
2439 );
2440 }
2441
2442 #[serial(cooldown_state)]
2443 #[tokio::test]
2444 async fn record_credit_exhaustion_billing_cycle_escalates() {
2445 super::reset_global_state().await;
2446 let store = test_store().await;
2447 *cooldown_store().lock().await = Some(store.clone());
2448
2449 let agent = "test_billing_cycle_escalation";
2450
2451 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2453 let remaining_1 = {
2454 let map = cooldowns().lock().unwrap();
2455 let entry = map.get(agent).expect("cooldown entry should exist");
2456 entry.cooldown_until - chrono::Utc::now().timestamp()
2457 };
2458 assert!(
2459 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2460 .contains(&remaining_1),
2461 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2462 );
2463
2464 {
2466 let mut map = cooldowns().lock().unwrap();
2467 map.remove(agent);
2468 }
2469
2470 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2472 let remaining_2 = {
2473 let map = cooldowns().lock().unwrap();
2474 let entry = map.get(agent).expect("cooldown entry should exist");
2475 entry.cooldown_until - chrono::Utc::now().timestamp()
2476 };
2477 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2479 remaining_2 >= expected_2 - 5,
2480 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2481 );
2482
2483 {
2485 let mut map = cooldowns().lock().unwrap();
2486 map.remove(agent);
2487 }
2488
2489 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2491 let remaining_3 = {
2492 let map = cooldowns().lock().unwrap();
2493 let entry = map.get(agent).expect("cooldown entry should exist");
2494 entry.cooldown_until - chrono::Utc::now().timestamp()
2495 };
2496 assert!(
2497 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2498 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2499 );
2500 assert!(
2501 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2502 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2503 );
2504 }
2505
2506 #[serial(cooldown_state)]
2507 #[tokio::test]
2508 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2509 super::reset_global_state().await;
2510 let store = test_store().await;
2511 *cooldown_store().lock().await = Some(store.clone());
2512
2513 let agent = "test_billing_cycle_agent_capped";
2514
2515 for _ in 0..5 {
2516 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2517 }
2518
2519 let remaining = {
2520 let map = cooldowns().lock().unwrap();
2521 let entry = map.get(agent).expect("cooldown entry should exist");
2522 entry.cooldown_until - chrono::Utc::now().timestamp()
2523 };
2524 assert!(
2525 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2526 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2527 );
2528 assert!(
2529 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2530 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2531 );
2532 }
2533
2534 #[serial(cooldown_state)]
2535 #[test]
2536 fn credit_exhaustion_backoff_constants() {
2537 assert_eq!(
2539 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2540 CREDIT_BACKOFF_BASE_SECS
2541 );
2542 assert_eq!(
2544 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2545 ORG_BACKOFF_BASE_SECS
2546 );
2547 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2549 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2550 }
2551
2552 #[serial(cooldown_state)]
2553 #[test]
2554 fn compute_backoff_long_cap_without_extended_stays_capped() {
2555 assert_eq!(
2556 compute_backoff(
2557 3,
2558 BILLING_CYCLE_COOLDOWN_SECS,
2559 BILLING_CYCLE_MAX_SECS,
2560 false
2561 ),
2562 BILLING_CYCLE_MAX_SECS,
2563 "7-day capped paths should stop at the documented max"
2564 );
2565 assert_eq!(
2566 compute_backoff(
2567 9,
2568 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2569 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2570 false,
2571 ),
2572 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2573 "persistent model failures should not extend past 7 days"
2574 );
2575 }
2576
2577 #[serial(cooldown_state)]
2580 #[test]
2581 fn mark_and_check_agent_degraded() {
2582 let agent = "test_degraded_mark";
2583 assert!(!is_agent_degraded(agent));
2584 mark_agent_degraded(agent);
2585 assert!(is_agent_degraded(agent));
2586 clear_agent_degraded(agent);
2587 assert!(!is_agent_degraded(agent));
2588 }
2589
2590 #[serial(cooldown_state)]
2591 #[tokio::test]
2592 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2593 super::reset_global_state().await;
2594 let store = test_store().await;
2595 let agent = "test_sync_external_clear";
2596 let future = chrono::Utc::now().timestamp() + 600;
2597
2598 init_cooldown_store(store.clone()).await;
2600 set_cooldown_async(agent, future, "test").await;
2601 assert!(is_agent_in_cooldown(agent));
2602
2603 store
2607 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2608 .await
2609 .expect("kv_set tombstone");
2610 assert!(
2611 is_agent_in_cooldown(agent),
2612 "in-memory cooldown should still be present before sync"
2613 );
2614
2615 sync_from_kv(&store).await;
2616
2617 assert!(
2618 !is_agent_in_cooldown(agent),
2619 "sync_from_kv must drop in-memory entries the KV tombstoned"
2620 );
2621 }
2622
2623 #[serial(cooldown_state)]
2624 #[tokio::test]
2625 async fn sync_from_kv_closes_github_circuit_when_row_cleared() {
2626 super::reset_global_state().await;
2627 let store = test_store().await;
2628 {
2629 *github_5xx_circuit_open_lock() = Some(chrono::Utc::now().timestamp() + 600);
2630 github_5xx_timestamps_lock().push(chrono::Utc::now().timestamp());
2631 }
2632 assert!(is_github_circuit_open());
2633
2634 store
2635 .kv_set("cooldown:github:5xx", "0")
2636 .await
2637 .expect("kv_set tombstone");
2638 sync_from_kv(&store).await;
2639
2640 assert!(!is_github_circuit_open());
2641 assert!(github_5xx_timestamps_lock().is_empty());
2642 }
2643
2644 #[serial(cooldown_state)]
2645 #[tokio::test]
2646 async fn sync_from_kv_keeps_5xx_window_when_expired_row_and_breaker_closed() {
2647 super::reset_global_state().await;
2648 let store = test_store().await;
2649 let now = chrono::Utc::now().timestamp();
2650 {
2651 github_5xx_timestamps_lock().push(now);
2652 github_5xx_timestamps_lock().push(now);
2653 }
2654
2655 store
2656 .kv_set("cooldown:github:5xx", &(now - 600).to_string())
2657 .await
2658 .expect("kv_set past");
2659 sync_from_kv(&store).await;
2660
2661 assert!(!is_github_circuit_open());
2662 assert_eq!(github_5xx_timestamps_lock().len(), 2);
2663
2664 github_5xx_timestamps_lock().clear();
2665 }
2666
2667 #[serial(cooldown_state)]
2668 #[tokio::test]
2669 async fn sync_from_kv_keeps_github_circuit_open_for_future_row() {
2670 super::reset_global_state().await;
2671 let store = test_store().await;
2672 let future = chrono::Utc::now().timestamp() + 600;
2673 *github_5xx_circuit_open_lock() = Some(future);
2674
2675 store
2676 .kv_set("cooldown:github:5xx", &future.to_string())
2677 .await
2678 .expect("kv_set future");
2679 sync_from_kv(&store).await;
2680
2681 assert!(is_github_circuit_open());
2682
2683 *github_5xx_circuit_open_lock() = None;
2684 }
2685
2686 #[serial(cooldown_state)]
2687 #[tokio::test]
2688 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2689 super::reset_global_state().await;
2693 let store = test_store().await;
2694 let agent = "test_sync_race_window";
2695 let future = chrono::Utc::now().timestamp() + 600;
2696
2697 set_cooldown_in_memory(agent, future, "raced");
2699 assert!(is_agent_in_cooldown(agent));
2700
2701 sync_from_kv(&store).await;
2702
2703 assert!(
2704 is_agent_in_cooldown(agent),
2705 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2706 );
2707 }
2708
2709 #[serial(cooldown_state)]
2710 #[tokio::test]
2711 async fn sync_from_kv_adopts_kv_timestamp() {
2712 super::reset_global_state().await;
2715 let store = test_store().await;
2716 let agent = "test_sync_adopt_kv";
2717 let original = chrono::Utc::now().timestamp() + 600;
2718 let updated = chrono::Utc::now().timestamp() + 1200;
2719
2720 init_cooldown_store(store.clone()).await;
2721 set_cooldown_async(agent, original, "initial").await;
2722
2723 store
2724 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2725 .await
2726 .expect("kv_set updated");
2727
2728 sync_from_kv(&store).await;
2729
2730 let until = cooldown_until(agent).expect("cooldown still active");
2731 assert_eq!(
2732 until, updated,
2733 "sync_from_kv must adopt KV's authoritative timestamp"
2734 );
2735 }
2736
2737 #[serial(cooldown_state)]
2738 #[tokio::test]
2739 async fn refresh_degraded_agents_marks_cooled_agent() {
2740 super::reset_global_state().await;
2741 let store = test_store().await;
2742 let agent = "test_refresh_cooled";
2743
2744 record_agent_failure_with_message(agent, "").await;
2746 assert!(is_agent_in_cooldown(agent));
2747
2748 let agents = vec![agent.to_string()];
2750 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2751
2752 assert!(
2753 is_agent_degraded(agent),
2754 "agent in cooldown 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_no_models_agent() {
2764 super::reset_global_state().await;
2765 let store = test_store().await;
2766 let agent = "test_refresh_no_models";
2767
2768 let agents = vec![agent.to_string()];
2770 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2771
2772 assert!(
2773 is_agent_degraded(agent),
2774 "agent with no available models should be marked degraded"
2775 );
2776
2777 clear_agent_degraded(agent);
2779 }
2780
2781 #[serial(cooldown_state)]
2782 #[tokio::test]
2783 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2784 super::reset_global_state().await;
2785 let store = test_store().await;
2786 let agent = "test_refresh_rate_limited";
2787
2788 for _ in 0..4 {
2790 store.record_rate_limit(agent, "rate", None).await.unwrap();
2791 }
2792
2793 let agents = vec![agent.to_string()];
2794 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2796
2797 assert!(
2798 is_agent_degraded(agent),
2799 "agent with rate limit count >= threshold should be marked degraded"
2800 );
2801
2802 clear_agent_degraded(agent);
2804 }
2805
2806 #[serial(cooldown_state)]
2807 #[tokio::test]
2808 async fn refresh_degraded_agents_clears_healthy_agent() {
2809 super::reset_global_state().await;
2810 let store = test_store().await;
2811 let agent = "test_refresh_healthy";
2812
2813 mark_agent_degraded(agent);
2815 assert!(is_agent_degraded(agent));
2816
2817 let agents = vec![agent.to_string()];
2819 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2820
2821 assert!(
2822 !is_agent_degraded(agent),
2823 "healthy agent should have degraded flag cleared"
2824 );
2825 }
2826
2827 #[serial(cooldown_state)]
2828 #[tokio::test]
2829 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2830 super::reset_global_state().await;
2831 let store = test_store().await;
2832 let agent = "test_refresh_expired_cooldown";
2833
2834 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2838 let mut map = cooldowns().lock().unwrap();
2839 map.insert(
2840 agent.to_string(),
2841 CooldownEntry {
2842 cooldown_until: expired_ts,
2843 reason: "expired_test".to_string(),
2844 },
2845 );
2846 }
2847
2848 assert!(
2850 !is_agent_in_cooldown(agent),
2851 "expired cooldown should not be active"
2852 );
2853
2854 mark_agent_degraded(agent);
2856 assert!(is_agent_degraded(agent));
2857
2858 let agents = vec![agent.to_string()];
2860 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2861
2862 assert!(
2863 !is_agent_degraded(agent),
2864 "agent with only expired cooldown should NOT be marked degraded"
2865 );
2866
2867 clear_agent_degraded(agent);
2869 {
2870 let mut map = cooldowns().lock().unwrap();
2871 map.remove(agent);
2872 }
2873 }
2874
2875 #[serial(cooldown_state)]
2876 #[test]
2877 fn compute_backoff_grows_exponentially() {
2878 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2880 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2881 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2883 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2885 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2887 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2889 }
2890
2891 #[serial(cooldown_state)]
2892 #[test]
2893 fn compute_backoff_extended_tier_1_applies_24h() {
2894 assert_eq!(
2900 compute_backoff(5, 300, 14400, true),
2901 14400,
2902 "first at max, not extended"
2903 );
2904 assert_eq!(
2905 compute_backoff(6, 300, 14400, true),
2906 14400,
2907 "times_at_max=1, not extended"
2908 );
2909 assert_eq!(
2910 compute_backoff(7, 300, 14400, true),
2911 86400,
2912 "times_at_max=2, extended 6x"
2913 );
2914 assert_eq!(
2915 compute_backoff(8, 300, 14400, true),
2916 86400,
2917 "times_at_max=3, still extended 6x"
2918 );
2919 }
2920
2921 #[serial(cooldown_state)]
2922 #[test]
2923 fn compute_backoff_extended_tier_2_applies_48h() {
2924 assert_eq!(
2927 compute_backoff(9, 300, 14400, true),
2928 86400,
2929 "times_at_max=4, still tier 1"
2930 );
2931 assert_eq!(
2932 compute_backoff(10, 300, 14400, true),
2933 172800,
2934 "times_at_max=5, extended 12x"
2935 );
2936 assert_eq!(
2937 compute_backoff(15, 300, 14400, true),
2938 172800,
2939 "times_at_max=10, still tier 2"
2940 );
2941 assert_eq!(
2942 compute_backoff(20, 300, 14400, true),
2943 172800,
2944 "times_at_max=15, still tier 2"
2945 );
2946 }
2947
2948 #[serial(cooldown_state)]
2949 #[test]
2950 fn compute_backoff_extended_tier_with_different_base_max() {
2951 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2958 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2959 assert_eq!(
2960 compute_backoff(7, 60, 3600, true),
2961 21600,
2962 "times_at_max=2, extended 6x"
2963 );
2964 }
2965
2966 #[serial(cooldown_state)]
2967 #[test]
2968 fn compute_backoff_extended_tier_persistent_fails_example() {
2969 assert_eq!(
2973 compute_backoff(21, 300, 14400, true),
2974 172800,
2975 "21 failures should get 48h extended cooldown"
2976 );
2977
2978 assert_eq!(
2980 compute_backoff(36, 300, 14400, true),
2981 172800,
2982 "36 failures should also get 48h extended cooldown"
2983 );
2984 }
2985
2986 #[serial(cooldown_state)]
2987 #[test]
2988 fn compute_backoff_extended_tier_zero_base() {
2989 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2991 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2992 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2993 }
2994
2995 #[serial(cooldown_state)]
2996 #[test]
2997 fn detect_credit_exhaustion_billing_cycle() {
2998 assert_eq!(
2999 detect_credit_exhaustion("billing cycle"),
3000 Some(CreditExhaustionReason::BillingCycleExhausted)
3001 );
3002 assert_eq!(
3003 detect_credit_exhaustion("quota refreshed next cycle"),
3004 Some(CreditExhaustionReason::BillingCycleExhausted)
3005 );
3006 assert_eq!(
3007 detect_credit_exhaustion(
3008 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
3009 ),
3010 Some(CreditExhaustionReason::BillingCycleExhausted)
3011 );
3012 assert_eq!(
3013 detect_credit_exhaustion("monthly quota exceeded"),
3014 Some(CreditExhaustionReason::BillingCycleExhausted)
3015 );
3016 assert_eq!(
3017 detect_credit_exhaustion(
3018 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
3019 ),
3020 Some(CreditExhaustionReason::BillingCycleExhausted)
3021 );
3022 }
3023
3024 #[serial(cooldown_state)]
3025 #[test]
3026 fn detect_credit_exhaustion_billing_cycle_minimax() {
3027 assert_eq!(
3028 detect_credit_exhaustion(
3029 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
3030 ),
3031 Some(CreditExhaustionReason::BillingCycleExhausted)
3032 );
3033 assert_eq!(
3034 detect_credit_exhaustion("Upgrade your Token Plan"),
3035 Some(CreditExhaustionReason::BillingCycleExhausted)
3036 );
3037 assert_eq!(
3038 detect_credit_exhaustion("Token Plan usage limit reached"),
3039 Some(CreditExhaustionReason::BillingCycleExhausted)
3040 );
3041 assert_eq!(
3042 detect_credit_exhaustion("purchase Credits for more usage"),
3043 Some(CreditExhaustionReason::BillingCycleExhausted)
3044 );
3045 assert_eq!(
3046 detect_credit_exhaustion(
3047 "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"
3048 ),
3049 Some(CreditExhaustionReason::BillingCycleExhausted)
3050 );
3051 }
3052
3053 #[serial(cooldown_state)]
3054 #[test]
3055 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
3056 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
3061 this billing cycle. Your quota will be refreshed soon.";
3062 assert_eq!(
3063 detect_credit_exhaustion_outcome(msg),
3064 Some("billing_cycle_exhausted")
3065 );
3066 }
3067
3068 #[serial(cooldown_state)]
3069 #[test]
3070 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
3071 assert_eq!(
3072 detect_credit_exhaustion_outcome("out of credits"),
3073 Some("credit_exhausted")
3074 );
3075 assert_eq!(
3076 detect_credit_exhaustion_outcome("organization has been disabled"),
3077 Some("org_level_disabled")
3078 );
3079 assert_eq!(
3080 detect_credit_exhaustion_outcome("monthly quota exceeded"),
3081 Some("billing_cycle_exhausted")
3082 );
3083 }
3084
3085 #[test]
3086 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
3087 assert_eq!(
3088 detect_credit_exhaustion_outcome("429 Too Many Requests"),
3089 None
3090 );
3091 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
3092 }
3093
3094 #[serial(cooldown_state)]
3095 #[test]
3096 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
3097 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
3100 assert_eq!(
3101 detect_credit_exhaustion(msg),
3102 Some(CreditExhaustionReason::OrgLevelDisabled)
3103 );
3104 }
3105
3106 #[serial(cooldown_state)]
3107 #[tokio::test]
3108 async fn refresh_degraded_agents_below_threshold_not_degraded() {
3109 super::reset_global_state().await;
3110 let store = test_store().await;
3111 let agent = "test_refresh_below_threshold";
3112
3113 for _ in 0..2 {
3115 store.record_rate_limit(agent, "rate", None).await.unwrap();
3116 }
3117
3118 let agents = vec![agent.to_string()];
3119 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3121
3122 assert!(
3123 !is_agent_degraded(agent),
3124 "agent with rate limit count < threshold should not be degraded"
3125 );
3126 }
3127
3128 #[serial(cooldown_state)]
3129 #[tokio::test]
3130 async fn refresh_degraded_agents_ignores_error_limit_type() {
3131 super::reset_global_state().await;
3132 let store = test_store().await;
3133 let agent = "test_refresh_error_rows";
3134
3135 for _ in 0..3 {
3136 store.record_rate_limit(agent, "error", 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 error rows in 6 hours must not mark an agent degraded"
3145 );
3146 }
3147
3148 #[serial(cooldown_state)]
3149 #[tokio::test]
3150 async fn refresh_degraded_agents_marks_three_rate_rows() {
3151 super::reset_global_state().await;
3152 let store = test_store().await;
3153 let agent = "test_refresh_three_rate_rows";
3154
3155 for _ in 0..3 {
3156 store.record_rate_limit(agent, "rate", None).await.unwrap();
3157 }
3158
3159 let agents = vec![agent.to_string()];
3160 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3161
3162 assert!(
3163 is_agent_degraded(agent),
3164 "three rate rows in 6 hours must mark an agent degraded"
3165 );
3166 assert_eq!(
3167 agent_degraded_reason(agent).as_deref(),
3168 Some("rate limit threshold exceeded")
3169 );
3170 let listed = list_degraded_agents();
3171 assert!(
3172 listed
3173 .iter()
3174 .any(|(a, r)| a == agent && r == "rate limit threshold exceeded"),
3175 "orch cooldown list must show the degraded reason, got {listed:?}"
3176 );
3177
3178 clear_agent_degraded(agent);
3179 }
3180
3181 #[serial(cooldown_state)]
3182 #[tokio::test]
3183 async fn refresh_degraded_agents_logs_reason_change() {
3184 super::reset_global_state().await;
3185 let store = test_store().await;
3186 let agent = "test_refresh_reason_change";
3187 *cooldown_store().lock().await = Some(store.clone());
3188
3189 set_agent_cooldown(agent, 600).await;
3190 let agents = vec![agent.to_string()];
3191 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3192 assert_eq!(
3193 agent_degraded_reason(agent).as_deref(),
3194 Some("agent in cooldown")
3195 );
3196
3197 clear_cooldown(agent, &store).await;
3198 let after_clear = sqlite_datetime_offset(60);
3199 for _ in 0..3 {
3200 store
3201 .record_rate_limit_at(agent, "rate", None, &after_clear)
3202 .await
3203 .unwrap();
3204 }
3205 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3206 assert_eq!(
3207 agent_degraded_reason(agent).as_deref(),
3208 Some("rate limit threshold exceeded"),
3209 "degraded reason must update when the hold changes"
3210 );
3211
3212 clear_agent_degraded(agent);
3213 }
3214
3215 #[serial(cooldown_state)]
3216 #[tokio::test]
3217 async fn clear_un_degrades_on_next_health_check() {
3218 super::reset_global_state().await;
3219 let store = test_store().await;
3220 let agent = "test_clear_un_degrades";
3221 let before = sqlite_datetime_offset(-30);
3222 let after = sqlite_datetime_offset(60);
3223
3224 for _ in 0..3 {
3225 store
3226 .record_rate_limit_at(agent, "rate", None, &before)
3227 .await
3228 .unwrap();
3229 }
3230
3231 let agents = vec![agent.to_string()];
3232 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3233 assert!(
3234 is_agent_degraded(agent),
3235 "three rate rows must mark an agent degraded"
3236 );
3237
3238 clear_cooldown(agent, &store).await;
3239 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3240 assert!(
3241 !is_agent_degraded(agent),
3242 "clear must un-degrade on the next health check"
3243 );
3244
3245 store
3246 .record_rate_limit_at(agent, "rate", None, &after)
3247 .await
3248 .unwrap();
3249 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3250 assert!(
3251 !is_agent_degraded(agent),
3252 "one rate row after the mark must not degrade again"
3253 );
3254
3255 for _ in 0..2 {
3256 store
3257 .record_rate_limit_at(agent, "rate", None, &after)
3258 .await
3259 .unwrap();
3260 }
3261 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3262 assert!(
3263 is_agent_degraded(agent),
3264 "three rate rows after the mark must degrade again"
3265 );
3266
3267 clear_agent_degraded(agent);
3268 }
3269
3270 #[serial(cooldown_state)]
3271 #[tokio::test]
3272 async fn clear_all_un_degrades_rate_limited_agent() {
3273 super::reset_global_state().await;
3274 let store = test_store().await;
3275 let agent = "test_clear_all_un_degrades";
3276 let before = sqlite_datetime_offset(-30);
3277
3278 for _ in 0..3 {
3279 store
3280 .record_rate_limit_at(agent, "rate", None, &before)
3281 .await
3282 .unwrap();
3283 }
3284
3285 let agents = vec![agent.to_string()];
3286 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3287 assert!(is_agent_degraded(agent));
3288
3289 clear_cooldown("*", &store).await;
3290 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3291 assert!(
3292 !is_agent_degraded(agent),
3293 "clear --all must un-degrade an agent held only by rate rows"
3294 );
3295 }
3296
3297 #[serial(cooldown_state)]
3300 #[tokio::test]
3301 async fn repeated_agent_failures_escalate_backoff() {
3302 super::reset_global_state().await;
3303 let store = test_store().await;
3304 *cooldown_store().lock().await = Some(store.clone());
3305
3306 let agent = "test_escalation_agent";
3307
3308 record_agent_failure_with_message(agent, "").await;
3310 let remaining_1 = {
3311 let map = cooldowns().lock().unwrap();
3312 let entry = map.get(agent).expect("should have cooldown");
3313 entry.cooldown_until - chrono::Utc::now().timestamp()
3314 };
3315 assert!(
3316 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
3317 "first failure should be ~5 min, got {remaining_1}s"
3318 );
3319
3320 {
3322 let mut map = cooldowns().lock().unwrap();
3323 map.remove(agent);
3324 }
3325
3326 record_agent_failure_with_message(agent, "").await;
3328 let remaining_2 = {
3329 let map = cooldowns().lock().unwrap();
3330 let entry = map.get(agent).expect("should have cooldown");
3331 entry.cooldown_until - chrono::Utc::now().timestamp()
3332 };
3333 assert!(
3334 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
3335 "second failure should be ~15 min, got {remaining_2}s"
3336 );
3337 }
3338
3339 #[serial(cooldown_state)]
3340 #[tokio::test]
3341 async fn vendor_past_retry_falls_back_to_backoff() {
3342 super::reset_global_state().await;
3343 let store = test_store().await;
3347 *cooldown_store().lock().await = Some(store.clone());
3348
3349 let agent = "test_vendor_past_retry";
3350
3351 {
3353 let mut map = cooldowns().lock().unwrap();
3354 map.remove(agent);
3355 }
3356
3357 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
3359
3360 record_agent_failure_with_message(agent, past_msg).await;
3361
3362 assert!(
3366 is_agent_in_cooldown(agent),
3367 "agent should be in cooldown after failure"
3368 );
3369
3370 let remaining = {
3371 let map = cooldowns().lock().unwrap();
3372 let entry = map.get(agent).expect("cooldown entry should exist");
3373 entry.cooldown_until - chrono::Utc::now().timestamp()
3374 };
3375
3376 assert!(
3377 remaining >= BACKOFF_BASE_SECS - 5,
3378 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3379 );
3380
3381 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3383 let stored = store
3384 .kv_get(&kv_key)
3385 .await
3386 .unwrap()
3387 .expect("failure count should be stored");
3388 assert_eq!(stored, "1");
3389 }
3390
3391 #[serial(cooldown_state)]
3392 #[tokio::test]
3393 async fn record_agent_success_resets_backoff() {
3394 super::reset_global_state().await;
3395 let store = test_store().await;
3396 *cooldown_store().lock().await = Some(store.clone());
3397
3398 let agent = "test_success_reset_agent";
3399 let model = "test_model";
3400
3401 for _ in 0..3 {
3403 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3404 }
3405 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3407 let count: u32 = store
3408 .kv_get(&kv_key)
3409 .await
3410 .unwrap()
3411 .unwrap()
3412 .parse()
3413 .unwrap();
3414 assert_eq!(count, 3);
3415
3416 record_agent_success(agent, model).await;
3418
3419 let count_after: u32 = store
3420 .kv_get(&kv_key)
3421 .await
3422 .unwrap()
3423 .unwrap()
3424 .parse()
3425 .unwrap();
3426 assert_eq!(count_after, 0, "success should reset failure count to 0");
3427
3428 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3430 let model_count: u32 = store
3431 .kv_get(&model_kv_key)
3432 .await
3433 .unwrap()
3434 .unwrap()
3435 .parse()
3436 .unwrap();
3437 assert_eq!(
3438 model_count, 0,
3439 "success should reset model failure count to 0"
3440 );
3441 }
3442
3443 #[serial(cooldown_state)]
3444 #[tokio::test]
3445 async fn record_agent_success_without_model_skips_empty_model_key() {
3446 super::reset_global_state().await;
3447 let store = test_store().await;
3448 *cooldown_store().lock().await = Some(store.clone());
3449
3450 let agent = "test_success_without_model";
3451
3452 record_agent_success(agent, "").await;
3453
3454 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3455 let agent_count = store
3456 .kv_get(&agent_kv_key)
3457 .await
3458 .unwrap()
3459 .expect("agent failure count should be stored");
3460 assert_eq!(agent_count, "0");
3461
3462 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3463 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3464 assert!(
3465 empty_model_count.is_none(),
3466 "success without model should not create an empty model failure count key"
3467 );
3468 }
3469
3470 #[serial(cooldown_state)]
3471 #[tokio::test]
3472 async fn clear_cooldown_resets_failure_counts() {
3473 super::reset_global_state().await;
3474 let store = test_store().await;
3475 *cooldown_store().lock().await = Some(store.clone());
3476
3477 let agent = "test_clear_fc_agent";
3478
3479 for _ in 0..5 {
3481 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3482 }
3483 set_agent_cooldown(agent, 3600).await;
3485
3486 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3488 let count: u32 = store
3489 .kv_get(&kv_key)
3490 .await
3491 .unwrap()
3492 .unwrap()
3493 .parse()
3494 .unwrap();
3495 assert_eq!(count, 5);
3496
3497 clear_cooldown(agent, &store).await;
3499
3500 let count_after: u32 = store
3502 .kv_get(&kv_key)
3503 .await
3504 .unwrap()
3505 .unwrap()
3506 .parse()
3507 .unwrap();
3508 assert_eq!(
3509 count_after, 0,
3510 "clear_cooldown should reset failure count to 0"
3511 );
3512
3513 assert!(
3515 !is_agent_in_cooldown(agent),
3516 "agent should not be in cooldown after clear"
3517 );
3518 }
3519
3520 #[serial(cooldown_state)]
3521 #[tokio::test]
3522 async fn clear_all_cooldowns_resets_all_failure_counts() {
3523 super::reset_global_state().await;
3524 let store = test_store().await;
3525 *cooldown_store().lock().await = Some(store.clone());
3526
3527 let agents = ["test_clear_all_a", "test_clear_all_b"];
3528 for agent in &agents {
3529 let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
3530 set_agent_cooldown(agent, 3600).await;
3531 }
3532
3533 clear_cooldown("*", &store).await;
3534
3535 for agent in &agents {
3536 assert!(!is_agent_in_cooldown(agent));
3537 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3538 let count: u32 = store
3539 .kv_get(&kv_key)
3540 .await
3541 .unwrap()
3542 .unwrap()
3543 .parse()
3544 .unwrap();
3545 assert_eq!(
3546 count, 0,
3547 "clear --all should reset failure count for {agent}"
3548 );
3549 }
3550 }
3551
3552 #[serial(cooldown_state)]
3553 #[test]
3554 fn compute_backoff_saturates_on_large_counts() {
3555 assert_eq!(
3557 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3558 BACKOFF_MAX_SECS * 12
3559 );
3560 assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
3561 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
3562 }
3563
3564 #[serial(cooldown_state)]
3565 #[test]
3566 fn compute_backoff_zero_base_returns_zero() {
3567 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3568 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3569 }
3570
3571 #[serial(cooldown_state)]
3574 #[tokio::test]
3575 async fn github_circuit_stays_closed_below_threshold() {
3576 super::reset_global_state().await;
3577 {
3579 let mut open = github_5xx_circuit_open().lock().unwrap();
3580 *open = None;
3581 let mut ts = github_5xx_timestamps().lock().unwrap();
3582 ts.clear();
3583 let mut map = cooldown_lock();
3587 map.remove("github:5xx");
3588 }
3589
3590 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3592 record_github_5xx().await;
3593 }
3594
3595 assert!(
3596 !is_github_circuit_open(),
3597 "circuit should stay closed below threshold"
3598 );
3599 assert_eq!(github_circuit_remaining_secs(), 0);
3600 }
3601
3602 #[serial(cooldown_state)]
3603 #[tokio::test]
3604 async fn github_circuit_trips_at_threshold() {
3605 super::reset_global_state().await;
3606 {
3608 let mut open = github_5xx_circuit_open_lock();
3609 *open = None;
3610 let mut ts = github_5xx_timestamps_lock();
3611 ts.clear();
3612 let mut map = cooldown_lock();
3613 map.remove("github:5xx");
3614 }
3615
3616 for _ in 0..GITHUB_5XX_THRESHOLD {
3618 record_github_5xx().await;
3619 }
3620
3621 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3622 assert!(
3623 github_circuit_remaining_secs() > 0,
3624 "remaining should be > 0 when open"
3625 );
3626 assert!(
3627 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3628 "remaining should not exceed cooldown duration"
3629 );
3630
3631 {
3633 let mut open = github_5xx_circuit_open().lock().unwrap();
3634 *open = None;
3635 }
3636 }
3637
3638 #[serial(cooldown_state)]
3639 #[tokio::test]
3640 async fn github_circuit_cannot_double_trip() {
3641 super::reset_global_state().await;
3642 {
3644 let mut open = github_5xx_circuit_open_lock();
3645 *open = None;
3646 let mut ts = github_5xx_timestamps_lock();
3647 ts.clear();
3648 let mut map = cooldown_lock();
3649 map.remove("github:5xx");
3650 }
3651
3652 for _ in 0..GITHUB_5XX_THRESHOLD {
3654 record_github_5xx().await;
3655 }
3656 assert!(is_github_circuit_open());
3657
3658 let first_remaining = github_circuit_remaining_secs();
3659
3660 for _ in 0..GITHUB_5XX_THRESHOLD {
3662 record_github_5xx().await;
3663 }
3664
3665 let second_remaining = github_circuit_remaining_secs();
3666 assert!(
3667 second_remaining <= first_remaining,
3668 "additional errors should not extend cooldown"
3669 );
3670
3671 {
3673 let mut open = github_5xx_circuit_open().lock().unwrap();
3674 *open = None;
3675 }
3676 }
3677
3678 #[serial(cooldown_state)]
3679 #[tokio::test]
3680 async fn github_circuit_recovers_after_cooldown() {
3681 super::reset_global_state().await;
3682 {
3684 let mut open = github_5xx_circuit_open_lock();
3685 *open = None;
3686 let mut ts = github_5xx_timestamps_lock();
3687 ts.clear();
3688 let mut map = cooldown_lock();
3689 map.remove("github:5xx");
3690 }
3691
3692 {
3694 let mut open = github_5xx_circuit_open().lock().unwrap();
3695 *open = Some(chrono::Utc::now().timestamp() - 1);
3696 }
3697
3698 assert!(
3700 !is_github_circuit_open(),
3701 "circuit should auto-recover after cooldown"
3702 );
3703 }
3704
3705 #[serial(cooldown_state)]
3706 #[tokio::test]
3707 async fn github_circuit_sliding_window_ages_out() {
3708 super::reset_global_state().await;
3709 {
3711 let mut open = github_5xx_circuit_open_lock();
3712 *open = None;
3713 let mut ts = github_5xx_timestamps_lock();
3714 ts.clear();
3715 let mut map = cooldown_lock();
3716 map.remove("github:5xx");
3717 }
3718
3719 {
3721 let mut ts = github_5xx_timestamps().lock().unwrap();
3722 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3723 for _ in 0..GITHUB_5XX_THRESHOLD {
3724 ts.push(old);
3725 }
3726 }
3727
3728 record_github_5xx().await;
3730
3731 assert!(
3732 !is_github_circuit_open(),
3733 "old errors outside window should not trip circuit"
3734 );
3735 }
3736}