Add observers for initial exec-server connection attempts (#52724)

## What changed

Expose `Environment::observe_connection_attempts` and `ConnectionAttemptOutcome` to report elapsed time and success, failure, or cancellation for initial connections. Measure from the end of provisioning through client installation, including connection retries.

The first registered observer wins. Past attempts and reconnections are omitted.

## Testing

Extend tests to assert failure followed by success when startup is retried, and a single cancellation notification when an environment is replaced during startup.

GitOrigin-RevId: 88300f38dcf9428b824c5ff36aa1b40d7d466b6f
This commit is contained in:
aaronl-openai
2026-10-09 23:28:38 +00:00
committed by copyberry
parent 2dc0791c7d
commit 7efd184c24
7 changed files with 93 additions and 2 deletions
+24
View File
@@ -390,6 +390,8 @@ impl Drop for PendingProcessStartSession {
}
type ConnectionResult = Result<ExecServerClient, Arc<ExecServerError>>;
pub(super) type ConnectionObserver =
Arc<dyn Fn(Duration, crate::ConnectionAttemptOutcome) + Send + Sync>;
#[derive(Clone)]
pub(crate) struct LazyRemoteExecServerClient {
@@ -402,6 +404,7 @@ pub(crate) struct LazyRemoteExecServerClient {
current_client: Arc<StdMutex<Option<ExecServerClient>>>,
reconnect: Arc<StdMutex<Option<Arc<ConnectionAttempt>>>>,
refresh_lock: Arc<Mutex<()>>,
pub(super) connection_observer: Arc<OnceLock<ConnectionObserver>>,
environment_connection_state_tx: watch::Sender<EnvironmentConnectionState>,
}
@@ -418,6 +421,7 @@ impl LazyRemoteExecServerClient {
current_client: Arc::new(StdMutex::new(None)),
reconnect: Arc::new(StdMutex::new(None)),
refresh_lock: Arc::new(Mutex::new(())),
connection_observer: Default::default(),
environment_connection_state_tx: watch::channel(
EnvironmentConnectionState::Disconnected,
)
@@ -2298,6 +2302,7 @@ mod tests {
codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
/*started*/ &mut None,
)
.await
.expect("stdio transport should connect");
@@ -3236,6 +3241,17 @@ mod tests {
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
assert!(
client
.connection_observer
.set(Arc::new(move |_, outcome| {
observed_tx
.send(outcome)
.expect("record connection outcome");
}))
.is_ok()
);
let failed_startup = match client.get().await {
Ok(_) => panic!("initial connection should fail"),
Err(error) => error,
@@ -3261,6 +3277,14 @@ mod tests {
.await
.expect("server should observe replacement initialization");
assert_eq!(
std::iter::from_fn(|| observed_rx.try_recv().ok()).collect::<Vec<_>>(),
[
crate::ConnectionAttemptOutcome::Failure,
crate::ConnectionAttemptOutcome::Success
],
);
drop(first);
drop(second);
drop(client);
@@ -251,6 +251,7 @@ impl LazyRemoteExecServerClient {
reconnect: std::sync::Arc::new(std::sync::Mutex::new(None)),
refresh_lock: std::sync::Arc::new(tokio::sync::Mutex::new(())),
environment_connection_state_tx,
connection_observer: Default::default(),
}
}
+30 -2
View File
@@ -23,6 +23,7 @@ use std::sync::Mutex as StdMutex;
use futures::future::BoxFuture;
use tokio::sync::OnceCell;
use tokio::sync::watch;
use tokio::time::Instant;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
@@ -33,6 +34,7 @@ use super::ExecServerError;
use super::Inner;
use super::LazyRemoteExecServerClient;
use super::fail_all_in_flight_work;
use crate::ConnectionAttemptOutcome;
use crate::EnvironmentConnectionState;
use crate::NoiseChannelPublicKey;
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
@@ -50,6 +52,20 @@ pub(super) struct ConnectionAttempt {
pub(super) transport: Option<ExecServerTransportParams>,
}
struct ConnectionObservation<'a> {
started: Option<Instant>,
outcome: ConnectionAttemptOutcome,
observer: Option<&'a super::ConnectionObserver>,
}
impl Drop for ConnectionObservation<'_> {
fn drop(&mut self) {
if let (Some(started), Some(observer)) = (self.started, self.observer) {
observer(started.elapsed(), self.outcome);
}
}
}
// 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 {
@@ -225,11 +241,22 @@ impl LazyRemoteExecServerClient {
"missing transport params for lazy exec-server connection".to_string(),
))
})?;
let client = tokio::select! {
let mut observation = ConnectionObservation {
started: None,
outcome: ConnectionAttemptOutcome::Cancelled,
observer: (self.cached_client().is_none() && attempt.transport.is_none())
.then(|| self.connection_observer.get())
.flatten(),
};
let result = 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)?,
result = ExecServerClient::connect_for_transport(transport.clone(), self.http_client_factory.clone(), &mut observation.started) => result,
};
let client = result.map_err(|error| {
observation.outcome = ConnectionAttemptOutcome::Failure;
Arc::new(error)
})?;
// Cancellation can race with a completed handshake. Recheck before attaching
// state or installing the client, under the same lock used by refresh.
{
@@ -242,6 +269,7 @@ impl LazyRemoteExecServerClient {
self.environment_connection_state_tx.clone(),
);
*current = Some(client.clone());
observation.outcome = ConnectionAttemptOutcome::Success;
return Ok(client);
}
}
@@ -331,6 +331,7 @@ impl ExecServerClient {
pub(crate) async fn connect_for_transport(
transport_params: ExecServerTransportParams,
http_client_factory: HttpClientFactory,
started: &mut Option<Instant>,
) -> Result<Self, ExecServerError> {
let (transport_params, deferred_readiness) = match transport_params {
ExecServerTransportParams::Deferred(deferred) => {
@@ -362,6 +363,7 @@ impl ExecServerClient {
None
};
*started = Some(Instant::now());
let websocket = match transport_params {
ExecServerTransportParams::Deferred(_) => {
return Err(ExecServerError::Protocol(
@@ -126,6 +126,7 @@ impl SequenceNoiseConnectProvider {
codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
/*started*/ &mut None,
)
.await
}
+34
View File
@@ -73,6 +73,14 @@ pub enum EnvironmentConnectionState {
Disconnected,
}
/// Outcome of an initial connection operation, including its internal retries.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConnectionAttemptOutcome {
Success,
Failure,
Cancelled,
}
/// Owns the execution/filesystem environments available to the Codex runtime.
///
/// `EnvironmentManager` is a shared registry for concrete environments. Its
@@ -907,6 +915,17 @@ impl Environment {
})
}
/// Observes initial connection attempts from the end of provisioning through client installation.
/// Register before connecting. The first observer wins; past attempts and reconnections are omitted.
pub fn observe_connection_attempts(
&self,
observer: impl Fn(std::time::Duration, ConnectionAttemptOutcome) + Send + Sync + 'static,
) {
if let Some(client) = &self.remote_client {
let _ = client.connection_observer.set(Arc::new(observer));
}
}
/// Subscribes to the current connection state for this remote environment.
pub fn subscribe_connection_state(
&self,
@@ -1850,6 +1869,16 @@ mod tests {
let environment = manager
.get_environment("executor-a")
.expect("first remote environment");
let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel();
environment.observe_connection_attempts(move |_, outcome| {
observed_tx
.send(outcome)
.expect("record cancelled connection");
});
let _connection = timeout(Duration::from_secs(1), first_listener.accept())
.await
.expect("connection should start")
.expect("accept initial connection");
let startup_abort = environment
.startup_task
.lock()
@@ -1880,6 +1909,11 @@ mod tests {
})
.await
.expect("replacing the environment should cancel its startup task");
assert_eq!(
observed_rx.try_recv(),
Ok(super::ConnectionAttemptOutcome::Cancelled)
);
assert!(observed_rx.try_recv().is_err());
}
#[tokio::test]
+1
View File
@@ -113,6 +113,7 @@ pub use environment::CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
pub use environment::ConnectionAttemptOutcome;
pub use environment::Environment;
pub use environment::EnvironmentConnectionState;
pub use environment::EnvironmentManager;