use chrono::TimeZone;
use once_cell::sync::Lazy;
use regex::Regex;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
pub const AGENT_COOLDOWN_SECS: i64 = 30 * 60;
pub const BACKOFF_BASE_SECS: i64 = 5 * 60;
pub const BACKOFF_MAX_SECS: i64 = 4 * 60 * 60;
pub const PERSISTENT_MODEL_BACKOFF_BASE_SECS: i64 = 4 * 60 * 60;
pub const PERSISTENT_MODEL_BACKOFF_MAX_SECS: i64 = 7 * 24 * 60 * 60;
pub const CREDIT_BACKOFF_BASE_SECS: i64 = 60 * 60;
pub const ORG_BACKOFF_BASE_SECS: i64 = 2 * 60 * 60;
pub const CREDIT_BACKOFF_MAX_SECS: i64 = 8 * 60 * 60;
pub const BILLING_CYCLE_COOLDOWN_SECS: i64 = 24 * 60 * 60;
pub const BILLING_CYCLE_MAX_SECS: i64 = 7 * 24 * 60 * 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CreditExhaustionReason {
OutOfCredits,
OrgLevelDisabled,
BillingCycleExhausted,
}
impl CreditExhaustionReason {
pub fn outcome_str(self) -> &'static str {
match self {
CreditExhaustionReason::OutOfCredits => "credit_exhausted",
CreditExhaustionReason::OrgLevelDisabled => "org_level_disabled",
CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
}
}
}
pub const SILENCE_AGENT_COOLDOWN_SECS: u64 = 120;
pub const SILENCE_COUNT_THRESHOLD: usize = 5;
pub const SILENCE_COUNT_WINDOW_SECS: i64 = 24 * 60 * 60;
pub const SILENCE_EXTENDED_COOLDOWN_SECS: u64 = 4 * 60 * 60;
const KV_PREFIX: &str = "cooldown:";
const SILENCE_COUNT_PREFIX: &str = "silence_count:";
const FAILURE_COUNT_PREFIX: &str = "failure_count:";
const CREDIT_FAILURE_COUNT_PREFIX: &str = "credit_failure_count:";
const PERSISTENT_FAILURE_COUNT_PREFIX: &str = "persistent_failure_count:";
struct CooldownEntry {
cooldown_until: i64,
reason: String,
}
pub struct SilenceCountResult {
pub count: usize,
pub extended_cooldown_applied: bool,
}
fn cooldowns() -> &'static Mutex<HashMap<String, CooldownEntry>> {
static COOLDOWNS: OnceLock<Mutex<HashMap<String, CooldownEntry>>> = OnceLock::new();
COOLDOWNS.get_or_init(|| Mutex::new(HashMap::new()))
}
fn cooldown_store() -> &'static tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>> {
static STORE: OnceLock<tokio::sync::Mutex<Option<Arc<crate::store::TaskStore>>>> =
OnceLock::new();
STORE.get_or_init(|| tokio::sync::Mutex::new(None))
}
fn parse_cooldown_value(value: &str) -> Option<(i64, String)> {
if let Ok(serde_json::Value::Object(map)) = serde_json::from_str(value) {
let until = map.get("until")?.as_i64()?;
let reason = map
.get("reason")
.and_then(|r| r.as_str())
.unwrap_or("persisted")
.to_string();
return Some((until, reason));
}
value
.parse::<i64>()
.ok()
.map(|until| (until, "persisted".to_string()))
}
pub async fn init_cooldown_store(store: Arc<crate::store::TaskStore>) {
match store.kv_list_prefix(KV_PREFIX).await {
Ok(rows) => {
let now = chrono::Utc::now().timestamp();
let mut map = cooldown_lock();
let mut loaded = 0;
for (key, value) in rows {
let cooldown_key = key.trim_start_matches(KV_PREFIX);
if let Some((cooldown_until, reason)) = parse_cooldown_value(&value) {
if now < cooldown_until {
map.insert(
cooldown_key.to_string(),
CooldownEntry {
cooldown_until,
reason,
},
);
loaded += 1;
}
}
}
if loaded > 0 {
tracing::info!(loaded, "loaded persisted cooldowns from KV store");
}
}
Err(e) => {
tracing::warn!(err = %e, "failed to load cooldowns from KV store");
}
}
if let Ok(Some(raw)) = store.kv_get("cooldown:github:5xx").await {
if let Ok(cooldown_until) = raw.parse::<i64>() {
let now = chrono::Utc::now().timestamp();
if now < cooldown_until {
let mut open = github_5xx_circuit_open_lock();
*open = Some(cooldown_until);
let remaining = cooldown_until - now;
tracing::warn!(
remaining_secs = remaining,
"GitHub 5xx circuit breaker restored from KV (still open)"
);
}
}
}
*cooldown_store().lock().await = Some(store);
}
pub async fn sync_from_kv(store: &Arc<crate::store::TaskStore>) {
let rows = match store.kv_list_prefix(KV_PREFIX).await {
Ok(rows) => rows,
Err(e) => {
tracing::warn!(err = %e, "cooldown sync_from_kv: KV list failed");
return;
}
};
let now = chrono::Utc::now().timestamp();
let mut explicit_kv: HashMap<String, Option<(i64, String)>> =
HashMap::with_capacity(rows.len());
for (kv_key, value) in &rows {
let cooldown_key = kv_key.trim_start_matches(KV_PREFIX);
if let Some((ts, reason)) = parse_cooldown_value(value) {
let normalized = if ts > now { Some((ts, reason)) } else { None };
explicit_kv.insert(cooldown_key.to_string(), normalized);
}
}
let mut map = cooldown_lock();
let mut dropped = 0usize;
let mut updated = 0usize;
map.retain(|key, _| {
if matches!(explicit_kv.get(key), Some(None)) {
dropped += 1;
false
} else {
true
}
});
for (key, normalized) in explicit_kv {
if let Some((ts, reason)) = normalized {
match map.get_mut(&key) {
Some(entry) => {
if entry.cooldown_until != ts {
entry.cooldown_until = ts;
updated += 1;
}
}
None => {
map.insert(
key,
CooldownEntry {
cooldown_until: ts,
reason,
},
);
updated += 1;
}
}
}
}
if dropped > 0 || updated > 0 {
tracing::info!(
dropped,
updated,
"cooldown sync_from_kv: reconciled in-memory map with KV"
);
}
}
const EXTENDED_TIER_1_THRESHOLD: u64 = 2;
const EXTENDED_TIER_1_MULTIPLIER: i64 = 6;
const EXTENDED_TIER_2_THRESHOLD: u64 = 5;
const EXTENDED_TIER_2_MULTIPLIER: i64 = 12;
pub fn compute_backoff(count: u64, base: i64, max: i64, allow_extended: bool) -> i64 {
if count <= 1 {
return base;
}
let exp = count.saturating_sub(1).min(u32::MAX as u64) as u32;
let factor = 3_i64.saturating_pow(exp);
let standard = base.saturating_mul(factor).min(max);
let first_max = if base <= 0 {
0
} else {
let mut accumulated = base;
let mut first_max_count = 1u64;
while accumulated < max && first_max_count < count {
accumulated = accumulated.saturating_mul(3);
first_max_count += 1;
}
first_max_count
};
let times_at_max = count.saturating_sub(first_max);
if !allow_extended {
return standard;
}
if standard == max && times_at_max >= EXTENDED_TIER_2_THRESHOLD {
max.saturating_mul(EXTENDED_TIER_2_MULTIPLIER)
} else if standard == max && times_at_max >= EXTENDED_TIER_1_THRESHOLD {
max.saturating_mul(EXTENDED_TIER_1_MULTIPLIER)
} else {
standard
}
}
async fn read_and_increment_failure_count(
store_opt: &Option<Arc<crate::store::TaskStore>>,
key: &str,
) -> u64 {
read_and_increment_failure_count_with_prefix(store_opt, FAILURE_COUNT_PREFIX, key).await
}
async fn read_and_increment_failure_count_with_prefix(
store_opt: &Option<Arc<crate::store::TaskStore>>,
prefix: &str,
key: &str,
) -> u64 {
let kv_key = format!("{prefix}{key}");
if let Some(store) = store_opt {
match store.kv_increment(&kv_key).await {
Ok(n) => n,
Err(e) => {
tracing::warn!(key = %kv_key, err = %e, "failed to increment failure count — treating as high count for safe backoff");
u64::MAX
}
}
} else {
1
}
}
pub async fn record_agent_success(agent_name: &str, model: &str) {
let store_opt = cooldown_store().lock().await.clone();
if let Some(store) = store_opt {
let agent_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}");
let credit_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{agent_name}");
if let Err(e) = store.kv_set(&agent_key, "0").await {
tracing::warn!(key = agent_key, err = %e, "failed to reset failure count");
}
if !model.is_empty() {
let model_key = format!("{FAILURE_COUNT_PREFIX}{agent_name}:{model}");
if let Err(e) = store.kv_set(&model_key, "0").await {
tracing::warn!(key = model_key, err = %e, "failed to reset failure count");
}
}
if let Err(e) = store.kv_set(&credit_key, "0").await {
tracing::warn!(key = credit_key, err = %e, "failed to reset failure count");
}
}
}
pub async fn record_agent_failure_with_message(agent_name: &str, error_message: &str) {
if let Some(cooldown_until) = parse_retry_at(error_message) {
let now = chrono::Utc::now().timestamp();
if cooldown_until > now {
set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
return;
}
tracing::debug!(
agent = agent_name,
cooldown_until,
"vendor retry_at is in the past — using backoff instead"
);
}
let store_opt = cooldown_store().lock().await.clone();
let count = read_and_increment_failure_count(&store_opt, agent_name).await;
let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
let cooldown_until =
chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
set_cooldown_async(agent_name, cooldown_until, "agent_error").await;
}
pub fn detect_credit_exhaustion_outcome(error_message: &str) -> Option<&'static str> {
detect_credit_exhaustion(error_message).map(CreditExhaustionReason::outcome_str)
}
pub fn detect_credit_exhaustion(error_message: &str) -> Option<CreditExhaustionReason> {
let lower = error_message.to_lowercase();
let overage_patterns = [
"overagedisabledreason",
"overage_disabled_reason",
"overage disabled reason",
];
if overage_patterns.iter().any(|p| lower.contains(p)) {
if lower.contains("out_of_credits") || lower.contains("outofcredits") {
return Some(CreditExhaustionReason::OutOfCredits);
}
if lower.contains("org_level_disabled") || lower.contains("orgdisabled") {
return Some(CreditExhaustionReason::OrgLevelDisabled);
}
}
let billing_patterns = [
"credit balance too low",
"credit balance insufficient",
"insufficient credit balance",
"billing quota exceeded",
"billing limit exceeded",
"organization has been disabled",
"org disabled",
];
if billing_patterns.iter().any(|p| lower.contains(p)) {
return Some(CreditExhaustionReason::OrgLevelDisabled);
}
let billing_cycle_patterns = [
"billing cycle",
"quota refreshed next cycle",
"monthly quota",
"quota will be refreshed",
"next billing cycle",
"refreshed in the next cycle",
"weekly/monthly limit exhausted",
"limit will reset at",
"upgrade to pro", "upgrade your token plan", "token plan usage limit", "purchase credits for more", ];
if billing_cycle_patterns.iter().any(|p| lower.contains(p)) {
return Some(CreditExhaustionReason::BillingCycleExhausted);
}
if lower.contains("out of credits")
|| lower.contains("outofcredits")
|| lower.contains("insufficient funds")
|| lower.contains("no credits remaining")
|| lower.contains("insufficient_quota")
|| lower.contains("quota exceeded")
|| lower.contains("insufficient balance")
|| lower.contains("please recharge")
|| lower.contains("no resource package")
|| lower.contains("balance exhausted")
|| lower.contains("402 payment required")
{
return Some(CreditExhaustionReason::OutOfCredits);
}
None
}
pub async fn record_credit_exhaustion(agent_name: &str, reason: CreditExhaustionReason) {
let reason_str = match reason {
CreditExhaustionReason::OutOfCredits => "credit_exhaustion_out_of_credits",
CreditExhaustionReason::OrgLevelDisabled => "credit_exhaustion_org_level_disabled",
CreditExhaustionReason::BillingCycleExhausted => "billing_cycle_exhausted",
};
let store_opt = cooldown_store().lock().await.clone();
let count = read_and_increment_failure_count_with_prefix(
&store_opt,
CREDIT_FAILURE_COUNT_PREFIX,
agent_name,
)
.await;
let (base, max) = match reason {
CreditExhaustionReason::OutOfCredits => (CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS),
CreditExhaustionReason::OrgLevelDisabled => {
(ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS)
}
CreditExhaustionReason::BillingCycleExhausted => {
(BILLING_CYCLE_COOLDOWN_SECS, BILLING_CYCLE_MAX_SECS)
}
};
let cooldown_secs = compute_backoff(
count,
base,
max,
!matches!(reason, CreditExhaustionReason::BillingCycleExhausted),
);
let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
set_cooldown_async(agent_name, cooldown_until, reason_str).await;
tracing::warn!(
agent = agent_name,
reason = reason_str,
attempt = count,
cooldown_secs,
"credit exhaustion detected: applying exponential agent-wide cooldown"
);
}
pub async fn record_model_failure(agent_name: &str, model: &str) {
let key = format!("{agent_name}:{model}");
let store_opt = cooldown_store().lock().await.clone();
let count = read_and_increment_failure_count(&store_opt, &key).await;
let base = crate::engine::router::config::get_agent_backoff_base(agent_name);
let cooldown_until =
chrono::Utc::now().timestamp() + compute_backoff(count, base, BACKOFF_MAX_SECS, true);
set_cooldown_async(&key, cooldown_until, "model_error").await;
}
pub async fn record_rate_limit(agent_name: &str, model: Option<&str>, error_message: &str) {
if let Some(window_secs) = parse_relative_usage_window(error_message) {
if let Some(m) = model {
let cooldown_until = chrono::Utc::now().timestamp() + window_secs as i64;
set_model_cooldown(agent_name, m, window_secs, "rate_limit_usage_window").await;
tracing::info!(
agent = agent_name,
model = m,
window_secs,
cooldown_until,
"relative usage window detected: applying model-specific cooldown"
);
return;
}
}
record_agent_failure_with_message(agent_name, error_message).await;
if let Some(m) = model {
record_model_failure(agent_name, m).await;
}
}
pub async fn record_persistent_model_failure(agent_name: &str, model: &str) {
let key = format!("{agent_name}:{model}");
let store_opt = cooldown_store().lock().await.clone();
let count = read_and_increment_failure_count_with_prefix(
&store_opt,
PERSISTENT_FAILURE_COUNT_PREFIX,
&key,
)
.await;
let cooldown_secs = compute_backoff(
count,
PERSISTENT_MODEL_BACKOFF_BASE_SECS,
PERSISTENT_MODEL_BACKOFF_MAX_SECS,
false,
);
let cooldown_until = chrono::Utc::now().timestamp() + cooldown_secs;
set_cooldown_async(&key, cooldown_until, "persistent_model_error").await;
}
pub async fn set_model_cooldown(agent_name: &str, model: &str, duration_secs: u64, reason: &str) {
let key = format!("{agent_name}:{model}");
let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
set_cooldown_async(&key, cooldown_until, reason).await;
}
pub async fn set_model_cooldown_until(agent_name: &str, model: &str, until: i64) {
let key = format!("{agent_name}:{model}");
set_cooldown_async(&key, until, "vendor_reset_at").await;
}
pub async fn set_agent_cooldown(agent_name: &str, duration_secs: u64) {
let cooldown_until = chrono::Utc::now().timestamp() + duration_secs as i64;
set_cooldown_async(agent_name, cooldown_until, "silence_agent_cooldown").await;
}
pub async fn record_silence_detection(agent_name: &str, model: &str) -> Option<SilenceCountResult> {
let store_opt = cooldown_store().lock().await.clone();
let store = match store_opt {
Some(store) => store,
None => {
tracing::debug!(
agent = agent_name,
model,
"skipping silence count record (no KV store)"
);
return None;
}
};
let key = format!("{SILENCE_COUNT_PREFIX}{agent_name}:{model}");
let now = chrono::Utc::now().timestamp();
let window_start = now - SILENCE_COUNT_WINDOW_SECS;
let mut timestamps = match store.kv_get(&key).await {
Ok(Some(raw)) => match serde_json::from_str::<Vec<i64>>(&raw) {
Ok(ts) => ts,
Err(e) => {
tracing::warn!(
kv_key = %key,
err = %e,
"failed to parse silence timestamps from KV — resetting window"
);
Vec::new()
}
},
Ok(None) => Vec::new(),
Err(err) => {
tracing::warn!(
kv_key = %key,
err = %err,
"failed to load silence count from KV"
);
Vec::new()
}
};
timestamps.retain(|ts| *ts >= window_start);
timestamps.push(now);
let count = timestamps.len();
let mut extended_cooldown_applied = false;
if count >= SILENCE_COUNT_THRESHOLD {
set_model_cooldown(
agent_name,
model,
SILENCE_EXTENDED_COOLDOWN_SECS,
"silence_detected",
)
.await;
extended_cooldown_applied = true;
timestamps.clear();
}
let value = serde_json::to_string(×tamps).unwrap_or_else(|_| "[]".to_string());
if let Err(err) = store.kv_set(&key, &value).await {
tracing::warn!(kv_key = key, err = %err, "failed to persist silence count");
}
Some(SilenceCountResult {
count,
extended_cooldown_applied,
})
}
pub fn is_model_in_cooldown(agent_name: &str, model: &str) -> bool {
let key = format!("{agent_name}:{model}");
is_in_cooldown(&key)
}
pub fn is_agent_in_cooldown(agent_name: &str) -> bool {
is_in_cooldown(agent_name)
}
pub fn cooldown_until(key: &str) -> Option<i64> {
let map = cooldown_lock();
map.get(key).and_then(|entry| {
let now = chrono::Utc::now().timestamp();
if now < entry.cooldown_until {
Some(entry.cooldown_until)
} else {
None
}
})
}
pub fn cooldown_reason(key: &str) -> Option<String> {
let map = cooldown_lock();
map.get(key).and_then(|entry| {
let now = chrono::Utc::now().timestamp();
if now < entry.cooldown_until {
Some(entry.reason.clone())
} else {
None
}
})
}
pub fn list_all_cooldowns() -> Vec<(String, i64, String)> {
let now = chrono::Utc::now().timestamp();
let map = cooldown_lock();
let mut result: Vec<(String, i64, String)> = map
.iter()
.filter(|(_, entry)| now < entry.cooldown_until)
.map(|(key, entry)| (key.clone(), entry.cooldown_until, entry.reason.clone()))
.collect();
result.sort_by(|a, b| a.0.cmp(&b.0));
result
}
pub async fn clear_cooldown(key: &str, store: &Arc<crate::store::TaskStore>) {
if key == "*" {
let keys: Vec<String> = {
let mut map = cooldown_lock();
let keys: Vec<String> = map.keys().cloned().collect();
map.clear();
keys
};
for k in &keys {
let kv_key = format!("{KV_PREFIX}{k}");
if let Err(e) = store.kv_set(&kv_key, "0").await {
tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
}
let fc_key = format!("{FAILURE_COUNT_PREFIX}{k}");
if let Err(e) = store.kv_set(&fc_key, "0").await {
tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
}
let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{k}");
if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
}
let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{k}");
if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
}
}
if let Ok(fc_entries) = store.kv_list_prefix(FAILURE_COUNT_PREFIX).await {
for (fc_key, _) in fc_entries {
if let Err(e) = store.kv_set(&fc_key, "0").await {
tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
}
}
}
if let Ok(cfc_entries) = store.kv_list_prefix(CREDIT_FAILURE_COUNT_PREFIX).await {
for (cfc_key, _) in cfc_entries {
if let Err(e) = store.kv_set(&cfc_key, "0").await {
tracing::warn!(key = cfc_key, err = %e, "failed to reset credit failure count");
}
}
}
if let Ok(pfc_entries) = store.kv_list_prefix(PERSISTENT_FAILURE_COUNT_PREFIX).await {
for (pfc_key, _) in pfc_entries {
if let Err(e) = store.kv_set(&pfc_key, "0").await {
tracing::warn!(key = pfc_key, err = %e, "failed to reset persistent failure count");
}
}
}
tracing::info!(
count = keys.len(),
"cleared all cooldowns and failure counts"
);
} else {
{
let mut map = cooldown_lock();
map.remove(key);
}
let kv_key = format!("{KV_PREFIX}{key}");
if let Err(e) = store.kv_set(&kv_key, "0").await {
tracing::warn!(key = kv_key, err = %e, "failed to clear cooldown");
}
let fc_key = format!("{FAILURE_COUNT_PREFIX}{key}");
if let Err(e) = store.kv_set(&fc_key, "0").await {
tracing::warn!(key = fc_key, err = %e, "failed to reset failure count");
}
let credit_fc_key = format!("{CREDIT_FAILURE_COUNT_PREFIX}{key}");
if let Err(e) = store.kv_set(&credit_fc_key, "0").await {
tracing::warn!(key = credit_fc_key, err = %e, "failed to reset credit failure count");
}
let persistent_fc_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{key}");
if let Err(e) = store.kv_set(&persistent_fc_key, "0").await {
tracing::warn!(key = persistent_fc_key, err = %e, "failed to reset persistent failure count");
}
tracing::info!(key, "cleared cooldown and failure count");
}
}
fn set_cooldown_in_memory(key: &str, cooldown_until: i64, reason: &str) -> bool {
let mut map = cooldown_lock();
if let Some(existing) = map.get(key) {
if existing.cooldown_until >= cooldown_until {
return false;
}
}
map.insert(
key.to_string(),
CooldownEntry {
cooldown_until,
reason: reason.to_string(),
},
);
true
}
async fn set_cooldown_async(key: &str, cooldown_until: i64, reason: &str) -> bool {
if !set_cooldown_in_memory(key, cooldown_until, reason) {
return false;
}
let store_opt = cooldown_store().lock().await.clone();
if let Some(store) = store_opt {
let kv_key = format!("{KV_PREFIX}{key}");
let value = serde_json::json!({"until": cooldown_until, "reason": reason}).to_string();
if let Err(e) = store.kv_set(&kv_key, &value).await {
tracing::warn!(
kv_key,
err = %e,
"failed to persist cooldown to KV store — rolling back in-memory entry"
);
let mut map = cooldown_lock();
if map.get(key).map(|e| e.cooldown_until) == Some(cooldown_until) {
map.remove(key);
}
return false;
}
}
true
}
fn cooldown_lock() -> std::sync::MutexGuard<'static, HashMap<String, CooldownEntry>> {
cooldowns().lock().unwrap_or_else(|e| {
tracing::error!("cooldown map lock poisoned — recovering with partial state; backoff and degraded-agent decisions may be incorrect");
e.into_inner()
})
}
fn degraded_agents_lock() -> std::sync::MutexGuard<'static, std::collections::HashSet<String>> {
degraded_agents().lock().unwrap_or_else(|e| {
tracing::error!("degraded agents lock poisoned — recovering with partial state; health check decisions may be incorrect");
e.into_inner()
})
}
fn github_5xx_timestamps_lock() -> std::sync::MutexGuard<'static, Vec<i64>> {
github_5xx_timestamps().lock().unwrap_or_else(|e| {
tracing::error!("github 5xx timestamps lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
e.into_inner()
})
}
fn github_5xx_circuit_open_lock() -> std::sync::MutexGuard<'static, Option<i64>> {
github_5xx_circuit_open().lock().unwrap_or_else(|e| {
tracing::error!("github 5xx circuit open lock poisoned — recovering with partial state; circuit breaker may operate incorrectly");
e.into_inner()
})
}
fn is_in_cooldown(key: &str) -> bool {
let map = cooldown_lock();
if let Some(entry) = map.get(key) {
let now = chrono::Utc::now().timestamp();
return now < entry.cooldown_until;
}
false
}
pub const DEFAULT_HEALTH_CHECK_WINDOW_HOURS: u32 = 6;
pub const DEFAULT_DEGRADED_RATE_LIMIT_THRESHOLD: i64 = 3;
fn degraded_agents() -> &'static Mutex<std::collections::HashSet<String>> {
static DEGRADED: OnceLock<Mutex<std::collections::HashSet<String>>> = OnceLock::new();
DEGRADED.get_or_init(|| Mutex::new(std::collections::HashSet::new()))
}
pub fn is_agent_degraded(agent: &str) -> bool {
let set = degraded_agents_lock();
set.contains(agent)
}
pub fn mark_agent_degraded(agent: &str) {
let mut set = degraded_agents_lock();
set.insert(agent.to_string());
}
pub fn clear_agent_degraded(agent: &str) {
let mut set = degraded_agents_lock();
set.remove(agent);
}
pub async fn refresh_degraded_agents(
store: &Arc<crate::store::TaskStore>,
available_agents: &[String],
model_checker: &(dyn Fn(&str) -> bool + Send + Sync),
window_hours: u32,
threshold: i64,
) {
clear_expired_cooldowns();
let start = chrono::Utc::now();
let counts = match store.recent_rate_limit_counts(window_hours).await {
Ok(c) => {
let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
tracing::info!(
duration_ms = dur_ms,
window_hours,
"recent_rate_limit_counts query completed"
);
c
}
Err(e) => {
let dur_ms = (chrono::Utc::now() - start).num_milliseconds();
tracing::warn!(
err = %e,
duration_ms = dur_ms,
"failed to query recent rate limit counts for health check"
);
return;
}
};
for agent in available_agents {
let in_cooldown = is_agent_in_cooldown(agent);
let no_models = !model_checker(agent);
let rate_limit_count = counts.get(agent.as_str()).copied().unwrap_or(0);
let over_threshold = rate_limit_count >= threshold;
let degraded = in_cooldown || no_models || over_threshold;
if degraded {
if !is_agent_degraded(agent) {
let reason = if in_cooldown {
"agent in cooldown"
} else if no_models {
"all models cooled"
} else {
"rate limit threshold exceeded"
};
tracing::warn!(
agent,
reason,
rate_limit_count,
window_hours,
threshold,
"pre-emptive health check: marking agent as degraded"
);
}
mark_agent_degraded(agent);
} else {
if is_agent_degraded(agent) {
tracing::info!(
agent,
rate_limit_count,
"pre-emptive health check: agent recovered, clearing degraded flag"
);
}
clear_agent_degraded(agent);
}
}
}
pub fn clear_expired_cooldowns() {
let mut map = cooldown_lock();
let now = chrono::Utc::now().timestamp();
map.retain(|_key, entry| now < entry.cooldown_until);
}
const GITHUB_5XX_WINDOW_SECS: i64 = 120;
const GITHUB_5XX_THRESHOLD: usize = 5;
const GITHUB_5XX_COOLDOWN_SECS: i64 = 180;
fn github_5xx_timestamps() -> &'static Mutex<Vec<i64>> {
static TS: OnceLock<Mutex<Vec<i64>>> = OnceLock::new();
TS.get_or_init(|| Mutex::new(Vec::new()))
}
fn github_5xx_circuit_open() -> &'static Mutex<Option<i64>> {
static OPEN: OnceLock<Mutex<Option<i64>>> = OnceLock::new();
OPEN.get_or_init(|| Mutex::new(None))
}
pub async fn record_github_5xx() {
let now = chrono::Utc::now().timestamp();
let window_start = now - GITHUB_5XX_WINDOW_SECS;
let should_trip = {
let mut ts = github_5xx_timestamps_lock();
ts.retain(|t| *t >= window_start);
ts.push(now);
ts.len() >= GITHUB_5XX_THRESHOLD
};
if should_trip {
let cooldown_until = {
let mut open = github_5xx_circuit_open_lock();
if open.is_none() {
let cd = now + GITHUB_5XX_COOLDOWN_SECS;
*open = Some(cd);
tracing::warn!(
errors_in_window = GITHUB_5XX_THRESHOLD,
window_secs = GITHUB_5XX_WINDOW_SECS,
cooldown_secs = GITHUB_5XX_COOLDOWN_SECS,
"GitHub 5xx circuit breaker OPEN — throttling non-critical work"
);
Some(cd)
} else {
None
}
};
if let Some(cd) = cooldown_until {
let store_opt = cooldown_store().lock().await.clone();
if let Some(store) = store_opt {
let _ = store.kv_set("cooldown:github:5xx", &cd.to_string()).await;
}
set_agent_cooldown("github:5xx", GITHUB_5XX_COOLDOWN_SECS as u64).await;
}
let mut ts = github_5xx_timestamps_lock();
ts.clear();
}
}
pub fn is_github_circuit_open() -> bool {
{
let mut open = github_5xx_circuit_open_lock();
if let Some(until) = *open {
let now = chrono::Utc::now().timestamp();
if now >= until {
*open = None;
tracing::info!("GitHub 5xx circuit breaker CLOSED — resuming normal operations");
return false;
}
return true;
}
}
if let Some(until) = cooldown_until("github:5xx") {
let mut open = github_5xx_circuit_open_lock();
if open.is_none() {
*open = Some(until);
let now = chrono::Utc::now().timestamp();
if now < until {
let remaining = until - now;
tracing::warn!(
remaining_secs = remaining,
"GitHub 5xx circuit breaker OPEN (discovered via generic cooldown) — throttling non-critical work"
);
}
}
return true;
}
false
}
pub fn github_circuit_remaining_secs() -> u64 {
{
let open = github_5xx_circuit_open_lock();
if let Some(until) = *open {
let now = chrono::Utc::now().timestamp();
if now < until {
return (until - now) as u64;
}
return 0;
}
}
if let Some(until) = cooldown_until("github:5xx") {
let now = chrono::Utc::now().timestamp();
if now < until {
return (until - now) as u64;
}
}
0
}
pub(crate) fn parse_retry_at(error_message: &str) -> Option<i64> {
if error_message.is_empty() {
return None;
}
let lower = error_message.to_lowercase();
for retry_marker in [
"try again at ",
"reset at ",
"resets at ",
"resets ",
"reset ",
] {
if let Some(idx) = lower.find(retry_marker) {
let date_str = &lower[idx + retry_marker.len()..];
let date_str = date_str
.split(['.', '\n'])
.next()
.unwrap_or(date_str)
.trim();
let cleaned = date_str
.replace("st,", ",")
.replace("nd,", ",")
.replace("rd,", ",")
.replace("th,", ",");
for fmt in &[
"%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", ] {
if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(&cleaned, fmt) {
use chrono::TimeZone;
let utc = match chrono::Local.from_local_datetime(&dt) {
chrono::LocalResult::Single(local_dt) => {
local_dt.with_timezone(&chrono::Utc)
}
chrono::LocalResult::Ambiguous(earliest, _) => {
earliest.with_timezone(&chrono::Utc)
}
chrono::LocalResult::None => {
dt.and_utc()
}
};
tracing::info!(
retry_at = %utc,
raw = date_str,
marker = retry_marker,
"parsed rate limit retry-at date"
);
return Some(utc.timestamp());
}
}
tracing::debug!(
raw = date_str,
marker = retry_marker,
"could not parse retry-at date"
);
let time_only = date_str
.split_once('(')
.map(|(t, _)| t)
.unwrap_or(date_str)
.trim();
let mut compact = time_only.replace(" am", "am").replace(" pm", "pm");
let split = compact.rfind(' ').map_or(0, |i| i + 1);
let last = &compact[split..];
if !last.contains(':')
&& (last.ends_with("am") || last.ends_with("pm"))
&& last.len() > 2
{
let at = split + last.len() - 2;
compact.insert_str(at, ":00");
}
let md = compact.replace(" at ", " ");
let now_local = chrono::Local::now();
let this_year = chrono::Datelike::year(&now_local);
for year in [this_year, this_year + 1] {
let with_year = format!("{year} {md}");
if let Ok(dt) =
chrono::NaiveDateTime::parse_from_str(&with_year, "%Y %b %d %I:%M%p")
{
use chrono::TimeZone;
if let Some(local_dt) = chrono::Local.from_local_datetime(&dt).earliest() {
if local_dt > now_local {
return Some(local_dt.timestamp());
}
}
}
}
let time_only = compact.as_str();
for tf in &["%I:%M %p", "%I:%M%p", "%I %p", "%H:%M"] {
if let Ok(t) = chrono::NaiveTime::parse_from_str(time_only, tf) {
let local_now = chrono::Local::now();
let today = local_now.date_naive();
let candidate = chrono::NaiveDateTime::new(today, t);
let candidate_local: chrono::DateTime<chrono::Local> =
match chrono::Local.from_local_datetime(&candidate) {
chrono::LocalResult::Single(dt) => dt,
chrono::LocalResult::Ambiguous(earliest, _) => earliest,
chrono::LocalResult::None => {
let utc_dt = candidate.and_utc();
return Some(utc_dt.timestamp());
}
};
let chosen = if candidate_local <= local_now {
let tomorrow = today.succ_opt().unwrap_or(today);
let ndt = chrono::NaiveDateTime::new(tomorrow, t);
match chrono::Local.from_local_datetime(&ndt) {
chrono::LocalResult::Single(dt) => dt.with_timezone(&chrono::Utc),
chrono::LocalResult::Ambiguous(earliest, _) => {
earliest.with_timezone(&chrono::Utc)
}
chrono::LocalResult::None => ndt.and_utc(),
}
} else {
candidate_local.with_timezone(&chrono::Utc)
};
tracing::info!(
retry_at = %chosen,
raw = date_str,
marker = retry_marker,
"parsed rate limit retry-at time-only; interpreted as next occurrence"
);
return Some(chosen.timestamp());
}
}
}
}
None
}
pub(crate) fn parse_relative_usage_window(error_message: &str) -> Option<u64> {
if error_message.is_empty() {
return None;
}
let lower = error_message.to_lowercase();
let has_context = lower.contains("usage limit")
|| lower.contains("quota")
|| lower.contains("window")
|| lower.contains("limit");
if !has_context {
return None;
}
static WINDOW_RE: Lazy<Regex> = Lazy::new(|| {
Regex::new(
r"(?i)(?:^|\s|[\(\)\.,;:!?])(?:(\d+)\s*-?\s*(hour|hours|hr|hrs|minute|minutes|min|mins|day|days)|(hourly|daily|weekly|monthly))(?:$|\s|[\(\)\.,;:!?])"
)
.expect("usage window regex is valid")
});
let mut max_secs: Option<u64> = None;
for caps in WINDOW_RE.captures_iter(&lower) {
let secs = if let Some(num) = caps.get(1) {
let n: u64 = num.as_str().parse().ok()?;
let unit = caps.get(2)?.as_str().to_lowercase();
let multiplier = match unit.as_str() {
"hour" | "hours" | "hr" | "hrs" => 3600,
"minute" | "minutes" | "min" | "mins" => 60,
"day" | "days" => 86400,
_ => return None,
};
n.checked_mul(multiplier)?
} else if let Some(word) = caps.get(3) {
match word.as_str().to_lowercase().as_str() {
"hourly" => 3600,
"daily" => 86400,
"weekly" => 604800,
"monthly" => 2592000,
_ => return None,
}
} else {
continue;
};
max_secs = Some(max_secs.map_or(secs, |m| m.max(secs)));
}
max_secs
}
#[cfg(test)]
pub(crate) async fn reset_global_state() {
cooldowns().lock().unwrap().clear();
*cooldown_store().lock().await = None;
github_5xx_timestamps().lock().unwrap().clear();
*github_5xx_circuit_open().lock().unwrap() = None;
degraded_agents().lock().unwrap().clear();
crate::engine::runner::agents::opencode::reset_model_caches_for_test();
}
#[cfg(test)]
mod tests {
use super::*;
use serial_test::serial;
async fn test_store() -> Arc<crate::store::TaskStore> {
Arc::new(
crate::store::TaskStore::open_memory()
.await
.expect("in-memory store"),
)
}
#[serial(cooldown_state)]
#[tokio::test]
async fn init_loads_unexpired_cooldowns_from_kv() {
super::reset_global_state().await;
let store = test_store().await;
let cooldown_until = chrono::Utc::now().timestamp() + 3600;
store
.kv_set(
&format!("{KV_PREFIX}kimi:k2p5"),
&cooldown_until.to_string(),
)
.await
.unwrap();
let expired = chrono::Utc::now().timestamp() - 1;
store
.kv_set(
&format!("{KV_PREFIX}opencode:old-free"),
&expired.to_string(),
)
.await
.unwrap();
{
let mut map = cooldown_lock();
map.clear();
}
init_cooldown_store(store).await;
assert!(
is_model_in_cooldown("kimi", "k2p5"),
"unexpired cooldown should be loaded"
);
assert!(
!is_model_in_cooldown("opencode", "old-free"),
"expired cooldown should not be loaded"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_model_failure_writes_to_kv() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
record_model_failure("testagent_persist", "testmodel_persist").await;
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let kv_key = format!("{KV_PREFIX}testagent_persist:testmodel_persist");
let stored = store.kv_get(&kv_key).await.unwrap();
assert!(stored.is_some(), "model failure should be persisted to KV");
let (ts, reason) = parse_cooldown_value(&stored.unwrap()).expect("cooldown value");
let now = chrono::Utc::now().timestamp();
assert!(ts > now, "persisted timestamp should be in the future");
assert_eq!(reason, "model_error");
}
#[serial(cooldown_state)]
#[tokio::test]
async fn cooldown_reason_survives_restart_via_kv() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_reason_restart";
let future = chrono::Utc::now().timestamp() + 600;
init_cooldown_store(store.clone()).await;
set_cooldown_async(agent, future, "billing_cycle_exhausted").await;
{
let mut map = cooldown_lock();
map.clear();
}
init_cooldown_store(store).await;
let entries = list_all_cooldowns();
let entry = entries.iter().find(|(key, ..)| key == agent);
assert_eq!(
entry.map(|(_, _, reason)| reason.as_str()),
Some("billing_cycle_exhausted"),
"cooldown reason must survive a restart, not collapse to 'persisted'"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_persistent_model_failure_escalates_and_increments_failure_count() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "testagent_persistent";
let model = "testmodel_persistent";
let key = format!("{agent}:{model}");
record_persistent_model_failure(agent, model).await;
let now = chrono::Utc::now().timestamp();
let first_until = cooldown_until(&key).expect("first cooldown should exist");
let first_remaining = first_until.saturating_sub(now);
assert!(
first_remaining >= PERSISTENT_MODEL_BACKOFF_BASE_SECS - 5,
"first persistent model cooldown should be ~4h, got {first_remaining}s"
);
record_persistent_model_failure(agent, model).await;
let now2 = chrono::Utc::now().timestamp();
let second_until = cooldown_until(&key).expect("second cooldown should exist");
let second_remaining = second_until.saturating_sub(now2);
assert!(
second_remaining >= (PERSISTENT_MODEL_BACKOFF_BASE_SECS * 3) - 5,
"second persistent model cooldown should be ~12h, got {second_remaining}s"
);
let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
let stored_count = store
.kv_get(&count_key)
.await
.expect("kv read should succeed")
.expect("failure count should be written");
assert_eq!(stored_count, "2");
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_agent_success_does_not_reset_persistent_failure_count() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "testagent_persistent_success";
let model = "testmodel_persistent_success";
let key = format!("{agent}:{model}");
record_persistent_model_failure(agent, model).await;
record_persistent_model_failure(agent, model).await;
let count_key = format!("{PERSISTENT_FAILURE_COUNT_PREFIX}{agent}:{model}");
let stored_count = store
.kv_get(&count_key)
.await
.expect("kv read should succeed")
.expect("failure count should be written");
assert_eq!(stored_count, "2");
record_agent_success(agent, model).await;
let stored_count_after_success = store
.kv_get(&count_key)
.await
.expect("kv read should succeed")
.expect("persistent failure count should still exist");
assert_eq!(
stored_count_after_success, "2",
"persistent failure count must survive an intervening success"
);
record_persistent_model_failure(agent, model).await;
let now = chrono::Utc::now().timestamp();
let third_until = cooldown_until(&key).expect("third cooldown should exist");
let third_remaining = third_until.saturating_sub(now);
assert!(
third_remaining > PERSISTENT_MODEL_BACKOFF_BASE_SECS,
"third persistent model failure should escalate past the 4h base, got {third_remaining}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_persistent_model_failure_stays_capped_at_7_days() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "testagent_persistent_capped";
let model = "testmodel_persistent_capped";
let key = format!("{agent}:{model}");
for _ in 0..9 {
record_persistent_model_failure(agent, model).await;
}
let remaining = cooldown_until(&key)
.expect("cooldown should exist")
.saturating_sub(chrono::Utc::now().timestamp());
assert!(
remaining >= PERSISTENT_MODEL_BACKOFF_MAX_SECS - 5,
"persistent model cooldown should stay capped at 7 days, got {remaining}s"
);
assert!(
remaining <= PERSISTENT_MODEL_BACKOFF_MAX_SECS + 5,
"persistent model cooldown should not exceed 7 days, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_codex_format() {
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.";
let ts = parse_retry_at(msg);
assert!(ts.is_some(), "should parse codex retry-at date");
use chrono::TimeZone;
let naive = chrono::NaiveDateTime::new(
chrono::NaiveDate::from_ymd_opt(2026, 3, 26).unwrap(),
chrono::NaiveTime::from_hms_opt(5, 55, 0).unwrap(),
);
let expected_utc = match chrono::Local.from_local_datetime(&naive) {
chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
chrono::LocalResult::None => naive.and_utc(),
};
assert_eq!(
ts.unwrap(),
expected_utc.timestamp(),
"timestamp should reflect local-timezone interpretation of the vendor date"
);
let dt = chrono::DateTime::from_timestamp(ts.unwrap(), 0).unwrap();
use chrono::Datelike;
assert_eq!(dt.month(), 3);
assert_eq!(dt.day(), 26);
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_hour_only_and_month_day() {
let now = chrono::Utc::now().timestamp();
for msg in [
"You've hit your session limit · resets 8am (America/Sao_Paulo)",
"You've hit your session limit · resets 6pm (America/Sao_Paulo)",
"You've hit your session limit · resets 12am (America/Sao_Paulo)",
"You've hit your weekly limit · resets Jun 9 at 1am (America/Sao_Paulo)",
] {
let ts = parse_retry_at(msg).unwrap_or_else(|| panic!("should parse: {msg}"));
assert!(ts > now, "{msg}");
}
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_weekly_resets_without_year() {
use chrono::Datelike;
let t = chrono::Local::now() + chrono::Duration::days(2);
let msg = format!(
"You've hit your weekly limit · resets {} {} at 1am (America/Sao_Paulo)",
t.format("%b"),
t.day()
);
let ts = parse_retry_at(&msg).expect("should parse year-less reset date");
let remaining = ts - chrono::Utc::now().timestamp();
assert!(
remaining > 3600 && remaining < 3 * 86400,
"expected ~2d retry-at, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_no_date() {
assert!(parse_retry_at("generic rate limit error").is_none());
assert!(parse_retry_at("").is_none());
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_glm_reset_format_with_seconds() {
let msg = "Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48";
let ts = parse_retry_at(msg).expect("should parse GLM reset-at date");
use chrono::TimeZone;
let naive = chrono::NaiveDateTime::new(
chrono::NaiveDate::from_ymd_opt(2026, 4, 23).expect("valid date"),
chrono::NaiveTime::from_hms_opt(18, 33, 48).expect("valid time"),
);
let expected_utc = match chrono::Local.from_local_datetime(&naive) {
chrono::LocalResult::Single(local_dt) => local_dt.with_timezone(&chrono::Utc),
chrono::LocalResult::Ambiguous(earliest, _) => earliest.with_timezone(&chrono::Utc),
chrono::LocalResult::None => naive.and_utc(),
};
assert_eq!(
ts,
expected_utc.timestamp(),
"timestamp should reflect local-timezone interpretation of vendor reset date"
);
}
#[serial(cooldown_state)]
#[test]
fn parse_retry_at_billing_cycle_without_date_returns_none() {
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.";
assert!(
parse_retry_at(msg).is_none(),
"billing cycle without date should return None"
);
}
#[serial(cooldown_state)]
#[test]
fn parse_relative_usage_window_extracts_kimi_five_hour_window() {
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.";
assert_eq!(
parse_relative_usage_window(msg),
Some(5 * 3600),
"should extract the 5-hour window as seconds"
);
}
#[serial(cooldown_state)]
#[test]
fn parse_relative_usage_window_handles_common_durations() {
assert_eq!(
parse_relative_usage_window("You've reached your hourly usage limit."),
Some(3600)
);
assert_eq!(
parse_relative_usage_window("You've reached your daily usage limit."),
Some(86400)
);
assert_eq!(
parse_relative_usage_window("You've reached your weekly limit."),
Some(604800)
);
assert_eq!(
parse_relative_usage_window("You've reached your 30-minute usage limit."),
Some(1800)
);
assert_eq!(
parse_relative_usage_window("You've reached your 2-day quota."),
Some(172800)
);
}
#[serial(cooldown_state)]
#[test]
fn parse_relative_usage_window_returns_none_without_usage_context() {
assert!(parse_relative_usage_window("429 Too Many Requests").is_none());
assert!(parse_relative_usage_window("rate limit exceeded").is_none());
assert!(parse_relative_usage_window("").is_none());
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_rate_limit_uses_model_specific_window_cooldown() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_record_rl_window_agent";
let model = "test_record_rl_window_model";
let msg = "You've reached your 5-hour usage limit. Your quota will reset when the current 5-hour window expires.";
record_rate_limit(agent, Some(model), msg).await;
assert!(
is_model_in_cooldown(agent, model),
"model should be cooled for the relative usage window"
);
assert!(
!is_agent_in_cooldown(agent),
"agent should not be cooled for a model-scoped quota"
);
let key = format!("{agent}:{model}");
let remaining = cooldown_until(&key)
.expect("model cooldown should exist")
.saturating_sub(chrono::Utc::now().timestamp());
assert!(
remaining >= 5 * 3600 - 5,
"model cooldown should be at least ~5 hours, got {remaining}s"
);
assert_eq!(
cooldown_reason(&key).as_deref(),
Some("rate_limit_usage_window"),
"relative usage-window cooldown must be labeled as a rate limit, not silence"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_rate_limit_without_window_falls_back_to_agent_and_model() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_record_rl_fallback_agent";
let model = "test_record_rl_fallback_model";
let msg = "429 Too Many Requests";
record_rate_limit(agent, Some(model), msg).await;
assert!(
is_agent_in_cooldown(agent),
"agent should be cooled when no relative window is detected"
);
assert!(
is_model_in_cooldown(agent, model),
"model should be cooled when no relative window is detected"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn agent_cooldown_persists_to_kv() {
super::reset_global_state().await;
let agent = "test_agent_persist_check";
record_agent_failure_with_message(agent, "").await;
assert!(is_agent_in_cooldown(agent));
}
#[serial(cooldown_state)]
#[tokio::test]
async fn cooldown_never_shortens() {
super::reset_global_state().await;
let agent = "test_agent_no_shorten";
set_agent_cooldown(agent, 24 * 60 * 60).await;
let initial = {
let map = cooldowns().lock().unwrap();
map.get(agent)
.expect("cooldown entry should exist")
.cooldown_until
};
record_agent_failure_with_message(agent, "rate limit").await;
let after = {
let map = cooldowns().lock().unwrap();
map.get(agent)
.expect("cooldown entry should exist")
.cooldown_until
};
assert_eq!(
initial, after,
"cooldown should not be shortened by later failures"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn silence_agent_cooldown_is_short_lived() {
super::reset_global_state().await;
let agent = "test_silence_agent_cd";
assert!(!is_agent_in_cooldown(agent));
set_agent_cooldown(agent, SILENCE_AGENT_COOLDOWN_SECS).await;
assert!(is_agent_in_cooldown(agent));
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("should have cooldown entry");
let now = chrono::Utc::now().timestamp();
let remaining = entry.cooldown_until - now;
assert!(
remaining <= SILENCE_AGENT_COOLDOWN_SECS as i64,
"silence agent cooldown should be <= {SILENCE_AGENT_COOLDOWN_SECS}s, got {remaining}s"
);
assert!(remaining > 0, "cooldown should still be active");
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_silence_detection_applies_extended_cooldown() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_silence_count_agent";
let model = "test_silence_count_model";
assert!(!is_model_in_cooldown(agent, model));
for _ in 0..SILENCE_COUNT_THRESHOLD {
let result = record_silence_detection(agent, model).await;
assert!(result.is_some());
}
assert!(is_model_in_cooldown(agent, model));
let key = format!("{agent}:{model}");
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(&key).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= SILENCE_EXTENDED_COOLDOWN_SECS as i64 - 5,
"extended cooldown should be applied, got {remaining}s"
);
let kv_key = format!("{SILENCE_COUNT_PREFIX}{agent}:{model}");
let stored = store.kv_get(&kv_key).await.unwrap().unwrap();
assert_eq!(stored, "[]");
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_out_of_credits() {
let msg = "Error: API error - {\"error\":{\"message\":\"You exceeded your current quota\",\"type\":\"insufficient_quota\",\"param\":null,\"code\":\"insufficient_quota\"}}";
assert_eq!(
detect_credit_exhaustion(msg),
Some(CreditExhaustionReason::OutOfCredits)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_out_of_credits_variations() {
assert_eq!(
detect_credit_exhaustion("out of credits"),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("outofcredits"),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("no credits remaining"),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("insufficient funds"),
Some(CreditExhaustionReason::OutOfCredits)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_org_level_disabled() {
let msg = "Error: API error - {\"error\":{\"message\":\"Organization has been disabled\",\"type\":\"org_disabled\",\"code\":\"overageDisabledReason:org_level_disabled\"}}";
assert_eq!(
detect_credit_exhaustion(msg),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_org_level_disabled_variations() {
assert_eq!(
detect_credit_exhaustion("organization has been disabled"),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
assert_eq!(
detect_credit_exhaustion("org disabled"),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
assert_eq!(
detect_credit_exhaustion("billing quota exceeded"),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
assert_eq!(
detect_credit_exhaustion("credit balance too low"),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_overage_disabled_reason_parsing() {
let msg = "overageDisabledReason: out_of_credits";
assert_eq!(
detect_credit_exhaustion(msg),
Some(CreditExhaustionReason::OutOfCredits)
);
let msg = "overageDisabledReason: org_level_disabled";
assert_eq!(
detect_credit_exhaustion(msg),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_glm_insufficient_balance() {
assert_eq!(
detect_credit_exhaustion(
"API Error: Request rejected (429) · Insufficient balance or no resource package. Please recharge."
),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("Insufficient balance"),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("please recharge your account"),
Some(CreditExhaustionReason::OutOfCredits)
);
assert_eq!(
detect_credit_exhaustion("no resource package available"),
Some(CreditExhaustionReason::OutOfCredits)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_none_for_regular_rate_limit() {
assert!(detect_credit_exhaustion("rate limit exceeded").is_none());
assert!(detect_credit_exhaustion("too many requests").is_none());
assert!(detect_credit_exhaustion("429 Too Many Requests").is_none());
assert!(detect_credit_exhaustion("you've hit your usage limit").is_none());
assert!(detect_credit_exhaustion(
"You've hit your usage limit, try again in a few minutes."
)
.is_none());
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_credit_exhaustion_applies_agent_cooldown() {
super::reset_global_state().await;
let agent = "test_credit_exhaust_agent";
assert!(!is_agent_in_cooldown(agent));
record_credit_exhaustion(agent, CreditExhaustionReason::OutOfCredits).await;
assert!(is_agent_in_cooldown(agent));
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= CREDIT_BACKOFF_BASE_SECS - 5,
"first out_of_credits cooldown should be ~1 hour, got {remaining}s"
);
assert!(
remaining <= CREDIT_BACKOFF_MAX_SECS,
"cooldown should not exceed cap of 8h, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_credit_exhaustion_org_level_applies_longer_cooldown() {
super::reset_global_state().await;
let agent = "test_org_disabled_agent";
assert!(!is_agent_in_cooldown(agent));
record_credit_exhaustion(agent, CreditExhaustionReason::OrgLevelDisabled).await;
assert!(is_agent_in_cooldown(agent));
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= ORG_BACKOFF_BASE_SECS - 5,
"first org-level disabled cooldown should be ~2 hours, got {remaining}s"
);
assert!(
remaining <= CREDIT_BACKOFF_MAX_SECS,
"cooldown should not exceed cap of 8h, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_credit_exhaustion_billing_cycle_applies_24h() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_billing_cycle_agent";
assert!(!is_agent_in_cooldown(agent));
record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
assert!(is_agent_in_cooldown(agent));
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= BILLING_CYCLE_COOLDOWN_SECS - 5,
"first billing cycle cooldown should be ~24 hours, got {remaining}s"
);
assert!(
remaining <= BILLING_CYCLE_MAX_SECS,
"cooldown should not exceed cap of 7 days, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_credit_exhaustion_billing_cycle_escalates() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_billing_cycle_escalation";
record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
let remaining_1 = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
(BILLING_CYCLE_COOLDOWN_SECS - 5..=BILLING_CYCLE_COOLDOWN_SECS + 5)
.contains(&remaining_1),
"first billing cycle cooldown should be ~24h, got {remaining_1}s"
);
{
let mut map = cooldowns().lock().unwrap();
map.remove(agent);
}
record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
let remaining_2 = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
let expected_2 = BILLING_CYCLE_COOLDOWN_SECS * 3; assert!(
remaining_2 >= expected_2 - 5,
"second billing cycle cooldown should be ~72h, got {remaining_2}s"
);
{
let mut map = cooldowns().lock().unwrap();
map.remove(agent);
}
record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
let remaining_3 = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining_3 >= BILLING_CYCLE_MAX_SECS - 5,
"third billing cycle cooldown should be capped at 7 days, got {remaining_3}s"
);
assert!(
remaining_3 <= BILLING_CYCLE_MAX_SECS + 5,
"third billing cycle cooldown should not exceed 7 days, got {remaining_3}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_credit_exhaustion_billing_cycle_stays_capped_after_repeats() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_billing_cycle_agent_capped";
for _ in 0..5 {
record_credit_exhaustion(agent, CreditExhaustionReason::BillingCycleExhausted).await;
}
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= BILLING_CYCLE_MAX_SECS - 5,
"billing cycle cooldown should stay capped at 7 days, got {remaining}s"
);
assert!(
remaining <= BILLING_CYCLE_MAX_SECS + 5,
"billing cycle cooldown should not exceed 7 days, got {remaining}s"
);
}
#[serial(cooldown_state)]
#[test]
fn credit_exhaustion_backoff_constants() {
assert_eq!(
compute_backoff(1, CREDIT_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
CREDIT_BACKOFF_BASE_SECS
);
assert_eq!(
compute_backoff(1, ORG_BACKOFF_BASE_SECS, CREDIT_BACKOFF_MAX_SECS, true),
ORG_BACKOFF_BASE_SECS
);
assert_eq!(BILLING_CYCLE_COOLDOWN_SECS, 24 * 60 * 60);
assert_eq!(BILLING_CYCLE_MAX_SECS, 7 * 24 * 60 * 60);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_long_cap_without_extended_stays_capped() {
assert_eq!(
compute_backoff(
3,
BILLING_CYCLE_COOLDOWN_SECS,
BILLING_CYCLE_MAX_SECS,
false
),
BILLING_CYCLE_MAX_SECS,
"7-day capped paths should stop at the documented max"
);
assert_eq!(
compute_backoff(
9,
PERSISTENT_MODEL_BACKOFF_BASE_SECS,
PERSISTENT_MODEL_BACKOFF_MAX_SECS,
false,
),
PERSISTENT_MODEL_BACKOFF_MAX_SECS,
"persistent model failures should not extend past 7 days"
);
}
#[serial(cooldown_state)]
#[test]
fn mark_and_check_agent_degraded() {
let agent = "test_degraded_mark";
assert!(!is_agent_degraded(agent));
mark_agent_degraded(agent);
assert!(is_agent_degraded(agent));
clear_agent_degraded(agent);
assert!(!is_agent_degraded(agent));
}
#[serial(cooldown_state)]
#[tokio::test]
async fn sync_from_kv_drops_in_memory_entry_cleared_externally() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_sync_external_clear";
let future = chrono::Utc::now().timestamp() + 600;
init_cooldown_store(store.clone()).await;
set_cooldown_async(agent, future, "test").await;
assert!(is_agent_in_cooldown(agent));
store
.kv_set(&format!("{KV_PREFIX}{agent}"), "0")
.await
.expect("kv_set tombstone");
assert!(
is_agent_in_cooldown(agent),
"in-memory cooldown should still be present before sync"
);
sync_from_kv(&store).await;
assert!(
!is_agent_in_cooldown(agent),
"sync_from_kv must drop in-memory entries the KV tombstoned"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn sync_from_kv_keeps_in_memory_entry_when_kv_row_missing() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_sync_race_window";
let future = chrono::Utc::now().timestamp() + 600;
set_cooldown_in_memory(agent, future, "raced");
assert!(is_agent_in_cooldown(agent));
sync_from_kv(&store).await;
assert!(
is_agent_in_cooldown(agent),
"sync_from_kv must NOT drop entries whose KV row is missing — that is the set_cooldown_async race window"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn sync_from_kv_adopts_kv_timestamp() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_sync_adopt_kv";
let original = chrono::Utc::now().timestamp() + 600;
let updated = chrono::Utc::now().timestamp() + 1200;
init_cooldown_store(store.clone()).await;
set_cooldown_async(agent, original, "initial").await;
store
.kv_set(&format!("{KV_PREFIX}{agent}"), &updated.to_string())
.await
.expect("kv_set updated");
sync_from_kv(&store).await;
let until = cooldown_until(agent).expect("cooldown still active");
assert_eq!(
until, updated,
"sync_from_kv must adopt KV's authoritative timestamp"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_marks_cooled_agent() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_cooled";
record_agent_failure_with_message(agent, "").await;
assert!(is_agent_in_cooldown(agent));
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
assert!(
is_agent_degraded(agent),
"agent in cooldown should be marked degraded"
);
clear_agent_degraded(agent);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_marks_no_models_agent() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_no_models";
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| false, 6, 3).await;
assert!(
is_agent_degraded(agent),
"agent with no available models should be marked degraded"
);
clear_agent_degraded(agent);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_marks_rate_limited_agent() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_rate_limited";
for _ in 0..4 {
store
.record_rate_limit(agent, "rate_limit", None)
.await
.unwrap();
}
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
assert!(
is_agent_degraded(agent),
"agent with rate limit count >= threshold should be marked degraded"
);
clear_agent_degraded(agent);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_clears_healthy_agent() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_healthy";
mark_agent_degraded(agent);
assert!(is_agent_degraded(agent));
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
assert!(
!is_agent_degraded(agent),
"healthy agent should have degraded flag cleared"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_ignores_expired_cooldowns() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_expired_cooldown";
let expired_ts = chrono::Utc::now().timestamp() - 100; {
let mut map = cooldowns().lock().unwrap();
map.insert(
agent.to_string(),
CooldownEntry {
cooldown_until: expired_ts,
reason: "expired_test".to_string(),
},
);
}
assert!(
!is_agent_in_cooldown(agent),
"expired cooldown should not be active"
);
mark_agent_degraded(agent);
assert!(is_agent_degraded(agent));
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
assert!(
!is_agent_degraded(agent),
"agent with only expired cooldown should NOT be marked degraded"
);
clear_agent_degraded(agent);
{
let mut map = cooldowns().lock().unwrap();
map.remove(agent);
}
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_grows_exponentially() {
assert_eq!(compute_backoff(0, 300, 14400, true), 300);
assert_eq!(compute_backoff(1, 300, 14400, true), 300);
assert_eq!(compute_backoff(2, 300, 14400, true), 900);
assert_eq!(compute_backoff(3, 300, 14400, true), 2700);
assert_eq!(compute_backoff(4, 300, 14400, true), 8100);
assert_eq!(compute_backoff(5, 300, 14400, true), 14400);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_extended_tier_1_applies_24h() {
assert_eq!(
compute_backoff(5, 300, 14400, true),
14400,
"first at max, not extended"
);
assert_eq!(
compute_backoff(6, 300, 14400, true),
14400,
"times_at_max=1, not extended"
);
assert_eq!(
compute_backoff(7, 300, 14400, true),
86400,
"times_at_max=2, extended 6x"
);
assert_eq!(
compute_backoff(8, 300, 14400, true),
86400,
"times_at_max=3, still extended 6x"
);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_extended_tier_2_applies_48h() {
assert_eq!(
compute_backoff(9, 300, 14400, true),
86400,
"times_at_max=4, still tier 1"
);
assert_eq!(
compute_backoff(10, 300, 14400, true),
172800,
"times_at_max=5, extended 12x"
);
assert_eq!(
compute_backoff(15, 300, 14400, true),
172800,
"times_at_max=10, still tier 2"
);
assert_eq!(
compute_backoff(20, 300, 14400, true),
172800,
"times_at_max=15, still tier 2"
);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_extended_tier_with_different_base_max() {
assert_eq!(compute_backoff(5, 60, 3600, true), 3600, "first at max");
assert_eq!(compute_backoff(6, 60, 3600, true), 3600, "times_at_max=1");
assert_eq!(
compute_backoff(7, 60, 3600, true),
21600,
"times_at_max=2, extended 6x"
);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_extended_tier_persistent_fails_example() {
assert_eq!(
compute_backoff(21, 300, 14400, true),
172800,
"21 failures should get 48h extended cooldown"
);
assert_eq!(
compute_backoff(36, 300, 14400, true),
172800,
"36 failures should also get 48h extended cooldown"
);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_extended_tier_zero_base() {
assert_eq!(compute_backoff(1, 0, 14400, true), 0);
assert_eq!(compute_backoff(5, 0, 14400, true), 0);
assert_eq!(compute_backoff(10, 0, 14400, true), 0);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_billing_cycle() {
assert_eq!(
detect_credit_exhaustion("billing cycle"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion("quota refreshed next cycle"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion(
"You've reached your usage limit for this billing cycle. Your quota will be refreshed in the next cycle."
),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion("monthly quota exceeded"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion(
"Weekly/Monthly Limit Exhausted. Your limit will reset at 2026-04-23 18:33:48"
),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_billing_cycle_minimax() {
assert_eq!(
detect_credit_exhaustion(
"minimax rate limit: API Error: Request rejected (429) · Token Plan usage limit reached: Upgrade your Token Plan or purchase Credits for more usage."
),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion("Upgrade your Token Plan"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion("Token Plan usage limit reached"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion("purchase Credits for more usage"),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
assert_eq!(
detect_credit_exhaustion(
"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"
),
Some(CreditExhaustionReason::BillingCycleExhausted)
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_outcome_kimi_billing_cycle() {
let msg = "Failed to authenticate. API Error: 403 You've reached your usage limit for \
this billing cycle. Your quota will be refreshed soon.";
assert_eq!(
detect_credit_exhaustion_outcome(msg),
Some("billing_cycle_exhausted")
);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_outcome_maps_all_reasons() {
assert_eq!(
detect_credit_exhaustion_outcome("out of credits"),
Some("credit_exhausted")
);
assert_eq!(
detect_credit_exhaustion_outcome("organization has been disabled"),
Some("org_level_disabled")
);
assert_eq!(
detect_credit_exhaustion_outcome("monthly quota exceeded"),
Some("billing_cycle_exhausted")
);
}
#[test]
fn detect_credit_exhaustion_outcome_returns_none_for_unrelated_errors() {
assert_eq!(
detect_credit_exhaustion_outcome("429 Too Many Requests"),
None
);
assert_eq!(detect_credit_exhaustion_outcome("connection reset"), None);
}
#[serial(cooldown_state)]
#[test]
fn detect_credit_exhaustion_overage_takes_priority_over_billing_cycle() {
let msg = "overageDisabledReason: org_level_disabled for billing cycle";
assert_eq!(
detect_credit_exhaustion(msg),
Some(CreditExhaustionReason::OrgLevelDisabled)
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn refresh_degraded_agents_below_threshold_not_degraded() {
super::reset_global_state().await;
let store = test_store().await;
let agent = "test_refresh_below_threshold";
for _ in 0..2 {
store
.record_rate_limit(agent, "rate_limit", None)
.await
.unwrap();
}
let agents = vec![agent.to_string()];
refresh_degraded_agents(&store, &agents, &|_| true, 6, 3).await;
assert!(
!is_agent_degraded(agent),
"agent with rate limit count < threshold should not be degraded"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn repeated_agent_failures_escalate_backoff() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_escalation_agent";
record_agent_failure_with_message(agent, "").await;
let remaining_1 = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("should have cooldown");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
(BACKOFF_BASE_SECS - 5..=BACKOFF_BASE_SECS + 5).contains(&remaining_1),
"first failure should be ~5 min, got {remaining_1}s"
);
{
let mut map = cooldowns().lock().unwrap();
map.remove(agent);
}
record_agent_failure_with_message(agent, "").await;
let remaining_2 = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("should have cooldown");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining_2 >= BACKOFF_BASE_SECS * 3 - 5,
"second failure should be ~15 min, got {remaining_2}s"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn vendor_past_retry_falls_back_to_backoff() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_vendor_past_retry";
{
let mut map = cooldowns().lock().unwrap();
map.remove(agent);
}
let past_msg = "You've hit the limit; try again at Jan 01, 2000 01:00 AM.";
record_agent_failure_with_message(agent, past_msg).await;
assert!(
is_agent_in_cooldown(agent),
"agent should be in cooldown after failure"
);
let remaining = {
let map = cooldowns().lock().unwrap();
let entry = map.get(agent).expect("cooldown entry should exist");
entry.cooldown_until - chrono::Utc::now().timestamp()
};
assert!(
remaining >= BACKOFF_BASE_SECS - 5,
"expected backoff of ~{BACKOFF_BASE_SECS}s when vendor timestamp is past, got {remaining}s"
);
let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
let stored = store
.kv_get(&kv_key)
.await
.unwrap()
.expect("failure count should be stored");
assert_eq!(stored, "1");
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_agent_success_resets_backoff() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_success_reset_agent";
let model = "test_model";
for _ in 0..3 {
let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
}
let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
let count: u32 = store
.kv_get(&kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(count, 3);
record_agent_success(agent, model).await;
let count_after: u32 = store
.kv_get(&kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(count_after, 0, "success should reset failure count to 0");
let model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:{model}");
let model_count: u32 = store
.kv_get(&model_kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(
model_count, 0,
"success should reset model failure count to 0"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn record_agent_success_without_model_skips_empty_model_key() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_success_without_model";
record_agent_success(agent, "").await;
let agent_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
let agent_count = store
.kv_get(&agent_kv_key)
.await
.unwrap()
.expect("agent failure count should be stored");
assert_eq!(agent_count, "0");
let empty_model_kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}:");
let empty_model_count = store.kv_get(&empty_model_kv_key).await.unwrap();
assert!(
empty_model_count.is_none(),
"success without model should not create an empty model failure count key"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn clear_cooldown_resets_failure_counts() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agent = "test_clear_fc_agent";
for _ in 0..5 {
let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
}
set_agent_cooldown(agent, 3600).await;
let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
let count: u32 = store
.kv_get(&kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(count, 5);
clear_cooldown(agent, &store).await;
let count_after: u32 = store
.kv_get(&kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(
count_after, 0,
"clear_cooldown should reset failure count to 0"
);
assert!(
!is_agent_in_cooldown(agent),
"agent should not be in cooldown after clear"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn clear_all_cooldowns_resets_all_failure_counts() {
super::reset_global_state().await;
let store = test_store().await;
*cooldown_store().lock().await = Some(store.clone());
let agents = ["test_clear_all_a", "test_clear_all_b"];
for agent in &agents {
let _ = read_and_increment_failure_count(&Some(store.clone()), agent).await;
set_agent_cooldown(agent, 3600).await;
}
clear_cooldown("*", &store).await;
for agent in &agents {
assert!(!is_agent_in_cooldown(agent));
let kv_key = format!("{FAILURE_COUNT_PREFIX}{agent}");
let count: u32 = store
.kv_get(&kv_key)
.await
.unwrap()
.unwrap()
.parse()
.unwrap();
assert_eq!(
count, 0,
"clear --all should reset failure count for {agent}"
);
}
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_saturates_on_large_counts() {
assert_eq!(
compute_backoff(100, BACKOFF_BASE_SECS, BACKOFF_MAX_SECS, true),
BACKOFF_MAX_SECS * 12
);
assert_eq!(compute_backoff(u32::MAX as u64, 300, 14400, true), 172800);
assert_eq!(compute_backoff(u64::MAX, 300, 14400, true), 172800);
}
#[serial(cooldown_state)]
#[test]
fn compute_backoff_zero_base_returns_zero() {
assert_eq!(compute_backoff(1, 0, 14400, true), 0);
assert_eq!(compute_backoff(5, 0, 14400, true), 0);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn github_circuit_stays_closed_below_threshold() {
super::reset_global_state().await;
{
let mut open = github_5xx_circuit_open().lock().unwrap();
*open = None;
let mut ts = github_5xx_timestamps().lock().unwrap();
ts.clear();
let mut map = cooldown_lock();
map.remove("github:5xx");
}
for _ in 0..GITHUB_5XX_THRESHOLD - 1 {
record_github_5xx().await;
}
assert!(
!is_github_circuit_open(),
"circuit should stay closed below threshold"
);
assert_eq!(github_circuit_remaining_secs(), 0);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn github_circuit_trips_at_threshold() {
super::reset_global_state().await;
{
let mut open = github_5xx_circuit_open_lock();
*open = None;
let mut ts = github_5xx_timestamps_lock();
ts.clear();
let mut map = cooldown_lock();
map.remove("github:5xx");
}
for _ in 0..GITHUB_5XX_THRESHOLD {
record_github_5xx().await;
}
assert!(is_github_circuit_open(), "circuit should trip at threshold");
assert!(
github_circuit_remaining_secs() > 0,
"remaining should be > 0 when open"
);
assert!(
github_circuit_remaining_secs() <= GITHUB_5XX_COOLDOWN_SECS as u64,
"remaining should not exceed cooldown duration"
);
{
let mut open = github_5xx_circuit_open().lock().unwrap();
*open = None;
}
}
#[serial(cooldown_state)]
#[tokio::test]
async fn github_circuit_cannot_double_trip() {
super::reset_global_state().await;
{
let mut open = github_5xx_circuit_open_lock();
*open = None;
let mut ts = github_5xx_timestamps_lock();
ts.clear();
let mut map = cooldown_lock();
map.remove("github:5xx");
}
for _ in 0..GITHUB_5XX_THRESHOLD {
record_github_5xx().await;
}
assert!(is_github_circuit_open());
let first_remaining = github_circuit_remaining_secs();
for _ in 0..GITHUB_5XX_THRESHOLD {
record_github_5xx().await;
}
let second_remaining = github_circuit_remaining_secs();
assert!(
second_remaining <= first_remaining,
"additional errors should not extend cooldown"
);
{
let mut open = github_5xx_circuit_open().lock().unwrap();
*open = None;
}
}
#[serial(cooldown_state)]
#[tokio::test]
async fn github_circuit_recovers_after_cooldown() {
super::reset_global_state().await;
{
let mut open = github_5xx_circuit_open_lock();
*open = None;
let mut ts = github_5xx_timestamps_lock();
ts.clear();
let mut map = cooldown_lock();
map.remove("github:5xx");
}
{
let mut open = github_5xx_circuit_open().lock().unwrap();
*open = Some(chrono::Utc::now().timestamp() - 1);
}
assert!(
!is_github_circuit_open(),
"circuit should auto-recover after cooldown"
);
}
#[serial(cooldown_state)]
#[tokio::test]
async fn github_circuit_sliding_window_ages_out() {
super::reset_global_state().await;
{
let mut open = github_5xx_circuit_open_lock();
*open = None;
let mut ts = github_5xx_timestamps_lock();
ts.clear();
let mut map = cooldown_lock();
map.remove("github:5xx");
}
{
let mut ts = github_5xx_timestamps().lock().unwrap();
let old = chrono::Utc::now().timestamp() - GITHUB_5XX_WINDOW_SECS - 10;
for _ in 0..GITHUB_5XX_THRESHOLD {
ts.push(old);
}
}
record_github_5xx().await;
assert!(
!is_github_circuit_open(),
"old errors outside window should not trip circuit"
);
}
}