codex/codex-rs/core/tests/suite/retry_after.rs

use anyhow::Result;
use codex_api::ApiError;
use codex_api::map_api_error;
use codex_client::RetryOn;
use codex_client::RetryPolicy;
use codex_client::run_with_retry;
use codex_http_client::Request;
use codex_http_client::RetryAfter;
use codex_http_client::TransportError;
use codex_login::CodexAuth;
use codex_models_manager::bundled_models_response;
use codex_protocol::protocol::CodexErrorInfo;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::turn_input::TurnInputRequest;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::TestCodex;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use http::Method;
use http::StatusCode;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::sync::Mutex;
use std::time::Duration;
use std::time::Instant;
use tokio::net::TcpSocket;
use tokio::sync::mpsc;
use tracing::Event;
use tracing::Subscriber;
use tracing::dispatcher::DefaultGuard;
use tracing::field::Field;
use tracing::field::Visit;
use tracing::span::Attributes;
use tracing::span::Id;
use tracing_subscriber::Layer;
use tracing_subscriber::layer::Context;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;

const FIRST_RETRY_MIN_DELAY: Duration = Duration::from_millis(180);
const FIRST_RETRY_MAX_DELAY: Duration = Duration::from_millis(220);
const SECOND_RETRY_MIN_DELAY: Duration = Duration::from_millis(360);
const SECOND_RETRY_MAX_DELAY: Duration = Duration::from_millis(440);

#[derive(Debug, PartialEq, Eq)]
pub(super) struct RetryTelemetryEvent {
    attempt: u64,
    pub(super) delay: Duration,
    pub(super) layer: String,
    pub(super) operation: String,
}

#[derive(Default)]
struct RetryTelemetryVisitor {
    name: Option<String>,
    attempt: Option<u64>,
    delay_ms: Option<u64>,
    layer: Option<String>,
    operation: Option<String>,
    is_responses_request: bool,
}

impl Visit for RetryTelemetryVisitor {
    fn record_u64(&mut self, field: &Field, value: u64) {
        match field.name() {
            "retry.attempt" => self.attempt = Some(value),
            "retry.delay_ms" => self.delay_ms = Some(value),
            _ => {}
        }
    }

    fn record_i64(&mut self, field: &Field, value: i64) {
        if let Ok(value) = u64::try_from(value) {
            self.record_u64(field, value);
        }
    }

    fn record_str(&mut self, field: &Field, value: &str) {
        match field.name() {
            "event.name" => self.name = Some(value.to_string()),
            "retry.layer" => self.layer = Some(value.to_string()),
            "retry.operation" => self.operation = Some(value.to_string()),
            "message" => {
                self.is_responses_request = value
                    .split_once(": ")
                    .and_then(|(request, _)| request.split_whitespace().nth(2))
                    .is_some_and(|url| url.ends_with("/responses"));
            }
            _ => {}
        }
    }

    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
        let value = format!("{value:?}");
        self.record_str(field, value.trim_matches('"'));
    }
}

struct RetryTelemetryLayer {
    events: mpsc::UnboundedSender<RetryTelemetryEvent>,
    resumptions: mpsc::UnboundedSender<Duration>,
    last_request: Mutex<Option<Instant>>,
    pending_retry: Mutex<Option<Instant>>,
}

impl RetryTelemetryLayer {
    fn record_request_after_retry(&self) {
        let now = Instant::now();
        let mut last_request = self
            .last_request
            .lock()
            .expect("last request should not be poisoned");
        let started = self
            .pending_retry
            .lock()
            .expect("pending retry should not be poisoned")
            .take();
        if let Some(started) = started {
            let _ = self.resumptions.send(now.duration_since(started));
        }
        *last_request = Some(now);
    }
}

impl<S> Layer<S> for RetryTelemetryLayer
where
    S: Subscriber,
{
    fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
        if event.metadata().target() == "codex_http_client::transport" {
            let mut visitor = RetryTelemetryVisitor::default();
            event.record(&mut visitor);
            if visitor.is_responses_request {
                self.record_request_after_retry();
            }
            return;
        }

        if event.metadata().target() != "codex_otel.trace_safe" {
            return;
        }

        let mut visitor = RetryTelemetryVisitor::default();
        event.record(&mut visitor);
        if visitor.name.as_deref() != Some("codex.retry") {
            return;
        }

        let retry = RetryTelemetryEvent {
            attempt: visitor
                .attempt
                .expect("retry event should include an attempt"),
            delay: Duration::from_millis(
                visitor
                    .delay_ms
                    .expect("retry event should include its selected delay"),
            ),
            layer: visitor
                .layer
                .expect("retry event should identify its layer"),
            operation: visitor
                .operation
                .expect("retry event should identify its operation"),
        };
        // The deadline is captured before retry telemetry. Starting at the failed request
        // gives a lower bound that still holds if the test is descheduled before telemetry.
        let started = *self
            .last_request
            .lock()
            .expect("last request should not be poisoned");
        *self
            .pending_retry
            .lock()
            .expect("pending retry should not be poisoned") = started;
        let _ = self.events.send(retry);
    }

    fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: Context<'_, S>) {
        if matches!(
            attributes.metadata().name(),
            "responses_websocket.connect" | "responses_websocket.stream_request"
        ) {
            self.record_request_after_retry();
        }
    }
}

pub(super) struct RetryTelemetryCapture {
    events: mpsc::UnboundedReceiver<RetryTelemetryEvent>,
    resumptions: mpsc::UnboundedReceiver<Duration>,
    _subscriber: DefaultGuard,
    _interest_cache_guard: tracing::Dispatch,
}

impl RetryTelemetryCapture {
    pub(super) fn install() -> Self {
        let (sender, events) = mpsc::unbounded_channel();
        let (resumptions_sender, resumptions) = mpsc::unbounded_channel();
        // Avoid caching no interest when another test first reaches a shared callsite.
        let interest_cache_guard =
            tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default());
        let subscriber = tracing_subscriber::registry()
            .with(RetryTelemetryLayer {
                events: sender,
                resumptions: resumptions_sender,
                last_request: Mutex::new(None),
                pending_retry: Mutex::new(None),
            })
            .set_default();
        tracing::callsite::rebuild_interest_cache();

        Self {
            events,
            resumptions,
            _subscriber: subscriber,
            _interest_cache_guard: interest_cache_guard,
        }
    }

    pub(super) async fn next_retry(&mut self) -> RetryTelemetryEvent {
        tokio::time::timeout(Duration::from_secs(10), self.events.recv())
            .await
            .expect("timed out waiting for retry telemetry")
            .expect("retry telemetry subscriber should remain installed")
    }
}

async fn wait_for_retry(
    telemetry: &mut RetryTelemetryCapture,
    retry: &RetryTelemetryEvent,
) -> Duration {
    let elapsed = tokio::time::timeout(
        retry.delay + Duration::from_secs(10),
        telemetry.resumptions.recv(),
    )
    .await
    .expect("timed out waiting for the request after a retry")
    .expect("retry should start another request after its sleep");
    assert!(
        elapsed >= retry.delay,
        "{} {} retry waited {elapsed:?}, less than its selected {:?} delay",
        retry.layer,
        retry.operation,
        retry.delay
    );
    elapsed
}

async fn submit_user_input(test: &TestCodex, text: &str) -> Result<()> {
    test.codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: text.to_string(),
            text_elements: Vec::new(),
        }]))
        .await?;
    Ok(())
}

async fn wait_for_turn_completion(test: &TestCodex) {
    let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await
    else {
        unreachable!("predicate guarantees a turn complete event");
    };
    assert_eq!(completed.error, None, "turn should complete successfully");
}

/// HTTP overloads use server advice within the request and stream retry budgets.
#[tokio::test(flavor = "current_thread")]
async fn responses_http_uses_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            ResponseTemplate::new(503)
                .insert_header("Retry-After", "1")
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
            ResponseTemplate::new(503)
                .insert_header("Retry-After", "1")
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("recovered"),
                responses::ev_completed("recovered"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(1);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "retry the upstream overload").await?;
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "http".into(),
            operation: "request".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    wait_for_turn_completion(&test).await;

    assert_eq!(response_mock.requests().len(), 3);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    Ok(())
}

/// A 429 with server advice retries in the sampling loop and completes without a terminal error.
#[tokio::test(flavor = "current_thread")]
async fn responses_http_429_uses_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            ResponseTemplate::new(429)
                .insert_header("Retry-After", "1")
                .set_body_json(json!({ "error": { "code": "rate_limit_exceeded" } })),
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("recovered"),
                responses::ev_completed("recovered"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "retry the rate-limited request").await?;
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    loop {
        match wait_for_event(&test.codex, |_| true).await {
            EventMsg::Error(error) => panic!("unexpected terminal error: {error:?}"),
            EventMsg::TurnComplete(completed) => {
                assert_eq!(completed.error, None);
                break;
            }
            _ => {}
        }
    }
    assert_eq!(response_mock.requests().len(), 2);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    Ok(())
}

/// Check backoff without a live server or tracing events coordinating the retry loop.
#[tokio::test(start_paused = true)]
async fn http_retry_backoff_exhausts_attempts() {
    let attempts = Mutex::new(Vec::new());
    let result = run_with_retry(
        RetryPolicy {
            max_attempts: 2,
            base_delay: Duration::from_millis(200),
            retry_on: RetryOn {
                retry_429: false,
                retry_5xx: true,
                retry_transport: false,
            },
        },
        || Request::new(Method::POST, "http://localhost/v1/responses".into()),
        |_, attempt| {
            attempts
                .lock()
                .expect("retry attempts should not be poisoned")
                .push((attempt, tokio::time::Instant::now()));
            std::future::ready(Err::<(), _>(TransportError::Http {
                retry_after: None,
                status: StatusCode::SERVICE_UNAVAILABLE,
                url: None,
                headers: None,
                body: None,
            }))
        },
    )
    .await;

    assert!(
        matches!(result, Err(TransportError::Http { status, .. }) if status == StatusCode::SERVICE_UNAVAILABLE)
    );
    let attempts = attempts
        .into_inner()
        .expect("retry attempts should not be poisoned");
    assert_eq!(
        attempts
            .iter()
            .map(|(attempt, _)| *attempt)
            .collect::<Vec<_>>(),
        vec![0, 1, 2]
    );
    assert!(
        (FIRST_RETRY_MIN_DELAY..=FIRST_RETRY_MAX_DELAY).contains(&(attempts[1].1 - attempts[0].1))
    );
    assert!(
        (SECOND_RETRY_MIN_DELAY..=SECOND_RETRY_MAX_DELAY)
            .contains(&(attempts[2].1 - attempts[1].1))
    );
}

/// Exhausting HTTP retries and mapping the error preserve the last server deadline.
#[tokio::test(start_paused = true)]
async fn exhausted_http_retries_preserve_deadline_through_error_mapping() {
    use tokio::time::Instant;

    let started = Instant::now();
    let transport_error = run_with_retry(
        RetryPolicy {
            max_attempts: 1,
            base_delay: Duration::from_millis(200),
            retry_on: RetryOn {
                retry_429: false,
                retry_5xx: true,
                retry_transport: false,
            },
        },
        || Request::new(Method::POST, "http://localhost/v1/responses".into()),
        |_, attempt| {
            let elapsed = Instant::now() - started;
            assert_eq!(elapsed.as_secs(), attempt * 3);
            let seconds = if attempt == 0 { 3 } else { 10 };
            std::future::ready(Err::<(), _>(TransportError::Http {
                status: StatusCode::INTERNAL_SERVER_ERROR,
                url: None,
                headers: None,
                body: None,
                retry_after: RetryAfter::from_delay(Duration::from_secs(seconds)),
            }))
        },
    )
    .await
    .unwrap_err();
    assert_eq!(Instant::now() - started, Duration::from_secs(3));
    tokio::time::advance(Duration::from_secs(4)).await;
    let outer_error = map_api_error(ApiError::Transport(transport_error));
    let retry_after = outer_error.retry_after().expect("retry advice");
    assert_eq!(retry_after.deadline(), started + Duration::from_secs(13));
    assert_eq!(
        outer_error.server_retry_delay(),
        Some(Duration::from_secs(6))
    );
    tokio::time::advance(Duration::from_secs(6)).await;
    assert_eq!(outer_error.retry_after(), Some(retry_after));
    assert_eq!(outer_error.server_retry_delay(), Some(Duration::ZERO));
}

/// Retry advice changes delays but never extends either retry budget or its displayed count.
#[test_case::test_case(None, 2, 3; "without_advice")]
#[test_case::test_case(Some("0"), 0, 3; "advice_without_stream_retries")]
#[test_case::test_case(Some("0"), 2, 9; "advice_with_stream_retries")]
#[tokio::test(flavor = "current_thread")]
async fn responses_http_overload_respects_retry_limits(
    retry_after: Option<&str>,
    stream_max_retries: u64,
    expected_requests: usize,
) -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = responses::start_mock_server().await;
    let mut overload = ResponseTemplate::new(503)
        .set_body_json(json!({ "error": { "code": "server_is_overloaded" } }));
    if let Some(retry_after) = retry_after {
        overload = overload.insert_header("Retry-After", retry_after);
    }
    Mock::given(method("POST"))
        .and(path("/v1/responses"))
        .respond_with(overload)
        .mount(&server)
        .await;
    let test = test_codex()
        .with_config(move |config| {
            config.model_provider.request_max_retries = Some(2);
            config.model_provider.stream_max_retries = Some(stream_max_retries);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "reject the disabled model").await?;
    let mut error_events = 0;
    let mut reconnect_messages = Vec::new();
    tokio::time::timeout(Duration::from_secs(10), async {
        loop {
            match wait_for_event(&test.codex, |_| true).await {
                EventMsg::Error(error) => {
                    error_events += 1;
                    assert_eq!(
                        error.codex_error_info,
                        Some(CodexErrorInfo::ServerOverloaded)
                    );
                    assert_eq!(
                        error.message,
                        "Selected model is at capacity. Please try a different model."
                    );
                }
                EventMsg::StreamError(error) => reconnect_messages.push(error.message),
                EventMsg::TurnComplete(event) => {
                    assert_eq!(
                        event.error.and_then(|error| error.codex_error_info),
                        Some(CodexErrorInfo::ServerOverloaded)
                    );
                    break;
                }
                _ => {}
            }
        }
    })
    .await
    .expect("overload retries should finish the turn within 10 seconds");

    assert_eq!(error_events, 1);
    let expected_messages = if retry_after.is_some() {
        (1..=stream_max_retries)
            .map(|retry| format!("Reconnecting... {retry}/{stream_max_retries}"))
            .collect::<Vec<_>>()
    } else {
        Vec::new()
    };
    assert_eq!(reconnect_messages, expected_messages);
    let request_count = server
        .received_requests()
        .await
        .expect("mock server should record requests")
        .into_iter()
        .filter(|request| request.url.path() == "/v1/responses")
        .count();
    assert_eq!(
        request_count, expected_requests,
        "overload should exhaust the configured retry budgets"
    );

    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn responses_http_504_preserves_status_after_stream_retry() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = responses::start_mock_server().await;
    let response_mock =
        responses::mount_response_sequence(&server, vec![ResponseTemplate::new(504); 2]).await;
    let test = test_codex()
        .with_config(|config| {
            config.service_tier = Some("flex".to_string());
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "surface the timeout after one retry").await?;
    assert_terminal_failure(
        &test,
        CodexErrorInfo::HttpConnectionFailed {
            http_status_code: Some(504),
        },
        /*expected_retries*/ 1,
    )
    .await?;
    assert_eq!(response_mock.requests().len(), 2);
    Ok(())
}

/// Remote compaction v2 uses the upstream header before another HTTP request.
#[tokio::test(flavor = "current_thread")]
async fn compact_v2_uses_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("seed"),
                responses::ev_completed("seed"),
            ])),
            ResponseTemplate::new(503)
                .insert_header("Retry-After", "1")
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
            responses::sse_response(responses::sse(vec![
                json!({
                    "type": "response.output_item.done",
                    "item": {
                        "type": "compaction",
                        "encrypted_content": "RETRIED_COMPACTION_SUMMARY",
                    }
                }),
                responses::ev_completed("compacted"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(1);
            config.model_provider.stream_max_retries = Some(0);
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_turn("seed history for compaction").await?;

    test.codex.submit(Op::Compact).await?;
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "http".into(),
            operation: "request".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    wait_for_turn_completion(&test).await;

    let requests = response_mock.requests();
    assert_eq!(requests.len(), 3);
    for request in &requests[1..] {
        assert_eq!(request.path(), "/v1/responses");
        assert!(
            !request.inputs_of_type("compaction_trigger").is_empty(),
            "expected a remote compaction v2 request"
        );
    }
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    Ok(())
}

/// Remote compaction v2 stream failures retry without using the enclosing response header.
#[tokio::test(flavor = "current_thread")]
async fn compact_v2_stream_failure_uses_local_backoff_despite_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("seed"),
                responses::ev_completed("seed"),
            ])),
            responses::sse_response(responses::sse_failed(
                "rate-limited",
                "rate_limit_exceeded",
                "Rate limit exceeded.",
            ))
            .insert_header("Retry-After", "1"),
            responses::sse_response(responses::sse(vec![
                json!({
                    "type": "response.output_item.done",
                    "item": {
                        "type": "compaction",
                        "encrypted_content": "RETRIED_COMPACTION_SUMMARY",
                    }
                }),
                responses::ev_completed("compacted"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_turn("seed history for compaction").await?;

    test.codex.submit(Op::Compact).await?;
    let retry = telemetry.next_retry().await;
    assert!((FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&retry.delay));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "remote_compaction_v2".into(),
        }
    );
    wait_for_retry(&mut telemetry, &retry).await;
    wait_for_turn_completion(&test).await;

    let requests = response_mock.requests();
    assert_eq!(requests.len(), 3);
    for request in &requests[1..] {
        assert_eq!(request.path(), "/v1/responses");
        assert!(
            !request.inputs_of_type("compaction_trigger").is_empty(),
            "expected a remote compaction v2 request"
        );
    }
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// Headerless remote compaction stream rate limits exhaust retries before one terminal error.
#[tokio::test(flavor = "current_thread")]
async fn compact_v2_stream_failure_without_retry_after_exhausts_stream_retries() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("seed"),
                responses::ev_completed("seed"),
            ])),
            responses::sse_response(responses::sse_failed(
                "rate-limited",
                "rate_limit_exceeded",
                "Rate limit exceeded.",
            )),
            responses::sse_response(responses::sse_failed(
                "still-rate-limited",
                "rate_limit_exceeded",
                "Rate limit exceeded.",
            )),
        ],
    )
    .await;
    let test = test_codex()
        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_turn("seed history for compaction").await?;

    test.codex.submit(Op::Compact).await?;
    let retry = telemetry.next_retry().await;
    assert!((FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&retry.delay));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "remote_compaction_v2".into(),
        }
    );
    wait_for_retry(&mut telemetry, &retry).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::RateLimitExceeded)
                );
                assert!(error.message.contains("Rate limit exceeded."));
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::RateLimitExceeded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 1);
    let requests = response_mock.requests();
    assert_eq!(requests.len(), 3);
    for request in &requests[1..] {
        assert_eq!(request.path(), "/v1/responses");
        assert!(
            !request.inputs_of_type("compaction_trigger").is_empty(),
            "expected a remote compaction v2 request"
        );
    }
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// Remote compaction honors message-only advice within its budget, with headers taking precedence.
#[test_case::test_case(Some("1"), "Try again in 4 seconds."; "nested_header")]
#[test_case::test_case(None, "Try again in 1 second."; "plaintext")]
#[tokio::test(flavor = "current_thread")]
async fn compact_v2_rate_limit_uses_server_advice(
    error_header: Option<&str>,
    message: &str,
) -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let failed = responses::sse_response(responses::sse(vec![json!({
        "type": "response.failed",
        "response": {"error": {
            "code": "rate_limit_exceeded",
            "message": message,
            "headers": {"Retry-After": error_header}
        }}
    })]));
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("seed"),
                responses::ev_completed("seed"),
            ])),
            failed,
            responses::sse_response(responses::sse(vec![
                json!({
                    "type": "response.output_item.done",
                    "item": {
                        "type": "compaction",
                        "encrypted_content": "RETRIED_COMPACTION_SUMMARY",
                    }
                }),
                responses::ev_completed("compacted"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_turn("seed history for compaction").await?;

    test.codex.submit(Op::Compact).await?;
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "remote_compaction_v2".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    wait_for_turn_completion(&test).await;

    let requests = response_mock.requests();
    assert_eq!(requests.len(), 3);
    for request in &requests[1..] {
        assert_eq!(request.path(), "/v1/responses");
        assert!(
            !request.inputs_of_type("compaction_trigger").is_empty(),
            "expected a remote compaction v2 request"
        );
    }
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// Headerless remote compaction v2 overloads exhaust request retries before one terminal error.
#[tokio::test(flavor = "current_thread")]
async fn compact_v2_overload_without_retry_after_exhausts_request_retries() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("seed"),
                responses::ev_completed("seed"),
            ])),
            ResponseTemplate::new(503)
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
            ResponseTemplate::new(503)
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
            ResponseTemplate::new(503)
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
        ],
    )
    .await;
    let test = test_codex()
        .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(2);
            config.model_provider.stream_max_retries = Some(2);
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_turn("seed history for compaction").await?;

    test.codex.submit(Op::Compact).await?;
    let first_retry = telemetry.next_retry().await;
    assert!((FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&first_retry.delay));
    assert_eq!(
        first_retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: first_retry.delay,
            layer: "http".into(),
            operation: "request".into(),
        }
    );
    wait_for_retry(&mut telemetry, &first_retry).await;
    let second_retry = telemetry.next_retry().await;
    assert!((SECOND_RETRY_MIN_DELAY..SECOND_RETRY_MAX_DELAY).contains(&second_retry.delay));
    assert_eq!(
        second_retry,
        RetryTelemetryEvent {
            attempt: 2,
            delay: second_retry.delay,
            layer: "http".into(),
            operation: "request".into(),
        }
    );
    wait_for_retry(&mut telemetry, &second_retry).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::ServerOverloaded)
                );
                assert!(
                    error
                        .message
                        .contains("Selected model is at capacity. Please try a different model.")
                );
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::ServerOverloaded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 0);
    let requests = response_mock.requests();
    assert_eq!(
        requests.len(),
        4,
        "expected a seed request and three remote compaction v2 attempts"
    );
    for request in &requests[1..] {
        assert_eq!(request.path(), "/v1/responses");
        assert!(
            !request.inputs_of_type("compaction_trigger").is_empty(),
            "expected a remote compaction v2 request"
        );
    }
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// SSE failures currently retry with local backoff instead of the enclosing response header.
#[tokio::test(flavor = "current_thread")]
async fn sse_failure_uses_local_backoff_despite_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse(vec![
                json!({
                    "type": "error",
                    "error": {
                        "type": "tokens",
                        "code": "rate_limit_exceeded",
                        "message": "Rate limit exceeded."
                    }
                }),
                json!({
                    "type": "response.failed",
                    "response": {
                        "id": "rate-limited",
                        "status": "failed",
                        "error": {
                            "code": "rate_limit_exceeded",
                            "message": "Rate limit exceeded."
                        }
                    }
                }),
            ]))
            .insert_header("Retry-After", "1"),
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("recovered"),
                responses::ev_completed("recovered"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "retry the rate-limited stream").await?;
    let retry = telemetry.next_retry().await;
    assert!((FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&retry.delay));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    wait_for_retry(&mut telemetry, &retry).await;
    wait_for_turn_completion(&test).await;

    assert_eq!(response_mock.requests().len(), 2);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    Ok(())
}

/// Headerless sampled stream rate limits exhaust retries before one terminal error.
#[test_case::test_case("rate_limit_exceeded"; "rate_limit")]
#[test_case::test_case("slow_down"; "slow_down")]
#[tokio::test(flavor = "current_thread")]
async fn sse_failure_without_retry_after_exhausts_stream_retries(code: &str) -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            responses::sse_response(responses::sse_failed(
                "rate-limited",
                code,
                "Rate limit exceeded.",
            )),
            responses::sse_response(responses::sse_failed(
                "still-rate-limited",
                code,
                "Rate limit exceeded.",
            )),
        ],
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "exhaust the headerless rate-limited stream").await?;
    let retry = telemetry.next_retry().await;
    assert!(
        (FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&retry.delay),
        "{retry:?}",
    );
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    wait_for_retry(&mut telemetry, &retry).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::RateLimitExceeded)
                );
                assert!(error.message.contains("Rate limit exceeded."));
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::RateLimitExceeded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 1);
    assert_eq!(response_mock.requests().len(), 2);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// Both sources survive to the agent's retry loop, and valid headers always beat message advice.
#[test_case::test_case("rate_limit_exceeded", Some("1"), "Rate limit exceeded. Try again in 4 seconds."; "azure_style_header")]
#[test_case::test_case("rate_limit_exceeded", None, "Rate limit exceeded. Try again in 1 second."; "azure_style_plaintext")]
#[test_case::test_case("slow_down", Some("1"), "Please try again in 4.000s."; "slow_down_fractional_header")]
#[test_case::test_case("slow_down", None, "Please try again in 1.000s."; "slow_down_fractional_plaintext")]
#[test_case::test_case("rate_limit_exceeded", Some("bad"), "Try again in 1 second."; "malformed_header")]
#[test_case::test_case("server_is_overloaded", Some("1"), "This model is overloaded."; "overload_header")]
#[tokio::test(flavor = "current_thread")]
async fn sse_failure_uses_server_advice(
    code: &str,
    error_header: Option<&str>,
    message: &str,
) -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let failed = responses::sse_response(responses::sse(vec![json!({
        "type": "response.failed",
        "response": {"error": {
            "code": code,
            "message": message,
            "headers": {"Retry-After": error_header}
        }}
    })]));
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            failed,
            responses::sse_response(responses::sse(vec![
                responses::ev_response_created("recovered"),
                responses::ev_completed("recovered"),
            ])),
        ],
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "retry after the server's advice").await?;
    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    wait_for_turn_completion(&test).await;

    assert_eq!(response_mock.requests().len(), 2);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// A streamed backend overload remains terminal despite an enclosing retry header.
#[tokio::test(flavor = "current_thread")]
async fn sse_overload_with_retry_after_is_terminal() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_response_once(
        &server,
        responses::sse_response(responses::sse_failed(
            "disabled-model",
            "server_is_overloaded",
            "This model is disabled.",
        ))
        .insert_header("Retry-After", "1"),
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(2);
            config.model_provider.stream_max_retries = Some(2);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "reject the streamed overload despite retry advice").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::ServerOverloaded)
                );
                assert_eq!(
                    error.message,
                    "Selected model is at capacity. Please try a different model."
                );
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::ServerOverloaded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 0);
    assert_eq!(response_mock.requests().len(), 1);
    let request_count = server
        .received_requests()
        .await
        .expect("mock server should record requests")
        .into_iter()
        .filter(|request| request.url.path() == "/v1/responses")
        .count();
    assert_eq!(request_count, 1, "streamed overload must not retry");
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// A streamed backend overload without retry advice must complete with one terminal error.
#[tokio::test(flavor = "current_thread")]
async fn sse_overload_without_retry_after_is_terminal() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_mock_server().await;
    let response_mock = responses::mount_sse_once(
        &server,
        responses::sse_failed(
            "disabled-model",
            "server_is_overloaded",
            "This model is disabled.",
        ),
    )
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(2);
            config.model_provider.stream_max_retries = Some(2);
        })
        .build_with_auto_env(&server)
        .await?;

    submit_user_input(&test, "reject the streamed overload").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::ServerOverloaded)
                );
                assert_eq!(
                    error.message,
                    "Selected model is at capacity. Please try a different model."
                );
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::ServerOverloaded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 0);
    assert_eq!(response_mock.requests().len(), 1);
    let request_count = server
        .received_requests()
        .await
        .expect("mock server should record requests")
        .into_iter()
        .filter(|request| request.url.path() == "/v1/responses")
        .count();
    assert_eq!(request_count, 1, "headerless SSE overload must not retry");
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

async fn assert_terminal_failure(
    test: &TestCodex,
    expected: CodexErrorInfo,
    expected_retries: usize,
) -> Result<()> {
    let mut errors = Vec::new();
    let mut completions = Vec::new();
    let mut retries = 0;
    loop {
        match wait_for_event(&test.codex, |_| true).await {
            EventMsg::Error(error) => errors.push(error.codex_error_info),
            EventMsg::StreamError(_) => retries += 1,
            EventMsg::TurnComplete(event) => {
                completions.push(event.error.and_then(|error| error.codex_error_info));
                test.codex.submit(Op::Shutdown).await?;
            }
            EventMsg::ShutdownComplete => break,
            _ => {}
        }
    }
    assert_eq!(errors, vec![Some(expected.clone())]);
    assert_eq!(completions, vec![Some(expected)]);
    assert_eq!(retries, expected_retries);
    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn http_and_sse_flex_unavailable_are_terminal() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let error = json!({
        "type": "resource_unavailable", "code": "flex_unavailable",
        "message": "Flex capacity unavailable."
    });
    for response in [
        ResponseTemplate::new(429)
            .insert_header("Retry-After", "300")
            .set_body_json(json!({"error": error})),
        responses::sse_response(responses::sse(vec![
            json!({"type": "error", "error": error}),
        ])),
        responses::sse_response(responses::sse_failed(
            "failed",
            "flex_unavailable",
            "No capacity",
        )),
    ] {
        let server = responses::start_mock_server().await;
        let response_mock = responses::mount_response_once(&server, response).await;
        let test = test_codex()
            .with_config(|config| {
                config.service_tier = Some("flex".to_string());
                config.model_provider.request_max_retries = Some(2);
                config.model_provider.stream_max_retries = Some(2);
            })
            .build_with_auto_env(&server)
            .await?;

        submit_user_input(&test, "surface the Flex capacity failure").await?;
        assert_terminal_failure(
            &test,
            CodexErrorInfo::FlexUnavailable,
            /*expected_retries*/ 0,
        )
        .await?;
        assert_eq!(
            response_mock.single_request().body_json()["service_tier"],
            json!("flex")
        );
        assert_eq!(
            server
                .received_requests()
                .await
                .unwrap()
                .iter()
                .filter(|request| request.url.path() == "/v1/responses")
                .count(),
            1
        );
    }
    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn websocket_streamed_flex_unavailable_is_terminal() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let error = json!({
        "type": "resource_unavailable", "code": "flex_unavailable",
        "message": "Flex capacity unavailable."
    });
    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![
            vec![
                responses::ev_response_created("prewarm"),
                responses::ev_completed("prewarm"),
            ],
            vec![event],
        ]])
        .await;
        let test = test_codex()
            .with_config(|config| {
                config.service_tier = Some("flex".to_string());
                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);
            })
            .build_with_websocket_server(&server)
            .await?;

        submit_user_input(&test, "surface the Flex capacity failure").await?;
        assert_terminal_failure(
            &test,
            CodexErrorInfo::FlexUnavailable,
            /*expected_retries*/ 0,
        )
        .await?;
        let requests = server.single_connection();
        assert_eq!(requests.len(), 2, "prewarm and failed turn");
        assert_eq!(requests[1].body_json()["service_tier"], json!("flex"));
        server.shutdown().await;
    }
    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<()>
{
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let bootstrap_server = responses::start_mock_server().await;
    let unavailable_socket = TcpSocket::new_v4()?;
    unavailable_socket.bind("127.0.0.1:0".parse()?)?;
    let unavailable_address = unavailable_socket.local_addr()?;

    let test = test_codex()
        .with_config(move |config| {
            config.model_provider.base_url = Some(format!("http://{unavailable_address}/v1"));
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(0);
            config.model_provider.supports_websockets = false;
        })
        .build_with_auto_env(&bootstrap_server)
        .await?;

    submit_user_input(&test, "recover after repeated network failures").await?;

    let first_retry = telemetry.next_retry().await;
    assert_eq!(
        first_retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: Duration::from_secs(5),
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    wait_for_retry(&mut telemetry, &first_retry).await;

    let second_retry = telemetry.next_retry().await;
    assert_eq!(
        second_retry,
        RetryTelemetryEvent {
            attempt: 2,
            delay: Duration::from_secs(10),
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );

    let recovered_server = MockServer::builder()
        .listener(unavailable_socket.listen(/*backlog*/ 128)?.into_std()?)
        .start()
        .await;
    let response_mock = responses::mount_sse_once(
        &recovered_server,
        responses::sse(vec![
            responses::ev_response_created("recovered"),
            responses::ev_completed("recovered"),
        ]),
    )
    .await;

    wait_for_retry(&mut telemetry, &second_retry).await;
    wait_for_turn_completion(&test).await;

    assert_eq!(response_mock.requests().len(), 1);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );

    Ok(())
}

/// Retryable websocket errors reconnect after the delay reported by retry telemetry.
#[tokio::test(flavor = "current_thread")]
async fn websocket_connection_limit_retries_with_local_backoff() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_websocket_server(vec![
        vec![
            vec![
                responses::ev_response_created("prewarm"),
                responses::ev_completed("prewarm"),
            ],
            vec![json!({
                "type": "error",
                "status": 400,
                "error": {
                    "type": "invalid_request_error",
                    "code": "websocket_connection_limit_reached",
                    "message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."
                }
            })],
        ],
        vec![vec![
            responses::ev_response_created("recovered"),
            responses::ev_completed("recovered"),
        ]],
    ])
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
        .await;
    assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
    submit_user_input(&test, "retry after reaching the websocket connection limit").await?;
    let retry = telemetry.next_retry().await;
    assert!((FIRST_RETRY_MIN_DELAY..FIRST_RETRY_MAX_DELAY).contains(&retry.delay));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    wait_for_retry(&mut telemetry, &retry).await;
    wait_for_turn_completion(&test).await;

    let connections = server.connections();
    assert_eq!(connections.len(), 2);
    let request_count: usize = connections.iter().map(Vec::len).sum();
    assert_eq!(request_count, 3);
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    server.shutdown().await;

    Ok(())
}

/// WebSocket failed-event headers take precedence over plaintext during an allowed retry.
#[tokio::test(flavor = "current_thread")]
async fn websocket_failed_event_uses_header_over_plaintext_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_websocket_server(vec![
        vec![
            vec![
                responses::ev_response_created("prewarm"),
                responses::ev_completed("prewarm"),
            ],
            vec![json!({
                "type": "response.failed",
                "response": {
                    "error": {
                        "code": "rate_limit_exceeded",
                        "message": "Rate limit exceeded. Try again in 4 seconds.",
                        "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_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
        .await;
    assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
    submit_user_input(&test, "retry the rate-limited websocket stream").await?;

    let retry = telemetry.next_retry().await;
    assert!(retry.delay <= Duration::from_secs(1));
    assert_eq!(
        retry,
        RetryTelemetryEvent {
            attempt: 1,
            delay: retry.delay,
            layer: "stream".into(),
            operation: "sampling".into(),
        }
    );
    assert!(wait_for_retry(&mut telemetry, &retry).await >= Duration::from_secs(1));
    wait_for_turn_completion(&test).await;

    let connections = server.connections();
    assert_eq!(
        connections.iter().map(Vec::len).collect::<Vec<_>>(),
        vec![2, 1]
    );
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    server.shutdown().await;

    Ok(())
}

/// Advised WebSocket failures reach HTTP after the existing retry budget. HTTP
/// recovery uses its own retry budget and honors the preceding advice deadlines.
#[tokio::test(flavor = "current_thread")]
async fn websocket_upgrade_rejection_uses_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = responses::start_mock_server().await;
    let prewarm_rejection = Mock::given(method("GET"))
        .and(path("/v1/responses"))
        .respond_with(ResponseTemplate::new(/*s*/ 503))
        .up_to_n_times(/*n*/ 1)
        .expect(/*r*/ 1)
        .mount_as_scoped(&server)
        .await;
    Mock::given(method("GET"))
        .and(path("/v1/responses"))
        .respond_with(
            ResponseTemplate::new(/*s*/ 503)
                .insert_header("Retry-After", "1")
                .set_body_json(json!({ "error": { "code": "server_is_overloaded" } })),
        )
        .mount(&server)
        .await;
    let http_overload = ResponseTemplate::new(/*s*/ 503)
        .insert_header("Retry-After", "0")
        .set_body_json(json!({ "error": { "code": "server_is_overloaded" } }));
    let response_mock = responses::mount_response_sequence(
        &server,
        vec![
            http_overload,
            ResponseTemplate::new(/*s*/ 200)
                .insert_header("content-type", "text/event-stream")
                .set_body_string(responses::sse(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.supports_websockets = true;
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_auto_env(&server)
        .await?;

    // Prewarm is best effort; start the advised failures in the sampling retry loop.
    tokio::time::timeout(
        Duration::from_secs(10),
        prewarm_rejection.wait_until_satisfied(),
    )
    .await?;
    drop(prewarm_rejection);
    let start = Instant::now();
    submit_user_input(&test, "retry the rejected websocket upgrade").await?;
    wait_for_turn_completion(&test).await;

    // Both WebSocket advice deadlines apply, including the one before transport fallback.
    assert!(start.elapsed() >= Duration::from_secs(2));
    let requests = server.received_requests().await.unwrap_or_default();
    let methods = requests
        .iter()
        .filter(|request| request.url.path() == "/v1/responses")
        .map(|request| request.method.as_str())
        .collect::<Vec<_>>();
    assert_eq!(methods, vec!["GET", "GET", "GET", "POST", "POST"]);
    assert_eq!(response_mock.requests().len(), 2);
    Ok(())
}

/// Wrapped websocket rate limits honor nested advice.
#[tokio::test(flavor = "current_thread")]
async fn websocket_error_uses_nested_retry_after() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    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![
            vec![
                responses::ev_response_created("prewarm"),
                responses::ev_completed("prewarm"),
            ],
            vec![error],
        ],
        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);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
        .await;
    assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
    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)
    );
    server.shutdown().await;

    Ok(())
}

/// 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(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_websocket_server(vec![vec![
        vec![
            responses::ev_response_created("prewarm"),
            responses::ev_completed("prewarm"),
        ],
        vec![json!({
            "type": "error",
            "status": 429,
            "error": {
                "type": "rate_limit_error",
                "code": "rate_limit_exceeded",
                "message": "Rate limit exceeded."
            }
        })],
    ]])
    .await;
    let test = test_codex()
        .with_config(|config| {
            config.model_provider.request_max_retries = Some(0);
            config.model_provider.stream_max_retries = Some(1);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .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 headerless 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");
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    server.shutdown().await;

    Ok(())
}

/// 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_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![
            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![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", "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(3);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0)
        .await;
    assert_eq!(warmup.body_json()["generate"].as_bool(), Some(false));
    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;
    let mut fallback_warning_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::ServerOverloaded)
                );
                assert_eq!(
                    error.message,
                    "Selected model is at capacity. Please try a different model."
                );
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::Warning(warning)
                if warning.message.contains("Falling back from WebSockets") =>
            {
                fallback_warning_events += 1;
            }
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::ServerOverloaded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    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, 5,
        "expected prewarm, three advised errors, and a terminal error"
    );
    assert_eq!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    server.shutdown().await;

    Ok(())
}

/// Headerless websocket overloads must neither reconnect nor fall back to HTTP.
#[tokio::test(flavor = "current_thread")]
async fn websocket_overload_without_retry_after_is_terminal() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let mut telemetry = RetryTelemetryCapture::install();
    let server = responses::start_websocket_server(vec![vec![
        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."
            }
        })],
    ]])
    .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);
        })
        .build_with_websocket_server(&server)
        .await?;

    let warmup = server
        .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").await?;

    let mut error_events = 0;
    let mut stream_error_events = 0;
    let mut fallback_warning_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::ServerOverloaded)
                );
                assert_eq!(
                    error.message,
                    "Selected model is at capacity. Please try a different model."
                );
            }
            EventMsg::StreamError(_) => stream_error_events += 1,
            EventMsg::Warning(warning)
                if warning.message.contains("Falling back from WebSockets") =>
            {
                fallback_warning_events += 1;
            }
            EventMsg::TurnComplete(event) => {
                assert_eq!(
                    event.error.and_then(|error| error.codex_error_info),
                    Some(CodexErrorInfo::ServerOverloaded)
                );
                break;
            }
            _ => {}
        }
    }

    assert_eq!(error_events, 1);
    assert_eq!(stream_error_events, 0);
    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!(
        telemetry.events.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    );
    server.shutdown().await;

    Ok(())
}