diff --git a/lib/crates/fabro-workflow/src/handler/mod.rs b/lib/crates/fabro-workflow/src/handler/mod.rs index ec93ff93b..59ec4e801 100644 --- a/lib/crates/fabro-workflow/src/handler/mod.rs +++ b/lib/crates/fabro-workflow/src/handler/mod.rs @@ -74,29 +74,7 @@ impl EngineServices { /// Bridge the core executor's atomic cancel flag to sandbox command cancellation. pub fn sandbox_cancel_token(&self) -> Option { - let cancel_requested = self.cancel_requested.clone()?; - let token = CancellationToken::new(); - - if cancel_requested.load(Ordering::Relaxed) { - token.cancel(); - return Some(token); - } - - let token_clone = token.clone(); - tokio::spawn(async move { - loop { - if token_clone.is_cancelled() { - return; - } - if cancel_requested.load(Ordering::Relaxed) { - token_clone.cancel(); - return; - } - time::sleep(Duration::from_millis(10)).await; - } - }); - - Some(token) + sandbox_cancel_token(self.cancel_requested.clone()) } /// Run lifecycle hooks and return the merged decision. @@ -150,6 +128,34 @@ impl EngineServices { } } +pub(crate) fn sandbox_cancel_token( + cancel_requested: Option>, +) -> Option { + let cancel_requested = cancel_requested?; + let token = CancellationToken::new(); + + if cancel_requested.load(Ordering::Relaxed) { + token.cancel(); + return Some(token); + } + + let token_clone = token.clone(); + tokio::spawn(async move { + loop { + if token_clone.is_cancelled() { + return; + } + if cancel_requested.load(Ordering::Relaxed) { + token_clone.cancel(); + return; + } + time::sleep(Duration::from_millis(10)).await; + } + }); + + Some(token) +} + /// The handler interface for node execution. #[async_trait] pub trait Handler: Send + Sync { diff --git a/lib/crates/fabro-workflow/src/pipeline/initialize.rs b/lib/crates/fabro-workflow/src/pipeline/initialize.rs index a03fd9bc7..de09870f7 100644 --- a/lib/crates/fabro-workflow/src/pipeline/initialize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/initialize.rs @@ -589,16 +589,24 @@ pub async fn initialize( index, }); let cmd_start = Instant::now(); + let cancel_token = + crate::handler::sandbox_cancel_token(options.run_options.cancel_token.clone()); let result = sandbox .exec_command( command, options.lifecycle.setup_command_timeout_ms, None, None, - None, + cancel_token.clone(), ) .await .map_err(|e| FabroError::engine(format!("Setup command failed: {e}")))?; + if let Some(token) = &cancel_token { + if token.is_cancelled() { + return Err(FabroError::Cancelled); + } + token.cancel(); + } let duration_ms = crate::millis_u64(cmd_start.elapsed()); if result.exit_code != 0 { options.emitter.emit(&Event::SetupFailed {