Files
zesdex/src/app/workflow/engine.rs
T

494 lines
18 KiB
Rust

//! Workflow engine: interprets `ScriptPrimitive` values (agent, parallel,
//! pipeline, phase) by spawning subagents, collecting results, and
//! managing concurrency.
//!
//! Key design points:
//! - `Parallel` branches run concurrently (capped by semaphore) — this is
//! the main advantage over single-turn chat.
//! - `Pipeline` branches run sequentially so each stage sees findings from
//! the previous one.
//! - `run_workflow_tracked` accepts a `LiveState` callback that receives
//! real-time agent status updates for the TUI panel.
//! - Findings (inter-agent notes) are scoped per invocation via an
//! `Arc<Mutex<Vec<String>>>` threaded through `execute_primitive` and
//! `spawn_single_agent` rather than a global static, preventing data
//! leaks between concurrent workflow runs.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::script::{ScriptPrimitive, WorkflowScript};
/// The lifecycle state of an agent within a workflow run.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AgentState {
Idle,
Running,
Completed,
Failed,
}
/// Timestamped status of one workflow agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStatus {
pub state: AgentState,
pub started_at: Option<i64>,
pub completed_at: Option<i64>,
pub error: Option<String>,
}
/// A single agent tracked within a workflow run.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowAgent {
pub id: String,
pub name: String,
pub status: AgentStatus,
}
/// Orchestrator for running workflow scripts: holds agent roster and a
/// shared finding accumulator visible to all pipeline stages.
#[derive(Debug, Clone)]
pub struct WorkflowEngine {
pub agents: Vec<WorkflowAgent>,
pub findings: Vec<String>,
}
impl WorkflowEngine {
/// Create an empty workflow engine with no agents or findings.
pub fn new() -> Self {
WorkflowEngine {
agents: Vec::new(),
findings: Vec::new(),
}
}
}
/// Shared live state used by `run_workflow_tracked` to push real-time
/// agent status updates into the TUI's `WorkflowEngine`.
///
/// The closure receives `(agent_id, new_status)` and should update the
/// corresponding agent in `AppStateRest::workflow_engine`.
pub type LiveStateFn = Arc<dyn Fn(String, AgentStatus) + Send + Sync>;
/// Spawn a single synchronous subagent with the given prompt, passing it
/// any findings from earlier sibling agents. Updates live state before and
/// after to reflect Running → Completed/Failed transitions.
///
/// Flow: push agent as `Running` → build SubagentContext with prompt +
/// findings preamble, linking the `workflow_findings` Arc so the subagent's
/// `note_finding` tool pushes into the same vec → call `run_subagent`
/// (draining the event channel into a consumer so events are not blocked)
/// → push `Completed` or `Failed`.
///
/// Why: the `workflow_findings` Arc is shared by all agents within the same
/// `execute_primitive` scope, so pipeline stages can pass data between each
/// other while different workflow invocations remain isolated.
///
/// When `timeout_ms` is `Some`, the subagent is killed (abandoned on a
/// separate thread) if it does not complete within the deadline, preventing
/// a stuck stage from blocking the entire pipeline forever.
///
/// Return: the agent's text output, or an error on failure.
fn spawn_single_agent(
agent_id: &str,
agent_name: &str,
prompt: &str,
findings_snapshot: Vec<String>,
findings: &Arc<Mutex<Vec<String>>>,
live: Option<&LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
timeout_ms: Option<u64>,
) -> anyhow::Result<String> {
use crate::app::subagent::context::build_subagent_context;
use crate::app::subagent::engine::run_subagent;
use crate::app::subagent::spawn::AgentDefinition;
let started_at = chrono::Utc::now().timestamp_millis();
// Notify UI: this agent is now running
if let Some(f) = live {
f(agent_id.to_string(), AgentStatus {
state: AgentState::Running,
started_at: Some(started_at),
completed_at: None,
error: None,
});
}
let def = AgentDefinition::new(agent_name.to_string(), "coder".to_string())
.with_max_steps(50);
let mut ctx = build_subagent_context(def);
ctx.session_dir = session_dir.to_path_buf();
ctx.workspaces = workspaces.to_vec();
let findings_section = if findings_snapshot.is_empty() {
String::new()
} else {
format!(
"\n\nFindings from sibling agents in this workflow run:\n{}",
findings_snapshot
.iter()
.enumerate()
.map(|(i, f)| format!("{}. {}", i + 1, f))
.collect::<Vec<_>>()
.join("\n")
)
};
ctx.system_prompt = format!("{}{}", prompt, findings_section);
// 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());
// Create an mpsc channel and drain events in a background thread so
// run_subagent's blocking_send never blocks (previously the _rx was
// dropped immediately, which would cause blocking_send to panic/fail
// on a closed channel).
let (tx, rx) = tokio::sync::mpsc::channel(64);
let _drain_thread = std::thread::spawn(move || {
// Drain all events so run_subagent's blocking_send never blocks.
// Individual SubagentEvent items are not surfaced to the TUI —
// the live state callbacks above handle coarse-grained Running /
// Completed / Failed status. ToolCall / ToolResult / StepCompleted
// events are traced at debug level for observability.
use crate::app::subagent::event::SubagentEvent;
let mut rx = rx;
while let Some(event) = rx.blocking_recv() {
match &event {
SubagentEvent::ToolCall { _tool, _args } => {
tracing::debug!("[subagent] tool call: {}", _tool);
}
SubagentEvent::ToolResult { _tool, .. } => {
tracing::debug!("[subagent] tool result: {}", _tool);
}
SubagentEvent::StepCompleted { _step, .. } => {
tracing::trace!("[subagent] step {} completed", _step);
}
SubagentEvent::StepFailed { _step, _error } => {
tracing::warn!("[subagent] step {} failed: {}", _step, _error);
}
SubagentEvent::Completed { .. } => {
tracing::debug!("[subagent] completed");
}
}
}
});
// 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 '{}' timed out after {}ms",
agent_name, timeout,
)),
}
} else {
run_subagent(ctx, tx)
};
let completed_at = chrono::Utc::now().timestamp_millis();
// Notify UI: agent completed or failed
if let Some(f) = live {
match &result {
Ok(_) => f(agent_id.to_string(), AgentStatus {
state: AgentState::Completed,
started_at: Some(started_at),
completed_at: Some(completed_at),
error: None,
}),
Err(e) => f(agent_id.to_string(), AgentStatus {
state: AgentState::Failed,
started_at: Some(started_at),
completed_at: Some(completed_at),
error: Some(e.to_string()),
}),
}
}
result
}
type ParallelResult = (usize, anyhow::Result<Vec<String>>);
/// Recursively execute a `ScriptPrimitive` tree, respecting an overall
/// concurrency cap for parallel branches.
///
/// Flow: match the primitive →
/// `Agent` → `spawn_single_agent`
/// `Parallel` → spawn threads up to `concurrency_cap` (semaphore-gated),
/// collect results in submission order
/// `Pipeline` → execute stages sequentially; findings flow between stages
/// `Phase` → recurse (pass-through wrapper)
///
/// Why: `Parallel` uses OS threads + a semaphore so the main async event
/// loop remains responsive. `Pipeline` is sequential so each stage sees
/// findings deposited by the previous one. Findings are scoped to an
/// `Arc<Mutex<Vec<String>>>` rather than a global static, so concurrent
/// workflow runs are isolated from each other.
///
/// `timeout_ms` propagates to individual agents so that no single agent
/// can block the entire workflow beyond the configured deadline.
///
/// Return: a `Vec<String>` of all agent outputs (or error strings) in
/// the order they were submitted.
pub fn execute_primitive(
primitive: &ScriptPrimitive,
args: &HashMap<String, String>,
concurrency_cap: usize,
continue_on_error: bool,
live: Option<&LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
findings: &Arc<Mutex<Vec<String>>>,
timeout_ms: Option<u64>,
) -> anyhow::Result<Vec<String>> {
match primitive {
ScriptPrimitive::Agent(prompt) => {
let resolved = resolve_template(prompt, args);
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) {
Ok(text) => Ok(vec![text]),
Err(e) => {
if continue_on_error {
Ok(vec![format!("agent error: {}", e)])
} else {
Err(e)
}
}
}
}
ScriptPrimitive::Parallel(scripts) => {
// All branches run concurrently, capped by semaphore.
// This is the primary advantage over single-turn chat: multiple
// independent subagents work simultaneously.
// Each branch shares the same `findings` Arc so note_finding
// calls within any branch are visible to all other branches.
let semaphore = Arc::new(Semaphore::new(concurrency_cap.max(1)));
let results: Arc<Mutex<Vec<ParallelResult>>> =
Arc::new(Mutex::new(Vec::new()));
let handles: Vec<_> = scripts
.iter()
.enumerate()
.map(|(idx, script)| {
let script = script.clone();
let args = args.clone();
let sem = Arc::clone(&semaphore);
let results = Arc::clone(&results);
let cap = concurrency_cap;
let live_clone = live.cloned();
let session_dir = session_dir.to_path_buf();
let workspaces = workspaces.to_vec();
let findings = Arc::clone(findings);
let to = timeout_ms;
std::thread::spawn(move || {
let _permit = sem.acquire();
let result = execute_primitive(
&script, &args, cap, continue_on_error,
live_clone.as_ref(),
&session_dir,
&workspaces,
&findings,
to,
);
if let Ok(mut locked) = results.lock() {
locked.push((idx, result));
}
})
})
.collect();
for handle in handles {
let _ = handle.join();
}
let mut locked = results.lock().map_err(|_| anyhow::anyhow!("parallel results lock poisoned"))?;
locked.sort_by_key(|(idx, _)| *idx);
let mut all = Vec::new();
for (_, res) in locked.drain(..) {
match res {
Ok(outputs) => all.extend(outputs),
Err(e) => all.push(format!("agent error: {}", e)),
}
}
Ok(all)
}
ScriptPrimitive::Pipeline(scripts) => {
// Sequential: each stage runs only after the previous completes.
//
// 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) {
Ok(outputs) => all.extend(outputs),
Err(e) => {
if continue_on_error {
all.push(format!("pipeline stage {} error: {}", idx, e));
} else {
return Err(e);
}
}
}
}
Ok(all)
}
ScriptPrimitive::Phase { name: _name, script } => {
execute_primitive(script, args, concurrency_cap, continue_on_error, live, session_dir, workspaces, findings, timeout_ms)
}
}
}
/// Run a `WorkflowScript` with the given template arguments and produce a
/// summary string. Uses no live-state callback.
///
/// Return: a human-readable summary string.
pub fn run_workflow(
script: &WorkflowScript,
args: &HashMap<String, String>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
) -> anyhow::Result<String> {
run_workflow_tracked(script, args, None, session_dir, workspaces)
}
/// Run a `WorkflowScript` with real-time live-state callbacks so the TUI
/// panel updates as each agent transitions between Idle/Running/Done/Failed.
///
/// Flow: create an empty findings Arc (scoped to this invocation) → cap
/// concurrency to 8 → call `execute_primitive` with the live callback and
/// findings → format results.
///
/// Why: findings are scoped to an `Arc<Mutex<Vec<String>>>` rather than a
/// global static, so concurrent `run_workflow_tracked` calls from different
/// spawn_agents invocations remain fully isolated.
///
/// Return: a human-readable summary string.
pub fn run_workflow_tracked(
script: &WorkflowScript,
args: &HashMap<String, String>,
live: Option<LiveStateFn>,
session_dir: &std::path::Path,
workspaces: &[std::path::PathBuf],
) -> anyhow::Result<String> {
let concurrency_cap = if script.options.max_concurrency > 0 {
script.options.max_concurrency.min(10) // allow up to 10 parallel agents
} else {
10
};
let live_ref = live.as_ref();
let findings = Arc::new(Mutex::new(Vec::new()));
let results = execute_primitive(
&script.script, args, concurrency_cap,
script.options.continue_on_error, live_ref,
session_dir, workspaces, &findings,
script.options.timeout_ms,
)?;
let summary = if results.is_empty() {
"workflow completed with no output".to_string()
} else {
format!(
"workflow '{}' completed. {} agent result(s):\n{}",
script.name,
results.len(),
results
.iter()
.enumerate()
.map(|(i, r)| format!("[{}] {}", i + 1, r.lines().next().unwrap_or(r)))
.collect::<Vec<_>>()
.join("\n")
)
};
Ok(summary)
}
/// Simple template engine: replace `{{key}}` placeholders with values
/// from `args`.
///
/// Why: a structured template engine is unnecessary for the limited
/// use-case; this is intentionally simple and safe.
fn resolve_template(template: &str, args: &HashMap<String, String>) -> String {
let mut result = template.to_string();
for (key, value) in args {
result = result.replace(&format!("{{{{{}}}}}", key), value);
}
result
}
/// A counting semaphore built from a `Mutex` + `Condvar`.
///
/// Used by `execute_primitive` to cap concurrent parallel branches.
///
/// Panic-safety: if a thread panics while holding a permit, the Mutex
/// becomes poisoned. Both `acquire` and the `Drop` implementation recover
/// from poisoned mutexes by discarding the poison, ensuring the semaphore
/// remains usable after a thread panic.
struct Semaphore {
count: Mutex<usize>,
condvar: std::sync::Condvar,
}
impl Semaphore {
fn new(count: usize) -> Self {
Semaphore {
count: Mutex::new(count),
condvar: std::sync::Condvar::new(),
}
}
fn acquire(&self) -> SemaphoreGuard<'_> {
let mut count = self.count.lock().unwrap_or_else(|e| {
tracing::warn!("[semaphore] mutex poisoned in acquire, recovering");
e.into_inner()
});
while *count == 0 {
count = self.condvar.wait(count).unwrap_or_else(|e| {
tracing::warn!("[semaphore] mutex poisoned in wait, recovering");
e.into_inner()
});
}
*count -= 1;
SemaphoreGuard { sem: self }
}
}
struct SemaphoreGuard<'a> {
sem: &'a Semaphore,
}
impl<'a> Drop for SemaphoreGuard<'a> {
fn drop(&mut self) {
let mut count = self.sem.count.lock().unwrap_or_else(|e| {
tracing::warn!("[semaphore] mutex poisoned in drop, recovering");
e.into_inner()
});
*count += 1;
self.sem.condvar.notify_one();
}
}