386 lines
15 KiB
Rust
386 lines
15 KiB
Rust
//! 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<SessionManager>, monitor: Arc<UsageMonitor>) {
|
|
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<f64>) {
|
|
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<UsageMonitor>) {
|
|
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<UsageMonitor>,
|
|
manager: &SessionManager,
|
|
owed: &OwedResume,
|
|
) -> Option<UsageSnapshot> {
|
|
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<f64> {
|
|
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<UsageWindow>) -> 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
|
|
);
|
|
}
|
|
}
|