Fix code-mode session admission retries for proxy HTTP 503 responses (#61153)

## Why

Proxy HTTP 503 responses can contain plain-text or HTML bodies instead of gRPC frames. These responses should reach the existing session admission retry logic as `Unavailable` errors.

## What changed

Strip body data from HTTP 503 responses to `OpenSession`, preserving headers and trailers so tonic can honor an explicit gRPC status or infer `Unavailable` when none is present. Keep this handling scoped to session admission and retain the existing five-attempt retry limit.

## Testing

Extend admission tests to cover bounded retries for plain-text 503 responses, recovery after an HTML 503 response, and an explicit `Internal` gRPC status that prevents retries.

GitOrigin-RevId: c996e13506464e8b65bae06172324bcbadb1f394

codex-rs/code-mode-host/tests/grpc/admission_tests.rs

@@ -2,6 +2,7 @@
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 
+use axum::http::StatusCode;
 use codex_code_mode::GrpcSessionOpenError;
 use pretty_assertions::assert_eq;
 
@@ -9,7 +10,22 @@ use super::*;
 
 #[tokio::test]
 async fn session_admission_retries_only_transient_failures_and_is_bounded() -> Result<()> {
-    for (code, failures, expected_attempts, succeeds) in [
+    let proxy_response = axum::http::Response::builder()
+        .status(StatusCode::SERVICE_UNAVAILABLE)
+        .header(
+            axum::http::header::CONTENT_TYPE,
+            "text/plain; charset=utf-8",
+        )
+        .body("no healthy upstream")?;
+    let mut explicit_status_response = proxy_response.clone();
+    explicit_status_response
+        .headers_mut()
+        .insert(tonic::Status::GRPC_STATUS, "13".parse()?);
+    let html_proxy_response = axum::http::Response::builder()
+        .status(StatusCode::SERVICE_UNAVAILABLE)
+        .header(axum::http::header::CONTENT_TYPE, "text/html")
+        .body("<html>no healthy upstream</html>")?;
+    let cases = [
         (Code::Unavailable, 1, 2, true),
         (Code::ResourceExhausted, 1, 2, true),
         (Code::Unavailable, usize::MAX, 5, false),
@@ -17,7 +33,29 @@ async fn session_admission_retries_only_transient_failures_and_is_bounded() -> R
         (Code::DeadlineExceeded, 1, 1, false),
         (Code::InvalidArgument, 1, 1, false),
         (Code::Unknown, 1, 1, false),
-    ] {
+    ]
+    .into_iter()
+    .map(|(code, failures, expected_attempts, succeeds)| {
+        (
+            tonic::Status::new(code, "transient admission test").into_http::<&'static str>(),
+            code,
+            failures,
+            expected_attempts,
+            succeeds,
+        )
+    })
+    .chain([
+        (proxy_response, Code::Unavailable, usize::MAX, 5, false),
+        (html_proxy_response, Code::Unavailable, 1, 2, true),
+        (
+            explicit_status_response,
+            Code::Internal,
+            usize::MAX,
+            1,
+            false,
+        ),
+    ]);
+    for (response, code, failures, expected_attempts, succeeds) in cases {
         let attempts = Arc::new(AtomicUsize::new(0));
         let count = Arc::clone(&attempts);
         let app = axum::Router::new()
@@ -27,12 +65,12 @@ async fn session_admission_retries_only_transient_failures_and_is_bounded() -> R
             .layer(axum::middleware::from_fn(
                 move |request: axum::extract::Request, next: axum::middleware::Next| {
                     let count = Arc::clone(&count);
+                    let response = response.clone();
                     async move {
                         if request.uri().path().ends_with("/OpenSession")
                             && count.fetch_add(1, Ordering::SeqCst) < failures
                         {
-                            return tonic::Status::new(code, "transient admission test")
-                                .into_http::<axum::body::Body>();
+                            return response.map(axum::body::Body::from);
                         }
                         next.run(request).await
                     }

codex-rs/code-mode/src/grpc_session/mod.rs

@@ -59,6 +59,9 @@ mod transport;
 
 type GrpcClient = CodeModeHostClient<GrpcTransport>;
 
+#[derive(Clone)]
+struct SessionAdmission;
+
 const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
 const ROUTE_HEADER: &str = "x-code-mode-route";
 // Bound result-delivery RPCs, not running tools. A full completion burst can
@@ -184,6 +187,7 @@ impl GrpcCodeModeSessionProvider {
                         }
                     },
                 });
+                request.extensions_mut().insert(SessionAdmission);
                 inject_span_traceparent(&mut request, &open_session_span);
                 let result = async {
                     let response = client.as_ref().clone().open_session(request).await?;

codex-rs/code-mode/src/grpc_session/transport.rs

@@ -6,8 +6,10 @@ use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
 use codex_http_client::ClientRouteClass;
 use codex_http_client::HttpClientFactory;
 use http_body_util::BodyExt;
+use reqwest::StatusCode;
 use tonic::Status;
 use tonic::body::Body;
+use tonic::codegen::Bytes;
 use tonic::codegen::http::Request;
 use tonic::codegen::http::Response;
 use tonic::codegen::http::Uri;
@@ -18,6 +20,7 @@ use tower::service_fn;
 use tower::util::BoxCloneSyncService;
 
 use super::GrpcClient;
+use super::SessionAdmission;
 
 pub(super) type GrpcTransport = BoxCloneSyncService<Request<Body>, Response<Body>, Status>;
 
@@ -118,6 +121,7 @@ impl SharedTransport {
             let transport = service_fn(move |request: Request<Body>| {
                 let client = client.clone();
                 async move {
+                    let admission = request.extensions().get::<SessionAdmission>().is_some();
                     let request =
                         request.map(|body| reqwest::Body::wrap_stream(body.into_data_stream()));
                     let request = reqwest::Request::try_from(request)
@@ -141,6 +145,13 @@ impl SharedTransport {
                             status
                         })?
                         .into_http_response();
+                    if admission && response.status() == StatusCode::SERVICE_UNAVAILABLE {
+                        // A 503 cannot grant a lease. Leave headers and trailers for tonic
+                        // to read the gRPC status, or infer Unavailable when it is absent.
+                        return Ok(response.map(|body| {
+                            Body::new(body.map_frame(|frame| frame.map_data(|_| Bytes::new())))
+                        }));
+                    }
                     Ok::<_, Status>(response.map(Body::new))
                 }
             });