From 0f9e31204fd835a8e27e53668ff1521c19272a71 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Thu, 26 Feb 2026 13:37:08 -0500 Subject: [PATCH] Add TOML-based task config for `attractor run` Support `attractor run task.toml` as an alternative to `.dot` files. Auto-detects format by file extension. TOML bundles pipeline config (graph path, model, setup commands, working directory) into a single file with precedence: CLI flag > TOML > DOT graph attrs > defaults. Co-Authored-By: Claude Opus 4.6 Entire-Checkpoint: 8744d57d6053 --- Cargo.lock | 1 + Cargo.toml | 1 + crates/attractor/Cargo.toml | 1 + crates/attractor/src/cli/mod.rs | 5 +- crates/attractor/src/cli/run.rs | 61 ++++++- crates/attractor/src/cli/task_config.rs | 230 ++++++++++++++++++++++++ 6 files changed, 289 insertions(+), 10 deletions(-) create mode 100644 crates/attractor/src/cli/task_config.rs diff --git a/Cargo.lock b/Cargo.lock index 12de16e15..f7606a43e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -215,6 +215,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "toml", "tower", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index 9d1a9f3f8..f82520f50 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,3 +33,4 @@ git2 = "0.19" walkdir = "2" regex = "1" aho-corasick = "1" +toml = "0.8" diff --git a/crates/attractor/Cargo.toml b/crates/attractor/Cargo.toml index 96ae8796b..090d0010e 100644 --- a/crates/attractor/Cargo.toml +++ b/crates/attractor/Cargo.toml @@ -38,6 +38,7 @@ async-trait.workspace = true futures.workspace = true chrono = { workspace = true, features = ["serde"] } nom = "7" +toml.workspace = true dirs = "6" dialoguer.workspace = true axum = { version = "0.8", optional = true } diff --git a/crates/attractor/src/cli/mod.rs b/crates/attractor/src/cli/mod.rs index 1e8ed9663..86dbce6b9 100644 --- a/crates/attractor/src/cli/mod.rs +++ b/crates/attractor/src/cli/mod.rs @@ -2,6 +2,7 @@ pub mod backend; pub mod run; #[cfg(feature = "server")] pub mod serve; +pub mod task_config; pub mod validate; use std::path::Path; @@ -24,7 +25,7 @@ pub struct Cli { #[derive(Subcommand)] pub enum Command { - /// Launch a pipeline from a .dot file + /// Launch a pipeline from a .dot or .toml task file Run(RunArgs), /// Parse and validate a pipeline without executing Validate(ValidateArgs), @@ -35,7 +36,7 @@ pub enum Command { #[derive(Args)] pub struct RunArgs { - /// Path to the .dot pipeline file + /// Path to a .dot pipeline file or .toml task config pub pipeline: PathBuf, /// Log/artifact directory diff --git a/crates/attractor/src/cli/run.rs b/crates/attractor/src/cli/run.rs index ed737e74d..15ba1aa79 100644 --- a/crates/attractor/src/cli/run.rs +++ b/crates/attractor/src/cli/run.rs @@ -17,6 +17,7 @@ use crate::pipeline::PipelineBuilder; use crate::validation::Severity; use super::backend::AgentBackend; +use super::task_config; use super::{compute_stage_cost, format_cost, format_duration_human, format_event_detail, format_event_summary, format_tokens_human, print_diagnostics, read_dot_file, RunArgs}; /// Accumulates token usage and cost across all pipeline stages. @@ -37,8 +38,33 @@ struct CostAccumulator { /// /// Returns an error if the pipeline cannot be read, parsed, validated, or executed. pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Result<()> { + // 0. Load task config if TOML, resolve DOT path, run setup + let (dot_path, task_cfg) = if args.pipeline.extension().is_some_and(|ext| ext == "toml") { + let cfg = task_config::load_task_config(&args.pipeline)?; + let dot = task_config::resolve_graph_path(&args.pipeline, &cfg.graph); + (dot, Some(cfg)) + } else { + (args.pipeline.clone(), None) + }; + + if let Some(ref cfg) = task_cfg { + if let Some(ref dir) = cfg.directory { + std::env::set_current_dir(dir) + .map_err(|e| anyhow::anyhow!("Failed to set working directory to {dir}: {e}"))?; + } + if let Some(ref setup) = cfg.setup { + let cwd = std::env::current_dir()?; + eprintln!( + "{dim}Running {} setup command(s)…{reset}", + setup.commands.len(), + dim = styles.dim, reset = styles.reset, + ); + task_config::run_setup(setup, &cwd).await?; + } + } + // 1. Parse and validate pipeline - let source = read_dot_file(&args.pipeline)?; + let source = read_dot_file(&dot_path)?; let (graph, diagnostics) = PipelineBuilder::new().prepare(&source)?; eprintln!( @@ -74,6 +100,11 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu tokio::fs::create_dir_all(&logs_dir).await?; tokio::fs::write(logs_dir.join("graph.dot"), &source).await?; tokio::fs::write(logs_dir.join("run.pid"), std::process::id().to_string()).await?; + if args.pipeline.extension().is_some_and(|ext| ext == "toml") { + if let Ok(toml_contents) = tokio::fs::read(&args.pipeline).await { + tokio::fs::write(logs_dir.join("task.toml"), toml_contents).await?; + } + } if args.verbose >= 1 { eprintln!( @@ -208,16 +239,30 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu } }; - let provider = args.provider.or_else(|| { - graph - .attrs - .get("default_provider") - .and_then(|v| v.as_str()) - .map(String::from) - }); + let toml_model = task_cfg + .as_ref() + .and_then(|c| c.llm.as_ref()) + .and_then(|l| l.model.clone()); + let toml_provider = task_cfg + .as_ref() + .and_then(|c| c.llm.as_ref()) + .and_then(|l| l.provider.clone()); + + // Precedence: CLI flag > TOML > DOT graph attrs > defaults + let provider = args + .provider + .or(toml_provider) + .or_else(|| { + graph + .attrs + .get("default_provider") + .and_then(|v| v.as_str()) + .map(String::from) + }); let model = args .model + .or(toml_model) .or_else(|| { graph .attrs diff --git a/crates/attractor/src/cli/task_config.rs b/crates/attractor/src/cli/task_config.rs new file mode 100644 index 000000000..0df77f10d --- /dev/null +++ b/crates/attractor/src/cli/task_config.rs @@ -0,0 +1,230 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{bail, Context}; +use serde::Deserialize; + +const SUPPORTED_VERSION: u32 = 1; + +#[derive(Debug, Deserialize)] +pub struct TaskConfig { + pub version: u32, + pub task: String, + pub graph: String, + pub directory: Option, + pub llm: Option, + pub setup: Option, +} + +#[derive(Debug, Deserialize)] +pub struct LlmConfig { + pub model: Option, + pub provider: Option, +} + +#[derive(Debug, Deserialize)] +pub struct SetupConfig { + pub commands: Vec, + pub timeout_ms: Option, +} + +/// Load and validate a task config from a TOML file. +/// +/// The `graph` path in the returned config is resolved relative to the +/// TOML file's parent directory. +pub fn load_task_config(path: &Path) -> anyhow::Result { + let contents = + std::fs::read_to_string(path).with_context(|| format!("Failed to read {}", path.display()))?; + let config = parse_task_config(&contents)?; + + Ok(config) +} + +/// Resolve the graph path relative to the TOML file's parent directory. +pub fn resolve_graph_path(toml_path: &Path, graph: &str) -> PathBuf { + let graph_path = Path::new(graph); + if graph_path.is_absolute() { + graph_path.to_path_buf() + } else { + toml_path + .parent() + .unwrap_or(Path::new(".")) + .join(graph_path) + } +} + +fn parse_task_config(contents: &str) -> anyhow::Result { + let config: TaskConfig = + toml::from_str(contents).context("Failed to parse task config TOML")?; + + if config.version != SUPPORTED_VERSION { + bail!( + "Unsupported task config version {}. Only version {SUPPORTED_VERSION} is supported.", + config.version + ); + } + + Ok(config) +} + +/// Run setup commands sequentially in the given directory. +/// +/// Each command gets the full `timeout_ms` budget. Commands are executed +/// via `sh -c` so shell features (pipes, redirects, etc.) work. +pub async fn run_setup(setup: &SetupConfig, directory: &Path) -> anyhow::Result<()> { + let timeout = std::time::Duration::from_millis(setup.timeout_ms.unwrap_or(300_000)); + + for cmd in &setup.commands { + let fut = tokio::process::Command::new("sh") + .arg("-c") + .arg(cmd) + .current_dir(directory) + .output(); + + let output = tokio::time::timeout(timeout, fut) + .await + .with_context(|| format!("Setup command timed out after {}ms: {cmd}", timeout.as_millis()))? + .with_context(|| format!("Failed to execute setup command: {cmd}"))?; + + if !output.status.success() { + let code = output + .status + .code() + .map_or("unknown".to_string(), |c| c.to_string()); + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("Setup command failed (exit code {code}): {cmd}\n{stderr}"); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_minimal_toml() { + let toml = r#" +version = 1 +task = "Run tests" +graph = "pipeline.dot" +"#; + let config = parse_task_config(toml).unwrap(); + assert_eq!(config.version, 1); + assert_eq!(config.task, "Run tests"); + assert_eq!(config.graph, "pipeline.dot"); + assert!(config.directory.is_none()); + assert!(config.llm.is_none()); + assert!(config.setup.is_none()); + } + + #[test] + fn parse_full_toml() { + let toml = r#" +version = 1 +task = "Full workflow" +graph = "pipeline.dot" +directory = "/tmp/repo" + +[llm] +model = "claude-haiku" +provider = "anthropic" + +[setup] +commands = ["pip install -r requirements.txt", "npm install"] +timeout_ms = 60000 +"#; + let config = parse_task_config(toml).unwrap(); + assert_eq!(config.task, "Full workflow"); + assert_eq!(config.directory.as_deref(), Some("/tmp/repo")); + + let llm = config.llm.unwrap(); + assert_eq!(llm.model.as_deref(), Some("claude-haiku")); + assert_eq!(llm.provider.as_deref(), Some("anthropic")); + + let setup = config.setup.unwrap(); + assert_eq!(setup.commands.len(), 2); + assert_eq!(setup.timeout_ms, Some(60000)); + } + + #[test] + fn unsupported_version_rejected() { + let toml = r#" +version = 2 +task = "x" +graph = "p.dot" +"#; + let err = parse_task_config(toml).unwrap_err(); + assert!( + err.to_string().contains("Unsupported task config version 2"), + "unexpected error: {err}" + ); + } + + #[test] + fn graph_path_resolved_relative_to_toml() { + let toml_path = Path::new("/tmp/sub/task.toml"); + let resolved = resolve_graph_path(toml_path, "p.dot"); + assert_eq!(resolved, PathBuf::from("/tmp/sub/p.dot")); + } + + #[test] + fn graph_path_absolute_unchanged() { + let toml_path = Path::new("/tmp/sub/task.toml"); + let resolved = resolve_graph_path(toml_path, "/other/pipeline.dot"); + assert_eq!(resolved, PathBuf::from("/other/pipeline.dot")); + } + + #[test] + fn missing_required_fields() { + let no_task = r#" +version = 1 +graph = "p.dot" +"#; + assert!(parse_task_config(no_task).is_err()); + + let no_graph = r#" +version = 1 +task = "x" +"#; + assert!(parse_task_config(no_graph).is_err()); + } + + #[tokio::test] + async fn run_setup_succeeds() { + let dir = tempfile::tempdir().unwrap(); + let setup = SetupConfig { + commands: vec!["echo hello".to_string()], + timeout_ms: None, + }; + run_setup(&setup, dir.path()).await.unwrap(); + } + + #[tokio::test] + async fn run_setup_fails_on_nonzero_exit() { + let dir = tempfile::tempdir().unwrap(); + let setup = SetupConfig { + commands: vec!["false".to_string()], + timeout_ms: None, + }; + let err = run_setup(&setup, dir.path()).await.unwrap_err(); + assert!( + err.to_string().contains("exit code"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn run_setup_timeout() { + let dir = tempfile::tempdir().unwrap(); + let setup = SetupConfig { + commands: vec!["sleep 10".to_string()], + timeout_ms: Some(100), + }; + let err = run_setup(&setup, dir.path()).await.unwrap_err(); + assert!( + err.to_string().contains("timed out"), + "unexpected error: {err}" + ); + } +}