Allow provisioned executors more time to come online (#48575)

## Why

A provisioned executor can still be resuming after readiness is reported. Initial connection attempts can exhaust the ordinary registry retry limits before the executor comes online.

## What changed

Retry `environment_offline` registry responses during initial Noise rendezvous connections for provisioned environments until a fixed five-minute deadline, starting after provisioning succeeds. Keep the existing retry limits for other registry errors and stop on permanent failures.

## Testing

Add tests covering the five-minute deadline, ordinary limits for other errors and stalled requests, and termination on a later permanent error. Verify that readiness and info requests remain pending through an extended offline period and succeed using the same environment handle once the executor comes online.

GitOrigin-RevId: e22211499d9ec2f3590e75224eb39e2e4bf3538d

codex-rs/exec-server/src/client.rs

@@ -150,7 +150,6 @@ mod recovery;
 #[path = "client_refresh.rs"]
 mod refresh;
 pub(crate) use connection_failure::can_retry_connection_attempt;
-#[cfg(test)]
 pub(crate) use recovery::is_environment_offline_error;
 pub(crate) use recovery::is_retryable_recovery_error;
 

codex-rs/exec-server/src/client_refresh_tests.rs

@@ -1,5 +1,8 @@
 use std::sync::Arc;
 use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use anyhow::Result;
@@ -49,6 +52,8 @@ struct Registry {
     target: Mutex<Target>,
     next_lookup: Mutex<Option<oneshot::Receiver<()>>>,
     lookup_started: Notify,
+    offline: AtomicBool,
+    calls: AtomicUsize,
 }
 
 impl Registry {
@@ -65,6 +70,14 @@ impl NoiseRendezvousConnectProvider for Registry {
         _: NoiseChannelPublicKey,
     ) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
         Box::pin(async move {
+            self.calls.fetch_add(1, Ordering::Relaxed);
+            if self.offline.load(Ordering::Relaxed) {
+                return Err(ExecServerError::EnvironmentRegistryHttp {
+                    status: http::StatusCode::CONFLICT,
+                    code: Some("environment_offline".to_owned()),
+                    message: "executor is still starting".to_owned(),
+                });
+            }
             let target = self.target.lock().unwrap().clone();
             let block = self.next_lookup.lock().unwrap().take();
             if let Some(block) = block {
@@ -203,6 +216,8 @@ impl Executor {
             target: Mutex::new(self.target.clone()),
             next_lookup: Mutex::new(None),
             lookup_started: Notify::new(),
+            offline: AtomicBool::new(false),
+            calls: AtomicUsize::new(0),
         });
         let client = LazyRemoteExecServerClient::new(
             ExecServerTransportParams::NoiseRendezvous {
@@ -635,6 +650,63 @@ async fn superseded_refresh_lookup_does_not_retire_a_newer_session() -> Result<(
     Ok(())
 }
 
+#[tokio::test]
+async fn provisioned_environment_waits_for_offline_executor_on_the_same_handle() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (_, registry) = executor.client()?;
+    registry.offline.store(true, Ordering::Relaxed);
+    let manager = crate::EnvironmentManager::from_snapshot(
+        crate::environment_provider::EnvironmentProviderSnapshot {
+            environments: Vec::new(),
+            default: crate::environment_provider::EnvironmentDefault::Disabled,
+            include_local: false,
+        },
+        /*local_runtime_paths*/ None,
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )?;
+    let environment = manager
+        .report_environment_provisioning_status(
+            "environment".to_owned(),
+            Ok(crate::EnvironmentReadyInfo::default()),
+            registry.clone(),
+        )?
+        .expect("provisioned environment");
+    let ready = environment.wait_until_ready();
+    let info = environment.info();
+    tokio::pin!(ready, info);
+    assert!(futures::poll!(ready.as_mut()).is_pending());
+    assert!(futures::poll!(info.as_mut()).is_pending());
+    assert_eq!(registry.calls.load(Ordering::Relaxed), 1);
+
+    // A restore can outlast both the ordinary registry retry count and deadline.
+    for _ in 0..35 {
+        tokio::time::advance(Duration::from_secs(1)).await;
+        assert!(futures::poll!(ready.as_mut()).is_pending());
+        assert!(futures::poll!(info.as_mut()).is_pending());
+    }
+    assert!(registry.calls.load(Ordering::Relaxed) > 5);
+    let reported = manager
+        .report_environment_provisioning_status(
+            "environment".to_owned(),
+            Ok(crate::EnvironmentReadyInfo::default()),
+            registry.clone(),
+        )?
+        .expect("same provisioned environment");
+    assert!(Arc::ptr_eq(&environment, &reported));
+    assert!(futures::poll!(ready.as_mut()).is_pending());
+    assert!(futures::poll!(info.as_mut()).is_pending());
+
+    registry.offline.store(false, Ordering::Relaxed);
+    tokio::time::advance(Duration::from_secs(8)).await;
+    let ((), _) = tokio::try_join!(ready, info)?;
+    assert_eq!(
+        environment.cached_executor_registration_id(),
+        Some(executor.target.registration.clone())
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn environment_refresh_preserves_environment_and_filesystem_handles() -> Result<()> {
     let _clock = freeze_clock();

codex-rs/exec-server/src/client_transport.rs

@@ -31,6 +31,7 @@ use crate::ExecServerClient;
 use crate::ExecServerError;
 use crate::client::NoiseInitializeContext;
 use crate::client::accepted::AcceptedConnectionSource;
+use crate::client::is_environment_offline_error;
 use crate::client::is_retryable_registry_error;
 use crate::client::registry_recovery_retry_delay;
 use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
@@ -56,6 +57,7 @@ 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);
+const PROVISIONED_ENVIRONMENT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5 * 60);
 
 pub(crate) async fn connect_websocket_request(
     request: http::Request<()>,
@@ -337,7 +339,7 @@ impl ExecServerClient {
             transport_params => (transport_params, None),
         };
 
-        if let Some(mut readiness) = deferred_readiness {
+        let provisioning_deadline = if let Some(mut readiness) = deferred_readiness {
             let provisioning_result = readiness
                 .wait_for(Option::is_some)
                 .await
@@ -355,7 +357,10 @@ impl ExecServerClient {
                     )
                 })?;
             provisioning_result.map_err(ExecServerError::ProvisioningFailed)?;
-        }
+            Some(Instant::now() + PROVISIONED_ENVIRONMENT_CONNECT_TIMEOUT)
+        } else {
+            None
+        };
 
         let websocket = match transport_params {
             ExecServerTransportParams::Deferred(_) => {
@@ -379,6 +384,7 @@ impl ExecServerClient {
                     &provider,
                     &identity,
                     http_client_factory.clone(),
+                    provisioning_deadline,
                 )
                 .await?;
                 let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous {
@@ -431,6 +437,7 @@ impl ExecServerClient {
         provider: &Arc<dyn NoiseRendezvousConnectProvider>,
         identity: &NoiseChannelIdentity,
         http_client_factory: HttpClientFactory,
+        provisioning_deadline: Option<Instant>,
     ) -> Result<(ReadyNoiseRendezvousConnection, crate::NoiseChannelPublicKey), ExecServerError>
     {
         let open_connection = |bundle: NoiseRendezvousConnectBundle| {
@@ -464,14 +471,22 @@ impl ExecServerClient {
         loop {
             let bundle = match result {
                 Ok(bundle) => bundle,
-                Err(error)
-                    if is_retryable_registry_error(&error)
-                        && retries < INITIAL_REGISTRY_MAX_RETRIES =>
-                {
+                Err(error) => {
+                    // A provisioned executor may be resuming after an earlier ready
+                    // report. Only offline responses get the longer, fixed deadline.
+                    let retry_deadline = match provisioning_deadline {
+                        Some(deadline) if is_environment_offline_error(&error) => deadline,
+                        _ if is_retryable_registry_error(&error)
+                            && retries < INITIAL_REGISTRY_MAX_RETRIES =>
+                        {
+                            deadline
+                        }
+                        _ => return Err(error),
+                    };
                     // Session resumption owns its separate recovery deadline.
                     let delay = registry_recovery_retry_delay(&retry_key, retries);
                     retries += 1;
-                    result = match timeout_at(deadline, async {
+                    result = match timeout_at(retry_deadline, async {
                         sleep(delay).await;
                         connect_bundle().await
                     })
@@ -482,7 +497,6 @@ impl ExecServerClient {
                     };
                     continue;
                 }
-                Err(error) => return Err(error),
             };
             let executor_public_key = bundle.executor_public_key.clone();
             match open_connection(bundle).await {

codex-rs/exec-server/src/client_transport_tests.rs

@@ -24,6 +24,7 @@ use super::ExecServerReconnectStrategy;
 use super::INITIAL_REGISTRY_MAX_RETRIES;
 use super::INITIAL_REGISTRY_OPERATION_TIMEOUT;
 use super::INITIAL_REGISTRY_REQUEST_TIMEOUT;
+use super::PROVISIONED_ENVIRONMENT_CONNECT_TIMEOUT;
 use crate::ExecServerError;
 use crate::NoiseChannelIdentity;
 use crate::NoiseChannelPublicKey;
@@ -33,7 +34,9 @@ use crate::NoiseRendezvousConnectProvider;
 use crate::client::NoiseInitializeContext;
 use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
 use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;
+use crate::client_api::Deferred;
 use crate::client_api::ExecServerClientConnectOptions;
+use crate::client_api::ExecServerTransportParams;
 use crate::connection::JsonRpcConnection;
 use crate::noise_channel::PendingResponderHandshake;
 use crate::noise_channel::noise_channel_prologue;
@@ -101,10 +104,31 @@ impl SequenceNoiseConnectProvider {
             codex_http_client::HttpClientFactory::new(
                 codex_http_client::OutboundProxyPolicy::ReqwestDefault,
             ),
+            /*provisioning_deadline*/ None,
         )
         .await
         .map(|(ready, _)| (ready.connection, ready.options))
     }
+
+    async fn connect_provisioned(
+        self: &Arc<Self>,
+        identity: &NoiseChannelIdentity,
+    ) -> Result<ExecServerClient, ExecServerError> {
+        let (_readiness_tx, readiness) = tokio::sync::watch::channel(Some(Ok(())));
+        ExecServerClient::connect_for_transport(
+            ExecServerTransportParams::Deferred(Box::new(Deferred {
+                readiness,
+                transport: ExecServerTransportParams::NoiseRendezvous {
+                    provider: self.clone(),
+                    identity: identity.clone(),
+                },
+            })),
+            codex_http_client::HttpClientFactory::new(
+                codex_http_client::OutboundProxyPolicy::ReqwestDefault,
+            ),
+        )
+        .await
+    }
 }
 
 impl NoiseRendezvousConnectProvider for SequenceNoiseConnectProvider {
@@ -455,6 +479,124 @@ async fn initial_noise_connection_does_not_retry_permanent_registry_errors() ->
     Ok(())
 }
 
+#[tokio::test(start_paused = true)]
+async fn provisioned_noise_connection_bounds_offline_retries() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    // Enough responses to reach the deadline even at the shortest retry interval.
+    for _ in 0..1024 {
+        sequence.push_error(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ));
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+
+    let error = sequence
+        .connect_provisioned(&identity)
+        .await
+        .err()
+        .expect("a provisioned executor that stays offline must time out");
+
+    assert!(crate::client::is_environment_offline_error(&error));
+    assert_eq!(started.elapsed(), PROVISIONED_ENVIRONMENT_CONNECT_TIMEOUT);
+    let requests = sequence.requested_keys().len();
+    assert!(requests > INITIAL_REGISTRY_MAX_RETRIES as usize + 1);
+    sequence.assert_requested_identity(&identity, requests);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn provisioned_noise_connection_keeps_other_registry_retry_limits() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    for _ in 0..=INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_error(registry_error(
+            http::StatusCode::SERVICE_UNAVAILABLE,
+            "unavailable",
+        ));
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+
+    let error = sequence
+        .connect_provisioned(&identity)
+        .await
+        .err()
+        .expect("only an offline executor receives the provisioning grace");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp {
+            status: http::StatusCode::SERVICE_UNAVAILABLE,
+            ..
+        }
+    ));
+    assert!(started.elapsed() <= INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    assert!(sequence.requested_keys().len() <= INITIAL_REGISTRY_MAX_RETRIES as usize + 1);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn provisioned_noise_connection_bounds_a_stalled_retry_request() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    sequence.push_error(registry_error(
+        http::StatusCode::CONFLICT,
+        "environment_offline",
+    ));
+    for _ in 0..INITIAL_REGISTRY_MAX_RETRIES {
+        sequence.push_pending();
+    }
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+    let error = sequence
+        .connect_provisioned(&identity)
+        .await
+        .err()
+        .expect("a stalled registry request must keep the ordinary timeout");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout()
+    ));
+    assert_eq!(started.elapsed(), INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    let requests = sequence.requested_keys().len();
+    assert!((2..=3).contains(&requests));
+    sequence.assert_requested_identity(&identity, requests);
+    Ok(())
+}
+
+#[tokio::test(start_paused = true)]
+async fn provisioned_noise_connection_stops_on_a_late_permanent_error() -> Result<()> {
+    let sequence = Arc::new(SequenceNoiseConnectProvider::default());
+    for _ in 0..10 {
+        sequence.push_error(registry_error(
+            http::StatusCode::CONFLICT,
+            "environment_offline",
+        ));
+    }
+    sequence.push_error(registry_error(http::StatusCode::FORBIDDEN, "forbidden"));
+    let identity = NoiseChannelIdentity::generate()?;
+    let started = tokio::time::Instant::now();
+
+    let error = sequence
+        .connect_provisioned(&identity)
+        .await
+        .err()
+        .expect("a permanent error must stop provisioning retries");
+
+    assert!(matches!(
+        error,
+        ExecServerError::EnvironmentRegistryHttp {
+            status: http::StatusCode::FORBIDDEN,
+            code: Some(code),
+            ..
+        } if code == "forbidden"
+    ));
+    sequence.assert_requested_identity(&identity, /*requests*/ 11);
+    assert!(started.elapsed() > INITIAL_REGISTRY_OPERATION_TIMEOUT);
+    Ok(())
+}
+
 #[tokio::test]
 async fn noise_session_resume_rejects_a_changed_noise_key_before_connecting() -> Result<()> {
     let sequence = Arc::new(SequenceNoiseConnectProvider::default());