Added back the logic to handle rate-limit errors when using API key (#3070)
A previous PR removed this when adding rate-limit errors for the ChatGPT auth path.
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
use std::io::BufRead;
|
use std::io::BufRead;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
use std::sync::OnceLock;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
@@ -7,6 +8,7 @@ use codex_login::AuthManager;
|
|||||||
use codex_login::AuthMode;
|
use codex_login::AuthMode;
|
||||||
use eventsource_stream::Eventsource;
|
use eventsource_stream::Eventsource;
|
||||||
use futures::prelude::*;
|
use futures::prelude::*;
|
||||||
|
use regex_lite::Regex;
|
||||||
use reqwest::StatusCode;
|
use reqwest::StatusCode;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
@@ -53,6 +55,8 @@ struct ErrorResponse {
|
|||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct Error {
|
struct Error {
|
||||||
r#type: Option<String>,
|
r#type: Option<String>,
|
||||||
|
#[allow(dead_code)]
|
||||||
|
code: Option<String>,
|
||||||
message: Option<String>,
|
message: Option<String>,
|
||||||
|
|
||||||
// Optional fields available on "usage_limit_reached" and "usage_not_included" errors
|
// Optional fields available on "usage_limit_reached" and "usage_not_included" errors
|
||||||
@@ -564,8 +568,9 @@ async fn process_sse<S>(
|
|||||||
if let Some(error) = error {
|
if let Some(error) = error {
|
||||||
match serde_json::from_value::<Error>(error.clone()) {
|
match serde_json::from_value::<Error>(error.clone()) {
|
||||||
Ok(error) => {
|
Ok(error) => {
|
||||||
|
let delay = try_parse_retry_after(&error);
|
||||||
let message = error.message.unwrap_or_default();
|
let message = error.message.unwrap_or_default();
|
||||||
response_error = Some(CodexErr::Stream(message, None));
|
response_error = Some(CodexErr::Stream(message, delay));
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
debug!("failed to parse ErrorResponse: {e}");
|
debug!("failed to parse ErrorResponse: {e}");
|
||||||
@@ -651,6 +656,40 @@ async fn stream_from_fixture(
|
|||||||
Ok(ResponseStream { rx_event })
|
Ok(ResponseStream { rx_event })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn rate_limit_regex() -> &'static Regex {
|
||||||
|
static RE: OnceLock<Regex> = OnceLock::new();
|
||||||
|
|
||||||
|
#[expect(clippy::unwrap_used)]
|
||||||
|
RE.get_or_init(|| Regex::new(r"Please try again in (\d+(?:\.\d+)?)(s|ms)").unwrap())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_parse_retry_after(err: &Error) -> Option<Duration> {
|
||||||
|
if err.code != Some("rate_limit_exceeded".to_string()) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
// parse the Please try again in 1.898s format using regex
|
||||||
|
let re = rate_limit_regex();
|
||||||
|
if let Some(message) = &err.message
|
||||||
|
&& let Some(captures) = re.captures(message)
|
||||||
|
{
|
||||||
|
let seconds = captures.get(1);
|
||||||
|
let unit = captures.get(2);
|
||||||
|
|
||||||
|
if let (Some(value), Some(unit)) = (seconds, unit) {
|
||||||
|
let value = value.as_str().parse::<f64>().ok()?;
|
||||||
|
let unit = unit.as_str();
|
||||||
|
|
||||||
|
if unit == "s" {
|
||||||
|
return Some(Duration::from_secs_f64(value));
|
||||||
|
} else if unit == "ms" {
|
||||||
|
return Some(Duration::from_millis(value as u64));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -871,7 +910,7 @@ mod tests {
|
|||||||
msg,
|
msg,
|
||||||
"Rate limit reached for gpt-5 in organization org-AAA on tokens per min (TPM): Limit 30000, Used 22999, Requested 12528. Please try again in 11.054s. Visit https://platform.openai.com/account/rate-limits to learn more."
|
"Rate limit reached for gpt-5 in organization org-AAA on tokens per min (TPM): Limit 30000, Used 22999, Requested 12528. Please try again in 11.054s. Visit https://platform.openai.com/account/rate-limits to learn more."
|
||||||
);
|
);
|
||||||
assert_eq!(*delay, None);
|
assert_eq!(*delay, Some(Duration::from_secs_f64(11.054)));
|
||||||
}
|
}
|
||||||
other => panic!("unexpected second event: {other:?}"),
|
other => panic!("unexpected second event: {other:?}"),
|
||||||
}
|
}
|
||||||
@@ -975,4 +1014,31 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_try_parse_retry_after() {
|
||||||
|
let err = Error {
|
||||||
|
r#type: None,
|
||||||
|
message: Some("Rate limit reached for gpt-5 in organization org- on tokens per min (TPM): Limit 1, Used 1, Requested 19304. Please try again in 28ms. Visit https://platform.openai.com/account/rate-limits to learn more.".to_string()),
|
||||||
|
code: Some("rate_limit_exceeded".to_string()),
|
||||||
|
plan_type: None,
|
||||||
|
resets_in_seconds: None
|
||||||
|
};
|
||||||
|
|
||||||
|
let delay = try_parse_retry_after(&err);
|
||||||
|
assert_eq!(delay, Some(Duration::from_millis(28)));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_try_parse_retry_after_no_delay() {
|
||||||
|
let err = Error {
|
||||||
|
r#type: None,
|
||||||
|
message: Some("Rate limit reached for gpt-5 in organization <ORG> on tokens per min (TPM): Limit 30000, Used 6899, Requested 24050. Please try again in 1.898s. Visit https://platform.openai.com/account/rate-limits to learn more.".to_string()),
|
||||||
|
code: Some("rate_limit_exceeded".to_string()),
|
||||||
|
plan_type: None,
|
||||||
|
resets_in_seconds: None
|
||||||
|
};
|
||||||
|
let delay = try_parse_retry_after(&err);
|
||||||
|
assert_eq!(delay, Some(Duration::from_secs_f64(1.898)));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user