From e325e3acd9ab64cd287a2c4d6cd7a7cebb639618 Mon Sep 17 00:00:00 2001 From: hesham-oai Date: Thu, 27 Aug 2026 23:17:34 +0000 Subject: [PATCH] Retry confirmed remote registration conflicts (#41219) ## Why Remote environment registration can return an explicit `503` `registration_conflict` after the registry finishes its own recovery attempts. Retrying ambiguous failures is unsafe because a timed-out request may still have replaced a newer registration. ## What changed - Retry only explicit `registration_conflict` responses with jittered backoff, both during initial registration and after a rejected relay connection. - Apply the existing connection timeout to the complete registration attempt, including response-body diagnostics. - Keep retry waits and in-flight requests cancellable while preserving the Noise identity and initialized session across successful recovery. ## Testing Added unit coverage for retry eligibility and timeout handling, plus relay tests for identity and session preservation and prompt shutdown during registration. GitOrigin-RevId: 374e6fdac4755f377cbea9400c76a228b790945a --- codex-rs/exec-server/src/remote.rs | 53 ++++- .../src/remote/registration_retry.rs | 54 +++++ .../src/remote/registration_retry_tests.rs | 178 +++++++++++++++ codex-rs/exec-server/tests/relay.rs | 3 + .../tests/relay/registration_retry_tests.rs | 210 ++++++++++++++++++ 5 files changed, 487 insertions(+), 11 deletions(-) create mode 100644 codex-rs/exec-server/src/remote/registration_retry.rs create mode 100644 codex-rs/exec-server/src/remote/registration_retry_tests.rs create mode 100644 codex-rs/exec-server/tests/relay/registration_retry_tests.rs diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 8836e4065a..cb4282c04a 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -15,6 +15,7 @@ use http::HeaderValue; use http::StatusCode; use serde::Deserialize; use tokio::time::sleep; +use tokio::time::timeout_at; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; use tracing::info; @@ -52,6 +53,8 @@ use crate::trace_context::current_trace_context_headers; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; +mod registration_retry; + #[derive(Clone)] struct EnvironmentRegistryClient { base_url: String, @@ -132,6 +135,7 @@ impl EnvironmentRegistryClient { environment_id: &str, executor_public_key: &NoiseChannelPublicKey, ) -> Result { + let deadline = tokio::time::Instant::now() + self.connect_timeout; let url = endpoint_url( &self.base_url, &format!("/cloud/environment/{environment_id}/register"), @@ -140,16 +144,40 @@ impl EnvironmentRegistryClient { security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), executor_public_key: executor_public_key.clone(), }; - let response = self - .http - .post(url) - .headers(self.resolve_auth_headers().await?) - .headers(current_trace_context_headers()) - .json(&body) - .send() - .await?; + let response = timeout_at(deadline, async { + self.http + .post(url) + .headers(self.resolve_auth_headers().await?) + .headers(current_trace_context_headers()) + .json(&body) + .send() + .await + .map_err(ExecServerError::EnvironmentRegistryRequest) + }) + .await + .unwrap_or_else(|_| { + Err(ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + )) + })?; + let status = response.status(); + // Read diagnostics within the same attempt budget, preserving a known error status. let response: EnvironmentRegistryRegistrationResponse = - self.parse_json_response(response).await?; + timeout_at(deadline, self.parse_json_response(response)) + .await + .unwrap_or_else(|_| { + Err(match status { + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => { + environment_registry_auth_error(status, "response body timed out") + } + status if !status.is_success() => { + environment_registry_http_error(status, "response body timed out") + } + _ => ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Timeout, + ), + }) + })?; if response.environment_id != environment_id { return Err(ExecServerError::Protocol( "environment registry returned a different environment id".to_string(), @@ -616,7 +644,7 @@ async fn run_remote_environment_connections( })?; let mut backoff = Duration::from_secs(1); let mut response = client - .register_environment(&config.environment_id, &identity.public_key()) + .register_environment_with_retry(&config.environment_id, &identity.public_key()) .await?; loop { @@ -674,7 +702,10 @@ async fn run_remote_environment_connections( if registration_rejected { config.telemetry.remote_reconnect("registration_rejected"); response = client - .register_environment(&config.environment_id, &identity.public_key()) + .register_environment_with_retry( + &config.environment_id, + &identity.public_key(), + ) .await?; } else { config.telemetry.remote_reconnect("connect_failed"); diff --git a/codex-rs/exec-server/src/remote/registration_retry.rs b/codex-rs/exec-server/src/remote/registration_retry.rs new file mode 100644 index 0000000000..1a2dc989da --- /dev/null +++ b/codex-rs/exec-server/src/remote/registration_retry.rs @@ -0,0 +1,54 @@ +//! Retry only explicit registration conflicts after the registry's write-retry loop has finished. +//! Ambiguous failures must not be replayed: a timed-out request can still replace a newer registration. +//! The enclosing remote-transport future owns cancellation; retries spawn no background work. + +use http::StatusCode; +use tokio::time::sleep; +use tracing::warn; + +use super::EnvironmentRegistryClient; +use crate::EnvironmentRegistryRegistrationResponse; +use crate::ExecServerError; +use crate::NoiseChannelPublicKey; +use crate::client::registry_recovery_retry_delay; + +impl EnvironmentRegistryClient { + pub(super) async fn register_environment_with_retry( + &self, + environment_id: &str, + executor_public_key: &NoiseChannelPublicKey, + ) -> Result { + // Competing executors for the same environment must not retry in lockstep. + let retry_key = uuid::Uuid::new_v4().to_string(); + let mut attempt = 0_u32; + loop { + match self + .register_environment(environment_id, executor_public_key) + .await + { + Ok(response) => return Ok(response), + Err(ExecServerError::EnvironmentRegistryHttp { status, code, .. }) + if status == StatusCode::SERVICE_UNAVAILABLE + && code.as_deref() == Some("registration_conflict") => + { + let delay = registry_recovery_retry_delay(&retry_key, attempt); + attempt = attempt.saturating_add(1); + // Do not log response bodies or transport errors: they can contain credentials. + warn!( + noise_event = "registration", + noise_outcome = "retry", + retry_attempt = attempt, + retry_delay_ms = delay.as_millis() as u64, + "Noise executor retrying registry registration conflict" + ); + sleep(delay).await; + } + Err(error) => return Err(error), + } + } + } +} + +#[cfg(test)] +#[path = "registration_retry_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/remote/registration_retry_tests.rs b/codex-rs/exec-server/src/remote/registration_retry_tests.rs new file mode 100644 index 0000000000..2c280c4947 --- /dev/null +++ b/codex-rs/exec-server/src/remote/registration_retry_tests.rs @@ -0,0 +1,178 @@ +//! Only confirmed conflicts permit replay; ambiguous registration results remain terminal. + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_api::AuthProvider; +use codex_http_client::RouteAwareRequestError; +use http::HeaderMap; +use pretty_assertions::assert_eq; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::time::advance; +use tokio::time::sleep; +use tokio::time::timeout; +use tokio_util::task::AbortOnDropHandle; +use tracing::instrument::WithSubscriber; +use tracing_subscriber::prelude::*; + +use super::EnvironmentRegistryClient; +use super::ExecServerError; +use crate::NoiseChannelIdentity; + +const ENVIRONMENT_ID: &str = "registration-retry-test"; +const CONFLICT_BODY: &str = + r#"{"error":{"code":"registration_conflict","message":"registration unavailable"}}"#; +const ERROR_BODY: &str = + r#"{"error":{"code":"registration_denied","message":"registration unavailable"}}"#; +const SUCCESS_BODY: &str = r#"{"environment_id":"registration-retry-test","url":"ws://localhost/relay","security_profile":"noise_hybrid_ik_v1","executor_registration_id":"committed-registration"}"#; + +#[derive(Debug, Default)] +struct RegistrationAuthProvider { + calls: AtomicUsize, +} + +impl AuthProvider for RegistrationAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { + self.calls.fetch_add(1, Ordering::Relaxed); + Box::pin(async { Ok(HeaderMap::new()) }) + } +} + +struct RetryObserved(Arc); + +impl tracing_subscriber::Layer for RetryObserved { + fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) { + if event.metadata().target() == "codex_exec_server::remote::registration_retry" + && event.metadata().fields().field("retry_attempt").is_some() + { + self.0.notify_one(); + } + } +} + +#[test_case::test_case(200, "not JSON", Duration::ZERO; "malformed_success")] +#[test_case::test_case(200, SUCCESS_BODY, Duration::from_secs(60); "committed_success_body_timeout")] +#[test_case::test_case(503, CONFLICT_BODY, Duration::ZERO; "confirmed_conflict_backoff_is_cancellable")] +#[test_case::test_case(503, CONFLICT_BODY, Duration::from_secs(60); "conflict_code_not_received")] +#[test_case::test_case(503, "not JSON", Duration::ZERO; "malformed_unavailable")] +#[test_case::test_case(503, "{}", Duration::ZERO; "missing_conflict_code")] +#[test_case::test_case(503, ERROR_BODY, Duration::ZERO; "different_error_code")] +#[test_case::test_case(502, CONFLICT_BODY, Duration::ZERO; "gateway_error")] +#[test_case::test_case(408, CONFLICT_BODY, Duration::ZERO; "request_timeout")] +#[test_case::test_case(429, CONFLICT_BODY, Duration::ZERO; "too_many_requests")] +#[test_case::test_case(401, ERROR_BODY, Duration::from_secs(60); "unauthorized_stalled_body")] +#[test_case::test_case(403, ERROR_BODY, Duration::from_secs(60); "forbidden_stalled_body")] +#[test_case::test_case(404, ERROR_BODY, Duration::from_secs(60); "environment_deleted_stalled_body")] +#[test_case::test_case(401, ERROR_BODY, Duration::from_millis(50); "delayed_unauthorized_details")] +#[test_case::test_case(403, ERROR_BODY, Duration::from_millis(50); "delayed_forbidden_details")] +#[test_case::test_case(404, ERROR_BODY, Duration::from_millis(50); "delayed_error_details")] +#[tokio::test] +async fn registration_requires_a_confirmed_conflict_before_replay( + status: u16, + body: &'static str, + body_delay: Duration, +) -> anyhow::Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let registry_url = format!("http://{}", listener.local_addr()?); + let _server = AbortOnDropHandle::new(tokio::spawn(async move { + let (mut stream, _) = listener.accept().await?; + // Drain the complete request before responding so unread bytes cannot cause a reset. + let mut reader = BufReader::new(&mut stream); + let mut content_length = 0; + loop { + let mut line = String::new(); + anyhow::ensure!(reader.read_line(&mut line).await? > 0, "missing headers"); + if line == "\r\n" { + break; + } + if let Some((name, value)) = line.split_once(':') + && name.eq_ignore_ascii_case("content-length") + { + content_length = value.trim().parse::()?; + } + } + anyhow::ensure!(content_length <= 16_384, "registration request too large"); + reader.read_exact(&mut vec![0; content_length]).await?; + drop(reader); + + let headers = format!( + "HTTP/1.1 {status} Registration Result\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).await?; + sleep(body_delay).await; + stream.write_all(body.as_bytes()).await?; + std::future::pending::<()>().await; + anyhow::Ok(()) + })); + let auth = Arc::new(RegistrationAuthProvider::default()); + let mut client = EnvironmentRegistryClient::new(registry_url, auth.clone())?; + client.connect_timeout = Duration::from_millis(500); + let key = NoiseChannelIdentity::generate()?.public_key(); + let retry = Arc::new(Notify::new()); + let subscriber = tracing_subscriber::registry().with(RetryObserved(retry.clone())); + let mut registration = Box::pin( + client + .register_environment_with_retry(ENVIRONMENT_ID, &key) + .with_subscriber(subscriber), + ); + + if status == 503 && body == CONFLICT_BODY && body_delay.is_zero() { + timeout(Duration::from_secs(1), async { + tokio::select! { + _ = retry.notified() => anyhow::Ok(()), + _ = &mut registration => anyhow::bail!("a confirmed conflict must enter backoff"), + } + }) + .await??; + // The retry event fires in the same poll that enters backoff, without a timing guess. + drop(registration); + tokio::time::pause(); + advance(Duration::from_secs(60)).await; + tokio::task::yield_now().await; + assert_eq!(auth.calls.load(Ordering::Relaxed), 1); + return Ok(()); + } + + let error = timeout(Duration::from_secs(1), registration) + .await? + .expect_err("an unconfirmed registration outcome must not be replayed"); + assert_eq!(auth.calls.load(Ordering::Relaxed), 1); + + match error { + ExecServerError::Json(_) if status == 200 && body_delay.is_zero() => {} + ExecServerError::EnvironmentRegistryRequest(RouteAwareRequestError::Timeout) + if status == 200 && body_delay > client.connect_timeout => {} + ExecServerError::EnvironmentRegistryAuth(message) if matches!(status, 401 | 403) => { + if body_delay < client.connect_timeout { + assert!(message.ends_with(": registration unavailable")); + } + } + ExecServerError::EnvironmentRegistryHttp { + status: actual, + code, + message, + } if !matches!(status, 200 | 401 | 403) => { + let expected_code = match (body_delay < client.connect_timeout, body) { + (true, ERROR_BODY) => Some("registration_denied"), + (true, CONFLICT_BODY) => Some("registration_conflict"), + _ => None, + }; + assert_eq!((actual.as_u16(), code.as_deref()), (status, expected_code)); + if expected_code.is_some() { + assert_eq!(message, "registration unavailable"); + } + } + _ => anyhow::bail!("unexpected registration failure kind"), + } + Ok(()) +} diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index f7dc181ed6..e4e44c2a28 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -3,6 +3,9 @@ mod common; #[path = "common/relay.rs"] mod relay_support; +#[path = "relay/registration_retry_tests.rs"] +mod registration_retry; + use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; diff --git a/codex-rs/exec-server/tests/relay/registration_retry_tests.rs b/codex-rs/exec-server/tests/relay/registration_retry_tests.rs new file mode 100644 index 0000000000..99455a6cc7 --- /dev/null +++ b/codex-rs/exec-server/tests/relay/registration_retry_tests.rs @@ -0,0 +1,210 @@ +//! Confirmed registration conflicts retry without replacing the Noise identity or session handler. + +use std::time::Duration; + +use codex_exec_server::ExecServerClient; +use codex_exec_server::NoiseChannelIdentity; +use codex_exec_server::NoiseRendezvousConnectArgs; +use pretty_assertions::assert_eq; +use tokio::sync::oneshot; +use tokio_tungstenite::accept_hdr_async; +use tokio_tungstenite::tungstenite::handshake::server::Request; +use tokio_tungstenite::tungstenite::handshake::server::Response; + +use super::*; + +type RemoteTask = AbortOnDropHandle>; +type RelayTask = AbortOnDropHandle>; + +struct RegistryFixture { + registry: MockServer, + listener: TcpListener, + requests: mpsc::UnboundedReceiver, +} + +impl RegistryFixture { + async fn new(statuses: Vec, response_delay: Duration) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + let (request_tx, requests) = mpsc::unbounded_channel(); + let attempts = AtomicUsize::new(0); + Mock::given(method("POST")) + .and(path(format!("/cloud/environment/{ENVIRONMENT_ID}/register"))) + .respond_with(move |_: &wiremock::Request| { + let attempt = attempts.fetch_add(1, Ordering::SeqCst) + 1; + let _ = request_tx.send(attempt); + let status = statuses[(attempt - 1).min(statuses.len() - 1)]; + let mut response = ResponseTemplate::new(status).set_delay(response_delay); + if status == 200 { + response = response.set_body_json(serde_json::json!({ + "environment_id": ENVIRONMENT_ID, + "url": format!("{rendezvous_url}/relay?role=environment®istration={attempt}"), + "security_profile": "noise_hybrid_ik_v1", + "executor_registration_id": format!("registration-{attempt}"), + })); + } else { + response = response.set_body_json(serde_json::json!({ + "error": {"code": "registration_conflict", "message": "registration conflicted"}, + })); + } + response + }) + .mount(®istry) + .await; + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/validate" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "valid": true, + }))) + .mount(®istry) + .await; + Ok(Self { + registry, + listener, + requests, + }) + } + + fn start(&self) -> Result<(oneshot::Sender<()>, RemoteTask)> { + let config = RemoteEnvironmentConfig::new( + self.registry.uri(), + ENVIRONMENT_ID.to_string(), + static_registry_auth_provider(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + )?; + let (codex_exe, sandbox_exe) = common::current_test_binary_helper_paths()?; + let runtime_paths = ExecServerRuntimePaths::new(codex_exe, sandbox_exe)?; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let task = AbortOnDropHandle::new(tokio::spawn( + codex_exec_server::run_remote_environment_until_shutdown( + config, + runtime_paths, + async move { + let _ = shutdown_rx.await; + }, + ), + )); + Ok((shutdown_tx, task)) + } + + async fn connect( + &self, + registration_id: &str, + harness_identity: &NoiseChannelIdentity, + resume_session_id: Option, + ) -> Result<(ExecServerClient, RelayTask)> { + let environment = accept_websocket(&self.listener, "environment").await?; + let args = NoiseRendezvousConnectArgs { + bundle: NoiseRendezvousConnectBundle { + websocket_url: format!("ws://{}/relay?role=harness", self.listener.local_addr()?), + environment_id: ENVIRONMENT_ID.to_string(), + executor_registration_id: registration_id.to_string(), + executor_public_key: registered_executor_public_key(&self.registry).await?, + harness_key_authorization: HARNESS_KEY_AUTHORIZATION.to_string(), + }, + harness_identity: harness_identity.clone(), + client_name: "registration-retry-test".to_string(), + connect_timeout: TEST_TIMEOUT, + initialize_timeout: TEST_TIMEOUT, + resume_session_id, + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + }; + let client = AbortOnDropHandle::new(tokio::spawn(async move { + ExecServerClient::connect_noise_rendezvous(args).await + })); + let harness = accept_websocket(&self.listener, "harness").await?; + let relay = AbortOnDropHandle::new(tokio::spawn(proxy_relay_frames( + environment, + harness, + Arc::new(Mutex::new(Vec::new())), + ))); + Ok((timeout(TEST_TIMEOUT, client).await???, relay)) + } + + async fn registered_keys(&self) -> Result> { + self.registry + .received_requests() + .await + .context("registry should retain requests")? + .iter() + .filter(|request| request.url.path().ends_with("/register")) + .map(|request| { + let body: serde_json::Value = serde_json::from_slice(&request.body)?; + Ok(serde_json::from_value(body["executor_public_key"].clone())?) + }) + .collect() + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn registration_retries_preserve_noise_identity_and_initialized_session() -> Result<()> { + let fixture = RegistryFixture::new(vec![503, 200, 503, 503, 200], Duration::ZERO).await?; + let (shutdown, remote) = fixture.start()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let (client, first_relay) = fixture + .connect( + "registration-2", + &harness_identity, + /*resume_session_id*/ None, + ) + .await?; + let session_id = client.session_id().context("initialized session ID")?; + let environment_info = client.force_environment_info().await?; + let key = registered_executor_public_key(&fixture.registry).await?; + assert_eq!(fixture.registered_keys().await?, vec![key.clone(); 2]); + + first_relay.abort(); + let _ = first_relay.await; + drop(client); + let (socket, _) = timeout(TEST_TIMEOUT, fixture.listener.accept()).await??; + let rejected = http::Response::builder() + .status(http::StatusCode::UNAUTHORIZED) + .body(Some("expired registration".to_string()))?; + let stale_url = accept_hdr_async(socket, |request: &Request, _: Response| { + assert_eq!( + request + .uri() + .path_and_query() + .map(http::uri::PathAndQuery::as_str), + Some("/relay?role=environment®istration=2") + ); + Err(rejected) + }); + assert!(timeout(TEST_TIMEOUT, stale_url).await?.is_err()); + + let (resumed, second_relay) = fixture + .connect( + "registration-5", + &harness_identity, + Some(session_id.clone()), + ) + .await?; + assert_eq!(resumed.session_id(), Some(session_id)); + assert_eq!(resumed.force_environment_info().await?, environment_info); + assert_eq!(fixture.registered_keys().await?, vec![key; 5]); + + assert!(shutdown.send(()).is_ok(), "remote task exited"); + timeout(TEST_TIMEOUT, remote).await???; + drop(resumed); + second_relay.abort(); + let _ = second_relay.await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shutdown_interrupts_in_flight_registration() -> Result<()> { + let mut fixture = RegistryFixture::new(vec![503], Duration::from_secs(60)).await?; + let (shutdown, remote) = fixture.start()?; + assert_eq!( + timeout(TEST_TIMEOUT, fixture.requests.recv()).await?, + Some(1) + ); + assert!(shutdown.send(()).is_ok(), "remote task exited"); + timeout(Duration::from_secs(1), remote) + .await + .context("shutdown must not wait for the registration request timeout")???; + Ok(()) +}