diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 2b874bc11f..368662ff1d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -390,6 +390,8 @@ impl Drop for PendingProcessStartSession { } type ConnectionResult = Result>; +pub(super) type ConnectionObserver = + Arc; #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { @@ -402,6 +404,7 @@ pub(crate) struct LazyRemoteExecServerClient { current_client: Arc>>, reconnect: Arc>>>, refresh_lock: Arc>, + pub(super) connection_observer: Arc>, environment_connection_state_tx: watch::Sender, } @@ -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::>(), + [ + crate::ConnectionAttemptOutcome::Failure, + crate::ConnectionAttemptOutcome::Success + ], + ); + drop(first); drop(second); drop(client); diff --git a/codex-rs/exec-server/src/client/accepted.rs b/codex-rs/exec-server/src/client/accepted.rs index ae6d5f84e7..61f12219fe 100644 --- a/codex-rs/exec-server/src/client/accepted.rs +++ b/codex-rs/exec-server/src/client/accepted.rs @@ -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(), } } diff --git a/codex-rs/exec-server/src/client_refresh.rs b/codex-rs/exec-server/src/client_refresh.rs index 3579cb4f72..43cae92820 100644 --- a/codex-rs/exec-server/src/client_refresh.rs +++ b/codex-rs/exec-server/src/client_refresh.rs @@ -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, } +struct ConnectionObservation<'a> { + started: Option, + 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); } } diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index beb47a1c37..85a56ca2cc 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -331,6 +331,7 @@ impl ExecServerClient { pub(crate) async fn connect_for_transport( transport_params: ExecServerTransportParams, http_client_factory: HttpClientFactory, + started: &mut Option, ) -> Result { 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( diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index fe1966f8fc..8858aed896 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -126,6 +126,7 @@ impl SequenceNoiseConnectProvider { codex_http_client::HttpClientFactory::new( codex_http_client::OutboundProxyPolicy::ReqwestDefault, ), + /*started*/ &mut None, ) .await } diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 9ffbb0b838..277f218069 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -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] diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index eaf3cd60d6..adca2b271c 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -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;