iris/ai-app
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rename setups and add provider reauthentication

Browse Source
This commit is contained in:
iris committed 2026-09-12 22:56:43 -04:00
1 parent e9a0f1b9da
commit 7d9df5d572
36 files changed
+1866 -684
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/usage.rs
+196 -74
View File
@@ -16,11 +16,11 @@
//! behind the same snapshot shape.
//!
//! **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.
//! runs wherever its machine 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
//! read through the session `Transport`, one snapshot per machine that offers
//! that provider.
//!
//! The token is read *to* the backend and the HTTP call is made from here, so
@@ -36,7 +36,7 @@ use std::time::{Duration, Instant};
use serde::Serialize;
use serde_json::Value;
use crate::config::SetupConfig;
use crate::config::MachineConfig;
use crate::session::transport::{Launch, Transport};
const USAGE_URL: &str = "https://api.anthropic.com/api/oauth/usage";
@@ -72,12 +72,12 @@ pub struct UsageWindow {
/// 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
/// Named 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.
/// last, so a perfectly healthy machine read as broken.
#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(tag = "state", rename_all = "camelCase")]
pub enum UsageState {
@@ -88,6 +88,13 @@ pub enum UsageState {
NotLoggedIn,
/// The machine could not be asked at all.
Unreachable { detail: String },
/// An explicit Claude login is waiting for its browser code. Kept apart
/// from a failure because nothing is broken while somebody is completing
/// the operation on the phone.
Authenticating,
/// Credentials exist, but the CLI could not renew them. Unlike a generic
/// provider failure, this has a direct action the phone can offer.
LoginRequired { detail: String },
/// The machine is logged in, but the usage endpoint did not answer.
Failed { detail: String },
}
@@ -98,10 +105,10 @@ 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,
pub machine: String,
/// That machine's current label, resolved when the snapshot is built, so
/// renaming a setup renames it here too.
pub setup_name: String,
/// renaming a machine renames it here too.
pub machine_name: String,
/// The provider's billing pool, when it exposes more than one. Codex uses
/// this to keep regular usage separate from its Luna reserve allowance.
#[serde(skip_serializing_if = "Option::is_none")]
@@ -149,8 +156,8 @@ pub trait UsageProvider: Send + Sync {
/// 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,
pub machine: String,
pub machine_name: String,
/// How to reach that machine. `Here` for the backend's own.
pub transport: Transport,
/// The CLI to run there, for the one thing this asks of it: refreshing its
@@ -168,8 +175,8 @@ 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(),
machine: self.machine.clone(),
machine_name: self.machine_name.clone(),
limit_id: None,
limit_name: None,
state,
@@ -234,8 +241,8 @@ impl UsageProvider for ClaudeUsage {
/// app-server protocol. The CLI owns authentication and token refresh; this
/// process never opens or copies its credentials.
pub struct CodexUsage {
pub setup: String,
pub setup_name: String,
pub machine: String,
pub machine_name: String,
pub transport: Transport,
pub program: String,
}
@@ -244,8 +251,8 @@ impl CodexUsage {
fn snapshot(&self, state: UsageState, windows: Vec<UsageWindow>) -> UsageSnapshot {
UsageSnapshot {
provider: CODEX.to_string(),
setup: self.setup.clone(),
setup_name: self.setup_name.clone(),
machine: self.machine.clone(),
machine_name: self.machine_name.clone(),
limit_id: None,
limit_name: None,
state,
@@ -283,7 +290,7 @@ impl UsageProvider for CodexUsage {
Err(err) => {
return vec![self.snapshot(
UsageState::Unreachable {
detail: format!("couldn't ask Codex on {}: {err:#}", self.setup_name),
detail: format!("couldn't ask Codex on {}: {err:#}", self.machine_name),
},
Vec::new(),
)];
@@ -446,7 +453,7 @@ impl ClaudeUsage {
return Err(UsageState::Failed {
detail: format!(
"the Claude login on {} has expired, and `{} doctor` couldn't be run there to refresh it: {err:#}",
self.setup_name, self.program
self.machine_name, self.program
),
});
}
@@ -463,10 +470,10 @@ impl ClaudeUsage {
/// A login the CLI could not renew: the one state here somebody has to act
/// on, so it says where and what to run.
fn still_expired(&self) -> UsageState {
UsageState::Failed {
UsageState::LoginRequired {
detail: format!(
"the Claude login on {} has expired and could not be refreshed; run `{} /login` there",
self.setup_name, self.program
"The Claude login on {} has expired and could not be refreshed.",
self.machine_name
),
}
}
@@ -727,6 +734,8 @@ fn describe(state: &UsageState, windows: &[UsageWindow]) -> String {
},
UsageState::NotLoggedIn => "nobody is logged in on this machine".to_string(),
UsageState::Unreachable { detail } => format!("machine unreachable ({detail})"),
UsageState::Authenticating => "sign-in is in progress".to_string(),
UsageState::LoginRequired { detail } => format!("sign-in required ({detail})"),
UsageState::Failed { detail } => format!("the meter failed ({detail})"),
}
}
@@ -734,8 +743,8 @@ fn describe(state: &UsageState, windows: &[UsageWindow]) -> String {
/// The fixture, as a provider, so it travels the same route and the same
/// cache as a real meter rather than being spliced in at the screen.
struct EchoUsage {
setup: String,
setup_name: String,
machine: String,
machine_name: String,
fixture: Fixture,
}
@@ -759,8 +768,8 @@ impl UsageProvider for EchoUsage {
));
vec![UsageSnapshot {
provider: self.name().to_string(),
setup: self.setup.clone(),
setup_name: self.setup_name.clone(),
machine: self.machine.clone(),
machine_name: self.machine_name.clone(),
limit_id: None,
limit_name: None,
state,
@@ -778,7 +787,7 @@ impl UsageProvider for EchoUsage {
/// Which paid services a machine can be asked about.
///
/// Derived from what the setup says it can run, so a machine with no Claude
/// Derived from what the machine 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.
///
@@ -787,9 +796,9 @@ impl UsageProvider for EchoUsage {
/// one of these rows by that same name: two lists that disagreed would leave a
/// session looking for a snapshot nothing produces. A second service later is a
/// name there and an impl beside [`ClaudeUsage`], not a screen.
fn providers_for(setup: &SetupConfig, fixture: &Fixture) -> Vec<Box<dyn UsageProvider>> {
fn providers_for(machine: &MachineConfig, fixture: &Fixture) -> Vec<Box<dyn UsageProvider>> {
let mut found: Vec<Box<dyn UsageProvider>> = Vec::new();
for provider in &setup.providers {
for provider in &machine.providers {
let Some(name) = provider.kind.usage_provider() else {
continue;
};
@@ -801,23 +810,23 @@ fn providers_for(setup: &SetupConfig, fixture: &Fixture) -> Vec<Box<dyn UsagePro
}
match name {
CLAUDE => found.push(Box::new(ClaudeUsage {
setup: setup.id.clone(),
setup_name: setup.name.clone(),
transport: Transport::for_setup(setup),
machine: machine.id.clone(),
machine_name: machine.name.clone(),
transport: Transport::for_machine(machine),
program: provider.program().to_string(),
})),
CODEX => found.push(Box::new(CodexUsage {
setup: setup.id.clone(),
setup_name: setup.name.clone(),
transport: Transport::for_setup(setup),
machine: machine.id.clone(),
machine_name: machine.name.clone(),
transport: Transport::for_machine(machine),
program: provider.program().to_string(),
})),
// Nothing at all until a test has asked for something: an
// echo session costs nothing, so the honest answer is no row
// rather than a row saying zero.
ECHO if fixture.is_set() => found.push(Box::new(EchoUsage {
setup: setup.id.clone(),
setup_name: setup.name.clone(),
machine: machine.id.clone(),
machine_name: machine.name.clone(),
fixture: fixture.clone(),
})),
_ => {}
@@ -829,16 +838,22 @@ fn providers_for(setup: &SetupConfig, fixture: &Fixture) -> Vec<Box<dyn UsagePro
/// 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 --
/// fixed at startup -- machines 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, Vec<UsageSnapshot>)>;
type ProviderGate = Arc<Mutex<()>>;
#[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>,
/// One gate per machine and metered provider. The cache lock cannot cover
/// a network call without making an unreachable machine stall every other
/// one, but releasing it used to let concurrent `/usage` requests launch
/// two token refreshers against the same rotating credential.
gates: Mutex<HashMap<(String, &'static str), ProviderGate>>,
/// The invented meter an echo session can put up; empty unless one
/// has. Shared with the session layer, which is where the command
/// that sets it is typed -- see [`Fixture`].
@@ -849,37 +864,116 @@ impl UsageMonitor {
pub fn new(fixture: Fixture) -> Self {
Self {
cache: Mutex::new(Cached::new()),
gates: Mutex::new(HashMap::new()),
fixture,
}
}
fn gate_for(&self, key: &(String, &'static str)) -> ProviderGate {
self.gates
.lock()
.unwrap()
.entry(key.clone())
.or_default()
.clone()
}
fn cached(
&self,
key: &(String, &'static str),
machine_name: &str,
max_age: Option<Duration>,
) -> Option<Vec<UsageSnapshot>> {
let cache = self.cache.lock().unwrap();
let (fetched, snapshots) = cache.get(key)?;
if max_age.is_some_and(|age| fetched.elapsed() >= age) {
return None;
}
let mut snapshots = snapshots.clone();
for snapshot in &mut snapshots {
snapshot.machine_name = machine_name.to_string();
}
Some(snapshots)
}
/// Prevents automatic refresh from overlapping an explicit login. The
/// fresh cache entry closes the small gap before the login worker acquires
/// the same gate, while an attempt lasting beyond the normal cache window
/// is protected by the worker holding it.
pub(crate) fn claude_authentication_started(&self, machine: &MachineConfig) -> ProviderGate {
let key = (machine.id.clone(), CLAUDE);
let gate = self.gate_for(&key);
let guard = gate.lock().unwrap();
self.cache.lock().unwrap().insert(
key,
(
Instant::now(),
vec![UsageSnapshot {
provider: CLAUDE.to_string(),
machine: machine.id.clone(),
machine_name: machine.name.clone(),
limit_id: None,
limit_name: None,
state: UsageState::Authenticating,
windows: Vec::new(),
fetched_at: crate::session::now(),
}],
),
);
drop(guard);
gate
}
/// Makes the first read after a login ask the provider instead of keeping
/// the pre-login state for the ordinary cache interval.
pub(crate) fn claude_authentication_finished(&self, machine: &str) {
self.cache
.lock()
.unwrap()
.remove(&(machine.to_string(), CLAUDE));
}
/// One or more snapshots per machine that offers a paid service, in the
/// order the machines are configured and the provider reports them.
///
/// Blocking -- call via `spawn_blocking`. Takes the setups rather than
/// Blocking -- call via `spawn_blocking`. Takes the machines 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> {
pub fn snapshots(&self, machines: &[MachineConfig]) -> Vec<UsageSnapshot> {
let mut fresh = Vec::new();
for setup in setups {
for provider in providers_for(setup, &self.fixture) {
let key = (setup.id.clone(), provider.name());
if let Some((fetched, snapshot)) = self.cache.lock().unwrap().get(&key)
&& fetched.elapsed() < provider.poll_interval()
for machine in machines {
for provider in providers_for(machine, &self.fixture) {
let key = (machine.id.clone(), provider.name());
if let Some(snapshots) =
self.cached(&key, &machine.name, Some(provider.poll_interval()))
{
fresh.extend(snapshots);
continue;
}
let gate = self.gate_for(&key);
let _guard = match gate.try_lock() {
Ok(guard) => guard,
Err(std::sync::TryLockError::WouldBlock) => {
// A refresh already in flight can keep serving the last
// measured answer. An explicit login seeds its own
// authenticating answer before taking the gate.
if let Some(snapshots) = self.cached(&key, &machine.name, None) {
fresh.extend(snapshots);
continue;
}
// The first-ever fetch has no honest stale answer. Wait
// for its one producer, then recheck below.
gate.lock().unwrap()
}
Err(std::sync::TryLockError::Poisoned(err)) => err.into_inner(),
};
if let Some(snapshots) =
self.cached(&key, &machine.name, Some(provider.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 snapshots = snapshot.clone();
for snapshot in &mut snapshots {
snapshot.setup_name = setup.name.clone();
}
fresh.extend(snapshots);
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 snapshots = provider.fetch();
self.cache
.lock()
@@ -890,11 +984,14 @@ impl UsageMonitor {
}
// 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();
machines.iter().map(|machine| machine.id.as_str()).collect();
self.cache
.lock()
.unwrap()
.retain(|(setup, _), _| live.contains(setup.as_str()));
.retain(|(machine, _), _| live.contains(machine.as_str()));
self.gates.lock().unwrap().retain(|(machine, _), gate| {
live.contains(machine.as_str()) || Arc::strong_count(gate) > 1
});
fresh
}
}
@@ -904,6 +1001,31 @@ mod tests {
use super::*;
use crate::config::DriverKind;
#[test]
fn refresh_gates_are_shared_per_machine_and_provider_only() {
let monitor = UsageMonitor::new(Fixture::default());
let claude_here = monitor.gate_for(&("here".to_string(), CLAUDE));
let same = monitor.gate_for(&("here".to_string(), CLAUDE));
let claude_there = monitor.gate_for(&("there".to_string(), CLAUDE));
let codex_here = monitor.gate_for(&("here".to_string(), CODEX));
let held = claude_here.lock().unwrap();
assert!(
matches!(same.try_lock(), Err(std::sync::TryLockError::WouldBlock)),
"the same machine/provider pair must have one credential writer"
);
assert!(
claude_there.try_lock().is_ok(),
"another machine must not wait"
);
assert!(
codex_here.try_lock().is_ok(),
"another provider must not wait"
);
drop(held);
assert!(same.try_lock().is_ok());
}
#[test]
fn a_refusal_the_endpoint_answered_is_not_reported_as_an_unreachable_one() {
assert!(why(&ureq::Error::StatusCode(500)).contains("HTTP 500"));
@@ -913,9 +1035,9 @@ mod tests {
#[test]
fn an_expired_login_says_where_to_log_in_rather_than_naming_the_network() {
let provider = ClaudeUsage {
setup: "far".to_string(),
setup_name: "somewhere else".to_string(),
transport: Transport::for_setup(&unreachable_setup()),
machine: "far".to_string(),
machine_name: "somewhere else".to_string(),
transport: Transport::for_machine(&unreachable_machine()),
program: "/opt/claude".to_string(),
};
// The machine cannot be reached, so the refresh attempt fails there
@@ -934,10 +1056,10 @@ mod tests {
"the token must never be quoted back"
);
let UsageState::Failed { detail } = provider.still_expired() else {
panic!("still expired is a fault");
let UsageState::LoginRequired { detail } = provider.still_expired() else {
panic!("still expired must offer a new login");
};
assert!(detail.contains("/opt/claude /login"), "{detail}");
assert!(detail.contains("somewhere else"), "{detail}");
assert!(!detail.contains("unreachable"), "{detail}");
}
@@ -967,10 +1089,10 @@ mod tests {
assert_eq!(windows[3].resets_at, None);
}
/// A setup naming a machine that cannot be dialled, so nothing here touches
/// A machine naming a machine that cannot be dialled, so nothing here touches
/// the network beyond ssh failing to resolve it.
fn unreachable_setup() -> SetupConfig {
SetupConfig {
fn unreachable_machine() -> MachineConfig {
MachineConfig {
id: "far".to_string(),
name: "somewhere else".to_string(),
ssh: Some(crate::config::SshConfig {
@@ -993,9 +1115,9 @@ mod tests {
#[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()),
machine: "far".to_string(),
machine_name: "somewhere else".to_string(),
transport: Transport::for_machine(&unreachable_machine()),
program: "claude".to_string(),
};
let snapshot = provider.fetch().remove(0);
@@ -1007,8 +1129,8 @@ mod tests {
"{:?}",
snapshot.state
);
assert_eq!(snapshot.setup, "far");
assert_eq!(snapshot.setup_name, "somewhere else");
assert_eq!(snapshot.machine, "far");
assert_eq!(snapshot.machine_name, "somewhere else");
assert!(snapshot.windows.is_empty());
}
@@ -1043,7 +1165,7 @@ mod tests {
#[test]
fn only_machines_that_can_run_claude_are_asked_about_it() {
let mut echo_only = unreachable_setup();
let mut echo_only = unreachable_machine();
echo_only.providers = vec![crate::config::ProviderConfig {
name: "echo".to_string(),
kind: DriverKind::Echo,
@@ -1056,7 +1178,7 @@ mod tests {
// otherwise there is no meter to report.
let unset = Fixture::new();
assert!(providers_for(&echo_only, &unset).is_empty());
assert_eq!(providers_for(&unreachable_setup(), &unset).len(), 1);
assert_eq!(providers_for(&unreachable_machine(), &unset).len(), 1);
// And with one set, that machine has exactly the invented meter
// -- under the name the session's `usageProvider` will name.
@@ -1142,8 +1264,8 @@ mod tests {
#[test]
fn keeps_codex_reserve_as_a_named_pool() {
let provider = CodexUsage {
setup: "local".to_string(),
setup_name: "this machine".to_string(),
machine: "local".to_string(),
machine_name: "this machine".to_string(),
transport: Transport::Here,
program: "codex".to_string(),
};