diff --git a/codex-rs/exec-server/BUILD.bazel b/codex-rs/exec-server/BUILD.bazel index db2230799f66..a93c14990df8 100644 --- a/codex-rs/exec-server/BUILD.bazel +++ b/codex-rs/exec-server/BUILD.bazel @@ -9,6 +9,10 @@ exports_files( codex_rust_crate( name = "exec-server", crate_name = "codex_exec_server", + crate_srcs = glob([ + "src/**/*.rs", + "tests/unit/**/*.rs", + ]), deps_extra = [ "@crates//:opentelemetry", "@crates//:opentelemetry_sdk", diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 20053588e1c5..bc441d9fe4f4 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -7,6 +7,11 @@ license.workspace = true [lib] doctest = false +# Compiled through #[path] in src/client.rs; cargo-shear 1.11.2 does not normalize +# the ../tests path and incorrectly reports this linked unit-test module as unlinked. +[package.metadata.cargo-shear] +ignored-paths = ["tests/unit/client_provisioning_tests.rs"] + [lints] workspace = true diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index eb76b6f373df..f33b21c137fd 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -129,6 +129,9 @@ use codex_http_client::HttpClientFactory; pub(crate) mod accepted; pub(crate) mod http_client; mod network_policy_audit; +#[cfg(test)] +#[path = "../tests/unit/client_provisioning_tests.rs"] +mod provisioning_tests; #[path = "client_recovery.rs"] mod recovery; #[path = "client_refresh.rs"] @@ -426,6 +429,13 @@ impl LazyRemoteExecServerClient { } pub(crate) async fn status(&self) -> crate::EnvironmentObservedStatus { + if let Some(ExecServerTransportParams::Deferred(deferred)) = &self.transport_params + && let Some(Err(error)) = deferred.readiness.borrow().as_ref() + { + return crate::EnvironmentObservedStatus::Disconnected { + error: ExecServerError::ProvisioningFailed(error.clone()).to_string(), + }; + } // Fail-fast lookup preserves the non-mutating contract: never start or recover a client. let client = match self.fail_fast().get().await { Ok(client) => client, @@ -498,23 +508,34 @@ impl LazyRemoteExecServerClient { } async fn initial_client(&self) -> Result { - if self.can_reconnect() + let result = if self.can_reconnect() && (self.startup.cancelled.is_cancelled() || self.startup.result.get().is_some_and(|result| { result .as_ref() .is_err_and(|error| recovery::is_retryable_recovery_error(error)) - })) + })) { + Box::pin(self.reconnect()).await + } else { + self.startup + .result + .get_or_init(|| self.connect_once(&self.startup)) + .await + .clone() + .map_err(ExecServerError::ConnectionAttempt) + }; + // Ready may arrive before an older attempt publishes its provisioning failure. + if let Err(ExecServerError::ConnectionAttempt(error)) = &result + && matches!(error.as_ref(), ExecServerError::ProvisioningFailed(_)) + && matches!( + &self.transport_params, + Some(ExecServerTransportParams::Deferred(deferred)) + if matches!(*deferred.readiness.borrow(), Some(Ok(()))) + ) { return Box::pin(self.reconnect()).await; } - - self.startup - .result - .get_or_init(|| self.connect_once(&self.startup)) - .await - .clone() - .map_err(ExecServerError::ConnectionAttempt) + result } async fn reconnect(&self) -> Result { @@ -618,6 +639,8 @@ pub enum ExecServerError { Closed, #[error("{0}")] Disconnected(String), + #[error("environment unavailable: {0}")] + ProvisioningFailed(String), #[error("failed to serialize or deserialize exec-server JSON: {0}")] Json(#[from] serde_json::Error), #[error("HTTP request failed: {0}")] diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 25bc54819ec7..8bb5a6e72063 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -818,7 +818,8 @@ pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool { is_transport_closed_error(error) || matches!( error, - ExecServerError::WebSocketConnectTimeout { .. } + ExecServerError::ProvisioningFailed(_) + | ExecServerError::WebSocketConnectTimeout { .. } | ExecServerError::WebSocketConnect { .. } | ExecServerError::InitializeTimedOut { .. } ) diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index bf1f9fa706e3..766f2b0d3258 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -187,9 +187,7 @@ impl ExecServerClient { .to_string(), ) })?; - provisioning_result.map_err(|message| { - ExecServerError::Disconnected(format!("environment unavailable: {message}")) - })?; + provisioning_result.map_err(ExecServerError::ProvisioningFailed)?; } let (websocket_url, connect_timeout, initialize_timeout) = match transport_params { diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 0cb6df392a32..a63938be58ef 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -395,9 +395,9 @@ impl EnvironmentManager { /// Ordinary environments are ignored. A provisioned environment keeps the same `Arc` from /// Pending through Ready or Failed, and is created if the report arrives first. /// - /// Ready updates capability roots. Failed keeps the first error. Repeating the same result is - /// allowed, but changing between Ready and Failed is rejected. Invalid Ready information fails - /// an existing Pending environment but does not create a missing environment. + /// Ready updates capability roots and can recover a failed provisioning attempt. Failed keeps + /// the first error until a Ready report arrives; a late failure cannot replace Ready. Invalid + /// Ready information fails an existing Pending environment but does not create a missing one. /// /// This only updates provisioning. The connection starts when the environment is selected. pub fn report_environment_provisioning_status( @@ -835,31 +835,19 @@ impl Environment { return Ok(()); }; let mut transition_error = None; - provisioning_status_tx.send_if_modified(|current| match current.as_ref() { - Some(Err(error)) => { - transition_error = Some(ExecServerError::Protocol(format!( - "environment `{environment_id}` provisioning already failed: {error}" - ))); - false - } - None => { - if let Err(error) = validate_environment_ready_info(environment_id, &ready_info) { + provisioning_status_tx.send_if_modified(|current| { + if let Err(error) = validate_environment_ready_info(environment_id, &ready_info) { + let pending = current.is_none(); + if pending { *current = Some(Err(error.to_string())); - transition_error = Some(error); - } else { - self.ready_info.store(Some(Arc::new(ready_info.clone()))); - *current = Some(Ok(())); } - true - } - Some(Ok(())) => { - if let Err(error) = validate_environment_ready_info(environment_id, &ready_info) { - transition_error = Some(error); - } else { - self.ready_info.store(Some(Arc::new(ready_info.clone()))); - } - false + transition_error = Some(error); + return pending; } + self.ready_info.store(Some(Arc::new(ready_info.clone()))); + let was_ready = matches!(current, Some(Ok(()))); + *current = Some(Ok(())); + !was_ready }); transition_error.map_or(Ok(()), Err) diff --git a/codex-rs/exec-server/tests/deferred_environment.rs b/codex-rs/exec-server/tests/deferred_environment.rs index d08b9ab5e622..123e55e8b730 100644 --- a/codex-rs/exec-server/tests/deferred_environment.rs +++ b/codex-rs/exec-server/tests/deferred_environment.rs @@ -136,7 +136,7 @@ async fn ordinary_environment_ignores_provisioning_reports() -> anyhow::Result<( } #[tokio::test] -async fn failure_before_materialization_is_terminal_without_connecting() -> anyhow::Result<()> { +async fn failure_before_materialization_is_reported_without_connecting() -> anyhow::Result<()> { let manager = environment_manager_without_environments(); let provider = Arc::new(FailingNoiseConnectProvider::default()); @@ -153,8 +153,15 @@ async fn failure_before_materialization_is_terminal_without_connecting() -> anyh )?; assert!(Arc::ptr_eq(&failed, &materialized)); + assert_eq!( + failed.status().await, + codex_exec_server::EnvironmentObservedStatus::Disconnected { + error: "environment unavailable: provisioning failed".to_string(), + } + ); let error = failed.wait_until_ready().await.unwrap_err(); assert!(error.to_string().ends_with("provisioning failed")); + assert!(failed.startup_finished()); assert_eq!(provider.calls(), 0); Ok(()) } @@ -183,7 +190,7 @@ async fn failure_releases_the_existing_pending_environment_without_connecting() } #[tokio::test] -async fn repeated_failure_preserves_the_first_error_and_rejects_ready() -> anyhow::Result<()> { +async fn repeated_failure_preserves_the_first_error_until_ready() -> anyhow::Result<()> { let manager = environment_manager_without_environments(); let provider = Arc::new(FailingNoiseConnectProvider::default()); let failed = manager @@ -203,18 +210,39 @@ async fn repeated_failure_preserves_the_first_error_and_rejects_ready() -> anyho .expect("repeated failure should be idempotent"); assert!(Arc::ptr_eq(&failed, &repeated)); - let error = manager + let error = failed.wait_until_ready().await.unwrap_err(); + assert!(error.to_string().ends_with("first failure")); + assert_eq!(provider.calls(), 0); + let invalid_ready_error = manager .report_environment_provisioning_status( "tools".to_string(), - Ok(ready_info("selected-root", "tools")?), + Ok(ready_info("selected-root", "other")?), provider.clone(), ) .unwrap_err(); - assert!(error.to_string().contains("first failure")); + assert!(matches!(invalid_ready_error, ExecServerError::Protocol(_))); + assert_eq!( + failed.wait_until_ready().await.unwrap_err().to_string(), + error.to_string() + ); assert!(failed.selected_capability_roots().is_empty()); - let error = failed.wait_until_ready().await.unwrap_err(); - assert!(error.to_string().ends_with("first failure")); assert_eq!(provider.calls(), 0); + let selected = ready_info("selected-root", "tools")?; + let ready = manager + .report_environment_provisioning_status( + "tools".to_string(), + Ok(selected.clone()), + provider.clone(), + )? + .expect("successful provisioning should recover the same environment"); + assert!(Arc::ptr_eq(&failed, &ready)); + assert_eq!( + failed.selected_capability_roots(), + selected.selected_capability_roots + ); + let error = failed.wait_until_ready().await.unwrap_err(); + assert!(error.to_string().contains("test Noise provider called")); + assert_eq!(provider.calls(), 1); Ok(()) } @@ -319,19 +347,19 @@ async fn invalid_ready_report_fails_the_provisioning_gate() -> anyhow::Result<() assert!(environment.selected_capability_roots().is_empty()); assert_eq!(provider.calls(), 0); - let later_ready_error = manager + let selected = ready_info("selected-root", "tools")?; + let reported = manager .report_environment_provisioning_status( "tools".to_string(), - Ok(ready_info("selected-root", "tools")?), + Ok(selected.clone()), provider.clone(), - ) - .unwrap_err(); - assert!( - later_ready_error - .to_string() - .contains("provisioning already failed") + )? + .expect("a corrected ready report should recover provisioning"); + assert!(Arc::ptr_eq(&environment, &reported)); + assert_eq!( + environment.selected_capability_roots(), + selected.selected_capability_roots ); - assert!(environment.selected_capability_roots().is_empty()); assert_eq!(provider.calls(), 0); Ok(()) } diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index ae56123bdfe3..d37f76331a20 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -94,7 +94,7 @@ impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[serial_test::serial] -async fn pending_noise_environment_connects_and_reconnects_after_ready_report() -> Result<()> { +async fn failed_noise_environment_recovers_and_reconnects_after_ready_report() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let rendezvous_address = listener.local_addr()?; let environment_rendezvous_url = @@ -194,21 +194,46 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report() executor_public_key: registered_executor_public_key(®istry).await?, calls: AtomicUsize::new(0), }); - let manager = EnvironmentManager::without_environments(http_client_factory); + let manager = Arc::new(EnvironmentManager::without_environments( + http_client_factory, + )); let environment = manager .materialize_pending_noise_environment(ENVIRONMENT_ID.to_string(), provider.clone())?; let mut connection_state = environment .subscribe_connection_state() .context("remote environment connection state")?; - assert_eq!(provider.calls(), 0); + let capability_root = TempDir::new()?; + let skill_file = capability_root.path().join("SKILL.md"); + let skill_contents = b"# Recovered capability\n"; + std::fs::write(&skill_file, skill_contents)?; let selected_capability_roots = vec![SelectedCapabilityRoot { id: "executor-plugin".to_string(), location: CapabilityRootLocation::Environment { environment_id: ENVIRONMENT_ID.to_string(), - path: PathUri::parse("file:///plugins/executor-plugin")?, + path: PathUri::from_host_native_path(capability_root.path())?, }, }]; + assert!( + manager + .resolve_selected_capability_roots(&selected_capability_roots, &HashMap::new()) + .await + .is_empty() + ); + assert_eq!(provider.calls(), 0); + manager.report_environment_provisioning_status( + ENVIRONMENT_ID.to_string(), + Err("first provisioning attempt failed".to_string()), + provider.clone(), + )?; + timeout(TEST_TIMEOUT, async { + while !environment.startup_finished() { + tokio::task::yield_now().await; + } + }) + .await + .expect("failed capability startup should record its completion"); + assert_eq!(provider.calls(), 0); let reported = manager .report_environment_provisioning_status( ENVIRONMENT_ID.to_string(), @@ -220,9 +245,15 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report() .context("ready report should apply to the pending environment")?; assert!(Arc::ptr_eq(&environment, &reported)); assert_eq!(provider.calls(), 0); - let initial_info = tokio::spawn({ - let environment = Arc::clone(&environment); - async move { environment.info().await } + // Capability resolution must retry on its own, without an explicit connection call. + let resolved_roots = tokio::spawn({ + let manager = Arc::clone(&manager); + let selected_capability_roots = selected_capability_roots.clone(); + async move { + manager + .resolve_selected_capability_roots(&selected_capability_roots, &HashMap::new()) + .await + } }); let harness_websocket = accept_websocket(&listener, "harness").await?; assert_eq!( @@ -234,9 +265,31 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report() harness_websocket, Arc::new(Mutex::new(Vec::new())), )); - let initial_info = timeout(TEST_TIMEOUT, initial_info) + let resolved_roots = timeout(TEST_TIMEOUT, resolved_roots) .await - .context("pending Noise environment should become ready")???; + .context("capability resolution should recover after Ready")??; + assert_eq!( + resolved_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect::>(), + selected_capability_roots + ); + let [resolved_root] = resolved_roots.as_slice() else { + anyhow::bail!("the recovered capability root should resolve"); + }; + assert!(Arc::ptr_eq(resolved_root.environment(), &environment)); + let recovered_skill = resolved_root + .environment() + .get_filesystem() + .read_file( + &PathUri::from_host_native_path(skill_file)?, + Default::default(), + /*sandbox*/ None, + ) + .await?; + assert_eq!(recovered_skill, skill_contents.to_vec()); + let initial_info = environment.info().await?; assert_eq!( environment.selected_capability_roots(), selected_capability_roots diff --git a/codex-rs/exec-server/tests/unit/client_provisioning_tests.rs b/codex-rs/exec-server/tests/unit/client_provisioning_tests.rs new file mode 100644 index 000000000000..b4281c625ed2 --- /dev/null +++ b/codex-rs/exec-server/tests/unit/client_provisioning_tests.rs @@ -0,0 +1,76 @@ +//! Exercises provisioning recovery while startup is publishing its previous failure. + +use super::*; +use crate::client_api::Deferred; +use crate::noise_channel::NoiseChannelIdentity; +use futures::poll; +use tokio::sync::oneshot; + +#[test_case::test_case(false; "initial startup")] +#[test_case::test_case(true; "subsequent attempt")] +#[tokio::test] +async fn caller_after_ready_retries_an_unpublished_provisioning_failure(reconnecting: bool) { + let (readiness, readiness_rx) = watch::channel(Some(Err("first failure".to_string()))); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::Deferred(Box::new(Deferred { + readiness: readiness_rx, + transport: ExecServerTransportParams::NoiseRendezvous { + provider: Arc::new(FailingProvider), + identity: NoiseChannelIdentity::generate().expect("Noise identity"), + }, + })), + HttpClientFactory::new(codex_http_client::OutboundProxyPolicy::ReqwestDefault), + ); + let (failed_tx, failed_rx) = oneshot::channel(); + let (publish_tx, publish_rx) = oneshot::channel(); + let attempt = if reconnecting { + assert!(client.wait_until_ready().await.is_err()); + let attempt = Arc::new(ConnectionAttempt::default()); + *client.reconnect.lock().expect("reconnect lock") = Some(Arc::clone(&attempt)); + attempt + } else { + Arc::clone(&client.startup) + }; + let startup_client = client.clone(); + let startup = tokio::spawn(async move { + attempt + .result + .get_or_init(|| async { + let result = startup_client.connect_once(&attempt).await; + assert!(result.is_err()); + failed_tx.send(()).expect("failure observed"); + // Suspend at the publication boundary, as another executor thread could. + publish_rx.await.expect("publish startup result"); + result + }) + .await + .clone() + }); + failed_rx.await.expect("startup consumed the Failed report"); + readiness.send_replace(Some(Ok(()))); + + let mut after_ready = Box::pin(client.wait_until_ready()); + assert!(poll!(&mut after_ready).is_pending()); + publish_tx.send(()).expect("release startup"); + let error = after_ready.await.unwrap_err(); + assert!( + error.to_string().contains("provider reached after Ready"), + "post-Ready caller should retry the stale failure: {error}" + ); + assert!(startup.await.expect("startup task").is_err()); +} + +struct FailingProvider; + +impl crate::NoiseRendezvousConnectProvider for FailingProvider { + fn connect_bundle( + &self, + _: crate::NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(ExecServerError::Protocol( + "provider reached after Ready".to_string(), + )) + }) + } +}