Compare commits

...

16 Commits

Author SHA1 Message Date
kevin zhao
c7d8c38d2f plan type reconciliation 2025-11-17 15:04:52 -08:00
kevin zhao
6259f91e5d adding drop logic 2025-11-17 14:29:00 -08:00
kevin zhao
189afb052f del config.toml 2025-11-17 13:23:40 -08:00
kevin zhao
82f49f4b4e simplify prefetch_rate_limits 2025-11-17 13:23:05 -08:00
kevin zhao
58ce6bcfb5 fix comment 2025-11-17 13:22:02 -08:00
kevin zhao
296c98b159 . 2025-11-17 13:20:02 -08:00
kevin zhao
82639ec81b . 2025-11-17 13:19:13 -08:00
kevin zhao
c8ecc49441 undo diff 2025-11-17 13:12:37 -08:00
kevin zhao
6343be1100 background fetch 2025-11-17 13:08:09 -08:00
kevin zhao
21ae112767 rate limit prefetching 2025-11-17 12:44:26 -08:00
kevin zhao
a1db68171c delete config.toml 2025-11-17 11:38:41 -08:00
kevin zhao
fc6aa8406a simplifying 2025-11-17 11:37:11 -08:00
kevin zhao
7ffdcecb8f fetching even if timeout 2025-11-17 11:35:51 -08:00
kevin zhao
f8e3c57975 remove unecessary config.toml + fmt 2025-11-17 11:01:18 -08:00
kevin zhao
bfa6fbc700 unecessary diff 2025-11-17 10:56:14 -08:00
kevin zhao
57c3632a9e fetching rate limits when calling /status 2025-11-17 10:47:28 -08:00
9 changed files with 166 additions and 7 deletions

1
codex-rs/Cargo.lock generated
View File

@@ -1470,6 +1470,7 @@ dependencies = [
"codex-ansi-escape",
"codex-app-server-protocol",
"codex-arg0",
"codex-backend-client",
"codex-common",
"codex-core",
"codex-feedback",

View File

@@ -967,6 +967,7 @@ impl CodexMessageProcessor {
client
.get_rate_limits()
.await
.map(|status| status.snapshot)
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to fetch codex rate limits: {err}"),

View File

@@ -1,5 +1,6 @@
use crate::types::CodeTaskDetailsResponse;
use crate::types::PaginatedListTaskListItem;
use crate::types::PlanType;
use crate::types::RateLimitStatusPayload;
use crate::types::RateLimitWindowSnapshot;
use crate::types::TurnAttemptsSiblingTurnsResponse;
@@ -44,6 +45,12 @@ pub struct Client {
path_style: PathStyle,
}
#[derive(Clone, Debug)]
pub struct RateLimitStatus {
pub plan_type: String,
pub snapshot: RateLimitSnapshot,
}
impl Client {
pub fn new(base_url: impl Into<String>) -> Result<Self> {
let mut base_url = base_url.into();
@@ -155,7 +162,7 @@ impl Client {
}
}
pub async fn get_rate_limits(&self) -> Result<RateLimitSnapshot> {
pub async fn get_rate_limits(&self) -> Result<RateLimitStatus> {
let url = match self.path_style {
PathStyle::CodexApi => format!("{}/api/codex/usage", self.base_url),
PathStyle::ChatGptApi => format!("{}/wham/usage", self.base_url),
@@ -163,7 +170,9 @@ impl Client {
let req = self.http.get(&url).headers(self.headers());
let (body, ct) = self.exec_request(req, "GET", &url).await?;
let payload: RateLimitStatusPayload = self.decode_json(&url, &ct, &body)?;
Ok(Self::rate_limit_snapshot_from_payload(payload))
let plan_type = Self::plan_type_to_string(payload.plan_type);
let snapshot = Self::rate_limit_snapshot_from_payload(payload);
Ok(RateLimitStatus { plan_type, snapshot })
}
pub async fn list_tasks(
@@ -288,6 +297,22 @@ impl Client {
}
}
fn plan_type_to_string(plan_type: PlanType) -> String {
match plan_type {
PlanType::Free => "free",
PlanType::Go => "go",
PlanType::Plus => "plus",
PlanType::Pro => "pro",
PlanType::Team => "team",
PlanType::Business => "business",
PlanType::Education => "education",
PlanType::Quorum => "quorum",
PlanType::Enterprise => "enterprise",
PlanType::Edu => "edu",
}
.to_string()
}
fn map_rate_limit_window(
window: Option<Option<Box<RateLimitWindowSnapshot>>>,
) -> Option<RateLimitWindow> {

View File

@@ -2,6 +2,7 @@ mod client;
pub mod types;
pub use client::Client;
pub use client::RateLimitStatus;
pub use types::CodeTaskDetailsResponse;
pub use types::CodeTaskDetailsResponseExt;
pub use types::PaginatedListTaskListItem;

View File

@@ -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",

View File

@@ -496,6 +496,12 @@ impl App {
AppEvent::FileSearchResult { query, matches } => {
self.chat_widget.apply_file_search_result(query, matches);
}
AppEvent::RateLimitSnapshotFetched(status) => {
self.chat_widget
.on_rate_limit_snapshot(Some(status.snapshot.clone()));
self.chat_widget
.reconcile_plan_type_from_usage(status.plan_type);
}
AppEvent::UpdateReasoningEffort(effort) => {
self.on_update_reasoning_effort(effort);
}

View File

@@ -1,5 +1,6 @@
use std::path::PathBuf;
use codex_backend_client::RateLimitStatus;
use codex_common::approval_presets::ApprovalPreset;
use codex_common::model_presets::ModelPreset;
use codex_core::protocol::ConversationPathResponseEvent;
@@ -41,6 +42,9 @@ pub(crate) enum AppEvent {
matches: Vec<FileMatch>,
},
/// Result of refreshing rate limits
RateLimitSnapshotFetched(RateLimitStatus),
/// Result of computing a `/diff` command.
DiffResult(String),

View File

@@ -2,7 +2,11 @@ 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_backend_client::RateLimitStatus;
use codex_core::config::Config;
use codex_core::config::types::Notifications;
use codex_core::git_info::current_branch_name;
@@ -62,6 +66,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 +121,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 +261,8 @@ pub(crate) struct ChatWidget {
rate_limit_snapshot: Option<RateLimitSnapshotDisplay>,
rate_limit_warnings: RateLimitWarningState,
rate_limit_switch_prompt: RateLimitSwitchPromptState,
rate_limit_poller: Option<JoinHandle<()>>,
last_plan_refresh_attempt: Option<String>,
// Stream lifecycle controller
stream_controller: Option<StreamController>,
running_commands: HashMap<String, RunningCommand>,
@@ -494,7 +502,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
@@ -547,6 +555,51 @@ impl ChatWidget {
self.rate_limit_snapshot = None;
}
}
pub(crate) fn reconcile_plan_type_from_usage(&mut self, backend_plan: String) {
let trimmed = backend_plan.trim();
if trimmed.is_empty() {
return;
}
let plan_lower = trimmed.to_ascii_lowercase();
let Some(auth) = self.auth_manager.auth() else {
self.last_plan_refresh_attempt = None;
return;
};
if auth.mode != AuthMode::ChatGPT {
self.last_plan_refresh_attempt = None;
return;
}
let current_plan = auth
.raw_plan_type()
.map(|plan| plan.to_ascii_lowercase());
if current_plan
.as_deref()
.is_some_and(|plan| plan == plan_lower)
{
self.last_plan_refresh_attempt = None;
return;
}
if self
.last_plan_refresh_attempt
.as_deref()
.is_some_and(|attempt| attempt == plan_lower)
{
return;
}
self.last_plan_refresh_attempt = Some(plan_lower);
let auth_manager = self.auth_manager.clone();
tokio::spawn(async move {
if let Err(err) = auth_manager.refresh_token().await {
debug!(?err, "failed to refresh auth after plan mismatch");
}
});
}
/// Finalize any active exec as failed and stop/clear running UI state.
fn finalize_turn(&mut self) {
// Ensure any spinner is replaced by a red ✗ and flushed into history.
@@ -1034,7 +1087,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 +1111,8 @@ impl ChatWidget {
rate_limit_snapshot: None,
rate_limit_warnings: RateLimitWarningState::default(),
rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
rate_limit_poller: None,
last_plan_refresh_attempt: None,
stream_controller: None,
running_commands: HashMap::new(),
task_complete_pending: false,
@@ -1076,7 +1131,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 +1160,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 +1184,8 @@ impl ChatWidget {
rate_limit_snapshot: None,
rate_limit_warnings: RateLimitWarningState::default(),
rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
rate_limit_poller: None,
last_plan_refresh_attempt: None,
stream_controller: None,
running_commands: HashMap::new(),
task_complete_pending: false,
@@ -1143,7 +1204,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 +1802,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(status) = fetch_rate_limits(base_url.clone(), auth.clone()).await {
app_event_tx.send(AppEvent::RateLimitSnapshotFetched(status));
}
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);
@@ -2774,6 +2871,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 +2995,22 @@ fn extract_first_bold(s: &str) -> Option<String> {
None
}
async fn fetch_rate_limits(base_url: String, auth: CodexAuth) -> Option<RateLimitStatus> {
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,

View File

@@ -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,