From 1c68331738978544ddb11d6bbc2155450155b016 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Sat, 12 Sep 2026 11:44:38 -0600 Subject: [PATCH] feat(TASK-008): add shared transient matcher and retry support to agent nodes Hoist the llm node's transient-error matcher to graph::is_transient_error, extend it with broken-pipe / connection-error / error-sending-request signatures, and give AgentNode an opt-in max_attempts retry loop with llm-parity semantics: contexts applied before the transient check, fresh wall-clock budget per attempt, peer identity (assignment + registry) re-armed before each retry, retirement exactly once after the final attempt, and structured extraction outside the loop. --- Cargo.toml | 1 + src/graph/agent.rs | 399 ++++++++++++++++++++++++++++++++++++++--- src/graph/executor.rs | 1 + src/graph/llm.rs | 34 +--- src/graph/map.rs | 1 + src/graph/mod.rs | 54 ++++++ src/graph/types.rs | 42 ++++- src/graph/validator.rs | 1 + 8 files changed, 474 insertions(+), 59 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f0e05bf..c5368c7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -146,6 +146,7 @@ ctor = "1.0.13" pretty_assertions = "1.4.0" rmcp = { version = "3.1.2", features = ["server"] } serial_test = "3" +tokio = { version = "1.34.0", features = ["test-util"] } [[bin]] name = "coyote" diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 03b9830..85a6a0b 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -1,3 +1,4 @@ +use super::is_transient_error; use super::state::StateManager; use super::state_updates; use super::structured; @@ -5,10 +6,14 @@ 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 crate::supervisor::mailbox::{Inbox, PeerRegistry}; +use anyhow::{Context, Error, Result, anyhow}; +use log::{debug, warn}; use serde_json::Value; use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; use tokio::time::timeout; const DEFAULT_TIMEOUT_SECS: u64 = 300; @@ -34,31 +39,26 @@ impl AgentNodeExecutor { } let secs = node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); - let bound = wall_clock(secs); + let agent_name = node.agent.clone(); - // 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 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 @@ -77,6 +77,106 @@ impl AgentNodeExecutor { } } +/// 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 +/// solver ("implementation of `Send` is not general enough"). +type AttemptFuture<'a> = Pin> + Send + 'a>>; + +/// The per-attempt runner the retry loop drives. The attempt future may only +/// borrow the ctx it is handed, so runners move owned clones of everything +/// else into the future. +type AttemptRunner<'f> = dyn for<'a> FnMut(&'a mut RequestContext) -> AttemptFuture<'a> + Send + 'f; + +/// Identity funnel that pins a closure to the runner's higher-ranked +/// signature, so inline closures at call sites infer the right lifetimes. +fn attempt_runner(f: F) -> F +where + F: for<'a> FnMut(&'a mut RequestContext) -> AttemptFuture<'a> + Send, +{ + f +} + +fn boxed_attempt<'a>(fut: impl Future> + Send + 'a) -> AttemptFuture<'a> { + Box::pin(fut) +} + +/// Runs the node's agent up to `node.max_attempts` times, retrying only +/// failures that `is_transient_error` recognizes. +/// +/// `run_agent_for_graph` takes the peer identity off the ctx and never +/// restores it, so BOTH halves — the (id, inbox) assignment and the registry +/// — are captured up front and re-armed before every retry; re-arming only +/// the assignment would give a retried attempt its identity back but no +/// roster and a broken `agent__send_message`. A frontier peer is retired +/// exactly once, after the final attempt, the moment the agent stops — not +/// after extraction. Chains re-arm the same identity for later steps and +/// leave retirement to the chain runner. +async fn run_with_retries( + node: &AgentNode, + parent_ctx: &mut RequestContext, + retire_peer_on_return: bool, + run_attempt: &mut AttemptRunner<'_>, +) -> Result { + let assignment = parent_ctx.peer_assignment.clone(); + let registry = parent_ctx.peer_registry.clone(); + + let result = retry_transient(node, parent_ctx, &assignment, ®istry, run_attempt).await; + + if retire_peer_on_return && let (Some(registry), Some((id, _))) = (®istry, &assignment) { + registry.mark_finished(id); + } + + result +} + +async fn retry_transient( + node: &AgentNode, + parent_ctx: &mut RequestContext, + assignment: &Option<(String, Arc)>, + registry: &Option>, + run_attempt: &mut AttemptRunner<'_>, +) -> Result { + let mut last_err: Option = None; + for attempt in 1..=node.max_attempts { + if attempt > 1 { + if let Some(assignment) = assignment { + parent_ctx.peer_assignment = Some(assignment.clone()); + } + if let Some(registry) = registry { + parent_ctx.peer_registry = Some(Arc::clone(registry)); + } + } + match run_attempt(parent_ctx).await { + Ok(out) => return Ok(out), + Err(e) if is_transient_error(&e) && attempt < node.max_attempts => { + warn!("agent node attempt {attempt} failed (transient): {e:#}; retrying"); + last_err = Some(e); + } + Err(e) => return Err(e), + } + } + Err(last_err.unwrap_or_else(|| anyhow!("agent node exhausted retries"))) +} + +/// Applies the wall-clock bound and the human-readable failure contexts to a +/// single attempt. The contexts land here — before the caller inspects the +/// error — because `tokio::time::error::Elapsed` renders as "deadline has +/// elapsed", which the transient matcher would not recognize as a timeout. +async fn bounded_attempt( + agent_name: &str, + secs: u64, + fut: impl Future>, +) -> Result { + let raw_result = match wall_clock(secs) { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + }; + raw_result + .with_context(|| format!("Agent '{agent_name}' timed out after {secs}s"))? + .with_context(|| format!("Agent '{agent_name}' failed")) +} + /// 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 @@ -135,6 +235,7 @@ mod tests { state_updates: updates, output_schema: None, timeout: None, + max_attempts: 1, inputs: None, teammates: false, } @@ -184,6 +285,254 @@ mod tests { assert!(!registry.is_finished(&id)); } + fn plain_ctx() -> RequestContext { + RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd) + } + + fn retryable_node(max_attempts: u32) -> AgentNode { + let mut node = node_with("hi", None); + node.max_attempts = max_attempts; + node + } + + #[tokio::test] + async fn run_with_retries_retries_transient_failure_and_succeeds() { + let mut ctx = plain_ctx(); + let node = retryable_node(2); + let mut attempts = 0u32; + + let out = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + let result = if attempts == 1 { + Err(anyhow!("error sending request for url")) + } else { + Ok("recovered".to_string()) + }; + boxed_attempt(async move { result }) + }), + ) + .await + .unwrap(); + + assert_eq!(out, "recovered"); + assert_eq!(attempts, 2); + } + + #[tokio::test] + async fn run_with_retries_does_not_retry_non_transient_failures() { + let mut ctx = plain_ctx(); + let node = retryable_node(3); + let mut attempts = 0u32; + + let err = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + boxed_attempt(async move { Err::(anyhow!("Unknown model 'foo'")) }) + }), + ) + .await + .expect_err("non-transient failure must propagate immediately"); + + assert_eq!(attempts, 1); + assert!(format!("{err:#}").contains("Unknown model 'foo'")); + } + + #[tokio::test] + async fn run_with_retries_stops_at_max_attempts_and_keeps_failure_context() { + let mut ctx = plain_ctx(); + let node = retryable_node(2); + let mut attempts = 0u32; + + let err = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + boxed_attempt(bounded_attempt("test_agent", 0, async { + Err(anyhow!("connection error: unexpected end of stream")) + })) + }), + ) + .await + .expect_err("exhausted retries must propagate the last error"); + + assert_eq!(attempts, 2); + let chain = format!("{err:#}"); + assert!(chain.contains("Agent 'test_agent' failed"), "{chain}"); + assert!(chain.contains("connection error"), "{chain}"); + } + + #[tokio::test(start_paused = true)] + async fn run_with_retries_retries_a_timeout_and_succeeds() { + let mut ctx = plain_ctx(); + let node = retryable_node(2); + let mut attempts = 0u32; + + let out = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + if attempts == 1 { + boxed_attempt(bounded_attempt( + "test_agent", + 1, + std::future::pending::>(), + )) + } else { + boxed_attempt(async move { Ok("recovered".to_string()) }) + } + }), + ) + .await + .unwrap(); + + assert_eq!(out, "recovered"); + assert_eq!(attempts, 2); + } + + #[tokio::test(start_paused = true)] + async fn bounded_attempt_timeout_context_is_transient() { + let err = bounded_attempt("test_agent", 1, std::future::pending::>()) + .await + .expect_err("pending future must hit the wall clock"); + + let chain = format!("{err:#}"); + assert!( + chain.contains("Agent 'test_agent' timed out after 1s"), + "{chain}" + ); + assert!(is_transient_error(&err)); + } + + #[tokio::test] + async fn run_with_retries_rearms_the_same_peer_identity_for_each_retry() { + let (mut ctx, registry, id) = ctx_at_max_depth_with_peer(); + let expected_inbox = Arc::clone(&ctx.peer_assignment.as_ref().unwrap().1); + let node = retryable_node(2); + let mut attempts = 0u32; + let mut observations: Vec<(String, bool, bool)> = Vec::new(); + + let out = run_with_retries( + &node, + &mut ctx, + true, + &mut attempt_runner(|ctx| { + attempts += 1; + // Mirror run_agent_for_graph: consume both identity halves. + let (peer_id, inbox) = ctx.peer_assignment.take().expect("assignment armed"); + let reg = ctx.peer_registry.take().expect("registry armed"); + observations.push(( + peer_id, + Arc::ptr_eq(&inbox, &expected_inbox), + Arc::ptr_eq(®, ®istry), + )); + assert!( + !registry.is_finished(&id), + "peer must not be retired between attempts" + ); + let result = if attempts == 1 { + Err(anyhow!("stream closed because of a broken pipe")) + } else { + Ok("done".to_string()) + }; + boxed_attempt(async move { result }) + }), + ) + .await + .unwrap(); + + assert_eq!(out, "done"); + assert_eq!(observations.len(), 2); + for (peer_id, same_inbox, same_registry) in &observations { + assert_eq!(peer_id, &id); + assert!(same_inbox, "retry must reuse the same inbox Arc"); + assert!(same_registry, "retry must reuse the same registry Arc"); + } + assert!(registry.is_finished(&id)); + } + + #[tokio::test] + async fn run_with_retries_retires_frontier_peer_after_the_final_failed_attempt() { + let (mut ctx, registry, id) = ctx_at_max_depth_with_peer(); + let node = retryable_node(2); + let mut attempts = 0u32; + + let err = run_with_retries( + &node, + &mut ctx, + true, + &mut attempt_runner(|ctx| { + attempts += 1; + ctx.peer_assignment.take(); + ctx.peer_registry.take(); + assert!( + !registry.is_finished(&id), + "peer must not be retired between attempts" + ); + boxed_attempt(async move { Err::(anyhow!("Connection reset by peer")) }) + }), + ) + .await + .expect_err("both attempts fail"); + + assert_eq!(attempts, 2); + assert!(registry.is_finished(&id)); + assert!(format!("{err:#}").contains("Connection reset")); + } + + #[tokio::test] + async fn run_with_retries_leaves_peer_live_when_not_retiring() { + let (mut ctx, registry, id) = ctx_at_max_depth_with_peer(); + let node = retryable_node(1); + + run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|ctx| { + ctx.peer_assignment.take(); + ctx.peer_registry.take(); + boxed_attempt(async move { Ok("done".to_string()) }) + }), + ) + .await + .unwrap(); + + assert!(!registry.is_finished(&id)); + } + + #[tokio::test] + async fn run_with_retries_zero_attempts_reports_exhausted_retries() { + let mut ctx = plain_ctx(); + let node = retryable_node(0); + let mut attempts = 0u32; + + let err = run_with_retries( + &node, + &mut ctx, + false, + &mut attempt_runner(|_ctx| { + attempts += 1; + boxed_attempt(async move { Ok("never".to_string()) }) + }), + ) + .await + .expect_err("zero attempts cannot succeed"); + + assert_eq!(attempts, 0); + assert!(format!("{err:#}").contains("agent node exhausted retries")); + } + 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 127aa95..1367c88 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -827,6 +827,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + max_attempts: 1, inputs: None, teammates, }), diff --git a/src/graph/llm.rs b/src/graph/llm.rs index a409469..7623c27 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -2,7 +2,7 @@ use super::state::StateManager; use super::state_updates; use super::structured; use super::types::LlmNode; -use super::wall_clock; +use super::{is_transient_error, wall_clock}; use crate::client::{Model, ModelType, call_chat_completions}; use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ @@ -250,7 +250,7 @@ async fn run_with_retries( for attempt in 1..=node.max_attempts { match run_chat_loop(node, prompt, ctx, abort).await { Ok(out) => return Ok(out), - Err(e) if is_transient(&e) && attempt < node.max_attempts => { + Err(e) if is_transient_error(&e) && attempt < node.max_attempts => { warn!("llm node attempt {attempt} failed (transient): {e}; retrying"); last_err = Some(e); } @@ -449,16 +449,6 @@ fn validate_tools_subset( Ok(()) } -fn is_transient(err: &Error) -> bool { - let s = format!("{err:#}"); - s.contains("timed out") - || s.contains("rate limit") - || s.contains("429") - || s.contains("Connection reset") - || s.contains("Connection refused") - || s.contains("produced no output") -} - fn apply_state_updates_with_output( node: &LlmNode, state_manager: &mut StateManager, @@ -724,26 +714,6 @@ mod tests { assert!(mcp.is_empty()); } - #[test] - fn is_transient_matches_expected_signatures() { - assert!(is_transient(&anyhow!("request timed out after 30s"))); - assert!(is_transient(&anyhow!("rate limit reached"))); - assert!(is_transient(&anyhow!("429 too many requests"))); - assert!(is_transient(&anyhow!("Connection reset by peer"))); - assert!(is_transient(&anyhow!("Connection refused"))); - assert!(is_transient(&anyhow!("llm produced no output"))); - } - - #[test] - fn is_transient_rejects_non_transient_errors() { - assert!(!is_transient(&anyhow!("Unknown model 'foo'"))); - assert!(!is_transient(&anyhow!( - "llm node references unknown tool 'bad'" - ))); - 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()); diff --git a/src/graph/map.rs b/src/graph/map.rs index 29ad7de..6d74543 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -464,6 +464,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + max_attempts: 1, inputs: None, teammates, }) diff --git a/src/graph/mod.rs b/src/graph/mod.rs index 3b1d72d..1d7ae9e 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -34,6 +34,22 @@ pub(crate) fn wall_clock(secs: u64) -> Option { (secs != 0).then(|| Duration::from_secs(secs)) } +/// Whether an error looks like a transient transport/provider failure worth +/// retrying. Providers and HTTP stacks surface these as rendered text rather +/// than typed errors, so the whole context chain is substring-matched. +pub(crate) fn is_transient_error(err: &anyhow::Error) -> bool { + let s = format!("{err:#}"); + s.contains("timed out") + || s.contains("rate limit") + || s.contains("429") + || s.contains("Connection reset") + || s.contains("Connection refused") + || s.contains("produced no output") + || s.contains("broken pipe") + || s.contains("connection error") + || s.contains("error sending request") +} + pub const MAX_STATE_SIZE_BYTES: usize = 32 * 1024; pub(in crate::graph) fn type_name(value: &Value) -> &'static str { @@ -50,6 +66,7 @@ pub(in crate::graph) fn type_name(value: &Value) -> &'static str { #[cfg(test)] mod tests { use super::*; + use anyhow::anyhow; #[test] fn wall_clock_zero_is_no_bound() { @@ -60,4 +77,41 @@ mod tests { fn wall_clock_nonzero_is_that_many_seconds() { assert_eq!(wall_clock(7), Some(Duration::from_secs(7))); } + + #[test] + fn is_transient_error_matches_expected_signatures() { + assert!(is_transient_error(&anyhow!("request timed out after 30s"))); + assert!(is_transient_error(&anyhow!("rate limit reached"))); + assert!(is_transient_error(&anyhow!("429 too many requests"))); + assert!(is_transient_error(&anyhow!("Connection reset by peer"))); + assert!(is_transient_error(&anyhow!("Connection refused"))); + assert!(is_transient_error(&anyhow!("llm produced no output"))); + assert!(is_transient_error(&anyhow!( + "stream closed because of a broken pipe" + ))); + assert!(is_transient_error(&anyhow!( + "connection error: unexpected end of stream" + ))); + assert!(is_transient_error(&anyhow!( + "error sending request for url" + ))); + } + + #[test] + fn is_transient_error_sees_through_context_chains() { + let err = anyhow!("stream closed because of a broken pipe") + .context("Failed to call chat-completions api") + .context("Agent 'domain-reviewer' failed"); + assert!(is_transient_error(&err)); + } + + #[test] + fn is_transient_error_rejects_non_transient_errors() { + assert!(!is_transient_error(&anyhow!("Unknown model 'foo'"))); + assert!(!is_transient_error(&anyhow!( + "llm node references unknown tool 'bad'" + ))); + assert!(!is_transient_error(&anyhow!("hit max_iterations"))); + assert!(!is_transient_error(&anyhow!("authentication failed"))); + } } diff --git a/src/graph/types.rs b/src/graph/types.rs index 6e00d04..b69d442 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 = "default_max_attempts")] + pub max_attempts: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] pub inputs: Option>, @@ -321,7 +324,7 @@ pub struct LlmNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub fallback: Option, - #[serde(default = "default_llm_max_attempts")] + #[serde(default = "default_max_attempts")] pub max_attempts: u32, #[serde(default = "default_llm_max_iterations")] @@ -349,7 +352,7 @@ pub struct LlmNode { pub skill_instructions: Option, } -fn default_llm_max_attempts() -> u32 { +fn default_max_attempts() -> u32 { 1 } @@ -640,6 +643,41 @@ nodes: } } + #[test] + fn agent_node_max_attempts_parses_and_defaults_to_one() { + let yaml = r#" +name: g +start: a +nodes: + a: + id: a + type: agent + agent: helper + prompt: hi + max_attempts: 2 + 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.max_attempts, 2), + _ => panic!("expected Agent variant"), + } + match &graph.get_node("b").unwrap().node_type { + NodeType::Agent(n) => assert_eq!(n.max_attempts, 1), + _ => 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 08188fe..c190b1f 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -1924,6 +1924,7 @@ mod tests { state_updates: None, output_schema: None, timeout: None, + max_attempts: 1, inputs: None, teammates: false, }),