1650 lines
49 KiB
Rust
1650 lines
49 KiB
Rust
use super::executor::{StepContext, StepResult, step};
|
|
use super::state::StateManager;
|
|
use super::types::{ConcurrencyCap, Graph, MapNode, Node, NodeType};
|
|
use super::validator::branch_subgraph;
|
|
use crate::config::{RenderMode, RequestContext};
|
|
use crate::graph::type_name;
|
|
use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id};
|
|
use anyhow::{Context, Result, anyhow, bail};
|
|
use futures_util::future::{BoxFuture, join_all};
|
|
use serde_json::Value;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::Arc;
|
|
use tokio::sync::Semaphore;
|
|
|
|
pub(super) struct MapNodeExecutor;
|
|
|
|
impl MapNodeExecutor {
|
|
/// Boxed because item chains call back into `step`, which calls this;
|
|
/// the concrete return type breaks the otherwise-recursive opaque future.
|
|
pub(super) fn execute<'a>(
|
|
node: &'a MapNode,
|
|
state: &'a mut StateManager,
|
|
ctx: &'a mut RequestContext,
|
|
step_ctx: &'a StepContext<'a>,
|
|
node_id: &'a str,
|
|
) -> BoxFuture<'a, Result<()>> {
|
|
Box::pin(run_map(node, state, ctx, step_ctx, node_id))
|
|
}
|
|
}
|
|
|
|
async fn run_map(
|
|
node: &MapNode,
|
|
state: &mut StateManager,
|
|
ctx: &mut RequestContext,
|
|
step_ctx: &StepContext<'_>,
|
|
node_id: &str,
|
|
) -> Result<()> {
|
|
let over_value = state
|
|
.interpolate_raw(&node.over)
|
|
.with_context(|| format!("map node '{node_id}': evaluating `over` template"))?;
|
|
|
|
let items = over_value.as_array().ok_or_else(|| {
|
|
anyhow!(
|
|
"map node '{}': `over` template '{}' must resolve to an array, got {}",
|
|
node_id,
|
|
node.over,
|
|
type_name(&over_value)
|
|
)
|
|
})?;
|
|
let items = items.clone();
|
|
|
|
if step_ctx.graph.get_node(&node.branch).is_none() {
|
|
bail!(
|
|
"map node '{node_id}': branch '{}' not found in graph",
|
|
node.branch
|
|
);
|
|
}
|
|
let subgraph = Arc::new(branch_subgraph(&step_ctx.graph, &node.branch));
|
|
|
|
let max_conc = resolve_max_concurrency(node, state, step_ctx.max_concurrency, node_id)?;
|
|
let semaphore = Arc::new(Semaphore::new(max_conc));
|
|
let mut sub_tasks = Vec::with_capacity(items.len());
|
|
|
|
let peers = provision_map_peers(&step_ctx.graph, &subgraph, &node.branch, items.len());
|
|
|
|
for (idx, item) in items.iter().enumerate() {
|
|
let item = item.clone();
|
|
let as_name = node.as_name.clone();
|
|
let mut sub_state = state.fork_for_branch_state();
|
|
// The chain must write the result itself; an inherited parent
|
|
// value would otherwise be collected as if the item produced it.
|
|
sub_state.state_mut().remove(&node.output_key);
|
|
sub_state.state_mut().set(as_name, item);
|
|
let mut sub_ctx = ctx.fork_for_branch();
|
|
sub_ctx.render_mode = RenderMode::Silent;
|
|
let item_peers: Option<(Arc<PeerRegistry>, PeerAssignment)> = peers
|
|
.as_ref()
|
|
.map(|(registry, assignments)| (Arc::clone(registry), assignments[idx].clone()));
|
|
let script_clone = step_ctx.script_executor.clone();
|
|
let graph = Arc::clone(&step_ctx.graph);
|
|
let subgraph = Arc::clone(&subgraph);
|
|
let entry = node.branch.clone();
|
|
let map_id = node_id.to_string();
|
|
let max_concurrency = step_ctx.max_concurrency;
|
|
let sem = semaphore.clone();
|
|
let abort = step_ctx.abort_signal.clone();
|
|
|
|
let task = tokio::spawn(async move {
|
|
let _permit = sem
|
|
.acquire()
|
|
.await
|
|
.expect("map semaphore should not be closed");
|
|
let mut state = sub_state;
|
|
let mut ctx = sub_ctx;
|
|
let step_ctx = StepContext {
|
|
graph,
|
|
script_executor: &script_clone,
|
|
max_concurrency,
|
|
abort_signal: &abort,
|
|
branch_mode: true,
|
|
};
|
|
let chain = ItemChain {
|
|
map_id: &map_id,
|
|
entry: &entry,
|
|
subgraph: &subgraph,
|
|
idx,
|
|
peers: item_peers.as_ref(),
|
|
};
|
|
let result = run_item_chain(&chain, &mut state, &mut ctx, &step_ctx).await;
|
|
(idx, state, result)
|
|
});
|
|
sub_tasks.push(task);
|
|
}
|
|
|
|
let joined = join_all(sub_tasks).await;
|
|
|
|
// Collect outputs keyed by input index so order is preserved regardless of finish order.
|
|
let mut outputs: HashMap<usize, Value> = HashMap::new();
|
|
for join_result in joined {
|
|
let (idx, sub_state, exec_result) =
|
|
join_result.map_err(|e| anyhow!("map sub-branch panicked: {e}"))?;
|
|
|
|
exec_result?;
|
|
|
|
let output_value = sub_state
|
|
.state()
|
|
.get(&node.output_key)
|
|
.cloned()
|
|
.ok_or_else(|| {
|
|
anyhow!(
|
|
"map node '{node_id}': sub-branch [{idx}] did not write output_key '{}'; \
|
|
write it via state_updates, output_schema, or a script's JSON output \
|
|
(the parent's value is not inherited inside a map branch)",
|
|
node.output_key
|
|
)
|
|
})?;
|
|
|
|
outputs.insert(idx, output_value);
|
|
}
|
|
|
|
let mut collected = Vec::with_capacity(items.len());
|
|
for idx in 0..items.len() {
|
|
let value = outputs.remove(&idx).ok_or_else(|| {
|
|
anyhow!("map node '{node_id}': internal error: missing result for sub-branch [{idx}]")
|
|
})?;
|
|
collected.push(value);
|
|
}
|
|
|
|
state
|
|
.state_mut()
|
|
.set(node.collect_into.clone(), Value::Array(collected));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// One item's slice of a map fan-out: where its chain starts, which nodes
|
|
/// it may visit, and (when the fan-out is a team) the identity it speaks as.
|
|
struct ItemChain<'a> {
|
|
map_id: &'a str,
|
|
entry: &'a str,
|
|
subgraph: &'a HashSet<String>,
|
|
idx: usize,
|
|
peers: Option<&'a (Arc<PeerRegistry>, PeerAssignment)>,
|
|
}
|
|
|
|
/// Runs one item's chain to completion on its forked state and ctx. Owns the
|
|
/// item's teammate identity for the whole chain: arms it only before
|
|
/// `teammates: true` agent steps and retires it when the chain exits,
|
|
/// successfully or not.
|
|
async fn run_item_chain(
|
|
chain: &ItemChain<'_>,
|
|
state: &mut StateManager,
|
|
ctx: &mut RequestContext,
|
|
step_ctx: &StepContext<'_>,
|
|
) -> Result<()> {
|
|
let result = run_chain_steps(chain, state, ctx, step_ctx).await;
|
|
if let Some((registry, (id, _))) = chain.peers {
|
|
registry.mark_finished(id);
|
|
}
|
|
result
|
|
}
|
|
|
|
async fn run_chain_steps(
|
|
chain: &ItemChain<'_>,
|
|
state: &mut StateManager,
|
|
ctx: &mut RequestContext,
|
|
step_ctx: &StepContext<'_>,
|
|
) -> Result<()> {
|
|
let mut current = chain.entry.to_string();
|
|
let mut step_no = 0usize;
|
|
loop {
|
|
step_no += 1;
|
|
let next = run_chain_step(chain, state, ctx, step_ctx, ¤t)
|
|
.await
|
|
.with_context(|| {
|
|
format!(
|
|
"map node '{}': sub-branch [{}] failed at node '{current}' (step {step_no})",
|
|
chain.map_id, chain.idx
|
|
)
|
|
})?;
|
|
match next {
|
|
Some(target) => current = target,
|
|
None => return Ok(()),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Runs a single node of the chain and returns the node to visit next, or
|
|
/// `None` once the chain has ended.
|
|
async fn run_chain_step(
|
|
chain: &ItemChain<'_>,
|
|
state: &mut StateManager,
|
|
ctx: &mut RequestContext,
|
|
step_ctx: &StepContext<'_>,
|
|
current: &str,
|
|
) -> Result<Option<String>> {
|
|
let ItemChain { map_id, idx, .. } = *chain;
|
|
if step_ctx.abort_signal.aborted() {
|
|
bail!("map sub-branch [{idx}] aborted");
|
|
}
|
|
let node = step_ctx.graph.get_node(current).ok_or_else(|| {
|
|
anyhow!("map node '{map_id}': sub-branch [{idx}] routed to unknown node '{current}'")
|
|
})?;
|
|
if !matches!(
|
|
node.node_type,
|
|
NodeType::Llm(_) | NodeType::Agent(_) | NodeType::Rag(_) | NodeType::Script(_)
|
|
) {
|
|
bail!(
|
|
"map branch '{current}' has type that cannot run inside a map \
|
|
(validator should have caught this; internal error)"
|
|
);
|
|
}
|
|
|
|
state.state_mut().visit_node(current);
|
|
let visits = state.state().loop_count(current);
|
|
let max_loops = step_ctx.graph.settings.max_loop_iterations;
|
|
if visits > max_loops {
|
|
bail!(
|
|
"node '{current}' visited {visits} times in map branch [{idx}] \
|
|
(max_loop_iterations={max_loops})"
|
|
);
|
|
}
|
|
|
|
if let Some((registry, assignment)) = chain.peers
|
|
&& wants_peer_identity(node)
|
|
{
|
|
ctx.peer_registry = Some(Arc::clone(registry));
|
|
ctx.peer_assignment = Some(assignment.clone());
|
|
}
|
|
|
|
match step(node, state, ctx, step_ctx, current).await? {
|
|
StepResult::Continue(targets) => match targets.as_slice() {
|
|
[] => {
|
|
debug!(
|
|
"[graph:{}] map '{map_id}' [{idx}] {current} → END",
|
|
step_ctx.graph_name()
|
|
);
|
|
Ok(None)
|
|
}
|
|
[target] => {
|
|
if !chain.subgraph.contains(target) {
|
|
let mut branch_nodes: Vec<&str> =
|
|
chain.subgraph.iter().map(String::as_str).collect();
|
|
branch_nodes.sort_unstable();
|
|
bail!(
|
|
"map node '{map_id}': sub-branch [{idx}] routed to '{target}' which is \
|
|
outside the branch subgraph rooted at '{}' (branch nodes: {}). Script \
|
|
`_next` targets inside a map branch must stay within the branch.",
|
|
chain.entry,
|
|
branch_nodes.join(", ")
|
|
);
|
|
}
|
|
debug!(
|
|
"[graph:{}] map '{map_id}' [{idx}] {current} → {target}",
|
|
step_ctx.graph_name()
|
|
);
|
|
Ok(Some(target.clone()))
|
|
}
|
|
many => bail!(
|
|
"map node '{map_id}': sub-branch [{idx}] node '{current}' fanned out to {many:?}; \
|
|
a map branch must route to a single node"
|
|
),
|
|
},
|
|
StepResult::End(_) => bail!(
|
|
"map node '{map_id}': sub-branch [{idx}] reached end node '{current}' inside a map branch"
|
|
),
|
|
}
|
|
}
|
|
|
|
fn wants_peer_identity(node: &Node) -> bool {
|
|
matches!(&node.node_type, NodeType::Agent(a) if a.teammates)
|
|
}
|
|
|
|
/// A templated cap is resolved against the parent state right before the
|
|
/// fan-out. Scripts often emit numbers as strings, so a numeric string is
|
|
/// accepted; anything else — including 0 — is the author's bug and surfaces
|
|
/// as an error rather than being clamped.
|
|
fn resolve_max_concurrency(
|
|
node: &MapNode,
|
|
state: &StateManager,
|
|
default: usize,
|
|
node_id: &str,
|
|
) -> Result<usize> {
|
|
let template = match &node.max_concurrency {
|
|
None => return Ok(default.max(1)),
|
|
Some(ConcurrencyCap::Fixed(n)) => return Ok((*n).max(1)),
|
|
Some(ConcurrencyCap::Template(t)) => t,
|
|
};
|
|
|
|
let value = state
|
|
.interpolate_raw(template)
|
|
.with_context(|| format!("map node '{node_id}': evaluating `max_concurrency` template"))?;
|
|
|
|
let resolved = match &value {
|
|
Value::Number(n) => n.as_u64().and_then(|n| usize::try_from(n).ok()),
|
|
Value::String(s) => s.trim().parse::<usize>().ok(),
|
|
_ => None,
|
|
};
|
|
|
|
match resolved {
|
|
Some(n) if n >= 1 => Ok(n),
|
|
_ => Err(anyhow!(
|
|
"map node '{node_id}': max_concurrency template \"{template}\" resolved to \
|
|
{value} ({}); expected a positive integer",
|
|
type_name(&value)
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Pre-provision one teammate identity per item before any chain starts, so a
|
|
/// fast sibling can message one still waiting on the semaphore (the message
|
|
/// queues in the pre-created inbox). The identity lives for the item's whole
|
|
/// chain; every `teammates: true` agent step in it speaks as that item. A
|
|
/// single-item fan-out has no peers, so it gets no registry.
|
|
fn provision_map_peers(
|
|
graph: &Graph,
|
|
subgraph: &HashSet<String>,
|
|
entry: &str,
|
|
item_count: usize,
|
|
) -> Option<(Arc<PeerRegistry>, Vec<PeerAssignment>)> {
|
|
if item_count < 2 {
|
|
return None;
|
|
}
|
|
let flagged_agent = |node_id: &str| match graph.get_node(node_id).map(|n| &n.node_type) {
|
|
Some(NodeType::Agent(a)) if a.teammates => Some(a.agent.as_str()),
|
|
_ => None,
|
|
};
|
|
let agent_name = flagged_agent(entry).or_else(|| {
|
|
graph
|
|
.nodes
|
|
.keys()
|
|
.filter(|id| subgraph.contains(*id))
|
|
.find_map(|id| flagged_agent(id))
|
|
})?;
|
|
|
|
let registry = Arc::new(PeerRegistry::new());
|
|
let assignments = (0..item_count)
|
|
.map(|i| {
|
|
let id = graph_agent_id(agent_name);
|
|
let inbox = Arc::new(Inbox::new());
|
|
registry.insert(id.clone(), format!("{entry}[{i}]"), inbox.clone());
|
|
(id, inbox)
|
|
})
|
|
.collect();
|
|
Some((registry, assignments))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::super::types::{AgentNode, EndNode, GraphSettings, Node};
|
|
use super::*;
|
|
use indexmap::IndexMap;
|
|
|
|
fn agent_branch(teammates: bool) -> NodeType {
|
|
NodeType::Agent(AgentNode {
|
|
agent: "worker".into(),
|
|
prompt: "p".into(),
|
|
state_updates: None,
|
|
output_schema: None,
|
|
timeout: None,
|
|
teammates,
|
|
})
|
|
}
|
|
|
|
fn graph_with_branch(node_type: NodeType) -> Graph {
|
|
let mut nodes: IndexMap<String, Node> = IndexMap::new();
|
|
nodes.insert(
|
|
"shards".into(),
|
|
Node {
|
|
id: "shards".into(),
|
|
description: String::new(),
|
|
node_type,
|
|
next: None,
|
|
},
|
|
);
|
|
|
|
Graph {
|
|
name: "t".into(),
|
|
description: String::new(),
|
|
version: "1.0".into(),
|
|
model: None,
|
|
temperature: None,
|
|
top_p: None,
|
|
reasoning_effort: None,
|
|
max_concurrent_jobs: None,
|
|
can_spawn_agents: None,
|
|
max_concurrent_agents: None,
|
|
max_agent_depth: None,
|
|
global_tools: Vec::new(),
|
|
mcp_servers: Vec::new(),
|
|
mcp_tools: None,
|
|
skills_enabled: None,
|
|
enabled_skills: None,
|
|
inject_skill_instructions: None,
|
|
skill_instructions: None,
|
|
conversation_starters: Vec::new(),
|
|
variables: Vec::new(),
|
|
settings: GraphSettings::default(),
|
|
initial_state: HashMap::new(),
|
|
reducers: HashMap::new(),
|
|
start: "shards".into(),
|
|
nodes,
|
|
}
|
|
}
|
|
|
|
fn provision(
|
|
node_type: NodeType,
|
|
n: usize,
|
|
) -> Option<(Arc<PeerRegistry>, Vec<PeerAssignment>)> {
|
|
let graph = graph_with_branch(node_type);
|
|
let subgraph = HashSet::from(["shards".to_string()]);
|
|
provision_map_peers(&graph, &subgraph, "shards", n)
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_registers_every_branch() {
|
|
let (registry, assignments) = provision(agent_branch(true), 3).expect("should provision");
|
|
|
|
assert_eq!(assignments.len(), 3);
|
|
let roster = registry.roster();
|
|
assert_eq!(roster.len(), 3);
|
|
for (i, (id, inbox)) in assignments.iter().enumerate() {
|
|
assert_eq!(roster[i].0, *id);
|
|
assert_eq!(roster[i].1, format!("shards[{i}]"));
|
|
let resolved = registry.get(id).expect("assigned id should resolve");
|
|
assert!(Arc::ptr_eq(&resolved, inbox));
|
|
}
|
|
assert_ne!(assignments[0].0, assignments[1].0);
|
|
assert_ne!(assignments[1].0, assignments[2].0);
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_single_item_gets_no_registry() {
|
|
assert!(provision(agent_branch(true), 1).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_unflagged_gets_no_registry() {
|
|
assert!(provision(agent_branch(false), 3).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_non_agent_branch_gets_no_registry() {
|
|
let branch = NodeType::End(EndNode {
|
|
output: String::new(),
|
|
state_updates: None,
|
|
});
|
|
assert!(provision(branch, 3).is_none());
|
|
}
|
|
|
|
fn map_with_cap(cap: Option<ConcurrencyCap>) -> MapNode {
|
|
MapNode {
|
|
over: "{{items}}".into(),
|
|
as_name: "item".into(),
|
|
branch: "br".into(),
|
|
output_key: "output".into(),
|
|
collect_into: "results".into(),
|
|
max_concurrency: cap,
|
|
}
|
|
}
|
|
|
|
fn state_with(key: &str, value: Value) -> StateManager {
|
|
StateManager::new(HashMap::from([(key.to_string(), value)]))
|
|
}
|
|
|
|
fn resolve(cap: Option<ConcurrencyCap>, state: &StateManager) -> Result<usize> {
|
|
resolve_max_concurrency(&map_with_cap(cap), state, 8, "m")
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_fixed_uses_literal() {
|
|
let state = StateManager::new(HashMap::new());
|
|
assert_eq!(resolve(Some(ConcurrencyCap::Fixed(4)), &state).unwrap(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_none_falls_back_to_step_default() {
|
|
let state = StateManager::new(HashMap::new());
|
|
assert_eq!(resolve(None, &state).unwrap(), 8);
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_accepts_integer_value() {
|
|
let state = state_with("budget", Value::from(3));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
assert_eq!(resolve(cap, &state).unwrap(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_accepts_numeric_string() {
|
|
let state = state_with("budget", Value::from("3"));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
assert_eq!(resolve(cap, &state).unwrap(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_rejects_non_numeric_string() {
|
|
let state = state_with("budget", Value::from("high"));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
let msg = resolve(cap, &state).unwrap_err().to_string();
|
|
assert!(msg.contains("map node 'm'"), "{msg}");
|
|
assert!(msg.contains("\"{{budget}}\""), "{msg}");
|
|
assert!(msg.contains("resolved to \"high\" (string)"), "{msg}");
|
|
assert!(msg.contains("expected a positive integer"), "{msg}");
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_rejects_zero() {
|
|
let state = state_with("budget", Value::from(0));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
let msg = resolve(cap, &state).unwrap_err().to_string();
|
|
assert!(msg.contains("resolved to 0 (number)"), "{msg}");
|
|
assert!(msg.contains("expected a positive integer"), "{msg}");
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_rejects_non_integer_number() {
|
|
let state = state_with("budget", Value::from(2.5));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
let msg = resolve(cap, &state).unwrap_err().to_string();
|
|
assert!(msg.contains("resolved to 2.5 (number)"), "{msg}");
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_rejects_array() {
|
|
let state = state_with("budget", Value::Array(vec![Value::from(3)]));
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
let msg = resolve(cap, &state).unwrap_err().to_string();
|
|
assert!(msg.contains("(array)"), "{msg}");
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_max_concurrency_template_missing_key_names_map_node() {
|
|
let state = StateManager::new(HashMap::new());
|
|
let cap = Some(ConcurrencyCap::Template("{{budget}}".into()));
|
|
let err = resolve(cap, &state).unwrap_err();
|
|
let chain = format!("{err:#}");
|
|
assert!(
|
|
chain.contains("map node 'm': evaluating `max_concurrency` template"),
|
|
"{chain}"
|
|
);
|
|
assert!(chain.contains("budget"), "{chain}");
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
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, 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()
|
|
}
|
|
|
|
struct TestWorkspace {
|
|
dir: PathBuf,
|
|
}
|
|
|
|
impl TestWorkspace {
|
|
fn new() -> Self {
|
|
let dir = temp_file("-graph-map-", "");
|
|
fs::create_dir_all(&dir).unwrap();
|
|
Self { dir }
|
|
}
|
|
|
|
fn write_script(&self, name: &str, contents: &str) {
|
|
fs::write(self.dir.join(name), contents).unwrap();
|
|
}
|
|
|
|
fn write_py(&self, name: &str, body: &str) {
|
|
self.write_script(
|
|
name,
|
|
&format!(
|
|
"#!/usr/bin/env python3\nimport os, json\n\
|
|
state = json.loads(os.environ.get(\"GRAPH_STATE\", \"{{}}\"))\n{body}\n"
|
|
),
|
|
);
|
|
}
|
|
}
|
|
|
|
impl Drop for TestWorkspace {
|
|
fn drop(&mut self) {
|
|
let _ = fs::remove_dir_all(&self.dir);
|
|
}
|
|
}
|
|
|
|
fn make_ctx() -> RequestContext {
|
|
RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd)
|
|
}
|
|
|
|
async fn run_graph(yaml: &str, ws: &TestWorkspace) -> Result<String> {
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
let mut ctx = make_ctx();
|
|
GraphExecutor::new(graph, &ws.dir)
|
|
.execute(&mut ctx, create_abort_signal())
|
|
.await
|
|
}
|
|
|
|
fn collected(result: &str) -> Vec<Value> {
|
|
serde_json::from_str::<Value>(result)
|
|
.unwrap_or_else(|_| panic!("expected JSON array, got: {result}"))
|
|
.as_array()
|
|
.expect("collected results should be an array")
|
|
.clone()
|
|
}
|
|
|
|
fn error_chain(result: Result<String>) -> String {
|
|
match result {
|
|
Ok(out) => panic!("expected failure, got output: {out}"),
|
|
Err(e) => format!("{e:#}"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_two_script_steps_route_via_next() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py(
|
|
"step_one.py",
|
|
r#"print(json.dumps({"draft": state["item"] * 2}))"#,
|
|
);
|
|
ws.write_py(
|
|
"step_two.py",
|
|
r#"print(json.dumps({"output": state["draft"] + 1}))"#,
|
|
);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1, 2, 3]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: step_one
|
|
collect_into: results
|
|
next: done
|
|
step_one:
|
|
type: script
|
|
script: step_one.py
|
|
next: step_two
|
|
step_two:
|
|
type: script
|
|
script: step_two.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
|
|
assert_eq!(collected(&result), vec![json!(3), json!(5), json!(7)]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_two_script_steps_pass_validation_before_run() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py(
|
|
"step_one.py",
|
|
r#"print(json.dumps({"draft": state["item"] * 2}))"#,
|
|
);
|
|
ws.write_py(
|
|
"step_two.py",
|
|
r#"print(json.dumps({"output": state["draft"] + 1}))"#,
|
|
);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: true
|
|
initial_state:
|
|
items: [1, 2, 3]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: step_one
|
|
collect_into: results
|
|
next: done
|
|
step_one:
|
|
type: script
|
|
script: step_one.py
|
|
next: step_two
|
|
step_two:
|
|
type: script
|
|
script: step_two.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("validation rejected a runnable chain: {e:#}"));
|
|
|
|
assert_eq!(collected(&result), vec![json!(3), json!(5), json!(7)]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_script_next_loop_is_bounded_by_max_loop_iterations() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("looper.py", r#"print(json.dumps({"_next": "looper"}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
max_loop_iterations: 3
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: looper
|
|
collect_into: results
|
|
next: done
|
|
looper:
|
|
type: script
|
|
script: looper.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let chain = error_chain(run_graph(yaml, &ws).await);
|
|
|
|
assert!(
|
|
chain.contains("map node 'fan_out': sub-branch [0] failed at node 'looper' (step 4)"),
|
|
"{chain}"
|
|
);
|
|
assert!(
|
|
chain.contains(
|
|
"node 'looper' visited 4 times in map branch [0] (max_loop_iterations=3)"
|
|
),
|
|
"{chain}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_next_outside_subgraph_fails_item() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("gate.py", r#"print(json.dumps({"_next": "done"}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: gate
|
|
collect_into: results
|
|
next: done
|
|
gate:
|
|
type: script
|
|
script: gate.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let chain = error_chain(run_graph(yaml, &ws).await);
|
|
|
|
assert!(
|
|
chain.contains(
|
|
"routed to 'done' which is outside the branch subgraph rooted at 'gate' \
|
|
(branch nodes: gate)"
|
|
),
|
|
"{chain}"
|
|
);
|
|
assert!(
|
|
chain.contains(
|
|
"Script `_next` targets inside a map branch must stay within the branch."
|
|
),
|
|
"{chain}"
|
|
);
|
|
assert!(chain.contains("failed at node 'gate' (step 1)"), "{chain}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_not_writing_output_key_errors_with_hint() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("worker.py", r#"print(json.dumps({"other": 1}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1]
|
|
output: stale
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: worker
|
|
collect_into: results
|
|
next: done
|
|
worker:
|
|
type: script
|
|
script: worker.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let chain = error_chain(run_graph(yaml, &ws).await);
|
|
|
|
assert!(
|
|
chain.contains("sub-branch [0] did not write output_key 'output'"),
|
|
"{chain}"
|
|
);
|
|
assert!(
|
|
chain.contains("the parent's value is not inherited inside a map branch"),
|
|
"{chain}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_output_equal_to_parent_value_is_collected() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("worker.py", r#"print(json.dumps({"output": 7}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [7, 7]
|
|
output: 7
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: worker
|
|
collect_into: results
|
|
next: done
|
|
worker:
|
|
type: script
|
|
script: worker.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
|
|
assert_eq!(collected(&result), vec![json!(7), json!(7)]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_state_updates_reading_output_key_sees_empty_string() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("worker.py", r#"print(json.dumps({"draft": "d"}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: ["a"]
|
|
summary: parent
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: worker
|
|
output_key: summary
|
|
collect_into: results
|
|
next: done
|
|
worker:
|
|
type: script
|
|
script: worker.py
|
|
state_updates:
|
|
summary: "{{summary}}|{{draft}}"
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
|
|
assert_eq!(collected(&result), vec![json!("|d")]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_pre_check_rejects_approval_node() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("gate.py", r#"print(json.dumps({}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: gate
|
|
collect_into: results
|
|
next: done
|
|
gate:
|
|
type: script
|
|
script: gate.py
|
|
next: ask
|
|
ask:
|
|
type: approval
|
|
question: "ok?"
|
|
options: ["yes", "no"]
|
|
routes:
|
|
"yes": gate
|
|
on_other: gate
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let chain = error_chain(run_graph(yaml, &ws).await);
|
|
|
|
assert!(
|
|
chain.contains("map branch 'ask' has type that cannot run inside a map"),
|
|
"{chain}"
|
|
);
|
|
assert!(chain.contains("failed at node 'ask' (step 2)"), "{chain}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_script_error_routes_to_fallback() {
|
|
if !cmd_available("python3") || !cmd_available("bash") {
|
|
eprintln!("skipping: python3 or bash not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_script("worker.sh", "#!/bin/bash\necho boom >&2\nexit 1\n");
|
|
ws.write_py(
|
|
"recover.py",
|
|
r#"print(json.dumps({"output": "recovered"}))"#,
|
|
);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1, 2]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: worker
|
|
collect_into: results
|
|
next: done
|
|
worker:
|
|
type: script
|
|
script: worker.sh
|
|
fallback: recover
|
|
recover:
|
|
type: script
|
|
script: recover.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
|
|
assert_eq!(
|
|
collected(&result),
|
|
vec![json!("recovered"), json!("recovered")]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn map_chain_llm_failure_routes_to_fallback() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let ws = TestWorkspace::new();
|
|
ws.write_py("recover.py", r#"print(json.dumps({"output": "fb"}))"#);
|
|
|
|
let yaml = r#"
|
|
name: chain
|
|
start: fan_out
|
|
settings:
|
|
validate_before_run: false
|
|
initial_state:
|
|
items: [1]
|
|
nodes:
|
|
fan_out:
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: check
|
|
collect_into: results
|
|
next: done
|
|
check:
|
|
type: llm
|
|
prompt: "hi"
|
|
fallback: recover
|
|
recover:
|
|
type: script
|
|
script: recover.py
|
|
done:
|
|
type: end
|
|
output: "{{results}}"
|
|
"#;
|
|
let result = run_graph(yaml, &ws)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("executor failed: {e:#}"));
|
|
|
|
assert_eq!(collected(&result), vec![json!("fb")]);
|
|
}
|
|
|
|
struct Harness {
|
|
ws: TestWorkspace,
|
|
graph: Arc<Graph>,
|
|
abort: AbortSignal,
|
|
}
|
|
|
|
impl Harness {
|
|
fn new(yaml: &str) -> Self {
|
|
Self {
|
|
ws: TestWorkspace::new(),
|
|
graph: Arc::new(serde_yaml::from_str(yaml).unwrap()),
|
|
abort: create_abort_signal(),
|
|
}
|
|
}
|
|
|
|
fn subgraph(&self, entry: &str) -> HashSet<String> {
|
|
branch_subgraph(&self.graph, entry)
|
|
}
|
|
|
|
fn provision(&self, entry: &str, n: usize) -> (Arc<PeerRegistry>, Vec<PeerAssignment>) {
|
|
provision_map_peers(&self.graph, &self.subgraph(entry), entry, n)
|
|
.expect("subgraph with a flagged agent should provision peers")
|
|
}
|
|
|
|
async fn run(
|
|
&self,
|
|
entry: &str,
|
|
peers: Option<&(Arc<PeerRegistry>, PeerAssignment)>,
|
|
state: &mut StateManager,
|
|
ctx: &mut RequestContext,
|
|
) -> Result<()> {
|
|
let script = ScriptExecutor::new(&self.ws.dir);
|
|
let step_ctx = StepContext {
|
|
graph: Arc::clone(&self.graph),
|
|
script_executor: &script,
|
|
max_concurrency: 4,
|
|
abort_signal: &self.abort,
|
|
branch_mode: true,
|
|
};
|
|
let subgraph = self.subgraph(entry);
|
|
let chain = ItemChain {
|
|
map_id: "fan_out",
|
|
entry,
|
|
subgraph: &subgraph,
|
|
idx: 0,
|
|
peers,
|
|
};
|
|
run_item_chain(&chain, state, ctx, &step_ctx).await
|
|
}
|
|
}
|
|
|
|
fn item_state(item: Value) -> StateManager {
|
|
StateManager::new(HashMap::from([("item".to_string(), item)]))
|
|
}
|
|
|
|
fn silent_ctx() -> RequestContext {
|
|
let mut ctx = make_ctx();
|
|
ctx.render_mode = RenderMode::Silent;
|
|
ctx
|
|
}
|
|
|
|
fn manual_peer(label: &str) -> (Arc<PeerRegistry>, PeerAssignment) {
|
|
let registry = Arc::new(PeerRegistry::new());
|
|
let id = graph_agent_id(label);
|
|
let inbox = Arc::new(Inbox::new());
|
|
registry.insert(id.clone(), format!("{label}[0]"), Arc::clone(&inbox));
|
|
(registry, (id, inbox))
|
|
}
|
|
|
|
fn unwrap_err_chain(result: Result<()>) -> String {
|
|
match result {
|
|
Ok(()) => panic!("expected the chain to fail"),
|
|
Err(e) => format!("{e:#}"),
|
|
}
|
|
}
|
|
|
|
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
|
|
nodes:
|
|
doubler:
|
|
type: script
|
|
script: doubler.py
|
|
"#;
|
|
|
|
const GATE_WORKER_FINISH_GRAPH: &str = r#"
|
|
name: t
|
|
start: gate
|
|
nodes:
|
|
gate:
|
|
type: script
|
|
script: gate.py
|
|
next: worker
|
|
worker:
|
|
type: agent
|
|
agent: no-such-agent
|
|
prompt: "p"
|
|
teammates: true
|
|
next: finish
|
|
finish:
|
|
type: script
|
|
script: finish.py
|
|
"#;
|
|
|
|
const FLAGGED_WORKER_GRAPH: &str = r#"
|
|
name: t
|
|
start: worker
|
|
nodes:
|
|
worker:
|
|
type: agent
|
|
agent: no-such-agent
|
|
prompt: "p"
|
|
teammates: true
|
|
"#;
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_single_script_branch_final_fork_state() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let h = Harness::new(DOUBLER_GRAPH);
|
|
h.ws.write_py(
|
|
"doubler.py",
|
|
r#"print(json.dumps({"output": state["item"] * 2}))"#,
|
|
);
|
|
let mut state = item_state(json!(3));
|
|
let mut ctx = silent_ctx();
|
|
|
|
h.run("doubler", None, &mut state, &mut ctx)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("chain failed: {e:#}"));
|
|
|
|
assert_eq!(
|
|
*state.state().data(),
|
|
HashMap::from([
|
|
("item".to_string(), json!(3)),
|
|
("output".to_string(), json!(6)),
|
|
])
|
|
);
|
|
assert_eq!(state.state().loop_count("doubler"), 1);
|
|
assert_eq!(state.state().current_node(), Some("doubler"));
|
|
assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_aborted_before_first_step_touches_nothing() {
|
|
let h = Harness::new(DOUBLER_GRAPH);
|
|
h.ws.write_py(
|
|
"doubler.py",
|
|
r#"print(json.dumps({"output": state["item"] * 2}))"#,
|
|
);
|
|
let peers = manual_peer("doubler");
|
|
let mut state = item_state(json!(3));
|
|
let mut ctx = silent_ctx();
|
|
h.abort.set_ctrlc();
|
|
|
|
let chain = unwrap_err_chain(h.run("doubler", Some(&peers), &mut state, &mut ctx).await);
|
|
|
|
assert!(chain.contains("map sub-branch [0] aborted"), "{chain}");
|
|
assert!(
|
|
chain.contains("failed at node 'doubler' (step 1)"),
|
|
"{chain}"
|
|
);
|
|
assert_eq!(state.state().loop_count("doubler"), 0);
|
|
assert!(state.state().get("output").is_none());
|
|
assert!(peers.0.is_finished(&peers.1.0));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_marks_identity_finished_on_success_without_arming_script_steps() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let h = Harness::new(GATE_WORKER_FINISH_GRAPH);
|
|
h.ws.write_py("gate.py", r#"print(json.dumps({"_next": "finish"}))"#);
|
|
h.ws.write_py("finish.py", r#"print(json.dumps({"output": "ok"}))"#);
|
|
let (registry, assignments) = h.provision("gate", 2);
|
|
let peers = (Arc::clone(®istry), assignments[0].clone());
|
|
let mut state = item_state(json!(0));
|
|
let mut ctx = silent_ctx();
|
|
|
|
h.run("gate", Some(&peers), &mut state, &mut ctx)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("chain failed: {e:#}"));
|
|
|
|
assert_eq!(state.state().get("output"), Some(&json!("ok")));
|
|
assert!(registry.is_finished(&assignments[0].0));
|
|
assert!(!registry.is_finished(&assignments[1].0));
|
|
// Fork-time arming would have left these Some: script steps never hold the identity.
|
|
assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_marks_identity_finished_when_agent_step_errors() {
|
|
let h = Harness::new(FLAGGED_WORKER_GRAPH);
|
|
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("Agent 'no-such-agent' failed"), "{chain}");
|
|
assert!(registry.is_finished(&assignments[0].0));
|
|
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(®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") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let h = Harness::new(
|
|
r#"
|
|
name: t
|
|
start: prep
|
|
nodes:
|
|
prep:
|
|
type: script
|
|
script: prep.py
|
|
next: worker
|
|
worker:
|
|
type: agent
|
|
agent: no-such-agent
|
|
prompt: "p"
|
|
teammates: true
|
|
"#,
|
|
);
|
|
h.ws.write_py("prep.py", r#"print(json.dumps({"draft": "x"}))"#);
|
|
let (registry, assignments) = h.provision("prep", 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("prep", Some(&peers), &mut state, &mut ctx).await);
|
|
|
|
assert!(
|
|
chain.contains("failed at node 'worker' (step 2)"),
|
|
"{chain}"
|
|
);
|
|
assert_eq!(state.state().get("draft"), Some(&json!("x")));
|
|
assert!(registry.is_finished(&assignments[0].0));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_never_arms_unflagged_agent() {
|
|
let h = Harness::new(
|
|
r#"
|
|
name: t
|
|
start: worker
|
|
nodes:
|
|
worker:
|
|
type: agent
|
|
agent: no-such-agent
|
|
prompt: "p"
|
|
teammates: false
|
|
"#,
|
|
);
|
|
let peers = manual_peer("worker");
|
|
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("Agent 'no-such-agent' failed"), "{chain}");
|
|
assert!(peers.0.is_finished(&peers.1.0));
|
|
}
|
|
|
|
#[test]
|
|
fn wants_peer_identity_only_for_flagged_agent_nodes() {
|
|
let graph: Graph = serde_yaml::from_str(
|
|
r#"
|
|
name: t
|
|
start: s
|
|
nodes:
|
|
s:
|
|
type: script
|
|
script: s.py
|
|
next: l
|
|
l:
|
|
type: llm
|
|
prompt: "p"
|
|
next: plain
|
|
plain:
|
|
type: agent
|
|
agent: a
|
|
prompt: "p"
|
|
next: flagged
|
|
flagged:
|
|
type: agent
|
|
agent: a
|
|
prompt: "p"
|
|
teammates: true
|
|
next: e
|
|
e:
|
|
type: end
|
|
output: ""
|
|
"#,
|
|
)
|
|
.unwrap();
|
|
let node = |id: &str| graph.get_node(id).unwrap();
|
|
|
|
assert!(!wants_peer_identity(node("s")));
|
|
assert!(!wants_peer_identity(node("l")));
|
|
assert!(!wants_peer_identity(node("plain")));
|
|
assert!(!wants_peer_identity(node("e")));
|
|
assert!(wants_peer_identity(node("flagged")));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn run_item_chain_llm_step_leaves_ctx_role_and_scopes_unchanged() {
|
|
if !cmd_available("python3") {
|
|
eprintln!("skipping: python3 not available");
|
|
return;
|
|
}
|
|
let h = Harness::new(
|
|
r#"
|
|
name: t
|
|
start: check
|
|
nodes:
|
|
check:
|
|
type: llm
|
|
prompt: "hi"
|
|
fallback: finish
|
|
finish:
|
|
type: script
|
|
script: finish.py
|
|
"#,
|
|
);
|
|
h.ws.write_py("finish.py", r#"print(json.dumps({"output": "ok"}))"#);
|
|
let mut state = item_state(json!(0));
|
|
let mut ctx = silent_ctx();
|
|
ctx.role = Some(Role::new("marker", "x"));
|
|
ctx.node_job_scope = Some(vec!["j".to_string()]);
|
|
let mcp_tools = Some(("n".to_string(), IndexMap::new()));
|
|
ctx.active_node_mcp_tools = mcp_tools.clone();
|
|
|
|
h.run("check", None, &mut state, &mut ctx)
|
|
.await
|
|
.unwrap_or_else(|e| panic!("chain failed: {e:#}"));
|
|
|
|
assert_eq!(state.state().get("output"), Some(&json!("ok")));
|
|
assert_eq!(ctx.role.as_ref().map(Role::name), Some("marker"));
|
|
assert_eq!(ctx.node_job_scope, Some(vec!["j".to_string()]));
|
|
assert_eq!(ctx.active_node_mcp_tools, mcp_tools);
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_flagged_agent_anywhere_in_subgraph() {
|
|
let h = Harness::new(GATE_WORKER_FINISH_GRAPH);
|
|
|
|
let (registry, assignments) = h.provision("gate", 2);
|
|
|
|
assert_eq!(assignments.len(), 2);
|
|
let roster = registry.roster();
|
|
assert_eq!(roster[0].1, "gate[0]");
|
|
assert_eq!(roster[1].1, "gate[1]");
|
|
for (id, _) in &assignments {
|
|
assert!(id.starts_with("graph_agent_no-such-agent_"), "{id}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn provision_map_peers_ignores_flagged_agent_outside_subgraph() {
|
|
let h = Harness::new(GATE_WORKER_FINISH_GRAPH);
|
|
let subgraph = HashSet::from(["gate".to_string()]);
|
|
|
|
assert!(provision_map_peers(&h.graph, &subgraph, "gate", 2).is_none());
|
|
}
|
|
}
|