diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 002ebfc..f1e89e4 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -21,9 +21,18 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::Semaphore; +/// Test-only hook invoked inside a frontier branch task right after its +/// teammate identity is retired, with the super-step's registry and the +/// retired peer id. Runs before the task returns, so an observer sees the +/// registry as siblings still in flight see it. +#[cfg(test)] +type FrontierObserver = Arc; + pub struct GraphExecutor { graph: Graph, base_dir: PathBuf, + #[cfg(test)] + frontier_observer: Option, } impl GraphExecutor { @@ -31,9 +40,17 @@ impl GraphExecutor { Self { graph, base_dir: base_dir.into(), + #[cfg(test)] + frontier_observer: None, } } + #[cfg(test)] + fn with_frontier_observer(mut self, observer: FrontierObserver) -> Self { + self.frontier_observer = Some(observer); + self + } + pub async fn execute( self, ctx: &mut RequestContext, @@ -58,7 +75,12 @@ impl GraphExecutor { ctx: &mut RequestContext, abort_signal: AbortSignal, ) -> Result { - let GraphExecutor { graph, base_dir } = self; + let GraphExecutor { + graph, + base_dir, + #[cfg(test)] + frontier_observer, + } = self; if graph.settings.validate_before_run { let mut validator = GraphValidator::new(&base_dir); @@ -194,6 +216,8 @@ impl GraphExecutor { let current = node_id.clone(); let sem_clone = semaphore.clone(); let abort_clone = abort_signal.clone(); + #[cfg(test)] + let observer = frontier_observer.clone(); let task = tokio::spawn(async move { let _permit = sem_clone @@ -228,6 +252,10 @@ impl GraphExecutor { let result = step(&node, &mut state, &mut ctx, &step_ctx, ¤t).await; if let (Some(registry), Some(id)) = (®istry_for_task, &peer_id) { registry.mark_finished(id); + #[cfg(test)] + if let Some(observe) = &observer { + observe(registry, id); + } } let elapsed = node_start.elapsed(); match &result { @@ -795,6 +823,7 @@ mod integration_tests { #[cfg(unix)] use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor, notification}; use crate::utils::{create_abort_signal, temp_file}; + use parking_lot::Mutex; use std::fs; #[cfg(unix)] use std::mem; @@ -1007,6 +1036,106 @@ nodes: ); } + /// Two flagged agent nodes share one super-step under a single permit, so + /// the first branch retires its identity while its sibling has not even + /// started. The observer runs inside the branch task, before that task + /// returns, so the snapshot it takes is what a still-running sibling + /// would see if it messaged the finished node. + #[tokio::test] + async fn frontier_branch_is_finished_before_super_step_join() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n"); + + let yaml = r#" +name: frontier_finish_test +start: dispatcher +settings: + max_concurrency: 1 + validate_before_run: false +nodes: + dispatcher: + type: script + script: dispatcher.sh + state_updates: {} + next: [worker_a, worker_b] + worker_a: + type: agent + agent: no-such-agent + prompt: "p" + teammates: true + next: join + worker_b: + type: agent + agent: no-such-agent + prompt: "p" + teammates: true + next: join + join: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + + // Each observation: (retired peer id, every roster (label, is_finished)). + type RosterSnapshot = Vec<(String, bool)>; + let observations: Arc>> = Arc::default(); + let sink = Arc::clone(&observations); + let observer: FrontierObserver = Arc::new(move |registry, id| { + let snapshot = registry + .roster() + .into_iter() + .map(|(peer_id, label)| (label, registry.is_finished(&peer_id))) + .collect(); + sink.lock().push((id.to_string(), snapshot)); + }); + + let result = GraphExecutor::new(graph, &ws.dir) + .with_frontier_observer(observer) + .execute(&mut ctx, abort) + .await; + + assert!(result.is_err(), "both agent branches fail at Agent::init"); + let observations = observations.lock(); + assert_eq!( + observations.len(), + 2, + "one observation per flagged branch: {observations:?}" + ); + + let (_, first) = &observations[0]; + let finished_first: Vec<&str> = first + .iter() + .filter(|(_, finished)| *finished) + .map(|(label, _)| label.as_str()) + .collect(); + assert_eq!( + finished_first.len(), + 1, + "the first branch to complete is retired while its sibling is still \ + unfinished, so this read happened before the super-step join: {first:?}" + ); + assert!( + first.iter().all(|(label, _)| label.starts_with("worker_")), + "{first:?}" + ); + + let (_, second) = &observations[1]; + assert!( + second.iter().all(|(_, finished)| *finished), + "both identities are retired once the last branch completes: {second:?}" + ); + assert_ne!( + observations[0].0, observations[1].0, + "each branch retires its own identity exactly once" + ); + } + #[tokio::test] async fn multi_end_in_super_step_is_rejected() { if !cmd_available("bash") { diff --git a/src/graph/map.rs b/src/graph/map.rs index f3a1cd5..e61eb1d 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -568,12 +568,16 @@ mod chain_tests { use super::super::executor::GraphExecutor; use super::super::script::ScriptExecutor; use super::*; + use crate::config::paths; use crate::config::{AppState, Role, WorkingMode}; - use crate::utils::{AbortSignal, create_abort_signal, temp_file}; + use crate::utils::{AbortSignal, create_abort_signal, get_env_name, temp_file}; use indexmap::IndexMap; use serde_json::json; + use serial_test::serial; + use std::env; use std::fs; use std::path::PathBuf; + use std::time::{SystemTime, UNIX_EPOCH}; fn cmd_available(name: &str) -> bool { which::which(name).is_ok() @@ -1171,6 +1175,89 @@ nodes: } } + struct TestConfigDirGuard { + key: String, + previous: Option, + path: PathBuf, + } + + impl TestConfigDirGuard { + fn new() -> Self { + let key = get_env_name("config_dir"); + let previous = env::var_os(&key); + let unique = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let path = env::temp_dir().join(format!("coyote-graph-map-tests-{unique}")); + fs::create_dir_all(&path).unwrap(); + unsafe { + env::set_var(&key, &path); + } + Self { + key, + previous, + path, + } + } + } + + impl Drop for TestConfigDirGuard { + fn drop(&mut self) { + if let Some(previous) = &self.previous { + unsafe { + env::set_var(&self.key, previous); + } + } else { + unsafe { + env::remove_var(&self.key); + } + } + let _ = fs::remove_dir_all(&self.path); + } + } + + const PROBE_AGENT: &str = "timeout-probe"; + + /// Materializes a graph agent under the guarded config dir that + /// `Agent::init` can load offline: no global tools, variables, MCP servers, + /// or rag nodes, so nothing in `run_agent_for_graph` bails before the + /// agent's own graph starts executing. Its single script node sleeps for + /// `hold_secs` and then writes `output`, so the whole run needs no LLM and + /// the agent future can be held open for as long as a test needs. + fn materialize_probe_agent(hold_secs: f64) { + let graph_path = paths::agent_graph_file(PROBE_AGENT); + let agent_dir = graph_path.parent().unwrap(); + fs::create_dir_all(agent_dir).unwrap(); + fs::write( + agent_dir.join("hold.py"), + format!( + "#!/usr/bin/env python3\nimport json, time\ntime.sleep({hold_secs})\n\ + print(json.dumps({{\"output\": \"held\"}}))\n" + ), + ) + .unwrap(); + fs::write( + &graph_path, + format!( + r#" +name: {PROBE_AGENT} +start: hold +nodes: + hold: + type: script + script: hold.py + timeout: 30 + next: done + done: + type: end + output: "{{{{output}}}}" +"# + ), + ) + .unwrap(); + } + const DOUBLER_GRAPH: &str = r#" name: t start: doubler @@ -1308,6 +1395,94 @@ nodes: assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); } + fn flagged_probe_graph(timeout: Option) -> String { + let timeout_line = timeout + .map(|t| format!(" timeout: {t}\n")) + .unwrap_or_default(); + format!( + r#" +name: t +start: worker +nodes: + worker: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + teammates: true + state_updates: + output: "{{{{output}}}}" +{timeout_line}"# + ) + } + + /// Control for the timeout test below: the same materialized agent, with + /// the default timeout and no hold, runs its own graph to completion + /// offline and hands its output back through the chain. That pins the + /// recipe, so a `timed out` error in the sibling test cannot be an init + /// failure in disguise. + #[tokio::test] + #[serial] + async fn run_item_chain_probe_agent_runs_its_graph_offline() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let h = Harness::new(&flagged_probe_graph(None)); + let peers = manual_peer("worker"); + let mut state = item_state(json!(0)); + let mut ctx = silent_ctx(); + + h.run("worker", Some(&peers), &mut state, &mut ctx) + .await + .unwrap_or_else(|e| panic!("chain failed: {e:#}")); + + assert_eq!(state.state().get("output"), Some(&json!("held"))); + assert!(peers.0.is_finished(&peers.1.0)); + assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); + } + + /// `tokio::time::timeout` polls the agent future first and a zero-duration + /// sleep still goes through the time driver, so `Elapsed` only fires if the + /// agent future is still pending when the timer is next polled. The + /// materialized agent's graph holds its script step open for seconds, so + /// the future is dropped mid-flight and the chain runner alone retires + /// the identity. + #[tokio::test] + #[serial] + async fn run_item_chain_marks_identity_finished_when_agent_step_times_out() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(5.0); + let h = Harness::new(&flagged_probe_graph(Some(0))); + let (registry, assignments) = h.provision("worker", 2); + let peers = (Arc::clone(®istry), assignments[0].clone()); + let mut state = item_state(json!(0)); + let mut ctx = silent_ctx(); + + let chain = unwrap_err_chain(h.run("worker", Some(&peers), &mut state, &mut ctx).await); + + assert!( + chain.contains("failed at node 'worker' (step 1)"), + "{chain}" + ); + assert!( + chain.contains(&format!("Agent '{PROBE_AGENT}' timed out after 0s")), + "{chain}" + ); + assert!( + !chain.contains(&format!("Agent '{PROBE_AGENT}' failed")), + "{chain}" + ); + assert!(registry.is_finished(&assignments[0].0)); + assert!(!registry.is_finished(&assignments[1].0)); + assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); + } + #[tokio::test] async fn run_item_chain_reaches_flagged_agent_after_script_step() { if !cmd_available("python3") {