1use chrono::TimeZone;
19use regex::Regex;
20use std::collections::HashMap;
21use std::sync::{Arc, Mutex, OnceLock};
22
23pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
29
30pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
32
33pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
35
36pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
38
39pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
41
42pub const TRANSIENT_MODEL_BACKOFF_BASE_SECS: i64 = 5 * 60;
44
45pub const TRANSIENT_MODEL_BACKOFF_MAX_SECS: i64 = 15 * 60;
47
48pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
50
51pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
53
54pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
57
58pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
66
67pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub enum CreditExhaustionReason {
77 OutOfCredits,
79 OrgLevelDisabled,
81 BillingCycleExhausted,
83}
84
85impl CreditExhaustionReason {
86 pub fn outcome_str(self) -> &'static str {
90 match self {
91 CreditExhaustionReason::OutOfCredits => "credit_exhausted",
92 CreditExhaustionReason::OrgLevelDisabled => "org_level_disabled",
93 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
94 }
95 }
96}
97
98pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
105
106pub const SILENCE_COUNT_THRESHOLD: usize = 5;
108
109pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
111
112pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
114
115const KV_PREFIX: &str = "cooldown:";
117
118const SILENCE_COUNT_PREFIX: &str = "silence_count:";
120
121const FAILURE_COUNT_PREFIX: &str = "failure_count:";
123
124const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
127
128const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
132
133const RATE_LIMITS_CLEARED_AT_PREFIX: &str = "rate_limits_cleared_at:";
135
136fn agent_from_key(key: &str) -> &str {
138 key.split_once(':').map_or(key, |(agent, _)| agent)
139}
140
141struct CooldownEntry {
142 cooldown_until: i64,
144 reason: String,
145}
146
147#[derive(Debug)]
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 && now < cooldown_until
210 {
211 map.insert(
212 cooldown_key.to_string(),
213 CooldownEntry {
214 cooldown_until,
215 reason,
216 },
217 );
218 loaded += 1;
219 }
220 }
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 && let Ok(cooldown_until) = raw.parse::<i64>()
235 {
236 let now = chrono::Utc::now().timestamp();
237 if now < cooldown_until {
238 let mut open = github_5xx_circuit_open_lock();
239 *open = Some(cooldown_until);
240 let remaining = cooldown_until - now;
241 tracing::warn!(
242 remaining_secs = remaining,
243 "GitHub 5xx circuit breaker restored from KV (still open)"
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 if let Some(entry) = map.get_mut(&key) {
322 if entry.cooldown_until != ts {
323 entry.cooldown_until = ts;
324 updated += 1;
325 }
326 } else {
327 map.insert(
328 key,
329 CooldownEntry {
330 cooldown_until: ts,
331 reason,
332 },
333 );
334 updated += 1;
335 }
336 }
337 }
338
339 if dropped > 0 || updated > 0 {
340 tracing::info!(
341 dropped,
342 updated,
343 "cooldown sync_from_kv: reconciled in-memory map with KV"
344 );
345 }
346}
347
348const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
352
353const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
355
356const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
360
361const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
363
364pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
383 if count <= 1 {
384 return base;
385 }
386 let exp = u32::try_from(count.saturating_sub(1)).unwrap_or(u32::MAX);
387 let factor = 3_i64.saturating_pow(exp);
388 let standard = base.saturating_mul(factor).min(max);
389
390 let first_max = if base <= 0 {
399 0
400 } else {
401 let mut accumulated = base;
403 let mut first_max_count = 1u64;
404 while accumulated < max && first_max_count < count {
405 accumulated = accumulated.saturating_mul(3);
406 first_max_count += 1;
407 }
408 first_max_count
409 };
410
411 let times_at_max = count.saturating_sub(first_max);
413
414 if !allow_extended {
415 return standard;
416 }
417
418 if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
419 max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
420 } else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
421 max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
422 } else {
423 standard
424 }
425}
426
427async fn read_and_increment_failure_count(
436 store_opt: Option<&Arc<crate::store::TaskStore>>,
437 key: &str,
438) -> u64 {
439 read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
440}
441
442async fn read_and_increment_failure_count_with_prefix(
447 store_opt: Option<&Arc<crate::store::TaskStore>>,
448 prefix: &str,
449 key: &str,
450) -> u64 {
451 let kv_key = format!("{prefix}{key}");
452 if let Some(store) = store_opt {
453 match store.kv_increment(&kv_key).await {
454 Ok(n) => n,
455 Err(e) => {
456 tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
457 u64::MAX
458 }
459 }
460 } else {
461 1
462 }
463}
464
465pub async fn record_agent_success(agent_name: &str, model: &str) {
471 let store_opt = cooldown_store().lock().await.clone();
472 if let Some(store) = store_opt {
473 let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
474 let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
475 if let Err(e) = store.kv_set(&agent_key, "0").await {
476 tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
477 }
478 if !model.is_empty() {
479 let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
480 if let Err(e) = store.kv_set(&model_key, "0").await {
481 tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
482 }
483 }
484 if let Err(e) = store.kv_set(&credit_key, "0").await {
485 tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
486 }
487 }
488}
489
490pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
496 if let Some(cooldown_until) = parse_retry_at(error_message) {
501 let now = chrono::Utc::now().timestamp();
502 if cooldown_until > now {
503 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
504 return;
505 }
506 tracing::debug!(
507 agent = agent_name,
508 cooldown_until,
509 "vendor retry_at is in the past — using backoff instead"
510 );
511 }
512
513 let store_opt = cooldown_store().lock().await.clone();
514 let count = read_and_increment_failure_count(store_opt.as_ref(), agent_name).await;
515 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
516 let cooldown_until =
517 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
518 set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
519}
520
521pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
530 detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
531}
532
533pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
538 let lower = error_message.to_lowercase();
539
540 let overage_patterns = [
541 "overagedisabledreason",
542 "overage_disabled_reason",
543 "overage disabled reason",
544 ];
545
546 if overage_patterns.iter().any(|p| lower.contains(p)) {
547 if lower.contains("out_of_credits") || lower.contains("outofcredits") {
548 return Some(CreditExhaustionReason::OutOfCredits);
549 }
550 if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
551 return Some(CreditExhaustionReason::OrgLevelDisabled);
552 }
553 }
554
555 let billing_patterns = [
556 "credit balance too low",
557 "credit balance insufficient",
558 "insufficient credit balance",
559 "billing quota exceeded",
560 "billing limit exceeded",
561 "organization has been disabled",
562 "org disabled",
563 ];
564
565 if billing_patterns.iter().any(|p| lower.contains(p)) {
566 return Some(CreditExhaustionReason::OrgLevelDisabled);
567 }
568
569 let billing_cycle_patterns = [
572 "billing cycle",
573 "quota refreshed next cycle",
574 "monthly quota",
575 "quota will be refreshed",
576 "next billing cycle",
577 "refreshed in the next cycle",
578 "weekly/monthly limit exhausted",
579 "limit will reset at",
580 "upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
585
586 if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
587 return Some(CreditExhaustionReason::BillingCycleExhausted);
588 }
589
590 if lower.contains("out of credits")
591 || lower.contains("outofcredits")
592 || lower.contains("insufficient funds")
593 || lower.contains("no credits remaining")
594 || lower.contains("insufficient_quota")
595 || lower.contains("quota exceeded")
596 || lower.contains("insufficient balance")
597 || lower.contains("please recharge")
598 || lower.contains("no resource package")
599 || lower.contains("balance exhausted")
600 || lower.contains("402 payment required")
601 {
602 return Some(CreditExhaustionReason::OutOfCredits);
603 }
604
605 None
606}
607
608pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
621 let reason_str = match reason {
622 CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
623 CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
624 CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
625 };
626
627 let store_opt = cooldown_store().lock().await.clone();
628 let count = read_and_increment_failure_count_with_prefix(
633 store_opt.as_ref(),
634 CREDIT_FAILURE_COUNT_PREFIX,
635 agent_name,
636 )
637 .await;
638
639 let (base, max) = match reason {
640 CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
641 CreditExhaustionReason::OrgLevelDisabled => {
642 (ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
643 }
644 CreditExhaustionReason::BillingCycleExhausted => {
645 (BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
646 }
647 };
648 let cooldown_secs = compute_backoff(
649 count,
650 base,
651 max,
652 !matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
653 );
654 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
655 set_cooldown_async(agent_name, cooldown_until, reason_str).await;
656 tracing::warn!(
657 agent = agent_name,
658 reason = reason_str,
659 attempt = count,
660 cooldown_secs,
661 "credit exhaustion detected: applying exponential agent-wide cooldown"
662 );
663}
664
665pub async fn record_model_failure(agent_name: &str, model: &str) {
669 let key = format!("{agent_name}:{model}");
670 let store_opt = cooldown_store().lock().await.clone();
671 let count = read_and_increment_failure_count(store_opt.as_ref(), &key).await;
672 let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
673 let cooldown_until =
674 chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
675 set_cooldown_async(&key, cooldown_until, "model_error").await;
676}
677
678pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
686 if let Some(window_secs) = parse_relative_usage_window(error_message)
687 && let Some(m) = model
688 {
689 let cooldown_until = chrono::Utc::now()
690 .timestamp()
691 .saturating_add(i64::try_from(window_secs).unwrap_or(i64::MAX));
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 record_agent_failure_with_message(agent_name, error_message).await;
704 if let Some(m) = model {
705 record_model_failure(agent_name, m).await;
706 }
707}
708
709pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
721 let key = format!("{agent_name}:{model}");
722 let store_opt = cooldown_store().lock().await.clone();
723 let count = read_and_increment_failure_count_with_prefix(
724 store_opt.as_ref(),
725 PERSISTENT_FAILURE_COUNT_PREFIX,
726 &key,
727 )
728 .await;
729 let cooldown_secs = compute_backoff(
730 count,
731 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
732 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
733 false,
734 );
735 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
736 set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
737}
738
739pub async fn record_model_unavailable(agent_name: &str, model: &str, message: &str) {
744 if is_model_permanently_gone(message) {
745 record_persistent_model_failure(agent_name, model).await;
746 } else {
747 record_model_failure(agent_name, model).await;
748 }
749}
750
751fn is_model_permanently_gone(message: &str) -> bool {
752 let lower = message.to_lowercase();
753 lower.contains("not found")
754 || lower.contains("not supported")
755 || lower.contains("not available in your country")
756 || lower.contains("not available in your region")
757}
758
759pub async fn record_transient_model_failure(agent_name: &str, model: &str) {
765 let key = format!("{agent_name}:{model}");
766 let store_opt = cooldown_store().lock().await.clone();
767 let count = read_and_increment_failure_count(store_opt.as_ref(), &key).await;
768 let cooldown_secs = compute_backoff(
769 count,
770 TRANSIENT_MODEL_BACKOFF_BASE_SECS,
771 TRANSIENT_MODEL_BACKOFF_MAX_SECS,
772 false,
773 );
774 let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
775 set_cooldown_async(&key, cooldown_until, "transient_model_error").await;
776}
777
778pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
785 let key = format!("{agent_name}:{model}");
786 let cooldown_until = chrono::Utc::now()
787 .timestamp()
788 .saturating_add(i64::try_from(duration_secs).unwrap_or(i64::MAX));
789 set_cooldown_async(&key, cooldown_until, reason).await;
790}
791
792pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
794 let key = format!("{agent_name}:{model}");
795 set_cooldown_async(&key, until, "vendor_reset_at").await;
796}
797
798pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
805 let cooldown_until = chrono::Utc::now()
806 .timestamp()
807 .saturating_add(i64::try_from(duration_secs).unwrap_or(i64::MAX));
808 set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
809}
810
811pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
814 let store_opt = cooldown_store().lock().await.clone();
815 let Some(store) = store_opt else {
816 tracing::debug!(
817 agent = agent_name,
818 model,
819 "skipping silence count record (no KV store)"
820 );
821 return None;
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 u64::try_from(until - now).unwrap_or(0);
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 u64::try_from(until - now).unwrap_or(0);
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.split_once('(').map_or(date_str, |(t, _)| t).trim();
1558 let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
1560 let split = compact.rfind(' ').map_or(0, |i| i + 1);
1561 let last = &compact[split..];
1562 if !last.contains(':')
1563 && (last.ends_with("am") || last.ends_with("pm"))
1564 && last.len() > 2
1565 {
1566 let at = split + last.len() - 2;
1567 compact.insert_str(at, ":00");
1568 }
1569 let md = compact.replace(" at ", " ");
1571 let now_local = chrono::Local::now();
1572 let this_year = chrono::Datelike::year(&now_local);
1573 for year in [this_year, this_year + 1] {
1574 let with_year = format!("{year} {md}");
1575 if let Ok(dt) =
1576 chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
1577 {
1578 use chrono::TimeZone;
1579 if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest()
1580 && local_dt > now_local
1581 {
1582 return Some(local_dt.timestamp());
1583 }
1584 }
1585 }
1586 let time_only = compact.as_str();
1587 for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
1588 if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
1589 let local_now = chrono::Local::now();
1592 let today = local_now.date_naive();
1593 let candidate = chrono::NaiveDateTime::new(today, t);
1594 let candidate_local: chrono::DateTime<chrono::Local> =
1597 match chrono::Local.from_local_datetime(&candidate) {
1598 chrono::LocalResult::Single(dt) => dt,
1599 chrono::LocalResult::Ambiguous(earliest, _) => earliest,
1600 chrono::LocalResult::None => {
1601 let utc_dt = candidate.and_utc();
1603 return Some(utc_dt.timestamp());
1604 }
1605 };
1606 let chosen = if candidate_local <= local_now {
1607 let tomorrow = today.succ_opt().unwrap_or(today);
1609 let ndt = chrono::NaiveDateTime::new(tomorrow, t);
1610 match chrono::Local.from_local_datetime(&ndt) {
1611 chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
1612 chrono::LocalResult::Ambiguous(earliest, _) => {
1613 earliest.with_timezone(&chrono::Utc)
1614 }
1615 chrono::LocalResult::None => ndt.and_utc(),
1616 }
1617 } else {
1618 candidate_local.with_timezone(&chrono::Utc)
1619 };
1620 tracing::info!(
1621 retry_at = %chosen,
1622 raw = date_str,
1623 marker = retry_marker,
1624 "parsed rate limit retry-at time-only; interpreted as next occurrence"
1625 );
1626 return Some(chosen.timestamp());
1627 }
1628 }
1629 }
1630 }
1631
1632 None
1636}
1637
1638pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
1649 static WINDOW_RE: std::sync::LazyLock<Regex> = std::sync::LazyLock::new(|| {
1650 Regex::new(
1651 r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
1652 )
1653 .expect("usage window regex is valid")
1654 });
1655
1656 if error_message.is_empty() {
1657 return None;
1658 }
1659
1660 let lower = error_message.to_lowercase();
1661 let has_context = lower.contains("usage limit")
1662 || lower.contains("quota")
1663 || lower.contains("window")
1664 || lower.contains("limit");
1665 if !has_context {
1666 return None;
1667 }
1668
1669 let mut max_secs: Option<u64> = None;
1670 for caps in WINDOW_RE.captures_iter(&lower) {
1671 let secs = if let Some(num) = caps.get(1) {
1672 let n: u64 = num.as_str().parse().ok()?;
1673 let unit = caps.get(2)?.as_str().to_lowercase();
1674 let multiplier = match unit.as_str() {
1675 "hour" | "hours" | "hr" | "hrs" => 3600,
1676 "minute" | "minutes" | "min" | "mins" => 60,
1677 "day" | "days" => 86400,
1678 _ => return None,
1679 };
1680 n.checked_mul(multiplier)?
1681 } else if let Some(word) = caps.get(3) {
1682 match word.as_str().to_lowercase().as_str() {
1683 "hourly" => 3600,
1684 "daily" => 86400,
1685 "weekly" => 604_800,
1686 "monthly" => 2_592_000,
1687 _ => return None,
1688 }
1689 } else {
1690 continue;
1691 };
1692 max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
1693 }
1694
1695 max_secs
1696}
1697
1698#[cfg(test)]
1699pub(crate) async fn reset_global_state() {
1700 cooldowns().lock().unwrap().clear();
1701 *cooldown_store().lock().await = None;
1702 github_5xx_timestamps().lock().unwrap().clear();
1703 *github_5xx_circuit_open().lock().unwrap() = None;
1704 degraded_agents().lock().unwrap().clear();
1705 crate::engine::runner::agents::opencode::reset_model_caches_for_test();
1706}
1707
1708#[cfg(test)]
1709mod tests {
1710 use super::*;
1711 use serial_test::serial;
1712
1713 async fn test_store() -> Arc<crate::store::TaskStore> {
1714 Arc::new(
1715 crate::store::TaskStore::open_memory()
1716 .await
1717 .expect("in-memory store"),
1718 )
1719 }
1720
1721 fn sqlite_datetime_offset(secs: i64) -> String {
1722 (chrono::Utc::now() + chrono::Duration::seconds(secs))
1723 .format("%Y-%m-%d %H:%M:%S")
1724 .to_string()
1725 }
1726
1727 #[serial(cooldown_state)]
1728 #[tokio::test]
1729 async fn init_loads_unexpired_cooldowns_from_kv() {
1730 super::reset_global_state().await;
1731 let store = test_store().await;
1732 let cooldown_until = chrono::Utc::now().timestamp() + 3600; store
1735 .kv_set(
1736 &format!("{KV_PREFIX}kimi:k2p5"),
1737 &cooldown_until.to_string(),
1738 )
1739 .await
1740 .unwrap();
1741
1742 let expired = chrono::Utc::now().timestamp() - 1;
1744 store
1745 .kv_set(
1746 &format!("{KV_PREFIX}opencode:old-free"),
1747 &expired.to_string(),
1748 )
1749 .await
1750 .unwrap();
1751
1752 {
1753 let mut map = cooldown_lock();
1754 map.clear();
1755 }
1756
1757 init_cooldown_store(store).await;
1758
1759 assert!(
1760 is_model_in_cooldown("kimi", "k2p5"),
1761 "unexpired cooldown should be loaded"
1762 );
1763 assert!(
1764 !is_model_in_cooldown("opencode", "old-free"),
1765 "expired cooldown should not be loaded"
1766 );
1767 }
1768
1769 #[serial(cooldown_state)]
1770 #[tokio::test]
1771 async fn record_model_failure_writes_to_kv() {
1772 super::reset_global_state().await;
1773 let store = test_store().await;
1774
1775 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
1776
1777 record_model_failure("testagent_persist", "testmodel_persist").await;
1778
1779 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1781
1782 let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
1783 let stored = store.kv_get(&kv_key).await.unwrap();
1784 assert!(stored.is_some(), "model failure should be persisted to KV");
1785 let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
1786 let now = chrono::Utc::now().timestamp();
1787 assert!(ts > now, "persisted timestamp should be in the future");
1789 assert_eq!(reason, "model_error");
1790 }
1791
1792 #[serial(cooldown_state)]
1793 #[tokio::test]
1794 async fn cooldown_reason_survives_restart_via_kv() {
1795 super::reset_global_state().await;
1796 let store = test_store().await;
1797 let agent = "test_reason_restart";
1798 let future = chrono::Utc::now().timestamp() + 600;
1799
1800 init_cooldown_store(std::sync::Arc::clone(&store)).await;
1801 set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
1802
1803 {
1805 let mut map = cooldown_lock();
1806 map.clear();
1807 }
1808 init_cooldown_store(store).await;
1809
1810 let entries = list_all_cooldowns();
1811 let entry = entries.iter().find(|(key, ..)| key == agent);
1812 assert_eq!(
1813 entry.map(|(_, _, reason)| reason.as_str()),
1814 Some("billing_cycle_exhausted"),
1815 "cooldown reason must survive a restart, not collapse to 'persisted'"
1816 );
1817 }
1818
1819 #[serial(cooldown_state)]
1820 #[tokio::test]
1821 async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
1822 super::reset_global_state().await;
1823 let store = test_store().await;
1824 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
1825
1826 let agent = "testagent_persistent";
1827 let model = "testmodel_persistent";
1828 let key = format!("{agent}:{model}");
1829
1830 record_persistent_model_failure(agent, model).await;
1832 let now = chrono::Utc::now().timestamp();
1833 let first_until = cooldown_until(&key).expect("first cooldown should exist");
1834 let first_remaining = first_until.saturating_sub(now);
1835 assert!(
1836 first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
1837 "first persistent model cooldown should be ~4h, got {first_remaining}s"
1838 );
1839
1840 record_persistent_model_failure(agent, model).await;
1842 let now2 = chrono::Utc::now().timestamp();
1843 let second_until = cooldown_until(&key).expect("second cooldown should exist");
1844 let second_remaining = second_until.saturating_sub(now2);
1845 assert!(
1846 second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
1847 "second persistent model cooldown should be ~12h, got {second_remaining}s"
1848 );
1849
1850 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1851 let stored_count = store
1852 .kv_get(&count_key)
1853 .await
1854 .expect("kv read should succeed")
1855 .expect("failure count should be written");
1856 assert_eq!(stored_count, "2");
1857 }
1858
1859 #[serial(cooldown_state)]
1860 #[tokio::test]
1861 async fn record_agent_success_does_not_reset_persistent_failure_count() {
1862 super::reset_global_state().await;
1863 let store = test_store().await;
1864 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
1865
1866 let agent = "testagent_persistent_success";
1867 let model = "testmodel_persistent_success";
1868 let key = format!("{agent}:{model}");
1869
1870 record_persistent_model_failure(agent, model).await;
1872 record_persistent_model_failure(agent, model).await;
1873
1874 let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
1875 let stored_count = store
1876 .kv_get(&count_key)
1877 .await
1878 .expect("kv read should succeed")
1879 .expect("failure count should be written");
1880 assert_eq!(stored_count, "2");
1881
1882 record_agent_success(agent, model).await;
1886
1887 let stored_count_after_success = store
1888 .kv_get(&count_key)
1889 .await
1890 .expect("kv read should succeed")
1891 .expect("persistent failure count should still exist");
1892 assert_eq!(
1893 stored_count_after_success, "2",
1894 "persistent failure count must survive an intervening success"
1895 );
1896
1897 record_persistent_model_failure(agent, model).await;
1899 let now = chrono::Utc::now().timestamp();
1900 let third_until = cooldown_until(&key).expect("third cooldown should exist");
1901 let third_remaining = third_until.saturating_sub(now);
1902 assert!(
1903 third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
1904 "third persistent model failure should escalate past the 4h base, got {third_remaining}s"
1905 );
1906 }
1907
1908 #[serial(cooldown_state)]
1909 #[tokio::test]
1910 async fn record_persistent_model_failure_stays_capped_at_7_days() {
1911 super::reset_global_state().await;
1912 let store = test_store().await;
1913 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
1914
1915 let agent = "testagent_persistent_capped";
1916 let model = "testmodel_persistent_capped";
1917 let key = format!("{agent}:{model}");
1918
1919 for _ in 0..9 {
1920 record_persistent_model_failure(agent, model).await;
1921 }
1922
1923 let remaining = cooldown_until(&key)
1924 .expect("cooldown should exist")
1925 .saturating_sub(chrono::Utc::now().timestamp());
1926 assert!(
1927 remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
1928 "persistent model cooldown should stay capped at 7 days, got {remaining}s"
1929 );
1930 assert!(
1931 remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
1932 "persistent model cooldown should not exceed 7 days, got {remaining}s"
1933 );
1934 }
1935
1936 #[serial(cooldown_state)]
1937 #[test]
1938 fn parse_retry_at_codex_format() {
1939 use chrono::Datelike;
1940 use chrono::TimeZone;
1941
1942 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.";
1943 let ts = parse_retry_at(msg);
1944 assert!(ts.is_some(), "should parse codex retry-at date");
1945
1946 let naive = chrono::NaiveDateTime::new(
1949 chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
1950 chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
1951 );
1952 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
1953 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
1954 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
1955 chrono::LocalResult::None => naive.and_utc(),
1956 };
1957 assert_eq!(
1958 ts.unwrap(),
1959 expected_utc.timestamp(),
1960 "timestamp should reflect local-timezone interpretation of the vendor date"
1961 );
1962
1963 let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
1965 assert_eq!(dt.month(), 3);
1966 assert_eq!(dt.day(), 26);
1967 }
1968
1969 #[serial(cooldown_state)]
1970 #[test]
1971 fn parse_retry_at_hour_only_and_month_day() {
1972 let now = chrono::Utc::now().timestamp();
1973 for msg in [
1974 "You've hit your session limit · resets 8am (America/Sao_Paulo)",
1975 "You've hit your session limit · resets 6pm (America/Sao_Paulo)",
1976 "You've hit your session limit · resets 12am (America/Sao_Paulo)",
1977 "You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
1978 ] {
1979 let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
1980 assert!(ts > now, "{msg}");
1981 }
1982 }
1983
1984 #[serial(cooldown_state)]
1985 #[test]
1986 fn parse_retry_at_weekly_resets_without_year() {
1987 use chrono::Datelike;
1988 let t = chrono::Local::now() + chrono::Duration::days(2);
1989 let msg = format!(
1990 "You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
1991 t.format("%b"),
1992 t.day()
1993 );
1994 let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
1995 let remaining = ts - chrono::Utc::now().timestamp();
1996 assert!(
1997 remaining > 3600 && remaining < 3 * 86400,
1998 "expected ~2d retry-at, got {remaining}s"
1999 );
2000 }
2001
2002 #[serial(cooldown_state)]
2003 #[test]
2004 fn parse_retry_at_no_date() {
2005 assert!(parse_retry_at("generic rate limit error").is_none());
2006 assert!(parse_retry_at("").is_none());
2007 }
2008
2009 #[serial(cooldown_state)]
2010 #[test]
2011 fn parse_retry_at_glm_reset_format_with_seconds() {
2012 use chrono::TimeZone;
2013
2014 let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
2015 let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
2016
2017 let naive = chrono::NaiveDateTime::new(
2018 chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
2019 chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
2020 );
2021 let expected_utc = match chrono::Local.from_local_datetime(&naive) {
2022 chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
2023 chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
2024 chrono::LocalResult::None => naive.and_utc(),
2025 };
2026
2027 assert_eq!(
2028 ts,
2029 expected_utc.timestamp(),
2030 "timestamp should reflect local-timezone interpretation of vendor reset date"
2031 );
2032 }
2033
2034 #[serial(cooldown_state)]
2035 #[test]
2036 fn parse_retry_at_billing_cycle_without_date_returns_none() {
2037 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.";
2040 assert!(
2041 parse_retry_at(msg).is_none(),
2042 "billing cycle without date should return None"
2043 );
2044 }
2045
2046 #[serial(cooldown_state)]
2047 #[test]
2048 fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
2049 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.";
2050 assert_eq!(
2051 parse_relative_usage_window(msg),
2052 Some(5 * 3600),
2053 "should extract the 5-hour window as seconds"
2054 );
2055 }
2056
2057 #[serial(cooldown_state)]
2058 #[test]
2059 fn parse_relative_usage_window_handles_common_durations() {
2060 assert_eq!(
2061 parse_relative_usage_window("You've reached your hourly usage limit."),
2062 Some(3600)
2063 );
2064 assert_eq!(
2065 parse_relative_usage_window("You've reached your daily usage limit."),
2066 Some(86400)
2067 );
2068 assert_eq!(
2069 parse_relative_usage_window("You've reached your weekly limit."),
2070 Some(604_800)
2071 );
2072 assert_eq!(
2073 parse_relative_usage_window("You've reached your 30-minute usage limit."),
2074 Some(1800)
2075 );
2076 assert_eq!(
2077 parse_relative_usage_window("You've reached your 2-day quota."),
2078 Some(172_800)
2079 );
2080 }
2081
2082 #[serial(cooldown_state)]
2083 #[test]
2084 fn parse_relative_usage_window_returns_none_without_usage_context() {
2085 assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
2088 assert!(parse_relative_usage_window("rate limit exceeded").is_none());
2089 assert!(parse_relative_usage_window("").is_none());
2090 }
2091
2092 #[serial(cooldown_state)]
2093 #[tokio::test]
2094 async fn record_rate_limit_uses_model_specific_window_cooldown() {
2095 super::reset_global_state().await;
2096 let store = test_store().await;
2097 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2098
2099 let agent = "test_record_rl_window_agent";
2100 let model = "test_record_rl_window_model";
2101 let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
2102
2103 record_rate_limit(agent, Some(model), msg).await;
2104
2105 assert!(
2106 is_model_in_cooldown(agent, model),
2107 "model should be cooled for the relative usage window"
2108 );
2109 assert!(
2110 !is_agent_in_cooldown(agent),
2111 "agent should not be cooled for a model-scoped quota"
2112 );
2113
2114 let key = format!("{agent}:{model}");
2115 let remaining = cooldown_until(&key)
2116 .expect("model cooldown should exist")
2117 .saturating_sub(chrono::Utc::now().timestamp());
2118 assert!(
2119 remaining >= 5 * 3600 - 5,
2120 "model cooldown should be at least ~5 hours, got {remaining}s"
2121 );
2122 assert_eq!(
2123 cooldown_reason(&key).as_deref(),
2124 Some("rate_limit_usage_window"),
2125 "relative usage-window cooldown must be labeled as a rate limit, not silence"
2126 );
2127 }
2128
2129 #[serial(cooldown_state)]
2130 #[tokio::test]
2131 async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
2132 super::reset_global_state().await;
2133 let store = test_store().await;
2134 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2135
2136 let agent = "test_record_rl_fallback_agent";
2137 let model = "test_record_rl_fallback_model";
2138 let msg = "429 Too Many Requests";
2139
2140 record_rate_limit(agent, Some(model), msg).await;
2141
2142 assert!(
2143 is_agent_in_cooldown(agent),
2144 "agent should be cooled when no relative window is detected"
2145 );
2146 assert!(
2147 is_model_in_cooldown(agent, model),
2148 "model should be cooled when no relative window is detected"
2149 );
2150 }
2151
2152 #[serial(cooldown_state)]
2153 #[tokio::test]
2154 async fn agent_cooldown_persists_to_kv() {
2155 super::reset_global_state().await;
2156 let agent = "test_agent_persist_check";
2158 record_agent_failure_with_message(agent, "").await;
2159 assert!(is_agent_in_cooldown(agent));
2160 }
2161
2162 #[serial(cooldown_state)]
2163 #[tokio::test]
2164 async fn cooldown_never_shortens() {
2165 super::reset_global_state().await;
2166 let agent = "test_agent_no_shorten";
2167 set_agent_cooldown(agent, 24 * 60 * 60).await;
2168 let initial = {
2169 let map = cooldowns().lock().unwrap();
2170 map.get(agent)
2171 .expect("cooldown entry should exist")
2172 .cooldown_until
2173 };
2174
2175 record_agent_failure_with_message(agent, "rate limit").await;
2177 let after = {
2178 let map = cooldowns().lock().unwrap();
2179 map.get(agent)
2180 .expect("cooldown entry should exist")
2181 .cooldown_until
2182 };
2183 assert_eq!(
2184 initial, after,
2185 "cooldown should not be shortened by later failures"
2186 );
2187 }
2188
2189 #[serial(cooldown_state)]
2190 #[tokio::test]
2191 async fn silence_agent_cooldown_is_short_lived() {
2192 super::reset_global_state().await;
2193 let agent = "test_silence_agent_cd";
2194 assert!(!is_agent_in_cooldown(agent));
2195
2196 set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
2197 assert!(is_agent_in_cooldown(agent));
2198
2199 let map = cooldowns().lock().unwrap();
2201 let entry = map.get(agent).expect("should have cooldown entry");
2202 let now = chrono::Utc::now().timestamp();
2203 let remaining = entry.cooldown_until - now;
2204 assert!(
2205 remaining <= i64::try_from(SILENCE_AGENT_COOLDOWN_SECS).unwrap(),
2206 "silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
2207 );
2208 assert!(remaining > 0, "cooldown should still be active");
2209 }
2210
2211 #[serial(cooldown_state)]
2212 #[tokio::test]
2213 async fn record_silence_detection_applies_extended_cooldown() {
2214 super::reset_global_state().await;
2215 let store = test_store().await;
2216 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2217
2218 let agent = "test_silence_count_agent";
2219 let model = "test_silence_count_model";
2220 assert!(!is_model_in_cooldown(agent, model));
2221
2222 for _ in 0..SILENCE_COUNT_THRESHOLD {
2223 let result = record_silence_detection(agent, model).await;
2224 assert!(result.is_some());
2225 }
2226
2227 assert!(is_model_in_cooldown(agent, model));
2228 let key = format!("{agent}:{model}");
2229 let remaining = {
2230 let map = cooldowns().lock().unwrap();
2231 let entry = map.get(&key).expect("cooldown entry should exist");
2232 entry.cooldown_until - chrono::Utc::now().timestamp()
2233 };
2234 assert!(
2235 remaining >= i64::try_from(SILENCE_EXTENDED_COOLDOWN_SECS).unwrap() - 5,
2236 "extended cooldown should be applied, got {remaining}s"
2237 );
2238
2239 let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
2240 let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
2241 assert_eq!(stored, "[]");
2242 }
2243
2244 #[serial(cooldown_state)]
2245 #[test]
2246 fn detect_credit_exhaustion_out_of_credits() {
2247 let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
2248 assert_eq!(
2249 detect_credit_exhaustion(msg),
2250 Some(CreditExhaustionReason::OutOfCredits)
2251 );
2252 }
2253
2254 #[serial(cooldown_state)]
2255 #[test]
2256 fn detect_credit_exhaustion_out_of_credits_variations() {
2257 assert_eq!(
2258 detect_credit_exhaustion("out of credits"),
2259 Some(CreditExhaustionReason::OutOfCredits)
2260 );
2261 assert_eq!(
2262 detect_credit_exhaustion("outofcredits"),
2263 Some(CreditExhaustionReason::OutOfCredits)
2264 );
2265 assert_eq!(
2266 detect_credit_exhaustion("no credits remaining"),
2267 Some(CreditExhaustionReason::OutOfCredits)
2268 );
2269 assert_eq!(
2270 detect_credit_exhaustion("insufficient funds"),
2271 Some(CreditExhaustionReason::OutOfCredits)
2272 );
2273 }
2274
2275 #[serial(cooldown_state)]
2276 #[test]
2277 fn detect_credit_exhaustion_org_level_disabled() {
2278 let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
2279 assert_eq!(
2280 detect_credit_exhaustion(msg),
2281 Some(CreditExhaustionReason::OrgLevelDisabled)
2282 );
2283 }
2284
2285 #[serial(cooldown_state)]
2286 #[test]
2287 fn detect_credit_exhaustion_org_level_disabled_variations() {
2288 assert_eq!(
2289 detect_credit_exhaustion("organization has been disabled"),
2290 Some(CreditExhaustionReason::OrgLevelDisabled)
2291 );
2292 assert_eq!(
2293 detect_credit_exhaustion("org disabled"),
2294 Some(CreditExhaustionReason::OrgLevelDisabled)
2295 );
2296 assert_eq!(
2297 detect_credit_exhaustion("billing quota exceeded"),
2298 Some(CreditExhaustionReason::OrgLevelDisabled)
2299 );
2300 assert_eq!(
2301 detect_credit_exhaustion("credit balance too low"),
2302 Some(CreditExhaustionReason::OrgLevelDisabled)
2303 );
2304 }
2305
2306 #[serial(cooldown_state)]
2307 #[test]
2308 fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
2309 let msg = "overageDisabledReason: out_of_credits";
2310 assert_eq!(
2311 detect_credit_exhaustion(msg),
2312 Some(CreditExhaustionReason::OutOfCredits)
2313 );
2314
2315 let msg = "overageDisabledReason: org_level_disabled";
2316 assert_eq!(
2317 detect_credit_exhaustion(msg),
2318 Some(CreditExhaustionReason::OrgLevelDisabled)
2319 );
2320 }
2321
2322 #[serial(cooldown_state)]
2323 #[test]
2324 fn detect_credit_exhaustion_glm_insufficient_balance() {
2325 assert_eq!(
2327 detect_credit_exhaustion(
2328 "API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
2329 ),
2330 Some(CreditExhaustionReason::OutOfCredits)
2331 );
2332 assert_eq!(
2333 detect_credit_exhaustion("Insufficient balance"),
2334 Some(CreditExhaustionReason::OutOfCredits)
2335 );
2336 assert_eq!(
2337 detect_credit_exhaustion("please recharge your account"),
2338 Some(CreditExhaustionReason::OutOfCredits)
2339 );
2340 assert_eq!(
2341 detect_credit_exhaustion("no resource package available"),
2342 Some(CreditExhaustionReason::OutOfCredits)
2343 );
2344 }
2345
2346 #[serial(cooldown_state)]
2347 #[test]
2348 fn detect_credit_exhaustion_none_for_regular_rate_limit() {
2349 assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
2350 assert!(detect_credit_exhaustion("too many requests").is_none());
2351 assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
2352 assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
2353 assert!(
2354 detect_credit_exhaustion("You've hit your usage limit, try again in a few minutes.")
2355 .is_none()
2356 );
2357 }
2358
2359 #[serial(cooldown_state)]
2360 #[tokio::test]
2361 async fn record_credit_exhaustion_applies_agent_cooldown() {
2362 super::reset_global_state().await;
2363 let agent = "test_credit_exhaust_agent";
2364 assert!(!is_agent_in_cooldown(agent));
2365
2366 record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
2367 assert!(is_agent_in_cooldown(agent));
2368
2369 let remaining = {
2370 let map = cooldowns().lock().unwrap();
2371 let entry = map.get(agent).expect("cooldown entry should exist");
2372 entry.cooldown_until - chrono::Utc::now().timestamp()
2373 };
2374 assert!(
2376 remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
2377 "first out_of_credits cooldown should be ~1 hour, got {remaining}s"
2378 );
2379 assert!(
2380 remaining <= CREDIT_BACKOFF_MAX_SECS,
2381 "cooldown should not exceed cap of 8h, got {remaining}s"
2382 );
2383 }
2384
2385 #[serial(cooldown_state)]
2386 #[tokio::test]
2387 async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
2388 super::reset_global_state().await;
2389 let agent = "test_org_disabled_agent";
2390 assert!(!is_agent_in_cooldown(agent));
2391
2392 record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
2393 assert!(is_agent_in_cooldown(agent));
2394
2395 let remaining = {
2396 let map = cooldowns().lock().unwrap();
2397 let entry = map.get(agent).expect("cooldown entry should exist");
2398 entry.cooldown_until - chrono::Utc::now().timestamp()
2399 };
2400 assert!(
2402 remaining >= ORG_BACKOFF_BASE_SECS - 5,
2403 "first org-level disabled cooldown should be ~2 hours, got {remaining}s"
2404 );
2405 assert!(
2406 remaining <= CREDIT_BACKOFF_MAX_SECS,
2407 "cooldown should not exceed cap of 8h, got {remaining}s"
2408 );
2409 }
2410
2411 #[serial(cooldown_state)]
2412 #[tokio::test]
2413 async fn record_credit_exhaustion_billing_cycle_applies_24h() {
2414 super::reset_global_state().await;
2415 let store = test_store().await;
2416 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2417
2418 let agent = "test_billing_cycle_agent";
2419 assert!(!is_agent_in_cooldown(agent));
2420
2421 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2422 assert!(is_agent_in_cooldown(agent));
2423
2424 let remaining = {
2425 let map = cooldowns().lock().unwrap();
2426 let entry = map.get(agent).expect("cooldown entry should exist");
2427 entry.cooldown_until - chrono::Utc::now().timestamp()
2428 };
2429 assert!(
2431 remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
2432 "first billing cycle cooldown should be ~24 hours, got {remaining}s"
2433 );
2434 assert!(
2435 remaining <= BILLING_CYCLE_MAX_SECS,
2436 "cooldown should not exceed cap of 7 days, got {remaining}s"
2437 );
2438 }
2439
2440 #[serial(cooldown_state)]
2441 #[tokio::test]
2442 async fn record_credit_exhaustion_billing_cycle_escalates() {
2443 super::reset_global_state().await;
2444 let store = test_store().await;
2445 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2446
2447 let agent = "test_billing_cycle_escalation";
2448
2449 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2451 let remaining_1 = {
2452 let map = cooldowns().lock().unwrap();
2453 let entry = map.get(agent).expect("cooldown entry should exist");
2454 entry.cooldown_until - chrono::Utc::now().timestamp()
2455 };
2456 assert!(
2457 (BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
2458 .contains(&remaining_1),
2459 "first billing cycle cooldown should be ~24h, got {remaining_1}s"
2460 );
2461
2462 {
2464 let mut map = cooldowns().lock().unwrap();
2465 map.remove(agent);
2466 }
2467
2468 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2470 let remaining_2 = {
2471 let map = cooldowns().lock().unwrap();
2472 let entry = map.get(agent).expect("cooldown entry should exist");
2473 entry.cooldown_until - chrono::Utc::now().timestamp()
2474 };
2475 let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
2477 remaining_2 >= expected_2 - 5,
2478 "second billing cycle cooldown should be ~72h, got {remaining_2}s"
2479 );
2480
2481 {
2483 let mut map = cooldowns().lock().unwrap();
2484 map.remove(agent);
2485 }
2486
2487 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2489 let remaining_3 = {
2490 let map = cooldowns().lock().unwrap();
2491 let entry = map.get(agent).expect("cooldown entry should exist");
2492 entry.cooldown_until - chrono::Utc::now().timestamp()
2493 };
2494 assert!(
2495 remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
2496 "third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
2497 );
2498 assert!(
2499 remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
2500 "third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
2501 );
2502 }
2503
2504 #[serial(cooldown_state)]
2505 #[tokio::test]
2506 async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
2507 super::reset_global_state().await;
2508 let store = test_store().await;
2509 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
2510
2511 let agent = "test_billing_cycle_agent_capped";
2512
2513 for _ in 0..5 {
2514 record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
2515 }
2516
2517 let remaining = {
2518 let map = cooldowns().lock().unwrap();
2519 let entry = map.get(agent).expect("cooldown entry should exist");
2520 entry.cooldown_until - chrono::Utc::now().timestamp()
2521 };
2522 assert!(
2523 remaining >= BILLING_CYCLE_MAX_SECS - 5,
2524 "billing cycle cooldown should stay capped at 7 days, got {remaining}s"
2525 );
2526 assert!(
2527 remaining <= BILLING_CYCLE_MAX_SECS + 5,
2528 "billing cycle cooldown should not exceed 7 days, got {remaining}s"
2529 );
2530 }
2531
2532 #[serial(cooldown_state)]
2533 #[test]
2534 fn credit_exhaustion_backoff_constants() {
2535 assert_eq!(
2537 compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2538 CREDIT_BACKOFF_BASE_SECS
2539 );
2540 assert_eq!(
2542 compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
2543 ORG_BACKOFF_BASE_SECS
2544 );
2545 assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
2547 assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
2548 }
2549
2550 #[serial(cooldown_state)]
2551 #[test]
2552 fn compute_backoff_long_cap_without_extended_stays_capped() {
2553 assert_eq!(
2554 compute_backoff(
2555 3,
2556 BILLING_CYCLE_COOLDOWN_SECS,
2557 BILLING_CYCLE_MAX_SECS,
2558 false
2559 ),
2560 BILLING_CYCLE_MAX_SECS,
2561 "7-day capped paths should stop at the documented max"
2562 );
2563 assert_eq!(
2564 compute_backoff(
2565 9,
2566 PERSISTENT_MODEL_BACKOFF_BASE_SECS,
2567 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2568 false,
2569 ),
2570 PERSISTENT_MODEL_BACKOFF_MAX_SECS,
2571 "persistent model failures should not extend past 7 days"
2572 );
2573 }
2574
2575 #[serial(cooldown_state)]
2578 #[test]
2579 fn mark_and_check_agent_degraded() {
2580 let agent = "test_degraded_mark";
2581 assert!(!is_agent_degraded(agent));
2582 mark_agent_degraded(agent);
2583 assert!(is_agent_degraded(agent));
2584 clear_agent_degraded(agent);
2585 assert!(!is_agent_degraded(agent));
2586 }
2587
2588 #[serial(cooldown_state)]
2589 #[tokio::test]
2590 async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
2591 super::reset_global_state().await;
2592 let store = test_store().await;
2593 let agent = "test_sync_external_clear";
2594 let future = chrono::Utc::now().timestamp() + 600;
2595
2596 init_cooldown_store(std::sync::Arc::clone(&store)).await;
2598 set_cooldown_async(agent, future, "test").await;
2599 assert!(is_agent_in_cooldown(agent));
2600
2601 store
2605 .kv_set(&format!("{KV_PREFIX}{agent}"), "0")
2606 .await
2607 .expect("kv_set tombstone");
2608 assert!(
2609 is_agent_in_cooldown(agent),
2610 "in-memory cooldown should still be present before sync"
2611 );
2612
2613 sync_from_kv(&store).await;
2614
2615 assert!(
2616 !is_agent_in_cooldown(agent),
2617 "sync_from_kv must drop in-memory entries the KV tombstoned"
2618 );
2619 }
2620
2621 #[serial(cooldown_state)]
2622 #[tokio::test]
2623 async fn sync_from_kv_closes_github_circuit_when_row_cleared() {
2624 super::reset_global_state().await;
2625 let store = test_store().await;
2626 {
2627 *github_5xx_circuit_open_lock() = Some(chrono::Utc::now().timestamp() + 600);
2628 github_5xx_timestamps_lock().push(chrono::Utc::now().timestamp());
2629 }
2630 assert!(is_github_circuit_open());
2631
2632 store
2633 .kv_set("cooldown:github:5xx", "0")
2634 .await
2635 .expect("kv_set tombstone");
2636 sync_from_kv(&store).await;
2637
2638 assert!(!is_github_circuit_open());
2639 assert!(github_5xx_timestamps_lock().is_empty());
2640 }
2641
2642 #[serial(cooldown_state)]
2643 #[tokio::test]
2644 async fn sync_from_kv_keeps_5xx_window_when_expired_row_and_breaker_closed() {
2645 super::reset_global_state().await;
2646 let store = test_store().await;
2647 let now = chrono::Utc::now().timestamp();
2648 {
2649 github_5xx_timestamps_lock().push(now);
2650 github_5xx_timestamps_lock().push(now);
2651 }
2652
2653 store
2654 .kv_set("cooldown:github:5xx", &(now - 600).to_string())
2655 .await
2656 .expect("kv_set past");
2657 sync_from_kv(&store).await;
2658
2659 assert!(!is_github_circuit_open());
2660 assert_eq!(github_5xx_timestamps_lock().len(), 2);
2661
2662 github_5xx_timestamps_lock().clear();
2663 }
2664
2665 #[serial(cooldown_state)]
2666 #[tokio::test]
2667 async fn sync_from_kv_keeps_github_circuit_open_for_future_row() {
2668 super::reset_global_state().await;
2669 let store = test_store().await;
2670 let future = chrono::Utc::now().timestamp() + 600;
2671 *github_5xx_circuit_open_lock() = Some(future);
2672
2673 store
2674 .kv_set("cooldown:github:5xx", &future.to_string())
2675 .await
2676 .expect("kv_set future");
2677 sync_from_kv(&store).await;
2678
2679 assert!(is_github_circuit_open());
2680
2681 *github_5xx_circuit_open_lock() = None;
2682 }
2683
2684 #[serial(cooldown_state)]
2685 #[tokio::test]
2686 async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
2687 super::reset_global_state().await;
2691 let store = test_store().await;
2692 let agent = "test_sync_race_window";
2693 let future = chrono::Utc::now().timestamp() + 600;
2694
2695 set_cooldown_in_memory(agent, future, "raced");
2697 assert!(is_agent_in_cooldown(agent));
2698
2699 sync_from_kv(&store).await;
2700
2701 assert!(
2702 is_agent_in_cooldown(agent),
2703 "sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
2704 );
2705 }
2706
2707 #[serial(cooldown_state)]
2708 #[tokio::test]
2709 async fn sync_from_kv_adopts_kv_timestamp() {
2710 super::reset_global_state().await;
2713 let store = test_store().await;
2714 let agent = "test_sync_adopt_kv";
2715 let original = chrono::Utc::now().timestamp() + 600;
2716 let updated = chrono::Utc::now().timestamp() + 1200;
2717
2718 init_cooldown_store(std::sync::Arc::clone(&store)).await;
2719 set_cooldown_async(agent, original, "initial").await;
2720
2721 store
2722 .kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
2723 .await
2724 .expect("kv_set updated");
2725
2726 sync_from_kv(&store).await;
2727
2728 let until = cooldown_until(agent).expect("cooldown still active");
2729 assert_eq!(
2730 until, updated,
2731 "sync_from_kv must adopt KV's authoritative timestamp"
2732 );
2733 }
2734
2735 #[serial(cooldown_state)]
2736 #[tokio::test]
2737 async fn refresh_degraded_agents_marks_cooled_agent() {
2738 super::reset_global_state().await;
2739 let store = test_store().await;
2740 let agent = "test_refresh_cooled";
2741
2742 record_agent_failure_with_message(agent, "").await;
2744 assert!(is_agent_in_cooldown(agent));
2745
2746 let agents = vec![agent.to_string()];
2748 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2749
2750 assert!(
2751 is_agent_degraded(agent),
2752 "agent in cooldown should be marked degraded"
2753 );
2754
2755 clear_agent_degraded(agent);
2757 }
2758
2759 #[serial(cooldown_state)]
2760 #[tokio::test]
2761 async fn refresh_degraded_agents_marks_no_models_agent() {
2762 super::reset_global_state().await;
2763 let store = test_store().await;
2764 let agent = "test_refresh_no_models";
2765
2766 let agents = vec![agent.to_string()];
2768 refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
2769
2770 assert!(
2771 is_agent_degraded(agent),
2772 "agent with no available models should be marked degraded"
2773 );
2774
2775 clear_agent_degraded(agent);
2777 }
2778
2779 #[serial(cooldown_state)]
2780 #[tokio::test]
2781 async fn refresh_degraded_agents_marks_rate_limited_agent() {
2782 super::reset_global_state().await;
2783 let store = test_store().await;
2784 let agent = "test_refresh_rate_limited";
2785
2786 for _ in 0..4 {
2788 store.record_rate_limit(agent, "rate", None).await.unwrap();
2789 }
2790
2791 let agents = vec![agent.to_string()];
2792 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2794
2795 assert!(
2796 is_agent_degraded(agent),
2797 "agent with rate limit count >= threshold should be marked degraded"
2798 );
2799
2800 clear_agent_degraded(agent);
2802 }
2803
2804 #[serial(cooldown_state)]
2805 #[tokio::test]
2806 async fn refresh_degraded_agents_clears_healthy_agent() {
2807 super::reset_global_state().await;
2808 let store = test_store().await;
2809 let agent = "test_refresh_healthy";
2810
2811 mark_agent_degraded(agent);
2813 assert!(is_agent_degraded(agent));
2814
2815 let agents = vec![agent.to_string()];
2817 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2818
2819 assert!(
2820 !is_agent_degraded(agent),
2821 "healthy agent should have degraded flag cleared"
2822 );
2823 }
2824
2825 #[serial(cooldown_state)]
2826 #[tokio::test]
2827 async fn refresh_degraded_agents_ignores_expired_cooldowns() {
2828 super::reset_global_state().await;
2829 let store = test_store().await;
2830 let agent = "test_refresh_expired_cooldown";
2831
2832 let expired_ts = chrono::Utc::now().timestamp() - 100; {
2836 let mut map = cooldowns().lock().unwrap();
2837 map.insert(
2838 agent.to_string(),
2839 CooldownEntry {
2840 cooldown_until: expired_ts,
2841 reason: "expired_test".to_string(),
2842 },
2843 );
2844 }
2845
2846 assert!(
2848 !is_agent_in_cooldown(agent),
2849 "expired cooldown should not be active"
2850 );
2851
2852 mark_agent_degraded(agent);
2854 assert!(is_agent_degraded(agent));
2855
2856 let agents = vec![agent.to_string()];
2858 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
2859
2860 assert!(
2861 !is_agent_degraded(agent),
2862 "agent with only expired cooldown should NOT be marked degraded"
2863 );
2864
2865 clear_agent_degraded(agent);
2867 {
2868 let mut map = cooldowns().lock().unwrap();
2869 map.remove(agent);
2870 }
2871 }
2872
2873 #[serial(cooldown_state)]
2874 #[test]
2875 fn compute_backoff_grows_exponentially() {
2876 assert_eq!(compute_backoff(0, 300, 14400, true), 300);
2878 assert_eq!(compute_backoff(1, 300, 14400, true), 300);
2879 assert_eq!(compute_backoff(2, 300, 14400, true), 900);
2881 assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
2883 assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
2885 assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
2887 }
2888
2889 #[serial(cooldown_state)]
2890 #[test]
2891 fn compute_backoff_extended_tier_1_applies_24h() {
2892 assert_eq!(
2898 compute_backoff(5, 300, 14400, true),
2899 14400,
2900 "first at max, not extended"
2901 );
2902 assert_eq!(
2903 compute_backoff(6, 300, 14400, true),
2904 14400,
2905 "times_at_max=1, not extended"
2906 );
2907 assert_eq!(
2908 compute_backoff(7, 300, 14400, true),
2909 86400,
2910 "times_at_max=2, extended 6x"
2911 );
2912 assert_eq!(
2913 compute_backoff(8, 300, 14400, true),
2914 86400,
2915 "times_at_max=3, still extended 6x"
2916 );
2917 }
2918
2919 #[serial(cooldown_state)]
2920 #[test]
2921 fn compute_backoff_extended_tier_2_applies_48h() {
2922 assert_eq!(
2925 compute_backoff(9, 300, 14400, true),
2926 86400,
2927 "times_at_max=4, still tier 1"
2928 );
2929 assert_eq!(
2930 compute_backoff(10, 300, 14400, true),
2931 172_800,
2932 "times_at_max=5, extended 12x"
2933 );
2934 assert_eq!(
2935 compute_backoff(15, 300, 14400, true),
2936 172_800,
2937 "times_at_max=10, still tier 2"
2938 );
2939 assert_eq!(
2940 compute_backoff(20, 300, 14400, true),
2941 172_800,
2942 "times_at_max=15, still tier 2"
2943 );
2944 }
2945
2946 #[serial(cooldown_state)]
2947 #[test]
2948 fn compute_backoff_extended_tier_with_different_base_max() {
2949 assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
2956 assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
2957 assert_eq!(
2958 compute_backoff(7, 60, 3600, true),
2959 21600,
2960 "times_at_max=2, extended 6x"
2961 );
2962 }
2963
2964 #[serial(cooldown_state)]
2965 #[test]
2966 fn compute_backoff_extended_tier_persistent_fails_example() {
2967 assert_eq!(
2971 compute_backoff(21, 300, 14400, true),
2972 172_800,
2973 "21 failures should get 48h extended cooldown"
2974 );
2975
2976 assert_eq!(
2978 compute_backoff(36, 300, 14400, true),
2979 172_800,
2980 "36 failures should also get 48h extended cooldown"
2981 );
2982 }
2983
2984 #[serial(cooldown_state)]
2985 #[test]
2986 fn compute_backoff_extended_tier_zero_base() {
2987 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
2989 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
2990 assert_eq!(compute_backoff(10, 0, 14400, true), 0);
2991 }
2992
2993 #[serial(cooldown_state)]
2994 #[test]
2995 fn detect_credit_exhaustion_billing_cycle() {
2996 assert_eq!(
2997 detect_credit_exhaustion("billing cycle"),
2998 Some(CreditExhaustionReason::BillingCycleExhausted)
2999 );
3000 assert_eq!(
3001 detect_credit_exhaustion("quota refreshed next cycle"),
3002 Some(CreditExhaustionReason::BillingCycleExhausted)
3003 );
3004 assert_eq!(
3005 detect_credit_exhaustion(
3006 "You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
3007 ),
3008 Some(CreditExhaustionReason::BillingCycleExhausted)
3009 );
3010 assert_eq!(
3011 detect_credit_exhaustion("monthly quota exceeded"),
3012 Some(CreditExhaustionReason::BillingCycleExhausted)
3013 );
3014 assert_eq!(
3015 detect_credit_exhaustion(
3016 "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
3017 ),
3018 Some(CreditExhaustionReason::BillingCycleExhausted)
3019 );
3020 }
3021
3022 #[serial(cooldown_state)]
3023 #[test]
3024 fn detect_credit_exhaustion_billing_cycle_minimax() {
3025 assert_eq!(
3026 detect_credit_exhaustion(
3027 "minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
3028 ),
3029 Some(CreditExhaustionReason::BillingCycleExhausted)
3030 );
3031 assert_eq!(
3032 detect_credit_exhaustion("Upgrade your Token Plan"),
3033 Some(CreditExhaustionReason::BillingCycleExhausted)
3034 );
3035 assert_eq!(
3036 detect_credit_exhaustion("Token Plan usage limit reached"),
3037 Some(CreditExhaustionReason::BillingCycleExhausted)
3038 );
3039 assert_eq!(
3040 detect_credit_exhaustion("purchase Credits for more usage"),
3041 Some(CreditExhaustionReason::BillingCycleExhausted)
3042 );
3043 assert_eq!(
3044 detect_credit_exhaustion(
3045 "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"
3046 ),
3047 Some(CreditExhaustionReason::BillingCycleExhausted)
3048 );
3049 }
3050
3051 #[serial(cooldown_state)]
3052 #[test]
3053 fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
3054 let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
3059 this billing cycle. Your quota will be refreshed soon.";
3060 assert_eq!(
3061 detect_credit_exhaustion_outcome(msg),
3062 Some("billing_cycle_exhausted")
3063 );
3064 }
3065
3066 #[serial(cooldown_state)]
3067 #[test]
3068 fn detect_credit_exhaustion_outcome_maps_all_reasons() {
3069 assert_eq!(
3070 detect_credit_exhaustion_outcome("out of credits"),
3071 Some("credit_exhausted")
3072 );
3073 assert_eq!(
3074 detect_credit_exhaustion_outcome("organization has been disabled"),
3075 Some("org_level_disabled")
3076 );
3077 assert_eq!(
3078 detect_credit_exhaustion_outcome("monthly quota exceeded"),
3079 Some("billing_cycle_exhausted")
3080 );
3081 }
3082
3083 #[test]
3084 fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
3085 assert_eq!(
3086 detect_credit_exhaustion_outcome("429 Too Many Requests"),
3087 None
3088 );
3089 assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
3090 }
3091
3092 #[serial(cooldown_state)]
3093 #[test]
3094 fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
3095 let msg = "overageDisabledReason: org_level_disabled for billing cycle";
3098 assert_eq!(
3099 detect_credit_exhaustion(msg),
3100 Some(CreditExhaustionReason::OrgLevelDisabled)
3101 );
3102 }
3103
3104 #[serial(cooldown_state)]
3105 #[tokio::test]
3106 async fn refresh_degraded_agents_below_threshold_not_degraded() {
3107 super::reset_global_state().await;
3108 let store = test_store().await;
3109 let agent = "test_refresh_below_threshold";
3110
3111 for _ in 0..2 {
3113 store.record_rate_limit(agent, "rate", None).await.unwrap();
3114 }
3115
3116 let agents = vec![agent.to_string()];
3117 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3119
3120 assert!(
3121 !is_agent_degraded(agent),
3122 "agent with rate limit count < threshold should not be degraded"
3123 );
3124 }
3125
3126 #[serial(cooldown_state)]
3127 #[tokio::test]
3128 async fn refresh_degraded_agents_ignores_error_limit_type() {
3129 super::reset_global_state().await;
3130 let store = test_store().await;
3131 let agent = "test_refresh_error_rows";
3132
3133 for _ in 0..3 {
3134 store.record_rate_limit(agent, "error", None).await.unwrap();
3135 }
3136
3137 let agents = vec![agent.to_string()];
3138 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3139
3140 assert!(
3141 !is_agent_degraded(agent),
3142 "three error rows in 6 hours must not mark an agent degraded"
3143 );
3144 }
3145
3146 #[serial(cooldown_state)]
3147 #[tokio::test]
3148 async fn refresh_degraded_agents_marks_three_rate_rows() {
3149 super::reset_global_state().await;
3150 let store = test_store().await;
3151 let agent = "test_refresh_three_rate_rows";
3152
3153 for _ in 0..3 {
3154 store.record_rate_limit(agent, "rate", None).await.unwrap();
3155 }
3156
3157 let agents = vec![agent.to_string()];
3158 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3159
3160 assert!(
3161 is_agent_degraded(agent),
3162 "three rate rows in 6 hours must mark an agent degraded"
3163 );
3164 assert_eq!(
3165 agent_degraded_reason(agent).as_deref(),
3166 Some("rate limit threshold exceeded")
3167 );
3168 let listed = list_degraded_agents();
3169 assert!(
3170 listed
3171 .iter()
3172 .any(|(a, r)| a == agent && r == "rate limit threshold exceeded"),
3173 "orch cooldown list must show the degraded reason, got {listed:?}"
3174 );
3175
3176 clear_agent_degraded(agent);
3177 }
3178
3179 #[serial(cooldown_state)]
3180 #[tokio::test]
3181 async fn refresh_degraded_agents_logs_reason_change() {
3182 super::reset_global_state().await;
3183 let store = test_store().await;
3184 let agent = "test_refresh_reason_change";
3185 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3186
3187 set_agent_cooldown(agent, 600).await;
3188 let agents = vec![agent.to_string()];
3189 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3190 assert_eq!(
3191 agent_degraded_reason(agent).as_deref(),
3192 Some("agent in cooldown")
3193 );
3194
3195 clear_cooldown(agent, &store).await;
3196 let after_clear = sqlite_datetime_offset(60);
3197 for _ in 0..3 {
3198 store
3199 .record_rate_limit_at(agent, "rate", None, &after_clear)
3200 .await
3201 .unwrap();
3202 }
3203 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3204 assert_eq!(
3205 agent_degraded_reason(agent).as_deref(),
3206 Some("rate limit threshold exceeded"),
3207 "degraded reason must update when the hold changes"
3208 );
3209
3210 clear_agent_degraded(agent);
3211 }
3212
3213 #[serial(cooldown_state)]
3214 #[tokio::test]
3215 async fn clear_un_degrades_on_next_health_check() {
3216 super::reset_global_state().await;
3217 let store = test_store().await;
3218 let agent = "test_clear_un_degrades";
3219 let before = sqlite_datetime_offset(-30);
3220 let after = sqlite_datetime_offset(60);
3221
3222 for _ in 0..3 {
3223 store
3224 .record_rate_limit_at(agent, "rate", None, &before)
3225 .await
3226 .unwrap();
3227 }
3228
3229 let agents = vec![agent.to_string()];
3230 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3231 assert!(
3232 is_agent_degraded(agent),
3233 "three rate rows must mark an agent degraded"
3234 );
3235
3236 clear_cooldown(agent, &store).await;
3237 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3238 assert!(
3239 !is_agent_degraded(agent),
3240 "clear must un-degrade on the next health check"
3241 );
3242
3243 store
3244 .record_rate_limit_at(agent, "rate", None, &after)
3245 .await
3246 .unwrap();
3247 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3248 assert!(
3249 !is_agent_degraded(agent),
3250 "one rate row after the mark must not degrade again"
3251 );
3252
3253 for _ in 0..2 {
3254 store
3255 .record_rate_limit_at(agent, "rate", None, &after)
3256 .await
3257 .unwrap();
3258 }
3259 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3260 assert!(
3261 is_agent_degraded(agent),
3262 "three rate rows after the mark must degrade again"
3263 );
3264
3265 clear_agent_degraded(agent);
3266 }
3267
3268 #[serial(cooldown_state)]
3269 #[tokio::test]
3270 async fn clear_all_un_degrades_rate_limited_agent() {
3271 super::reset_global_state().await;
3272 let store = test_store().await;
3273 let agent = "test_clear_all_un_degrades";
3274 let before = sqlite_datetime_offset(-30);
3275
3276 for _ in 0..3 {
3277 store
3278 .record_rate_limit_at(agent, "rate", None, &before)
3279 .await
3280 .unwrap();
3281 }
3282
3283 let agents = vec![agent.to_string()];
3284 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3285 assert!(is_agent_degraded(agent));
3286
3287 clear_cooldown("*", &store).await;
3288 refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
3289 assert!(
3290 !is_agent_degraded(agent),
3291 "clear --all must un-degrade an agent held only by rate rows"
3292 );
3293 }
3294
3295 #[serial(cooldown_state)]
3298 #[tokio::test]
3299 async fn repeated_agent_failures_escalate_backoff() {
3300 super::reset_global_state().await;
3301 let store = test_store().await;
3302 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3303
3304 let agent = "test_escalation_agent";
3305
3306 record_agent_failure_with_message(agent, "").await;
3308 let remaining_1 = {
3309 let map = cooldowns().lock().unwrap();
3310 let entry = map.get(agent).expect("should have cooldown");
3311 entry.cooldown_until - chrono::Utc::now().timestamp()
3312 };
3313 assert!(
3314 (BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
3315 "first failure should be ~5 min, got {remaining_1}s"
3316 );
3317
3318 {
3320 let mut map = cooldowns().lock().unwrap();
3321 map.remove(agent);
3322 }
3323
3324 record_agent_failure_with_message(agent, "").await;
3326 let remaining_2 = {
3327 let map = cooldowns().lock().unwrap();
3328 let entry = map.get(agent).expect("should have cooldown");
3329 entry.cooldown_until - chrono::Utc::now().timestamp()
3330 };
3331 assert!(
3332 remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
3333 "second failure should be ~15 min, got {remaining_2}s"
3334 );
3335 }
3336
3337 #[serial(cooldown_state)]
3338 #[tokio::test]
3339 async fn vendor_past_retry_falls_back_to_backoff() {
3340 super::reset_global_state().await;
3341 let store = test_store().await;
3345 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3346
3347 let agent = "test_vendor_past_retry";
3348
3349 {
3351 let mut map = cooldowns().lock().unwrap();
3352 map.remove(agent);
3353 }
3354
3355 let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
3357
3358 record_agent_failure_with_message(agent, past_msg).await;
3359
3360 assert!(
3364 is_agent_in_cooldown(agent),
3365 "agent should be in cooldown after failure"
3366 );
3367
3368 let remaining = {
3369 let map = cooldowns().lock().unwrap();
3370 let entry = map.get(agent).expect("cooldown entry should exist");
3371 entry.cooldown_until - chrono::Utc::now().timestamp()
3372 };
3373
3374 assert!(
3375 remaining >= BACKOFF_BASE_SECS - 5,
3376 "expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
3377 );
3378
3379 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3381 let stored = store
3382 .kv_get(&kv_key)
3383 .await
3384 .unwrap()
3385 .expect("failure count should be stored");
3386 assert_eq!(stored, "1");
3387 }
3388
3389 #[serial(cooldown_state)]
3390 #[tokio::test]
3391 async fn record_agent_success_resets_backoff() {
3392 super::reset_global_state().await;
3393 let store = test_store().await;
3394 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3395
3396 let agent = "test_success_reset_agent";
3397 let model = "test_model";
3398
3399 for _ in 0..3 {
3401 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3402 }
3403 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3405 let count: u32 = store
3406 .kv_get(&kv_key)
3407 .await
3408 .unwrap()
3409 .unwrap()
3410 .parse()
3411 .unwrap();
3412 assert_eq!(count, 3);
3413
3414 record_agent_success(agent, model).await;
3416
3417 let count_after: u32 = store
3418 .kv_get(&kv_key)
3419 .await
3420 .unwrap()
3421 .unwrap()
3422 .parse()
3423 .unwrap();
3424 assert_eq!(count_after, 0, "success should reset failure count to 0");
3425
3426 let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
3428 let model_count: u32 = store
3429 .kv_get(&model_kv_key)
3430 .await
3431 .unwrap()
3432 .unwrap()
3433 .parse()
3434 .unwrap();
3435 assert_eq!(
3436 model_count, 0,
3437 "success should reset model failure count to 0"
3438 );
3439 }
3440
3441 #[serial(cooldown_state)]
3442 #[tokio::test]
3443 async fn record_agent_success_without_model_skips_empty_model_key() {
3444 super::reset_global_state().await;
3445 let store = test_store().await;
3446 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3447
3448 let agent = "test_success_without_model";
3449
3450 record_agent_success(agent, "").await;
3451
3452 let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3453 let agent_count = store
3454 .kv_get(&agent_kv_key)
3455 .await
3456 .unwrap()
3457 .expect("agent failure count should be stored");
3458 assert_eq!(agent_count, "0");
3459
3460 let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
3461 let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
3462 assert!(
3463 empty_model_count.is_none(),
3464 "success without model should not create an empty model failure count key"
3465 );
3466 }
3467
3468 #[serial(cooldown_state)]
3469 #[tokio::test]
3470 async fn clear_cooldown_resets_failure_counts() {
3471 super::reset_global_state().await;
3472 let store = test_store().await;
3473 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3474
3475 let agent = "test_clear_fc_agent";
3476
3477 for _ in 0..5 {
3479 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3480 }
3481 set_agent_cooldown(agent, 3600).await;
3483
3484 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3486 let count: u32 = store
3487 .kv_get(&kv_key)
3488 .await
3489 .unwrap()
3490 .unwrap()
3491 .parse()
3492 .unwrap();
3493 assert_eq!(count, 5);
3494
3495 clear_cooldown(agent, &store).await;
3497
3498 let count_after: u32 = store
3500 .kv_get(&kv_key)
3501 .await
3502 .unwrap()
3503 .unwrap()
3504 .parse()
3505 .unwrap();
3506 assert_eq!(
3507 count_after, 0,
3508 "clear_cooldown should reset failure count to 0"
3509 );
3510
3511 assert!(
3513 !is_agent_in_cooldown(agent),
3514 "agent should not be in cooldown after clear"
3515 );
3516 }
3517
3518 #[serial(cooldown_state)]
3519 #[tokio::test]
3520 async fn clear_all_cooldowns_resets_all_failure_counts() {
3521 super::reset_global_state().await;
3522 let store = test_store().await;
3523 *cooldown_store().lock().await = Some(std::sync::Arc::clone(&store));
3524
3525 let agents = ["test_clear_all_a", "test_clear_all_b"];
3526 for agent in &agents {
3527 let _ = read_and_increment_failure_count(Some(&store), agent).await;
3528 set_agent_cooldown(agent, 3600).await;
3529 }
3530
3531 clear_cooldown("*", &store).await;
3532
3533 for agent in &agents {
3534 assert!(!is_agent_in_cooldown(agent));
3535 let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
3536 let count: u32 = store
3537 .kv_get(&kv_key)
3538 .await
3539 .unwrap()
3540 .unwrap()
3541 .parse()
3542 .unwrap();
3543 assert_eq!(
3544 count, 0,
3545 "clear --all should reset failure count for {agent}"
3546 );
3547 }
3548 }
3549
3550 #[serial(cooldown_state)]
3551 #[test]
3552 fn compute_backoff_saturates_on_large_counts() {
3553 assert_eq!(
3555 compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
3556 BACKOFF_MAX_SECS * 12
3557 );
3558 assert_eq!(
3559 compute_backoff(u64::from(u32::MAX), 300, 14400, true),
3560 172_800
3561 );
3562 assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172_800);
3563 }
3564
3565 #[serial(cooldown_state)]
3566 #[test]
3567 fn compute_backoff_zero_base_returns_zero() {
3568 assert_eq!(compute_backoff(1, 0, 14400, true), 0);
3569 assert_eq!(compute_backoff(5, 0, 14400, true), 0);
3570 }
3571
3572 #[serial(cooldown_state)]
3575 #[tokio::test]
3576 async fn github_circuit_stays_closed_below_threshold() {
3577 super::reset_global_state().await;
3578 {
3580 let mut open = github_5xx_circuit_open().lock().unwrap();
3581 *open = None;
3582 let mut ts = github_5xx_timestamps().lock().unwrap();
3583 ts.clear();
3584 let mut map = cooldown_lock();
3588 map.remove("github:5xx");
3589 }
3590
3591 for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
3593 record_github_5xx().await;
3594 }
3595
3596 assert!(
3597 !is_github_circuit_open(),
3598 "circuit should stay closed below threshold"
3599 );
3600 assert_eq!(github_circuit_remaining_secs(), 0);
3601 }
3602
3603 #[serial(cooldown_state)]
3604 #[tokio::test]
3605 async fn github_circuit_trips_at_threshold() {
3606 super::reset_global_state().await;
3607 {
3609 let mut open = github_5xx_circuit_open_lock();
3610 *open = None;
3611 let mut ts = github_5xx_timestamps_lock();
3612 ts.clear();
3613 let mut map = cooldown_lock();
3614 map.remove("github:5xx");
3615 }
3616
3617 for _ in 0..GITHUB_5XX_THRESHOLD {
3619 record_github_5xx().await;
3620 }
3621
3622 assert!(is_github_circuit_open(), "circuit should trip at threshold");
3623 assert!(
3624 github_circuit_remaining_secs() > 0,
3625 "remaining should be > 0 when open"
3626 );
3627 assert!(
3628 github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
3629 "remaining should not exceed cooldown duration"
3630 );
3631
3632 {
3634 let mut open = github_5xx_circuit_open().lock().unwrap();
3635 *open = None;
3636 }
3637 }
3638
3639 #[serial(cooldown_state)]
3640 #[tokio::test]
3641 async fn github_circuit_cannot_double_trip() {
3642 super::reset_global_state().await;
3643 {
3645 let mut open = github_5xx_circuit_open_lock();
3646 *open = None;
3647 let mut ts = github_5xx_timestamps_lock();
3648 ts.clear();
3649 let mut map = cooldown_lock();
3650 map.remove("github:5xx");
3651 }
3652
3653 for _ in 0..GITHUB_5XX_THRESHOLD {
3655 record_github_5xx().await;
3656 }
3657 assert!(is_github_circuit_open());
3658
3659 let first_remaining = github_circuit_remaining_secs();
3660
3661 for _ in 0..GITHUB_5XX_THRESHOLD {
3663 record_github_5xx().await;
3664 }
3665
3666 let second_remaining = github_circuit_remaining_secs();
3667 assert!(
3668 second_remaining <= first_remaining,
3669 "additional errors should not extend cooldown"
3670 );
3671
3672 {
3674 let mut open = github_5xx_circuit_open().lock().unwrap();
3675 *open = None;
3676 }
3677 }
3678
3679 #[serial(cooldown_state)]
3680 #[tokio::test]
3681 async fn github_circuit_recovers_after_cooldown() {
3682 super::reset_global_state().await;
3683 {
3685 let mut open = github_5xx_circuit_open_lock();
3686 *open = None;
3687 let mut ts = github_5xx_timestamps_lock();
3688 ts.clear();
3689 let mut map = cooldown_lock();
3690 map.remove("github:5xx");
3691 }
3692
3693 {
3695 let mut open = github_5xx_circuit_open().lock().unwrap();
3696 *open = Some(chrono::Utc::now().timestamp() - 1);
3697 }
3698
3699 assert!(
3701 !is_github_circuit_open(),
3702 "circuit should auto-recover after cooldown"
3703 );
3704 }
3705
3706 #[serial(cooldown_state)]
3707 #[tokio::test]
3708 async fn github_circuit_sliding_window_ages_out() {
3709 super::reset_global_state().await;
3710 {
3712 let mut open = github_5xx_circuit_open_lock();
3713 *open = None;
3714 let mut ts = github_5xx_timestamps_lock();
3715 ts.clear();
3716 let mut map = cooldown_lock();
3717 map.remove("github:5xx");
3718 }
3719
3720 {
3722 let mut ts = github_5xx_timestamps().lock().unwrap();
3723 let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
3724 for _ in 0..GITHUB_5XX_THRESHOLD {
3725 ts.push(old);
3726 }
3727 }
3728
3729 record_github_5xx().await;
3731
3732 assert!(
3733 !is_github_circuit_open(),
3734 "old errors outside window should not trip circuit"
3735 );
3736 }
3737}