1use serde::{Deserialize, Serialize};
37
38use super::account::CredentialRef;
39use super::provider::OAuthProviderId;
40
41pub mod cred_serde {
44 use super::CredentialRef;
45 use serde::{Deserialize, Deserializer, Serialize, Serializer};
46
47 pub fn serialize<S: Serializer>(c: &CredentialRef, s: S) -> Result<S::Ok, S::Error> {
48 c.storage_key().serialize(s)
49 }
50
51 pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<CredentialRef, D::Error> {
52 let raw = String::deserialize(d)?;
53 CredentialRef::parse_storage_key(&raw)
54 .ok_or_else(|| serde::de::Error::custom(format!("invalid credential key '{raw}'")))
55 }
56}
57
58pub const WEEKLY_WINDOW_MIN_MS: u64 = 6 * 24 * 60 * 60 * 1000;
62pub const WEEKLY_WINDOW_MAX_MS: u64 = 8 * 24 * 60 * 60 * 1000;
64pub const MAX_FUTURE_OBSERVATION_SKEW_MS: u64 = 5 * 60 * 1000;
67pub const DEFAULT_EXHAUSTED_AT_PERCENT: f64 = 100.0;
69pub const DEFAULT_URGENT_HORIZON_MS: u64 = 24 * 60 * 60 * 1000;
73pub const UNANCHORED_TOLERANCE_MS: u64 = 5 * 60 * 1000;
78pub const UNANCHORED_MAX_USED_PERCENT: f64 = 0.5;
82
83#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
90pub struct WindowLimit {
91 pub id: String,
93 #[serde(default, skip_serializing_if = "Option::is_none")]
95 pub duration_ms: Option<u64>,
96 #[serde(default, skip_serializing_if = "Option::is_none")]
98 pub used_percent: Option<f64>,
99 #[serde(default, skip_serializing_if = "Option::is_none")]
104 pub limit_reached: Option<bool>,
105 #[serde(default, skip_serializing_if = "Option::is_none")]
107 pub resets_at_ms: Option<u64>,
108 #[serde(default, skip_serializing_if = "Option::is_none")]
110 pub models: Option<Vec<String>>,
111}
112
113#[derive(Debug, Clone, PartialEq)]
115pub enum WindowVerdict {
116 Headroom { used_percent: Option<f64> },
118 Exhausted { resets_at_ms: Option<u64> },
120 Unknown,
122}
123
124pub fn model_matches(entry: &str, requested: &str) -> bool {
129 let entry = entry.trim().to_ascii_lowercase();
130 let req = requested.trim().to_ascii_lowercase();
131 if entry.is_empty() || req.is_empty() {
132 return false;
133 }
134 if entry == req {
135 return true;
136 }
137 req.split(['-', '_', '/', ':', '.']).any(|tok| tok == entry)
138}
139
140impl WindowLimit {
141 pub fn applies_to(&self, model: Option<&str>) -> bool {
143 match (&self.models, model) {
144 (None, _) => true,
145 (Some(_), None) => true,
148 (Some(list), Some(m)) => list.iter().any(|x| model_matches(x, m)),
149 }
150 }
151
152 pub fn is_weekly(&self) -> bool {
154 matches!(self.duration_ms, Some(d) if (WEEKLY_WINDOW_MIN_MS..=WEEKLY_WINDOW_MAX_MS).contains(&d))
155 }
156
157 pub fn verdict(&self, exhausted_at_percent: f64) -> WindowVerdict {
162 if !valid_threshold(exhausted_at_percent) {
163 return WindowVerdict::Unknown;
164 }
165 if self.limit_reached == Some(true) {
166 return WindowVerdict::Exhausted {
167 resets_at_ms: self.resets_at_ms,
168 };
169 }
170 let used = match self.used_percent {
171 Some(p) if p > 100.0 && p <= super::usage::OVERRUN_PERCENT_MAX => Some(100.0),
173 Some(p) if !p.is_finite() || !(0.0..=100.0).contains(&p) => {
174 return WindowVerdict::Unknown
175 }
176 other => other,
177 };
178 match used {
179 Some(p) if p >= exhausted_at_percent => WindowVerdict::Exhausted {
180 resets_at_ms: self.resets_at_ms,
181 },
182 Some(p) => WindowVerdict::Headroom {
183 used_percent: Some(p),
184 },
185 None if self.limit_reached == Some(false) => {
186 WindowVerdict::Headroom { used_percent: None }
187 }
188 None => WindowVerdict::Unknown,
189 }
190 }
191}
192
193pub fn valid_threshold(percent: f64) -> bool {
195 percent.is_finite() && (0.0..=100.0).contains(&percent)
196}
197
198#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
201#[serde(rename_all = "snake_case")]
202pub enum ModelState {
203 Available,
204 Exhausted,
205 Unknown,
206}
207
208#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
210pub struct ModelAvailability {
211 pub model: String,
212 pub state: ModelState,
213}
214
215#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
217#[serde(tag = "kind", rename_all = "snake_case")]
218pub enum QuotaObservation {
219 Ok {
221 windows: Vec<WindowLimit>,
222 #[serde(default, skip_serializing_if = "Option::is_none")]
224 models: Option<Vec<ModelAvailability>>,
225 },
226 AuthError,
228 Unsupported,
230 Malformed,
232 Unknown,
234}
235
236#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
238pub struct AccountCapacity {
239 #[serde(with = "cred_serde")]
240 pub credential: CredentialRef,
241 #[serde(default, skip_serializing_if = "Option::is_none")]
243 pub observed_at_ms: Option<u64>,
244 pub observation: QuotaObservation,
245 #[serde(default, skip_serializing_if = "Option::is_none")]
247 pub cooldown_until_ms: Option<u64>,
248 #[serde(default, skip_serializing_if = "Option::is_none")]
252 pub failure: Option<ObservationFailure>,
253}
254
255impl AccountCapacity {
256 pub fn unobserved(credential: CredentialRef) -> Self {
258 Self {
259 credential,
260 observed_at_ms: None,
261 observation: QuotaObservation::Unknown,
262 cooldown_until_ms: None,
263 failure: None,
264 }
265 }
266
267 pub fn failed(credential: CredentialRef, failure: ObservationFailure) -> Self {
269 Self {
270 observation: failure.observation(),
271 failure: Some(failure),
272 ..Self::unobserved(credential)
273 }
274 }
275}
276
277#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
283#[serde(rename_all = "snake_case")]
284pub enum AccountFailureKind {
285 CoolingDown,
287 Exhausted,
289 ModelUnavailable,
291 UsageRateLimited,
293 UsageTimeout,
294 UsageUpstream,
296 UsageTransport,
297 UsageMalformed,
299 UsageUnauthorized,
301 RefreshTransient,
304 RefreshLockBusy,
306 RefreshOther,
309 RefreshRejected,
311 RefreshActionRequired,
313 CredentialMissing,
315 CredentialStore,
317 AuthError,
319 Unsupported,
321 QuotaUnknown,
323 StaleSnapshot,
324 WindowUnknown,
325 ModelUnknown,
326 NoLimitEvidence,
327 Excluded,
329 NotPresent,
330 InvalidLabel,
331 #[serde(other)]
334 Unknown,
335}
336
337impl AccountFailureKind {
338 pub fn as_str(self) -> &'static str {
339 match self {
340 Self::CoolingDown => "cooling_down",
341 Self::Exhausted => "exhausted",
342 Self::ModelUnavailable => "model_unavailable",
343 Self::UsageRateLimited => "usage_rate_limited",
344 Self::UsageTimeout => "usage_timeout",
345 Self::UsageUpstream => "usage_upstream",
346 Self::UsageTransport => "usage_transport",
347 Self::UsageMalformed => "usage_malformed",
348 Self::UsageUnauthorized => "usage_unauthorized",
349 Self::RefreshTransient => "refresh_transient",
350 Self::RefreshLockBusy => "refresh_lock_busy",
351 Self::RefreshOther => "refresh_other",
352 Self::RefreshRejected => "refresh_rejected",
353 Self::RefreshActionRequired => "refresh_action_required",
354 Self::CredentialMissing => "credential_missing",
355 Self::CredentialStore => "credential_store",
356 Self::AuthError => "auth_error",
357 Self::Unsupported => "unsupported",
358 Self::QuotaUnknown => "quota_unknown",
359 Self::StaleSnapshot => "stale_snapshot",
360 Self::WindowUnknown => "window_unknown",
361 Self::ModelUnknown => "model_unknown",
362 Self::NoLimitEvidence => "no_limit_evidence",
363 Self::Excluded => "excluded",
364 Self::NotPresent => "not_present",
365 Self::InvalidLabel => "invalid_label",
366 Self::Unknown => "unknown",
367 }
368 }
369
370 pub fn needs_login(self) -> bool {
372 matches!(
373 self,
374 Self::UsageUnauthorized
375 | Self::RefreshRejected
376 | Self::CredentialMissing
377 | Self::AuthError
378 )
379 }
380
381 pub fn is_transient(self) -> bool {
384 matches!(
385 self,
386 Self::UsageRateLimited
387 | Self::UsageTimeout
388 | Self::UsageUpstream
389 | Self::UsageTransport
390 | Self::RefreshTransient
391 | Self::RefreshLockBusy
392 | Self::RefreshOther
393 )
394 }
395
396 pub fn is_limited(self) -> bool {
398 matches!(
399 self,
400 Self::CoolingDown | Self::Exhausted | Self::ModelUnavailable
401 )
402 }
403
404 pub fn needs_action(self) -> bool {
407 matches!(self, Self::RefreshActionRequired | Self::CredentialStore)
408 }
409}
410
411impl std::fmt::Display for AccountFailureKind {
412 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
413 f.write_str(self.as_str())
414 }
415}
416
417#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
421pub struct ObservationFailure {
422 pub kind: AccountFailureKind,
423 #[serde(default, skip_serializing_if = "Option::is_none")]
424 pub http_status: Option<u16>,
425 #[serde(default, skip_serializing_if = "Option::is_none")]
426 pub retry_after_s: Option<u64>,
427}
428
429impl ObservationFailure {
430 pub fn new(kind: AccountFailureKind) -> Self {
431 Self {
432 kind,
433 http_status: None,
434 retry_after_s: None,
435 }
436 }
437
438 pub fn with_status(mut self, status: u16) -> Self {
439 self.http_status = Some(status);
440 self
441 }
442
443 pub fn with_retry_after(mut self, secs: Option<u64>) -> Self {
444 self.retry_after_s = secs;
445 self
446 }
447
448 pub fn observation(&self) -> QuotaObservation {
450 match self.kind {
451 k if k.needs_login() => QuotaObservation::AuthError,
452 AccountFailureKind::UsageMalformed => QuotaObservation::Malformed,
453 AccountFailureKind::Unsupported => QuotaObservation::Unsupported,
454 _ => QuotaObservation::Unknown,
455 }
456 }
457}
458
459impl std::fmt::Display for ObservationFailure {
460 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
461 write!(f, "{}", self.kind)?;
462 if let Some(status) = self.http_status {
463 write!(f, ", HTTP {status}")?;
464 }
465 if let Some(secs) = self.retry_after_s {
466 write!(f, ", retry after {secs}s")?;
467 }
468 Ok(())
469 }
470}
471
472#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
474#[serde(rename_all = "snake_case")]
475pub enum Strategy {
476 PreferenceOrder,
479 LowestUtilization,
483 SoonestReset,
490}
491
492#[derive(Debug, Clone)]
494pub struct SelectionRequest<'a> {
495 pub provider: OAuthProviderId,
496 pub model: Option<&'a str>,
498 pub now_ms: u64,
499 pub max_snapshot_age_ms: u64,
501 pub preference: &'a [String],
503 pub explicit_account: Option<&'a str>,
505 pub exclude: &'a [CredentialRef],
507 pub strategy: Strategy,
508 pub exhausted_at_percent: f64,
510 pub current: Option<&'a CredentialRef>,
514 pub urgent_horizon_ms: u64,
517 pub sticky: bool,
523 pub grace: &'a [AccountCapacity],
533 pub grace_max_age_ms: u64,
536}
537
538impl<'a> SelectionRequest<'a> {
539 pub fn new(provider: OAuthProviderId, now_ms: u64, max_snapshot_age_ms: u64) -> Self {
543 Self {
544 provider,
545 model: None,
546 now_ms,
547 max_snapshot_age_ms,
548 preference: &[],
549 explicit_account: None,
550 exclude: &[],
551 strategy: Strategy::LowestUtilization,
552 exhausted_at_percent: DEFAULT_EXHAUSTED_AT_PERCENT,
553 current: None,
554 urgent_horizon_ms: DEFAULT_URGENT_HORIZON_MS,
555 sticky: false,
556 grace: &[],
557 grace_max_age_ms: super::usage_reads::USAGE_READ_GRACE.as_millis() as u64,
558 }
559 }
560
561 pub fn soonest_reset(provider: OAuthProviderId, now_ms: u64, max_snapshot_age_ms: u64) -> Self {
565 Self {
566 strategy: Strategy::SoonestReset,
567 sticky: true,
568 ..Self::new(provider, now_ms, max_snapshot_age_ms)
569 }
570 }
571}
572
573#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
577#[serde(tag = "reason", rename_all = "snake_case")]
578pub enum RejectReason {
579 Stale {
581 #[serde(default, skip_serializing_if = "Option::is_none")]
582 age_ms: Option<u64>,
583 },
584 Malformed,
586 Unknown,
587 Unsupported,
588 AuthError,
589 CoolingDown {
590 until_ms: u64,
591 },
592 Exhausted {
593 window: String,
594 #[serde(default, skip_serializing_if = "Option::is_none")]
595 resets_at_ms: Option<u64>,
596 },
597 WindowUnknown {
599 window: String,
600 },
601 NoLimitEvidence,
603 ModelUnavailable {
604 model: String,
605 },
606 ModelUnknown {
608 model: String,
609 },
610 Excluded,
612 NotPresent,
614 InvalidLabel {
616 label: String,
617 },
618 Outranked {
622 #[serde(default)]
623 rank: u32,
624 },
625}
626
627impl std::fmt::Display for RejectReason {
628 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
629 match self {
630 Self::Stale { age_ms: Some(a) } => write!(f, "stale snapshot ({}s old)", a / 1000),
631 Self::Stale { age_ms: None } => write!(f, "no snapshot"),
632 Self::Malformed => write!(f, "malformed snapshot"),
633 Self::Unknown => write!(f, "quota unknown"),
634 Self::Unsupported => write!(f, "usage unsupported"),
635 Self::AuthError => write!(f, "auth error"),
636 Self::CoolingDown { until_ms } => write!(f, "cooling down until {until_ms}"),
637 Self::Exhausted {
638 window,
639 resets_at_ms: Some(r),
640 } => write!(f, "window '{window}' exhausted, resets at {r}"),
641 Self::Exhausted { window, .. } => write!(f, "window '{window}' exhausted"),
642 Self::WindowUnknown { window } => write!(f, "window '{window}' has no evidence"),
643 Self::NoLimitEvidence => write!(f, "no applicable limit evidence"),
644 Self::ModelUnavailable { model } => write!(f, "model '{model}' unavailable"),
645 Self::ModelUnknown { model } => write!(f, "model '{model}' availability unknown"),
646 Self::Excluded => write!(f, "excluded"),
647 Self::NotPresent => write!(f, "account not present"),
648 Self::InvalidLabel { label } => write!(f, "invalid account label '{label}'"),
649 Self::Outranked { rank } => write!(f, "eligible; outranked (rank {rank})"),
650 }
651 }
652}
653
654#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
656pub struct Rejection {
657 #[serde(with = "cred_serde")]
658 pub credential: CredentialRef,
659 pub reason: RejectReason,
660 #[serde(default, skip_serializing_if = "Option::is_none")]
665 pub cause: Option<ObservationFailure>,
666 #[serde(default, skip_serializing_if = "Option::is_none")]
671 pub usable_after_ms: Option<u64>,
672}
673
674impl Rejection {
675 pub fn describe(&self) -> String {
677 match &self.cause {
678 Some(cause) => format!("{} ({cause})", self.reason),
679 None => self.reason.to_string(),
680 }
681 }
682}
683
684#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
686pub struct AccountFailure {
687 pub account: String,
689 pub kind: AccountFailureKind,
690 #[serde(default, skip_serializing_if = "Option::is_none")]
691 pub http_status: Option<u16>,
692 #[serde(default, skip_serializing_if = "Option::is_none")]
693 pub retry_after_s: Option<u64>,
694 #[serde(default, skip_serializing_if = "Option::is_none")]
696 pub until_ms: Option<u64>,
697 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
700 pub cooling_down: bool,
701}
702
703impl AccountFailure {
704 pub fn from_rejection(rejection: &Rejection) -> Option<Self> {
709 use AccountFailureKind as K;
710 let mut out = Self {
711 account: rejection.credential.account.label_str().to_string(),
712 kind: K::QuotaUnknown,
713 http_status: None,
714 retry_after_s: None,
715 until_ms: rejection.usable_after_ms,
716 cooling_down: matches!(rejection.reason, RejectReason::CoolingDown { .. }),
717 };
718 if let RejectReason::Outranked { .. } = rejection.reason {
719 return None;
720 }
721 if let Some(cause) = &rejection.cause {
722 out.kind = cause.kind;
723 out.http_status = cause.http_status;
724 out.retry_after_s = cause.retry_after_s;
725 return Some(out);
726 }
727 out.kind = match &rejection.reason {
728 RejectReason::Outranked { .. } => return None,
729 RejectReason::CoolingDown { .. } => K::CoolingDown,
730 RejectReason::Exhausted { .. } => K::Exhausted,
731 RejectReason::Stale { .. } => K::StaleSnapshot,
732 RejectReason::Malformed => K::UsageMalformed,
733 RejectReason::Unknown => K::QuotaUnknown,
734 RejectReason::Unsupported => K::Unsupported,
735 RejectReason::AuthError => K::AuthError,
736 RejectReason::WindowUnknown { .. } => K::WindowUnknown,
737 RejectReason::NoLimitEvidence => K::NoLimitEvidence,
738 RejectReason::ModelUnavailable { .. } => K::ModelUnavailable,
739 RejectReason::ModelUnknown { .. } => K::ModelUnknown,
740 RejectReason::Excluded => K::Excluded,
741 RejectReason::NotPresent => K::NotPresent,
742 RejectReason::InvalidLabel { .. } => K::InvalidLabel,
743 };
744 Some(out)
745 }
746
747 pub fn is_limited(&self) -> bool {
750 self.cooling_down || self.kind.is_limited()
751 }
752
753 pub fn is_live(&self) -> bool {
756 !self.is_limited() && self.kind != AccountFailureKind::Excluded
757 }
758
759 pub fn log_form(&self) -> String {
762 let mut out = format!("{}:{}", self.account, self.kind);
763 if let Some(status) = self.http_status {
764 out.push_str(&format!(":{status}"));
765 }
766 if self.cooling_down && self.kind != AccountFailureKind::CoolingDown {
767 out.push_str(":cooling");
768 }
769 out
770 }
771}
772
773#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
775#[serde(rename_all = "snake_case")]
776pub enum SelectionFailureKind {
777 NoCredentials,
779 UsageReadFailed,
782 CredentialRejected,
784 AllCoolingDown,
786 QuotaUnknown,
788 SeatChurn,
790 RemoteNoCapacity,
793 ActionRequired,
797 CredentialStore,
799 #[serde(other)]
801 Unknown,
802}
803
804impl SelectionFailureKind {
805 pub fn as_str(self) -> &'static str {
806 match self {
807 Self::NoCredentials => "no_credentials",
808 Self::UsageReadFailed => "usage_read_failed",
809 Self::CredentialRejected => "credential_rejected",
810 Self::AllCoolingDown => "all_cooling_down",
811 Self::QuotaUnknown => "quota_unknown",
812 Self::SeatChurn => "seat_churn",
813 Self::RemoteNoCapacity => "remote_no_capacity",
814 Self::ActionRequired => "action_required",
815 Self::CredentialStore => "credential_store",
816 Self::Unknown => "unknown",
817 }
818 }
819}
820
821impl std::fmt::Display for SelectionFailureKind {
822 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
823 f.write_str(self.as_str())
824 }
825}
826
827#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
830pub struct SelectionFailure {
831 pub kind: SelectionFailureKind,
832 #[serde(default)]
833 pub accounts: Vec<AccountFailure>,
834 #[serde(default, skip_serializing_if = "Option::is_none")]
836 pub earliest_reset_ms: Option<u64>,
837 #[serde(default, skip_serializing_if = "Option::is_none")]
840 pub usage_read_retry_ms: Option<u64>,
841}
842
843impl SelectionFailure {
844 pub fn new(kind: SelectionFailureKind) -> Self {
846 Self {
847 kind,
848 accounts: Vec::new(),
849 earliest_reset_ms: None,
850 usage_read_retry_ms: None,
851 }
852 }
853
854 pub fn from_rejections(rejections: &[Rejection], now_ms: u64) -> Self {
870 use SelectionFailureKind as K;
871 let accounts: Vec<_> = rejections
872 .iter()
873 .filter_map(AccountFailure::from_rejection)
874 .collect();
875 let live: Vec<_> = accounts.iter().filter(|a| a.is_live()).collect();
879 let kind = if accounts.is_empty() {
880 K::NoCredentials
881 } else if live.iter().any(|a| a.kind.is_transient()) {
882 K::UsageReadFailed
883 } else if live.is_empty() && accounts.iter().any(AccountFailure::is_limited) {
884 K::AllCoolingDown
885 } else if live.is_empty() {
886 K::QuotaUnknown
887 } else if live.iter().all(|a| a.kind.needs_login()) {
888 K::CredentialRejected
889 } else if live
890 .iter()
891 .all(|a| a.kind.needs_login() || a.kind.needs_action())
892 {
893 K::ActionRequired
894 } else {
895 K::QuotaUnknown
896 };
897 let earliest_reset_ms = accounts
898 .iter()
899 .filter(|a| a.is_limited() && !a.kind.needs_login() && !a.kind.needs_action())
900 .filter_map(|a| a.until_ms)
901 .filter(|t| *t > now_ms)
902 .min();
903 Self {
904 kind,
905 accounts,
906 earliest_reset_ms,
907 usage_read_retry_ms: None,
908 }
909 }
910
911 pub fn transient_retry_after_s(&self) -> Option<u64> {
916 let waits: Option<Vec<u64>> = self
917 .accounts
918 .iter()
919 .filter(|a| a.is_live() && a.kind.is_transient())
920 .map(|a| a.retry_after_s)
921 .collect();
922 waits?.into_iter().min()
923 }
924
925 pub fn login_labels(&self) -> Vec<&str> {
927 self.accounts
928 .iter()
929 .filter(|a| a.kind.needs_login())
930 .map(|a| a.account.as_str())
931 .collect()
932 }
933
934 pub fn action_labels(&self) -> Vec<&str> {
936 self.accounts
937 .iter()
938 .filter(|a| a.kind.needs_action())
939 .map(|a| a.account.as_str())
940 .collect()
941 }
942
943 pub fn accounts_field(&self) -> String {
945 self.accounts
946 .iter()
947 .map(AccountFailure::log_form)
948 .collect::<Vec<_>>()
949 .join(",")
950 }
951}
952
953#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
955#[serde(tag = "outcome", rename_all = "snake_case")]
956pub enum Selection {
957 Selected {
958 #[serde(with = "cred_serde")]
959 credential: CredentialRef,
960 #[serde(default, skip_serializing_if = "Option::is_none")]
962 utilization: Option<f64>,
963 #[serde(default)]
967 tier: u8,
968 #[serde(default, skip_serializing_if = "Option::is_none")]
970 budget_reset_ms: Option<u64>,
971 rejections: Vec<Rejection>,
973 },
974 NoCapacity {
975 rejections: Vec<Rejection>,
976 #[serde(default, skip_serializing_if = "Option::is_none")]
980 earliest_reset_ms: Option<u64>,
981 },
982}
983
984impl Selection {
985 pub fn selected(&self) -> Option<&CredentialRef> {
986 match self {
987 Self::Selected { credential, .. } => Some(credential),
988 Self::NoCapacity { .. } => None,
989 }
990 }
991}
992
993#[derive(Debug, Clone, PartialEq)]
997pub struct Eligible {
998 pub utilization: Option<f64>,
1000 pub budget_reset_ms: Option<u64>,
1004 pub anchored: Option<bool>,
1009 pub tier: u8,
1017 pub graced: bool,
1020}
1021
1022fn tier_for(
1024 anchored: Option<bool>,
1025 budget_reset_ms: Option<u64>,
1026 now_ms: u64,
1027 urgent_horizon_ms: u64,
1028) -> u8 {
1029 match (anchored, budget_reset_ms) {
1030 (_, None) => 4,
1031 (Some(false), Some(_)) => 2,
1032 (_, Some(reset)) if reset.saturating_sub(now_ms) <= urgent_horizon_ms => 1,
1033 (_, Some(_)) => 3,
1034 }
1035}
1036
1037pub fn evaluate(
1040 req: &SelectionRequest<'_>,
1041 cap: &AccountCapacity,
1042) -> Result<Eligible, RejectReason> {
1043 if req
1044 .exclude
1045 .iter()
1046 .any(|c| c.storage_key() == cap.credential.storage_key())
1047 {
1048 return Err(RejectReason::Excluded);
1049 }
1050 if let Some(until) = cap.cooldown_until_ms {
1051 if until > req.now_ms {
1052 return Err(RejectReason::CoolingDown { until_ms: until });
1053 }
1054 }
1055 let (windows, models) = match &cap.observation {
1056 QuotaObservation::Ok { windows, models } => (windows, models),
1057 QuotaObservation::AuthError => return Err(RejectReason::AuthError),
1058 QuotaObservation::Unsupported => return Err(RejectReason::Unsupported),
1059 QuotaObservation::Malformed => return Err(RejectReason::Malformed),
1060 QuotaObservation::Unknown => return evaluate_graced(req, cap),
1061 };
1062 let observed = cap
1063 .observed_at_ms
1064 .ok_or(RejectReason::Stale { age_ms: None })?;
1065 if observed > req.now_ms.saturating_add(MAX_FUTURE_OBSERVATION_SKEW_MS) {
1066 return Err(RejectReason::Malformed);
1067 }
1068 let age = req.now_ms.saturating_sub(observed);
1069 if age > req.max_snapshot_age_ms {
1070 return Err(RejectReason::Stale { age_ms: Some(age) });
1071 }
1072 if let (Some(model), Some(list)) = (req.model, models) {
1073 for m in list.iter().filter(|m| model_matches(&m.model, model)) {
1074 match m.state {
1075 ModelState::Available => {}
1076 ModelState::Exhausted => {
1077 return Err(RejectReason::ModelUnavailable {
1078 model: model.to_string(),
1079 })
1080 }
1081 ModelState::Unknown => {
1082 return Err(RejectReason::ModelUnknown {
1083 model: model.to_string(),
1084 })
1085 }
1086 }
1087 }
1088 }
1089 let mut applicable = 0usize;
1090 let mut utilization: Option<f64> = None;
1091 for w in windows.iter().filter(|w| w.applies_to(req.model)) {
1092 applicable += 1;
1093 match w.verdict(req.exhausted_at_percent) {
1094 WindowVerdict::Headroom { used_percent } => {
1095 if let Some(p) = used_percent {
1096 utilization = Some(utilization.map_or(p, |u| u.max(p)));
1097 }
1098 }
1099 WindowVerdict::Exhausted { resets_at_ms } => {
1100 return Err(RejectReason::Exhausted {
1101 window: w.id.clone(),
1102 resets_at_ms,
1103 });
1104 }
1105 WindowVerdict::Unknown => {
1106 return Err(RejectReason::WindowUnknown {
1107 window: w.id.clone(),
1108 });
1109 }
1110 }
1111 }
1112 if applicable == 0 {
1113 return Err(RejectReason::NoLimitEvidence);
1114 }
1115 let budget = budget_window(windows, req.model);
1119 let budget_reset_ms = budget
1120 .and_then(|w| w.resets_at_ms)
1121 .filter(|t| *t > req.now_ms);
1122 let anchored = match (budget, budget_reset_ms) {
1123 (Some(w), Some(_)) => is_unanchored(w, observed).map(|u| !u),
1124 _ => None,
1125 };
1126 let tier = tier_for(anchored, budget_reset_ms, req.now_ms, req.urgent_horizon_ms);
1127 Ok(Eligible {
1128 utilization,
1129 budget_reset_ms,
1130 anchored,
1131 tier,
1132 graced: false,
1133 })
1134}
1135
1136fn evaluate_graced(
1142 req: &SelectionRequest<'_>,
1143 cap: &AccountCapacity,
1144) -> Result<Eligible, RejectReason> {
1145 let key = cap.credential.storage_key();
1146 let Some(good) = req
1147 .grace
1148 .iter()
1149 .find(|g| g.credential.storage_key() == key)
1150 .filter(|g| matches!(g.observation, QuotaObservation::Ok { .. }))
1151 else {
1152 return Err(RejectReason::Unknown);
1153 };
1154 let young = good
1157 .observed_at_ms
1158 .is_some_and(|t| req.now_ms.saturating_sub(t) < req.grace_max_age_ms);
1159 if !young {
1160 return Err(RejectReason::Unknown);
1161 }
1162 let graced_req = SelectionRequest {
1163 max_snapshot_age_ms: req.grace_max_age_ms,
1164 grace: &[],
1165 ..req.clone()
1166 };
1167 let reading = AccountCapacity {
1168 cooldown_until_ms: cap.cooldown_until_ms,
1169 ..good.clone()
1170 };
1171 evaluate(&graced_req, &reading)
1172 .map(|e| Eligible { graced: true, ..e })
1173 .map_err(|_| RejectReason::Unknown)
1174}
1175
1176fn utilization_cmp(a: Option<f64>, b: Option<f64>) -> std::cmp::Ordering {
1178 a.is_none().cmp(&b.is_none()).then_with(|| {
1179 a.unwrap_or(0.0)
1180 .partial_cmp(&b.unwrap_or(0.0))
1181 .unwrap_or(std::cmp::Ordering::Equal)
1182 })
1183}
1184
1185fn preference_rank(req: &SelectionRequest<'_>, cred: &CredentialRef) -> usize {
1186 let label = cred.account.label_str();
1187 req.preference
1188 .iter()
1189 .position(|p| p == label)
1190 .unwrap_or(usize::MAX)
1191}
1192
1193fn observation_cause(reason: &RejectReason, cap: &AccountCapacity) -> Option<ObservationFailure> {
1198 match reason {
1199 RejectReason::Unknown
1200 | RejectReason::AuthError
1201 | RejectReason::Malformed
1202 | RejectReason::Unsupported
1203 | RejectReason::CoolingDown { .. } => cap.failure.clone(),
1204 _ => None,
1205 }
1206}
1207
1208fn usable_after_ms(req: &SelectionRequest<'_>, cap: &AccountCapacity) -> Option<u64> {
1214 let mut latest = cap.cooldown_until_ms.filter(|t| *t > req.now_ms);
1215 if let QuotaObservation::Ok { windows, .. } = &cap.observation {
1216 for w in windows.iter().filter(|w| w.applies_to(req.model)) {
1217 if let WindowVerdict::Exhausted { resets_at_ms } = w.verdict(req.exhausted_at_percent) {
1218 let reset = resets_at_ms.filter(|t| *t > req.now_ms)?;
1219 latest = Some(latest.map_or(reset, |l| l.max(reset)));
1220 }
1221 }
1222 }
1223 latest
1224}
1225
1226pub fn select(req: &SelectionRequest<'_>, candidates: &[AccountCapacity]) -> Selection {
1230 let mut rejections: Vec<Rejection> = Vec::new();
1231 let mut eligible: Vec<(&AccountCapacity, Eligible)> = Vec::new();
1232 let mut explicit_seen = false;
1233
1234 for cap in candidates
1235 .iter()
1236 .filter(|c| c.credential.provider == req.provider)
1237 {
1238 if let Some(label) = req.explicit_account {
1239 if cap.credential.account.label_str() != label {
1240 continue;
1241 }
1242 explicit_seen = true;
1243 }
1244 match evaluate(req, cap) {
1245 Ok(e) => eligible.push((cap, e)),
1246 Err(reason) => rejections.push(Rejection {
1247 credential: cap.credential.clone(),
1248 cause: observation_cause(&reason, cap),
1249 usable_after_ms: usable_after_ms(req, cap),
1250 reason,
1251 }),
1252 }
1253 }
1254
1255 if let (Some(label), false) = (req.explicit_account, explicit_seen) {
1256 let (credential, reason) = match super::account::Account::parse(label) {
1260 Ok(account) => (
1261 CredentialRef::new(req.provider, account),
1262 RejectReason::NotPresent,
1263 ),
1264 Err(_) => (
1265 CredentialRef::default_for(req.provider),
1266 RejectReason::InvalidLabel {
1267 label: label.to_string(),
1268 },
1269 ),
1270 };
1271 return Selection::NoCapacity {
1272 rejections: vec![Rejection {
1273 credential,
1274 reason,
1275 cause: None,
1276 usable_after_ms: None,
1277 }],
1278 earliest_reset_ms: None,
1279 };
1280 }
1281
1282 eligible.sort_by(|(a, ea), (b, eb)| {
1286 let ka = a.credential.storage_key();
1287 let kb = b.credential.storage_key();
1288 let pa = preference_rank(req, &a.credential);
1289 let pb = preference_rank(req, &b.credential);
1290 ea.graced.cmp(&eb.graced).then_with(|| match req.strategy {
1291 Strategy::PreferenceOrder => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
1292 Strategy::LowestUtilization => utilization_cmp(ea.utilization, eb.utilization)
1293 .then_with(|| pa.cmp(&pb))
1294 .then_with(|| ka.cmp(&kb)),
1295 Strategy::SoonestReset => ea
1296 .tier
1297 .cmp(&eb.tier)
1298 .then_with(|| match ea.tier {
1299 1 | 3 => ea.budget_reset_ms.cmp(&eb.budget_reset_ms),
1301 2 => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
1303 _ => utilization_cmp(ea.utilization, eb.utilization),
1305 })
1306 .then_with(|| utilization_cmp(ea.utilization, eb.utilization))
1307 .then_with(|| pa.cmp(&pb))
1308 .then_with(|| ka.cmp(&kb)),
1309 })
1310 });
1311
1312 if req.sticky && req.strategy == Strategy::SoonestReset {
1319 if let Some(current) = req.current {
1320 let key = current.storage_key();
1321 let at = eligible
1322 .iter()
1323 .position(|(c, _)| c.credential.storage_key() == key);
1324 if let Some(i) = at.filter(|i| *i > 0) {
1325 let best = &eligible[0].1;
1326 let cur = &eligible[i].1;
1327 let switch = best.tier < cur.tier
1328 || (best.tier == 1
1329 && cur.tier == 1
1330 && best.budget_reset_ms < cur.budget_reset_ms);
1331 if !switch {
1332 eligible[..=i].rotate_right(1);
1333 }
1334 }
1335 }
1336 }
1337
1338 match eligible.first() {
1339 Some((cap, e)) => {
1340 let winner = cap.credential.clone();
1341 for (i, (other, _)) in eligible.iter().enumerate().skip(1) {
1342 rejections.push(Rejection {
1343 credential: other.credential.clone(),
1344 reason: RejectReason::Outranked {
1345 rank: (i + 1) as u32,
1346 },
1347 cause: None,
1348 usable_after_ms: None,
1349 });
1350 }
1351 Selection::Selected {
1352 credential: winner,
1353 utilization: e.utilization,
1354 tier: e.tier,
1355 budget_reset_ms: e.budget_reset_ms,
1356 rejections,
1357 }
1358 }
1359 None => {
1360 let earliest_reset_ms = rejections
1361 .iter()
1362 .filter_map(|r| match &r.reason {
1363 RejectReason::Exhausted {
1364 resets_at_ms: Some(t),
1365 ..
1366 } if *t > req.now_ms => Some(*t),
1367 _ => None,
1368 })
1369 .min();
1370 Selection::NoCapacity {
1371 rejections,
1372 earliest_reset_ms,
1373 }
1374 }
1375 }
1376}
1377
1378#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1382#[serde(rename_all = "snake_case")]
1383pub enum FailureKind {
1384 QuotaExhaustedPreOutput,
1387 RateLimited,
1389 Other,
1391}
1392
1393#[derive(Debug, Clone, PartialEq, Eq)]
1395pub struct FailureEvidence {
1396 pub kind: FailureKind,
1397 pub output_started: bool,
1399 pub tool_activity: bool,
1401 pub failovers_so_far: u32,
1403}
1404
1405pub const MAX_FAILOVERS: u32 = 1;
1407
1408#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1409#[serde(tag = "decision", rename_all = "snake_case")]
1410pub enum FailoverDecision {
1411 Failover {
1412 #[serde(with = "cred_serde")]
1413 credential: CredentialRef,
1414 },
1415 NoFailover {
1416 reason: String,
1417 },
1418}
1419
1420pub fn failover(
1426 req: &SelectionRequest<'_>,
1427 failed: &CredentialRef,
1428 evidence: &FailureEvidence,
1429 candidates: &[AccountCapacity],
1430) -> FailoverDecision {
1431 let no = |reason: &str| FailoverDecision::NoFailover {
1432 reason: reason.to_string(),
1433 };
1434 if evidence.output_started {
1435 return no("output already started; never replay on another account");
1436 }
1437 if evidence.tool_activity {
1438 return no("tool activity occurred; never replay on another account");
1439 }
1440 if evidence.failovers_so_far >= MAX_FAILOVERS {
1441 return no("failover budget exhausted");
1442 }
1443 if evidence.kind != FailureKind::QuotaExhaustedPreOutput {
1444 return no("not a recognized pre-output quota failure");
1445 }
1446 if req.explicit_account.is_some() {
1447 return no("explicit account requested; no silent fallback");
1448 }
1449 let mut exclude: Vec<CredentialRef> = req.exclude.to_vec();
1450 exclude.push(failed.clone());
1451 let retry = SelectionRequest {
1452 exclude: &exclude,
1453 ..req.clone()
1454 };
1455 match select(&retry, candidates) {
1456 Selection::Selected { credential, .. } => FailoverDecision::Failover { credential },
1457 Selection::NoCapacity { .. } => no("no other eligible account"),
1458 }
1459}
1460
1461fn pick_weekly<'w>(windows: impl Iterator<Item = &'w WindowLimit>) -> Option<&'w WindowLimit> {
1465 let mut weekly: Vec<&WindowLimit> = windows.filter(|w| w.is_weekly()).collect();
1466 weekly.sort_by(|a, b| {
1467 a.models
1468 .is_some()
1469 .cmp(&b.models.is_some())
1470 .then_with(|| a.id.cmp(&b.id))
1471 });
1472 weekly.into_iter().next()
1473}
1474
1475pub fn weekly_window(windows: &[WindowLimit]) -> Option<&WindowLimit> {
1481 pick_weekly(windows.iter())
1482}
1483
1484pub fn budget_window<'w>(
1494 windows: &'w [WindowLimit],
1495 model: Option<&str>,
1496) -> Option<&'w WindowLimit> {
1497 let applicable = || windows.iter().filter(|w| w.applies_to(model));
1498 if let Some(weekly) = pick_weekly(applicable()) {
1499 return Some(weekly);
1500 }
1501 applicable()
1502 .filter(|w| w.duration_ms.is_some())
1503 .min_by(|a, b| {
1504 b.duration_ms
1505 .cmp(&a.duration_ms)
1506 .then_with(|| a.models.is_some().cmp(&b.models.is_some()))
1507 .then_with(|| a.id.cmp(&b.id))
1508 })
1509}
1510
1511pub fn is_unanchored(window: &WindowLimit, observed_at_ms: u64) -> Option<bool> {
1526 let (Some(duration), Some(reset)) = (window.duration_ms, window.resets_at_ms) else {
1527 return None;
1528 };
1529 let near_zero = matches!(
1530 window.used_percent,
1531 Some(p) if (0.0..=UNANCHORED_MAX_USED_PERCENT).contains(&p)
1532 );
1533 if !near_zero {
1534 return Some(false);
1535 }
1536 let expected = observed_at_ms.saturating_add(duration);
1537 Some(expected.abs_diff(reset) <= UNANCHORED_TOLERANCE_MS)
1538}
1539
1540#[cfg(test)]
1541mod tests {
1542 use super::*;
1543 use crate::auth::account::Account;
1544
1545 const HOUR: u64 = 3_600_000;
1546 const DAY: u64 = 24 * HOUR;
1547 const NOW: u64 = 1_800_000_000_000;
1548
1549 fn cred(label: &str) -> CredentialRef {
1550 CredentialRef::new(OAuthProviderId::OpenAiCodex, Account::parse(label).unwrap())
1551 }
1552
1553 fn win(id: &str, dur: u64, used: f64, resets: u64) -> WindowLimit {
1554 WindowLimit {
1555 id: id.into(),
1556 duration_ms: Some(dur),
1557 used_percent: Some(used),
1558 limit_reached: Some(used >= 100.0),
1559 resets_at_ms: Some(resets),
1560 models: None,
1561 }
1562 }
1563
1564 fn ok(label: &str, windows: Vec<WindowLimit>) -> AccountCapacity {
1565 AccountCapacity {
1566 credential: cred(label),
1567 observed_at_ms: Some(NOW - 1000),
1568 observation: QuotaObservation::Ok {
1569 windows,
1570 models: None,
1571 },
1572 cooldown_until_ms: None,
1573 failure: None,
1574 }
1575 }
1576
1577 fn req() -> SelectionRequest<'static> {
1578 SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
1579 }
1580
1581 #[test]
1582 fn window_verdicts_fail_closed() {
1583 let mut w = win("primary", 7 * DAY, 20.0, NOW + DAY);
1584 assert_eq!(
1585 w.verdict(100.0),
1586 WindowVerdict::Headroom {
1587 used_percent: Some(20.0)
1588 }
1589 );
1590 w.used_percent = Some(100.0);
1591 w.limit_reached = Some(false);
1592 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1593 w.used_percent = Some(50.0);
1594 w.limit_reached = Some(true);
1595 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1596 w.used_percent = None;
1597 w.limit_reached = None;
1598 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1599 w.limit_reached = Some(false);
1600 assert_eq!(
1601 w.verdict(100.0),
1602 WindowVerdict::Headroom { used_percent: None }
1603 );
1604 w.used_percent = Some(f64::NAN);
1605 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1606 w.used_percent = Some(140.0);
1608 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1609 w.used_percent = Some(1000.5);
1610 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1611 let saved = w.limit_reached;
1613 w.limit_reached = Some(true);
1614 for p in [1500.0, f64::NAN, -3.0] {
1615 w.used_percent = Some(p);
1616 assert!(
1617 matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }),
1618 "{p}"
1619 );
1620 }
1621 w.limit_reached = saved;
1622 w.used_percent = Some(-1.0);
1623 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1624 w.used_percent = Some(1.0);
1626 assert_eq!(w.verdict(f64::NAN), WindowVerdict::Unknown);
1627 assert_eq!(w.verdict(101.0), WindowVerdict::Unknown);
1628 assert_eq!(w.verdict(-0.5), WindowVerdict::Unknown);
1629 assert_eq!(w.verdict(f64::INFINITY), WindowVerdict::Unknown);
1630 let r = SelectionRequest {
1631 exhausted_at_percent: f64::NAN,
1632 ..SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, HOUR)
1633 };
1634 assert!(select(
1635 &r,
1636 &[ok("a", vec![win("primary", 7 * DAY, 1.0, NOW + DAY)])]
1637 )
1638 .selected()
1639 .is_none());
1640 }
1641
1642 #[test]
1643 fn weekly_is_from_duration_not_name() {
1644 let primary_5h = win("primary", 5 * HOUR, 0.0, NOW + HOUR);
1645 let primary_week = win("primary", 7 * DAY, 0.0, NOW + DAY);
1646 let no_dur = WindowLimit {
1647 duration_ms: None,
1648 ..win("seven_day", 0, 0.0, NOW)
1649 };
1650 assert!(!primary_5h.is_weekly());
1651 assert!(primary_week.is_weekly());
1652 assert!(!no_dur.is_weekly());
1653 assert_eq!(
1654 weekly_window(&[primary_5h.clone(), primary_week.clone()]).map(|w| w.duration_ms),
1655 Some(Some(7 * DAY))
1656 );
1657 assert!(weekly_window(&[primary_5h, no_dur]).is_none());
1658 let scoped = WindowLimit {
1660 id: "aaa_scoped".into(),
1661 models: Some(vec!["m".into()]),
1662 ..primary_week.clone()
1663 };
1664 assert_eq!(
1665 weekly_window(&[scoped, primary_week]).map(|w| w.id.as_str()),
1666 Some("primary")
1667 );
1668 }
1669
1670 #[test]
1671 fn selects_lowest_utilization_deterministically() {
1672 let a = ok("a", vec![win("primary", 7 * DAY, 60.0, NOW + DAY)]);
1673 let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1674 let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1675 let sel = select(&req(), &[a.clone(), c.clone(), b.clone()]);
1676 assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@b");
1677 let sel2 = select(&req(), &[c, b, a]);
1679 assert_eq!(sel2.selected().unwrap().storage_key(), "openai-codex@b");
1680 }
1681
1682 #[test]
1683 fn preference_breaks_ties_and_orders() {
1684 let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1685 let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1686 let pref = vec!["c".to_string(), "b".to_string()];
1687 let r = SelectionRequest {
1688 preference: &pref,
1689 ..req()
1690 };
1691 assert_eq!(
1692 select(&r, &[b.clone(), c.clone()])
1693 .selected()
1694 .unwrap()
1695 .storage_key(),
1696 "openai-codex@c"
1697 );
1698 let a = ok("a", vec![win("primary", 7 * DAY, 90.0, NOW + DAY)]);
1699 let pref2 = vec!["a".to_string()];
1700 let r2 = SelectionRequest {
1701 preference: &pref2,
1702 strategy: Strategy::PreferenceOrder,
1703 ..req()
1704 };
1705 assert_eq!(
1706 select(&r2, &[b, c, a]).selected().unwrap().storage_key(),
1707 "openai-codex@a"
1708 );
1709 }
1710
1711 #[test]
1715 fn graced_readings_are_bounded_rank_after_fresh_and_keep_a_sticky_seat() {
1716 let grace_ms = super::super::usage_reads::USAGE_READ_GRACE.as_millis() as u64;
1717 let failed = |label: &str| AccountCapacity::unobserved(cred(label));
1718 let good = |label: &str, age: u64, used: f64, reset: u64| AccountCapacity {
1719 observed_at_ms: Some(NOW - age),
1720 ..ok(label, vec![win("primary", 7 * DAY, used, reset)])
1721 };
1722 let fresh_req = || SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 60_000);
1723 let winner = |sel: Selection| match sel {
1724 Selection::Selected { credential, .. } => credential.account.label_str().to_string(),
1725 Selection::NoCapacity { rejections, .. } => panic!("no capacity: {rejections:?}"),
1726 };
1727 let one = |grace: &[AccountCapacity], cap: &AccountCapacity| {
1728 evaluate(&SelectionRequest { grace, ..fresh_req() }, cap)
1729 };
1730
1731 let young = [good("a", grace_ms - 60_000, 20.0, NOW + DAY)];
1732 assert!(one(&young, &failed("a")).unwrap().graced);
1733 let old = [good("a", grace_ms + 60_000, 20.0, NOW + DAY)];
1734 assert_eq!(one(&old, &failed("a")), Err(RejectReason::Unknown), "past the grace age");
1735 let edge = [good("a", grace_ms, 20.0, NOW + DAY)];
1737 assert_eq!(one(&edge, &failed("a")), Err(RejectReason::Unknown), "at the grace age");
1738 let inside = [good("a", grace_ms - 1, 20.0, NOW + DAY)];
1739 assert!(one(&inside, &failed("a")).unwrap().graced, "just under it");
1740 let exhausted = [good("a", 60_000, 100.0, NOW + HOUR)];
1741 assert_eq!(one(&exhausted, &failed("a")), Err(RejectReason::Unknown), "exhausted window");
1742 let cooling = AccountCapacity {
1743 cooldown_until_ms: Some(NOW + HOUR),
1744 ..failed("a")
1745 };
1746 assert!(matches!(one(&young, &cooling), Err(RejectReason::CoolingDown { .. })));
1747 let auth = AccountCapacity {
1748 observation: QuotaObservation::AuthError,
1749 ..failed("a")
1750 };
1751 assert_eq!(one(&young, &auth), Err(RejectReason::AuthError), "only a failed read is graced");
1752
1753 let grace = [good("a", 60_000, 20.0, NOW + HOUR)];
1755 let b = good("b", 1000, 20.0, NOW + 3 * DAY);
1756 let anew = SelectionRequest { grace: &grace, ..fresh_req() };
1757 assert_eq!(winner(select(&anew, &[failed("a"), b.clone()])), "b");
1758 assert_eq!(winner(select(&anew, &[failed("a")])), "a", "alone, the graced seat serves");
1759 let current = cred("a");
1761 let sticky = SelectionRequest {
1762 current: Some(¤t),
1763 ..anew.clone()
1764 };
1765 assert_eq!(winner(select(&sticky, &[failed("a"), b])), "a");
1766
1767 let mut switched = 0;
1772 for a_reset in [NOW + HOUR, NOW + 2 * DAY, NOW + 10 * HOUR] {
1773 for b_reset in [NOW + HOUR / 2, NOW + 3 * DAY, NOW + 12 * HOUR, NOW + 5 * DAY] {
1774 let a_good = good("a", 1000, 20.0, a_reset);
1775 let b = good("b", 1000, 20.0, b_reset);
1776 let graced = [a_good.clone()];
1777 let on_grace = SelectionRequest {
1778 grace: &graced,
1779 current: Some(¤t),
1780 ..fresh_req()
1781 };
1782 let on_fresh = SelectionRequest {
1783 current: Some(¤t),
1784 ..fresh_req()
1785 };
1786 let g = winner(select(&on_grace, &[failed("a"), b.clone()]));
1787 let f = winner(select(&on_fresh, &[a_good, b]));
1788 assert_eq!(g, f, "a resets at {a_reset}, b at {b_reset}");
1789 switched += usize::from(g == "b");
1790 }
1791 }
1792 assert!(switched > 0 && switched < 12, "the cases cover staying and switching: {switched}");
1793 }
1794
1795 #[test]
1796 fn stale_unknown_malformed_auth_are_not_capacity() {
1797 let mut stale = ok("s", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1798 stale.observed_at_ms = Some(NOW - 2 * HOUR);
1799 let mut never = ok("n", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1800 never.observed_at_ms = None;
1801 let mut future = ok("f", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1802 future.observed_at_ms = Some(NOW + HOUR);
1803 let mut auth = ok("x", vec![]);
1804 auth.observation = QuotaObservation::AuthError;
1805 let mut unk = ok("u", vec![]);
1806 unk.observation = QuotaObservation::Unknown;
1807 let mut mal = ok("m", vec![]);
1808 mal.observation = QuotaObservation::Malformed;
1809 let mut uns = ok("v", vec![]);
1810 uns.observation = QuotaObservation::Unsupported;
1811 let none = ok("e", vec![]);
1812 let sel = select(&req(), &[stale, never, future, auth, unk, mal, uns, none]);
1813 let Selection::NoCapacity {
1814 rejections,
1815 earliest_reset_ms,
1816 } = sel
1817 else {
1818 panic!("expected no capacity");
1819 };
1820 assert_eq!(earliest_reset_ms, None);
1821 let reasons: Vec<(String, RejectReason)> = rejections
1822 .into_iter()
1823 .map(|r| (r.credential.account.label_str().to_string(), r.reason))
1824 .collect();
1825 assert!(reasons.contains(&(
1826 "s".into(),
1827 RejectReason::Stale {
1828 age_ms: Some(2 * HOUR)
1829 }
1830 )));
1831 assert!(reasons.contains(&("n".into(), RejectReason::Stale { age_ms: None })));
1832 assert!(reasons.contains(&("f".into(), RejectReason::Malformed)));
1833 assert!(reasons.contains(&("x".into(), RejectReason::AuthError)));
1834 assert!(reasons.contains(&("u".into(), RejectReason::Unknown)));
1835 assert!(reasons.contains(&("m".into(), RejectReason::Malformed)));
1836 assert!(reasons.contains(&("v".into(), RejectReason::Unsupported)));
1837 assert!(reasons.contains(&("e".into(), RejectReason::NoLimitEvidence)));
1838 }
1839
1840 #[test]
1841 fn exhausted_cooldown_and_unknown_window_excluded() {
1842 let exhausted = ok("x", vec![win("primary", 7 * DAY, 100.0, NOW + 2 * DAY)]);
1843 let mut cooling = ok("c", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1844 cooling.cooldown_until_ms = Some(NOW + HOUR);
1845 let mut cooled = ok("d", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1846 cooled.cooldown_until_ms = Some(NOW - 1);
1847 let unknown_window = ok(
1848 "w",
1849 vec![WindowLimit {
1850 used_percent: None,
1851 limit_reached: None,
1852 ..win("primary", 7 * DAY, 0.0, NOW + DAY)
1853 }],
1854 );
1855 let past_exhausted = ok("p", vec![win("primary", 7 * DAY, 100.0, NOW - DAY)]);
1856 let sel = select(
1857 &req(),
1858 &[exhausted, cooling, unknown_window, past_exhausted],
1859 );
1860 let Selection::NoCapacity {
1861 earliest_reset_ms, ..
1862 } = &sel
1863 else {
1864 panic!("no capacity expected");
1865 };
1866 assert_eq!(*earliest_reset_ms, Some(NOW + 2 * DAY));
1868 let sel = select(&req(), &[cooled]);
1869 assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@d");
1870 }
1871
1872 #[test]
1873 fn all_applicable_limits_including_model_scoped_count() {
1874 let model_win = WindowLimit {
1875 id: "weekly_sonnet".into(),
1876 models: Some(vec!["sonnet".into()]),
1877 ..win("weekly_sonnet", 7 * DAY, 100.0, NOW + DAY)
1878 };
1879 let five_h = win("five_hour", 5 * HOUR, 99.0, NOW + HOUR);
1880 let acct = ok(
1881 "a",
1882 vec![
1883 win("seven_day", 7 * DAY, 10.0, NOW + DAY),
1884 five_h,
1885 model_win,
1886 ],
1887 );
1888 let r = SelectionRequest {
1890 model: Some("sonnet"),
1891 ..req()
1892 };
1893 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
1894 let r = SelectionRequest {
1896 model: Some("opus"),
1897 ..req()
1898 };
1899 match select(&r, std::slice::from_ref(&acct)) {
1900 Selection::Selected { utilization, .. } => assert_eq!(utilization, Some(99.0)),
1901 other => panic!("unexpected {other:?}"),
1902 }
1903 assert!(select(&req(), std::slice::from_ref(&acct))
1905 .selected()
1906 .is_none());
1907 let mut unavailable = acct;
1909 if let QuotaObservation::Ok { models, .. } = &mut unavailable.observation {
1910 *models = Some(vec![ModelAvailability {
1911 model: "opus".into(),
1912 state: ModelState::Exhausted,
1913 }]);
1914 }
1915 let r = SelectionRequest {
1916 model: Some("opus"),
1917 ..req()
1918 };
1919 let sel = select(&r, &[unavailable.clone()]);
1920 match sel {
1921 Selection::NoCapacity { rejections, .. } => assert_eq!(
1922 rejections[0].reason,
1923 RejectReason::ModelUnavailable {
1924 model: "opus".into()
1925 }
1926 ),
1927 other => panic!("unexpected {other:?}"),
1928 }
1929 let mut unknown = unavailable;
1931 if let QuotaObservation::Ok { models, .. } = &mut unknown.observation {
1932 *models = Some(vec![ModelAvailability {
1933 model: "opus".into(),
1934 state: ModelState::Unknown,
1935 }]);
1936 }
1937 match select(&r, &[unknown]) {
1938 Selection::NoCapacity { rejections, .. } => assert_eq!(
1939 rejections[0].reason,
1940 RejectReason::ModelUnknown {
1941 model: "opus".into()
1942 }
1943 ),
1944 other => panic!("unexpected {other:?}"),
1945 }
1946 }
1947
1948 #[test]
1949 fn explicit_account_never_falls_back() {
1950 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
1951 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1952 let r = SelectionRequest {
1953 explicit_account: Some("a"),
1954 ..req()
1955 };
1956 let sel = select(&r, &[a.clone(), b.clone()]);
1957 assert!(sel.selected().is_none());
1958 let r = SelectionRequest {
1959 explicit_account: Some("zzz"),
1960 ..req()
1961 };
1962 match select(&r, &[a.clone(), b.clone()]) {
1963 Selection::NoCapacity { rejections, .. } => {
1964 assert_eq!(rejections.len(), 1);
1965 assert_eq!(rejections[0].reason, RejectReason::NotPresent);
1966 assert_eq!(rejections[0].credential.storage_key(), "openai-codex@zzz");
1967 }
1968 other => panic!("unexpected {other:?}"),
1969 }
1970 let r = SelectionRequest {
1971 explicit_account: Some("Bad Label"),
1972 ..req()
1973 };
1974 match select(&r, &[a, b]) {
1975 Selection::NoCapacity { rejections, .. } => assert_eq!(
1976 rejections[0].reason,
1977 RejectReason::InvalidLabel {
1978 label: "Bad Label".into()
1979 }
1980 ),
1981 other => panic!("unexpected {other:?}"),
1982 }
1983 }
1984
1985 #[test]
1986 fn other_providers_and_exclusions_are_ignored() {
1987 let mut other = ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1988 other.credential =
1989 CredentialRef::new(OAuthProviderId::Anthropic, Account::parse("a").unwrap());
1990 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1991 let excl = vec![cred("b")];
1992 let r = SelectionRequest {
1993 exclude: &excl,
1994 ..req()
1995 };
1996 let sel = select(&r, &[other, b]);
1997 match sel {
1998 Selection::NoCapacity { rejections, .. } => {
1999 assert_eq!(rejections.len(), 1);
2000 assert_eq!(rejections[0].reason, RejectReason::Excluded);
2001 }
2002 other => panic!("unexpected {other:?}"),
2003 }
2004 }
2005
2006 #[test]
2007 fn failover_is_bounded_and_never_after_output() {
2008 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2009 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
2010 let c = ok("c", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2011 let cands = [a, b, c];
2012 let ev = FailureEvidence {
2013 kind: FailureKind::QuotaExhaustedPreOutput,
2014 output_started: false,
2015 tool_activity: false,
2016 failovers_so_far: 0,
2017 };
2018 match failover(&req(), &cred("a"), &ev, &cands) {
2019 FailoverDecision::Failover { credential } => {
2020 assert_eq!(credential.storage_key(), "openai-codex@b")
2021 }
2022 other => panic!("unexpected {other:?}"),
2023 }
2024 let ev2 = FailureEvidence {
2026 failovers_so_far: 1,
2027 ..ev.clone()
2028 };
2029 assert!(matches!(
2030 failover(&req(), &cred("b"), &ev2, &cands),
2031 FailoverDecision::NoFailover { .. }
2032 ));
2033 let ev3 = FailureEvidence {
2035 output_started: true,
2036 ..ev.clone()
2037 };
2038 assert!(matches!(
2039 failover(&req(), &cred("a"), &ev3, &cands),
2040 FailoverDecision::NoFailover { .. }
2041 ));
2042 let ev4 = FailureEvidence {
2044 tool_activity: true,
2045 ..ev.clone()
2046 };
2047 assert!(matches!(
2048 failover(&req(), &cred("a"), &ev4, &cands),
2049 FailoverDecision::NoFailover { .. }
2050 ));
2051 let ev5 = FailureEvidence {
2053 kind: FailureKind::RateLimited,
2054 ..ev.clone()
2055 };
2056 assert!(matches!(
2057 failover(&req(), &cred("a"), &ev5, &cands),
2058 FailoverDecision::NoFailover { .. }
2059 ));
2060 let r = SelectionRequest {
2062 explicit_account: Some("a"),
2063 ..req()
2064 };
2065 assert!(matches!(
2066 failover(&r, &cred("a"), &ev, &cands),
2067 FailoverDecision::NoFailover { .. }
2068 ));
2069 let only = [ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)])];
2072 assert!(matches!(
2073 failover(&req(), &cred("a"), &ev, &only),
2074 FailoverDecision::NoFailover { .. }
2075 ));
2076 }
2077
2078 #[test]
2079 fn unknown_utilization_sorts_after_known() {
2080 let proven = ok(
2081 "p",
2082 vec![WindowLimit {
2083 used_percent: None,
2084 limit_reached: Some(false),
2085 ..win("primary", 7 * DAY, 0.0, NOW + DAY)
2086 }],
2087 );
2088 let known = ok("k", vec![win("primary", 7 * DAY, 95.0, NOW + DAY)]);
2089 assert_eq!(
2090 select(&req(), &[proven.clone(), known])
2091 .selected()
2092 .unwrap()
2093 .storage_key(),
2094 "openai-codex@k"
2095 );
2096 assert_eq!(
2097 select(&req(), &[proven]).selected().unwrap().storage_key(),
2098 "openai-codex@p"
2099 );
2100 }
2101
2102 #[test]
2103 fn model_matching_uses_family_tokens() {
2104 assert!(model_matches("sonnet", "claude-sonnet-4-6"));
2105 assert!(model_matches("Sonnet", "CLAUDE-SONNET-4-6"));
2106 assert!(model_matches("gpt-6-astra", "gpt-6-astra"));
2107 assert!(!model_matches("opus", "claude-sonnet-4-6"));
2108 assert!(!model_matches("", "x") && !model_matches("x", ""));
2109 let w = WindowLimit {
2110 models: Some(vec!["sonnet".into()]),
2111 ..win("seven_day_sonnet", 7 * DAY, 100.0, NOW + DAY)
2112 };
2113 assert!(w.applies_to(Some("claude-sonnet-4-6")));
2114 assert!(!w.applies_to(Some("claude-opus-4-6")));
2115 assert!(w.applies_to(None));
2116 let acct = ok("a", vec![win("seven_day", 7 * DAY, 1.0, NOW + DAY), w]);
2117 let r = SelectionRequest {
2118 model: Some("claude-sonnet-4-6"),
2119 ..req()
2120 };
2121 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
2122 let r = SelectionRequest {
2123 model: Some("claude-opus-4-6"),
2124 ..req()
2125 };
2126 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_some());
2127 }
2128
2129 #[test]
2130 fn serde_round_trip_is_secret_free() {
2131 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2132 let sel = select(&req(), &[a]);
2133 let json = serde_json::to_string(&sel).unwrap();
2134 assert!(json.contains("no_capacity"));
2135 let back: Selection = serde_json::from_str(&json).unwrap();
2136 assert_eq!(back, sel);
2137 }
2138
2139 fn sr_req() -> SelectionRequest<'static> {
2144 SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
2145 }
2146
2147 const OBSERVED: u64 = NOW - 1000;
2149
2150 fn anchored(used: f64, resets: u64) -> WindowLimit {
2153 win("primary", 7 * DAY, used, resets)
2154 }
2155
2156 fn unanchored() -> WindowLimit {
2159 win("primary", 7 * DAY, 0.0, OBSERVED + 7 * DAY)
2160 }
2161
2162 fn no_reset(used: f64) -> WindowLimit {
2164 WindowLimit {
2165 resets_at_ms: None,
2166 ..win("primary", 7 * DAY, used, 0)
2167 }
2168 }
2169
2170 fn label_of(sel: &Selection) -> &str {
2171 sel.selected().expect("a selection").account.label_str()
2172 }
2173
2174 fn ranks(sel: &Selection) -> Vec<(String, u32)> {
2175 let (Selection::Selected { rejections, .. } | Selection::NoCapacity { rejections, .. }) =
2176 sel;
2177 rejections
2178 .iter()
2179 .filter_map(|r| match r.reason {
2180 RejectReason::Outranked { rank } => {
2181 Some((r.credential.account.label_str().to_string(), rank))
2182 }
2183 _ => None,
2184 })
2185 .collect()
2186 }
2187
2188 #[test]
2189 fn soonest_reset_beats_lower_utilization_where_lowest_utilization_does_not() {
2190 let a = ok("a", vec![anchored(60.0, NOW + 2 * DAY)]);
2192 let b = ok("b", vec![anchored(10.0, NOW + 5 * DAY)]);
2193 let cands = [b.clone(), a.clone()];
2194 match select(&sr_req(), &cands) {
2195 Selection::Selected {
2196 credential,
2197 tier,
2198 budget_reset_ms,
2199 utilization,
2200 ..
2201 } => {
2202 assert_eq!(credential.storage_key(), "openai-codex@a");
2203 assert_eq!(tier, 3);
2204 assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
2205 assert_eq!(utilization, Some(60.0));
2206 }
2207 other => panic!("unexpected {other:?}"),
2208 }
2209 assert_eq!(label_of(&select(&req(), &cands)), "b");
2211 let r = req();
2213 assert_eq!(r.strategy, Strategy::LowestUtilization);
2214 assert!(!r.sticky && r.current.is_none());
2215 assert_eq!(r.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
2216 let s = sr_req();
2217 assert_eq!(s.strategy, Strategy::SoonestReset);
2218 assert!(s.sticky && s.current.is_none());
2219 assert_eq!(s.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
2220 }
2221
2222 #[test]
2223 fn tiers_order_urgent_then_unanchored_then_far_then_no_reset() {
2224 let urgent = ok("u", vec![anchored(50.0, NOW + 6 * HOUR)]);
2225 let sliding = ok("s", vec![unanchored()]);
2226 let far = ok("f", vec![anchored(0.0, NOW + 5 * DAY)]);
2227 let none = ok("n", vec![no_reset(0.0)]);
2228 let r = sr_req();
2229 for (cap, tier, anch) in [
2230 (&urgent, 1u8, Some(true)),
2231 (&sliding, 2, Some(false)),
2232 (&far, 3, Some(true)),
2233 (&none, 4, None),
2234 ] {
2235 let e = evaluate(&r, cap).unwrap();
2236 assert_eq!(e.tier, tier, "{}", cap.credential.storage_key());
2237 assert_eq!(e.anchored, anch, "{}", cap.credential.storage_key());
2238 }
2239 assert_eq!(evaluate(&req(), &sliding).unwrap().tier, 2);
2241 let cands = [none.clone(), far.clone(), sliding.clone(), urgent.clone()];
2243 let sel = select(&r, &cands);
2244 assert_eq!(label_of(&sel), "u");
2245 assert_eq!(
2246 ranks(&sel),
2247 vec![
2248 ("s".to_string(), 2),
2249 ("f".to_string(), 3),
2250 ("n".to_string(), 4)
2251 ]
2252 );
2253 assert_eq!(
2254 label_of(&select(&r, &[none.clone(), far.clone(), sliding])),
2255 "s"
2256 );
2257 assert_eq!(label_of(&select(&r, &[none.clone(), far])), "f");
2258 assert_eq!(label_of(&select(&r, &[none])), "n");
2259 let short = SelectionRequest {
2262 urgent_horizon_ms: HOUR,
2263 ..sr_req()
2264 };
2265 assert_eq!(evaluate(&short, &urgent).unwrap().tier, 3);
2266 assert_eq!(
2267 label_of(&select(&short, &[urgent, ok("s", vec![unanchored()])])),
2268 "s"
2269 );
2270 let edge = ok("e", vec![anchored(0.0, NOW + DEFAULT_URGENT_HORIZON_MS)]);
2272 assert_eq!(evaluate(&sr_req(), &edge).unwrap().tier, 1);
2273 }
2274
2275 #[test]
2276 fn within_tier_orders_by_reset_preference_or_utilization() {
2277 let a = ok("a", vec![anchored(90.0, NOW + 2 * HOUR)]);
2279 let b = ok("b", vec![anchored(1.0, NOW + 3 * HOUR)]);
2280 assert_eq!(label_of(&select(&sr_req(), &[b, a])), "a");
2281 let s1 = ok("s1", vec![unanchored()]);
2283 let s2 = ok("s2", vec![unanchored()]);
2284 assert_eq!(
2285 label_of(&select(&sr_req(), &[s2.clone(), s1.clone()])),
2286 "s1"
2287 );
2288 let pref = vec!["s2".to_string()];
2289 let r = SelectionRequest {
2290 preference: &pref,
2291 ..sr_req()
2292 };
2293 assert_eq!(label_of(&select(&r, &[s1, s2])), "s2");
2294 let n1 = ok("n1", vec![no_reset(40.0)]);
2296 let n2 = ok("n2", vec![no_reset(5.0)]);
2297 let n3 = ok(
2298 "n3",
2299 vec![WindowLimit {
2300 used_percent: None,
2301 limit_reached: Some(false),
2302 ..no_reset(0.0)
2303 }],
2304 );
2305 let sel = select(&sr_req(), &[n3, n1, n2]);
2306 assert_eq!(label_of(&sel), "n2");
2307 assert_eq!(
2308 ranks(&sel),
2309 vec![("n1".to_string(), 2), ("n3".to_string(), 3)]
2310 );
2311 let t1 = ok("t1", vec![anchored(30.0, NOW + 3 * DAY)]);
2313 let t2 = ok("t2", vec![anchored(20.0, NOW + 3 * DAY)]);
2314 assert_eq!(label_of(&select(&sr_req(), &[t1, t2])), "t2");
2315 }
2316
2317 #[test]
2318 fn exhausted_soonest_seat_is_skipped_to_next_soonest() {
2319 let x = ok("x", vec![anchored(100.0, NOW + HOUR)]);
2320 let a = ok("a", vec![anchored(20.0, NOW + 2 * DAY)]);
2321 let b = ok("b", vec![anchored(20.0, NOW + 3 * DAY)]);
2322 let sel = select(&sr_req(), &[b, x.clone(), a]);
2323 assert_eq!(label_of(&sel), "a");
2324 let Selection::Selected { rejections, .. } = &sel else {
2325 panic!("expected selection");
2326 };
2327 assert!(rejections
2328 .iter()
2329 .any(|r| r.credential.account.label_str() == "x"
2330 && r.reason
2331 == RejectReason::Exhausted {
2332 window: "primary".into(),
2333 resets_at_ms: Some(NOW + HOUR),
2334 }));
2335 let y = ok("y", vec![anchored(100.0, NOW + 2 * DAY)]);
2337 let past = ok("p", vec![anchored(100.0, NOW - HOUR)]);
2338 match select(&sr_req(), &[y, past, x]) {
2339 Selection::NoCapacity {
2340 earliest_reset_ms, ..
2341 } => assert_eq!(earliest_reset_ms, Some(NOW + HOUR)),
2342 other => panic!("unexpected {other:?}"),
2343 }
2344 }
2345
2346 #[test]
2347 fn five_hour_throttle_rejects_then_reenters_after_reset() {
2348 let throttled = ok(
2349 "t",
2350 vec![
2351 win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
2352 win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
2353 ],
2354 );
2355 let other = ok("o", vec![win("seven_day", 7 * DAY, 10.0, NOW + 4 * DAY)]);
2356 let sel = select(&sr_req(), &[throttled, other.clone()]);
2357 assert_eq!(label_of(&sel), "o");
2358 let Selection::Selected { rejections, .. } = &sel else {
2359 panic!("expected selection");
2360 };
2361 assert_eq!(
2362 rejections[0].reason,
2363 RejectReason::Exhausted {
2364 window: "five_hour".into(),
2365 resets_at_ms: Some(NOW + HOUR),
2366 }
2367 );
2368 let now2 = NOW + HOUR + 1;
2371 let mut throttled2 = ok(
2372 "t",
2373 vec![
2374 win("five_hour", 5 * HOUR, 0.0, now2 + 5 * HOUR),
2375 win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
2376 ],
2377 );
2378 throttled2.observed_at_ms = Some(now2 - 1000);
2379 let mut other2 = other;
2380 other2.observed_at_ms = Some(now2 - 1000);
2381 let r = SelectionRequest {
2382 now_ms: now2,
2383 ..sr_req()
2384 };
2385 match select(&r, &[other2, throttled2]) {
2386 Selection::Selected {
2387 credential,
2388 tier,
2389 budget_reset_ms,
2390 ..
2391 } => {
2392 assert_eq!(credential.storage_key(), "openai-codex@t");
2393 assert_eq!(tier, 3);
2394 assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
2395 }
2396 other => panic!("unexpected {other:?}"),
2397 }
2398 }
2399
2400 #[test]
2401 fn is_unanchored_classifies_sliding_resets_only() {
2402 let obs = NOW - 90_000;
2403 let codex_idle = win("primary", 7 * DAY, 0.0, obs + 7 * DAY);
2405 assert_eq!(is_unanchored(&codex_idle, obs), Some(true));
2406 let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY + 60_000);
2408 assert_eq!(is_unanchored(&skewed, obs), Some(true));
2409 let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY - 60_000);
2410 assert_eq!(is_unanchored(&skewed, obs), Some(true));
2411 let beyond = win(
2413 "primary",
2414 7 * DAY,
2415 0.0,
2416 obs + 7 * DAY + UNANCHORED_TOLERANCE_MS + 1,
2417 );
2418 assert_eq!(is_unanchored(&beyond, obs), Some(false));
2419 let anthropic = win("seven_day", 7 * DAY, 0.0, obs + 5 * DAY);
2421 assert_eq!(is_unanchored(&anthropic, obs), Some(false));
2422 let used = win("primary", 7 * DAY, 40.0, obs + 7 * DAY);
2424 assert_eq!(is_unanchored(&used, obs), Some(false));
2425 let half = win(
2427 "primary",
2428 7 * DAY,
2429 UNANCHORED_MAX_USED_PERCENT,
2430 obs + 7 * DAY,
2431 );
2432 assert_eq!(is_unanchored(&half, obs), Some(true));
2433 let unreported = WindowLimit {
2435 used_percent: None,
2436 limit_reached: Some(false),
2437 ..codex_idle.clone()
2438 };
2439 assert_eq!(is_unanchored(&unreported, obs), Some(false));
2440 let no_dur = WindowLimit {
2442 duration_ms: None,
2443 ..codex_idle.clone()
2444 };
2445 assert_eq!(is_unanchored(&no_dur, obs), None);
2446 let no_reset = WindowLimit {
2447 resets_at_ms: None,
2448 ..codex_idle
2449 };
2450 assert_eq!(is_unanchored(&no_reset, obs), None);
2451 }
2452
2453 #[test]
2454 fn budget_window_prefers_weekly_then_longest() {
2455 let five_h = win("five_hour", 5 * HOUR, 0.0, NOW + HOUR);
2456 let one_h = win("one_hour", HOUR, 0.0, NOW + HOUR);
2457 let thirty_d = win("thirty_day", 30 * DAY, 0.0, NOW + 20 * DAY);
2458 let weekly_wide = win("weekly", 7 * DAY, 0.0, NOW + 3 * DAY);
2459 let weekly_scoped = WindowLimit {
2460 id: "aaa_weekly_scoped".into(),
2461 models: Some(vec!["m".into()]),
2462 ..weekly_wide.clone()
2463 };
2464 let all = [
2466 thirty_d.clone(),
2467 weekly_scoped.clone(),
2468 five_h.clone(),
2469 weekly_wide.clone(),
2470 ];
2471 assert_eq!(
2472 budget_window(&all, Some("m")).map(|w| w.id.as_str()),
2473 Some("weekly")
2474 );
2475 assert_eq!(
2476 budget_window(&all, None).map(|w| w.id.as_str()),
2477 Some("weekly")
2478 );
2479 let no_wide = [thirty_d.clone(), weekly_scoped.clone(), five_h.clone()];
2480 assert_eq!(
2481 budget_window(&no_wide, Some("m")).map(|w| w.id.as_str()),
2482 Some("aaa_weekly_scoped")
2483 );
2484 assert_eq!(
2486 budget_window(&no_wide, Some("other")).map(|w| w.id.as_str()),
2487 Some("thirty_day")
2488 );
2489 assert_eq!(
2491 budget_window(&[five_h.clone(), one_h.clone()], None).map(|w| w.id.as_str()),
2492 Some("five_hour")
2493 );
2494 let b5 = WindowLimit {
2496 id: "b5".into(),
2497 ..five_h.clone()
2498 };
2499 let a5_scoped = WindowLimit {
2500 id: "a5".into(),
2501 models: Some(vec!["m".into()]),
2502 ..five_h.clone()
2503 };
2504 assert_eq!(
2505 budget_window(&[b5.clone(), a5_scoped.clone()], Some("m")).map(|w| w.id.as_str()),
2506 Some("b5")
2507 );
2508 let a5 = WindowLimit {
2509 id: "a5".into(),
2510 ..five_h.clone()
2511 };
2512 assert_eq!(
2513 budget_window(&[b5, a5], None).map(|w| w.id.as_str()),
2514 Some("a5")
2515 );
2516 let no_dur = WindowLimit {
2518 duration_ms: None,
2519 ..five_h
2520 };
2521 assert!(budget_window(&[no_dur.clone(), no_dur], None).is_none());
2522 assert!(budget_window(&[], None).is_none());
2523 assert_eq!(weekly_window(&all).map(|w| w.id.as_str()), Some("weekly"));
2525 }
2526
2527 #[test]
2528 fn past_reset_with_headroom_is_not_reset_evidence() {
2529 let stale_reset = ok("p", vec![anchored(20.0, NOW - 1)]);
2532 let e = evaluate(&sr_req(), &stale_reset).unwrap();
2533 assert_eq!(e.budget_reset_ms, None);
2534 assert_eq!(e.anchored, None);
2535 assert_eq!(e.tier, 4);
2536 let far = ok("f", vec![anchored(80.0, NOW + 6 * DAY)]);
2537 assert_eq!(label_of(&select(&sr_req(), &[stale_reset, far])), "f");
2538 }
2539
2540 #[test]
2541 fn sticky_keeps_current_within_tier_3() {
2542 let cur = cred("cur");
2543 let current = ok("cur", vec![anchored(30.0, NOW + 5 * DAY)]);
2544 let sooner = ok("soon", vec![anchored(10.0, NOW + 2 * DAY)]);
2545 let r = SelectionRequest {
2546 current: Some(&cur),
2547 ..sr_req()
2548 };
2549 let sel = select(&r, &[sooner.clone(), current.clone()]);
2550 assert_eq!(label_of(&sel), "cur");
2551 assert_eq!(ranks(&sel), vec![("soon".to_string(), 2)]);
2552 assert_eq!(
2554 label_of(&select(&sr_req(), &[sooner.clone(), current.clone()])),
2555 "soon"
2556 );
2557 let r = SelectionRequest { sticky: false, ..r };
2559 assert_eq!(label_of(&select(&r, &[sooner, current])), "soon");
2560 }
2561
2562 #[test]
2563 fn sticky_yields_to_strictly_higher_tier() {
2564 let cur = cred("cur");
2565 let current = ok("cur", vec![anchored(10.0, NOW + 5 * DAY)]);
2566 let urgent = ok("urg", vec![anchored(90.0, NOW + 6 * HOUR)]);
2567 let r = SelectionRequest {
2568 current: Some(&cur),
2569 ..sr_req()
2570 };
2571 assert_eq!(label_of(&select(&r, &[current.clone(), urgent])), "urg");
2572 let sliding = ok("sld", vec![unanchored()]);
2574 assert_eq!(label_of(&select(&r, &[current, sliding])), "sld");
2575 let cur_sliding = ok("cur", vec![unanchored()]);
2578 let a_sliding = ok("a", vec![unanchored()]);
2579 let far = ok("far", vec![anchored(0.0, NOW + 2 * DAY)]);
2580 assert_eq!(label_of(&select(&r, &[a_sliding, far, cur_sliding])), "cur");
2581 }
2582
2583 #[test]
2584 fn sticky_tier1_yields_only_to_strictly_sooner_reset() {
2585 let cur = cred("cur");
2586 let current = ok("cur", vec![anchored(50.0, NOW + 10 * HOUR)]);
2587 let r = SelectionRequest {
2588 current: Some(&cur),
2589 ..sr_req()
2590 };
2591 let sooner = ok("soon", vec![anchored(80.0, NOW + 6 * HOUR)]);
2592 assert_eq!(label_of(&select(&r, &[current.clone(), sooner])), "soon");
2593 let equal = ok("eq", vec![anchored(1.0, NOW + 10 * HOUR)]);
2596 let sel = select(&r, &[equal.clone(), current.clone()]);
2597 assert_eq!(label_of(&sel), "cur");
2598 assert_eq!(ranks(&sel), vec![("eq".to_string(), 2)]);
2599 let later = ok("late", vec![anchored(0.0, NOW + 12 * HOUR)]);
2601 assert_eq!(label_of(&select(&r, &[later, current])), "cur");
2602 }
2603
2604 #[test]
2605 fn sticky_never_keeps_ineligible_current_and_only_applies_to_soonest_reset() {
2606 let cur = cred("cur");
2607 let other = ok("o", vec![anchored(50.0, NOW + 5 * DAY)]);
2608 let r = SelectionRequest {
2609 current: Some(&cur),
2610 ..sr_req()
2611 };
2612 let exhausted = ok("cur", vec![anchored(100.0, NOW + DAY)]);
2614 assert_eq!(label_of(&select(&r, &[exhausted, other.clone()])), "o");
2615 let mut cooling = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2617 cooling.cooldown_until_ms = Some(NOW + HOUR);
2618 assert_eq!(label_of(&select(&r, &[cooling, other.clone()])), "o");
2619 let mut stale = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2621 stale.observed_at_ms = Some(NOW - 2 * HOUR);
2622 assert_eq!(label_of(&select(&r, &[stale, other.clone()])), "o");
2623 let excl = vec![cur.clone()];
2625 let r_excl = SelectionRequest {
2626 exclude: &excl,
2627 ..r.clone()
2628 };
2629 let healthy = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2630 assert_eq!(
2631 label_of(&select(&r_excl, &[healthy.clone(), other.clone()])),
2632 "o"
2633 );
2634 assert_eq!(label_of(&select(&r, std::slice::from_ref(&other))), "o");
2636 let busy_current = ok("cur", vec![anchored(90.0, NOW + 5 * DAY)]);
2638 let r_lu = SelectionRequest {
2639 strategy: Strategy::LowestUtilization,
2640 sticky: true,
2641 current: Some(&cur),
2642 ..req()
2643 };
2644 assert_eq!(
2645 label_of(&select(&r_lu, &[busy_current.clone(), other.clone()])),
2646 "o"
2647 );
2648 let pref = vec!["o".to_string()];
2649 let r_po = SelectionRequest {
2650 strategy: Strategy::PreferenceOrder,
2651 sticky: true,
2652 current: Some(&cur),
2653 preference: &pref,
2654 ..req()
2655 };
2656 assert_eq!(label_of(&select(&r_po, &[busy_current, other])), "o");
2657 }
2658
2659 #[test]
2660 fn failover_under_soonest_reset_never_repicks_sticky_failed_seat() {
2661 let cur = cred("a");
2662 let a = ok("a", vec![anchored(0.0, NOW + 2 * DAY)]);
2664 let b = ok("b", vec![anchored(50.0, NOW + 5 * DAY)]);
2665 let r = SelectionRequest {
2666 current: Some(&cur),
2667 ..sr_req()
2668 };
2669 assert_eq!(label_of(&select(&r, &[a.clone(), b.clone()])), "a");
2670 let ev = FailureEvidence {
2671 kind: FailureKind::QuotaExhaustedPreOutput,
2672 output_started: false,
2673 tool_activity: false,
2674 failovers_so_far: 0,
2675 };
2676 match failover(&r, &cur, &ev, &[a.clone(), b]) {
2677 FailoverDecision::Failover { credential } => {
2678 assert_eq!(credential.storage_key(), "openai-codex@b")
2679 }
2680 other => panic!("unexpected {other:?}"),
2681 }
2682 assert!(matches!(
2683 failover(&r, &cur, &ev, &[a]),
2684 FailoverDecision::NoFailover { .. }
2685 ));
2686 }
2687
2688 #[test]
2689 fn explicit_account_under_soonest_reset_never_falls_back() {
2690 let a = ok("a", vec![anchored(100.0, NOW + HOUR)]);
2691 let b = ok("b", vec![anchored(0.0, NOW + 6 * HOUR)]);
2692 let cur = cred("b");
2693 let r = SelectionRequest {
2695 explicit_account: Some("a"),
2696 current: Some(&cur),
2697 ..sr_req()
2698 };
2699 match select(&r, &[a.clone(), b.clone()]) {
2700 Selection::NoCapacity { rejections, .. } => {
2701 assert_eq!(rejections.len(), 1);
2702 assert_eq!(rejections[0].credential.storage_key(), "openai-codex@a");
2703 }
2704 other => panic!("unexpected {other:?}"),
2705 }
2706 let r = SelectionRequest {
2708 explicit_account: Some("b"),
2709 ..sr_req()
2710 };
2711 match select(&r, &[a, b]) {
2712 Selection::Selected {
2713 credential,
2714 rejections,
2715 tier,
2716 ..
2717 } => {
2718 assert_eq!(credential.storage_key(), "openai-codex@b");
2719 assert!(rejections.is_empty());
2720 assert_eq!(tier, 1);
2721 }
2722 other => panic!("unexpected {other:?}"),
2723 }
2724 }
2725
2726 fn permutations<T: Clone>(items: &[T]) -> Vec<Vec<T>> {
2728 fn go<T: Clone>(k: usize, a: &mut Vec<T>, out: &mut Vec<Vec<T>>) {
2729 if k <= 1 {
2730 out.push(a.clone());
2731 return;
2732 }
2733 go(k - 1, a, out);
2734 for i in 0..k - 1 {
2735 if k.is_multiple_of(2) {
2736 a.swap(i, k - 1);
2737 } else {
2738 a.swap(0, k - 1);
2739 }
2740 go(k - 1, a, out);
2741 }
2742 }
2743 let mut a = items.to_vec();
2744 let mut out = Vec::new();
2745 go(a.len(), &mut a, &mut out);
2746 out
2747 }
2748
2749 #[test]
2750 fn soonest_reset_is_deterministic_under_candidate_shuffle() {
2751 let cur = cred("t3b");
2752 let seats = vec![
2753 ok("t1", vec![anchored(70.0, NOW + 3 * HOUR)]),
2754 ok("t2", vec![unanchored()]),
2755 ok("t3a", vec![anchored(5.0, NOW + 2 * DAY)]),
2756 ok("t3b", vec![anchored(40.0, NOW + 4 * DAY)]),
2757 ok("x", vec![anchored(100.0, NOW + DAY)]),
2758 ];
2759 let r = SelectionRequest {
2760 current: Some(&cur),
2761 ..sr_req()
2762 };
2763 let perms = permutations(&seats);
2764 assert_eq!(perms.len(), 120);
2765 let baseline = select(&r, &seats);
2766 assert_eq!(label_of(&baseline), "t1");
2767 let normalize = |sel: &Selection| {
2768 let Selection::Selected {
2769 credential,
2770 tier,
2771 budget_reset_ms,
2772 utilization,
2773 rejections,
2774 } = sel
2775 else {
2776 panic!("expected selection");
2777 };
2778 let mut rej: Vec<(String, RejectReason)> = rejections
2779 .iter()
2780 .map(|r| (r.credential.storage_key(), r.reason.clone()))
2781 .collect();
2782 rej.sort_by(|a, b| a.0.cmp(&b.0));
2783 (
2784 credential.storage_key(),
2785 *tier,
2786 *budget_reset_ms,
2787 *utilization,
2788 rej,
2789 )
2790 };
2791 let expected = normalize(&baseline);
2792 assert_eq!(
2793 ranks(&baseline),
2794 vec![
2795 ("t2".to_string(), 2),
2796 ("t3a".to_string(), 3),
2797 ("t3b".to_string(), 4)
2798 ]
2799 );
2800 for p in perms {
2801 assert_eq!(normalize(&select(&r, &p)), expected);
2802 }
2803 }
2804
2805 #[test]
2806 fn outranked_ranks_are_two_to_n_in_ranking_order() {
2807 let seats = [
2808 ok("c", vec![anchored(0.0, NOW + 4 * DAY)]),
2809 ok("a", vec![anchored(0.0, NOW + 2 * DAY)]),
2810 ok("d", vec![anchored(0.0, NOW + 5 * DAY)]),
2811 ok("b", vec![anchored(0.0, NOW + 3 * DAY)]),
2812 ];
2813 let sel = select(&sr_req(), &seats);
2814 assert_eq!(label_of(&sel), "a");
2815 assert_eq!(
2816 ranks(&sel),
2817 vec![
2818 ("b".to_string(), 2),
2819 ("c".to_string(), 3),
2820 ("d".to_string(), 4)
2821 ]
2822 );
2823 let lu = [
2825 ok("hi", vec![anchored(80.0, NOW + DAY)]),
2826 ok("lo", vec![anchored(10.0, NOW + DAY)]),
2827 ok("mid", vec![anchored(50.0, NOW + DAY)]),
2828 ];
2829 let sel = select(&req(), &lu);
2830 assert_eq!(label_of(&sel), "lo");
2831 assert_eq!(
2832 ranks(&sel),
2833 vec![("mid".to_string(), 2), ("hi".to_string(), 3)]
2834 );
2835 assert_eq!(
2836 RejectReason::Outranked { rank: 3 }.to_string(),
2837 "eligible; outranked (rank 3)"
2838 );
2839 }
2840
2841 #[test]
2842 fn soonest_reset_serde_is_backward_compatible() {
2843 assert_eq!(
2844 serde_json::to_string(&Strategy::SoonestReset).unwrap(),
2845 "\"soonest_reset\""
2846 );
2847 assert_eq!(
2848 serde_json::from_str::<Strategy>("\"soonest_reset\"").unwrap(),
2849 Strategy::SoonestReset
2850 );
2851 let a = ok("a", vec![anchored(10.0, NOW + 2 * HOUR)]);
2852 let b = ok("b", vec![anchored(10.0, NOW + 3 * HOUR)]);
2853 let sel = select(&sr_req(), &[a, b]);
2854 let json = serde_json::to_string(&sel).unwrap();
2855 assert!(json.contains("\"tier\":1"));
2856 assert!(json.contains("\"budget_reset_ms\""));
2857 assert!(json.contains("\"rank\":2"));
2858 assert_eq!(serde_json::from_str::<Selection>(&json).unwrap(), sel);
2859 let old = r#"{"outcome":"selected","credential":"openai-codex@a","utilization":10.0,"rejections":[{"credential":"openai-codex@b","reason":{"reason":"outranked"}}]}"#;
2861 match serde_json::from_str::<Selection>(old).unwrap() {
2862 Selection::Selected {
2863 tier,
2864 budget_reset_ms,
2865 rejections,
2866 ..
2867 } => {
2868 assert_eq!(tier, 0);
2869 assert_eq!(budget_reset_ms, None);
2870 assert_eq!(rejections[0].reason, RejectReason::Outranked { rank: 0 });
2871 }
2872 other => panic!("unexpected {other:?}"),
2873 }
2874 }
2875
2876 fn failed(label: &str, failure: ObservationFailure) -> AccountCapacity {
2879 AccountCapacity::failed(cred(label), failure)
2880 }
2881
2882 fn failure_of(candidates: &[AccountCapacity]) -> SelectionFailure {
2883 match select(&req(), candidates) {
2884 Selection::NoCapacity { rejections, .. } => {
2885 SelectionFailure::from_rejections(&rejections, NOW)
2886 }
2887 other => panic!("expected no capacity, got {other:?}"),
2888 }
2889 }
2890
2891 #[test]
2892 fn observation_failure_maps_to_a_never_capacity_observation() {
2893 use AccountFailureKind as K;
2894 for (kind, obs) in [
2895 (K::UsageRateLimited, QuotaObservation::Unknown),
2896 (K::RefreshTransient, QuotaObservation::Unknown),
2897 (K::UsageUnauthorized, QuotaObservation::AuthError),
2898 (K::RefreshRejected, QuotaObservation::AuthError),
2899 (K::CredentialMissing, QuotaObservation::AuthError),
2900 (K::UsageMalformed, QuotaObservation::Malformed),
2901 (K::Unsupported, QuotaObservation::Unsupported),
2902 (K::RefreshActionRequired, QuotaObservation::Unknown),
2903 ] {
2904 assert_eq!(ObservationFailure::new(kind).observation(), obs, "{kind}");
2905 }
2906 let f = ObservationFailure::new(K::UsageRateLimited)
2907 .with_status(429)
2908 .with_retry_after(Some(120));
2909 assert_eq!(
2910 f.to_string(),
2911 "usage_rate_limited, HTTP 429, retry after 120s"
2912 );
2913 assert_eq!(
2914 serde_json::to_value(&f).unwrap(),
2915 serde_json::json!({"kind": "usage_rate_limited", "http_status": 429, "retry_after_s": 120})
2916 );
2917 let old: Rejection = serde_json::from_str(
2919 r#"{"credential":"openai-codex@a","reason":{"reason":"unknown"}}"#,
2920 )
2921 .unwrap();
2922 assert_eq!(old.cause, None);
2923 assert_eq!(old.describe(), "quota unknown");
2924 }
2925
2926 #[test]
2927 fn rejection_keeps_the_observation_cause_beside_a_cooldown() {
2928 use AccountFailureKind as K;
2929 let limited = ObservationFailure::new(K::UsageRateLimited).with_status(429);
2930 let mut cooling = failed("b", limited.clone());
2931 cooling.cooldown_until_ms = Some(NOW + HOUR);
2932 match select(&req(), &[failed("a", limited.clone()), cooling]) {
2933 Selection::NoCapacity { rejections, .. } => {
2934 assert_eq!(rejections[0].reason, RejectReason::Unknown);
2935 assert_eq!(rejections[0].cause.as_ref(), Some(&limited));
2936 assert_eq!(
2937 rejections[0].describe(),
2938 "quota unknown (usage_rate_limited, HTTP 429)"
2939 );
2940 assert!(matches!(
2943 rejections[1].reason,
2944 RejectReason::CoolingDown { .. }
2945 ));
2946 assert_eq!(rejections[1].cause.as_ref(), Some(&limited));
2947 }
2948 other => panic!("unexpected {other:?}"),
2949 }
2950 }
2951
2952 #[test]
2956 fn a_cooling_accounts_failure_keeps_its_read_kind_and_counts_as_limited() {
2957 use AccountFailureKind as K;
2958 use SelectionFailureKind as S;
2959 let mut refused = failed(
2960 "a",
2961 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
2962 );
2963 refused.cooldown_until_ms = Some(NOW + 2 * HOUR);
2964 let mut throttled = failed(
2965 "b",
2966 ObservationFailure::new(K::UsageRateLimited)
2967 .with_status(429)
2968 .with_retry_after(Some(90)),
2969 );
2970 throttled.cooldown_until_ms = Some(NOW + HOUR);
2971 let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
2972 plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
2973
2974 let f = failure_of(&[refused, throttled, plain]);
2975 assert_eq!(f.kind, S::AllCoolingDown, "{f:?}");
2976 let a = &f.accounts[0];
2977 assert_eq!(
2978 (a.kind, a.http_status, a.cooling_down, a.until_ms),
2979 (K::UsageUnauthorized, Some(401), true, Some(NOW + 2 * HOUR))
2980 );
2981 let b = &f.accounts[1];
2982 assert_eq!(
2983 (b.kind, b.http_status, b.retry_after_s, b.cooling_down),
2984 (K::UsageRateLimited, Some(429), Some(90), true)
2985 );
2986 let c = &f.accounts[2];
2987 assert_eq!((c.kind, c.cooling_down), (K::CoolingDown, true));
2988 assert_eq!(f.login_labels(), vec!["a"]);
2989 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
2990 assert_eq!(
2991 f.accounts_field(),
2992 "a:usage_unauthorized:401:cooling,b:usage_rate_limited:429:cooling,c:cooling_down"
2993 );
2994 let v = serde_json::to_value(&f).unwrap();
2996 assert_eq!(v["accounts"][0]["cooling_down"], true);
2997 let g = failure_of(&[failed("d", ObservationFailure::new(K::UsageTimeout))]);
2998 assert!(serde_json::to_value(&g).unwrap()["accounts"][0]
2999 .get("cooling_down")
3000 .is_none());
3001 }
3002
3003 #[test]
3010 fn the_headline_skips_a_cooling_accounts_read_failure_which_stays_per_account() {
3011 use AccountFailureKind as K;
3012 use SelectionFailureKind as S;
3013 let throttled = || {
3014 let mut c = failed(
3015 "claude5",
3016 ObservationFailure::new(K::UsageRateLimited)
3017 .with_status(429)
3018 .with_retry_after(Some(90)),
3019 );
3020 c.cooldown_until_ms = Some(NOW + HOUR);
3021 c
3022 };
3023 let c5_kept = |f: &SelectionFailure| {
3024 let c5 = &f.accounts[0];
3025 assert_eq!(
3026 (
3027 c5.account.as_str(),
3028 c5.kind,
3029 c5.http_status,
3030 c5.retry_after_s,
3031 c5.cooling_down,
3032 c5.until_ms
3033 ),
3034 (
3035 "claude5",
3036 K::UsageRateLimited,
3037 Some(429),
3038 Some(90),
3039 true,
3040 Some(NOW + HOUR)
3041 ),
3042 "{f:?}"
3043 );
3044 };
3045
3046 let mut plain = ok("claude6", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3047 plain.cooldown_until_ms = Some(NOW + 2 * HOUR);
3048 let a = failure_of(&[throttled(), plain]);
3049 assert_eq!(a.kind, S::AllCoolingDown, "{a:?}");
3050 c5_kept(&a);
3051 assert_eq!(
3052 a.accounts_field(),
3053 "claude5:usage_rate_limited:429:cooling,claude6:cooling_down"
3054 );
3055 assert_eq!(a.earliest_reset_ms, Some(NOW + HOUR));
3056 assert!(a.login_labels().is_empty(), "{a:?}");
3057
3058 let refused = failed(
3059 "claude2",
3060 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
3061 );
3062 let b = failure_of(&[throttled(), refused]);
3063 assert_eq!(b.kind, S::CredentialRejected, "{b:?}");
3064 c5_kept(&b);
3065 assert_eq!(
3066 b.accounts_field(),
3067 "claude5:usage_rate_limited:429:cooling,claude2:usage_unauthorized:401"
3068 );
3069 assert_eq!(b.earliest_reset_ms, Some(NOW + HOUR));
3070 assert_eq!(b.login_labels(), vec!["claude2"]);
3071 }
3072
3073 #[test]
3077 fn an_accounts_until_is_its_latest_exhausted_reset_and_cooldown() {
3078 let later = ok(
3080 "a",
3081 vec![
3082 win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
3083 win("seven_day", 7 * DAY, 100.0, NOW + 3 * DAY),
3084 ],
3085 );
3086 let earlier = ok(
3088 "b",
3089 vec![
3090 win("seven_day", 7 * DAY, 100.0, NOW + 2 * DAY),
3091 win("five_hour", 5 * HOUR, 100.0, NOW + 2 * HOUR),
3092 ],
3093 );
3094 let mut cooled = ok("c", vec![win("five_hour", 5 * HOUR, 100.0, NOW + HOUR)]);
3096 cooled.cooldown_until_ms = Some(NOW + 4 * HOUR);
3097 let mut unknown = ok("d", vec![win("seven_day", 7 * DAY, 100.0, NOW + DAY)]);
3099 if let QuotaObservation::Ok { windows, .. } = &mut unknown.observation {
3100 windows.push(WindowLimit {
3101 resets_at_ms: None,
3102 ..win("five_hour", 5 * HOUR, 100.0, 0)
3103 });
3104 }
3105
3106 let until = |f: &SelectionFailure, label: &str| {
3107 f.accounts
3108 .iter()
3109 .find(|a| a.account == label)
3110 .unwrap()
3111 .until_ms
3112 };
3113 let f = failure_of(&[later.clone(), earlier.clone(), cooled.clone(), unknown]);
3114 assert_eq!(until(&f, "a"), Some(NOW + 3 * DAY));
3115 assert_eq!(until(&f, "b"), Some(NOW + 2 * DAY));
3116 assert_eq!(until(&f, "c"), Some(NOW + 4 * HOUR));
3117 assert_eq!(until(&f, "d"), None);
3118 assert_eq!(f.earliest_reset_ms, Some(NOW + 4 * HOUR));
3119 assert_eq!(
3120 failure_of(&[later, earlier]).earliest_reset_ms,
3121 Some(NOW + 2 * DAY)
3122 );
3123 match select(&req(), &[cooled]) {
3126 Selection::NoCapacity { rejections, .. } => {
3127 assert!(matches!(
3128 rejections[0].reason,
3129 RejectReason::CoolingDown { .. }
3130 ));
3131 }
3132 other => panic!("unexpected {other:?}"),
3133 }
3134 }
3135
3136 #[test]
3141 fn the_earliest_reset_skips_accounts_that_need_a_login_or_an_action() {
3142 use AccountFailureKind as K;
3143 let mut refused = failed(
3144 "a",
3145 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
3146 );
3147 refused.cooldown_until_ms = Some(NOW + HOUR);
3148 let mut broken = failed("b", ObservationFailure::new(K::CredentialStore));
3149 broken.cooldown_until_ms = Some(NOW + 2 * HOUR);
3150 let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3151 plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
3152
3153 let f = failure_of(&[refused.clone(), broken.clone(), plain]);
3154 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3155 assert_eq!(f.earliest_reset_ms, Some(NOW + 3 * HOUR));
3156 assert_eq!(f.accounts[0].until_ms, Some(NOW + HOUR));
3157 assert_eq!(f.accounts[1].until_ms, Some(NOW + 2 * HOUR));
3158 assert_eq!(failure_of(&[refused, broken]).earliest_reset_ms, None);
3159 }
3160
3161 #[test]
3164 fn action_required_and_store_failures_get_their_own_kind() {
3165 use AccountFailureKind as K;
3166 use SelectionFailureKind as S;
3167 let action = ObservationFailure::new(K::RefreshActionRequired);
3168 let store = ObservationFailure::new(K::CredentialStore);
3169 let refused = ObservationFailure::new(K::UsageUnauthorized).with_status(401);
3170 let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3171 cooling.cooldown_until_ms = Some(NOW + HOUR);
3172
3173 for f in [
3174 failure_of(&[failed("a", action.clone())]),
3175 failure_of(&[failed("a", store.clone())]),
3176 failure_of(&[
3177 failed("a", action.clone()),
3178 failed("b", store.clone()),
3179 cooling.clone(),
3180 ]),
3181 failure_of(&[failed("a", action.clone()), failed("b", refused.clone())]),
3182 ] {
3183 assert_eq!(f.kind, S::ActionRequired, "{f:?}");
3184 }
3185 assert!(K::RefreshActionRequired.needs_action());
3186 assert!(K::CredentialStore.needs_action());
3187 assert!(
3188 !K::RefreshActionRequired.needs_login() && !K::RefreshActionRequired.is_transient()
3189 );
3190 let f = failure_of(&[failed("a", action.clone()), failed("b", store.clone())]);
3191 assert_eq!(f.action_labels(), vec!["a", "b"]);
3192 assert!(f.login_labels().is_empty());
3193 let f = failure_of(&[
3195 failed("a", action),
3196 failed("b", ObservationFailure::new(K::UsageTimeout)),
3197 ]);
3198 assert_eq!(f.kind, S::UsageReadFailed);
3199 assert_eq!(f.action_labels(), vec!["a"]);
3200 assert_eq!(S::ActionRequired.as_str(), "action_required");
3201 assert_eq!(S::CredentialStore.as_str(), "credential_store");
3202 }
3203
3204 #[test]
3205 fn selection_failure_kind_priority_and_login_labels() {
3206 use AccountFailureKind as K;
3207 use SelectionFailureKind as S;
3208 let transient = ObservationFailure::new(K::UsageUpstream).with_status(503);
3209 let rejected = ObservationFailure::new(K::RefreshRejected);
3210 let exhausted = ok("x", vec![win("primary", 5 * HOUR, 100.0, NOW + 2 * HOUR)]);
3211 let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3212 cooling.cooldown_until_ms = Some(NOW + HOUR);
3213
3214 assert_eq!(failure_of(&[]).kind, S::NoCredentials);
3215
3216 let f = failure_of(&[
3218 failed("a", rejected.clone()),
3219 failed("b", transient.clone()),
3220 ]);
3221 assert_eq!(f.kind, S::UsageReadFailed);
3222 assert_eq!(f.login_labels(), vec!["a"]);
3223 assert_eq!(
3224 f.accounts_field(),
3225 "a:refresh_rejected,b:usage_upstream:503"
3226 );
3227
3228 let f = failure_of(&[failed("a", rejected.clone()), cooling.clone()]);
3230 assert_eq!(f.kind, S::CredentialRejected);
3231 assert_eq!(f.login_labels(), vec!["a"]);
3232 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3233
3234 let f = failure_of(&[exhausted.clone(), cooling.clone()]);
3237 assert_eq!(f.kind, S::AllCoolingDown);
3238 assert_eq!(f.accounts_field(), "x:exhausted,c:cooling_down");
3239 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3240 assert!(f.login_labels().is_empty());
3241
3242 let no_evidence = ok("n", vec![]);
3245 let f = failure_of(&[failed("a", rejected), no_evidence]);
3246 assert_eq!(f.kind, S::QuotaUnknown);
3247 assert_eq!(f.login_labels(), vec!["a"]);
3248
3249 for kind in [
3251 K::RefreshLockBusy,
3252 K::RefreshOther,
3253 K::UsageTimeout,
3254 K::UsageTransport,
3255 ] {
3256 let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
3257 assert_eq!(f.kind, S::UsageReadFailed, "{kind}");
3258 }
3259 for kind in [K::RefreshActionRequired, K::CredentialStore] {
3262 let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
3263 assert_eq!(f.kind, S::ActionRequired, "{kind}");
3264 assert!(f.login_labels().is_empty(), "{kind}");
3265 }
3266 let f = failure_of(&[failed("a", ObservationFailure::new(K::UsageMalformed))]);
3267 assert_eq!(f.kind, S::QuotaUnknown);
3268 assert!(f.login_labels().is_empty());
3269 let stale_exhausted = ok("y", vec![win("primary", 5 * HOUR, 100.0, NOW - 1)]);
3271 assert_eq!(failure_of(&[stale_exhausted]).earliest_reset_ms, None);
3272 }
3273 #[test]
3278 fn an_excluded_duplicate_seat_does_not_pick_the_headline() {
3279 use AccountFailureKind as K;
3280 let reset = NOW + 2 * HOUR;
3281 let spent = ok("claude1", vec![win("primary", 5 * HOUR, 100.0, reset)]);
3282 let dup = ok("claude1-dup", vec![win("primary", 5 * HOUR, 100.0, reset)]);
3283 let exclude = [cred("claude1-dup")];
3284 let request = SelectionRequest {
3285 exclude: &exclude,
3286 ..req()
3287 };
3288 let f = match select(&request, &[spent, dup]) {
3289 Selection::NoCapacity { rejections, .. } => {
3290 SelectionFailure::from_rejections(&rejections, NOW)
3291 }
3292 other => panic!("expected no capacity, got {other:?}"),
3293 };
3294 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3295 assert_eq!(f.earliest_reset_ms, Some(reset));
3296 assert_eq!(f.accounts[1].kind, K::Excluded);
3297
3298 let mut cooling = ok("a", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3300 cooling.cooldown_until_ms = Some(NOW + HOUR);
3301 let dup = failed("b", ObservationFailure::new(K::UsageRateLimited));
3302 let exclude = [cred("b")];
3303 let request = SelectionRequest {
3304 exclude: &exclude,
3305 ..req()
3306 };
3307 let f = match select(&request, &[cooling, dup]) {
3308 Selection::NoCapacity { rejections, .. } => {
3309 SelectionFailure::from_rejections(&rejections, NOW)
3310 }
3311 other => panic!("expected no capacity, got {other:?}"),
3312 };
3313 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3314 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3315
3316 let only = ok("x", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3319 let exclude = [cred("x")];
3320 let request = SelectionRequest {
3321 exclude: &exclude,
3322 ..req()
3323 };
3324 let f = match select(&request, &[only]) {
3325 Selection::NoCapacity { rejections, .. } => {
3326 SelectionFailure::from_rejections(&rejections, NOW)
3327 }
3328 other => panic!("expected no capacity, got {other:?}"),
3329 };
3330 assert_eq!(f.kind, SelectionFailureKind::QuotaUnknown, "{f:?}");
3331 }
3332}