diff --git a/crates/zesdex-backend/src/app/runtime/actions/turn.rs b/crates/zesdex-backend/src/app/runtime/actions/turn.rs index 2932169..ed00187 100644 --- a/crates/zesdex-backend/src/app/runtime/actions/turn.rs +++ b/crates/zesdex-backend/src/app/runtime/actions/turn.rs @@ -306,9 +306,7 @@ pub(super) fn run_agent_turn( // Check abort after pipeline completes, before entering main loop. // This catches the case where the user pressed Esc during the pipeline // phase, which previously ran unchecked for minutes at a time. - if tc - .abort_flag - .load(std::sync::atomic::Ordering::SeqCst) + if crate::app::util::abort::is_aborted_direct(&tc.abort_flag) { push_event(&events_q, TurnEvent::Error("Generation aborted by user".to_string())); return Ok(()); @@ -326,9 +324,7 @@ pub(super) fn run_agent_turn( // Skip message compaction if abort was requested — the non-streaming // LLM call for summarization would block without checking abort_flag. - let wire_msgs = if !tc - .abort_flag - .load(std::sync::atomic::Ordering::SeqCst) + let wire_msgs = if !crate::app::util::abort::is_aborted_direct(&tc.abort_flag) && crate::app::runtime::context::shaping::should_shape( token_estimate, max_wire_tokens, @@ -371,7 +367,7 @@ pub(super) fn run_agent_turn( Some(tc.temperature), tc.max_tokens, |event| -> bool { - if tc.abort_flag.load(std::sync::atomic::Ordering::SeqCst) { + if crate::app::util::abort::is_aborted_direct(&tc.abort_flag) { return false; } if let Ok(mut q) = events_q.lock() { @@ -425,7 +421,7 @@ pub(super) fn run_agent_turn( Ok((msg, u)) => (msg, u.or(usage)), Err(e) => { // If abort was requested, return immediately. - if tc.abort_flag.load(std::sync::atomic::Ordering::SeqCst) + if crate::app::util::abort::is_aborted_direct(&tc.abort_flag) || e.to_string().contains("aborted") { push_event(&events_q, TurnEvent::Error( @@ -552,9 +548,7 @@ pub(super) fn run_agent_turn( }); for (tool_call, tool_name, args, output, is_error, is_edit) in results_vec { - if tc - .abort_flag - .load(std::sync::atomic::Ordering::SeqCst) + if crate::app::util::abort::is_aborted_direct(&tc.abort_flag) { push_event(&events_q, TurnEvent::Error( "Turn aborted by user".to_string(), diff --git a/crates/zesdex-backend/src/app/subagent/engine.rs b/crates/zesdex-backend/src/app/subagent/engine.rs index cac5e66..9c1f9af 100644 --- a/crates/zesdex-backend/src/app/subagent/engine.rs +++ b/crates/zesdex-backend/src/app/subagent/engine.rs @@ -117,11 +117,7 @@ pub fn run_subagent( for step in 0..ctx.max_steps { // Check abort flag before each LLM call so a stuck subagent can // be cancelled from the parent (mirrors main agent behaviour). - if ctx - .abort_flag - .as_ref() - .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) - { + if crate::app::util::abort::is_aborted(&ctx.abort_flag) { let _ = tx.blocking_send(SubagentEvent::StepFailed { step, error: "subagent aborted by parent".to_string(), @@ -154,11 +150,7 @@ pub fn run_subagent( Some(4096), |event| -> bool { // Check abort on every SSE event for responsive cancellation. - if ctx - .abort_flag - .as_ref() - .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) - { + if crate::app::util::abort::is_aborted(&ctx.abort_flag) { return false; // signals provider to abort } match event { @@ -194,10 +186,7 @@ pub fn run_subagent( Ok(result) => break result, Err(e) => { let err_str = e.to_string(); - let is_abort = ctx - .abort_flag - .as_ref() - .is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) + let is_abort = crate::app::util::abort::is_aborted(&ctx.abort_flag) || err_str.contains("aborted"); if is_abort || !should_retry_subagent_step(&err_str) || step_attempt >= max_step_retries { @@ -279,7 +268,7 @@ pub fn run_subagent( for tool_call in &tool_calls { let handle = s.spawn(move || { // Check abort flag before each tool execution - if ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) { + if crate::app::util::abort::is_aborted(&ctx.abort_flag) { return (tool_call, Err(anyhow::anyhow!("subagent aborted by parent during tool execution"))); } diff --git a/crates/zesdex-backend/src/app/util/abort.rs b/crates/zesdex-backend/src/app/util/abort.rs new file mode 100644 index 0000000..ac22ca1 --- /dev/null +++ b/crates/zesdex-backend/src/app/util/abort.rs @@ -0,0 +1,26 @@ +//! Shared abort-flag checks. +//! +//! The two variants (Option> and bare AtomicBool) are +//! used across the agent runtime, subagent, workflow engine, and provider. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +/// Check whether an optional abort flag has been signalled. +pub fn is_aborted(flag: &Option>) -> bool { + flag.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) +} + +/// Check whether a bare abort flag has been signalled. +pub fn is_aborted_direct(flag: &AtomicBool) -> bool { + flag.load(Ordering::SeqCst) +} + +/// Check whether an optional borrowed abort flag has been signalled. +/// +/// This variant handles the `Option<&AtomicBool>` pattern used in +/// service/provider.rs where the flag is passed as a by-value optional +/// reference rather than an `Arc`. +pub fn is_aborted_ref(flag: Option<&AtomicBool>) -> bool { + flag.is_some_and(|f| f.load(Ordering::SeqCst)) +} diff --git a/crates/zesdex-backend/src/app/util/mod.rs b/crates/zesdex-backend/src/app/util/mod.rs index c02e7fc..0ab54e3 100644 --- a/crates/zesdex-backend/src/app/util/mod.rs +++ b/crates/zesdex-backend/src/app/util/mod.rs @@ -1,3 +1,4 @@ //! Utility modules for shared helpers. +pub mod abort; pub mod backoff; diff --git a/crates/zesdex-backend/src/app/workflow/engine/mod.rs b/crates/zesdex-backend/src/app/workflow/engine/mod.rs index b882c72..5f79179 100644 --- a/crates/zesdex-backend/src/app/workflow/engine/mod.rs +++ b/crates/zesdex-backend/src/app/workflow/engine/mod.rs @@ -25,7 +25,7 @@ pub(crate) use primitives::PrimitiveCtx; use serde::{Deserialize, Serialize}; use std::sync::{ - atomic::{AtomicBool, Ordering}, + atomic::AtomicBool, Arc, Mutex, }; use std::time::Duration; @@ -376,10 +376,7 @@ fn spawn_single_agent(sp: SpawnCtx<'_>) -> anyhow::Result { }); // Check abort before even starting the subagent. - if sp - .abort_flag - .as_ref() - .is_some_and(|f| f.load(Ordering::SeqCst)) + if crate::app::util::abort::is_aborted(sp.abort_flag) { anyhow::bail!("subagent '{}' aborted before start", sp.agent_name); } @@ -405,9 +402,7 @@ fn spawn_single_agent(sp: SpawnCtx<'_>) -> anyhow::Result { for attempt in 1..=2 { // Don't retry if aborted. - if bg_abort_thread - .as_ref() - .is_some_and(|f| f.load(Ordering::SeqCst)) + if crate::app::util::abort::is_aborted(&bg_abort_thread) { let _ = done_tx.send(Err(anyhow::anyhow!( "subagent '{bg_name_thread}' aborted by user" @@ -453,7 +448,7 @@ fn spawn_single_agent(sp: SpawnCtx<'_>) -> anyhow::Result { "subagent '{bg_name}' timed out after {timeout}ms", )); } - if bg_abort.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) { + if crate::app::util::abort::is_aborted(&bg_abort) { break Err(anyhow::anyhow!("subagent '{bg_name}' aborted by user")); } } @@ -462,7 +457,7 @@ fn spawn_single_agent(sp: SpawnCtx<'_>) -> anyhow::Result { if let Ok(r) = done_rx.recv_timeout(poll_interval) { break r; } - if bg_abort.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) { + if crate::app::util::abort::is_aborted(&bg_abort) { break Err(anyhow::anyhow!("subagent '{bg_name}' aborted by user")); } } diff --git a/crates/zesdex-backend/src/service/provider.rs b/crates/zesdex-backend/src/service/provider.rs index 5722d99..cee03f7 100644 --- a/crates/zesdex-backend/src/service/provider.rs +++ b/crates/zesdex-backend/src/service/provider.rs @@ -198,7 +198,7 @@ impl LlmClient { // Check abort before each retry so user cancellation is // responsive even during a long non-streaming backoff chain. - if abort_flag.is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) { + if crate::app::util::abort::is_aborted_ref(abort_flag) { anyhow::bail!("aborted"); } @@ -373,7 +373,7 @@ impl LlmClient { if meaningful_content { // Check abort before entering the blocking non-streaming // call — otherwise the fallback ignores user cancellation. - if abort_flag.is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) { + if crate::app::util::abort::is_aborted_ref(abort_flag) { return Err(anyhow::anyhow!("aborted")); } tracing::warn!(