test(graph): pin teammate retirement on agent timeout and frontier completion

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<dyn Fn(&PeerRegistry, &str) + Send + Sync>;
+
 pub struct GraphExecutor {
     graph: Graph,
     base_dir: PathBuf,
+    #[cfg(test)]
+    frontier_observer: Option<FrontierObserver>,
 }
 
 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<String> {
-        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, &current).await;
                     if let (Some(registry), Some(id)) = (&registry_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<Mutex<Vec<(String, RosterSnapshot)>>> = 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") {

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<std::ffi::OsString>,
+        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<u64>) -> 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(&registry), 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") {