From 5663754f62adb9b7701f6ef76f8b86be21dae275 Mon Sep 17 00:00:00 2001 From: richardopenai Date: Thu, 20 Aug 2026 18:34:43 +0000 Subject: [PATCH] Retry transient registry failures during initial exec connection (#39777) ## What changed - Retry transient environment-registry failures while opening the initial Noise rendezvous connection, including timeouts, interrupted response bodies, retryable HTTP statuses, and temporarily offline environments. - Bound registry attempts with per-request and overall timeouts, exponential backoff, and a retry limit while returning permanent registry errors immediately. - Preserve the existing one-time bundle refresh after an unauthorized WebSocket handshake, with a fresh registry retry window. - Treat malformed complete JSON responses as terminal rather than retryable transport failures. ## Testing - Cover transient and permanent registry responses, stalled requests, retry limits and deadlines, interrupted bodies, malformed JSON, unauthorized bundle refresh, and session-resume behavior. GitOrigin-RevId: 5929274103702ad477e7a638ab489319a43a93fa --- codex-rs/exec-server/src/client.rs | 4 + codex-rs/exec-server/src/client_recovery.rs | 31 +- .../exec-server/src/client_recovery_tests.rs | 26 +- codex-rs/exec-server/src/client_transport.rs | 82 +++- .../exec-server/src/client_transport_tests.rs | 407 +++++++++++++++++- codex-rs/exec-server/src/remote.rs | 87 +++- 6 files changed, 587 insertions(+), 50 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index af0d47514c..a686606210 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -129,7 +129,11 @@ pub(crate) mod http_client; mod network_policy_audit; #[path = "client_recovery.rs"] mod recovery; +#[cfg(test)] +pub(crate) use recovery::is_environment_offline_error; pub(crate) use recovery::is_retryable_recovery_error; +pub(crate) use recovery::is_retryable_registry_error; +pub(crate) use recovery::registry_recovery_retry_delay; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 1d6144f9c2..97724180a1 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -825,33 +825,44 @@ pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool { ) } -fn is_retryable_registry_error(error: &ExecServerError) -> bool { +pub(crate) fn is_retryable_registry_error(error: &ExecServerError) -> bool { matches!( error, ExecServerError::EnvironmentRegistryRequest(error) - if error.is_connect() || error.is_timeout() + if error.is_connect() + || error.is_timeout() + || error.is_body() + || matches!( + error, + codex_http_client::RouteAwareRequestError::Request(error) + if error.is_decode() + ) ) || matches!( error, - ExecServerError::EnvironmentRegistryHttp { status, code, .. } + ExecServerError::EnvironmentRegistryHttp { status, .. } if status.is_server_error() || *status == http::StatusCode::REQUEST_TIMEOUT || *status == http::StatusCode::TOO_MANY_REQUESTS - // TODO: Replace this coarse retry with an explicit registry/presence - // recovery FSM so `environment_offline` is retried only while the - // executor is expected to reconnect. - || (*status == http::StatusCode::CONFLICT - && code.as_deref() == Some("environment_offline")) + ) || is_environment_offline_error(error) +} + +pub(crate) fn is_environment_offline_error(error: &ExecServerError) -> bool { + matches!( + error, + ExecServerError::EnvironmentRegistryHttp { status, code, .. } + if *status == http::StatusCode::CONFLICT + && code.as_deref() == Some("environment_offline") ) } -fn registry_recovery_retry_delay(session_id: &str, attempt: u32) -> Duration { +pub(crate) fn registry_recovery_retry_delay(retry_key: &str, attempt: u32) -> Duration { let multiplier = 1_u32.checked_shl(attempt.min(4)).unwrap_or(u32::MAX); let base_delay = REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL .saturating_mul(multiplier) .min(REGISTRY_RECOVERY_MAX_RETRY_INTERVAL); let base_millis = base_delay.as_millis() as u64; let mut hasher = DefaultHasher::new(); - session_id.hash(&mut hasher); + retry_key.hash(&mut hasher); attempt.hash(&mut hasher); Duration::from_millis(base_millis + hasher.finish() % (base_millis / 2 + 1)) diff --git a/codex-rs/exec-server/src/client_recovery_tests.rs b/codex-rs/exec-server/src/client_recovery_tests.rs index 545154728b..4abc939d32 100644 --- a/codex-rs/exec-server/src/client_recovery_tests.rs +++ b/codex-rs/exec-server/src/client_recovery_tests.rs @@ -37,13 +37,31 @@ fn registry_recovery_retry_delay_exponentially_backs_off_and_caps() { #[test] fn recovery_retries_transient_registry_errors() { - let error = registry_error(http::StatusCode::TOO_MANY_REQUESTS, /*code*/ None); + for status in [ + http::StatusCode::REQUEST_TIMEOUT, + http::StatusCode::TOO_MANY_REQUESTS, + http::StatusCode::INTERNAL_SERVER_ERROR, + http::StatusCode::BAD_GATEWAY, + http::StatusCode::SERVICE_UNAVAILABLE, + ] { + let error = registry_error(status, /*code*/ None); + + assert!(is_retryable_registry_error(&error)); + assert!(is_retryable_recovery_error(&error)); + assert!(is_retryable_recovery_error( + &ExecServerError::ConnectionAttempt(Arc::new(error)) + )); + } +} + +#[test] +fn recovery_retries_registry_request_timeouts() { + let error = ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + ); assert!(is_retryable_registry_error(&error)); assert!(is_retryable_recovery_error(&error)); - assert!(is_retryable_recovery_error( - &ExecServerError::ConnectionAttempt(Arc::new(error)) - )); } #[test] diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index b4a75bb6b3..62fff714fd 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -5,7 +5,10 @@ use std::time::Duration; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Command; +use tokio::time::Instant; +use tokio::time::sleep; use tokio::time::timeout; +use tokio::time::timeout_at; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::warn; @@ -18,6 +21,8 @@ use codex_websocket_client::WebSocketTlsMode; use crate::ExecServerClient; use crate::ExecServerError; +use crate::client::is_retryable_registry_error; +use crate::client::registry_recovery_retry_delay; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT; use crate::client_api::ExecServerClientConnectOptions; @@ -37,6 +42,9 @@ use crate::relay::harness_connection_from_websocket; use crate::trace_context::current_rendezvous_headers; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; +const INITIAL_REGISTRY_MAX_RETRIES: u32 = 4; +const INITIAL_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(6); +const INITIAL_REGISTRY_OPERATION_TIMEOUT: Duration = Duration::from_secs(14); /// Reopens the transport for one logical exec-server client session. /// @@ -199,23 +207,65 @@ impl ExecServerClient { http_client_factory: http_client_factory.clone(), }) }; - let bundle = provider.connect_bundle(identity.public_key()).await?; - match open_connection(bundle).await { - Err(error) - if matches!( - &error, - ExecServerError::WebSocketConnect { source, .. } - if matches!( - source, - tokio_tungstenite::tungstenite::Error::Http(response) - if response.status().as_u16() == 401 - ) - ) => - { - let bundle = provider.connect_bundle(identity.public_key()).await?; - open_connection(bundle).await + let mut deadline = Instant::now() + INITIAL_REGISTRY_OPERATION_TIMEOUT; + let retry_key = uuid::Uuid::new_v4().to_string(); + let mut retries = 0; + let mut refreshed_unauthorized_bundle = false; + let connect_bundle = || async { + timeout( + INITIAL_REGISTRY_REQUEST_TIMEOUT, + provider.connect_bundle(identity.public_key()), + ) + .await + .unwrap_or_else(|_| { + Err(ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + )) + }) + }; + let mut result = connect_bundle().await; + loop { + let bundle = match result { + Ok(bundle) => bundle, + Err(error) + if is_retryable_registry_error(&error) + && retries < INITIAL_REGISTRY_MAX_RETRIES => + { + // Session resumption owns its separate recovery deadline. + let delay = registry_recovery_retry_delay(&retry_key, retries); + retries += 1; + result = match timeout_at(deadline, async { + sleep(delay).await; + connect_bundle().await + }) + .await + { + Ok(result) => result, + Err(_) => return Err(error), + }; + continue; + } + Err(error) => return Err(error), + }; + match open_connection(bundle).await { + Err(error) + if !refreshed_unauthorized_bundle + && matches!( + &error, + ExecServerError::WebSocketConnect { source, .. } + if matches!( + source, + tokio_tungstenite::tungstenite::Error::Http(response) + if response.status().as_u16() == 401 + ) + ) => + { + refreshed_unauthorized_bundle = true; + deadline = Instant::now() + INITIAL_REGISTRY_OPERATION_TIMEOUT; + result = connect_bundle().await; + } + result => return result, } - result => result, } } diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index cce85b3331..846dab9bc4 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use std::sync::Mutex; use anyhow::Result; +use futures::FutureExt; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use tokio::io::AsyncReadExt; @@ -11,22 +12,36 @@ use tokio::net::TcpListener; use tokio_tungstenite::accept_async; use super::ExecServerClient; +use super::ExecServerReconnectStrategy; +use super::INITIAL_REGISTRY_MAX_RETRIES; +use super::INITIAL_REGISTRY_OPERATION_TIMEOUT; +use super::INITIAL_REGISTRY_REQUEST_TIMEOUT; use crate::ExecServerError; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; use crate::NoiseRendezvousConnectBundle; use crate::NoiseRendezvousConnectProvider; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT; struct SequenceNoiseConnectProvider { - bundles: Mutex>, + bundles: + Mutex>>>, returned_urls: Mutex>, + requested_keys: Mutex>, } impl SequenceNoiseConnectProvider { - fn new(bundles: Vec) -> Self { + fn new(bundles: Vec>) -> Self { Self { - bundles: Mutex::new(bundles.into()), + bundles: Mutex::new( + bundles + .into_iter() + .map(|bundle| futures::future::ready(bundle).boxed()) + .collect(), + ), returned_urls: Mutex::new(Vec::new()), + requested_keys: Mutex::new(Vec::new()), } } @@ -41,21 +56,25 @@ impl SequenceNoiseConnectProvider { impl NoiseRendezvousConnectProvider for SequenceNoiseConnectProvider { fn connect_bundle( &self, - _: NoiseChannelPublicKey, + harness_public_key: NoiseChannelPublicKey, ) -> BoxFuture<'_, Result> { - let result = self + self.requested_keys.lock().unwrap().push(harness_public_key); + let response = self .bundles .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .pop_front() - .ok_or_else(|| ExecServerError::Protocol("test Noise provider exhausted".to_string())); - if let Ok(bundle) = &result { - self.returned_urls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(bundle.websocket_url.clone()); - } - Box::pin(async move { result }) + .expect("test Noise provider exhausted"); + Box::pin(async move { + let result = response.await; + if let Ok(bundle) = &result { + self.returned_urls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(bundle.websocket_url.clone()); + } + result + }) } } @@ -69,6 +88,329 @@ fn test_bundle(websocket_url: String) -> Result { }) } +fn registry_error(status: http::StatusCode, code: &str) -> ExecServerError { + ExecServerError::EnvironmentRegistryHttp { + status, + code: Some(code.to_string()), + message: "registry unavailable".to_string(), + } +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_bounds_offline_retries() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new( + (0..=INITIAL_REGISTRY_MAX_RETRIES) + .map(|_| { + Err(registry_error( + http::StatusCode::CONFLICT, + "environment_offline", + )) + }) + .collect(), + )); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + let started = tokio::time::Instant::now(); + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("offline retries must end"); + + assert!(crate::client::is_environment_offline_error(&error)); + let requested_keys = sequence.requested_keys.lock().unwrap(); + assert!((4..=INITIAL_REGISTRY_MAX_RETRIES as usize + 1).contains(&requested_keys.len())); + assert_eq!( + *requested_keys, + vec![identity.public_key(); requested_keys.len()] + ); + assert!(started.elapsed() <= INITIAL_REGISTRY_OPERATION_TIMEOUT); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_bounds_a_stalled_retry_request() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![Err( + registry_error(http::StatusCode::CONFLICT, "environment_offline"), + )])); + sequence + .bundles + .lock() + .unwrap() + .extend((0..INITIAL_REGISTRY_MAX_RETRIES).map(|_| { + futures::future::pending::>() + .boxed() + })); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + let started = tokio::time::Instant::now(); + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("stalled retry must time out"); + + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout() + )); + assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT); + let requested_keys = sequence.requested_keys.lock().unwrap(); + assert!((2..=3).contains(&requested_keys.len())); + assert_eq!( + *requested_keys, + vec![identity.public_key(); requested_keys.len()] + ); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_bounds_a_stalled_initial_request() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![])); + sequence + .bundles + .lock() + .unwrap() + .extend((0..=INITIAL_REGISTRY_MAX_RETRIES).map(|_| { + futures::future::pending::>() + .boxed() + })); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + let started = tokio::time::Instant::now(); + + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("stalled initial request must time out"); + + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout() + )); + assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT); + let requested_keys = sequence.requested_keys.lock().unwrap(); + assert!((2..=3).contains(&requested_keys.len())); + assert_eq!( + *requested_keys, + vec![identity.public_key(); requested_keys.len()] + ); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_retries_a_stalled_initial_request() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![])); + sequence.bundles.lock().unwrap().extend([ + futures::future::pending::>().boxed(), + futures::future::ready(Err(registry_error( + http::StatusCode::FORBIDDEN, + "forbidden", + ))) + .boxed(), + ]); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + let started = tokio::time::Instant::now(); + + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("terminal response must stop the retry sequence"); + + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryHttp { + status: http::StatusCode::FORBIDDEN, + .. + } + )); + assert!(started.elapsed() >= INITIAL_REGISTRY_REQUEST_TIMEOUT); + assert!(started.elapsed() < INITIAL_REGISTRY_OPERATION_TIMEOUT); + assert_eq!( + *sequence.requested_keys.lock().unwrap(), + vec![identity.public_key(); 2] + ); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_retries_transient_registry_statuses() -> Result<()> { + for status in [ + http::StatusCode::REQUEST_TIMEOUT, + http::StatusCode::TOO_MANY_REQUESTS, + http::StatusCode::INTERNAL_SERVER_ERROR, + http::StatusCode::BAD_GATEWAY, + http::StatusCode::SERVICE_UNAVAILABLE, + ] { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![ + Err(registry_error(status, "temporarily_unavailable")), + Err(registry_error(http::StatusCode::FORBIDDEN, "forbidden")), + ])); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("terminal response must stop the retry sequence"); + + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryHttp { + status: http::StatusCode::FORBIDDEN, + .. + } + )); + assert_eq!( + *sequence.requested_keys.lock().unwrap(), + vec![identity.public_key(); 2], + "registry status {status} should be retried" + ); + } + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_retries_registry_request_timeouts() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![ + Err(ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + )), + Err(registry_error(http::StatusCode::FORBIDDEN, "forbidden")), + ])); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("terminal response must stop the retry sequence"); + + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryHttp { + status: http::StatusCode::FORBIDDEN, + .. + } + )); + assert_eq!( + *sequence.requested_keys.lock().unwrap(), + vec![identity.public_key(); 2] + ); + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn initial_noise_connection_does_not_retry_permanent_registry_errors() -> Result<()> { + for (status, code) in [ + (http::StatusCode::UNAUTHORIZED, "unauthorized"), + (http::StatusCode::FORBIDDEN, "forbidden"), + (http::StatusCode::BAD_REQUEST, "bad_request"), + (http::StatusCode::NOT_FOUND, "environment_not_found"), + (http::StatusCode::CONFLICT, "registration_conflict"), + (http::StatusCode::CONFLICT, "route_unavailable"), + ] { + // A terminal error must also stop a retry sequence already in progress. + for initial_offline in [false, true] { + let mut responses = Vec::new(); + if initial_offline { + responses.push(Err(registry_error( + http::StatusCode::CONFLICT, + "environment_offline", + ))); + } + responses.push(Err(registry_error(status, code))); + let sequence = Arc::new(SequenceNoiseConnectProvider::new(responses)); + let provider: Arc = sequence.clone(); + let identity = NoiseChannelIdentity::generate()?; + let error = ExecServerClient::open_initial_noise_rendezvous_connection( + &provider, + &identity, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .err() + .expect("other errors must propagate"); + assert!( + matches!(error, ExecServerError::EnvironmentRegistryHttp { status: actual_status, code: Some(actual_code), .. } if actual_status == status && actual_code == code) + ); + assert_eq!( + sequence.requested_keys.lock().unwrap().len(), + 1 + usize::from(initial_offline) + ); + } + } + Ok(()) +} + +#[tokio::test(start_paused = true)] +async fn noise_session_resume_leaves_offline_retries_to_recovery() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![Err( + registry_error(http::StatusCode::CONFLICT, "environment_offline"), + )])); + let identity = NoiseChannelIdentity::generate()?; + let strategy = ExecServerReconnectStrategy::NoiseRendezvous { + provider: sequence.clone(), + identity: identity.clone(), + client_name: "test".to_string(), + connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + }; + let started = tokio::time::Instant::now(); + let error = strategy + .resume("session") + .await + .err() + .expect("resume must return the offline error"); + assert!(crate::client::is_environment_offline_error(&error)); + assert_eq!(started.elapsed(), std::time::Duration::ZERO); + assert_eq!( + *sequence.requested_keys.lock().unwrap(), + vec![identity.public_key()] + ); + Ok(()) +} + #[tokio::test] async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake() -> Result<()> { let unauthorized_listener = TcpListener::bind("127.0.0.1:0").await?; @@ -92,10 +434,37 @@ async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake( let _websocket = accept_async(socket).await?; anyhow::Ok(()) }); - let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![ - test_bundle(unauthorized_url.clone())?, - test_bundle(accepted_url.clone())?, - ])); + let sequence = Arc::new(SequenceNoiseConnectProvider::new(vec![])); + let unauthorized_bundle = test_bundle(unauthorized_url.clone())?; + let accepted_bundle = test_bundle(accepted_url.clone())?; + sequence.bundles.lock().unwrap().extend([ + async { + tokio::time::pause(); + Err(registry_error( + http::StatusCode::CONFLICT, + "environment_offline", + )) + } + .boxed(), + async move { + tokio::time::resume(); + Ok(unauthorized_bundle) + } + .boxed(), + async { + tokio::time::pause(); + Err(registry_error( + http::StatusCode::CONFLICT, + "environment_offline", + )) + } + .boxed(), + async move { + tokio::time::resume(); + Ok(accepted_bundle) + } + .boxed(), + ]); let provider: Arc = sequence.clone(); let identity = NoiseChannelIdentity::generate()?; @@ -112,6 +481,10 @@ async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake( sequence.returned_urls(), vec![unauthorized_url, accepted_url] ); + assert_eq!( + *sequence.requested_keys.lock().unwrap(), + vec![identity.public_key(); 4] + ); unauthorized_server.await??; accepted_server.await??; Ok(()) diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index e5c85f8919..8836e4065a 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -250,10 +250,11 @@ impl EnvironmentRegistryClient { R: for<'de> Deserialize<'de>, { if response.status().is_success() { - return response - .json::() + let body = response + .text() .await - .map_err(|error| ExecServerError::EnvironmentRegistryRequest(error.into())); + .map_err(|error| ExecServerError::EnvironmentRegistryRequest(error.into()))?; + return serde_json::from_str(&body).map_err(ExecServerError::Json); } let status = response.status(); @@ -817,6 +818,7 @@ mod tests { use opentelemetry::trace::TracerProvider as _; use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; + use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tracing::Instrument; @@ -1039,6 +1041,85 @@ mod tests { )); } + #[tokio::test] + async fn connect_environment_retries_interrupted_registry_response_bodies() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("registry listener should bind"); + let registry_url = format!( + "http://{}", + listener + .local_addr() + .expect("registry listener should have an address") + ); + tokio::spawn(async move { + let (mut stream, _) = listener + .accept() + .await + .expect("registry request should connect"); + let mut request = [0_u8; 4096]; + let bytes_read = stream + .read(&mut request) + .await + .expect("registry request should arrive before the response"); + assert_ne!(bytes_read, 0, "registry request should not be empty"); + stream + .write_all( + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 256\r\n\r\n{", + ) + .await + .expect("registry response headers should write"); + stream + .shutdown() + .await + .expect("registry connection should close"); + }); + let client = EnvironmentRegistryClient::new(registry_url, static_registry_auth_provider()) + .expect("client"); + let harness_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); + + let error = client + .connect_environment("environment-requested", harness_public_key) + .await + .err() + .expect("interrupted response body must fail"); + + assert!( + crate::client::is_retryable_registry_error(&error), + "interrupted registry response body should be retryable: {error:?}" + ); + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryRequest(_) + )); + } + + #[tokio::test] + async fn connect_environment_does_not_retry_malformed_successful_responses() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-requested/connect")) + .respond_with(ResponseTemplate::new(200).set_body_string("{")) + .mount(&server) + .await; + let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) + .expect("client"); + let harness_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); + + let error = client + .connect_environment("environment-requested", harness_public_key) + .await + .err() + .expect("malformed response must fail"); + + assert!(!crate::client::is_retryable_registry_error(&error)); + assert!(matches!(error, ExecServerError::Json(_))); + } + #[tokio::test] async fn register_environment_does_not_follow_redirects_with_auth_headers() { let server = MockServer::start().await;