diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 63ed561c4a..bd758cdd9c 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -131,11 +131,14 @@ pub(crate) mod http_client; mod network_policy_audit; #[path = "client_recovery.rs"] mod recovery; +#[path = "client_refresh.rs"] +mod refresh; #[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; +use refresh::ConnectionAttempt; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); @@ -264,6 +267,7 @@ struct Inner { // Keep admission shared while recovered transports finish older requests. rpc_inbound_request_slots: Arc, session_id: OnceLock, + retired: CancellationToken, /// Caches metadata from initialization or the first successful info request for this client's lifetime. environment_info: OnceCell, reconnect_strategy: Option, @@ -347,7 +351,6 @@ impl Drop for PendingProcessStartSession { } type ConnectionResult = Result>; -type ConnectionAttempt = OnceCell; #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { @@ -359,6 +362,7 @@ pub(crate) struct LazyRemoteExecServerClient { // The latest successful client, replaced whenever reconnecting succeeds. current_client: Arc>>, reconnect: Arc>>>, + refresh_lock: Arc>, environment_connection_state_tx: watch::Sender, } @@ -371,9 +375,10 @@ impl LazyRemoteExecServerClient { transport_params: Some(transport_params), http_client_factory, recovery_policy: RecoveryPolicy::Wait, - startup: Arc::new(ConnectionAttempt::new()), + startup: Arc::new(ConnectionAttempt::default()), current_client: Arc::new(StdMutex::new(None)), reconnect: Arc::new(StdMutex::new(None)), + refresh_lock: Arc::new(Mutex::new(())), environment_connection_state_tx: watch::channel( EnvironmentConnectionState::Disconnected, ) @@ -406,14 +411,15 @@ impl LazyRemoteExecServerClient { } pub(crate) fn startup_finished(&self) -> bool { - self.startup.get().is_some() + // Explicit refresh can install the first client without polling startup. + self.cached_client().is_some() || self.startup.result.get().is_some() } pub(crate) fn readiness_result(&self) -> Option> { if let Some(client) = self.cached_client() { return client.readiness_result(); } - self.startup.get().and_then(|result| match result { + self.startup.result.get().and_then(|result| match result { Ok(client) => client.readiness_result(), Err(error) => Some(Err(ExecServerError::ConnectionAttempt(Arc::clone(error)))), }) @@ -425,7 +431,7 @@ impl LazyRemoteExecServerClient { Ok(client) => client, Err(error) => { // Without a completed startup attempt, there is no exec-server connection to probe. - if self.cached_client().is_none() && self.startup.get().is_none() { + if self.cached_client().is_none() && self.startup.result.get().is_none() { return crate::EnvironmentObservedStatus::Pending; } // A known connection failure is reported without retrying it as part of status. @@ -458,7 +464,7 @@ impl LazyRemoteExecServerClient { if matches!(self.recovery_policy, RecoveryPolicy::FailFast) { let client = match self.cached_client() { Some(client) => client, - None => match self.startup.get() { + None => match self.startup.result.get() { Some(Ok(client)) => client.clone(), Some(Err(error)) => { return Err(ExecServerError::ConnectionAttempt(Arc::clone(error))); @@ -492,28 +498,23 @@ 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) + if self.can_reconnect() + && (self.startup.cancelled.is_cancelled() + || self.startup.result.get().is_some_and(|result| { + result + .as_ref() + .is_err_and(|error| 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 { - Ok(client) => { - let mut current_client = self - .current_client - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if current_client.is_none() { - *current_client = Some(client.clone()); - } - Ok(client.clone()) - } - Err(error) => Err(ExecServerError::ConnectionAttempt(Arc::clone(error))), - } + self.startup + .result + .get_or_init(|| self.connect_once(&self.startup)) + .await + .clone() + .map_err(ExecServerError::ConnectionAttempt) } async fn reconnect(&self) -> Result { @@ -527,20 +528,12 @@ impl LazyRemoteExecServerClient { return Ok(client); } reconnect - .get_or_insert_with(|| Arc::new(ConnectionAttempt::new())) + .get_or_insert_with(|| Arc::new(ConnectionAttempt::default())) .clone() }; let result = attempt - .get_or_init(|| async { - let result = self.connect_once().await; - if let Ok(client) = &result { - *self - .current_client - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(client.clone()); - } - result - }) + .result + .get_or_init(|| self.connect_once(&attempt)) .await; let mut reconnect = self .reconnect @@ -578,26 +571,6 @@ impl LazyRemoteExecServerClient { ) ) } - - #[tracing::instrument(name = "codex.exec_server.remote.connect", skip_all)] - async fn connect_once(&self) -> ConnectionResult { - let transport_params = self.transport_params.as_ref().ok_or_else(|| { - Arc::new(ExecServerError::Protocol( - "missing transport params for lazy exec-server connection".to_string(), - )) - })?; - let result = ExecServerClient::connect_for_transport( - transport_params.clone(), - self.http_client_factory.clone(), - ) - .await - .map_err(Arc::new); - if let Ok(client) = &result { - client - .attach_environment_connection_state(self.environment_connection_state_tx.clone()); - } - result - } } impl HttpClient for LazyRemoteExecServerClient { @@ -780,7 +753,7 @@ impl ExecServerClient { // TODO: Remove after app-server migrates off this call. pub async fn force_environment_info(&self) -> Result { let rpc_client = self.rpc_client().await?; - map_rpc_call_result( + self.map_rpc_call_result( rpc_client .call_with_timeout(ENVIRONMENT_INFO_METHOD, &(), ENVIRONMENT_INFO_TIMEOUT) .await, @@ -797,7 +770,7 @@ impl ExecServerClient { pub async fn environment_status(&self) -> Result { // Health checks only reuse an existing RPC connection and never initiate recovery. let rpc_client = self.rpc_client_without_recovery()?; - map_rpc_call_result( + self.map_rpc_call_result( rpc_client .call_with_timeout(ENVIRONMENT_STATUS_METHOD, &(), ENVIRONMENT_STATUS_TIMEOUT) .await, @@ -1045,6 +1018,12 @@ impl ExecServerClient { .with_current_subscriber(), ); let result = result_rx.await; + // The response task may have queued a session before retirement. + if self.inner.retired.is_cancelled() { + return Err(ExecServerError::Disconnected( + "exec-server executor was replaced".to_string(), + )); + } if matches!(&result, Ok(Ok(_))) { pending_start.armed = false; let _ = result_received_tx.send(()); @@ -1130,6 +1109,7 @@ impl ExecServerClient { http_body_stream_next_id: AtomicU64::new(1), rpc_inbound_request_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_SERVER_CALLS)), session_id, + retired: CancellationToken::new(), environment_info: OnceCell::new(), reconnect_strategy, }); @@ -1170,7 +1150,7 @@ impl ExecServerClient { P: serde::Serialize, T: serde::de::DeserializeOwned, { - map_rpc_call_result(rpc_client.call(method, params).await) + self.map_rpc_call_result(rpc_client.call(method, params).await) } async fn call_for_cleanup(&self, method: &str, params: &P) -> Result @@ -1179,19 +1159,29 @@ impl ExecServerClient { T: serde::de::DeserializeOwned, { let rpc_client = self.inner.rpc_client().await?; - map_rpc_call_result(rpc_client.call_for_cleanup(method, params).await) + self.map_rpc_call_result(rpc_client.call_for_cleanup(method, params).await) } -} -fn map_rpc_call_result(result: Result) -> Result { - result.map_err(|error| { - let error = ExecServerError::from(error); - if is_transport_closed_error(&error) { - ExecServerError::Disconnected(disconnected_message(/*reason*/ None)) - } else { - error + fn map_rpc_call_result( + &self, + result: Result, + ) -> Result { + // Explicit retirement rejects late responses. Ordinary EOF still preserves + // responses received before disconnect, as ordered by the RPC reader. + if self.inner.retired.is_cancelled() { + return Err(ExecServerError::Disconnected( + "exec-server executor was replaced".to_string(), + )); } - }) + result.map_err(|error| { + let error = ExecServerError::from(error); + if is_transport_closed_error(&error) { + ExecServerError::Disconnected(disconnected_message(/*reason*/ None)) + } else { + error + } + }) + } } async fn cleanup_process_start( diff --git a/codex-rs/exec-server/src/client/accepted.rs b/codex-rs/exec-server/src/client/accepted.rs index 535eecb694..c7853221a0 100644 --- a/codex-rs/exec-server/src/client/accepted.rs +++ b/codex-rs/exec-server/src/client/accepted.rs @@ -243,9 +243,13 @@ impl LazyRemoteExecServerClient { transport_params: None, http_client_factory, recovery_policy: super::RecoveryPolicy::Wait, - startup: std::sync::Arc::new(OnceCell::new_with(Some(Ok(client)))), - current_client: std::sync::Arc::new(std::sync::Mutex::new(None)), + startup: std::sync::Arc::new(super::ConnectionAttempt { + result: OnceCell::new_with(Some(Ok(client.clone()))), + ..Default::default() + }), + current_client: std::sync::Arc::new(std::sync::Mutex::new(Some(client))), reconnect: std::sync::Arc::new(std::sync::Mutex::new(None)), + refresh_lock: std::sync::Arc::new(tokio::sync::Mutex::new(())), environment_connection_state_tx, } } diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 0c0999cb18..25bc54819e 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -351,7 +351,11 @@ impl Inner { self.notify_connection_changed(); let inner = Arc::clone(self); tokio::spawn(async move { - inner.recover(disconnect_message).await; + tokio::select! { + biased; + _ = inner.retired.cancelled() => {}, + _ = inner.recover(disconnect_message) => {}, + } }); } diff --git a/codex-rs/exec-server/src/client_refresh.rs b/codex-rs/exec-server/src/client_refresh.rs new file mode 100644 index 0000000000..9918227b71 --- /dev/null +++ b/codex-rs/exec-server/src/client_refresh.rs @@ -0,0 +1,269 @@ +//! Explicit connection refresh after a planned executor replacement. +//! +//! Ordinary recovery tries to resume the same executor session after a transient +//! disconnect. Replacement needs a fresh session, without waiting for old recovery +//! to give up. The caller supplies the ordering: register the replacement first, +//! then refresh. Executor identity stays inside this connection layer. +//! +//! Flow: fresh registry lookup -> reuse or retire session -> connect if needed -> +//! live status probe. The lazy client and its `Environment` remain the same objects; +//! only the underlying `ExecServerClient` may change. The public caller contract is on +//! `Environment::refresh_connection`. +//! +//! Two races determine the synchronization here. A client installed during the +//! lookup makes that lookup stale, so refresh checks again. A connection attempt +//! cancelled by refresh must never install later. Cancellation and installation +//! synchronize on the `current_client` lock; acquire `reconnect` first when both are needed. +//! `refresh_lock` serializes only explicit refreshes; ordinary connection and recovery +//! work can continue concurrently. Retired sessions cannot publish environment state. + +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use futures::future::BoxFuture; +use tokio::sync::OnceCell; +use tokio::sync::watch; +use tokio::time::timeout; +use tokio_util::sync::CancellationToken; + +use super::ConnectionResult; +use super::ConnectionStatus; +use super::ExecServerClient; +use super::ExecServerError; +use super::Inner; +use super::LazyRemoteExecServerClient; +use super::fail_all_in_flight_work; +use crate::EnvironmentConnectionState; +use crate::NoiseChannelPublicKey; +use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; +use crate::client_api::ExecServerTransportParams; +use crate::client_api::NoiseRendezvousConnectBundle; +use crate::client_api::NoiseRendezvousConnectProvider; +use crate::client_transport::ExecServerReconnectStrategy; + +/// Shared startup/reconnect result plus cancellation for work superseded by refresh. +/// The optional transport carries a refresh lookup's bundle into the normal connector. +#[derive(Default)] +pub(super) struct ConnectionAttempt { + pub(super) result: OnceCell, + pub(super) cancelled: CancellationToken, + pub(super) transport: Option, +} + +// Use the compared bundle intact for the first connection: address, key and authorization +// belong together. Later lookups, including authorization refresh, use the real provider. +struct PrefetchedConnectProvider { + bundle: StdMutex>, + provider: Arc, +} + +impl NoiseRendezvousConnectProvider for PrefetchedConnectProvider { + fn connect_bundle( + &self, + harness_public_key: NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let bundle = self + .bundle + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + match bundle { + Some(bundle) => Ok(bundle), + None => self.provider.connect_bundle(harness_public_key).await, + } + }) + } +} + +impl LazyRemoteExecServerClient { + #[expect( + clippy::await_holding_invalid_type, + reason = "serialize explicit refreshes, not ordinary connection or recovery attempts" + )] + pub(crate) async fn refresh_connection(&self) -> Result<(), ExecServerError> { + let _refresh = self.refresh_lock.lock().await; + let (previous, attempt) = loop { + let observed = self.cached_client(); + let mut transport = self.transport_params.clone().ok_or_else(|| { + ExecServerError::Protocol( + "connection refresh requires a Noise registry".to_string(), + ) + })?; + let target = match &mut transport { + ExecServerTransportParams::Deferred(deferred) => &mut deferred.transport, + transport => transport, + }; + let ExecServerTransportParams::NoiseRendezvous { provider, identity } = target else { + return Err(ExecServerError::Protocol( + "connection refresh requires a Noise registry".to_string(), + )); + }; + // This lookup is independent of the old session and its recovery deadline. + let bundle = timeout( + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + provider.connect_bundle(identity.public_key()), + ) + .await + .map_err(|_| { + ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + ) + })??; + let executor_public_key = bundle.executor_public_key.clone(); + *provider = Arc::new(PrefetchedConnectProvider { + bundle: StdMutex::new(Some(bundle)), + provider: Arc::clone(provider), + }); + + let mut reconnect = self + .reconnect + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let current = self + .current_client + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // An ordinary connection may have finished during the lookup. Re-read the + // registry rather than retire a newer client using a superseded response. + if !match (&observed, &*current) { + (Some(observed), Some(current)) => Arc::ptr_eq(&observed.inner, ¤t.inner), + (None, None) => true, + _ => false, + } { + continue; + } + // is_disconnected means terminally failed, not temporarily recovering. + // Preserve same-executor recovery; the final probe fails fast if still recovering. + if current.as_ref().is_some_and(|client| { + !client.is_disconnected() + && matches!( + client.inner.reconnect_strategy.as_ref(), + Some(ExecServerReconnectStrategy::NoiseRendezvous { + executor_public_key: key, .. + }) if key == &executor_public_key + ) + }) { + break (current.clone(), None); + } + // Cancellation and connection installation use the same lock. A late + // handshake cannot install a client after its attempt has been superseded. + self.startup.cancelled.cancel(); + if let Some(attempt) = reconnect.as_ref() { + attempt.cancelled.cancel(); + } + self.environment_connection_state_tx + .send_replace(EnvironmentConnectionState::Disconnected); + let attempt = Arc::new(ConnectionAttempt { + transport: Some(transport), + ..Default::default() + }); + *reconnect = Some(Arc::clone(&attempt)); + break (current.clone(), Some(attempt)); + }; + let client = match attempt { + Some(attempt) => { + if let Some(previous) = previous { + previous.inner.retire().await; + } + let result = attempt + .result + .get_or_init(|| self.connect_once(&attempt)) + .await + .clone(); + let mut reconnect = self + .reconnect + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if reconnect + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &attempt)) + { + *reconnect = None; + } + result.map_err(ExecServerError::ConnectionAttempt)? + } + None => previous.ok_or_else(|| { + ExecServerError::Protocol("current executor session is missing".to_string()) + })?, + }; + // Metadata may be cached; readiness requires a live, non-recovering probe. + client.environment_status().await.map(drop) + } + + #[tracing::instrument(name = "codex.exec_server.remote.connect", skip_all)] + pub(super) fn connect_once<'a>( + &'a self, + attempt: &'a ConnectionAttempt, + ) -> BoxFuture<'a, ConnectionResult> { + // Keep the transport future out of every caller's async layout, including + // the CLI entry point, which otherwise exceeds rustc's query-depth limit. + Box::pin(async move { + let transport = attempt + .transport + .as_ref() + .or(self.transport_params.as_ref()) + .ok_or_else(|| { + Arc::new(ExecServerError::Protocol( + "missing transport params for lazy exec-server connection".to_string(), + )) + })?; + let client = tokio::select! { + biased; + _ = attempt.cancelled.cancelled() => return Err(Arc::new(ExecServerError::Disconnected("connection attempt was superseded".to_string()))), + result = ExecServerClient::connect_for_transport(transport.clone(), self.http_client_factory.clone()) => result.map_err(Arc::new)?, + }; + // Cancellation can race with a completed handshake. Recheck before attaching + // state or installing the client, under the same lock used by refresh. + { + let mut current = self + .current_client + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !attempt.cancelled.is_cancelled() { + client.attach_environment_connection_state( + self.environment_connection_state_tx.clone(), + ); + *current = Some(client.clone()); + return Ok(client); + } + } + client.inner.retire().await; + Err(Arc::new(ExecServerError::Disconnected( + "connection attempt was superseded".to_string(), + ))) + }) + } +} + +impl Inner { + async fn retire(self: &Arc) { + let message = "exec-server executor was replaced".to_string(); + let rpc_client = { + let mut connection = self + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Detach before a later transport completion can publish stale state. + connection.environment_connection_state_tx = + watch::channel(EnvironmentConnectionState::Disconnected).0; + let rpc_client = match &connection.status { + ConnectionStatus::Connected(client) => Some(Arc::clone(client)), + ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None, + }; + self.retired.cancel(); + connection.set_status(ConnectionStatus::Failed(message.clone())); + rpc_client + }; + self.connection_changed.send_replace(()); + // Drain pending RPCs before stream cleanup, which may wait for other work. + if let Some(rpc_client) = rpc_client { + rpc_client.close_transport().await; + } + fail_all_in_flight_work(self, message).await; + } +} + +#[cfg(test)] +#[path = "client_refresh_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/client_refresh_tests.rs b/codex-rs/exec-server/src/client_refresh_tests.rs new file mode 100644 index 0000000000..34dab7bb14 --- /dev/null +++ b/codex-rs/exec-server/src/client_refresh_tests.rs @@ -0,0 +1,620 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use anyhow::Result; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::sync::oneshot; +use tokio::task::JoinSet; +use tokio_util::task::AbortOnDropHandle; + +use super::*; +use crate::ExecServerRuntimePaths; +use crate::NoiseChannelIdentity; +use crate::ProcessId; +use crate::relay::HarnessKeyValidator; +use crate::relay::run_multiplexed_environment; +use crate::server::ConnectionProcessor; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; + +// These tests exercise real sockets, not keepalive deadlines. The shared unit-test +// Pong timeout is only 100 ms and can expire during Noise handshakes under load. +// A blocking task prevents paused time from auto-advancing while socket I/O is +// pending; dropping the returned sender releases it without polling or sleeping. +fn freeze_clock() -> std::sync::mpsc::Sender<()> { + tokio::time::pause(); + let (guard, dropped) = std::sync::mpsc::channel(); + tokio::task::spawn_blocking(move || { + let _ = dropped.recv(); + }); + guard +} + +#[derive(Clone)] +struct Target { + url: String, + identity: NoiseChannelIdentity, + registration: String, +} + +struct Registry { + target: Mutex, + next_lookup: Mutex>>, + lookup_started: Notify, +} + +impl Registry { + fn block_next_lookup(&self) -> oneshot::Sender<()> { + let (tx, rx) = oneshot::channel(); + *self.next_lookup.lock().unwrap() = Some(rx); + tx + } +} + +impl NoiseRendezvousConnectProvider for Registry { + fn connect_bundle( + &self, + _: NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let target = self.target.lock().unwrap().clone(); + let block = self.next_lookup.lock().unwrap().take(); + if let Some(block) = block { + self.lookup_started.notify_one(); + block + .await + .map_err(|_| ExecServerError::Protocol("test lookup failed".to_owned()))?; + } + Ok(NoiseRendezvousConnectBundle { + websocket_url: target.url, + environment_id: "environment".to_owned(), + executor_registration_id: target.registration, + executor_public_key: target.identity.public_key(), + harness_key_authorization: "authorization".to_owned(), + }) + }) + } +} + +#[derive(Clone, Default)] +struct Validator { + handshake: Option>, + started: Arc, +} + +impl HarnessKeyValidator for Validator { + async fn validate_harness_key( + &self, + _: &NoiseChannelPublicKey, + _: &str, + ) -> Result<(), ExecServerError> { + if let Some(handshake) = &self.handshake { + self.started.notify_one(); + handshake.notified().await; + } + Ok(()) + } +} + +struct Executor { + target: Target, + _server: AbortOnDropHandle<()>, +} + +impl Executor { + async fn start(validator: Validator) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let target = Target { + url: format!("ws://{}", listener.local_addr()?), + identity: NoiseChannelIdentity::generate()?, + registration: uuid::Uuid::new_v4().to_string(), + }; + let executor = target.clone(); + let processor = ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?); + let server = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + socket = listener.accept() => { + let (socket, _) = socket.unwrap(); + let executor = executor.clone(); + let processor = processor.clone(); + let validator = validator.clone(); + connections.spawn(async move { + let socket = tokio_tungstenite::accept_async(socket).await.unwrap(); + run_multiplexed_environment(socket, processor, "environment".to_owned(), executor.registration, executor.identity, validator).await; + }); + } + _ = connections.join_next(), if !connections.is_empty() => {} + } + } + }); + Ok(Self { + target, + _server: AbortOnDropHandle::new(server), + }) + } + + fn client(&self) -> Result<(LazyRemoteExecServerClient, Arc)> { + let registry = Arc::new(Registry { + target: Mutex::new(self.target.clone()), + next_lookup: Mutex::new(None), + lookup_started: Notify::new(), + }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::NoiseRendezvous { + provider: registry.clone(), + identity: NoiseChannelIdentity::generate()?, + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + Ok((client, registry)) + } +} + +async fn disconnect(client: &ExecServerClient) { + let rpc_client = { + let connection = client.inner.connection.lock().unwrap(); + let ConnectionStatus::Connected(rpc_client) = &connection.status else { + panic!("expected connected session") + }; + Arc::clone(rpc_client) + }; + rpc_client.close_transport().await; +} + +#[tokio::test] +async fn refresh_cancels_old_recovery_and_connects_without_resuming_old_session() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let original = client.get().await?; + let process = original + .register_session(&ProcessId::from("old-process")) + .await?; + let blocked_recovery = registry.block_next_lookup(); + disconnect(&original).await; + registry.lookup_started.notified().await; + *registry.target.lock().unwrap() = new.target.clone(); + let refresh = client.refresh_connection(); + tokio::pin!(refresh); + let started = tokio::time::Instant::now(); + assert!(futures::poll!(refresh.as_mut()).is_pending()); + assert!(original.inner.retired.is_cancelled()); + assert!(original.is_disconnected()); + assert_eq!(started.elapsed(), Duration::ZERO); + refresh.await?; + let replacement = client.get().await?; + assert_ne!(original.session_id(), replacement.session_id()); + assert_eq!( + *client.environment_connection_state_tx.borrow(), + EnvironmentConnectionState::Connected + ); + assert!(matches!( + process.write(b"never replay".to_vec()).await, + Err(ExecServerError::Disconnected(_)) + )); + // Releasing an old registry response cannot reinstall or disconnect the replacement. + let _ = blocked_recovery.send(()); + tokio::task::yield_now().await; + assert!(Arc::ptr_eq(&client.get().await?.inner, &replacement.inner)); + replacement.environment_status().await?; + Ok(()) +} + +#[tokio::test] +async fn refresh_retires_a_still_connected_old_executor() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let original = client.get().await?; + *registry.target.lock().unwrap() = new.target.clone(); + client.refresh_connection().await?; + assert!(original.is_disconnected()); + assert_ne!(original.session_id(), client.get().await?.session_id()); + Ok(()) +} + +#[tokio::test] +async fn refresh_preserves_a_current_session_across_registration_renewal() -> Result<()> { + let _clock = freeze_clock(); + let executor = Executor::start(Validator::default()).await?; + let (client, registry) = executor.client()?; + let original = client.get().await?; + registry.target.lock().unwrap().registration = "renewed-registration".to_owned(); + let concurrent = client.clone(); + let concurrent = tokio::spawn(async move { concurrent.refresh_connection().await }); + client.refresh_connection().await?; + concurrent.await??; + assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner)); + original.environment_status().await?; + Ok(()) +} + +#[tokio::test] +async fn refresh_cancels_a_stalled_initial_lookup() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let release = registry.block_next_lookup(); + let initial = client.get(); + tokio::pin!(initial); + assert!(futures::poll!(initial.as_mut()).is_pending()); + *registry.target.lock().unwrap() = new.target.clone(); + client.refresh_connection().await?; + assert!(initial.await.is_err()); + let _ = release.send(()); + client.get().await?.environment_status().await?; + Ok(()) +} + +#[tokio::test] +async fn refresh_cancels_a_stalled_noise_handshake() -> Result<()> { + let _clock = freeze_clock(); + let validator = Validator { + handshake: Some(Arc::new(Notify::new())), + ..Default::default() + }; + let old = Executor::start(validator.clone()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let connecting = client.clone(); + let initial = tokio::spawn(async move { connecting.get().await }); + validator.started.notified().await; + *registry.target.lock().unwrap() = new.target.clone(); + client.refresh_connection().await?; + assert!(initial.await?.is_err()); + validator.handshake.unwrap().notify_one(); + client.get().await?.environment_status().await?; + assert_eq!( + *client.environment_connection_state_tx.borrow(), + EnvironmentConnectionState::Connected + ); + Ok(()) +} + +#[tokio::test] +async fn failed_refresh_lookup_leaves_the_existing_session_usable() -> Result<()> { + let _clock = freeze_clock(); + let executor = Executor::start(Validator::default()).await?; + let (client, registry) = executor.client()?; + let original = client.get().await?; + drop(registry.block_next_lookup()); + assert!(client.refresh_connection().await.is_err()); + assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner)); + original.environment_status().await?; + Ok(()) +} + +#[tokio::test] +async fn failed_replacement_connection_keeps_old_handles_retired_and_get_retries() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let original = client.get().await?; + let process = original + .register_session(&ProcessId::from("old-process")) + .await?; + + // The registry knows the replacement, but its endpoint drops the new connection. + let unavailable = TcpListener::bind("127.0.0.1:0").await?; + let target = Target { + url: format!("ws://{}", unavailable.local_addr()?), + ..new.target.clone() + }; + let _rejected_connection = AbortOnDropHandle::new(tokio::spawn(async move { + drop(unavailable.accept().await.unwrap()); + })); + *registry.target.lock().unwrap() = target; + assert!(matches!( + client.refresh_connection().await, + Err(ExecServerError::ConnectionAttempt(_)) + )); + assert!(original.inner.retired.is_cancelled()); + assert!(matches!( + original.environment_status().await, + Err(ExecServerError::Disconnected(_)) + )); + assert!(matches!( + process.write(b"never replay".to_vec()).await, + Err(ExecServerError::Disconnected(_)) + )); + assert_eq!( + *client.environment_connection_state_tx.borrow(), + EnvironmentConnectionState::Disconnected + ); + + // A later caller retries against the registry instead of reviving the retired client. + *registry.target.lock().unwrap() = new.target.clone(); + let replacement = client.get().await?; + assert_ne!(original.session_id(), replacement.session_id()); + replacement.environment_status().await?; + assert_eq!( + *client.environment_connection_state_tx.borrow(), + EnvironmentConnectionState::Connected + ); + assert!(matches!( + process.write(b"still retired".to_vec()).await, + Err(ExecServerError::Disconnected(_)) + )); + Ok(()) +} + +#[tokio::test] +async fn superseded_refresh_lookup_does_not_retire_a_newer_session() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (client, registry) = old.client()?; + let original = client.get().await?; + let release = registry.block_next_lookup(); + let refreshing = client.refresh_connection(); + tokio::pin!(refreshing); + assert!(futures::poll!(refreshing.as_mut()).is_pending()); + *registry.target.lock().unwrap() = new.target.clone(); + original.inner.retire().await; + let replacement = client.get().await?; + release.send(()).unwrap(); + refreshing.await?; + assert!(Arc::ptr_eq(&replacement.inner, &client.get().await?.inner)); + assert!(!replacement.inner.retired.is_cancelled()); + replacement.environment_status().await?; + Ok(()) +} + +#[tokio::test] +async fn environment_refresh_preserves_environment_and_filesystem_handles() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let new = Executor::start(Validator::default()).await?; + let (_, registry) = old.client()?; + 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 + .materialize_pending_noise_environment("environment".to_owned(), registry.clone())?; + manager.report_environment_provisioning_status( + "environment".to_owned(), + Ok(crate::EnvironmentReadyInfo { + selected_capability_roots: Vec::new(), + }), + registry.clone(), + )?; + environment.info().await?; + let filesystem = environment.get_filesystem(); + *registry.target.lock().unwrap() = new.target.clone(); + environment.refresh_connection().await?; + assert!(Arc::ptr_eq( + &environment, + &manager.get_environment("environment").unwrap() + )); + assert!(Arc::ptr_eq(&filesystem, &environment.get_filesystem())); + environment.info().await?; + Ok(()) +} + +#[tokio::test] +async fn refresh_does_not_accept_cached_metadata_during_recovery() -> Result<()> { + let _clock = freeze_clock(); + let executor = Executor::start(Validator::default()).await?; + let (client, registry) = executor.client()?; + let original = client.get().await?; + original.environment_info().await?; + let blocked_recovery = registry.block_next_lookup(); + disconnect(&original).await; + registry.lookup_started.notified().await; + let started = tokio::time::Instant::now(); + assert!(matches!( + client.refresh_connection().await, + Err(ExecServerError::Disconnected(_)) + )); + assert_eq!(started.elapsed(), Duration::ZERO); + assert!(!original.inner.retired.is_cancelled()); + drop(blocked_recovery); + Ok(()) +} + +#[tokio::test] +async fn refresh_before_startup_marks_startup_finished() -> Result<()> { + let _clock = freeze_clock(); + let executor = Executor::start(Validator::default()).await?; + let (client, _) = executor.client()?; + assert!(!client.startup_finished()); + client.refresh_connection().await?; + assert!(client.startup_finished()); + assert!(matches!(client.readiness_result(), Some(Ok(())))); + Ok(()) +} + +struct ControlledRpc { + client: ExecServerClient, + requests: tokio::sync::mpsc::Receiver, + responses: tokio::sync::mpsc::Sender, +} + +async fn controlled_rpc() -> Result { + use crate::connection::JsonRpcConnection; + use crate::connection::JsonRpcConnectionEvent; + use crate::connection::JsonRpcTransport; + use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCResponse; + let (outgoing_tx, mut requests) = tokio::sync::mpsc::channel(/*buffer*/ 8); + let (responses, incoming_rx) = tokio::sync::mpsc::channel(/*buffer*/ 8); + let connection = JsonRpcConnection { + outgoing_tx, + incoming_rx, + disconnected_rx: tokio::sync::watch::channel(/*init*/ false).1, + task_handles: Vec::new(), + transport: JsonRpcTransport::Plain, + }; + let connecting = ExecServerClient::connect(connection, /*options*/ Default::default()); + tokio::pin!(connecting); + assert!(futures::poll!(connecting.as_mut()).is_pending()); + let Some(JSONRPCMessage::Request(initialize)) = requests.recv().await else { + anyhow::bail!("expected initialize request"); + }; + responses + .send(JsonRpcConnectionEvent::message(JSONRPCMessage::Response( + JSONRPCResponse { + id: initialize.id, + result: serde_json::json!({"sessionId": "controlled-session"}), + }, + ))) + .await?; + let client = connecting.await?; + assert!(matches!( + requests.recv().await, + Some(JSONRPCMessage::Notification(_)) + )); + Ok(ControlledRpc { + client, + requests, + responses, + }) +} + +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "hold stream cleanup pending to exercise retirement ordering" +)] +async fn retirement_rejects_pending_mutation_before_stream_cleanup() -> Result<()> { + use crate::connection::JsonRpcConnectionEvent; + use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCResponse; + for response_queued in [false, true] { + let mut rpc = controlled_rpc().await?; + let call = rpc.client.fs_remove(crate::protocol::FsRemoveParams { + path: "file:///retired-file".parse()?, + recursive: None, + force: None, + follow_symlinks: None, + sandbox: None, + }); + tokio::pin!(call); + assert!(futures::poll!(call.as_mut()).is_pending()); + let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else { + anyhow::bail!("expected filesystem request"); + }; + let response = JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::json!({}), + }); + if response_queued { + rpc.responses + .send(JsonRpcConnectionEvent::message(response.clone())) + .await?; + let transport = rpc.client.rpc_client_without_recovery()?; + tokio::time::timeout(Duration::from_secs(5), async { + while transport.pending_request_count().await != 0 { + tokio::task::yield_now().await; + } + }) + .await?; + } + // Hold stream cleanup. Cover both a late response and one already queued + // for the caller when retirement begins; closing the socket alone misses the latter. + let streams = rpc.client.inner.http_body_streams_write_lock.lock().await; + let retirement = rpc.client.inner.retire(); + tokio::pin!(retirement); + assert!(futures::poll!(retirement.as_mut()).is_pending()); + if !response_queued { + rpc.responses + .send(JsonRpcConnectionEvent::message(response)) + .await?; + } + assert!(matches!(call.await, Err(ExecServerError::Disconnected(_)))); + drop(streams); + retirement.await; + } + Ok(()) +} + +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "hold stream cleanup pending to exercise retirement ordering" +)] +async fn retirement_rejects_pending_process_start_before_stream_cleanup() -> Result<()> { + use crate::connection::JsonRpcConnectionEvent; + use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCResponse; + for response_queued in [false, true] { + let mut rpc = controlled_rpc().await?; + let process_id = ProcessId::from("retired-process"); + let call = rpc.client.start_process( + crate::protocol::ExecParams { + process_id: process_id.clone(), + argv: vec!["unused".to_owned()], + cwd: "file:///".parse()?, + shell_snapshot: None, + env_policy: None, + env: Default::default(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }, + /*network_policy_decider*/ None, + ); + tokio::pin!(call); + assert!(futures::poll!(call.as_mut()).is_pending()); + let Some(JSONRPCMessage::Request(request)) = rpc.requests.recv().await else { + anyhow::bail!("expected process start request"); + }; + let response = JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::json!({"processId": "retired-process"}), + }); + if response_queued { + rpc.responses + .send(JsonRpcConnectionEvent::message(response.clone())) + .await?; + let state = rpc + .client + .inner + .get_session(&process_id) + .expect("pending process"); + // The start task has a second response channel to the original caller. + tokio::time::timeout(Duration::from_secs(5), async { + while !state.recoverable.load(std::sync::atomic::Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await?; + } + let streams = rpc.client.inner.http_body_streams_write_lock.lock().await; + let retirement = rpc.client.inner.retire(); + tokio::pin!(retirement); + assert!(futures::poll!(retirement.as_mut()).is_pending()); + if !response_queued { + rpc.responses + .send(JsonRpcConnectionEvent::message(response)) + .await?; + } + assert!(matches!(call.await, Err(ExecServerError::Disconnected(_)))); + drop(streams); + retirement.await; + } + Ok(()) +} diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index add01f69d4..bf1f9fa706 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -103,6 +103,8 @@ pub(crate) enum ExecServerReconnectStrategy { Accepted(AcceptedConnectionSource), WebSocket(RemoteExecServerConnectArgs), NoiseRendezvous { + // The executor that created the session, not the latest recovery lookup. + executor_public_key: crate::NoiseChannelPublicKey, provider: Arc, identity: NoiseChannelIdentity, client_name: String, @@ -126,6 +128,7 @@ impl ExecServerReconnectStrategy { Ok(ReconnectAttempt::new(connection, args.into())) } Self::NoiseRendezvous { + executor_public_key: _, provider, identity, client_name, @@ -201,20 +204,22 @@ impl ExecServerClient { initialize_timeout, } => (websocket_url, connect_timeout, initialize_timeout), ExecServerTransportParams::NoiseRendezvous { provider, identity } => { + let (connection, options, executor_public_key) = + Self::open_initial_noise_rendezvous_connection( + &provider, + &identity, + http_client_factory.clone(), + ) + .await?; let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous { - provider: Arc::clone(&provider), - identity: identity.clone(), + executor_public_key, + provider, + identity, client_name: ENVIRONMENT_CLIENT_NAME.to_string(), connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, - http_client_factory: http_client_factory.clone(), - }; - let (connection, options) = Self::open_initial_noise_rendezvous_connection( - &provider, - &identity, http_client_factory, - ) - .await?; + }; return Self::connect_with_recovery(connection, options, Some(reconnect_strategy)) .await; } @@ -247,7 +252,14 @@ impl ExecServerClient { provider: &Arc, identity: &NoiseChannelIdentity, http_client_factory: HttpClientFactory, - ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { + ) -> Result< + ( + JsonRpcConnection, + ExecServerClientConnectOptions, + crate::NoiseChannelPublicKey, + ), + ExecServerError, + > { let open_connection = |bundle: NoiseRendezvousConnectBundle| { Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { bundle, @@ -299,6 +311,7 @@ impl ExecServerClient { } Err(error) => return Err(error), }; + let executor_public_key = bundle.executor_public_key.clone(); match open_connection(bundle).await { Err(error) if !refreshed_unauthorized_bundle @@ -317,7 +330,10 @@ impl ExecServerClient { retries = 0; result = connect_bundle().await; } - result => return result, + result => { + return result + .map(|(connection, options)| (connection, options, executor_public_key)); + } } } } diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 411dfc0e8d..362de81654 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -85,6 +85,7 @@ impl SequenceNoiseConnectProvider { ), ) .await + .map(|(connection, options, _)| (connection, options)) } } @@ -340,6 +341,7 @@ async fn noise_session_resume_leaves_offline_retries_to_recovery() -> Result<()> )); let identity = NoiseChannelIdentity::generate()?; let strategy = ExecServerReconnectStrategy::NoiseRendezvous { + executor_public_key: NoiseChannelIdentity::generate()?.public_key(), provider: sequence.clone(), identity: identity.clone(), client_name: "test".to_string(), diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index e3f69de24f..953f8d4da2 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -902,6 +902,47 @@ impl Environment { } } + /// Refresh the connection to the executor currently registered for this environment. + /// + /// # Caller contract + /// + /// Call after a planned replacement has registered and become available under the + /// same environment ID. This method does not provision or destroy executors, or + /// wait for the registry to identify a particular replacement. It requires a remote + /// Noise registry-backed environment; other environment types return an error. + /// + /// # Session behavior + /// + /// A fresh registry lookup determines whether the current session can be reused. + /// A changed executor key, or a failed or missing session, causes a fresh connection + /// without resuming the old session. Retirement cancels old recovery, fails its + /// outstanding work and process handles, and never replays commands. The environment + /// object and filesystem handle remain usable through the new connection. + /// A matching executor key preserves a session that has not failed, including one + /// that is recovering; the live readiness check rejects a recovering connection. + /// + /// # Completion and errors + /// + /// Success means the selected connection answered a live status RPC, not merely that + /// metadata was cached. Refresh bypasses the old session's recovery deadline, but + /// registry lookup, connection, and status RPC timeouts still apply. If the initial + /// registry lookup fails, refresh leaves the old session untouched; errors after + /// retirement do not restore it. Ordinary disconnect recovery is unchanged unless + /// refresh retires the session. + #[tracing::instrument( + name = "exec_server.environment.refresh_connection", + skip_all, + fields(remote = self.is_remote()) + )] + pub async fn refresh_connection(&self) -> Result<(), ExecServerError> { + let client = self.remote_client.as_ref().ok_or_else(|| { + ExecServerError::Protocol( + "connection refresh requires a remote environment".to_string(), + ) + })?; + client.refresh_connection().await + } + /// Fetches uncached metadata, connecting or waiting for recovery as needed. // TODO: Remove after app-server migrates off of force_environment_info. #[tracing::instrument( @@ -1023,7 +1064,7 @@ impl Environment { } } - /// Returns whether the initial startup attempt has completed. + /// Returns whether startup has completed, including a first connection made by refresh. pub fn startup_finished(&self) -> bool { self.remote_client .as_ref()