From 064b2d9bca2b5ba2f9f399ffefd06a75d97852f0 Mon Sep 17 00:00:00 2001 From: Sean Huang Date: Tue, 10 Nov 2026 01:57:05 +0000 Subject: [PATCH] Preserve typed gRPC code-mode session admission errors (#60496) ## Why Session-opening failures were returned as strings, preventing callers from inspecting the terminal gRPC admission status directly. ## What changed - Add `GrpcCodeModeSessionProvider::open_session` returning `GrpcSessionOpenError`, distinguishing admission status codes, admission timeouts, and other setup failures. - Keep string errors for the existing `CodeModeSessionProvider` trait API. - Re-export `ProcessOwnedGrpcCodeModeSessionProvider` from `codex-core`. ## Testing Extend admission tests to assert terminal status codes, cover exhausted `ResourceExhausted` retries, and verify that `DeadlineExceeded` is not retried. GitOrigin-RevId: 9d7f2e733e12b8bb7def368fa5fc5a09ea2a374e --- .../tests/grpc/admission_tests.rs | 17 ++++--- .../code-mode/src/grpc_session/deadline.rs | 34 +++++++++++++- codex-rs/code-mode/src/grpc_session/mod.rs | 44 ++++++++++++++----- .../code-mode/src/grpc_session/reconnect.rs | 20 +++++---- codex-rs/code-mode/src/lib.rs | 1 + codex-rs/core/src/lib.rs | 1 + 6 files changed, 88 insertions(+), 29 deletions(-) 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 b9bc41646e..6c4669a388 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 codex_code_mode::GrpcSessionOpenError; use pretty_assertions::assert_eq; use super::*; @@ -12,6 +13,8 @@ async fn session_admission_retries_only_transient_failures_and_is_bounded() -> R (Code::Unavailable, 1, 2, true), (Code::ResourceExhausted, 1, 2, true), (Code::Unavailable, usize::MAX, 5, false), + (Code::ResourceExhausted, usize::MAX, 5, false), + (Code::DeadlineExceeded, 1, 1, false), (Code::InvalidArgument, 1, 1, false), (Code::Unknown, 1, 1, false), ] { @@ -43,13 +46,13 @@ async fn session_admission_retries_only_transient_failures_and_is_bounded() -> R .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() - ); + let result = timeout(TEST_TIMEOUT, provider.open_session(Default::default())).await?; + let actual_code = match result.as_ref() { + Ok(_) => None, + Err(GrpcSessionOpenError::AdmissionStatus { code, .. }) => Some(*code), + Err(error) => anyhow::bail!("unexpected session open error: {error}"), + }; + assert_eq!(actual_code, (!succeeds).then_some(code)); assert_eq!(attempts.load(Ordering::SeqCst), expected_attempts); if let Ok(session) = result { let response = execute( diff --git a/codex-rs/code-mode/src/grpc_session/deadline.rs b/codex-rs/code-mode/src/grpc_session/deadline.rs index 4afc7e58f7..12b31e5822 100644 --- a/codex-rs/code-mode/src/grpc_session/deadline.rs +++ b/codex-rs/code-mode/src/grpc_session/deadline.rs @@ -1,4 +1,5 @@ use std::fmt::Display; +use std::fmt::Formatter; use std::future::Future; use std::time::Duration; @@ -7,6 +8,35 @@ use super::SessionInner; const TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60); const MAX_ERROR_BYTES: usize = 512; +/// Terminal failure while establishing a new gRPC code-mode session. +#[derive(Debug, PartialEq, Eq)] +pub enum GrpcSessionOpenError { + /// Final status from OpenSession or the first read of its lease, with the bounded diagnostic. + AdmissionStatus { code: tonic::Code, message: String }, + /// The client deadline elapsed while waiting for admission, including retries. + AdmissionTimeout(String), + /// Setup failed outside admission, for example during transport setup or tool subscription. + Other(String), +} + +impl Display for GrpcSessionOpenError { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::AdmissionStatus { message, .. } + | Self::AdmissionTimeout(message) + | Self::Other(message) => f.write_str(message), + } + } +} + +impl std::error::Error for GrpcSessionOpenError {} + +impl> From for GrpcSessionOpenError { + fn from(message: T) -> Self { + Self::Other(message.into()) + } +} + pub(super) async fn startup( operation: &str, request: impl Future>, @@ -85,7 +115,7 @@ pub(super) fn failure(operation: &str, error: impl Display) -> String { message } -async fn enforce( +pub(super) async fn enforce( operation: &str, runtime_timeout: Duration, request: impl Future>, @@ -100,7 +130,7 @@ async fn enforce( } } -enum RequestError { +pub(super) enum RequestError { Failed(E), TimedOut(String), } diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index eec382d9a6..dda02b71d7 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -27,6 +27,7 @@ use codex_code_mode_protocol::grpc; use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use futures::TryFutureExt; use rand::Rng; use tokio::sync::Semaphore; use tokio::sync::watch; @@ -35,6 +36,8 @@ use tokio_util::task::TaskTracker; use tonic::transport::Channel; use tracing::Instrument; +pub use self::deadline::GrpcSessionOpenError; +use self::deadline::RequestError; use self::operations::WaitSlot; use self::state::ClosedCell; use self::state::SessionState; @@ -134,12 +137,22 @@ impl GrpcCodeModeSessionProvider { } } + /// Opens a new session and preserves the terminal admission status for the caller. + pub async fn open_session( + &self, + options: CodeModeSessionOptions, + ) -> Result, GrpcSessionOpenError> { + let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options)); + session.initialize().await?; + Ok(session) + } + #[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)] async fn open_binding( &self, options: &CodeModeSessionOptions, pinned_mode: Option, - ) -> Result, String> { + ) -> Result, GrpcSessionOpenError> { let client = deadline::startup("transport connection", self.transport.client()).await?; let requested_mode = pinned_mode.unwrap_or(options.output_truncation); let limits = &options.cell_execution_limits; @@ -155,7 +168,8 @@ 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 (lease, first, route) = deadline::startup("session opening", async { + let operation = "session opening"; + let (lease, first, route) = deadline::enforce(operation, Duration::ZERO, async { let mut retries = 0; loop { let mut request = tonic::Request::new(grpc::OpenSessionRequest { @@ -211,9 +225,16 @@ impl GrpcCodeModeSessionProvider { } }) .instrument(open_session_span.clone()) - .await?; + .await + .map_err(|error| match error { + RequestError::Failed(status) => GrpcSessionOpenError::AdmissionStatus { + code: status.code(), + message: deadline::failure(operation, status), + }, + RequestError::TimedOut(reason) => GrpcSessionOpenError::AdmissionTimeout(reason), + })?; let Some(grpc::session_event::Event::Opened(opened)) = first.event else { - return Err("gRPC code-mode session lease omitted its opening event".to_string()); + return Err("gRPC code-mode session lease omitted its opening event".into()); }; validate_identifier(&opened.session_id, "session ID")?; let host_mode = @@ -230,13 +251,13 @@ impl GrpcCodeModeSessionProvider { { return Err(format!( "gRPC code-mode host selected {output_truncation:?} output truncation, requested {requested_mode:?}" - )); + ).into()); } if pinned_mode.is_some_and(|previous| previous != output_truncation) { // Dropping the opening lease releases this incompatible session before subscribing. return Err( "gRPC code-mode host changed the session output truncation mode on reconnect" - .to_string(), + .into(), ); } @@ -277,7 +298,7 @@ impl GrpcCodeModeSessionProvider { Ok(response) => response, Err(error) => { let _ = wait_for_watch(inner.request_shutdown()).await; - return Err(error); + return Err(error.into()); } }; inner.spawn_tool_subscription(response.into_inner()); @@ -313,11 +334,10 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider { &'a self, options: CodeModeSessionOptions, ) -> CodeModeSessionProviderFuture<'a> { - Box::pin(async move { - let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options)); - session.initialize().await?; - Ok(session as _) - }) + Box::pin( + self.open_session(options) + .map_err(|error| error.to_string()), + ) } } diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index 9b39a445f4..f00b1a5742 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -20,6 +20,7 @@ use tokio_util::sync::CancellationToken; use super::GrpcCodeModeSession; use super::GrpcCodeModeSessionProvider; +use super::GrpcSessionOpenError; use super::SHUTDOWN_ERROR; use super::generation; use super::generation::GenerationDelegate; @@ -66,7 +67,7 @@ impl ReconnectableSession { } } - pub(super) async fn initialize(&self) -> Result<(), String> { + pub(super) async fn initialize(&self) -> Result<(), GrpcSessionOpenError> { self.inner.get_or_open_binding().await.map(|_| ()) } } @@ -151,7 +152,10 @@ impl Drop for ReconnectableSession { impl ReconnectInner { async fn binding_for_request(&self, generation: Option) -> Result { - let binding = self.get_or_open_binding().await?; + let binding = self + .get_or_open_binding() + .await + .map_err(|error| error.to_string())?; if generation != Some(binding.generation) { return Err(concat!( "The exec runtime was replaced. Stored values and running cells were lost. ", @@ -162,9 +166,9 @@ impl ReconnectInner { Ok(binding) } - async fn get_or_open_binding(&self) -> Result { + async fn get_or_open_binding(&self) -> Result { if self.shutdown_requested.is_cancelled() { - return Err(SHUTDOWN_ERROR.to_string()); + return Err(SHUTDOWN_ERROR.into()); } if let Some(binding) = self.live_binding() { return Ok(binding); @@ -173,13 +177,13 @@ impl ReconnectInner { let _opening_permit = tokio::select! { biased; _ = self.shutdown_requested.cancelled() => { - return Err(SHUTDOWN_ERROR.to_string()); + return Err(SHUTDOWN_ERROR.into()); } permit = self.opening_permit.acquire() => permit .map_err(|_| "gRPC code-mode session opening coordinator closed".to_string())?, }; if self.shutdown_requested.is_cancelled() { - return Err(SHUTDOWN_ERROR.to_string()); + return Err(SHUTDOWN_ERROR.into()); } if let Some(binding) = self.live_binding() { return Ok(binding); @@ -196,7 +200,7 @@ impl ReconnectInner { let session = tokio::select! { biased; _ = self.shutdown_requested.cancelled() => { - return Err(SHUTDOWN_ERROR.to_string()); + return Err(SHUTDOWN_ERROR.into()); } session = self.provider.open_binding(&self.options, self.output_truncation.get().copied()) => session?, }; @@ -217,7 +221,7 @@ impl ReconnectInner { }; if !published { let _ = wait_for_watch(binding.session.inner.request_shutdown()).await; - return Err(SHUTDOWN_ERROR.to_string()); + return Err(SHUTDOWN_ERROR.into()); } Ok(binding) } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 8bed201d33..9ad685a48c 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -3,6 +3,7 @@ mod remote_session; pub use codex_code_mode_protocol::*; pub use grpc_session::GrpcCodeModeSessionProvider; +pub use grpc_session::GrpcSessionOpenError; pub use remote_session::DisabledCodeModeSessionProvider; pub use remote_session::ProcessOwnedCodeModeSession; pub use remote_session::ProcessOwnedCodeModeSessionProvider; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f95157512e..ec127194f3 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -273,6 +273,7 @@ pub use client::X_CODEX_TURN_METADATA_HEADER; pub use client_common::Prompt; pub use client_common::ResponseEvent; pub use client_common::ResponseStream; +pub use code_mode_host::ProcessOwnedGrpcCodeModeSessionProvider; pub use codex_prompts::REVIEW_PROMPT; pub use codex_runtime_agent_instance::resolve_installation_id; pub use compact::content_items_to_text;