feat(TASK-009): add fallback routing to agent nodes

Agent nodes gain an optional fallback target with llm-node parity: on any
failure (retries exhausted, hard error, or structured-extraction failure)
the failure chain is written into state as 'Agent node failed: <chain>'
and the chain routes to the fallback node. Without a fallback the error
propagates unchanged. Validator collects agent fallback edges for
missing-target checks, reachability, and map-branch locality.
This commit is contained in:
Dark-Alex-17 committed 2026-09-14 10:50:33 -06:00
1 parent 1c68331738
commit be0c5462ed
5 files changed
+416 -57

No files matched your search

+204 -49
View File
@@ -18,65 +18,111 @@ use tokio::time::timeout;
const DEFAULT_TIMEOUT_SECS: u64 = 300; const DEFAULT_TIMEOUT_SECS: u64 = 300;
#[derive(Debug, PartialEq, Eq)]
pub(super) enum AgentExecutionOutcome {
Continue(String),
FellBack(String),
}
pub struct AgentNodeExecutor; pub struct AgentNodeExecutor;
impl AgentNodeExecutor { impl AgentNodeExecutor {
pub async fn execute( pub(super) async fn execute(
node: &AgentNode, node: &AgentNode,
state_manager: &mut StateManager, state_manager: &mut StateManager,
parent_ctx: &mut RequestContext, parent_ctx: &mut RequestContext,
retire_peer_on_return: bool, retire_peer_on_return: bool,
) -> Result<String> { ) -> Result<AgentExecutionOutcome> {
let prompt = state_manager let result = run(node, state_manager, parent_ctx, retire_peer_on_return).await;
.interpolate(&node.prompt) outcome_from(node, state_manager, result)
.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)
} }
} }
/// 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: <chain>"` (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<String>,
) -> Result<AgentExecutionOutcome> {
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<String> {
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 /// One boxed attempt against the ctx. Boxed (rather than an opaque
/// `AsyncFnMut` future) because agent nodes run inside `tokio::spawn`ed map /// `AsyncFnMut` future) because agent nodes run inside `tokio::spawn`ed map
/// branches, where higher-ranked opaque futures trip the `Send` auto-trait /// branches, where higher-ranked opaque futures trip the `Send` auto-trait
@@ -236,6 +282,7 @@ mod tests {
output_schema: None, output_schema: None,
timeout: None, timeout: None,
max_attempts: 1, max_attempts: 1,
fallback: None,
inputs: None, inputs: None,
teammates: false, teammates: false,
} }
@@ -533,6 +580,114 @@ mod tests {
assert!(format!("{err:#}").contains("agent node exhausted retries")); assert!(format!("{err:#}").contains("agent node exhausted retries"));
} }
fn failure_capture_updates() -> HashMap<String, String> {
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::<String, _>(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 { fn node_with_inputs(pairs: &[(&str, &str)]) -> AgentNode {
let mut node = node_with("hi", None); let mut node = node_with("hi", None);
node.inputs = Some( node.inputs = Some(
+10 -3
View File
@@ -1,4 +1,4 @@
use super::agent::AgentNodeExecutor; use super::agent::{AgentExecutionOutcome, AgentNodeExecutor};
use super::llm::{LlmExecutionOutcome, LlmNodeExecutor}; use super::llm::{LlmExecutionOutcome, LlmNodeExecutor};
use super::logging::{GraphLogger, narrate_node_complete, narrate_node_failed}; use super::logging::{GraphLogger, narrate_node_complete, narrate_node_failed};
use super::map::MapNodeExecutor; use super::map::MapNodeExecutor;
@@ -574,8 +574,14 @@ pub(super) async fn step(
) -> Result<StepResult> { ) -> Result<StepResult> {
match &node.node_type { match &node.node_type {
NodeType::Agent(agent_node) => { NodeType::Agent(agent_node) => {
AgentNodeExecutor::execute(agent_node, state, ctx, !step_ctx.branch_mode).await?; let outcome =
let targets = static_next_targets(node, current, "agent", step_ctx.branch_mode)?; 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)) Ok(StepResult::Continue(targets))
} }
NodeType::Script(script_node) => { NodeType::Script(script_node) => {
@@ -828,6 +834,7 @@ mod tests {
output_schema: None, output_schema: None,
timeout: None, timeout: None,
max_attempts: 1, max_attempts: 1,
fallback: None,
inputs: None, inputs: None,
teammates, teammates,
}), }),
+50
View File
@@ -465,6 +465,7 @@ mod tests {
output_schema: None, output_schema: None,
timeout: None, timeout: None,
max_attempts: 1, max_attempts: 1,
fallback: None,
inputs: None, inputs: None,
teammates, teammates,
}) })
@@ -1609,6 +1610,55 @@ nodes:
assert_eq!(collected(&result), vec![json!("fb")]); 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 { struct Harness {
ws: TestWorkspace, ws: TestWorkspace,
graph: Arc<Graph>, graph: Arc<Graph>,
+38
View File
@@ -243,6 +243,9 @@ pub struct AgentNode {
#[serde(default = "default_max_attempts")] #[serde(default = "default_max_attempts")]
pub max_attempts: u32, pub max_attempts: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub fallback: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub inputs: Option<HashMap<String, String>>, pub inputs: Option<HashMap<String, String>>,
@@ -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] #[test]
fn agent_node_inputs_parse_optionally_and_empty_map_is_some() { fn agent_node_inputs_parse_optionally_and_empty_map_is_some() {
let yaml = r#" let yaml = r#"
+114 -5
View File
@@ -719,6 +719,7 @@ impl GraphValidator {
.filter_map(|(_, n)| match &n.node_type { .filter_map(|(_, n)| match &n.node_type {
NodeType::Script(s) => s.fallback.as_deref(), NodeType::Script(s) => s.fallback.as_deref(),
NodeType::Llm(l) => l.fallback.as_deref(), NodeType::Llm(l) => l.fallback.as_deref(),
NodeType::Agent(a) => a.fallback.as_deref(),
_ => None, _ => None,
}) })
.collect(); .collect();
@@ -982,12 +983,17 @@ fn declared_targets(node: &Node) -> Vec<(String, &'static str)> {
out.push((t.clone(), "llm 'fallback'")); out.push((t.clone(), "llm 'fallback'"));
} }
} }
NodeType::Agent(a) => {
if let Some(t) = &a.fallback {
out.push((t.clone(), "agent 'fallback'"));
}
}
NodeType::Map(m) => { NodeType::Map(m) => {
out.push((m.branch.clone(), "map 'branch'")); out.push((m.branch.clone(), "map 'branch'"));
} }
// `agent`/`input`/`rag` route only via `next` (already collected // `input`/`rag` route only via `next` (already collected above);
// above); `end` is terminal. No type-specific routing edges to add. // `end` is terminal. No type-specific routing edges to add.
NodeType::Agent(_) | NodeType::Input(_) | NodeType::Rag(_) | NodeType::End(_) => {} NodeType::Input(_) | NodeType::Rag(_) | NodeType::End(_) => {}
} }
out out
} }
@@ -1056,8 +1062,8 @@ fn main_flow_reachable(graph: &Graph) -> HashSet<String> {
} }
/// Nodes reachable from `entry` over the edges a chain can follow at run /// 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 /// time: `next` targets and script/llm/agent `fallback`s. Approval routes
/// nested map's `branch` are not followed. Neither may appear inside a /// and a nested map's `branch` are not followed. Neither may appear inside a
/// branch, and the validator reports them separately. /// branch, and the validator reports them separately.
pub(super) fn branch_subgraph(graph: &Graph, entry: &str) -> HashSet<String> { pub(super) fn branch_subgraph(graph: &Graph, entry: &str) -> HashSet<String> {
let mut reachable: HashSet<String> = HashSet::new(); let mut reachable: HashSet<String> = HashSet::new();
@@ -1082,6 +1088,7 @@ pub(super) fn branch_subgraph(graph: &Graph, entry: &str) -> HashSet<String> {
match &node.node_type { match &node.node_type {
NodeType::Script(s) => edges.extend(s.fallback.as_ref()), NodeType::Script(s) => edges.extend(s.fallback.as_ref()),
NodeType::Llm(l) => edges.extend(l.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 { for next in edges {
@@ -1925,6 +1932,7 @@ mod tests {
output_schema: None, output_schema: None,
timeout: None, timeout: None,
max_attempts: 1, max_attempts: 1,
fallback: None,
inputs: None, inputs: None,
teammates: false, 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] #[test]
fn detects_two_node_cycle() { fn detects_two_node_cycle() {
let mut a = end_node("a"); 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] #[test]
fn map_branch_deep_approval_node_errors_on_the_map() { fn map_branch_deep_approval_node_errors_on_the_map() {
let map = map_node_basic("m", "br", Some("end")); let map = map_node_basic("m", "br", Some("end"));