From de23ba8ef0ad3f8eea5d75398d2318820cd995c9 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 12 Jun 2026 00:46:12 -0700 Subject: [PATCH] [codex] support pending environment handles --- codex-rs/exec-server/src/client.rs | 224 ++++++++++++-- codex-rs/exec-server/src/environment.rs | 273 +++++++++++++++--- .../exec-server/src/environment_provider.rs | 6 +- codex-rs/exec-server/src/environment_toml.rs | 2 +- 4 files changed, 430 insertions(+), 75 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 978c502b7d..5c49aaf0a0 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -200,51 +200,141 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { - transport_params: ExecServerTransportParams, - client: Arc>>, + state: Arc>, + transport_revision_tx: watch::Sender, connect_lock: Arc, } +struct LazyRemoteExecServerClientState { + transport_params: Option>, + client: Option, +} + impl LazyRemoteExecServerClient { pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self { + Self::with_transport(Some(transport_params)) + } + + pub(crate) fn pending() -> Self { + Self::with_transport(/*transport_params*/ None) + } + + fn with_transport(transport_params: Option) -> Self { + let (transport_revision_tx, _transport_revision_rx) = watch::channel(0); Self { - transport_params, - client: Arc::new(StdMutex::new(None)), + state: Arc::new(StdMutex::new(LazyRemoteExecServerClientState { + transport_params: transport_params.map(Arc::new), + client: None, + })), + transport_revision_tx, connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)), } } pub(crate) async fn get(&self) -> Result { - if let Some(client) = self.connected_client() { - return Ok(client); - } - - let _connect_permit = self.connect_lock.acquire().await.map_err(|_| { - ExecServerError::Protocol("exec-server connect lock closed".to_string()) - })?; - if let Some(client) = self.connected_client() { - return Ok(client); - } - - let next_client = match self.cached_client() { - Some(_client) - if matches!( - &self.transport_params, - ExecServerTransportParams::WebSocketUrl { .. } - ) => - { - ExecServerClient::connect_for_transport(self.transport_params.clone()).await? + loop { + if let Some(client) = self.connected_client() { + return Ok(client); } - Some(client) => return Ok(client), - None => ExecServerClient::connect_for_transport(self.transport_params.clone()).await?, - }; - let mut cached_client = self - .client + let _connect_permit = self.connect_lock.acquire().await.map_err(|_| { + ExecServerError::Protocol("exec-server connect lock closed".to_string()) + })?; + if let Some(client) = self.connected_client() { + return Ok(client); + } + + let transport_params = self.wait_for_transport().await?; + if let Some(client) = self.cached_client() + && matches!( + transport_params.as_ref(), + ExecServerTransportParams::StdioCommand { .. } + ) + { + return Ok(client); + } + let mut transport_revision_rx = self.transport_revision_tx.subscribe(); + if !self.is_current_transport(&transport_params) { + continue; + } + let next_client = tokio::select! { + result = ExecServerClient::connect_for_transport(transport_params.as_ref().clone()) => { + match result { + Ok(client) => client, + Err(_err) if !self.is_current_transport(&transport_params) => continue, + Err(err) => return Err(err), + } + } + changed = transport_revision_rx.changed() => { + changed.map_err(|_| ExecServerError::Closed)?; + continue; + } + }; + + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .transport_params + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &transport_params)) + { + state.client = Some(next_client.clone()); + return Ok(next_client); + } + } + } + + pub(crate) fn set_transport(&self, transport_params: ExecServerTransportParams) { + let mut state = self + .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - *cached_client = Some(next_client.clone()); - Ok(next_client) + if state + .transport_params + .as_deref() + .is_some_and(|current| current == &transport_params) + { + return; + } + state.transport_params = Some(Arc::new(transport_params)); + state.client = None; + drop(state); + self.transport_revision_tx + .send_modify(|revision| *revision += 1); + } + + pub(crate) fn transport_params(&self) -> Option> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .transport_params + .clone() + } + + pub(crate) async fn wait_for_transport( + &self, + ) -> Result, ExecServerError> { + let mut transport_revision_rx = self.transport_revision_tx.subscribe(); + loop { + if let Some(transport_params) = self.transport_params() { + return Ok(transport_params); + } + transport_revision_rx + .changed() + .await + .map_err(|_| ExecServerError::Closed)?; + } + } + + fn is_current_transport(&self, transport_params: &Arc) -> bool { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .transport_params + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, transport_params)) } fn connected_client(&self) -> Option { @@ -253,9 +343,10 @@ impl LazyRemoteExecServerClient { } fn cached_client(&self) -> Option { - self.client + self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) + .client .clone() } } @@ -1589,6 +1680,77 @@ mod tests { server.await.expect("server task should finish"); } + #[tokio::test] + async fn remote_websocket_client_cancels_stale_connection_after_transport_update() { + let stale_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("stale listener should bind"); + let stale_websocket_url = format!( + "ws://{}", + stale_listener + .local_addr() + .expect("stale listener should have address") + ); + let (stale_accepted_tx, stale_accepted_rx) = oneshot::channel(); + let (release_stale_tx, release_stale_rx) = oneshot::channel(); + let stale_server = tokio::spawn(async move { + let _stale_websocket = accept_websocket(&stale_listener).await; + stale_accepted_tx + .send(()) + .expect("stale connection should signal"); + let _ = release_stale_rx.await; + }); + + let ready_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("ready listener should bind"); + let ready_websocket_url = format!( + "ws://{}", + ready_listener + .local_addr() + .expect("ready listener should have address") + ); + let ready_server = tokio::spawn(async move { + let mut websocket = accept_websocket(&ready_listener).await; + complete_websocket_initialize( + &mut websocket, + "ready-session", + /*expected_resume_session_id*/ None, + ) + .await; + }); + + let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl { + websocket_url: stale_websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(60), + }); + let connecting_client = client.clone(); + let connect_task = tokio::spawn(async move { connecting_client.get().await }); + stale_accepted_rx + .await + .expect("stale connection should be accepted"); + + client.set_transport(ExecServerTransportParams::WebSocketUrl { + websocket_url: ready_websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + }); + + let connected = timeout(Duration::from_secs(1), connect_task) + .await + .expect("updated transport should not wait for the stale timeout") + .expect("connection task should not panic") + .expect("updated transport should connect"); + assert_eq!(connected.session_id().as_deref(), Some("ready-session")); + + release_stale_tx + .send(()) + .expect("stale server should be released"); + stale_server.await.expect("stale server should finish"); + ready_server.await.expect("ready server should finish"); + } + #[tokio::test] async fn wake_notifications_do_not_block_other_sessions() { let (client_stdin, server_reader) = duplex(1 << 20); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index bc100b9422..1b4d28cbfa 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -31,7 +31,7 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// Owns the execution/filesystem environments available to the Codex runtime. /// -/// `EnvironmentManager` is a shared registry for concrete environments. Its +/// `EnvironmentManager` is a shared registry for environments. Its /// default constructor preserves the legacy `CODEX_EXEC_SERVER_URL` behavior /// while configured construction accepts a provider-supplied snapshot. /// @@ -251,6 +251,25 @@ impl EnvironmentManager { .cloned() } + /// Registers a named remote environment before its exec-server endpoint is available. + pub fn register_pending_environment( + &self, + environment_id: String, + ) -> Result, ExecServerError> { + validate_environment_id(&environment_id)?; + let mut environments = self + .environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(Arc::clone( + environments.entry(environment_id).or_insert_with(|| { + Arc::new(Environment::pending_remote( + self.local_runtime_paths.clone(), + )) + }), + )) + } + /// Adds or replaces a named remote environment without changing the /// manager's default environment selection. pub fn upsert_environment( @@ -258,11 +277,7 @@ impl EnvironmentManager { environment_id: String, exec_server_url: String, ) -> Result<(), ExecServerError> { - if environment_id.is_empty() { - return Err(ExecServerError::Protocol( - "environment id cannot be empty".to_string(), - )); - } + validate_environment_id(&environment_id)?; let (exec_server_url, disabled) = normalize_exec_server_url(Some(exec_server_url)); if disabled { return Err(ExecServerError::Protocol( @@ -274,24 +289,47 @@ impl EnvironmentManager { "remote environment requires an exec-server url".to_string(), )); }; - let environment = - Environment::remote_inner(exec_server_url, self.local_runtime_paths.clone()); - self.environments + let mut environments = self + .environments .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(environment_id, Arc::new(environment)); + .unwrap_or_else(std::sync::PoisonError::into_inner); + match environments.get(&environment_id) { + Some(existing) if existing.is_remote() => existing + .set_remote_transport(ExecServerTransportParams::websocket_url(exec_server_url)), + _ => { + environments.insert( + environment_id, + Arc::new(Environment::remote_inner( + exec_server_url, + self.local_runtime_paths.clone(), + )), + ); + } + } Ok(()) } } -/// Concrete execution/filesystem environment selected for a session. +fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> { + if environment_id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id cannot be empty".to_string(), + )); + } + if environment_id == LOCAL_ENVIRONMENT_ID { + return Err(ExecServerError::Protocol(format!( + "environment id `{LOCAL_ENVIRONMENT_ID}` is reserved for EnvironmentManager" + ))); + } + Ok(()) +} + +/// Execution/filesystem environment selected for a session. /// -/// This bundles the selected backend metadata together with the local runtime -/// paths used by filesystem helpers. +/// The handle remains stable while its concrete backend is supplied or replaced. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, - remote_transport: Option, + remote_client: Option, info_provider: Arc, exec_backend: Arc, filesystem: Arc, @@ -332,8 +370,7 @@ impl Environment { /// Builds a test-only local environment without configured sandbox helper paths. pub fn default_for_tests() -> Self { Self { - exec_server_url: None, - remote_transport: None, + remote_client: None, info_provider: Arc::new(LocalEnvironmentInfoProvider), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), @@ -346,7 +383,9 @@ impl Environment { impl std::fmt::Debug for Environment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Environment") - .field("exec_server_url", &self.exec_server_url) + .field("is_remote", &self.is_remote()) + .field("is_ready", &self.is_ready()) + .field("exec_server_url", &self.exec_server_url()) .finish_non_exhaustive() } } @@ -389,8 +428,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { - exec_server_url: None, - remote_transport: None, + remote_client: None, info_provider: Arc::new(LocalEnvironmentInfoProvider), exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( @@ -401,6 +439,10 @@ impl Environment { } } + fn pending_remote(local_runtime_paths: Option) -> Self { + Self::remote_with_client(LazyRemoteExecServerClient::pending(), local_runtime_paths) + } + pub(crate) fn remote_inner( exec_server_url: String, local_runtime_paths: Option, @@ -415,21 +457,22 @@ impl Environment { remote_transport: ExecServerTransportParams, local_runtime_paths: Option, ) -> Self { - let exec_server_url = match &remote_transport { - ExecServerTransportParams::WebSocketUrl { - websocket_url: exec_server_url, - .. - } => Some(exec_server_url.clone()), - ExecServerTransportParams::StdioCommand { .. } => None, - }; - let client = LazyRemoteExecServerClient::new(remote_transport.clone()); + Self::remote_with_client( + LazyRemoteExecServerClient::new(remote_transport), + local_runtime_paths, + ) + } + + fn remote_with_client( + client: LazyRemoteExecServerClient, + local_runtime_paths: Option, + ) -> Self { let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url, - remote_transport: Some(remote_transport), + remote_client: Some(client.clone()), info_provider: Arc::new(RemoteEnvironmentInfoProvider::new(client.clone())), exec_backend, filesystem, @@ -439,12 +482,26 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.remote_transport.is_some() + self.remote_client.is_some() } - /// Returns the remote exec-server URL when this environment is remote. - pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() + pub fn is_ready(&self) -> bool { + self.remote_client + .as_ref() + .is_none_or(|client| client.transport_params().is_some()) + } + + /// Returns the currently configured remote exec-server URL. + pub fn exec_server_url(&self) -> Option { + self.remote_client + .as_ref()? + .transport_params() + .and_then(|transport_params| match transport_params.as_ref() { + ExecServerTransportParams::WebSocketUrl { websocket_url, .. } => { + Some(websocket_url.clone()) + } + ExecServerTransportParams::StdioCommand { .. } => None, + }) } pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { @@ -467,6 +524,20 @@ impl Environment { pub fn get_filesystem(&self) -> Arc { Arc::clone(&self.filesystem) } + + fn set_remote_transport(&self, remote_transport: ExecServerTransportParams) { + 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 { @@ -489,6 +560,7 @@ impl From for ShellInfo { #[cfg(test)] mod tests { use std::sync::Arc; + use std::time::Duration; use super::Environment; use super::EnvironmentManager; @@ -570,7 +642,10 @@ mod tests { Some(REMOTE_ENVIRONMENT_ID) ); assert!(environment.is_remote()); - assert_eq!(environment.exec_server_url(), Some("ws://127.0.0.1:8765")); + assert_eq!( + environment.exec_server_url().as_deref(), + Some("ws://127.0.0.1:8765") + ); assert!(Arc::ptr_eq( &environment, &manager @@ -760,7 +835,7 @@ mod tests { assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths)); let manager = EnvironmentManager::create_for_tests( - environment.exec_server_url().map(str::to_owned), + environment.exec_server_url(), Some( environment .local_runtime_paths() @@ -816,16 +891,96 @@ mod tests { #[tokio::test] async fn environment_manager_upserts_named_remote_environment() { + let reserved_listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("test exec-server port should bind"); + let exec_server_address = reserved_listener + .local_addr() + .expect("test exec-server should have an address"); + drop(reserved_listener); + let exec_server_url = format!("ws://{exec_server_address}"); + let exec_server_url_for_task = exec_server_url.clone(); + let exec_server_task = tokio::spawn(async move { + crate::run_main(&exec_server_url_for_task, test_runtime_paths()).await + }); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if tokio::net::TcpStream::connect(exec_server_address) + .await + .is_ok() + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("test exec-server should start"); + let manager = EnvironmentManager::without_environments(); + let pending = manager + .register_pending_environment("executor-a".to_string()) + .expect("pending environment"); + assert!(pending.is_remote() && !pending.is_ready()); + let exec_backend = pending.get_exec_backend(); + let http_client = pending.get_http_client(); + let filesystem = pending.get_filesystem(); + let waiting_exec_backend = Arc::clone(&exec_backend); + let exec_waiter = tokio::spawn(async move { + waiting_exec_backend + .start(crate::ExecParams { + process_id: ProcessId::from("pending-environment-proc"), + argv: vec!["true".to_string()], + cwd: std::env::current_dir().expect("read current dir"), + env_policy: None, + env: Default::default(), + tty: false, + pipe_stdin: false, + arg0: None, + }) + .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(), "ws://127.0.0.1:8765".to_string()) + .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") + .expect("pending capability task should not panic") + .expect("pending capability should execute"); + assert_eq!( + started.process.process_id().as_str(), + "pending-environment-proc" + ); let first = manager .get_environment("executor-a") .expect("first remote environment"); + assert!(Arc::ptr_eq(&pending, &first) && first.is_ready()); assert!(first.is_remote()); - assert_eq!(first.exec_server_url(), Some("ws://127.0.0.1:8765")); + assert_eq!( + first.exec_server_url().as_deref(), + Some(exec_server_url.as_str()) + ); assert_eq!(manager.default_environment_id(), None); manager @@ -835,8 +990,44 @@ mod tests { .get_environment("executor-a") .expect("second remote environment"); assert!(second.is_remote()); - assert_eq!(second.exec_server_url(), Some("ws://127.0.0.1:9876")); - assert!(!Arc::ptr_eq(&first, &second)); + assert_eq!( + second.exec_server_url().as_deref(), + Some("ws://127.0.0.1:9876") + ); + assert!(Arc::ptr_eq(&first, &second)); + assert!(Arc::ptr_eq(&exec_backend, &second.get_exec_backend())); + assert!(Arc::ptr_eq(&http_client, &second.get_http_client())); + assert!(Arc::ptr_eq(&filesystem, &second.get_filesystem())); + + exec_server_task.abort(); + let _ = exec_server_task.await; + } + + #[tokio::test] + async fn environment_manager_rejects_reserved_dynamic_environment_id() { + let manager = EnvironmentManager::default_for_tests(); + + let errors = vec![ + manager + .register_pending_environment(LOCAL_ENVIRONMENT_ID.to_string()) + .expect_err("pending local environment should fail") + .to_string(), + manager + .upsert_environment( + LOCAL_ENVIRONMENT_ID.to_string(), + "ws://127.0.0.1:8765".to_string(), + ) + .expect_err("remote local environment should fail") + .to_string(), + ]; + + assert_eq!( + errors, + vec![ + "exec-server protocol error: environment id `local` is reserved for EnvironmentManager", + "exec-server protocol error: environment id `local` is reserved for EnvironmentManager", + ] + ); } #[tokio::test] diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index ffb5cbd116..53b5e0aa35 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -178,7 +178,7 @@ mod tests { let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; assert!(remote_environment.is_remote()); assert_eq!( - remote_environment.exec_server_url(), + remote_environment.exec_server_url().as_deref(), Some("ws://127.0.0.1:8765") ); assert_eq!( @@ -194,7 +194,9 @@ mod tests { let environments: HashMap<_, _> = snapshot.environments.into_iter().collect(); assert_eq!( - environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), + environments[REMOTE_ENVIRONMENT_ID] + .exec_server_url() + .as_deref(), Some("ws://127.0.0.1:8765") ); } diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 87d5bea8b3..16dd396b9f 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -390,7 +390,7 @@ mod tests { assert!(include_local); assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID)); assert_eq!( - environments["devbox"].exec_server_url(), + environments["devbox"].exec_server_url().as_deref(), Some("ws://127.0.0.1:8765") ); assert!(environments["ssh-dev"].is_remote());