openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add explicit remote executor connection refresh (#40710)

Browse Source 
## Why

Planned executor replacement needs a fresh session without waiting for the old
session's transient-disconnect recovery to finish.

## What changed

- Add `Environment::refresh_connection` for remote Noise registry-backed
  environments. It performs a fresh registry lookup, reuses a healthy session
  when the executor identity is unchanged, and connects to a replacement when
  it has changed.
- Retire superseded sessions and connection attempts so they cannot publish
  stale state, accept late RPC results, or replay outstanding work.
- Preserve the existing environment and filesystem handles while replacing the
  underlying client, and require a live status probe before refresh succeeds.

## Testing

Add coverage for replacement and session reuse, recovery and connection races,
lookup and handshake failures, handle preservation, and late RPC responses.

GitOrigin-RevId: f1d11208cbfe8af8feb25f6b6b8100da82169a99
This commit is contained in:
Rasmus Rygaard
2026-08-25 20:52:36 +00:00
committed by copyberry
parent 5ca4175295
commit eb49f491c6
8 changed files with 1029 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
codex-rs/exec-server/src/client.rs
View File

@@ -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<Semaphore>,
session_id: OnceLock<String>,
retired: CancellationToken,
/// Caches metadata from initialization or the first successful info request for this client's lifetime.
environment_info: OnceCell<EnvironmentInfo>,
reconnect_strategy: Option<ExecServerReconnectStrategy>,
@@ -347,7 +351,6 @@ impl Drop for PendingProcessStartSession {
}
type ConnectionResult = Result<ExecServerClient, Arc<ExecServerError>>;
type ConnectionAttempt = OnceCell<ConnectionResult>;
#[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<StdMutex<Option<ExecServerClient>>>,
reconnect: Arc<StdMutex<Option<Arc<ConnectionAttempt>>>>,
refresh_lock: Arc<Mutex<()>>,
environment_connection_state_tx: watch::Sender<EnvironmentConnectionState>,
}
@@ -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<Result<(), ExecServerError>> {
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<ExecServerClient, ExecServerError> {
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<ExecServerClient, ExecServerError> {
@@ -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<EnvironmentInfo, ExecServerError> {
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<EnvironmentStatus, ExecServerError> {
// 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<P, T>(&self, method: &str, params: &P) -> Result<T, ExecServerError>
@@ -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<T>(result: Result<T, RpcCallError>) -> Result<T, ExecServerError> {
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<T>(
&self,
result: Result<T, RpcCallError>,
) -> Result<T, ExecServerError> {
// 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(

8
codex-rs/exec-server/src/client/accepted.rs
View File

@@ -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,
}
}

6
codex-rs/exec-server/src/client_recovery.rs
View File

@@ -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) => {},
}
});
}

269
codex-rs/exec-server/src/client_refresh.rs Normal file
View File

@@ -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<ConnectionResult>,
pub(super) cancelled: CancellationToken,
pub(super) transport: Option<ExecServerTransportParams>,
}
// 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<Option<NoiseRendezvousConnectBundle>>,
provider: Arc<dyn NoiseRendezvousConnectProvider>,
}
impl NoiseRendezvousConnectProvider for PrefetchedConnectProvider {
fn connect_bundle(
&self,
harness_public_key: NoiseChannelPublicKey,
) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
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, &current.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<Self>) {
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;

620
codex-rs/exec-server/src/client_refresh_tests.rs Normal file
View File

@@ -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<Target>,
next_lookup: Mutex<Option<oneshot::Receiver<()>>>,
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<NoiseRendezvousConnectBundle, ExecServerError>> {
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<Arc<Notify>>,
started: Arc<Notify>,
}
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<Self> {
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<Registry>)> {
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<codex_exec_server_protocol::JSONRPCMessage>,
responses: tokio::sync::mpsc::Sender<crate::connection::JsonRpcConnectionEvent>,
}
async fn controlled_rpc() -> Result<ControlledRpc> {
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(())
}

38
codex-rs/exec-server/src/client_transport.rs
View File

@@ -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<dyn NoiseRendezvousConnectProvider>,
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<dyn NoiseRendezvousConnectProvider>,
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));
}
}
}
}

2
codex-rs/exec-server/src/client_transport_tests.rs
View File

@@ -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(),

43
codex-rs/exec-server/src/environment.rs
View File

@@ -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()