pitboard-core 0.5.1

//! How often pitboard is allowed to ask Anthropic about an account.
//!
//! `status` asked about every enrolled account plus the live login on every run, with no
//! memory of having just asked, and the menu bar app asked the same questions every five
//! minutes, on every wake, and on every panel open, from a process that knew nothing about
//! the command line's. Nothing honoured `Retry-After`: a 429 became a stale row and the
//! identical request went out on the next tick. Two accounts and a running app is on the
//! order of six hundred authenticated requests a day that nobody asked for.
//!
//! Beyond the cost, that is the part of pitboard's behaviour that reads least like a person
//! switching between their own accounts, which is the one appearance this project cannot
//! afford.
//!
//! So there is a floor, and it is derived rather than chosen: an account may be asked again
//! once its tightest limit could have moved by one percentage point. A five-hour window
//! moves at most 100% in five hours, so that is three minutes; a weekly window, a hundred.
//! Whatever the floor is, it is far below what changes anybody's decision, and it collapses
//! the repeated asking to one request per account per few minutes however many front ends
//! are running.
//!
//! Nothing here is secret. It is names, times and counts.

use crate::context::Context;
use crate::usage::Snapshot;
use crate::{atomic, home};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// How long a window runs, which is what makes the floor a derivation rather than a
/// preference.
///
/// The reading says, where the service said: OpenAI states every window's length, and
/// pitboard derives Anthropic's from its kind when it reads the answer. A reading
/// remembered from before pitboard kept the length has only its kind, and Anthropic's
/// kinds are the only ones that could have been remembered then.
///
/// Until this read the length, it knew only `five_hour` and `seven_day`, and Anthropic has
/// been answering `session` and `weekly_all`. Every account was asked about once a minute,
/// the unknown-window fallback, against the three minutes this module promises.
fn window_seconds(window: &crate::usage::Window) -> Option<i64> {
    window
        .length_seconds
        .or_else(|| crate::usage::anthropic_window_length(&window.kind))
}

/// One percentage point of the tightest window this account has, in seconds. The floor.
///
/// With no reading to derive it from, a minute: enough to collapse a burst of front ends
/// starting together, short enough that nobody notices.
const UNKNOWN_FLOOR: i64 = 60;

pub fn floor_for(reading: Option<&Snapshot>) -> i64 {
    let shortest = reading
        .map(|s| s.windows.as_slice())
        .unwrap_or_default()
        .iter()
        .filter_map(window_seconds)
        .min();
    shortest.map_or(UNKNOWN_FLOOR, |seconds| (seconds / 100).max(1))
}

/// Anthropic asked for less traffic, or could not be reached. Both mean waiting, and how
/// long is the difference between them.
const RATE_LIMITED_FIRST: i64 = 60;
const RATE_LIMITED_MOST: i64 = 3600;
const UNREACHABLE_FIRST: i64 = 30;
const UNREACHABLE_MOST: i64 = 900;

/// What is known about asking one account.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct Record {
    /// When Anthropic last answered about this account, in epoch seconds.
    #[serde(default)]
    pub answered_at: i64,
    /// Not before this, in epoch seconds. Set by a refusal, cleared by an answer.
    #[serde(default)]
    pub held_until: i64,
    /// Refusals in a row, which is what makes the wait grow.
    #[serde(default)]
    pub refusals: u32,
}

/// One account's line in the ledger: its record, and why its wait was set.
///
/// The reason is kept beside the record rather than in it. `Record` is public and can be
/// written as a literal, so a field added to it breaks every such literal, which a patch
/// release must not do. The file is the same either way: the reason is one more key.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct Entry {
    #[serde(flatten)]
    record: Record,
    /// Absent in a line written before pitboard kept it, whose wait is then told by its
    /// length, as it always was.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    held_for: Option<Reason>,
}

/// Why an account is being held off, which decides whether asking for it may go through.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
enum Reason {
    /// The service asked for less traffic.
    RateLimited,
    /// The service could not be reached, and pitboard chose the wait.
    Unreachable,
}

impl Entry {
    fn held(&self) -> Held {
        let record = &self.record;
        match self.held_for {
            Some(Reason::RateLimited) => Held::RateLimited,
            Some(Reason::Unreachable) => Held::Unreachable,
            // No unreachable wait runs longer than its cap, so a longer one was asked for.
            None if record.refusals > 0
                && record.held_until - record.answered_at > UNREACHABLE_MOST =>
            {
                Held::RateLimited
            }
            None => Held::Unreachable,
        }
    }
}

type Ledger = HashMap<String, Entry>;

fn path(ctx: &Context) -> PathBuf {
    home::dir(ctx).join("asking.json")
}

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()
}

fn save(ctx: &Context, ledger: &Ledger) {
    if let Ok(body) = serde_json::to_string(ledger) {
        let _ = atomic::write(&path(ctx), body.as_bytes(), atomic::Perms::Secret);
    }
}

/// Why an account is not being asked about right now.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Held {
    /// Asked recently enough that the answer cannot have moved by a point.
    Fresh,
    /// Anthropic asked for less traffic.
    RateLimited,
    /// Anthropic could not be reached, and trying again at once helps nobody.
    Unreachable,
}

/// Whether to ask about this account now.
///
/// `None` means ask. `Some(held)` means do not, and says which of the three reasons it is,
/// so a front end can show the right thing and a person can tell a quiet answer from a
/// refused one.
///
/// `forced` is a person asking, with `status --fresh` or the app's Refresh. It goes past the
/// floor and past a wait pitboard chose after failing to reach the service, since asking is
/// how somebody says the network is back. It does not go past a wait the service asked for:
/// that request is the traffic the service asked not to get, and it would answer with the
/// same refusal.
pub fn may_ask(
    ctx: &Context,
    account_uuid: &str,
    last_reading: Option<&Snapshot>,
    forced: bool,
) -> Option<Held> {
    let ledger = load(ctx);
    let entry = ledger.get(account_uuid)?;
    let record = &entry.record;
    let now = ctx.now();
    if record.held_until > now {
        let held = entry.held();
        return (!forced || held == Held::RateLimited).then_some(held);
    }
    (!forced && now - record.answered_at < floor_for(last_reading)).then_some(Held::Fresh)
}

/// What one request turned out to be, for recording afterwards.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Outcome {
    Answered,
    /// Anthropic asked for less traffic. `Some` is what it said to wait, in seconds.
    RateLimited(Option<i64>),
    /// Anthropic could not be reached, or answered in a way that trying again might fix.
    Unreachable,
}

/// Record what came back, for one account or for several.
///
/// Several at a time on purpose. Every front end asks about all its accounts at once, on a
/// thread each, and a record-as-you-go would have each thread read this file, change one
/// entry and write the whole thing back, so the last writer would erase what the others
/// learned. Which is exactly what happened: two accounts asked together, one budgeted and
/// one not, forever.
pub fn record(ctx: &Context, outcomes: &[(String, Outcome)]) {
    if outcomes.is_empty() {
        return;
    }
    let now = ctx.now();
    let mut ledger = load(ctx);
    for (account_uuid, outcome) in outcomes {
        let entry = ledger.entry(account_uuid.clone()).or_default();
        let record = &mut entry.record;
        match outcome {
            Outcome::Answered => {
                record.answered_at = now;
                record.held_until = 0;
                record.refusals = 0;
                entry.held_for = None;
            }
            Outcome::RateLimited(retry_after) => {
                hold(
                    record,
                    now,
                    RATE_LIMITED_FIRST,
                    RATE_LIMITED_MOST,
                    *retry_after,
                );
                entry.held_for = Some(Reason::RateLimited);
            }
            Outcome::Unreachable => {
                hold(record, now, UNREACHABLE_FIRST, UNREACHABLE_MOST, None);
                entry.held_for = Some(Reason::Unreachable);
            }
        }
    }
    save(ctx, &ledger);
}

/// `retry_after` is what Anthropic said to wait, where it said anything, and is believed
/// over anything pitboard would pick.
fn hold(record: &mut Record, now: i64, first: i64, most: i64, retry_after: Option<i64>) {
    record.refusals = record.refusals.saturating_add(1);
    let backoff = retry_after
        .filter(|seconds| *seconds > 0)
        .unwrap_or_else(|| doubling(first, most, record.refusals - 1));
    record.held_until = now + backoff;
}

/// Doubling from `first`, capped at `most`. `before` is how many refusals came before this
/// one, so the first wait is `first` itself.
///
/// No jitter, because these waits are per account and per machine: there is no herd here to
/// spread out, and a jittered wait is one a person cannot predict from what doctor told
/// them.
fn doubling(first: i64, most: i64, before: u32) -> i64 {
    let shifted = first.saturating_mul(1i64 << before.min(16));
    shifted.clamp(first, most)
}

/// Every account currently being held off, for `doctor` to report.
pub fn holds(ctx: &Context) -> Vec<(String, i64)> {
    let now = ctx.now();
    let mut held: Vec<(String, i64)> = load(ctx)
        .into_iter()
        .filter(|(_, e)| e.record.held_until > now)
        .map(|(uuid, e)| (uuid, e.record.held_until - now))
        .collect();
    held.sort();
    held
}

/// Drop what is known about an account nobody is enrolled as any more.
pub fn forget(ctx: &Context, account_uuid: &str) {
    let mut ledger = load(ctx);
    if ledger.remove(account_uuid).is_some() {
        save(ctx, &ledger);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::time::{Clock, FixedClock};
    use crate::usage::{Source, Window};
    use std::sync::Arc;

    const NOW: i64 = 1_760_000_000;

    struct Scratch(PathBuf);
    impl Drop for Scratch {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    fn machine(name: &str) -> (Context, Arc<FixedClock>, Scratch) {
        let root = std::env::temp_dir().join(format!(
            "pitboard-budget-{name}-{}-{:?}",
            std::process::id(),
            std::thread::current().id()
        ));
        let _ = std::fs::remove_dir_all(&root);
        let clock = Arc::new(FixedClock::at(NOW));
        let ctx = Context::new(root.clone())
            .with_pitboard_home(root.clone())
            .with_clock(Arc::clone(&clock) as Arc<dyn Clock>);
        home::ensure(&ctx).expect("a home");
        (ctx, clock, Scratch(root))
    }

    fn reading(kinds: &[&str]) -> Snapshot {
        Snapshot {
            observed_at: Some(NOW),
            account_uuid: None,
            source: Source::Live,
            windows: kinds
                .iter()
                .map(|kind| Window {
                    kind: (*kind).to_string(),
                    scope: None,
                    percent: 10.0,
                    resets_at: None,
                    is_active: true,
                    severity: None,
                    length_seconds: None,
                })
                .collect(),
        }
    }

    /// Anthropic names its windows `session` and `weekly_all`. The floor knew only the older
    /// `five_hour` and `seven_day`, so every account was asked about once a minute against
    /// the three the module promises.
    #[test]
    fn the_floor_knows_the_windows_anthropic_actually_sends() {
        assert_eq!(floor_for(Some(&reading(&["session"]))), 180);
        assert_eq!(
            floor_for(Some(&reading(&["weekly_all", "weekly_scoped"]))),
            6048
        );
    }

    /// A window whose length the service stated is timed by that, whatever it is called:
    /// OpenAI's are measured in seconds and named after their length.
    #[test]
    fn a_stated_length_wins_over_the_name() {
        let mut stated = reading(&["primary"]);
        stated.windows[0].length_seconds = Some(18_000);
        assert_eq!(floor_for(Some(&stated)), 180);
    }

    /// The floor is a derivation, not a preference: how long the tightest limit takes to
    /// move by one percentage point.
    #[test]
    fn the_floor_comes_from_the_window_rather_than_from_taste() {
        assert_eq!(floor_for(Some(&reading(&["five_hour"]))), 180);
        assert_eq!(floor_for(Some(&reading(&["seven_day"]))), 6048);
        assert_eq!(
            floor_for(Some(&reading(&["seven_day", "five_hour"]))),
            180,
            "the tightest window is the one that decides"
        );
        assert_eq!(floor_for(None), UNKNOWN_FLOOR);
        assert_eq!(
            floor_for(Some(&reading(&["something_new"]))),
            UNKNOWN_FLOOR,
            "a window kind pitboard does not know is not a reason to ask forever"
        );
    }

    #[test]
    fn an_account_nobody_has_asked_about_is_asked_about() {
        let (ctx, _clock, _s) = machine("first");
        assert_eq!(may_ask(&ctx, "acc", None, false), None);
    }

    #[test]
    fn asking_again_inside_the_floor_serves_what_is_already_known() {
        let (ctx, clock, _s) = machine("floor");
        let five_hour = reading(&["five_hour"]);
        record(&ctx, &[("acc".into(), Outcome::Answered)]);

        assert_eq!(
            may_ask(&ctx, "acc", Some(&five_hour), false),
            Some(Held::Fresh)
        );
        clock.advance(179);
        assert_eq!(
            may_ask(&ctx, "acc", Some(&five_hour), false),
            Some(Held::Fresh)
        );
        clock.advance(2);
        assert_eq!(may_ask(&ctx, "acc", Some(&five_hour), false), None);
    }

    #[test]
    fn asking_for_it_goes_past_the_floor() {
        let (ctx, _clock, _s) = machine("forced");
        record(&ctx, &[("acc".into(), Outcome::Answered)]);
        assert_eq!(
            may_ask(&ctx, "acc", Some(&reading(&["five_hour"])), true),
            None
        );
    }

    /// `status --fresh` and the app's Refresh do not ask through a wait the service asked
    /// for, as the 0.3.0 changelog and doctor's advice say. They asked through it before.
    #[test]
    fn asking_for_it_keeps_a_wait_the_service_asked_for() {
        let (ctx, clock, _s) = machine("forced-rate-limited");
        record(&ctx, &[("acc".into(), Outcome::RateLimited(Some(300)))]);
        assert_eq!(may_ask(&ctx, "acc", None, true), Some(Held::RateLimited));
        clock.advance(301);
        assert_eq!(may_ask(&ctx, "acc", None, true), None);
    }

    /// A wait pitboard chose after failing to reach the service is its own guess, and
    /// somebody asking is how it learns the network is back.
    #[test]
    fn asking_for_it_tries_an_unreachable_service_again() {
        let (ctx, _clock, _s) = machine("forced-unreachable");
        record(&ctx, &[("acc".into(), Outcome::Unreachable)]);
        assert_eq!(may_ask(&ctx, "acc", None, false), Some(Held::Unreachable));
        assert_eq!(may_ask(&ctx, "acc", None, true), None);
    }

    /// The reason was told by the wait's length alone, and an account never answered
    /// has a long way back to its last answer: one that could not be reached was said
    /// to be rate limited.
    #[test]
    fn an_unreachable_service_is_not_called_rate_limiting() {
        let (ctx, _clock, _s) = machine("never-answered");
        record(&ctx, &[("acc".into(), Outcome::Unreachable)]);
        assert_eq!(may_ask(&ctx, "acc", None, false), Some(Held::Unreachable));
    }

    /// A record an older pitboard wrote has no reason, and is told by its length as it
    /// was then: longer than any unreachable wait means asked for.
    #[test]
    fn a_wait_an_older_pitboard_recorded_is_told_by_its_length() {
        let (ctx, _clock, _s) = machine("older");
        let written = format!(
            r#"{{"asked":{{"answered_at":{a},"held_until":{h},"refusals":1}},"short":{{"answered_at":{a},"held_until":{s},"refusals":1}}}}"#,
            a = NOW - 10,
            h = NOW + 3000,
            s = NOW + 20,
        );
        std::fs::write(path(&ctx), written).unwrap();
        assert_eq!(may_ask(&ctx, "asked", None, true), Some(Held::RateLimited));
        assert_eq!(may_ask(&ctx, "short", None, false), Some(Held::Unreachable));
        assert_eq!(may_ask(&ctx, "short", None, true), None);
    }

    /// What Anthropic said to wait is believed over anything pitboard would pick.
    #[test]
    fn a_retry_after_is_taken_at_its_word() {
        let (ctx, clock, _s) = machine("retry-after");
        record(&ctx, &[("acc".into(), Outcome::RateLimited(Some(300)))]);

        assert!(may_ask(&ctx, "acc", None, false).is_some());
        clock.advance(299);
        assert!(may_ask(&ctx, "acc", None, false).is_some());
        clock.advance(2);
        assert_eq!(may_ask(&ctx, "acc", None, false), None);
    }

    #[test]
    fn refusals_in_a_row_wait_longer_each_time_up_to_a_cap() {
        let (ctx, clock, _s) = machine("doubling");
        let waits: Vec<i64> = (0..8)
            .map(|_| {
                record(&ctx, &[("acc".into(), Outcome::RateLimited(None))]);
                let held = holds(&ctx);
                clock.advance(held.first().map_or(0, |(_, left)| *left));
                held.first().map_or(0, |(_, left)| *left)
            })
            .collect();
        assert_eq!(waits[0], RATE_LIMITED_FIRST);
        assert!(waits[1] > waits[0]);
        assert!(
            waits.iter().all(|w| *w <= RATE_LIMITED_MOST),
            "the wait has a ceiling: {waits:?}"
        );
        assert_eq!(*waits.last().expect("eight of them"), RATE_LIMITED_MOST);
    }

    #[test]
    fn an_answer_clears_a_wait_and_starts_the_floor_again() {
        let (ctx, _clock, _s) = machine("cleared");
        record(&ctx, &[("acc".into(), Outcome::RateLimited(Some(3600)))]);
        assert!(!holds(&ctx).is_empty());

        record(&ctx, &[("acc".into(), Outcome::Answered)]);
        assert!(holds(&ctx).is_empty());
        assert_eq!(may_ask(&ctx, "acc", None, false), Some(Held::Fresh));
    }

    #[test]
    fn what_is_known_about_an_account_goes_when_the_account_does() {
        let (ctx, _clock, _s) = machine("forget");
        record(&ctx, &[("acc".into(), Outcome::RateLimited(Some(3600)))]);
        forget(&ctx, "acc");
        assert!(holds(&ctx).is_empty());
        assert_eq!(may_ask(&ctx, "acc", None, false), None);
    }
}