diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index bf13a4e117..c2af0e3e85 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/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"] diff --git a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs new file mode 100644 index 0000000000..5dcad3df6c --- /dev/null +++ b/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::(); + } + 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(()) +} diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 5c58592e22..9dc478bbae 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/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()); diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index 9b32d005b0..0f9b336da1 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/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, Response, io::Error>; +pub(super) type GrpcTransport = BoxCloneSyncService, Response, 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)