openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Track executor registrations across connection refresh and recovery (#46555)

Browse Source 
## Why

An executor can renew its registration while retaining its Noise identity. Connection refresh must distinguish registrations even when their keys match, and a stale registry lookup must not retire a session that recovery has already renewed.

## What changed

- Track the installed session's registration and expose it through `Environment::cached_executor_registration_id()` without connecting. Publish registration changes only when the new connection is installed.
- Require both registration and executor key to match before refresh reuses a session. Retry stale lookups and prevent recovery from installing a connection after retirement.
- Pin session recovery to the original Noise key while allowing registration renewal to preserve the session and running processes.
- Expose `CodexThread::active_turn_environment_selections()` so hosts can authorize steering against the active turn's selections, including starting and failed environments, independently of later settings updates.

## Testing

Add regression coverage for registration replacement, renewal with a running process, stale refresh lookups, missing sessions, and rejection of changed Noise keys. Extend active-turn tests to check selection snapshots across settings updates, deferred startup, and turn completion.

GitOrigin-RevId: 4395bc6c57f137691887e6585f52dcc5b958101c
This commit is contained in:
viyatb-oai
2026-09-19 01:27:02 +00:00
committed by copyberry
parent a633ebc124
commit 328feb0c29
9 changed files with 445 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex_thread.rs
+23
View File
@@ -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<Vec<TurnEnvironmentSelection>> {
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,
codex-rs/core/tests/suite/remote_env.rs
+30
View File
@@ -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<_>>(),
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]
codex-rs/exec-server/src/client.rs
+20
View File
@@ -291,6 +291,8 @@ struct Inner {
struct ConnectionState {
status: ConnectionStatus,
// Publish registration renewal together with the recovered transport.
executor_registration_id: Option<String>,
active_process_starts: usize,
environment_connection_state_tx: watch::Sender<EnvironmentConnectionState>,
}
@@ -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<String> {
self.cached_client()
.and_then(|client| client.executor_registration_id())
}
fn cached_client(&self) -> Option<ExecServerClient> {
self.current_client
.lock()
@@ -693,6 +701,15 @@ pub enum ExecServerError {
}
impl ExecServerClient {
fn executor_registration_id(&self) -> Option<String> {
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<EnvironmentConnectionState>,
@@ -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,
codex-rs/exec-server/src/client_recovery.rs
+22 -6
View File
@@ -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<RpcClient>) -> bool {
fn install_recovered_client(
&self,
rpc_client: Arc<RpcClient>,
executor_registration_id: Option<String>,
) -> 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<Self>,
session_id: &str,
) -> Result<(Arc<RpcClient>, Option<tokio::sync::OwnedSemaphorePermit>), ExecServerError> {
) -> Result<
(
Arc<RpcClient>,
Option<tokio::sync::OwnedSemaphorePermit>,
Option<String>,
),
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(
codex-rs/exec-server/src/client_refresh.rs
+25 -8
View File
@@ -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, &current.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.
codex-rs/exec-server/src/client_refresh_tests.rs
+273 -11
View File
@@ -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<Notify>,
}
impl NoiseStreamHandler for NotifyingProcessor {
type Incoming = <ConnectionProcessor as NoiseStreamHandler>::Incoming;
type Outgoing = <ConnectionProcessor as NoiseStreamHandler>::Outgoing;
fn decode(payload: Bytes) -> Result<Self::Incoming, ExecServerError> {
ConnectionProcessor::decode(payload)
}
fn encode(message: Self::Outgoing) -> Result<NoiseOutboundMessage, ExecServerError> {
ConnectionProcessor::encode(message)
}
async fn run_connection(
self,
connection: NoiseStreamConnection<Self::Incoming, Self::Outgoing>,
) {
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<Notify>,
_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();
Self::start_with_identity(validator, NoiseChannelIdentity::generate()?).await
}
async fn start_with_identity(
validator: Validator,
identity: NoiseChannelIdentity,
) -> Result<Self> {
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<Self> {
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();
codex-rs/exec-server/src/client_transport.rs
+13 -3
View File
@@ -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<dyn NoiseRendezvousConnectProvider>,
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",
codex-rs/exec-server/src/client_transport_tests.rs
+24
View File
@@ -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());
codex-rs/exec-server/src/environment.rs
+15 -2
View File
@@ -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<String> {
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