diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 696822f845..81fc3110b2 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -3,6 +3,7 @@ use crate::config::ConstraintResult; use crate::context::ContextualUserFragment; use crate::context::GuardianReviewEvidence; use crate::elicitation::ElicitationRegistration; +use crate::environment_selection::TurnEnvironmentState; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::new_submission_id; @@ -517,6 +518,28 @@ impl CodexThread { self.session.inject_if_running(items).await } + /// Environment selections captured by the active turn, before later settings updates. + /// Includes environments that are still starting or have failed. Hosts use this snapshot + /// to authorize steering against every executor that the active turn selected. + pub async fn active_turn_environment_selections( + &self, + ) -> Option> { + let active = self.session.active_turn.lock().await; + let task = active.as_ref()?.task.as_ref()?; + Some( + task.turn_context + .environments + .environments + .iter() + .map(|environment| match environment { + TurnEnvironmentState::Ready(environment) => environment.selection(), + TurnEnvironmentState::Starting(environment) => environment.selection.clone(), + TurnEnvironmentState::Failed { selection, .. } => selection.clone(), + }) + .collect(), + ) + } + /// Captures a regular turn only after its input is recorded. The caller must flush the rollout. pub async fn interrupted_turn( &self, diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 9bf677af28..f1e04c4dbe 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -907,6 +907,7 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( let test = builder.build(&server).await?; let initial_cwd = test.config.cwd.clone(); let initial_environments = test.codex.environment_selections().await; + assert_eq!(test.codex.active_turn_environment_selections().await, None); let next_workspace = TempDir::new()?; let next_cwd = next_workspace.path().abs(); let next_environments = @@ -924,6 +925,11 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( }) .await; + assert_eq!( + test.codex.active_turn_environment_selections().await, + Some(initial_environments.clone()) + ); + let preview = test .codex .preview_thread_settings_overrides(CodexThreadSettingsOverrides { @@ -954,6 +960,10 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( test.codex.environment_selections().await, next_environments.environments ); + assert_eq!( + test.codex.active_turn_environment_selections().await, + Some(initial_environments) + ); let snapshot = test.codex.config_snapshot().await; assert_eq!( snapshot.environment_selections(), @@ -978,6 +988,7 @@ async fn settings_update_does_not_retarget_active_turn_environment() -> Result<( matches!(event, EventMsg::TurnComplete(_)) }) .await; + assert_eq!(test.codex.active_turn_environment_selections().await, None); test.submit_turn("start the next turn").await?; let request_texts = response_mock @@ -1105,6 +1116,19 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() }) .await; + let active_environments = test + .codex + .active_turn_environment_selections() + .await + .context("active turn environments")?; + assert_eq!( + active_environments + .iter() + .map(|selection| selection.environment_id.as_str()) + .collect::>(), + vec![REMOTE_ENVIRONMENT_ID, "local"] + ); + let requests = response_mock.requests(); let initial_context = requests[1] .message_input_texts("user") @@ -1164,6 +1188,10 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() } }); core_test_support::wait_for_mcp_server(&test.codex, "deferred").await?; + assert_eq!( + test.codex.active_turn_environment_selections().await, + Some(active_environments) + ); test.codex .submit(Op::UserInputAnswer { id: request.turn_id, @@ -1182,6 +1210,8 @@ async fn deferred_executor_promotes_primary_environment_when_startup_completes() }) .await; + assert_eq!(test.codex.active_turn_environment_selections().await, None); + let requests = response_mock.requests(); assert!( requests[2] diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 99d0255493..eced178795 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -291,6 +291,8 @@ struct Inner { struct ConnectionState { status: ConnectionStatus, + // Publish registration renewal together with the recovered transport. + executor_registration_id: Option, active_process_starts: usize, environment_connection_state_tx: watch::Sender, } @@ -347,6 +349,7 @@ pub struct ExecServerClient { /// `timeout_for_error` keeps diagnostics tied to the caller's configured /// budget after readiness has consumed part of that budget. pub(crate) struct NoiseInitializeContext { + pub(crate) executor_registration_id: String, pub(crate) span: tracing::Span, pub(crate) timeout_for_error: Duration, } @@ -598,6 +601,11 @@ impl LazyRemoteExecServerClient { .filter(|client| !client.is_disconnected()) } + pub(crate) fn cached_executor_registration_id(&self) -> Option { + self.cached_client() + .and_then(|client| client.executor_registration_id()) + } + fn cached_client(&self) -> Option { self.current_client .lock() @@ -693,6 +701,15 @@ pub enum ExecServerError { } impl ExecServerClient { + fn executor_registration_id(&self) -> Option { + self.inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .executor_registration_id + .clone() + } + fn attach_environment_connection_state( &self, state_tx: watch::Sender, @@ -1216,6 +1233,9 @@ impl ExecServerClient { let inner = Arc::new(Inner { connection: StdMutex::new(ConnectionState { status: ConnectionStatus::Connected(Arc::clone(&rpc_client)), + executor_registration_id: noise_context + .as_ref() + .map(|context| context.executor_registration_id.clone()), active_process_starts: 0, environment_connection_state_tx: watch::channel( EnvironmentConnectionState::Connected, diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 6b4cf61df1..f3363fd3ce 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -394,8 +394,8 @@ impl Inner { let mut registry_retry_attempt = 0; let last_error = loop { match timeout_at(deadline, self.resume_once(&session_id)).await { - Ok(Ok((rpc_client, _attempt))) => { - if !rpc_client.is_disconnected() && self.install_recovered_client(rpc_client) { + Ok(Ok((rpc_client, _attempt, executor_registration_id))) => { + if self.install_recovered_client(rpc_client, executor_registration_id) { return; } } @@ -444,17 +444,23 @@ impl Inner { } } - fn install_recovered_client(&self, rpc_client: Arc) -> bool { + fn install_recovered_client( + &self, + rpc_client: Arc, + executor_registration_id: Option, + ) -> bool { let installed = { let mut connection = self .connection .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if !matches!(connection.status, ConnectionStatus::Recovering) + if self.retired.is_cancelled() + || !matches!(connection.status, ConnectionStatus::Recovering) || rpc_client.is_disconnected() { false } else { + connection.executor_registration_id = executor_registration_id; connection.set_status(ConnectionStatus::Connected(rpc_client)); true } @@ -472,13 +478,23 @@ impl Inner { async fn resume_once( self: &Arc, session_id: &str, - ) -> Result<(Arc, Option), ExecServerError> { + ) -> Result< + ( + Arc, + Option, + Option, + ), + ExecServerError, + > { let reconnect_strategy = self .reconnect_strategy .as_ref() .ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?; let attempt = reconnect_strategy.resume(session_id).await?; let (connection, options, attempt_permit, noise_context) = attempt.into_parts(); + let executor_registration_id = noise_context + .as_ref() + .map(|context| context.executor_registration_id.clone()); let (rpc_client, events_rx) = RpcClient::new(connection); let rpc_client = Arc::new(rpc_client); let client = ExecServerClient { @@ -495,7 +511,7 @@ impl Inner { .await?; self.recover_processes(&rpc_client).await?; - Ok((rpc_client, attempt_permit)) + Ok((rpc_client, attempt_permit, executor_registration_id)) } async fn recover_processes( diff --git a/codex-rs/exec-server/src/client_refresh.rs b/codex-rs/exec-server/src/client_refresh.rs index 9918227b71..e9e31dfcfc 100644 --- a/codex-rs/exec-server/src/client_refresh.rs +++ b/codex-rs/exec-server/src/client_refresh.rs @@ -85,6 +85,9 @@ impl LazyRemoteExecServerClient { let _refresh = self.refresh_lock.lock().await; let (previous, attempt) = loop { let observed = self.cached_client(); + let observed_registration = observed + .as_ref() + .and_then(ExecServerClient::executor_registration_id); let mut transport = self.transport_params.clone().ok_or_else(|| { ExecServerError::Protocol( "connection refresh requires a Noise registry".to_string(), @@ -111,6 +114,7 @@ impl LazyRemoteExecServerClient { ) })??; let executor_public_key = bundle.executor_public_key.clone(); + let executor_registration_id = bundle.executor_registration_id.clone(); *provider = Arc::new(PrefetchedConnectProvider { bundle: StdMutex::new(Some(bundle)), provider: Arc::clone(provider), @@ -124,8 +128,8 @@ impl LazyRemoteExecServerClient { .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. + // A new client makes the lookup stale. Recovery can also renew the + // registration on the same client, so compare it under the connection lock. if !match (&observed, &*current) { (Some(observed), Some(current)) => Arc::ptr_eq(&observed.inner, ¤t.inner), (None, None) => true, @@ -133,18 +137,31 @@ impl LazyRemoteExecServerClient { } { 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() + if let Some(client) = current.as_ref() { + let connection = client + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if connection.executor_registration_id != observed_registration { + continue; + } + // Preserve same-executor recovery; the final probe fails fast if still recovering. + if !matches!(connection.status, ConnectionStatus::Failed(_)) + && connection.executor_registration_id.as_ref() + == Some(&executor_registration_id) && matches!( client.inner.reconnect_strategy.as_ref(), Some(ExecServerReconnectStrategy::NoiseRendezvous { executor_public_key: key, .. }) if key == &executor_public_key ) - }) { - break (current.clone(), None); + { + break (current.clone(), None); + } + // Claim retirement before recovery can install another registration. + // Transport shutdown and pending-work cleanup happen outside these locks. + client.inner.retired.cancel(); } // Cancellation and connection installation use the same lock. A late // handshake cannot install a client after its attempt has been superseded. diff --git a/codex-rs/exec-server/src/client_refresh_tests.rs b/codex-rs/exec-server/src/client_refresh_tests.rs index 4e45a85ec5..4beee04dee 100644 --- a/codex-rs/exec-server/src/client_refresh_tests.rs +++ b/codex-rs/exec-server/src/client_refresh_tests.rs @@ -3,6 +3,7 @@ use std::sync::Mutex; use std::time::Duration; use anyhow::Result; +use bytes::Bytes; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use tokio::net::TcpListener; @@ -15,6 +16,9 @@ use super::*; use crate::ExecServerRuntimePaths; use crate::NoiseChannelIdentity; use crate::ProcessId; +use crate::noise_relay::stream_handler::NoiseOutboundMessage; +use crate::noise_relay::stream_handler::NoiseStreamConnection; +use crate::noise_relay::stream_handler::NoiseStreamHandler; use crate::relay::HarnessKeyValidator; use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; @@ -100,24 +104,74 @@ impl HarnessKeyValidator for Validator { } } +#[derive(Clone)] +struct NotifyingProcessor { + processor: ConnectionProcessor, + detached: Arc, +} + +impl NoiseStreamHandler for NotifyingProcessor { + type Incoming = ::Incoming; + type Outgoing = ::Outgoing; + + fn decode(payload: Bytes) -> Result { + ConnectionProcessor::decode(payload) + } + + fn encode(message: Self::Outgoing) -> Result { + ConnectionProcessor::encode(message) + } + + async fn run_connection( + self, + connection: NoiseStreamConnection, + ) { + NoiseStreamHandler::run_connection(self.processor, connection).await; + // Connection cleanup has detached the session before recovery can resume it. + self.detached.notify_one(); + } +} + struct Executor { target: Target, + processor: ConnectionProcessor, + detached: Arc, _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(); + Self::start_with_identity(validator, NoiseChannelIdentity::generate()?).await + } + + async fn start_with_identity( + validator: Validator, + identity: NoiseChannelIdentity, + ) -> Result { let processor = ConnectionProcessor::new(ExecServerRuntimePaths::new( std::env::current_exe()?, /*codex_linux_sandbox_exe*/ None, )?); + Self::start_with_processor(validator, identity, processor).await + } + + async fn start_with_processor( + validator: Validator, + identity: NoiseChannelIdentity, + processor: ConnectionProcessor, + ) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let target = Target { + url: format!("ws://{}", listener.local_addr()?), + identity, + registration: uuid::Uuid::new_v4().to_string(), + }; + let executor = target.clone(); + let detached = Arc::new(Notify::new()); + let server_processor = NotifyingProcessor { + processor: processor.clone(), + detached: Arc::clone(&detached), + }; let server = tokio::spawn(async move { let mut connections = JoinSet::new(); loop { @@ -125,7 +179,7 @@ impl Executor { socket = listener.accept() => { let (socket, _) = socket.unwrap(); let executor = executor.clone(); - let processor = processor.clone(); + let processor = server_processor.clone(); let validator = validator.clone(); connections.spawn(async move { let socket = tokio_tungstenite::accept_async(socket).await.unwrap(); @@ -138,6 +192,8 @@ impl Executor { }); Ok(Self { target, + processor, + detached, _server: AbortOnDropHandle::new(server), }) } @@ -225,12 +281,11 @@ async fn refresh_retires_a_still_connected_old_executor() -> Result<()> { } #[tokio::test] -async fn refresh_preserves_a_current_session_across_registration_renewal() -> Result<()> { +async fn refresh_preserves_a_current_registration() -> Result<()> { let _clock = freeze_clock(); let executor = Executor::start(Validator::default()).await?; - let (client, registry) = executor.client()?; + 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?; @@ -240,6 +295,213 @@ async fn refresh_preserves_a_current_session_across_registration_renewal() -> Re Ok(()) } +#[tokio::test] +async fn registration_snapshot_changes_only_after_same_key_replacement_is_installed() -> Result<()> +{ + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let validator = Validator { + handshake: Some(Arc::new(Notify::new())), + ..Default::default() + }; + let new = Executor::start_with_identity(validator.clone(), old.target.identity.clone()).await?; + let (client, registry) = old.client()?; + let environment = + crate::Environment::remote_with_client(client.clone(), /*local_runtime_paths*/ None); + assert_eq!(environment.cached_executor_registration_id(), None); + let original = client.get().await?; + assert_eq!( + environment.cached_executor_registration_id(), + Some(old.target.registration.clone()) + ); + *registry.target.lock().unwrap() = new.target.clone(); + let refreshing = environment.clone(); + let refresh = tokio::spawn(async move { refreshing.refresh_connection().await }); + validator.started.notified().await; + assert_eq!( + environment.cached_executor_registration_id(), + Some(old.target.registration.clone()) + ); + validator.handshake.as_ref().unwrap().notify_one(); + refresh.await??; + assert_eq!( + environment.cached_executor_registration_id(), + Some(new.target.registration.clone()) + ); + assert_ne!(original.session_id(), client.get().await?.session_id()); + assert!(original.is_disconnected()); + Ok(()) +} + +#[cfg(unix)] +#[tokio::test] +async fn registration_renewal_recovers_the_session_and_running_process() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let validator = Validator { + handshake: Some(Arc::new(Notify::new())), + ..Default::default() + }; + let renewed = Executor::start_with_processor( + validator.clone(), + old.target.identity.clone(), + old.processor.clone(), + ) + .await?; + let (client, registry) = old.client()?; + let environment = + crate::Environment::remote_with_client(client.clone(), /*local_runtime_paths*/ None); + let original = client.get().await?; + let session_id = original.session_id(); + let process = + original + .start_process( + crate::protocol::ExecParams { + metadata: Default::default(), + process_id: ProcessId::from("renewed-process"), + argv: vec![ + "/bin/sh".to_owned(), + "-c".to_owned(), + "IFS= read line; printf 'from-stdin:%s\\n' \"$line\"".to_owned(), + ], + cwd: codex_utils_path_uri::PathUri::from_host_native_path( + std::env::current_dir()?, + )?, + shell_snapshot: None, + env_policy: None, + env: Default::default(), + tty: false, + pipe_stdin: true, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }, + /*network_policy_decider*/ None, + ) + .await?; + let mut events = process.subscribe_events(); + *registry.target.lock().unwrap() = renewed.target.clone(); + disconnect(&original).await; + validator.started.notified().await; + assert_eq!( + environment.cached_executor_registration_id(), + Some(old.target.registration.clone()) + ); + old.detached.notified().await; + validator.handshake.as_ref().unwrap().notify_one(); + original.inner.rpc_client().await?; + assert_eq!( + environment.cached_executor_registration_id(), + Some(renewed.target.registration.clone()) + ); + let recovered = client.get().await?; + assert!(Arc::ptr_eq(&original.inner, &recovered.inner)); + assert_eq!(recovered.session_id(), session_id); + assert_eq!( + process.write(b"after-renewal\n".to_vec()).await?.status, + crate::protocol::WriteStatus::Accepted + ); + let mut output = Vec::new(); + let mut exit_code = None; + loop { + match events.recv().await? { + crate::process::ExecProcessEvent::Output(chunk) => { + output.extend_from_slice(&chunk.chunk.into_inner()); + } + crate::process::ExecProcessEvent::Exited { + exit_code: code, .. + } => { + exit_code = Some(code); + } + crate::process::ExecProcessEvent::Closed { .. } => break, + crate::process::ExecProcessEvent::Failed(message) => { + anyhow::bail!("process failed after registration renewal: {message}"); + } + } + } + assert_eq!( + (String::from_utf8(output)?, exit_code), + ("from-stdin:after-renewal\n".to_owned(), Some(0)) + ); + Ok(()) +} + +#[tokio::test] +async fn same_key_replacement_cannot_resume_a_missing_session() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let replacement = + Executor::start_with_identity(Validator::default(), old.target.identity.clone()).await?; + let (client, registry) = old.client()?; + let environment = + crate::Environment::remote_with_client(client.clone(), /*local_runtime_paths*/ None); + let original = client.get().await?; + let process = original + .register_session(&ProcessId::from("old-process")) + .await?; + *registry.target.lock().unwrap() = replacement.target.clone(); + disconnect(&original).await; + assert!(matches!( + original.inner.rpc_client().await, + Err(ExecServerError::Disconnected(_)) + )); + assert_eq!( + environment.cached_executor_registration_id(), + Some(old.target.registration.clone()) + ); + assert!(matches!( + process.write(b"never replay".to_vec()).await, + Err(ExecServerError::Disconnected(_)) + )); + assert!(matches!( + original.environment_status().await, + Err(ExecServerError::Disconnected(_)) + )); + Ok(()) +} + +#[tokio::test] +async fn stale_refresh_lookup_does_not_retire_a_renewed_session() -> Result<()> { + let _clock = freeze_clock(); + let old = Executor::start(Validator::default()).await?; + let renewed = Executor::start_with_processor( + Validator::default(), + old.target.identity.clone(), + old.processor.clone(), + ) + .await?; + let (client, registry) = old.client()?; + let environment = + crate::Environment::remote_with_client(client.clone(), /*local_runtime_paths*/ None); + 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() = renewed.target.clone(); + let recovery = registry.block_next_lookup(); + disconnect(&original).await; + old.detached.notified().await; + recovery.send(()).unwrap(); + original.inner.rpc_client().await?; + assert_eq!( + environment.cached_executor_registration_id(), + Some(renewed.target.registration.clone()) + ); + release.send(()).unwrap(); + refreshing.await?; + assert!(Arc::ptr_eq(&original.inner, &client.get().await?.inner)); + assert!(!original.inner.retired.is_cancelled()); + assert_eq!( + environment.cached_executor_registration_id(), + Some(renewed.target.registration.clone()) + ); + original.environment_status().await?; + Ok(()) +} + #[tokio::test] async fn refresh_cancels_a_stalled_initial_lookup() -> Result<()> { let _clock = freeze_clock(); diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index bbdf769038..0ec478adcf 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -174,6 +174,7 @@ pub(crate) struct ReconnectAttempt { } struct OpenNoiseRendezvousConnection { + executor_registration_id: String, connection: JsonRpcConnection, options: ExecServerClientConnectOptions, handshake_ready: tokio::sync::oneshot::Receiver<()>, @@ -254,7 +255,7 @@ pub(crate) enum ExecServerReconnectStrategy { http_headers: HeaderMap, }, NoiseRendezvous { - // The executor that created the session, not the latest recovery lookup. + // Registration can renew; recovery still pins the executor's Noise identity. executor_public_key: crate::NoiseChannelPublicKey, provider: Arc, identity: NoiseChannelIdentity, @@ -280,7 +281,7 @@ impl ExecServerReconnectStrategy { Ok(ReconnectAttempt::new(connection, args.into())) } Self::NoiseRendezvous { - executor_public_key: _, + executor_public_key, provider, identity, client_name, @@ -289,6 +290,13 @@ impl ExecServerReconnectStrategy { http_client_factory, } => { let bundle = provider.connect_bundle(identity.public_key()).await?; + if &bundle.executor_public_key != executor_public_key { + return Err(ExecServerError::Protocol( + "executor key changed during session recovery".to_string(), + )); + } + // An expired rendezvous URL can make the same executor register again. + // Initialization must still resume the original session on that executor. let opened = ExecServerClient::open_noise_rendezvous_connection( NoiseRendezvousConnectArgs { bundle, @@ -660,13 +668,14 @@ impl ExecServerClient { NoiseHarnessConnectionArgs { connection_label, environment_id, - executor_registration_id, + executor_registration_id: executor_registration_id.clone(), identity: harness_identity, responder_public_key: executor_public_key, harness_key_authorization, }, ); Ok(OpenNoiseRendezvousConnection { + executor_registration_id, connection: connection.connection, options: ExecServerClientConnectOptions { client_name, @@ -710,6 +719,7 @@ impl ExecServerClient { // startup parent while making its two child operations visible. let initialize_timeout = connection.options.initialize_timeout; let noise_context = NoiseInitializeContext { + executor_registration_id: connection.executor_registration_id, span: tracing::info_span!( "codex.exec_server.request", otel.kind = "client", diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 63345c227a..6f66c13d32 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -239,6 +239,7 @@ async fn deferred_initialize_timeout_reports_configured_budget() { }, /*reconnect_strategy*/ None, NoiseInitializeContext { + executor_registration_id: "registration".to_string(), span: tracing::info_span!("codex.exec_server.request"), timeout_for_error: configured_timeout, }, @@ -454,6 +455,29 @@ async fn initial_noise_connection_does_not_retry_permanent_registry_errors() -> Ok(()) } +#[tokio::test] +async fn noise_session_resume_rejects_a_changed_noise_key_before_connecting() -> Result<()> { + let sequence = Arc::new(SequenceNoiseConnectProvider::default()); + let bundle = test_bundle("ws://must-not-connect.invalid".to_string())?; + let executor_public_key = NoiseChannelIdentity::generate()?.public_key(); + sequence.push_response(futures::future::ready(Ok(bundle))); + let strategy = ExecServerReconnectStrategy::NoiseRendezvous { + executor_public_key, + provider: sequence, + identity: NoiseChannelIdentity::generate()?, + client_name: "test".to_string(), + connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + }; + assert!( + matches!(strategy.resume("session").await, Err(ExecServerError::Protocol(message)) if message == "executor key changed during session recovery") + ); + Ok(()) +} + #[tokio::test(start_paused = true)] async fn noise_session_resume_leaves_offline_retries_to_recovery() -> Result<()> { let sequence = Arc::new(SequenceNoiseConnectProvider::default()); diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index de77ff3efc..c2cd7384ac 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -908,6 +908,19 @@ impl Environment { } } + /// Registration on the installed Noise session, without connecting or refreshing it. + /// + /// Returns `None` before a session is installed and for non-Noise environments. + /// A registry lookup does not change this value until the new connection is + /// installed. Recovery may renew the registration while preserving the session. + /// Callers that authorize asynchronously must check the snapshot again before + /// admitting work. + pub fn cached_executor_registration_id(&self) -> Option { + self.remote_client + .as_ref() + .and_then(LazyRemoteExecServerClient::cached_executor_registration_id) + } + /// Refresh the connection to the executor currently registered for this environment. /// /// # Caller contract @@ -920,11 +933,11 @@ impl Environment { /// # 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 + /// A changed registration or 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 + /// A matching registration and executor key preserve a session that has not failed, including one /// that is recovering; the live readiness check rejects a recovering connection. /// /// # Completion and errors