Show separate Codex usage pools

This commit is contained in:
iris committed 2026-09-08 00:41:06 -04:00
1 parent 7ee88dfd9c
commit 00538cc19b
4 files changed
+118 -27

No files matched your search

@@ -891,6 +891,10 @@ data class UsageSnapshot(
val setup: String, val setup: String,
/** That machine's current label. */ /** That machine's current label. */
val setupName: String, val setupName: String,
/** Provider-specific billing pool, such as Codex's regular or Luna Reserve pool. */
val limitId: String?,
/** Provider-specific human-facing pool name, when supplied. */
val limitName: String?,
/** /**
* What came back: "ok", "notLoggedIn", "unreachable" or "failed". * What came back: "ok", "notLoggedIn", "unreachable" or "failed".
* *
@@ -912,6 +916,8 @@ fun fetchUsage(settings: ServerSettings): List<UsageSnapshot> =
provider = snapshot.getString("provider"), provider = snapshot.getString("provider"),
setup = snapshot.optString("setup"), setup = snapshot.optString("setup"),
setupName = snapshot.optString("setupName"), setupName = snapshot.optString("setupName"),
limitId = snapshot.optString("limitId").ifEmpty { null },
limitName = snapshot.optString("limitName").ifEmpty { null },
// Unknown to an older backend, and unknown is not "fine": defaulting to "ok" would // Unknown to an older backend, and unknown is not "fine": defaulting to "ok" would
// draw an empty card as a healthy one. // draw an empty card as a healthy one.
state = snapshot.optString("state").ifEmpty { "failed" }, state = snapshot.optString("state").ifEmpty { "failed" },
@@ -108,7 +108,7 @@ private fun UsageBody(state: LoadState<List<UsageSnapshot>>) {
// read as a section of their own. Small and quiet, because the numbers // read as a section of their own. Small and quiet, because the numbers
// below are what somebody opened this to see. // below are what somebody opened this to see.
Text( Text(
"${snapshot.setupName.ifEmpty { snapshot.setup }} · ${snapshot.provider}", usageSectionTitle(snapshot),
style = MaterialTheme.typography.bodySmall, style = MaterialTheme.typography.bodySmall,
color = MaterialTheme.colorScheme.onSurfaceVariant, color = MaterialTheme.colorScheme.onSurfaceVariant,
) )
@@ -127,6 +127,20 @@ private fun UsageBody(state: LoadState<List<UsageSnapshot>>) {
} }
} }
private fun usageSectionTitle(snapshot: UsageSnapshot): String {
val machine = snapshot.setupName.ifEmpty { snapshot.setup }
val provider = snapshot.provider
val pool =
if (provider == "codex" && snapshot.limitId != "codex") {
when (snapshot.limitName) {
"gpt-reserve" -> "Luna Reserve"
null -> snapshot.limitId
else -> snapshot.limitName
}
} else null
return listOfNotNull(machine, provider, pool).joinToString(" · ")
}
/** /**
* Anything other than numbers: why this machine has none. * Anything other than numbers: why this machine has none.
* *
+2
View File
@@ -269,6 +269,8 @@ mod tests {
provider: crate::usage::CLAUDE.to_string(), provider: crate::usage::CLAUDE.to_string(),
setup: "local".to_string(), setup: "local".to_string(),
setup_name: "this machine".to_string(), setup_name: "this machine".to_string(),
limit_id: None,
limit_name: None,
state, state,
windows, windows,
fetched_at: 0.0, fetched_at: 0.0,
+95 -26
View File
@@ -99,6 +99,12 @@ pub struct UsageSnapshot {
/// That machine's current label, resolved when the snapshot is built, so /// That machine's current label, resolved when the snapshot is built, so
/// renaming a setup renames it here too. /// renaming a setup renames it here too.
pub setup_name: String, pub setup_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")]
pub limit_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub limit_name: Option<String>,
#[serde(flatten)] #[serde(flatten)]
pub state: UsageState, pub state: UsageState,
pub windows: Vec<UsageWindow>, pub windows: Vec<UsageWindow>,
@@ -120,8 +126,9 @@ pub const ECHO: &str = "echo";
pub trait UsageProvider: Send + Sync { pub trait UsageProvider: Send + Sync {
fn name(&self) -> &'static str; fn name(&self) -> &'static str;
/// Blocking -- call off the async workers. /// Blocking -- call off the async workers. A provider may expose more than
fn fetch(&self) -> UsageSnapshot; /// one billing pool, so each answer is a separately labeled snapshot.
fn fetch(&self) -> Vec<UsageSnapshot>;
/// How long an answer from this one may be reused. /// How long an answer from this one may be reused.
/// ///
/// A property of the provider rather than of the cache, because what /// A property of the provider rather than of the cache, because what
@@ -160,6 +167,8 @@ impl ClaudeUsage {
provider: self.name().to_string(), provider: self.name().to_string(),
setup: self.setup.clone(), setup: self.setup.clone(),
setup_name: self.setup_name.clone(), setup_name: self.setup_name.clone(),
limit_id: None,
limit_name: None,
state, state,
windows, windows,
fetched_at: crate::session::now(), fetched_at: crate::session::now(),
@@ -199,22 +208,22 @@ impl UsageProvider for ClaudeUsage {
CLAUDE CLAUDE
} }
fn fetch(&self) -> UsageSnapshot { fn fetch(&self) -> Vec<UsageSnapshot> {
let token = match self.access_token() { let token = match self.access_token() {
Ok(token) => token, Ok(token) => token,
Err(state) => return self.snapshot(state, Vec::new()), Err(state) => return vec![self.snapshot(state, Vec::new())],
}; };
let body = match self.call(&token) { let body = match self.call(&token) {
Ok(body) => body, Ok(body) => body,
Err(Refused::Other(detail)) => { Err(Refused::Other(detail)) => {
return self.snapshot(UsageState::Failed { detail }, Vec::new()); return vec![self.snapshot(UsageState::Failed { detail }, Vec::new())];
} }
Err(Refused::Unauthorized) => match self.after_cli_refresh(&token) { Err(Refused::Unauthorized) => match self.after_cli_refresh(&token) {
Ok(body) => body, Ok(body) => body,
Err(state) => return self.snapshot(state, Vec::new()), Err(state) => return vec![self.snapshot(state, Vec::new())],
}, },
}; };
self.snapshot(UsageState::Ok, parse_windows(&body)) vec![self.snapshot(UsageState::Ok, parse_windows(&body))]
} }
} }
@@ -234,6 +243,8 @@ impl CodexUsage {
provider: CODEX.to_string(), provider: CODEX.to_string(),
setup: self.setup.clone(), setup: self.setup.clone(),
setup_name: self.setup_name.clone(), setup_name: self.setup_name.clone(),
limit_id: None,
limit_name: None,
state, state,
windows, windows,
fetched_at: crate::session::now(), fetched_at: crate::session::now(),
@@ -246,7 +257,7 @@ impl UsageProvider for CodexUsage {
CODEX CODEX
} }
fn fetch(&self) -> UsageSnapshot { fn fetch(&self) -> Vec<UsageSnapshot> {
let launch = Launch::new( let launch = Launch::new(
&self.program, &self.program,
vec!["app-server".to_string(), "--stdio".to_string()], vec!["app-server".to_string(), "--stdio".to_string()],
@@ -267,12 +278,12 @@ impl UsageProvider for CodexUsage {
{ {
Ok(answer) => answer, Ok(answer) => answer,
Err(err) => { Err(err) => {
return self.snapshot( return vec![self.snapshot(
UsageState::Unreachable { UsageState::Unreachable {
detail: format!("couldn't ask Codex on {}: {err:#}", self.setup_name), detail: format!("couldn't ask Codex on {}: {err:#}", self.setup_name),
}, },
Vec::new(), Vec::new(),
); )];
} }
}; };
if let Some(error) = answer.pointer("/error/message").and_then(Value::as_str) { if let Some(error) = answer.pointer("/error/message").and_then(Value::as_str) {
@@ -285,17 +296,45 @@ impl UsageProvider for CodexUsage {
detail: error.to_string(), detail: error.to_string(),
} }
}; };
return self.snapshot(state, Vec::new()); return vec![self.snapshot(state, Vec::new())];
} }
let Some(limits) = answer.pointer("/result/rateLimits") else { let Some(limits) = answer.pointer("/result/rateLimits") else {
return self.snapshot( return vec![self.snapshot(
UsageState::Failed { UsageState::Failed {
detail: "Codex returned no rate-limit snapshot".to_string(), detail: "Codex returned no rate-limit snapshot".to_string(),
}, },
Vec::new(), Vec::new(),
); )];
}; };
self.snapshot(UsageState::Ok, parse_codex_windows(limits)) let mut snapshots = vec![self.snapshot_for_limit("codex", None, limits)];
if let Some(pools) = answer
.pointer("/result/rateLimitsByLimitId")
.and_then(Value::as_object)
{
for (key, pool) in pools {
let limit_id = pool.get("limitId").and_then(Value::as_str).unwrap_or(key);
if limit_id == "codex" {
continue;
}
let limit_name = pool.get("limitName").and_then(Value::as_str);
snapshots.push(self.snapshot_for_limit(limit_id, limit_name, pool));
}
}
snapshots
}
}
impl CodexUsage {
fn snapshot_for_limit(
&self,
limit_id: &str,
limit_name: Option<&str>,
limits: &Value,
) -> UsageSnapshot {
let mut snapshot = self.snapshot(UsageState::Ok, parse_codex_windows(limits));
snapshot.limit_id = Some(limit_id.to_string());
snapshot.limit_name = limit_name.map(str::to_string);
snapshot
} }
} }
@@ -686,7 +725,7 @@ impl UsageProvider for EchoUsage {
ECHO ECHO
} }
fn fetch(&self) -> UsageSnapshot { fn fetch(&self) -> Vec<UsageSnapshot> {
let (state, windows) = self let (state, windows) = self
.fixture .fixture
.read() .read()
@@ -699,14 +738,16 @@ impl UsageProvider for EchoUsage {
}, },
Vec::new(), Vec::new(),
)); ));
UsageSnapshot { vec![UsageSnapshot {
provider: self.name().to_string(), provider: self.name().to_string(),
setup: self.setup.clone(), setup: self.setup.clone(),
setup_name: self.setup_name.clone(), setup_name: self.setup_name.clone(),
limit_id: None,
limit_name: None,
state, state,
windows, windows,
fetched_at: crate::session::now(), fetched_at: crate::session::now(),
} }]
} }
/// Read from memory, and set by somebody who is about to look at the /// Read from memory, and set by somebody who is about to look at the
@@ -772,7 +813,7 @@ fn providers_for(setup: &SetupConfig, fixture: &Fixture) -> Vec<Box<dyn UsagePro
/// fixed at startup -- setups are added, renamed and removed from the phone -- /// 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 /// and a positional cache would hand one machine's numbers to another the
/// moment the list shifted. /// moment the list shifted.
type Cached = HashMap<(String, &'static str), (Instant, UsageSnapshot)>; type Cached = HashMap<(String, &'static str), (Instant, Vec<UsageSnapshot>)>;
#[derive(Default)] #[derive(Default)]
/// The cache in front of whatever machines exist: at most one real fetch per /// The cache in front of whatever machines exist: at most one real fetch per
@@ -793,8 +834,8 @@ impl UsageMonitor {
} }
} }
/// One snapshot per machine that offers a paid service, in the order the /// One or more snapshots per machine that offers a paid service, in the
/// machines are configured. /// 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 setups rather than
/// holding the manager, so this module stays below the session layer rather /// holding the manager, so this module stays below the session layer rather
@@ -810,20 +851,22 @@ impl UsageMonitor {
// Cached numbers, but the machine's *name* is read fresh: a // Cached numbers, but the machine's *name* is read fresh: a
// rename should show immediately rather than waiting out a // rename should show immediately rather than waiting out a
// poll interval it has nothing to do with. // poll interval it has nothing to do with.
let mut snapshot = snapshot.clone(); let mut snapshots = snapshot.clone();
for snapshot in &mut snapshots {
snapshot.setup_name = setup.name.clone(); snapshot.setup_name = setup.name.clone();
fresh.push(snapshot); }
fresh.extend(snapshots);
continue; continue;
} }
// Fetched without the lock held: this makes a network call per // Fetched without the lock held: this makes a network call per
// machine, and holding the cache across them would serialise // machine, and holding the cache across them would serialise
// every phone asking for the screen behind the slowest ssh. // every phone asking for the screen behind the slowest ssh.
let snapshot = provider.fetch(); let snapshots = provider.fetch();
self.cache self.cache
.lock() .lock()
.unwrap() .unwrap()
.insert(key, (Instant::now(), snapshot.clone())); .insert(key, (Instant::now(), snapshots.clone()));
fresh.push(snapshot); fresh.extend(snapshots);
} }
} }
// Machines that have gone away should not keep their numbers alive. // Machines that have gone away should not keep their numbers alive.
@@ -936,7 +979,7 @@ mod tests {
transport: Transport::for_setup(&unreachable_setup()), transport: Transport::for_setup(&unreachable_setup()),
program: "claude".to_string(), program: "claude".to_string(),
}; };
let snapshot = provider.fetch(); let snapshot = provider.fetch().remove(0);
// The distinction the old single `error` string could not make: this // The distinction the old single `error` string could not make: this
// machine was never reached, which is not the same as a machine that // machine was never reached, which is not the same as a machine that
// answered and has nobody logged in. // answered and has nobody logged in.
@@ -1074,4 +1117,30 @@ mod tests {
.is_some_and(|at| at.ends_with('Z')) .is_some_and(|at| at.ends_with('Z'))
); );
} }
#[test]
fn keeps_codex_reserve_as_a_named_pool() {
let provider = CodexUsage {
setup: "local".to_string(),
setup_name: "this machine".to_string(),
transport: Transport::Here,
program: "codex".to_string(),
};
let snapshot = provider.snapshot_for_limit(
"base_model_inference",
Some("gpt-reserve"),
&serde_json::json!({
"primary": {
"usedPercent": 12,
"windowDurationMins": 10080,
"resetsAt": 1789446343_i64
},
"secondary": null
}),
);
assert_eq!(snapshot.limit_id.as_deref(), Some("base_model_inference"));
assert_eq!(snapshot.limit_name.as_deref(), Some("gpt-reserve"));
assert_eq!(snapshot.windows.len(), 1);
assert_eq!(snapshot.windows[0].label, "Weekly");
}
} }