diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 85a6a0b..974d8b8 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -18,65 +18,111 @@ use tokio::time::timeout; const DEFAULT_TIMEOUT_SECS: u64 = 300; +#[derive(Debug, PartialEq, Eq)] +pub(super) enum AgentExecutionOutcome { + Continue(String), + FellBack(String), +} + pub struct AgentNodeExecutor; impl AgentNodeExecutor { - pub async fn execute( + pub(super) async fn execute( 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 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 secs = node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); - let agent_name = node.agent.clone(); - - let mut run_attempt = attempt_runner(move |ctx| { - let agent_name = agent_name.clone(); - let prompt = prompt.clone(); - let graph_inputs = graph_inputs.clone(); - boxed_attempt(async move { - bounded_attempt( - &agent_name, - secs, - run_agent_for_graph(ctx, &agent_name, &prompt, graph_inputs), - ) - .await - }) - }); - let raw = - run_with_retries(node, parent_ctx, retire_peer_on_return, &mut run_attempt).await?; - - // Extraction stays outside the retry loop: by this point the agent - // itself succeeded, so an extraction failure is never retried. - let output_value = match &node.output_schema { - Some(schema) => structured::extract(&raw, schema, parent_ctx) - .await - .with_context(|| { - format!( - "Agent '{}' output failed structured-output extraction", - node.agent - ) - })?, - None => Value::String(raw.clone()), - }; - - apply_state_updates(node, state_manager, &output_value); - - Ok(raw) + ) -> Result { + let result = run(node, state_manager, parent_ctx, retire_peer_on_return).await; + outcome_from(node, state_manager, result) } } +/// Turns the node run's final result into a routing outcome. With a +/// `fallback` declared, any failure — retries exhausted, hard error, or +/// extraction failure — is written into state as +/// `"Agent node failed: "` (through the node's own state_updates / +/// output-schema path, so the fallback node can interpolate it) and routes +/// to the fallback. Without one the error propagates unchanged; unlike llm +/// nodes there is no teaching bail, because agent nodes always failed loudly +/// and callers assert on the existing "Agent 'X' failed" chains. +fn outcome_from( + node: &AgentNode, + state_manager: &mut StateManager, + result: Result, +) -> Result { + match result { + Ok(raw) => Ok(AgentExecutionOutcome::Continue(raw)), + Err(e) => match &node.fallback { + Some(fb) => { + warn!("agent node failed, routing to fallback '{fb}': {e:#}"); + state_updates::apply( + state_manager, + &Value::String(format!("Agent node failed: {e:#}")), + node.output_schema.is_some(), + node.state_updates.as_ref(), + ); + Ok(AgentExecutionOutcome::FellBack(fb.clone())) + } + None => Err(e), + }, + } +} + +async fn run( + 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 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 secs = node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); + let agent_name = node.agent.clone(); + + let mut run_attempt = attempt_runner(move |ctx| { + let agent_name = agent_name.clone(); + let prompt = prompt.clone(); + let graph_inputs = graph_inputs.clone(); + boxed_attempt(async move { + bounded_attempt( + &agent_name, + secs, + run_agent_for_graph(ctx, &agent_name, &prompt, graph_inputs), + ) + .await + }) + }); + let raw = run_with_retries(node, parent_ctx, retire_peer_on_return, &mut run_attempt).await?; + + // Extraction stays outside the retry loop: by this point the agent + // itself succeeded, so an extraction failure is never retried. + let output_value = match &node.output_schema { + Some(schema) => structured::extract(&raw, schema, parent_ctx) + .await + .with_context(|| { + format!( + "Agent '{}' output failed structured-output extraction", + node.agent + ) + })?, + None => Value::String(raw.clone()), + }; + + apply_state_updates(node, state_manager, &output_value); + + Ok(raw) +} + /// One boxed attempt against the ctx. Boxed (rather than an opaque /// `AsyncFnMut` future) because agent nodes run inside `tokio::spawn`ed map /// branches, where higher-ranked opaque futures trip the `Send` auto-trait @@ -236,6 +282,7 @@ mod tests { output_schema: None, timeout: None, max_attempts: 1, + fallback: None, inputs: None, teammates: false, } @@ -533,6 +580,114 @@ mod tests { assert!(format!("{err:#}").contains("agent node exhausted retries")); } + fn failure_capture_updates() -> HashMap { + let mut u = HashMap::new(); + u.insert("failure".into(), "{{output}}".into()); + u + } + + #[tokio::test] + async fn execute_failure_with_fallback_routes_and_records_the_failure() { + let mut ctx = plain_ctx(); + ctx.current_depth = default_max_agent_depth(); + let mut node = node_with("hi", Some(failure_capture_updates())); + node.fallback = Some("recover".into()); + let mut state = manager_with(&[]); + + let outcome = AgentNodeExecutor::execute(&node, &mut state, &mut ctx, false) + .await + .unwrap(); + + assert_eq!(outcome, AgentExecutionOutcome::FellBack("recover".into())); + let failure = state + .state() + .get("failure") + .and_then(Value::as_str) + .expect("failure text must land in state for the fallback node") + .to_string(); + assert!(failure.starts_with("Agent node failed: "), "{failure}"); + assert!(failure.contains("Agent 'test_agent' failed"), "{failure}"); + assert!(failure.contains("Max agent depth exceeded"), "{failure}"); + } + + fn extraction_error() -> Error { + anyhow!("no JSON object found in output") + .context("Agent 'test_agent' output failed structured-output extraction") + } + + #[test] + fn outcome_from_success_is_continue() { + let node = node_with("hi", None); + let mut state = manager_with(&[]); + + let outcome = outcome_from(&node, &mut state, Ok("raw".into())).unwrap(); + + assert_eq!(outcome, AgentExecutionOutcome::Continue("raw".into())); + } + + #[test] + fn outcome_from_extraction_failure_with_fallback_falls_back() { + let mut node = node_with("hi", Some(failure_capture_updates())); + node.fallback = Some("recover".into()); + let mut state = manager_with(&[]); + + let outcome = outcome_from(&node, &mut state, Err(extraction_error())).unwrap(); + + assert_eq!(outcome, AgentExecutionOutcome::FellBack("recover".into())); + let failure = state + .state() + .get("failure") + .and_then(Value::as_str) + .unwrap(); + assert!( + failure.contains("structured-output extraction"), + "{failure}" + ); + assert!(failure.contains("no JSON object found"), "{failure}"); + } + + #[test] + fn outcome_from_failure_without_fallback_propagates_the_error_unchanged() { + let node = node_with("hi", Some(failure_capture_updates())); + let mut state = manager_with(&[]); + + let err = outcome_from(&node, &mut state, Err(extraction_error())) + .expect_err("no fallback must propagate"); + + assert_eq!( + format!("{err:#}"), + "Agent 'test_agent' output failed structured-output extraction: \ + no JSON object found in output" + ); + assert!(state.state().get("failure").is_none()); + } + + #[tokio::test] + async fn fallback_engages_only_after_transient_retries_exhaust() { + let mut ctx = plain_ctx(); + let mut node = retryable_node(2); + node.fallback = Some("recover".into()); + let mut attempts = 0u32; + + let result = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + boxed_attempt(async move { + Err::(anyhow!("connection error: reset by peer")) + }) + }), + ) + .await; + let mut state = manager_with(&[]); + let outcome = outcome_from(&node, &mut state, result).unwrap(); + + assert_eq!(attempts, 2, "fallback must not preempt the retry budget"); + assert_eq!(outcome, AgentExecutionOutcome::FellBack("recover".into())); + } + fn node_with_inputs(pairs: &[(&str, &str)]) -> AgentNode { let mut node = node_with("hi", None); node.inputs = Some( diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 1367c88..805516d 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -1,4 +1,4 @@ -use super::agent::AgentNodeExecutor; +use super::agent::{AgentExecutionOutcome, AgentNodeExecutor}; use super::llm::{LlmExecutionOutcome, LlmNodeExecutor}; use super::logging::{GraphLogger, narrate_node_complete, narrate_node_failed}; use super::map::MapNodeExecutor; @@ -574,8 +574,14 @@ pub(super) async fn step( ) -> Result { match &node.node_type { NodeType::Agent(agent_node) => { - AgentNodeExecutor::execute(agent_node, state, ctx, !step_ctx.branch_mode).await?; - let targets = static_next_targets(node, current, "agent", step_ctx.branch_mode)?; + let outcome = + AgentNodeExecutor::execute(agent_node, state, ctx, !step_ctx.branch_mode).await?; + let targets = match outcome { + AgentExecutionOutcome::Continue(_) => { + static_next_targets(node, current, "agent", step_ctx.branch_mode)? + } + AgentExecutionOutcome::FellBack(target) => vec![target], + }; Ok(StepResult::Continue(targets)) } NodeType::Script(script_node) => { @@ -828,6 +834,7 @@ mod tests { output_schema: None, timeout: None, max_attempts: 1, + fallback: None, inputs: None, teammates, }), diff --git a/src/graph/map.rs b/src/graph/map.rs index 6d74543..2f982fc 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -465,6 +465,7 @@ mod tests { output_schema: None, timeout: None, max_attempts: 1, + fallback: None, inputs: None, teammates, }) @@ -1609,6 +1610,55 @@ nodes: assert_eq!(collected(&result), vec![json!("fb")]); } + #[tokio::test] + async fn map_chain_agent_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": "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: agent + agent: no-such-agent + 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!("recovered"), json!("recovered")] + ); + } + struct Harness { ws: TestWorkspace, graph: Arc, diff --git a/src/graph/types.rs b/src/graph/types.rs index b69d442..6a8c02f 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -243,6 +243,9 @@ pub struct AgentNode { #[serde(default = "default_max_attempts")] pub max_attempts: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub fallback: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub inputs: Option>, @@ -678,6 +681,41 @@ nodes: } } + #[test] + fn agent_node_fallback_parses_and_defaults_to_none() { + let yaml = r#" +name: g +start: a +nodes: + a: + id: a + type: agent + agent: helper + prompt: hi + fallback: b + next: b + b: + id: b + type: agent + agent: helper + prompt: hi + next: e + e: + id: e + type: end + output: done +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + match &graph.get_node("a").unwrap().node_type { + NodeType::Agent(n) => assert_eq!(n.fallback.as_deref(), Some("b")), + _ => panic!("expected Agent variant"), + } + match &graph.get_node("b").unwrap().node_type { + NodeType::Agent(n) => assert!(n.fallback.is_none()), + _ => panic!("expected Agent variant"), + } + } + #[test] fn agent_node_inputs_parse_optionally_and_empty_map_is_some() { let yaml = r#" diff --git a/src/graph/validator.rs b/src/graph/validator.rs index c190b1f..0fb3fde 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -719,6 +719,7 @@ impl GraphValidator { .filter_map(|(_, n)| match &n.node_type { NodeType::Script(s) => s.fallback.as_deref(), NodeType::Llm(l) => l.fallback.as_deref(), + NodeType::Agent(a) => a.fallback.as_deref(), _ => None, }) .collect(); @@ -982,12 +983,17 @@ fn declared_targets(node: &Node) -> Vec<(String, &'static str)> { out.push((t.clone(), "llm 'fallback'")); } } + NodeType::Agent(a) => { + if let Some(t) = &a.fallback { + out.push((t.clone(), "agent 'fallback'")); + } + } NodeType::Map(m) => { out.push((m.branch.clone(), "map 'branch'")); } - // `agent`/`input`/`rag` route only via `next` (already collected - // above); `end` is terminal. No type-specific routing edges to add. - NodeType::Agent(_) | NodeType::Input(_) | NodeType::Rag(_) | NodeType::End(_) => {} + // `input`/`rag` route only via `next` (already collected above); + // `end` is terminal. No type-specific routing edges to add. + NodeType::Input(_) | NodeType::Rag(_) | NodeType::End(_) => {} } out } @@ -1056,8 +1062,8 @@ fn main_flow_reachable(graph: &Graph) -> HashSet { } /// 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 +/// time: `next` targets and script/llm/agent `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(); @@ -1082,6 +1088,7 @@ pub(super) fn branch_subgraph(graph: &Graph, entry: &str) -> HashSet { match &node.node_type { NodeType::Script(s) => edges.extend(s.fallback.as_ref()), NodeType::Llm(l) => edges.extend(l.fallback.as_ref()), + NodeType::Agent(a) => edges.extend(a.fallback.as_ref()), _ => {} } for next in edges { @@ -1925,6 +1932,7 @@ mod tests { output_schema: None, timeout: None, max_attempts: 1, + fallback: None, inputs: None, teammates: false, }), @@ -2018,6 +2026,51 @@ mod tests { ); } + #[test] + fn flags_missing_agent_fallback_target() { + let mut a = agent_node("a", "helper", Some("end")); + if let NodeType::Agent(ref mut an) = a.node_type { + an.fallback = Some("ghost".into()); + } + let graph = graph_with(vec![("a", a), ("end", end_node("end"))], "a"); + + let result = validator().validate(&graph); + + assert!(!result.is_valid()); + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("non-existent node 'ghost'") + && e.message.contains("agent 'fallback'")) + ); + } + + #[test] + fn agent_fallback_target_is_not_flagged_unreachable() { + let mut a = agent_node("a", "helper", Some("end")); + if let NodeType::Agent(ref mut an) = a.node_type { + an.fallback = Some("recover".into()); + } + let recover = agent_node("recover", "helper", Some("end")); + let graph = graph_with( + vec![("a", a), ("recover", recover), ("end", end_node("end"))], + "a", + ); + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|w| w.node_id.as_deref() == Some("recover") + && w.message.contains("unreachable")), + "fallback target 'recover' incorrectly marked unreachable: {:?}", + result.warnings + ); + } + #[test] fn detects_two_node_cycle() { let mut a = end_node("a"); @@ -2857,6 +2910,62 @@ mod tests { ); } + #[test] + fn map_branch_agent_fallback_escaping_the_branch_gets_hint() { + let map = map_node_basic("m", "br", Some("end")); + let mut branch = agent_node("br", "helper", None); + if let NodeType::Agent(ref mut an) = branch.node_type { + an.fallback = Some("end".into()); + } + 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_agent_fallback_to_branch_local_node_passes() { + let map = map_node_basic("m", "br", Some("end")); + let mut branch = agent_node("br", "helper", None); + if let NodeType::Agent(ref mut an) = branch.node_type { + an.fallback = Some("recover".into()); + } + let recover = agent_node("recover", "helper", None); + let graph = graph_with( + vec![ + ("m", map), + ("br", branch), + ("recover", recover), + ("end", end_node("end")), + ], + "m", + ); + + let result = validator().validate(&graph); + + // Agent-existence errors on 'br'/'recover' are expected (no agent + // dirs in the test cwd); branch-locality errors land on the map. + assert!( + !result + .errors + .iter() + .any(|e| e.node_id.as_deref() == Some("m")), + "a branch-local agent fallback is allowed: {:?}", + result.errors + ); + } + #[test] fn map_branch_deep_approval_node_errors_on_the_map() { let map = map_node_basic("m", "br", Some("end"));