use serde::{Deserialize, Serialize};
use super::account::CredentialRef;
use super::provider::OAuthProviderId;
pub mod cred_serde {
use super::CredentialRef;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<S: Serializer>(c: &CredentialRef, s: S) -> Result<S::Ok, S::Error> {
c.storage_key().serialize(s)
}
pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<CredentialRef, D::Error> {
let raw = String::deserialize(d)?;
CredentialRef::parse_storage_key(&raw)
.ok_or_else(|| serde::de::Error::custom(format!("invalid credential key '{raw}'")))
}
}
pub const WEEKLY_WINDOW_MIN_MS: u64 = 6 * 24 * 60 * 60 * 1000;
pub const WEEKLY_WINDOW_MAX_MS: u64 = 8 * 24 * 60 * 60 * 1000;
pub const MAX_FUTURE_OBSERVATION_SKEW_MS: u64 = 5 * 60 * 1000;
pub const DEFAULT_EXHAUSTED_AT_PERCENT: f64 = 100.0;
pub const DEFAULT_URGENT_HORIZON_MS: u64 = 24 * 60 * 60 * 1000;
pub const UNANCHORED_TOLERANCE_MS: u64 = 5 * 60 * 1000;
pub const UNANCHORED_MAX_USED_PERCENT: f64 = 0.5;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WindowLimit {
pub id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub used_percent: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit_reached: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub resets_at_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub models: Option<Vec<String>>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum WindowVerdict {
Headroom { used_percent: Option<f64> },
Exhausted { resets_at_ms: Option<u64> },
Unknown,
}
pub fn model_matches(entry: &str, requested: &str) -> bool {
let entry = entry.trim().to_ascii_lowercase();
let req = requested.trim().to_ascii_lowercase();
if entry.is_empty() || req.is_empty() {
return false;
}
if entry == req {
return true;
}
req.split(['-', '_', '/', ':', '.']).any(|tok| tok == entry)
}
impl WindowLimit {
pub fn applies_to(&self, model: Option<&str>) -> bool {
match (&self.models, model) {
(None, _) => true,
(Some(_), None) => true,
(Some(list), Some(m)) => list.iter().any(|x| model_matches(x, m)),
}
}
pub fn is_weekly(&self) -> bool {
matches!(self.duration_ms, Some(d) if (WEEKLY_WINDOW_MIN_MS..=WEEKLY_WINDOW_MAX_MS).contains(&d))
}
pub fn verdict(&self, exhausted_at_percent: f64) -> WindowVerdict {
if !valid_threshold(exhausted_at_percent) {
return WindowVerdict::Unknown;
}
if self.limit_reached == Some(true) {
return WindowVerdict::Exhausted {
resets_at_ms: self.resets_at_ms,
};
}
let used = match self.used_percent {
Some(p) if p > 100.0 && p <= super::usage::OVERRUN_PERCENT_MAX => Some(100.0),
Some(p) if !p.is_finite() || !(0.0..=100.0).contains(&p) => {
return WindowVerdict::Unknown
}
other => other,
};
match used {
Some(p) if p >= exhausted_at_percent => WindowVerdict::Exhausted {
resets_at_ms: self.resets_at_ms,
},
Some(p) => WindowVerdict::Headroom {
used_percent: Some(p),
},
None if self.limit_reached == Some(false) => {
WindowVerdict::Headroom { used_percent: None }
}
None => WindowVerdict::Unknown,
}
}
}
pub fn valid_threshold(percent: f64) -> bool {
percent.is_finite() && (0.0..=100.0).contains(&percent)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelState {
Available,
Exhausted,
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelAvailability {
pub model: String,
pub state: ModelState,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum QuotaObservation {
Ok {
windows: Vec<WindowLimit>,
#[serde(default, skip_serializing_if = "Option::is_none")]
models: Option<Vec<ModelAvailability>>,
},
AuthError,
Unsupported,
Malformed,
Unknown,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AccountCapacity {
#[serde(with = "cred_serde")]
pub credential: CredentialRef,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at_ms: Option<u64>,
pub observation: QuotaObservation,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cooldown_until_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub failure: Option<ObservationFailure>,
}
impl AccountCapacity {
pub fn unobserved(credential: CredentialRef) -> Self {
Self {
credential,
observed_at_ms: None,
observation: QuotaObservation::Unknown,
cooldown_until_ms: None,
failure: None,
}
}
pub fn failed(credential: CredentialRef, failure: ObservationFailure) -> Self {
Self {
observation: failure.observation(),
failure: Some(failure),
..Self::unobserved(credential)
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AccountFailureKind {
CoolingDown,
Exhausted,
ModelUnavailable,
UsageRateLimited,
UsageTimeout,
UsageUpstream,
UsageTransport,
UsageMalformed,
UsageUnauthorized,
RefreshTransient,
RefreshLockBusy,
RefreshOther,
RefreshRejected,
RefreshActionRequired,
CredentialMissing,
CredentialStore,
AuthError,
Unsupported,
QuotaUnknown,
StaleSnapshot,
WindowUnknown,
ModelUnknown,
NoLimitEvidence,
Excluded,
NotPresent,
InvalidLabel,
#[serde(other)]
Unknown,
}
impl AccountFailureKind {
pub fn as_str(self) -> &'static str {
match self {
Self::CoolingDown => "cooling_down",
Self::Exhausted => "exhausted",
Self::ModelUnavailable => "model_unavailable",
Self::UsageRateLimited => "usage_rate_limited",
Self::UsageTimeout => "usage_timeout",
Self::UsageUpstream => "usage_upstream",
Self::UsageTransport => "usage_transport",
Self::UsageMalformed => "usage_malformed",
Self::UsageUnauthorized => "usage_unauthorized",
Self::RefreshTransient => "refresh_transient",
Self::RefreshLockBusy => "refresh_lock_busy",
Self::RefreshOther => "refresh_other",
Self::RefreshRejected => "refresh_rejected",
Self::RefreshActionRequired => "refresh_action_required",
Self::CredentialMissing => "credential_missing",
Self::CredentialStore => "credential_store",
Self::AuthError => "auth_error",
Self::Unsupported => "unsupported",
Self::QuotaUnknown => "quota_unknown",
Self::StaleSnapshot => "stale_snapshot",
Self::WindowUnknown => "window_unknown",
Self::ModelUnknown => "model_unknown",
Self::NoLimitEvidence => "no_limit_evidence",
Self::Excluded => "excluded",
Self::NotPresent => "not_present",
Self::InvalidLabel => "invalid_label",
Self::Unknown => "unknown",
}
}
pub fn needs_login(self) -> bool {
matches!(
self,
Self::UsageUnauthorized
| Self::RefreshRejected
| Self::CredentialMissing
| Self::AuthError
)
}
pub fn is_transient(self) -> bool {
matches!(
self,
Self::UsageRateLimited
| Self::UsageTimeout
| Self::UsageUpstream
| Self::UsageTransport
| Self::RefreshTransient
| Self::RefreshLockBusy
| Self::RefreshOther
)
}
pub fn is_limited(self) -> bool {
matches!(
self,
Self::CoolingDown | Self::Exhausted | Self::ModelUnavailable
)
}
pub fn needs_action(self) -> bool {
matches!(self, Self::RefreshActionRequired | Self::CredentialStore)
}
}
impl std::fmt::Display for AccountFailureKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ObservationFailure {
pub kind: AccountFailureKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub http_status: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retry_after_s: Option<u64>,
}
impl ObservationFailure {
pub fn new(kind: AccountFailureKind) -> Self {
Self {
kind,
http_status: None,
retry_after_s: None,
}
}
pub fn with_status(mut self, status: u16) -> Self {
self.http_status = Some(status);
self
}
pub fn with_retry_after(mut self, secs: Option<u64>) -> Self {
self.retry_after_s = secs;
self
}
pub fn observation(&self) -> QuotaObservation {
match self.kind {
k if k.needs_login() => QuotaObservation::AuthError,
AccountFailureKind::UsageMalformed => QuotaObservation::Malformed,
AccountFailureKind::Unsupported => QuotaObservation::Unsupported,
_ => QuotaObservation::Unknown,
}
}
}
impl std::fmt::Display for ObservationFailure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.kind)?;
if let Some(status) = self.http_status {
write!(f, ", HTTP {status}")?;
}
if let Some(secs) = self.retry_after_s {
write!(f, ", retry after {secs}s")?;
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Strategy {
PreferenceOrder,
LowestUtilization,
SoonestReset,
}
#[derive(Debug, Clone)]
pub struct SelectionRequest<'a> {
pub provider: OAuthProviderId,
pub model: Option<&'a str>,
pub now_ms: u64,
pub max_snapshot_age_ms: u64,
pub preference: &'a [String],
pub explicit_account: Option<&'a str>,
pub exclude: &'a [CredentialRef],
pub strategy: Strategy,
pub exhausted_at_percent: f64,
pub current: Option<&'a CredentialRef>,
pub urgent_horizon_ms: u64,
pub sticky: bool,
pub grace: &'a [AccountCapacity],
pub grace_max_age_ms: u64,
}
impl<'a> SelectionRequest<'a> {
pub fn new(provider: OAuthProviderId, now_ms: u64, max_snapshot_age_ms: u64) -> Self {
Self {
provider,
model: None,
now_ms,
max_snapshot_age_ms,
preference: &[],
explicit_account: None,
exclude: &[],
strategy: Strategy::LowestUtilization,
exhausted_at_percent: DEFAULT_EXHAUSTED_AT_PERCENT,
current: None,
urgent_horizon_ms: DEFAULT_URGENT_HORIZON_MS,
sticky: false,
grace: &[],
grace_max_age_ms: super::usage_reads::USAGE_READ_GRACE.as_millis() as u64,
}
}
pub fn soonest_reset(provider: OAuthProviderId, now_ms: u64, max_snapshot_age_ms: u64) -> Self {
Self {
strategy: Strategy::SoonestReset,
sticky: true,
..Self::new(provider, now_ms, max_snapshot_age_ms)
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "reason", rename_all = "snake_case")]
pub enum RejectReason {
Stale {
#[serde(default, skip_serializing_if = "Option::is_none")]
age_ms: Option<u64>,
},
Malformed,
Unknown,
Unsupported,
AuthError,
CoolingDown {
until_ms: u64,
},
Exhausted {
window: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
resets_at_ms: Option<u64>,
},
WindowUnknown {
window: String,
},
NoLimitEvidence,
ModelUnavailable {
model: String,
},
ModelUnknown {
model: String,
},
Excluded,
NotPresent,
InvalidLabel {
label: String,
},
Outranked {
#[serde(default)]
rank: u32,
},
}
impl std::fmt::Display for RejectReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Stale { age_ms: Some(a) } => write!(f, "stale snapshot ({}s old)", a / 1000),
Self::Stale { age_ms: None } => write!(f, "no snapshot"),
Self::Malformed => write!(f, "malformed snapshot"),
Self::Unknown => write!(f, "quota unknown"),
Self::Unsupported => write!(f, "usage unsupported"),
Self::AuthError => write!(f, "auth error"),
Self::CoolingDown { until_ms } => write!(f, "cooling down until {until_ms}"),
Self::Exhausted {
window,
resets_at_ms: Some(r),
} => write!(f, "window '{window}' exhausted, resets at {r}"),
Self::Exhausted { window, .. } => write!(f, "window '{window}' exhausted"),
Self::WindowUnknown { window } => write!(f, "window '{window}' has no evidence"),
Self::NoLimitEvidence => write!(f, "no applicable limit evidence"),
Self::ModelUnavailable { model } => write!(f, "model '{model}' unavailable"),
Self::ModelUnknown { model } => write!(f, "model '{model}' availability unknown"),
Self::Excluded => write!(f, "excluded"),
Self::NotPresent => write!(f, "account not present"),
Self::InvalidLabel { label } => write!(f, "invalid account label '{label}'"),
Self::Outranked { rank } => write!(f, "eligible; outranked (rank {rank})"),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Rejection {
#[serde(with = "cred_serde")]
pub credential: CredentialRef,
pub reason: RejectReason,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cause: Option<ObservationFailure>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usable_after_ms: Option<u64>,
}
impl Rejection {
pub fn describe(&self) -> String {
match &self.cause {
Some(cause) => format!("{} ({cause})", self.reason),
None => self.reason.to_string(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AccountFailure {
pub account: String,
pub kind: AccountFailureKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub http_status: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub retry_after_s: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub until_ms: Option<u64>,
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub cooling_down: bool,
}
impl AccountFailure {
pub fn from_rejection(rejection: &Rejection) -> Option<Self> {
use AccountFailureKind as K;
let mut out = Self {
account: rejection.credential.account.label_str().to_string(),
kind: K::QuotaUnknown,
http_status: None,
retry_after_s: None,
until_ms: rejection.usable_after_ms,
cooling_down: matches!(rejection.reason, RejectReason::CoolingDown { .. }),
};
if let RejectReason::Outranked { .. } = rejection.reason {
return None;
}
if let Some(cause) = &rejection.cause {
out.kind = cause.kind;
out.http_status = cause.http_status;
out.retry_after_s = cause.retry_after_s;
return Some(out);
}
out.kind = match &rejection.reason {
RejectReason::Outranked { .. } => return None,
RejectReason::CoolingDown { .. } => K::CoolingDown,
RejectReason::Exhausted { .. } => K::Exhausted,
RejectReason::Stale { .. } => K::StaleSnapshot,
RejectReason::Malformed => K::UsageMalformed,
RejectReason::Unknown => K::QuotaUnknown,
RejectReason::Unsupported => K::Unsupported,
RejectReason::AuthError => K::AuthError,
RejectReason::WindowUnknown { .. } => K::WindowUnknown,
RejectReason::NoLimitEvidence => K::NoLimitEvidence,
RejectReason::ModelUnavailable { .. } => K::ModelUnavailable,
RejectReason::ModelUnknown { .. } => K::ModelUnknown,
RejectReason::Excluded => K::Excluded,
RejectReason::NotPresent => K::NotPresent,
RejectReason::InvalidLabel { .. } => K::InvalidLabel,
};
Some(out)
}
pub fn is_limited(&self) -> bool {
self.cooling_down || self.kind.is_limited()
}
pub fn is_live(&self) -> bool {
!self.is_limited() && self.kind != AccountFailureKind::Excluded
}
pub fn log_form(&self) -> String {
let mut out = format!("{}:{}", self.account, self.kind);
if let Some(status) = self.http_status {
out.push_str(&format!(":{status}"));
}
if self.cooling_down && self.kind != AccountFailureKind::CoolingDown {
out.push_str(":cooling");
}
out
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SelectionFailureKind {
NoCredentials,
UsageReadFailed,
CredentialRejected,
AllCoolingDown,
QuotaUnknown,
SeatChurn,
RemoteNoCapacity,
ActionRequired,
CredentialStore,
#[serde(other)]
Unknown,
}
impl SelectionFailureKind {
pub fn as_str(self) -> &'static str {
match self {
Self::NoCredentials => "no_credentials",
Self::UsageReadFailed => "usage_read_failed",
Self::CredentialRejected => "credential_rejected",
Self::AllCoolingDown => "all_cooling_down",
Self::QuotaUnknown => "quota_unknown",
Self::SeatChurn => "seat_churn",
Self::RemoteNoCapacity => "remote_no_capacity",
Self::ActionRequired => "action_required",
Self::CredentialStore => "credential_store",
Self::Unknown => "unknown",
}
}
}
impl std::fmt::Display for SelectionFailureKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SelectionFailure {
pub kind: SelectionFailureKind,
#[serde(default)]
pub accounts: Vec<AccountFailure>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub earliest_reset_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage_read_retry_ms: Option<u64>,
}
impl SelectionFailure {
pub fn new(kind: SelectionFailureKind) -> Self {
Self {
kind,
accounts: Vec::new(),
earliest_reset_ms: None,
usage_read_retry_ms: None,
}
}
pub fn from_rejections(rejections: &[Rejection], now_ms: u64) -> Self {
use SelectionFailureKind as K;
let accounts: Vec<_> = rejections
.iter()
.filter_map(AccountFailure::from_rejection)
.collect();
let live: Vec<_> = accounts.iter().filter(|a| a.is_live()).collect();
let kind = if accounts.is_empty() {
K::NoCredentials
} else if live.iter().any(|a| a.kind.is_transient()) {
K::UsageReadFailed
} else if live.is_empty() && accounts.iter().any(AccountFailure::is_limited) {
K::AllCoolingDown
} else if live.is_empty() {
K::QuotaUnknown
} else if live.iter().all(|a| a.kind.needs_login()) {
K::CredentialRejected
} else if live
.iter()
.all(|a| a.kind.needs_login() || a.kind.needs_action())
{
K::ActionRequired
} else {
K::QuotaUnknown
};
let earliest_reset_ms = accounts
.iter()
.filter(|a| a.is_limited() && !a.kind.needs_login() && !a.kind.needs_action())
.filter_map(|a| a.until_ms)
.filter(|t| *t > now_ms)
.min();
Self {
kind,
accounts,
earliest_reset_ms,
usage_read_retry_ms: None,
}
}
pub fn transient_retry_after_s(&self) -> Option<u64> {
let waits: Option<Vec<u64>> = self
.accounts
.iter()
.filter(|a| a.is_live() && a.kind.is_transient())
.map(|a| a.retry_after_s)
.collect();
waits?.into_iter().min()
}
pub fn login_labels(&self) -> Vec<&str> {
self.accounts
.iter()
.filter(|a| a.kind.needs_login())
.map(|a| a.account.as_str())
.collect()
}
pub fn action_labels(&self) -> Vec<&str> {
self.accounts
.iter()
.filter(|a| a.kind.needs_action())
.map(|a| a.account.as_str())
.collect()
}
pub fn accounts_field(&self) -> String {
self.accounts
.iter()
.map(AccountFailure::log_form)
.collect::<Vec<_>>()
.join(",")
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "outcome", rename_all = "snake_case")]
pub enum Selection {
Selected {
#[serde(with = "cred_serde")]
credential: CredentialRef,
#[serde(default, skip_serializing_if = "Option::is_none")]
utilization: Option<f64>,
#[serde(default)]
tier: u8,
#[serde(default, skip_serializing_if = "Option::is_none")]
budget_reset_ms: Option<u64>,
rejections: Vec<Rejection>,
},
NoCapacity {
rejections: Vec<Rejection>,
#[serde(default, skip_serializing_if = "Option::is_none")]
earliest_reset_ms: Option<u64>,
},
}
impl Selection {
pub fn selected(&self) -> Option<&CredentialRef> {
match self {
Self::Selected { credential, .. } => Some(credential),
Self::NoCapacity { .. } => None,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct Eligible {
pub utilization: Option<f64>,
pub budget_reset_ms: Option<u64>,
pub anchored: Option<bool>,
pub tier: u8,
pub graced: bool,
}
fn tier_for(
anchored: Option<bool>,
budget_reset_ms: Option<u64>,
now_ms: u64,
urgent_horizon_ms: u64,
) -> u8 {
match (anchored, budget_reset_ms) {
(_, None) => 4,
(Some(false), Some(_)) => 2,
(_, Some(reset)) if reset.saturating_sub(now_ms) <= urgent_horizon_ms => 1,
(_, Some(_)) => 3,
}
}
pub fn evaluate(
req: &SelectionRequest<'_>,
cap: &AccountCapacity,
) -> Result<Eligible, RejectReason> {
if req
.exclude
.iter()
.any(|c| c.storage_key() == cap.credential.storage_key())
{
return Err(RejectReason::Excluded);
}
if let Some(until) = cap.cooldown_until_ms {
if until > req.now_ms {
return Err(RejectReason::CoolingDown { until_ms: until });
}
}
let (windows, models) = match &cap.observation {
QuotaObservation::Ok { windows, models } => (windows, models),
QuotaObservation::AuthError => return Err(RejectReason::AuthError),
QuotaObservation::Unsupported => return Err(RejectReason::Unsupported),
QuotaObservation::Malformed => return Err(RejectReason::Malformed),
QuotaObservation::Unknown => return evaluate_graced(req, cap),
};
let observed = cap
.observed_at_ms
.ok_or(RejectReason::Stale { age_ms: None })?;
if observed > req.now_ms.saturating_add(MAX_FUTURE_OBSERVATION_SKEW_MS) {
return Err(RejectReason::Malformed);
}
let age = req.now_ms.saturating_sub(observed);
if age > req.max_snapshot_age_ms {
return Err(RejectReason::Stale { age_ms: Some(age) });
}
if let (Some(model), Some(list)) = (req.model, models) {
for m in list.iter().filter(|m| model_matches(&m.model, model)) {
match m.state {
ModelState::Available => {}
ModelState::Exhausted => {
return Err(RejectReason::ModelUnavailable {
model: model.to_string(),
})
}
ModelState::Unknown => {
return Err(RejectReason::ModelUnknown {
model: model.to_string(),
})
}
}
}
}
let mut applicable = 0usize;
let mut utilization: Option<f64> = None;
for w in windows.iter().filter(|w| w.applies_to(req.model)) {
applicable += 1;
match w.verdict(req.exhausted_at_percent) {
WindowVerdict::Headroom { used_percent } => {
if let Some(p) = used_percent {
utilization = Some(utilization.map_or(p, |u| u.max(p)));
}
}
WindowVerdict::Exhausted { resets_at_ms } => {
return Err(RejectReason::Exhausted {
window: w.id.clone(),
resets_at_ms,
});
}
WindowVerdict::Unknown => {
return Err(RejectReason::WindowUnknown {
window: w.id.clone(),
});
}
}
}
if applicable == 0 {
return Err(RejectReason::NoLimitEvidence);
}
let budget = budget_window(windows, req.model);
let budget_reset_ms = budget
.and_then(|w| w.resets_at_ms)
.filter(|t| *t > req.now_ms);
let anchored = match (budget, budget_reset_ms) {
(Some(w), Some(_)) => is_unanchored(w, observed).map(|u| !u),
_ => None,
};
let tier = tier_for(anchored, budget_reset_ms, req.now_ms, req.urgent_horizon_ms);
Ok(Eligible {
utilization,
budget_reset_ms,
anchored,
tier,
graced: false,
})
}
fn evaluate_graced(
req: &SelectionRequest<'_>,
cap: &AccountCapacity,
) -> Result<Eligible, RejectReason> {
let key = cap.credential.storage_key();
let Some(good) = req
.grace
.iter()
.find(|g| g.credential.storage_key() == key)
.filter(|g| matches!(g.observation, QuotaObservation::Ok { .. }))
else {
return Err(RejectReason::Unknown);
};
let young = good
.observed_at_ms
.is_some_and(|t| req.now_ms.saturating_sub(t) < req.grace_max_age_ms);
if !young {
return Err(RejectReason::Unknown);
}
let graced_req = SelectionRequest {
max_snapshot_age_ms: req.grace_max_age_ms,
grace: &[],
..req.clone()
};
let reading = AccountCapacity {
cooldown_until_ms: cap.cooldown_until_ms,
..good.clone()
};
evaluate(&graced_req, &reading)
.map(|e| Eligible { graced: true, ..e })
.map_err(|_| RejectReason::Unknown)
}
fn utilization_cmp(a: Option<f64>, b: Option<f64>) -> std::cmp::Ordering {
a.is_none().cmp(&b.is_none()).then_with(|| {
a.unwrap_or(0.0)
.partial_cmp(&b.unwrap_or(0.0))
.unwrap_or(std::cmp::Ordering::Equal)
})
}
fn preference_rank(req: &SelectionRequest<'_>, cred: &CredentialRef) -> usize {
let label = cred.account.label_str();
req.preference
.iter()
.position(|p| p == label)
.unwrap_or(usize::MAX)
}
fn observation_cause(reason: &RejectReason, cap: &AccountCapacity) -> Option<ObservationFailure> {
match reason {
RejectReason::Unknown
| RejectReason::AuthError
| RejectReason::Malformed
| RejectReason::Unsupported
| RejectReason::CoolingDown { .. } => cap.failure.clone(),
_ => None,
}
}
fn usable_after_ms(req: &SelectionRequest<'_>, cap: &AccountCapacity) -> Option<u64> {
let mut latest = cap.cooldown_until_ms.filter(|t| *t > req.now_ms);
if let QuotaObservation::Ok { windows, .. } = &cap.observation {
for w in windows.iter().filter(|w| w.applies_to(req.model)) {
if let WindowVerdict::Exhausted { resets_at_ms } = w.verdict(req.exhausted_at_percent) {
let reset = resets_at_ms.filter(|t| *t > req.now_ms)?;
latest = Some(latest.map_or(reset, |l| l.max(reset)));
}
}
}
latest
}
pub fn select(req: &SelectionRequest<'_>, candidates: &[AccountCapacity]) -> Selection {
let mut rejections: Vec<Rejection> = Vec::new();
let mut eligible: Vec<(&AccountCapacity, Eligible)> = Vec::new();
let mut explicit_seen = false;
for cap in candidates
.iter()
.filter(|c| c.credential.provider == req.provider)
{
if let Some(label) = req.explicit_account {
if cap.credential.account.label_str() != label {
continue;
}
explicit_seen = true;
}
match evaluate(req, cap) {
Ok(e) => eligible.push((cap, e)),
Err(reason) => rejections.push(Rejection {
credential: cap.credential.clone(),
cause: observation_cause(&reason, cap),
usable_after_ms: usable_after_ms(req, cap),
reason,
}),
}
}
if let (Some(label), false) = (req.explicit_account, explicit_seen) {
let (credential, reason) = match super::account::Account::parse(label) {
Ok(account) => (
CredentialRef::new(req.provider, account),
RejectReason::NotPresent,
),
Err(_) => (
CredentialRef::default_for(req.provider),
RejectReason::InvalidLabel {
label: label.to_string(),
},
),
};
return Selection::NoCapacity {
rejections: vec![Rejection {
credential,
reason,
cause: None,
usable_after_ms: None,
}],
earliest_reset_ms: None,
};
}
eligible.sort_by(|(a, ea), (b, eb)| {
let ka = a.credential.storage_key();
let kb = b.credential.storage_key();
let pa = preference_rank(req, &a.credential);
let pb = preference_rank(req, &b.credential);
ea.graced.cmp(&eb.graced).then_with(|| match req.strategy {
Strategy::PreferenceOrder => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
Strategy::LowestUtilization => utilization_cmp(ea.utilization, eb.utilization)
.then_with(|| pa.cmp(&pb))
.then_with(|| ka.cmp(&kb)),
Strategy::SoonestReset => ea
.tier
.cmp(&eb.tier)
.then_with(|| match ea.tier {
1 | 3 => ea.budget_reset_ms.cmp(&eb.budget_reset_ms),
2 => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
_ => utilization_cmp(ea.utilization, eb.utilization),
})
.then_with(|| utilization_cmp(ea.utilization, eb.utilization))
.then_with(|| pa.cmp(&pb))
.then_with(|| ka.cmp(&kb)),
})
});
if req.sticky && req.strategy == Strategy::SoonestReset {
if let Some(current) = req.current {
let key = current.storage_key();
let at = eligible
.iter()
.position(|(c, _)| c.credential.storage_key() == key);
if let Some(i) = at.filter(|i| *i > 0) {
let best = &eligible[0].1;
let cur = &eligible[i].1;
let switch = best.tier < cur.tier
|| (best.tier == 1
&& cur.tier == 1
&& best.budget_reset_ms < cur.budget_reset_ms);
if !switch {
eligible[..=i].rotate_right(1);
}
}
}
}
match eligible.first() {
Some((cap, e)) => {
let winner = cap.credential.clone();
for (i, (other, _)) in eligible.iter().enumerate().skip(1) {
rejections.push(Rejection {
credential: other.credential.clone(),
reason: RejectReason::Outranked {
rank: (i + 1) as u32,
},
cause: None,
usable_after_ms: None,
});
}
Selection::Selected {
credential: winner,
utilization: e.utilization,
tier: e.tier,
budget_reset_ms: e.budget_reset_ms,
rejections,
}
}
None => {
let earliest_reset_ms = rejections
.iter()
.filter_map(|r| match &r.reason {
RejectReason::Exhausted {
resets_at_ms: Some(t),
..
} if *t > req.now_ms => Some(*t),
_ => None,
})
.min();
Selection::NoCapacity {
rejections,
earliest_reset_ms,
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FailureKind {
QuotaExhaustedPreOutput,
RateLimited,
Other,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FailureEvidence {
pub kind: FailureKind,
pub output_started: bool,
pub tool_activity: bool,
pub failovers_so_far: u32,
}
pub const MAX_FAILOVERS: u32 = 1;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "decision", rename_all = "snake_case")]
pub enum FailoverDecision {
Failover {
#[serde(with = "cred_serde")]
credential: CredentialRef,
},
NoFailover {
reason: String,
},
}
pub fn failover(
req: &SelectionRequest<'_>,
failed: &CredentialRef,
evidence: &FailureEvidence,
candidates: &[AccountCapacity],
) -> FailoverDecision {
let no = |reason: &str| FailoverDecision::NoFailover {
reason: reason.to_string(),
};
if evidence.output_started {
return no("output already started; never replay on another account");
}
if evidence.tool_activity {
return no("tool activity occurred; never replay on another account");
}
if evidence.failovers_so_far >= MAX_FAILOVERS {
return no("failover budget exhausted");
}
if evidence.kind != FailureKind::QuotaExhaustedPreOutput {
return no("not a recognized pre-output quota failure");
}
if req.explicit_account.is_some() {
return no("explicit account requested; no silent fallback");
}
let mut exclude: Vec<CredentialRef> = req.exclude.to_vec();
exclude.push(failed.clone());
let retry = SelectionRequest {
exclude: &exclude,
..req.clone()
};
match select(&retry, candidates) {
Selection::Selected { credential, .. } => FailoverDecision::Failover { credential },
Selection::NoCapacity { .. } => no("no other eligible account"),
}
}
fn pick_weekly<'w>(windows: impl Iterator<Item = &'w WindowLimit>) -> Option<&'w WindowLimit> {
let mut weekly: Vec<&WindowLimit> = windows.filter(|w| w.is_weekly()).collect();
weekly.sort_by(|a, b| {
a.models
.is_some()
.cmp(&b.models.is_some())
.then_with(|| a.id.cmp(&b.id))
});
weekly.into_iter().next()
}
pub fn weekly_window(windows: &[WindowLimit]) -> Option<&WindowLimit> {
pick_weekly(windows.iter())
}
pub fn budget_window<'w>(
windows: &'w [WindowLimit],
model: Option<&str>,
) -> Option<&'w WindowLimit> {
let applicable = || windows.iter().filter(|w| w.applies_to(model));
if let Some(weekly) = pick_weekly(applicable()) {
return Some(weekly);
}
applicable()
.filter(|w| w.duration_ms.is_some())
.min_by(|a, b| {
b.duration_ms
.cmp(&a.duration_ms)
.then_with(|| a.models.is_some().cmp(&b.models.is_some()))
.then_with(|| a.id.cmp(&b.id))
})
}
pub fn is_unanchored(window: &WindowLimit, observed_at_ms: u64) -> Option<bool> {
let (Some(duration), Some(reset)) = (window.duration_ms, window.resets_at_ms) else {
return None;
};
let near_zero = matches!(
window.used_percent,
Some(p) if (0.0..=UNANCHORED_MAX_USED_PERCENT).contains(&p)
);
if !near_zero {
return Some(false);
}
let expected = observed_at_ms.saturating_add(duration);
Some(expected.abs_diff(reset) <= UNANCHORED_TOLERANCE_MS)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::account::Account;
const HOUR: u64 = 3_600_000;
const DAY: u64 = 24 * HOUR;
const NOW: u64 = 1_800_000_000_000;
fn cred(label: &str) -> CredentialRef {
CredentialRef::new(OAuthProviderId::OpenAiCodex, Account::parse(label).unwrap())
}
fn win(id: &str, dur: u64, used: f64, resets: u64) -> WindowLimit {
WindowLimit {
id: id.into(),
duration_ms: Some(dur),
used_percent: Some(used),
limit_reached: Some(used >= 100.0),
resets_at_ms: Some(resets),
models: None,
}
}
fn ok(label: &str, windows: Vec<WindowLimit>) -> AccountCapacity {
AccountCapacity {
credential: cred(label),
observed_at_ms: Some(NOW - 1000),
observation: QuotaObservation::Ok {
windows,
models: None,
},
cooldown_until_ms: None,
failure: None,
}
}
fn req() -> SelectionRequest<'static> {
SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
}
#[test]
fn window_verdicts_fail_closed() {
let mut w = win("primary", 7 * DAY, 20.0, NOW + DAY);
assert_eq!(
w.verdict(100.0),
WindowVerdict::Headroom {
used_percent: Some(20.0)
}
);
w.used_percent = Some(100.0);
w.limit_reached = Some(false);
assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
w.used_percent = Some(50.0);
w.limit_reached = Some(true);
assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
w.used_percent = None;
w.limit_reached = None;
assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
w.limit_reached = Some(false);
assert_eq!(
w.verdict(100.0),
WindowVerdict::Headroom { used_percent: None }
);
w.used_percent = Some(f64::NAN);
assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
w.used_percent = Some(140.0);
assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
w.used_percent = Some(1000.5);
assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
let saved = w.limit_reached;
w.limit_reached = Some(true);
for p in [1500.0, f64::NAN, -3.0] {
w.used_percent = Some(p);
assert!(
matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }),
"{p}"
);
}
w.limit_reached = saved;
w.used_percent = Some(-1.0);
assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
w.used_percent = Some(1.0);
assert_eq!(w.verdict(f64::NAN), WindowVerdict::Unknown);
assert_eq!(w.verdict(101.0), WindowVerdict::Unknown);
assert_eq!(w.verdict(-0.5), WindowVerdict::Unknown);
assert_eq!(w.verdict(f64::INFINITY), WindowVerdict::Unknown);
let r = SelectionRequest {
exhausted_at_percent: f64::NAN,
..SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, HOUR)
};
assert!(select(
&r,
&[ok("a", vec![win("primary", 7 * DAY, 1.0, NOW + DAY)])]
)
.selected()
.is_none());
}
#[test]
fn weekly_is_from_duration_not_name() {
let primary_5h = win("primary", 5 * HOUR, 0.0, NOW + HOUR);
let primary_week = win("primary", 7 * DAY, 0.0, NOW + DAY);
let no_dur = WindowLimit {
duration_ms: None,
..win("seven_day", 0, 0.0, NOW)
};
assert!(!primary_5h.is_weekly());
assert!(primary_week.is_weekly());
assert!(!no_dur.is_weekly());
assert_eq!(
weekly_window(&[primary_5h.clone(), primary_week.clone()]).map(|w| w.duration_ms),
Some(Some(7 * DAY))
);
assert!(weekly_window(&[primary_5h, no_dur]).is_none());
let scoped = WindowLimit {
id: "aaa_scoped".into(),
models: Some(vec!["m".into()]),
..primary_week.clone()
};
assert_eq!(
weekly_window(&[scoped, primary_week]).map(|w| w.id.as_str()),
Some("primary")
);
}
#[test]
fn selects_lowest_utilization_deterministically() {
let a = ok("a", vec![win("primary", 7 * DAY, 60.0, NOW + DAY)]);
let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
let sel = select(&req(), &[a.clone(), c.clone(), b.clone()]);
assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@b");
let sel2 = select(&req(), &[c, b, a]);
assert_eq!(sel2.selected().unwrap().storage_key(), "openai-codex@b");
}
#[test]
fn preference_breaks_ties_and_orders() {
let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
let pref = vec!["c".to_string(), "b".to_string()];
let r = SelectionRequest {
preference: &pref,
..req()
};
assert_eq!(
select(&r, &[b.clone(), c.clone()])
.selected()
.unwrap()
.storage_key(),
"openai-codex@c"
);
let a = ok("a", vec![win("primary", 7 * DAY, 90.0, NOW + DAY)]);
let pref2 = vec!["a".to_string()];
let r2 = SelectionRequest {
preference: &pref2,
strategy: Strategy::PreferenceOrder,
..req()
};
assert_eq!(
select(&r2, &[b, c, a]).selected().unwrap().storage_key(),
"openai-codex@a"
);
}
#[test]
fn graced_readings_are_bounded_rank_after_fresh_and_keep_a_sticky_seat() {
let grace_ms = super::super::usage_reads::USAGE_READ_GRACE.as_millis() as u64;
let failed = |label: &str| AccountCapacity::unobserved(cred(label));
let good = |label: &str, age: u64, used: f64, reset: u64| AccountCapacity {
observed_at_ms: Some(NOW - age),
..ok(label, vec![win("primary", 7 * DAY, used, reset)])
};
let fresh_req = || SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 60_000);
let winner = |sel: Selection| match sel {
Selection::Selected { credential, .. } => credential.account.label_str().to_string(),
Selection::NoCapacity { rejections, .. } => panic!("no capacity: {rejections:?}"),
};
let one = |grace: &[AccountCapacity], cap: &AccountCapacity| {
evaluate(&SelectionRequest { grace, ..fresh_req() }, cap)
};
let young = [good("a", grace_ms - 60_000, 20.0, NOW + DAY)];
assert!(one(&young, &failed("a")).unwrap().graced);
let old = [good("a", grace_ms + 60_000, 20.0, NOW + DAY)];
assert_eq!(one(&old, &failed("a")), Err(RejectReason::Unknown), "past the grace age");
let edge = [good("a", grace_ms, 20.0, NOW + DAY)];
assert_eq!(one(&edge, &failed("a")), Err(RejectReason::Unknown), "at the grace age");
let inside = [good("a", grace_ms - 1, 20.0, NOW + DAY)];
assert!(one(&inside, &failed("a")).unwrap().graced, "just under it");
let exhausted = [good("a", 60_000, 100.0, NOW + HOUR)];
assert_eq!(one(&exhausted, &failed("a")), Err(RejectReason::Unknown), "exhausted window");
let cooling = AccountCapacity {
cooldown_until_ms: Some(NOW + HOUR),
..failed("a")
};
assert!(matches!(one(&young, &cooling), Err(RejectReason::CoolingDown { .. })));
let auth = AccountCapacity {
observation: QuotaObservation::AuthError,
..failed("a")
};
assert_eq!(one(&young, &auth), Err(RejectReason::AuthError), "only a failed read is graced");
let grace = [good("a", 60_000, 20.0, NOW + HOUR)];
let b = good("b", 1000, 20.0, NOW + 3 * DAY);
let anew = SelectionRequest { grace: &grace, ..fresh_req() };
assert_eq!(winner(select(&anew, &[failed("a"), b.clone()])), "b");
assert_eq!(winner(select(&anew, &[failed("a")])), "a", "alone, the graced seat serves");
let current = cred("a");
let sticky = SelectionRequest {
current: Some(¤t),
..anew.clone()
};
assert_eq!(winner(select(&sticky, &[failed("a"), b])), "a");
let mut switched = 0;
for a_reset in [NOW + HOUR, NOW + 2 * DAY, NOW + 10 * HOUR] {
for b_reset in [NOW + HOUR / 2, NOW + 3 * DAY, NOW + 12 * HOUR, NOW + 5 * DAY] {
let a_good = good("a", 1000, 20.0, a_reset);
let b = good("b", 1000, 20.0, b_reset);
let graced = [a_good.clone()];
let on_grace = SelectionRequest {
grace: &graced,
current: Some(¤t),
..fresh_req()
};
let on_fresh = SelectionRequest {
current: Some(¤t),
..fresh_req()
};
let g = winner(select(&on_grace, &[failed("a"), b.clone()]));
let f = winner(select(&on_fresh, &[a_good, b]));
assert_eq!(g, f, "a resets at {a_reset}, b at {b_reset}");
switched += usize::from(g == "b");
}
}
assert!(switched > 0 && switched < 12, "the cases cover staying and switching: {switched}");
}
#[test]
fn stale_unknown_malformed_auth_are_not_capacity() {
let mut stale = ok("s", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
stale.observed_at_ms = Some(NOW - 2 * HOUR);
let mut never = ok("n", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
never.observed_at_ms = None;
let mut future = ok("f", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
future.observed_at_ms = Some(NOW + HOUR);
let mut auth = ok("x", vec![]);
auth.observation = QuotaObservation::AuthError;
let mut unk = ok("u", vec![]);
unk.observation = QuotaObservation::Unknown;
let mut mal = ok("m", vec![]);
mal.observation = QuotaObservation::Malformed;
let mut uns = ok("v", vec![]);
uns.observation = QuotaObservation::Unsupported;
let none = ok("e", vec![]);
let sel = select(&req(), &[stale, never, future, auth, unk, mal, uns, none]);
let Selection::NoCapacity {
rejections,
earliest_reset_ms,
} = sel
else {
panic!("expected no capacity");
};
assert_eq!(earliest_reset_ms, None);
let reasons: Vec<(String, RejectReason)> = rejections
.into_iter()
.map(|r| (r.credential.account.label_str().to_string(), r.reason))
.collect();
assert!(reasons.contains(&(
"s".into(),
RejectReason::Stale {
age_ms: Some(2 * HOUR)
}
)));
assert!(reasons.contains(&("n".into(), RejectReason::Stale { age_ms: None })));
assert!(reasons.contains(&("f".into(), RejectReason::Malformed)));
assert!(reasons.contains(&("x".into(), RejectReason::AuthError)));
assert!(reasons.contains(&("u".into(), RejectReason::Unknown)));
assert!(reasons.contains(&("m".into(), RejectReason::Malformed)));
assert!(reasons.contains(&("v".into(), RejectReason::Unsupported)));
assert!(reasons.contains(&("e".into(), RejectReason::NoLimitEvidence)));
}
#[test]
fn exhausted_cooldown_and_unknown_window_excluded() {
let exhausted = ok("x", vec![win("primary", 7 * DAY, 100.0, NOW + 2 * DAY)]);
let mut cooling = ok("c", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
cooling.cooldown_until_ms = Some(NOW + HOUR);
let mut cooled = ok("d", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
cooled.cooldown_until_ms = Some(NOW - 1);
let unknown_window = ok(
"w",
vec![WindowLimit {
used_percent: None,
limit_reached: None,
..win("primary", 7 * DAY, 0.0, NOW + DAY)
}],
);
let past_exhausted = ok("p", vec![win("primary", 7 * DAY, 100.0, NOW - DAY)]);
let sel = select(
&req(),
&[exhausted, cooling, unknown_window, past_exhausted],
);
let Selection::NoCapacity {
earliest_reset_ms, ..
} = &sel
else {
panic!("no capacity expected");
};
assert_eq!(*earliest_reset_ms, Some(NOW + 2 * DAY));
let sel = select(&req(), &[cooled]);
assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@d");
}
#[test]
fn all_applicable_limits_including_model_scoped_count() {
let model_win = WindowLimit {
id: "weekly_sonnet".into(),
models: Some(vec!["sonnet".into()]),
..win("weekly_sonnet", 7 * DAY, 100.0, NOW + DAY)
};
let five_h = win("five_hour", 5 * HOUR, 99.0, NOW + HOUR);
let acct = ok(
"a",
vec![
win("seven_day", 7 * DAY, 10.0, NOW + DAY),
five_h,
model_win,
],
);
let r = SelectionRequest {
model: Some("sonnet"),
..req()
};
assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
let r = SelectionRequest {
model: Some("opus"),
..req()
};
match select(&r, std::slice::from_ref(&acct)) {
Selection::Selected { utilization, .. } => assert_eq!(utilization, Some(99.0)),
other => panic!("unexpected {other:?}"),
}
assert!(select(&req(), std::slice::from_ref(&acct))
.selected()
.is_none());
let mut unavailable = acct;
if let QuotaObservation::Ok { models, .. } = &mut unavailable.observation {
*models = Some(vec![ModelAvailability {
model: "opus".into(),
state: ModelState::Exhausted,
}]);
}
let r = SelectionRequest {
model: Some("opus"),
..req()
};
let sel = select(&r, &[unavailable.clone()]);
match sel {
Selection::NoCapacity { rejections, .. } => assert_eq!(
rejections[0].reason,
RejectReason::ModelUnavailable {
model: "opus".into()
}
),
other => panic!("unexpected {other:?}"),
}
let mut unknown = unavailable;
if let QuotaObservation::Ok { models, .. } = &mut unknown.observation {
*models = Some(vec![ModelAvailability {
model: "opus".into(),
state: ModelState::Unknown,
}]);
}
match select(&r, &[unknown]) {
Selection::NoCapacity { rejections, .. } => assert_eq!(
rejections[0].reason,
RejectReason::ModelUnknown {
model: "opus".into()
}
),
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn explicit_account_never_falls_back() {
let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
let r = SelectionRequest {
explicit_account: Some("a"),
..req()
};
let sel = select(&r, &[a.clone(), b.clone()]);
assert!(sel.selected().is_none());
let r = SelectionRequest {
explicit_account: Some("zzz"),
..req()
};
match select(&r, &[a.clone(), b.clone()]) {
Selection::NoCapacity { rejections, .. } => {
assert_eq!(rejections.len(), 1);
assert_eq!(rejections[0].reason, RejectReason::NotPresent);
assert_eq!(rejections[0].credential.storage_key(), "openai-codex@zzz");
}
other => panic!("unexpected {other:?}"),
}
let r = SelectionRequest {
explicit_account: Some("Bad Label"),
..req()
};
match select(&r, &[a, b]) {
Selection::NoCapacity { rejections, .. } => assert_eq!(
rejections[0].reason,
RejectReason::InvalidLabel {
label: "Bad Label".into()
}
),
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn other_providers_and_exclusions_are_ignored() {
let mut other = ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
other.credential =
CredentialRef::new(OAuthProviderId::Anthropic, Account::parse("a").unwrap());
let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
let excl = vec![cred("b")];
let r = SelectionRequest {
exclude: &excl,
..req()
};
let sel = select(&r, &[other, b]);
match sel {
Selection::NoCapacity { rejections, .. } => {
assert_eq!(rejections.len(), 1);
assert_eq!(rejections[0].reason, RejectReason::Excluded);
}
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn failover_is_bounded_and_never_after_output() {
let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
let c = ok("c", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
let cands = [a, b, c];
let ev = FailureEvidence {
kind: FailureKind::QuotaExhaustedPreOutput,
output_started: false,
tool_activity: false,
failovers_so_far: 0,
};
match failover(&req(), &cred("a"), &ev, &cands) {
FailoverDecision::Failover { credential } => {
assert_eq!(credential.storage_key(), "openai-codex@b")
}
other => panic!("unexpected {other:?}"),
}
let ev2 = FailureEvidence {
failovers_so_far: 1,
..ev.clone()
};
assert!(matches!(
failover(&req(), &cred("b"), &ev2, &cands),
FailoverDecision::NoFailover { .. }
));
let ev3 = FailureEvidence {
output_started: true,
..ev.clone()
};
assert!(matches!(
failover(&req(), &cred("a"), &ev3, &cands),
FailoverDecision::NoFailover { .. }
));
let ev4 = FailureEvidence {
tool_activity: true,
..ev.clone()
};
assert!(matches!(
failover(&req(), &cred("a"), &ev4, &cands),
FailoverDecision::NoFailover { .. }
));
let ev5 = FailureEvidence {
kind: FailureKind::RateLimited,
..ev.clone()
};
assert!(matches!(
failover(&req(), &cred("a"), &ev5, &cands),
FailoverDecision::NoFailover { .. }
));
let r = SelectionRequest {
explicit_account: Some("a"),
..req()
};
assert!(matches!(
failover(&r, &cred("a"), &ev, &cands),
FailoverDecision::NoFailover { .. }
));
let only = [ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)])];
assert!(matches!(
failover(&req(), &cred("a"), &ev, &only),
FailoverDecision::NoFailover { .. }
));
}
#[test]
fn unknown_utilization_sorts_after_known() {
let proven = ok(
"p",
vec![WindowLimit {
used_percent: None,
limit_reached: Some(false),
..win("primary", 7 * DAY, 0.0, NOW + DAY)
}],
);
let known = ok("k", vec![win("primary", 7 * DAY, 95.0, NOW + DAY)]);
assert_eq!(
select(&req(), &[proven.clone(), known])
.selected()
.unwrap()
.storage_key(),
"openai-codex@k"
);
assert_eq!(
select(&req(), &[proven]).selected().unwrap().storage_key(),
"openai-codex@p"
);
}
#[test]
fn model_matching_uses_family_tokens() {
assert!(model_matches("sonnet", "claude-sonnet-4-6"));
assert!(model_matches("Sonnet", "CLAUDE-SONNET-4-6"));
assert!(model_matches("gpt-6-astra", "gpt-6-astra"));
assert!(!model_matches("opus", "claude-sonnet-4-6"));
assert!(!model_matches("", "x") && !model_matches("x", ""));
let w = WindowLimit {
models: Some(vec!["sonnet".into()]),
..win("seven_day_sonnet", 7 * DAY, 100.0, NOW + DAY)
};
assert!(w.applies_to(Some("claude-sonnet-4-6")));
assert!(!w.applies_to(Some("claude-opus-4-6")));
assert!(w.applies_to(None));
let acct = ok("a", vec![win("seven_day", 7 * DAY, 1.0, NOW + DAY), w]);
let r = SelectionRequest {
model: Some("claude-sonnet-4-6"),
..req()
};
assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
let r = SelectionRequest {
model: Some("claude-opus-4-6"),
..req()
};
assert!(select(&r, std::slice::from_ref(&acct)).selected().is_some());
}
#[test]
fn serde_round_trip_is_secret_free() {
let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
let sel = select(&req(), &[a]);
let json = serde_json::to_string(&sel).unwrap();
assert!(json.contains("no_capacity"));
let back: Selection = serde_json::from_str(&json).unwrap();
assert_eq!(back, sel);
}
fn sr_req() -> SelectionRequest<'static> {
SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
}
const OBSERVED: u64 = NOW - 1000;
fn anchored(used: f64, resets: u64) -> WindowLimit {
win("primary", 7 * DAY, used, resets)
}
fn unanchored() -> WindowLimit {
win("primary", 7 * DAY, 0.0, OBSERVED + 7 * DAY)
}
fn no_reset(used: f64) -> WindowLimit {
WindowLimit {
resets_at_ms: None,
..win("primary", 7 * DAY, used, 0)
}
}
fn label_of(sel: &Selection) -> &str {
sel.selected().expect("a selection").account.label_str()
}
fn ranks(sel: &Selection) -> Vec<(String, u32)> {
let (Selection::Selected { rejections, .. } | Selection::NoCapacity { rejections, .. }) =
sel;
rejections
.iter()
.filter_map(|r| match r.reason {
RejectReason::Outranked { rank } => {
Some((r.credential.account.label_str().to_string(), rank))
}
_ => None,
})
.collect()
}
#[test]
fn soonest_reset_beats_lower_utilization_where_lowest_utilization_does_not() {
let a = ok("a", vec![anchored(60.0, NOW + 2 * DAY)]);
let b = ok("b", vec![anchored(10.0, NOW + 5 * DAY)]);
let cands = [b.clone(), a.clone()];
match select(&sr_req(), &cands) {
Selection::Selected {
credential,
tier,
budget_reset_ms,
utilization,
..
} => {
assert_eq!(credential.storage_key(), "openai-codex@a");
assert_eq!(tier, 3);
assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
assert_eq!(utilization, Some(60.0));
}
other => panic!("unexpected {other:?}"),
}
assert_eq!(label_of(&select(&req(), &cands)), "b");
let r = req();
assert_eq!(r.strategy, Strategy::LowestUtilization);
assert!(!r.sticky && r.current.is_none());
assert_eq!(r.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
let s = sr_req();
assert_eq!(s.strategy, Strategy::SoonestReset);
assert!(s.sticky && s.current.is_none());
assert_eq!(s.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
}
#[test]
fn tiers_order_urgent_then_unanchored_then_far_then_no_reset() {
let urgent = ok("u", vec![anchored(50.0, NOW + 6 * HOUR)]);
let sliding = ok("s", vec![unanchored()]);
let far = ok("f", vec![anchored(0.0, NOW + 5 * DAY)]);
let none = ok("n", vec![no_reset(0.0)]);
let r = sr_req();
for (cap, tier, anch) in [
(&urgent, 1u8, Some(true)),
(&sliding, 2, Some(false)),
(&far, 3, Some(true)),
(&none, 4, None),
] {
let e = evaluate(&r, cap).unwrap();
assert_eq!(e.tier, tier, "{}", cap.credential.storage_key());
assert_eq!(e.anchored, anch, "{}", cap.credential.storage_key());
}
assert_eq!(evaluate(&req(), &sliding).unwrap().tier, 2);
let cands = [none.clone(), far.clone(), sliding.clone(), urgent.clone()];
let sel = select(&r, &cands);
assert_eq!(label_of(&sel), "u");
assert_eq!(
ranks(&sel),
vec![
("s".to_string(), 2),
("f".to_string(), 3),
("n".to_string(), 4)
]
);
assert_eq!(
label_of(&select(&r, &[none.clone(), far.clone(), sliding])),
"s"
);
assert_eq!(label_of(&select(&r, &[none.clone(), far])), "f");
assert_eq!(label_of(&select(&r, &[none])), "n");
let short = SelectionRequest {
urgent_horizon_ms: HOUR,
..sr_req()
};
assert_eq!(evaluate(&short, &urgent).unwrap().tier, 3);
assert_eq!(
label_of(&select(&short, &[urgent, ok("s", vec![unanchored()])])),
"s"
);
let edge = ok("e", vec![anchored(0.0, NOW + DEFAULT_URGENT_HORIZON_MS)]);
assert_eq!(evaluate(&sr_req(), &edge).unwrap().tier, 1);
}
#[test]
fn within_tier_orders_by_reset_preference_or_utilization() {
let a = ok("a", vec![anchored(90.0, NOW + 2 * HOUR)]);
let b = ok("b", vec![anchored(1.0, NOW + 3 * HOUR)]);
assert_eq!(label_of(&select(&sr_req(), &[b, a])), "a");
let s1 = ok("s1", vec![unanchored()]);
let s2 = ok("s2", vec![unanchored()]);
assert_eq!(
label_of(&select(&sr_req(), &[s2.clone(), s1.clone()])),
"s1"
);
let pref = vec!["s2".to_string()];
let r = SelectionRequest {
preference: &pref,
..sr_req()
};
assert_eq!(label_of(&select(&r, &[s1, s2])), "s2");
let n1 = ok("n1", vec![no_reset(40.0)]);
let n2 = ok("n2", vec![no_reset(5.0)]);
let n3 = ok(
"n3",
vec![WindowLimit {
used_percent: None,
limit_reached: Some(false),
..no_reset(0.0)
}],
);
let sel = select(&sr_req(), &[n3, n1, n2]);
assert_eq!(label_of(&sel), "n2");
assert_eq!(
ranks(&sel),
vec![("n1".to_string(), 2), ("n3".to_string(), 3)]
);
let t1 = ok("t1", vec![anchored(30.0, NOW + 3 * DAY)]);
let t2 = ok("t2", vec![anchored(20.0, NOW + 3 * DAY)]);
assert_eq!(label_of(&select(&sr_req(), &[t1, t2])), "t2");
}
#[test]
fn exhausted_soonest_seat_is_skipped_to_next_soonest() {
let x = ok("x", vec![anchored(100.0, NOW + HOUR)]);
let a = ok("a", vec![anchored(20.0, NOW + 2 * DAY)]);
let b = ok("b", vec![anchored(20.0, NOW + 3 * DAY)]);
let sel = select(&sr_req(), &[b, x.clone(), a]);
assert_eq!(label_of(&sel), "a");
let Selection::Selected { rejections, .. } = &sel else {
panic!("expected selection");
};
assert!(rejections
.iter()
.any(|r| r.credential.account.label_str() == "x"
&& r.reason
== RejectReason::Exhausted {
window: "primary".into(),
resets_at_ms: Some(NOW + HOUR),
}));
let y = ok("y", vec![anchored(100.0, NOW + 2 * DAY)]);
let past = ok("p", vec![anchored(100.0, NOW - HOUR)]);
match select(&sr_req(), &[y, past, x]) {
Selection::NoCapacity {
earliest_reset_ms, ..
} => assert_eq!(earliest_reset_ms, Some(NOW + HOUR)),
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn five_hour_throttle_rejects_then_reenters_after_reset() {
let throttled = ok(
"t",
vec![
win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
],
);
let other = ok("o", vec![win("seven_day", 7 * DAY, 10.0, NOW + 4 * DAY)]);
let sel = select(&sr_req(), &[throttled, other.clone()]);
assert_eq!(label_of(&sel), "o");
let Selection::Selected { rejections, .. } = &sel else {
panic!("expected selection");
};
assert_eq!(
rejections[0].reason,
RejectReason::Exhausted {
window: "five_hour".into(),
resets_at_ms: Some(NOW + HOUR),
}
);
let now2 = NOW + HOUR + 1;
let mut throttled2 = ok(
"t",
vec![
win("five_hour", 5 * HOUR, 0.0, now2 + 5 * HOUR),
win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
],
);
throttled2.observed_at_ms = Some(now2 - 1000);
let mut other2 = other;
other2.observed_at_ms = Some(now2 - 1000);
let r = SelectionRequest {
now_ms: now2,
..sr_req()
};
match select(&r, &[other2, throttled2]) {
Selection::Selected {
credential,
tier,
budget_reset_ms,
..
} => {
assert_eq!(credential.storage_key(), "openai-codex@t");
assert_eq!(tier, 3);
assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
}
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn is_unanchored_classifies_sliding_resets_only() {
let obs = NOW - 90_000;
let codex_idle = win("primary", 7 * DAY, 0.0, obs + 7 * DAY);
assert_eq!(is_unanchored(&codex_idle, obs), Some(true));
let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY + 60_000);
assert_eq!(is_unanchored(&skewed, obs), Some(true));
let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY - 60_000);
assert_eq!(is_unanchored(&skewed, obs), Some(true));
let beyond = win(
"primary",
7 * DAY,
0.0,
obs + 7 * DAY + UNANCHORED_TOLERANCE_MS + 1,
);
assert_eq!(is_unanchored(&beyond, obs), Some(false));
let anthropic = win("seven_day", 7 * DAY, 0.0, obs + 5 * DAY);
assert_eq!(is_unanchored(&anthropic, obs), Some(false));
let used = win("primary", 7 * DAY, 40.0, obs + 7 * DAY);
assert_eq!(is_unanchored(&used, obs), Some(false));
let half = win(
"primary",
7 * DAY,
UNANCHORED_MAX_USED_PERCENT,
obs + 7 * DAY,
);
assert_eq!(is_unanchored(&half, obs), Some(true));
let unreported = WindowLimit {
used_percent: None,
limit_reached: Some(false),
..codex_idle.clone()
};
assert_eq!(is_unanchored(&unreported, obs), Some(false));
let no_dur = WindowLimit {
duration_ms: None,
..codex_idle.clone()
};
assert_eq!(is_unanchored(&no_dur, obs), None);
let no_reset = WindowLimit {
resets_at_ms: None,
..codex_idle
};
assert_eq!(is_unanchored(&no_reset, obs), None);
}
#[test]
fn budget_window_prefers_weekly_then_longest() {
let five_h = win("five_hour", 5 * HOUR, 0.0, NOW + HOUR);
let one_h = win("one_hour", HOUR, 0.0, NOW + HOUR);
let thirty_d = win("thirty_day", 30 * DAY, 0.0, NOW + 20 * DAY);
let weekly_wide = win("weekly", 7 * DAY, 0.0, NOW + 3 * DAY);
let weekly_scoped = WindowLimit {
id: "aaa_weekly_scoped".into(),
models: Some(vec!["m".into()]),
..weekly_wide.clone()
};
let all = [
thirty_d.clone(),
weekly_scoped.clone(),
five_h.clone(),
weekly_wide.clone(),
];
assert_eq!(
budget_window(&all, Some("m")).map(|w| w.id.as_str()),
Some("weekly")
);
assert_eq!(
budget_window(&all, None).map(|w| w.id.as_str()),
Some("weekly")
);
let no_wide = [thirty_d.clone(), weekly_scoped.clone(), five_h.clone()];
assert_eq!(
budget_window(&no_wide, Some("m")).map(|w| w.id.as_str()),
Some("aaa_weekly_scoped")
);
assert_eq!(
budget_window(&no_wide, Some("other")).map(|w| w.id.as_str()),
Some("thirty_day")
);
assert_eq!(
budget_window(&[five_h.clone(), one_h.clone()], None).map(|w| w.id.as_str()),
Some("five_hour")
);
let b5 = WindowLimit {
id: "b5".into(),
..five_h.clone()
};
let a5_scoped = WindowLimit {
id: "a5".into(),
models: Some(vec!["m".into()]),
..five_h.clone()
};
assert_eq!(
budget_window(&[b5.clone(), a5_scoped.clone()], Some("m")).map(|w| w.id.as_str()),
Some("b5")
);
let a5 = WindowLimit {
id: "a5".into(),
..five_h.clone()
};
assert_eq!(
budget_window(&[b5, a5], None).map(|w| w.id.as_str()),
Some("a5")
);
let no_dur = WindowLimit {
duration_ms: None,
..five_h
};
assert!(budget_window(&[no_dur.clone(), no_dur], None).is_none());
assert!(budget_window(&[], None).is_none());
assert_eq!(weekly_window(&all).map(|w| w.id.as_str()), Some("weekly"));
}
#[test]
fn past_reset_with_headroom_is_not_reset_evidence() {
let stale_reset = ok("p", vec![anchored(20.0, NOW - 1)]);
let e = evaluate(&sr_req(), &stale_reset).unwrap();
assert_eq!(e.budget_reset_ms, None);
assert_eq!(e.anchored, None);
assert_eq!(e.tier, 4);
let far = ok("f", vec![anchored(80.0, NOW + 6 * DAY)]);
assert_eq!(label_of(&select(&sr_req(), &[stale_reset, far])), "f");
}
#[test]
fn sticky_keeps_current_within_tier_3() {
let cur = cred("cur");
let current = ok("cur", vec![anchored(30.0, NOW + 5 * DAY)]);
let sooner = ok("soon", vec![anchored(10.0, NOW + 2 * DAY)]);
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
let sel = select(&r, &[sooner.clone(), current.clone()]);
assert_eq!(label_of(&sel), "cur");
assert_eq!(ranks(&sel), vec![("soon".to_string(), 2)]);
assert_eq!(
label_of(&select(&sr_req(), &[sooner.clone(), current.clone()])),
"soon"
);
let r = SelectionRequest { sticky: false, ..r };
assert_eq!(label_of(&select(&r, &[sooner, current])), "soon");
}
#[test]
fn sticky_yields_to_strictly_higher_tier() {
let cur = cred("cur");
let current = ok("cur", vec![anchored(10.0, NOW + 5 * DAY)]);
let urgent = ok("urg", vec![anchored(90.0, NOW + 6 * HOUR)]);
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
assert_eq!(label_of(&select(&r, &[current.clone(), urgent])), "urg");
let sliding = ok("sld", vec![unanchored()]);
assert_eq!(label_of(&select(&r, &[current, sliding])), "sld");
let cur_sliding = ok("cur", vec![unanchored()]);
let a_sliding = ok("a", vec![unanchored()]);
let far = ok("far", vec![anchored(0.0, NOW + 2 * DAY)]);
assert_eq!(label_of(&select(&r, &[a_sliding, far, cur_sliding])), "cur");
}
#[test]
fn sticky_tier1_yields_only_to_strictly_sooner_reset() {
let cur = cred("cur");
let current = ok("cur", vec![anchored(50.0, NOW + 10 * HOUR)]);
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
let sooner = ok("soon", vec![anchored(80.0, NOW + 6 * HOUR)]);
assert_eq!(label_of(&select(&r, &[current.clone(), sooner])), "soon");
let equal = ok("eq", vec![anchored(1.0, NOW + 10 * HOUR)]);
let sel = select(&r, &[equal.clone(), current.clone()]);
assert_eq!(label_of(&sel), "cur");
assert_eq!(ranks(&sel), vec![("eq".to_string(), 2)]);
let later = ok("late", vec![anchored(0.0, NOW + 12 * HOUR)]);
assert_eq!(label_of(&select(&r, &[later, current])), "cur");
}
#[test]
fn sticky_never_keeps_ineligible_current_and_only_applies_to_soonest_reset() {
let cur = cred("cur");
let other = ok("o", vec![anchored(50.0, NOW + 5 * DAY)]);
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
let exhausted = ok("cur", vec![anchored(100.0, NOW + DAY)]);
assert_eq!(label_of(&select(&r, &[exhausted, other.clone()])), "o");
let mut cooling = ok("cur", vec![anchored(0.0, NOW + DAY)]);
cooling.cooldown_until_ms = Some(NOW + HOUR);
assert_eq!(label_of(&select(&r, &[cooling, other.clone()])), "o");
let mut stale = ok("cur", vec![anchored(0.0, NOW + DAY)]);
stale.observed_at_ms = Some(NOW - 2 * HOUR);
assert_eq!(label_of(&select(&r, &[stale, other.clone()])), "o");
let excl = vec![cur.clone()];
let r_excl = SelectionRequest {
exclude: &excl,
..r.clone()
};
let healthy = ok("cur", vec![anchored(0.0, NOW + DAY)]);
assert_eq!(
label_of(&select(&r_excl, &[healthy.clone(), other.clone()])),
"o"
);
assert_eq!(label_of(&select(&r, std::slice::from_ref(&other))), "o");
let busy_current = ok("cur", vec![anchored(90.0, NOW + 5 * DAY)]);
let r_lu = SelectionRequest {
strategy: Strategy::LowestUtilization,
sticky: true,
current: Some(&cur),
..req()
};
assert_eq!(
label_of(&select(&r_lu, &[busy_current.clone(), other.clone()])),
"o"
);
let pref = vec!["o".to_string()];
let r_po = SelectionRequest {
strategy: Strategy::PreferenceOrder,
sticky: true,
current: Some(&cur),
preference: &pref,
..req()
};
assert_eq!(label_of(&select(&r_po, &[busy_current, other])), "o");
}
#[test]
fn failover_under_soonest_reset_never_repicks_sticky_failed_seat() {
let cur = cred("a");
let a = ok("a", vec![anchored(0.0, NOW + 2 * DAY)]);
let b = ok("b", vec![anchored(50.0, NOW + 5 * DAY)]);
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
assert_eq!(label_of(&select(&r, &[a.clone(), b.clone()])), "a");
let ev = FailureEvidence {
kind: FailureKind::QuotaExhaustedPreOutput,
output_started: false,
tool_activity: false,
failovers_so_far: 0,
};
match failover(&r, &cur, &ev, &[a.clone(), b]) {
FailoverDecision::Failover { credential } => {
assert_eq!(credential.storage_key(), "openai-codex@b")
}
other => panic!("unexpected {other:?}"),
}
assert!(matches!(
failover(&r, &cur, &ev, &[a]),
FailoverDecision::NoFailover { .. }
));
}
#[test]
fn explicit_account_under_soonest_reset_never_falls_back() {
let a = ok("a", vec![anchored(100.0, NOW + HOUR)]);
let b = ok("b", vec![anchored(0.0, NOW + 6 * HOUR)]);
let cur = cred("b");
let r = SelectionRequest {
explicit_account: Some("a"),
current: Some(&cur),
..sr_req()
};
match select(&r, &[a.clone(), b.clone()]) {
Selection::NoCapacity { rejections, .. } => {
assert_eq!(rejections.len(), 1);
assert_eq!(rejections[0].credential.storage_key(), "openai-codex@a");
}
other => panic!("unexpected {other:?}"),
}
let r = SelectionRequest {
explicit_account: Some("b"),
..sr_req()
};
match select(&r, &[a, b]) {
Selection::Selected {
credential,
rejections,
tier,
..
} => {
assert_eq!(credential.storage_key(), "openai-codex@b");
assert!(rejections.is_empty());
assert_eq!(tier, 1);
}
other => panic!("unexpected {other:?}"),
}
}
fn permutations<T: Clone>(items: &[T]) -> Vec<Vec<T>> {
fn go<T: Clone>(k: usize, a: &mut Vec<T>, out: &mut Vec<Vec<T>>) {
if k <= 1 {
out.push(a.clone());
return;
}
go(k - 1, a, out);
for i in 0..k - 1 {
if k.is_multiple_of(2) {
a.swap(i, k - 1);
} else {
a.swap(0, k - 1);
}
go(k - 1, a, out);
}
}
let mut a = items.to_vec();
let mut out = Vec::new();
go(a.len(), &mut a, &mut out);
out
}
#[test]
fn soonest_reset_is_deterministic_under_candidate_shuffle() {
let cur = cred("t3b");
let seats = vec![
ok("t1", vec![anchored(70.0, NOW + 3 * HOUR)]),
ok("t2", vec![unanchored()]),
ok("t3a", vec![anchored(5.0, NOW + 2 * DAY)]),
ok("t3b", vec![anchored(40.0, NOW + 4 * DAY)]),
ok("x", vec![anchored(100.0, NOW + DAY)]),
];
let r = SelectionRequest {
current: Some(&cur),
..sr_req()
};
let perms = permutations(&seats);
assert_eq!(perms.len(), 120);
let baseline = select(&r, &seats);
assert_eq!(label_of(&baseline), "t1");
let normalize = |sel: &Selection| {
let Selection::Selected {
credential,
tier,
budget_reset_ms,
utilization,
rejections,
} = sel
else {
panic!("expected selection");
};
let mut rej: Vec<(String, RejectReason)> = rejections
.iter()
.map(|r| (r.credential.storage_key(), r.reason.clone()))
.collect();
rej.sort_by(|a, b| a.0.cmp(&b.0));
(
credential.storage_key(),
*tier,
*budget_reset_ms,
*utilization,
rej,
)
};
let expected = normalize(&baseline);
assert_eq!(
ranks(&baseline),
vec![
("t2".to_string(), 2),
("t3a".to_string(), 3),
("t3b".to_string(), 4)
]
);
for p in perms {
assert_eq!(normalize(&select(&r, &p)), expected);
}
}
#[test]
fn outranked_ranks_are_two_to_n_in_ranking_order() {
let seats = [
ok("c", vec![anchored(0.0, NOW + 4 * DAY)]),
ok("a", vec![anchored(0.0, NOW + 2 * DAY)]),
ok("d", vec![anchored(0.0, NOW + 5 * DAY)]),
ok("b", vec![anchored(0.0, NOW + 3 * DAY)]),
];
let sel = select(&sr_req(), &seats);
assert_eq!(label_of(&sel), "a");
assert_eq!(
ranks(&sel),
vec![
("b".to_string(), 2),
("c".to_string(), 3),
("d".to_string(), 4)
]
);
let lu = [
ok("hi", vec![anchored(80.0, NOW + DAY)]),
ok("lo", vec![anchored(10.0, NOW + DAY)]),
ok("mid", vec![anchored(50.0, NOW + DAY)]),
];
let sel = select(&req(), &lu);
assert_eq!(label_of(&sel), "lo");
assert_eq!(
ranks(&sel),
vec![("mid".to_string(), 2), ("hi".to_string(), 3)]
);
assert_eq!(
RejectReason::Outranked { rank: 3 }.to_string(),
"eligible; outranked (rank 3)"
);
}
#[test]
fn soonest_reset_serde_is_backward_compatible() {
assert_eq!(
serde_json::to_string(&Strategy::SoonestReset).unwrap(),
"\"soonest_reset\""
);
assert_eq!(
serde_json::from_str::<Strategy>("\"soonest_reset\"").unwrap(),
Strategy::SoonestReset
);
let a = ok("a", vec![anchored(10.0, NOW + 2 * HOUR)]);
let b = ok("b", vec![anchored(10.0, NOW + 3 * HOUR)]);
let sel = select(&sr_req(), &[a, b]);
let json = serde_json::to_string(&sel).unwrap();
assert!(json.contains("\"tier\":1"));
assert!(json.contains("\"budget_reset_ms\""));
assert!(json.contains("\"rank\":2"));
assert_eq!(serde_json::from_str::<Selection>(&json).unwrap(), sel);
let old = r#"{"outcome":"selected","credential":"openai-codex@a","utilization":10.0,"rejections":[{"credential":"openai-codex@b","reason":{"reason":"outranked"}}]}"#;
match serde_json::from_str::<Selection>(old).unwrap() {
Selection::Selected {
tier,
budget_reset_ms,
rejections,
..
} => {
assert_eq!(tier, 0);
assert_eq!(budget_reset_ms, None);
assert_eq!(rejections[0].reason, RejectReason::Outranked { rank: 0 });
}
other => panic!("unexpected {other:?}"),
}
}
fn failed(label: &str, failure: ObservationFailure) -> AccountCapacity {
AccountCapacity::failed(cred(label), failure)
}
fn failure_of(candidates: &[AccountCapacity]) -> SelectionFailure {
match select(&req(), candidates) {
Selection::NoCapacity { rejections, .. } => {
SelectionFailure::from_rejections(&rejections, NOW)
}
other => panic!("expected no capacity, got {other:?}"),
}
}
#[test]
fn observation_failure_maps_to_a_never_capacity_observation() {
use AccountFailureKind as K;
for (kind, obs) in [
(K::UsageRateLimited, QuotaObservation::Unknown),
(K::RefreshTransient, QuotaObservation::Unknown),
(K::UsageUnauthorized, QuotaObservation::AuthError),
(K::RefreshRejected, QuotaObservation::AuthError),
(K::CredentialMissing, QuotaObservation::AuthError),
(K::UsageMalformed, QuotaObservation::Malformed),
(K::Unsupported, QuotaObservation::Unsupported),
(K::RefreshActionRequired, QuotaObservation::Unknown),
] {
assert_eq!(ObservationFailure::new(kind).observation(), obs, "{kind}");
}
let f = ObservationFailure::new(K::UsageRateLimited)
.with_status(429)
.with_retry_after(Some(120));
assert_eq!(
f.to_string(),
"usage_rate_limited, HTTP 429, retry after 120s"
);
assert_eq!(
serde_json::to_value(&f).unwrap(),
serde_json::json!({"kind": "usage_rate_limited", "http_status": 429, "retry_after_s": 120})
);
let old: Rejection = serde_json::from_str(
r#"{"credential":"openai-codex@a","reason":{"reason":"unknown"}}"#,
)
.unwrap();
assert_eq!(old.cause, None);
assert_eq!(old.describe(), "quota unknown");
}
#[test]
fn rejection_keeps_the_observation_cause_beside_a_cooldown() {
use AccountFailureKind as K;
let limited = ObservationFailure::new(K::UsageRateLimited).with_status(429);
let mut cooling = failed("b", limited.clone());
cooling.cooldown_until_ms = Some(NOW + HOUR);
match select(&req(), &[failed("a", limited.clone()), cooling]) {
Selection::NoCapacity { rejections, .. } => {
assert_eq!(rejections[0].reason, RejectReason::Unknown);
assert_eq!(rejections[0].cause.as_ref(), Some(&limited));
assert_eq!(
rejections[0].describe(),
"quota unknown (usage_rate_limited, HTTP 429)"
);
assert!(matches!(
rejections[1].reason,
RejectReason::CoolingDown { .. }
));
assert_eq!(rejections[1].cause.as_ref(), Some(&limited));
}
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn a_cooling_accounts_failure_keeps_its_read_kind_and_counts_as_limited() {
use AccountFailureKind as K;
use SelectionFailureKind as S;
let mut refused = failed(
"a",
ObservationFailure::new(K::UsageUnauthorized).with_status(401),
);
refused.cooldown_until_ms = Some(NOW + 2 * HOUR);
let mut throttled = failed(
"b",
ObservationFailure::new(K::UsageRateLimited)
.with_status(429)
.with_retry_after(Some(90)),
);
throttled.cooldown_until_ms = Some(NOW + HOUR);
let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
let f = failure_of(&[refused, throttled, plain]);
assert_eq!(f.kind, S::AllCoolingDown, "{f:?}");
let a = &f.accounts[0];
assert_eq!(
(a.kind, a.http_status, a.cooling_down, a.until_ms),
(K::UsageUnauthorized, Some(401), true, Some(NOW + 2 * HOUR))
);
let b = &f.accounts[1];
assert_eq!(
(b.kind, b.http_status, b.retry_after_s, b.cooling_down),
(K::UsageRateLimited, Some(429), Some(90), true)
);
let c = &f.accounts[2];
assert_eq!((c.kind, c.cooling_down), (K::CoolingDown, true));
assert_eq!(f.login_labels(), vec!["a"]);
assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
assert_eq!(
f.accounts_field(),
"a:usage_unauthorized:401:cooling,b:usage_rate_limited:429:cooling,c:cooling_down"
);
let v = serde_json::to_value(&f).unwrap();
assert_eq!(v["accounts"][0]["cooling_down"], true);
let g = failure_of(&[failed("d", ObservationFailure::new(K::UsageTimeout))]);
assert!(serde_json::to_value(&g).unwrap()["accounts"][0]
.get("cooling_down")
.is_none());
}
#[test]
fn the_headline_skips_a_cooling_accounts_read_failure_which_stays_per_account() {
use AccountFailureKind as K;
use SelectionFailureKind as S;
let throttled = || {
let mut c = failed(
"claude5",
ObservationFailure::new(K::UsageRateLimited)
.with_status(429)
.with_retry_after(Some(90)),
);
c.cooldown_until_ms = Some(NOW + HOUR);
c
};
let c5_kept = |f: &SelectionFailure| {
let c5 = &f.accounts[0];
assert_eq!(
(
c5.account.as_str(),
c5.kind,
c5.http_status,
c5.retry_after_s,
c5.cooling_down,
c5.until_ms
),
(
"claude5",
K::UsageRateLimited,
Some(429),
Some(90),
true,
Some(NOW + HOUR)
),
"{f:?}"
);
};
let mut plain = ok("claude6", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
plain.cooldown_until_ms = Some(NOW + 2 * HOUR);
let a = failure_of(&[throttled(), plain]);
assert_eq!(a.kind, S::AllCoolingDown, "{a:?}");
c5_kept(&a);
assert_eq!(
a.accounts_field(),
"claude5:usage_rate_limited:429:cooling,claude6:cooling_down"
);
assert_eq!(a.earliest_reset_ms, Some(NOW + HOUR));
assert!(a.login_labels().is_empty(), "{a:?}");
let refused = failed(
"claude2",
ObservationFailure::new(K::UsageUnauthorized).with_status(401),
);
let b = failure_of(&[throttled(), refused]);
assert_eq!(b.kind, S::CredentialRejected, "{b:?}");
c5_kept(&b);
assert_eq!(
b.accounts_field(),
"claude5:usage_rate_limited:429:cooling,claude2:usage_unauthorized:401"
);
assert_eq!(b.earliest_reset_ms, Some(NOW + HOUR));
assert_eq!(b.login_labels(), vec!["claude2"]);
}
#[test]
fn an_accounts_until_is_its_latest_exhausted_reset_and_cooldown() {
let later = ok(
"a",
vec![
win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
win("seven_day", 7 * DAY, 100.0, NOW + 3 * DAY),
],
);
let earlier = ok(
"b",
vec![
win("seven_day", 7 * DAY, 100.0, NOW + 2 * DAY),
win("five_hour", 5 * HOUR, 100.0, NOW + 2 * HOUR),
],
);
let mut cooled = ok("c", vec![win("five_hour", 5 * HOUR, 100.0, NOW + HOUR)]);
cooled.cooldown_until_ms = Some(NOW + 4 * HOUR);
let mut unknown = ok("d", vec![win("seven_day", 7 * DAY, 100.0, NOW + DAY)]);
if let QuotaObservation::Ok { windows, .. } = &mut unknown.observation {
windows.push(WindowLimit {
resets_at_ms: None,
..win("five_hour", 5 * HOUR, 100.0, 0)
});
}
let until = |f: &SelectionFailure, label: &str| {
f.accounts
.iter()
.find(|a| a.account == label)
.unwrap()
.until_ms
};
let f = failure_of(&[later.clone(), earlier.clone(), cooled.clone(), unknown]);
assert_eq!(until(&f, "a"), Some(NOW + 3 * DAY));
assert_eq!(until(&f, "b"), Some(NOW + 2 * DAY));
assert_eq!(until(&f, "c"), Some(NOW + 4 * HOUR));
assert_eq!(until(&f, "d"), None);
assert_eq!(f.earliest_reset_ms, Some(NOW + 4 * HOUR));
assert_eq!(
failure_of(&[later, earlier]).earliest_reset_ms,
Some(NOW + 2 * DAY)
);
match select(&req(), &[cooled]) {
Selection::NoCapacity { rejections, .. } => {
assert!(matches!(
rejections[0].reason,
RejectReason::CoolingDown { .. }
));
}
other => panic!("unexpected {other:?}"),
}
}
#[test]
fn the_earliest_reset_skips_accounts_that_need_a_login_or_an_action() {
use AccountFailureKind as K;
let mut refused = failed(
"a",
ObservationFailure::new(K::UsageUnauthorized).with_status(401),
);
refused.cooldown_until_ms = Some(NOW + HOUR);
let mut broken = failed("b", ObservationFailure::new(K::CredentialStore));
broken.cooldown_until_ms = Some(NOW + 2 * HOUR);
let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
let f = failure_of(&[refused.clone(), broken.clone(), plain]);
assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
assert_eq!(f.earliest_reset_ms, Some(NOW + 3 * HOUR));
assert_eq!(f.accounts[0].until_ms, Some(NOW + HOUR));
assert_eq!(f.accounts[1].until_ms, Some(NOW + 2 * HOUR));
assert_eq!(failure_of(&[refused, broken]).earliest_reset_ms, None);
}
#[test]
fn action_required_and_store_failures_get_their_own_kind() {
use AccountFailureKind as K;
use SelectionFailureKind as S;
let action = ObservationFailure::new(K::RefreshActionRequired);
let store = ObservationFailure::new(K::CredentialStore);
let refused = ObservationFailure::new(K::UsageUnauthorized).with_status(401);
let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
cooling.cooldown_until_ms = Some(NOW + HOUR);
for f in [
failure_of(&[failed("a", action.clone())]),
failure_of(&[failed("a", store.clone())]),
failure_of(&[
failed("a", action.clone()),
failed("b", store.clone()),
cooling.clone(),
]),
failure_of(&[failed("a", action.clone()), failed("b", refused.clone())]),
] {
assert_eq!(f.kind, S::ActionRequired, "{f:?}");
}
assert!(K::RefreshActionRequired.needs_action());
assert!(K::CredentialStore.needs_action());
assert!(
!K::RefreshActionRequired.needs_login() && !K::RefreshActionRequired.is_transient()
);
let f = failure_of(&[failed("a", action.clone()), failed("b", store.clone())]);
assert_eq!(f.action_labels(), vec!["a", "b"]);
assert!(f.login_labels().is_empty());
let f = failure_of(&[
failed("a", action),
failed("b", ObservationFailure::new(K::UsageTimeout)),
]);
assert_eq!(f.kind, S::UsageReadFailed);
assert_eq!(f.action_labels(), vec!["a"]);
assert_eq!(S::ActionRequired.as_str(), "action_required");
assert_eq!(S::CredentialStore.as_str(), "credential_store");
}
#[test]
fn selection_failure_kind_priority_and_login_labels() {
use AccountFailureKind as K;
use SelectionFailureKind as S;
let transient = ObservationFailure::new(K::UsageUpstream).with_status(503);
let rejected = ObservationFailure::new(K::RefreshRejected);
let exhausted = ok("x", vec![win("primary", 5 * HOUR, 100.0, NOW + 2 * HOUR)]);
let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
cooling.cooldown_until_ms = Some(NOW + HOUR);
assert_eq!(failure_of(&[]).kind, S::NoCredentials);
let f = failure_of(&[
failed("a", rejected.clone()),
failed("b", transient.clone()),
]);
assert_eq!(f.kind, S::UsageReadFailed);
assert_eq!(f.login_labels(), vec!["a"]);
assert_eq!(
f.accounts_field(),
"a:refresh_rejected,b:usage_upstream:503"
);
let f = failure_of(&[failed("a", rejected.clone()), cooling.clone()]);
assert_eq!(f.kind, S::CredentialRejected);
assert_eq!(f.login_labels(), vec!["a"]);
assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
let f = failure_of(&[exhausted.clone(), cooling.clone()]);
assert_eq!(f.kind, S::AllCoolingDown);
assert_eq!(f.accounts_field(), "x:exhausted,c:cooling_down");
assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
assert!(f.login_labels().is_empty());
let no_evidence = ok("n", vec![]);
let f = failure_of(&[failed("a", rejected), no_evidence]);
assert_eq!(f.kind, S::QuotaUnknown);
assert_eq!(f.login_labels(), vec!["a"]);
for kind in [
K::RefreshLockBusy,
K::RefreshOther,
K::UsageTimeout,
K::UsageTransport,
] {
let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
assert_eq!(f.kind, S::UsageReadFailed, "{kind}");
}
for kind in [K::RefreshActionRequired, K::CredentialStore] {
let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
assert_eq!(f.kind, S::ActionRequired, "{kind}");
assert!(f.login_labels().is_empty(), "{kind}");
}
let f = failure_of(&[failed("a", ObservationFailure::new(K::UsageMalformed))]);
assert_eq!(f.kind, S::QuotaUnknown);
assert!(f.login_labels().is_empty());
let stale_exhausted = ok("y", vec![win("primary", 5 * HOUR, 100.0, NOW - 1)]);
assert_eq!(failure_of(&[stale_exhausted]).earliest_reset_ms, None);
}
#[test]
fn an_excluded_duplicate_seat_does_not_pick_the_headline() {
use AccountFailureKind as K;
let reset = NOW + 2 * HOUR;
let spent = ok("claude1", vec![win("primary", 5 * HOUR, 100.0, reset)]);
let dup = ok("claude1-dup", vec![win("primary", 5 * HOUR, 100.0, reset)]);
let exclude = [cred("claude1-dup")];
let request = SelectionRequest {
exclude: &exclude,
..req()
};
let f = match select(&request, &[spent, dup]) {
Selection::NoCapacity { rejections, .. } => {
SelectionFailure::from_rejections(&rejections, NOW)
}
other => panic!("expected no capacity, got {other:?}"),
};
assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
assert_eq!(f.earliest_reset_ms, Some(reset));
assert_eq!(f.accounts[1].kind, K::Excluded);
let mut cooling = ok("a", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
cooling.cooldown_until_ms = Some(NOW + HOUR);
let dup = failed("b", ObservationFailure::new(K::UsageRateLimited));
let exclude = [cred("b")];
let request = SelectionRequest {
exclude: &exclude,
..req()
};
let f = match select(&request, &[cooling, dup]) {
Selection::NoCapacity { rejections, .. } => {
SelectionFailure::from_rejections(&rejections, NOW)
}
other => panic!("expected no capacity, got {other:?}"),
};
assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
let only = ok("x", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
let exclude = [cred("x")];
let request = SelectionRequest {
exclude: &exclude,
..req()
};
let f = match select(&request, &[only]) {
Selection::NoCapacity { rejections, .. } => {
SelectionFailure::from_rejections(&rejections, NOW)
}
other => panic!("expected no capacity, got {other:?}"),
};
assert_eq!(f.kind, SelectionFailureKind::QuotaUnknown, "{f:?}");
}
}