From 104b0daf4cc51581de04f03d6b727cdb16f9b6c3 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 13 Jul 2026 09:09:23 +0700 Subject: [PATCH] feat: implement abort mechanism for workflows and subagents --- src/app/runtime/actions/mod.rs | 16 ++++- src/app/workflow/company.rs | 7 ++- src/app/workflow/engine.rs | 106 ++++++++++++++++++++++++--------- src/tool/spawn.rs | 4 ++ src/tool/workflow.rs | 3 + 5 files changed, 104 insertions(+), 32 deletions(-) diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index af20489..bce9b5f 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -682,8 +682,9 @@ pub fn apply_action(state: &mut AppStateRest, action: Action) { }); let args: HashMap = HashMap::new(); + let no_abort: Option> = 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; diff --git a/src/app/workflow/company.rs b/src/app/workflow/company.rs index 3daf1e5..fb9a876 100644 --- a/src/app/workflow/company.rs +++ b/src/app/workflow/company.rs @@ -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>>>, + abort_flag: &Option>, ) -> anyhow::Result { let divisions = division::all_divisions(); let mut pipeline_scripts: Vec = 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>>>, + abort_flag: &Option>, ) -> anyhow::Result { 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() diff --git a/src/app/workflow/engine.rs b/src/app/workflow/engine.rs index 903ed7f..9357164 100644 --- a/src/app/workflow/engine.rs +++ b/src/app/workflow/engine.rs @@ -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>>, + abort_flag: &Option>, 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::>(); - 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::>(); + 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, concurrency_cap: usize, continue_on_error: bool, + abort_flag: &Option>, 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::(); - 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 { - 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, + abort_flag: &Option>, 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, )?; diff --git a/src/tool/spawn.rs b/src/tool/spawn.rs index 4dd5618..aa57ebd 100644 --- a/src/tool/spawn.rs +++ b/src/tool/spawn.rs @@ -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>> = Arc::new(Mutex::new(Vec::new())); + let no_abort: Option> = 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>> = Arc::new(Mutex::new(Vec::new())); + let no_abort: Option> = 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, diff --git a/src/tool/workflow.rs b/src/tool/workflow.rs index 9505111..4d44e28 100644 --- a/src/tool/workflow.rs +++ b/src/tool/workflow.rs @@ -186,6 +186,7 @@ impl Tool for CompanyPipeline { .and_then(|v| v.as_str()) .unwrap_or("full"); + let no_abort: Option> = 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, ) } }