openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-31 19:28:49 +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 21:23:17 +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
codex-rs/exec-server/src/client.rs
+58 -68
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(