fix(core): narrow capacity retry ownership

codex-rs/codex-api/src/api_bridge.rs

@@ -38,6 +38,9 @@ pub fn map_api_error(err: ApiError) -> CodexErr {
         }
         ApiError::InvalidRequest { message } => CodexErr::InvalidRequest(message),
         ApiError::CyberPolicy { message } => CodexErr::CyberPolicy { message },
+        ApiError::Transport(ref transport) if is_server_overloaded_transport_error(transport) => {
+            CodexErr::ServerOverloaded
+        }
         ApiError::Transport(transport) => match transport {
             TransportError::Http {
                 status,
@@ -47,19 +50,6 @@ pub fn map_api_error(err: ApiError) -> CodexErr {
             } => {
                 let body_text = body.unwrap_or_default();
 
-                if status == http::StatusCode::SERVICE_UNAVAILABLE
-                    && let Ok(value) = serde_json::from_str::<serde_json::Value>(&body_text)
-                    && matches!(
-                        value
-                            .get("error")
-                            .and_then(|error| error.get("code"))
-                            .and_then(serde_json::Value::as_str),
-                        Some("server_is_overloaded" | "slow_down")
-                    )
-                {
-                    return CodexErr::ServerOverloaded;
-                }
-
                 if status == http::StatusCode::BAD_REQUEST {
                     if let Ok(parsed) = serde_json::from_str::<Value>(&body_text)
                         && let Some(error) = parsed.get("error")
@@ -141,6 +131,28 @@ pub fn map_api_error(err: ApiError) -> CodexErr {
     }
 }
 
+pub(crate) fn is_server_overloaded_transport_error(err: &TransportError) -> bool {
+    let TransportError::Http {
+        status,
+        body: Some(body),
+        ..
+    } = err
+    else {
+        return false;
+    };
+    *status == http::StatusCode::SERVICE_UNAVAILABLE
+        && serde_json::from_str::<serde_json::Value>(body)
+            .ok()
+            .is_some_and(|value| {
+                matches!(
+                    value
+                        .pointer("/error/code")
+                        .and_then(serde_json::Value::as_str),
+                    Some("server_is_overloaded" | "slow_down")
+                )
+            })
+}
+
 const ACTIVE_LIMIT_HEADER: &str = "x-codex-active-limit";
 const REQUEST_ID_HEADER: &str = "x-request-id";
 const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";

codex-rs/codex-api/src/api_bridge_tests.rs

@@ -26,6 +26,36 @@ fn map_api_error_maps_server_overloaded_from_503_body() {
     assert!(matches!(err, CodexErr::ServerOverloaded));
 }
 
+#[test]
+fn server_overload_transport_detection_is_narrow() {
+    let error = |status, body: &str| TransportError::Http {
+        status,
+        url: None,
+        headers: None,
+        body: Some(body.to_string()),
+    };
+    for code in ["server_is_overloaded", "slow_down"] {
+        let body = serde_json::json!({ "error": { "code": code } }).to_string();
+        assert!(is_server_overloaded_transport_error(&error(
+            http::StatusCode::SERVICE_UNAVAILABLE,
+            &body,
+        )));
+    }
+
+    assert!(!is_server_overloaded_transport_error(&error(
+        http::StatusCode::SERVICE_UNAVAILABLE,
+        "unavailable",
+    )));
+    assert!(!is_server_overloaded_transport_error(&error(
+        http::StatusCode::SERVICE_UNAVAILABLE,
+        r#"{"error":{"code":"other"}}"#,
+    )));
+    assert!(!is_server_overloaded_transport_error(&error(
+        http::StatusCode::INTERNAL_SERVER_ERROR,
+        r#"{"error":{"code":"server_is_overloaded"}}"#,
+    )));
+}
+
 #[test]
 fn map_api_error_maps_cloudflare_blocked_response_to_user_message() {
     let mut headers = HeaderMap::new();

codex-rs/codex-api/src/endpoint/responses.rs

@@ -26,6 +26,7 @@ use tracing::instrument;
 pub struct ResponsesClient<T: HttpTransport> {
     session: EndpointSession<T>,
     sse_telemetry: Option<Arc<dyn SseTelemetry>>,
+    defer_server_overloaded_retries: bool,
 }
 
 #[derive(Default)]
@@ -43,6 +44,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
         Self {
             session: EndpointSession::new(transport, provider, auth),
             sse_telemetry: None,
+            defer_server_overloaded_retries: false,
         }
     }
 
@@ -54,9 +56,16 @@ impl<T: HttpTransport> ResponsesClient<T> {
         Self {
             session: self.session.with_request_telemetry(request),
             sse_telemetry: sse,
+            defer_server_overloaded_retries: self.defer_server_overloaded_retries,
         }
     }
 
+    /// Return capacity-coded 503s to a caller that applies a slower retry policy.
+    pub fn with_deferred_server_overload_retries(mut self, defer: bool) -> Self {
+        self.defer_server_overloaded_retries = defer;
+        self
+    }
+
     #[instrument(
         name = "responses.stream_request",
         level = "info",
@@ -93,8 +102,14 @@ impl<T: HttpTransport> ResponsesClient<T> {
             insert_header(&mut headers, "x-openai-subagent", &subagent);
         }
 
-        self.stream_encoded(body, headers, compression, turn_state)
-            .await
+        self.stream_encoded(
+            body,
+            headers,
+            compression,
+            turn_state,
+            self.defer_server_overloaded_retries,
+        )
+        .await
     }
 
     fn path() -> &'static str {
@@ -121,8 +136,14 @@ impl<T: HttpTransport> ResponsesClient<T> {
     ) -> Result<ResponseStream, ApiError> {
         let body = EncodedJsonBody::encode(&body)
             .map_err(|e| ApiError::Stream(format!("failed to encode responses request: {e}")))?;
-        self.stream_encoded(body, extra_headers, compression, turn_state)
-            .await
+        self.stream_encoded(
+            body,
+            extra_headers,
+            compression,
+            turn_state,
+            self.defer_server_overloaded_retries,
+        )
+        .await
     }
 
     async fn stream_encoded(
@@ -131,6 +152,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
         extra_headers: HeaderMap,
         compression: Compression,
         turn_state: Option<Arc<OnceLock<String>>>,
+        defer_server_overloaded_retries: bool,
     ) -> Result<ResponseStream, ApiError> {
         let request_compression = match compression {
             Compression::None => RequestCompression::None,
@@ -144,6 +166,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
                 Self::path(),
                 extra_headers,
                 Some(body),
+                defer_server_overloaded_retries,
                 |req| {
                     req.headers.insert(
                         http::header::ACCEPT,

codex-rs/codex-api/src/endpoint/session.rs

@@ -1,3 +1,4 @@
+use crate::api_bridge::is_server_overloaded_transport_error;
 use crate::auth::SharedAuthProvider;
 use crate::error::ApiError;
 use crate::provider::Provider;
@@ -107,6 +108,7 @@ impl<T: HttpTransport> EndpointSession<T> {
                     transport.execute(req).await
                 }
             },
+            |_| true,
         )
         .await?;
 
@@ -125,6 +127,7 @@ impl<T: HttpTransport> EndpointSession<T> {
         path: &str,
         extra_headers: HeaderMap,
         body: Option<EncodedJsonBody>,
+        defer_server_overloaded_retries: bool,
         configure: C,
     ) -> Result<StreamResponse, ApiError>
     where
@@ -148,6 +151,7 @@ impl<T: HttpTransport> EndpointSession<T> {
                     transport.stream(req).await
                 }
             },
+            |err| !defer_server_overloaded_retries || !is_server_overloaded_transport_error(err),
         )
         .await?;
 

codex-rs/codex-api/src/telemetry.rs

@@ -5,7 +5,7 @@ use codex_client::Response;
 use codex_client::RetryPolicy;
 use codex_client::StreamResponse;
 use codex_client::TransportError;
-use codex_client::run_with_retry;
+use codex_client::run_with_retry_if;
 use http::StatusCode;
 use std::future::Future;
 use std::sync::Arc;
@@ -70,6 +70,7 @@ pub(crate) async fn run_with_request_telemetry<T, F, Fut>(
     telemetry: Option<Arc<dyn RequestTelemetry>>,
     make_request: impl FnMut() -> Request,
     send: F,
+    retry_if: impl Fn(&TransportError) -> bool,
 ) -> Result<T, TransportError>
 where
     T: WithStatus,
@@ -78,21 +79,26 @@ where
 {
     // Wraps `run_with_retry` to attach per-attempt request telemetry for both
     // unary and streaming HTTP calls.
-    run_with_retry(policy, make_request, move |req, attempt| {
-        let telemetry = telemetry.clone();
-        let send = send.clone();
-        async move {
-            let start = Instant::now();
-            let result = send(req).await;
-            if let Some(t) = telemetry.as_ref() {
-                let (status, err) = match &result {
-                    Ok(resp) => (Some(resp.status()), None),
-                    Err(err) => (http_status(err), Some(err)),
-                };
-                t.on_request(attempt, status, err, start.elapsed());
+    run_with_retry_if(
+        policy,
+        make_request,
+        move |req, attempt| {
+            let telemetry = telemetry.clone();
+            let send = send.clone();
+            async move {
+                let start = Instant::now();
+                let result = send(req).await;
+                if let Some(t) = telemetry.as_ref() {
+                    let (status, err) = match &result {
+                        Ok(resp) => (Some(resp.status()), None),
+                        Err(err) => (http_status(err), Some(err)),
+                    };
+                    t.on_request(attempt, status, err, start.elapsed());
+                }
+                result
             }
-            result
-        }
-    })
+        },
+        retry_if,
+    )
     .await
 }