agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support multiple rate limits (#11260)

Browse Source 
Added multi-limit support end-to-end by carrying limit_name in
rate-limit snapshots and handling multiple buckets instead of only
codex.
Extended /usage client parsing to consume additional_rate_limits
Updated TUI /status and in-memory state to store/render per-limit
snapshots
Extended app-server rate-limit read response: kept rate_limits and added
rate_limits_by_name.
Adjusted usage-limit error messaging for non-default codex limit buckets
This commit is contained in:
xl-openai
2026-02-10 20:09:31 -08:00
committed by GitHub
Unverified
parent 641d5268fa
commit fdd0cd1de9
36 changed files with 1435 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/api_bridge.rs
+35 -3
View File
@@ -4,7 +4,7 @@ use codex_api::AuthProvider as ApiAuthProvider;
use codex_api::TransportError;
use codex_api::error::ApiError;
use codex_api::rate_limits::parse_promo_message;
use codex_api::rate_limits::parse_rate_limit;
use codex_api::rate_limits::parse_rate_limit_for_limit;
use http::HeaderMap;
use serde::Deserialize;
@@ -71,7 +71,10 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr {
if let Ok(err) = serde_json::from_str::<UsageErrorResponse>(&body_text) {
if err.error.error_type.as_deref() == Some("usage_limit_reached") {
let rate_limits = headers.as_ref().and_then(parse_rate_limit);
let limit_id = extract_header(headers.as_ref(), ACTIVE_LIMIT_HEADER);
let rate_limits = headers.as_ref().and_then(|map| {
parse_rate_limit_for_limit(map, limit_id.as_deref())
});
let promo_message = headers.as_ref().and_then(parse_promo_message);
let resets_at = err
.error
@@ -80,8 +83,9 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr {
return CodexErr::UsageLimitReached(UsageLimitReachedError {
plan_type: err.error.plan_type,
resets_at,
rate_limits,
rate_limits: rate_limits.map(Box::new),
promo_message,
limit_name: limit_id,
});
} else if err.error.error_type.as_deref() == Some("usage_not_included") {
return CodexErr::UsageNotIncluded;
@@ -117,6 +121,7 @@ pub(crate) fn map_api_error(err: ApiError) -> CodexErr {
const MODEL_CAP_MODEL_HEADER: &str = "x-codex-model-cap-model";
const MODEL_CAP_RESET_AFTER_HEADER: &str = "x-codex-model-cap-reset-after-seconds";
const ACTIVE_LIMIT_HEADER: &str = "x-codex-active-limit";
const REQUEST_ID_HEADER: &str = "x-request-id";
const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";
const CF_RAY_HEADER: &str = "cf-ray";
@@ -152,6 +157,33 @@ mod tests {
assert_eq!(model_cap.model, "boomslang");
assert_eq!(model_cap.reset_after_seconds, Some(120));
}
#[test]
fn map_api_error_maps_usage_limit_limit_name_header() {
let mut headers = HeaderMap::new();
headers.insert(
ACTIVE_LIMIT_HEADER,
http::HeaderValue::from_static("codex_other"),
);
let body = serde_json::json!({
"error": {
"type": "usage_limit_reached",
"plan_type": "pro",
}
})
.to_string();
let err = map_api_error(ApiError::Transport(TransportError::Http {
status: StatusCode::TOO_MANY_REQUESTS,
url: Some("http://example.com/v1/responses".to_string()),
headers: Some(headers),
body: Some(body),
}));
let CodexErr::UsageLimitReached(usage_limit) = err else {
panic!("expected CodexErr::UsageLimitReached, got {err:?}");
};
assert_eq!(usage_limit.limit_name.as_deref(), Some("codex_other"));
}
}
fn extract_request_tracking_id(headers: Option<&HeaderMap>) -> Option<String> {
codex-rs/core/src/codex.rs
+13 -1
View File
@@ -4280,7 +4280,7 @@ async fn run_sampling_request(
Err(CodexErr::UsageLimitReached(e)) => {
let rate_limits = e.rate_limits.clone();
if let Some(rate_limits) = rate_limits {
sess.update_rate_limits(&turn_context, rate_limits).await;
sess.update_rate_limits(&turn_context, *rate_limits).await;
}
return Err(CodexErr::UsageLimitReached(e));
}
@@ -5823,6 +5823,8 @@ mod tests {
let mut state = SessionState::new(session_configuration);
let initial = RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 10.0,
window_minutes: Some(15),
@@ -5839,6 +5841,8 @@ mod tests {
state.set_rate_limits(initial.clone());
let update = RateLimitSnapshot {
limit_id: Some("codex_other".to_string()),
limit_name: Some("codex_other".to_string()),
primary: Some(RateLimitWindow {
used_percent: 40.0,
window_minutes: Some(30),
@@ -5857,6 +5861,8 @@ mod tests {
assert_eq!(
state.latest_rate_limits,
Some(RateLimitSnapshot {
limit_id: Some("codex_other".to_string()),
limit_name: Some("codex_other".to_string()),
primary: update.primary.clone(),
secondary: update.secondary,
credits: initial.credits,
@@ -5906,6 +5912,8 @@ mod tests {
let mut state = SessionState::new(session_configuration);
let initial = RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 15.0,
window_minutes: Some(20),
@@ -5926,6 +5934,8 @@ mod tests {
state.set_rate_limits(initial.clone());
let update = RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 35.0,
window_minutes: Some(25),
@@ -5940,6 +5950,8 @@ mod tests {
assert_eq!(
state.latest_rate_limits,
Some(RateLimitSnapshot {
limit_id: Some("codex".to_string()),
limit_name: None,
primary: update.primary,
secondary: update.secondary,
credits: initial.credits,
codex-rs/core/src/error.rs
+63 -14
View File
@@ -409,12 +409,23 @@ impl std::fmt::Display for RetryLimitReachedError {
pub struct UsageLimitReachedError {
pub(crate) plan_type: Option<PlanType>,
pub(crate) resets_at: Option<DateTime<Utc>>,
pub(crate) rate_limits: Option<RateLimitSnapshot>,
pub(crate) rate_limits: Option<Box<RateLimitSnapshot>>,
pub(crate) promo_message: Option<String>,
pub(crate) limit_name: Option<String>,
}
impl std::fmt::Display for UsageLimitReachedError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if let Some(limit_name) = self.limit_name.as_deref()
&& !limit_name.eq_ignore_ascii_case("codex")
{
return write!(
f,
"You've hit your usage limit for {limit_name}.{}",
retry_suffix(self.resets_at.as_ref())
);
}
if let Some(promo_message) = &self.promo_message {
return write!(
f,
@@ -699,6 +710,8 @@ mod tests {
.unwrap()
.timestamp();
RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 50.0,
window_minutes: Some(60),
@@ -728,8 +741,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Plus)),
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -875,8 +889,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Free)),
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -889,8 +904,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Go)),
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -903,8 +919,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: None,
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -921,8 +938,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Team)),
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!(
"You've hit your usage limit. To get more access now, send a request to your admin or try again at {expected_time}."
@@ -936,8 +954,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Business)),
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -950,8 +969,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Enterprise)),
resets_at: None,
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
assert_eq!(
err.to_string(),
@@ -968,8 +988,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Pro)),
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!(
"You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}."
@@ -978,6 +999,29 @@ mod tests {
});
}
#[test]
fn usage_limit_reached_error_hides_upsell_for_non_codex_limit_name() {
let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
let resets_at = base + ChronoDuration::hours(1);
with_now_override(base, move || {
let expected_time = format_retry_timestamp(&resets_at);
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Plus)),
resets_at: Some(resets_at),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: Some(
"Visit https://chatgpt.com/codex/settings/usage to purchase more credits"
.to_string(),
),
limit_name: Some("codex_other".to_string()),
};
let expected = format!(
"You've hit your usage limit for codex_other. Try again at {expected_time}."
);
assert_eq!(err.to_string(), expected);
});
}
#[test]
fn usage_limit_reached_includes_minutes_when_available() {
let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
@@ -987,8 +1031,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: None,
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!("You've hit your usage limit. Try again at {expected_time}.");
assert_eq!(err.to_string(), expected);
@@ -1096,8 +1141,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: Some(PlanType::Known(KnownPlan::Plus)),
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!(
"You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}."
@@ -1116,8 +1162,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: None,
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!("You've hit your usage limit. Try again at {expected_time}.");
assert_eq!(err.to_string(), expected);
@@ -1133,8 +1180,9 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: None,
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: None,
limit_name: None,
};
let expected = format!("You've hit your usage limit. Try again at {expected_time}.");
assert_eq!(err.to_string(), expected);
@@ -1150,10 +1198,11 @@ mod tests {
let err = UsageLimitReachedError {
plan_type: None,
resets_at: Some(resets_at),
rate_limits: Some(rate_limit_snapshot()),
rate_limits: Some(Box::new(rate_limit_snapshot())),
promo_message: Some(
"To continue using Codex, start a free trial of <PLAN> today".to_string(),
),
limit_name: None,
};
let expected = format!(
"You've hit your usage limit. To continue using Codex, start a free trial of <PLAN> today, or try again at {expected_time}."
codex-rs/core/src/state/session.rs
+134 -1
View File
@@ -170,11 +170,21 @@ impl SessionState {
}
}
// Sometimes new snapshots don't include credits or plan information.
// Merge partial rate-limit updates: new fields overwrite existing values;
// missing fields retain prior values. If `limit_id` is absent everywhere,
// default it to `"codex"`.
fn merge_rate_limit_fields(
previous: Option<&RateLimitSnapshot>,
mut snapshot: RateLimitSnapshot,
) -> RateLimitSnapshot {
if snapshot.limit_id.is_none() {
snapshot.limit_id = previous
.and_then(|prior| prior.limit_id.clone())
.or_else(|| Some("codex".to_string()));
}
if snapshot.limit_name.is_none() {
snapshot.limit_name = previous.and_then(|prior| prior.limit_name.clone());
}
if snapshot.credits.is_none() {
snapshot.credits = previous.and_then(|prior| prior.credits.clone());
}
@@ -188,6 +198,7 @@ fn merge_rate_limit_fields(
mod tests {
use super::*;
use crate::codex::make_session_configuration_for_tests;
use crate::protocol::RateLimitWindow;
use pretty_assertions::assert_eq;
#[tokio::test]
@@ -258,4 +269,126 @@ mod tests {
assert_eq!(state.get_mcp_tool_selection(), None);
}
#[tokio::test]
async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
state.set_rate_limits(RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 12.0,
window_minutes: Some(60),
resets_at: Some(100),
}),
secondary: None,
credits: None,
plan_type: None,
});
assert_eq!(
state
.latest_rate_limits
.as_ref()
.and_then(|v| v.limit_id.clone()),
Some("codex".to_string())
);
}
#[tokio::test]
async fn set_rate_limits_preserves_previous_limit_id_when_missing() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
state.set_rate_limits(RateLimitSnapshot {
limit_id: Some("codex_other".to_string()),
limit_name: Some("codex_other".to_string()),
primary: Some(RateLimitWindow {
used_percent: 20.0,
window_minutes: Some(60),
resets_at: Some(200),
}),
secondary: None,
credits: None,
plan_type: None,
});
state.set_rate_limits(RateLimitSnapshot {
limit_id: None,
limit_name: None,
primary: Some(RateLimitWindow {
used_percent: 30.0,
window_minutes: Some(60),
resets_at: Some(300),
}),
secondary: None,
credits: None,
plan_type: None,
});
assert_eq!(
state
.latest_rate_limits
.as_ref()
.and_then(|v| v.limit_id.clone()),
Some("codex_other".to_string())
);
}
#[tokio::test]
async fn set_rate_limits_accepts_new_limit_id_bucket() {
let session_configuration = make_session_configuration_for_tests().await;
let mut state = SessionState::new(session_configuration);
state.set_rate_limits(RateLimitSnapshot {
limit_id: Some("codex".to_string()),
limit_name: Some("codex".to_string()),
primary: Some(RateLimitWindow {
used_percent: 10.0,
window_minutes: Some(60),
resets_at: Some(100),
}),
secondary: None,
credits: Some(crate::protocol::CreditsSnapshot {
has_credits: true,
unlimited: false,
balance: Some("50".to_string()),
}),
plan_type: Some(codex_protocol::account::PlanType::Plus),
});
state.set_rate_limits(RateLimitSnapshot {
limit_id: Some("codex_other".to_string()),
limit_name: Some("codex_other".to_string()),
primary: Some(RateLimitWindow {
used_percent: 30.0,
window_minutes: Some(120),
resets_at: Some(200),
}),
secondary: None,
credits: None,
plan_type: None,
});
assert_eq!(
state.latest_rate_limits,
Some(RateLimitSnapshot {
limit_id: Some("codex_other".to_string()),
limit_name: Some("codex_other".to_string()),
primary: Some(RateLimitWindow {
used_percent: 30.0,
window_minutes: Some(120),
resets_at: Some(200),
}),
secondary: None,
credits: Some(crate::protocol::CreditsSnapshot {
has_credits: true,
unlimited: false,
balance: Some("50".to_string()),
}),
plan_type: Some(codex_protocol::account::PlanType::Plus),
})
);
}
}