From 99bb2a8080e6e59339ff8ed710f60960c90f627f Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 12 Jun 2026 16:23:38 -0700 Subject: [PATCH] [exec-server] clarify pending environment lifecycle --- codex-rs/exec-server/src/client.rs | 8 +++++++ codex-rs/exec-server/src/environment.rs | 31 ++++++------------------- 2 files changed, 15 insertions(+), 24 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 5c49aaf0a0..e9a03875f7 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -206,6 +206,7 @@ pub(crate) struct LazyRemoteExecServerClient { } struct LazyRemoteExecServerClientState { + // None means we do not know how to connect to this environment yet. transport_params: Option>, client: Option, } @@ -220,6 +221,7 @@ impl LazyRemoteExecServerClient { } fn with_transport(transport_params: Option) -> Self { + // The revision channel wakes pending callers and cancels stale connection attempts. let (transport_revision_tx, _transport_revision_rx) = watch::channel(0); Self { state: Arc::new(StdMutex::new(LazyRemoteExecServerClientState { @@ -237,6 +239,7 @@ impl LazyRemoteExecServerClient { return Ok(client); } + // Only one caller creates and publishes a connection; followers recheck the cache. let _connect_permit = self.connect_lock.acquire().await.map_err(|_| { ExecServerError::Protocol("exec-server connect lock closed".to_string()) })?; @@ -244,6 +247,7 @@ impl LazyRemoteExecServerClient { return Ok(client); } + // Wait here until upsert provides the environment's connection details. let transport_params = self.wait_for_transport().await?; if let Some(client) = self.cached_client() && matches!( @@ -257,6 +261,7 @@ impl LazyRemoteExecServerClient { if !self.is_current_transport(&transport_params) { continue; } + // Stop waiting on an obsolete endpoint as soon as a replacement arrives. let next_client = tokio::select! { result = ExecServerClient::connect_for_transport(transport_params.as_ref().clone()) => { match result { @@ -271,6 +276,7 @@ impl LazyRemoteExecServerClient { } }; + // Publish only if this connection still belongs to the current transport. let mut state = self .state .lock() @@ -301,6 +307,7 @@ impl LazyRemoteExecServerClient { state.transport_params = Some(Arc::new(transport_params)); state.client = None; drop(state); + // Notify after releasing the state lock so woken callers can inspect it immediately. self.transport_revision_tx .send_modify(|revision| *revision += 1); } @@ -318,6 +325,7 @@ impl LazyRemoteExecServerClient { ) -> Result, ExecServerError> { let mut transport_revision_rx = self.transport_revision_tx.subscribe(); loop { + // This only waits for connection details; it does not connect yet. if let Some(transport_params) = self.transport_params() { return Ok(transport_params); } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 1b4d28cbfa..9ecb01733e 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -261,6 +261,7 @@ impl EnvironmentManager { .environments .write() .unwrap_or_else(std::sync::PoisonError::into_inner); + // Reuse the same Arc so early turn and capability handles observe later updates. Ok(Arc::clone( environments.entry(environment_id).or_insert_with(|| { Arc::new(Environment::pending_remote( @@ -294,6 +295,7 @@ impl EnvironmentManager { .write() .unwrap_or_else(std::sync::PoisonError::into_inner); match environments.get(&environment_id) { + // Update the shared client inside the stable Environment handle. Some(existing) if existing.is_remote() => existing .set_remote_transport(ExecServerTransportParams::websocket_url(exec_server_url)), _ => { @@ -329,6 +331,7 @@ fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> /// The handle remains stable while its concrete backend is supplied or replaced. #[derive(Clone)] pub struct Environment { + // None is local; Some is remote and may still be waiting for its transport. remote_client: Option, info_provider: Arc, exec_backend: Arc, @@ -440,6 +443,7 @@ impl Environment { } fn pending_remote(local_runtime_paths: Option) -> Self { + // Build capability wrappers now; their first operation waits inside the lazy client. Self::remote_with_client(LazyRemoteExecServerClient::pending(), local_runtime_paths) } @@ -467,6 +471,7 @@ impl Environment { client: LazyRemoteExecServerClient, local_runtime_paths: Option, ) -> Self { + // These clones share one internal state, so every capability sees transport updates. let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); @@ -485,6 +490,7 @@ impl Environment { self.remote_client.is_some() } + /// Ready means we know how to connect; the connection is opened only when needed. pub fn is_ready(&self) -> bool { self.remote_client .as_ref() @@ -526,18 +532,11 @@ impl Environment { } fn set_remote_transport(&self, remote_transport: ExecServerTransportParams) { + // Mutate the shared client rather than replacing the Environment held by existing turns. if let Some(client) = &self.remote_client { client.set_transport(remote_transport); } } - - #[cfg(test)] - async fn wait_until_ready(&self) -> Result<(), ExecServerError> { - if let Some(client) = &self.remote_client { - client.wait_for_transport().await?; - } - Ok(()) - } } impl EnvironmentInfo { @@ -941,28 +940,12 @@ mod tests { }) .await }); - let waiting_environment = Arc::clone(&pending); - let ready_waiter = tokio::spawn(async move { - waiting_environment - .wait_until_ready() - .await - .map(|()| waiting_environment.exec_server_url()) - }); tokio::task::yield_now().await; assert!(!exec_waiter.is_finished()); manager .upsert_environment("executor-a".to_string(), exec_server_url.clone()) .expect("remote environment"); - assert_eq!( - tokio::time::timeout(Duration::from_secs(1), ready_waiter) - .await - .expect("ready waiter should be notified") - .expect("ready waiter should not panic") - .expect("environment should become ready") - .as_deref(), - Some(exec_server_url.as_str()) - ); let started = tokio::time::timeout(Duration::from_secs(2), exec_waiter) .await .expect("pending capability should wake")