306 lines
11 KiB
Rust
306 lines
11 KiB
Rust
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use crate::session::{LimitHit, OwedResume, SessionManager, now};
|
|
use crate::usage::{UsageMonitor, UsageSnapshot, UsageState};
|
|
|
|
const TICK: Duration = Duration::from_secs(60);
|
|
|
|
const NEARLY: f64 = 300.0;
|
|
|
|
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;
|
|
|
|
/// Something has to bound it, or a machine that can never be asked -- an
|
|
/// unplugged laptop, a setup 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;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum Step {
|
|
Send,
|
|
/// Ask again at this epoch second.
|
|
WaitUntil(f64),
|
|
GiveUp,
|
|
}
|
|
|
|
/// 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;
|
|
}
|
|
}
|
|
|
|
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:#}"),
|
|
}
|
|
}
|
|
|
|
async fn sweep(manager: &SessionManager, monitor: &Arc<UsageMonitor>) {
|
|
let at = now();
|
|
for owed in manager.owed_resumes() {
|
|
if owed.scheduled.at - at > NEARLY {
|
|
continue;
|
|
}
|
|
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.setup,
|
|
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 => {
|
|
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.setup),
|
|
};
|
|
if let Err(err) = manager.abandon_resume(&owed.session_id, &why) {
|
|
tracing::error!("couldn't clear {}'s resume: {err:#}", owed.session_id);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn snapshot_for(
|
|
monitor: Arc<UsageMonitor>,
|
|
manager: &SessionManager,
|
|
owed: &OwedResume,
|
|
) -> Option<UsageSnapshot> {
|
|
let setups: Vec<_> = manager
|
|
.setups()
|
|
.into_iter()
|
|
.filter(|setup| setup.id == owed.setup)
|
|
.collect();
|
|
if setups.is_empty() {
|
|
return None;
|
|
}
|
|
let provider = owed.provider;
|
|
tokio::task::spawn_blocking(move || {
|
|
monitor
|
|
.snapshots(&setups)
|
|
.into_iter()
|
|
.find(|snapshot| snapshot.provider == provider)
|
|
})
|
|
.await
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
pub fn decide(snapshot: Option<&UsageSnapshot>, owed: &OwedResume, at: f64) -> Step {
|
|
let step = match snapshot {
|
|
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 {
|
|
match spent
|
|
.iter()
|
|
.filter_map(|window| epoch_of(window.resets_at.as_deref()))
|
|
.min_by(f64::total_cmp)
|
|
{
|
|
Some(resets) => Step::WaitUntil(resets),
|
|
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 {
|
|
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(),
|
|
setup: "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(),
|
|
setup: "local".to_string(),
|
|
setup_name: "this machine".to_string(),
|
|
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,
|
|
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() {
|
|
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:?}"
|
|
);
|
|
}
|
|
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
|
|
);
|
|
let clear = snapshot(UsageState::Ok, vec![window(3.0, None)]);
|
|
assert_eq!(
|
|
decide(Some(&clear), &owed(at - GIVE_UP - 1.0), at),
|
|
Step::Send
|
|
);
|
|
}
|
|
}
|