MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement abort mechanism for workflows and subagents

Browse Source
This commit is contained in:
asepharyana
2026-07-13 09:09:23 +07:00
parent 65647ce517
commit 104b0daf4c
5 changed files with 104 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/app/runtime/actions/mod.rs
+15 -1
View File
@@ -682,8 +682,9 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) {
});
let args: HashMap<String, String> = HashMap::new();
let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
let result = crate::app::workflow::engine::run_workflow_tracked(
&wf, &args, Some(&live), &session_dir, &workspace_roots,
&wf, &args, &no_abort, Some(&live), &session_dir, &workspace_roots,
);
let (kind, message) = match result {
@@ -1082,12 +1083,14 @@ fn run_agent_turn(
});
}
let pipeline_abort = Some(tc.abort_flag.clone());
let pipeline_result = if use_full {
crate::app::workflow::company::run_company_pipeline(
user_request,
&tc.edit_log_session_dir,
&tc.workspace_roots,
Some(events_q),
&pipeline_abort,
)
} else {
crate::app::workflow::company::run_company_pipeline_quick(
@@ -1095,6 +1098,7 @@ fn run_agent_turn(
&tc.edit_log_session_dir,
&tc.workspace_roots,
Some(events_q),
&pipeline_abort,
)
};
@@ -1127,6 +1131,16 @@ fn run_agent_turn(
tracing::debug!("[ceo] pipeline not triggered — handling directly");
}
// 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 let Ok(mut q) = events_q.lock() {
q.push_back(TurnEvent::Error("Generation aborted by user".to_string()));
}
return Ok(());
}
let mut turn_step = 0usize;
let mut todo_retry_count = 0usize;
src/app/workflow/company.rs
+5 -2
View File
@@ -22,7 +22,7 @@
use std::collections::HashMap;
use std::fmt::Write;
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Mutex, atomic::AtomicBool};
use crate::app::workflow::engine::{execute_primitive, LiveStateFn, AgentStatus};
use crate::app::workflow::script::{ScriptPrimitive, ScriptOptions, WorkflowScript};
use crate::app::subagent::division;
@@ -45,6 +45,7 @@ pub fn run_company_pipeline(
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
turn_events: Option<&Arc<Mutex<std::collections::VecDeque<crate::app::state::runtime::TurnEvent>>>>,
abort_flag: &Option<Arc<AtomicBool>>,
) -> anyhow::Result<String> {
let divisions = division::all_divisions();
let mut pipeline_scripts: Vec<ScriptPrimitive> = Vec::with_capacity(divisions.len());
@@ -102,6 +103,7 @@ pub fn run_company_pipeline(
&args,
1,
true,
abort_flag,
live_ref,
session_dir,
workspaces,
@@ -126,6 +128,7 @@ pub fn run_company_pipeline_quick(
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
turn_events: Option<&Arc<Mutex<std::collections::VecDeque<crate::app::state::runtime::TurnEvent>>>>,
abort_flag: &Option<Arc<AtomicBool>>,
) -> anyhow::Result<String> {
let divisions = division::all_divisions();
// Only use first 3 divisions for quick pipeline: Strategy, Engineering, Quality
@@ -174,7 +177,7 @@ pub fn run_company_pipeline_quick(
let results = execute_primitive(
&wf.script, &args, 1, true,
live.as_ref(), session_dir, workspaces, &findings, None,
abort_flag, live.as_ref(), session_dir, workspaces, &findings, None,
)?;
let all_findings = findings.lock()
src/app/workflow/engine.rs
+77 -29
View File
@@ -15,7 +15,7 @@
//! leaks between concurrent workflow runs.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Mutex, atomic::{AtomicBool, Ordering}};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::script::{ScriptPrimitive, WorkflowScript};
@@ -105,6 +105,7 @@ fn spawn_single_agent(
prompt: &str,
findings_snapshot: &[String],
findings: &Arc<Mutex<Vec<String>>>,
abort_flag: &Option<Arc<AtomicBool>>,
live: Option<&LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
@@ -157,9 +158,7 @@ fn spawn_single_agent(
// Link the shared findings Arc so note_finding calls within this
// subagent write into the same vec visible to sibling agents.
ctx.workflow_findings = Some(findings.clone());
// Abort flag stays None by default — the parent can set it to abort
// long-running agents. No abort mechanism is wired yet at this level;
// future work can expose a kill-switch per agent via the live callback.
ctx.abort_flag = abort_flag.clone();
// Create an mpsc channel and drain events in a background thread.
// The drain thread also pushes intra-division progress updates to the
@@ -221,27 +220,59 @@ fn spawn_single_agent(
}
});
// Enforce timeout by running subagent on a separate thread and
// waiting with a deadline. If the deadline expires, the thread is
// abandoned (Rust threads cannot be forcibly killed, but we proceed
// without waiting for it — the drain thread will drop when tx is
// dropped on thread exit).
let result = if let Some(timeout) = timeout_ms {
let (done_tx, done_rx) = std::sync::mpsc::channel::<anyhow::Result<String>>();
let timeout_ctx = ctx;
let timeout_tx = tx;
std::thread::spawn(move || {
let _ = done_tx.send(run_subagent(&timeout_ctx, &timeout_tx));
});
match done_rx.recv_timeout(Duration::from_millis(timeout)) {
Ok(r) => r,
Err(_) => Err(anyhow::anyhow!(
"subagent '{agent_name}' timed out after {timeout}ms",
)),
// Check abort before even starting the subagent.
if abort_flag.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) {
anyhow::bail!("subagent '{agent_name}' aborted before start");
}
// Run subagent on a separate thread so the abort flag can be polled.
// If abort is requested while the subagent is running, we abandon the
// thread (Rust threads cannot be forcibly killed) and return early.
let (done_tx, done_rx) = std::sync::mpsc::channel::<anyhow::Result<String>>();
let bg_ctx = ctx;
let bg_tx = tx;
let bg_name = agent_name.to_string();
let bg_abort = abort_flag.clone();
std::thread::spawn(move || {
let _ = done_tx.send(run_subagent(&bg_ctx, &bg_tx));
});
let poll_interval = Duration::from_millis(500);
let result = if let Some(timeout) = timeout_ms {
let deadline = Duration::from_millis(timeout);
let mut elapsed = Duration::ZERO;
loop {
match done_rx.recv_timeout(poll_interval) {
Ok(r) => break r,
Err(_) => {
elapsed += poll_interval;
if elapsed >= deadline {
break Err(anyhow::anyhow!(
"subagent '{bg_name}' timed out after {timeout}ms",
));
}
if bg_abort.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) {
break Err(anyhow::anyhow!(
"subagent '{bg_name}' aborted by user",
));
}
}
}
} else {
run_subagent(&ctx, &tx)
};
}
} else {
loop {
match done_rx.recv_timeout(poll_interval) {
Ok(r) => break r,
Err(_) => {
if bg_abort.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) {
break Err(anyhow::anyhow!(
"subagent '{bg_name}' aborted by user",
));
}
}
}
}
};
let completed_at = chrono::Utc::now().timestamp_millis();
@@ -305,6 +336,7 @@ pub fn execute_primitive(
args: &HashMap<String, String>,
concurrency_cap: usize,
continue_on_error: bool,
abort_flag: &Option<Arc<AtomicBool>>,
live: Option<&LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
@@ -317,7 +349,7 @@ pub fn execute_primitive(
let findings_snapshot = findings.lock().map(|f| f.clone()).unwrap_or_default();
let agent_id = uuid::Uuid::new_v4().to_string();
let agent_name = resolved.chars().take(40).collect::<String>();
match spawn_single_agent(&agent_id, &agent_name, &resolved, &findings_snapshot, findings, live, session_dir, workspaces, timeout_ms) {
match spawn_single_agent(&agent_id, &agent_name, &resolved, &findings_snapshot, findings, abort_flag, live, session_dir, workspaces, timeout_ms) {
Ok(text) => Ok(vec![text]),
Err(e) => {
if continue_on_error {
@@ -348,6 +380,7 @@ pub fn execute_primitive(
let sem = Arc::clone(&semaphore);
let results = Arc::clone(&results);
let cap = concurrency_cap;
let abort = abort_flag.clone();
let live_clone = live.cloned();
let session_dir = session_dir.to_path_buf();
let workspaces = workspaces.to_vec();
@@ -358,6 +391,7 @@ pub fn execute_primitive(
let _permit = sem.acquire();
let result = execute_primitive(
&script, &args, cap, continue_on_error,
&abort,
live_clone.as_ref(),
&session_dir,
&workspaces,
@@ -390,13 +424,26 @@ pub fn execute_primitive(
ScriptPrimitive::Pipeline(scripts) => {
// Sequential: each stage runs only after the previous completes.
//
// Abort is checked between stages so the user can cancel the
// pipeline immediately when moving to the next division, rather
// than having to wait for the current subagent to finish.
//
// Why: parallel execution defeats the purpose of a pipeline whose
// stages are supposed to build on each other's output. Findings
// written by stage N are visible to stage N+1 through the shared
// `findings` Arc (same isolation scope as parent).
let mut all = Vec::new();
for (idx, script) in scripts.iter().enumerate() {
match execute_primitive(script, args, concurrency_cap, continue_on_error, live, session_dir, workspaces, findings, timeout_ms) {
// Check abort before each pipeline stage so we don't
// launch the next division after the user cancelled.
if abort_flag.as_ref().is_some_and(|f| f.load(Ordering::SeqCst)) {
if continue_on_error {
all.push(format!("pipeline aborted at stage {idx}"));
break;
}
anyhow::bail!("pipeline aborted by user at stage {idx}");
}
match execute_primitive(script, args, concurrency_cap, continue_on_error, abort_flag, live, session_dir, workspaces, findings, timeout_ms) {
Ok(outputs) => all.extend(outputs),
Err(e) => {
if continue_on_error {
@@ -411,7 +458,7 @@ pub fn execute_primitive(
}
ScriptPrimitive::Phase { name: _name, script } => {
execute_primitive(script, args, concurrency_cap, continue_on_error, live, session_dir, workspaces, findings, timeout_ms)
execute_primitive(script, args, concurrency_cap, continue_on_error, abort_flag, live, session_dir, workspaces, findings, timeout_ms)
}
}
}
@@ -426,7 +473,7 @@ pub fn run_workflow(
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
) -> anyhow::Result<String> {
run_workflow_tracked(script, args, None, session_dir, workspaces)
run_workflow_tracked(script, args, &None, None, session_dir, workspaces)
}
/// Run a `WorkflowScript` with real-time live-state callbacks so the TUI
@@ -444,6 +491,7 @@ pub fn run_workflow(
pub fn run_workflow_tracked(
script: &WorkflowScript,
args: &HashMap<String, String>,
abort_flag: &Option<Arc<AtomicBool>>,
live: Option<&LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
@@ -457,7 +505,7 @@ pub fn run_workflow_tracked(
let findings = Arc::new(Mutex::new(Vec::new()));
let results = execute_primitive(
&script.script, args, concurrency_cap,
script.options.continue_on_error, live,
script.options.continue_on_error, abort_flag, live,
session_dir, workspaces, &findings,
script.options.timeout_ms,
)?;
src/tool/spawn.rs
+4
View File
@@ -106,11 +106,13 @@ impl Tool for SpawnAgents {
// by this tool call are isolated from any other concurrent
// spawn_agents or workflow_run invocations.
let findings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
let results = crate::app::workflow::engine::execute_primitive(
&wf.script,
&HashMap::new(),
max_concurrency,
true,
&no_abort,
live.as_ref(),
&ctx.session_dir,
&ctx.workspaces,
@@ -195,11 +197,13 @@ impl Tool for SpawnPipeline {
// Per-invocation findings scope isolates this pipeline from any
// other concurrent spawn_agents / spawn_pipeline / workflow_run.
let findings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
let results = crate::app::workflow::engine::execute_primitive(
&wf.script,
&HashMap::new(),
1,
false,
&no_abort,
live.as_ref(),
&ctx.session_dir,
&ctx.workspaces,
src/tool/workflow.rs
+3
View File
@@ -186,6 +186,7 @@ impl Tool for CompanyPipeline {
.and_then(|v| v.as_str())
.unwrap_or("full");
let no_abort: Option<std::sync::Arc<std::sync::atomic::AtomicBool>> = None;
match mode {
"quick" => {
crate::app::workflow::company::run_company_pipeline_quick(
@@ -193,6 +194,7 @@ impl Tool for CompanyPipeline {
&ctx.session_dir,
&ctx.workspaces,
ctx.turn_events.as_ref(),
&no_abort,
)
}
_ => {
@@ -201,6 +203,7 @@ impl Tool for CompanyPipeline {
&ctx.session_dir,
&ctx.workspaces,
ctx.turn_events.as_ref(),
&no_abort,
)
}
}