diff --git a/Cargo.lock b/Cargo.lock index 2917f2f5..2dd52ee2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -777,6 +777,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -837,6 +846,29 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1005,6 +1037,15 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_users" version = "0.5.3" @@ -1191,6 +1232,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "serde" version = "1.0.229" @@ -1500,17 +1547,21 @@ dependencies = [ "chrono", "dotenvy", "futures", + "parking_lot", "serde", "serde_json", + "tempfile", "thiserror", "tinyagents-definition", "tinyagents-graph", "tinyagents-harness", "tinyagents-registry", "tinyagents-runtime", + "tinyagents-session", "tinyinference-llm", "tinytools", "tokio", + "uuid", ] [[package]] diff --git a/crates/tinyagents-harness/src/lib.rs b/crates/tinyagents-harness/src/lib.rs index 447f0eef..d1f1c918 100644 --- a/crates/tinyagents-harness/src/lib.rs +++ b/crates/tinyagents-harness/src/lib.rs @@ -96,6 +96,7 @@ pub mod token_estimation; pub mod tool; #[cfg(feature = "builtin-tools")] pub mod tools; +pub mod workspace; /// Re-exported vendor crates. Downstream consumers should reach these /// dependencies' types through these re-exports (e.g. diff --git a/crates/tinyagents-harness/src/workspace/README.md b/crates/tinyagents-harness/src/workspace/README.md new file mode 100644 index 00000000..ecabd2ef --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/README.md @@ -0,0 +1,73 @@ +# harness::workspace + +Workspace isolation and sandbox hooks for tools that run over real files or +command executors. + +## Why this exists + +Application-specific worktree/sandbox providers need a common seam: a +`tinytools::WorkspaceDescriptor` tells a tool which filesystem root it may +touch, and a [`WorkspaceIsolation`] provider trait prepares/tears down a +per-agent environment. TinyAgents does not own any concrete isolation policy +— it owns the interface, plus two providers: a trivial single-shared-root +default and a real git-worktree-backed implementation. + +## Public surface + +- [`WorkspaceIsolation`] — the provider trait: `prepare(run_id, agent) -> + WorkspaceDescriptor` and `cleanup(&descriptor)`. +- [`prepare_workspace`] / [`cleanup_workspace`] — free functions that drive a + `WorkspaceIsolation` provider *and* emit the corresponding + `AgentEvent::WorkspacePrepared` / `AgentEvent::WorkspaceCleanup` on the + run's event sink, so isolation setup/teardown is observable. +- [`SharedRootWorkspace`] — the trivial provider: scopes every agent to one + shared root without per-agent copying (`prepare` is a descriptor + construction, `cleanup` is a no-op). A sensible default and a test double. +- [`enforce_workspace_path`] — the fail-closed path gate a tool calls + *before* touching a path: emits `AgentEvent::WorkspaceViolation` and returns + a validation error when the path is outside every allowed root. +- From `git` (git-worktree-backed isolation): + - [`GitWorktreeIsolation`] — a real `WorkspaceIsolation` implementation: + each run gets its own `git worktree` checkout under + `/.claude/worktrees/`. Builder methods: + `with_base_ref`, `with_sandbox`, `with_trusted_root`. + - [`GitWorktreeBaseRef`] — which ref a new worktree branches from (`Head` or + `Fresh`, i.e. the repo's default branch). + - [`GitWorktreeStatus`] — a snapshot (`path`, `branch`, `is_dirty`, + `changed_files`) of one worktree's state. + - [`GitWorktreeError`] — errors from the worktree manager (`NotAGitRepo`, + `DirtyRefused`, `GitFailed`, `Io`). + - Standalone functions usable without the `WorkspaceIsolation` trait: + `create_git_worktree`, `list_git_worktrees`, `git_worktree_status`, + `git_worktree_diff_summary`, `remove_git_worktree`, + `detect_worktree_overlaps` (finds files touched by more than one sibling + worker — useful for merge-conflict-avoidance heuristics). + - `GIT_WORKTREE_SUBDIR` — the fixed subdirectory (`.claude/worktrees`) + worktrees are created under. + +## Files + +| File | Role | +| --- | --- | +| `mod.rs` | `WorkspaceIsolation` driving helpers (`prepare_workspace`/`cleanup_workspace`) and `SharedRootWorkspace`. | +| `types.rs` | The `WorkspaceIsolation` trait. | +| `policy.rs` | `enforce_workspace_path`, the fail-closed path gate (paired with the descriptor's lexical `allows` check, which lives in `tinytools`). | +| `git.rs` | `GitWorktreeIsolation` and the standalone git-worktree management functions. | +| `git/test.rs` | Tests for the git-worktree isolation provider (spawns real `git` subprocesses against a temp repo). | +| `test.rs` | Tests for the descriptor-allows and `SharedRootWorkspace`/event-emission hooks. | + +## Operational constraints + +- `WorkspaceDescriptor`'s allowed-root policy and its lexical `allows` check + live in `tinytools`, which owns the tool vocabulary; this module only + supplies the event-emitting wrapper (`enforce_workspace_path`) and the + isolation providers. +- `GitWorktreeIsolation::cleanup` refuses to remove a dirty worktree unless + forced — `remove_git_worktree`'s `force` parameter is `false` in the + `WorkspaceIsolation::cleanup` path, so uncommitted work is never silently + discarded by the isolation lifecycle. +- Run ids are sanitized (`sanitize_run_id`) into a filesystem- and + git-ref-safe slug before being used in a worktree path or branch name. +- All git operations shell out to the `git` binary via `std::process::Command` + — there is no `libgit2`/`gix` dependency. `git.rs`'s private `git`/`git_raw` + helpers are the only two call sites that invoke it. diff --git a/crates/tinyagents-harness/src/workspace/git.rs b/crates/tinyagents-harness/src/workspace/git.rs new file mode 100644 index 00000000..37d208b1 --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/git.rs @@ -0,0 +1,444 @@ +//! Git-worktree backed workspace isolation. +//! +//! This provider gives each agent run its own checkout under +//! `/.claude/worktrees/`, so parallel edit-capable workers can +//! operate without sharing one mutable filesystem root. It intentionally owns +//! only generic git/workspace behavior; host applications remain responsible +//! for product-specific audit, cleanup policy, and merge UX. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::workspace::WorkspaceIsolation; +use crate::{Result, TinyAgentsError}; +use tinytools::SandboxMode; +use tinytools::WorkspaceDescriptor; + +/// Directory, relative to the repository root, where isolated worktrees are +/// created. +pub const GIT_WORKTREE_SUBDIR: &str = ".claude/worktrees"; + +/// Which ref a new worktree branches from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GitWorktreeBaseRef { + /// Branch off the repo's current `HEAD`. + Head, + /// Branch off the repository's default branch (`origin/HEAD`, then local + /// `HEAD`, then `main`). + Fresh, +} + +impl GitWorktreeBaseRef { + /// Parses a user/config string. Unknown or empty values default to + /// [`GitWorktreeBaseRef::Head`]. + pub fn parse(value: Option<&str>) -> Self { + match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() { + Some("fresh") => Self::Fresh, + _ => Self::Head, + } + } + + /// Stable lowercase label for logs and policy ids. + pub fn as_str(self) -> &'static str { + match self { + Self::Head => "head", + Self::Fresh => "fresh", + } + } +} + +/// Snapshot of a single git worktree's state. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct GitWorktreeStatus { + /// Absolute path to the worktree checkout. + pub path: PathBuf, + /// Checked-out branch, or `(detached HEAD)` for detached worktrees. + pub branch: Option, + /// Whether staged, unstaged, or untracked changes are present. + pub is_dirty: bool, + /// Changed files relative to the worktree root. + pub changed_files: Vec, +} + +/// [`WorkspaceIsolation`] implementation backed by `git worktree`. +#[derive(Debug, Clone)] +pub struct GitWorktreeIsolation { + repo_root: PathBuf, + base_ref: GitWorktreeBaseRef, + sandbox: SandboxMode, + trusted_roots: Vec, +} + +impl GitWorktreeIsolation { + /// Creates an isolation provider rooted at a git repository. + pub fn new(repo_root: impl Into) -> Self { + Self { + repo_root: repo_root.into(), + base_ref: GitWorktreeBaseRef::Head, + sandbox: SandboxMode::Inherit, + trusted_roots: Vec::new(), + } + } + + /// Selects which ref newly prepared worktrees branch from. + pub fn with_base_ref(mut self, base_ref: GitWorktreeBaseRef) -> Self { + self.base_ref = base_ref; + self + } + + /// Advertises the sandbox expectation on prepared descriptors. + pub fn with_sandbox(mut self, sandbox: SandboxMode) -> Self { + self.sandbox = sandbox; + self + } + + /// Adds an extra root tools may touch alongside the isolated checkout. + pub fn with_trusted_root(mut self, root: impl Into) -> Self { + self.trusted_roots.push(root.into()); + self + } +} + +#[async_trait] +impl WorkspaceIsolation for GitWorktreeIsolation { + async fn prepare(&self, run_id: &str, agent: Option<&str>) -> Result { + let status = create_git_worktree(&self.repo_root, run_id, self.base_ref) + .map_err(|err| TinyAgentsError::Tool(err.to_string()))?; + let policy_id = match agent { + Some(agent) if !agent.is_empty() => format!("git.worktree:{agent}:{run_id}"), + _ => format!("git.worktree:{run_id}"), + }; + let mut descriptor = WorkspaceDescriptor::new(status.path) + .with_policy_id(policy_id) + .with_sandbox(self.sandbox); + for root in &self.trusted_roots { + descriptor = descriptor.with_trusted_root(root.clone()); + } + Ok(descriptor) + } + + async fn cleanup(&self, descriptor: &WorkspaceDescriptor) -> Result<()> { + remove_git_worktree(&self.repo_root, &descriptor.root, false) + .map_err(|err| TinyAgentsError::Tool(err.to_string())) + } +} + +/// Errors surfaced by the git worktree manager. +#[derive(Debug, thiserror::Error)] +pub enum GitWorktreeError { + /// The supplied path is not inside a git work tree. + #[error("path is not inside a git repository: {0}")] + NotAGitRepo(PathBuf), + + /// Dirty worktrees are not removed unless `force = true`. + #[error("worktree is dirty and force=false; refusing to remove: {0}")] + DirtyRefused(PathBuf), + + /// A git command exited unsuccessfully. + #[error("git command `{command}` failed: {stderr}")] + GitFailed { command: String, stderr: String }, + + /// Spawning git or creating the worktree parent directory failed. + #[error("io error running git: {0}")] + Io(#[from] std::io::Error), +} + +/// Shorthand for a result carrying [`GitWorktreeError`]. +type GitResult = std::result::Result; + +/// Creates an isolated worktree for `run_id` and returns its status snapshot. +pub fn create_git_worktree( + repo_root: &Path, + run_id: &str, + base_ref: GitWorktreeBaseRef, +) -> GitResult { + let repo_top = validate_repo_root(repo_root)?; + let run_slug = sanitize_run_id(run_id); + let worktree_path = repo_top.join(GIT_WORKTREE_SUBDIR).join(&run_slug); + let branch = format!("worker/{run_slug}"); + let base = match base_ref { + GitWorktreeBaseRef::Head => "HEAD".to_string(), + GitWorktreeBaseRef::Fresh => resolve_fresh_base(&repo_top), + }; + + if let Some(parent) = worktree_path.parent() { + std::fs::create_dir_all(parent)?; + } + + let worktree = worktree_path.to_string_lossy().to_string(); + git( + &repo_top, + &["worktree", "add", "-b", &branch, &worktree, &base], + )?; + match git_worktree_status(&repo_top, &worktree_path) { + Ok(status) => Ok(status), + Err(error) => { + // Do not leave a registered checkout behind when validation of a + // newly-created worktree fails. + let _ = git(&repo_top, &["worktree", "remove", "--force", &worktree]); + let _ = git(&repo_top, &["branch", "-D", &branch]); + Err(error) + } + } +} + +/// Lists worktrees registered on the repository at `repo_root`. +pub fn list_git_worktrees(repo_root: &Path) -> GitResult> { + let repo_top = validate_repo_root(repo_root)?; + let porcelain = git(&repo_top, &["worktree", "list", "--porcelain"])?; + let mut out = Vec::new(); + let mut cur_path: Option = None; + let mut cur_branch: Option = None; + + let mut flush = |path: &mut Option, branch: &mut Option| -> GitResult<()> { + if let Some(path) = path.take() { + let (is_dirty, changed_files) = dirty_state(&path)?; + out.push(GitWorktreeStatus { + path, + branch: branch.take(), + is_dirty, + changed_files, + }); + } else { + *branch = None; + } + Ok(()) + }; + + for line in porcelain.lines() { + if let Some(rest) = line.strip_prefix("worktree ") { + flush(&mut cur_path, &mut cur_branch)?; + cur_path = Some(PathBuf::from(rest.trim())); + } else if let Some(rest) = line.strip_prefix("branch ") { + let trimmed = rest.trim(); + cur_branch = Some( + trimmed + .strip_prefix("refs/heads/") + .unwrap_or(trimmed) + .to_string(), + ); + } else if line.trim() == "detached" { + cur_branch = Some("(detached HEAD)".to_string()); + } + } + flush(&mut cur_path, &mut cur_branch)?; + Ok(out) +} + +/// Returns branch, dirty, and changed-file status for one worktree. +pub fn git_worktree_status(repo_root: &Path, worktree_path: &Path) -> GitResult { + let repo_top = validate_repo_root(repo_root)?; + let worktree_path = if worktree_path.is_absolute() { + worktree_path.to_path_buf() + } else { + repo_top.join(worktree_path) + }; + if !worktree_path.exists() { + return Err(GitWorktreeError::NotAGitRepo(worktree_path)); + } + let branch = git(&worktree_path, &["rev-parse", "--abbrev-ref", "HEAD"]) + .ok() + .map(|branch| { + if branch == "HEAD" { + "(detached HEAD)".to_string() + } else { + branch + } + }); + let (is_dirty, changed_files) = dirty_state(&worktree_path)?; + Ok(GitWorktreeStatus { + path: worktree_path, + branch, + is_dirty, + changed_files, + }) +} + +/// Human-readable diff stat of working changes vs `HEAD`, including untracked +/// files. +pub fn git_worktree_diff_summary(repo_root: &Path, worktree_path: &Path) -> GitResult { + validate_repo_root(repo_root)?; + let stat = git(worktree_path, &["diff", "HEAD", "--stat"])?; + let untracked = git( + worktree_path, + &["ls-files", "--others", "--exclude-standard"], + )?; + let mut parts = Vec::new(); + if !stat.is_empty() { + parts.push(stat); + } + if !untracked.is_empty() { + parts.push( + untracked + .lines() + .map(|line| format!(" {line} (untracked)")) + .collect::>() + .join("\n"), + ); + } + Ok(parts.join("\n")) +} + +/// Removes a worktree. Dirty worktrees are refused unless `force = true`. +pub fn remove_git_worktree(repo_root: &Path, worktree_path: &Path, force: bool) -> GitResult<()> { + let repo_top = validate_repo_root(repo_root)?; + let worktree_path = if worktree_path.is_absolute() { + worktree_path.to_path_buf() + } else { + repo_top.join(worktree_path) + }; + let (is_dirty, _) = dirty_state(&worktree_path)?; + if is_dirty && !force { + return Err(GitWorktreeError::DirtyRefused(worktree_path)); + } + + let worktree = worktree_path.to_string_lossy().to_string(); + let mut args = vec!["worktree", "remove", &worktree]; + if force { + args.push("--force"); + } + git(&repo_top, &args)?; + Ok(()) +} + +/// Detects changed files touched by more than one sibling worker. +pub fn detect_worktree_overlaps( + per_worker: &[(String, Vec)], +) -> std::collections::BTreeMap> { + use std::collections::{BTreeMap, BTreeSet}; + + let mut by_file: BTreeMap> = BTreeMap::new(); + for (worker_id, files) in per_worker { + let mut seen = BTreeSet::new(); + for file in files { + if seen.insert(file.clone()) { + by_file + .entry(file.clone()) + .or_default() + .push(worker_id.clone()); + } + } + } + + by_file + .into_iter() + .filter_map(|(file, mut workers)| { + workers.sort(); + workers.dedup(); + (workers.len() > 1).then_some((file, workers)) + }) + .collect() +} + +/// Runs `git` in `cwd` with `args`, returning trimmed stdout on success or a +/// [`GitWorktreeError::GitFailed`] carrying stderr on a non-zero exit. +fn git(cwd: &Path, args: &[&str]) -> GitResult { + let output = Command::new("git").current_dir(cwd).args(args).output()?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + return Err(GitWorktreeError::GitFailed { + command: format!("git {}", args.join(" ")), + stderr, + }); + } + Ok(String::from_utf8_lossy(&output.stdout).trim().to_string()) +} + +/// Like [`git`], but returns raw (untrimmed) stdout — used where leading +/// column characters (e.g. `git status --porcelain` status codes) are +/// significant. +fn git_raw(cwd: &Path, args: &[&str]) -> GitResult { + let output = Command::new("git").current_dir(cwd).args(args).output()?; + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + return Err(GitWorktreeError::GitFailed { + command: format!("git {}", args.join(" ")), + stderr, + }); + } + Ok(String::from_utf8_lossy(&output.stdout).to_string()) +} + +/// Confirms `repo_root` is inside a git work tree and returns the repo's +/// top-level directory (`git rev-parse --show-toplevel`), which every +/// worktree operation anchors on regardless of which subdirectory of the repo +/// `repo_root` names. +fn validate_repo_root(repo_root: &Path) -> GitResult { + if !repo_root.exists() { + return Err(GitWorktreeError::NotAGitRepo(repo_root.to_path_buf())); + } + let inside = git(repo_root, &["rev-parse", "--is-inside-work-tree"]) + .map_err(|_| GitWorktreeError::NotAGitRepo(repo_root.to_path_buf()))?; + if inside.trim() != "true" { + return Err(GitWorktreeError::NotAGitRepo(repo_root.to_path_buf())); + } + let top = git(repo_root, &["rev-parse", "--show-toplevel"]) + .map_err(|_| GitWorktreeError::NotAGitRepo(repo_root.to_path_buf()))?; + Ok(PathBuf::from(top.trim())) +} + +/// Resolves the ref [`GitWorktreeBaseRef::Fresh`] branches from: the remote's +/// default branch (`origin/HEAD`) if set, else the local `HEAD`'s symbolic +/// name, else the literal `"main"` as a last resort. +fn resolve_fresh_base(repo_top: &Path) -> String { + if let Ok(sym) = git( + repo_top, + &["symbolic-ref", "--short", "refs/remotes/origin/HEAD"], + ) && !sym.is_empty() + { + return sym; + } + if let Ok(head) = git(repo_top, &["symbolic-ref", "--short", "HEAD"]) + && !head.is_empty() + { + return head; + } + "main".to_string() +} + +/// Parses `git status --porcelain` for `worktree_path` into a dirty flag plus +/// the sorted, deduplicated list of changed paths (rename entries `a -> b` are +/// reduced to `b`). +fn dirty_state(worktree_path: &Path) -> GitResult<(bool, Vec)> { + let porcelain = git_raw(worktree_path, &["status", "--porcelain"])?; + let mut changed = Vec::new(); + for line in porcelain.lines() { + if line.len() > 3 { + let path = line[3..].trim_end(); + let path = path.rsplit(" -> ").next().unwrap_or(path); + changed.push(PathBuf::from(path)); + } + } + changed.sort(); + changed.dedup(); + Ok((!changed.is_empty(), changed)) +} + +/// Maps `run_id` to a filesystem- and git-ref-safe slug: non-alphanumeric +/// characters (other than `_`/`-`) become `-`, leading/trailing `-` is +/// trimmed, and an empty result falls back to `"worker"`. +fn sanitize_run_id(run_id: &str) -> String { + let cleaned: String = run_id + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' { + ch + } else { + '-' + } + }) + .collect(); + let trimmed = cleaned.trim_matches('-'); + if trimmed.is_empty() { + "worker".to_string() + } else { + trimmed.to_string() + } +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/workspace/git/test.rs b/crates/tinyagents-harness/src/workspace/git/test.rs new file mode 100644 index 00000000..15de3487 --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/git/test.rs @@ -0,0 +1,258 @@ +//! Tests for the git-worktree isolation provider. + +use super::*; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +static TEMP_DIR_COUNTER: AtomicU64 = AtomicU64::new(0); + +struct TempDir { + path: PathBuf, +} + +impl TempDir { + fn new() -> Self { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock") + .as_nanos(); + let counter = TEMP_DIR_COUNTER.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "tinyagents-worktree-test-{}-{nanos}-{counter}", + std::process::id() + )); + std::fs::create_dir_all(&path).expect("create temp dir"); + Self { path } + } + + fn path(&self) -> &Path { + &self.path + } +} + +impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } +} + +fn git_available() -> bool { + Command::new("git") + .arg("--version") + .output() + .map(|output| output.status.success()) + .unwrap_or(false) +} + +fn run(dir: &Path, args: &[&str]) { + let output = Command::new("git") + .current_dir(dir) + .args(args) + .output() + .expect("git invocation"); + assert!( + output.status.success(), + "git {:?} failed: {}", + args, + String::from_utf8_lossy(&output.stderr) + ); +} + +fn init_repo() -> (TempDir, PathBuf) { + let tmp = TempDir::new(); + let root = tmp.path().to_path_buf(); + run(&root, &["init", "-b", "main"]); + run(&root, &["config", "user.email", "test@example.com"]); + run(&root, &["config", "user.name", "Test User"]); + std::fs::write(root.join("README.md"), "hello\n").unwrap(); + run(&root, &["add", "README.md"]); + run(&root, &["commit", "-m", "initial"]); + (tmp, root) +} + +#[test] +fn validate_repo_root_rejects_non_repo() { + if !git_available() { + return; + } + let tmp = TempDir::new(); + let err = create_git_worktree(tmp.path(), "run-1", GitWorktreeBaseRef::Head).unwrap_err(); + assert!(matches!(err, GitWorktreeError::NotAGitRepo(_))); +} + +#[test] +fn create_then_status_reports_clean_worktree() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = + create_git_worktree(&root, "run-1", GitWorktreeBaseRef::Head).expect("create worktree"); + assert!(status.path.exists()); + assert_eq!(status.branch.as_deref(), Some("worker/run-1")); + assert!(!status.is_dirty); + assert!(status.changed_files.is_empty()); + assert!(status.path.ends_with(Path::new(".claude/worktrees/run-1"))); +} + +#[test] +fn list_includes_created_worktrees() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + create_git_worktree(&root, "run-a", GitWorktreeBaseRef::Head).expect("create a"); + create_git_worktree(&root, "run-b", GitWorktreeBaseRef::Fresh).expect("create b"); + + let all = list_git_worktrees(&root).expect("list worktrees"); + assert!(all.len() >= 3, "expected main + two worktrees, got {all:?}"); + let branches: Vec<_> = all + .iter() + .filter_map(|worktree| worktree.branch.clone()) + .collect(); + assert!(branches.iter().any(|branch| branch == "worker/run-a")); + assert!(branches.iter().any(|branch| branch == "worker/run-b")); +} + +#[test] +fn status_detects_dirty_changes() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = create_git_worktree(&root, "run-dirty", GitWorktreeBaseRef::Head).expect("create"); + std::fs::write(status.path.join("README.md"), "changed\n").unwrap(); + std::fs::write(status.path.join("new.txt"), "fresh\n").unwrap(); + + let status = git_worktree_status(&root, &status.path).expect("status"); + assert!(status.is_dirty); + let names: Vec<_> = status + .changed_files + .iter() + .map(|path| path.to_string_lossy().to_string()) + .collect(); + assert!(names.iter().any(|name| name.contains("README.md"))); + assert!(names.iter().any(|name| name.contains("new.txt"))); +} + +#[test] +fn diff_summary_lists_tracked_and_untracked() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = create_git_worktree(&root, "run-diff", GitWorktreeBaseRef::Head).expect("create"); + std::fs::write(status.path.join("README.md"), "changed body\n").unwrap(); + std::fs::write(status.path.join("brand_new.txt"), "x\n").unwrap(); + + let summary = git_worktree_diff_summary(&root, &status.path).expect("diff"); + assert!(summary.contains("README.md")); + assert!(summary.contains("brand_new.txt") && summary.contains("untracked")); +} + +#[test] +fn remove_refuses_dirty_without_force() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = create_git_worktree(&root, "run-keep", GitWorktreeBaseRef::Head).expect("create"); + std::fs::write(status.path.join("README.md"), "dirty\n").unwrap(); + + let err = remove_git_worktree(&root, &status.path, false).expect_err("must refuse dirty"); + assert!(matches!(err, GitWorktreeError::DirtyRefused(_))); + assert!(status.path.exists()); +} + +#[test] +fn remove_force_deletes_dirty_worktree() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = create_git_worktree(&root, "run-force", GitWorktreeBaseRef::Head).expect("create"); + std::fs::write(status.path.join("README.md"), "dirty\n").unwrap(); + + remove_git_worktree(&root, &status.path, true).expect("force remove"); + assert!(!status.path.exists()); +} + +#[test] +fn remove_clean_worktree_succeeds() { + if !git_available() { + return; + } + let (_tmp, root) = init_repo(); + let status = create_git_worktree(&root, "run-clean", GitWorktreeBaseRef::Head).expect("create"); + + remove_git_worktree(&root, &status.path, false).expect("clean remove"); + assert!(!status.path.exists()); +} + +#[test] +fn base_ref_parse_defaults_to_head() { + assert_eq!(GitWorktreeBaseRef::parse(None), GitWorktreeBaseRef::Head); + assert_eq!( + GitWorktreeBaseRef::parse(Some("head")), + GitWorktreeBaseRef::Head + ); + assert_eq!( + GitWorktreeBaseRef::parse(Some(" Fresh ")), + GitWorktreeBaseRef::Fresh + ); + assert_eq!( + GitWorktreeBaseRef::parse(Some("garbage")), + GitWorktreeBaseRef::Head + ); +} + +#[test] +fn sanitize_run_id_strips_unsafe_chars() { + assert_eq!(sanitize_run_id("sub-1234"), "sub-1234"); + assert_eq!(sanitize_run_id("a/b\\c"), "a-b-c"); + assert_eq!(sanitize_run_id("///"), "worker"); + assert_eq!(sanitize_run_id(""), "worker"); +} + +#[test] +fn detect_overlaps_flags_shared_files() { + let per_worker = vec![ + ( + "w1".to_string(), + vec![PathBuf::from("src/a.rs"), PathBuf::from("src/b.rs")], + ), + ( + "w2".to_string(), + vec![PathBuf::from("src/b.rs"), PathBuf::from("src/c.rs")], + ), + ("w3".to_string(), vec![PathBuf::from("src/c.rs")]), + ]; + + let overlaps = detect_worktree_overlaps(&per_worker); + assert_eq!(overlaps.len(), 2); + assert_eq!( + overlaps.get(&PathBuf::from("src/b.rs")).unwrap(), + &vec!["w1".to_string(), "w2".to_string()] + ); + assert_eq!( + overlaps.get(&PathBuf::from("src/c.rs")).unwrap(), + &vec!["w2".to_string(), "w3".to_string()] + ); +} + +#[test] +fn detect_overlaps_empty_when_disjoint_or_duplicate_within_one_worker() { + let disjoint = vec![ + ("w1".to_string(), vec![PathBuf::from("a.rs")]), + ("w2".to_string(), vec![PathBuf::from("b.rs")]), + ]; + assert!(detect_worktree_overlaps(&disjoint).is_empty()); + + let duplicate = vec![( + "w1".to_string(), + vec![PathBuf::from("a.rs"), PathBuf::from("a.rs")], + )]; + assert!(detect_worktree_overlaps(&duplicate).is_empty()); +} diff --git a/crates/tinyagents-harness/src/workspace/mod.rs b/crates/tinyagents-harness/src/workspace/mod.rs new file mode 100644 index 00000000..cbf422de --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/mod.rs @@ -0,0 +1,108 @@ +//! Workspace isolation and sandbox hooks for tools that run over real files or +//! command executors. +//! +//! See [`tinytools::WorkspaceDescriptor`] for the allowed-root policy a tool +//! reads from its execution context) and the [`WorkspaceIsolation`] provider +//! trait (per-agent environment preparation/cleanup). This module ships one +//! trivial provider, [`SharedRootWorkspace`], which scopes every agent to a +//! single shared root without copying — a sensible default and a test double. +//! Application-specific worktree/sandbox providers implement +//! [`WorkspaceIsolation`] themselves. + +mod git; +mod policy; +mod types; + +pub use git::*; +pub use policy::enforce_workspace_path; +pub use types::*; + +use std::path::PathBuf; + +use async_trait::async_trait; + +use crate::Result; +use crate::events::{AgentEvent, EventSink}; +use tinytools::{SandboxMode, WorkspaceDescriptor}; + +/// Prepares a per-agent environment through `isolation` and emits an +/// [`AgentEvent::WorkspacePrepared`] on the run's event sink so late observers +/// and journals record the isolation setup. Returns the descriptor to thread +/// into the run via [`RunContext::with_workspace`][crate::context::RunContext::with_workspace]. +/// +/// A preparation failure is propagated to the caller (there is no partial +/// environment to clean up); the paired teardown is +/// [`cleanup_workspace`]. +pub async fn prepare_workspace( + isolation: &dyn WorkspaceIsolation, + events: &EventSink, + run_id: &str, + agent: Option<&str>, +) -> Result { + let descriptor = isolation.prepare(run_id, agent).await?; + events.emit(AgentEvent::WorkspacePrepared { + policy_id: descriptor.policy_id.clone(), + root: descriptor.root.display().to_string(), + }); + Ok(descriptor) +} + +/// Tears down a previously prepared environment through `isolation` and emits an +/// [`AgentEvent::WorkspaceCleanup`] (with `error` set when cleanup fails) so the +/// teardown is observable. The cleanup result is returned unchanged. +pub async fn cleanup_workspace( + isolation: &dyn WorkspaceIsolation, + events: &EventSink, + descriptor: &WorkspaceDescriptor, +) -> Result<()> { + let result = isolation.cleanup(descriptor).await; + events.emit(AgentEvent::WorkspaceCleanup { + policy_id: descriptor.policy_id.clone(), + error: result.as_ref().err().map(|e| e.to_string()), + }); + result +} + +/// A [`WorkspaceIsolation`] provider that scopes every agent to one shared root +/// without creating per-agent copies. +/// +/// `prepare` returns a descriptor rooted at the shared directory (tagged with +/// the run id as the policy identity) and `cleanup` is a no-op, since nothing +/// per-agent was created. +#[derive(Clone, Debug)] +pub struct SharedRootWorkspace { + root: PathBuf, + sandbox: SandboxMode, +} + +impl SharedRootWorkspace { + /// Creates a provider scoping agents to `root`. + pub fn new(root: impl Into) -> Self { + Self { + root: root.into(), + sandbox: SandboxMode::Inherit, + } + } + + /// Sets the sandbox mode advertised on prepared descriptors. + pub fn with_sandbox(mut self, sandbox: SandboxMode) -> Self { + self.sandbox = sandbox; + self + } +} + +#[async_trait] +impl WorkspaceIsolation for SharedRootWorkspace { + async fn prepare(&self, run_id: &str, _agent: Option<&str>) -> Result { + Ok(WorkspaceDescriptor::new(self.root.clone()) + .with_policy_id(run_id) + .with_sandbox(self.sandbox)) + } + + async fn cleanup(&self, _descriptor: &WorkspaceDescriptor) -> Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/workspace/policy.rs b/crates/tinyagents-harness/src/workspace/policy.rs new file mode 100644 index 00000000..c6652c37 --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/policy.rs @@ -0,0 +1,76 @@ +//! The fail-closed path gate for a [`WorkspaceDescriptor`]. +//! +//! The descriptor and its lexical `allows` check live in `tinytools`, which +//! owns the tool vocabulary. What stays here is the half that needs this +//! crate: emitting a [`WorkspaceViolation`][crate::events::AgentEvent::WorkspaceViolation] +//! and returning this crate's error type. It is a free function rather than an +//! inherent method because the descriptor is now a foreign type. + +use std::path::{Path, PathBuf}; + +use tinytools::WorkspaceDescriptor; + +use crate::Result; +use crate::events::{AgentEvent, EventSink}; + +/// Fail-closed path gate to call *before* a tool touches `path`. +/// +/// When the path is outside every allowed root, emits a +/// [`AgentEvent::WorkspaceViolation`] on `events` and returns a validation +/// error so the caller blocks the operation. Returns `Ok(())` when the path is +/// allowed. +/// +/// # Errors +/// +/// Returns [`TinyAgentsError::Validation`][crate::error::TinyAgentsError::Validation] +/// when `path` lies outside the descriptor's root and trusted roots. +pub fn enforce_workspace_path( + workspace: &WorkspaceDescriptor, + path: &Path, + events: &EventSink, +) -> Result<()> { + if workspace.allows(path) && resolved_path_is_allowed(workspace, path) { + return Ok(()); + } + let rendered = path.display().to_string(); + events.emit(AgentEvent::WorkspaceViolation { + path: rendered.clone(), + }); + Err(crate::error::TinyAgentsError::Validation(format!( + "path `{rendered}` is outside the allowed workspace roots" + ))) +} + +/// Resolves existing filesystem components before checking containment so an +/// in-workspace symlink cannot redirect a tool to an untrusted root. For a new +/// path, its nearest existing ancestor is resolved and the remaining components +/// are appended; the eventual open must still use no-follow semantics where a +/// host supports them to close the final TOCTOU window. +fn resolved_path_is_allowed(workspace: &WorkspaceDescriptor, path: &Path) -> bool { + let candidate = canonicalize_with_missing_tail(path); + let Some(candidate) = candidate else { + return false; + }; + std::iter::once(&workspace.root) + .chain(workspace.trusted_roots.iter()) + .filter_map(|root| std::fs::canonicalize(root).ok()) + .any(|root| candidate.starts_with(root)) +} + +fn canonicalize_with_missing_tail(path: &Path) -> Option { + let mut candidate = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir().ok()?.join(path) + }; + let mut missing = Vec::new(); + while !candidate.exists() { + missing.push(candidate.file_name()?.to_os_string()); + candidate = candidate.parent()?.to_path_buf(); + } + let mut resolved = std::fs::canonicalize(candidate).ok()?; + for component in missing.iter().rev() { + resolved.push(component); + } + Some(resolved) +} diff --git a/crates/tinyagents-harness/src/workspace/test.rs b/crates/tinyagents-harness/src/workspace/test.rs new file mode 100644 index 00000000..bc076916 --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/test.rs @@ -0,0 +1,125 @@ +//! Tests for workspace isolation hooks. + +use super::*; +use std::path::Path; +use tinytools::SandboxMode; + +#[test] +fn descriptor_allows_paths_under_root_and_trusted_roots() { + let ws = WorkspaceDescriptor::new("/work/agent-a") + .with_trusted_root("/shared/cache") + .with_sandbox(SandboxMode::Required) + .with_policy_id("run-1"); + + assert!(ws.allows(Path::new("/work/agent-a/src/main.rs"))); + assert!(ws.allows(Path::new("/shared/cache/blob"))); + // Outside every root. + assert!(!ws.allows(Path::new("/etc/passwd"))); + // Escape attempt via `..` is normalized and rejected. + assert!(!ws.allows(Path::new("/work/agent-a/../agent-b/secret"))); + assert_eq!(ws.sandbox, SandboxMode::Required); + assert_eq!(ws.policy_id, "run-1"); +} + +#[test] +fn relative_root_rejects_leading_parent_sibling_escape() { + // Relative root: a candidate that walks up past the root and re-descends + // into a same-named sibling must be rejected, not admitted by a leading + // `..` collapsing back onto the root name. + let ws = WorkspaceDescriptor::new("ws"); + + assert!(ws.allows(Path::new("ws/src/main.rs"))); + // `ws/../../ws/secret` resolves outside the anchored `ws` root. + assert!(!ws.allows(Path::new("ws/../../ws/secret"))); + // A bare leading `..` escape is rejected. + assert!(!ws.allows(Path::new("../ws/secret"))); + assert!(!ws.allows(Path::new("../evil"))); +} + +#[tokio::test] +async fn shared_root_workspace_prepares_and_cleans_up() { + let provider = SharedRootWorkspace::new("/work").with_sandbox(SandboxMode::Disabled); + let descriptor = provider.prepare("run-42", Some("worker")).await.unwrap(); + + assert_eq!(descriptor.root, std::path::PathBuf::from("/work")); + assert_eq!(descriptor.policy_id, "run-42"); + assert_eq!(descriptor.sandbox, SandboxMode::Disabled); + assert!(descriptor.allows(Path::new("/work/output.txt"))); + + // Cleanup is a no-op for a shared root. + provider.cleanup(&descriptor).await.unwrap(); +} + +#[tokio::test] +async fn prepare_and_cleanup_helpers_emit_lifecycle_events() { + use crate::events::{AgentEvent, EventSink, RecordingListener}; + use std::sync::Arc; + + let events = EventSink::new(); + let recorder = Arc::new(RecordingListener::new()); + events.subscribe(recorder.clone()); + + let provider = SharedRootWorkspace::new("/work"); + let descriptor = prepare_workspace(&provider, &events, "run-7", Some("worker")) + .await + .unwrap(); + cleanup_workspace(&provider, &events, &descriptor) + .await + .unwrap(); + + let kinds: Vec<_> = recorder.events().iter().map(|r| r.event.kind()).collect(); + assert_eq!(kinds, vec!["workspace.prepared", "workspace.cleanup"]); + // The prepared event carries the policy id and root for audit. + match &recorder.events()[0].event { + AgentEvent::WorkspacePrepared { policy_id, root } => { + assert_eq!(policy_id, "run-7"); + assert_eq!(root, "/work"); + } + other => panic!("expected WorkspacePrepared, got {other:?}"), + } +} + +#[test] +fn enforce_blocks_unsafe_paths_and_emits_violation() { + use crate::events::{EventSink, RecordingListener}; + use std::sync::Arc; + + let events = EventSink::new(); + let recorder = Arc::new(RecordingListener::new()); + events.subscribe(recorder.clone()); + + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("agent-a"); + std::fs::create_dir_all(&root).unwrap(); + let ws = WorkspaceDescriptor::new(&root); + // Allowed path passes silently with no event. + enforce_workspace_path(&ws, &root.join("out.txt"), &events).unwrap(); + assert!(recorder.is_empty()); + + // Unsafe path fails closed and emits a violation. + let err = enforce_workspace_path(&ws, Path::new("/etc/passwd"), &events) + .expect_err("path outside root must be blocked"); + assert!(err.to_string().contains("outside the allowed workspace")); + assert_eq!(recorder.events()[0].event.kind(), "workspace.violation"); +} + +#[test] +fn run_context_workspace_threads_into_tool_execution_context() { + use crate::CallId; + use crate::context::{RunConfig, RunContext}; + use crate::tool::ToolExecutionContext; + + let ws = WorkspaceDescriptor::new("/work/agent-a").with_policy_id("run-9"); + let ctx: RunContext = RunContext::new(RunConfig::new("run-9"), ()).with_workspace(ws.clone()); + + let tool_ctx = ToolExecutionContext::from_run_context(&ctx, CallId::new("call-9")); + assert_eq!(tool_ctx.workspace, Some(ws)); +} + +#[test] +fn descriptor_serializes_for_audit() { + let ws = WorkspaceDescriptor::new("/work").with_policy_id("p1"); + let json = serde_json::to_value(&ws).unwrap(); + let back: WorkspaceDescriptor = serde_json::from_value(json).unwrap(); + assert_eq!(ws, back); +} diff --git a/crates/tinyagents-harness/src/workspace/types.rs b/crates/tinyagents-harness/src/workspace/types.rs new file mode 100644 index 00000000..7090ca12 --- /dev/null +++ b/crates/tinyagents-harness/src/workspace/types.rs @@ -0,0 +1,31 @@ +//! Workspace isolation and sandbox types. +//! +//! These are the SDK-owned, application-policy-neutral hooks agents use when +//! their tools run over real files or command executors: a +//! [`tinytools::WorkspaceDescriptor`] tells a tool which filesystem root it may touch, and +//! a [`WorkspaceIsolation`] provider prepares and tears down per-agent +//! worktrees/sandboxes. TinyAgents does not own any concrete policy; it owns the +//! interface so parallel agents can be isolated consistently. + +use async_trait::async_trait; + +use crate::Result; + +/// Prepares and tears down per-agent execution environments. +/// +/// Implementations create a worktree/sandbox for one agent run and clean it up +/// afterward. The returned [`WorkspaceDescriptor`](tinytools::WorkspaceDescriptor) is what the run threads into +/// tool execution contexts. +#[async_trait] +pub trait WorkspaceIsolation: Send + Sync { + /// Prepares an environment for `run_id` (optionally on behalf of a named + /// `agent`). + async fn prepare( + &self, + run_id: &str, + agent: Option<&str>, + ) -> Result; + + /// Cleans up a previously prepared environment. + async fn cleanup(&self, descriptor: &tinytools::WorkspaceDescriptor) -> Result<()>; +} diff --git a/crates/tinyagents-orchestration/Cargo.toml b/crates/tinyagents-orchestration/Cargo.toml index 29f9b332..1157dc5b 100644 --- a/crates/tinyagents-orchestration/Cargo.toml +++ b/crates/tinyagents-orchestration/Cargo.toml @@ -15,11 +15,15 @@ chrono = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } thiserror = "2" +parking_lot = "0.12" +tinyagents-graph = { path = "../tinyagents-graph", version = "2.1.2", default-features = false } tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2", default-features = false } +tinyagents-session = { path = "../tinyagents-session", version = "2.1.2", default-features = false } tinyagents-runtime = { path = "../tinyagents-runtime", version = "2.1.2", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.3.0" } tokio = { workspace = true, features = ["sync", "rt", "macros"] } +uuid = { version = "1", features = ["v4"] } [features] default = [] @@ -30,6 +34,7 @@ tracing = [ [dev-dependencies] dotenvy = "0.15" futures = { workspace = true } +tempfile = "3" tinyagents-definition = { path = "../tinyagents-definition", version = "2.1.2" } tinyagents-graph = { path = "../tinyagents-graph", version = "2.1.2" } tinyagents-registry = { path = "../tinyagents-registry", version = "2.1.2" } diff --git a/crates/tinyagents-orchestration/src/lib.rs b/crates/tinyagents-orchestration/src/lib.rs index ccaa5b01..3b2c9f3f 100644 --- a/crates/tinyagents-orchestration/src/lib.rs +++ b/crates/tinyagents-orchestration/src/lib.rs @@ -10,6 +10,8 @@ //! depend on this composition layer. pub mod subagent; +pub mod teams; +pub mod workflow; #[cfg(test)] mod boundary_tests { diff --git a/crates/tinyagents-orchestration/src/teams/README.md b/crates/tinyagents-orchestration/src/teams/README.md new file mode 100644 index 00000000..60a91baa --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/README.md @@ -0,0 +1,67 @@ +# `orchestration::teams` — durable agent-team coordination + +A **team** is a durable group of worker agents (members) who claim and +complete collaborative tasks. The module enforces team structure invariants +(no duplicate member names, valid task dependencies) and manages the event +log that records every state change: members added, tasks created, claims, +completions, and inter-member messages. + +## Public surface + +- **`TeamService`** — the main API: creates teams, adds members, + creates tasks, handles claims and completions, and routes messages. Generic + over a durable ledger. +- **`TeamLedger`** — the trait for durable team state: CRUD operations for + teams, members, tasks, and events. Implementations handle all persistence. +- **`SessionTeamLedger`** — built-in ledger backed by `tinyagents-session`'s + run ledger at a caller-supplied workspace root. +- **`NewMember`, `TeamView`, `MemberShutdown`** — data types for snapshots and + operations. +- **`TeamError`** — validation errors: duplicate members, unknown members, + cyclic task dependencies, etc. +- **`run_member_graph`** — executes a member's work as a generic DAG (execute → + complete/fail → done), bridging host worker callbacks and the graph layer. + +## Design and invariants + +- **Durable state at every step:** Team members, tasks, events, and watermarks + are all durably persisted via the ledger. Hosts can use any storage backend + that implements [`TeamLedger`]. +- **Event-sourced messaging:** All team events (messages, member adds, task + completions) are logged in the durable run event stream. Members read their + undelivered messages and advance a delivery watermark — making repeated + reads idempotent and safe for retries. +- **Dependency-aware tasks:** Tasks form a directed acyclic graph with + dependencies. The service validates this at creation time (no cycles, no + self-dependencies, no unknown member/task references). +- **Atomic task claims:** Task claiming uses atomic compare-and-swap to ensure + exactly one member claims each task. Completion is similarly atomic, + advancing task status and optionally requiring evidence. +- **Watermark-based delivery:** Members have a message delivery watermark + (sequence number) that advances durably as messages are read. Hosts can call + message-delivery functions repeatedly without duplicating messages. + +## File map + +- **`service.rs`** — [`TeamService`] implementation: team/member/task CRUD, + validation, and coordination. [`TeamLedger`] trait definition. + [`SessionTeamLedger`] — session run-ledger adapter. +- **`types.rs`** — data types: [`NewMember`], [`TeamView`], [`MemberShutdown`], + [`TeamError`]. +- **`graph.rs`** — member worker graph: `execute` → `complete`/`fail` → `done` + DAG. Host supplies `run_worker`, `on_complete`, `on_failed` callbacks. + [`MemberOutcome`] routes on success/failure. +- **`runtime.rs`** — message delivery and prompt composition: + `deliver_pending_messages()` reads undelivered messages and advances + watermarks. `build_member_prompt()` composes the stable worker prompt from + task and delivered messages. +- **`tests.rs`** — tests for team creation, member lifecycle, task coordination, + concurrency, and message delivery. + +## Relationship to other modules + +- **Depends on:** `tinyagents-graph` (DAG execution), `tinyagents-session` (run + ledger), `tinyagents-harness` (error types). +- **Used by:** `orchestration::workflow` (to model multi-agent phases), + host code (for team management). +- **Integration:** The workflow engine can model a phase as a team of agents. diff --git a/crates/tinyagents-orchestration/src/teams/graph.rs b/crates/tinyagents-orchestration/src/teams/graph.rs new file mode 100644 index 00000000..d9097f13 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/graph.rs @@ -0,0 +1,174 @@ +//! Member worker graph execution: a generic three-step DAG that invokes a host +//! worker, routes on success/failure, and calls back to the host. +//! +//! This module owns the structure of member execution as seen by the graph +//! layer (entry → execute → complete/fail → done), while letting the host +//! supply the actual work (`run_worker`) and post-work effects (`on_complete` +//! and `on_failed`). It is the bridge between durable team coordination +//! (membership, tasks, event delivery) and the graph's lifecycle and event +//! streaming. + +use std::future::Future; +use std::sync::Arc; + +use anyhow::Result; +use tinyagents_graph::export::GraphTopology; +use tinyagents_graph::stream::GraphEventSink; +use tinyagents_graph::{ + ClosureStateReducer, Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, +}; + +/// Terminal classification of a host worker run. +/// +/// Returned by a host's `run_worker` callback to signal whether the member +/// completed its work successfully or failed. The member graph routes on this +/// outcome and calls the appropriate host callback (`on_complete` or +/// `on_failed`). +pub enum MemberOutcome { + /// Member completed successfully; carries the output to be recorded. + Completed { output: String }, + /// Member failed; carries the failure reason to be recorded. + Failed { reason: String }, +} + +#[derive(Clone, Default)] +struct MemberState { + payload: Option, +} + +enum MemberUpdate { + Payload(String), + Noop, +} + +fn graph_err(error: anyhow::Error) -> tinyagents_harness::TinyAgentsError { + tinyagents_harness::TinyAgentsError::Graph(error.to_string()) +} + +/// Run the generic complete-or-fail member graph with host supplied effects. +/// +/// `event_sink` is optional because observability belongs to the embedding host; +/// when supplied it receives the graph executor's lifecycle events unchanged. +pub async fn run_member_graph( + event_sink: Option>, + run_worker: W, + on_complete: C, + on_failed: F, +) -> Result<()> +where + W: Fn() -> WF + Clone + Send + Sync + 'static, + WF: Future> + Send + 'static, + C: Fn(String) -> CF + Clone + Send + Sync + 'static, + CF: Future> + Send + 'static, + F: Fn(String) -> FF + Clone + Send + Sync + 'static, + FF: Future> + Send + 'static, +{ + let mut graph = build_member_graph(run_worker, on_complete, on_failed)?; + if let Some(event_sink) = event_sink { + graph = graph.with_event_sink(event_sink); + } + graph + .run(MemberState::default()) + .await + .map_err(|error| anyhow::anyhow!("member graph run failed: {error}"))?; + Ok(()) +} + +fn build_member_graph( + run_worker: W, + on_complete: C, + on_failed: F, +) -> Result> +where + W: Fn() -> WF + Clone + Send + Sync + 'static, + WF: Future> + Send + 'static, + C: Fn(String) -> CF + Clone + Send + Sync + 'static, + CF: Future> + Send + 'static, + F: Fn(String) -> FF + Clone + Send + Sync + 'static, + FF: Future> + Send + 'static, +{ + let mut builder = GraphBuilder::::new().set_reducer( + ClosureStateReducer::new(|mut state: MemberState, update: MemberUpdate| { + if let MemberUpdate::Payload(payload) = update { + state.payload = Some(payload); + } + Ok(state) + }), + ); + builder = builder.add_node( + "execute", + move |_state: MemberState, _context: NodeContext| { + let run_worker = run_worker.clone(); + async move { + match run_worker().await.map_err(graph_err)? { + MemberOutcome::Completed { output } => Ok(NodeResult::Command( + Command::default() + .with_update(MemberUpdate::Payload(output)) + .with_goto(["complete"]), + )), + MemberOutcome::Failed { reason } => Ok(NodeResult::Command( + Command::default() + .with_update(MemberUpdate::Payload(reason)) + .with_goto(["fail"]), + )), + } + } + }, + ); + builder = builder.add_node( + "complete", + move |state: MemberState, _context: NodeContext| { + let on_complete = on_complete.clone(); + async move { + on_complete(state.payload.unwrap_or_default()) + .await + .map_err(graph_err)?; + Ok(NodeResult::Update(MemberUpdate::Noop)) + } + }, + ); + builder = builder.add_node("fail", move |state: MemberState, _context: NodeContext| { + let on_failed = on_failed.clone(); + async move { + on_failed(state.payload.unwrap_or_default()) + .await + .map_err(graph_err)?; + Ok(NodeResult::Update(MemberUpdate::Noop)) + } + }); + builder + .add_node( + "done", + |_state: MemberState, _context: NodeContext| async move { + Ok(NodeResult::Update(MemberUpdate::Noop)) + }, + ) + .add_edge("complete", "done") + .add_edge("fail", "done") + .set_entry("execute") + .mark_command_routing("execute") + .set_finish("done") + .compile() + .map_err(|error| anyhow::anyhow!("member graph compile failed: {error}")) +} + +/// Structure-only view of the generic member execution graph. +/// +/// Returns a topology that reflects the nodes and edges (entry, execute, +/// complete, fail, done, routing rules) without running any real worker or +/// calling effects. Used for introspection and documentation. +pub fn member_graph_topology() -> Result { + Ok(build_member_graph( + || async { + Ok(MemberOutcome::Completed { + output: String::new(), + }) + }, + |_| async { Ok(()) }, + |_| async { Ok(()) }, + )? + .topology()) +} + +#[cfg(test)] +mod tests; diff --git a/crates/tinyagents-orchestration/src/teams/graph/tests.rs b/crates/tinyagents-orchestration/src/teams/graph/tests.rs new file mode 100644 index 00000000..c0fdd55d --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/graph/tests.rs @@ -0,0 +1,100 @@ +//! Tests for member worker graph execution. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use tinyagents_graph::stream::CollectingSink; + +use super::*; + +#[tokio::test] +async fn member_graph_routes_completed_and_failed_workers() { + let complete = Arc::new(AtomicBool::new(false)); + let failed = Arc::new(AtomicBool::new(false)); + let complete_flag = complete.clone(); + let failed_flag = failed.clone(); + run_member_graph( + None, + || async { + Ok(MemberOutcome::Completed { + output: "done".into(), + }) + }, + move |output| { + let complete = complete_flag.clone(); + async move { + assert_eq!(output, "done"); + complete.store(true, Ordering::SeqCst); + Ok(()) + } + }, + move |_| { + let failed = failed_flag.clone(); + async move { + failed.store(true, Ordering::SeqCst); + Ok(()) + } + }, + ) + .await + .unwrap(); + assert!(complete.load(Ordering::SeqCst)); + assert!(!failed.load(Ordering::SeqCst)); + + let failed = Arc::new(AtomicBool::new(false)); + let failed_flag = failed.clone(); + run_member_graph( + None, + || async { + Ok(MemberOutcome::Failed { + reason: "boom".into(), + }) + }, + |_| async { Ok(()) }, + move |reason| { + let failed = failed_flag.clone(); + async move { + assert_eq!(reason, "boom"); + failed.store(true, Ordering::SeqCst); + Ok(()) + } + }, + ) + .await + .unwrap(); + assert!(failed.load(Ordering::SeqCst)); +} + +#[tokio::test] +async fn worker_engine_errors_propagate() { + let result = run_member_graph( + None, + || async { Err(anyhow::anyhow!("worker unavailable")) }, + |_| async { Ok(()) }, + |_| async { Ok(()) }, + ) + .await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn injected_event_sink_observes_member_graph_lifecycle() { + let sink = Arc::new(CollectingSink::new()); + run_member_graph( + Some(sink.clone()), + || async { + Ok(MemberOutcome::Completed { + output: "done".into(), + }) + }, + |_| async { Ok(()) }, + |_| async { Ok(()) }, + ) + .await + .unwrap(); + + assert!( + !sink.is_empty(), + "an injected sink must receive the member graph lifecycle" + ); +} diff --git a/crates/tinyagents-orchestration/src/teams/mod.rs b/crates/tinyagents-orchestration/src/teams/mod.rs new file mode 100644 index 00000000..8194f094 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/mod.rs @@ -0,0 +1,31 @@ +//! Durable, dependency-aware agent-team composition. +//! +//! A **team** is a durable group of worker agents (members) who claim and +//! complete collaborative tasks. This module owns team creation, member +//! lifecycle, task assignment and completion, and the event log that records +//! all durable state changes. +//! +//! [`TeamService`] is the primary API: it validates team structure, manages +//! member and task persistence (via the [`TeamLedger`] trait), and enforces +//! coordination invariants (no duplicate names, no cycles in task dependencies, +//! no dangling member or task references). `runtime` handles the per-member +//! details: reading undelivered messages from the event log and composing the +//! prompt a worker should receive. `graph` executes a member's work as a +//! generic execute → complete/fail → done DAG, bridging the graph layer and +//! durable team state. + +mod graph; +mod runtime; +mod service; +mod types; + +pub use graph::{MemberOutcome, member_graph_topology, run_member_graph}; +pub use runtime::{ + DeliveredMessages, EVENT_PAGE_SIZE, MESSAGE_DELIVERED_EVENT, TEAM_MESSAGE_EVENT, + build_member_prompt, deliver_pending_messages, drain_run_events, truncate_chars, +}; +pub use service::{SessionTeamLedger, TeamLedger, TeamService, claimable_task}; +pub use types::{LEAD_SENDER, MemberShutdown, NewMember, TeamError, TeamView}; + +#[cfg(test)] +mod tests; diff --git a/crates/tinyagents-orchestration/src/teams/runtime.rs b/crates/tinyagents-orchestration/src/teams/runtime.rs new file mode 100644 index 00000000..32ba92e1 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/runtime.rs @@ -0,0 +1,150 @@ +//! Host-neutral mechanics used when a durable team member starts work. +//! +//! Hosts decide when and how to run a worker. This module only reads the +//! durable event log, advances delivery watermarks, and produces the prompt +//! that a worker should receive at its start boundary. + +use anyhow::Result; +use serde_json::json; +use tinyagents_session::run_ledger::{ + AgentTeamTask, RunEvent, RunEventAppend, RunEventListRequest, +}; + +use super::TeamLedger; + +/// Event type used for durable lead and teammate messages. +pub const TEAM_MESSAGE_EVENT: &str = "team_message"; +/// Event type used to record that a member consumed messages through a sequence. +pub const MESSAGE_DELIVERED_EVENT: &str = "team_message_delivered"; +/// Maximum page accepted by the session run ledger. +pub const EVENT_PAGE_SIZE: u32 = 1_000; + +/// The result of reading one member's undelivered team messages. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeliveredMessages { + /// Message bodies in durable sequence order. + pub messages: Vec, + /// The highest delivered message sequence, when a watermark was written. + pub up_to_sequence: Option, +} + +/// Compose the stable worker prompt from the claimed task and delivered messages. +pub fn build_member_prompt(task: &AgentTeamTask, messages: &[String]) -> String { + let mut prompt = format!("You are a teammate on an agent team. Task: {}", task.title); + if let Some(objective) = task + .objective + .as_deref() + .filter(|value| !value.trim().is_empty()) + { + prompt.push_str("\n\nObjective:\n"); + prompt.push_str(objective.trim()); + } + if !messages.is_empty() { + prompt.push_str("\n\nMessages from your lead / teammates:\n"); + for message in messages { + prompt.push_str("- "); + prompt.push_str(message); + prompt.push('\n'); + } + } + prompt.push_str("\n\nComplete the task and report what you did."); + prompt +} + +/// Drain every durable event in ascending sequence order. +/// +/// The session ledger caps each request, so callers must use this rather than +/// assuming that one event-list request represents a complete team history. +pub fn drain_run_events(ledger: &L, team_id: &str) -> Result> { + let mut events = Vec::new(); + let mut after_sequence = None; + loop { + let page = ledger.list_events(&RunEventListRequest { + run_id: team_id.to_string(), + after_sequence, + limit: Some(EVENT_PAGE_SIZE), + })?; + let exhausted = page.len() < EVENT_PAGE_SIZE as usize; + after_sequence = page.last().map(|event| event.sequence); + events.extend(page); + if exhausted || after_sequence.is_none() { + return Ok(events); + } + } +} + +/// Read messages for a member, then durably advance that member's watermark. +/// +/// Direct messages and broadcasts are selected; messages for other members +/// stay untouched. The watermark is append-only, which makes repeated calls +/// idempotent while retaining the full event history for audit and replay. +pub fn deliver_pending_messages( + ledger: &L, + team_id: &str, + member_id: &str, +) -> Result { + let events = drain_run_events(ledger, team_id)?; + let watermark = events + .iter() + .filter(|event| event.event_type == MESSAGE_DELIVERED_EVENT) + .filter(|event| { + event + .payload + .get("memberId") + .and_then(|value| value.as_str()) + == Some(member_id) + }) + .filter_map(|event| { + event + .payload + .get("upToSeq") + .and_then(|value| value.as_u64()) + }) + .max() + .unwrap_or_default(); + + let mut up_to_sequence = watermark; + let mut messages = Vec::new(); + for event in &events { + if event.event_type != TEAM_MESSAGE_EVENT || event.sequence <= watermark { + continue; + } + let recipient = event.payload.get("to").and_then(|value| value.as_str()); + if recipient.is_none() || recipient == Some(member_id) { + if let Some(content) = event + .payload + .get("content") + .and_then(|value| value.as_str()) + { + messages.push(content.to_string()); + } + up_to_sequence = up_to_sequence.max(event.sequence); + } + } + + let delivered_up_to = (up_to_sequence > watermark).then_some(up_to_sequence); + if let Some(up_to_sequence) = delivered_up_to { + ledger.append_event(RunEventAppend { + run_id: team_id.to_string(), + event_type: MESSAGE_DELIVERED_EVENT.to_string(), + payload: json!({ "memberId": member_id, "upToSeq": up_to_sequence }), + })?; + } + Ok(DeliveredMessages { + messages, + up_to_sequence: delivered_up_to, + }) +} + +/// Truncate text at a Unicode character boundary and append an ellipsis. +pub fn truncate_chars(value: &str, max_chars: usize) -> String { + if value.chars().count() <= max_chars { + return value.to_string(); + } + let mut output: String = value.chars().take(max_chars).collect(); + output.push('…'); + output +} + +#[cfg(test)] +mod tests; diff --git a/crates/tinyagents-orchestration/src/teams/runtime/tests.rs b/crates/tinyagents-orchestration/src/teams/runtime/tests.rs new file mode 100644 index 00000000..7b4bde3c --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/runtime/tests.rs @@ -0,0 +1,84 @@ +//! Tests for member prompt composition and message delivery. + +use tempfile::TempDir; + +use super::*; +use crate::teams::{NewMember, SessionTeamLedger, TeamService}; + +fn team(dir: &TempDir) -> (SessionTeamLedger, String, String) { + let ledger = SessionTeamLedger::new(dir.path()); + let service = TeamService::new(ledger.clone()); + let team = service + .create_team( + "lead", + None, + None, + &[NewMember { + name: "member".into(), + agent_id: None, + }], + ) + .unwrap(); + (ledger, team.team.id, team.members[0].id.clone()) +} + +#[test] +fn delivery_selects_direct_and_broadcast_messages_once() { + let dir = TempDir::new().unwrap(); + let (ledger, team_id, member_id) = team(&dir); + let service = TeamService::new(ledger.clone()); + service + .message_member(&team_id, None, Some(&member_id), "direct", None) + .unwrap(); + service + .message_member(&team_id, None, None, "broadcast", None) + .unwrap(); + let first = deliver_pending_messages(&ledger, &team_id, &member_id).unwrap(); + assert_eq!(first.messages, ["direct", "broadcast"]); + assert!(first.up_to_sequence.is_some()); + assert!( + deliver_pending_messages(&ledger, &team_id, &member_id) + .unwrap() + .messages + .is_empty() + ); +} + +#[test] +fn delivery_pages_past_the_session_ledger_cap() { + let dir = TempDir::new().unwrap(); + let (ledger, team_id, member_id) = team(&dir); + for index in 0..=EVENT_PAGE_SIZE { + ledger + .append_event(tinyagents_session::run_ledger::RunEventAppend { + run_id: team_id.clone(), + event_type: "noise".into(), + payload: serde_json::json!({ "index": index }), + }) + .unwrap(); + } + let service = TeamService::new(ledger.clone()); + service + .message_member(&team_id, None, Some(&member_id), "late", None) + .unwrap(); + assert_eq!( + deliver_pending_messages(&ledger, &team_id, &member_id) + .unwrap() + .messages, + ["late"] + ); +} + +#[test] +fn prompt_and_truncation_preserve_text_boundaries() { + let dir = TempDir::new().unwrap(); + let (ledger, team_id, _) = team(&dir); + let service = TeamService::new(ledger); + let task = service + .assign_task(&team_id, "Ship", Some(" Build it "), None, &[]) + .unwrap(); + let prompt = build_member_prompt(&task, &["coordinate".into()]); + assert!(prompt.contains("Build it")); + assert!(prompt.contains("coordinate")); + assert_eq!(truncate_chars("aébc", 2), "aé…"); +} diff --git a/crates/tinyagents-orchestration/src/teams/service.rs b/crates/tinyagents-orchestration/src/teams/service.rs new file mode 100644 index 00000000..0a2cc453 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/service.rs @@ -0,0 +1,557 @@ +//! Team service and durable ledger: the public API for team management. +//! +//! [`TeamService`] validates team structure, manages member and task +//! persistence (via [`TeamLedger`]), and enforces coordination invariants. +//! [`SessionTeamLedger`] provides a built-in `tinyagents-session` backend; +//! hosts can supply their own ledger implementation for testing or custom +//! storage. + +use std::collections::HashSet; +use std::path::{Path, PathBuf}; + +use anyhow::{Result, anyhow}; +use chrono::Utc; +use serde_json::json; +use tinyagents_graph::dag::{DagNode, has_cycle}; +use tinyagents_session::run_ledger::{ + self, AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, + AgentTeamMemberStatus, AgentTeamMemberUpsert, AgentTeamStatus, AgentTeamTask, + AgentTeamTaskStatus, AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, CompletionOutcome, + RunEvent, RunEventAppend, RunEventListRequest, +}; +use uuid::Uuid; + +use super::runtime::drain_run_events; +use super::{LEAD_SENDER, MemberShutdown, NewMember, TEAM_MESSAGE_EVENT, TeamError, TeamView}; + +/// Durable team state required by [`TeamService`]. +/// +/// The trait intentionally mirrors the session run-ledger operations rather +/// than introducing another task store. Hosts can substitute a fake ledger in +/// tests or select their own durable implementation. +pub trait TeamLedger: Send + Sync { + /// Upserts a team row: creates if absent, updates if present. + fn upsert_team(&self, upsert: AgentTeamUpsert) -> Result; + /// Retrieves a team by id. + fn get_team(&self, id: &str) -> Result>; + /// Lists teams according to the request filters. + fn list_teams(&self, request: &AgentTeamListRequest) -> Result; + /// Upserts a team member: creates if absent, updates if present. + fn upsert_member(&self, upsert: AgentTeamMemberUpsert) -> Result; + /// Lists members of a team in creation order. + fn list_members(&self, team_id: &str) -> Result>; + /// Lists tasks assigned to a team in creation order. + fn list_tasks(&self, team_id: &str) -> Result>; + /// Upserts a task: creates if absent, updates if present. + fn upsert_task(&self, upsert: AgentTeamTaskUpsert) -> Result; + /// Attempts to claim a task for a member using an atomic CAS operation. + /// Returns the claim outcome (success, already claimed, not found). + fn claim_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + claim_token: &str, + ) -> Result; + /// Records a task completion with optional evidence, atomically advancing + /// the task status and optionally validating evidence before transition. + fn complete_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + evidence: &[String], + require_evidence: bool, + ) -> Result; + /// Stops a member and releases its claimed tasks, returning the stopped + /// member and the released task ids. + fn shutdown_member( + &self, + team_id: &str, + member_id: &str, + ) -> Result)>>; + /// Appends an event to the durable run event log. + fn append_event(&self, event: RunEventAppend) -> Result; + /// Lists events from the run event log according to the request filters. + fn list_events(&self, request: &RunEventListRequest) -> Result>; +} + +/// [`TeamLedger`] backed by `tinyagents-session`'s run ledger at a caller +/// supplied workspace root. It makes no workspace or host policy decision. +/// +/// Delegates all operations to the session run-ledger functions, projecting +/// the workspace path into each call. This is the default implementation when +/// using TinyAgents' built-in session storage. +#[derive(Debug, Clone)] +pub struct SessionTeamLedger { + workspace_dir: PathBuf, +} + +impl SessionTeamLedger { + /// Creates a ledger wrapping the session layer at the specified workspace. + pub fn new(workspace_dir: impl Into) -> Self { + Self { + workspace_dir: workspace_dir.into(), + } + } + + /// Returns the workspace directory used for all ledger operations. + pub fn workspace_dir(&self) -> &Path { + &self.workspace_dir + } +} + +impl TeamLedger for SessionTeamLedger { + fn upsert_team(&self, upsert: AgentTeamUpsert) -> Result { + Ok(run_ledger::upsert_agent_team(&self.workspace_dir, upsert)?) + } + fn get_team(&self, id: &str) -> Result> { + Ok(run_ledger::get_agent_team(&self.workspace_dir, id)?) + } + fn list_teams(&self, request: &AgentTeamListRequest) -> Result { + Ok(run_ledger::list_agent_teams(&self.workspace_dir, request)?) + } + fn upsert_member(&self, upsert: AgentTeamMemberUpsert) -> Result { + Ok(run_ledger::upsert_agent_team_member( + &self.workspace_dir, + upsert, + )?) + } + fn list_members(&self, team_id: &str) -> Result> { + Ok(run_ledger::list_agent_team_members( + &self.workspace_dir, + team_id, + )?) + } + fn list_tasks(&self, team_id: &str) -> Result> { + Ok(run_ledger::list_agent_team_tasks( + &self.workspace_dir, + team_id, + )?) + } + fn upsert_task(&self, upsert: AgentTeamTaskUpsert) -> Result { + Ok(run_ledger::upsert_agent_team_task( + &self.workspace_dir, + upsert, + )?) + } + fn claim_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + claim_token: &str, + ) -> Result { + Ok(run_ledger::claim_agent_team_task( + &self.workspace_dir, + team_id, + task_id, + member_id, + claim_token, + )?) + } + fn complete_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + evidence: &[String], + require_evidence: bool, + ) -> Result { + Ok(run_ledger::complete_agent_team_task( + &self.workspace_dir, + team_id, + task_id, + member_id, + evidence, + require_evidence, + )?) + } + fn shutdown_member( + &self, + team_id: &str, + member_id: &str, + ) -> Result)>> { + Ok(run_ledger::shutdown_agent_team_member( + &self.workspace_dir, + team_id, + member_id, + )?) + } + fn append_event(&self, event: RunEventAppend) -> Result { + Ok(run_ledger::append_run_event(&self.workspace_dir, event)?) + } + fn list_events(&self, request: &RunEventListRequest) -> Result> { + Ok(run_ledger::list_recent_run_events(&self.workspace_dir, request)?.events) + } +} + +/// Host-neutral service for durable, dependency-aware agent teams. +/// +/// Provides a high-level API for team management: creating teams, adding +/// members, creating and claiming tasks, composing prompts, and shutting down +/// members. All mutations are durably persisted via a caller-supplied +/// [`TeamLedger`]; the service enforces coordination invariants (no duplicate +/// member names, valid task dependencies, no cycles). +/// +/// Generic over the ledger to allow hosts to inject their own storage +/// implementation or a test double. +#[derive(Debug, Clone)] +pub struct TeamService { + ledger: L, +} + +impl TeamService { + /// Creates a service wrapping the provided [`TeamLedger`]. + pub fn new(ledger: L) -> Self { + Self { ledger } + } + + /// Returns a reference to the underlying ledger. + pub fn ledger(&self) -> &L { + &self.ledger + } +} + +impl TeamService { + pub fn create_team( + &self, + lead_agent_id: &str, + parent_thread_id: Option<&str>, + summary: Option<&str>, + members: &[NewMember], + ) -> Result { + let mut seen = HashSet::new(); + for member in members { + if !seen.insert(member.name.as_str()) { + return Err(anyhow!(TeamError::DuplicateMemberName { + name: member.name.clone() + })); + } + } + let team_id = format!("team-{}", Uuid::new_v4().simple()); + self.ledger.upsert_team(AgentTeamUpsert { + id: team_id.clone(), + parent_thread_id: parent_thread_id.map(str::to_string), + lead_agent_id: lead_agent_id.to_string(), + status: AgentTeamStatus::Active, + summary: summary.map(str::to_string), + created_at: None, + closed_at: None, + })?; + for member in members { + self.ledger.upsert_member(AgentTeamMemberUpsert { + id: format!("member-{}", Uuid::new_v4().simple()), + team_id: team_id.clone(), + name: member.name.clone(), + agent_id: member.agent_id.clone(), + member_status: AgentTeamMemberStatus::Pending, + current_task_id: None, + worker_thread_id: None, + run_id: None, + created_at: None, + })?; + } + self.team_view(&team_id) + } + + pub fn list_teams(&self, request: &AgentTeamListRequest) -> Result { + self.ledger.list_teams(request) + } + + pub fn get_team(&self, team_id: &str) -> Result> { + if self.ledger.get_team(team_id)?.is_some() { + self.team_view(team_id).map(Some) + } else { + Ok(None) + } + } + + pub fn assign_task( + &self, + team_id: &str, + title: &str, + objective: Option<&str>, + owner_member_id: Option<&str>, + depends_on: &[String], + ) -> Result { + let team = self + .ledger + .get_team(team_id)? + .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; + if team.status == AgentTeamStatus::Closed { + return Err(anyhow!("team is closed: {team_id}")); + } + let existing = self.ledger.list_tasks(team_id)?; + if let Some(owner) = owner_member_id + && !self.ledger.list_members(team_id)?.iter().any(|member| { + member.id == owner && member.member_status != AgentTeamMemberStatus::Stopped + }) + { + return Err(anyhow!(TeamError::UnknownMember { + member_id: owner.to_string() + })); + } + let task_id = format!("task-{}", Uuid::new_v4().simple()); + validate_dependencies(&task_id, depends_on, &existing)?; + self.ledger.upsert_task(AgentTeamTaskUpsert { + id: task_id, + team_id: team_id.to_string(), + title: title.to_string(), + objective: objective.map(str::to_string), + status: AgentTeamTaskStatus::Todo, + owner_member_id: owner_member_id.map(str::to_string), + depends_on: depends_on.to_vec(), + gate_status: None, + gate_reason: None, + evidence: vec![], + source_run_id: None, + order_index: existing.len() as i64, + created_at: None, + }) + } + + pub fn claim_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + claim_token: &str, + ) -> Result { + self.ensure_team_active(team_id)?; + self.ensure_member(team_id, member_id)?; + self.ledger + .claim_task(team_id, task_id, member_id, claim_token) + } + + pub fn message_member( + &self, + team_id: &str, + from_member_id: Option<&str>, + to_member_id: Option<&str>, + content: &str, + visibility: Option<&str>, + ) -> Result { + self.ledger + .get_team(team_id)? + .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; + if let Some(from) = from_member_id { + self.ensure_member(team_id, from)?; + } + if let Some(to) = to_member_id { + self.ensure_member(team_id, to)?; + } + self.ledger.append_event(RunEventAppend { + run_id: team_id.to_string(), + event_type: TEAM_MESSAGE_EVENT.to_string(), + payload: json!({"from": from_member_id.unwrap_or(LEAD_SENDER), "to": to_member_id, + "content": content, "visibility": visibility.unwrap_or("team")}), + }) + } + + pub fn list_messages(&self, team_id: &str, limit: Option) -> Result> { + Ok(drain_run_events(&self.ledger, team_id)? + .into_iter() + .filter(|event| event.event_type == TEAM_MESSAGE_EVENT) + .take(limit.unwrap_or(u32::MAX) as usize) + .collect()) + } + + pub fn complete_task( + &self, + team_id: &str, + task_id: &str, + member_id: &str, + evidence: &[String], + require_evidence: bool, + ) -> Result { + self.ensure_member(team_id, member_id)?; + self.ledger + .complete_task(team_id, task_id, member_id, evidence, require_evidence) + } + + pub fn shutdown_member(&self, team_id: &str, member_id: &str) -> Result { + self.ledger + .shutdown_member(team_id, member_id)? + .map(|(member, released_task_ids)| MemberShutdown { + member, + released_task_ids, + }) + .ok_or_else(|| { + anyhow!(TeamError::UnknownMember { + member_id: member_id.to_string() + }) + }) + } + + pub fn close_team(&self, team_id: &str, summary: Option<&str>) -> Result { + let existing = self + .ledger + .get_team(team_id)? + .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; + self.ledger.upsert_team(AgentTeamUpsert { + id: team_id.to_string(), + parent_thread_id: existing.parent_thread_id, + lead_agent_id: existing.lead_agent_id, + status: AgentTeamStatus::Closed, + summary: summary.map(str::to_string), + created_at: Some(existing.created_at), + closed_at: Some(Utc::now()), + }) + } + + fn team_view(&self, team_id: &str) -> Result { + let team = self + .ledger + .get_team(team_id)? + .ok_or_else(|| anyhow!("team missing after creation: {team_id}"))?; + Ok(TeamView { + team, + members: self.ledger.list_members(team_id)?, + tasks: self.ledger.list_tasks(team_id)?, + }) + } + + fn ensure_member(&self, team_id: &str, member_id: &str) -> Result<()> { + if self.ledger.list_members(team_id)?.iter().any(|member| { + member.id == member_id && member.member_status != AgentTeamMemberStatus::Stopped + }) { + Ok(()) + } else { + Err(anyhow!(TeamError::UnknownMember { + member_id: member_id.to_string() + })) + } + } + + fn ensure_team_active(&self, team_id: &str) -> Result<()> { + let team = self + .ledger + .get_team(team_id)? + .ok_or_else(|| anyhow!("unknown team: {team_id}"))?; + if team.status == AgentTeamStatus::Active { + Ok(()) + } else { + Err(anyhow!("team is closed: {team_id}")) + } + } +} + +fn validate_dependencies( + new_task_id: &str, + depends_on: &[String], + existing: &[AgentTeamTask], +) -> Result<()> { + let known: HashSet<&str> = existing.iter().map(|task| task.id.as_str()).collect(); + for dependency in depends_on { + if dependency == new_task_id { + return Err(anyhow!(TeamError::SelfDependency { + task_id: new_task_id.to_string() + })); + } + if !known.contains(dependency.as_str()) { + return Err(anyhow!(TeamError::UnknownDependency { + depends_on: dependency.clone() + })); + } + } + if has_task_cycle(new_task_id, depends_on, existing) { + return Err(anyhow!(TeamError::CyclicDependency)); + } + Ok(()) +} + +fn has_task_cycle(new_task_id: &str, depends_on: &[String], existing: &[AgentTeamTask]) -> bool { + let mut nodes: Vec> = existing + .iter() + .map(|task| DagNode::new(task.id.as_str(), task.depends_on.iter().map(String::as_str))) + .collect(); + nodes.push(DagNode::new( + new_task_id, + depends_on.iter().map(String::as_str), + )); + has_cycle(&nodes) +} + +/// Select the next task a member may claim without making any policy decision. +pub fn claimable_task<'a>( + tasks: &'a [AgentTeamTask], + member_id: &str, +) -> Option<&'a AgentTeamTask> { + let done: HashSet<&str> = tasks + .iter() + .filter(|task| task.status == AgentTeamTaskStatus::Done) + .map(|task| task.id.as_str()) + .collect(); + tasks.iter().find(|task| { + matches!( + task.status, + AgentTeamTaskStatus::Todo | AgentTeamTaskStatus::Ready + ) && task.claimed_by_member_id.is_none() + && task + .owner_member_id + .as_deref() + .map(|owner| owner == member_id) + .unwrap_or(true) + && task + .depends_on + .iter() + .all(|dependency| done.contains(dependency.as_str())) + }) +} + +#[cfg(test)] +mod dependency_tests { + use chrono::Utc; + use tinyagents_session::run_ledger::AgentTeamTaskStatus; + + use super::*; + + fn task(id: &str, depends_on: &[&str]) -> AgentTeamTask { + let now = Utc::now(); + AgentTeamTask { + id: id.to_string(), + team_id: "team".to_string(), + title: id.to_string(), + objective: None, + status: AgentTeamTaskStatus::Todo, + owner_member_id: None, + claimed_by_member_id: None, + claim_token: None, + depends_on: depends_on + .iter() + .map(|dependency| (*dependency).to_string()) + .collect(), + gate_status: "pending".to_string(), + gate_reason: None, + evidence: Vec::new(), + source_run_id: None, + order_index: 0, + created_at: now, + updated_at: now, + } + } + + #[test] + fn rejects_self_dependency() { + let error = + validate_dependencies("task-self", &["task-self".to_string()], &[]).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + TeamError::SelfDependency { + task_id: "task-self".to_string() + } + ); + } + + #[test] + fn rejects_dependency_cycle() { + let existing = vec![task("task-a", &["task-new"])]; + let error = + validate_dependencies("task-new", &["task-a".to_string()], &existing).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + TeamError::CyclicDependency + ); + } +} diff --git a/crates/tinyagents-orchestration/src/teams/tests.rs b/crates/tinyagents-orchestration/src/teams/tests.rs new file mode 100644 index 00000000..c191e110 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/tests.rs @@ -0,0 +1,390 @@ +//! Tests for team service: creation, member lifecycle, task coordination, +//! and messaging flow. + +use std::sync::{Arc, Barrier, Mutex}; +use std::thread; + +use chrono::Utc; +use tempfile::TempDir; +use tinyagents_session::run_ledger::{ + AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, AgentTeamMemberStatus, + AgentTeamMemberUpsert, AgentTeamStatus, AgentTeamTask, AgentTeamTaskUpsert, AgentTeamUpsert, + ClaimOutcome, CompletionOutcome, RunEvent, RunEventAppend, RunEventListRequest, +}; + +use super::*; + +fn service(dir: &TempDir) -> TeamService { + TeamService::new(SessionTeamLedger::new(dir.path())) +} + +fn solo_team(service: &TeamService) -> (String, String) { + let view = service + .create_team( + "lead", + None, + None, + &[NewMember { + name: "alice".into(), + agent_id: None, + }], + ) + .unwrap(); + (view.team.id, view.members[0].id.clone()) +} + +fn team_error(error: anyhow::Error) -> TeamError { + error.downcast::().unwrap() +} + +#[test] +fn rejects_duplicate_members_and_unknown_dependencies() { + let dir = TempDir::new().unwrap(); + let service = service(&dir); + let error = service + .create_team( + "lead", + None, + None, + &[ + NewMember { + name: "alice".into(), + agent_id: None, + }, + NewMember { + name: "alice".into(), + agent_id: None, + }, + ], + ) + .unwrap_err(); + assert_eq!( + team_error(error), + TeamError::DuplicateMemberName { + name: "alice".into() + } + ); + + let (team_id, _) = solo_team(&service); + let error = service + .assign_task(&team_id, "task", None, None, &["missing".into()]) + .unwrap_err(); + assert_eq!( + team_error(error), + TeamError::UnknownDependency { + depends_on: "missing".into() + } + ); +} + +#[test] +fn task_claim_completion_and_quality_gate_are_durable() { + let dir = TempDir::new().unwrap(); + let service = service(&dir); + let (team_id, member_id) = solo_team(&service); + let task = service + .assign_task(&team_id, "ship", None, None, &[]) + .unwrap(); + assert!(matches!( + service + .claim_task(&team_id, &task.id, &member_id, "claim-1") + .unwrap(), + ClaimOutcome::Claimed(_) + )); + assert!(matches!( + service + .complete_task(&team_id, &task.id, &member_id, &[], true) + .unwrap(), + CompletionOutcome::GateFailed { .. } + )); + let after_failed_gate = service.get_team(&team_id).unwrap().unwrap(); + let durable_task = after_failed_gate + .tasks + .iter() + .find(|candidate| candidate.id == task.id) + .unwrap(); + assert_eq!( + durable_task.status, + tinyagents_session::run_ledger::AgentTeamTaskStatus::InProgress + ); + assert_eq!(durable_task.gate_status, "failed"); + assert!(matches!( + service + .complete_task(&team_id, &task.id, &member_id, &["proof".into()], true) + .unwrap(), + CompletionOutcome::Completed(_) + )); +} + +#[test] +fn racing_claims_have_one_winner_and_one_already_claimed_loser() { + let dir = TempDir::new().unwrap(); + let service = service(&dir); + let view = service + .create_team( + "lead", + None, + None, + &[ + NewMember { + name: "alice".into(), + agent_id: None, + }, + NewMember { + name: "bob".into(), + agent_id: None, + }, + ], + ) + .unwrap(); + let task = service + .assign_task(&view.team.id, "race", None, None, &[]) + .unwrap(); + let barrier = Arc::new(Barrier::new(2)); + let mut handles = Vec::new(); + for (member, token) in [ + (view.members[0].id.clone(), "alice-token"), + (view.members[1].id.clone(), "bob-token"), + ] { + let service = service.clone(); + let team_id = view.team.id.clone(); + let task_id = task.id.clone(); + let barrier = barrier.clone(); + handles.push(thread::spawn(move || { + barrier.wait(); + service + .claim_task(&team_id, &task_id, &member, token) + .unwrap() + })); + } + let outcomes: Vec<_> = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect(); + assert_eq!( + outcomes + .iter() + .filter(|outcome| matches!(outcome, ClaimOutcome::Claimed(_))) + .count(), + 1 + ); + assert_eq!( + outcomes + .iter() + .filter(|outcome| matches!(outcome, ClaimOutcome::AlreadyClaimed)) + .count(), + 1 + ); +} + +#[test] +fn completion_rejects_non_claimants_and_owner_mismatches() { + let dir = TempDir::new().unwrap(); + let service = service(&dir); + let view = service + .create_team( + "lead", + None, + None, + &[ + NewMember { + name: "alice".into(), + agent_id: None, + }, + NewMember { + name: "bob".into(), + agent_id: None, + }, + ], + ) + .unwrap(); + let alice = &view.members[0].id; + let bob = &view.members[1].id; + + let claimed_by_alice = service + .assign_task(&view.team.id, "alice work", None, None, &[]) + .unwrap(); + service + .claim_task(&view.team.id, &claimed_by_alice.id, alice, "alice-token") + .unwrap(); + assert!(matches!( + service + .complete_task( + &view.team.id, + &claimed_by_alice.id, + bob, + &["proof".into()], + false + ) + .unwrap(), + CompletionOutcome::NotClaimed + )); + + let owned_by_alice = service + .assign_task(&view.team.id, "owned work", None, Some(alice), &[]) + .unwrap(); + service + .claim_task(&view.team.id, &owned_by_alice.id, bob, "bob-token") + .unwrap(); + let outcome = service + .complete_task( + &view.team.id, + &owned_by_alice.id, + bob, + &["proof".into()], + false, + ) + .unwrap(); + assert!(matches!( + outcome, + CompletionOutcome::GateFailed { ref reasons } + if reasons.iter().any(|reason| reason.contains("owned by")) + )); +} + +#[test] +fn messages_remain_ordered_and_member_shutdown_releases_work() { + let dir = TempDir::new().unwrap(); + let service = service(&dir); + let view = service + .create_team( + "lead", + None, + None, + &[ + NewMember { + name: "alice".into(), + agent_id: None, + }, + NewMember { + name: "bob".into(), + agent_id: None, + }, + ], + ) + .unwrap(); + let team_id = view.team.id; + let alice = view.members[0].id.clone(); + let bob = view.members[1].id.clone(); + service + .message_member(&team_id, Some(&alice), Some(&bob), "first", None) + .unwrap(); + service + .message_member(&team_id, None, Some(&alice), "second", None) + .unwrap(); + let messages = service.list_messages(&team_id, None).unwrap(); + assert_eq!( + messages + .iter() + .map(|event| event.payload["content"].as_str()) + .collect::>(), + vec![Some("first"), Some("second")] + ); + assert_eq!(messages[1].payload["from"], LEAD_SENDER); + + let task = service + .assign_task(&team_id, "ship", None, None, &[]) + .unwrap(); + service + .claim_task(&team_id, &task.id, &alice, "claim-1") + .unwrap(); + let shutdown = service.shutdown_member(&team_id, &alice).unwrap(); + assert_eq!(shutdown.released_task_ids, vec![task.id]); + assert_eq!( + shutdown.member.member_status, + AgentTeamMemberStatus::Stopped + ); +} + +#[test] +fn fake_ledger_exercises_member_validation_without_session_storage() { + let ledger = FakeLedger::default(); + let service = TeamService::new(ledger.clone()); + ledger.teams.lock().unwrap().push(team("team-1")); + let error = service + .claim_task("team-1", "task-1", "unknown", "token") + .unwrap_err(); + assert_eq!( + team_error(error), + TeamError::UnknownMember { + member_id: "unknown".into() + } + ); +} + +fn team(id: &str) -> AgentTeam { + AgentTeam { + id: id.into(), + parent_thread_id: None, + lead_agent_id: "lead".into(), + status: AgentTeamStatus::Active, + summary: None, + created_at: Utc::now(), + updated_at: Utc::now(), + closed_at: None, + } +} + +#[derive(Clone, Default)] +struct FakeLedger { + teams: Arc>>, +} + +impl TeamLedger for FakeLedger { + fn upsert_team(&self, upsert: AgentTeamUpsert) -> anyhow::Result { + Ok(team(&upsert.id)) + } + fn get_team(&self, id: &str) -> anyhow::Result> { + Ok(self + .teams + .lock() + .unwrap() + .iter() + .find(|team| team.id == id) + .cloned()) + } + fn list_teams(&self, _: &AgentTeamListRequest) -> anyhow::Result { + Ok(AgentTeamListResponse { + teams: self.teams.lock().unwrap().clone(), + count: self.teams.lock().unwrap().len(), + }) + } + fn upsert_member(&self, _: AgentTeamMemberUpsert) -> anyhow::Result { + unreachable!() + } + fn list_members(&self, _: &str) -> anyhow::Result> { + Ok(vec![]) + } + fn list_tasks(&self, _: &str) -> anyhow::Result> { + Ok(vec![]) + } + fn upsert_task(&self, _: AgentTeamTaskUpsert) -> anyhow::Result { + unreachable!() + } + fn claim_task(&self, _: &str, _: &str, _: &str, _: &str) -> anyhow::Result { + unreachable!() + } + fn complete_task( + &self, + _: &str, + _: &str, + _: &str, + _: &[String], + _: bool, + ) -> anyhow::Result { + unreachable!() + } + fn shutdown_member( + &self, + _: &str, + _: &str, + ) -> anyhow::Result)>> { + Ok(None) + } + fn append_event(&self, _: RunEventAppend) -> anyhow::Result { + unreachable!() + } + fn list_events(&self, _: &RunEventListRequest) -> anyhow::Result> { + Ok(vec![]) + } +} diff --git a/crates/tinyagents-orchestration/src/teams/types.rs b/crates/tinyagents-orchestration/src/teams/types.rs new file mode 100644 index 00000000..6f2c9e64 --- /dev/null +++ b/crates/tinyagents-orchestration/src/teams/types.rs @@ -0,0 +1,91 @@ +//! Data types for team coordination: members, tasks, shutdown results, and +//! validation errors. + +use serde::Serialize; +use tinyagents_session::run_ledger::{AgentTeam, AgentTeamMember, AgentTeamTask}; + +/// Sentinel sender for a lead or user message rather than a member row. +/// +/// Used in the event log to distinguish lead/user messages from member-to-member +/// messages. All events with this sender should be treated as external input +/// rather than team member output. +pub const LEAD_SENDER: &str = "lead"; + +/// One member supplied when a team is created. +/// +/// Carries only the identity and (optional) agent-id hint needed for +/// initialization; the full persistent [`AgentTeamMember`] row is created +/// by the ledger. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NewMember { + /// Human-readable member name for prompts and display. + pub name: String, + /// Optional reference to an agent definition in the host registry. When + /// supplied, the host may use it to resolve member prompts and tool access. + pub agent_id: Option, +} + +/// A durable team and the member/task rows needed to render it. +/// +/// Projects a complete team state at one point in time: the team metadata, +/// its members, and its tasks. Used for snapshots, UI display, and audit logs. +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct TeamView { + /// The team metadata and status. + pub team: AgentTeam, + /// All members of the team, in creation order. + pub members: Vec, + /// All tasks the team has been assigned, in creation order. + pub tasks: Vec, +} + +/// Result of stopping a member and releasing its active claims. +/// +/// Returned when a running member is shut down. Records the stopped member +/// and the task ids it had claimed (now released and available for +/// reassignment). +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct MemberShutdown { + /// The member that was shut down. + pub member: AgentTeamMember, + /// Task ids that member had claimed, now released. + pub released_task_ids: Vec, +} + +/// Coordination validation errors that are independent of a host or storage backend. +/// +/// These errors detect structural issues in team definition or task dependency +/// graphs that do not require access to a host's agent registry or persistence +/// layer. They are deterministic and stable across runs. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase", tag = "kind", content = "detail")] +pub enum TeamError { + /// A team member name appears more than once. + DuplicateMemberName { name: String }, + /// A task references a member id that is not registered. + UnknownMember { member_id: String }, + /// A task depends on itself. + SelfDependency { task_id: String }, + /// Task dependencies form a cycle. + CyclicDependency, + /// A task depends on another task that is not registered. + UnknownDependency { depends_on: String }, +} + +impl std::fmt::Display for TeamError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::DuplicateMemberName { name } => write!(f, "duplicate member name: {name}"), + Self::UnknownMember { member_id } => write!(f, "unknown member: {member_id}"), + Self::SelfDependency { task_id } => write!(f, "task {task_id} cannot depend on itself"), + Self::CyclicDependency => write!(f, "dependency cycle detected"), + Self::UnknownDependency { depends_on } => { + write!(f, "unknown dependency: {depends_on}") + } + } + } +} + +impl std::error::Error for TeamError {} diff --git a/crates/tinyagents-orchestration/src/workflow/README.md b/crates/tinyagents-orchestration/src/workflow/README.md new file mode 100644 index 00000000..9fb17b96 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/README.md @@ -0,0 +1,79 @@ +# `orchestration::workflow` — durable phase DAG execution + +A **workflow** is a directed acyclic graph of phases, each with associated +agents, dependencies, and concurrency limits. The [`WorkflowEngine`] schedules +phases topologically, spawns bounded-concurrent child tasks per phase, +collects results, and manages state transitions durably. + +## Public surface + +- **`WorkflowEngine`** — the execution engine: schedules + runnable phases, spawns child tasks, handles concurrency limits, collects + results, and persists state. Generic over a host-supplied executor. +- **`WorkflowExecutor`** — the trait for host-supplied work: create and + monitor child tasks, cancel them, and query their status. +- **`WorkflowStore`, `SessionWorkflowStore`** — durable run state: loads, + saves, claims (lease), and compare-and-swap updates. [`SessionWorkflowStore`] + wraps the session run ledger. +- **`WorkflowDefinition`, `WorkflowPhase`** — the declarative phase DAG: phase + names, descriptions, agent ids, dependencies, concurrency settings. +- **`PhaseStatus`** — phase state machine: Pending → Running → (Completed | + Failed). Interrupted phases can reset to Pending for retry. +- **State projection functions** — query and mutate the JSON phase-state + document: `phase_status()`, `next_runnable_phase()`, `all_phases_completed()`, + `reset_running_phases()`, `phase_prompt()`, `synthesize_summary()`. + +## Design and invariants + +- **Durable at every boundary:** Workflow runs and phase states are durably + persisted via [`WorkflowStore`]. The engine is safe to interrupt and resume. +- **Topological scheduling:** Phases are scheduled in dependency order. The + engine finds the next runnable phase (all dependencies met), preventing + partial execution and cycles. +- **Bounded parallelism:** The engine respects `default_concurrency` (agents + per phase in parallel) and `max_children` (total spawned children at once). + Phase work is fan-out (parallel agents) + fan-in (collect results). +- **Deterministic retry:** If the engine is interrupted while running phases, + all running phases are reset to Pending (outputs cleared). Completed phases + remain immutable and are never retried. +- **Result aggregation:** Upstream outputs (from dependency phases) are + collected and passed to downstream phases' prompts, allowing workflows to + reason over prior results. +- **JSON phase state:** The phase-state document is JSON (a BTreeMap of phase + names to `{ status, outputs, reason, ... }`). This projection is durable, + queryable, and renderable in UIs. + +## File map + +- **`engine.rs`** — [`WorkflowEngine`] implementation: phase scheduling, child + task spawning and monitoring, state persistence, concurrency enforcement. + [`WorkflowStore`], [`SessionWorkflowStore`], [`WorkflowExecutor`] trait. +- **`types.rs`** — data types: [`WorkflowDefinition`], [`WorkflowPhase`], + [`DefinitionError`], [`WorkflowDefinitionListResponse`]. +- **`state.rs`** — phase-state projection: JSON document queries and mutations. + [`PhaseStatus`], `next_runnable_phase()`, `phase_prompt()`, + `synthesize_summary()`, etc. +- **`graph.rs`** — scheduler DAG: models workflow phases as a directed graph + for topological sorting and dependency resolution. +- **`validate.rs`** — structural validation: no duplicate phases, valid + dependencies, no cycles, valid concurrency settings, etc. +- **`tests.rs`** — tests for scheduling, phase transitions, concurrency, + result aggregation, and interruption/retry behavior. + +## Relationship to other modules + +- **Depends on:** `tinyagents-graph` (DAG validation), `tinyagents-session` + (run ledger), `tinyagents-harness` (cancellation, error types). +- **Used by:** host orchestration logic (workflow management and execution). +- **Integration:** Can model multi-agent phases as teams (via + `orchestration::teams`). + +## Typical usage + +1. Define a [`WorkflowDefinition`] with phases and dependencies. +2. Create a [`WorkflowEngine`] with a host-supplied [`WorkflowExecutor`] and + [`WorkflowStore`]. +3. Call `engine.run()`: the engine schedules phases, spawns bounded child + tasks, collects results, and persists state. +4. On interruption, retry `engine.run()`: running phases reset to Pending; + completed phases remain done. diff --git a/crates/tinyagents-orchestration/src/workflow/engine.rs b/crates/tinyagents-orchestration/src/workflow/engine.rs new file mode 100644 index 00000000..4b83f844 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/engine.rs @@ -0,0 +1,882 @@ +//! Workflow execution engine: orchestrates phase scheduling, task spawning, +//! and result collection. +//! +//! [`WorkflowEngine`] drives the scheduler, manages child task creation and +//! claim-based assignment, handles phase state persistence and transitions, +//! and enforces bounded concurrency. It is generic over a host-supplied +//! [`WorkflowExecutor`] that creates and monitors actual work. + +use std::fmt; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use chrono::Utc; +use serde_json::{Value, json}; +use tinyagents_graph::GraphEventSink; +use tinyagents_graph::parallel::{FailurePolicy, ParallelOptions, map_reduce}; +use tinyagents_harness::CancellationToken; +use tinyagents_session::run_ledger::{ + WorkflowLeaseClaim, WorkflowRun, WorkflowRunStatus, WorkflowRunUpsert, + compare_and_swap_workflow_run, get_workflow_run, renew_workflow_run_lease, + try_claim_workflow_run, upsert_workflow_run, +}; + +use super::state::{ + PhaseStatus, all_phases_completed, init_phase_states, next_runnable_phase, phase_prompt, + reset_running_phases, set_phase_reason, set_phase_status, synthesize_summary, upstream_outputs, +}; +use super::{WorkflowDefinition, WorkflowPhase}; + +/// Error returned by a host child executor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OrchestrationError(pub String); + +impl fmt::Display for OrchestrationError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(formatter) + } +} + +impl std::error::Error for OrchestrationError {} + +impl From for OrchestrationError { + fn from(error: anyhow::Error) -> Self { + Self(error.to_string()) + } +} + +impl From for OrchestrationError { + fn from(error: tinyagents_harness::TinyAgentsError) -> Self { + Self(error.to_string()) + } +} + +/// Durable workflow rows supplied by the session owner. +pub trait WorkflowStore: Send + Sync { + fn load(&self, id: &str) -> Result, OrchestrationError>; + fn upsert(&self, row: WorkflowRunUpsert) -> Result; + fn claim( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result; + fn compare_and_swap( + &self, + row: WorkflowRunUpsert, + expected_revision: u64, + owner: &str, + lease_for: Duration, + ) -> Result, OrchestrationError>; + fn renew(&self, id: &str, owner: &str, lease_for: Duration) + -> Result; +} + +/// `tinyagents-session` run-ledger adapter with a caller-selected workspace. +#[derive(Debug, Clone)] +pub struct SessionWorkflowStore { + workspace_dir: PathBuf, +} + +impl SessionWorkflowStore { + pub fn new(workspace_dir: impl Into) -> Self { + Self { + workspace_dir: workspace_dir.into(), + } + } + + pub fn workspace_dir(&self) -> &Path { + &self.workspace_dir + } +} + +impl WorkflowStore for SessionWorkflowStore { + fn load(&self, id: &str) -> Result, OrchestrationError> { + get_workflow_run(&self.workspace_dir, id).map_err(OrchestrationError::from) + } + + fn upsert(&self, row: WorkflowRunUpsert) -> Result { + upsert_workflow_run(&self.workspace_dir, row).map_err(OrchestrationError::from) + } + + fn claim( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result { + try_claim_workflow_run( + &self.workspace_dir, + id, + owner, + chrono::Duration::from_std(lease_for) + .map_err(|error| OrchestrationError(error.to_string()))?, + ) + .map_err(OrchestrationError::from) + } + + fn compare_and_swap( + &self, + row: WorkflowRunUpsert, + expected_revision: u64, + owner: &str, + lease_for: Duration, + ) -> Result, OrchestrationError> { + compare_and_swap_workflow_run( + &self.workspace_dir, + row, + expected_revision, + owner, + chrono::Duration::from_std(lease_for) + .map_err(|error| OrchestrationError(error.to_string()))?, + ) + .map_err(OrchestrationError::from) + } + + fn renew( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result { + renew_workflow_run_lease( + &self.workspace_dir, + id, + owner, + chrono::Duration::from_std(lease_for) + .map_err(|error| OrchestrationError(error.to_string()))?, + ) + .map_err(OrchestrationError::from) + } +} + +/// One host-authorized child invocation. +#[derive(Debug, Clone, PartialEq)] +pub struct WorkflowChildRequest { + pub run_id: String, + pub phase: String, + pub agent_id: String, + pub index_in_phase: usize, + pub prompt: String, +} + +/// One child's terminal result. `output` is retained verbatim in phase state. +#[derive(Debug, Clone, PartialEq)] +pub struct WorkflowChildResult { + pub child_id: String, + pub output: Value, +} + +/// Called by a host immediately after it has created a real child. Registering +/// before waiting makes the child visible to a concurrent cancellation request +/// even when the worker is still in flight. +pub trait WorkflowChildRegistration: Send + Sync { + fn register(&self, child_id: String) -> Result<(), OrchestrationError>; +} + +/// Host-owned execution and cancellation mechanism. +#[async_trait] +pub trait WorkflowExecutor: Send + Sync { + async fn execute( + &self, + request: WorkflowChildRequest, + cancel: CancellationToken, + registration: Arc, + ) -> Result; + + async fn cancel_children(&self, child_ids: &[String]); +} + +/// Generic durable workflow engine. It never creates tasks: hosts decide where +/// work runs and which authorization context is in force via [`WorkflowExecutor`]. +pub struct WorkflowEngine { + store: Arc, + executor: Arc, + event_sink: Option>, + event_seq: AtomicU64, + lease_for: Duration, +} + +const WORKFLOW_LEASE: Duration = Duration::from_secs(10 * 60); + +struct PersistRequest { + phase_states: Value, + child_run_ids: Vec, + status: WorkflowRunStatus, + summary: Option, + terminal: bool, +} + +struct PhaseRegistration { + store: Arc, + owner: String, + run: parking_lot::Mutex, + phase_states: Value, + lease_for: Duration, +} + +impl PhaseRegistration { + fn current(&self) -> WorkflowRun { + self.run.lock().clone() + } +} + +impl WorkflowChildRegistration for PhaseRegistration { + fn register(&self, child_id: String) -> Result<(), OrchestrationError> { + let mut run = self.run.lock(); + if run.child_run_ids.iter().any(|known| known == &child_id) { + return Ok(()); + } + let mut children = run.child_run_ids.clone(); + children.push(child_id); + let Some(updated) = self.store.compare_and_swap( + WorkflowRunUpsert { + id: run.id.clone(), + definition_id: run.definition_id.clone(), + parent_thread_id: run.parent_thread_id.clone(), + input: run.input.clone(), + phase_states: self.phase_states.clone(), + child_run_ids: children, + status: WorkflowRunStatus::Running, + summary: None, + started_at: Some(run.started_at), + completed_at: None, + }, + run.revision, + &self.owner, + self.lease_for, + )? + else { + return Err(OrchestrationError( + "workflow lease lost while registering child".into(), + )); + }; + *run = updated; + Ok(()) + } +} + +impl WorkflowEngine +where + S: WorkflowStore + 'static, + E: WorkflowExecutor + 'static, +{ + pub fn new(store: Arc, executor: Arc) -> Self { + Self { + store, + executor, + event_sink: None, + event_seq: AtomicU64::new(0), + lease_for: WORKFLOW_LEASE, + } + } + + /// Attach an optional host sink for ordinary graph lifecycle tracing. + pub fn with_event_sink(mut self, sink: Arc) -> Self { + self.event_sink = Some(sink); + self + } + + /// Override the driver lease for deterministic tests or hosts with a + /// deliberately shorter failure-detection window. + pub fn with_lease_duration(mut self, lease_for: Duration) -> Self { + self.lease_for = lease_for.max(Duration::from_millis(3)); + self + } + + /// Initialise a durable run before the host schedules [`Self::drive`]. + pub fn initialise( + &self, + id: String, + definition: &WorkflowDefinition, + input: Value, + parent_thread_id: Option, + ) -> Result { + self.store.upsert(WorkflowRunUpsert { + id, + definition_id: definition.id.clone(), + parent_thread_id, + input, + phase_states: init_phase_states(definition), + child_run_ids: Vec::new(), + status: WorkflowRunStatus::Running, + summary: None, + started_at: None, + completed_at: None, + }) + } + + /// Drive a run to a terminal state. Completed phases are never executed + /// again, so a host may safely call this after process restart or resume. + pub async fn drive( + &self, + run_id: &str, + definition: &WorkflowDefinition, + cancel: CancellationToken, + ) -> Result<(), OrchestrationError> { + // A driver lease is acquired before looking for runnable work. This + // is deliberately separate from the in-process cancellation token: + // resume can race in another process, and only the durable lease + // prevents both drivers from spawning the same phase. + let owner = uuid::Uuid::new_v4().to_string(); + let mut run = match self.store.claim(run_id, &owner, self.lease_for)? { + WorkflowLeaseClaim::Acquired(run) => run, + WorkflowLeaseClaim::Busy(_) => return Ok(()), + WorkflowLeaseClaim::Missing => { + return Err(OrchestrationError(format!( + "workflow run {run_id} vanished before start" + ))); + } + }; + // A crashed owner can leave the durable phase marked `running`. That + // marker is intentionally not runnable, so reclaiming the lease must + // turn it back into retryable work before scheduling. The claim above + // fences every prior owner; this CAS is the new owner's durable + // recovery transition rather than a read/modify/write race. + if run.phase_states.as_object().is_some_and(|phases| { + phases + .values() + .any(|phase| phase.get("status").and_then(Value::as_str) == Some("running")) + }) { + // The previous owner may have crashed after registering remote + // children. Fence side effects before making its phase runnable. + self.executor.cancel_children(&run.child_run_ids).await; + let mut phase_states = run.phase_states.clone(); + reset_running_phases( + &mut phase_states, + "workflow owner expired; phase will retry after lease takeover", + ); + run = self.persist( + &run, + PersistRequest { + phase_states, + child_run_ids: Vec::new(), + status: WorkflowRunStatus::Running, + summary: None, + terminal: false, + }, + &owner, + )?; + } + self.emit( + run_id, + tinyagents_graph::GraphEvent::RunStarted { + run_id: tinyagents_harness::ids::RunId::new(run_id), + }, + ); + // Registered children from an interrupted attempt are historical, not + // part of the retry's spawn budget. + let mut total_spawned = 0; + + loop { + if cancel.is_cancelled() { + self.executor.cancel_children(&run.child_run_ids).await; + let mut phase_states = run.phase_states.clone(); + reset_running_phases( + &mut phase_states, + "workflow interrupted; phase will retry on resume", + ); + if let Err(error) = self.persist( + &run, + PersistRequest { + phase_states, + child_run_ids: run.child_run_ids.clone(), + status: WorkflowRunStatus::Interrupted, + summary: None, + terminal: false, + }, + &owner, + ) { + if self.owner_lost(run_id, &owner) + || self.emit_recorded_terminal(run_id, total_spawned as usize) + { + return Ok(()); + } + self.finish_failed(run_id, error.to_string()); + return Err(error); + } + self.finish_cancelled(run_id); + return Ok(()); + } + let Some(phase) = next_runnable_phase(definition, &run.phase_states).cloned() else { + if all_phases_completed(definition, &run.phase_states) { + if let Err(error) = self.persist( + &run, + PersistRequest { + phase_states: run.phase_states.clone(), + child_run_ids: run.child_run_ids.clone(), + status: WorkflowRunStatus::Completed, + summary: synthesize_summary(definition, &run.phase_states), + terminal: true, + }, + &owner, + ) { + if self.owner_lost(run_id, &owner) { + return Ok(()); + } + self.finish_failed(run_id, error.to_string()); + return Err(error); + } + self.finish_completed(run_id, total_spawned as usize); + } else { + let reason = "no runnable phase (dependency deadlock)".to_owned(); + if let Err(error) = self.persist( + &run, + PersistRequest { + phase_states: run.phase_states.clone(), + child_run_ids: run.child_run_ids.clone(), + status: WorkflowRunStatus::Failed, + summary: Some(reason.clone()), + terminal: true, + }, + &owner, + ) { + if self.owner_lost(run_id, &owner) { + return Ok(()); + } + self.finish_failed(run_id, error.to_string()); + return Err(error); + } + self.finish_failed(run_id, reason); + } + return Ok(()); + }; + self.emit( + run_id, + tinyagents_graph::GraphEvent::NodeStarted { + node: tinyagents_harness::ids::NodeId::new("run_phase"), + step: total_spawned as usize + 1, + }, + ); + let phase_result = self + .run_phase( + &run, + definition, + &phase, + total_spawned, + cancel.clone(), + &owner, + ) + .await; + let (updated, spawned) = match phase_result { + Ok(result) => result, + Err(error) => { + // A host stop/resume fences this owner with a revision CAS. + // Do not turn that intentional hand-off into a stale + // failure event or overwrite the newer durable state. + if self.owner_lost(run_id, &owner) + || self.emit_recorded_terminal(run_id, total_spawned as usize) + { + return Ok(()); + } + self.finish_failed(run_id, error.to_string()); + return Err(error); + } + }; + run = updated; + self.emit( + run_id, + tinyagents_graph::GraphEvent::NodeCompleted { + node: tinyagents_harness::ids::NodeId::new("run_phase"), + step: total_spawned as usize + 1, + }, + ); + total_spawned += spawned; + if run.status != WorkflowRunStatus::Running { + match run.status { + WorkflowRunStatus::Completed => { + self.finish_completed(run_id, total_spawned as usize) + } + WorkflowRunStatus::Interrupted | WorkflowRunStatus::Cancelled => { + self.finish_cancelled(run_id) + } + WorkflowRunStatus::Failed => self.finish_failed( + run_id, + run.summary + .clone() + .unwrap_or_else(|| "workflow phase failed".to_owned()), + ), + WorkflowRunStatus::Pending | WorkflowRunStatus::Running => {} + } + return Ok(()); + } + } + } + + async fn run_phase( + &self, + run: &WorkflowRun, + definition: &WorkflowDefinition, + phase: &WorkflowPhase, + total_spawned: u32, + cancel: CancellationToken, + owner: &str, + ) -> Result<(WorkflowRun, u32), OrchestrationError> { + let mut phase_states = run.phase_states.clone(); + let mut child_ids = run.child_run_ids.clone(); + set_phase_status(&mut phase_states, &phase.name, PhaseStatus::Running, None); + let running = self.persist( + run, + PersistRequest { + phase_states: phase_states.clone(), + child_run_ids: child_ids.clone(), + status: WorkflowRunStatus::Running, + summary: None, + terminal: false, + }, + owner, + )?; + + let budget = definition.max_children.saturating_sub(total_spawned) as usize; + if budget == 0 { + return self.fail_phase( + &running, + &mut phase_states, + child_ids, + phase, + format!( + "max_children cap ({}) reached before phase '{}' completed", + definition.max_children, phase.name + ), + owner, + ); + } + let capacity = phase.agent_ids.len().min(budget); + let capped = capacity != phase.agent_ids.len(); + let upstream = upstream_outputs(phase, &phase_states); + let requests = phase.agent_ids[..capacity] + .iter() + .enumerate() + .map(|(index_in_phase, agent_id)| WorkflowChildRequest { + run_id: run.id.clone(), + phase: phase.name.clone(), + agent_id: agent_id.clone(), + index_in_phase, + prompt: phase_prompt(&run.input, phase, index_in_phase, &upstream), + }) + .collect::>(); + let registration = Arc::new(PhaseRegistration { + store: self.store.clone(), + owner: owner.to_owned(), + run: parking_lot::Mutex::new(running.clone()), + phase_states: phase_states.clone(), + lease_for: self.lease_for, + }); + let executor = self.executor.clone(); + let worker_cancel = cancel.clone(); + let worker_registration = registration.clone(); + let outcomes = map_reduce( + requests, + ParallelOptions::default() + .with_max_concurrency(definition.default_concurrency as usize) + .with_failure_policy(FailurePolicy::CollectAll) + .with_cancellation(cancel.clone()), + move |_index, request| { + let executor = executor.clone(); + let cancel = worker_cancel.clone(); + let registration = worker_registration.clone(); + async move { + executor + .execute(request, cancel, registration) + .await + .map_err(|error| { + tinyagents_harness::TinyAgentsError::Graph(error.to_string()) + }) + } + }, + ); + tokio::pin!(outcomes); + let renew_every = self.lease_for.div_f32(3.0).max(Duration::from_millis(1)); + let mut heartbeat = tokio::time::interval(renew_every); + // Ignore interval's eager first tick: `claim` has just installed this + // lease, so renewals begin only while a child can be in flight. + heartbeat.tick().await; + let outcomes = loop { + tokio::select! { + outcomes = &mut outcomes => break outcomes, + _ = heartbeat.tick() => { + let renewed = self.store.renew(&run.id, owner, self.lease_for); + if !matches!(renewed, Ok(true)) { + cancel.cancel(); + let children = registration.current().child_run_ids; + self.executor.cancel_children(&children).await; + if let Err(error) = renewed { + return Err(OrchestrationError(format!( + "workflow lease renewal errored; cancelled registered children: {error}" + ))); + } + return Err(OrchestrationError( + "workflow lease renewal failed; cancelled registered children".to_owned(), + )); + } + } + } + }; + let outcomes = match outcomes { + Ok(outcomes) => outcomes, + Err(tinyagents_harness::TinyAgentsError::Cancelled) => { + let children = registration.current().child_run_ids; + self.executor.cancel_children(&children).await; + reset_running_phases( + &mut phase_states, + "workflow interrupted; phase will retry on resume", + ); + let updated = self.persist( + ®istration.current(), + PersistRequest { + phase_states, + child_run_ids: children, + status: WorkflowRunStatus::Interrupted, + summary: None, + terminal: false, + }, + owner, + )?; + return Ok((updated, 0)); + } + Err(error) => return Err(OrchestrationError(error.to_string())), + }; + child_ids = registration.current().child_run_ids; + let mut outputs = Vec::new(); + let mut failure = None; + let mut spawned = 0_u32; + for outcome in outcomes.outcomes { + match outcome.result { + Ok(result) => { + spawned += 1; + // The executor registered the real id before it could + // await completion. Keep older executors harmlessly + // compatible by accepting an already-present id only. + if !child_ids.iter().any(|id| id == &result.child_id) { + child_ids.push(result.child_id.clone()); + } + outputs.push(json!({ + // Preserve the persisted/RPC v1 projection while the + // v2 metadata remains lossless for engine consumers. + "agentId": phase.agent_ids[outcome.index], + "output": render_compat_output(&result.output), + "metadata": { "version": 2, "rawOutput": result.output }, + })); + } + Err(error) if failure.is_none() => failure = Some(error), + Err(_) => {} + } + } + if cancel.is_cancelled() { + let children = registration.current().child_run_ids; + self.executor.cancel_children(&children).await; + reset_running_phases( + &mut phase_states, + "workflow interrupted; phase will retry on resume", + ); + let updated = self.persist( + ®istration.current(), + PersistRequest { + phase_states, + child_run_ids: children, + status: WorkflowRunStatus::Interrupted, + summary: None, + terminal: false, + }, + owner, + )?; + return Ok((updated, 0)); + } + if let Some(reason) = failure.or_else(|| { + capped.then(|| { + format!( + "max_children cap ({}) reached before phase '{}' completed", + definition.max_children, phase.name + ) + }) + }) { + return self.fail_phase( + ®istration.current(), + &mut phase_states, + child_ids, + phase, + reason, + owner, + ); + } + set_phase_status( + &mut phase_states, + &phase.name, + PhaseStatus::Completed, + Some(Value::Array(outputs)), + ); + let updated = self.persist( + ®istration.current(), + PersistRequest { + phase_states, + child_run_ids: child_ids, + status: WorkflowRunStatus::Running, + summary: None, + terminal: false, + }, + owner, + )?; + Ok((updated, spawned)) + } + + fn fail_phase( + &self, + run: &WorkflowRun, + phase_states: &mut Value, + child_ids: Vec, + phase: &WorkflowPhase, + reason: String, + owner: &str, + ) -> Result<(WorkflowRun, u32), OrchestrationError> { + set_phase_status( + phase_states, + &phase.name, + PhaseStatus::Failed, + Some(json!([])), + ); + set_phase_reason(phase_states, &phase.name, &reason); + let updated = self.persist( + run, + PersistRequest { + phase_states: phase_states.clone(), + child_run_ids: child_ids, + status: WorkflowRunStatus::Failed, + summary: Some(reason), + terminal: true, + }, + owner, + )?; + Ok((updated, 0)) + } + + fn persist( + &self, + run: &WorkflowRun, + request: PersistRequest, + owner: &str, + ) -> Result { + self.store + .compare_and_swap( + WorkflowRunUpsert { + id: run.id.clone(), + definition_id: run.definition_id.clone(), + parent_thread_id: run.parent_thread_id.clone(), + input: run.input.clone(), + phase_states: request.phase_states, + child_run_ids: request.child_run_ids, + status: request.status, + summary: request.summary, + started_at: Some(run.started_at), + completed_at: request.terminal.then(Utc::now), + }, + run.revision, + owner, + self.lease_for, + )? + .ok_or_else(|| { + OrchestrationError("workflow lease lost before durable state transition".to_owned()) + }) + } + + fn emit(&self, run_id: &str, event: tinyagents_graph::GraphEvent) { + if let Some(sink) = &self.event_sink { + sink.emit(tinyagents_graph::GraphEventEnvelope { + run_id: tinyagents_harness::ids::RunId::new(run_id), + task_id: None, + ns: Vec::new(), + seq: self.event_seq.fetch_add(1, Ordering::Relaxed), + event, + }); + } + } + + fn finish_completed(&self, run_id: &str, steps: usize) { + self.emit( + run_id, + tinyagents_graph::GraphEvent::RunCompleted { + run_id: tinyagents_harness::ids::RunId::new(run_id), + steps, + }, + ); + self.flush_terminal_events(); + } + + fn finish_failed(&self, run_id: &str, error: String) { + self.emit( + run_id, + tinyagents_graph::GraphEvent::RunFailed { + run_id: tinyagents_harness::ids::RunId::new(run_id), + error, + }, + ); + self.flush_terminal_events(); + } + + fn finish_cancelled(&self, run_id: &str) { + self.emit( + run_id, + tinyagents_graph::GraphEvent::RunCancelled { + run_id: tinyagents_harness::ids::RunId::new(run_id), + }, + ); + self.flush_terminal_events(); + } + + fn flush_terminal_events(&self) { + if let Some(sink) = &self.event_sink { + sink.flush(); + } + } + + /// A lifecycle hand-off or lease takeover has fenced this driver. It must + /// not manufacture a terminal graph event for the replacement owner. + fn owner_lost(&self, run_id: &str, owner: &str) -> bool { + self.store + .load(run_id) + .ok() + .flatten() + .is_some_and(|current| { + current.lease_owner.as_deref() != Some(owner) + || current + .lease_expires_at + .is_none_or(|expires| expires <= Utc::now()) + }) + } + + /// Returns true after emitting the terminal event already committed by a + /// newer lifecycle owner. This is the stale-driver escape hatch: it never + /// writes, so a stop/resume hand-off cannot be overwritten by its loser. + fn emit_recorded_terminal(&self, run_id: &str, steps: usize) -> bool { + let Ok(Some(current)) = self.store.load(run_id) else { + return false; + }; + if !current.status.is_terminal() { + return false; + } + match current.status { + WorkflowRunStatus::Completed => self.finish_completed(run_id, steps), + WorkflowRunStatus::Interrupted | WorkflowRunStatus::Cancelled => { + self.finish_cancelled(run_id) + } + WorkflowRunStatus::Failed => self.finish_failed( + run_id, + current + .summary + .unwrap_or_else(|| "workflow phase failed".to_owned()), + ), + WorkflowRunStatus::Pending | WorkflowRunStatus::Running => return false, + } + true + } +} + +fn render_compat_output(output: &Value) -> String { + match output { + Value::String(text) => text.clone(), + _ => serde_json::to_string(output) + .unwrap_or_else(|_| "".to_owned()), + } +} diff --git a/crates/tinyagents-orchestration/src/workflow/graph.rs b/crates/tinyagents-orchestration/src/workflow/graph.rs new file mode 100644 index 00000000..ab28b748 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/graph.rs @@ -0,0 +1,78 @@ +//! Workflow scheduler DAG: phase scheduling, dependency resolution, and +//! topological ordering. +//! +//! Builds a directed acyclic graph representing the workflow's phase +//! dependencies, computes runnable phases, and projects the schedule into +//! the graph layer for topology introspection. + +use anyhow::{Result, anyhow}; +use tinyagents_graph::export::GraphTopology; +use tinyagents_graph::recursion::RecursionPolicy; +use tinyagents_graph::{ + ClosureStateReducer, Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, +}; + +#[derive(Clone, Default)] +pub(crate) struct SchedulerState; + +pub(crate) enum SchedulerUpdate { + Noop, +} + +/// Structure-only *preview* of the scheduler topology. +/// +/// `WorkflowEngine` is the effectful scheduler. This helper deliberately does +/// not execute that engine; it only supplies a stable topology to diagnostic +/// UIs, so callers must never present it as the graph that ran a workflow. +pub fn scheduler_topology_preview() -> Result { + Ok(build_scheduler_graph(1)?.topology()) +} + +pub(crate) fn build_scheduler_graph( + phase_count: usize, +) -> Result> { + let mut builder = GraphBuilder::::new().set_reducer( + ClosureStateReducer::new(|state: SchedulerState, update| { + match update { + SchedulerUpdate::Noop => {} + } + Ok(state) + }), + ); + // Effects are supplied at invocation time through the execution context. + // These nodes only make the validated dispatch/run/done topology inspectable; + // `WorkflowEngine` installs its effectful variants below. + builder = builder.add_node( + "dispatch", + |_state: SchedulerState, _context: NodeContext| async move { + Ok(NodeResult::Command(Command::default().with_goto(["done"]))) + }, + ); + let graph = builder + .add_node( + "run_phase", + |_state: SchedulerState, _context: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["dispatch"]), + )) + }, + ) + .add_node( + "done", + |_state: SchedulerState, _context: NodeContext| async move { + Ok(NodeResult::Update(SchedulerUpdate::Noop)) + }, + ) + .set_entry("dispatch") + .mark_command_routing("dispatch") + .mark_command_routing("run_phase") + .set_finish("done") + .compile() + .map_err(|error| anyhow!("workflow scheduler graph compile failed: {error}"))? + .with_recursion_policy(RecursionPolicy { + max_visits_per_node: Some(phase_count + 2), + max_total_steps: (phase_count + 1) * 3 + 16, + ..RecursionPolicy::default() + }); + Ok(graph) +} diff --git a/crates/tinyagents-orchestration/src/workflow/lower.rs b/crates/tinyagents-orchestration/src/workflow/lower.rs deleted file mode 100644 index 9bd8f3bd..00000000 --- a/crates/tinyagents-orchestration/src/workflow/lower.rs +++ /dev/null @@ -1,405 +0,0 @@ -//! Lowers a [`WorkflowDefinition`] to a `tinyagents_graph::CompiledGraph` -//! (Phase 4 of `docs/runtime-comparison/feature-gaps.md`; see -//! `docs/runtime-comparison/code-review-graph.md` I12/R6). -//! -//! Two distinct things live here, on purpose: -//! -//! - [`lowered_topology`]: a pure, *never-executed* structural export — one -//! graph node per phase, `depends_on` expressed as literal -//! [`GraphBuilder::add_waiting_edge`] barriers — used only so a host or a -//! test can inspect "the DAG this workflow defines" (mirrors the -//! `scheduler_topology_preview`/`build_scheduler_graph` split already in -//! `workflow::graph`, whose doc makes the same "structure-only preview, -//! never the graph that ran the workflow" distinction). -//! - [`lower_workflow`]: the *executable* lowering [`WorkflowEngine::drive`] -//! actually runs under the `graph-workflows` feature. It builds a small -//! `dispatch -> -> dispatch -> ... ` graph — one real node per -//! phase plus a `dispatch` router — rather than literal per-phase -//! waiting-edge topology. -//! -//! ## Why the executable graph doesn't just run the waiting-edge topology -//! -//! [`WorkflowEngine::run_phase`] persists a phase's Running/Completed/Failed -//! transition with an optimistic compare-and-swap keyed to the durable -//! run's revision. Every branch active in one graph superstep is handed the -//! *same* pre-step state snapshot (this holds for both sequential and -//! parallel supersteps — see `tinyagents_graph::compiled`'s module docs), so -//! two phases that became ready in the same superstep (e.g. two independent -//! phases that both depend only on a common upstream phase) would both -//! start `run_phase` from the same stale revision and race the same CAS. -//! Only one would win; the loser's `persist` would report a spurious -//! "lease lost" failure even though the lease is fine — a real correctness -//! hazard, not a cosmetic one. -//! -//! The `dispatch` router node sidesteps this by construction: it is the -//! *only* thing that decides which phase runs next -//! ([`next_runnable_phase`], the same function the legacy scheduler uses), -//! and it always routes to exactly one phase node before looping back to -//! itself. Exactly one phase node is ever active in any given superstep, so -//! `run_phase`'s CAS is never raced and can be reused completely unchanged. -//! `depends_on` ordering is therefore enforced the same way it always was — -//! algorithmically, by `next_runnable_phase` — not by the executable -//! graph's own edges (which is why [`lowered_topology`]'s literal -//! waiting-edge shape is a separate, non-executed structure). -//! -//! Each phase's own agent fan-out (bounded by `WorkflowDefinition`'s single -//! `default_concurrency`, since [`WorkflowPhase`] carries no per-phase -//! override) still goes through `run_phase`'s existing -//! `tinyagents_graph::parallel::map_reduce` call — the graph crate's own -//! bounded fan-out primitive — completely unchanged. - -use std::collections::HashSet; -use std::sync::Arc; - -use tinyagents_graph::export::GraphTopology; -use tinyagents_graph::recursion::RecursionPolicy; -use tinyagents_graph::{ - ClosureStateReducer, Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, START, -}; -use tinyagents_harness::CancellationToken; -use tinyagents_harness::TinyAgentsError; -use tinyagents_session::run_ledger::{WorkflowRun, WorkflowRunStatus}; - -use super::WorkflowDefinition; -use super::engine::{ - OrchestrationError, PersistRequest, WorkflowEngine, WorkflowExecutor, WorkflowStore, -}; -use super::state::{ - all_phases_completed, next_runnable_phase, reset_running_phases, synthesize_summary, -}; - -/// The lowered graph's `State`/`Update` type: the durable run row plus a -/// running count of how many children this drive has spawned so far (the -/// `max_children` budget). An overwrite reducer (`Update == State`) is safe -/// here because — per the module doc — exactly one node is ever active per -/// superstep, so there is never a sibling update to merge. -#[derive(Clone)] -pub(crate) struct SchedulerState { - pub(crate) run: WorkflowRun, - pub(crate) total_spawned: u32, -} - -/// Builds the *executable* lowered graph: a `dispatch` router plus one node -/// per phase in `definition`. See the module doc for why this shape (rather -/// than literal per-phase waiting edges) is what actually runs. -pub(crate) fn lower_workflow( - engine: Arc>, - definition: Arc, - run_id: String, - owner: String, - cancel: CancellationToken, -) -> Result, OrchestrationError> -where - S: WorkflowStore + 'static, - E: WorkflowExecutor + 'static, -{ - let mut builder = GraphBuilder::::overwrite(); - - { - let engine = engine.clone(); - let definition = definition.clone(); - let run_id = run_id.clone(); - let owner = owner.clone(); - let cancel = cancel.clone(); - builder = builder.add_node( - "dispatch", - move |state: SchedulerState, _ctx: NodeContext| { - let engine = engine.clone(); - let definition = definition.clone(); - let run_id = run_id.clone(); - let owner = owner.clone(); - let cancel = cancel.clone(); - async move { dispatch(engine, definition, run_id, owner, cancel, state).await } - }, - ); - } - builder = builder - .set_entry("dispatch") - .mark_command_routing("dispatch"); - - for phase in &definition.phases { - let engine = engine.clone(); - let definition = definition.clone(); - let run_id = run_id.clone(); - let owner = owner.clone(); - let cancel = cancel.clone(); - let phase = phase.clone(); - let node_id = phase.name.clone(); - builder = builder - .add_node( - node_id.clone(), - move |state: SchedulerState, _ctx: NodeContext| { - let engine = engine.clone(); - let definition = definition.clone(); - let run_id = run_id.clone(); - let owner = owner.clone(); - let cancel = cancel.clone(); - let phase = phase.clone(); - async move { - run_phase_node(engine, definition, run_id, owner, cancel, phase, state) - .await - } - }, - ) - .mark_command_routing(node_id); - } - - let phase_count = definition.phases.len(); - let graph = builder - .compile() - .map_err(|error| OrchestrationError(format!("workflow graph lowering failed: {error}")))? - .with_recursion_policy(RecursionPolicy { - max_visits_per_node: Some(phase_count + 2), - max_total_steps: (phase_count + 1) * 3 + 16, - ..RecursionPolicy::default() - }); - Ok(graph) -} - -/// The `dispatch` node: picks the next runnable phase -/// ([`next_runnable_phase`], identical to the legacy scheduler), or settles -/// the run when nothing is left to run — either every phase completed, or a -/// cancellation was observed with no phase in flight, or no phase is -/// runnable and the workflow is not done (a `depends_on` deadlock). -async fn dispatch( - engine: Arc>, - definition: Arc, - run_id: String, - owner: String, - cancel: CancellationToken, - state: SchedulerState, -) -> tinyagents_graph::Result> -where - S: WorkflowStore + 'static, - E: WorkflowExecutor + 'static, -{ - if cancel.is_cancelled() { - let mut phase_states = state.run.phase_states.clone(); - reset_running_phases( - &mut phase_states, - "workflow interrupted; phase will retry on resume", - ); - return match engine - .persist( - &state.run, - PersistRequest { - phase_states, - child_run_ids: state.run.child_run_ids.clone(), - status: WorkflowRunStatus::Interrupted, - summary: None, - terminal: false, - }, - &owner, - ) - .await - { - Ok(updated) => { - engine.finish_cancelled(&run_id); - Ok(NodeResult::Update(SchedulerState { - run: updated, - ..state - })) - } - Err(error) => settle_infra_error(&engine, &run_id, &owner, state, error).await, - }; - } - - if let Some(phase) = next_runnable_phase(&definition, &state.run.phase_states) { - let target = phase.name.clone(); - return Ok(NodeResult::Command( - Command::default().with_update(state).with_goto([target]), - )); - } - - if all_phases_completed(&definition, &state.run.phase_states) { - let summary = synthesize_summary(&definition, &state.run.phase_states); - return match engine - .persist( - &state.run, - PersistRequest { - phase_states: state.run.phase_states.clone(), - child_run_ids: state.run.child_run_ids.clone(), - status: WorkflowRunStatus::Completed, - summary, - terminal: true, - }, - &owner, - ) - .await - { - Ok(updated) => { - engine.finish_completed(&run_id, state.total_spawned as usize); - Ok(NodeResult::Update(SchedulerState { - run: updated, - ..state - })) - } - Err(error) => settle_infra_error(&engine, &run_id, &owner, state, error).await, - }; - } - - let reason = "no runnable phase (dependency deadlock)".to_owned(); - match engine - .persist( - &state.run, - PersistRequest { - phase_states: state.run.phase_states.clone(), - child_run_ids: state.run.child_run_ids.clone(), - status: WorkflowRunStatus::Failed, - summary: Some(reason.clone()), - terminal: true, - }, - &owner, - ) - .await - { - Ok(updated) => { - engine.finish_failed(&run_id, reason); - Ok(NodeResult::Update(SchedulerState { - run: updated, - ..state - })) - } - Err(error) => settle_infra_error(&engine, &run_id, &owner, state, error).await, - } -} - -/// One phase's node body: reuses [`WorkflowEngine::run_phase`] unchanged, -/// then either loops back to `dispatch` (the phase completed and the run is -/// still `Running`) or stops (a failure/interrupt already durably -/// persisted, and its terminal event already emitted, by `run_phase` -/// itself). -#[allow(clippy::too_many_arguments)] -async fn run_phase_node( - engine: Arc>, - definition: Arc, - run_id: String, - owner: String, - cancel: CancellationToken, - phase: super::WorkflowPhase, - state: SchedulerState, -) -> tinyagents_graph::Result> -where - S: WorkflowStore + 'static, - E: WorkflowExecutor + 'static, -{ - match engine - .run_phase( - &state.run, - &definition, - &phase, - state.total_spawned, - cancel, - &owner, - ) - .await - { - Ok((updated, spawned)) => { - let next = SchedulerState { - run: updated, - total_spawned: state.total_spawned + spawned, - }; - if next.run.status == WorkflowRunStatus::Running { - Ok(NodeResult::Command( - Command::default().with_update(next).with_goto(["dispatch"]), - )) - } else { - // `run_phase` already persisted this transition and it is a - // terminal one (Failed/Interrupted/Cancelled) — surface the - // matching event and stop the loop (no `goto`). - match next.run.status { - WorkflowRunStatus::Interrupted | WorkflowRunStatus::Cancelled => { - engine.finish_cancelled(&run_id) - } - WorkflowRunStatus::Failed => engine.finish_failed( - &run_id, - next.run - .summary - .clone() - .unwrap_or_else(|| "workflow phase failed".to_owned()), - ), - WorkflowRunStatus::Completed | WorkflowRunStatus::Pending => {} - WorkflowRunStatus::Running => unreachable!("handled above"), - } - Ok(NodeResult::Update(next)) - } - } - Err(error) => settle_infra_error(&engine, &run_id, &owner, state, error).await, - } -} - -/// A `run_phase`/`persist` call failed for infrastructure reasons (not an -/// ordinary phase outcome). Mirrors `drive_legacy`'s own -/// `owner_lost`/`emit_recorded_terminal` escape hatches exactly: a -/// stop/resume hand-off or a lease takeover must not manufacture a stale -/// terminal event *or* a hard error for a driver that has already been -/// fenced — that case stops the loop silently (`Ok`, no `goto`), matching -/// `drive_legacy`'s `return Ok(())`. A genuine, unfenced infrastructure -/// failure instead propagates as an `Err`, matching `drive_legacy`'s -/// `return Err(error)`; `WorkflowEngine::drive_via_graph`'s single -/// catch-all around `graph.run(..)` is what emits `finish_failed` for it -/// (once, regardless of which node's `Err` bubbled up). -async fn settle_infra_error( - engine: &Arc>, - run_id: &str, - owner: &str, - state: SchedulerState, - error: OrchestrationError, -) -> tinyagents_graph::Result> -where - S: WorkflowStore + 'static, - E: WorkflowExecutor + 'static, -{ - if engine.owner_lost(run_id, owner).await - || engine - .emit_recorded_terminal(run_id, state.total_spawned as usize) - .await - { - return Ok(NodeResult::Update(state)); - } - // A genuine, unfenced infrastructure failure: propagate it and let - // `WorkflowEngine::drive_via_graph`'s single catch-all emit - // `finish_failed` exactly once, rather than emitting it here too. - Err(TinyAgentsError::Graph(error.to_string())) -} - -/// A pure, never-executed structural export: one node per phase, with -/// `depends_on` expressed as literal [`GraphBuilder::add_waiting_edge`] -/// barriers — "the DAG this workflow defines", not the graph that actually -/// runs it (see the module doc). Requires a single root phase (a phase with -/// an empty `depends_on`); `tinyagents_graph`'s topology only records one -/// `entry` node, so a definition with several independent root phases would -/// lose all but one of them here. -pub fn lowered_topology( - definition: &WorkflowDefinition, -) -> Result { - let mut builder = GraphBuilder::<(), ()>::new() - .set_reducer(ClosureStateReducer::new(|state: (), _update: ()| Ok(state))); - for phase in &definition.phases { - builder = builder.add_node( - phase.name.clone(), - |state: (), _ctx: NodeContext| async move { Ok(NodeResult::Update(state)) }, - ); - } - let depended_on: HashSet<&str> = definition - .phases - .iter() - .flat_map(|phase| phase.depends_on.iter().map(String::as_str)) - .collect(); - for phase in &definition.phases { - if phase.depends_on.is_empty() { - builder = builder.add_edge(START, phase.name.clone()); - } else { - for dependency in &phase.depends_on { - builder = builder.add_waiting_edge(dependency.clone(), phase.name.clone()); - } - } - if !depended_on.contains(phase.name.as_str()) { - builder = builder.set_finish(phase.name.clone()); - } - } - let graph = builder.compile().map_err(|error| { - OrchestrationError(format!("workflow topology lowering failed: {error}")) - })?; - Ok(graph.topology()) -} diff --git a/crates/tinyagents-orchestration/src/workflow/mod.rs b/crates/tinyagents-orchestration/src/workflow/mod.rs new file mode 100644 index 00000000..64a1d1b6 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/mod.rs @@ -0,0 +1,29 @@ +//! Durable, host-neutral workflow definitions and execution. +//! +//! The workflow engine owns phase scheduling, bounded fan-out, cancellation, +//! resume semantics, and the JSON phase-state projection. A host supplies a +//! [`WorkflowStore`] and [`WorkflowExecutor`]; therefore this module has no +//! knowledge of credentials, model selection, policy, progress, or RPC. + +mod engine; +mod graph; +mod state; +mod types; +mod validate; + +pub use engine::{ + OrchestrationError, SessionWorkflowStore, WorkflowChildRegistration, WorkflowChildRequest, + WorkflowChildResult, WorkflowEngine, WorkflowExecutor, WorkflowStore, +}; +pub use graph::scheduler_topology_preview; +pub use state::{ + PhaseStatus, all_phases_completed, init_phase_states, next_runnable_phase, phase_prompt, + phase_status, reset_running_phases, synthesize_summary, upstream_outputs, +}; +pub use types::{ + DefinitionError, WorkflowDefinition, WorkflowDefinitionListResponse, WorkflowPhase, +}; +pub use validate::{validate_agents, validate_structure}; + +#[cfg(test)] +mod tests; diff --git a/crates/tinyagents-orchestration/src/workflow/state.rs b/crates/tinyagents-orchestration/src/workflow/state.rs new file mode 100644 index 00000000..70a37d7d --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/state.rs @@ -0,0 +1,259 @@ +//! Phase state projection and advancement for workflow runs. +//! +//! This module owns the JSON phase-state document (a BTreeMap of phase names +//! to status + outputs) that the workflow engine persists. It provides queries +//! (what phases are runnable?) and mutations (mark complete, reset after +//! interruption) over this state without touching the underlying persistence layer. + +use serde_json::{Value, json}; + +use super::{WorkflowDefinition, WorkflowPhase}; + +/// Durable status of one phase in a workflow run. +/// +/// Phases transition: Pending → Running → (Completed | Failed). A phase +/// interrupted while running can be reset to Pending for retry; completed +/// phases are immutable. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PhaseStatus { + /// Phase has not yet started. + Pending, + /// Phase is currently executing (likely in a child). + Running, + /// Phase completed successfully; immutable. + Completed, + /// Phase failed; may be retried by resetting to Pending. + Failed, +} + +impl PhaseStatus { + /// Returns the string representation of this status (used in JSON). + pub const fn as_str(self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Running => "running", + Self::Completed => "completed", + Self::Failed => "failed", + } + } +} + +/// Initializes a phase-state document from a workflow definition. +/// +/// All phases start in the Pending status with empty outputs. The returned +/// JSON structure maps phase names to `{ status, outputs, ... }` objects. +pub fn init_phase_states(definition: &WorkflowDefinition) -> Value { + Value::Object( + definition + .phases + .iter() + .map(|phase| { + ( + phase.name.clone(), + json!({ "status": "pending", "outputs": [] }), + ) + }) + .collect(), + ) +} + +/// Queries the current status string of a phase, or None if the phase is unknown. +pub fn phase_status<'a>(phase_states: &'a Value, name: &str) -> Option<&'a str> { + phase_states.get(name)?.get("status")?.as_str() +} + +pub(crate) fn set_phase_status( + phase_states: &mut Value, + name: &str, + status: PhaseStatus, + outputs: Option, +) { + let Some(entries) = phase_states.as_object_mut() else { + return; + }; + let entry = entries + .entry(name.to_owned()) + .or_insert_with(|| json!({ "status": "pending", "outputs": [] })); + if let Some(object) = entry.as_object_mut() { + object.insert("status".to_owned(), json!(status.as_str())); + if let Some(outputs) = outputs { + object.insert("outputs".to_owned(), outputs); + } + } +} + +pub(crate) fn set_phase_reason(phase_states: &mut Value, name: &str, reason: &str) { + if let Some(object) = phase_states.get_mut(name).and_then(Value::as_object_mut) { + object.insert("reason".to_owned(), json!(reason)); + } +} + +/// Make an interrupted phase runnable again. A stopped phase may have spawned +/// children whose results were never durably collected; retrying the whole +/// phase is the only safe, deterministic recovery. Completed phases remain +/// immutable and are never retried. +pub fn reset_running_phases(phase_states: &mut Value, reason: &str) { + let Some(phases) = phase_states.as_object_mut() else { + return; + }; + for entry in phases.values_mut() { + let Some(state) = entry.as_object_mut() else { + continue; + }; + if state.get("status").and_then(Value::as_str) == Some("running") { + state.insert("status".to_owned(), json!(PhaseStatus::Pending.as_str())); + state.insert("outputs".to_owned(), json!([])); + state.insert("reason".to_owned(), json!(reason)); + } + } +} + +/// Finds the next phase that should run: a phase that is not already +/// completed, running, or failed, and all of whose dependencies are completed. +/// +/// Returns the first such phase in definition order, or `None` if no phase is +/// ready (either all are done, or some have unmet dependencies). +pub fn next_runnable_phase<'a>( + definition: &'a WorkflowDefinition, + phase_states: &Value, +) -> Option<&'a WorkflowPhase> { + definition.phases.iter().find(|phase| { + !matches!( + phase_status(phase_states, &phase.name), + Some("completed" | "running" | "failed") + ) && phase + .depends_on + .iter() + .all(|dependency| phase_status(phase_states, dependency) == Some("completed")) + }) +} + +/// Checks whether all phases in the workflow have completed. +pub fn all_phases_completed(definition: &WorkflowDefinition, phase_states: &Value) -> bool { + definition + .phases + .iter() + .all(|phase| phase_status(phase_states, &phase.name) == Some("completed")) +} + +/// Collects outputs from all upstream (dependency) phases for a given phase. +/// Filters out empty and null outputs to yield only meaningful results. +pub fn upstream_outputs(phase: &WorkflowPhase, phase_states: &Value) -> Vec { + phase + .depends_on + .iter() + .flat_map(|dependency| { + phase_states + .get(dependency) + .and_then(|entry| entry.get("outputs")) + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(move |item| { + durable_output(item) + .filter(|output| match output { + Value::Null => false, + Value::String(text) => !text.trim().is_empty(), + _ => true, + }) + .map(|output| json!({ "phase": dependency, "output": output })) + }) + }) + .collect() +} + +/// Composes the prompt for a worker in a phase. +/// +/// Includes the phase name and description, the input question, the worker's +/// index (if multiple workers), and relevant upstream outputs from dependencies. +pub fn phase_prompt( + input: &Value, + phase: &WorkflowPhase, + index: usize, + upstream: &[Value], +) -> String { + let question = input + .get("question") + .or_else(|| input.get("input")) + .and_then(Value::as_str) + .map(str::to_owned) + .unwrap_or_else(|| input.to_string()); + let mut prompt = format!( + "Workflow phase: {}\n{}\n\nInput:\n{}\n", + phase.name, phase.description, question + ); + if phase.agent_ids.len() > 1 { + prompt.push_str(&format!( + "\n(You are worker #{} in this phase.)\n", + index + 1 + )); + } + if !upstream.is_empty() { + prompt.push_str("\nContext from prior phases:\n"); + for item in upstream { + if let (Some(source), Some(output)) = ( + item.get("phase").and_then(Value::as_str), + item.get("output"), + ) { + prompt.push_str(&format!("- [{source}] {}\n", render_output(output))); + } + } + } + prompt +} + +/// Composes a workflow summary from all final phase outputs. +/// +/// Returns `None` if all phases are complete and there are no outputs; +/// otherwise returns a formatted summary of all non-empty outputs in phase order. +pub fn synthesize_summary(definition: &WorkflowDefinition, phase_states: &Value) -> Option { + let outputs_for = |name: &str| { + phase_states + .get(name) + .and_then(|entry| entry.get("outputs")) + .and_then(Value::as_array) + .map(|outputs| { + outputs + .iter() + .filter_map(|output| durable_output(output).map(|value| render_output(&value))) + .filter(|output| !output.trim().is_empty() && output != "null") + .collect::>() + .join("\n") + }) + .filter(|summary| !summary.trim().is_empty()) + }; + outputs_for("synthesize").or_else(|| { + definition + .phases + .iter() + .rev() + .find_map(|phase| outputs_for(&phase.name)) + }) +} + +/// The public/RPC projection remains `{ agentId, output: String }` for +/// compatibility. New rows carry the exact result in `metadata.rawOutput` so +/// future phases retain arbitrary JSON without changing the old wire shape. +fn durable_output(item: &Value) -> Option { + item.get("metadata") + .and_then(|metadata| metadata.get("version")) + .and_then(Value::as_u64) + .filter(|version| *version >= 2) + .and_then(|_| { + item.get("metadata") + .and_then(|metadata| metadata.get("rawOutput")) + }) + .cloned() + .or_else(|| item.get("output").cloned()) +} + +/// Preserve every JSON output in prompt context and summaries. JSON object's +/// map ordering is canonical under serde_json's default map implementation, +/// so repeated resume/synthesis renders the same bytes rather than silently +/// discarding structured child results. +fn render_output(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + _ => serde_json::to_string(value).unwrap_or_else(|_| "".to_owned()), + } +} diff --git a/crates/tinyagents-orchestration/src/workflow/tests.rs b/crates/tinyagents-orchestration/src/workflow/tests.rs new file mode 100644 index 00000000..5c0efea5 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/tests.rs @@ -0,0 +1,853 @@ +//! Tests for workflow execution: scheduling, phase transitions, concurrency, +//! and error handling. + +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use chrono::Utc; +use parking_lot::Mutex; +use serde_json::json; +use tinyagents_graph::CollectingSink; +use tinyagents_harness::CancellationToken; +use tinyagents_session::run_ledger::{ + WorkflowLeaseClaim, WorkflowRun, WorkflowRunStatus, WorkflowRunUpsert, +}; + +use super::state::set_phase_status; +use super::*; + +fn definition() -> WorkflowDefinition { + WorkflowDefinition { + id: "test".into(), + name: "Test".into(), + description: "test workflow".into(), + phases: vec![ + WorkflowPhase { + name: "plan".into(), + description: "plan".into(), + agent_ids: vec!["planner".into()], + depends_on: vec![], + }, + WorkflowPhase { + name: "research".into(), + description: "research".into(), + agent_ids: vec!["researcher".into(), "researcher".into()], + depends_on: vec!["plan".into()], + }, + WorkflowPhase { + name: "synthesize".into(), + description: "synthesize".into(), + agent_ids: vec!["writer".into()], + depends_on: vec!["research".into()], + }, + ], + default_concurrency: 2, + max_children: 8, + extensions: BTreeMap::new(), + } +} + +#[derive(Default)] +struct MemoryStore(Mutex>); + +impl WorkflowStore for MemoryStore { + fn load(&self, id: &str) -> Result, OrchestrationError> { + Ok(self.0.lock().get(id).cloned()) + } + + fn upsert(&self, update: WorkflowRunUpsert) -> Result { + let now = Utc::now(); + let prior = self.0.lock().get(&update.id).cloned(); + let row = WorkflowRun { + id: update.id.clone(), + definition_id: update.definition_id, + parent_thread_id: update.parent_thread_id, + input: update.input, + phase_states: update.phase_states, + child_run_ids: update.child_run_ids, + status: update.status, + summary: update + .summary + .or_else(|| prior.as_ref().and_then(|row| row.summary.clone())), + started_at: update + .started_at + .unwrap_or_else(|| prior.as_ref().map(|row| row.started_at).unwrap_or(now)), + updated_at: now, + completed_at: update + .completed_at + .or_else(|| prior.as_ref().and_then(|row| row.completed_at)), + revision: prior.as_ref().map_or(0, |row| row.revision + 1), + lease_owner: prior.as_ref().and_then(|row| row.lease_owner.clone()), + lease_expires_at: prior.as_ref().and_then(|row| row.lease_expires_at), + }; + self.0.lock().insert(row.id.clone(), row.clone()); + Ok(row) + } + + fn claim( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result { + let mut rows = self.0.lock(); + let Some(row) = rows.get_mut(id) else { + return Ok(WorkflowLeaseClaim::Missing); + }; + let now = Utc::now(); + if row + .lease_owner + .as_deref() + .is_some_and(|current| current != owner) + && row.lease_expires_at.is_some_and(|until| until > now) + { + return Ok(WorkflowLeaseClaim::Busy(row.clone())); + } + row.lease_owner = Some(owner.to_owned()); + row.lease_expires_at = chrono::Duration::from_std(lease_for) + .ok() + .map(|duration| now + duration); + row.revision += 1; + Ok(WorkflowLeaseClaim::Acquired(row.clone())) + } + + fn compare_and_swap( + &self, + update: WorkflowRunUpsert, + expected_revision: u64, + owner: &str, + lease_for: Duration, + ) -> Result, OrchestrationError> { + let mut rows = self.0.lock(); + let Some(prior) = rows.get(&update.id).cloned() else { + return Ok(None); + }; + if prior.revision != expected_revision || prior.lease_owner.as_deref() != Some(owner) { + return Ok(None); + } + let now = Utc::now(); + let row = WorkflowRun { + id: update.id, + definition_id: update.definition_id, + parent_thread_id: update.parent_thread_id, + input: update.input, + phase_states: update.phase_states, + child_run_ids: update.child_run_ids, + status: update.status, + summary: update.summary.or(prior.summary), + started_at: update.started_at.unwrap_or(prior.started_at), + updated_at: now, + completed_at: update.completed_at.or(prior.completed_at), + revision: prior.revision + 1, + lease_owner: (!update.status.is_terminal()).then(|| owner.to_owned()), + lease_expires_at: (!update.status.is_terminal()) + .then(|| { + chrono::Duration::from_std(lease_for) + .ok() + .map(|duration| now + duration) + }) + .flatten(), + }; + rows.insert(row.id.clone(), row.clone()); + Ok(Some(row)) + } + + fn renew( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result { + let mut rows = self.0.lock(); + let Some(row) = rows.get_mut(id) else { + return Ok(false); + }; + let now = Utc::now(); + if row.lease_owner.as_deref() != Some(owner) + || row.lease_expires_at.is_none_or(|expires| expires <= now) + { + return Ok(false); + } + row.lease_expires_at = chrono::Duration::from_std(lease_for) + .ok() + .map(|duration| now + duration); + Ok(true) + } +} + +#[derive(Default)] +struct RenewFailStore(MemoryStore); + +impl WorkflowStore for RenewFailStore { + fn load(&self, id: &str) -> Result, OrchestrationError> { + self.0.load(id) + } + + fn upsert(&self, update: WorkflowRunUpsert) -> Result { + self.0.upsert(update) + } + + fn claim( + &self, + id: &str, + owner: &str, + lease_for: Duration, + ) -> Result { + self.0.claim(id, owner, lease_for) + } + + fn compare_and_swap( + &self, + update: WorkflowRunUpsert, + expected_revision: u64, + owner: &str, + lease_for: Duration, + ) -> Result, OrchestrationError> { + self.0 + .compare_and_swap(update, expected_revision, owner, lease_for) + } + + fn renew( + &self, + _id: &str, + _owner: &str, + _lease_for: Duration, + ) -> Result { + Ok(false) + } +} + +#[derive(Default)] +struct FakeExecutor { + calls: Mutex>, + active: AtomicUsize, + peak: AtomicUsize, + fail_agent: Mutex>, + cancelled: AtomicUsize, +} + +#[derive(Default)] +struct BlockingExecutor { + started: tokio::sync::Notify, + cancelled: Mutex>, + calls: AtomicUsize, +} + +#[async_trait] +impl WorkflowExecutor for BlockingExecutor { + async fn execute( + &self, + request: WorkflowChildRequest, + cancel: CancellationToken, + registration: Arc, + ) -> Result { + let id = format!("live-{}-{}", request.phase, request.index_in_phase); + registration.register(id.clone())?; + self.calls.fetch_add(1, Ordering::SeqCst); + self.started.notify_one(); + cancel.cancelled().await; + Err(OrchestrationError("cancelled while child was live".into())) + } + + async fn cancel_children(&self, child_ids: &[String]) { + self.cancelled.lock().extend(child_ids.iter().cloned()); + } +} + +#[async_trait] +impl WorkflowExecutor for FakeExecutor { + async fn execute( + &self, + request: WorkflowChildRequest, + cancel: CancellationToken, + registration: Arc, + ) -> Result { + if cancel.is_cancelled() { + return Err(OrchestrationError("cancelled".into())); + } + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.peak.fetch_max(active, Ordering::SeqCst); + self.calls.lock().push(request.clone()); + registration.register(format!("{}-{}", request.phase, request.index_in_phase))?; + tokio::task::yield_now().await; + self.active.fetch_sub(1, Ordering::SeqCst); + if self.fail_agent.lock().as_deref() == Some(request.agent_id.as_str()) { + return Err(OrchestrationError("child failure".into())); + } + Ok(WorkflowChildResult { + child_id: format!("{}-{}", request.phase, request.index_in_phase), + output: json!(format!("{} output", request.phase)), + }) + } + + async fn cancel_children(&self, _child_ids: &[String]) { + self.cancelled.fetch_add(1, Ordering::SeqCst); + } +} + +fn engine() -> ( + Arc, + Arc, + WorkflowEngine, +) { + let store = Arc::new(MemoryStore::default()); + let executor = Arc::new(FakeExecutor::default()); + let engine = WorkflowEngine::new(store.clone(), executor.clone()); + (store, executor, engine) +} + +#[test] +fn structural_validation_covers_invalid_definitions() { + let mut empty = definition(); + empty.phases.clear(); + assert_eq!(validate_structure(&empty), vec![DefinitionError::NoPhases]); + let mut bad = definition(); + bad.default_concurrency = 0; + bad.max_children = 0; + bad.phases[1].name = "plan".into(); + bad.phases[2].depends_on = vec!["missing".into()]; + let errors = validate_structure(&bad); + assert!( + errors + .iter() + .any(|error| matches!(error, DefinitionError::DuplicatePhase { .. })) + ); + assert!( + errors + .iter() + .any(|error| matches!(error, DefinitionError::UnknownDependency { .. })) + ); + assert!( + errors + .iter() + .any(|error| matches!(error, DefinitionError::InvalidConcurrency { .. })) + ); + let mut cyclic = definition(); + cyclic.phases[0].depends_on = vec!["synthesize".into()]; + assert!(validate_structure(&cyclic).contains(&DefinitionError::CyclicDependency)); +} + +#[test] +fn scheduler_topology_preview_exposes_dispatch_run_and_done() { + let topology = scheduler_topology_preview().expect("topology"); + let nodes = topology + .nodes + .iter() + .map(|node| node.id.as_str()) + .collect::>(); + assert!(nodes.contains(&"dispatch") && nodes.contains(&"run_phase") && nodes.contains(&"done")); +} + +#[tokio::test] +async fn engine_runs_in_deterministic_dependency_order_and_threads_context() { + let (store, executor, engine) = engine(); + let def = definition(); + engine + .initialise("run".into(), &def, json!({"question":"q"}), None) + .unwrap(); + engine + .drive("run", &def, CancellationToken::new()) + .await + .unwrap(); + let run = store.load("run").unwrap().unwrap(); + assert_eq!(run.status, WorkflowRunStatus::Completed); + let calls = executor.calls.lock(); + assert_eq!( + calls + .iter() + .map(|call| call.phase.as_str()) + .collect::>(), + vec!["plan", "research", "research", "synthesize"] + ); + assert!( + calls + .last() + .unwrap() + .prompt + .contains("Context from prior phases") + ); + assert!(run.summary.unwrap().contains("synthesize output")); +} + +#[tokio::test] +async fn engine_respects_concurrency_global_cap_and_partial_failure() { + let (store, executor, engine) = engine(); + let mut def = definition(); + def.default_concurrency = 1; + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + engine + .drive("run", &def, CancellationToken::new()) + .await + .unwrap(); + assert!(executor.peak.load(Ordering::SeqCst) <= 1); + let mut cap = definition(); + cap.max_children = 2; + engine + .initialise("cap".into(), &cap, json!("q"), None) + .unwrap(); + engine + .drive("cap", &cap, CancellationToken::new()) + .await + .unwrap(); + assert_eq!( + store.load("cap").unwrap().unwrap().status, + WorkflowRunStatus::Failed + ); + *executor.fail_agent.lock() = Some("planner".into()); + let failing = definition(); + engine + .initialise("failed".into(), &failing, json!("q"), None) + .unwrap(); + engine + .drive("failed", &failing, CancellationToken::new()) + .await + .unwrap(); + assert_eq!( + store.load("failed").unwrap().unwrap().status, + WorkflowRunStatus::Failed + ); +} + +#[tokio::test] +async fn cancellation_and_resume_do_not_repeat_completed_phases() { + let (store, executor, engine) = engine(); + let def = definition(); + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + let cancel = CancellationToken::new(); + cancel.cancel(); + engine.drive("run", &def, cancel).await.unwrap(); + assert_eq!( + store.load("run").unwrap().unwrap().status, + WorkflowRunStatus::Interrupted + ); + let mut states = init_phase_states(&def); + set_phase_status( + &mut states, + "plan", + PhaseStatus::Completed, + Some(json!([{ "output": "already" }])), + ); + let run = store.load("run").unwrap().unwrap(); + store + .upsert(WorkflowRunUpsert { + id: run.id, + definition_id: run.definition_id, + parent_thread_id: run.parent_thread_id, + input: run.input, + phase_states: states, + child_run_ids: vec!["old".into()], + status: WorkflowRunStatus::Running, + summary: None, + started_at: Some(run.started_at), + completed_at: None, + }) + .unwrap(); + engine + .drive("run", &def, CancellationToken::new()) + .await + .unwrap(); + assert!( + !executor + .calls + .lock() + .iter() + .any(|call| call.phase == "plan") + ); + assert_eq!( + store.load("run").unwrap().unwrap().status, + WorkflowRunStatus::Completed + ); +} + +#[tokio::test] +async fn cancellation_after_workers_start_cancels_durably_registered_children() { + let store = Arc::new(MemoryStore::default()); + let executor = Arc::new(BlockingExecutor::default()); + let engine = Arc::new(WorkflowEngine::new(store.clone(), executor.clone())); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "live".into(), + description: "live".into(), + agent_ids: vec!["one".into(), "two".into()], + depends_on: vec![], + }]; + def.default_concurrency = 2; + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + let cancel = CancellationToken::new(); + let drive = { + let engine = engine.clone(); + let def = def.clone(); + let cancel = cancel.clone(); + tokio::spawn(async move { engine.drive("run", &def, cancel).await }) + }; + executor.started.notified().await; + // The notification occurs only after register(), so this races precisely + // the old orphan window between real child spawn and ledger persistence. + cancel.cancel(); + drive.await.unwrap().unwrap(); + let run = store.load("run").unwrap().unwrap(); + assert_eq!(run.status, WorkflowRunStatus::Interrupted); + assert_eq!( + run.phase_states["live"]["status"], + json!("pending"), + "an interrupted in-flight phase must be retryable on resume" + ); + assert!(!run.child_run_ids.is_empty()); + let cancelled = executor.cancelled.lock().clone(); + for id in &run.child_run_ids { + assert!(cancelled.contains(id), "missing cancellation for {id}"); + } +} + +#[tokio::test] +async fn concurrent_drives_acquire_one_lease_and_do_not_duplicate_children() { + let (store, executor, engine) = engine(); + let def = definition(); + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + let engine = Arc::new(engine); + let first = { + let engine = engine.clone(); + let def = def.clone(); + tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) + }; + let second = { + let engine = engine.clone(); + let def = def.clone(); + tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) + }; + first.await.unwrap().unwrap(); + second.await.unwrap().unwrap(); + assert_eq!( + store.load("run").unwrap().unwrap().status, + WorkflowRunStatus::Completed + ); + assert_eq!(executor.calls.lock().len(), 4, "one driver owns all phases"); +} + +#[tokio::test] +async fn expired_owner_takeover_resets_running_phase_and_retries_once() { + let (store, executor, engine) = engine(); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "recover".into(), + description: "recover".into(), + agent_ids: vec!["worker".into()], + depends_on: vec![], + }]; + engine + .initialise("takeover".into(), &def, json!("q"), None) + .unwrap(); + + // Simulate a process death after the owner has persisted `running`, but + // before it can complete or reset the phase. + let old = match store + .claim("takeover", "crashed-owner", Duration::from_millis(3)) + .unwrap() + { + WorkflowLeaseClaim::Acquired(run) => run, + other => panic!("expected lease, got {other:?}"), + }; + store + .compare_and_swap( + WorkflowRunUpsert { + id: old.id.clone(), + definition_id: old.definition_id.clone(), + parent_thread_id: old.parent_thread_id.clone(), + input: old.input.clone(), + phase_states: json!({ + "recover": {"status": "running", "outputs": []} + }), + child_run_ids: vec!["recover-0".into()], + status: WorkflowRunStatus::Running, + summary: None, + started_at: Some(old.started_at), + completed_at: None, + }, + old.revision, + "crashed-owner", + Duration::from_millis(3), + ) + .unwrap() + .expect("crashed owner persists phase start"); + tokio::time::sleep(Duration::from_millis(10)).await; + + engine + .drive("takeover", &def, CancellationToken::new()) + .await + .expect("new owner retries reclaimed phase"); + + let run = store.load("takeover").unwrap().expect("run remains"); + assert_eq!(run.status, WorkflowRunStatus::Completed); + assert_eq!(run.phase_states["recover"]["status"], json!("completed")); + assert_eq!( + executor + .calls + .lock() + .iter() + .filter(|call| call.phase == "recover") + .count(), + 1, + "lease takeover must schedule the reclaimed phase once" + ); + assert_eq!( + run.child_run_ids + .iter() + .filter(|id| id.as_str() == "recover-0") + .count(), + 1, + "the retry must not duplicate a durably registered child id" + ); +} + +#[tokio::test] +async fn heartbeat_renews_a_short_lease_while_a_child_is_running() { + let store = Arc::new(MemoryStore::default()); + let executor = Arc::new(BlockingExecutor::default()); + let engine = Arc::new( + WorkflowEngine::new(store.clone(), executor.clone()) + .with_lease_duration(Duration::from_millis(30)), + ); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "live".into(), + description: "live".into(), + agent_ids: vec!["one".into()], + depends_on: vec![], + }]; + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + let cancel = CancellationToken::new(); + let first = { + let engine = engine.clone(); + let def = def.clone(); + let cancel = cancel.clone(); + tokio::spawn(async move { engine.drive("run", &def, cancel).await }) + }; + executor.started.notified().await; + tokio::time::sleep(Duration::from_millis(75)).await; + // This exceeds the original lease, so it only remains busy if the + // in-flight driver's heartbeat kept renewing it. + engine + .drive("run", &def, CancellationToken::new()) + .await + .unwrap(); + assert_eq!(executor.calls.load(Ordering::SeqCst), 1); + cancel.cancel(); + first.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn lost_heartbeat_cancels_registered_children_and_fails_closed() { + let store = Arc::new(RenewFailStore::default()); + let executor = Arc::new(BlockingExecutor::default()); + let engine = Arc::new( + WorkflowEngine::new(store.clone(), executor.clone()) + .with_lease_duration(Duration::from_millis(30)), + ); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "live".into(), + description: "live".into(), + agent_ids: vec!["one".into()], + depends_on: vec![], + }]; + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + let drive = { + let engine = engine.clone(); + let def = def.clone(); + tokio::spawn(async move { engine.drive("run", &def, CancellationToken::new()).await }) + }; + executor.started.notified().await; + let error = drive + .await + .unwrap() + .expect_err("renewal loss must fail closed"); + assert!(error.0.contains("lease renewal failed")); + assert!( + executor + .cancelled + .lock() + .iter() + .any(|id| id == "live-live-0"), + "the child registered before waiting must be cancelled" + ); +} + +#[tokio::test] +async fn terminal_events_are_truthful_and_flushed() { + let (store, executor, engine) = engine(); + let sink = Arc::new(CollectingSink::new()); + let engine = engine.with_event_sink(sink.clone()); + let def = definition(); + engine + .initialise("ok".into(), &def, json!("q"), None) + .unwrap(); + engine + .drive("ok", &def, CancellationToken::new()) + .await + .unwrap(); + assert!( + sink.events() + .iter() + .any(|event| matches!(event, tinyagents_graph::GraphEvent::RunCompleted { .. })) + ); + + *executor.fail_agent.lock() = Some("planner".into()); + engine + .initialise("failed".into(), &def, json!("q"), None) + .unwrap(); + engine + .drive("failed", &def, CancellationToken::new()) + .await + .unwrap(); + assert!(sink.events().iter().any(|event| matches!( + event, + tinyagents_graph::GraphEvent::RunFailed { error, .. } if error.contains("child failure") + ))); + assert_eq!( + store.load("failed").unwrap().unwrap().status, + WorkflowRunStatus::Failed + ); +} + +#[tokio::test] +async fn fenced_driver_exits_silently_when_a_replacement_is_running() { + let store = Arc::new(RenewFailStore::default()); + let executor = Arc::new(BlockingExecutor::default()); + let sink = Arc::new(CollectingSink::new()); + let engine = Arc::new( + WorkflowEngine::new(store.clone(), executor.clone()) + .with_lease_duration(Duration::from_millis(30)) + .with_event_sink(sink.clone()), + ); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "live".into(), + description: "live".into(), + agent_ids: vec!["one".into()], + depends_on: vec![], + }]; + engine + .initialise("handoff".into(), &def, json!("q"), None) + .unwrap(); + let first = { + let engine = engine.clone(); + let def = def.clone(); + tokio::spawn(async move { + engine + .drive("handoff", &def, CancellationToken::new()) + .await + }) + }; + executor.started.notified().await; + + // A replacement owner acquires after expiry before the old driver's + // heartbeat sees its loss. The old loop must not flush a false failure. + let current = store.load("handoff").unwrap().unwrap(); + store + .compare_and_swap( + WorkflowRunUpsert { + id: current.id.clone(), + definition_id: current.definition_id.clone(), + parent_thread_id: current.parent_thread_id.clone(), + input: current.input.clone(), + phase_states: current.phase_states.clone(), + child_run_ids: current.child_run_ids.clone(), + status: WorkflowRunStatus::Running, + summary: None, + started_at: Some(current.started_at), + completed_at: None, + }, + current.revision, + current.lease_owner.as_deref().unwrap(), + Duration::from_millis(3), + ) + .unwrap() + .unwrap(); + tokio::time::sleep(Duration::from_millis(5)).await; + match store + .claim("handoff", "replacement", Duration::from_secs(1)) + .unwrap() + { + WorkflowLeaseClaim::Acquired(run) => { + assert_eq!(run.lease_owner.as_deref(), Some("replacement")); + } + other => panic!("expected replacement lease, got {other:?}"), + } + first.await.unwrap().expect("fenced driver exits cleanly"); + assert!( + !sink + .events() + .iter() + .any(|event| matches!(event, tinyagents_graph::GraphEvent::RunFailed { .. })), + "a fenced driver must not report the active replacement as failed" + ); +} + +#[test] +fn structured_outputs_are_preserved_in_context_and_summary() { + let def = definition(); + let mut states = init_phase_states(&def); + set_phase_status( + &mut states, + "plan", + PhaseStatus::Completed, + Some(json!([{ + "output": { "claims": ["a", "b"], "score": 7 } + }])), + ); + let upstream = upstream_outputs(&def.phases[1], &states); + let prompt = phase_prompt(&json!("q"), &def.phases[1], 0, &upstream); + assert!(prompt.contains(r#"{"claims":["a","b"],"score":7}"#)); + set_phase_status( + &mut states, + "synthesize", + PhaseStatus::Completed, + Some(json!([{ + "output": { "answer": "kept" } + }])), + ); + assert_eq!( + synthesize_summary(&def, &states).as_deref(), + Some(r#"{"answer":"kept"}"#) + ); +} + +#[tokio::test] +async fn output_wire_shape_remains_compatible_while_json_stays_lossless() { + let (store, _executor, engine) = engine(); + let mut def = definition(); + def.phases = vec![WorkflowPhase { + name: "only".into(), + description: "only".into(), + agent_ids: vec!["planner".into()], + depends_on: vec![], + }]; + engine + .initialise("run".into(), &def, json!("q"), None) + .unwrap(); + engine + .drive("run", &def, CancellationToken::new()) + .await + .unwrap(); + let output = store.load("run").unwrap().unwrap().phase_states["only"]["outputs"][0].clone(); + assert_eq!(output["agentId"], json!("planner")); + assert!(output["output"].is_string()); + assert_eq!(output["metadata"]["version"], json!(2)); + assert_eq!(output["metadata"]["rawOutput"], json!("only output")); +} diff --git a/crates/tinyagents-orchestration/src/workflow/types.rs b/crates/tinyagents-orchestration/src/workflow/types.rs new file mode 100644 index 00000000..bda9d287 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/types.rs @@ -0,0 +1,87 @@ +//! Workflow definition types: phases, definitions, and validation errors. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// One phase of a declarative workflow. +/// +/// A phase specifies a set of agents to work on it concurrently and its +/// dependencies (other phases that must complete first). Agents run in +/// parallel within a phase; the phase itself is the unit of scheduling. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowPhase { + /// The human-readable name of this phase. + pub name: String, + /// A description of what this phase does. + pub description: String, + /// Agent ids that should work on this phase in parallel. + pub agent_ids: Vec, + /// Phase names this phase depends on (must complete first). + pub depends_on: Vec, +} + +/// A host-neutral declarative phase DAG. +/// +/// Defines a workflow as a directed acyclic graph of phases, each with +/// associated agents, concurrency limits, and dependencies. The orchestration +/// engine schedules phases topologically and manages bounded parallelism via +/// `default_concurrency` and `max_children`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowDefinition { + /// Unique identifier for this workflow. + pub id: String, + /// Human-readable workflow name. + pub name: String, + /// Description of what the workflow does. + pub description: String, + /// All phases in the workflow. + pub phases: Vec, + /// Default concurrency limit for agents in phases that don't specify one. + pub default_concurrency: u32, + /// Maximum number of child tasks that can be spawned concurrently. + pub max_children: u32, + /// Host-defined wire metadata which the orchestration engine never reads. + #[serde(flatten, default)] + pub extensions: BTreeMap, +} + +/// List response used by hosts that expose a workflow catalog. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowDefinitionListResponse { + /// The workflow definitions in this response. + pub definitions: Vec, + /// Total count of workflows in the host's catalog. + pub count: usize, +} + +/// A structural or host-supplied lookup problem in a workflow definition. +/// +/// These errors detect issues that prevent a workflow from running: +/// unknown agent references, missing dependencies, cycles, or invalid +/// concurrency settings. They are host-independent validation issues. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum DefinitionError { + /// A phase references an agent the host does not know about. + UnknownAgent { phase: String, agent_id: String }, + /// A phase depends on another phase that is not defined. + UnknownDependency { phase: String, depends_on: String }, + /// Multiple phases have the same name. + DuplicatePhase { name: String }, + /// A phase has no agents assigned to it. + EmptyPhase { phase: String }, + /// The phase dependencies form a cycle. + CyclicDependency, + /// The workflow has no phases. + NoPhases, + /// The concurrency settings are invalid (e.g. default > max). + InvalidConcurrency { + default_concurrency: u32, + max_children: u32, + }, +} diff --git a/crates/tinyagents-orchestration/src/workflow/validate.rs b/crates/tinyagents-orchestration/src/workflow/validate.rs new file mode 100644 index 00000000..fb525441 --- /dev/null +++ b/crates/tinyagents-orchestration/src/workflow/validate.rs @@ -0,0 +1,71 @@ +//! Workflow definition validation: structural checks and error detection. +//! +//! Validates a workflow definition for structural issues (missing phases, +//! duplicate names, cycles, invalid concurrency) and host-specific issues +//! (unknown agents). Errors are deterministic and host-independent. + +use tinyagents_graph::dag::{DagIssue, DagNode, validate_dag}; + +use super::{DefinitionError, WorkflowDefinition}; + +/// Validate properties that do not require a host agent registry. +pub fn validate_structure(definition: &WorkflowDefinition) -> Vec { + if definition.phases.is_empty() { + return vec![DefinitionError::NoPhases]; + } + + let mut errors = definition + .phases + .iter() + .filter(|phase| phase.agent_ids.is_empty()) + .map(|phase| DefinitionError::EmptyPhase { + phase: phase.name.clone(), + }) + .collect::>(); + let nodes = definition + .phases + .iter() + .map(|phase| { + DagNode::new( + phase.name.as_str(), + phase.depends_on.iter().map(String::as_str), + ) + }) + .collect::>(); + errors.extend(validate_dag(&nodes).into_iter().map(|issue| match issue { + DagIssue::DuplicateNode { id } => DefinitionError::DuplicatePhase { name: id }, + DagIssue::UnknownDependency { node, depends_on } => DefinitionError::UnknownDependency { + phase: node, + depends_on, + }, + DagIssue::Cycle => DefinitionError::CyclicDependency, + })); + if definition.default_concurrency == 0 || definition.max_children == 0 { + errors.push(DefinitionError::InvalidConcurrency { + default_concurrency: definition.default_concurrency, + max_children: definition.max_children, + }); + } + errors +} + +/// Validate agent identifiers using a host-owned registry lookup. +pub fn validate_agents(definition: &WorkflowDefinition, is_known: F) -> Vec +where + F: Fn(&str) -> bool, +{ + definition + .phases + .iter() + .flat_map(|phase| { + phase + .agent_ids + .iter() + .filter(|agent_id| !is_known(agent_id)) + .map(|agent_id| DefinitionError::UnknownAgent { + phase: phase.name.clone(), + agent_id: agent_id.clone(), + }) + }) + .collect() +}