github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 16:34:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Honor Retry-After in WebSocket error events (#51440)

Browse Source 
## Why

WebSocket rate-limit and overload errors ignored `Retry-After` advice and ended turns even when the server supplied a retry delay.

## What changed

Extract retry advice from nested `error.headers`, falling back to top-level `headers` when nested advice is missing or invalid. Use the shared HTTP header validation and retry-delay parser, and preserve the delay when mapping WebSocket errors.

## Testing

Add coverage for nested-header precedence, malformed headers, HTTP header validation, and recovery after an advised rate limit. Verify that overloads retry only while advice is present, without falling back to HTTP, and that headerless rate limits, usage limits, and `flex_unavailable` errors remain terminal.

GitOrigin-RevId: 61add0c36197e10402d6d3c325332aac72a5a521
This commit is contained in:
Adam Perry @ OpenAI
2026-10-06 18:17:48 +00:00
committed by copyberry
parent bc3ebcb77d
commit f6cf05af1d
2 changed files with 174 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/endpoint/responses_websocket.rs
+43 -2
View File
@@ -584,6 +584,8 @@ fn map_ws_error(err: WsError, url: &Url) -> ApiError {
struct WrappedWebsocketError {
code: Option<String>,
message: Option<String>,
#[serde(default)]
headers: Option<Value>,
}
#[derive(Debug, Deserialize)]
@@ -616,6 +618,16 @@ fn map_wrapped_websocket_error_event(
headers,
..
} = event;
let retry_after = [
error
.as_ref()
.and_then(|error| error.headers.as_ref())
.and_then(Value::as_object),
headers.as_ref(),
]
.into_iter()
.flatten()
.find_map(|headers| RetryAfter::from_headers(&json_headers_to_http_headers(headers)));
if let Some(error) = error.as_ref()
&& let Some(code) = error.code.as_deref()
@@ -632,7 +644,7 @@ fn map_wrapped_websocket_error_event(
.message
.clone()
.unwrap_or_else(|| fallback_message.to_string()),
retry_after: None,
retry_after,
});
}
@@ -646,7 +658,7 @@ fn map_wrapped_websocket_error_event(
url: None,
headers: headers.as_ref().map(json_headers_to_http_headers),
body: Some(original_payload),
retry_after: None,
retry_after,
}))
}
@@ -1051,6 +1063,35 @@ mod tests {
assert!(body.contains("The usage limit has been reached"));
}
/// Websocket error advice uses the shared HTTP value validation and duplicate handling.
#[tokio::test(start_paused = true)]
async fn wrapped_websocket_retry_after_uses_http_header_validation() {
for (error_headers, expected_seconds) in [
(json!("invalid"), 12),
(json!({"retry-after": "\n5\n"}), 12),
(json!({"retry-after": "\t5\t"}), 5),
(json!({"Retry-After": "5", "retry-after": "30"}), 30),
] {
let payload = json!({
"type": "error",
"status": 429,
"error": {"code": "rate_limit_exceeded", "headers": error_headers},
"headers": {"retry-after": "12"}
})
.to_string();
let wrapped = parse_wrapped_websocket_error_event(&payload).unwrap();
let Some(ApiError::Transport(TransportError::Http { retry_after, .. })) =
map_wrapped_websocket_error_event(wrapped, payload)
else {
panic!("expected a websocket HTTP error");
};
assert_eq!(
retry_after,
RetryAfter::from_delay(Duration::from_secs(expected_seconds))
);
}
}
#[test]
fn wrapped_websocket_usage_limit_preserves_optional_window() {
for (window, expected) in [(Some(10080), Some(10080)), (None, None)] {
codex-rs/core/tests/suite/retry_after.rs
+131 -69
View File
@@ -1508,6 +1508,7 @@ async fn websocket_streamed_flex_unavailable_is_terminal() -> Result<()> {
});
for event in [
json!({"type": "error", "status": 429, "headers": {"retry-after": "300"}, "error": error}),
json!({"type": "error", "error": {"code": "flex_unavailable", "headers": {"retry-after": "300"}}}),
json!({"type": "response.failed", "response": {"error": error}}),
] {
let server = responses::start_websocket_server(vec![vec![
@@ -1544,6 +1545,45 @@ async fn websocket_streamed_flex_unavailable_is_terminal() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "current_thread")]
async fn websocket_usage_limit_with_malformed_nested_headers_is_terminal() -> Result<()> {
skip_if_no_network!(Ok(()));
let error = json!({
"type": "error",
"status": 429,
"headers": {"retry-after": "300"},
"error": {"type": "usage_limit_reached", "headers": "invalid"}
});
let server = responses::start_websocket_server(vec![vec![
vec![
responses::ev_response_created("prewarm"),
responses::ev_completed("prewarm"),
],
vec![error],
]])
.await;
let test = test_codex()
.with_config(|config| {
config.model_catalog = Some(bundled_models_response().expect("bundled models"));
config.model_provider.request_max_retries = Some(2);
config.model_provider.stream_max_retries = Some(2);
})
.build_with_websocket_server(&server)
.await?;
submit_user_input(&test, "surface the usage limit").await?;
assert_terminal_failure(
&test,
CodexErrorInfo::UsageLimitExceeded,
/*expected_retries*/ 0,
)
.await?;
assert_eq!(server.single_connection().len(), 2);
server.shutdown().await;
Ok(())
}
/// Network reconnects keep their own attempt count without consuming stream retry budget.
#[tokio::test(flavor = "current_thread")]
async fn connection_failures_increment_retry_telemetry_without_consuming_retry_budget() -> Result<()>
@@ -1828,32 +1868,40 @@ async fn websocket_upgrade_rejection_uses_retry_after() -> Result<()> {
Ok(())
}
// TODO(anp) respect Retry-After
/// Nested websocket retry headers are currently ignored, leaving rate-limit errors terminal.
/// Wrapped websocket rate limits honor nested advice.
#[tokio::test(flavor = "current_thread")]
async fn websocket_rate_limit_with_nested_retry_after_is_terminal() -> Result<()> {
async fn websocket_error_uses_nested_retry_after() -> Result<()> {
skip_if_no_network!(Ok(()));
let mut telemetry = RetryTelemetryCapture::install();
let server = responses::start_websocket_server(vec![vec![
let error = json!({
"type": "error",
"status": 429,
"headers": {"Retry-After": "5"},
"error": {
"type": "rate_limit_error",
"code": "rate_limit_exceeded",
"message": "Please wait.",
"headers": {"retry-after": 1}
}
});
let server = responses::start_websocket_server(vec![
vec![
responses::ev_response_created("prewarm"),
responses::ev_completed("prewarm"),
vec![
responses::ev_response_created("prewarm"),
responses::ev_completed("prewarm"),
],
vec![error],
],
vec![json!({
"type": "error",
"status": 429,
"error": {
"type": "rate_limit_error",
"code": "rate_limit_exceeded",
"message": "Rate limit exceeded.",
"headers": { "Retry-After": "1" }
}
})],
]])
vec![vec![
responses::ev_response_created("recovered"),
responses::ev_completed("recovered"),
]],
])
.await;
let test = test_codex()
.with_config(|config| {
config.model_catalog = Some(bundled_models_response().expect("bundled models"));
config.model_provider.request_max_retries = Some(0);
config.model_provider.stream_max_retries = Some(1);
})
@@ -1864,39 +1912,27 @@ async fn websocket_rate_limit_with_nested_retry_after_is_terminal() -> Result<()
.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
.await;
assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
submit_user_input(&test, "surface the websocket rate limit").await?;
let mut error_events = 0;
let mut stream_error_events = 0;
loop {
match wait_for_event(&test.codex, |_| true).await {
EventMsg::Error(error) => {
error_events += 1;
assert_eq!(
error.codex_error_info,
Some(CodexErrorInfo::ResponseTooManyFailedAttempts {
http_status_code: Some(429),
})
);
}
EventMsg::StreamError(_) => stream_error_events += 1,
EventMsg::TurnComplete(event) => {
assert_eq!(
event.error.and_then(|error| error.codex_error_info),
Some(CodexErrorInfo::ResponseTooManyFailedAttempts {
http_status_code: Some(429),
})
);
break;
}
_ => {}
}
}
assert_eq!(error_events, 1);
assert_eq!(stream_error_events, 0);
let request_count: usize = server.connections().iter().map(Vec::len).sum();
assert_eq!(request_count, 2, "expected only prewarm and terminal error");
submit_user_input(&test, "recover from the websocket error").await?;
let retry = telemetry.next_retry().await;
assert!(retry.delay <= Duration::from_secs(1));
assert_eq!(
(
retry.attempt,
retry.layer.as_str(),
retry.operation.as_str()
),
(1, "stream", "sampling")
);
assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
wait_for_turn_completion(&test).await;
assert_eq!(
server
.connections()
.iter()
.map(Vec::len)
.collect::<Vec<_>>(),
vec![2, 1]
);
assert_eq!(
telemetry.events.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
@@ -1906,7 +1942,7 @@ async fn websocket_rate_limit_with_nested_retry_after_is_terminal() -> Result<()
Ok(())
}
/// Headerless websocket rate limits complete with the same terminal error as nested headers.
/// Headerless websocket rate limits remain terminal.
#[tokio::test(flavor = "current_thread")]
async fn websocket_rate_limit_without_retry_after_is_terminal() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -1982,28 +2018,45 @@ async fn websocket_rate_limit_without_retry_after_is_terminal() -> Result<()> {
Ok(())
}
// TODO(anp) respect Retry-After
/// Websocket overloads remain terminal despite a nested retry header.
/// When an overload stops giving advice, it is terminal even if the retry limit
/// has been reached. The transport must not switch after that terminal error.
#[tokio::test(flavor = "current_thread")]
async fn websocket_overload_with_nested_retry_after_is_terminal() -> Result<()> {
async fn websocket_overload_retries_only_while_it_has_retry_after() -> Result<()> {
skip_if_no_network!(Ok(()));
let mut telemetry = RetryTelemetryCapture::install();
let server = responses::start_websocket_server(vec![vec![
let server = responses::start_websocket_server(vec![
vec![
responses::ev_response_created("prewarm"),
responses::ev_completed("prewarm"),
vec![
responses::ev_response_created("prewarm"),
responses::ev_completed("prewarm"),
],
vec![json!({
"type": "error",
"status": 503,
"error": {
"code": "server_is_overloaded",
"message": "This model is disabled.",
"headers": { "Retry-After": "1" }
}
})],
],
vec![json!({
vec![vec![json!({
"type": "error",
"status": 503,
"error": {
"code": "server_is_overloaded",
"message": "This model is disabled.",
"headers": { "Retry-After": "1" }
}
})],
]])
"error": {"code": "server_is_overloaded", "headers": {"Retry-After": "1"}}
})]],
vec![vec![json!({
"type": "error",
"status": 503,
"error": {"code": "server_is_overloaded", "headers": {"Retry-After": "1"}}
})]],
vec![vec![json!({
"type": "error",
"status": 503,
"error": {"code": "server_is_overloaded", "message": "This model is disabled."}
})]],
])
.await;
let test = test_codex()
.with_config(|config| {
@@ -2011,7 +2064,7 @@ async fn websocket_overload_with_nested_retry_after_is_terminal() -> Result<()>
config.model_catalog =
Some(bundled_models_response().expect("bundled models.json should parse"));
config.model_provider.request_max_retries = Some(2);
config.model_provider.stream_max_retries = Some(2);
config.model_provider.stream_max_retries = Some(3);
})
.build_with_websocket_server(&server)
.await?;
@@ -2020,7 +2073,13 @@ async fn websocket_overload_with_nested_retry_after_is_terminal() -> Result<()>
.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
.await;
assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
submit_user_input(&test, "reject the websocket overload despite retry advice").await?;
submit_user_input(&test, "retry the websocket overload only when advised").await?;
for attempt in 1..=3 {
let retry = telemetry.next_retry().await;
assert_eq!(retry.attempt, attempt);
assert!(retry.delay <= Duration::from_secs(1));
assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
}
let mut error_events = 0;
let mut stream_error_events = 0;
@@ -2056,13 +2115,16 @@ async fn websocket_overload_with_nested_retry_after_is_terminal() -> Result<()>
}
assert_eq!(error_events, 1);
assert_eq!(stream_error_events, 0);
assert_eq!(stream_error_events, 2 + usize::from(cfg!(debug_assertions)));
assert_eq!(
fallback_warning_events, 0,
"websocket must not fall back to HTTP"
);
let request_count: usize = server.connections().iter().map(Vec::len).sum();
assert_eq!(request_count, 2, "expected only prewarm and terminal error");
assert_eq!(
request_count, 5,
"expected prewarm, three advised errors, and a terminal error"
);
assert_eq!(
telemetry.events.try_recv(),
Err(mpsc::error::TryRecvError::Empty)