diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 333a7832e2..1ebde3302f 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -331,7 +331,7 @@ pub(crate) struct LazyRemoteExecServerClient { pub(crate) transport_params: ExecServerTransportParams, http_client_factory: HttpClientFactory, recovery_policy: RecoveryPolicy, - // Saves the first startup result so callers share it and failures remain final. + // Saves the first startup result so callers share it; retryable failures use reconnect. startup: Arc, // The latest successful client, replaced whenever reconnecting succeeds. current_client: Arc>>, @@ -469,6 +469,13 @@ impl LazyRemoteExecServerClient { } async fn initial_client(&self) -> Result { + if let Some(Err(error)) = self.startup.get() + && self.can_reconnect() + && recovery::is_retryable_recovery_error(error) + { + return Box::pin(self.reconnect()).await; + } + // The first caller starts the work; every other caller waits for that same result. let result = self.startup.get_or_init(|| self.connect_once()).await; match result { @@ -2681,6 +2688,77 @@ mod tests { assert!(Arc::ptr_eq(&first, &second)); } + #[tokio::test] + async fn retryable_startup_failure_does_not_burn_environment() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (mut failed_startup, _) = listener.accept().await.expect("startup should arrive"); + failed_startup + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + + let mut replacement = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut replacement, + "replacement-session", + /*expected_resume_session_id*/ None, + ) + .await; + replacement_initialized_tx + .send(()) + .expect("replacement initialization should be observed"); + timeout(Duration::from_secs(1), replacement.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + + let failed_startup = match client.get().await { + Ok(_) => panic!("initial connection should fail"), + Err(error) => error, + }; + assert!(matches!( + failed_startup, + super::ExecServerError::ConnectionAttempt(_) + )); + assert!(client.startup_finished()); + assert!(matches!( + client.status().await, + EnvironmentObservedStatus::Disconnected { .. } + )); + + let (ready, first, second) = + tokio::join!(client.wait_until_ready(), client.get(), client.get()); + ready.expect("later readiness check should retry startup"); + let first = first.expect("first waiter should receive the replacement client"); + let second = second.expect("second waiter should receive the same replacement client"); + assert_eq!(first.session_id().as_deref(), Some("replacement-session")); + assert!(Arc::ptr_eq(&first.inner, &second.inner)); + replacement_initialized_rx + .await + .expect("server should observe replacement initialization"); + + drop(first); + drop(second); + drop(client); + server.await.expect("server task should finish"); + } + #[tokio::test] async fn failed_reconnect_does_not_burn_environment() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 4318840ab5..6588de2109 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -770,7 +770,7 @@ impl ExecServerClient { } } -fn is_retryable_recovery_error(error: &ExecServerError) -> bool { +pub(super) fn is_retryable_recovery_error(error: &ExecServerError) -> bool { is_transport_closed_error(error) || matches!( error, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 9b5f914841..32a669ecef 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -961,14 +961,14 @@ impl Environment { } } - /// Returns whether initial startup has either succeeded or permanently failed. + /// Returns whether the initial startup attempt has completed. pub fn startup_finished(&self) -> bool { self.remote_client .as_ref() .is_none_or(LazyRemoteExecServerClient::startup_finished) } - /// Waits for initial startup. A failed startup is never attempted again. + /// Waits for initial startup, retrying a previous transient failure when possible. pub async fn wait_until_ready(&self) -> Result<(), ExecServerError> { match &self.remote_client { Some(client) => client.wait_until_ready().await,