From f6cf05af1dc701322664bb3f22ee4dafbc733ffc Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 6 Oct 2026 17:44:36 +0000 Subject: [PATCH] Honor Retry-After in WebSocket error events (#51440) ## 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 --- .../src/endpoint/responses_websocket.rs | 45 +++- codex-rs/core/tests/suite/retry_after.rs | 200 ++++++++++++------ 2 files changed, 174 insertions(+), 71 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index af7fca9c41..4e041ab0a5 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -584,6 +584,8 @@ fn map_ws_error(err: WsError, url: &Url) -> ApiError { struct WrappedWebsocketError { code: Option, message: Option, + #[serde(default)] + headers: Option, } #[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)] { diff --git a/codex-rs/core/tests/suite/retry_after.rs b/codex-rs/core/tests/suite/retry_after.rs index bd15f5ac76..bceca6f5b4 100644 --- a/codex-rs/core/tests/suite/retry_after.rs +++ b/codex-rs/core/tests/suite/retry_after.rs @@ -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![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)