From 4e2dfd22306985f846115758bbf7f9a59351bb29 Mon Sep 17 00:00:00 2001 From: pansul47-oai Date: Thu, 29 Oct 2026 00:32:43 +0000 Subject: [PATCH] Retry transient environment startup failures on the next turn (#57866) ## Why A transient environment startup failure persisted across turns when the selected environment and workspace were unchanged, preventing commands from recovering after the executor became available. ## What changed Retry failed environment resolution at the start of a new turn when the error is retryable and the workspace remains selected. Reuse healthy or pending resolutions and preserve provisioning failures until the selection is updated. ## Testing Add coverage for command execution recovering on an unchanged next turn after endpoint availability is restored. Extend provisioning failure tests to verify that subsequent turns receive the original error without restarting resolution. GitOrigin-RevId: b80af60990f30d277804160c9c39c286024c7667 --- codex-rs/core/src/environment_selection.rs | 71 ++++++- codex-rs/core/src/session/environment.rs | 8 +- .../tests/suite/remote_env_failure_tests.rs | 187 +++++++++++++++++- 3 files changed, 252 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 3700d594ec..ce7bf9c8f8 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -308,6 +308,39 @@ struct ThreadEnvironmentsState { environments: Vec, } +#[derive(Clone, Copy)] +enum StartupRetry<'a> { + AfterSelectionUpdate, + OnNewTurn(&'a [TurnEnvironmentSelection]), +} + +impl StartupRetry<'_> { + fn should_restart(self, environment: &SelectedTurnEnvironment) -> bool { + if matches!( + environment.selection.config, + EnvironmentConfigState::Failed(_) + ) { + return matches!(self, Self::AfterSelectionUpdate); + } + // Avoid cloning a resolved environment on every healthy turn. + if environment.resolution.peek().is_some_and(Result::is_ok) { + return false; + } + match environment.resolution.clone().now_or_never() { + Some(Err(error)) => match self { + Self::AfterSelectionUpdate => true, + Self::OnNewTurn(selected) => { + error.is_retryable_preparation_error() + && selected + .iter() + .any(|s| s.has_same_workspace(&environment.selection)) + } + }, + Some(Ok(_)) | None => false, + } + } +} + impl ThreadEnvironments { pub(crate) fn new( environment_manager: Arc, @@ -449,6 +482,27 @@ impl ThreadEnvironments { /// Updates the selected list before waking work that was waiting on the previous selection. pub(crate) fn update_selections(&self, environments: &[TurnEnvironmentSelection]) { + self.apply_selections(environments, StartupRetry::AfterSelectionUpdate); + } + + /// A new turn can retry a failed connection even when it keeps the same workspace. + pub(crate) fn activate_selections(&self, environments: &[TurnEnvironmentSelection]) { + let retry = StartupRetry::OnNewTurn(environments); + if environments != self.selections() { + self.update_selections(environments); + } else if self + .state + .lock() + .unwrap_or_else(PoisonError::into_inner) + .environments + .iter() + .any(|environment| retry.should_restart(environment)) + { + self.apply_selections(environments, retry); + } + } + + fn apply_selections(&self, environments: &[TurnEnvironmentSelection], retry: StartupRetry<'_>) { let bindings = self .configured .lock() @@ -460,9 +514,16 @@ impl ThreadEnvironments { environments, &bindings, self.connection_event_tx.get(), + retry, ); // Retry configured bindings with their own settings, keeping turn-only changes temporary. - self.update_configured_selections(&self.configured_snapshot_now().all_selections()); + self.update_selections_with_bindings( + &self.configured, + &self.configured_snapshot_now().all_selections(), + &[], + /*tx_event*/ None, + retry, + ); } /// Retains configured bindings without changing the active turn's attachments. @@ -472,6 +533,7 @@ impl ThreadEnvironments { environments, &[], /*tx_event*/ None, + StartupRetry::AfterSelectionUpdate, ); } @@ -481,6 +543,7 @@ impl ThreadEnvironments { environments: &[TurnEnvironmentSelection], bindings: &[SelectedTurnEnvironment], tx_event: Option<&Sender>, + retry: StartupRetry<'_>, ) { let mut state = state.lock().unwrap_or_else(PoisonError::into_inner); let ThreadEnvironmentsState { @@ -526,11 +589,7 @@ impl ThreadEnvironments { let (selected_environment, config_origin) = resolve_selection_config(selected_environment.clone(), thread_defaults); if let Some(environment) = previous { - let failed = - matches!( - environment.selection.config, - EnvironmentConfigState::Failed(_) - ) || matches!(environment.resolution.clone().now_or_never(), Some(Err(_))); + let failed = retry.should_restart(environment); let mut environment = if current_environment.is_some() { environment.clone() } else { diff --git a/codex-rs/core/src/session/environment.rs b/codex-rs/core/src/session/environment.rs index 09dd1c68d1..3820b8b7d8 100644 --- a/codex-rs/core/src/session/environment.rs +++ b/codex-rs/core/src/session/environment.rs @@ -214,11 +214,9 @@ impl Session { self.services .turn_environments .set_active_thread_defaults(configuration.inferred_environment_config()); - if environments != self.services.turn_environments.selections() { - self.services - .turn_environments - .update_selections(&environments); - } + self.services + .turn_environments + .activate_selections(&environments); } self.services.turn_environments.snapshot() }; diff --git a/codex-rs/core/tests/suite/remote_env_failure_tests.rs b/codex-rs/core/tests/suite/remote_env_failure_tests.rs index d50143eec6..0a5d424223 100644 --- a/codex-rs/core/tests/suite/remote_env_failure_tests.rs +++ b/codex-rs/core/tests/suite/remote_env_failure_tests.rs @@ -1,18 +1,49 @@ -//! Provisioning failures reach the model without terminating its turn. +//! Environment startup failures reach the model, and transient failures can recover next turn. use super::*; use codex_protocol::protocol::TurnEnvironmentRequest; use codex_protocol::protocol::TurnEnvironmentRequests; +use core_test_support::exec_server::ExecServerProcess; use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicBool; use test_case::test_case; +use tokio::task::JoinSet; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Id; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context as LayerContext; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; + +struct EnvironmentResolutionCounter(Arc); + +impl Layer for EnvironmentResolutionCounter { + fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: LayerContext<'_, S>) { + if attributes.metadata().target() == "codex_core::environment_selection" + && attributes.metadata().name() == "environments.resolve" + { + self.0.fetch_add(1, Ordering::SeqCst); + } + } +} #[test_case(true; "failed_before_turn")] #[test_case(false; "failed_while_waiting")] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Keep spawned tasks on the thread with the scoped tracing subscriber. +#[tokio::test(flavor = "current_thread")] async fn provisioning_failure_reaches_model_and_turn_continues( fail_before_turn: bool, ) -> Result<()> { const REASON: &str = "This repository is empty. Push an initial commit, then retry."; const CALL_ID: &str = "wait-for-failed-environment"; + const SECOND_CALL_ID: &str = "wait-for-failed-environment-again"; + let resolutions = Arc::new(AtomicUsize::new(0)); + // Keep concurrent tests without a subscriber from caching this callsite as disabled. + let _interest_cache_guard = + tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default()); + let _tracing = tracing_subscriber::registry() + .with(EnvironmentResolutionCounter(Arc::clone(&resolutions))) + .set_default(); let server = start_mock_server().await; let response_mock = mount_sse_sequence( &server, @@ -35,6 +66,21 @@ async fn provisioning_failure_reaches_model_and_turn_continues( ), ev_completed("resp-2"), ]), + sse(vec![ + ev_function_call( + SECOND_CALL_ID, + "wait_for_environment", + &json!({ + "environment_id": REMOTE_ENVIRONMENT_ID, + }) + .to_string(), + ), + ev_completed("resp-3"), + ]), + sse(vec![ + ev_assistant_message("message-2", "The repository still needs an initial commit."), + ev_completed("resp-4"), + ]), ], ) .await; @@ -42,6 +88,7 @@ async fn provisioning_failure_reaches_model_and_turn_continues( assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); }); let test = expect_startup(builder.build(&server)).await; + let resolutions_before_selection = resolutions.load(Ordering::SeqCst); let manager = test.thread_manager.environment_manager(); let provider = Arc::new(FailingNoiseConnectProvider::default()); manager.materialize_pending_noise_environment( @@ -80,7 +127,7 @@ async fn provisioning_failure_reaches_model_and_turn_continues( manager.report_environment_provisioning_status( REMOTE_ENVIRONMENT_ID.to_string(), Err(REASON.to_string()), - provider, + provider.clone(), )?; } wait_for_event(&test.codex, |event| { @@ -104,5 +151,139 @@ async fn provisioning_failure_reaches_model_and_turn_continues( let context = context_request.message_input_texts("user").join("\n"); assert!(context.contains("failed"), "{context}"); assert!(context.contains(REASON), "{context}"); + + let initial_resolutions = resolutions.load(Ordering::SeqCst); + assert!(initial_resolutions > resolutions_before_selection); + let selections = test.codex.environment_selections().await; + test.submit_text_turn("check the repository again").await?; + let requests = response_mock.requests(); + assert_eq!(requests.len(), 4); + let (output, _) = requests[3] + .function_call_output_content_and_success(SECOND_CALL_ID) + .context("the unchanged next turn should receive the original failure")?; + assert!( + output.as_deref().is_some_and(|text| text.contains(REASON)), + "{output:?}" + ); + assert_eq!(test.codex.environment_selections().await, selections); + // A retry can fail with the same provisioning error before invoking the provider. + // Observe attachment resolution itself so preserving the error alone cannot pass. + assert_eq!(resolutions.load(Ordering::SeqCst), initial_resolutions); + assert_eq!(provider.calls.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unchanged_next_turn_recovers_from_transient_environment_startup_failure() -> Result<()> { + const FIRST_CALL_ID: &str = "exec-before-recovery"; + const SECOND_CALL_ID: &str = "exec-after-recovery"; + const MARKER: &str = "recovered.txt"; + + let exec_server = ExecServerProcess::start().await?; + let upstream_address = url::Url::parse(&exec_server.websocket_url)? + .socket_addrs(|| None)? + .into_iter() + .next() + .context("exec-server should have a loopback address")?; + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let available = Arc::new(AtomicBool::new(false)); + let proxy_available = Arc::clone(&available); + let mut proxy_tasks = JoinSet::new(); + proxy_tasks.spawn(async move { + while let Ok((mut incoming, _)) = listener.accept().await { + if !proxy_available.load(Ordering::SeqCst) { + // Fail the handshake until the executor becomes available between turns. + continue; + } + let mut outgoing = TcpStream::connect(upstream_address).await?; + let _ = tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await; + } + Ok::<_, std::io::Error>(()) + }); + + let server = start_mock_server().await; + let arguments = json!({ + "cmd": format!("echo recovered > {MARKER}"), + "login": false, + "yield_time_ms": 10_000, + "environment_id": REMOTE_ENVIRONMENT_ID, + }) + .to_string(); + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call(FIRST_CALL_ID, "exec_command", &arguments), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "The environment is unavailable."), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_function_call(SECOND_CALL_ID, "exec_command", &arguments), + ev_completed("resp-3"), + ]), + sse(vec![ + ev_assistant_message("msg-2", "The command succeeded."), + ev_completed("resp-4"), + ]), + ], + ) + .await; + let mut builder = test_codex().with_config(|config| { + assert!(config.features.disable(Feature::DeferredExecutor).is_ok()); + }); + let test = builder.build(&server).await?; + let manager = test.thread_manager.environment_manager(); + manager.upsert_environment( + REMOTE_ENVIRONMENT_ID.to_string(), + websocket_url, + /*connect_timeout*/ None, + )?; + let environment = manager + .get_environment(REMOTE_ENVIRONMENT_ID) + .context("remote environment should be registered")?; + let selection = TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: PathUri::from_abs_path(&test.config.cwd), + workspace_roots: vec![PathUri::from_abs_path(&test.config.cwd)], + config: EnvironmentConfigState::FromThread, + selected_capability_roots: Default::default(), + }; + test.submit_turn_with_environments("write the marker", Some(vec![selection])) + .await?; + let selections = test.codex.environment_selections().await; + let marker_path = test.config.cwd.join(MARKER); + let first_output = response_mock + .function_call_output_text(FIRST_CALL_ID) + .context("model should receive the initial tool failure")?; + assert!( + first_output.contains("No usable execution environment is available."), + "{first_output}" + ); + assert!(!marker_path.as_path().exists()); + + // Restore only endpoint availability. The next turn keeps the same environment and + // workspace, without refreshing the transport or resubmitting environment settings. + available.store(true, Ordering::SeqCst); + test.submit_text_turn("try writing the marker again") + .await?; + let second_output = response_mock + .function_call_output_text(SECOND_CALL_ID) + .context("model should receive the recovered command result")?; + assert!( + second_output.contains("Process exited with code 0"), + "the unchanged next turn should execute successfully: {second_output}" + ); + assert!(fs::metadata(marker_path)?.len() > 0); + assert_eq!(test.codex.environment_selections().await, selections); + assert!(Arc::ptr_eq( + &environment, + &manager + .get_environment(REMOTE_ENVIRONMENT_ID) + .context("remote environment should remain registered")?, + )); Ok(()) }