//! Auto-resume: picking a session back up when its account's usage limit //! lifts. //! //! Off unless a session was switched to it, because this spends quota the //! moment quota exists and does it while nobody is watching. What it does is //! narrow on purpose: it sends one message -- "continue" unless something else //! was typed -- to a session that stopped because the account ran out, and //! then it is done. There is no retry loop around the conversation itself. //! //! **The schedule is a plan to ask, never a plan to send.** A reset time is //! the one thing here that cannot be trusted: the dialect's is a hint written //! when the turn failed, the endpoint's moves when the window moves, and both //! are wrong across the case this exists for -- a limit that lifts later than //! it said. So the wait ends in a *question* to [`crate::usage`], and only an //! answer that says the limits no longer apply sends anything. Every other //! answer, including one that cannot be got at all, becomes a new wait. //! //! This is the top layer: it holds the session manager and the usage monitor //! and neither holds it. That is what lets the decision below be a pure //! function of a snapshot and a clock, which is the whole of what is worth //! testing here. use std::sync::Arc; use std::time::Duration; use crate::session::{LimitHit, OwedResume, SessionManager, now}; use crate::usage::{UsageMonitor, UsageSnapshot, UsageState}; /// How often to look at the schedule. Coarse deliberately: a wait measured in /// hours does not deserve a fine-grained clock, and the meter behind it is /// cached for three minutes anyway. const TICK: Duration = Duration::from_secs(60); /// How close to a scheduled check is close enough to ask the meter. Anything /// further out is left alone, so a session waiting five hours costs nothing /// until the last few minutes of it. const NEARLY: f64 = 300.0; /// How long to wait after an answer that decided nothing -- the machine could /// not be asked, or it says the limit is still on with no reset time. const BACKOFF: f64 = 300.0; /// The least time to wait before asking again, whatever a reset time says. A /// window that claims to reset in the past would otherwise be asked about on /// every tick. const AT_LEAST: f64 = 60.0; /// How long after the limit was hit to stop waiting. /// /// Something has to bound it, or a machine that can never be asked -- an /// unplugged laptop, a machine somebody edited away -- is retried for ever with /// nothing on screen saying so. A day is past the longest window Claude /// reports, so reaching this means the wait was never going to end on its own. const GIVE_UP: f64 = 24.0 * 60.0 * 60.0; /// The percentage at which a window is spent. The API counts up to 100, so /// this is an equality in all but name; written as a threshold because a /// figure arriving slightly over is a full window, not a corrupt one. const SPENT: f64 = 100.0; /// What to do about one owed resume, having asked the meter. #[derive(Debug, Clone, Copy, PartialEq)] pub enum Step { /// The limits no longer apply: send the message. Send, /// Ask again at this epoch second. WaitUntil(f64), /// This has been waiting longer than anything real would take. GiveUp, } /// Runs the schedule until the server stops. /// /// Two things wake it: the tick, and a session reporting that it has just run /// out. The second is not an optimisation -- a limit hit is what *creates* a /// schedule, and a tick that happened a moment before it would otherwise leave /// the session unrecorded until the next one. pub async fn run(manager: Arc, monitor: Arc) { let mut limits = manager.subscribe_limits(); loop { tokio::select! { _ = tokio::time::sleep(TICK) => {} hit = limits.recv() => match hit { Ok(LimitHit { session_id, resets_at }) => note(&manager, &session_id, resets_at), // Lagged: some reports were dropped, and a session that hit a // limit while this was busy has no schedule. Nothing is lost // for good -- the sweep below reads the config, and the // session will report again the next time it is poked -- but // it is worth saying, because until then that session waits // for a person. Err(tokio::sync::broadcast::error::RecvError::Lagged(missed)) => { tracing::warn!("auto-resume missed {missed} limit reports"); } Err(tokio::sync::broadcast::error::RecvError::Closed) => return, }, } sweep(&manager, &monitor).await; } } /// Records a limit against the session that hit it, if it is one that resumes. pub(crate) fn note(manager: &SessionManager, session_id: &str, resets_at: Option) { match manager.note_limit(session_id, resets_at) { Ok(true) => tracing::info!("session {session_id} hit its usage limit; auto-resume is on"), Ok(false) => {} Err(err) => tracing::error!("couldn't schedule a resume for {session_id}: {err:#}"), } } /// One pass over everything owed a message. async fn sweep(manager: &SessionManager, monitor: &Arc) { let at = now(); for owed in manager.owed_resumes() { if owed.scheduled.at - at > NEARLY { continue; } // Asked per session rather than once for the whole sweep: the answer // is cached per machine and per meter, so several sessions on one // account share one fetch, and a machine nobody is waiting on is not // dialled at all. let snapshot = snapshot_for(Arc::clone(monitor), manager, &owed).await; match decide(snapshot.as_ref(), &owed, now()) { Step::Send => match manager.resume_now(&owed.session_id) { Ok(message) => tracing::info!( "the limit on {} has lifted; sent \"{message}\" to {}", owed.machine, owed.session_id ), Err(err) => { tracing::error!("couldn't resume {}: {err:#}", owed.session_id) } }, Step::WaitUntil(next) => { if let Err(err) = manager.reschedule_resume(&owed.session_id, next) { tracing::error!( "couldn't move {}'s resume to {next}: {err:#}", owed.session_id ); } } Step::GiveUp => { // About the machine rather than in the state's own words: the // detail is in the log, and what lands in the transcript has // to read on a phone. let why = match snapshot.as_ref().map(|snapshot| &snapshot.state) { Some(UsageState::Ok) => "the limit has not lifted in a day".to_string(), _ => format!("{} could not be asked for a day", owed.machine), }; if let Err(err) = manager.abandon_resume(&owed.session_id, &why) { tracing::error!("couldn't clear {}'s resume: {err:#}", owed.session_id); } } } } } /// The numbers for the machine and the meter this session is billed against, /// and `None` when nothing reports on it. /// /// Blocking work, so it goes to a blocking thread: the fetch behind it reads a /// credential file over ssh and then makes an HTTP call. async fn snapshot_for( monitor: Arc, manager: &SessionManager, owed: &OwedResume, ) -> Option { let machines: Vec<_> = manager .machines() .into_iter() .filter(|machine| machine.id == owed.machine) .collect(); if machines.is_empty() { return None; } let provider = owed.provider; tokio::task::spawn_blocking(move || { monitor .snapshots(&machines) .into_iter() .find(|snapshot| snapshot.provider == provider) }) .await .unwrap_or_default() } /// What one owed resume should do, given what the meter said and the time. /// /// A pure function of the two, which is what makes the rule inspectable: every /// answer that is not "the limits no longer apply" is a longer wait, and the /// only thing that ends the waiting other than success is the clock. /// /// The reset time comes from the *snapshot* rather than from the schedule, so /// a window that turns out to reset later than the dialect said pushes the /// check back, and one that resets sooner pulls it forward. That is the case /// the whole design is about: the first answer was a guess, this one is a /// measurement. pub fn decide(snapshot: Option<&UsageSnapshot>, owed: &OwedResume, at: f64) -> Step { let step = match snapshot { // The meter answered with numbers, which is the only answer that can // send anything. Some(snapshot) if snapshot.state == UsageState::Ok => { let spent: Vec<&crate::usage::UsageWindow> = snapshot .windows .iter() .filter(|window| window.percent >= SPENT) .collect(); if spent.is_empty() { Step::Send } else { // The earliest of the spent windows: it is the first moment // the situation can have changed, and if the others are still // full this comes straight back here. match spent .iter() .filter_map(|window| epoch_of(window.resets_at.as_deref())) .min_by(f64::total_cmp) { Some(resets) => Step::WaitUntil(resets), // Spent with no reset time anybody could read. Not a // reason to send: what is known is that the limit is on. None => Step::WaitUntil(at + BACKOFF), } } } // Logged out, unreachable, or the endpoint refused us -- and nothing // at all, which is a session whose machine or provider has gone. None // of them says the limit has lifted, and sending on any of them is // exactly the "inferred value presented as a measured one" this is // built to avoid. _ => Step::WaitUntil(at + BACKOFF), }; match step { // Waiting past the point where a real window would have reset means // whatever is wrong is not going to fix itself. Step::WaitUntil(_) if at - owed.scheduled.since > GIVE_UP => Step::GiveUp, Step::WaitUntil(next) => Step::WaitUntil(next.max(at + AT_LEAST)), other => other, } } /// An RFC-3339 timestamp as epoch seconds, and `None` for one that is absent /// or unreadable -- the same two answers the phone's countdown makes, kept /// apart from each other nowhere here because both mean "this cannot decide /// when to ask". fn epoch_of(resets_at: Option<&str>) -> Option { let text = resets_at?; time::OffsetDateTime::parse(text, &time::format_description::well_known::Rfc3339) .ok() .map(|at| at.unix_timestamp() as f64) } #[cfg(test)] mod tests { use super::*; use crate::config::ScheduledResume; use crate::usage::UsageWindow; fn owed(since: f64) -> OwedResume { OwedResume { session_id: "s1".to_string(), machine: "local".to_string(), provider: crate::usage::CLAUDE, scheduled: ScheduledResume { at: since, since }, } } fn snapshot(state: UsageState, windows: Vec) -> UsageSnapshot { UsageSnapshot { provider: crate::usage::CLAUDE.to_string(), machine: "local".to_string(), machine_name: "this machine".to_string(), limit_id: None, limit_name: None, state, windows, fetched_at: 0.0, } } fn window(percent: f64, resets_at: Option<&str>) -> UsageWindow { UsageWindow { kind: "session".to_string(), label: "5-hour window".to_string(), percent, duration_minutes: Some(300), resets_at: resets_at.map(str::to_string), active: true, } } #[test] fn a_meter_with_room_in_it_is_the_only_thing_that_sends() { let clear = snapshot(UsageState::Ok, vec![window(41.0, None)]); assert_eq!(decide(Some(&clear), &owed(0.0), 100.0), Step::Send); } #[test] fn a_window_still_spent_moves_the_check_to_its_own_reset_time() { // The case the feature exists for: the wait was scheduled for one // time, the limit is still on, and the endpoint now names another. let at = 1_788_546_972.0; let later = "2026-09-05T12:00:00+00:00"; let spent = snapshot(UsageState::Ok, vec![window(100.0, Some(later))]); assert_eq!( decide(Some(&spent), &owed(at - 60.0), at), Step::WaitUntil(epoch_of(Some(later)).expect("parses")) ); } #[test] fn a_reset_time_already_past_still_waits_a_little() { let at = 1_788_546_972.0; let spent = snapshot( UsageState::Ok, vec![window(100.0, Some("2020-01-01T00:00:00+00:00"))], ); assert_eq!( decide(Some(&spent), &owed(at - 60.0), at), Step::WaitUntil(at + AT_LEAST) ); } #[test] fn the_earliest_spent_window_is_the_one_worth_waiting_on() { let at = 1_788_546_972.0; let soon = "2026-09-05T12:00:00+00:00"; let far = "2026-09-09T12:00:00+00:00"; let mut weekly = window(100.0, Some(far)); weekly.kind = "weekly_all".to_string(); let spent = snapshot(UsageState::Ok, vec![window(100.0, Some(soon)), weekly]); assert_eq!( decide(Some(&spent), &owed(at - 60.0), at), Step::WaitUntil(epoch_of(Some(soon)).expect("parses")) ); } #[test] fn a_meter_that_could_not_be_asked_never_sends() { let at = 1_788_546_972.0; for state in [ UsageState::NotLoggedIn, UsageState::Unreachable { detail: "no route".to_string(), }, UsageState::Failed { detail: "429".to_string(), }, ] { let broken = snapshot(state.clone(), Vec::new()); assert_eq!( decide(Some(&broken), &owed(at - 60.0), at), Step::WaitUntil(at + BACKOFF), "{state:?}" ); } // And no snapshot at all -- a machine or provider edited away under a // session that was waiting on it. assert_eq!( decide(None, &owed(at - 60.0), at), Step::WaitUntil(at + BACKOFF) ); } #[test] fn waiting_longer_than_any_real_window_gives_up_rather_than_retrying_for_ever() { let at = 1_788_546_972.0; let broken = snapshot( UsageState::Unreachable { detail: "no route".to_string(), }, Vec::new(), ); assert_eq!( decide(Some(&broken), &owed(at - GIVE_UP - 1.0), at), Step::GiveUp ); // A meter that answers is still allowed to send on the same tick: the // ceiling bounds waiting, not resuming. let clear = snapshot(UsageState::Ok, vec![window(3.0, None)]); assert_eq!( decide(Some(&clear), &owed(at - GIVE_UP - 1.0), at), Step::Send ); } }