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)