Show separate Codex usage pools

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,

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<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub limit_name: Option<String>,
     #[serde(flatten)]
     pub state: UsageState,
     pub windows: Vec<UsageWindow>,
@@ -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<UsageSnapshot>;
     /// 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<UsageSnapshot> {
         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<UsageSnapshot> {
         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<UsageSnapshot> {
         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<Box<dyn UsagePro
 /// 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)>;
+type Cached = HashMap<(String, &'static str), (Instant, Vec<UsageSnapshot>)>;
 
 #[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");
+    }
 }