diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt index a102b21..f70bb30 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/Api.kt @@ -891,6 +891,10 @@ data class UsageSnapshot( val setup: String, /** That machine's current label. */ 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". * @@ -912,6 +916,8 @@ fun fetchUsage(settings: ServerSettings): List = provider = snapshot.getString("provider"), setup = snapshot.optString("setup"), 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 // draw an empty card as a healthy one. state = snapshot.optString("state").ifEmpty { "failed" }, diff --git a/app/androidApp/src/main/kotlin/com/example/aiapp/UsageDialog.kt b/app/androidApp/src/main/kotlin/com/example/aiapp/UsageDialog.kt index a95e071..88fab49 100644 --- a/app/androidApp/src/main/kotlin/com/example/aiapp/UsageDialog.kt +++ b/app/androidApp/src/main/kotlin/com/example/aiapp/UsageDialog.kt @@ -108,7 +108,7 @@ private fun UsageBody(state: LoadState>) { // read as a section of their own. Small and quiet, because the numbers // below are what somebody opened this to see. Text( - "${snapshot.setupName.ifEmpty { snapshot.setup }} · ${snapshot.provider}", + usageSectionTitle(snapshot), style = MaterialTheme.typography.bodySmall, color = MaterialTheme.colorScheme.onSurfaceVariant, ) @@ -127,6 +127,20 @@ private fun UsageBody(state: LoadState>) { } } +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. * diff --git a/server/src/resume.rs b/server/src/resume.rs index 2ef56d4..42947f7 100644 --- a/server/src/resume.rs +++ b/server/src/resume.rs @@ -269,6 +269,8 @@ mod tests { provider: crate::usage::CLAUDE.to_string(), setup: "local".to_string(), setup_name: "this machine".to_string(), + limit_id: None, + limit_name: None, state, windows, fetched_at: 0.0, diff --git a/server/src/usage.rs b/server/src/usage.rs index daa59e2..1f3482a 100644 --- a/server/src/usage.rs +++ b/server/src/usage.rs @@ -99,6 +99,12 @@ pub struct UsageSnapshot { /// That machine's current label, resolved when the snapshot is built, so /// renaming a setup renames it here too. 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, + #[serde(skip_serializing_if = "Option::is_none")] + pub limit_name: Option, #[serde(flatten)] pub state: UsageState, pub windows: Vec, @@ -120,8 +126,9 @@ pub const ECHO: &str = "echo"; pub trait UsageProvider: Send + Sync { fn name(&self) -> &'static str; - /// Blocking -- call off the async workers. - fn fetch(&self) -> UsageSnapshot; + /// Blocking -- call off the async workers. A provider may expose more than + /// one billing pool, so each answer is a separately labeled snapshot. + fn fetch(&self) -> Vec; /// How long an answer from this one may be reused. /// /// A property of the provider rather than of the cache, because what @@ -160,6 +167,8 @@ impl ClaudeUsage { provider: self.name().to_string(), setup: self.setup.clone(), setup_name: self.setup_name.clone(), + limit_id: None, + limit_name: None, state, windows, fetched_at: crate::session::now(), @@ -199,22 +208,22 @@ impl UsageProvider for ClaudeUsage { CLAUDE } - fn fetch(&self) -> UsageSnapshot { + fn fetch(&self) -> Vec { let token = match self.access_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) { Ok(body) => body, 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) { 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(), setup: self.setup.clone(), setup_name: self.setup_name.clone(), + limit_id: None, + limit_name: None, state, windows, fetched_at: crate::session::now(), @@ -246,7 +257,7 @@ impl UsageProvider for CodexUsage { CODEX } - fn fetch(&self) -> UsageSnapshot { + fn fetch(&self) -> Vec { let launch = Launch::new( &self.program, vec!["app-server".to_string(), "--stdio".to_string()], @@ -267,12 +278,12 @@ impl UsageProvider for CodexUsage { { Ok(answer) => answer, Err(err) => { - return self.snapshot( + return vec![self.snapshot( UsageState::Unreachable { detail: format!("couldn't ask Codex on {}: {err:#}", self.setup_name), }, Vec::new(), - ); + )]; } }; 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(), } }; - return self.snapshot(state, Vec::new()); + return vec![self.snapshot(state, Vec::new())]; } let Some(limits) = answer.pointer("/result/rateLimits") else { - return self.snapshot( + return vec![self.snapshot( UsageState::Failed { detail: "Codex returned no rate-limit snapshot".to_string(), }, 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 } - fn fetch(&self) -> UsageSnapshot { + fn fetch(&self) -> Vec { let (state, windows) = self .fixture .read() @@ -699,14 +738,16 @@ impl UsageProvider for EchoUsage { }, Vec::new(), )); - UsageSnapshot { + vec![UsageSnapshot { provider: self.name().to_string(), setup: self.setup.clone(), setup_name: self.setup_name.clone(), + limit_id: None, + limit_name: None, state, windows, fetched_at: crate::session::now(), - } + }] } /// 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; +type Cached = HashMap<(String, &'static str), (Instant, Vec)>; #[derive(Default)] /// 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 - /// machines are configured. + /// 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 /// 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 // 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); + 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 snapshot = provider.fetch(); + let snapshots = provider.fetch(); self.cache .lock() .unwrap() - .insert(key, (Instant::now(), snapshot.clone())); - fresh.push(snapshot); + .insert(key, (Instant::now(), snapshots.clone())); + fresh.extend(snapshots); } } // Machines that have gone away should not keep their numbers alive. @@ -936,7 +979,7 @@ mod tests { transport: Transport::for_setup(&unreachable_setup()), 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 // machine was never reached, which is not the same as a machine that // answered and has nobody logged in. @@ -1074,4 +1117,30 @@ mod tests { .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"); + } }