Retry transient gRPC code-mode session admission failures (#51185)

## Why

Transient admission errors or connection resets can prevent a code-mode session from opening before any execution starts.

## What changed

- Retry `OpenSession` and the initial lease message on `Unavailable` or `ResourceExhausted`, with at most three attempts and 100 ms and 200 ms delays within a single startup deadline.
- Classify HTTP connection and request failures as `Unavailable` so admission can recover from connection resets, while preserving error sources.
- Keep retries limited to initial admission so established executions are never replayed.

## Testing

Add gRPC integration tests covering transient status recovery, the retry limit, immediate failure for non-retryable statuses, and successful session execution after a connection reset.

GitOrigin-RevId: 3b654bcd377a5addf7c3ed58bb1c72a449b2f7f7

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

@@ -42,6 +42,8 @@ use tonic::Code;
 #[cfg(unix)]
 use tonic::transport::Server;
 
+#[path = "grpc/admission_tests.rs"]
+mod admission_tests;
 #[path = "support/host.rs"]
 mod host;
 #[path = "support/large_tool_delegate.rs"]

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

@@ -0,0 +1,108 @@
+//! Initial admission may retry transient failures; established work must not replay.
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use pretty_assertions::assert_eq;
+
+use super::*;
+
+#[tokio::test]
+async fn session_admission_retries_only_transient_failures_and_is_bounded() -> Result<()> {
+    for (code, failures, expected_attempts, succeeds) in [
+        (Code::Unavailable, 1, 2, true),
+        (Code::ResourceExhausted, 1, 2, true),
+        (Code::Unavailable, usize::MAX, 3, false),
+        (Code::InvalidArgument, 1, 1, false),
+        (Code::Unknown, 1, 1, false),
+    ] {
+        let attempts = Arc::new(AtomicUsize::new(0));
+        let count = Arc::clone(&attempts);
+        let app = axum::Router::new()
+            .fallback_service(grpc::code_mode_host_server::CodeModeHostServer::new(
+                codex_code_mode_host::GrpcCodeModeHost::new(),
+            ))
+            .layer(axum::middleware::from_fn(
+                move |request: axum::extract::Request, next: axum::middleware::Next| {
+                    let count = Arc::clone(&count);
+                    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>();
+                        }
+                        next.run(request).await
+                    }
+                },
+            ));
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+        let address = listener.local_addr()?;
+        let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+            axum::serve(listener, app)
+                .await
+                .expect("admission test server");
+        }));
+        let provider = GrpcCodeModeSessionProvider::new(format!("http://{address}"));
+        let result = timeout(TEST_TIMEOUT, provider.create_session()).await?;
+        assert_eq!(
+            result.is_ok(),
+            succeeds,
+            "status {code}: {}",
+            result.as_ref().err().cloned().unwrap_or_default()
+        );
+        assert_eq!(attempts.load(Ordering::SeqCst), expected_attempts);
+        if let Ok(session) = result {
+            let response = execute(
+                &session,
+                request("text(42);"),
+                Arc::new(NoopCodeModeSessionDelegate),
+            )
+            .await?;
+            assert_eq!(
+                response,
+                text_response("1", "42", response.code_mode_host_duration())
+            );
+            session.shutdown().await.map_err(anyhow::Error::msg)?;
+        }
+        drop(server);
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn session_admission_recovers_after_connection_reset() -> Result<()> {
+    use tokio::io::AsyncReadExt;
+
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let address = listener.local_addr()?;
+    let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+        // Drop the first connection with unread HTTP/2 bytes to send TCP RST.
+        // Serve the next connection normally; no timer or launcher mock is needed.
+        let (mut stream, _) = listener.accept().await?;
+        stream.read_exact(&mut [0]).await?;
+        drop(stream);
+        let app = axum::Router::new().fallback_service(
+            grpc::code_mode_host_server::CodeModeHostServer::new(
+                codex_code_mode_host::GrpcCodeModeHost::new(),
+            ),
+        );
+        axum::serve(listener, app).await
+    }));
+    let provider = GrpcCodeModeSessionProvider::new(format!("http://{address}"));
+    let session = timeout(TEST_TIMEOUT, provider.create_session())
+        .await?
+        .map_err(anyhow::Error::msg)?;
+    let response = execute(
+        &session,
+        request("text(42);"),
+        Arc::new(NoopCodeModeSessionDelegate),
+    )
+    .await?;
+    assert_eq!(
+        response,
+        text_response("1", "42", response.code_mode_host_duration())
+    );
+    session.shutdown().await.map_err(anyhow::Error::msg)?;
+    drop(server);
+    Ok(())
+}

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

@@ -112,21 +112,38 @@ impl GrpcCodeModeSessionProvider {
             || limits.max_heap_size_bytes.is_some())
         .then_some(limits);
         let open_session_span = tracing::info_span!("code_mode.grpc.open_session");
-        let mut open_session_request = tonic::Request::new(grpc::OpenSessionRequest {
-            cell_execution_limits,
-        });
-        inject_span_traceparent(&mut open_session_request, &open_session_span);
-        let (lease, first) = async {
-            let mut lease =
-                deadline::startup("session opening", client.open_session(open_session_request))
-                    .await?
-                    .into_inner();
-            let first = deadline::startup("session lease opening", lease.message())
-                .await?
-                .ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
-            Ok::<_, String>((lease, first))
-        }
-        .instrument(open_session_span)
+        let (lease, first) = deadline::startup("session opening", async {
+            let mut retries = 0;
+            loop {
+                let mut request = tonic::Request::new(grpc::OpenSessionRequest {
+                    cell_execution_limits,
+                });
+                inject_span_traceparent(&mut request, &open_session_span);
+                let result = async {
+                    let mut lease = client.open_session(request).await?.into_inner();
+                    let first = lease.message().await?.ok_or_else(|| {
+                        tonic::Status::unavailable("session lease ended before opening")
+                    })?;
+                    Ok::<_, tonic::Status>((lease, first))
+                }
+                .await;
+                let Err(error) = &result else { return result };
+                if retries == 2
+                    || !matches!(
+                        error.code(),
+                        tonic::Code::Unavailable | tonic::Code::ResourceExhausted
+                    )
+                {
+                    return result;
+                }
+                // No execution has started, and dropping a failed lease releases
+                // any partial admission. Never apply these retries to Execute.
+                retries += 1;
+                tracing::warn!(%error, retries, "retrying code-mode session admission");
+                tokio::time::sleep(Duration::from_millis(100 * retries)).await;
+            }
+        })
+        .instrument(open_session_span.clone())
         .await?;
         let Some(grpc::session_event::Event::Opened(opened)) = first.event else {
             return Err("gRPC code-mode session lease omitted its opening event".to_string());

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

@@ -1,10 +1,11 @@
-use std::io;
+use std::sync::Arc;
 
 use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
 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 tonic::Status;
 use tonic::body::Body;
 use tonic::codegen::http::Request;
 use tonic::codegen::http::Response;
@@ -17,7 +18,7 @@ use tower::util::BoxCloneSyncService;
 
 use super::GrpcClient;
 
-pub(super) type GrpcTransport = BoxCloneSyncService<Request<Body>, Response<Body>, io::Error>;
+pub(super) type GrpcTransport = BoxCloneSyncService<Request<Body>, Response<Body>, Status>;
 
 pub(super) enum SharedTransport {
     Url {
@@ -38,7 +39,7 @@ impl SharedTransport {
     }
 
     pub(super) fn with_channel(channel: Channel) -> Self {
-        let transport = channel.map_err(io::Error::other);
+        let transport = channel.map_err(|error| Status::from_error(Box::new(error)));
         Self::Connected(
             CodeModeHostClient::new(BoxCloneSyncService::new(transport))
                 .max_decoding_message_size(MAX_FRAME_BYTES)
@@ -71,7 +72,7 @@ impl SharedTransport {
             let channel = Endpoint::from_shared(endpoint.clone())
                 .map_err(|error| format!("invalid gRPC code-mode Unix socket endpoint: {error}"))?
                 .connect_lazy();
-            let transport = channel.map_err(io::Error::other);
+            let transport = channel.map_err(|error| Status::from_error(Box::new(error)));
             CodeModeHostClient::new(BoxCloneSyncService::new(transport))
         } else {
             let target = reqwest::Url::parse(endpoint)
@@ -104,13 +105,28 @@ impl SharedTransport {
                 async move {
                     let request =
                         request.map(|body| reqwest::Body::wrap_stream(body.into_data_stream()));
-                    let request = reqwest::Request::try_from(request).map_err(io::Error::other)?;
+                    let request = reqwest::Request::try_from(request)
+                        .map_err(|error| Status::invalid_argument(error.to_string()))?;
                     let response = client
                         .execute(request)
                         .await
-                        .map_err(io::Error::other)?
+                        .map_err(|error| {
+                            // Sending can fail on DNS or a reset connection before
+                            // headers arrive. Keep these distinct from RPC errors.
+                            let mut status = if matches!(
+                                &error,
+                                codex_http_client::HttpError::Request(error)
+                                    if error.is_connect() || error.is_request()
+                            ) {
+                                Status::unavailable(error.to_string())
+                            } else {
+                                Status::unknown(error.to_string())
+                            };
+                            status.set_source(Arc::new(error));
+                            status
+                        })?
                         .into_http_response();
-                    Ok::<_, io::Error>(response.map(Body::new))
+                    Ok::<_, Status>(response.map(Body::new))
                 }
             });
             CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin)