diff --git a/src/codex/accounts/model.rs b/src/codex/accounts/model.rs index 32f67768a..6ff9ad0da 100644 --- a/src/codex/accounts/model.rs +++ b/src/codex/accounts/model.rs @@ -29,6 +29,9 @@ pub struct Account { pub request_count: u64, pub empty_response_count: u64, pub window_request_count: u64, + pub window_input_tokens: u64, + pub window_output_tokens: u64, + pub window_cached_tokens: u64, pub window_started_at: Option>, pub window_reset_at: Option>, pub limit_window_seconds: Option, @@ -55,6 +58,9 @@ impl Account { request_count: 0, empty_response_count: 0, window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/accounts/pool.rs b/src/codex/accounts/pool.rs index 1a833eb97..900ecb0eb 100644 --- a/src/codex/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -160,6 +160,9 @@ impl AccountPool { account.request_count = 0; account.empty_response_count = 0; account.window_request_count = 0; + account.window_input_tokens = 0; + account.window_output_tokens = 0; + account.window_cached_tokens = 0; account.window_started_at = None; account.window_reset_at = None; true @@ -173,19 +176,51 @@ impl AccountPool { let Some(account) = self.accounts.get_mut(account_id) else { return false; }; + let final_cooldown_until = account + .quota_cooldown_until + .filter(|existing| *existing > cooldown_until) + .unwrap_or(cooldown_until); account.quota_limit_reached = true; - account.quota_cooldown_until = Some(cooldown_until); - account.window_reset_at = Some(cooldown_until); - if let Ok(seconds) = (cooldown_until - Utc::now()) + account.quota_cooldown_until = Some(final_cooldown_until); + account.window_reset_at = account + .window_reset_at + .filter(|existing| *existing > final_cooldown_until) + .or(Some(final_cooldown_until)); + if let Ok(seconds) = (final_cooldown_until - Utc::now()) .to_std() .map(|duration| duration.as_secs()) { - account.limit_window_seconds = Some(seconds); + account.limit_window_seconds.get_or_insert(seconds); } self.slots.remove(account_id); true } + pub fn sync_rate_limit_window( + &mut self, + account_id: &str, + new_reset_at: DateTime, + limit_window_seconds: Option, + ) -> bool { + self.sync_rate_limit_window_at(account_id, new_reset_at, limit_window_seconds, Utc::now()) + } + + pub fn record_window_token_usage( + &mut self, + account_id: &str, + input_tokens: u64, + output_tokens: u64, + cached_tokens: u64, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + account.window_input_tokens = account.window_input_tokens.saturating_add(input_tokens); + account.window_output_tokens = account.window_output_tokens.saturating_add(output_tokens); + account.window_cached_tokens = account.window_cached_tokens.saturating_add(cached_tokens); + true + } + pub fn set_cloudflare_cooldown_until( &mut self, account_id: &str, @@ -383,6 +418,27 @@ impl AccountPool { Some(account.clone()) } + fn sync_rate_limit_window_at( + &mut self, + account_id: &str, + new_reset_at: DateTime, + limit_window_seconds: Option, + now: DateTime, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + if should_reset_window_counters(account, new_reset_at, limit_window_seconds) { + reset_window_counters(account); + account.window_started_at = Some(now); + } + account.window_reset_at = Some(new_reset_at); + if let Some(limit_window_seconds) = limit_window_seconds { + account.limit_window_seconds = Some(limit_window_seconds); + } + true + } + fn cleanup_stale_slots(&mut self, now: DateTime) { let ttl = self.options.stale_slot_ttl; self.slots.retain(|_, slots| { @@ -458,7 +514,7 @@ fn refresh_quota_window(account: &mut Account, now: DateTime) { .is_some_and(|cooldown_until| now >= cooldown_until); if window_expired { - account.window_request_count = 0; + reset_window_counters(account); account.window_started_at = Some(now); account.window_reset_at = next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); @@ -470,6 +526,39 @@ fn refresh_quota_window(account: &mut Account, now: DateTime) { } } +fn should_reset_window_counters( + account: &Account, + new_reset_at: DateTime, + limit_window_seconds: Option, +) -> bool { + let Some(old_reset_at) = account.window_reset_at else { + return false; + }; + if old_reset_at == new_reset_at { + return false; + } + let drift = old_reset_at + .signed_duration_since(new_reset_at) + .num_seconds() + .unsigned_abs(); + let window_seconds = limit_window_seconds + .or(account.limit_window_seconds) + .unwrap_or(0); + let threshold = if window_seconds > 0 { + window_seconds / 2 + } else { + 3_600 + }; + drift >= threshold +} + +fn reset_window_counters(account: &mut Account) { + account.window_request_count = 0; + account.window_input_tokens = 0; + account.window_output_tokens = 0; + account.window_cached_tokens = 0; +} + fn next_window_reset_at( reset_at: Option>, limit_window_seconds: Option, diff --git a/src/codex/accounts/repository.rs b/src/codex/accounts/repository.rs index 18f97e61e..8ba935e55 100644 --- a/src/codex/accounts/repository.rs +++ b/src/codex/accounts/repository.rs @@ -359,6 +359,9 @@ impl AccountRepository { request_count: row.get::("usage_request_count").max(0) as u64, empty_response_count: row.get::("usage_empty_response_count").max(0) as u64, window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, @@ -560,15 +563,28 @@ impl AccountRepository { Ok(result.rows_affected() > 0) } + pub async fn get_quota_json( + &self, + account_id: &str, + ) -> AccountRepositoryResult> { + let row = sqlx::query("select quota_json from accounts where id = ?") + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + Ok(row.and_then(|row| row.get("quota_json"))) + } + pub async fn set_quota_cooldown_until( &self, id: &str, cooldown_until: DateTime, ) -> AccountRepositoryResult { + let cooldown_until = cooldown_until.to_rfc3339(); let result = sqlx::query( - "update accounts set quota_limit_reached = 1, quota_cooldown_until = ?, updated_at = ? where id = ?", + "update accounts set quota_limit_reached = 1, quota_cooldown_until = case when quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until else ? end, updated_at = ? where id = ?", ) - .bind(cooldown_until.to_rfc3339()) + .bind(&cooldown_until) + .bind(cooldown_until) .bind(Utc::now().to_rfc3339()) .bind(id) .execute(&self.pool) diff --git a/src/codex/accounts/service/import.rs b/src/codex/accounts/service/import.rs index b0f843074..65dfb5a3c 100644 --- a/src/codex/accounts/service/import.rs +++ b/src/codex/accounts/service/import.rs @@ -201,6 +201,9 @@ impl AccountService { request_count: 0, empty_response_count: 0, window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/accounts/service/runtime_pool.rs b/src/codex/accounts/service/runtime_pool.rs index 19f59796a..a32d09486 100644 --- a/src/codex/accounts/service/runtime_pool.rs +++ b/src/codex/accounts/service/runtime_pool.rs @@ -57,6 +57,9 @@ pub(super) fn pool_account_from_stored(account: StoredAccount) -> Account { request_count: 0, empty_response_count: 0, window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, window_started_at: None, window_reset_at: None, limit_window_seconds: None, diff --git a/src/codex/gateway/transport/client.rs b/src/codex/gateway/transport/client.rs index b8272ba75..4045bedc7 100644 --- a/src/codex/gateway/transport/client.rs +++ b/src/codex/gateway/transport/client.rs @@ -19,10 +19,11 @@ use crate::{ types::CodexResponsesRequest, usage::{extract_sse_usage, TokenUsage}, websocket::{ - create_response_via_websocket, create_response_via_websocket_stream, - create_response_via_websocket_stream_with_pool, transport_for_request, CodexTransport, - CodexWebSocketError, CodexWebSocketPool, CodexWebSocketPoolKey, - CodexWebSocketSseStream, CodexWebSocketStreamResponse, WebSocketSupportError, + append_rate_limit_updates, create_response_via_websocket, + create_response_via_websocket_stream, create_response_via_websocket_stream_with_pool, + transport_for_request, CodexTransport, CodexWebSocketError, CodexWebSocketPool, + CodexWebSocketPoolKey, CodexWebSocketSseStream, CodexWebSocketStreamResponse, + SharedRateLimitUpdates, WebSocketSupportError, }, }, }; @@ -97,6 +98,7 @@ pub struct CodexBackendWebSocketStream { pub turn_state: Option, pub set_cookie_headers: Vec, pub rate_limit_headers: Vec<(String, String)>, + pub rate_limit_updates: SharedRateLimitUpdates, } #[derive(Clone)] @@ -273,6 +275,7 @@ impl CodexBackendClient { turn_state: response.turn_state, set_cookie_headers: response.set_cookie_headers, rate_limit_headers: response.rate_limit_headers, + rate_limit_updates: response.rate_limit_updates, }) } @@ -537,7 +540,8 @@ async fn collect_websocket_stream_response( mut body_stream, turn_state, set_cookie_headers, - rate_limit_headers, + mut rate_limit_headers, + rate_limit_updates, } = response; let mut body = String::new(); while let Some(chunk) = body_stream.next().await { @@ -546,6 +550,7 @@ async fn collect_websocket_stream_response( if body.is_empty() { return Err(CodexWebSocketError::EmptyResponse); } + append_rate_limit_updates(&mut rate_limit_headers, &rate_limit_updates).await; Ok( crate::codex::gateway::transport::websocket::CodexWebSocketResponse { body, @@ -704,7 +709,14 @@ fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> { fn is_rate_limit_header(name: &str) -> bool { let name = name.to_ascii_lowercase(); - name == "retry-after" || name.contains("ratelimit") || name.contains("rate-limit") + name == "retry-after" + || name.contains("ratelimit") + || name.contains("rate-limit") + || name.starts_with("x-codex-primary-") + || name.starts_with("x-codex-secondary-") + || name.starts_with("x-codex-code-review-") + || name.starts_with("x-codex-review-") + || name.starts_with("x-code-review-") } fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option { diff --git a/src/codex/gateway/transport/mod.rs b/src/codex/gateway/transport/mod.rs index 038fae8a2..463ea73f3 100644 --- a/src/codex/gateway/transport/mod.rs +++ b/src/codex/gateway/transport/mod.rs @@ -1,5 +1,6 @@ pub mod client; pub mod headers; +pub mod rate_limits; pub mod sse; pub mod types; pub mod usage; diff --git a/src/codex/gateway/transport/rate_limits.rs b/src/codex/gateway/transport/rate_limits.rs new file mode 100644 index 000000000..96085edca --- /dev/null +++ b/src/codex/gateway/transport/rate_limits.rs @@ -0,0 +1,467 @@ +use std::{ + collections::BTreeMap, + time::{SystemTime, UNIX_EPOCH}, +}; + +use chrono::{DateTime, Duration, Utc}; +use serde_json::{json, Value}; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RateLimitWindow { + pub used_percent: f64, + pub window_minutes: Option, + pub reset_at: Option, +} + +impl RateLimitWindow { + pub fn limit_window_seconds(self) -> Option { + self.window_minutes + .and_then(|minutes| minutes.checked_mul(60)) + } + + pub fn reset_at_datetime(self) -> Option> { + let reset_at = self.reset_at?; + DateTime::::from_timestamp(reset_at, 0) + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RateLimitDetails { + pub allowed: Option, + pub limit_reached: Option, + pub primary: Option, + pub secondary: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ParsedRateLimits { + pub primary: Option, + pub secondary: Option, + pub code_review: Option, +} + +impl ParsedRateLimits { + pub fn primary_reset_at(self) -> Option> { + self.primary.and_then(RateLimitWindow::reset_at_datetime) + } + + pub fn primary_limit_window_seconds(self) -> Option { + self.primary.and_then(RateLimitWindow::limit_window_seconds) + } + + pub fn primary_limit_reached(self) -> bool { + self.primary + .is_some_and(|window| window.used_percent >= 100.0) + } +} + +pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option { + let mut normalized = BTreeMap::new(); + for (name, value) in headers { + normalized.insert(name.to_ascii_lowercase(), value.trim()); + } + + let primary = parse_window_from_lookup(&normalized, "x-codex-primary"); + let secondary = parse_window_from_lookup(&normalized, "x-codex-secondary"); + let code_review = parse_details_from_lookup(&normalized, "x-codex-code-review") + .or_else(|| parse_details_from_lookup(&normalized, "x-codex-review")) + .or_else(|| parse_details_from_lookup(&normalized, "x-code-review")); + + if primary.is_none() && secondary.is_none() && code_review.is_none() { + return None; + } + + Some(ParsedRateLimits { + primary, + secondary, + code_review, + }) +} + +pub fn parse_rate_limits_event(value: &Value) -> Option { + if value + .get("type") + .and_then(Value::as_str) + .is_some_and(|event_type| event_type != "codex.rate_limits") + { + return None; + } + + let details = value.get("rate_limits").and_then(parse_details_from_object); + let explicit_code_review = value + .get("code_review_rate_limits") + .and_then(parse_details_from_object) + .or_else(|| { + value + .get("code_review_rate_limit") + .and_then(parse_details_from_object) + }); + + let mut primary = details.and_then(|details| details.primary); + let mut secondary = details.and_then(|details| details.secondary); + let mut code_review = explicit_code_review; + if details.is_some() && is_review_limit_name(rate_limit_name(value)) { + code_review = code_review.or(details); + primary = None; + secondary = None; + } + + if primary.is_none() && secondary.is_none() && code_review.is_none() { + return None; + } + + Some(ParsedRateLimits { + primary, + secondary, + code_review, + }) +} + +pub fn parse_rate_limits_event_raw(raw: &str) -> Option { + serde_json::from_str::(raw) + .ok() + .and_then(|value| parse_rate_limits_event(&value)) +} + +pub fn rate_limit_quota( + rate_limits: &ParsedRateLimits, + plan_type: Option<&str>, + existing_quota: Option<&Value>, +) -> Value { + let mut quota = json!({ + "plan_type": plan_type.unwrap_or("unknown"), + "rate_limit": quota_window(rate_limits.primary, true), + "secondary_rate_limit": rate_limits.secondary.map(|window| quota_window(Some(window), false)).unwrap_or(Value::Null), + "code_review_rate_limit": rate_limits.code_review.map(code_review_quota_window).unwrap_or(Value::Null), + }); + + if let Some(credits) = existing_quota.and_then(|quota| quota.get("credits")) { + quota["credits"] = credits.clone(); + } + + quota +} + +pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(String, String)> { + let mut headers = Vec::new(); + push_window_headers(&mut headers, "x-codex-primary", rate_limits.primary); + push_window_headers(&mut headers, "x-codex-secondary", rate_limits.secondary); + if let Some(code_review) = rate_limits.code_review { + push_window_headers( + &mut headers, + "x-codex-code-review-primary", + code_review.primary, + ); + push_window_headers( + &mut headers, + "x-codex-code-review-secondary", + code_review.secondary, + ); + } + headers +} + +pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration { + let factor_bps = u64::from(factor_bps); + let variance = seconds.saturating_mul(factor_bps) / 10_000; + if variance == 0 { + return Duration::seconds(seconds.min(i64::MAX as u64) as i64); + } + let lower = seconds.saturating_sub(variance); + let span = variance.saturating_mul(2).saturating_add(1); + let offset = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| u64::from(duration.subsec_nanos()) % span) + .unwrap_or(variance); + Duration::seconds(lower.saturating_add(offset).min(i64::MAX as u64) as i64) +} + +fn push_window_headers( + headers: &mut Vec<(String, String)>, + prefix: &str, + window: Option, +) { + let Some(window) = window else { + return; + }; + headers.push(( + format!("{prefix}-used-percent"), + window.used_percent.to_string(), + )); + if let Some(window_minutes) = window.window_minutes { + headers.push(( + format!("{prefix}-window-minutes"), + window_minutes.to_string(), + )); + } + if let Some(reset_at) = window.reset_at { + headers.push((format!("{prefix}-reset-at"), reset_at.to_string())); + } +} + +fn parse_details_from_lookup( + headers: &BTreeMap, + prefix: &str, +) -> Option { + let primary = parse_window_from_lookup(headers, &format!("{prefix}-primary")); + let secondary = parse_window_from_lookup(headers, &format!("{prefix}-secondary")); + if primary.is_none() && secondary.is_none() { + return None; + } + Some(RateLimitDetails { + allowed: None, + limit_reached: None, + primary, + secondary, + }) +} + +fn parse_window_from_lookup( + headers: &BTreeMap, + prefix: &str, +) -> Option { + let used_percent = headers + .get(&format!("{prefix}-used-percent")) + .and_then(|value| parse_finite_percent(value))?; + let window_minutes = headers + .get(&format!("{prefix}-window-minutes")) + .and_then(|value| parse_positive_u64(value)); + let reset_at = headers + .get(&format!("{prefix}-reset-at")) + .and_then(|value| parse_positive_i64(value)); + + Some(RateLimitWindow { + used_percent, + window_minutes, + reset_at, + }) +} + +fn parse_details_from_object(value: &Value) -> Option { + let primary = value.get("primary").and_then(parse_window_from_object); + let secondary = value.get("secondary").and_then(parse_window_from_object); + if primary.is_none() && secondary.is_none() { + return None; + } + Some(RateLimitDetails { + allowed: value.get("allowed").and_then(Value::as_bool), + limit_reached: value.get("limit_reached").and_then(Value::as_bool), + primary, + secondary, + }) +} + +fn parse_window_from_object(value: &Value) -> Option { + let used_percent = value + .get("used_percent") + .and_then(Value::as_f64) + .filter(|value| value.is_finite())?; + let window_minutes = value.get("window_minutes").and_then(value_as_positive_u64); + let reset_at = value.get("reset_at").and_then(value_as_positive_i64); + + Some(RateLimitWindow { + used_percent, + window_minutes, + reset_at, + }) +} + +fn quota_window(window: Option, include_allowed: bool) -> Value { + let Some(window) = window else { + return Value::Null; + }; + let limit_reached = window.used_percent >= 100.0; + let mut value = json!({ + "used_percent": window.used_percent, + "remaining_percent": remaining_percent(window.used_percent), + "reset_at": window.reset_at, + "limit_window_seconds": window.limit_window_seconds(), + "limit_reached": limit_reached, + }); + if include_allowed { + value["allowed"] = Value::Bool(!limit_reached); + } + value +} + +fn code_review_quota_window(details: RateLimitDetails) -> Value { + let used_percent = details.primary.map(|window| window.used_percent); + let limit_reached = details + .limit_reached + .unwrap_or_else(|| used_percent.is_some_and(|used| used >= 100.0)); + json!({ + "allowed": details.allowed.unwrap_or(true), + "limit_reached": limit_reached, + "used_percent": used_percent, + "remaining_percent": used_percent.map(remaining_percent), + "reset_at": details.primary.and_then(|window| window.reset_at), + "limit_window_seconds": details.primary.and_then(RateLimitWindow::limit_window_seconds), + }) +} + +fn remaining_percent(used_percent: f64) -> i64 { + let used = used_percent.clamp(0.0, 100.0); + (100.0 - used).round() as i64 +} + +fn parse_finite_percent(value: &str) -> Option { + value + .trim() + .parse::() + .ok() + .filter(|value| value.is_finite()) +} + +fn parse_positive_u64(value: &str) -> Option { + value.trim().parse::().ok().filter(|value| *value > 0) +} + +fn parse_positive_i64(value: &str) -> Option { + value.trim().parse::().ok().filter(|value| *value > 0) +} + +fn value_as_positive_u64(value: &Value) -> Option { + value + .as_u64() + .filter(|value| *value > 0) + .or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok())) + .filter(|value| *value > 0) +} + +fn value_as_positive_i64(value: &Value) -> Option { + value + .as_i64() + .filter(|value| *value > 0) + .or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok())) + .filter(|value| *value > 0) +} + +fn rate_limit_name(value: &Value) -> Option<&str> { + value + .get("metered_limit_name") + .or_else(|| value.get("limit_name")) + .and_then(Value::as_str) +} + +fn is_review_limit_name(value: Option<&str>) -> bool { + let Some(value) = value else { + return false; + }; + let normalized = value.trim().to_ascii_lowercase().replace('-', "_"); + matches!( + normalized.as_str(), + "review" | "code_review" | "codex_review" | "codex_code_review" + ) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + parse_rate_limit_headers, parse_rate_limits_event, rate_limit_quota, RateLimitWindow, + }; + + #[test] + fn parse_rate_limit_headers_should_extract_primary_secondary_and_review_windows() { + let headers = vec![ + ( + "x-codex-primary-used-percent".to_string(), + "100".to_string(), + ), + ( + "x-codex-primary-window-minutes".to_string(), + "5".to_string(), + ), + ( + "x-codex-primary-reset-at".to_string(), + "1893456300".to_string(), + ), + ( + "x-codex-secondary-used-percent".to_string(), + "42.5".to_string(), + ), + ( + "x-codex-secondary-window-minutes".to_string(), + "10080".to_string(), + ), + ( + "x-codex-code-review-primary-used-percent".to_string(), + "80".to_string(), + ), + ( + "x-codex-code-review-primary-reset-at".to_string(), + "1893456600".to_string(), + ), + ]; + + let parsed = parse_rate_limit_headers(&headers).unwrap(); + + assert_eq!( + parsed.primary, + Some(RateLimitWindow { + used_percent: 100.0, + window_minutes: Some(5), + reset_at: Some(1_893_456_300), + }) + ); + assert_eq!(parsed.secondary.unwrap().window_minutes, Some(10080)); + assert_eq!( + parsed.code_review.unwrap().primary.unwrap().reset_at, + Some(1_893_456_600) + ); + } + + #[test] + fn parse_rate_limits_event_should_extract_internal_websocket_rate_limits() { + let event = json!({ + "type": "codex.rate_limits", + "rate_limits": { + "primary": { + "used_percent": 99.5, + "window_minutes": 300, + "reset_at": 1893456300 + }, + "secondary": { + "used_percent": 10, + "window_minutes": 10080, + "reset_at": 1894056000 + } + } + }); + + let parsed = parse_rate_limits_event(&event).unwrap(); + + assert_eq!(parsed.primary.unwrap().used_percent, 99.5); + assert_eq!(parsed.secondary.unwrap().reset_at, Some(1_894_056_000)); + } + + #[test] + fn rate_limit_quota_should_preserve_existing_credits_when_passive_data_lacks_credits() { + let headers = vec![ + ("x-codex-primary-used-percent".to_string(), "25".to_string()), + ( + "x-codex-primary-window-minutes".to_string(), + "5".to_string(), + ), + ( + "x-codex-primary-reset-at".to_string(), + "1893456300".to_string(), + ), + ]; + let existing = json!({ + "credits": { + "has_credits": true, + "unlimited": false, + "balance": 12 + } + }); + let parsed = parse_rate_limit_headers(&headers).unwrap(); + + let quota = rate_limit_quota(&parsed, Some("plus"), Some(&existing)); + + assert_eq!(quota["plan_type"], "plus"); + assert_eq!(quota["rate_limit"]["remaining_percent"], 75); + assert_eq!(quota["credits"]["balance"], 12); + } +} diff --git a/src/codex/gateway/transport/websocket.rs b/src/codex/gateway/transport/websocket.rs index 33761719e..b7178aa4e 100644 --- a/src/codex/gateway/transport/websocket.rs +++ b/src/codex/gateway/transport/websocket.rs @@ -21,7 +21,11 @@ use futures::{channel::mpsc, SinkExt, Stream, StreamExt}; use reqwest::{header::HeaderMap, StatusCode}; use serde_json::{json, Value}; -use crate::codex::gateway::transport::{sse::encode_sse_event, types::CodexResponsesRequest}; +use crate::codex::gateway::transport::{ + rate_limits::{parse_rate_limits_event_raw, rate_limits_to_header_pairs, ParsedRateLimits}, + sse::encode_sse_event, + types::CodexResponsesRequest, +}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexTransport { @@ -66,12 +70,14 @@ pub struct CodexWebSocketResponse { pub type CodexWebSocketSseStream = Pin> + Send>>; +pub type SharedRateLimitUpdates = Arc>>; pub struct CodexWebSocketStreamResponse { pub body_stream: CodexWebSocketSseStream, pub turn_state: Option, pub set_cookie_headers: Vec, pub rate_limit_headers: Vec<(String, String)>, + pub rate_limit_updates: SharedRateLimitUpdates, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -637,7 +643,8 @@ pub async fn create_response_via_websocket( mut body_stream, turn_state, set_cookie_headers, - rate_limit_headers, + mut rate_limit_headers, + rate_limit_updates, } = create_response_via_websocket_stream(base_url, request, headers).await?; let mut body = String::new(); @@ -647,6 +654,7 @@ pub async fn create_response_via_websocket( if body.is_empty() { return Err(CodexWebSocketError::EmptyResponse); } + append_rate_limit_updates(&mut rate_limit_headers, &rate_limit_updates).await; Ok(CodexWebSocketResponse { body, @@ -697,6 +705,7 @@ async fn create_response_via_websocket_stream_inner( ) -> Result { let mut pool_context = pool; let mut retry_stale_reuse = pool_context.is_some(); + let rate_limit_updates = Arc::new(Mutex::new(Vec::new())); loop { let mut active = acquire_websocket(base_url, headers.clone(), pool_context.clone()).await?; if let Err(error) = active @@ -746,6 +755,7 @@ async fn create_response_via_websocket_stream_inner( continue; }; if is_internal_websocket_event(&raw) { + capture_internal_rate_limit_event(&raw, &rate_limit_updates).await; continue; } if let Some(classified) = classify_ws_error_frame(&raw) { @@ -774,8 +784,9 @@ async fn create_response_via_websocket_stream_inner( if terminal { active.finish().await; } else { + let rate_limit_updates = rate_limit_updates.clone(); tokio::spawn(async move { - forward_websocket_as_sse(active, tx).await; + forward_websocket_as_sse(active, tx, rate_limit_updates).await; }); } return Ok(CodexWebSocketStreamResponse { @@ -783,6 +794,7 @@ async fn create_response_via_websocket_stream_inner( turn_state: metadata.turn_state, set_cookie_headers: metadata.set_cookie_headers, rate_limit_headers: metadata.rate_limit_headers, + rate_limit_updates, }); } } @@ -982,6 +994,7 @@ fn is_terminal_websocket_event(event: &str) -> bool { async fn forward_websocket_as_sse( mut active: ActiveWebSocket, tx: mpsc::UnboundedSender>, + rate_limit_updates: SharedRateLimitUpdates, ) { while let Some(message) = active.websocket.next().await { match message { @@ -991,6 +1004,7 @@ async fn forward_websocket_as_sse( continue; }; if is_internal_websocket_event(&raw) { + capture_internal_rate_limit_event(&raw, &rate_limit_updates).await; continue; } let event = websocket_event_type(&raw); @@ -1016,6 +1030,22 @@ async fn forward_websocket_as_sse( active.discard().await; } +async fn capture_internal_rate_limit_event(raw: &str, updates: &SharedRateLimitUpdates) { + if let Some(parsed) = parse_rate_limits_event_raw(raw) { + updates.lock().await.push(parsed); + } +} + +pub async fn append_rate_limit_updates( + headers: &mut Vec<(String, String)>, + updates: &SharedRateLimitUpdates, +) { + let updates = updates.lock().await; + for update in updates.iter() { + headers.extend(rate_limits_to_header_pairs(update)); + } +} + struct ClassifiedWebSocketError { status: StatusCode, connection_fatal: bool, @@ -1105,7 +1135,14 @@ fn rate_limit_headers(headers: &WsHeaderMap) -> Vec<(String, String)> { fn is_rate_limit_header(name: &str) -> bool { let name = name.to_ascii_lowercase(); - name == "retry-after" || name.contains("ratelimit") || name.contains("rate-limit") + name == "retry-after" + || name.contains("ratelimit") + || name.contains("rate-limit") + || name.starts_with("x-codex-primary-") + || name.starts_with("x-codex-secondary-") + || name.starts_with("x-codex-code-review-") + || name.starts_with("x-codex-review-") + || name.starts_with("x-code-review-") } fn retry_after_seconds(headers: &WsHeaderMap) -> Option { diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index bc1d5bf6a..afbc2bb25 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -39,10 +39,12 @@ use crate::{ build_reqwest_client, CodexBackendClient, CodexBackendStream, CodexBackendWebSocketStream, CodexClientError, CodexRequestContext, }, + rate_limits::{cooldown_with_jitter, parse_rate_limit_headers, rate_limit_quota}, types::CodexResponsesRequest, usage::{extract_sse_usage, TokenUsage}, websocket::{ - transport_for_request, CodexTransport, CodexWebSocketError, CodexWebSocketPool, + append_rate_limit_updates, transport_for_request, CodexTransport, CodexWebSocketError, + CodexWebSocketPool, SharedRateLimitUpdates, }, }, codex::logs::event::{EventLevel, EventLog}, @@ -354,8 +356,15 @@ async fn responses_http_sse_stream( .into_response(); } + let rate_limit_headers = stream_response.rate_limit_headers.clone(); let upstream = Box::pin(stream_response.response.bytes_stream()); - let audit = StreamAudit::new(deps, log_context, account.id); + let audit = StreamAudit::new( + deps, + log_context, + account.id, + account.plan_type, + rate_limit_headers, + ); use tokio::time::{interval, Duration}; const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); @@ -432,7 +441,7 @@ async fn send_codex_request_with_refresh_retry_deps( account: &Account, request_id: &str, ) -> Result { - match send_codex_request(deps, request, account, request_id).await { + let result = match send_codex_request(deps, request, account, request_id).await { Err(CodexClientError::Upstream { status, body, @@ -450,7 +459,17 @@ async fn send_codex_request_with_refresh_retry_deps( send_codex_request(deps, request, &refreshed, request_id).await } result => result, + }; + if let Ok(response) = &result { + apply_rate_limit_headers_with_deps( + deps, + &account.id, + account.plan_type.as_deref(), + &response.rate_limit_headers, + ) + .await; } + result } async fn send_codex_request( @@ -702,6 +721,63 @@ async fn record_response_affinity_with_deps( .await; } +async fn apply_rate_limit_headers_with_deps( + deps: &CodexUpstreamDependencies, + account_id: &str, + plan_type: Option<&str>, + rate_limit_headers: &[(String, String)], +) { + let Some(rate_limits) = parse_rate_limit_headers(rate_limit_headers) else { + return; + }; + + let existing_quota = existing_quota_json(deps, account_id).await; + let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); + if let Some(repo) = deps.account_repository.as_ref() { + if let Err(error) = repo.update_quota_json(account_id, "a.to_string()).await { + tracing::warn!( + error = %error, + account_id = %account_id, + "被动同步 quota 缓存失败" + ); + } + } + + let Some(reset_at) = rate_limits.primary_reset_at() else { + return; + }; + deps.account_pool.lock().await.sync_rate_limit_window( + account_id, + reset_at, + rate_limits.primary_limit_window_seconds(), + ); + if !rate_limits.primary_limit_reached() || reset_at <= Utc::now() { + return; + } + + if let Some(repo) = deps.account_repository.as_ref() { + if let Err(error) = repo.set_quota_cooldown_until(account_id, reset_at).await { + tracing::warn!( + error = %error, + account_id = %account_id, + cooldown_until = %reset_at, + "持久化被动 quota cooldown 失败" + ); + } + } + deps.account_pool + .lock() + .await + .mark_quota_limited_until(account_id, reset_at); + deps.websocket_pool.evict_account(account_id).await; +} + +async fn existing_quota_json(deps: &CodexUpstreamDependencies, account_id: &str) -> Option { + let repo = deps.account_repository.as_ref()?; + let raw = repo.get_quota_json(account_id).await.ok().flatten()?; + serde_json::from_str(&raw).ok() +} + async fn apply_upstream_retry_and_acquire_fallback_with_deps( deps: &CodexUpstreamDependencies, account: &Account, @@ -730,7 +806,7 @@ async fn apply_upstream_account_retry_with_deps( UpstreamAccountRetry::RateLimited { retry_after_seconds, } => { - let cooldown_until = Utc::now() + Duration::seconds(retry_after_seconds as i64); + let cooldown_until = Utc::now() + cooldown_with_jitter(retry_after_seconds, 2_000); if let Some(repo) = deps.account_repository.as_ref() { if let Err(error) = repo .set_quota_cooldown_until(&account.id, cooldown_until) @@ -918,14 +994,17 @@ async fn responses_websocket_stream( .into_response(); } - let audit = WebSocketStreamAudit::new( + let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account.id); + let audit = WebSocketStreamAudit { deps, - log_context, - account.id, + context: log_context, + account_slot, + account_plan_type: account.plan_type, request, - stream_response.turn_state, - stream_response.rate_limit_headers, - ); + turn_state: stream_response.turn_state, + rate_limit_headers: stream_response.rate_limit_headers, + rate_limit_updates: stream_response.rate_limit_updates, + }; let upstream = stream_response.body_stream; use tokio::time::{interval, Duration}; @@ -1046,6 +1125,8 @@ struct StreamAudit { deps: CodexUpstreamDependencies, context: CodexRequestLogContext, account_slot: AccountSlotGuard, + account_plan_type: Option, + rate_limit_headers: Vec<(String, String)>, } impl StreamAudit { @@ -1053,16 +1134,27 @@ impl StreamAudit { deps: CodexUpstreamDependencies, context: CodexRequestLogContext, account_id: String, + account_plan_type: Option, + rate_limit_headers: Vec<(String, String)>, ) -> Self { let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id); Self { deps, context, account_slot, + account_plan_type, + rate_limit_headers, } } async fn complete(&mut self, body: &[u8]) { + apply_rate_limit_headers_with_deps( + &self.deps, + &self.context.account_id, + self.account_plan_type.as_deref(), + &self.rate_limit_headers, + ) + .await; let body = String::from_utf8_lossy(body); let mut status = StatusCode::OK; let mut level = EventLevel::Info; @@ -1136,32 +1228,24 @@ struct WebSocketStreamAudit { deps: CodexUpstreamDependencies, context: CodexRequestLogContext, account_slot: AccountSlotGuard, + account_plan_type: Option, request: CodexResponsesRequest, turn_state: Option, rate_limit_headers: Vec<(String, String)>, + rate_limit_updates: SharedRateLimitUpdates, } impl WebSocketStreamAudit { - fn new( - deps: CodexUpstreamDependencies, - context: CodexRequestLogContext, - account_id: String, - request: CodexResponsesRequest, - turn_state: Option, - rate_limit_headers: Vec<(String, String)>, - ) -> Self { - let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id); - Self { - deps, - context, - account_slot, - request, - turn_state, - rate_limit_headers, - } - } - async fn complete(&mut self, body: &[u8]) { + let mut rate_limit_headers = self.rate_limit_headers.clone(); + append_rate_limit_updates(&mut rate_limit_headers, &self.rate_limit_updates).await; + apply_rate_limit_headers_with_deps( + &self.deps, + &self.context.account_id, + self.account_plan_type.as_deref(), + &rate_limit_headers, + ) + .await; let body = String::from_utf8_lossy(body).into_owned(); let usage_result = extract_sse_usage(&body); let response_usage = match &usage_result { diff --git a/src/codex/serving/dispatch/usage.rs b/src/codex/serving/dispatch/usage.rs index da74c9b98..3cb970a78 100644 --- a/src/codex/serving/dispatch/usage.rs +++ b/src/codex/serving/dispatch/usage.rs @@ -9,6 +9,12 @@ pub(super) async fn record_usage_with_deps( account_id: &str, usage: TokenUsage, ) -> Result<(), ()> { + deps.account_pool.lock().await.record_window_token_usage( + account_id, + usage.input_tokens, + usage.output_tokens, + usage.cached_tokens, + ); let Some(repo) = deps.account_repository.as_ref() else { return Ok(()); }; diff --git a/tests/codex_accounts/pool_scheduling.rs b/tests/codex_accounts/pool_scheduling.rs index bdda5b8d5..4465df1eb 100644 --- a/tests/codex_accounts/pool_scheduling.rs +++ b/tests/codex_accounts/pool_scheduling.rs @@ -78,6 +78,42 @@ fn account_pool_should_reuse_quota_limited_accounts_after_cooldown() { ); } +#[test] +fn account_pool_should_not_shorten_existing_quota_cooldown() { + let now = fixed_time(); + let mut pool = AccountPool::default(); + pool.insert(Account::test("limited", AccountStatus::Active)); + + pool.mark_quota_limited_until("limited", now + Duration::seconds(180)); + pool.mark_quota_limited_until("limited", now + Duration::seconds(60)); + + assert!(pool + .acquire_with(AccountAcquireRequest::new( + "gpt-5.5", + now + Duration::seconds(90) + )) + .is_none()); +} + +#[test] +fn account_pool_should_not_replace_known_window_length_with_cooldown_seconds() { + let now = Utc::now(); + let mut account = Account::test("limited", AccountStatus::Active); + account.limit_window_seconds = Some(300); + let mut pool = AccountPool::default(); + pool.insert(account); + + pool.mark_quota_limited_until("limited", now + Duration::seconds(60)); + let acquired = pool + .acquire_with(AccountAcquireRequest::new( + "gpt-5.5", + now + Duration::seconds(61), + )) + .unwrap(); + + assert_eq!(acquired.account.limit_window_seconds, Some(300)); +} + #[test] fn acquire_should_skip_accounts_with_expired_token_metadata() { let now = fixed_time(); @@ -380,6 +416,50 @@ fn acquire_should_mark_runtime_usage_window() { ); } +#[test] +fn sync_rate_limit_window_should_reset_window_counters_when_drift_crosses_threshold() { + let now = fixed_time(); + let mut account = Account::test("acct_a", AccountStatus::Active); + account.window_request_count = 5; + account.window_input_tokens = 100; + account.window_output_tokens = 40; + account.window_cached_tokens = 20; + account.window_reset_at = Some(now + Duration::seconds(60)); + account.limit_window_seconds = Some(60); + let mut pool = AccountPool::default(); + pool.insert(account); + + pool.sync_rate_limit_window("acct_a", now + Duration::seconds(121), Some(60)); + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert_eq!(acquired.account.window_request_count, 1); + assert_eq!(acquired.account.window_input_tokens, 0); + assert_eq!(acquired.account.window_output_tokens, 0); + assert_eq!(acquired.account.window_cached_tokens, 0); + assert_eq!( + acquired.account.window_reset_at, + Some(now + Duration::seconds(121)) + ); +} + +#[test] +fn record_window_token_usage_should_accumulate_runtime_window_tokens() { + let now = fixed_time(); + let mut pool = AccountPool::default(); + pool.insert(Account::test("acct_a", AccountStatus::Active)); + + pool.record_window_token_usage("acct_a", 7, 4, 2); + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert_eq!(acquired.account.window_input_tokens, 7); + assert_eq!(acquired.account.window_output_tokens, 4); + assert_eq!(acquired.account.window_cached_tokens, 2); +} + fn fixed_time() -> chrono::DateTime { Utc.with_ymd_and_hms(2026, 6, 11, 8, 0, 0).unwrap() } diff --git a/tests/codex_gateway/websocket.rs b/tests/codex_gateway/websocket.rs index c7988fa5c..4180fb68c 100644 --- a/tests/codex_gateway/websocket.rs +++ b/tests/codex_gateway/websocket.rs @@ -665,6 +665,88 @@ async fn websocket_success_should_capture_handshake_headers_and_rate_limit_heade .any(|(name, value)| { name == "x-ratelimit-remaining-requests" && value == "41" })); } +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_success_should_capture_internal_rate_limit_events_without_forwarding_them() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) + .await + .unwrap(); + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "codex.rate_limits", + "rate_limits": { + "primary": { + "used_percent": 100, + "window_minutes": 5, + "reset_at": 1893456300 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + websocket_completed_response("resp_ws_rate_limits", 1, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + + let mut request = base_request(); + request.use_websocket = true; + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), + ); + + let response = client + .create_response( + &request, + CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_ws_rate_limits", + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + installation_id: None, + session_id: None, + }, + ) + .await + .unwrap(); + + server.await.unwrap(); + assert!(!response.body.contains("codex.rate_limits")); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| { name == "x-codex-primary-used-percent" && value == "100" })); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| { name == "x-codex-primary-reset-at" && value == "1893456300" })); +} + #[tokio::test] #[expect( clippy::result_large_err, diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs index 96456576d..f572d203a 100644 --- a/tests/codex_serving/responses_http_sse.rs +++ b/tests/codex_serving/responses_http_sse.rs @@ -2,6 +2,7 @@ use axum::{ body::Body, http::{Request, StatusCode}, }; +use chrono::Utc; use serde_json::{json, Value}; use tower::ServiceExt; use wiremock::{ @@ -144,6 +145,66 @@ async fn v1_responses_should_use_imported_account_and_record_usage() { assert_eq!(event.4["usage"]["inputTokens"], 7); } +#[tokio::test] +async fn v1_responses_should_passively_cache_rate_limit_headers() { + let server = MockServer::start().await; + let reset_at = Utc::now().timestamp() + 300; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header("x-codex-primary-used-percent", "100") + .insert_header("x-codex-primary-window-minutes", "5") + .insert_header("x-codex-primary-reset-at", reset_at.to_string()) + .set_body_string(COMPLETED_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + sqlx::query( + r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}' where id = ?"#, + ) + .bind("acct_imported") + .execute(&imported.pool) + .await + .unwrap(); + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let stored: (String, i64, Option) = sqlx::query_as( + "select quota_json, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + let quota: Value = serde_json::from_str(&stored.0).unwrap(); + assert_eq!(quota["rate_limit"]["limit_reached"], true); + assert_eq!(quota["rate_limit"]["reset_at"], reset_at); + assert_eq!(quota["credits"]["balance"], 12); + assert_eq!(stored.1, 1); + assert!(stored.2.is_some()); +} + #[tokio::test] async fn v1_responses_should_record_empty_response_attempts() { let server = MockServer::start().await;