diff --git a/assets/agents/deep-research/graph.yaml b/assets/agents/deep-research/graph.yaml index dd611ed..4aa2b25 100644 --- a/assets/agents/deep-research/graph.yaml +++ b/assets/agents/deep-research/graph.yaml @@ -13,7 +13,7 @@ description: | node type (script, llm, rag, map, agent, input, approval, end) and both static fan-out and dynamic map fan-out. -version: "1.0" +version: '1.0' global_tools: - web_search_coyote.sh @@ -24,8 +24,8 @@ mcp_servers: - ddg-search conversation_starters: - - "How does HTTP/3 differ from HTTP/2?" - - "Summarize recent advances in solid-state battery chemistry" + - 'How does HTTP/3 differ from HTTP/2?' + - 'Summarize recent advances in solid-state battery chemistry' settings: max_loop_iterations: 40 @@ -34,15 +34,14 @@ settings: max_concurrency: 4 initial_state: - research_feedback: "" + research_feedback: '' research_attempts: 0 - local_context: "" - local_sources: "" + local_context: '' + local_sources: '' start: parse_request nodes: - parse_request: id: parse_request type: script @@ -52,10 +51,10 @@ nodes: ask_topic: id: ask_topic type: input - question: "What would you like me to research?" - validation: "len(input) > 0" + question: 'What would you like me to research?' + validation: 'len(input) > 0' state_updates: - topic: "{{input}}" + topic: '{{input}}' next: bootstrap_research bootstrap_research: @@ -77,7 +76,7 @@ nodes: are precise, self-contained queries (each one is sent on its own to a separate research worker, so they must be answerable without each other's context). - prompt: "Research topic: {{topic}}" + prompt: 'Research topic: {{topic}}' tools: [] output_schema: type: object @@ -99,22 +98,23 @@ nodes: type: rag documents: - ./knowledge/ - query: "{{topic}}" + query: '{{topic}}' top_k: 6 chunk_size: 1000 chunk_overlap: 100 state_updates: - local_context: "{{output.context}}" - local_sources: "{{output.sources}}" + local_context: '{{output.context}}' + local_sources: '{{output.sources}}' next: research_each_question research_each_question: id: research_each_question type: map - over: "{{questions}}" + over: '{{questions}}' as: question branch: research_one_question collect_into: question_findings + output_key: finding max_concurrency: 3 next: combine_findings @@ -151,7 +151,8 @@ nodes: - mcp:ddg-search max_iterations: 10 max_attempts: 2 - temperature: 0.1 + state_updates: + finding: '{{output}}' combine_findings: id: combine_findings @@ -177,7 +178,7 @@ nodes: - classify_source max_iterations: 15 state_updates: - source_assessment: "{{output}}" + source_assessment: '{{output}}' next: critique critique: @@ -217,7 +218,7 @@ nodes: {{source_assessment}} tools: [] state_updates: - critique: "{{output}}" + critique: '{{output}}' next: reflexion_gate reflexion_gate: @@ -242,7 +243,7 @@ nodes: Produce the final report following your instructions. timeout: 300 state_updates: - report: "{{output}}" + report: '{{output}}' next: verify_sources verify_sources: @@ -266,14 +267,14 @@ nodes: Accept this report? Pick "accept" or "reject", or type specific feedback to send the research back for another pass. options: - - "accept" - - "reject" + - 'accept' + - 'reject' routes: - "accept": end_accepted - "reject": end_rejected + 'accept': end_accepted + 'reject': end_rejected on_other: incorporate_feedback state_updates: - decision: "{{choice}}" + decision: '{{choice}}' incorporate_feedback: id: incorporate_feedback @@ -283,7 +284,7 @@ nodes: end_accepted: id: end_accepted type: end - output: "{{report}}" + output: '{{report}}' end_rejected: id: end_rejected diff --git a/assets/agents/finding-verifier/graph.yaml b/assets/agents/finding-verifier/graph.yaml index 472fe48..5e91756 100644 --- a/assets/agents/finding-verifier/graph.yaml +++ b/assets/agents/finding-verifier/graph.yaml @@ -27,9 +27,6 @@ settings: log_state_snapshots: true timeout: 900 -reducers: - output: overwrite - initial_state: findings_list: [] verdicts: [] @@ -95,6 +92,7 @@ nodes: as: finding branch: verify_one collect_into: verdicts + output_key: verdict max_concurrency: 4 next: done @@ -149,7 +147,8 @@ nodes: - fs_grep max_iterations: 8 max_attempts: 2 - temperature: 0.1 + state_updates: + verdict: '{{output}}' done: id: done diff --git a/assets/agents/librarian/graph.yaml b/assets/agents/librarian/graph.yaml index 13ab16d..696a0f0 100644 --- a/assets/agents/librarian/graph.yaml +++ b/assets/agents/librarian/graph.yaml @@ -32,9 +32,6 @@ settings: log_state_snapshots: true timeout: 600 -reducers: - output: overwrite - initial_state: language_ecosystem: 'general' doc_domain_hints: '' diff --git a/graph.example.yaml b/graph.example.yaml index 95b6514..70d62f3 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -116,9 +116,12 @@ variables: # --------------------------------------------------------------------------- settings: max_loop_iterations: 100 # Per-node visit cap. If one node id is entered more - # than this many times, execution aborts. Default 100. + # than this many times, execution aborts. Default 100; 0 = no cap. timeout: 600 # Optional wall-clock cap (seconds) on the whole run, - # checked between node transitions. + # checked between node transitions. 0 = no bound. + # Ctrl-C (or the enclosing turn aborting) ends the run + # promptly from anywhere: in-flight branches, tools, and + # child graph agents are cancelled with it. log_state_snapshots: true # Log state before each node (debug/trace). Default true. validate_before_run: true # Run the graph validator at startup. Default true. max_concurrency: 4 # Cap on simultaneously running branches in any @@ -242,7 +245,7 @@ nodes: - https://example.com/reference query: "{{topic}}" # Retrieval query (templated). Default: {{initial_prompt}}. top_k: 5 # Chunks to retrieve. Default = the KB's own top_k. - timeout: 120 # Retrieval timeout in seconds. Default 120. + timeout: 120 # Retrieval timeout in seconds. Default 120; 0 = no bound. # Knowledge-base build config (optional; used only when the KB is first # built). When embedding_model + chunk_size + chunk_overlap are all set, # the KB builds with no interactive prompts (works in non-interactive runs). @@ -311,7 +314,7 @@ nodes: id: synthesize type: script script: scripts/synthesize.py # Path relative to the agent directory - timeout: 30 # Seconds. Default 30. + timeout: 30 # Seconds. Default 30; 0 = no bound. state_updates: # Applied after the stdout JSON is merged decided_for: "{{topic}}" next: summarize # Default route if the script emits no `_next` @@ -333,7 +336,21 @@ nodes: prompt: | # User message sent to the child (templated) Research {{topic}} in depth. Existing context: {{context}} - timeout: 600 # Optional wall-clock cap, seconds. Default 300. + # inputs: # Optional. Only valid when `agent` is a GRAPH agent + # max_parallel: "{{n}}" # (has a graph.yaml; deep-research is one). Each value + # context: "Repo: {{r}}" # is a template against THIS graph's state. The resolved + # # map is overlaid on the child's `initial_state` before + # # its graph starts, so the child can read {{max_parallel}} + # # and {{context}} from its first node on. A lone {{key}} + # # keeps the value's JSON type (numbers stay numbers, lists + # # stay lists, null overlays the child's default); mixed text + # # becomes a string. `initial_prompt` is reserved (always + # # seeded from `prompt:`) and is rejected as an input key. + # # `inputs:` on a config-only agent is a load-time error, + # # and `inputs` values count as reads for the cross-branch + # # read check. See the wiki: Graph Agents -> "Passing state + # # into a child graph agent (inputs)". + timeout: 600 # Optional wall-clock cap, seconds. Default 300; 0 = no bound. teammates: false # Optional. When true and this node runs concurrently # with other flagged agent nodes (map fan-out branches # or the same super-step), the siblings can message @@ -343,7 +360,13 @@ nodes: # super-step frontier form two separate teams and cannot # message each other. Peers only span a single map # node's branches or a single super-step's flagged - # agent nodes. + # agent nodes. Inside a map branch chain the identity + # belongs to the ITEM, not the node: it lives for the + # item's whole chain, is handed to each `teammates: true` + # agent step in that chain (unflagged steps never see it), + # and is retired when the chain ends. A message delivered + # after the chain's last flagged agent step is accepted + # but never read. output_schema: # Optional. Same extraction as llm nodes type: object properties: @@ -357,9 +380,20 @@ nodes: next: review # Required for agent nodes # --- map node (Dynamic fan-out. Think: LangGraph's `Send` API) ---------------- - # Spawns one parallel sub-branch per item in `over`. Each sub-branch runs - # the node referenced by `branch:` with the item bound to `as:`. Outputs - # collect into the array named by `collect_into:`, preserving input order. + # Spawns one parallel sub-branch per item in `over`. Each sub-branch runs a + # per-item CHAIN of nodes on a private fork of the state: it starts at the + # node referenced by `branch:` (with the item bound to `as:`) and follows + # that node's `next:` / script `_next` / `fallback` edges until it reaches a + # node with no `next` and no `_next`. When every chain has finished, the + # value of `output_key` from each item's final fork state is collected into + # the array named by `collect_into:`, preserving input order. Everything + # else the chain wrote is fork-local scratch and is discarded at collect. + # + # Not supported inside a branch chain (validator errors): fan-out `next:` + # lists, another `map`, `end`, `approval`, `input`. A branch node must not + # also be reachable from `start` on the main flow. Per-item loops are script + # `_next` loops (bounded by settings.max_loop_iterations on the fork; `0` = unbounded); a + # declared `next`/`fallback` edge back to an earlier chain node is a cycle. # # Reach via `synthesize`'s `_next: subjects_map`. The producer is expected # to have written a list at `subjects` (e.g. an upstream LLM node with an @@ -371,31 +405,53 @@ nodes: # Empty list is allowed. It means no branches spawn, # and thus `collect_into` is written as []. as: subject # Required. Per-branch state key holding the - # current item. Read with {{subject}} inside - # the branch node's prompt. - branch: research_subject # Required. Node id to invoke per item. - # Must point to an llm | agent | rag | script - # node satisfying the map branch contract: - # - no `next:` (atomic, joined at map exit) - # - no `state_updates:` other than via the - # map's `collect_into` channel - # - no `output_schema:` (top-level merge - # would clash with collect_into) - # Validator enforces all three. + # current item. Read with {{subject}} from any + # node in the branch chain. Must differ from + # `output_key` (load-time error otherwise). + branch: research_subject # Required. ENTRY node of the per-item chain. + # Every node reachable from it via `next:` / + # `fallback:` is part of the chain and must be + # an llm | agent | rag | script node. Chain + # nodes may use `next:`, `state_updates:` + # (fork-local scratch such as a draft or a + # retry counter), `output_schema:`, and + # `fallback:` (honored per item; it must point + # at a node inside the chain, never at an + # `end_*` node). A script's `_next` must also + # stay inside the chain; routing out of it + # fails the item. collect_into: subject_findings # Required. State key for the array of # per-branch outputs, in input order # (not spawn-finish order). - max_concurrency: 3 # Optional per-map cap. Defaults to - # settings.max_concurrency above. - output_key: output # Optional. State key the branch's output - # appears under. Default "output". Useful - # only if the branch reads its own bound - # name back (rare). + max_concurrency: 3 # Optional cap on items in flight (one slot is + # held for an item's WHOLE chain). Defaults to + # settings.max_concurrency above. Accepts an + # integer or a template that is exactly one + # "{{key}}" (surrounding whitespace allowed, + # nothing else) and resolves to a positive integer, + # e.g. `max_concurrency: "{{research_budget}}"`. + output_key: subject_summary # Optional. State key whose value in the item's + # FINAL fork state becomes the item's result. + # Default "output"; must differ from `as:`. The + # key is removed from the fork before the chain + # starts, so the chain + # must write it (via state_updates, an + # output_schema property, or a script's JSON); + # an item that never writes it is an error, and + # a chain that reads it before writing it sees + # it absent (prompts fail, state_updates render + # "", scripts get no key) - seed per-item input + # through `as:` instead. Prefer a named key: with + # the default, the collector errors if the chain + # never writes it. next: aggregate_subjects # Where to go after all sub-branches finish. - # Branch node for subjects_map. Each invocation receives a different - # `subject` in state. The branch is "atomic", meaning it cannot route on - # its own; the surrounding `map` joins after all invocations finish. + # Entry node of subjects_map's per-item chain. Each item gets its own copy + # of this node with a different `subject` in state. With no `next:` it is + # also the LAST chain node, so the chain is a single step; give it `next:` + # (or emit `_next` from a script step) to extend the chain, e.g. + # research -> verify -> gate-script that re-emits `_next: research_subject` + # until the draft passes or a retry counter in fork state runs out. research_subject: id: research_subject type: llm @@ -403,8 +459,11 @@ nodes: prompt: "Research {{subject}}: pull the key facts and one citation." tools: - web_search_coyote - # No `next:`, `state_updates:`, or `output_schema:` here. Map branches - # have a strict contract (see `subjects_map.branch` comment). + # Recommended pattern: write the completion to a named key via + # `state_updates` and point the map's `output_key` at that same key. + # `{{output}}` is bound only while these templates render. + state_updates: + subject_summary: "{{output}}" # Aggregator that runs after the map joins. Reads the collected list. aggregate_subjects: @@ -434,9 +493,9 @@ nodes: temperature: 0.3 # Optional per-node sampling override reasoning_effort: null # Optional per-node reasoning effort override (e.g. low, medium, high) max_attempts: 2 # Retry count on transient errors only. Default 1. - max_iterations: 10 # Tool-call-loop turn cap. Default 10. + max_iterations: 10 # Tool-call-loop turn cap. Default 10; 0 = no cap. fallback: review # Route here if all attempts fail - timeout: 300 # Optional node wall-clock cap, seconds (unset = no timeout) + timeout: 300 # Optional node wall-clock cap, seconds (unset or 0 = no timeout) state_updates: research: "{{output}}" next: review # Required for llm nodes: the success route diff --git a/src/config/agent.rs b/src/config/agent.rs index 5dd3f64..c67cd49 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -633,6 +633,7 @@ impl Agent { Some(self.name().to_string()), tool_timeout, false, + None, )?; match value { Some(v) => Ok(v), diff --git a/src/config/request_context.rs b/src/config/request_context.rs index e596ad7..d5b36f9 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -8076,7 +8076,28 @@ mod tests { #[serial] fn bundled_graph_agents_parse_and_validate() { use crate::graph::GraphParser; - use crate::graph::validator::GraphValidator; + use crate::graph::validator::{GraphValidator, ValidationResult}; + use std::collections::BTreeMap; + + const UNREACHABLE: &str = "Node is unreachable from the start node via declared edges \ + (script `_next` routing is not analyzed)"; + + fn warning_lines(result: &ValidationResult) -> Vec { + let mut lines: Vec = result + .warnings + .iter() + .map(|w| format!("{}: {}", w.node_id.as_deref().unwrap_or("-"), w.message)) + .collect(); + lines.sort(); + lines + } + + fn unreachable(node_ids: &[&str]) -> Vec { + node_ids + .iter() + .map(|id| format!("{id}: {UNREACHABLE}")) + .collect() + } let _guard = TestConfigDirGuard::new(); @@ -8084,6 +8105,7 @@ mod tests { Skill::install_builtin_skills(false).unwrap(); let mut checked = Vec::new(); + let mut warnings_by_agent = BTreeMap::new(); for entry in std::fs::read_dir(paths::agents_data_dir()).unwrap() { let dir = entry.unwrap().path(); let graph_path = dir.join("graph.yaml"); @@ -8100,6 +8122,7 @@ mod tests { "graph.yaml for '{name}' failed validation: {:#?}", result.errors ); + warnings_by_agent.insert(name.clone(), warning_lines(&result)); checked.push(name); } checked.sort(); @@ -8109,6 +8132,89 @@ mod tests { "expected bundled graph agent '{expected}' to be checked; found {checked:?}" ); } + + // Warning parity with the pre-subgraph validator: the only warnings any + // bundled graph emits are the `_next`-routed unreachable nodes. A new + // rule that fires on a shipped asset (or a new bundled graph) must be + // recorded here deliberately. + let expected_warnings = BTreeMap::from([ + ( + "coder".to_string(), + unreachable(&[ + "analyze_request", + "end_rejected", + "end_success", + "fix_loop_gate", + "gate_approval", + "implement", + "route_complexity", + "route_review_result", + "self_review", + "verify_build", + "verify_tests", + ]), + ), + ("deep-research".to_string(), unreachable(&["ask_topic"])), + ("finding-verifier".to_string(), Vec::new()), + ("librarian".to_string(), Vec::new()), + ( + "step-runner".to_string(), + unreachable(&[ + "check_handoff", + "edge_case_sweep", + "end_blocked", + "end_rejected", + "end_success", + "fix_loop_gate", + "gate_blocked", + "gate_deviation", + "gate_user_review", + "get_revision", + "independent_review", + "revise_from_choice", + "route_review", + "route_sweep", + "verify_build", + "verify_format_lint", + "verify_tests", + "write_handoff", + ]), + ), + ]); + assert_eq!( + warnings_by_agent, expected_warnings, + "bundled graph validator warnings drifted from the recorded baseline" + ); + + // The shipped reference graph must load through the real parser and + // validate cleanly too; its `scripts/synthesize.py` only has to exist. + let base = _guard.path.join("example-base"); + create_dir_all(base.join("scripts")).unwrap(); + write(base.join("scripts").join("synthesize.py"), "").unwrap(); + let example = GraphParser::new(&base) + .load_from_string(include_str!("../../graph.example.yaml")) + .unwrap_or_else(|e| panic!("graph.example.yaml failed to parse: {e}")); + let result = GraphValidator::new(&base).validate(&example); + assert!( + result.errors.is_empty(), + "graph.example.yaml failed validation: {:#?}", + result.errors + ); + let example_warnings = warning_lines(&result); + assert!( + example_warnings.iter().all(|w| !w.contains("output_key")), + "graph.example.yaml's map must declare a writer for its output_key: {example_warnings:#?}" + ); + assert_eq!( + example_warnings, + unreachable(&[ + "aggregate_subjects", + "deep_dive", + "research_subject", + "subjects_map", + ]), + "graph.example.yaml validator warnings drifted from the recorded baseline" + ); } #[test] diff --git a/src/function/agents.rs b/src/function/agents.rs index a62b2c1..071bcb2 100644 --- a/src/function/agents.rs +++ b/src/function/agents.rs @@ -18,6 +18,7 @@ use indexmap::IndexMap; use log::{debug, warn}; use parking_lot::RwLock; use serde_json::{Value, json}; +use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -535,18 +536,57 @@ pub async fn handle_agent_tool( } pub fn run_child_agent( - mut child_ctx: RequestContext, + child_ctx: RequestContext, initial_input: Input, abort_signal: AbortSignal, +) -> Pin> + Send>> { + run_child_agent_with_graph_inputs(child_ctx, initial_input, abort_signal, None) +} + +/// Stops a child agent's own spawned subagents if the future running the +/// child is dropped mid-flight (parent timeout, abort) or the child fails +/// before reaching its own cleanup. Disarmed on the success paths, which +/// already cancel explicitly. +struct CancelOnDrop(Option>>); + +impl CancelOnDrop { + fn disarm(&mut self) { + self.0.take(); + } +} + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + if let Some(supervisor) = self.0.take() { + supervisor.read().cancel_recursive(); + } + } +} + +/// Runs a child agent to completion. When the child is a graph agent, +/// `graph_inputs` is overlaid on its `initial_state` before the graph +/// starts. Passing inputs to a config-only agent is rejected before the +/// child runs, since nothing could consume them. +fn run_child_agent_with_graph_inputs( + mut child_ctx: RequestContext, + initial_input: Input, + abort_signal: AbortSignal, + graph_inputs: Option>, ) -> Pin> + Send>> { Box::pin(async move { + let mut cancel_on_drop = CancelOnDrop(child_ctx.supervisor.clone()); if graph::active_agent_graph_name(&child_ctx).is_some() { - return graph::run_active_agent_graph( + let result = graph::run_active_agent_graph_with_inputs( &mut child_ctx, &initial_input.text(), abort_signal, + graph_inputs, ) .await; + if result.is_ok() { + cancel_on_drop.disarm(); + } + return result; } let mut accumulated_output = String::new(); @@ -606,6 +646,7 @@ pub fn run_child_agent( if let Some(supervisor) = child_ctx.supervisor.clone() { supervisor.read().cancel_recursive(); } + cancel_on_drop.disarm(); Ok(accumulated_output) }) @@ -685,10 +726,13 @@ fn effective_max_agent_depth(parent_ctx: &RequestContext) -> usize { /// output. This is similar to `handle_spawn` but runs the child agent in the /// current task (no tokio::spawn, no supervisor handle registration) so the /// graph executor can sequence agent nodes directly. +/// `graph_inputs` seeds the child's graph state and is only accepted when the +/// target is a graph agent. pub async fn run_agent_for_graph( parent_ctx: &mut RequestContext, agent_name: &str, prompt: &str, + graph_inputs: Option>, ) -> Result { let peer_assignment = parent_ctx.peer_assignment.take(); let assigned_peer = peer_assignment.is_some(); @@ -757,8 +801,22 @@ pub async fn run_agent_for_graph( Arc::clone(&child_inbox), agent_id.clone(), ); + child_ctx.session_abort = parent_ctx.session_abort.clone(); child_ctx.rag = agent.rag(); child_ctx.agent = Some(agent); + if graph_inputs.is_some() && graph::active_agent_graph_name(&child_ctx).is_none() { + bail!( + "agent '{agent_name}' has no graph.yaml; `inputs:` is only valid on agent nodes that target a graph agent" + ); + } + if graph_inputs + .as_ref() + .is_some_and(|inputs| inputs.contains_key("initial_prompt")) + { + bail!( + "agent '{agent_name}': `inputs.initial_prompt` is reserved (the dispatcher seeds it from `prompt:`)" + ); + } if assigned_peer { child_ctx.peer_registry = peer_registry.clone(); } @@ -786,13 +844,7 @@ pub async fn run_agent_for_graph( debug!("Spawning agent '{agent_name}' for graph node as '{agent_id}'"); - let result = run_child_agent(child_ctx, input, child_abort).await; - - if let Some(registry) = &peer_registry { - registry.mark_finished(&agent_id); - } - - result + run_child_agent_with_graph_inputs(child_ctx, input, child_abort, graph_inputs).await } async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[String]) -> Result<()> { @@ -1919,6 +1971,141 @@ mod tests { assert_eq!(effective_max_agent_depth(&ctx), 5); } + fn unique_agent_name(prefix: &str) -> String { + format!( + "{prefix}_{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + ) + } + + fn write_config_only_agent(agent_name: &str) { + let agent_dir = paths::agent_data_dir(agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("config.yaml"), + format!("name: {agent_name}\ninstructions: hi\n"), + ) + .unwrap(); + } + + fn write_graph_agent(agent_name: &str) { + let agent_dir = paths::agent_data_dir(agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("graph.yaml"), + format!( + "name: {agent_name}\nstart: done\nnodes:\n done:\n type: end\n output: done\n" + ), + ) + .unwrap(); + } + + fn run_agent_for_graph_error( + agent_name: &str, + graph_inputs: Option>, + ) -> String { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + let err = run_async(run_agent_for_graph( + &mut ctx, + agent_name, + "hi", + graph_inputs, + )) + .expect_err("run_agent_for_graph should reject the inputs"); + format!("{err:#}") + } + + #[test] + #[serial] + fn run_agent_for_graph_rejects_inputs_for_config_only_agent() { + let _guard = TestConfigDirGuard::new(); + let agent_name = unique_agent_name("test_inputs_plain_agent"); + write_config_only_agent(&agent_name); + + let chain = run_agent_for_graph_error(&agent_name, Some(HashMap::new())); + + assert_eq!( + chain, + format!( + "agent '{agent_name}' has no graph.yaml; `inputs:` is only valid on agent nodes that target a graph agent" + ) + ); + } + + #[test] + #[serial] + fn run_agent_for_graph_rejects_initial_prompt_input_on_graph_agent() { + let _guard = TestConfigDirGuard::new(); + let agent_name = unique_agent_name("test_inputs_graph_agent"); + write_graph_agent(&agent_name); + let inputs = HashMap::from([("initial_prompt".to_string(), json!("x"))]); + + let chain = run_agent_for_graph_error(&agent_name, Some(inputs)); + + assert_eq!( + chain, + format!( + "agent '{agent_name}': `inputs.initial_prompt` is reserved (the dispatcher seeds it from `prompt:`)" + ) + ); + } + + /// The child ctx must carry the parent's session flag: a session already + /// aborted before the child starts trips the child graph's own pre-check, + /// so it never reaches its (otherwise trivially successful) end node. + /// Without the inheritance the child sees `None` and returns "done". + #[test] + #[serial] + fn run_agent_for_graph_child_inherits_session_abort() { + let _guard = TestConfigDirGuard::new(); + let agent_name = unique_agent_name("test_session_abort_graph_agent"); + write_graph_agent(&agent_name); + let session = create_abort_signal(); + session.set_ctrlc(); + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.session_abort = Some(Arc::clone(&session)); + + let err = run_async(run_agent_for_graph(&mut ctx, &agent_name, "hi", None)) + .expect_err("the inherited session abort must end the child graph"); + + let chain = format!("{err:#}"); + assert!( + chain.contains(&format!("Graph execution failed for agent '{agent_name}'")), + "{chain}" + ); + assert!(chain.contains("aborted before super-step"), "{chain}"); + assert!( + Arc::ptr_eq(ctx.session_abort.as_ref().unwrap(), &session), + "the parent's own session signal is untouched" + ); + } + + #[test] + fn cancel_on_drop_guard_cancels_when_armed_and_not_when_disarmed() { + run_async(async { + let mut ctx = ctx_with_supervisor(4, 3); + let armed_abort = register_running_agent(&mut ctx, "a1", "explore"); + drop(CancelOnDrop(ctx.supervisor.clone())); + assert!( + armed_abort.aborted(), + "an armed guard cancels the supervisor's agents when dropped" + ); + + let mut ctx = ctx_with_supervisor(4, 3); + let disarmed_abort = register_running_agent(&mut ctx, "a2", "explore"); + let mut guard = CancelOnDrop(ctx.supervisor.clone()); + guard.disarm(); + drop(guard); + assert!( + !disarmed_abort.aborted(), + "a disarmed guard leaves the supervisor's agents running" + ); + }); + } + fn auto_continue_ctx() -> RequestContext { let app = AppState { config: Arc::new(AppConfig { diff --git a/src/function/mod.rs b/src/function/mod.rs index 7708b1f..d80ae25 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1684,6 +1684,7 @@ impl ToolCall { agent_name, ctx.app.config.tool_timeout, quiet, + ctx.session_abort.clone(), ) { Ok(Some(contents)) => serde_json::from_str(&contents) .ok() @@ -2293,6 +2294,7 @@ pub fn run_llm_function( agent_name: Option, tool_timeout: Option, quiet: bool, + abort: Option, ) -> Result> { let mut bin_dirs: Vec = vec![]; let mut command_name = cmd_name.clone(); @@ -2440,6 +2442,21 @@ pub fn run_llm_function( return Ok(Some(error_json.to_string())); } + if abort.as_ref().is_some_and(|a| a.aborted()) { + let _ = child.kill(); + let _ = child.wait(); + drop(stdout_thread); + drop(stderr_thread); + let tool_error_message = format!("Tool call '{command_name}' aborted"); + emit_tool_warning( + quiet, + &format!("⚠️ {tool_error_message} ⚠️"), + &tool_error_message, + ); + return Ok(Some( + json!({"tool_call_error": tool_error_message}).to_string(), + )); + } thread::sleep(Duration::from_millis(100)); }; let stdout_bytes = stdout_thread.join().unwrap_or_default(); @@ -3389,7 +3406,8 @@ mod tests { continue; } let name = path.file_stem().unwrap().to_string_lossy().to_string(); - let declarations = Functions::generate_declarations(&path) + let src = std::fs::read_to_string(&path).unwrap(); + let (_, declarations) = bash::build_bash_tool(&src, &name) .unwrap_or_else(|e| panic!("bundled tool '{name}' failed to parse: {e}")); assert!( !declarations.is_empty(), @@ -4561,6 +4579,7 @@ mod tests { None, None, false, + None, ) .unwrap() .expect("nonzero exit must return an error payload"); @@ -4589,6 +4608,7 @@ mod tests { None, None, true, + None, ) .unwrap() .expect("nonzero exit must return an error payload"); @@ -4604,6 +4624,38 @@ mod tests { assert_eq!(json["output"], "partial-output\n"); } + /// An already-aborted signal is noticed on the first wait-loop poll, so a + /// tool that would run for seconds is killed and reported within one poll + /// interval instead of running to completion or `tool_timeout`. + #[cfg(unix)] + #[test] + fn run_llm_function_aborts_running_tool() { + if which::which("bash").is_err() { + eprintln!("skipping: bash not available"); + return; + } + let abort = create_abort_signal(); + abort.set_ctrlc(); + + let started = Instant::now(); + let result = run_llm_function( + "bash".into(), + vec!["-c".into(), "sleep 5".into()], + HashMap::new(), + None, + None, + true, + Some(abort), + ) + .unwrap() + .expect("an aborted tool must return an error payload"); + let elapsed = started.elapsed(); + + assert!(elapsed < Duration::from_secs(1), "took {elapsed:?}"); + let json: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(json["tool_call_error"], "Tool call 'bash' aborted"); + } + #[test] #[serial] fn tool_timeout_secs_prefers_env_then_config_then_default() { diff --git a/src/graph/agent.rs b/src/graph/agent.rs index e5eb7d1..03b9830 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -1,14 +1,16 @@ use super::state::StateManager; +use super::state_updates; use super::structured; use super::types::AgentNode; +use super::wall_clock; use crate::config::RequestContext; use crate::function::agents::run_agent_for_graph; use anyhow::{Context, Result}; +use log::debug; use serde_json::Value; -use std::time::Duration; +use std::collections::HashMap; use tokio::time::timeout; -const OUTPUT_KEY: &str = "output"; const DEFAULT_TIMEOUT_SECS: u64 = 300; pub struct AgentNodeExecutor; @@ -18,26 +20,44 @@ impl AgentNodeExecutor { node: &AgentNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, + retire_peer_on_return: bool, ) -> Result { let prompt = state_manager .interpolate(&node.prompt) .with_context(|| format!("Failed to interpolate prompt for agent '{}'", node.agent))?; - let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS)); + let graph_inputs = resolve_inputs(node, state_manager)?; + if let Some(inputs) = &graph_inputs { + let mut keys: Vec<&String> = inputs.keys().collect(); + keys.sort_unstable(); + debug!("Agent '{}' graph inputs: {keys:?}", node.agent); + } - let raw = timeout( - timeout_dur, - run_agent_for_graph(parent_ctx, &node.agent, &prompt), - ) - .await - .with_context(|| { - format!( - "Agent '{}' timed out after {}s", - node.agent, - timeout_dur.as_secs() - ) - })? - .with_context(|| format!("Agent '{}' failed", node.agent))?; + let secs = node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); + let bound = wall_clock(secs); + + // run_agent_for_graph takes the identity off the ctx; keep a handle so + // a frontier peer is retired the moment the agent stops, not after + // extraction. Chains re-arm the same identity for later steps and + // leave retirement to the chain runner. + let peer = parent_ctx.peer_registry.clone().zip( + parent_ctx + .peer_assignment + .as_ref() + .map(|(id, _)| id.clone()), + ); + + let fut = run_agent_for_graph(parent_ctx, &node.agent, &prompt, graph_inputs); + let raw_result = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + }; + if retire_peer_on_return && let Some((registry, id)) = &peer { + registry.mark_finished(id); + } + let raw = raw_result + .with_context(|| format!("Agent '{}' timed out after {}s", node.agent, secs))? + .with_context(|| format!("Agent '{}' failed", node.agent))?; let output_value = match &node.output_schema { Some(schema) => structured::extract(&raw, schema, parent_ctx) @@ -57,46 +77,48 @@ impl AgentNodeExecutor { } } -fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) { - if node.output_schema.is_some() - && let Some(obj) = output.as_object() - { - for (k, v) in obj { - state_manager.state_mut().set(k.clone(), v.clone()); - } - } - - let Some(updates) = &node.state_updates else { - return; +/// Resolves each `inputs` template against the parent state. A lone +/// `{{key}}` yields the state value as-is (numbers, arrays, objects and +/// `null` all survive); anything else renders to a string. Every key is +/// strict: a missing reference fails the node before the child starts. +fn resolve_inputs( + node: &AgentNode, + state_manager: &StateManager, +) -> Result>> { + let Some(inputs) = &node.inputs else { + return Ok(None); }; - let prev_output = state_manager.state().get(OUTPUT_KEY).cloned(); - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), output.clone()); - - for (key, template) in updates { - let value = state_manager.interpolate_lenient(template); - state_manager - .state_mut() - .set(key.clone(), Value::String(value)); + let mut resolved = HashMap::with_capacity(inputs.len()); + for (key, template) in inputs { + let value = state_manager.interpolate_raw(template).with_context(|| { + format!( + "Failed to interpolate inputs.{key} for agent '{}'", + node.agent + ) + })?; + resolved.insert(key.clone(), value); } + Ok(Some(resolved)) +} - match prev_output { - Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v), - None => { - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), Value::Null); - } - } +fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, output: &Value) { + state_updates::apply( + state_manager, + output, + node.output_schema.is_some(), + node.state_updates.as_ref(), + ); } #[cfg(test)] mod tests { use super::super::types::AgentNode; use super::*; + use crate::config::{AppState, WorkingMode, default_max_agent_depth}; + use crate::supervisor::mailbox::{Inbox, PeerRegistry, graph_agent_id}; use serde_json::json; use std::collections::HashMap; + use std::sync::Arc; fn manager_with(pairs: &[(&str, Value)]) -> StateManager { let mut map = HashMap::new(); @@ -113,10 +135,117 @@ mod tests { state_updates: updates, output_schema: None, timeout: None, + inputs: None, teammates: false, } } + fn ctx_at_max_depth_with_peer() -> (RequestContext, Arc, String) { + let registry = Arc::new(PeerRegistry::new()); + let id = graph_agent_id("test_agent"); + let inbox = Arc::new(Inbox::new()); + registry.insert(id.clone(), "worker[0]".into(), Arc::clone(&inbox)); + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.current_depth = default_max_agent_depth(); + ctx.peer_registry = Some(Arc::clone(®istry)); + ctx.peer_assignment = Some((id.clone(), inbox)); + (ctx, registry, id) + } + + async fn execute_past_max_depth(ctx: &mut RequestContext, retire_peer_on_return: bool) { + let mut node = node_with("hi", None); + node.teammates = true; + let mut state = manager_with(&[]); + + let err = AgentNodeExecutor::execute(&node, &mut state, ctx, retire_peer_on_return) + .await + .expect_err("agent past max depth should fail before running"); + + let chain = format!("{err:#}"); + assert!(chain.contains("Agent 'test_agent' failed"), "{chain}"); + assert!(chain.contains("Max agent depth exceeded"), "{chain}"); + } + + #[tokio::test] + async fn execute_retires_peer_identity_on_frontier_when_agent_fails() { + let (mut ctx, registry, id) = ctx_at_max_depth_with_peer(); + + execute_past_max_depth(&mut ctx, true).await; + + assert!(registry.is_finished(&id)); + } + + #[tokio::test] + async fn execute_leaves_peer_identity_live_in_branch_mode_when_agent_fails() { + let (mut ctx, registry, id) = ctx_at_max_depth_with_peer(); + + execute_past_max_depth(&mut ctx, false).await; + + assert!(!registry.is_finished(&id)); + } + + fn node_with_inputs(pairs: &[(&str, &str)]) -> AgentNode { + let mut node = node_with("hi", None); + node.inputs = Some( + pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(), + ); + node + } + + #[test] + fn resolve_inputs_is_none_when_node_has_no_inputs() { + let node = node_with("hi", None); + let state = manager_with(&[("n", json!(3))]); + + assert_eq!(resolve_inputs(&node, &state).unwrap(), None); + } + + #[test] + fn resolve_inputs_lone_reference_keeps_the_array_value() { + let node = node_with_inputs(&[("items", "{{list}}")]); + let state = manager_with(&[("list", json!(["a", "b"]))]); + + let inputs = resolve_inputs(&node, &state).unwrap().unwrap(); + + assert_eq!(inputs.get("items"), Some(&json!(["a", "b"]))); + } + + #[test] + fn resolve_inputs_mixed_text_renders_to_a_string() { + let node = node_with_inputs(&[("label", "n={{n}}")]); + let state = manager_with(&[("n", json!(3))]); + + let inputs = resolve_inputs(&node, &state).unwrap().unwrap(); + + assert_eq!(inputs.get("label"), Some(&json!("n=3"))); + } + + #[test] + fn resolve_inputs_lone_reference_to_null_yields_null() { + let node = node_with_inputs(&[("x", "{{k}}")]); + let state = manager_with(&[("k", Value::Null)]); + + let inputs = resolve_inputs(&node, &state).unwrap().unwrap(); + + assert_eq!(inputs.get("x"), Some(&Value::Null)); + } + + #[test] + fn resolve_inputs_missing_key_errors_naming_the_input() { + let node = node_with_inputs(&[("width", "{{nope}}")]); + let state = manager_with(&[]); + + let err = resolve_inputs(&node, &state).expect_err("missing key must fail"); + + let chain = format!("{err:#}"); + assert!(chain.contains("inputs.width"), "{chain}"); + assert!(chain.contains("for agent 'test_agent'"), "{chain}"); + assert!(chain.contains("'nope' not found in state"), "{chain}"); + } + #[test] fn state_updates_use_output_placeholder() { let node = { @@ -162,7 +291,7 @@ mod tests { apply_state_updates(&node, &mut state, &json!("anything")); - assert_eq!(state.state().get("output"), Some(&Value::Null)); + assert!(state.state().get("output").is_none()); } #[test] diff --git a/src/graph/dispatch.rs b/src/graph/dispatch.rs index a65b36f..5cda4ec 100644 --- a/src/graph/dispatch.rs +++ b/src/graph/dispatch.rs @@ -5,6 +5,7 @@ use crate::utils::AbortSignal; use anyhow::{Context, Result, anyhow}; use log::info; use serde_json::Value; +use std::collections::HashMap; pub fn active_agent_graph_name(ctx: &RequestContext) -> Option { let name = ctx.agent.as_ref()?.name().to_string(); @@ -15,6 +16,18 @@ pub async fn run_active_agent_graph( ctx: &mut RequestContext, prompt: &str, abort_signal: AbortSignal, +) -> Result { + run_active_agent_graph_with_inputs(ctx, prompt, abort_signal, None).await +} + +/// Runs the active agent's graph with `inputs` overlaid on its +/// `initial_state` before the executor starts. `None` seeds exactly what +/// `run_active_agent_graph` does: the prompt under `initial_prompt`. +pub async fn run_active_agent_graph_with_inputs( + ctx: &mut RequestContext, + prompt: &str, + abort_signal: AbortSignal, + inputs: Option>, ) -> Result { let agent_name = active_agent_graph_name(ctx).ok_or_else(|| anyhow!("Active agent has no graph.yaml"))?; @@ -29,9 +42,7 @@ pub async fn run_active_agent_graph( .load_from_file(&graph_path) .with_context(|| format!("Failed to load graph.yaml for agent '{agent_name}'"))?; - graph - .initial_state - .insert("initial_prompt".into(), Value::String(prompt.to_string())); + seed_initial_state(&mut graph.initial_state, prompt, inputs); let executor = GraphExecutor::new(graph, agent_dir); let output = executor @@ -45,3 +56,91 @@ pub async fn run_active_agent_graph( Ok(output) } + +/// Seeds a child graph's `initial_state`, later writers winning: the graph's +/// own YAML defaults, then the caller's `inputs`, then `initial_prompt`, +/// which the dispatcher always owns. +pub fn seed_initial_state( + initial_state: &mut HashMap, + prompt: &str, + inputs: Option>, +) { + if let Some(inputs) = inputs { + initial_state.extend(inputs); + } + initial_state.insert("initial_prompt".into(), Value::String(prompt.to_string())); +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn yaml_defaults() -> HashMap { + HashMap::from([ + ("width".to_string(), json!(4)), + ("mode".to_string(), json!("fast")), + ]) + } + + #[test] + fn seed_initial_state_without_inputs_only_adds_initial_prompt() { + let mut state = yaml_defaults(); + + seed_initial_state(&mut state, "hello", None); + + let mut expected = yaml_defaults(); + expected.insert("initial_prompt".into(), json!("hello")); + assert_eq!(state, expected); + } + + #[test] + fn seed_initial_state_without_inputs_on_empty_state_inserts_exactly_one_key() { + let mut state = HashMap::new(); + + seed_initial_state(&mut state, "hello", None); + + assert_eq!( + state, + HashMap::from([("initial_prompt".to_string(), json!("hello"))]) + ); + } + + #[test] + fn seed_initial_state_inputs_override_yaml_defaults_and_keep_the_rest() { + let mut state = yaml_defaults(); + let inputs = HashMap::from([ + ("width".to_string(), json!(2)), + ("extra".to_string(), json!(["a", "b"])), + ]); + + seed_initial_state(&mut state, "hello", Some(inputs)); + + assert_eq!(state.get("width"), Some(&json!(2))); + assert_eq!(state.get("mode"), Some(&json!("fast"))); + assert_eq!(state.get("extra"), Some(&json!(["a", "b"]))); + assert_eq!(state.get("initial_prompt"), Some(&json!("hello"))); + assert_eq!(state.len(), 4); + } + + #[test] + fn seed_initial_state_dispatcher_prompt_beats_an_initial_prompt_input() { + let mut state = HashMap::from([("initial_prompt".to_string(), json!("from yaml"))]); + let inputs = HashMap::from([("initial_prompt".to_string(), json!("from inputs"))]); + + seed_initial_state(&mut state, "from dispatcher", Some(inputs)); + + assert_eq!(state.get("initial_prompt"), Some(&json!("from dispatcher"))); + assert_eq!(state.len(), 1); + } + + #[test] + fn seed_initial_state_null_input_overlays_a_default() { + let mut state = yaml_defaults(); + let inputs = HashMap::from([("width".to_string(), Value::Null)]); + + seed_initial_state(&mut state, "hello", Some(inputs)); + + assert_eq!(state.get("width"), Some(&Value::Null)); + } +} diff --git a/src/graph/executor.rs b/src/graph/executor.rs index aa854b5..e441e21 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -9,21 +9,33 @@ use super::state::StateManager; use super::types::{EndNode, Graph, Node, NodeType}; use super::user_interaction::{ApprovalNodeExecutor, InputNodeExecutor}; use super::validator::{AgentValidationContext, GraphValidator}; +use super::wall_clock; use crate::config::{RenderMode, RequestContext}; use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id}; -use crate::utils::AbortSignal; +use crate::utils::{AbortSignal, wait_abort_signal, wait_user_interrupt}; use anyhow::{Context, Result, anyhow, bail}; use futures_util::future::join_all; use serde_json::Value; use std::collections::{HashMap, HashSet}; +use std::future::Future; use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::Semaphore; +use tokio::task::{AbortHandle, JoinHandle}; + +/// 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 +43,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 +78,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); @@ -83,11 +108,36 @@ impl GraphExecutor { .unwrap_or_default(); let script_executor = ScriptExecutor::new(&base_dir).with_envs(agent_envs); let max_iterations = graph.settings.max_loop_iterations; - let graph_timeout = graph.settings.timeout.map(Duration::from_secs); + let graph_timeout = graph.settings.timeout.and_then(wall_clock); let max_concurrency = graph.settings.max_concurrency; + if max_concurrency > Semaphore::MAX_PERMITS { + bail!( + "Graph '{}': settings.max_concurrency {max_concurrency} exceeds the runtime limit of {}", + graph.name, + Semaphore::MAX_PERMITS + ); + } let graph = Arc::new(graph); let start = Instant::now(); + // Maps a user interrupt (SIGINT, or the enclosing turn aborting) onto + // this run's abort flag for as long as the run lives. Without a + // session signal (ACP, bare executor) SIGINT handling is unchanged. + let _bridge = match ctx.session_abort.clone() { + Some(session) if session.aborted() => { + abort_signal.set_ctrlc(); + None + } + Some(session) => { + let graph_abort = abort_signal.clone(); + Some(TaskCancelGuard::spawn_one(async move { + wait_user_interrupt(Some(&session)).await; + graph_abort.set_ctrlc(); + })) + } + None => None, + }; + let mut frontier: HashSet = HashSet::from([graph.start.clone()]); logger.graph_start(&graph.start, graph.nodes.len()); @@ -120,10 +170,11 @@ impl GraphExecutor { // Loop-count and visit tracking on live state, BEFORE forking. // This counts every entry to a node toward max_loop_iterations // regardless of how many parallel branches converged on it. + // A cap of 0 disables the check. for node_id in &frontier { state.state_mut().visit_node(node_id); let visits = state.state().loop_count(node_id); - if visits > max_iterations { + if max_iterations > 0 && visits > max_iterations { bail!( "Node '{}' visited {} times (max_loop_iterations={}). \ Possible infinite loop.", @@ -179,10 +230,13 @@ impl GraphExecutor { logger.node_start(&node, in_super_step); let branch_state = state.fork_for_branch_state(); let mut branch_ctx = ctx.fork_for_branch(); + let mut peer_id: Option = None; if let Some(assignment) = peer_assignments.remove(node_id) { + peer_id = Some(assignment.0.clone()); branch_ctx.peer_registry = peer_registry.clone(); branch_ctx.peer_assignment = Some(assignment); } + let registry_for_task = peer_registry.clone(); if in_super_step { branch_ctx.render_mode = RenderMode::Silent; } @@ -191,8 +245,24 @@ 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(); + // Retires the teammate identity however this task ends: + // dropped explicitly right after step() on the normal path, + // by unwinding on the abort return, task abort, or panic. + // Built before the spawn so a task aborted before its first + // poll still drops it with the future. + let retire = registry_for_task + .zip(peer_id) + .map(|(registry, id)| PeerRetireGuard { + registry, + id, + #[cfg(test)] + observer, + }); let task = tokio::spawn(async move { + let retire = retire; let _permit = sem_clone .acquire() .await @@ -216,12 +286,14 @@ impl GraphExecutor { let mut state = branch_state; let mut ctx = branch_ctx; let step_ctx = StepContext { - graph: graph_clone.as_ref(), + graph: Arc::clone(&graph_clone), script_executor: &script_exec_clone, max_concurrency, abort_signal: &abort_clone, + branch_mode: false, }; let result = step(&node, &mut state, &mut ctx, &step_ctx, ¤t).await; + drop(retire); let elapsed = node_start.elapsed(); match &result { Ok(StepResult::Continue(targets)) => { @@ -262,29 +334,34 @@ impl GraphExecutor { branch_tasks.push(task); } - let joined = match graph_timeout { - Some(t) => { - let remaining = t.saturating_sub(start.elapsed()); - let abort_handles: Vec<_> = branch_tasks - .iter() - .map(|task| task.abort_handle()) - .collect(); - match tokio::time::timeout(remaining, join_all(branch_tasks)).await { - Ok(joined) => joined, - Err(_) => { - for handle in abort_handles { - handle.abort(); - } - bail!( - "Graph '{}' timed out after {}s during super-step with frontier {:?}", - graph.name, - t.as_secs(), - sorted_frontier(&frontier) - ); - } + // Owns the branch tasks until the join completes: a timeout or + // abort bail drops it and cancels whatever is still in flight. + let _cancel = TaskCancelGuard::new(&branch_tasks); + let bounded_join = async { + match graph_timeout { + Some(t) => { + let remaining = t.saturating_sub(start.elapsed()); + tokio::time::timeout(remaining, join_all(branch_tasks)) + .await + .map_err(|_| { + anyhow!( + "Graph '{}' timed out after {}s during super-step with frontier {:?}", + graph.name, + t.as_secs(), + sorted_frontier(&frontier) + ) + }) } + None => Ok(join_all(branch_tasks).await), } - None => join_all(branch_tasks).await, + }; + let joined = tokio::select! { + joined = bounded_join => joined?, + _ = wait_abort_signal(&abort_signal) => bail!( + "Graph '{}' aborted during super-step with frontier {:?}", + graph.name, + sorted_frontier(&frontier) + ), }; let mut branch_writes: Vec = Vec::new(); @@ -397,11 +474,55 @@ fn provision_frontier_peers( (Some(registry), assignments) } +pub(super) struct TaskCancelGuard(Vec); + +impl TaskCancelGuard { + pub(super) fn new(tasks: &[JoinHandle]) -> Self { + Self(tasks.iter().map(|task| task.abort_handle()).collect()) + } + + pub(super) fn spawn_one(fut: F) -> Self + where + F: Future + Send + 'static, + { + Self(vec![tokio::spawn(fut).abort_handle()]) + } +} + +impl Drop for TaskCancelGuard { + fn drop(&mut self) { + for handle in &self.0 { + handle.abort(); + } + } +} + +/// Retires a teammate identity when dropped, so a branch or item task that +/// is aborted, panics, or returns early still leaves its peers seeing it as +/// finished. `mark_finished` is idempotent; an explicit earlier call is fine. +pub(super) struct PeerRetireGuard { + pub(super) registry: Arc, + pub(super) id: String, + #[cfg(test)] + pub(super) observer: Option, +} + +impl Drop for PeerRetireGuard { + fn drop(&mut self) { + self.registry.mark_finished(&self.id); + #[cfg(test)] + if let Some(observe) = &self.observer { + observe(&self.registry, &self.id); + } + } +} + pub(super) struct StepContext<'a> { - pub graph: &'a Graph, + pub graph: Arc, pub script_executor: &'a ScriptExecutor, pub max_concurrency: usize, pub abort_signal: &'a AbortSignal, + pub branch_mode: bool, } impl StepContext<'_> { @@ -410,7 +531,7 @@ impl StepContext<'_> { } } -enum StepResult { +pub(super) enum StepResult { // The set of next-node ids the executor should add to the next super-step's // frontier. A `Vec` of length 1 for sequential routing (default) and the // full target list for fan-out (`next: [a, b, ...]`). Dynamic single-route @@ -420,7 +541,7 @@ enum StepResult { End(String), } -async fn step( +pub(super) async fn step( node: &Node, state: &mut StateManager, ctx: &mut RequestContext, @@ -429,8 +550,8 @@ async fn step( ) -> Result { match &node.node_type { NodeType::Agent(agent_node) => { - AgentNodeExecutor::execute(agent_node, state, ctx).await?; - let targets = static_next_targets(node, current, "agent")?; + AgentNodeExecutor::execute(agent_node, state, ctx, !step_ctx.branch_mode).await?; + let targets = static_next_targets(node, current, "agent", step_ctx.branch_mode)?; Ok(StepResult::Continue(targets)) } NodeType::Script(script_node) => { @@ -452,7 +573,7 @@ async fn step( }; let targets = match dynamic { Some(n) => vec![n], - None => static_next_targets(node, current, "script")?, + None => static_next_targets(node, current, "script", step_ctx.branch_mode)?, }; Ok(StepResult::Continue(targets)) } @@ -466,28 +587,42 @@ async fn step( Ok(StepResult::Continue(vec![next])) } NodeType::Llm(llm_node) => { - let outcome = LlmNodeExecutor::execute(current, llm_node, state, ctx).await?; + let outcome = + LlmNodeExecutor::execute(current, llm_node, state, ctx, step_ctx.abort_signal) + .await?; let targets = match outcome { - LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?, + LlmExecutionOutcome::Continue => { + static_next_targets(node, current, "llm", step_ctx.branch_mode)? + } LlmExecutionOutcome::FellBack(target) => vec![target], }; Ok(StepResult::Continue(targets)) } NodeType::Rag(rag_node) => { RagNodeExecutor::execute(rag_node, current, state, ctx).await?; - let targets = static_next_targets(node, current, "rag")?; + let targets = static_next_targets(node, current, "rag", step_ctx.branch_mode)?; Ok(StepResult::Continue(targets)) } NodeType::End(end_node) => Ok(StepResult::End(resolve_end_output(end_node, state))), NodeType::Map(map_node) => { - let targets = static_next_targets(node, current, "map")?; + let targets = static_next_targets(node, current, "map", step_ctx.branch_mode)?; MapNodeExecutor::execute(map_node, state, ctx, step_ctx, current).await?; Ok(StepResult::Continue(targets)) } } } -fn static_next_targets(node: &Node, current: &str, kind: &str) -> Result> { +/// Inside a map branch a node without `next` simply ends the item's chain; on +/// the main flow it is a routing error. +fn static_next_targets( + node: &Node, + current: &str, + kind: &str, + branch_mode: bool, +) -> Result> { + if node.next.is_none() && branch_mode { + return Ok(Vec::new()); + } node.next .as_ref() .map(|t| t.as_slice().to_vec()) @@ -517,7 +652,7 @@ fn apply_simple_state_updates(updates: Option<&HashMap>, state: #[cfg(test)] mod tests { - use super::super::types::{AgentNode, GraphSettings}; + use super::super::types::{AgentNode, GraphSettings, NextTargets}; use super::*; use indexmap::IndexMap; use serde_json::json; @@ -612,6 +747,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates, }), next: None, @@ -730,20 +866,59 @@ mod tests { assert!(registry.is_none()); assert!(assignments.is_empty()); } + + #[test] + fn static_next_targets_branch_mode_treats_missing_next_as_terminal() { + let node = agent_node("n", false); + + let targets = static_next_targets(&node, "n", "agent", true).unwrap(); + + assert!(targets.is_empty()); + } + + #[test] + fn static_next_targets_main_flow_still_errors_without_next() { + let node = agent_node("n", false); + + let msg = static_next_targets(&node, "n", "agent", false) + .unwrap_err() + .to_string(); + + assert!( + msg.contains("agent node 'n' has no `next` and is not an end node"), + "{msg}" + ); + } + + #[test] + fn static_next_targets_returns_declared_targets_in_both_modes() { + let mut node = agent_node("n", false); + node.next = Some(NextTargets::Many(vec!["a".into(), "b".into()])); + + for branch_mode in [false, true] { + let targets = static_next_targets(&node, "n", "agent", branch_mode).unwrap(); + assert_eq!(targets, vec!["a".to_string(), "b".to_string()]); + } + } } #[cfg(test)] mod integration_tests { use super::*; + use crate::config::paths; use crate::config::{AppState, WorkingMode}; #[cfg(unix)] use crate::function::jobs::RingBuf; #[cfg(unix)] use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor, notification}; - use crate::utils::{create_abort_signal, temp_file}; + use crate::utils::{create_abort_signal, get_env_name, temp_file}; + use parking_lot::Mutex; + use serial_test::serial; + use std::env; use std::fs; #[cfg(unix)] use std::mem; + use std::time::{SystemTime, UNIX_EPOCH}; fn cmd_available(name: &str) -> bool { which::which(name).is_ok() @@ -834,6 +1009,63 @@ nodes: assert!(strs.contains(&"beta"), "missing 'beta' in {strs:?}"); } + /// Two parallel agent branches whose `state_updates` only read `{{output}}` + /// leave the key as they found it (absent), so the join sees no write to + /// `output` from either side and needs no reducer for it. + #[tokio::test] + #[serial] + async fn parallel_state_updates_branches_do_not_contend_on_output() { + if !cmd_available("bash") || !cmd_available("python3") { + eprintln!("skipping: bash or python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let ws = TestWorkspace::new(); + ws.write_script("dispatcher.sh", "#!/bin/bash\necho '{}'\n"); + + let yaml = format!( + r#" +name: t +settings: + validate_before_run: false +start: dispatcher +nodes: + dispatcher: + type: script + script: dispatcher.sh + state_updates: {{}} + next: [a, b] + a: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note_a: "{{{{output}}}}" + next: join + b: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note_b: "{{{{output}}}}" + next: join + join: + type: end + output: "{{{{note_a}}}}|{{{{note_b}}}}" +"# + ); + let graph: Graph = serde_yaml::from_str(&yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await + .unwrap_or_else(|e| panic!("executor failed: {e:#}")); + + assert_eq!(result, "held|held"); + } + #[tokio::test] async fn map_over_list_collects_outputs_in_input_order() { if !cmd_available("python3") { @@ -953,6 +1185,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") { @@ -1037,6 +1369,153 @@ nodes: assert!(err.contains("sleeper"), "error should name frontier: {err}"); } + #[tokio::test] + async fn settings_max_concurrency_above_semaphore_limit_errors_instead_of_panicking() { + let ws = TestWorkspace::new(); + let yaml = r#" +name: oversized_cap_test +start: done +settings: + validate_before_run: false +nodes: + done: + type: end + output: "done" +"#; + let mut graph: Graph = serde_yaml::from_str(yaml).unwrap(); + graph.settings.max_concurrency = Semaphore::MAX_PERMITS + 1; + let mut ctx = make_ctx(); + let err = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, create_abort_signal()) + .await + .expect_err("an oversized cap is an error, not a panic"); + let msg = format!("{err:#}"); + assert!(msg.contains("exceeds the runtime limit of"), "{msg}"); + } + + #[tokio::test] + async fn graph_timeout_zero_is_unbounded() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + ws.write_script("sleeper.sh", "#!/bin/bash\nsleep 1.5\necho '{}'\n"); + + let yaml = r#" +name: zero_timeout_test +start: sleeper +settings: + timeout: 0 +nodes: + sleeper: + type: script + script: sleeper.sh + state_updates: {} + next: done + done: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await; + + result.unwrap_or_else(|e| panic!("timeout: 0 should not bound the run: {e:#}")); + } + + /// Bash script that bumps `n` in state and re-enters `looper` while + /// `n < stop`; once `n == stop` it omits `_next` so the static + /// `next: done` edge is taken. + fn looper_script(stop: usize) -> String { + format!( + "#!/bin/bash\n\ + t=${{GRAPH_STATE#*'\"n\":'}}\n\ + n=${{t%%[!0-9]*}}\n\ + n=$((n + 1))\n\ + if (( n < {stop} )); then\n\ + printf '{{\"n\": %d, \"_next\": \"looper\"}}' \"$n\"\n\ + else\n\ + printf '{{\"n\": %d}}' \"$n\"\n\ + fi\n" + ) + } + + fn loop_cap_graph(max_loop_iterations: usize) -> Graph { + let yaml = format!( + r#" +name: loop_cap_test +start: looper +initial_state: + n: 0 +settings: + max_loop_iterations: {max_loop_iterations} +nodes: + looper: + type: script + script: looper.sh + next: done + done: + type: end + output: "{{{{n}}}}" +"# + ); + serde_yaml::from_str(&yaml).unwrap() + } + + #[tokio::test] + async fn max_loop_iterations_zero_is_unbounded() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + // Loop past the default cap so a fallback-to-default implementation + // would be caught, not just a lifted small cap. + let stop = crate::graph::DEFAULT_MAX_LOOP_ITERATIONS + 1; + let ws = TestWorkspace::new(); + ws.write_script("looper.sh", &looper_script(stop)); + + let graph = loop_cap_graph(0); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await + .unwrap_or_else(|e| panic!("max_loop_iterations: 0 should not cap visits: {e:#}")); + + assert_eq!(result, stop.to_string()); + } + + #[tokio::test] + async fn max_loop_iterations_nonzero_still_bounds() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + // Always re-enter `looper`: only the cap can end this run. + ws.write_script("looper.sh", "#!/bin/bash\necho '{\"_next\": \"looper\"}'\n"); + + let graph = loop_cap_graph(2); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await; + + assert!(result.is_err(), "expected the visit cap to abort the run"); + let err = format!("{:#}", result.unwrap_err()); + assert!( + err.contains( + "Node 'looper' visited 3 times (max_loop_iterations=2). Possible infinite loop." + ), + "error should report the visit cap: {err}" + ); + } + #[cfg(unix)] #[tokio::test] async fn background_job_survives_graph_node_execution() { @@ -1115,4 +1594,416 @@ nodes: assert_eq!(events[0].id, "job_bg"); assert_eq!(events[0].event, "job_completed"); } + + 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-executor-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; its single script node sleeps for + /// `hold_secs` and then writes `output`, so an agent node can be held + /// in flight without any LLM. + 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(); + } + + #[tokio::test] + async fn task_cancel_guard_aborts_owned_tasks() { + let tasks: Vec> = (0..3) + .map(|_| tokio::spawn(std::future::pending::<()>())) + .collect(); + + drop(TaskCancelGuard::new(&tasks)); + + for task in tasks { + let err = task + .await + .expect_err("an owned task is aborted when the guard drops"); + assert!(err.is_cancelled(), "{err}"); + } + } + + #[tokio::test] + async fn peer_retire_guard_fires_on_task_abort() { + let (registry, assignments) = provision_frontier_peers(vec![ + ("a".to_string(), "worker".to_string()), + ("b".to_string(), "worker".to_string()), + ]); + let registry = registry.expect("two flagged nodes should get a registry"); + let id_a = assignments["a"].0.clone(); + let id_b = assignments["b"].0.clone(); + let guard = PeerRetireGuard { + registry: Arc::clone(®istry), + id: id_a.clone(), + observer: None, + }; + let task = tokio::spawn(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + + task.abort(); + let err = task.await.expect_err("the task was aborted"); + + assert!(err.is_cancelled(), "{err}"); + assert!( + registry.is_finished(&id_a), + "the aborted task's identity is retired by the guard" + ); + assert!(!registry.is_finished(&id_b), "the sibling is untouched"); + } + + /// A task aborted before its first poll never runs its body, so a guard + /// built inside the future would never exist. Built outside and moved in, + /// it drops with the unpolled future and still retires the identity. + #[tokio::test] + async fn peer_retire_guard_fires_when_task_is_aborted_before_first_poll() { + let (registry, assignments) = provision_frontier_peers(vec![ + ("a".to_string(), "worker".to_string()), + ("b".to_string(), "worker".to_string()), + ]); + let registry = registry.expect("two flagged nodes should get a registry"); + let id_a = assignments["a"].0.clone(); + let id_b = assignments["b"].0.clone(); + + let outside = PeerRetireGuard { + registry: Arc::clone(®istry), + id: id_a.clone(), + observer: None, + }; + let inside_registry = Arc::clone(®istry); + let inside_id = id_b.clone(); + let task = tokio::spawn(async move { + let _outside = outside; + let _inside = PeerRetireGuard { + registry: inside_registry, + id: inside_id, + observer: None, + }; + std::future::pending::<()>().await; + }); + // No await between spawn and abort: on the current-thread test + // runtime the task has not been polled yet. + task.abort(); + let err = task.await.expect_err("the task was aborted"); + + assert!(err.is_cancelled(), "{err}"); + assert!(registry.is_finished(&id_a), "the moved-in guard retires"); + assert!( + !registry.is_finished(&id_b), + "a guard built inside never existed" + ); + } + + /// A session already aborted when the run starts trips the synchronous + /// pre-check: no bridge task, no super-step, the start script never runs. + #[tokio::test] + async fn graph_run_aborts_when_session_is_preset() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + let sentinel = ws.dir.join("ran"); + ws.write_script( + "mark.sh", + &format!("#!/bin/bash\ntouch '{}'\necho '{{}}'\n", sentinel.display()), + ); + + let yaml = r#" +name: preset_session_abort_test +start: mark +nodes: + mark: + type: script + script: mark.sh + state_updates: {} + next: done + done: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let session = create_abort_signal(); + session.set_ctrlc(); + ctx.session_abort = Some(session); + let graph_abort = create_abort_signal(); + + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, graph_abort.clone()) + .await; + + let err = format!( + "{:#}", + result.expect_err("a pre-set session aborts the run") + ); + assert!(err.contains("aborted before super-step"), "{err}"); + assert!( + graph_abort.aborted(), + "the session flag is mapped onto the graph's own abort" + ); + assert!(!sentinel.exists(), "the start script must never run"); + } + + /// The graph abort and the session abort are different signals and only + /// the session one is set, mid-run: the bridge is what must fire, and the + /// cancelled branch task must kill its script before the `&&` runs. + #[tokio::test] + async fn session_abort_mid_run_aborts_graph() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + let sentinel = ws.dir.join("finished"); + ws.write_script( + "hold.sh", + &format!( + "#!/bin/bash\nsleep 1.5 && touch '{}'\necho '{{}}'\n", + sentinel.display() + ), + ); + + let yaml = r#" +name: mid_run_session_abort_test +start: hold +nodes: + hold: + type: script + script: hold.sh + state_updates: {} + next: done + done: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let session = create_abort_signal(); + ctx.session_abort = Some(Arc::clone(&session)); + let graph_abort = create_abort_signal(); + assert!(!Arc::ptr_eq(&session, &graph_abort)); + + let trigger = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(300)).await; + session.set_ctrlc(); + }); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, graph_abort.clone()) + .await; + trigger.await.unwrap(); + + let err = format!("{:#}", result.expect_err("a session abort ends the run")); + assert!(err.contains("aborted"), "{err}"); + assert!(graph_abort.aborted(), "the bridge set the graph's abort"); + + tokio::time::sleep(Duration::from_millis(2500)).await; + assert!( + !sentinel.exists(), + "the cancelled branch's script was killed before its `&&` ran" + ); + } + + /// Two flagged agent nodes under one permit: the graph abort lands while + /// the first is holding its probe agent open and the second is still + /// waiting on the semaphore. Both tasks are cancelled by the abort, and + /// each guard retires its identity as the task unwinds. + #[tokio::test] + #[serial] + async fn aborted_frontier_agents_are_retired() { + if !cmd_available("bash") || !cmd_available("python3") { + eprintln!("skipping: bash or python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(5.0); + let ws = TestWorkspace::new(); + // The abort must land after the worker tasks exist (their retirement + // guards are created at spawn), i.e. after the dispatcher super-step + // has run. A fixed delay races bash start-up on slow runners, so the + // dispatcher leaves a marker and the trigger waits for it. + let dispatched = ws.dir.join("dispatched"); + ws.write_script( + "dispatcher.sh", + &format!( + "#!/bin/bash\necho '{{}}'\ntouch '{}'\n", + dispatched.display() + ), + ); + + let yaml = format!( + r#" +name: aborted_frontier_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: {PROBE_AGENT} + prompt: "p" + teammates: true + next: join + worker_b: + type: agent + agent: {PROBE_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 (peer id, label, is_finished)). + type RosterSnapshot = Vec<(String, 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)| { + let finished = registry.is_finished(&peer_id); + (peer_id, label, finished) + }) + .collect(); + sink.lock().push((id.to_string(), snapshot)); + }); + + let trigger = { + let abort = abort.clone(); + let dispatched = dispatched.clone(); + tokio::spawn(async move { + while !dispatched.exists() { + tokio::time::sleep(Duration::from_millis(20)).await; + } + tokio::time::sleep(Duration::from_millis(500)).await; + abort.set_ctrlc(); + }) + }; + let started = Instant::now(); + let result = GraphExecutor::new(graph, &ws.dir) + .with_frontier_observer(observer) + .execute(&mut ctx, abort) + .await; + trigger.await.unwrap(); + + let err = format!("{:#}", result.expect_err("the abort ends the run")); + assert!(err.contains("aborted"), "{err}"); + assert!( + started.elapsed() < Duration::from_secs(3), + "the run must not wait for the 5s probe hold: {:?}", + started.elapsed() + ); + + // Aborted tasks unwind on the runtime's next pass, not inside execute(). + tokio::time::sleep(Duration::from_millis(300)).await; + let observations = observations.lock(); + assert_eq!( + observations.len(), + 2, + "each cancelled branch retires its identity exactly once: {observations:?}" + ); + for (retired, snapshot) in observations.iter() { + assert!( + snapshot + .iter() + .any(|(peer_id, _, finished)| peer_id == retired && *finished), + "the retired id reads as finished in its own observation: {observations:?}" + ); + assert!( + snapshot + .iter() + .all(|(_, label, _)| label.starts_with("worker_")), + "{snapshot:?}" + ); + } + assert_ne!(observations[0].0, observations[1].0); + let (_, last) = &observations[1]; + assert!( + last.iter().all(|(_, _, finished)| *finished), + "both identities are retired once the last task unwinds: {last:?}" + ); + } } diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 0b9faae..a409469 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -1,6 +1,8 @@ use super::state::StateManager; +use super::state_updates; use super::structured; use super::types::LlmNode; +use super::wall_clock; use crate::client::{Model, ModelType, call_chat_completions}; use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ @@ -9,17 +11,14 @@ use crate::config::{ use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::function::jobs::reap_jobs; use crate::function::skill::skill_function_declarations; -use crate::utils::create_abort_signal; +use crate::utils::AbortSignal; use anyhow::{Context, Error, Result, anyhow, bail}; use log::warn; use serde_json::Value; use std::collections::HashSet; use std::sync::Arc; -use std::time::Duration; use tokio::time::timeout; -const OUTPUT_KEY: &str = "output"; - #[derive(Debug, PartialEq, Eq)] pub(super) enum LlmExecutionOutcome { Continue, @@ -34,8 +33,9 @@ impl LlmNodeExecutor { node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { - let result = run(node_id, node, state_manager, parent_ctx).await; + let result = run(node_id, node, state_manager, parent_ctx, abort).await; let (output, failure_reason) = match result { Ok(raw) => match &node.output_schema { Some(schema) => match structured::extract(&raw, schema, parent_ctx).await { @@ -84,6 +84,7 @@ async fn run( node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { let mut instructions: Option = match &node.instructions { Some(s) => Some( @@ -186,17 +187,12 @@ async fn run( node.mcp_tools.clone().map(|map| (node_id.to_string(), map)), ); parent_ctx.refresh_mcp_tool_filters(); - let result = match node.timeout { - Some(secs) => match timeout( - Duration::from_secs(secs), - run_with_retries(node, &prompt, parent_ctx), - ) - .await - { + let result = match node.timeout.and_then(wall_clock) { + Some(d) => match timeout(d, run_with_retries(node, &prompt, parent_ctx, abort)).await { Ok(r) => r, - Err(_) => Err(anyhow!("llm node timed out after {secs}s")), + Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())), }, - None => run_with_retries(node, &prompt, parent_ctx).await, + None => run_with_retries(node, &prompt, parent_ctx, abort).await, }; parent_ctx.role = saved_role; let node_jobs = @@ -248,10 +244,11 @@ async fn run_with_retries( node: &LlmNode, prompt: &str, ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { let mut last_err: Option = None; for attempt in 1..=node.max_attempts { - match run_chat_loop(node, prompt, ctx).await { + match run_chat_loop(node, prompt, ctx, abort).await { Ok(out) => return Ok(out), Err(e) if is_transient(&e) && attempt < node.max_attempts => { warn!("llm node attempt {attempt} failed (transient): {e}; retrying"); @@ -263,14 +260,28 @@ async fn run_with_retries( Err(last_err.unwrap_or_else(|| anyhow!("llm node exhausted retries"))) } -async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -> Result { - let abort = create_abort_signal(); +/// Whether `turn` (0-based) is the final turn the node's cap allows. A cap of 0 means no cap. +pub(crate) fn is_last_turn(turn: u32, max_iterations: u32) -> bool { + max_iterations > 0 && turn == max_iterations - 1 +} + +async fn run_chat_loop( + node: &LlmNode, + prompt: &str, + ctx: &mut RequestContext, + abort: &AbortSignal, +) -> Result { + let abort = abort.clone(); let app_cfg = Arc::clone(&ctx.app.config); let role_for_input = ctx.role.clone(); let mut input = Input::from_str(ctx, prompt, role_for_input)?; let mut accumulated = String::new(); - for turn in 0..node.max_iterations { + let mut turn: u32 = 0; + loop { + if abort.aborted() { + bail!("llm node aborted"); + } let client = input.create_client()?; ctx.before_chat_completion(&input)?; let (output, tool_results) = @@ -297,7 +308,7 @@ async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) - return Ok(accumulated); } GuardrailAction::Inject(prompt) => { - if turn + 1 == node.max_iterations { + if is_last_turn(turn, node.max_iterations) { bail!( "llm node hit max_iterations ({}) before LLM concluded", node.max_iterations @@ -305,22 +316,19 @@ async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) - } let role = ctx.role.clone(); input = Input::from_str(ctx, &prompt, role)?; - continue; } } + } else { + if is_last_turn(turn, node.max_iterations) { + bail!( + "llm node hit max_iterations ({}) before LLM concluded", + node.max_iterations + ); + } + input = input.merge_tool_results(output, tool_results); } - - if turn + 1 == node.max_iterations { - bail!( - "llm node hit max_iterations ({}) before LLM concluded", - node.max_iterations - ); - } - - input = input.merge_tool_results(output, tool_results); + turn = turn.saturating_add(1); } - - bail!("llm node ended without producing output") } fn build_inline_role( @@ -456,37 +464,12 @@ fn apply_state_updates_with_output( state_manager: &mut StateManager, output: &Value, ) { - if node.output_schema.is_some() - && let Some(obj) = output.as_object() - { - for (k, v) in obj { - state_manager.state_mut().set(k.clone(), v.clone()); - } - } - - let Some(updates) = &node.state_updates else { - return; - }; - let prev_output = state_manager.state().get(OUTPUT_KEY).cloned(); - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), output.clone()); - - for (key, template) in updates { - let value = state_manager.interpolate_lenient(template); - state_manager - .state_mut() - .set(key.clone(), Value::String(value)); - } - - match prev_output { - Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v), - None => { - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), Value::Null); - } - } + state_updates::apply( + state_manager, + output, + node.output_schema.is_some(), + node.state_updates.as_ref(), + ); } fn format_schema_hint(schema: &Value) -> String { @@ -499,10 +482,14 @@ fn format_schema_hint(schema: &Value) -> String { #[cfg(test)] mod tests { + use super::super::state_updates::OUTPUT_KEY; use super::super::types::*; use super::*; + use crate::config::{Agent, AgentConfig, AppState, WorkingMode}; + use crate::utils::create_abort_signal; use serde_json::json; use std::collections::HashMap; + use std::time::Duration; fn manager_with(pairs: &[(&str, Value)]) -> StateManager { let mut map = HashMap::new(); @@ -568,7 +555,7 @@ mod tests { apply_state_updates_with_output(&node, &mut state, &json!("anything")); - assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); + assert!(state.state().get(OUTPUT_KEY).is_none()); } #[test] @@ -756,4 +743,64 @@ mod tests { assert!(!is_transient(&anyhow!("hit max_iterations"))); assert!(!is_transient(&anyhow!("authentication failed"))); } + + #[test] + fn zero_timeout_resolves_to_no_bound() { + assert!(Some(0u64).and_then(wall_clock).is_none()); + assert_eq!( + Some(5u64).and_then(wall_clock), + Some(Duration::from_secs(5)) + ); + } + + #[test] + fn is_last_turn_bounded() { + assert!(is_last_turn(0, 1)); + assert!(!is_last_turn(0, 10)); + assert!(is_last_turn(9, 10)); + assert!(!is_last_turn(10, 10)); + } + + #[test] + fn is_last_turn_zero_cap_never_fires() { + assert!(!is_last_turn(0, 0)); + assert!(!is_last_turn(9, 0)); + assert!(!is_last_turn(u32::MAX, 0)); + } + + #[test] + fn is_last_turn_at_u32_max_does_not_overflow() { + assert!(is_last_turn(u32::MAX - 1, u32::MAX)); + assert!(!is_last_turn(u32::MAX, u32::MAX)); + } + + /// The turn-top abort check runs before any client is created, so an + /// already-aborted graph never issues a model call and the node fails + /// with the abort reason rather than a connection error. + #[tokio::test] + async fn run_chat_loop_bails_before_first_call_when_aborted() { + let node = node_with(None); + let mut state = manager_with(&[]); + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + let abort = create_abort_signal(); + abort.set_ctrlc(); + + let err = LlmNodeExecutor::execute("think", &node, &mut state, &mut ctx, &abort) + .await + .expect_err("a pre-set abort must fail the node"); + + let chain = format!("{err:#}"); + assert!( + chain.contains( + "LLM node failed and no fallback declared: LLM call failed: llm node aborted" + ), + "{chain}" + ); + assert_eq!( + state.state().get(OUTPUT_KEY), + None, + "no output is recorded for an aborted node without state_updates" + ); + } } diff --git a/src/graph/map.rs b/src/graph/map.rs index e541950..29ad7de 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -1,195 +1,461 @@ -use super::agent::AgentNodeExecutor; -use super::executor::StepContext; -use super::llm::LlmNodeExecutor; -use super::rag::RagNodeExecutor; -use super::state::StateManager; -use super::types::{MapNode, NodeType}; +use super::executor::{PeerRetireGuard, StepContext, StepResult, TaskCancelGuard, step}; +use super::state::{StateManager, is_lone_template}; +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}; -use futures_util::future::join_all; +use anyhow::{Context, Result, anyhow, bail}; +use futures_util::future::{BoxFuture, join_all}; use serde_json::Value; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::Arc; use tokio::sync::Semaphore; pub(super) struct MapNodeExecutor; impl MapNodeExecutor { - pub(super) async fn execute( - 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(); - - let branch_node = step_ctx - .graph - .get_node(&node.branch) - .ok_or_else(|| { - anyhow!( - "map node '{node_id}': branch '{}' not found in graph", - node.branch - ) - })? - .clone(); - - let max_conc = node - .max_concurrency - .unwrap_or(step_ctx.max_concurrency) - .max(1); - let semaphore = Arc::new(Semaphore::new(max_conc)); - let mut sub_tasks = Vec::with_capacity(items.len()); - - let peers = provision_map_peers(&branch_node.node_type, &node.branch, items.len()); - - for (idx, item) in items.iter().enumerate() { - let item = item.clone(); - let as_name = node.as_name.clone(); - let branch_clone = branch_node.clone(); - let mut sub_state = state.fork_for_branch_state(); - let mut sub_ctx = ctx.fork_for_branch(); - sub_ctx.render_mode = RenderMode::Silent; - if let Some((registry, assignments)) = &peers { - sub_ctx.peer_registry = Some(Arc::clone(registry)); - sub_ctx.peer_assignment = Some(assignments[idx].clone()); - } - let script_clone = step_ctx.script_executor.clone(); - let sub_branch_id = node.branch.clone(); - let sem = semaphore.clone(); - let abort = step_ctx.abort_signal.clone(); - - sub_state.state_mut().set(as_name, item); - - let task = tokio::spawn(async move { - let _permit = sem - .acquire() - .await - .expect("map semaphore should not be closed"); - if abort.aborted() { - return ( - idx, - sub_state, - Err(anyhow!("map sub-branch [{idx}] aborted")), - ); - } - let mut state = sub_state; - let mut ctx = sub_ctx; - - let exec_result: Result<()> = match &branch_clone.node_type { - NodeType::Llm(n) => { - LlmNodeExecutor::execute(&branch_clone.id, n, &mut state, &mut ctx) - .await - .map(|_| ()) - } - NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx) - .await - .map(|_| ()), - NodeType::Rag(n) => { - RagNodeExecutor::execute(n, &sub_branch_id, &mut state, &mut ctx).await - } - NodeType::Script(n) => script_clone.execute(n, &mut state).await.map(|_| ()), - _ => Err(anyhow!( - "map branch '{}' has type that cannot run inside a map \ - (validator should have caught this; internal error)", - branch_clone.id - )), - }; - - (idx, state, exec_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 = 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 - .with_context(|| format!("map node '{node_id}': sub-branch [{idx}] failed"))?; - - 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` '{}'", - 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(()) + /// 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)) } } -/// Pre-provision teammate identities before any branch starts, so a fast -/// sibling can message one still waiting on the semaphore (the message -/// queues in the pre-created inbox). A single-item fan-out has no peers, so -/// it gets no registry. +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 + ); + } + if node.as_name == node.output_key { + bail!( + "map node '{node_id}': `as` and `output_key` are both '{}'; the item binding \ + would be collected as the chain's result", + node.as_name + ); + } + 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, 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(); + + // Cancellation net behind the chain runner's own retirement: a task + // aborted before or during its chain still leaves the item's identity + // finished. Built before the spawn so a task aborted before its first + // poll drops it with the future. + let retire = item_peers + .as_ref() + .map(|(registry, assignment)| PeerRetireGuard { + registry: Arc::clone(registry), + id: assignment.0.clone(), + #[cfg(test)] + observer: None, + }); + let task = tokio::spawn(async move { + let _retire = retire; + 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); + } + + // Owns the item tasks across the join: if this future is dropped (the + // map's own task aborted), every in-flight chain is cancelled with it. + let _cancel = TaskCancelGuard::new(&sub_tasks); + 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 = 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, + idx: usize, + peers: Option<&'a (Arc, 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. Errors carry no map/item locator; the +/// caller's context supplies it. +async fn run_chain_step( + chain: &ItemChain<'_>, + state: &mut StateManager, + ctx: &mut RequestContext, + step_ctx: &StepContext<'_>, + current: &str, +) -> Result> { + let ItemChain { map_id, idx, .. } = *chain; + if step_ctx.abort_signal.aborted() { + bail!("aborted"); + } + let node = step_ctx + .graph + .get_node(current) + .ok_or_else(|| anyhow!("routed to unknown node '{current}'"))?; + let disallowed = match &node.node_type { + NodeType::Approval(_) => Some("an approval node"), + NodeType::Input(_) => Some("an input node"), + NodeType::End(_) => Some("an end node"), + NodeType::Map(_) => Some("a map node"), + NodeType::Agent(_) | NodeType::Llm(_) | NodeType::Rag(_) | NodeType::Script(_) => None, + }; + if let Some(type_phrase) = disallowed { + bail!( + "'{current}' is {type_phrase}; approval/input/end/map nodes cannot run inside a \ + map branch (enable settings.validate_before_run to catch this at load time)" + ); + } + + state.state_mut().visit_node(current); + let visits = state.state().loop_count(current); + let max_loops = step_ctx.graph.settings.max_loop_iterations; + if max_loops > 0 && visits > max_loops { + bail!("node '{current}' visited {visits} times (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()); + } + + let result = step(node, state, ctx, step_ctx, current).await; + ctx.peer_registry = None; + ctx.peer_assignment = None; + + match result? { + 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!( + "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!( + "node '{current}' fanned out to {many:?}; a map branch must route to a single node" + ), + }, + // The node-type pre-check above rejects End before step() runs. + StepResult::End(_) => bail!( + "internal error: step() returned End for '{current}' after the node-type pre-check" + ), + } +} + +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 is the author's bug and surfaces as an error +/// rather than being clamped. A literal 0 is rejected the same way, so a +/// graph loaded with validation off fails here instead of silently running +/// one item at a time, and a string that is not exactly one `{{key}}` is +/// rejected before interpolation for the same reason. Only the absent-cap default is clamped. +fn resolve_max_concurrency( + node: &MapNode, + state: &StateManager, + default: usize, + node_id: &str, +) -> Result { + let template = match &node.max_concurrency { + None => return Ok(default.max(1)), + Some(ConcurrencyCap::Fixed(0)) => { + bail!("map node '{node_id}': max_concurrency 0; expected a positive integer") + } + Some(ConcurrencyCap::Fixed(n)) => return bounded_cap(node_id, *n), + Some(ConcurrencyCap::Template(t)) => t, + }; + + if !is_lone_template(template) { + bail!( + "map node '{node_id}': max_concurrency \"{template}\" is a string but not a template; \ + write an integer or exactly one `{{{{key}}}}` with nothing around it" + ); + } + + 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::().ok(), + _ => None, + }; + + match resolved { + Some(n) if n >= 1 => bounded_cap(node_id, n), + _ => Err(anyhow!( + "map node '{node_id}': max_concurrency template \"{template}\" resolved to \ + {value} ({}); expected a positive integer", + type_name(&value) + )), + } +} + +/// `Semaphore::new` panics above `MAX_PERMITS`; a cap that large is a +/// configuration error, not a crash. +fn bounded_cap(node_id: &str, n: usize) -> Result { + if n > Semaphore::MAX_PERMITS { + bail!( + "map node '{node_id}': max_concurrency {n} exceeds the runtime limit of {}", + Semaphore::MAX_PERMITS + ); + } + Ok(n) +} + +/// 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. The identity is +/// named after the first `teammates: true` agent found by BFS from the entry, +/// so the pick follows the chain as a reader traces it, not YAML order. fn provision_map_peers( - branch: &NodeType, - branch_id: &str, + graph: &Graph, + subgraph: &HashSet, + entry: &str, item_count: usize, ) -> Option<(Arc, Vec)> { - match branch { - NodeType::Agent(n) if n.teammates && item_count >= 2 => { - let registry = Arc::new(PeerRegistry::new()); - let assignments = (0..item_count) - .map(|i| { - let id = graph_agent_id(&n.agent); - let inbox = Arc::new(Inbox::new()); - registry.insert(id.clone(), format!("{branch_id}[{i}]"), inbox.clone()); - (id, inbox) - }) - .collect(); - Some((registry, assignments)) - } - _ => None, + if item_count < 2 { + return None; } + let agent_name = first_flagged_agent_by_bfs(graph, subgraph, entry)?; + + 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)) +} + +fn first_flagged_agent_by_bfs<'g>( + graph: &'g Graph, + subgraph: &HashSet, + entry: &str, +) -> Option<&'g str> { + let mut visited: HashSet<&str> = HashSet::new(); + let mut queue: VecDeque<&str> = VecDeque::new(); + if subgraph.contains(entry) { + visited.insert(entry); + queue.push_back(entry); + } + while let Some(id) = queue.pop_front() { + let Some(node) = graph.get_node(id) else { + continue; + }; + if let NodeType::Agent(a) = &node.node_type + && a.teammates + { + return Some(a.agent.as_str()); + } + let mut edges: Vec<&String> = node + .next + .as_ref() + .map(|t| t.as_slice().iter().collect()) + .unwrap_or_default(); + match &node.node_type { + NodeType::Script(s) => edges.extend(s.fallback.as_ref()), + NodeType::Llm(l) => edges.extend(l.fallback.as_ref()), + _ => {} + } + for next in edges { + if subgraph.contains(next) && visited.insert(next) { + queue.push_back(next); + } + } + } + None } #[cfg(test)] mod tests { - use super::super::types::{AgentNode, EndNode}; + use super::super::types::{AgentNode, EndNode, GraphSettings, Node}; use super::*; + use indexmap::IndexMap; fn agent_branch(teammates: bool) -> NodeType { NodeType::Agent(AgentNode { @@ -198,14 +464,64 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates, }) } + fn graph_with_branch(node_type: NodeType) -> Graph { + let mut nodes: IndexMap = 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, Vec)> { + 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_map_peers(&agent_branch(true), "shards", 3).expect("should provision"); + let (registry, assignments) = provision(agent_branch(true), 3).expect("should provision"); assert_eq!(assignments.len(), 3); let roster = registry.roster(); @@ -222,12 +538,12 @@ mod tests { #[test] fn provision_map_peers_single_item_gets_no_registry() { - assert!(provision_map_peers(&agent_branch(true), "shards", 1).is_none()); + assert!(provision(agent_branch(true), 1).is_none()); } #[test] fn provision_map_peers_unflagged_gets_no_registry() { - assert!(provision_map_peers(&agent_branch(false), "shards", 3).is_none()); + assert!(provision(agent_branch(false), 3).is_none()); } #[test] @@ -236,6 +552,1892 @@ mod tests { output: String::new(), state_updates: None, }); - assert!(provision_map_peers(&branch, "shards", 3).is_none()); + assert!(provision(branch, 3).is_none()); + } + + #[tokio::test] + async fn map_item_peer_guard_retires_on_cancel() { + let (registry, assignments) = provision(agent_branch(true), 2).expect("should provision"); + let guard = PeerRetireGuard { + registry: Arc::clone(®istry), + id: assignments[0].0.clone(), + observer: None, + }; + let task = tokio::spawn(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + + task.abort(); + let err = task.await.expect_err("the item task was aborted"); + + assert!(err.is_cancelled(), "{err}"); + assert!( + registry.is_finished(&assignments[0].0), + "the cancelled item's identity is retired by the guard" + ); + assert!(!registry.is_finished(&assignments[1].0)); + } + + fn map_with_cap(cap: Option) -> 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, state: &StateManager) -> Result { + 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_fixed_zero_errors() { + let state = StateManager::new(HashMap::new()); + let msg = resolve(Some(ConcurrencyCap::Fixed(0)), &state) + .unwrap_err() + .to_string(); + assert_eq!( + msg, + "map node 'm': max_concurrency 0; expected a positive integer" + ); + } + + #[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}"); + } + + #[test] + fn resolve_max_concurrency_rejects_non_template_string() { + let state = state_with("budget", Value::from(3)); + let cap = Some(ConcurrencyCap::Template("4".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!( + msg.contains( + "is a string but not a template; write an integer or exactly one `{{key}}` \ + with nothing around it" + ), + "{msg}" + ); + assert!(msg.contains("map node 'm': max_concurrency \"4\""), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_fixed_above_semaphore_limit_errors() { + let state = StateManager::new(HashMap::new()); + let cap = Some(ConcurrencyCap::Fixed(Semaphore::MAX_PERMITS + 1)); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert_eq!( + msg, + format!( + "map node 'm': max_concurrency {} exceeds the runtime limit of {}", + Semaphore::MAX_PERMITS + 1, + Semaphore::MAX_PERMITS + ) + ); + // The limit itself is fine. + let cap = Some(ConcurrencyCap::Fixed(Semaphore::MAX_PERMITS)); + assert_eq!(resolve(cap, &state).unwrap(), Semaphore::MAX_PERMITS); + } + + #[test] + fn resolve_max_concurrency_template_above_semaphore_limit_errors() { + let state = state_with("budget", Value::from((Semaphore::MAX_PERMITS as u64) + 1)); + let cap = Some(ConcurrencyCap::Template("{{budget}}".into())); + let msg = resolve(cap, &state).unwrap_err().to_string(); + assert!(msg.contains("exceeds the runtime limit of"), "{msg}"); + } + + #[test] + fn resolve_max_concurrency_accepts_padded_template() { + let state = state_with("budget", Value::from(3)); + let cap = Some(ConcurrencyCap::Template(" {{budget}} ".into())); + assert_eq!(resolve(cap, &state).unwrap(), 3); + } +} + +#[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 { + 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 { + serde_json::from_str::(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 { + match result { + Ok(out) => panic!("expected failure, got output: {out}"), + Err(e) => format!("{e:#}"), + } + } + + /// Three item chains hold their scripts open; the graph abort lands + /// mid-fan-out. The map's frontier task is cancelled, which must cancel + /// every item task with it so the shells are killed before `&&` runs. + #[tokio::test] + async fn map_abort_cancels_item_chains() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + ws.write_script( + "hold.sh", + &format!( + "#!/bin/bash\n\ + t=${{GRAPH_STATE#*'\"item\":'}}\n\ + n=${{t%%[!0-9]*}}\n\ + sleep 1.5 && touch '{}/sentinel-'\"$n\"\n\ + echo '{{\"output\": 1}}'\n", + ws.dir.display() + ), + ); + + let yaml = r#" +name: map_abort_test +start: fan_out +settings: + validate_before_run: false +initial_state: + items: [1, 2, 3] +nodes: + fan_out: + type: map + over: "{{items}}" + as: item + branch: hold + collect_into: results + max_concurrency: 3 + next: done + hold: + type: script + script: hold.sh + done: + type: end + output: "{{results}}" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let trigger = { + let abort = abort.clone(); + tokio::spawn(async move { + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + abort.set_ctrlc(); + }) + }; + + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await; + trigger.await.unwrap(); + + let chain = error_chain(result); + assert!(chain.contains("aborted"), "{chain}"); + + tokio::time::sleep(std::time::Duration::from_millis(2500)).await; + for n in 1..=3 { + let sentinel = ws.dir.join(format!("sentinel-{n}")); + assert!( + !sentinel.exists(), + "item chain {n} kept running after the abort cancelled the map task" + ); + } + } + + /// With validation off, the collision is still caught before any item + /// forks: the branch never runs. + #[tokio::test] + async fn map_as_equal_output_key_fails_at_runtime() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + let sentinel = ws.dir.join("ran"); + ws.write_script( + "mark.sh", + &format!("#!/bin/bash\ntouch '{}'\necho '{{}}'\n", sentinel.display()), + ); + + let yaml = r#" +name: as_output_collision_test +start: fan_out +settings: + validate_before_run: false +initial_state: + items: [1, 2] +nodes: + fan_out: + type: map + over: "{{items}}" + as: output + branch: mark + collect_into: results + next: done + mark: + type: script + script: mark.sh + done: + type: end + output: "{{results}}" +"#; + let chain = error_chain(run_graph(yaml, &ws).await); + + assert!( + chain.contains( + "map node 'fan_out': `as` and `output_key` are both 'output'; the item binding \ + would be collected as the chain's result" + ), + "{chain}" + ); + assert!(!sentinel.exists(), "no item chain may start"); + } + + #[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)]); + } + + const RETRY_CHAIN_GRAPH: &str = r#" +name: retry-chain +start: fan_out +settings: + validate_before_run: true +initial_state: + items: ["a", "b", "c"] +nodes: + fan_out: + type: map + over: "{{items}}" + as: item + branch: draft + output_key: result + collect_into: results + next: done + draft: + type: script + script: draft.py + next: check + check: + type: script + script: check.py + done: + type: end + output: "{{results}}" +"#; + + fn write_retry_chain_scripts(ws: &TestWorkspace) { + ws.write_py( + "draft.py", + r##"attempts = state.get("attempts", 0) + 1 +print(json.dumps({"attempts": attempts, "draft": f"{state['item']}#{attempts}"}))"##, + ); + ws.write_py( + "check.py", + r#"if state["attempts"] < 2: + print(json.dumps({"_next": "draft"})) +else: + print(json.dumps({"result": {"item": state["item"], "draft": state["draft"], "attempts": state["attempts"]}}))"#, + ); + } + + fn retry_chain_results() -> Vec { + ["a", "b", "c"] + .iter() + .map(|item| json!({"item": item, "draft": format!("{item}#2"), "attempts": 2})) + .collect() + } + + #[tokio::test] + async fn map_chain_retries_via_script_next_until_check_accepts() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let ws = TestWorkspace::new(); + write_retry_chain_scripts(&ws); + + let result = run_graph(RETRY_CHAIN_GRAPH, &ws) + .await + .unwrap_or_else(|e| panic!("validated retry chain failed: {e:#}")); + + assert_eq!(collected(&result), retry_chain_results()); + } + + #[tokio::test] + async fn map_chain_retries_leave_parent_state_and_loop_counts_untouched() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let ws = TestWorkspace::new(); + write_retry_chain_scripts(&ws); + let graph: Arc = Arc::new(serde_yaml::from_str(RETRY_CHAIN_GRAPH).unwrap()); + let NodeType::Map(map_node) = &graph.get_node("fan_out").unwrap().node_type else { + panic!("fan_out should be a map node"); + }; + let script = ScriptExecutor::new(&ws.dir); + let abort = create_abort_signal(); + let step_ctx = StepContext { + graph: Arc::clone(&graph), + script_executor: &script, + max_concurrency: 4, + abort_signal: &abort, + branch_mode: false, + }; + let items = json!(["a", "b", "c"]); + let mut parent = StateManager::new(HashMap::from([("items".to_string(), items.clone())])); + let mut ctx = silent_ctx(); + + MapNodeExecutor::execute(map_node, &mut parent, &mut ctx, &step_ctx, "fan_out") + .await + .unwrap_or_else(|e| panic!("map failed: {e:#}")); + + // Only `collect_into` lands on the parent: the per-item `attempts`, + // `draft` and `result` writes were fork-local scratch. + assert_eq!( + *parent.state().data(), + HashMap::from([ + ("items".to_string(), items), + ("results".to_string(), Value::Array(retry_chain_results())), + ]) + ); + for node in ["fan_out", "draft", "check"] { + assert_eq!( + parent.state().loop_count(node), + 0, + "parent loop count for '{node}'" + ); + } + } + + #[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 (max_loop_iterations=3)"), + "{chain}" + ); + assert_eq!(chain.matches("sub-branch [").count(), 1, "{chain}"); + } + + #[tokio::test] + async fn map_chain_script_next_loop_with_zero_cap_runs_to_completion() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let ws = TestWorkspace::new(); + // Re-enter `looper` until n reaches 5; omitting `_next` ends the item. + ws.write_py( + "looper.py", + r#"n = state.get("n", 0) + 1 +out = {"n": n, "output": n} +if n < 5: + out["_next"] = "looper" +print(json.dumps(out))"#, + ); + + let yaml = r#" +name: chain +start: fan_out +settings: + max_loop_iterations: 0 + 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 result = run_graph(yaml, &ws).await; + + assert_eq!(collected(&result.unwrap()), vec![json!(5)]); + } + + #[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("map node 'fan_out': sub-branch [0] failed at node 'gate' (step 1)"), + "{chain}" + ); + assert_eq!(chain.matches("sub-branch [").count(), 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( + "'ask' is an approval node; approval/input/end/map nodes cannot run inside a \ + map branch (enable settings.validate_before_run to catch this at load time)" + ), + "{chain}" + ); + assert!(chain.contains("failed at node 'ask' (step 2)"), "{chain}"); + } + + #[tokio::test] + async fn map_chain_pre_check_rejects_end_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: finish + finish: + type: end + output: "done early" + done: + type: end + output: "{{results}}" +"#; + let chain = error_chain(run_graph(yaml, &ws).await); + + assert!( + chain.contains( + "'finish' is an end node; approval/input/end/map nodes cannot run inside a \ + map branch (enable settings.validate_before_run to catch this at load time)" + ), + "{chain}" + ); + assert!( + chain.contains("map node 'fan_out': sub-branch [0] failed at node 'finish' (step 2)"), + "{chain}" + ); + assert_eq!(chain.matches("sub-branch [").count(), 1, "{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, + 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 { + branch_subgraph(&self.graph, entry) + } + + fn provision(&self, entry: &str, n: usize) -> (Arc, Vec) { + 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, 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, 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, + 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 ECHO_INPUTS_AGENT: &str = "echo-inputs"; + + /// Materializes a graph agent whose only script echoes the `width` it + /// finds in `GRAPH_STATE` back as `output`. Its YAML default is + /// `width: 1`, so a parent that seeds `width` through `inputs:` is + /// observable against a parent that does not. + fn materialize_echo_inputs_agent() { + let graph_path = paths::agent_graph_file(ECHO_INPUTS_AGENT); + let agent_dir = graph_path.parent().unwrap(); + fs::create_dir_all(agent_dir).unwrap(); + fs::write( + agent_dir.join("echo.py"), + "#!/usr/bin/env python3\nimport os, json\n\ + state = json.loads(os.environ.get(\"GRAPH_STATE\", \"{}\"))\n\ + print(json.dumps({\"output\": str(state[\"width\"])}))\n", + ) + .unwrap(); + fs::write( + &graph_path, + format!( + r#" +name: {ECHO_INPUTS_AGENT} +start: echo +initial_state: + width: 1 +nodes: + echo: + type: script + script: echo.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("failed at node 'doubler' (step 1): aborted"), + "{chain}" + ); + assert_eq!(chain.matches("sub-branch [").count(), 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) -> 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()); + } + + /// An agent step whose `state_updates` name only other keys leaves the + /// default `output_key` unwritten, so the collector reports the missing + /// key instead of collecting a manufactured `null`. + #[tokio::test] + #[serial] + async fn map_branch_with_unrelated_state_updates_errors_instead_of_null() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let ws = TestWorkspace::new(); + let yaml = format!( + r#" +name: t +settings: + validate_before_run: false +initial_state: + items: [1] +start: fan_out +nodes: + fan_out: + type: map + over: "{{{{items}}}}" + as: item + branch: worker + collect_into: results + next: done + worker: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + state_updates: + note: "{{{{output}}}}" + 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}" + ); + } + + /// The materialized agent's graph holds its script step open for seconds, + /// so the 1s timer fires while that step is still pending: the agent + /// 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(1))); + 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 1s")), + "{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()); + } + + /// Positive twin of the timeout test above: with `timeout: 0` the agent + /// step has no wall-clock bound, so a hold longer than any accidental 0s + /// or 1s bound still runs to completion and hands its output back. + #[tokio::test] + #[serial] + async fn run_item_chain_probe_agent_with_zero_timeout_runs_to_completion() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(1.5); + let h = Harness::new(&flagged_probe_graph(Some(0))); + 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::test] + #[serial] + async fn run_item_chain_clears_peer_fields_after_flagged_agent_then_script() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(0.0); + let h = Harness::new(&format!( + r#" +name: t +start: worker +nodes: + worker: + type: agent + agent: {PROBE_AGENT} + prompt: "p" + teammates: true + state_updates: + output: "{{{{output}}}}" + next: finish + finish: + type: script + script: finish.py +"# + )); + h.ws.write_py( + "finish.py", + r#"print(json.dumps({"output": state["output"] + "+finished"}))"#, + ); + 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+finished"))); + assert!(peers.0.is_finished(&peers.1.0)); + assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); + } + + /// The prompt fails to interpolate before `run_agent_for_graph` can take + /// the peer fields, so only the chain runner's reset can clear them. + #[tokio::test] + async fn run_item_chain_clears_peer_fields_when_agent_step_fails_before_takeover() { + let h = Harness::new( + r#" +name: t +start: worker +nodes: + worker: + type: agent + agent: no-such-agent + prompt: "{{nope}}" + teammates: true +"#, + ); + 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("Failed to interpolate prompt for agent 'no-such-agent'"), + "{chain}" + ); + assert!(chain.contains("'nope' not found in state"), "{chain}"); + assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); + assert!(peers.0.is_finished(&peers.1.0)); + } + + fn echo_inputs_parent_graph(inputs_block: &str) -> String { + format!( + r#" +name: t +start: worker +nodes: + worker: + type: agent + agent: {ECHO_INPUTS_AGENT} + prompt: "p" +{inputs_block} state_updates: + output: "{{{{output}}}}" +"# + ) + } + + async fn run_echo_inputs_chain(parent_yaml: &str) -> StateManager { + let h = Harness::new(parent_yaml); + let mut state = StateManager::new(HashMap::from([("n".to_string(), json!(3))])); + let mut ctx = silent_ctx(); + + h.run("worker", None, &mut state, &mut ctx) + .await + .unwrap_or_else(|e| panic!("chain failed: {e:#}")); + + state + } + + /// The child's `width` default is overlaid by the parent's `inputs` + /// before the child graph starts, so its script sees the parent's value. + #[tokio::test] + #[serial] + async fn agent_node_inputs_seed_the_child_graph_state() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_echo_inputs_agent(); + let parent = echo_inputs_parent_graph(" inputs:\n width: \"{{n}}\"\n"); + + let state = run_echo_inputs_chain(&parent).await; + + assert_eq!(state.state().get("output"), Some(&json!("3"))); + } + + /// Control for the test above: the same child with no `inputs` on the + /// parent node keeps its YAML default, so an override in the sibling + /// test cannot be the child ignoring its own `initial_state`. + #[tokio::test] + #[serial] + async fn agent_node_without_inputs_leaves_the_child_graph_defaults() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_echo_inputs_agent(); + let parent = echo_inputs_parent_graph(""); + + let state = run_echo_inputs_chain(&parent).await; + + assert_eq!(state.state().get("output"), Some(&json!("1"))); + } + + #[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_picks_flagged_agent_nearest_to_entry_not_yaml_order() { + let h = Harness::new( + r#" +name: t +start: gate +nodes: + far: + type: agent + agent: far-agent + prompt: "p" + teammates: true + gate: + type: script + script: gate.py + next: near + near: + type: agent + agent: near-agent + prompt: "p" + teammates: true + next: far +"#, + ); + + let (_, assignments) = h.provision("gate", 2); + + for (id, _) in &assignments { + assert!(id.starts_with("graph_agent_near-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()); } } diff --git a/src/graph/mod.rs b/src/graph/mod.rs index b5355d8..3b1d72d 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -10,21 +10,30 @@ pub mod reducer; pub mod script; pub mod staging; pub mod state; +pub mod state_updates; pub mod structured; pub mod types; pub mod user_interaction; pub mod validator; -pub use dispatch::{active_agent_graph_name, run_active_agent_graph}; +pub use dispatch::{ + active_agent_graph_name, run_active_agent_graph, run_active_agent_graph_with_inputs, +}; pub use executor::GraphExecutor; pub use parser::{GraphParser, agent_has_graph}; use serde_json::Value; +use std::time::Duration; pub use types::{Graph, NodeType}; pub const GRAPH_SCHEMA_VERSION: &str = "1.0"; pub const DEFAULT_MAX_LOOP_ITERATIONS: usize = 100; +/// A timeout of `0` disables the wall-clock bound. +pub(crate) fn wall_clock(secs: u64) -> Option { + (secs != 0).then(|| Duration::from_secs(secs)) +} + pub const MAX_STATE_SIZE_BYTES: usize = 32 * 1024; pub(in crate::graph) fn type_name(value: &Value) -> &'static str { @@ -37,3 +46,18 @@ pub(in crate::graph) fn type_name(value: &Value) -> &'static str { Value::Object(_) => "object", } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wall_clock_zero_is_no_bound() { + assert!(wall_clock(0).is_none()); + } + + #[test] + fn wall_clock_nonzero_is_that_many_seconds() { + assert_eq!(wall_clock(7), Some(Duration::from_secs(7))); + } +} diff --git a/src/graph/rag.rs b/src/graph/rag.rs index 2aed523..fbba2aa 100644 --- a/src/graph/rag.rs +++ b/src/graph/rag.rs @@ -1,13 +1,13 @@ use super::state::StateManager; +use super::state_updates; use super::types::RagNode; +use super::wall_clock; use crate::config::RequestContext; use crate::utils::create_abort_signal; use anyhow::{Context, Result, anyhow}; use serde_json::{Map, Value}; -use std::time::Duration; use tokio::time::timeout; -const OUTPUT_KEY: &str = "output"; const DEFAULT_QUERY: &str = "{{initial_prompt}}"; const DEFAULT_RAG_TIMEOUT_SECS: u64 = 120; @@ -34,18 +34,16 @@ impl RagNodeExecutor { let top_k = node.top_k.unwrap_or_else(|| rag.configured_top_k()); let rerank = rag.configured_reranker(); - let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS)); + let secs = node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS); + let bound = wall_clock(secs); let abort = create_abort_signal(); - let (context, sources_str, _ids) = - timeout(timeout_dur, rag.search(&query, top_k, rerank, abort)) - .await - .with_context(|| { - format!( - "rag node '{node_id}' timed out after {}s", - timeout_dur.as_secs() - ) - })? - .with_context(|| format!("rag node '{node_id}' retrieval failed"))?; + let fut = rag.search(&query, top_k, rerank, abort); + let (context, sources_str, _ids) = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + } + .with_context(|| format!("rag node '{node_id}' timed out after {secs}s"))? + .with_context(|| format!("rag node '{node_id}' retrieval failed"))?; let output = build_rag_output(context, &sources_str); apply_state_updates(node, state_manager, &output); @@ -70,33 +68,14 @@ fn build_rag_output(context: String, sources_str: &str) -> Value { } fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: &Value) { - let Some(updates) = &node.state_updates else { - return; - }; - let prev_output = state_manager.state().get(OUTPUT_KEY).cloned(); - state_manager - .state_mut() - .set(OUTPUT_KEY.into(), output.clone()); - - for (key, template) in updates { - let value = state_manager.interpolate_lenient(template); - state_manager - .state_mut() - .set(key.clone(), Value::String(value)); - } - - match prev_output { - Some(v) => state_manager.state_mut().set(OUTPUT_KEY.into(), v), - None => state_manager - .state_mut() - .set(OUTPUT_KEY.into(), Value::Null), - } + state_updates::apply(state_manager, output, false, node.state_updates.as_ref()); } #[cfg(test)] mod tests { use super::*; use serde_json::json; + use std::time::Duration; #[test] fn build_rag_output_splits_bullet_sources_into_array() { @@ -126,4 +105,15 @@ mod tests { assert_eq!(out["sources"], json!(["plain/path"])); } + + /// Mirrors the resolution expression in `execute`. + fn resolve(timeout: Option) -> Option { + wall_clock(timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS)) + } + + #[test] + fn zero_timeout_resolves_to_no_bound() { + assert!(resolve(Some(0)).is_none()); + assert_eq!(resolve(None), Some(Duration::from_secs(120))); + } } diff --git a/src/graph/script.rs b/src/graph/script.rs index 688d6ad..a2d7363 100644 --- a/src/graph/script.rs +++ b/src/graph/script.rs @@ -1,5 +1,6 @@ use super::state::{StateManager, StateRepresentation}; use super::types::ScriptNode; +use super::wall_clock; use crate::config::paths; use crate::function::Language; use anyhow::{Context, Result, anyhow, bail}; @@ -8,7 +9,6 @@ use std::collections::HashMap; use std::env; use std::path::{Path, PathBuf}; use std::process::Stdio; -use std::time::Duration; use tokio::process::Command; use tokio::time::timeout; @@ -66,22 +66,25 @@ impl ScriptExecutor { } } - let timeout_dur = Duration::from_secs(node.timeout); - let output = timeout(timeout_dur, cmd.output()) - .await - .with_context(|| { - format!( - "Script '{}' timed out after {}s", - script_path.display(), - node.timeout - ) - })? - .with_context(|| { - format!( - "Failed to spawn script process for '{}'", - script_path.display() - ) - })?; + let bound = wall_clock(node.timeout); + let fut = cmd.output(); + let output = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + } + .with_context(|| { + format!( + "Script '{}' timed out after {}s", + script_path.display(), + node.timeout + ) + })? + .with_context(|| { + format!( + "Failed to spawn script process for '{}'", + script_path.display() + ) + })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); @@ -260,6 +263,33 @@ echo '{"quality": 0.85, "issues": 3, "_next": "approve"}' cleanup(&dir); } + #[tokio::test] + async fn zero_timeout_is_unbounded() { + if !cmd_available("bash") { + return; + } + let (dir, path) = write_script( + r#"#!/bin/bash +sleep 1.5 +echo '{"ok":true}' +"#, + "sh", + ); + let mut state = StateManager::new(HashMap::new()); + let executor = ScriptExecutor::new(&dir); + + executor + .execute( + &node_for(path.file_name().unwrap().to_str().unwrap(), 0), + &mut state, + ) + .await + .unwrap_or_else(|e| panic!("zero timeout should not bound the script: {e:#}")); + + assert_eq!(state.state().get("ok"), Some(&json!(true))); + cleanup(&dir); + } + #[tokio::test] async fn bash_script_can_read_state_from_env() { if !cmd_available("bash") { diff --git a/src/graph/state.rs b/src/graph/state.rs index e16f16f..6a1c21d 100644 --- a/src/graph/state.rs +++ b/src/graph/state.rs @@ -320,6 +320,26 @@ pub(super) fn template_root_keys(template: &str) -> Vec { .collect() } +// Whether `s` is exactly one `{{key}}` reference and nothing else (after the +// same trim `interpolate_raw` applies) whose path parses with the grammar +// `get_nested_value` uses, so this accepts exactly the strings the runtime +// can resolve as a pure reference. +pub(super) fn is_lone_template(s: &str) -> bool { + let s = s.trim(); + let Ok(Some(m)) = TEMPLATE_VAR_RE.captures(s) else { + return false; + }; + let whole = m.get(0).expect("group 0 always present"); + if whole.start() != 0 || whole.end() != s.len() { + return false; + } + let path = m.get(1).expect("template regex has one group").as_str(); + let mut parts = path.split('.'); + let first = parts.next().and_then(split_indices); + first.is_some_and(|(name, _)| !name.is_empty()) + && parts.all(|part| split_indices(part).is_some()) +} + fn value_to_string(value: &Value) -> String { match value { Value::String(s) => s.clone(), @@ -931,6 +951,60 @@ mod tests { assert_eq!(result, json!("v")); } + #[test] + fn is_lone_template_cases() { + for accepted in [ + "{{budget}}", + "{{cfg.limits[0]}}", + " {{budget}} ", + "\n{{k}}\t", + ] { + assert!( + is_lone_template(accepted), + "{accepted:?} should be a lone template" + ); + } + for rejected in [ + "", + "4", + "{{}}", + "{{ key }}", + "n={{k}}", + "{{k}}x", + "{{a}} {{b", + "{{a}}{{b}}", + "{{a-b}}", + "{{k", + ] { + assert!( + !is_lone_template(rejected), + "{rejected:?} must not be a lone template" + ); + } + } + + #[test] + fn is_lone_template_matches_runtime_path_grammar() { + for accepted in [ + "{{a.b[0][1].c}}", + "{{a[0][1]}}", + "{{a.}}", + "{{a..b}}", + "{{a.[0]}}", + ] { + assert!( + is_lone_template(accepted), + "{accepted:?} resolves at run time and must be accepted" + ); + } + for rejected in ["{{limits[foo]}}", "{{a[]}}", "{{a[0}}", "{{.a}}"] { + assert!( + !is_lone_template(rejected), + "{rejected:?} never resolves at run time and must be rejected" + ); + } + } + #[test] fn interpolate_raw_inner_spaces_treated_as_mixed() { let manager = manager_with(&[("k", json!("v"))]); diff --git a/src/graph/state_updates.rs b/src/graph/state_updates.rs new file mode 100644 index 0000000..a4ea300 --- /dev/null +++ b/src/graph/state_updates.rs @@ -0,0 +1,168 @@ +use super::state::StateManager; +use serde_json::Value; +use std::collections::HashMap; + +pub(super) const OUTPUT_KEY: &str = "output"; + +/// Merge `output_schema` top-level keys (when `has_schema`) and apply `state_updates` with +/// `{{output}}` bound to `output` for the duration of interpolation only. An explicit `output` +/// key in `updates` is honored; otherwise `output` is restored to its prior value, or removed +/// again if it was absent. +pub(super) fn apply( + state_manager: &mut StateManager, + output: &Value, + has_schema: bool, + updates: Option<&HashMap>, +) { + if has_schema && let Some(obj) = output.as_object() { + for (k, v) in obj { + state_manager.state_mut().set(k.clone(), v.clone()); + } + } + + let Some(updates) = updates else { + return; + }; + let prev_output = state_manager.state().get(OUTPUT_KEY).cloned(); + state_manager + .state_mut() + .set(OUTPUT_KEY.into(), output.clone()); + + let computed: Vec<(String, Value)> = updates + .iter() + .map(|(key, template)| { + ( + key.clone(), + Value::String(state_manager.interpolate_lenient(template)), + ) + }) + .collect(); + + match prev_output { + Some(prev) => state_manager.state_mut().set(OUTPUT_KEY.into(), prev), + None => { + state_manager.state_mut().remove(OUTPUT_KEY); + } + } + + for (key, value) in computed { + state_manager.state_mut().set(key, value); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn manager_with(pairs: &[(&str, Value)]) -> StateManager { + let mut map = HashMap::new(); + for (k, v) in pairs { + map.insert((*k).into(), v.clone()); + } + StateManager::new(map) + } + + fn updates(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect() + } + + #[test] + fn explicit_output_state_update_is_honored() { + let u = updates(&[("output", "{{output}}")]); + let mut state = manager_with(&[("output", json!("stale"))]); + + apply(&mut state, &json!("the completion"), false, Some(&u)); + + assert_eq!( + state.state().get(OUTPUT_KEY), + Some(&json!("the completion")) + ); + } + + #[test] + fn explicit_output_is_honored_alongside_other_updates() { + let u = updates(&[("output", "{{output}}"), ("summary", "got: {{output}}")]); + let mut state = manager_with(&[]); + + apply(&mut state, &json!("done"), false, Some(&u)); + + assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("done"))); + assert_eq!(state.state().get("summary"), Some(&json!("got: done"))); + } + + #[test] + fn no_updates_is_a_noop() { + let mut state = manager_with(&[("keep", json!(1))]); + + apply(&mut state, &json!("ignored"), false, None); + + assert_eq!(state.state().get("keep"), Some(&json!(1))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + + #[test] + fn output_is_restored_to_previous_value() { + let u = updates(&[("response", "{{output}}")]); + let mut state = manager_with(&[("output", json!("preserved"))]); + + apply(&mut state, &json!("new"), false, Some(&u)); + + assert_eq!(state.state().get("response"), Some(&json!("new"))); + assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!("preserved"))); + } + + #[test] + fn output_is_removed_when_absent() { + let u = updates(&[("response", "{{output}}")]); + let mut state = manager_with(&[]); + + apply(&mut state, &json!("new"), false, Some(&u)); + + assert_eq!(state.state().get("response"), Some(&json!("new"))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + + #[test] + fn explicit_null_prior_output_is_restored_as_null() { + let u = updates(&[("response", "{{output}}")]); + let mut state = manager_with(&[("output", json!(null))]); + + apply(&mut state, &json!("new"), false, Some(&u)); + + assert_eq!(state.state().get("response"), Some(&json!("new"))); + assert_eq!(state.state().get(OUTPUT_KEY), Some(&json!(null))); + } + + #[test] + fn schema_merges_top_level_keys_when_has_schema() { + let mut state = manager_with(&[]); + + apply(&mut state, &json!({"summary": "s", "score": 7}), true, None); + + assert_eq!(state.state().get("summary"), Some(&json!("s"))); + assert_eq!(state.state().get("score"), Some(&json!(7))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } + + #[test] + fn object_output_is_not_merged_without_schema() { + let u = updates(&[("ctx", "{{output.context}}")]); + let mut state = manager_with(&[]); + + apply( + &mut state, + &json!({"context": "c", "sources": ["a"]}), + false, + Some(&u), + ); + + assert!(state.state().get("context").is_none()); + assert!(state.state().get("sources").is_none()); + assert_eq!(state.state().get("ctx"), Some(&json!("c"))); + assert!(state.state().get(OUTPUT_KEY).is_none()); + } +} diff --git a/src/graph/types.rs b/src/graph/types.rs index 7af4483..6e00d04 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -240,6 +240,9 @@ pub struct AgentNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub timeout: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inputs: Option>, + #[serde(default)] pub teammates: bool, } @@ -426,13 +429,60 @@ pub struct MapNode { pub collect_into: String, #[serde(default, skip_serializing_if = "Option::is_none")] - pub max_concurrency: Option, + pub max_concurrency: Option, } fn default_map_output_key() -> String { "output".to_string() } +/// A map node's per-item concurrency cap: either a literal integer or a +/// `{{template}}` resolved against graph state when the map runs. +#[derive(Debug, Clone, Serialize, PartialEq)] +#[serde(untagged)] +pub enum ConcurrencyCap { + Fixed(usize), + Template(String), +} + +// Hand-written so a malformed value (`-1`, `4.5`, `true`) reports what was +// received and what is accepted, instead of serde's untagged-enum "did not +// match any variant". +impl<'de> Deserialize<'de> for ConcurrencyCap { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + struct CapVisitor; + + impl serde::de::Visitor<'_> for CapVisitor { + type Value = ConcurrencyCap; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a positive integer or a `{{template}}` string for max_concurrency") + } + + fn visit_u64(self, v: u64) -> Result { + usize::try_from(v) + .map(ConcurrencyCap::Fixed) + .map_err(|_| E::invalid_value(serde::de::Unexpected::Unsigned(v), &self)) + } + + fn visit_i64(self, v: i64) -> Result { + usize::try_from(v) + .map(ConcurrencyCap::Fixed) + .map_err(|_| E::invalid_value(serde::de::Unexpected::Signed(v), &self)) + } + + fn visit_str(self, v: &str) -> Result { + Ok(ConcurrencyCap::Template(v.to_owned())) + } + } + + deserializer.deserialize_any(CapVisitor) + } +} + #[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum Reducer { @@ -449,6 +499,7 @@ pub enum Reducer { #[derive(Debug, Clone, Default)] pub struct GraphState { data: HashMap, + #[cfg(test)] history: Vec, loop_counts: HashMap, } @@ -457,6 +508,7 @@ impl GraphState { pub fn new(initial: HashMap) -> Self { Self { data: initial, + #[cfg(test)] history: Vec::new(), loop_counts: HashMap::new(), } @@ -470,6 +522,10 @@ impl GraphState { self.data.insert(key, value); } + pub fn remove(&mut self, key: &str) -> Option { + self.data.remove(key) + } + pub fn merge(&mut self, json_obj: &serde_json::Map) { for (key, value) in json_obj { self.data.insert(key.clone(), value.clone()); @@ -481,6 +537,7 @@ impl GraphState { } pub fn visit_node(&mut self, node_id: &str) { + #[cfg(test)] self.history.push(node_id.to_string()); *self.loop_counts.entry(node_id.to_string()).or_insert(0) += 1; } @@ -583,6 +640,48 @@ nodes: } } + #[test] + fn agent_node_inputs_parse_optionally_and_empty_map_is_some() { + let yaml = r#" +name: g +start: with_inputs +nodes: + with_inputs: + id: with_inputs + type: agent + agent: helper + prompt: hi + inputs: + width: "{{n}}" + note: "Repo: {{repo}}" + next: without_inputs + without_inputs: + id: without_inputs + type: agent + agent: helper + prompt: hi + next: empty_inputs + empty_inputs: + id: empty_inputs + type: agent + agent: helper + prompt: hi + inputs: {} +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let inputs_of = |id: &str| match &graph.get_node(id).unwrap().node_type { + NodeType::Agent(n) => n.inputs.clone(), + _ => panic!("expected Agent variant"), + }; + + let with = inputs_of("with_inputs").expect("inputs should parse"); + assert_eq!(with.len(), 2); + assert_eq!(with.get("width").map(String::as_str), Some("{{n}}")); + assert_eq!(with.get("note").map(String::as_str), Some("Repo: {{repo}}")); + assert_eq!(inputs_of("without_inputs"), None); + assert_eq!(inputs_of("empty_inputs"), Some(HashMap::new())); + } + #[test] fn deserializes_mcp_tools_at_graph_and_node_level() { let yaml = r#" @@ -1232,7 +1331,7 @@ next: rank assert_eq!(map.branch, "research_subject"); assert_eq!(map.output_key, "research_result"); assert_eq!(map.collect_into, "research_results"); - assert_eq!(map.max_concurrency, Some(5)); + assert_eq!(map.max_concurrency, Some(ConcurrencyCap::Fixed(5))); } #[test] @@ -1256,6 +1355,78 @@ collect_into: results assert!(map.max_concurrency.is_none()); } + fn map_node_yaml_with_cap(cap: &str) -> String { + format!( + "id: fan_out\ntype: map\nover: \"{{{{items}}}}\"\nas: item\nbranch: process\n\ + collect_into: results\nmax_concurrency: {cap}\n" + ) + } + + fn parse_map_cap(cap: &str) -> Option { + let node: Node = serde_yaml::from_str(&map_node_yaml_with_cap(cap)).unwrap(); + match node.node_type { + NodeType::Map(m) => m.max_concurrency, + _ => panic!("expected Map variant"), + } + } + + #[test] + fn map_max_concurrency_integer_parses_as_fixed() { + assert_eq!(parse_map_cap("4"), Some(ConcurrencyCap::Fixed(4))); + } + + #[test] + fn map_max_concurrency_zero_parses_as_fixed_zero() { + assert_eq!(parse_map_cap("0"), Some(ConcurrencyCap::Fixed(0))); + } + + #[test] + fn map_max_concurrency_string_parses_as_template() { + assert_eq!( + parse_map_cap("\"{{budget}}\""), + Some(ConcurrencyCap::Template("{{budget}}".into())) + ); + } + + #[test] + fn map_max_concurrency_quoted_integer_parses_as_template() { + assert_eq!( + parse_map_cap("\"4\""), + Some(ConcurrencyCap::Template("4".into())) + ); + } + + #[test] + fn map_max_concurrency_malformed_values_name_accepted_forms() { + for bad in ["-1", "4.5", "true", "[4]", "{n: 4}"] { + let err = serde_yaml::from_str::(&map_node_yaml_with_cap(bad)) + .expect_err(&format!("max_concurrency: {bad} should fail to parse")); + let msg = err.to_string(); + assert!( + msg.contains("a positive integer or a `{{template}}` string for max_concurrency"), + "max_concurrency: {bad} produced an unhelpful error: {msg}" + ); + } + } + + #[test] + fn map_max_concurrency_negative_error_names_offending_value() { + let err = serde_yaml::from_str::(&map_node_yaml_with_cap("-1")).unwrap_err(); + assert!(err.to_string().contains("-1"), "{err}"); + } + + #[test] + fn concurrency_cap_serializes_as_bare_scalar() { + assert_eq!( + serde_json::to_value(ConcurrencyCap::Fixed(4)).unwrap(), + serde_json::json!(4) + ); + assert_eq!( + serde_json::to_value(ConcurrencyCap::Template("{{k}}".into())).unwrap(), + serde_json::json!("{{k}}") + ); + } + #[test] fn full_graph_with_all_new_phase_a_fields_parses() { let yaml = r#" @@ -1386,4 +1557,20 @@ nodes: assert!(!serde_yaml::to_string(&graph).unwrap().contains("driver")); } + + #[test] + fn graph_state_remove_returns_previous_value_and_drops_key() { + let mut state = GraphState::new(HashMap::from([("k".to_string(), json!(1))])); + + assert_eq!(state.remove("k"), Some(json!(1))); + assert_eq!(state.get("k"), None); + assert!(state.data().is_empty()); + } + + #[test] + fn graph_state_remove_missing_key_is_none() { + let mut state = GraphState::new(HashMap::new()); + + assert_eq!(state.remove("k"), None); + } } diff --git a/src/graph/validator.rs b/src/graph/validator.rs index b900e14..cf61276 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -1,5 +1,5 @@ -use super::state::template_root_keys; -use super::types::{Graph, Node, NodeType}; +use super::state::{is_lone_template, template_root_keys}; +use super::types::{ConcurrencyCap, Graph, NextTargets, Node, NodeType}; use crate::client::{Model, ModelType}; use crate::config; use crate::config::{Agent, AppConfig, paths}; @@ -8,6 +8,7 @@ use anyhow::{Result, bail}; use std::collections::{BTreeMap, HashSet, VecDeque}; use std::path::PathBuf; use std::sync::Arc; +use tokio::sync::Semaphore; #[derive(Debug, Clone)] pub struct ValidationError { @@ -176,8 +177,10 @@ impl GraphValidator { self.validate_llm_nodes(graph, &mut result); self.validate_llm_skills(graph, &mut result); self.validate_max_concurrency(graph, &mut result); + self.validate_max_concurrency_template(graph, &mut result); self.validate_orchestration_limits(graph, &mut result); - self.validate_map_branches(graph, &mut result); + self.validate_timeouts(graph, &mut result); + self.validate_map_subgraphs(graph, &mut result); self.validate_parallel_user_interaction(graph, &mut result); self.validate_parallel_writes(graph, &mut result); self.validate_parallel_reads(graph, &mut result); @@ -474,6 +477,34 @@ impl GraphValidator { a.agent ), )); + } else if has_config && has_graph { + result.error(ValidationError::with_node( + node_id, + format!( + "Agent '{}' has both config.yaml and graph.yaml; a graph agent is \ + defined by graph.yaml alone. Remove one of the two files", + a.agent + ), + )); + } + if let Some(inputs) = &a.inputs { + if has_config && !has_graph { + result.error(ValidationError::with_node( + node_id, + format!( + "agent node '{node_id}': `inputs:` requires agent '{}' to be a graph agent (no graph.yaml found)", + a.agent + ), + )); + } + if inputs.contains_key("initial_prompt") { + result.error(ValidationError::with_node( + node_id, + format!( + "agent node '{node_id}': `inputs.initial_prompt` is reserved (the dispatcher seeds it from `prompt:`)" + ), + )); + } } } } @@ -505,8 +536,8 @@ impl GraphValidator { // Parallel-execution validation. // // The v1 algorithm uses immediate-successor analysis only: a parallel group is the set of `next:` targets of a - // single fan-out node. Map nodes are checked separately by `validate_map_branches` (the branch is self-parallel, - // but enforcement comes from strict-mode rules on the branch node, not from group membership). Transitive parallel + // single fan-out node. Map branch subgraphs are checked separately by `validate_map_subgraphs` (per-item forks + // never race each other, so their nodes are not group members). Transitive parallel // groups (deeper fan-out chains) are a v2 enhancement; v1 over-reports rather than under-reports. A false positive // forces an unneeded reducer (mild annoyance); a false negative allows silent data races (catastrophic). fn validate_max_concurrency(&self, graph: &Graph, result: &mut ValidationResult) { @@ -516,10 +547,16 @@ impl GraphValidator { deadlock the executor", )); } + if graph.settings.max_concurrency > Semaphore::MAX_PERMITS { + result.error(ValidationError::new(format!( + "settings.max_concurrency must be <= {} (got {}); the runtime cannot allocate more permits", + Semaphore::MAX_PERMITS, graph.settings.max_concurrency + ))); + } for (node_id, node) in &graph.nodes { if let NodeType::Map(m) = &node.node_type - && let Some(0) = m.max_concurrency + && let Some(ConcurrencyCap::Fixed(0)) = m.max_concurrency { result.error(ValidationError::with_node( node_id, @@ -527,6 +564,35 @@ impl GraphValidator { would deadlock the executor", )); } + if let NodeType::Map(m) = &node.node_type + && let Some(ConcurrencyCap::Fixed(n)) = m.max_concurrency + && n > Semaphore::MAX_PERMITS + { + result.error(ValidationError::with_node( + node_id, + format!( + "map node's `max_concurrency` must be <= {} (got {n}); the runtime cannot allocate more permits", + Semaphore::MAX_PERMITS + ), + )); + } + } + } + + // A string cap that never interpolates would only fail at run time, after the + // map's inputs were already computed; catch the typo (`"4"` for `4`) at load. + fn validate_max_concurrency_template(&self, graph: &Graph, result: &mut ValidationResult) { + for (node_id, node) in &graph.nodes { + if let NodeType::Map(m) = &node.node_type + && let Some(ConcurrencyCap::Template(t)) = &m.max_concurrency + && !is_lone_template(t) + { + result.error(ValidationError::with_node( + node_id, + "map node's `max_concurrency` is a string but not a template; write \ + an integer or exactly one `{{key}}` with nothing around it", + )); + } } } @@ -547,107 +613,196 @@ impl GraphValidator { } } - fn validate_map_branches(&self, graph: &Graph, result: &mut ValidationResult) { + /// Explicit `0` opt-outs on `settings` and node bounds: a `timeout: 0` + /// disables the wall-clock bound, `max_loop_iterations: 0` disables the + /// per-node visit cap, and an llm node's `max_iterations: 0` disables its + /// tool-call-loop turn cap. All are legal; the warnings only make the + /// choice visible. + fn validate_timeouts(&self, graph: &Graph, result: &mut ValidationResult) { + if graph.settings.max_loop_iterations == 0 { + result.warning(ValidationError::new( + "settings.max_loop_iterations: 0 disables the per-node visit cap; a \ + non-converging loop runs until an enclosing timeout or abort", + )); + } + if graph.settings.timeout == Some(0) { + result.warning(ValidationError::new( + "settings.timeout: 0 disables the graph wall-clock bound; the run \ + ends only when the graph completes, is aborted, or another \ + limit stops it", + )); + } + for (node_id, node) in &graph.nodes { + if let NodeType::Llm(l) = &node.node_type + && l.max_iterations == 0 + { + result.warning(ValidationError::with_node( + node_id, + format!( + "max_iterations: 0 disables the turn cap for llm node '{node_id}'; it runs \ + until the model concludes, the context window fills, an enclosing timeout, or abort" + ), + )); + } + let kind = match &node.node_type { + NodeType::Agent(a) if a.timeout == Some(0) => "agent", + NodeType::Script(s) if s.timeout == 0 => "script", + NodeType::Llm(l) if l.timeout == Some(0) => "llm", + NodeType::Rag(r) if r.timeout == Some(0) => "rag", + _ => continue, + }; + result.warning(ValidationError::with_node( + node_id, + format!( + "timeout: 0 disables the wall-clock bound for {kind} node \ + '{node_id}'; it runs until it returns or an enclosing bound \ + (settings.timeout, an outer agent node's timeout, or abort) ends it" + ), + )); + } + } + + /// Rules for the per-item subgraph rooted at each map's `branch`. Every + /// error is attributed to the map node; the offending branch node is + /// named in the message. + fn validate_map_subgraphs(&self, graph: &Graph, result: &mut ValidationResult) { + let main_flow = main_flow_reachable(graph); + for (map_id, node) in &graph.nodes { let NodeType::Map(m) = &node.node_type else { continue; }; - let Some(branch) = graph.get_node(&m.branch) else { + if m.as_name == m.output_key { + result.error(ValidationError::with_node( + map_id, + format!( + "map node '{map_id}': `as` and `output_key` are both '{}'; the item \ + binding would be collected as the chain's result", + m.as_name + ), + )); + } + // A missing entry is reported by `validate_node_references`. + let mut member_ids: Vec = + branch_subgraph(graph, &m.branch).into_iter().collect(); + if member_ids.is_empty() { continue; - }; + } + member_ids.sort(); + let members: Vec<(&str, &Node)> = member_ids + .iter() + .filter_map(|id| graph.get_node(id).map(|n| (id.as_str(), n))) + .collect(); - match &branch.node_type { - NodeType::Approval(_) => { + let fallback_targets: HashSet<&str> = members + .iter() + .filter_map(|(_, n)| match &n.node_type { + NodeType::Script(s) => s.fallback.as_deref(), + NodeType::Llm(l) => l.fallback.as_deref(), + _ => None, + }) + .collect(); + + let mut has_disallowed_node = false; + for (id, member) in &members { + let disallowed = match &member.node_type { + NodeType::Approval(_) => Some( + "an approval node; approval/input nodes cannot run inside a parallel \ + map branch (the CLI would prompt the user N times concurrently)", + ), + NodeType::Input(_) => { + Some("an input node; input nodes cannot run inside a parallel map branch") + } + NodeType::End(_) => Some( + "an end node; map branches terminate via the map's collect \ + mechanism, not via end nodes", + ), + NodeType::Map(_) => { + Some("itself a map node; nested map fan-outs are not supported in v1") + } + NodeType::Agent(_) + | NodeType::Llm(_) + | NodeType::Rag(_) + | NodeType::Script(_) => None, + }; + if let Some(reason) = disallowed { + has_disallowed_node = true; + let mut message = if *id == m.branch { + format!("map node points to branch '{id}' which is {reason}") + } else { + format!( + "map node's branch subgraph (entry '{}') reaches '{id}', which is \ + {reason}", + m.branch + ) + }; + if fallback_targets.contains(id) { + message.push_str( + "; a branch's `fallback` must stay inside the branch subgraph; \ + `fallback` on a map branch was previously accepted but \ + never honored. Remove it or point it at a branch-local node", + ); + } + result.error(ValidationError::with_node(map_id, message)); + } + + if member.next.as_ref().is_some_and(NextTargets::is_fan_out) { result.error(ValidationError::with_node( map_id, format!( - "map node points to branch '{}' which is an approval node; \ - approval/input nodes cannot run inside a parallel map branch \ - (the CLI would prompt the user N times concurrently)", + "node '{id}' inside the branch subgraph (entry '{}') declares a \ + fan-out `next`; fan-out inside a map branch is not supported in \ + v1; use a nested map after the join", m.branch ), )); - continue; } - NodeType::Input(_) => { - result.error(ValidationError::with_node( - map_id, - format!( - "map node points to branch '{}' which is an input node; \ - input nodes cannot run inside a parallel map branch", - m.branch - ), - )); - continue; - } - NodeType::End(_) => { - result.error(ValidationError::with_node( - map_id, - format!( - "map node points to branch '{}' which is an end node; \ - map branches terminate via the map's collect mechanism, \ - not via end nodes", - m.branch - ), - )); - continue; - } - NodeType::Map(_) => { - result.error(ValidationError::with_node( - map_id, - format!( - "map node points to branch '{}' which is itself a map node; \ - nested map fan-outs are not supported in v1", - m.branch - ), - )); - continue; - } - _ => {} } - if branch.next.is_some() { + let shared: Vec<&str> = members + .iter() + .filter(|(id, _)| main_flow.contains(*id)) + .map(|(id, _)| *id) + .collect(); + if !shared.is_empty() { result.error(ValidationError::with_node( - m.branch.clone(), + map_id, format!( - "branch node '{}' has a `next` declared, but map branches must be \ - atomic (one node, one execution per item). Remove `next` or \ - restructure the workflow so any chaining happens after the map.", - m.branch + "branch subgraph (entry '{}') shares node(s) [{}] with the main flow \ + (reachable from start '{}' without entering a map branch); a branch \ + node that is also on the main path would run standalone without the \ + `as` binding and could route into the branch mid-chain. Give the \ + branch its own nodes, or move the shared step after the map.", + m.branch, + shared.join(", "), + graph.start ), )); } - if let Some(updates) = node_state_updates_keys(branch) { - for k in &updates { - if k != &m.output_key { - result.error(ValidationError::with_node( - m.branch.clone(), - format!( - "branch node '{}' writes state key '{}' via state_updates, \ - but map branches may only write through their `output_key` \ - ('{}'). Rename the write, or move the side effect outside \ - the map.", - m.branch, k, m.output_key - ), - )); - } - } + if has_disallowed_node { + continue; } - let schema_keys = output_schema_top_level_keys(branch); - if !schema_keys.is_empty() { - let mut keys_sorted: Vec = schema_keys.into_iter().collect(); - keys_sorted.sort(); - result.error(ValidationError::with_node( - m.branch.clone(), + // A script may write any key through its JSON output, so its + // presence makes the subgraph's write set unknowable. + let has_script = members + .iter() + .any(|(_, n)| matches!(n.node_type, NodeType::Script(_))); + if has_script { + continue; + } + let declares_output_key = members.iter().any(|(_, n)| { + node_state_updates_keys(n).is_some_and(|keys| keys.contains(&m.output_key)) + || output_schema_top_level_keys(n).contains(&m.output_key) + }); + if !declares_output_key { + result.warning(ValidationError::with_node( + map_id, format!( - "branch node '{}' has an `output_schema` with top-level \ - properties ({}); map branches must write only through their \ - `output_key` ('{}'). Remove `output_schema`, or use state_updates \ - to map the output explicitly.", - m.branch, - keys_sorted.join(", "), + "no node in the branch subgraph of map '{map_id}' declares a write to \ + output_key '{}'; the map will error at runtime if the chain does not \ + write it", m.output_key ), )); @@ -844,9 +999,83 @@ fn find_reachable_nodes(graph: &Graph) -> HashSet { reachable } +/// Nodes reachable from the start node without entering any map's branch: +/// the path the frontier itself walks. A map node contributes only its +/// `next` edges here, so per-item subgraphs can be checked for overlap +/// with the main flow. +fn main_flow_reachable(graph: &Graph) -> HashSet { + let mut reachable: HashSet = HashSet::new(); + let mut queue: VecDeque = VecDeque::new(); + + if !graph.has_node(&graph.start) { + return reachable; + } + + reachable.insert(graph.start.clone()); + queue.push_back(graph.start.clone()); + + while let Some(id) = queue.pop_front() { + let Some(node) = graph.get_node(&id) else { + continue; + }; + let edges: Vec = match &node.node_type { + NodeType::Map(_) => node + .next + .as_ref() + .map(|t| t.as_slice().to_vec()) + .unwrap_or_default(), + _ => outgoing_node_ids(node), + }; + for next in edges { + if graph.has_node(&next) && reachable.insert(next.clone()) { + queue.push_back(next); + } + } + } + reachable +} + +/// Nodes reachable from `entry` over the edges a chain can follow at run +/// time: `next` targets and script/llm `fallback`s. Approval routes and a +/// nested map's `branch` are not followed. Neither may appear inside a +/// branch, and the validator reports them separately. +pub(super) fn branch_subgraph(graph: &Graph, entry: &str) -> HashSet { + let mut reachable: HashSet = HashSet::new(); + let mut queue: VecDeque = VecDeque::new(); + + if !graph.has_node(entry) { + return reachable; + } + + reachable.insert(entry.to_string()); + queue.push_back(entry.to_string()); + + while let Some(id) = queue.pop_front() { + let Some(node) = graph.get_node(&id) else { + continue; + }; + let mut edges: Vec<&String> = node + .next + .as_ref() + .map(|t| t.as_slice().iter().collect()) + .unwrap_or_default(); + match &node.node_type { + NodeType::Script(s) => edges.extend(s.fallback.as_ref()), + NodeType::Llm(l) => edges.extend(l.fallback.as_ref()), + _ => {} + } + for next in edges { + if graph.has_node(next) && reachable.insert(next.clone()) { + queue.push_back(next.clone()); + } + } + } + reachable +} + // v1 parallel-group detection: only the immediate `next` targets of a fan-out node count as a parallel group. Map -// branches are handled separately by `validate_map_branches` (the branch's self-parallelism is checked via strict-mode -// rules on the branch node itself, not via group membership). +// branch subgraphs are handled separately by `validate_map_subgraphs` (per-item forks never race each other, so +// their nodes are not group members). // // Returns one HashSet per fan-out source; deeper transitive parallelism is intentionally out of scope for v1. fn compute_parallel_groups(graph: &Graph) -> Vec> { @@ -923,7 +1152,13 @@ fn primary_templated_fields(node: &Node) -> Vec { } v } - NodeType::Agent(n) => vec![n.prompt.clone()], + NodeType::Agent(n) => { + let mut v = vec![n.prompt.clone()]; + if let Some(inputs) = &n.inputs { + v.extend(inputs.values().cloned()); + } + v + } NodeType::Rag(n) => { vec![ n.query @@ -940,7 +1175,13 @@ fn primary_templated_fields(node: &Node) -> Vec { v } NodeType::End(n) => vec![n.output.clone()], - NodeType::Map(n) => vec![n.over.clone()], + NodeType::Map(n) => { + let mut v = vec![n.over.clone()]; + if let Some(ConcurrencyCap::Template(t)) = &n.max_concurrency { + v.push(t.clone()); + } + v + } NodeType::Script(_) => Vec::new(), } } @@ -1026,9 +1267,14 @@ fn detect_cycle_dfs( mod tests { use super::super::types::*; use super::*; + use crate::utils::get_env_name; use indexmap::IndexMap; + use serial_test::serial; use std::collections::HashMap; use std::env; + use std::fs; + use std::path::PathBuf; + use std::time::{SystemTime, UNIX_EPOCH}; fn graph_with(nodes: Vec<(&str, Node)>, start: &str) -> Graph { let mut map: IndexMap = IndexMap::new(); @@ -1658,6 +1904,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + inputs: None, teammates: false, }), next: next.map(NextTargets::from), @@ -2433,7 +2680,31 @@ mod tests { } #[test] - fn map_branch_cannot_have_next_declared() { + fn map_branch_with_in_subgraph_next_passes() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("br2")); + let br2 = llm_with_state_updates("br2", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("br2", br2), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.is_empty(), + "a branch chaining to a node inside its own subgraph is valid: {:?}", + result.errors + ); + } + + #[test] + fn map_branch_next_to_end_node_errors() { let map = map_node_basic("m", "br", Some("end")); let branch = llm_with_state_updates("br", &[("output", "{{output}}")], Some("somewhere")); let graph = graph_with( @@ -2452,10 +2723,38 @@ mod tests { result .errors .iter() - .any(|e| e.message.contains("has a `next` declared") - && e.message.contains("atomic") - && e.node_id.as_deref() == Some("br")), - "expected branch-has-next error: {:?}", + .any(|e| e.message.contains("'somewhere'") + && e.message.contains("end node") + && e.message.contains("collect mechanism") + && !e.message.contains("never honored") + && e.node_id.as_deref() == Some("m")), + "expected end-node-in-subgraph error attributed to the map: {:?}", + result.errors + ); + } + + #[test] + fn map_branch_next_escaping_to_main_flow_errors() { + let map = map_node_basic("m", "br", Some("after")); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], Some("after")); + let after = llm_node("after", None, Some("end")); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("after", after), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains("main flow") + && e.message.contains("[after") + && e.node_id.as_deref() == Some("m")), + "expected shared-with-main-flow error naming 'after': {:?}", result.errors ); } @@ -2481,7 +2780,7 @@ mod tests { } #[test] - fn map_branch_state_updates_wrong_key_errors() { + fn map_branch_state_updates_scratch_key_passes() { let map = map_node_basic("m", "br", Some("end")); let branch = llm_with_state_updates("br", &[("not_output", "{{output}}")], None); let graph = graph_with( @@ -2492,19 +2791,14 @@ mod tests { let result = validator().validate(&graph); assert!( - result - .errors - .iter() - .any(|e| e.message.contains("writes state key 'not_output'") - && e.message.contains("'output'") - && e.node_id.as_deref() == Some("br")), - "expected wrong-key error: {:?}", + result.errors.is_empty(), + "branch-local scratch writes are allowed: {:?}", result.errors ); } #[test] - fn map_branch_with_output_schema_errors() { + fn map_branch_with_output_schema_passes() { let map = map_node_basic("m", "br", Some("end")); let branch = llm_with_output_schema("br", &["foo", "bar"], None); let graph = graph_with( @@ -2515,17 +2809,308 @@ mod tests { let result = validator().validate(&graph); assert!( - result - .errors - .iter() - .any(|e| e.message.contains("output_schema") - && e.message.contains("top-level properties") - && e.node_id.as_deref() == Some("br")), - "expected output_schema-forbidden error: {:?}", + result.errors.is_empty(), + "an output_schema on a branch node is allowed: {:?}", result.errors ); } + #[test] + fn map_branch_disallowed_node_reached_via_fallback_gets_hint() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", Some("end"), None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains("'end'") + && e.message.contains("end node") + && e.message.contains("never honored") + && e.node_id.as_deref() == Some("m")), + "expected end-node error with the fallback hint: {:?}", + result.errors + ); + } + + #[test] + fn map_branch_deep_approval_node_errors_on_the_map() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("gate")); + let gate = approval_node("gate", &["yes"], &[("yes", "end")], "end"); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("gate", gate), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains("'gate'") + && e.message.contains("approval node") + && e.message.contains("map branch") + && e.node_id.as_deref() == Some("m")), + "expected deep approval error attributed to the map: {:?}", + result.errors + ); + } + + #[test] + fn map_branch_fan_out_inside_subgraph_errors() { + let map = map_node_basic("m", "br", Some("end")); + let mut branch = llm_node("br", None, None); + branch.next = Some(NextTargets::Many(vec!["x".into(), "y".into()])); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("x", llm_node("x", None, None)), + ("y", llm_node("y", None, None)), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains("'br'") + && e.message.contains("fan-out inside a map branch") + && e.node_id.as_deref() == Some("m")), + "expected fan-out-in-branch error: {:?}", + result.errors + ); + } + + #[test] + fn map_branch_sharing_node_with_main_flow_errors() { + let mut start = llm_node("s", None, None); + start.next = Some(NextTargets::Many(vec!["m".into(), "shared".into()])); + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("shared")); + let shared = llm_with_state_updates("shared", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![ + ("s", start), + ("m", map), + ("br", branch), + ("shared", shared), + ("end", end_node("end")), + ], + "s", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains("main flow") + && e.message.contains("[shared]") + && e.node_id.as_deref() == Some("m")), + "expected shared-node error naming only 'shared': {:?}", + result.errors + ); + } + + #[test] + fn two_maps_sharing_a_subgraph_pass() { + let first = map_node_basic("m1", "br", Some("m2")); + let second = map_node_basic("m2", "br", Some("end")); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![ + ("m1", first), + ("m2", second), + ("br", branch), + ("end", end_node("end")), + ], + "m1", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.is_empty(), + "two maps may share a branch subgraph: {:?}", + result.errors + ); + } + + fn output_key_warnings(result: &ValidationResult) -> Vec<&ValidationError> { + result + .warnings + .iter() + .filter(|w| w.message.contains("declares a write to output_key")) + .collect() + } + + #[test] + fn map_branch_without_declared_output_writer_warns() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!(result.errors.is_empty(), "{:?}", result.errors); + let warnings = output_key_warnings(&result); + assert_eq!(warnings.len(), 1, "{:?}", result.warnings); + assert_eq!(warnings[0].node_id.as_deref(), Some("m")); + assert!( + warnings[0] + .message + .contains("branch subgraph of map 'm' declares a write to output_key 'output'"), + "{}", + warnings[0].message + ); + } + + #[test] + fn map_branch_script_in_subgraph_suppresses_missing_writer_warning() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("finish")); + let finish = script_node("finish", "Cargo.toml", None); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("finish", finish), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!(result.errors.is_empty(), "{:?}", result.errors); + assert!( + output_key_warnings(&result).is_empty(), + "a script anywhere in the subgraph may write the key: {:?}", + result.warnings + ); + } + + #[test] + fn map_branch_state_updates_deeper_in_chain_suppress_missing_writer_warning() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("br2")); + let br2 = llm_with_state_updates("br2", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("br2", br2), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!(result.errors.is_empty(), "{:?}", result.errors); + assert!( + output_key_warnings(&result).is_empty(), + "state_updates naming output_key anywhere in the chain is a declared write: {:?}", + result.warnings + ); + } + + #[test] + fn map_branch_output_schema_naming_output_key_suppresses_missing_writer_warning() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_with_output_schema("br", &["output", "notes"], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!(result.errors.is_empty(), "{:?}", result.errors); + assert!( + output_key_warnings(&result).is_empty(), + "an output_schema property naming output_key is a declared write: {:?}", + result.warnings + ); + } + + #[test] + fn map_branch_three_node_chain_validates_clean() { + let map = map_node_basic("m", "a", Some("end")); + let a = llm_node("a", None, Some("b")); + let b = llm_with_state_updates("b", &[("draft", "{{output}}")], Some("c")); + let c = llm_with_state_updates("c", &[("output", "{{draft}}")], None); + let graph = graph_with( + vec![ + ("m", map), + ("a", a), + ("b", b), + ("c", c), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + assert!(result.errors.is_empty(), "{:?}", result.errors); + assert!( + output_key_warnings(&result).is_empty(), + "{:?}", + result.warnings + ); + } + + #[test] + fn main_flow_reachable_does_not_enter_map_branches() { + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_node("br", None, Some("br2")); + let br2 = llm_node("br2", Some("end"), None); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("br2", br2), + ("end", end_node("end")), + ], + "m", + ); + + assert_eq!(ids(&main_flow_reachable(&graph)), vec!["end", "m"]); + assert_eq!( + ids(&find_reachable_nodes(&graph)), + vec!["br", "br2", "end", "m"] + ); + } + + #[test] + fn main_flow_reachable_follows_approval_routes() { + let gate = approval_node("gate", &["yes"], &[("yes", "m")], "end"); + let map = map_node_basic("m", "br", Some("end")); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![ + ("gate", gate), + ("m", map), + ("br", branch), + ("end", end_node("end")), + ], + "gate", + ); + + assert_eq!(ids(&main_flow_reachable(&graph)), vec!["end", "gate", "m"]); + } + #[test] fn script_in_fan_out_without_state_updates_errors() { let a = script_node("worker_a", "Cargo.toml", None); @@ -2602,6 +3187,24 @@ mod tests { ); } + #[test] + fn settings_max_concurrency_above_semaphore_limit_errors() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.settings.max_concurrency = Semaphore::MAX_PERMITS + 1; + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message.contains(&format!( + "settings.max_concurrency must be <= {} (got {})", + Semaphore::MAX_PERMITS, + Semaphore::MAX_PERMITS + 1 + ))), + "expected graph-level max_concurrency upper-bound error: {:?}", + result.errors + ); + } + #[test] fn settings_max_concurrency_default_is_valid() { let graph = graph_with(vec![("e", end_node("e"))], "e"); @@ -2669,11 +3272,272 @@ mod tests { ); } + fn timeout_warnings(result: &ValidationResult) -> Vec<&ValidationError> { + result + .warnings + .iter() + .filter(|w| w.message.contains("disables the")) + .collect() + } + + #[test] + fn settings_timeout_zero_warns() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.settings.timeout = Some(0); + + let result = validator().validate(&graph); + + assert!(result.is_valid()); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id, None); + assert_eq!( + w[0].message, + "settings.timeout: 0 disables the graph wall-clock bound; the run ends only when \ + the graph completes, is aborted, or another limit stops it" + ); + } + + #[test] + fn agent_timeout_zero_warns() { + let mut a = agent_node("a", "worker", Some("end")); + if let NodeType::Agent(ref mut an) = a.node_type { + an.timeout = Some(0); + } + let graph = graph_with(vec![("a", a), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id.as_deref(), Some("a")); + assert_eq!( + w[0].message, + "timeout: 0 disables the wall-clock bound for agent node 'a'; it runs until it returns \ + or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it" + ); + } + + #[test] + fn script_timeout_zero_warns() { + let mut s = script_node("s", "does-not-exist.py", None); + if let NodeType::Script(ref mut sn) = s.node_type { + sn.timeout = 0; + } + s.next = Some("end".into()); + let graph = graph_with(vec![("s", s), ("end", end_node("end"))], "s"); + + let result = validator().validate(&graph); + + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id.as_deref(), Some("s")); + assert_eq!( + w[0].message, + "timeout: 0 disables the wall-clock bound for script node 's'; it runs until it returns \ + or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it" + ); + } + + #[test] + fn llm_timeout_zero_warns() { + let mut l = llm_node("l", None, Some("end")); + if let NodeType::Llm(ref mut ln) = l.node_type { + ln.timeout = Some(0); + } + let graph = graph_with(vec![("l", l), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id.as_deref(), Some("l")); + assert_eq!( + w[0].message, + "timeout: 0 disables the wall-clock bound for llm node 'l'; it runs until it returns \ + or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it" + ); + } + + #[test] + fn rag_timeout_zero_warns() { + let mut r = rag_node("r", &["docs/"], true); + if let NodeType::Rag(ref mut rn) = r.node_type { + rn.timeout = Some(0); + } + let graph = graph_with(vec![("r", r), ("end", end_node("end"))], "r"); + + let result = validator().validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id.as_deref(), Some("r")); + assert_eq!( + w[0].message, + "timeout: 0 disables the wall-clock bound for rag node 'r'; it runs until it returns \ + or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it" + ); + } + + #[test] + fn nonzero_and_unset_timeouts_do_not_warn() { + let mut a = agent_node("a", "worker", Some("s")); + if let NodeType::Agent(ref mut an) = a.node_type { + an.timeout = Some(600); + } + let mut s = script_node("s", "does-not-exist.py", None); + s.next = Some("l".into()); + let l = llm_node("l", None, Some("end")); + let mut graph = graph_with( + vec![("a", a), ("s", s), ("l", l), ("end", end_node("end"))], + "a", + ); + graph.settings.timeout = Some(30); + + let result = validator().validate(&graph); + + assert!( + timeout_warnings(&result).is_empty(), + "non-zero and unset timeouts should not warn: {:?}", + result.warnings + ); + } + + #[test] + fn max_loop_iterations_zero_warns() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.settings.max_loop_iterations = 0; + + let result = validator().validate(&graph); + + assert!(result.is_valid()); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert!(w[0].node_id.is_none()); + assert_eq!( + w[0].message, + "settings.max_loop_iterations: 0 disables the per-node visit cap; a non-converging \ + loop runs until an enclosing timeout or abort" + ); + } + + #[test] + fn max_loop_iterations_nonzero_does_not_warn() { + let mut graph = graph_with(vec![("e", end_node("e"))], "e"); + graph.settings.max_loop_iterations = 5; + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|w| w.message.contains("max_loop_iterations: 0")), + "non-zero cap should not warn: {:?}", + result.warnings + ); + } + + #[test] + fn llm_max_iterations_zero_warns() { + let mut l = llm_node("l", None, Some("end")); + if let NodeType::Llm(ref mut ln) = l.node_type { + ln.max_iterations = 0; + } + let graph = graph_with(vec![("l", l), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 1, "{:?}", result.warnings); + assert_eq!(w[0].node_id.as_deref(), Some("l")); + assert_eq!( + w[0].message, + "max_iterations: 0 disables the turn cap for llm node 'l'; it runs until the model \ + concludes, the context window fills, an enclosing timeout, or abort" + ); + } + + #[test] + fn llm_max_iterations_nonzero_does_not_warn() { + let graph = graph_with( + vec![ + ("l", llm_node("l", None, Some("end"))), + ("end", end_node("end")), + ], + "l", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|w| w.message.contains("max_iterations: 0")), + "non-zero cap should not warn: {:?}", + result.warnings + ); + } + + #[test] + fn llm_zero_cap_and_zero_timeout_warn_in_order() { + let mut l = llm_node("l", None, Some("end")); + if let NodeType::Llm(ref mut ln) = l.node_type { + ln.max_iterations = 0; + ln.timeout = Some(0); + } + let graph = graph_with(vec![("l", l), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + let w = timeout_warnings(&result); + assert_eq!(w.len(), 2, "{:?}", result.warnings); + assert!( + w[0].message + .starts_with("max_iterations: 0 disables the turn cap") + ); + assert!( + w[1].message + .starts_with("timeout: 0 disables the wall-clock bound") + ); + assert_eq!(w[0].node_id.as_deref(), Some("l")); + assert_eq!(w[1].node_id.as_deref(), Some("l")); + } + + #[test] + fn map_max_concurrency_above_semaphore_limit_errors() { + let mut map = map_node_basic("m", "br", Some("end")); + if let NodeType::Map(ref mut mm) = map.node_type { + mm.max_concurrency = Some(ConcurrencyCap::Fixed(Semaphore::MAX_PERMITS + 1)); + } + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e + .message + .contains(&format!("must be <= {}", Semaphore::MAX_PERMITS)) + && e.node_id.as_deref() == Some("m")), + "expected map max_concurrency upper-bound error: {:?}", + result.errors + ); + } + #[test] fn map_max_concurrency_zero_errors() { let mut map = map_node_basic("m", "br", Some("end")); if let NodeType::Map(ref mut mm) = map.node_type { - mm.max_concurrency = Some(0); + mm.max_concurrency = Some(ConcurrencyCap::Fixed(0)); } let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); let graph = graph_with( @@ -2714,6 +3578,226 @@ mod tests { ); } + fn map_with_cap(id: &str, branch: &str, next: Option<&str>, cap: &str) -> Node { + let mut map = map_node_basic(id, branch, next); + if let NodeType::Map(ref mut mm) = map.node_type { + mm.max_concurrency = Some(ConcurrencyCap::Template(cap.into())); + } + map + } + + fn map_with_as(id: &str, branch: &str, next: Option<&str>, as_name: &str) -> Node { + let mut map = map_node_basic(id, branch, next); + if let NodeType::Map(ref mut mm) = map.node_type { + mm.as_name = as_name.into(); + } + map + } + + #[test] + fn map_as_equal_output_key_errors() { + let map = map_with_as("m", "br", Some("end"), "output"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message + == "map node 'm': `as` and `output_key` are both 'output'; the item binding \ + would be collected as the chain's result" + && e.node_id.as_deref() == Some("m")), + "expected as/output_key collision error: {:?}", + result.errors + ); + } + + #[test] + fn map_as_distinct_from_output_key_passes() { + let map = map_with_as("m", "br", Some("end"), "item"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .errors + .iter() + .any(|e| e.message.contains("`as` and `output_key`")), + "distinct as/output_key must not error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_non_template_string_errors() { + let map = map_with_cap("m", "br", Some("end"), "4"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e + .message + .contains("`max_concurrency` is a string but not a template") + && e.node_id.as_deref() == Some("m")), + "expected non-template string cap error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_template_string_is_valid() { + let map = map_with_cap("m", "br", Some("end"), "{{k}}"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .errors + .iter() + .any(|e| e.message.contains("max_concurrency")), + "templated cap should not error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_unclosed_template_errors() { + let map = map_with_cap("m", "br", Some("end"), "{{k"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("is a string but not a template")), + "expected unclosed-template cap error: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_malformed_templates_error() { + for cap in ["{{}}", "{{ key }}", "n={{k}}", "{{a}} {{b", "{{a-b}}"] { + let map = map_with_cap("m", "br", Some("end"), cap); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message + == "map node's `max_concurrency` is a string but not a template; write \ + an integer or exactly one `{{key}}` with nothing around it" + && e.node_id.as_deref() == Some("m")), + "cap {cap:?} should fail the whole-string template check: {:?}", + result.errors + ); + } + } + + #[test] + fn map_max_concurrency_non_numeric_index_errors() { + let map = map_with_cap("m", "br", Some("end"), "{{limits[foo]}}"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e.message + == "map node's `max_concurrency` is a string but not a template; write \ + an integer or exactly one `{{key}}` with nothing around it" + && e.node_id.as_deref() == Some("m")), + "a non-numeric index can never resolve and must fail at load: {:?}", + result.errors + ); + } + + #[test] + fn map_max_concurrency_lone_template_forms_pass() { + for cap in ["{{budget}}", "{{cfg.limits[0]}}", " {{budget}} "] { + let map = map_with_cap("m", "br", Some("end"), cap); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let graph = graph_with( + vec![("m", map), ("br", branch), ("end", end_node("end"))], + "m", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .errors + .iter() + .any(|e| e.message.contains("is a string but not a template")), + "cap {cap:?} is a lone template and must pass: {:?}", + result.errors + ); + } + } + + #[test] + fn map_max_concurrency_template_reading_sibling_write_errors() { + let mut start = end_node("start"); + start.next = Some(NextTargets::Many(vec!["m".into(), "worker_b".into()])); + let map = map_with_cap("m", "br", Some("end"), "{{summary}}"); + let branch = llm_with_state_updates("br", &[("output", "{{output}}")], None); + let writer = llm_with_state_updates("worker_b", &[("summary", "static")], Some("end")); + let graph = graph_with( + vec![ + ("start", start), + ("m", map), + ("br", branch), + ("worker_b", writer), + ("end", end_node("end")), + ], + "start", + ); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("reads state key(s) `summary`") + && e.message.contains("'worker_b'") + && e.node_id.as_deref() == Some("m")), + "expected cross-branch read error for templated cap: {:?}", + result.errors + ); + } + fn llm_with_prompt(id: &str, prompt: &str, next: Option<&str>) -> Node { let mut node = llm_node(id, None, next); if let NodeType::Llm(ref mut n) = node.node_type { @@ -2849,4 +3933,293 @@ mod tests { result.errors ); } + + 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-validator-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); + } + } + + fn materialize_config_only_agent(name: &str) { + fs::create_dir_all(paths::agent_data_dir(name)).unwrap(); + fs::write(paths::agent_config_file(name), "").unwrap(); + } + + fn materialize_graph_agent(name: &str) { + fs::create_dir_all(paths::agent_data_dir(name)).unwrap(); + fs::write(paths::agent_graph_file(name), "").unwrap(); + } + + fn agent_node_with_inputs(id: &str, agent: &str, inputs: &[(&str, &str)]) -> Node { + let mut node = agent_node(id, agent, Some("end")); + if let NodeType::Agent(ref mut a) = node.node_type { + a.inputs = Some( + inputs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(), + ); + } + node + } + + fn inputs_requires_graph_agent_error(result: &ValidationResult, node_id: &str) -> bool { + result.errors.iter().any(|e| { + e.message.contains("`inputs:` requires agent") + && e.message + .contains("to be a graph agent (no graph.yaml found)") + && e.node_id.as_deref() == Some(node_id) + }) + } + + #[test] + #[serial] + fn agent_inputs_on_config_only_agent_errors() { + let _guard = TestConfigDirGuard::new(); + materialize_config_only_agent("plain-agent"); + let node = agent_node_with_inputs("a", "plain-agent", &[("x", "{{k}}")]); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| { + e.message + == "agent node 'a': `inputs:` requires agent 'plain-agent' to be a graph agent (no graph.yaml found)" + && e.node_id.as_deref() == Some("a") + }), + "{:?}", + result.errors + ); + } + + #[test] + #[serial] + fn agent_empty_inputs_on_config_only_agent_still_errors() { + let _guard = TestConfigDirGuard::new(); + materialize_config_only_agent("plain-agent"); + let node = agent_node_with_inputs("a", "plain-agent", &[]); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!( + inputs_requires_graph_agent_error(&result, "a"), + "{:?}", + result.errors + ); + } + + #[test] + #[serial] + fn agent_inputs_on_graph_agent_passes() { + let _guard = TestConfigDirGuard::new(); + materialize_graph_agent("graph-agent"); + let node = agent_node_with_inputs("a", "graph-agent", &[("x", "{{k}}")]); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!(result.is_valid(), "{:?}", result.errors); + assert!(!inputs_requires_graph_agent_error(&result, "a")); + } + + #[test] + #[serial] + fn agent_with_both_config_and_graph_errors() { + let _guard = TestConfigDirGuard::new(); + materialize_config_only_agent("mixed-agent"); + materialize_graph_agent("mixed-agent"); + let node = agent_node_with_inputs("a", "mixed-agent", &[("x", "{{k}}")]); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| e + .message + .starts_with("Agent 'mixed-agent' has both config.yaml and graph.yaml") + && e.node_id.as_deref() == Some("a")), + "{:?}", + result.errors + ); + } + + #[test] + #[serial] + fn agent_inputs_missing_agent_dir_gets_only_the_not_found_error() { + let _guard = TestConfigDirGuard::new(); + let node = agent_node_with_inputs("a", "ghost-agent", &[("x", "{{k}}")]); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("Agent 'ghost-agent' not found")), + "{:?}", + result.errors + ); + assert!(!inputs_requires_graph_agent_error(&result, "a")); + } + + #[test] + #[serial] + fn agent_inputs_initial_prompt_is_reserved() { + let _guard = TestConfigDirGuard::new(); + materialize_graph_agent("graph-agent"); + let node = agent_node_with_inputs( + "a", + "graph-agent", + &[("initial_prompt", "{{k}}"), ("x", "{{k}}")], + ); + let graph = graph_with(vec![("a", node), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!( + result.errors.iter().any(|e| { + e.message + == "agent node 'a': `inputs.initial_prompt` is reserved (the dispatcher seeds it from `prompt:`)" + && e.node_id.as_deref() == Some("a") + }), + "{:?}", + result.errors + ); + assert!(!inputs_requires_graph_agent_error(&result, "a")); + } + + #[test] + fn agent_inputs_reading_sibling_write_errors() { + let reader = agent_node_with_inputs("worker_a", "some-agent", &[("x", "{{k}}")]); + let writer = llm_with_state_updates("worker_b", &[("k", "static")], Some("end")); + let graph = fan_out_graph_with_two_workers(reader, writer); + + let result = validator().validate(&graph); + + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("reads state key(s) `k`") + && e.message.contains("'worker_b'") + && e.node_id.as_deref() == Some("worker_a")), + "expected cross-branch read error for agent inputs: {:?}", + result.errors + ); + } + + #[test] + fn agent_inputs_reading_upstream_key_passes() { + let reader = agent_node_with_inputs("worker_a", "some-agent", &[("x", "{{k}}")]); + let writer = llm_with_state_updates("worker_b", &[("other", "static")], Some("end")); + let graph = fan_out_graph_with_two_workers(reader, writer); + + let result = validator().validate(&graph); + + assert!( + !result + .errors + .iter() + .any(|e| e.message.contains("reads state key")), + "upstream `k` shouldn't trigger cross-branch read error: {:?}", + result.errors + ); + } + + fn ids(subgraph: &HashSet) -> Vec<&str> { + let mut ids: Vec<&str> = subgraph.iter().map(String::as_str).collect(); + ids.sort_unstable(); + ids + } + + #[test] + fn branch_subgraph_follows_next_and_fallback_edges() { + let mut a = script_node("a", "a.sh", Some("c")); + a.next = Some("b".to_string().into()); + let graph = graph_with( + vec![ + ("a", a), + ("b", llm_node("b", Some("d"), None)), + ("c", script_node("c", "c.sh", None)), + ("d", script_node("d", "d.sh", None)), + ("e", script_node("e", "e.sh", None)), + ], + "a", + ); + + assert_eq!(ids(&branch_subgraph(&graph, "a")), vec!["a", "b", "c", "d"]); + } + + #[test] + fn branch_subgraph_does_not_follow_approval_routes_or_nested_map_branch() { + let mut a = script_node("a", "a.sh", None); + a.next = Some(NextTargets::Many(vec!["m".into(), "p".into()])); + let graph = graph_with( + vec![ + ("a", a), + ("m", map_node_basic("m", "inner", None)), + ("inner", script_node("inner", "inner.sh", None)), + ("p", approval_node("p", &["yes"], &[("yes", "q")], "q")), + ("q", script_node("q", "q.sh", None)), + ], + "a", + ); + + assert_eq!(ids(&branch_subgraph(&graph, "a")), vec!["a", "m", "p"]); + } + + #[test] + fn branch_subgraph_of_unknown_entry_is_empty() { + let graph = graph_with(vec![("a", script_node("a", "a.sh", None))], "a"); + + assert!(branch_subgraph(&graph, "nope").is_empty()); + } + + #[test] + fn branch_subgraph_skips_undeclared_targets() { + let mut a = script_node("a", "a.sh", Some("ghost_fallback")); + a.next = Some("ghost_next".to_string().into()); + let graph = graph_with(vec![("a", a)], "a"); + + assert_eq!(ids(&branch_subgraph(&graph, "a")), vec!["a"]); + } } diff --git a/src/parsers/bash.rs b/src/parsers/bash.rs index 1547e11..040a156 100644 --- a/src/parsers/bash.rs +++ b/src/parsers/bash.rs @@ -18,23 +18,32 @@ pub fn generate_bash_declarations( .with_context(|| format!("Failed to load script at '{tool_file:?}'"))?; debug!("Building script at '{tool_file:?}'"); + let (build_script, declarations) = build_bash_tool(&src, file_name) + .with_context(|| format!("Failed to build script '{file_name}'"))?; + function::write_file_atomic(tools_file_path, &build_script, Some(0o755)) + .with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?; + + Ok(declarations) +} + +/// Builds the argc script and derives its declarations without touching the +/// filesystem, so callers can inspect a tool source without rewriting it. +pub fn build_bash_tool(src: &str, file_name: &str) -> Result<(String, Vec)> { let build_script = argc::build( - &src, + src, "", env::var("TERM_WIDTH").ok().and_then(|v| v.parse().ok()), )?; let build_script = allow_empty_required_values(&build_script); - function::write_file_atomic(tools_file_path, &build_script, Some(0o755)) - .with_context(|| format!("Failed to write built script to '{tools_file_path:?}'"))?; let command_value = argc::export(&build_script, file_name) - .with_context(|| format!("Failed to parse script at '{tool_file:?}'"))?; + .with_context(|| format!("Failed to parse script '{file_name}'"))?; if command_value.subcommands.is_empty() { let function_declaration = command_to_function_declaration(&command_value).ok_or_else(|| { anyhow::format_err!("Tool definition missing or empty description: {file_name}") })?; - Ok(vec![function_declaration]) + Ok((build_script, vec![function_declaration])) } else { let mut declarations = vec![]; for subcommand in &command_value.subcommands { @@ -54,7 +63,7 @@ pub fn generate_bash_declarations( } } - Ok(declarations) + Ok((build_script, declarations)) } } diff --git a/src/utils/spinner.rs b/src/utils/spinner.rs index a38a6d6..1f81bb2 100644 --- a/src/utils/spinner.rs +++ b/src/utils/spinner.rs @@ -2,6 +2,11 @@ use super::{AbortSignal, IS_STDOUT_TERMINAL, poll_abort_signal, wait_abort_signa use anyhow::{Result, bail}; use crossterm::{cursor, queue, style, terminal}; +#[cfg(test)] +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; use std::{ future::Future, io::{Write, stdout}, @@ -19,11 +24,22 @@ use tokio::{ pub struct SpinnerInner { index: usize, message: String, + #[cfg(test)] + cleared: Option>, } impl SpinnerInner { const DATA: [&'static str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; + #[cfg(test)] + fn observed(cleared: Arc) -> Self { + Self { + index: 0, + message: String::new(), + cleared: Some(cleared), + } + } + fn step(&mut self) -> Result<()> { if !*IS_STDOUT_TERMINAL || self.message.is_empty() { return Ok(()); @@ -50,10 +66,17 @@ impl SpinnerInner { } fn clear_message(&mut self) -> Result<()> { - if !*IS_STDOUT_TERMINAL || self.message.is_empty() { + if self.message.is_empty() { return Ok(()); } self.message.clear(); + #[cfg(test)] + if let Some(cleared) = &self.cleared { + cleared.fetch_add(1, Ordering::SeqCst); + } + if !*IS_STDOUT_TERMINAL { + return Ok(()); + } let mut writer = stdout(); queue!( writer, @@ -66,6 +89,17 @@ impl SpinnerInner { } } +/// Restores the terminal when the spinner goes away by any route, not just the +/// loop's normal exit: a future cancelled mid-spin (the surrounding task was +/// aborted) or an early `?` return would otherwise leave the last frame on +/// screen and the cursor hidden. `clear_message` is a no-op once the message +/// is gone, so a normal-path clear leaves nothing for the drop to redo. +impl Drop for SpinnerInner { + fn drop(&mut self) { + let _ = self.clear_message(); + } +} + #[derive(Clone)] pub struct Spinner(mpsc::UnboundedSender); @@ -215,3 +249,65 @@ async fn run_abortable_spinner( spinner.clear_message()?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn observed(message: &str) -> (SpinnerInner, Arc) { + let cleared = Arc::new(AtomicUsize::new(0)); + let mut spinner = SpinnerInner::observed(cleared.clone()); + spinner.set_message(message.to_string()).unwrap(); + (spinner, cleared) + } + + #[test] + fn drop_with_active_message_clears_once() { + let (spinner, cleared) = observed("Thinking"); + assert_eq!(cleared.load(Ordering::SeqCst), 0); + drop(spinner); + assert_eq!(cleared.load(Ordering::SeqCst), 1); + } + + #[test] + fn explicit_clear_then_drop_clears_exactly_once() { + let (mut spinner, cleared) = observed("Thinking"); + spinner.clear_message().unwrap(); + assert_eq!(cleared.load(Ordering::SeqCst), 1); + drop(spinner); + assert_eq!(cleared.load(Ordering::SeqCst), 1); + } + + #[test] + fn drop_without_message_is_silent() { + let (spinner, cleared) = observed(""); + + drop(spinner); + + assert_eq!(cleared.load(Ordering::SeqCst), 0); + } + + #[test] + fn set_message_replaces_and_clears_previous_once() { + let (mut spinner, cleared) = observed("Thinking"); + spinner.set_message("Fetching".to_string()).unwrap(); + assert_eq!(cleared.load(Ordering::SeqCst), 1); + drop(spinner); + assert_eq!(cleared.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn cancelled_task_restores_terminal_on_drop() { + let (spinner, cleared) = observed("Thinking"); + let task = tokio::spawn(async move { + let _spinner = spinner; + std::future::pending::<()>().await; + }); + tokio::task::yield_now().await; + assert_eq!(cleared.load(Ordering::SeqCst), 0); + + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert_eq!(cleared.load(Ordering::SeqCst), 1); + } +}