Skip to content

Honor Retry-After in WebSocket error events #51440

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copyberry merged 1 commit into from
Oct 6, 2026
Merged

Honor Retry-After in WebSocket error events #51440

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 43 additions & 2 deletions codex-rs/codex-api/src/endpoint/responses_websocket.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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()
Expand All @@ -632,7 +644,7 @@ fn map_wrapped_websocket_error_event(
.message
.clone()
.unwrap_or_else(|| fallback_message.to_string()),
retry_after: None,
retry_after,
});
}

Expand All @@ -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,
}))
}

Expand Down Expand Up @@ -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)] {
Expand Down
200 changes: 131 additions & 69 deletions codex-rs/core/tests/suite/retry_after.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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![
Expand Down Expand Up @@ -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<()>
Expand Down Expand Up @@ -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);
})
Expand All @@ -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)
Expand All @@ -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(()));
Expand Down Expand Up @@ -1982,36 +2018,53 @@ 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| {
// Capture inference retries without unrelated startup model-discovery retries.
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?;
Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand Down
Loading