use crate::context::Context;
use crate::provider::ProviderId;
use crate::service::Permit;
use crate::state::{Key, State};
use crate::status::Row;
use crate::usage::{Snapshot, Window, same_reset};
use crate::{atomic, home};
use serde::{Deserialize, Serialize};
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::path::PathBuf;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Threshold(u8);
impl Threshold {
pub const LOWEST: u8 = 50;
pub const HIGHEST: u8 = 99;
pub const DEFAULT: Threshold = Threshold(95);
pub fn new(percent: u8) -> Option<Threshold> {
(Threshold::LOWEST..=Threshold::HIGHEST)
.contains(&percent)
.then_some(Threshold(percent))
}
pub fn clamped(percent: i64) -> Threshold {
let lowest = i64::from(Threshold::LOWEST);
let highest = i64::from(Threshold::HIGHEST);
Threshold(u8::try_from(percent.clamp(lowest, highest)).unwrap_or(Threshold::DEFAULT.0))
}
pub const fn percent(self) -> u8 {
self.0
}
fn reached_by(self, window: &Window, now: i64) -> bool {
window.used(now) >= f64::from(self.0)
}
}
impl Default for Threshold {
fn default() -> Threshold {
Threshold::DEFAULT
}
}
pub const SETTLING_SECONDS: i64 = 300;
pub const RETRY_SECONDS: i64 = 60;
pub const RETRY_MOST_SECONDS: i64 = 900;
pub const ATTEMPTS: u32 = 3;
fn retry_after(waits: u32) -> i64 {
(RETRY_SECONDS << waits.saturating_sub(1).min(4)).min(RETRY_MOST_SECONDS)
}
#[derive(Debug, Clone, PartialEq)]
pub struct Plan {
pub from: Key,
pub from_uuid: String,
pub to: Key,
pub limit: Window,
}
impl Plan {
pub(crate) fn same(&self, other: &Plan) -> bool {
self.from == other.from
&& self.from_uuid == other.from_uuid
&& self.to == other.to
&& self.limit.same_limit(&other.limit)
&& same_reset(
self.limit.resets_at.unwrap_or(0),
other.limit.resets_at.unwrap_or(0),
)
}
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum Decision {
Stay,
Switch(Plan),
NoRoom {
from: Key,
limit: Window,
},
}
#[derive(Debug)]
pub enum Auto {
Idle,
Switched {
from: String,
to: String,
limit: Window,
adoption: crate::provider::Adoption,
},
NoRoom { from: String, limit: Window },
Skipped {
from: String,
limit: Window,
why: Skip,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Skip {
SwitchInterrupted,
CustomOauth,
Overridden(Vec<String>),
}
impl Skip {
pub fn code(&self) -> &'static str {
match self {
Skip::SwitchInterrupted => "switch_interrupted",
Skip::CustomOauth => "custom_oauth_endpoint",
Skip::Overridden(_) => "auth_overridden",
}
}
}
fn reached(reading: &Snapshot, threshold: Threshold, now: i64) -> Option<&Window> {
reading
.windows
.iter()
.filter(|window| threshold.reached_by(window, now))
.min_by(|a, b| {
b.used(now)
.total_cmp(&a.used(now))
.then_with(|| sooner(a.resets_at, b.resets_at))
})
}
fn sooner(a: Option<i64>, b: Option<i64>) -> Ordering {
match (a, b) {
(Some(a), Some(b)) => a.cmp(&b),
(Some(_), None) => Ordering::Less,
(None, Some(_)) => Ordering::Greater,
(None, None) => Ordering::Equal,
}
}
fn room(
candidate: &Row,
mine: &Snapshot,
limit: &Window,
threshold: Threshold,
now: i64,
) -> Option<(f64, f64)> {
let theirs = &candidate.usage.as_ref()?.windows;
if theirs
.iter()
.any(|window| threshold.reached_by(window, now))
{
return None;
}
let reports = |kind: &str| {
theirs
.iter()
.any(|window| window.scope.is_none() && crate::usage::limit_name(&window.kind) == kind)
};
let overall = reports("session") && reports("weekly_all");
for had in &mine.windows {
if theirs.iter().any(|window| window.same_limit(had)) {
continue;
}
if !(overall && had.scope.is_some() && !had.same_limit(limit)) {
return None;
}
}
let used = theirs
.iter()
.find(|window| window.same_limit(limit))?
.used(now);
let fullest = theirs
.iter()
.filter(|window| !window.same_limit(limit))
.map(|window| window.used(now))
.fold(0.0, f64::max);
Some((used, fullest))
}
pub(crate) fn decide(
state: &State,
rows: &[Row],
ledger: &Ledger,
threshold: Threshold,
now: i64,
) -> Decision {
let claude = |row: &&Row| row.provider == ProviderId::Claude && row.label.is_some();
let Some(current) = rows.iter().filter(claude).find(|row| row.signed_in) else {
return Decision::Stay;
};
let (Some(from), Some(reading)) = (current.key(), current.usage.as_ref()) else {
return Decision::Stay;
};
let Some(limit) = reached(reading, threshold, now) else {
return Decision::Stay;
};
let settling = state
.get(&from)
.and_then(|account| account.last_used_at)
.is_some_and(|at| (0..SETTLING_SECONDS).contains(&(now - at)));
let tried = ledger.tried(¤t.account_uuid, limit);
let waiting = tried.is_some_and(|tried| {
tried.switched || tried.attempts >= ATTEMPTS || now - tried.at < retry_after(tried.waits)
});
if settling || waiting {
return Decision::Stay;
}
let passed_over = tried.map_or(&[][..], |tried| tried.passed_over.as_slice());
let best = rows
.iter()
.filter(claude)
.filter(|row| row.switchable(now))
.filter(|row| !passed_over.contains(&row.account_uuid))
.filter_map(|row| Some((row, room(row, reading, limit, threshold, now)?)))
.min_by(|(_, (a, a_most)), (_, (b, b_most))| {
a.total_cmp(b).then_with(|| a_most.total_cmp(b_most))
});
match best.and_then(|(row, _)| row.key()) {
Some(to) => Decision::Switch(Plan {
from,
from_uuid: current.account_uuid.clone(),
to,
limit: limit.clone(),
}),
None => Decision::NoRoom {
from,
limit: limit.clone(),
},
}
}
#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct Ledger {
#[serde(default)]
limits: BTreeMap<String, Tried>,
}
#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
struct Tried {
resets_at: i64,
attempts: u32,
#[serde(default)]
waits: u32,
at: i64,
#[serde(default)]
switched: bool,
#[serde(default)]
passed_over: Vec<String>,
}
fn path(ctx: &Context) -> PathBuf {
home::dir(ctx).join("autoswitch.json")
}
fn entry(account_uuid: &str, limit: &Window) -> String {
[
account_uuid,
crate::usage::limit_name(&limit.kind),
limit.scope.as_deref().unwrap_or_default(),
]
.join("/")
}
impl Ledger {
pub(crate) fn load(ctx: &Context) -> Ledger {
std::fs::read_to_string(path(ctx))
.ok()
.and_then(|raw| serde_json::from_str(&raw).ok())
.unwrap_or_default()
}
pub(crate) fn save(&mut self, ctx: &Context, permit: Permit) -> crate::error::Result<()> {
let now = ctx.now();
self.limits.retain(|_, tried| {
tried.resets_at > now || (tried.resets_at == 0 && now - tried.at < 7 * 86_400)
});
let body = serde_json::to_vec(self).map_err(|e| crate::error::Error::StateWriteFailed {
path: path(ctx),
source: std::io::Error::other(e),
})?;
atomic::write(permit, &path(ctx), &body, atomic::Perms::Secret).map_err(|source| {
crate::error::Error::StateWriteFailed {
path: path(ctx),
source,
}
})
}
fn tried(&self, account_uuid: &str, limit: &Window) -> Option<&Tried> {
self.limits
.get(&entry(account_uuid, limit))
.filter(|tried| same_reset(tried.resets_at, limit.resets_at.unwrap_or(0)))
}
fn trying(&mut self, plan: &Plan) -> &mut Tried {
let resets_at = plan.limit.resets_at.unwrap_or(0);
let tried = self
.limits
.entry(entry(&plan.from_uuid, &plan.limit))
.or_default();
if !same_reset(tried.resets_at, resets_at) {
*tried = Tried {
resets_at,
..Tried::default()
};
}
tried
}
pub(crate) fn attempt(&mut self, plan: &Plan, now: i64) {
let tried = self.trying(plan);
tried.attempts += 1;
tried.at = now;
}
pub(crate) fn waited(&mut self, plan: &Plan) {
let tried = self.trying(plan);
tried.attempts = tried.attempts.saturating_sub(1);
tried.waits = tried.waits.saturating_add(1);
}
pub(crate) fn switched(&mut self, plan: &Plan) {
self.trying(plan).switched = true;
}
pub(crate) fn pass_over(&mut self, plan: &Plan, to_uuid: &str) {
let tried = self.trying(plan);
tried.waits = 0;
if !tried.passed_over.iter().any(|uuid| uuid == to_uuid) {
tried.passed_over.push(to_uuid.to_owned());
}
}
}
#[derive(Debug)]
pub(crate) enum Next {
Switch(Plan),
Say(Auto),
}
pub(crate) fn look(ctx: &Context, state: &State, threshold: Threshold) -> Next {
let rows = crate::status::gather_offline(ctx, state).rows;
let plan = match decide(state, &rows, &Ledger::load(ctx), threshold, ctx.now()) {
Decision::Stay => return Next::Say(Auto::Idle),
Decision::NoRoom { from, limit } => {
return Next::Say(Auto::NoRoom {
from: state.typed(&from),
limit,
});
}
Decision::Switch(plan) => plan,
};
let overridden = crate::provider::of(ProviderId::Claude).overridden_by(ctx);
let why = if crate::switch::interrupted(ctx) {
Skip::SwitchInterrupted
} else if crate::settings::custom_oauth(ctx) {
Skip::CustomOauth
} else if !overridden.is_empty() {
Skip::Overridden(overridden)
} else {
return Next::Switch(plan);
};
Next::Say(Auto::Skipped {
from: state.typed(&plan.from),
limit: plan.limit,
why,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::history::Runway;
use crate::state::Park;
use crate::switch::harness::{NOW, account, window};
use crate::usage::Source;
fn scoped(kind: &str, model: &str, percent: f64) -> Window {
Window {
scope: Some(model.into()),
..window(kind, percent)
}
}
fn row(label: &str, signed_in: bool, windows: Vec<Window>) -> Row {
Row {
provider: ProviderId::Claude,
label: Some(label.into()),
email: format!("{label}@example.com"),
account_uuid: label.into(),
signed_in,
parked: (!signed_in).then(|| Park {
service: format!("pitboard-park-{label}-1"),
parked_at: NOW,
refresh_fingerprint: label.into(),
access_expires_at: Some(NOW + 3600),
refresh_expires_at: Some(NOW + 30 * 86_400),
}),
usage: Some(Snapshot {
windows,
observed_at: Some(NOW),
account_uuid: Some(label.into()),
source: Source::Remembered,
}),
stale: None,
runway: Runway::Unknown,
}
}
fn state_of(rows: &[Row]) -> State {
let mut state = State::default();
for row in rows {
state.accounts.push(account(
row.label.as_deref().unwrap(),
&row.account_uuid,
row.parked.clone(),
));
}
state
}
fn decided(rows: &[Row]) -> Decision {
decide(
&state_of(rows),
rows,
&Ledger::default(),
Threshold::DEFAULT,
NOW,
)
}
fn to(decision: &Decision) -> Option<&str> {
match decision {
Decision::Switch(plan) => Some(plan.to.label.as_str()),
_ => None,
}
}
fn work_at(percent: f64) -> Row {
row(
"work",
true,
vec![window("session", percent), window("weekly_all", 20.0)],
)
}
fn spare() -> Row {
row(
"spare",
false,
vec![window("session", 10.0), window("weekly_all", 30.0)],
)
}
#[test]
fn a_limit_below_the_share_is_left_alone_and_one_at_it_switches() {
assert_eq!(decided(&[work_at(94.9), spare()]), Decision::Stay);
let decision = decided(&[work_at(95.0), spare()]);
assert_eq!(to(&decision), Some("spare"));
let Decision::Switch(plan) = decision else {
unreachable!()
};
assert_eq!(
(plan.from.label.as_str(), plan.from_uuid.as_str()),
("work", "work")
);
assert_eq!(plan.limit.kind, "session");
}
#[test]
fn every_limit_the_account_reports_can_reach_the_share() {
let weekly = row(
"work",
true,
vec![window("session", 10.0), window("weekly_all", 97.0)],
);
let opus = row(
"work",
true,
vec![
window("session", 10.0),
window("weekly_all", 10.0),
scoped("weekly_scoped", "Opus", 96.0),
],
);
let named_as_a_session_names_them = row(
"work",
true,
vec![window("five_hour", 99.0), window("seven_day", 10.0)],
);
for (case, work) in [
("weekly", weekly),
("opus", opus),
("five_hour", named_as_a_session_names_them),
] {
let spare = row(
"spare",
false,
vec![
window("session", 0.0),
window("weekly_all", 0.0),
scoped("weekly_scoped", "Opus", 0.0),
],
);
assert_eq!(to(&decided(&[work, spare])), Some("spare"), "{case}");
}
}
#[test]
fn the_limit_furthest_past_the_share_is_the_one_acted_on() {
let work = row(
"work",
true,
vec![window("session", 96.0), window("weekly_all", 98.0)],
);
let Decision::Switch(plan) = decided(&[work, spare()]) else {
panic!("a switch")
};
assert_eq!(plan.limit.kind, "weekly_all");
}
#[test]
fn a_limit_whose_reset_has_passed_counts_as_reset() {
let mut work = work_at(100.0);
work.usage.as_mut().unwrap().windows[0].resets_at = Some(NOW);
assert_eq!(decided(&[work, spare()]), Decision::Stay);
}
#[test]
fn an_account_in_use_with_no_reading_is_never_switched_from() {
let mut work = work_at(99.0);
work.usage = None;
assert_eq!(decided(&[work, spare()]), Decision::Stay);
}
#[test]
fn codex_is_never_switched() {
let mut work = work_at(99.0);
let mut other = spare();
work.provider = ProviderId::Codex;
other.provider = ProviderId::Codex;
assert_eq!(decided(&[work, other]), Decision::Stay);
assert_eq!(
decided(&[work_at(99.0), {
let mut codex = spare();
codex.provider = ProviderId::Codex;
codex
}]),
Decision::NoRoom {
from: Key::new(ProviderId::Claude, "work"),
limit: window("session", 99.0)
}
);
}
#[test]
fn an_account_with_any_limit_at_the_share_is_no_place_to_go() {
let full_weekly = || {
row(
"spare",
false,
vec![window("session", 0.0), window("weekly_all", 95.0)],
)
};
assert!(matches!(
decided(&[work_at(96.0), full_weekly()]),
Decision::NoRoom { .. }
));
let roomy = row(
"roomy",
false,
vec![window("session", 60.0), window("weekly_all", 10.0)],
);
assert_eq!(
to(&decided(&[work_at(96.0), full_weekly(), roomy])),
Some("roomy")
);
}
#[test]
fn an_account_that_cannot_be_switched_to_is_passed_over() {
let mut expired = spare();
expired.parked.as_mut().unwrap().refresh_expires_at = Some(NOW);
assert!(matches!(
decided(&[work_at(96.0), expired]),
Decision::NoRoom { .. }
));
}
#[test]
fn a_limit_the_other_account_does_not_report_counts_as_spent() {
let unknown_weekly = row("spare", false, vec![window("session", 0.0)]);
assert!(matches!(
decided(&[work_at(96.0), unknown_weekly]),
Decision::NoRoom { .. }
));
let mut unread = spare();
unread.usage = None;
assert!(matches!(
decided(&[work_at(96.0), unread]),
Decision::NoRoom { .. }
));
}
#[test]
fn a_limit_for_a_model_the_other_plan_does_not_have_is_not_counted_against_it() {
let work = row(
"work",
true,
vec![
window("session", 97.0),
window("weekly_all", 20.0),
scoped("weekly_scoped", "Opus", 40.0),
],
);
assert_eq!(to(&decided(&[work, spare()])), Some("spare"));
let opus_ran_out = row(
"work",
true,
vec![
window("session", 10.0),
window("weekly_all", 20.0),
scoped("weekly_scoped", "Opus", 97.0),
],
);
assert!(matches!(
decided(&[opus_ran_out, spare()]),
Decision::NoRoom { .. }
));
let session_only = row("spare", false, vec![window("session", 10.0)]);
let work = row(
"work",
true,
vec![
window("session", 97.0),
scoped("weekly_scoped", "Opus", 40.0),
],
);
assert!(matches!(
decided(&[work, session_only]),
Decision::NoRoom { .. }
));
}
#[test]
fn the_account_with_most_room_in_that_limit_is_chosen_then_the_least_full() {
let less = || {
row(
"less",
false,
vec![window("session", 50.0), window("weekly_all", 10.0)],
)
};
let more = row(
"more",
false,
vec![window("session", 20.0), window("weekly_all", 80.0)],
);
assert_eq!(to(&decided(&[work_at(96.0), less(), more])), Some("more"));
let tie = row(
"tie",
false,
vec![window("session", 50.0), window("weekly_all", 5.0)],
);
assert_eq!(to(&decided(&[work_at(96.0), less(), tie])), Some("tie"));
let same = row(
"same",
false,
vec![window("session", 50.0), window("weekly_all", 10.0)],
);
assert_eq!(
to(&decided(&[work_at(96.0), less(), same])),
Some("less"),
"the first of two alike"
);
}
#[test]
fn an_account_just_put_in_use_is_left_to_settle() {
let rows = [work_at(96.0), spare()];
let mut state = state_of(&rows);
let work = Key::new(ProviderId::Claude, "work");
state.used(&work, NOW - SETTLING_SECONDS + 1);
let ledger = Ledger::default();
assert_eq!(
decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW),
Decision::Stay
);
state.used(&work, NOW - SETTLING_SECONDS);
assert_eq!(
to(&decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW)),
Some("spare")
);
}
fn plan(rows: &[Row]) -> Plan {
match decided(rows) {
Decision::Switch(plan) => plan,
other => panic!("a switch, not {other:?}"),
}
}
#[test]
fn a_limit_switched_away_from_is_not_switched_away_from_again_before_it_resets() {
let rows = [work_at(96.0), spare()];
let state = state_of(&rows);
let mut ledger = Ledger::default();
let plan = plan(&rows);
ledger.attempt(&plan, NOW - 3600);
ledger.switched(&plan);
assert_eq!(
decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW),
Decision::Stay
);
let mut next_window = work_at(96.0);
next_window.usage.as_mut().unwrap().windows[0].resets_at = Some(NOW + 3 * 3600);
let rows = [next_window, spare()];
assert_eq!(
to(&decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW)),
Some("spare")
);
}
#[test]
fn a_failed_attempt_is_tried_again_after_a_minute_and_at_most_three_times() {
let rows = [work_at(96.0), spare()];
let state = state_of(&rows);
let mut ledger = Ledger::default();
let plan = plan(&rows);
ledger.attempt(&plan, NOW);
let at = |now| decide(&state, &rows, &ledger, Threshold::DEFAULT, now);
assert_eq!(at(NOW + RETRY_SECONDS - 1), Decision::Stay);
assert_eq!(to(&at(NOW + RETRY_SECONDS)), Some("spare"));
assert_eq!(at(NOW - 1), Decision::Stay, "a clock gone backwards waits");
ledger.attempt(&plan, NOW);
ledger.attempt(&plan, NOW);
let at = |now| decide(&state, &rows, &ledger, Threshold::DEFAULT, now);
assert_eq!(at(NOW + 3600), Decision::Stay);
}
#[test]
fn a_failure_trying_again_may_mend_waits_longer_each_time_and_is_not_counted() {
let rows = [work_at(96.0), spare()];
let state = state_of(&rows);
let mut ledger = Ledger::default();
let plan = plan(&rows);
let mut at = NOW;
for wait in [60, 120, 240, 480, 900, 900] {
ledger.attempt(&plan, at);
ledger.waited(&plan);
let decided = |now| decide(&state, &rows, &ledger, Threshold::DEFAULT, now);
assert_eq!(decided(at + wait - 1), Decision::Stay, "within {wait} s");
assert_eq!(to(&decided(at + wait)), Some("spare"), "after {wait} s");
at += wait;
}
}
#[test]
fn an_account_a_switch_was_refused_over_is_passed_over_for_the_next() {
let other = row(
"other",
false,
vec![window("session", 30.0), window("weekly_all", 30.0)],
);
let rows = [work_at(96.0), spare(), other];
let state = state_of(&rows);
let mut ledger = Ledger::default();
let plan = plan(&rows);
assert_eq!(plan.to.label, "spare");
ledger.attempt(&plan, NOW);
ledger.pass_over(&plan, "spare");
assert_eq!(
to(&decide(
&state,
&rows,
&ledger,
Threshold::DEFAULT,
NOW + RETRY_SECONDS
)),
Some("other")
);
}
#[test]
fn a_reset_a_second_off_is_the_same_window_and_a_minute_off_is_not() {
let rows = [work_at(96.0), spare()];
let mut ledger = Ledger::default();
let mut plan = plan(&rows);
let resets = plan.limit.resets_at.unwrap();
plan.limit.resets_at = Some(resets + 59);
ledger.attempt(&plan, NOW - 3600);
ledger.switched(&plan);
let state = state_of(&rows);
assert_eq!(
decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW),
Decision::Stay
);
plan.limit.resets_at = Some(resets + 61);
let mut ledger = Ledger::default();
ledger.attempt(&plan, NOW - 3600);
ledger.switched(&plan);
assert_eq!(
to(&decide(&state, &rows, &ledger, Threshold::DEFAULT, NOW)),
Some("spare")
);
}
#[test]
fn a_share_is_between_fifty_and_ninety_nine() {
assert_eq!(Threshold::new(49), None);
assert_eq!(Threshold::new(100), None);
assert_eq!(Threshold::new(50).map(Threshold::percent), Some(50));
assert_eq!(Threshold::clamped(120).percent(), 99);
assert_eq!(Threshold::clamped(-3).percent(), 50);
assert_eq!(Threshold::default().percent(), 95);
let lower = Threshold::new(80).unwrap();
let rows = [work_at(85.0), spare()];
assert_eq!(decided(&rows), Decision::Stay);
assert_eq!(
to(&decide(
&state_of(&rows),
&rows,
&Ledger::default(),
lower,
NOW
)),
Some("spare")
);
}
}