diff --git a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs b/codex-rs/code-mode-host/tests/grpc/admission_tests.rs index 6c4669a388..8eccd79858 100644 --- a/codex-rs/code-mode-host/tests/grpc/admission_tests.rs +++ b/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("no healthy upstream")?; + 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::(); + return response.map(axum::body::Body::from); } next.run(request).await } diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index dda02b71d7..66d10ac3c8 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -59,6 +59,9 @@ mod transport; type GrpcClient = CodeModeHostClient; +#[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?; diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index e34cfdb084..a94bbfaf1a 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/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, Response, Status>; @@ -118,6 +121,7 @@ impl SharedTransport { let transport = service_fn(move |request: Request| { let client = client.clone(); async move { + let admission = request.extensions().get::().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)) } });