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();
880 let kind = if accounts.is_empty() {
881 K::NoCredentials
882 } else if live.iter().any(|a| a.kind.is_transient()) {
883 K::UsageReadFailed
884 } else if live.is_empty() && accounts.iter().any(AccountFailure::is_limited) {
885 K::AllCoolingDown
886 } else if live.is_empty() {
887 K::QuotaUnknown
888 } else if live.iter().all(|a| a.kind.needs_login()) {
889 K::CredentialRejected
890 } else if live
891 .iter()
892 .all(|a| a.kind.needs_login() || a.kind.needs_action())
893 {
894 K::ActionRequired
895 } else {
896 K::QuotaUnknown
897 };
898 let earliest_reset_ms = accounts
899 .iter()
900 .filter(|a| a.is_limited() && !a.kind.needs_login() && !a.kind.needs_action())
901 .filter_map(|a| a.until_ms)
902 .filter(|t| *t > now_ms)
903 .min();
904 Self {
905 kind,
906 accounts,
907 earliest_reset_ms,
908 usage_read_retry_ms: None,
909 }
910 }
911
912 pub fn transient_retry_after_s(&self) -> Option<u64> {
917 let waits: Option<Vec<u64>> = self
918 .accounts
919 .iter()
920 .filter(|a| a.is_live() && a.kind.is_transient())
921 .map(|a| a.retry_after_s)
922 .collect();
923 waits?.into_iter().min()
924 }
925
926 pub fn login_labels(&self) -> Vec<&str> {
928 self.accounts
929 .iter()
930 .filter(|a| a.kind.needs_login())
931 .map(|a| a.account.as_str())
932 .collect()
933 }
934
935 pub fn action_labels(&self) -> Vec<&str> {
937 self.accounts
938 .iter()
939 .filter(|a| a.kind.needs_action())
940 .map(|a| a.account.as_str())
941 .collect()
942 }
943
944 pub fn accounts_field(&self) -> String {
946 self.accounts
947 .iter()
948 .map(AccountFailure::log_form)
949 .collect::<Vec<_>>()
950 .join(",")
951 }
952}
953
954#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
956#[serde(tag = "outcome", rename_all = "snake_case")]
957pub enum Selection {
958 Selected {
959 #[serde(with = "cred_serde")]
960 credential: CredentialRef,
961 #[serde(default, skip_serializing_if = "Option::is_none")]
963 utilization: Option<f64>,
964 #[serde(default)]
968 tier: u8,
969 #[serde(default, skip_serializing_if = "Option::is_none")]
971 budget_reset_ms: Option<u64>,
972 rejections: Vec<Rejection>,
974 },
975 NoCapacity {
976 rejections: Vec<Rejection>,
977 #[serde(default, skip_serializing_if = "Option::is_none")]
981 earliest_reset_ms: Option<u64>,
982 },
983}
984
985impl Selection {
986 pub fn selected(&self) -> Option<&CredentialRef> {
987 match self {
988 Self::Selected { credential, .. } => Some(credential),
989 Self::NoCapacity { .. } => None,
990 }
991 }
992}
993
994#[derive(Debug, Clone, PartialEq)]
998pub struct Eligible {
999 pub utilization: Option<f64>,
1001 pub budget_reset_ms: Option<u64>,
1005 pub anchored: Option<bool>,
1010 pub tier: u8,
1018 pub graced: bool,
1021}
1022
1023fn tier_for(
1025 anchored: Option<bool>,
1026 budget_reset_ms: Option<u64>,
1027 now_ms: u64,
1028 urgent_horizon_ms: u64,
1029) -> u8 {
1030 match (anchored, budget_reset_ms) {
1031 (_, None) => 4,
1032 (Some(false), Some(_)) => 2,
1033 (_, Some(reset)) if reset.saturating_sub(now_ms) <= urgent_horizon_ms => 1,
1034 (_, Some(_)) => 3,
1035 }
1036}
1037
1038pub fn evaluate(
1041 req: &SelectionRequest<'_>,
1042 cap: &AccountCapacity,
1043) -> Result<Eligible, RejectReason> {
1044 if let Some(until) = cap.cooldown_until_ms {
1049 if until > req.now_ms {
1050 return Err(RejectReason::CoolingDown { until_ms: until });
1051 }
1052 }
1053 if req
1054 .exclude
1055 .iter()
1056 .any(|c| c.storage_key() == cap.credential.storage_key())
1057 {
1058 return Err(RejectReason::Excluded);
1059 }
1060 let (windows, models) = match &cap.observation {
1061 QuotaObservation::Ok { windows, models } => (windows, models),
1062 QuotaObservation::AuthError => return Err(RejectReason::AuthError),
1063 QuotaObservation::Unsupported => return Err(RejectReason::Unsupported),
1064 QuotaObservation::Malformed => return Err(RejectReason::Malformed),
1065 QuotaObservation::Unknown => return evaluate_graced(req, cap),
1066 };
1067 let observed = cap
1068 .observed_at_ms
1069 .ok_or(RejectReason::Stale { age_ms: None })?;
1070 if observed > req.now_ms.saturating_add(MAX_FUTURE_OBSERVATION_SKEW_MS) {
1071 return Err(RejectReason::Malformed);
1072 }
1073 let age = req.now_ms.saturating_sub(observed);
1074 if age > req.max_snapshot_age_ms {
1075 return Err(RejectReason::Stale { age_ms: Some(age) });
1076 }
1077 if let (Some(model), Some(list)) = (req.model, models) {
1078 for m in list.iter().filter(|m| model_matches(&m.model, model)) {
1079 match m.state {
1080 ModelState::Available => {}
1081 ModelState::Exhausted => {
1082 return Err(RejectReason::ModelUnavailable {
1083 model: model.to_string(),
1084 })
1085 }
1086 ModelState::Unknown => {
1087 return Err(RejectReason::ModelUnknown {
1088 model: model.to_string(),
1089 })
1090 }
1091 }
1092 }
1093 }
1094 let mut applicable = 0usize;
1095 let mut utilization: Option<f64> = None;
1096 for w in windows.iter().filter(|w| w.applies_to(req.model)) {
1097 applicable += 1;
1098 match w.verdict(req.exhausted_at_percent) {
1099 WindowVerdict::Headroom { used_percent } => {
1100 if let Some(p) = used_percent {
1101 utilization = Some(utilization.map_or(p, |u| u.max(p)));
1102 }
1103 }
1104 WindowVerdict::Exhausted { resets_at_ms } => {
1105 return Err(RejectReason::Exhausted {
1106 window: w.id.clone(),
1107 resets_at_ms,
1108 });
1109 }
1110 WindowVerdict::Unknown => {
1111 return Err(RejectReason::WindowUnknown {
1112 window: w.id.clone(),
1113 });
1114 }
1115 }
1116 }
1117 if applicable == 0 {
1118 return Err(RejectReason::NoLimitEvidence);
1119 }
1120 let budget = budget_window(windows, req.model);
1124 let budget_reset_ms = budget
1125 .and_then(|w| w.resets_at_ms)
1126 .filter(|t| *t > req.now_ms);
1127 let anchored = match (budget, budget_reset_ms) {
1128 (Some(w), Some(_)) => is_unanchored(w, observed).map(|u| !u),
1129 _ => None,
1130 };
1131 let tier = tier_for(anchored, budget_reset_ms, req.now_ms, req.urgent_horizon_ms);
1132 Ok(Eligible {
1133 utilization,
1134 budget_reset_ms,
1135 anchored,
1136 tier,
1137 graced: false,
1138 })
1139}
1140
1141fn evaluate_graced(
1147 req: &SelectionRequest<'_>,
1148 cap: &AccountCapacity,
1149) -> Result<Eligible, RejectReason> {
1150 let key = cap.credential.storage_key();
1151 let Some(good) = req
1152 .grace
1153 .iter()
1154 .find(|g| g.credential.storage_key() == key)
1155 .filter(|g| matches!(g.observation, QuotaObservation::Ok { .. }))
1156 else {
1157 return Err(RejectReason::Unknown);
1158 };
1159 let young = good
1162 .observed_at_ms
1163 .is_some_and(|t| req.now_ms.saturating_sub(t) < req.grace_max_age_ms);
1164 if !young {
1165 return Err(RejectReason::Unknown);
1166 }
1167 let graced_req = SelectionRequest {
1168 max_snapshot_age_ms: req.grace_max_age_ms,
1169 grace: &[],
1170 ..req.clone()
1171 };
1172 let reading = AccountCapacity {
1173 cooldown_until_ms: cap.cooldown_until_ms,
1174 ..good.clone()
1175 };
1176 evaluate(&graced_req, &reading)
1177 .map(|e| Eligible { graced: true, ..e })
1178 .map_err(|_| RejectReason::Unknown)
1179}
1180
1181fn utilization_cmp(a: Option<f64>, b: Option<f64>) -> std::cmp::Ordering {
1183 a.is_none().cmp(&b.is_none()).then_with(|| {
1184 a.unwrap_or(0.0)
1185 .partial_cmp(&b.unwrap_or(0.0))
1186 .unwrap_or(std::cmp::Ordering::Equal)
1187 })
1188}
1189
1190fn preference_rank(req: &SelectionRequest<'_>, cred: &CredentialRef) -> usize {
1191 let label = cred.account.label_str();
1192 req.preference
1193 .iter()
1194 .position(|p| p == label)
1195 .unwrap_or(usize::MAX)
1196}
1197
1198fn observation_cause(reason: &RejectReason, cap: &AccountCapacity) -> Option<ObservationFailure> {
1203 match reason {
1204 RejectReason::Unknown
1205 | RejectReason::AuthError
1206 | RejectReason::Malformed
1207 | RejectReason::Unsupported
1208 | RejectReason::CoolingDown { .. } => cap.failure.clone(),
1209 _ => None,
1210 }
1211}
1212
1213fn usable_after_ms(req: &SelectionRequest<'_>, cap: &AccountCapacity) -> Option<u64> {
1219 let mut latest = cap.cooldown_until_ms.filter(|t| *t > req.now_ms);
1220 if let QuotaObservation::Ok { windows, .. } = &cap.observation {
1221 for w in windows.iter().filter(|w| w.applies_to(req.model)) {
1222 if let WindowVerdict::Exhausted { resets_at_ms } = w.verdict(req.exhausted_at_percent) {
1223 let reset = resets_at_ms.filter(|t| *t > req.now_ms)?;
1224 latest = Some(latest.map_or(reset, |l| l.max(reset)));
1225 }
1226 }
1227 }
1228 latest
1229}
1230
1231pub fn select(req: &SelectionRequest<'_>, candidates: &[AccountCapacity]) -> Selection {
1235 let mut rejections: Vec<Rejection> = Vec::new();
1236 let mut eligible: Vec<(&AccountCapacity, Eligible)> = Vec::new();
1237 let mut explicit_seen = false;
1238
1239 for cap in candidates
1240 .iter()
1241 .filter(|c| c.credential.provider == req.provider)
1242 {
1243 if let Some(label) = req.explicit_account {
1244 if cap.credential.account.label_str() != label {
1245 continue;
1246 }
1247 explicit_seen = true;
1248 }
1249 match evaluate(req, cap) {
1250 Ok(e) => eligible.push((cap, e)),
1251 Err(reason) => rejections.push(Rejection {
1252 credential: cap.credential.clone(),
1253 cause: observation_cause(&reason, cap),
1254 usable_after_ms: usable_after_ms(req, cap),
1255 reason,
1256 }),
1257 }
1258 }
1259
1260 if let (Some(label), false) = (req.explicit_account, explicit_seen) {
1261 let (credential, reason) = match super::account::Account::parse(label) {
1265 Ok(account) => (
1266 CredentialRef::new(req.provider, account),
1267 RejectReason::NotPresent,
1268 ),
1269 Err(_) => (
1270 CredentialRef::default_for(req.provider),
1271 RejectReason::InvalidLabel {
1272 label: label.to_string(),
1273 },
1274 ),
1275 };
1276 return Selection::NoCapacity {
1277 rejections: vec![Rejection {
1278 credential,
1279 reason,
1280 cause: None,
1281 usable_after_ms: None,
1282 }],
1283 earliest_reset_ms: None,
1284 };
1285 }
1286
1287 eligible.sort_by(|(a, ea), (b, eb)| {
1291 let ka = a.credential.storage_key();
1292 let kb = b.credential.storage_key();
1293 let pa = preference_rank(req, &a.credential);
1294 let pb = preference_rank(req, &b.credential);
1295 ea.graced.cmp(&eb.graced).then_with(|| match req.strategy {
1296 Strategy::PreferenceOrder => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
1297 Strategy::LowestUtilization => utilization_cmp(ea.utilization, eb.utilization)
1298 .then_with(|| pa.cmp(&pb))
1299 .then_with(|| ka.cmp(&kb)),
1300 Strategy::SoonestReset => ea
1301 .tier
1302 .cmp(&eb.tier)
1303 .then_with(|| match ea.tier {
1304 1 | 3 => ea.budget_reset_ms.cmp(&eb.budget_reset_ms),
1306 2 => pa.cmp(&pb).then_with(|| ka.cmp(&kb)),
1308 _ => utilization_cmp(ea.utilization, eb.utilization),
1310 })
1311 .then_with(|| utilization_cmp(ea.utilization, eb.utilization))
1312 .then_with(|| pa.cmp(&pb))
1313 .then_with(|| ka.cmp(&kb)),
1314 })
1315 });
1316
1317 if req.sticky && req.strategy == Strategy::SoonestReset {
1324 if let Some(current) = req.current {
1325 let key = current.storage_key();
1326 let at = eligible
1327 .iter()
1328 .position(|(c, _)| c.credential.storage_key() == key);
1329 if let Some(i) = at.filter(|i| *i > 0) {
1330 let best = &eligible[0].1;
1331 let cur = &eligible[i].1;
1332 let switch = best.tier < cur.tier
1333 || (best.tier == 1
1334 && cur.tier == 1
1335 && best.budget_reset_ms < cur.budget_reset_ms);
1336 if !switch {
1337 eligible[..=i].rotate_right(1);
1338 }
1339 }
1340 }
1341 }
1342
1343 match eligible.first() {
1344 Some((cap, e)) => {
1345 let winner = cap.credential.clone();
1346 for (i, (other, _)) in eligible.iter().enumerate().skip(1) {
1347 rejections.push(Rejection {
1348 credential: other.credential.clone(),
1349 reason: RejectReason::Outranked {
1350 rank: (i + 1) as u32,
1351 },
1352 cause: None,
1353 usable_after_ms: None,
1354 });
1355 }
1356 Selection::Selected {
1357 credential: winner,
1358 utilization: e.utilization,
1359 tier: e.tier,
1360 budget_reset_ms: e.budget_reset_ms,
1361 rejections,
1362 }
1363 }
1364 None => {
1365 let earliest_reset_ms = rejections
1366 .iter()
1367 .filter_map(|r| match &r.reason {
1368 RejectReason::Exhausted {
1369 resets_at_ms: Some(t),
1370 ..
1371 } if *t > req.now_ms => Some(*t),
1372 _ => None,
1373 })
1374 .min();
1375 Selection::NoCapacity {
1376 rejections,
1377 earliest_reset_ms,
1378 }
1379 }
1380 }
1381}
1382
1383#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1387#[serde(rename_all = "snake_case")]
1388pub enum FailureKind {
1389 QuotaExhaustedPreOutput,
1392 RateLimited,
1394 Other,
1396}
1397
1398#[derive(Debug, Clone, PartialEq, Eq)]
1400pub struct FailureEvidence {
1401 pub kind: FailureKind,
1402 pub output_started: bool,
1404 pub tool_activity: bool,
1406 pub failovers_so_far: u32,
1408}
1409
1410pub const MAX_FAILOVERS: u32 = 1;
1412
1413#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1414#[serde(tag = "decision", rename_all = "snake_case")]
1415pub enum FailoverDecision {
1416 Failover {
1417 #[serde(with = "cred_serde")]
1418 credential: CredentialRef,
1419 },
1420 NoFailover {
1421 reason: String,
1422 },
1423}
1424
1425pub fn failover(
1431 req: &SelectionRequest<'_>,
1432 failed: &CredentialRef,
1433 evidence: &FailureEvidence,
1434 candidates: &[AccountCapacity],
1435) -> FailoverDecision {
1436 let no = |reason: &str| FailoverDecision::NoFailover {
1437 reason: reason.to_string(),
1438 };
1439 if evidence.output_started {
1440 return no("output already started; never replay on another account");
1441 }
1442 if evidence.tool_activity {
1443 return no("tool activity occurred; never replay on another account");
1444 }
1445 if evidence.failovers_so_far >= MAX_FAILOVERS {
1446 return no("failover budget exhausted");
1447 }
1448 if evidence.kind != FailureKind::QuotaExhaustedPreOutput {
1449 return no("not a recognized pre-output quota failure");
1450 }
1451 if req.explicit_account.is_some() {
1452 return no("explicit account requested; no silent fallback");
1453 }
1454 let mut exclude: Vec<CredentialRef> = req.exclude.to_vec();
1455 exclude.push(failed.clone());
1456 let retry = SelectionRequest {
1457 exclude: &exclude,
1458 ..req.clone()
1459 };
1460 match select(&retry, candidates) {
1461 Selection::Selected { credential, .. } => FailoverDecision::Failover { credential },
1462 Selection::NoCapacity { .. } => no("no other eligible account"),
1463 }
1464}
1465
1466fn pick_weekly<'w>(windows: impl Iterator<Item = &'w WindowLimit>) -> Option<&'w WindowLimit> {
1470 let mut weekly: Vec<&WindowLimit> = windows.filter(|w| w.is_weekly()).collect();
1471 weekly.sort_by(|a, b| {
1472 a.models
1473 .is_some()
1474 .cmp(&b.models.is_some())
1475 .then_with(|| a.id.cmp(&b.id))
1476 });
1477 weekly.into_iter().next()
1478}
1479
1480pub fn weekly_window(windows: &[WindowLimit]) -> Option<&WindowLimit> {
1486 pick_weekly(windows.iter())
1487}
1488
1489pub fn budget_window<'w>(
1499 windows: &'w [WindowLimit],
1500 model: Option<&str>,
1501) -> Option<&'w WindowLimit> {
1502 let applicable = || windows.iter().filter(|w| w.applies_to(model));
1503 if let Some(weekly) = pick_weekly(applicable()) {
1504 return Some(weekly);
1505 }
1506 applicable()
1507 .filter(|w| w.duration_ms.is_some())
1508 .min_by(|a, b| {
1509 b.duration_ms
1510 .cmp(&a.duration_ms)
1511 .then_with(|| a.models.is_some().cmp(&b.models.is_some()))
1512 .then_with(|| a.id.cmp(&b.id))
1513 })
1514}
1515
1516pub fn is_unanchored(window: &WindowLimit, observed_at_ms: u64) -> Option<bool> {
1531 let (Some(duration), Some(reset)) = (window.duration_ms, window.resets_at_ms) else {
1532 return None;
1533 };
1534 let near_zero = matches!(
1535 window.used_percent,
1536 Some(p) if (0.0..=UNANCHORED_MAX_USED_PERCENT).contains(&p)
1537 );
1538 if !near_zero {
1539 return Some(false);
1540 }
1541 let expected = observed_at_ms.saturating_add(duration);
1542 Some(expected.abs_diff(reset) <= UNANCHORED_TOLERANCE_MS)
1543}
1544
1545#[cfg(test)]
1546mod tests {
1547 use super::*;
1548 use crate::auth::account::Account;
1549
1550 const HOUR: u64 = 3_600_000;
1551 const DAY: u64 = 24 * HOUR;
1552 const NOW: u64 = 1_800_000_000_000;
1553
1554 fn cred(label: &str) -> CredentialRef {
1555 CredentialRef::new(OAuthProviderId::OpenAiCodex, Account::parse(label).unwrap())
1556 }
1557
1558 fn win(id: &str, dur: u64, used: f64, resets: u64) -> WindowLimit {
1559 WindowLimit {
1560 id: id.into(),
1561 duration_ms: Some(dur),
1562 used_percent: Some(used),
1563 limit_reached: Some(used >= 100.0),
1564 resets_at_ms: Some(resets),
1565 models: None,
1566 }
1567 }
1568
1569 fn ok(label: &str, windows: Vec<WindowLimit>) -> AccountCapacity {
1570 AccountCapacity {
1571 credential: cred(label),
1572 observed_at_ms: Some(NOW - 1000),
1573 observation: QuotaObservation::Ok {
1574 windows,
1575 models: None,
1576 },
1577 cooldown_until_ms: None,
1578 failure: None,
1579 }
1580 }
1581
1582 fn req() -> SelectionRequest<'static> {
1583 SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
1584 }
1585
1586 #[test]
1587 fn window_verdicts_fail_closed() {
1588 let mut w = win("primary", 7 * DAY, 20.0, NOW + DAY);
1589 assert_eq!(
1590 w.verdict(100.0),
1591 WindowVerdict::Headroom {
1592 used_percent: Some(20.0)
1593 }
1594 );
1595 w.used_percent = Some(100.0);
1596 w.limit_reached = Some(false);
1597 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1598 w.used_percent = Some(50.0);
1599 w.limit_reached = Some(true);
1600 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1601 w.used_percent = None;
1602 w.limit_reached = None;
1603 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1604 w.limit_reached = Some(false);
1605 assert_eq!(
1606 w.verdict(100.0),
1607 WindowVerdict::Headroom { used_percent: None }
1608 );
1609 w.used_percent = Some(f64::NAN);
1610 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1611 w.used_percent = Some(140.0);
1613 assert!(matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }));
1614 w.used_percent = Some(1000.5);
1615 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1616 let saved = w.limit_reached;
1618 w.limit_reached = Some(true);
1619 for p in [1500.0, f64::NAN, -3.0] {
1620 w.used_percent = Some(p);
1621 assert!(
1622 matches!(w.verdict(100.0), WindowVerdict::Exhausted { .. }),
1623 "{p}"
1624 );
1625 }
1626 w.limit_reached = saved;
1627 w.used_percent = Some(-1.0);
1628 assert_eq!(w.verdict(100.0), WindowVerdict::Unknown);
1629 w.used_percent = Some(1.0);
1631 assert_eq!(w.verdict(f64::NAN), WindowVerdict::Unknown);
1632 assert_eq!(w.verdict(101.0), WindowVerdict::Unknown);
1633 assert_eq!(w.verdict(-0.5), WindowVerdict::Unknown);
1634 assert_eq!(w.verdict(f64::INFINITY), WindowVerdict::Unknown);
1635 let r = SelectionRequest {
1636 exhausted_at_percent: f64::NAN,
1637 ..SelectionRequest::new(OAuthProviderId::OpenAiCodex, NOW, HOUR)
1638 };
1639 assert!(select(
1640 &r,
1641 &[ok("a", vec![win("primary", 7 * DAY, 1.0, NOW + DAY)])]
1642 )
1643 .selected()
1644 .is_none());
1645 }
1646
1647 #[test]
1648 fn weekly_is_from_duration_not_name() {
1649 let primary_5h = win("primary", 5 * HOUR, 0.0, NOW + HOUR);
1650 let primary_week = win("primary", 7 * DAY, 0.0, NOW + DAY);
1651 let no_dur = WindowLimit {
1652 duration_ms: None,
1653 ..win("seven_day", 0, 0.0, NOW)
1654 };
1655 assert!(!primary_5h.is_weekly());
1656 assert!(primary_week.is_weekly());
1657 assert!(!no_dur.is_weekly());
1658 assert_eq!(
1659 weekly_window(&[primary_5h.clone(), primary_week.clone()]).map(|w| w.duration_ms),
1660 Some(Some(7 * DAY))
1661 );
1662 assert!(weekly_window(&[primary_5h, no_dur]).is_none());
1663 let scoped = WindowLimit {
1665 id: "aaa_scoped".into(),
1666 models: Some(vec!["m".into()]),
1667 ..primary_week.clone()
1668 };
1669 assert_eq!(
1670 weekly_window(&[scoped, primary_week]).map(|w| w.id.as_str()),
1671 Some("primary")
1672 );
1673 }
1674
1675 #[test]
1676 fn selects_lowest_utilization_deterministically() {
1677 let a = ok("a", vec![win("primary", 7 * DAY, 60.0, NOW + DAY)]);
1678 let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1679 let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1680 let sel = select(&req(), &[a.clone(), c.clone(), b.clone()]);
1681 assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@b");
1682 let sel2 = select(&req(), &[c, b, a]);
1684 assert_eq!(sel2.selected().unwrap().storage_key(), "openai-codex@b");
1685 }
1686
1687 #[test]
1688 fn preference_breaks_ties_and_orders() {
1689 let b = ok("b", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1690 let c = ok("c", vec![win("primary", 7 * DAY, 10.0, NOW + DAY)]);
1691 let pref = vec!["c".to_string(), "b".to_string()];
1692 let r = SelectionRequest {
1693 preference: &pref,
1694 ..req()
1695 };
1696 assert_eq!(
1697 select(&r, &[b.clone(), c.clone()])
1698 .selected()
1699 .unwrap()
1700 .storage_key(),
1701 "openai-codex@c"
1702 );
1703 let a = ok("a", vec![win("primary", 7 * DAY, 90.0, NOW + DAY)]);
1704 let pref2 = vec!["a".to_string()];
1705 let r2 = SelectionRequest {
1706 preference: &pref2,
1707 strategy: Strategy::PreferenceOrder,
1708 ..req()
1709 };
1710 assert_eq!(
1711 select(&r2, &[b, c, a]).selected().unwrap().storage_key(),
1712 "openai-codex@a"
1713 );
1714 }
1715
1716 #[test]
1720 fn graced_readings_are_bounded_rank_after_fresh_and_keep_a_sticky_seat() {
1721 let grace_ms = super::super::usage_reads::USAGE_READ_GRACE.as_millis() as u64;
1722 let failed = |label: &str| AccountCapacity::unobserved(cred(label));
1723 let good = |label: &str, age: u64, used: f64, reset: u64| AccountCapacity {
1724 observed_at_ms: Some(NOW - age),
1725 ..ok(label, vec![win("primary", 7 * DAY, used, reset)])
1726 };
1727 let fresh_req = || SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 60_000);
1728 let winner = |sel: Selection| match sel {
1729 Selection::Selected { credential, .. } => credential.account.label_str().to_string(),
1730 Selection::NoCapacity { rejections, .. } => panic!("no capacity: {rejections:?}"),
1731 };
1732 let one = |grace: &[AccountCapacity], cap: &AccountCapacity| {
1733 evaluate(&SelectionRequest { grace, ..fresh_req() }, cap)
1734 };
1735
1736 let young = [good("a", grace_ms - 60_000, 20.0, NOW + DAY)];
1737 assert!(one(&young, &failed("a")).unwrap().graced);
1738 let old = [good("a", grace_ms + 60_000, 20.0, NOW + DAY)];
1739 assert_eq!(one(&old, &failed("a")), Err(RejectReason::Unknown), "past the grace age");
1740 let edge = [good("a", grace_ms, 20.0, NOW + DAY)];
1742 assert_eq!(one(&edge, &failed("a")), Err(RejectReason::Unknown), "at the grace age");
1743 let inside = [good("a", grace_ms - 1, 20.0, NOW + DAY)];
1744 assert!(one(&inside, &failed("a")).unwrap().graced, "just under it");
1745 let exhausted = [good("a", 60_000, 100.0, NOW + HOUR)];
1746 assert_eq!(one(&exhausted, &failed("a")), Err(RejectReason::Unknown), "exhausted window");
1747 let cooling = AccountCapacity {
1748 cooldown_until_ms: Some(NOW + HOUR),
1749 ..failed("a")
1750 };
1751 assert!(matches!(one(&young, &cooling), Err(RejectReason::CoolingDown { .. })));
1752 let auth = AccountCapacity {
1753 observation: QuotaObservation::AuthError,
1754 ..failed("a")
1755 };
1756 assert_eq!(one(&young, &auth), Err(RejectReason::AuthError), "only a failed read is graced");
1757
1758 let grace = [good("a", 60_000, 20.0, NOW + HOUR)];
1760 let b = good("b", 1000, 20.0, NOW + 3 * DAY);
1761 let anew = SelectionRequest { grace: &grace, ..fresh_req() };
1762 assert_eq!(winner(select(&anew, &[failed("a"), b.clone()])), "b");
1763 assert_eq!(winner(select(&anew, &[failed("a")])), "a", "alone, the graced seat serves");
1764 let current = cred("a");
1766 let sticky = SelectionRequest {
1767 current: Some(¤t),
1768 ..anew.clone()
1769 };
1770 assert_eq!(winner(select(&sticky, &[failed("a"), b])), "a");
1771
1772 let mut switched = 0;
1777 for a_reset in [NOW + HOUR, NOW + 2 * DAY, NOW + 10 * HOUR] {
1778 for b_reset in [NOW + HOUR / 2, NOW + 3 * DAY, NOW + 12 * HOUR, NOW + 5 * DAY] {
1779 let a_good = good("a", 1000, 20.0, a_reset);
1780 let b = good("b", 1000, 20.0, b_reset);
1781 let graced = [a_good.clone()];
1782 let on_grace = SelectionRequest {
1783 grace: &graced,
1784 current: Some(¤t),
1785 ..fresh_req()
1786 };
1787 let on_fresh = SelectionRequest {
1788 current: Some(¤t),
1789 ..fresh_req()
1790 };
1791 let g = winner(select(&on_grace, &[failed("a"), b.clone()]));
1792 let f = winner(select(&on_fresh, &[a_good, b]));
1793 assert_eq!(g, f, "a resets at {a_reset}, b at {b_reset}");
1794 switched += usize::from(g == "b");
1795 }
1796 }
1797 assert!(switched > 0 && switched < 12, "the cases cover staying and switching: {switched}");
1798 }
1799
1800 #[test]
1801 fn stale_unknown_malformed_auth_are_not_capacity() {
1802 let mut stale = ok("s", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1803 stale.observed_at_ms = Some(NOW - 2 * HOUR);
1804 let mut never = ok("n", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1805 never.observed_at_ms = None;
1806 let mut future = ok("f", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1807 future.observed_at_ms = Some(NOW + HOUR);
1808 let mut auth = ok("x", vec![]);
1809 auth.observation = QuotaObservation::AuthError;
1810 let mut unk = ok("u", vec![]);
1811 unk.observation = QuotaObservation::Unknown;
1812 let mut mal = ok("m", vec![]);
1813 mal.observation = QuotaObservation::Malformed;
1814 let mut uns = ok("v", vec![]);
1815 uns.observation = QuotaObservation::Unsupported;
1816 let none = ok("e", vec![]);
1817 let sel = select(&req(), &[stale, never, future, auth, unk, mal, uns, none]);
1818 let Selection::NoCapacity {
1819 rejections,
1820 earliest_reset_ms,
1821 } = sel
1822 else {
1823 panic!("expected no capacity");
1824 };
1825 assert_eq!(earliest_reset_ms, None);
1826 let reasons: Vec<(String, RejectReason)> = rejections
1827 .into_iter()
1828 .map(|r| (r.credential.account.label_str().to_string(), r.reason))
1829 .collect();
1830 assert!(reasons.contains(&(
1831 "s".into(),
1832 RejectReason::Stale {
1833 age_ms: Some(2 * HOUR)
1834 }
1835 )));
1836 assert!(reasons.contains(&("n".into(), RejectReason::Stale { age_ms: None })));
1837 assert!(reasons.contains(&("f".into(), RejectReason::Malformed)));
1838 assert!(reasons.contains(&("x".into(), RejectReason::AuthError)));
1839 assert!(reasons.contains(&("u".into(), RejectReason::Unknown)));
1840 assert!(reasons.contains(&("m".into(), RejectReason::Malformed)));
1841 assert!(reasons.contains(&("v".into(), RejectReason::Unsupported)));
1842 assert!(reasons.contains(&("e".into(), RejectReason::NoLimitEvidence)));
1843 }
1844
1845 #[test]
1846 fn exhausted_cooldown_and_unknown_window_excluded() {
1847 let exhausted = ok("x", vec![win("primary", 7 * DAY, 100.0, NOW + 2 * DAY)]);
1848 let mut cooling = ok("c", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1849 cooling.cooldown_until_ms = Some(NOW + HOUR);
1850 let mut cooled = ok("d", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1851 cooled.cooldown_until_ms = Some(NOW - 1);
1852 let unknown_window = ok(
1853 "w",
1854 vec![WindowLimit {
1855 used_percent: None,
1856 limit_reached: None,
1857 ..win("primary", 7 * DAY, 0.0, NOW + DAY)
1858 }],
1859 );
1860 let past_exhausted = ok("p", vec![win("primary", 7 * DAY, 100.0, NOW - DAY)]);
1861 let sel = select(
1862 &req(),
1863 &[exhausted, cooling, unknown_window, past_exhausted],
1864 );
1865 let Selection::NoCapacity {
1866 earliest_reset_ms, ..
1867 } = &sel
1868 else {
1869 panic!("no capacity expected");
1870 };
1871 assert_eq!(*earliest_reset_ms, Some(NOW + 2 * DAY));
1873 let sel = select(&req(), &[cooled]);
1874 assert_eq!(sel.selected().unwrap().storage_key(), "openai-codex@d");
1875 }
1876
1877 #[test]
1878 fn all_applicable_limits_including_model_scoped_count() {
1879 let model_win = WindowLimit {
1880 id: "weekly_sonnet".into(),
1881 models: Some(vec!["sonnet".into()]),
1882 ..win("weekly_sonnet", 7 * DAY, 100.0, NOW + DAY)
1883 };
1884 let five_h = win("five_hour", 5 * HOUR, 99.0, NOW + HOUR);
1885 let acct = ok(
1886 "a",
1887 vec![
1888 win("seven_day", 7 * DAY, 10.0, NOW + DAY),
1889 five_h,
1890 model_win,
1891 ],
1892 );
1893 let r = SelectionRequest {
1895 model: Some("sonnet"),
1896 ..req()
1897 };
1898 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
1899 let r = SelectionRequest {
1901 model: Some("opus"),
1902 ..req()
1903 };
1904 match select(&r, std::slice::from_ref(&acct)) {
1905 Selection::Selected { utilization, .. } => assert_eq!(utilization, Some(99.0)),
1906 other => panic!("unexpected {other:?}"),
1907 }
1908 assert!(select(&req(), std::slice::from_ref(&acct))
1910 .selected()
1911 .is_none());
1912 let mut unavailable = acct;
1914 if let QuotaObservation::Ok { models, .. } = &mut unavailable.observation {
1915 *models = Some(vec![ModelAvailability {
1916 model: "opus".into(),
1917 state: ModelState::Exhausted,
1918 }]);
1919 }
1920 let r = SelectionRequest {
1921 model: Some("opus"),
1922 ..req()
1923 };
1924 let sel = select(&r, &[unavailable.clone()]);
1925 match sel {
1926 Selection::NoCapacity { rejections, .. } => assert_eq!(
1927 rejections[0].reason,
1928 RejectReason::ModelUnavailable {
1929 model: "opus".into()
1930 }
1931 ),
1932 other => panic!("unexpected {other:?}"),
1933 }
1934 let mut unknown = unavailable;
1936 if let QuotaObservation::Ok { models, .. } = &mut unknown.observation {
1937 *models = Some(vec![ModelAvailability {
1938 model: "opus".into(),
1939 state: ModelState::Unknown,
1940 }]);
1941 }
1942 match select(&r, &[unknown]) {
1943 Selection::NoCapacity { rejections, .. } => assert_eq!(
1944 rejections[0].reason,
1945 RejectReason::ModelUnknown {
1946 model: "opus".into()
1947 }
1948 ),
1949 other => panic!("unexpected {other:?}"),
1950 }
1951 }
1952
1953 #[test]
1954 fn explicit_account_never_falls_back() {
1955 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
1956 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1957 let r = SelectionRequest {
1958 explicit_account: Some("a"),
1959 ..req()
1960 };
1961 let sel = select(&r, &[a.clone(), b.clone()]);
1962 assert!(sel.selected().is_none());
1963 let r = SelectionRequest {
1964 explicit_account: Some("zzz"),
1965 ..req()
1966 };
1967 match select(&r, &[a.clone(), b.clone()]) {
1968 Selection::NoCapacity { rejections, .. } => {
1969 assert_eq!(rejections.len(), 1);
1970 assert_eq!(rejections[0].reason, RejectReason::NotPresent);
1971 assert_eq!(rejections[0].credential.storage_key(), "openai-codex@zzz");
1972 }
1973 other => panic!("unexpected {other:?}"),
1974 }
1975 let r = SelectionRequest {
1976 explicit_account: Some("Bad Label"),
1977 ..req()
1978 };
1979 match select(&r, &[a, b]) {
1980 Selection::NoCapacity { rejections, .. } => assert_eq!(
1981 rejections[0].reason,
1982 RejectReason::InvalidLabel {
1983 label: "Bad Label".into()
1984 }
1985 ),
1986 other => panic!("unexpected {other:?}"),
1987 }
1988 }
1989
1990 #[test]
1991 fn other_providers_and_exclusions_are_ignored() {
1992 let mut other = ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1993 other.credential =
1994 CredentialRef::new(OAuthProviderId::Anthropic, Account::parse("a").unwrap());
1995 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
1996 let excl = vec![cred("b")];
1997 let r = SelectionRequest {
1998 exclude: &excl,
1999 ..req()
2000 };
2001 let sel = select(&r, &[other, b]);
2002 match sel {
2003 Selection::NoCapacity { rejections, .. } => {
2004 assert_eq!(rejections.len(), 1);
2005 assert_eq!(rejections[0].reason, RejectReason::Excluded);
2006 }
2007 other => panic!("unexpected {other:?}"),
2008 }
2009 }
2010
2011 #[test]
2012 fn failover_is_bounded_and_never_after_output() {
2013 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2014 let b = ok("b", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)]);
2015 let c = ok("c", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2016 let cands = [a, b, c];
2017 let ev = FailureEvidence {
2018 kind: FailureKind::QuotaExhaustedPreOutput,
2019 output_started: false,
2020 tool_activity: false,
2021 failovers_so_far: 0,
2022 };
2023 match failover(&req(), &cred("a"), &ev, &cands) {
2024 FailoverDecision::Failover { credential } => {
2025 assert_eq!(credential.storage_key(), "openai-codex@b")
2026 }
2027 other => panic!("unexpected {other:?}"),
2028 }
2029 let ev2 = FailureEvidence {
2031 failovers_so_far: 1,
2032 ..ev.clone()
2033 };
2034 assert!(matches!(
2035 failover(&req(), &cred("b"), &ev2, &cands),
2036 FailoverDecision::NoFailover { .. }
2037 ));
2038 let ev3 = FailureEvidence {
2040 output_started: true,
2041 ..ev.clone()
2042 };
2043 assert!(matches!(
2044 failover(&req(), &cred("a"), &ev3, &cands),
2045 FailoverDecision::NoFailover { .. }
2046 ));
2047 let ev4 = FailureEvidence {
2049 tool_activity: true,
2050 ..ev.clone()
2051 };
2052 assert!(matches!(
2053 failover(&req(), &cred("a"), &ev4, &cands),
2054 FailoverDecision::NoFailover { .. }
2055 ));
2056 let ev5 = FailureEvidence {
2058 kind: FailureKind::RateLimited,
2059 ..ev.clone()
2060 };
2061 assert!(matches!(
2062 failover(&req(), &cred("a"), &ev5, &cands),
2063 FailoverDecision::NoFailover { .. }
2064 ));
2065 let r = SelectionRequest {
2067 explicit_account: Some("a"),
2068 ..req()
2069 };
2070 assert!(matches!(
2071 failover(&r, &cred("a"), &ev, &cands),
2072 FailoverDecision::NoFailover { .. }
2073 ));
2074 let only = [ok("a", vec![win("primary", 7 * DAY, 0.0, NOW + DAY)])];
2077 assert!(matches!(
2078 failover(&req(), &cred("a"), &ev, &only),
2079 FailoverDecision::NoFailover { .. }
2080 ));
2081 }
2082
2083 #[test]
2084 fn unknown_utilization_sorts_after_known() {
2085 let proven = ok(
2086 "p",
2087 vec![WindowLimit {
2088 used_percent: None,
2089 limit_reached: Some(false),
2090 ..win("primary", 7 * DAY, 0.0, NOW + DAY)
2091 }],
2092 );
2093 let known = ok("k", vec![win("primary", 7 * DAY, 95.0, NOW + DAY)]);
2094 assert_eq!(
2095 select(&req(), &[proven.clone(), known])
2096 .selected()
2097 .unwrap()
2098 .storage_key(),
2099 "openai-codex@k"
2100 );
2101 assert_eq!(
2102 select(&req(), &[proven]).selected().unwrap().storage_key(),
2103 "openai-codex@p"
2104 );
2105 }
2106
2107 #[test]
2108 fn model_matching_uses_family_tokens() {
2109 assert!(model_matches("sonnet", "claude-sonnet-4-6"));
2110 assert!(model_matches("Sonnet", "CLAUDE-SONNET-4-6"));
2111 assert!(model_matches("gpt-6-astra", "gpt-6-astra"));
2112 assert!(!model_matches("opus", "claude-sonnet-4-6"));
2113 assert!(!model_matches("", "x") && !model_matches("x", ""));
2114 let w = WindowLimit {
2115 models: Some(vec!["sonnet".into()]),
2116 ..win("seven_day_sonnet", 7 * DAY, 100.0, NOW + DAY)
2117 };
2118 assert!(w.applies_to(Some("claude-sonnet-4-6")));
2119 assert!(!w.applies_to(Some("claude-opus-4-6")));
2120 assert!(w.applies_to(None));
2121 let acct = ok("a", vec![win("seven_day", 7 * DAY, 1.0, NOW + DAY), w]);
2122 let r = SelectionRequest {
2123 model: Some("claude-sonnet-4-6"),
2124 ..req()
2125 };
2126 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_none());
2127 let r = SelectionRequest {
2128 model: Some("claude-opus-4-6"),
2129 ..req()
2130 };
2131 assert!(select(&r, std::slice::from_ref(&acct)).selected().is_some());
2132 }
2133
2134 #[test]
2135 fn serde_round_trip_is_secret_free() {
2136 let a = ok("a", vec![win("primary", 7 * DAY, 100.0, NOW + DAY)]);
2137 let sel = select(&req(), &[a]);
2138 let json = serde_json::to_string(&sel).unwrap();
2139 assert!(json.contains("no_capacity"));
2140 let back: Selection = serde_json::from_str(&json).unwrap();
2141 assert_eq!(back, sel);
2142 }
2143
2144 fn sr_req() -> SelectionRequest<'static> {
2149 SelectionRequest::soonest_reset(OAuthProviderId::OpenAiCodex, NOW, 30 * 60 * 1000)
2150 }
2151
2152 const OBSERVED: u64 = NOW - 1000;
2154
2155 fn anchored(used: f64, resets: u64) -> WindowLimit {
2158 win("primary", 7 * DAY, used, resets)
2159 }
2160
2161 fn unanchored() -> WindowLimit {
2164 win("primary", 7 * DAY, 0.0, OBSERVED + 7 * DAY)
2165 }
2166
2167 fn no_reset(used: f64) -> WindowLimit {
2169 WindowLimit {
2170 resets_at_ms: None,
2171 ..win("primary", 7 * DAY, used, 0)
2172 }
2173 }
2174
2175 fn label_of(sel: &Selection) -> &str {
2176 sel.selected().expect("a selection").account.label_str()
2177 }
2178
2179 fn ranks(sel: &Selection) -> Vec<(String, u32)> {
2180 let (Selection::Selected { rejections, .. } | Selection::NoCapacity { rejections, .. }) =
2181 sel;
2182 rejections
2183 .iter()
2184 .filter_map(|r| match r.reason {
2185 RejectReason::Outranked { rank } => {
2186 Some((r.credential.account.label_str().to_string(), rank))
2187 }
2188 _ => None,
2189 })
2190 .collect()
2191 }
2192
2193 #[test]
2194 fn soonest_reset_beats_lower_utilization_where_lowest_utilization_does_not() {
2195 let a = ok("a", vec![anchored(60.0, NOW + 2 * DAY)]);
2197 let b = ok("b", vec![anchored(10.0, NOW + 5 * DAY)]);
2198 let cands = [b.clone(), a.clone()];
2199 match select(&sr_req(), &cands) {
2200 Selection::Selected {
2201 credential,
2202 tier,
2203 budget_reset_ms,
2204 utilization,
2205 ..
2206 } => {
2207 assert_eq!(credential.storage_key(), "openai-codex@a");
2208 assert_eq!(tier, 3);
2209 assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
2210 assert_eq!(utilization, Some(60.0));
2211 }
2212 other => panic!("unexpected {other:?}"),
2213 }
2214 assert_eq!(label_of(&select(&req(), &cands)), "b");
2216 let r = req();
2218 assert_eq!(r.strategy, Strategy::LowestUtilization);
2219 assert!(!r.sticky && r.current.is_none());
2220 assert_eq!(r.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
2221 let s = sr_req();
2222 assert_eq!(s.strategy, Strategy::SoonestReset);
2223 assert!(s.sticky && s.current.is_none());
2224 assert_eq!(s.urgent_horizon_ms, DEFAULT_URGENT_HORIZON_MS);
2225 }
2226
2227 #[test]
2228 fn tiers_order_urgent_then_unanchored_then_far_then_no_reset() {
2229 let urgent = ok("u", vec![anchored(50.0, NOW + 6 * HOUR)]);
2230 let sliding = ok("s", vec![unanchored()]);
2231 let far = ok("f", vec![anchored(0.0, NOW + 5 * DAY)]);
2232 let none = ok("n", vec![no_reset(0.0)]);
2233 let r = sr_req();
2234 for (cap, tier, anch) in [
2235 (&urgent, 1u8, Some(true)),
2236 (&sliding, 2, Some(false)),
2237 (&far, 3, Some(true)),
2238 (&none, 4, None),
2239 ] {
2240 let e = evaluate(&r, cap).unwrap();
2241 assert_eq!(e.tier, tier, "{}", cap.credential.storage_key());
2242 assert_eq!(e.anchored, anch, "{}", cap.credential.storage_key());
2243 }
2244 assert_eq!(evaluate(&req(), &sliding).unwrap().tier, 2);
2246 let cands = [none.clone(), far.clone(), sliding.clone(), urgent.clone()];
2248 let sel = select(&r, &cands);
2249 assert_eq!(label_of(&sel), "u");
2250 assert_eq!(
2251 ranks(&sel),
2252 vec![
2253 ("s".to_string(), 2),
2254 ("f".to_string(), 3),
2255 ("n".to_string(), 4)
2256 ]
2257 );
2258 assert_eq!(
2259 label_of(&select(&r, &[none.clone(), far.clone(), sliding])),
2260 "s"
2261 );
2262 assert_eq!(label_of(&select(&r, &[none.clone(), far])), "f");
2263 assert_eq!(label_of(&select(&r, &[none])), "n");
2264 let short = SelectionRequest {
2267 urgent_horizon_ms: HOUR,
2268 ..sr_req()
2269 };
2270 assert_eq!(evaluate(&short, &urgent).unwrap().tier, 3);
2271 assert_eq!(
2272 label_of(&select(&short, &[urgent, ok("s", vec![unanchored()])])),
2273 "s"
2274 );
2275 let edge = ok("e", vec![anchored(0.0, NOW + DEFAULT_URGENT_HORIZON_MS)]);
2277 assert_eq!(evaluate(&sr_req(), &edge).unwrap().tier, 1);
2278 }
2279
2280 #[test]
2281 fn within_tier_orders_by_reset_preference_or_utilization() {
2282 let a = ok("a", vec![anchored(90.0, NOW + 2 * HOUR)]);
2284 let b = ok("b", vec![anchored(1.0, NOW + 3 * HOUR)]);
2285 assert_eq!(label_of(&select(&sr_req(), &[b, a])), "a");
2286 let s1 = ok("s1", vec![unanchored()]);
2288 let s2 = ok("s2", vec![unanchored()]);
2289 assert_eq!(
2290 label_of(&select(&sr_req(), &[s2.clone(), s1.clone()])),
2291 "s1"
2292 );
2293 let pref = vec!["s2".to_string()];
2294 let r = SelectionRequest {
2295 preference: &pref,
2296 ..sr_req()
2297 };
2298 assert_eq!(label_of(&select(&r, &[s1, s2])), "s2");
2299 let n1 = ok("n1", vec![no_reset(40.0)]);
2301 let n2 = ok("n2", vec![no_reset(5.0)]);
2302 let n3 = ok(
2303 "n3",
2304 vec![WindowLimit {
2305 used_percent: None,
2306 limit_reached: Some(false),
2307 ..no_reset(0.0)
2308 }],
2309 );
2310 let sel = select(&sr_req(), &[n3, n1, n2]);
2311 assert_eq!(label_of(&sel), "n2");
2312 assert_eq!(
2313 ranks(&sel),
2314 vec![("n1".to_string(), 2), ("n3".to_string(), 3)]
2315 );
2316 let t1 = ok("t1", vec![anchored(30.0, NOW + 3 * DAY)]);
2318 let t2 = ok("t2", vec![anchored(20.0, NOW + 3 * DAY)]);
2319 assert_eq!(label_of(&select(&sr_req(), &[t1, t2])), "t2");
2320 }
2321
2322 #[test]
2323 fn exhausted_soonest_seat_is_skipped_to_next_soonest() {
2324 let x = ok("x", vec![anchored(100.0, NOW + HOUR)]);
2325 let a = ok("a", vec![anchored(20.0, NOW + 2 * DAY)]);
2326 let b = ok("b", vec![anchored(20.0, NOW + 3 * DAY)]);
2327 let sel = select(&sr_req(), &[b, x.clone(), a]);
2328 assert_eq!(label_of(&sel), "a");
2329 let Selection::Selected { rejections, .. } = &sel else {
2330 panic!("expected selection");
2331 };
2332 assert!(rejections
2333 .iter()
2334 .any(|r| r.credential.account.label_str() == "x"
2335 && r.reason
2336 == RejectReason::Exhausted {
2337 window: "primary".into(),
2338 resets_at_ms: Some(NOW + HOUR),
2339 }));
2340 let y = ok("y", vec![anchored(100.0, NOW + 2 * DAY)]);
2342 let past = ok("p", vec![anchored(100.0, NOW - HOUR)]);
2343 match select(&sr_req(), &[y, past, x]) {
2344 Selection::NoCapacity {
2345 earliest_reset_ms, ..
2346 } => assert_eq!(earliest_reset_ms, Some(NOW + HOUR)),
2347 other => panic!("unexpected {other:?}"),
2348 }
2349 }
2350
2351 #[test]
2352 fn five_hour_throttle_rejects_then_reenters_after_reset() {
2353 let throttled = ok(
2354 "t",
2355 vec![
2356 win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
2357 win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
2358 ],
2359 );
2360 let other = ok("o", vec![win("seven_day", 7 * DAY, 10.0, NOW + 4 * DAY)]);
2361 let sel = select(&sr_req(), &[throttled, other.clone()]);
2362 assert_eq!(label_of(&sel), "o");
2363 let Selection::Selected { rejections, .. } = &sel else {
2364 panic!("expected selection");
2365 };
2366 assert_eq!(
2367 rejections[0].reason,
2368 RejectReason::Exhausted {
2369 window: "five_hour".into(),
2370 resets_at_ms: Some(NOW + HOUR),
2371 }
2372 );
2373 let now2 = NOW + HOUR + 1;
2376 let mut throttled2 = ok(
2377 "t",
2378 vec![
2379 win("five_hour", 5 * HOUR, 0.0, now2 + 5 * HOUR),
2380 win("seven_day", 7 * DAY, 30.0, NOW + 2 * DAY),
2381 ],
2382 );
2383 throttled2.observed_at_ms = Some(now2 - 1000);
2384 let mut other2 = other;
2385 other2.observed_at_ms = Some(now2 - 1000);
2386 let r = SelectionRequest {
2387 now_ms: now2,
2388 ..sr_req()
2389 };
2390 match select(&r, &[other2, throttled2]) {
2391 Selection::Selected {
2392 credential,
2393 tier,
2394 budget_reset_ms,
2395 ..
2396 } => {
2397 assert_eq!(credential.storage_key(), "openai-codex@t");
2398 assert_eq!(tier, 3);
2399 assert_eq!(budget_reset_ms, Some(NOW + 2 * DAY));
2400 }
2401 other => panic!("unexpected {other:?}"),
2402 }
2403 }
2404
2405 #[test]
2406 fn is_unanchored_classifies_sliding_resets_only() {
2407 let obs = NOW - 90_000;
2408 let codex_idle = win("primary", 7 * DAY, 0.0, obs + 7 * DAY);
2410 assert_eq!(is_unanchored(&codex_idle, obs), Some(true));
2411 let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY + 60_000);
2413 assert_eq!(is_unanchored(&skewed, obs), Some(true));
2414 let skewed = win("primary", 7 * DAY, 0.0, obs + 7 * DAY - 60_000);
2415 assert_eq!(is_unanchored(&skewed, obs), Some(true));
2416 let beyond = win(
2418 "primary",
2419 7 * DAY,
2420 0.0,
2421 obs + 7 * DAY + UNANCHORED_TOLERANCE_MS + 1,
2422 );
2423 assert_eq!(is_unanchored(&beyond, obs), Some(false));
2424 let anthropic = win("seven_day", 7 * DAY, 0.0, obs + 5 * DAY);
2426 assert_eq!(is_unanchored(&anthropic, obs), Some(false));
2427 let used = win("primary", 7 * DAY, 40.0, obs + 7 * DAY);
2429 assert_eq!(is_unanchored(&used, obs), Some(false));
2430 let half = win(
2432 "primary",
2433 7 * DAY,
2434 UNANCHORED_MAX_USED_PERCENT,
2435 obs + 7 * DAY,
2436 );
2437 assert_eq!(is_unanchored(&half, obs), Some(true));
2438 let unreported = WindowLimit {
2440 used_percent: None,
2441 limit_reached: Some(false),
2442 ..codex_idle.clone()
2443 };
2444 assert_eq!(is_unanchored(&unreported, obs), Some(false));
2445 let no_dur = WindowLimit {
2447 duration_ms: None,
2448 ..codex_idle.clone()
2449 };
2450 assert_eq!(is_unanchored(&no_dur, obs), None);
2451 let no_reset = WindowLimit {
2452 resets_at_ms: None,
2453 ..codex_idle
2454 };
2455 assert_eq!(is_unanchored(&no_reset, obs), None);
2456 }
2457
2458 #[test]
2459 fn budget_window_prefers_weekly_then_longest() {
2460 let five_h = win("five_hour", 5 * HOUR, 0.0, NOW + HOUR);
2461 let one_h = win("one_hour", HOUR, 0.0, NOW + HOUR);
2462 let thirty_d = win("thirty_day", 30 * DAY, 0.0, NOW + 20 * DAY);
2463 let weekly_wide = win("weekly", 7 * DAY, 0.0, NOW + 3 * DAY);
2464 let weekly_scoped = WindowLimit {
2465 id: "aaa_weekly_scoped".into(),
2466 models: Some(vec!["m".into()]),
2467 ..weekly_wide.clone()
2468 };
2469 let all = [
2471 thirty_d.clone(),
2472 weekly_scoped.clone(),
2473 five_h.clone(),
2474 weekly_wide.clone(),
2475 ];
2476 assert_eq!(
2477 budget_window(&all, Some("m")).map(|w| w.id.as_str()),
2478 Some("weekly")
2479 );
2480 assert_eq!(
2481 budget_window(&all, None).map(|w| w.id.as_str()),
2482 Some("weekly")
2483 );
2484 let no_wide = [thirty_d.clone(), weekly_scoped.clone(), five_h.clone()];
2485 assert_eq!(
2486 budget_window(&no_wide, Some("m")).map(|w| w.id.as_str()),
2487 Some("aaa_weekly_scoped")
2488 );
2489 assert_eq!(
2491 budget_window(&no_wide, Some("other")).map(|w| w.id.as_str()),
2492 Some("thirty_day")
2493 );
2494 assert_eq!(
2496 budget_window(&[five_h.clone(), one_h.clone()], None).map(|w| w.id.as_str()),
2497 Some("five_hour")
2498 );
2499 let b5 = WindowLimit {
2501 id: "b5".into(),
2502 ..five_h.clone()
2503 };
2504 let a5_scoped = WindowLimit {
2505 id: "a5".into(),
2506 models: Some(vec!["m".into()]),
2507 ..five_h.clone()
2508 };
2509 assert_eq!(
2510 budget_window(&[b5.clone(), a5_scoped.clone()], Some("m")).map(|w| w.id.as_str()),
2511 Some("b5")
2512 );
2513 let a5 = WindowLimit {
2514 id: "a5".into(),
2515 ..five_h.clone()
2516 };
2517 assert_eq!(
2518 budget_window(&[b5, a5], None).map(|w| w.id.as_str()),
2519 Some("a5")
2520 );
2521 let no_dur = WindowLimit {
2523 duration_ms: None,
2524 ..five_h
2525 };
2526 assert!(budget_window(&[no_dur.clone(), no_dur], None).is_none());
2527 assert!(budget_window(&[], None).is_none());
2528 assert_eq!(weekly_window(&all).map(|w| w.id.as_str()), Some("weekly"));
2530 }
2531
2532 #[test]
2533 fn past_reset_with_headroom_is_not_reset_evidence() {
2534 let stale_reset = ok("p", vec![anchored(20.0, NOW - 1)]);
2537 let e = evaluate(&sr_req(), &stale_reset).unwrap();
2538 assert_eq!(e.budget_reset_ms, None);
2539 assert_eq!(e.anchored, None);
2540 assert_eq!(e.tier, 4);
2541 let far = ok("f", vec![anchored(80.0, NOW + 6 * DAY)]);
2542 assert_eq!(label_of(&select(&sr_req(), &[stale_reset, far])), "f");
2543 }
2544
2545 #[test]
2546 fn sticky_keeps_current_within_tier_3() {
2547 let cur = cred("cur");
2548 let current = ok("cur", vec![anchored(30.0, NOW + 5 * DAY)]);
2549 let sooner = ok("soon", vec![anchored(10.0, NOW + 2 * DAY)]);
2550 let r = SelectionRequest {
2551 current: Some(&cur),
2552 ..sr_req()
2553 };
2554 let sel = select(&r, &[sooner.clone(), current.clone()]);
2555 assert_eq!(label_of(&sel), "cur");
2556 assert_eq!(ranks(&sel), vec![("soon".to_string(), 2)]);
2557 assert_eq!(
2559 label_of(&select(&sr_req(), &[sooner.clone(), current.clone()])),
2560 "soon"
2561 );
2562 let r = SelectionRequest { sticky: false, ..r };
2564 assert_eq!(label_of(&select(&r, &[sooner, current])), "soon");
2565 }
2566
2567 #[test]
2568 fn sticky_yields_to_strictly_higher_tier() {
2569 let cur = cred("cur");
2570 let current = ok("cur", vec![anchored(10.0, NOW + 5 * DAY)]);
2571 let urgent = ok("urg", vec![anchored(90.0, NOW + 6 * HOUR)]);
2572 let r = SelectionRequest {
2573 current: Some(&cur),
2574 ..sr_req()
2575 };
2576 assert_eq!(label_of(&select(&r, &[current.clone(), urgent])), "urg");
2577 let sliding = ok("sld", vec![unanchored()]);
2579 assert_eq!(label_of(&select(&r, &[current, sliding])), "sld");
2580 let cur_sliding = ok("cur", vec![unanchored()]);
2583 let a_sliding = ok("a", vec![unanchored()]);
2584 let far = ok("far", vec![anchored(0.0, NOW + 2 * DAY)]);
2585 assert_eq!(label_of(&select(&r, &[a_sliding, far, cur_sliding])), "cur");
2586 }
2587
2588 #[test]
2589 fn sticky_tier1_yields_only_to_strictly_sooner_reset() {
2590 let cur = cred("cur");
2591 let current = ok("cur", vec![anchored(50.0, NOW + 10 * HOUR)]);
2592 let r = SelectionRequest {
2593 current: Some(&cur),
2594 ..sr_req()
2595 };
2596 let sooner = ok("soon", vec![anchored(80.0, NOW + 6 * HOUR)]);
2597 assert_eq!(label_of(&select(&r, &[current.clone(), sooner])), "soon");
2598 let equal = ok("eq", vec![anchored(1.0, NOW + 10 * HOUR)]);
2601 let sel = select(&r, &[equal.clone(), current.clone()]);
2602 assert_eq!(label_of(&sel), "cur");
2603 assert_eq!(ranks(&sel), vec![("eq".to_string(), 2)]);
2604 let later = ok("late", vec![anchored(0.0, NOW + 12 * HOUR)]);
2606 assert_eq!(label_of(&select(&r, &[later, current])), "cur");
2607 }
2608
2609 #[test]
2610 fn sticky_never_keeps_ineligible_current_and_only_applies_to_soonest_reset() {
2611 let cur = cred("cur");
2612 let other = ok("o", vec![anchored(50.0, NOW + 5 * DAY)]);
2613 let r = SelectionRequest {
2614 current: Some(&cur),
2615 ..sr_req()
2616 };
2617 let exhausted = ok("cur", vec![anchored(100.0, NOW + DAY)]);
2619 assert_eq!(label_of(&select(&r, &[exhausted, other.clone()])), "o");
2620 let mut cooling = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2622 cooling.cooldown_until_ms = Some(NOW + HOUR);
2623 assert_eq!(label_of(&select(&r, &[cooling, other.clone()])), "o");
2624 let mut stale = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2626 stale.observed_at_ms = Some(NOW - 2 * HOUR);
2627 assert_eq!(label_of(&select(&r, &[stale, other.clone()])), "o");
2628 let excl = vec![cur.clone()];
2630 let r_excl = SelectionRequest {
2631 exclude: &excl,
2632 ..r.clone()
2633 };
2634 let healthy = ok("cur", vec![anchored(0.0, NOW + DAY)]);
2635 assert_eq!(
2636 label_of(&select(&r_excl, &[healthy.clone(), other.clone()])),
2637 "o"
2638 );
2639 assert_eq!(label_of(&select(&r, std::slice::from_ref(&other))), "o");
2641 let busy_current = ok("cur", vec![anchored(90.0, NOW + 5 * DAY)]);
2643 let r_lu = SelectionRequest {
2644 strategy: Strategy::LowestUtilization,
2645 sticky: true,
2646 current: Some(&cur),
2647 ..req()
2648 };
2649 assert_eq!(
2650 label_of(&select(&r_lu, &[busy_current.clone(), other.clone()])),
2651 "o"
2652 );
2653 let pref = vec!["o".to_string()];
2654 let r_po = SelectionRequest {
2655 strategy: Strategy::PreferenceOrder,
2656 sticky: true,
2657 current: Some(&cur),
2658 preference: &pref,
2659 ..req()
2660 };
2661 assert_eq!(label_of(&select(&r_po, &[busy_current, other])), "o");
2662 }
2663
2664 #[test]
2665 fn failover_under_soonest_reset_never_repicks_sticky_failed_seat() {
2666 let cur = cred("a");
2667 let a = ok("a", vec![anchored(0.0, NOW + 2 * DAY)]);
2669 let b = ok("b", vec![anchored(50.0, NOW + 5 * DAY)]);
2670 let r = SelectionRequest {
2671 current: Some(&cur),
2672 ..sr_req()
2673 };
2674 assert_eq!(label_of(&select(&r, &[a.clone(), b.clone()])), "a");
2675 let ev = FailureEvidence {
2676 kind: FailureKind::QuotaExhaustedPreOutput,
2677 output_started: false,
2678 tool_activity: false,
2679 failovers_so_far: 0,
2680 };
2681 match failover(&r, &cur, &ev, &[a.clone(), b]) {
2682 FailoverDecision::Failover { credential } => {
2683 assert_eq!(credential.storage_key(), "openai-codex@b")
2684 }
2685 other => panic!("unexpected {other:?}"),
2686 }
2687 assert!(matches!(
2688 failover(&r, &cur, &ev, &[a]),
2689 FailoverDecision::NoFailover { .. }
2690 ));
2691 }
2692
2693 #[test]
2694 fn explicit_account_under_soonest_reset_never_falls_back() {
2695 let a = ok("a", vec![anchored(100.0, NOW + HOUR)]);
2696 let b = ok("b", vec![anchored(0.0, NOW + 6 * HOUR)]);
2697 let cur = cred("b");
2698 let r = SelectionRequest {
2700 explicit_account: Some("a"),
2701 current: Some(&cur),
2702 ..sr_req()
2703 };
2704 match select(&r, &[a.clone(), b.clone()]) {
2705 Selection::NoCapacity { rejections, .. } => {
2706 assert_eq!(rejections.len(), 1);
2707 assert_eq!(rejections[0].credential.storage_key(), "openai-codex@a");
2708 }
2709 other => panic!("unexpected {other:?}"),
2710 }
2711 let r = SelectionRequest {
2713 explicit_account: Some("b"),
2714 ..sr_req()
2715 };
2716 match select(&r, &[a, b]) {
2717 Selection::Selected {
2718 credential,
2719 rejections,
2720 tier,
2721 ..
2722 } => {
2723 assert_eq!(credential.storage_key(), "openai-codex@b");
2724 assert!(rejections.is_empty());
2725 assert_eq!(tier, 1);
2726 }
2727 other => panic!("unexpected {other:?}"),
2728 }
2729 }
2730
2731 fn permutations<T: Clone>(items: &[T]) -> Vec<Vec<T>> {
2733 fn go<T: Clone>(k: usize, a: &mut Vec<T>, out: &mut Vec<Vec<T>>) {
2734 if k <= 1 {
2735 out.push(a.clone());
2736 return;
2737 }
2738 go(k - 1, a, out);
2739 for i in 0..k - 1 {
2740 if k.is_multiple_of(2) {
2741 a.swap(i, k - 1);
2742 } else {
2743 a.swap(0, k - 1);
2744 }
2745 go(k - 1, a, out);
2746 }
2747 }
2748 let mut a = items.to_vec();
2749 let mut out = Vec::new();
2750 go(a.len(), &mut a, &mut out);
2751 out
2752 }
2753
2754 #[test]
2755 fn soonest_reset_is_deterministic_under_candidate_shuffle() {
2756 let cur = cred("t3b");
2757 let seats = vec![
2758 ok("t1", vec![anchored(70.0, NOW + 3 * HOUR)]),
2759 ok("t2", vec![unanchored()]),
2760 ok("t3a", vec![anchored(5.0, NOW + 2 * DAY)]),
2761 ok("t3b", vec![anchored(40.0, NOW + 4 * DAY)]),
2762 ok("x", vec![anchored(100.0, NOW + DAY)]),
2763 ];
2764 let r = SelectionRequest {
2765 current: Some(&cur),
2766 ..sr_req()
2767 };
2768 let perms = permutations(&seats);
2769 assert_eq!(perms.len(), 120);
2770 let baseline = select(&r, &seats);
2771 assert_eq!(label_of(&baseline), "t1");
2772 let normalize = |sel: &Selection| {
2773 let Selection::Selected {
2774 credential,
2775 tier,
2776 budget_reset_ms,
2777 utilization,
2778 rejections,
2779 } = sel
2780 else {
2781 panic!("expected selection");
2782 };
2783 let mut rej: Vec<(String, RejectReason)> = rejections
2784 .iter()
2785 .map(|r| (r.credential.storage_key(), r.reason.clone()))
2786 .collect();
2787 rej.sort_by(|a, b| a.0.cmp(&b.0));
2788 (
2789 credential.storage_key(),
2790 *tier,
2791 *budget_reset_ms,
2792 *utilization,
2793 rej,
2794 )
2795 };
2796 let expected = normalize(&baseline);
2797 assert_eq!(
2798 ranks(&baseline),
2799 vec![
2800 ("t2".to_string(), 2),
2801 ("t3a".to_string(), 3),
2802 ("t3b".to_string(), 4)
2803 ]
2804 );
2805 for p in perms {
2806 assert_eq!(normalize(&select(&r, &p)), expected);
2807 }
2808 }
2809
2810 #[test]
2811 fn outranked_ranks_are_two_to_n_in_ranking_order() {
2812 let seats = [
2813 ok("c", vec![anchored(0.0, NOW + 4 * DAY)]),
2814 ok("a", vec![anchored(0.0, NOW + 2 * DAY)]),
2815 ok("d", vec![anchored(0.0, NOW + 5 * DAY)]),
2816 ok("b", vec![anchored(0.0, NOW + 3 * DAY)]),
2817 ];
2818 let sel = select(&sr_req(), &seats);
2819 assert_eq!(label_of(&sel), "a");
2820 assert_eq!(
2821 ranks(&sel),
2822 vec![
2823 ("b".to_string(), 2),
2824 ("c".to_string(), 3),
2825 ("d".to_string(), 4)
2826 ]
2827 );
2828 let lu = [
2830 ok("hi", vec![anchored(80.0, NOW + DAY)]),
2831 ok("lo", vec![anchored(10.0, NOW + DAY)]),
2832 ok("mid", vec![anchored(50.0, NOW + DAY)]),
2833 ];
2834 let sel = select(&req(), &lu);
2835 assert_eq!(label_of(&sel), "lo");
2836 assert_eq!(
2837 ranks(&sel),
2838 vec![("mid".to_string(), 2), ("hi".to_string(), 3)]
2839 );
2840 assert_eq!(
2841 RejectReason::Outranked { rank: 3 }.to_string(),
2842 "eligible; outranked (rank 3)"
2843 );
2844 }
2845
2846 #[test]
2847 fn soonest_reset_serde_is_backward_compatible() {
2848 assert_eq!(
2849 serde_json::to_string(&Strategy::SoonestReset).unwrap(),
2850 "\"soonest_reset\""
2851 );
2852 assert_eq!(
2853 serde_json::from_str::<Strategy>("\"soonest_reset\"").unwrap(),
2854 Strategy::SoonestReset
2855 );
2856 let a = ok("a", vec![anchored(10.0, NOW + 2 * HOUR)]);
2857 let b = ok("b", vec![anchored(10.0, NOW + 3 * HOUR)]);
2858 let sel = select(&sr_req(), &[a, b]);
2859 let json = serde_json::to_string(&sel).unwrap();
2860 assert!(json.contains("\"tier\":1"));
2861 assert!(json.contains("\"budget_reset_ms\""));
2862 assert!(json.contains("\"rank\":2"));
2863 assert_eq!(serde_json::from_str::<Selection>(&json).unwrap(), sel);
2864 let old = r#"{"outcome":"selected","credential":"openai-codex@a","utilization":10.0,"rejections":[{"credential":"openai-codex@b","reason":{"reason":"outranked"}}]}"#;
2866 match serde_json::from_str::<Selection>(old).unwrap() {
2867 Selection::Selected {
2868 tier,
2869 budget_reset_ms,
2870 rejections,
2871 ..
2872 } => {
2873 assert_eq!(tier, 0);
2874 assert_eq!(budget_reset_ms, None);
2875 assert_eq!(rejections[0].reason, RejectReason::Outranked { rank: 0 });
2876 }
2877 other => panic!("unexpected {other:?}"),
2878 }
2879 }
2880
2881 fn failed(label: &str, failure: ObservationFailure) -> AccountCapacity {
2884 AccountCapacity::failed(cred(label), failure)
2885 }
2886
2887 fn failure_of(candidates: &[AccountCapacity]) -> SelectionFailure {
2888 match select(&req(), candidates) {
2889 Selection::NoCapacity { rejections, .. } => {
2890 SelectionFailure::from_rejections(&rejections, NOW)
2891 }
2892 other => panic!("expected no capacity, got {other:?}"),
2893 }
2894 }
2895
2896 #[test]
2897 fn observation_failure_maps_to_a_never_capacity_observation() {
2898 use AccountFailureKind as K;
2899 for (kind, obs) in [
2900 (K::UsageRateLimited, QuotaObservation::Unknown),
2901 (K::RefreshTransient, QuotaObservation::Unknown),
2902 (K::UsageUnauthorized, QuotaObservation::AuthError),
2903 (K::RefreshRejected, QuotaObservation::AuthError),
2904 (K::CredentialMissing, QuotaObservation::AuthError),
2905 (K::UsageMalformed, QuotaObservation::Malformed),
2906 (K::Unsupported, QuotaObservation::Unsupported),
2907 (K::RefreshActionRequired, QuotaObservation::Unknown),
2908 ] {
2909 assert_eq!(ObservationFailure::new(kind).observation(), obs, "{kind}");
2910 }
2911 let f = ObservationFailure::new(K::UsageRateLimited)
2912 .with_status(429)
2913 .with_retry_after(Some(120));
2914 assert_eq!(
2915 f.to_string(),
2916 "usage_rate_limited, HTTP 429, retry after 120s"
2917 );
2918 assert_eq!(
2919 serde_json::to_value(&f).unwrap(),
2920 serde_json::json!({"kind": "usage_rate_limited", "http_status": 429, "retry_after_s": 120})
2921 );
2922 let old: Rejection = serde_json::from_str(
2924 r#"{"credential":"openai-codex@a","reason":{"reason":"unknown"}}"#,
2925 )
2926 .unwrap();
2927 assert_eq!(old.cause, None);
2928 assert_eq!(old.describe(), "quota unknown");
2929 }
2930
2931 #[test]
2932 fn rejection_keeps_the_observation_cause_beside_a_cooldown() {
2933 use AccountFailureKind as K;
2934 let limited = ObservationFailure::new(K::UsageRateLimited).with_status(429);
2935 let mut cooling = failed("b", limited.clone());
2936 cooling.cooldown_until_ms = Some(NOW + HOUR);
2937 match select(&req(), &[failed("a", limited.clone()), cooling]) {
2938 Selection::NoCapacity { rejections, .. } => {
2939 assert_eq!(rejections[0].reason, RejectReason::Unknown);
2940 assert_eq!(rejections[0].cause.as_ref(), Some(&limited));
2941 assert_eq!(
2942 rejections[0].describe(),
2943 "quota unknown (usage_rate_limited, HTTP 429)"
2944 );
2945 assert!(matches!(
2948 rejections[1].reason,
2949 RejectReason::CoolingDown { .. }
2950 ));
2951 assert_eq!(rejections[1].cause.as_ref(), Some(&limited));
2952 }
2953 other => panic!("unexpected {other:?}"),
2954 }
2955 }
2956
2957 #[test]
2961 fn a_cooling_accounts_failure_keeps_its_read_kind_and_counts_as_limited() {
2962 use AccountFailureKind as K;
2963 use SelectionFailureKind as S;
2964 let mut refused = failed(
2965 "a",
2966 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
2967 );
2968 refused.cooldown_until_ms = Some(NOW + 2 * HOUR);
2969 let mut throttled = failed(
2970 "b",
2971 ObservationFailure::new(K::UsageRateLimited)
2972 .with_status(429)
2973 .with_retry_after(Some(90)),
2974 );
2975 throttled.cooldown_until_ms = Some(NOW + HOUR);
2976 let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
2977 plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
2978
2979 let f = failure_of(&[refused, throttled, plain]);
2980 assert_eq!(f.kind, S::AllCoolingDown, "{f:?}");
2981 let a = &f.accounts[0];
2982 assert_eq!(
2983 (a.kind, a.http_status, a.cooling_down, a.until_ms),
2984 (K::UsageUnauthorized, Some(401), true, Some(NOW + 2 * HOUR))
2985 );
2986 let b = &f.accounts[1];
2987 assert_eq!(
2988 (b.kind, b.http_status, b.retry_after_s, b.cooling_down),
2989 (K::UsageRateLimited, Some(429), Some(90), true)
2990 );
2991 let c = &f.accounts[2];
2992 assert_eq!((c.kind, c.cooling_down), (K::CoolingDown, true));
2993 assert_eq!(f.login_labels(), vec!["a"]);
2994 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
2995 assert_eq!(
2996 f.accounts_field(),
2997 "a:usage_unauthorized:401:cooling,b:usage_rate_limited:429:cooling,c:cooling_down"
2998 );
2999 let v = serde_json::to_value(&f).unwrap();
3001 assert_eq!(v["accounts"][0]["cooling_down"], true);
3002 let g = failure_of(&[failed("d", ObservationFailure::new(K::UsageTimeout))]);
3003 assert!(serde_json::to_value(&g).unwrap()["accounts"][0]
3004 .get("cooling_down")
3005 .is_none());
3006 }
3007
3008 #[test]
3015 fn the_headline_skips_a_cooling_accounts_read_failure_which_stays_per_account() {
3016 use AccountFailureKind as K;
3017 use SelectionFailureKind as S;
3018 let throttled = || {
3019 let mut c = failed(
3020 "claude5",
3021 ObservationFailure::new(K::UsageRateLimited)
3022 .with_status(429)
3023 .with_retry_after(Some(90)),
3024 );
3025 c.cooldown_until_ms = Some(NOW + HOUR);
3026 c
3027 };
3028 let c5_kept = |f: &SelectionFailure| {
3029 let c5 = &f.accounts[0];
3030 assert_eq!(
3031 (
3032 c5.account.as_str(),
3033 c5.kind,
3034 c5.http_status,
3035 c5.retry_after_s,
3036 c5.cooling_down,
3037 c5.until_ms
3038 ),
3039 (
3040 "claude5",
3041 K::UsageRateLimited,
3042 Some(429),
3043 Some(90),
3044 true,
3045 Some(NOW + HOUR)
3046 ),
3047 "{f:?}"
3048 );
3049 };
3050
3051 let mut plain = ok("claude6", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3052 plain.cooldown_until_ms = Some(NOW + 2 * HOUR);
3053 let a = failure_of(&[throttled(), plain]);
3054 assert_eq!(a.kind, S::AllCoolingDown, "{a:?}");
3055 c5_kept(&a);
3056 assert_eq!(
3057 a.accounts_field(),
3058 "claude5:usage_rate_limited:429:cooling,claude6:cooling_down"
3059 );
3060 assert_eq!(a.earliest_reset_ms, Some(NOW + HOUR));
3061 assert!(a.login_labels().is_empty(), "{a:?}");
3062
3063 let refused = failed(
3064 "claude2",
3065 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
3066 );
3067 let b = failure_of(&[throttled(), refused]);
3068 assert_eq!(b.kind, S::CredentialRejected, "{b:?}");
3069 c5_kept(&b);
3070 assert_eq!(
3071 b.accounts_field(),
3072 "claude5:usage_rate_limited:429:cooling,claude2:usage_unauthorized:401"
3073 );
3074 assert_eq!(b.earliest_reset_ms, Some(NOW + HOUR));
3075 assert_eq!(b.login_labels(), vec!["claude2"]);
3076 }
3077
3078 #[test]
3082 fn an_accounts_until_is_its_latest_exhausted_reset_and_cooldown() {
3083 let later = ok(
3085 "a",
3086 vec![
3087 win("five_hour", 5 * HOUR, 100.0, NOW + HOUR),
3088 win("seven_day", 7 * DAY, 100.0, NOW + 3 * DAY),
3089 ],
3090 );
3091 let earlier = ok(
3093 "b",
3094 vec![
3095 win("seven_day", 7 * DAY, 100.0, NOW + 2 * DAY),
3096 win("five_hour", 5 * HOUR, 100.0, NOW + 2 * HOUR),
3097 ],
3098 );
3099 let mut cooled = ok("c", vec![win("five_hour", 5 * HOUR, 100.0, NOW + HOUR)]);
3101 cooled.cooldown_until_ms = Some(NOW + 4 * HOUR);
3102 let mut unknown = ok("d", vec![win("seven_day", 7 * DAY, 100.0, NOW + DAY)]);
3104 if let QuotaObservation::Ok { windows, .. } = &mut unknown.observation {
3105 windows.push(WindowLimit {
3106 resets_at_ms: None,
3107 ..win("five_hour", 5 * HOUR, 100.0, 0)
3108 });
3109 }
3110
3111 let until = |f: &SelectionFailure, label: &str| {
3112 f.accounts
3113 .iter()
3114 .find(|a| a.account == label)
3115 .unwrap()
3116 .until_ms
3117 };
3118 let f = failure_of(&[later.clone(), earlier.clone(), cooled.clone(), unknown]);
3119 assert_eq!(until(&f, "a"), Some(NOW + 3 * DAY));
3120 assert_eq!(until(&f, "b"), Some(NOW + 2 * DAY));
3121 assert_eq!(until(&f, "c"), Some(NOW + 4 * HOUR));
3122 assert_eq!(until(&f, "d"), None);
3123 assert_eq!(f.earliest_reset_ms, Some(NOW + 4 * HOUR));
3124 assert_eq!(
3125 failure_of(&[later, earlier]).earliest_reset_ms,
3126 Some(NOW + 2 * DAY)
3127 );
3128 match select(&req(), &[cooled]) {
3131 Selection::NoCapacity { rejections, .. } => {
3132 assert!(matches!(
3133 rejections[0].reason,
3134 RejectReason::CoolingDown { .. }
3135 ));
3136 }
3137 other => panic!("unexpected {other:?}"),
3138 }
3139 }
3140
3141 #[test]
3146 fn the_earliest_reset_skips_accounts_that_need_a_login_or_an_action() {
3147 use AccountFailureKind as K;
3148 let mut refused = failed(
3149 "a",
3150 ObservationFailure::new(K::UsageUnauthorized).with_status(401),
3151 );
3152 refused.cooldown_until_ms = Some(NOW + HOUR);
3153 let mut broken = failed("b", ObservationFailure::new(K::CredentialStore));
3154 broken.cooldown_until_ms = Some(NOW + 2 * HOUR);
3155 let mut plain = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3156 plain.cooldown_until_ms = Some(NOW + 3 * HOUR);
3157
3158 let f = failure_of(&[refused.clone(), broken.clone(), plain]);
3159 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3160 assert_eq!(f.earliest_reset_ms, Some(NOW + 3 * HOUR));
3161 assert_eq!(f.accounts[0].until_ms, Some(NOW + HOUR));
3162 assert_eq!(f.accounts[1].until_ms, Some(NOW + 2 * HOUR));
3163 assert_eq!(failure_of(&[refused, broken]).earliest_reset_ms, None);
3164 }
3165
3166 #[test]
3169 fn action_required_and_store_failures_get_their_own_kind() {
3170 use AccountFailureKind as K;
3171 use SelectionFailureKind as S;
3172 let action = ObservationFailure::new(K::RefreshActionRequired);
3173 let store = ObservationFailure::new(K::CredentialStore);
3174 let refused = ObservationFailure::new(K::UsageUnauthorized).with_status(401);
3175 let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3176 cooling.cooldown_until_ms = Some(NOW + HOUR);
3177
3178 for f in [
3179 failure_of(&[failed("a", action.clone())]),
3180 failure_of(&[failed("a", store.clone())]),
3181 failure_of(&[
3182 failed("a", action.clone()),
3183 failed("b", store.clone()),
3184 cooling.clone(),
3185 ]),
3186 failure_of(&[failed("a", action.clone()), failed("b", refused.clone())]),
3187 ] {
3188 assert_eq!(f.kind, S::ActionRequired, "{f:?}");
3189 }
3190 assert!(K::RefreshActionRequired.needs_action());
3191 assert!(K::CredentialStore.needs_action());
3192 assert!(
3193 !K::RefreshActionRequired.needs_login() && !K::RefreshActionRequired.is_transient()
3194 );
3195 let f = failure_of(&[failed("a", action.clone()), failed("b", store.clone())]);
3196 assert_eq!(f.action_labels(), vec!["a", "b"]);
3197 assert!(f.login_labels().is_empty());
3198 let f = failure_of(&[
3200 failed("a", action),
3201 failed("b", ObservationFailure::new(K::UsageTimeout)),
3202 ]);
3203 assert_eq!(f.kind, S::UsageReadFailed);
3204 assert_eq!(f.action_labels(), vec!["a"]);
3205 assert_eq!(S::ActionRequired.as_str(), "action_required");
3206 assert_eq!(S::CredentialStore.as_str(), "credential_store");
3207 }
3208
3209 #[test]
3210 fn selection_failure_kind_priority_and_login_labels() {
3211 use AccountFailureKind as K;
3212 use SelectionFailureKind as S;
3213 let transient = ObservationFailure::new(K::UsageUpstream).with_status(503);
3214 let rejected = ObservationFailure::new(K::RefreshRejected);
3215 let exhausted = ok("x", vec![win("primary", 5 * HOUR, 100.0, NOW + 2 * HOUR)]);
3216 let mut cooling = ok("c", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3217 cooling.cooldown_until_ms = Some(NOW + HOUR);
3218
3219 assert_eq!(failure_of(&[]).kind, S::NoCredentials);
3220
3221 let f = failure_of(&[
3223 failed("a", rejected.clone()),
3224 failed("b", transient.clone()),
3225 ]);
3226 assert_eq!(f.kind, S::UsageReadFailed);
3227 assert_eq!(f.login_labels(), vec!["a"]);
3228 assert_eq!(
3229 f.accounts_field(),
3230 "a:refresh_rejected,b:usage_upstream:503"
3231 );
3232
3233 let f = failure_of(&[failed("a", rejected.clone()), cooling.clone()]);
3235 assert_eq!(f.kind, S::CredentialRejected);
3236 assert_eq!(f.login_labels(), vec!["a"]);
3237 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3238
3239 let f = failure_of(&[exhausted.clone(), cooling.clone()]);
3242 assert_eq!(f.kind, S::AllCoolingDown);
3243 assert_eq!(f.accounts_field(), "x:exhausted,c:cooling_down");
3244 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3245 assert!(f.login_labels().is_empty());
3246
3247 let no_evidence = ok("n", vec![]);
3250 let f = failure_of(&[failed("a", rejected), no_evidence]);
3251 assert_eq!(f.kind, S::QuotaUnknown);
3252 assert_eq!(f.login_labels(), vec!["a"]);
3253
3254 for kind in [
3256 K::RefreshLockBusy,
3257 K::RefreshOther,
3258 K::UsageTimeout,
3259 K::UsageTransport,
3260 ] {
3261 let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
3262 assert_eq!(f.kind, S::UsageReadFailed, "{kind}");
3263 }
3264 for kind in [K::RefreshActionRequired, K::CredentialStore] {
3267 let f = failure_of(&[failed("a", ObservationFailure::new(kind))]);
3268 assert_eq!(f.kind, S::ActionRequired, "{kind}");
3269 assert!(f.login_labels().is_empty(), "{kind}");
3270 }
3271 let f = failure_of(&[failed("a", ObservationFailure::new(K::UsageMalformed))]);
3272 assert_eq!(f.kind, S::QuotaUnknown);
3273 assert!(f.login_labels().is_empty());
3274 let stale_exhausted = ok("y", vec![win("primary", 5 * HOUR, 100.0, NOW - 1)]);
3276 assert_eq!(failure_of(&[stale_exhausted]).earliest_reset_ms, None);
3277 }
3278 #[test]
3283 fn an_excluded_duplicate_seat_does_not_pick_the_headline() {
3284 use AccountFailureKind as K;
3285 let reset = NOW + 2 * HOUR;
3286 let spent = ok("claude1", vec![win("primary", 5 * HOUR, 100.0, reset)]);
3287 let dup = ok("claude1-dup", vec![win("primary", 5 * HOUR, 100.0, reset)]);
3288 let exclude = [cred("claude1-dup")];
3289 let request = SelectionRequest {
3290 exclude: &exclude,
3291 ..req()
3292 };
3293 let f = match select(&request, &[spent, dup]) {
3294 Selection::NoCapacity { rejections, .. } => {
3295 SelectionFailure::from_rejections(&rejections, NOW)
3296 }
3297 other => panic!("expected no capacity, got {other:?}"),
3298 };
3299 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3300 assert_eq!(f.earliest_reset_ms, Some(reset));
3301 assert_eq!(f.accounts[1].kind, K::Excluded);
3302
3303 let mut cooling = ok("a", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3305 cooling.cooldown_until_ms = Some(NOW + HOUR);
3306 let dup = failed("b", ObservationFailure::new(K::UsageRateLimited));
3307 let exclude = [cred("b")];
3308 let request = SelectionRequest {
3309 exclude: &exclude,
3310 ..req()
3311 };
3312 let f = match select(&request, &[cooling, dup]) {
3313 Selection::NoCapacity { rejections, .. } => {
3314 SelectionFailure::from_rejections(&rejections, NOW)
3315 }
3316 other => panic!("expected no capacity, got {other:?}"),
3317 };
3318 assert_eq!(f.kind, SelectionFailureKind::AllCoolingDown, "{f:?}");
3319 assert_eq!(f.earliest_reset_ms, Some(NOW + HOUR));
3320
3321 let only = ok("x", vec![win("primary", 5 * HOUR, 10.0, NOW + DAY)]);
3324 let exclude = [cred("x")];
3325 let request = SelectionRequest {
3326 exclude: &exclude,
3327 ..req()
3328 };
3329 let f = match select(&request, &[only]) {
3330 Selection::NoCapacity { rejections, .. } => {
3331 SelectionFailure::from_rejections(&rejections, NOW)
3332 }
3333 other => panic!("expected no capacity, got {other:?}"),
3334 };
3335 assert_eq!(f.kind, SelectionFailureKind::QuotaUnknown, "{f:?}");
3336 }
3337}