Show separate Codex usage pools
This commit is contained in:
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.
|
||||||
*
|
*
|
||||||
|
|||||||
@@ -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
@@ -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");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Reference in new issue
Block a user