ai-app/server/src/usage.rs

//! Usage-limit reporting -- the same numbers as Claude Code's `/usage`.
//!
//! Polls `https://api.anthropic.com/api/oauth/usage` with the OAuth access
//! token from Claude Code's local credential store. The endpoint is
//! undocumented and has changed before, so everything here is best-effort:
//! every field is optional, and failure degrades to an "unavailable" snapshot
//! with the reason, never an error that breaks the screen.
//!
//! Two rules learned from others hitting this endpoint: send `User-Agent:
//! claude-code/<version>` (without it, requests land in an aggressively
//! rate-limited bucket) and poll no more often than every 180 s. The cache
//! below enforces the latter across any number of phone refreshes; there is no
//! background poll at all.
//!
//! One [`UsageProvider`] per paid service, so a second service later is a new
//! impl behind the same snapshot shape, not a parallel screen.
//!
//! **Asked of the machine that spends the tokens, not of this one.** A session
//! runs wherever its setup says, so the account being billed is that machine's.
//! In the layout this project aims at, `ai-server` is on the host, the host has
//! no `claude` CLI, and the CLI machine is a remote -- so the one set of numbers
//! the screen could show would be an account with no sessions. Credentials are
//! read through the session `Transport`, one snapshot per setup that offers
//! Claude.
//!
//! The token is read *to* the backend and the HTTP call is made from here, so
//! the far machine needs nothing beyond a shell and the wire format stays in
//! one place. The cost is that a remote machine's token is in this process's
//! memory for the length of a fetch, which is the same trust the backend
//! already has over that machine.

use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use serde::Serialize;
use serde_json::Value;

use crate::config::{DriverKind, SetupConfig};
use crate::session::transport::{Launch, Transport};

const USAGE_URL: &str = "https://api.anthropic.com/api/oauth/usage";
const MIN_POLL_INTERVAL: Duration = Duration::from_secs(180);
/// Matched to the CLI version the wire formats were pinned against.
const USER_AGENT: &str = "claude-code/2.1.237";

/// One rate-limit window, as the phone renders it: a labeled bar.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageWindow {
    /// The API's own word for which window this is -- `session` for the
    /// five-hour one, `weekly_all`, `weekly_scoped`, or whatever new kind it
    /// starts sending.
    ///
    /// Carried beside the label because a caller that wants one particular
    /// window has to ask for it without matching on display text: the label is
    /// written for a person and would silently select nothing the day it
    /// changes.
    pub kind: String,
    pub label: String,
    /// 0-100.
    pub percent: f64,
    /// ISO-8601, as the API sends it; absent for windows that never reset.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resets_at: Option<String>,
    /// Whether this window is currently the binding one.
    pub active: bool,
}

/// What came back when a machine was asked about its limits.
///
/// Four answers rather than a flag and a message, because the screen has to
/// treat them differently. "Nobody is logged in here" is a machine working
/// exactly as configured, while "I could not reach it" is a fault worth
/// chasing, and "the endpoint refused me" says nothing about the machine at
/// all. Collapsing them into one `error` string made the first look like the
/// last, so a perfectly healthy setup read as broken.
#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(tag = "state", rename_all = "camelCase")]
pub enum UsageState {
    /// Numbers were fetched; `windows` has them.
    Ok,
    /// The machine answered and has no Claude credentials. A choice, not a
    /// fault: nothing to report and nothing to fix.
    NotLoggedIn,
    /// The machine could not be asked at all.
    Unreachable { detail: String },
    /// The machine is logged in, but the usage endpoint did not answer.
    Failed { detail: String },
}

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageSnapshot {
    pub provider: String,
    /// Which machine these are the numbers for. The point of the whole module:
    /// they belong to an account on a particular box.
    pub setup: String,
    /// That machine's current label, resolved when the snapshot is built, so
    /// renaming a setup renames it here too.
    pub setup_name: String,
    #[serde(flatten)]
    pub state: UsageState,
    pub windows: Vec<UsageWindow>,
    /// Epoch seconds the numbers were fetched (they can be up to the poll
    /// interval old).
    pub fetched_at: f64,
}

pub trait UsageProvider: Send + Sync {
    fn name(&self) -> &'static str;
    /// Blocking -- call off the async workers.
    fn fetch(&self) -> UsageSnapshot;
}

/// Reads the numbers behind Claude Code's `/usage` from one machine, using the
/// credentials that machine stores -- nothing to configure, and it reports on
/// exactly the account whose CLI runs the sessions there.
pub struct ClaudeUsage {
    pub setup: String,
    pub setup_name: String,
    /// How to reach that machine. `Here` for the backend's own.
    pub transport: Transport,
}

/// Where Claude Code keeps its credentials, as a shell word rather than a path:
/// `$HOME` is expanded by the shell on the machine being asked, which is the
/// only place that knows what it is.
const CREDENTIALS: &str = "$HOME/.claude/.credentials.json";

impl ClaudeUsage {
    fn snapshot(&self, state: UsageState, windows: Vec<UsageWindow>) -> UsageSnapshot {
        UsageSnapshot {
            provider: self.name().to_string(),
            setup: self.setup.clone(),
            setup_name: self.setup_name.clone(),
            state,
            windows,
            fetched_at: crate::session::now(),
        }
    }

    /// The machine's stored OAuth token, or which of the two ways of not having
    /// one this is. Read through `sh -c` so `$HOME` resolves on the far machine;
    /// a path built here would be this machine's home directory.
    fn access_token(&self) -> Result<String, UsageState> {
        let launch = Launch::new(
            "sh",
            vec!["-c".to_string(), format!("cat {CREDENTIALS}")],
            None,
        );
        let text = self
            .transport
            .capture_blocking(&launch)
            .map_err(|err| why_no_credentials(&format!("{err:#}")))?;
        serde_json::from_str::<Value>(&text)
            .ok()
            .and_then(|creds| {
                creds
                    .get("claudeAiOauth")?
                    .get("accessToken")?
                    .as_str()
                    .map(String::from)
            })
            // A file that exists but carries no token is the same situation as
            // no file: nobody has logged in here yet.
            .ok_or(UsageState::NotLoggedIn)
    }
}

impl UsageProvider for ClaudeUsage {
    fn name(&self) -> &'static str {
        "claude"
    }

    fn fetch(&self) -> UsageSnapshot {
        let token = match self.access_token() {
            Ok(token) => token,
            Err(state) => return self.snapshot(state, Vec::new()),
        };
        let text = match ureq::get(USAGE_URL)
            .header("Authorization", &format!("Bearer {token}"))
            .header("anthropic-beta", "oauth-2025-04-20")
            .header("User-Agent", USER_AGENT)
            .call()
            .and_then(|mut response| response.body_mut().read_to_string())
        {
            Ok(text) => text,
            Err(err) => {
                // The error string can embed the URL but never the token.
                return self.snapshot(
                    UsageState::Failed {
                        detail: format!("usage endpoint unreachable: {err}"),
                    },
                    Vec::new(),
                );
            }
        };
        let body: Value = match serde_json::from_str(&text) {
            Ok(body) => body,
            Err(err) => {
                return self.snapshot(
                    UsageState::Failed {
                        detail: format!("usage endpoint sent non-JSON: {err}"),
                    },
                    Vec::new(),
                );
            }
        };
        self.snapshot(UsageState::Ok, parse_windows(&body))
    }
}

/// Which kind of "no credentials" a failed read was.
///
/// The distinction is the point of having both states. `cat` failing because
/// the file is not there is a machine nobody has logged in on -- a decision
/// somebody made, with nothing to fix. Anything else is a machine this server
/// could not ask, which is a fault and reads as one.
///
/// Matched on the shell's own words rather than an exit status because there is
/// only one that survives being wrapped in `sh -c` and passed back through ssh.
fn why_no_credentials(detail: &str) -> UsageState {
    // "No such file or directory" is GNU and BSD coreutils; busybox says "can't
    // open". Anything unrecognised is treated as unreachable, which is the
    // answer that gets looked at rather than ignored.
    let missing = ["No such file", "no such file", "can't open", "cannot open"];
    if missing.iter().any(|phrase| detail.contains(phrase)) {
        UsageState::NotLoggedIn
    } else {
        UsageState::Unreachable {
            detail: detail.to_string(),
        }
    }
}

/// Pulls the `limits` array apart, defensively: entries with no percent are
/// skipped, and unknown kinds keep their raw name as the label rather than
/// being dropped -- a new window appearing should show up, not vanish.
fn parse_windows(body: &Value) -> Vec<UsageWindow> {
    let Some(limits) = body.get("limits").and_then(Value::as_array) else {
        return Vec::new();
    };
    limits
        .iter()
        .filter_map(|limit| {
            let percent = limit.get("percent")?.as_f64()?;
            let kind = limit
                .get("kind")
                .and_then(Value::as_str)
                .unwrap_or("unknown");
            let scope_model = limit
                .get("scope")
                .and_then(|scope| scope.get("model"))
                .and_then(|model| model.get("display_name"))
                .and_then(Value::as_str);
            let label = match (kind, scope_model) {
                ("session", _) => "5-hour window".to_string(),
                ("weekly_all", _) => "Weekly (all models)".to_string(),
                ("weekly_scoped", Some(model)) => format!("Weekly ({model})"),
                (other, Some(model)) => format!("{other} ({model})"),
                (other, None) => other.to_string(),
            };
            Some(UsageWindow {
                kind: kind.to_string(),
                label,
                percent,
                resets_at: limit
                    .get("resets_at")
                    .and_then(Value::as_str)
                    .map(String::from),
                active: limit
                    .get("is_active")
                    .and_then(Value::as_bool)
                    .unwrap_or(false),
            })
        })
        .collect()
}

/// Which paid services a machine can be asked about.
///
/// Derived from what the setup says it can run, so a machine with no Claude
/// provider is not asked about Claude limits -- it has none, and a row saying
/// so would be a fact about nothing. A second service later adds a branch here
/// and an impl beside [`ClaudeUsage`], not a screen.
fn providers_for(setup: &SetupConfig) -> Vec<Box<dyn UsageProvider>> {
    let mut found: Vec<Box<dyn UsageProvider>> = Vec::new();
    if setup
        .providers
        .iter()
        .any(|provider| provider.kind == DriverKind::ClaudeCli)
    {
        found.push(Box::new(ClaudeUsage {
            setup: setup.id.clone(),
            setup_name: setup.name.clone(),
            transport: Transport::for_setup(setup),
        }));
    }
    found
}

/// One machine's numbers for one service, and when they were fetched.
///
/// Keyed by the machine and the service rather than by position: the set is not
/// fixed at startup -- setups are added, renamed and removed from the phone --
/// and a positional cache would hand one machine's numbers to another the
/// moment the list shifted.
type Cached = HashMap<(String, &'static str), (Instant, UsageSnapshot)>;

#[derive(Default)]
/// The cache in front of whatever machines exist: at most one real fetch per
/// machine per service per [`MIN_POLL_INTERVAL`], however often the phone asks.
pub struct UsageMonitor {
    cache: Mutex<Cached>,
}

impl UsageMonitor {
    pub fn new() -> Self {
        Self::default()
    }

    /// One snapshot per machine that offers a paid service, in the order the
    /// machines are configured.
    ///
    /// Blocking -- call via `spawn_blocking`. Takes the setups rather than
    /// holding the manager, so this module stays below the session layer rather
    /// than reaching up into it.
    pub fn snapshots(&self, setups: &[SetupConfig]) -> Vec<UsageSnapshot> {
        let mut fresh = Vec::new();
        for setup in setups {
            for provider in providers_for(setup) {
                let key = (setup.id.clone(), provider.name());
                if let Some((fetched, snapshot)) = self.cache.lock().unwrap().get(&key)
                    && fetched.elapsed() < MIN_POLL_INTERVAL
                {
                    // Cached numbers, but the machine's *name* is read fresh: a
                    // rename should show immediately rather than waiting out a
                    // poll interval it has nothing to do with.
                    let mut snapshot = snapshot.clone();
                    snapshot.setup_name = setup.name.clone();
                    fresh.push(snapshot);
                    continue;
                }
                // Fetched without the lock held: this makes a network call per
                // machine, and holding the cache across them would serialise
                // every phone asking for the screen behind the slowest ssh.
                let snapshot = provider.fetch();
                self.cache
                    .lock()
                    .unwrap()
                    .insert(key, (Instant::now(), snapshot.clone()));
                fresh.push(snapshot);
            }
        }
        // Machines that have gone away should not keep their numbers alive.
        let live: std::collections::HashSet<&str> =
            setups.iter().map(|setup| setup.id.as_str()).collect();
        self.cache
            .lock()
            .unwrap()
            .retain(|(setup, _), _| live.contains(setup.as_str()));
        fresh
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_the_limits_array_defensively() {
        // Trimmed from a live 2026-08-24 response.
        let body: Value = serde_json::from_str(
            r#"{"limits":[
                {"kind":"session","group":"session","percent":70,"severity":"normal","resets_at":"2026-08-25T04:29:59+00:00","scope":null,"is_active":true},
                {"kind":"weekly_all","group":"weekly","percent":25,"resets_at":"2026-08-28T21:59:59+00:00","is_active":false},
                {"kind":"weekly_scoped","percent":15,"resets_at":"2026-08-28T21:59:59+00:00","scope":{"model":{"id":null,"display_name":"Fable"}},"is_active":false},
                {"kind":"mystery_new_window","percent":5},
                {"kind":"broken_entry_without_percent"}
            ]}"#,
        )
        .expect("json");
        let windows = parse_windows(&body);
        assert_eq!(windows.len(), 4);
        assert_eq!(windows[0].label, "5-hour window");
        assert_eq!(windows[0].percent, 70.0);
        assert!(windows[0].active);
        assert_eq!(windows[1].label, "Weekly (all models)");
        assert_eq!(windows[2].label, "Weekly (Fable)");
        // Unknown kinds surface under their raw name instead of vanishing.
        assert_eq!(windows[3].label, "mystery_new_window");
        assert_eq!(windows[3].resets_at, None);
    }

    /// A setup naming a machine that cannot be dialled, so nothing here touches
    /// the network beyond ssh failing to resolve it.
    fn unreachable_setup() -> SetupConfig {
        SetupConfig {
            id: "far".to_string(),
            name: "somewhere else".to_string(),
            ssh: Some(crate::config::SshConfig {
                address: "no-such-host.invalid".to_string(),
                port: None,
                identity_file: None,
                options: vec!["ConnectTimeout=1".to_string()],
                attachments_dir: None,
            }),
            providers: vec![crate::config::ProviderConfig {
                name: "claude-cli".to_string(),
                kind: DriverKind::ClaudeCli,
                command: None,
                models: vec![],
            }],
        }
    }

    #[test]
    fn a_machine_that_cannot_be_asked_says_so_rather_than_looking_logged_out() {
        let provider = ClaudeUsage {
            setup: "far".to_string(),
            setup_name: "somewhere else".to_string(),
            transport: Transport::for_setup(&unreachable_setup()),
        };
        let snapshot = provider.fetch();
        // The distinction the old single `error` string could not make: this
        // machine was never reached, which is not the same as a machine that
        // answered and has nobody logged in.
        assert!(
            matches!(snapshot.state, UsageState::Unreachable { .. }),
            "{:?}",
            snapshot.state
        );
        assert_eq!(snapshot.setup, "far");
        assert_eq!(snapshot.setup_name, "somewhere else");
        assert!(snapshot.windows.is_empty());
    }

    #[test]
    fn a_missing_credential_file_is_a_choice_and_anything_else_is_a_fault() {
        // What a real shell says when nobody has logged in on that machine.
        // Nothing to fix, so it must not read as an error.
        assert_eq!(
            why_no_credentials("cat: /home/x/.claude/.credentials.json: No such file or directory"),
            UsageState::NotLoggedIn
        );
        assert_eq!(
            why_no_credentials("cat: can't open '/home/x/.claude/.credentials.json'"),
            UsageState::NotLoggedIn
        );

        // What ssh says when the machine is not there. Worth chasing, and the
        // detail is carried so somebody can.
        let refused = why_no_credentials("ssh: connect to host vm port 22: Connection refused");
        assert!(
            matches!(&refused, UsageState::Unreachable { detail } if detail.contains("refused")),
            "{refused:?}"
        );

        // Anything unrecognised errs towards the state that gets looked at,
        // rather than silently claiming nobody is logged in.
        assert!(matches!(
            why_no_credentials("something nobody has seen before"),
            UsageState::Unreachable { .. }
        ));
    }

    #[test]
    fn only_machines_that_can_run_claude_are_asked_about_it() {
        let mut echo_only = unreachable_setup();
        echo_only.providers = vec![crate::config::ProviderConfig {
            name: "echo".to_string(),
            kind: DriverKind::Echo,
            command: None,
            models: vec![],
        }];
        // A machine with no Claude on it has no Claude limits, and a row
        // reporting on it would be a fact about nothing.
        assert!(providers_for(&echo_only).is_empty());
        assert_eq!(providers_for(&unreachable_setup()).len(), 1);
    }

    #[test]
    fn an_empty_or_alien_body_yields_no_windows() {
        assert!(parse_windows(&serde_json::json!({})).is_empty());
        assert!(parse_windows(&serde_json::json!({"limits": "what"})).is_empty());
    }
}