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

@@ -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<EnvironmentRegistryRegistrationResponse, ExecServerError> {
+        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<H: NoiseStreamHandler>(
     })?;
     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<H: NoiseStreamHandler>(
                 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");

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<EnvironmentRegistryRegistrationResponse, ExecServerError> {
+        // 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;

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<Notify>);
+
+impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> 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::<usize>()?;
+            }
+        }
+        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(())
+}

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;

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<Result<(), ExecServerError>>;
+type RelayTask = AbortOnDropHandle<Result<()>>;
+
+struct RegistryFixture {
+    registry: MockServer,
+    listener: TcpListener,
+    requests: mpsc::UnboundedReceiver<usize>,
+}
+
+impl RegistryFixture {
+    async fn new(statuses: Vec<u16>, response_delay: Duration) -> Result<Self> {
+        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&registration={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(&registry)
+            .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(&registry)
+            .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<String>,
+    ) -> 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<Vec<NoiseChannelPublicKey>> {
+        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&registration=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(())
+}