feat: show spend controls in tui status

This commit is contained in:
Edward Frazer
2026-05-15 23:45:26 -07:00
parent d322741267
commit ad1facce2b
4 changed files with 292 additions and 0 deletions
@@ -112,6 +112,7 @@ pub(super) fn snapshot(percent: f64) -> RateLimitSnapshot {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
}
@@ -531,6 +531,7 @@ async fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
unlimited: false,
balance: Some("17.5".to_string()),
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
}));
@@ -551,6 +552,7 @@ async fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
}));
@@ -590,6 +592,7 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() {
resets_at: None,
}),
credits: None,
spend_control: None,
plan_type: Some(PlanType::Plus),
rate_limit_reached_type: None,
}));
@@ -609,6 +612,7 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() {
resets_at: Some(234),
}),
credits: None,
spend_control: None,
plan_type: Some(PlanType::Pro),
rate_limit_reached_type: None,
}));
@@ -628,6 +632,7 @@ async fn rate_limit_snapshot_updates_and_retains_plan_type() {
resets_at: Some(567),
}),
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
}));
@@ -652,6 +657,7 @@ async fn rate_limit_snapshots_keep_separate_entries_per_limit_id() {
unlimited: false,
balance: Some("5.00".to_string()),
}),
spend_control: None,
plan_type: Some(PlanType::Pro),
rate_limit_reached_type: None,
}));
@@ -666,6 +672,7 @@ async fn rate_limit_snapshots_keep_separate_entries_per_limit_id() {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: Some(PlanType::Pro),
rate_limit_reached_type: None,
}));
@@ -719,6 +726,7 @@ async fn rate_limit_switch_prompt_skips_non_codex_limit() {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
}));
+186
View File
@@ -14,8 +14,12 @@ use chrono::Duration as ChronoDuration;
use chrono::Local;
use chrono::Utc;
use codex_app_server_protocol::CreditsSnapshot as CoreCreditsSnapshot;
use codex_app_server_protocol::GroupSpendControlLimitSnapshot;
use codex_app_server_protocol::RateLimitSnapshot;
use codex_app_server_protocol::RateLimitWindow;
use codex_app_server_protocol::SpendControlLimitSnapshot;
use codex_app_server_protocol::SpendControlLimitType;
use codex_app_server_protocol::SpendControlSnapshot;
const STATUS_LIMIT_BAR_SEGMENTS: usize = 20;
const STATUS_LIMIT_BAR_FILLED: &str = "█";
@@ -98,6 +102,8 @@ pub(crate) struct RateLimitSnapshotDisplay {
pub secondary: Option<RateLimitWindowDisplay>,
/// Optional credits metadata when available.
pub credits: Option<CreditsSnapshotDisplay>,
/// Optional spend-control metadata when available.
pub spend_control: Option<SpendControlSnapshotDisplay>,
}
/// Display-ready credits state extracted from protocol snapshots.
@@ -111,6 +117,28 @@ pub(crate) struct CreditsSnapshotDisplay {
pub balance: Option<String>,
}
/// Display-ready spend-control state extracted from protocol snapshots.
#[derive(Debug, Clone)]
pub(crate) struct SpendControlSnapshotDisplay {
/// Whether the backend says a spend control blocked the current request path.
pub reached: bool,
/// The concrete control type that was reached, when one is known.
pub reached_limit_type: Option<SpendControlLimitType>,
/// The most relevant active limit to summarize for `/status`.
pub active_limit: Option<SpendControlLimitDisplay>,
}
/// Display-ready view of one configured spend-control limit.
#[derive(Debug, Clone)]
pub(crate) struct SpendControlLimitDisplay {
/// Friendly control label for status output.
pub label: String,
/// Percent remaining according to the backend snapshot.
pub remaining_percent: f64,
/// Human-readable local reset time.
pub resets_at: Option<String>,
}
/// Converts a protocol snapshot into UI-friendly display data.
///
/// Pass the timestamp from the same observation point as `snapshot`; supplying a significantly
@@ -140,6 +168,9 @@ pub(crate) fn rate_limit_snapshot_display_for_limit(
.as_ref()
.map(|window| RateLimitWindowDisplay::from_window(window, captured_at)),
credits: snapshot.credits.as_ref().map(CreditsSnapshotDisplay::from),
spend_control: snapshot.spend_control.as_ref().map(|spend_control| {
SpendControlSnapshotDisplay::from_snapshot(spend_control, captured_at)
}),
}
}
@@ -153,6 +184,16 @@ impl From<&CoreCreditsSnapshot> for CreditsSnapshotDisplay {
}
}
impl SpendControlSnapshotDisplay {
fn from_snapshot(snapshot: &SpendControlSnapshot, captured_at: DateTime<Local>) -> Self {
Self {
reached: snapshot.reached,
reached_limit_type: snapshot.reached_limit_type,
active_limit: active_spend_control_limit(snapshot, captured_at),
}
}
}
/// Builds display rows from a snapshot and marks stale data by capture age.
///
/// Callers should pass `Local::now()` for `now` at render time; using a cached timestamp can make
@@ -268,6 +309,9 @@ pub(crate) fn compose_rate_limit_data_many(
{
rows.push(row);
}
if let Some(spend_control) = snapshot.spend_control.as_ref() {
rows.extend(spend_control_status_rows(spend_control));
}
}
if rows.is_empty() {
@@ -279,6 +323,107 @@ pub(crate) fn compose_rate_limit_data_many(
}
}
fn spend_control_status_rows(
spend_control: &SpendControlSnapshotDisplay,
) -> Vec<StatusRateLimitRow> {
let mut rows = Vec::new();
if let Some(active_limit) = spend_control.active_limit.as_ref() {
rows.push(StatusRateLimitRow {
label: active_limit.label.clone(),
value: StatusRateLimitValue::Window {
percent_used: (100.0 - active_limit.remaining_percent).clamp(0.0, 100.0),
resets_at: active_limit.resets_at.clone(),
},
});
}
if spend_control.reached
&& let Some(reached_limit_type) = spend_control.reached_limit_type
{
rows.push(StatusRateLimitRow {
label: "Spend control reached".to_string(),
value: StatusRateLimitValue::Text(spend_control_limit_type_label(reached_limit_type)),
});
}
rows
}
fn active_spend_control_limit(
snapshot: &SpendControlSnapshot,
captured_at: DateTime<Local>,
) -> Option<SpendControlLimitDisplay> {
snapshot
.individual_limit
.as_ref()
.map(|limit| spend_control_limit_display("Individual spend control", limit, captured_at))
.or_else(|| {
snapshot.group_default_limit.as_ref().map(|limit| {
group_spend_control_limit_display("Group default spend control", limit, captured_at)
})
})
.or_else(|| {
snapshot.workspace_default_limit.as_ref().map(|limit| {
spend_control_limit_display("Workspace default spend control", limit, captured_at)
})
})
.or_else(|| {
snapshot
.role_budget_limit
.as_ref()
.map(|limit| spend_control_limit_display("Role budget", limit, captured_at))
})
.or_else(|| {
snapshot.group_shared_limit.as_ref().map(|limit| {
group_spend_control_limit_display("Group shared spend pool", limit, captured_at)
})
})
.or_else(|| {
snapshot.workspace_shared_limit.as_ref().map(|limit| {
spend_control_limit_display("Workspace shared spend pool", limit, captured_at)
})
})
}
fn spend_control_limit_display(
label: &str,
limit: &SpendControlLimitSnapshot,
captured_at: DateTime<Local>,
) -> SpendControlLimitDisplay {
SpendControlLimitDisplay {
label: label.to_string(),
remaining_percent: f64::from(limit.remaining_percent),
resets_at: format_spend_control_reset(limit.reset_at, captured_at),
}
}
fn group_spend_control_limit_display(
label: &str,
limit: &GroupSpendControlLimitSnapshot,
captured_at: DateTime<Local>,
) -> SpendControlLimitDisplay {
SpendControlLimitDisplay {
label: format!("{label}: {}", limit.group_name),
remaining_percent: f64::from(limit.details.remaining_percent),
resets_at: format_spend_control_reset(limit.details.reset_at, captured_at),
}
}
fn format_spend_control_reset(reset_at: i32, captured_at: DateTime<Local>) -> Option<String> {
DateTime::<Utc>::from_timestamp(i64::from(reset_at), 0)
.map(|dt| dt.with_timezone(&Local))
.map(|dt| format_reset_timestamp(dt, captured_at))
}
fn spend_control_limit_type_label(limit_type: SpendControlLimitType) -> String {
match limit_type {
SpendControlLimitType::Individual => "Individual spend control".to_string(),
SpendControlLimitType::GroupDefault => "Group default spend control".to_string(),
SpendControlLimitType::WorkspaceDefault => "Workspace default spend control".to_string(),
SpendControlLimitType::RoleBudget => "Role budget".to_string(),
SpendControlLimitType::GroupShared => "Group shared spend pool".to_string(),
SpendControlLimitType::WorkspaceShared => "Workspace shared spend pool".to_string(),
}
}
/// Renders a fixed-width progress bar from remaining percentage.
///
/// This function expects a remaining value in the `0..=100` range and clamps out-of-range input.
@@ -349,9 +494,12 @@ mod tests {
use super::CreditsSnapshotDisplay;
use super::RateLimitSnapshotDisplay;
use super::RateLimitWindowDisplay;
use super::SpendControlLimitDisplay;
use super::SpendControlSnapshotDisplay;
use super::StatusRateLimitData;
use super::compose_rate_limit_data_many;
use chrono::Local;
use codex_app_server_protocol::SpendControlLimitType;
use pretty_assertions::assert_eq;
fn window(used_percent: f64) -> RateLimitWindowDisplay {
@@ -375,6 +523,7 @@ mod tests {
unlimited: false,
balance: Some("25".to_string()),
}),
spend_control: None,
};
let other = RateLimitSnapshotDisplay {
limit_name: "codex-other".to_string(),
@@ -386,6 +535,7 @@ mod tests {
unlimited: false,
balance: Some("99".to_string()),
}),
spend_control: None,
};
let rows = match compose_rate_limit_data_many(&[codex, other], now) {
@@ -423,6 +573,7 @@ mod tests {
window_minutes: None,
}),
credits: None,
spend_control: None,
};
let rows = match compose_rate_limit_data_many(&[other], now) {
@@ -439,4 +590,39 @@ mod tests {
]
);
}
#[test]
fn spend_control_rows_show_active_limit_and_reached_type() {
let now = Local::now();
let codex = RateLimitSnapshotDisplay {
limit_name: "codex".to_string(),
captured_at: now,
primary: None,
secondary: None,
credits: None,
spend_control: Some(SpendControlSnapshotDisplay {
reached: true,
reached_limit_type: Some(SpendControlLimitType::GroupShared),
active_limit: Some(SpendControlLimitDisplay {
label: "Group shared spend pool: Engineering".to_string(),
remaining_percent: 18.0,
resets_at: Some("soon".to_string()),
}),
}),
};
let rows = match compose_rate_limit_data_many(&[codex], now) {
StatusRateLimitData::Available(rows) => rows,
other => panic!("unexpected status: {other:?}"),
};
let labels: Vec<String> = rows.iter().map(|row| row.label.clone()).collect();
assert_eq!(
labels,
vec![
"Group shared spend pool: Engineering".to_string(),
"Spend control reached".to_string(),
]
);
}
}
+97
View File
@@ -19,11 +19,15 @@ use codex_app_server_protocol::FileSystemAccessMode;
use codex_app_server_protocol::FileSystemPath;
use codex_app_server_protocol::FileSystemSandboxEntry;
use codex_app_server_protocol::FileSystemSpecialPath;
use codex_app_server_protocol::GroupSpendControlLimitSnapshot;
use codex_app_server_protocol::PermissionProfile as AppServerPermissionProfile;
use codex_app_server_protocol::PermissionProfileFileSystemPermissions;
use codex_app_server_protocol::PermissionProfileNetworkPermissions;
use codex_app_server_protocol::RateLimitSnapshot;
use codex_app_server_protocol::RateLimitWindow;
use codex_app_server_protocol::SpendControlLimitSnapshot;
use codex_app_server_protocol::SpendControlLimitType;
use codex_app_server_protocol::SpendControlSnapshot;
use codex_config::LoaderOverrides;
use codex_model_provider_info::ModelProviderAwsAuthInfo;
use codex_model_provider_info::ModelProviderInfo;
@@ -236,6 +240,7 @@ async fn status_snapshot_includes_reasoning_details() {
resets_at: Some(reset_at_from(&captured_at, /*seconds*/ 1_200)),
}),
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -829,6 +834,87 @@ async fn status_snapshot_includes_monthly_limit() {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref());
let token_info = token_info_for(&model_slug, &config, &usage);
let composite = new_status_output(
&config,
account_display.as_ref(),
Some(&token_info),
&usage,
&None,
/*thread_name*/ None,
/*forked_from*/ None,
Some(&rate_display),
None,
captured_at,
&model_slug,
/*collaboration_mode*/ None,
/*reasoning_effort_override*/ None,
);
let mut rendered_lines = render_lines(&composite.display_lines(/*width*/ 80));
if cfg!(windows) {
for line in &mut rendered_lines {
*line = line.replace('\\', "/");
}
}
let sanitized = sanitize_directory(rendered_lines).join("\n");
assert_snapshot!(sanitized);
}
#[tokio::test]
async fn status_snapshot_includes_group_shared_spend_control() {
let temp_home = TempDir::new().expect("temp home");
let mut config = test_config(&temp_home).await;
config.model = Some("gpt-5.1-codex-max".to_string());
config.model_provider_id = "openai".to_string();
set_workspace_cwd(&mut config, test_path_buf("/workspace/tests").abs());
let account_display = test_status_account_display();
let usage = TokenUsage {
input_tokens: 800,
cached_input_tokens: 0,
output_tokens: 400,
reasoning_output_tokens: 0,
total_tokens: 1_200,
};
let captured_at = chrono::Local
.with_ymd_and_hms(2024, 5, 6, 7, 8, 9)
.single()
.expect("timestamp");
let snapshot = RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: None,
secondary: None,
credits: None,
spend_control: Some(SpendControlSnapshot {
reached: true,
reached_limit_type: Some(SpendControlLimitType::GroupShared),
individual_limit: None,
group_default_limit: None,
workspace_default_limit: None,
role_budget_limit: None,
group_shared_limit: Some(GroupSpendControlLimitSnapshot {
group_id: "group-1".to_string(),
group_name: "Engineering".to_string(),
details: SpendControlLimitSnapshot {
limit: "200".to_string(),
used: "190".to_string(),
remaining: "10".to_string(),
used_percent: 95,
remaining_percent: 5,
reset_after_seconds: 86_400,
reset_at: reset_at_from(&captured_at, /*seconds*/ 86_400) as i32,
},
}),
workspace_shared_limit: None,
}),
plan_type: None,
rate_limit_reached_type: None,
};
@@ -881,6 +967,7 @@ async fn status_snapshot_shows_unlimited_credits() {
unlimited: true,
balance: None,
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -931,6 +1018,7 @@ async fn status_snapshot_shows_positive_credits() {
unlimited: false,
balance: Some("12.5".to_string()),
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -981,6 +1069,7 @@ async fn status_snapshot_hides_zero_credits() {
unlimited: false,
balance: Some("0".to_string()),
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1029,6 +1118,7 @@ async fn status_snapshot_hides_when_has_no_credits_flag() {
unlimited: true,
balance: None,
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1135,6 +1225,7 @@ async fn status_snapshot_truncates_in_narrow_terminal() {
}),
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1300,6 +1391,7 @@ async fn status_snapshot_shows_refreshing_limits_notice() {
resets_at: Some(reset_at_from(&captured_at, /*seconds*/ 2_700)),
}),
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1371,6 +1463,7 @@ async fn status_snapshot_includes_credits_and_limits() {
unlimited: false,
balance: Some("37.5".to_string()),
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1425,6 +1518,7 @@ async fn status_snapshot_shows_unavailable_limits_message() {
primary: None,
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1482,6 +1576,7 @@ async fn status_snapshot_treats_refreshing_empty_limits_as_unavailable() {
primary: None,
secondary: None,
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1553,6 +1648,7 @@ async fn status_snapshot_shows_stale_limits_message() {
resets_at: Some(reset_at_from(&captured_at, /*seconds*/ 1_800)),
}),
credits: None,
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};
@@ -1624,6 +1720,7 @@ async fn status_snapshot_cached_limits_hide_credits_without_flag() {
unlimited: false,
balance: Some("80".to_string()),
}),
spend_control: None,
plan_type: None,
rate_limit_reached_type: None,
};