Skip to content
Merged
1 change: 1 addition & 0 deletions codex-rs/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions codex-rs/tui/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/tui/src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
4 changes: 4 additions & 0 deletions codex-rs/tui/src/app_event.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -41,6 +42,9 @@ pub(crate) enum AppEvent {
matches: Vec<FileMatch>,
},

/// Result of refreshing rate limits
RateLimitSnapshotFetched(RateLimitSnapshot),

/// Result of computing a `/diff` command.
DiffResult(String),

Expand Down
80 changes: 75 additions & 5 deletions codex-rs/tui/src/chatwidget.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -255,6 +260,7 @@ pub(crate) struct ChatWidget {
rate_limit_snapshot: Option<RateLimitSnapshotDisplay>,
rate_limit_warnings: RateLimitWarningState,
rate_limit_switch_prompt: RateLimitSwitchPromptState,
rate_limit_poller: Option<JoinHandle<()>>,
// Stream lifecycle controller
stream_controller: Option<StreamController>,
running_commands: HashMap<String, RunningCommand>,
Expand Down Expand Up @@ -494,7 +500,7 @@ impl ChatWidget {
}
}

fn on_rate_limit_snapshot(&mut self, snapshot: Option<RateLimitSnapshot>) {
pub(crate) fn on_rate_limit_snapshot(&mut self, snapshot: Option<RateLimitSnapshot>) {
if let Some(snapshot) = snapshot {
let warnings = self.rate_limit_warnings.take_warnings(
snapshot
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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).
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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) {
Expand Down Expand Up @@ -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<ModelPreset> {
let auth_mode = self.auth_manager.auth().map(|auth| auth.mode);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -2892,6 +2946,22 @@ fn extract_first_bold(s: &str) -> Option<String> {
None
}

async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option<RateLimitSnapshot> {
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,
Expand Down
1 change: 1 addition & 0 deletions codex-rs/tui/src/chatwidget/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading