diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index cca5ca0b43..2b0080e049 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1470,6 +1470,7 @@ dependencies = [ "codex-ansi-escape", "codex-app-server-protocol", "codex-arg0", + "codex-backend-client", "codex-common", "codex-core", "codex-feedback", diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index b524d8bfd4..b8b2ae2623 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -29,6 +29,7 @@ clap = { workspace = true, features = ["derive"] } codex-ansi-escape = { workspace = true } codex-app-server-protocol = { workspace = true } codex-arg0 = { workspace = true } +codex-backend-client = { workspace = true } codex-common = { workspace = true, features = [ "cli", "elapsed", diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 4aa295a323..9f996007f8 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -496,6 +496,9 @@ impl App { AppEvent::FileSearchResult { query, matches } => { self.chat_widget.apply_file_search_result(query, matches); } + AppEvent::RateLimitSnapshotFetched(snapshot) => { + self.chat_widget.on_rate_limit_snapshot(Some(snapshot)); + } AppEvent::UpdateReasoningEffort(effort) => { self.on_update_reasoning_effort(effort); } diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 39485faa93..0c3033c5c9 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -4,6 +4,7 @@ use codex_common::approval_presets::ApprovalPreset; use codex_common::model_presets::ModelPreset; use codex_core::protocol::ConversationPathResponseEvent; use codex_core::protocol::Event; +use codex_core::protocol::RateLimitSnapshot; use codex_file_search::FileMatch; use crate::bottom_pane::ApprovalRequest; @@ -41,6 +42,9 @@ pub(crate) enum AppEvent { matches: Vec, }, + /// Result of refreshing rate limits + RateLimitSnapshotFetched(RateLimitSnapshot), + /// Result of computing a `/diff` command. DiffResult(String), diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 79781f3352..83fea1e90e 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2,7 +2,10 @@ use std::collections::HashMap; use std::collections::VecDeque; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; +use codex_app_server_protocol::AuthMode; +use codex_backend_client::Client as BackendClient; use codex_core::config::Config; use codex_core::config::types::Notifications; use codex_core::git_info::current_branch_name; @@ -62,6 +65,7 @@ use ratatui::text::Line; use ratatui::widgets::Paragraph; use ratatui::widgets::Wrap; use tokio::sync::mpsc::UnboundedSender; +use tokio::task::JoinHandle; use tracing::debug; use crate::app_event::AppEvent; @@ -116,6 +120,7 @@ use codex_common::approval_presets::builtin_approval_presets; use codex_common::model_presets::ModelPreset; use codex_common::model_presets::builtin_model_presets; use codex_core::AuthManager; +use codex_core::CodexAuth; use codex_core::ConversationManager; use codex_core::protocol::AskForApproval; use codex_core::protocol::SandboxPolicy; @@ -255,6 +260,7 @@ pub(crate) struct ChatWidget { rate_limit_snapshot: Option, rate_limit_warnings: RateLimitWarningState, rate_limit_switch_prompt: RateLimitSwitchPromptState, + rate_limit_poller: Option>, // Stream lifecycle controller stream_controller: Option, running_commands: HashMap, @@ -494,7 +500,7 @@ impl ChatWidget { } } - fn on_rate_limit_snapshot(&mut self, snapshot: Option) { + pub(crate) fn on_rate_limit_snapshot(&mut self, snapshot: Option) { if let Some(snapshot) = snapshot { let warnings = self.rate_limit_warnings.take_warnings( snapshot @@ -1034,7 +1040,7 @@ impl ChatWidget { let placeholder = EXAMPLE_PROMPTS[rng.random_range(0..EXAMPLE_PROMPTS.len())].to_string(); let codex_op_tx = spawn_agent(config.clone(), app_event_tx.clone(), conversation_manager); - Self { + let mut widget = Self { app_event_tx: app_event_tx.clone(), frame_requester: frame_requester.clone(), codex_op_tx, @@ -1058,6 +1064,7 @@ impl ChatWidget { rate_limit_snapshot: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), + rate_limit_poller: None, stream_controller: None, running_commands: HashMap::new(), task_complete_pending: false, @@ -1076,7 +1083,11 @@ impl ChatWidget { last_rendered_width: std::cell::Cell::new(None), feedback, current_rollout_path: None, - } + }; + + widget.prefetch_rate_limits(); + + widget } /// Create a ChatWidget attached to an existing conversation (e.g., a fork). @@ -1101,7 +1112,7 @@ impl ChatWidget { let codex_op_tx = spawn_agent_from_existing(conversation, session_configured, app_event_tx.clone()); - Self { + let mut widget = Self { app_event_tx: app_event_tx.clone(), frame_requester: frame_requester.clone(), codex_op_tx, @@ -1125,6 +1136,7 @@ impl ChatWidget { rate_limit_snapshot: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), + rate_limit_poller: None, stream_controller: None, running_commands: HashMap::new(), task_complete_pending: false, @@ -1143,7 +1155,11 @@ impl ChatWidget { last_rendered_width: std::cell::Cell::new(None), feedback, current_rollout_path: None, - } + }; + + widget.prefetch_rate_limits(); + + widget } pub(crate) fn handle_key_event(&mut self, key_event: KeyEvent) { @@ -1737,6 +1753,38 @@ impl ChatWidget { Local::now(), )); } + fn stop_rate_limit_poller(&mut self) { + if let Some(handle) = self.rate_limit_poller.take() { + handle.abort(); + } + } + + fn prefetch_rate_limits(&mut self) { + self.stop_rate_limit_poller(); + + let Some(auth) = self.auth_manager.auth() else { + return; + }; + if auth.mode != AuthMode::ChatGPT { + return; + } + + let base_url = self.config.chatgpt_base_url.clone(); + let app_event_tx = self.app_event_tx.clone(); + + let handle = tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(60)); + + loop { + if let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth.clone()).await { + app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); + } + interval.tick().await; + } + }); + + self.rate_limit_poller = Some(handle); + } fn lower_cost_preset(&self) -> Option { let auth_mode = self.auth_manager.auth().map(|auth| auth.mode); @@ -2774,6 +2822,12 @@ impl ChatWidget { } } +impl Drop for ChatWidget { + fn drop(&mut self) { + self.stop_rate_limit_poller(); + } +} + impl Renderable for ChatWidget { fn render(&self, area: Rect, buf: &mut Buffer) { self.as_renderable().render(area, buf); @@ -2892,6 +2946,22 @@ fn extract_first_bold(s: &str) -> Option { None } +async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option { + match BackendClient::from_auth(base_url, &auth).await { + Ok(client) => match client.get_rate_limits().await { + Ok(snapshot) => Some(snapshot), + Err(err) => { + debug!(error = ?err, "failed to fetch rate limits from /usage"); + None + } + }, + Err(err) => { + debug!(error = ?err, "failed to construct backend client for rate limits"); + None + } + } +} + #[cfg(test)] pub(crate) fn show_review_commit_picker_with_entries( chat: &mut ChatWidget, diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index bdaae93353..c875862bbc 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -333,6 +333,7 @@ fn make_chatwidget_manual() -> ( rate_limit_snapshot: None, rate_limit_warnings: RateLimitWarningState::default(), rate_limit_switch_prompt: RateLimitSwitchPromptState::default(), + rate_limit_poller: None, stream_controller: None, running_commands: HashMap::new(), task_complete_pending: false,