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(()) }