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/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,
)?;