ai-app/server/src/resume.rs

//! 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 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;

/// 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.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 => {
                // 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.setup),
                };
                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 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()
}

/// 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(),
            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() {
        // 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
        );
    }
}