diff --git a/README.md b/README.md index 13ee8c1..198fada 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,8 @@ All live in [`examples/`](examples/): - **`basic_graph`** — a minimal typed state graph: `START`, nodes, edges, `END`. - **`complex_graph`** — conditional routing, fanout, and richer topology. - **`durable_graph`** — checkpoints, resume, and time-travel over supersteps. +- **`resilient_graph`** — node-level retry over transient failures, plus a + resumable failure checkpoint that `retry` restarts after an outage clears. - **`agent_loop_tools`** — the agent ↔ tool loop the harness runs. - **`orchestrator_subagents`** — **recursion in action:** an orchestrator agent that calls sub-agents as tools, with depth tracking and rolled-up usage. diff --git a/docs/modules/graph/fault-tolerance.md b/docs/modules/graph/fault-tolerance.md index 0fcca22..df16c3f 100644 --- a/docs/modules/graph/fault-tolerance.md +++ b/docs/modules/graph/fault-tolerance.md @@ -1,42 +1,80 @@ # Graph Error Handling And Fault Tolerance -Default behavior: - -- node error fails the run -- retry only if policy allows it -- timeout fails the node -- checkpoint remains at the last completed boundary -- pending writes from completed tasks remain available when the checkpointer - supports them - -Target behavior: - -- route node errors to a node-specific or default error handler -- expose handled errors in task events and checkpoint task metadata -- retry with backoff and jitter -- support optional nodes by explicit policy -- support cooperative drain/shutdown with a drain reason -- distinguish graph recursion errors from node failures - -Error types should distinguish: - -- missing start -- duplicate node -- missing node -- missing edge target -- duplicate route -- missing route -- invalid command target -- invalid parent graph command -- invalid send target -- recursion limit -- checkpoint required -- checkpoint missing -- interrupt resume mismatch -- reducer conflict -- invalid concurrent update -- node timeout -- node failure -- graph drained -- serialization failure -- checkpoint backend failure +The graph runtime survives transient network failures and makes a hard failure +restartable, using two opt-in mechanisms plus the existing checkpoint/resume +machinery. + +## Node retry (surviving network blips) + +`CompiledGraph::with_node_retry(RetryPolicy)` wraps every node handler. When a +handler fails with a [retryable] error — `TinyAgentsError::Model` or +`TinyAgentsError::Tool`, the transient class classified by +`harness::retry::is_retryable` — the executor re-runs the node **from its start** +up to the policy's `max_attempts`, emitting `GraphEvent::NodeRetryScheduled` +before each retry. Because a node is never suspended mid-flight, a retry rebuilds +the handler future and context from scratch, matching the durable-execution +model. + +Backoff between attempts is computed from the `RetryPolicy` (exponential, with +optional jitter) but only **slept on** when the policy opts in via +`RetryPolicy::with_backoff_sleep(true)`. The default is sleep-free so unit tests +stay deterministic; production callers enable it so retries wait a real, growing +delay. The same opt-in switch governs the harness model-call retry loop and +`RetryMiddleware`. + +Both the sequential and parallel runners drive handlers through one shared +`run_node_with_retry` helper, which also applies the per-node timeout — so a +retried attempt is still bounded by `node_timeout`. + +## Resumable failures (restarting / continuing after a hard failure) + +When a handler fails **beyond** the retry budget, or with a non-retryable error, +the run is not lost. On a checkpointed thread the executor: + +1. folds the updates of the branches that completed *before* the failing one + into committed state (partial parallel progress is preserved, not discarded); +2. persists a **failure-boundary checkpoint** whose `next_nodes` schedule the + failed node (and the not-yet-run tail of the step) for re-run, records the + successful nodes as `completed_tasks`, and stamps the rendered `error` and + `failed_node` into the checkpoint metadata; +3. records a `Failed` run status carrying that checkpoint id; +4. returns the error to the caller. + +Without a checkpointer/thread the run aborts immediately with the error, exactly +as before — the failure checkpoint is a no-op. + +### Restarting and continuing + +Because `resume`/`resume_from` replay from a checkpoint's `next_nodes`, a failed +run resumes exactly like an interrupted one: + +- `CompiledGraph::retry(thread)` — shorthand for `resume` with an empty command; + re-runs the failed node and the not-yet-run tail from the failure boundary. +- **Continue on user feedback** — inspect committed state with `get_state`, edit + it with `update_state` (attributed through the reducers), then `retry`/`resume`. + The re-run sees the edited state. + +See the `resilient_graph` example (`cargo run --example resilient_graph`) for +both mechanisms end to end. + +## Error taxonomy + +`TinyAgentsError` distinguishes structural/config errors (non-resumable) from +node failures (resumable on a checkpointed thread): + +- structural / not retried, not resumable: missing start, missing node, missing + edge target, missing route, invalid command/parent/send target, recursion + limit, node-visit limit, sub-agent depth, checkpoint required/missing, resume + mismatch, reducer conflict, invalid concurrent update, serialization failure, + checkpoint backend failure; +- transient / retryable and resumable: `Model`, `Tool` (and `Timeout` aborts the + node but leaves a resumable boundary like any other node failure). + +## Future work + +- Route node errors to a node-specific or default error-handler node. +- Cooperative drain/shutdown with a drain reason. +- Populate the checkpoint `pending_writes` list explicitly (today partial + progress is folded into committed state instead). + +[retryable]: ../../../src/harness/retry/mod.rs diff --git a/examples/resilient_graph.rs b/examples/resilient_graph.rs new file mode 100644 index 0000000..6fdb2bb --- /dev/null +++ b/examples/resilient_graph.rs @@ -0,0 +1,135 @@ +//! Surviving network blips and restarting a failed run. +//! +//! Two resilience primitives, on one durable graph: +//! +//! 1. **Node-level retry.** [`CompiledGraph::with_node_retry`] re-runs a node +//! from its start when its handler fails with a transient +//! ([retryable][tinyagents::harness::retry::is_retryable]) error — a model or +//! tool/network error. A `fetch` node here fails its first two attempts (a +//! simulated connectivity blip) and succeeds on the third, so the run +//! completes without any operator involvement. +//! +//! 2. **Resumable failure + restart.** A `commit` node fails *harder* than the +//! retry budget allows. On a checkpointed thread the run does not vanish: it +//! persists a resumable failure-boundary checkpoint (the failed node is +//! scheduled to re-run, earlier progress is preserved) and returns the error. +//! The "outage" then clears and [`CompiledGraph::retry`] restarts the run +//! from that checkpoint to completion. A run could equally be *continued on +//! user feedback* by editing state with `update_state` before `retry`. +//! +//! Run with: +//! +//! ```text +//! cargo run --example resilient_graph +//! ``` + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use tinyagents::harness::ids::ExecutionStatus; +use tinyagents::harness::retry::RetryPolicy; +use tinyagents::{ + GraphBuilder, InMemoryCheckpointer, NodeContext, NodeResult, Result, TinyAgentsError, +}; + +/// A tiny pipeline state: how far we got, plus an audit log. +#[derive(Clone, Debug, Default)] +struct Pipeline { + fetched: bool, + committed: bool, + log: Vec, +} + +#[tokio::main] +async fn main() -> Result<()> { + let checkpointer = Arc::new(InMemoryCheckpointer::::new()); + + // A flaky `fetch`: fails (retryable model error) its first two attempts, + // then succeeds — a transient connectivity blip the node-retry policy + // absorbs on its own. + let fetch_attempts = Arc::new(AtomicUsize::new(0)); + // A `commit` that stays down until the "outage" clears between runs. The + // outer restart, not the inner retry policy, is what recovers it. + let outage = Arc::new(AtomicUsize::new(0)); + + let fa = fetch_attempts.clone(); + let og = outage.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("fetch", move |mut state: Pipeline, _c: NodeContext| { + let fa = fa.clone(); + async move { + let n = fa.fetch_add(1, Ordering::SeqCst); + if n < 2 { + return Err(TinyAgentsError::Model(format!( + "connection reset (fetch attempt {})", + n + 1 + ))); + } + state.fetched = true; + state.log.push(format!("fetched on attempt {}", n + 1)); + Ok(NodeResult::Update(state)) + } + }) + .add_node("commit", move |mut state: Pipeline, _c: NodeContext| { + let og = og.clone(); + async move { + // Down while the outage flag is 0; healthy once it is raised. + if og.load(Ordering::SeqCst) == 0 { + return Err(TinyAgentsError::Tool( + "commit endpoint unreachable (outage)".into(), + )); + } + state.committed = true; + state.log.push("committed".into()); + Ok(NodeResult::Update(state)) + } + }) + .set_entry("fetch") + .add_edge("fetch", "commit") + .set_finish("commit") + .compile() + .unwrap() + // 1 try + 3 retries per node, absorbing the fetch blips. + .with_node_retry(RetryPolicy::default().with_max_attempts(4)) + .with_checkpointer(checkpointer.clone()); + + // ---- First run: fetch self-heals, but commit is fully down. ------------ + let thread = "orders-4711"; + let err = graph + .run_with_thread(thread, Pipeline::default()) + .await + .expect_err("commit is down, so the first run fails"); + println!("first run failed as expected: {err}"); + println!( + " fetch took {} attempts", + fetch_attempts.load(Ordering::SeqCst) + ); + + // The failure left a resumable checkpoint: `fetch`'s progress is committed + // and `commit` is scheduled to re-run. + let snapshot = graph + .get_state(thread, None) + .await? + .expect("a failure-boundary checkpoint exists"); + println!( + " durable state: fetched={}, committed={}", + snapshot.values.fetched, snapshot.values.committed + ); + println!(" will re-run on resume: {:?}", snapshot.next_nodes); + println!(" checkpoints on thread: {}", checkpointer.count(thread)); + + // ---- The outage clears; restart the run from the checkpoint. ----------- + outage.store(1, Ordering::SeqCst); + let resumed = graph.retry(thread).await?; + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + println!("\nrestarted to completion:"); + println!( + " fetched={}, committed={}", + resumed.state.fetched, resumed.state.committed + ); + for line in &resumed.state.log { + println!(" - {line}"); + } + + Ok(()) +} diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index f4d17d4..01ecb92 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -32,11 +32,37 @@ //! the fan-in / join: lower-index branches' updates are applied first. //! - The *lowest-index* branch that errors or interrupts is the step's terminal //! outcome. Updates produced by lower-index successful branches are still -//! applied/persisted; an error aborts the run, an interrupt persists a -//! checkpoint whose pending nodes are that branch and every later active node. +//! applied/persisted; an error persists a resumable failure boundary (see +//! below) and aborts, an interrupt persists a checkpoint whose pending nodes +//! are that branch and every later active node. //! - Because branches run on cloned snapshots and never share mutable state, //! concurrency is data-race free; the reducer alone resolves conflicting //! writes (deterministically, by index). +//! +//! ## Network resilience and resumable failures +//! +//! Two opt-in mechanisms make a run durable under transient failure and +//! restartable after a hard one: +//! +//! - **Node retry.** With +//! [`CompiledGraph::with_node_retry`], a node whose handler fails with a +//! [retryable][crate::harness::retry::is_retryable] error (a model or tool +//! error — the transient class) is re-run from its start up to the policy's +//! attempt cap, emitting +//! [`GraphEvent::NodeRetryScheduled`](crate::graph::stream::GraphEvent::NodeRetryScheduled) +//! and sleeping the opt-in backoff between attempts. A single network blip is +//! absorbed without touching the run. +//! - **Resumable failure.** When a handler fails beyond the retry budget (or the +//! error is non-retryable), the executor does not discard the step. On a +//! checkpointed thread it folds the branches that already completed into +//! committed state and persists a failure-boundary checkpoint whose +//! `next_nodes` schedule the failed node (and the not-yet-run tail) for a +//! later [`CompiledGraph::resume`]/[`CompiledGraph::retry`], with the error +//! and failed node stamped into the checkpoint metadata. The run then reports +//! `Failed` (carrying that checkpoint id) and returns the error. A caller can +//! restart it as-is, or continue on operator feedback by editing state with +//! [`CompiledGraph::update_state`] before resuming. Without a checkpointer the +//! run aborts immediately, exactly as before. mod types; @@ -48,7 +74,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; use crate::graph::builder::{ - Branch, BuilderNode, END, ForkId, NodeContext, NodeFuture, NodeMeta, START, + Branch, BuilderNode, END, ForkId, NodeContext, NodeFuture, NodeHandler, NodeMeta, START, }; use crate::graph::checkpoint::{ Checkpoint, CheckpointConfig, CheckpointTuple, Checkpointer, DurabilityMode, @@ -63,6 +89,7 @@ use crate::graph::stream::{GraphEvent, GraphEventSink}; use crate::harness::ids::{ CheckpointId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, ThreadId, }; +use crate::harness::retry::is_retryable; use crate::{Result, TinyAgentsError}; static SEQ: AtomicU64 = AtomicU64::new(0); @@ -110,6 +137,26 @@ struct StepRun { goto_map: HashMap>, /// The lowest-index branch interrupt, if any (its active-set index + value). interrupt: Option<(usize, Interrupt)>, + /// A node-handler failure that survived the node-retry policy, if any. When + /// set, `updates` still carries the updates of the branches that completed + /// *before* the failing branch, so the executor can fold that partial + /// progress into committed state and persist a resumable failure boundary. + failure: Option, +} + +/// A node-handler failure captured by a runner so the executor can persist a +/// resumable failure-boundary checkpoint instead of discarding partial progress. +struct StepFailure { + /// The node whose handler ultimately failed (after any retries). + failed_node: NodeId, + /// Nodes that completed successfully before the failure, in active-set order + /// (recorded as the checkpoint's completed tasks). + completed: Vec, + /// Nodes to re-run on resume: the failed node first, then the not-yet-folded + /// members of the step (recorded as the checkpoint's next nodes). + pending: Vec, + /// The escalated error. + error: TinyAgentsError, } /// One scheduled activation in the active set: the node plus an optional @@ -184,6 +231,7 @@ impl CompiledGraph { max_concurrency, node_timeout, durability: crate::graph::checkpoint::DurabilityMode::default(), + node_retry: None, } } @@ -232,6 +280,27 @@ impl CompiledGraph { self } + /// Sets the per-node [`RetryPolicy`] applied around every node handler. + /// + /// Opt-in network resilience for the graph: when a node handler fails with a + /// [retryable][crate::harness::retry::is_retryable] error (a model or tool + /// error — the transient class), the executor re-runs the node from its + /// start up to the policy's attempt cap, emitting a + /// [`GraphEvent::NodeRetryScheduled`](crate::graph::stream::GraphEvent::NodeRetryScheduled) + /// before each retry. Backoff between attempts is slept on only when the + /// policy opts in via + /// [`RetryPolicy::with_backoff_sleep`](crate::harness::retry::RetryPolicy::with_backoff_sleep). + /// + /// Non-retryable errors, and retryable errors once attempts are exhausted, + /// escalate — and, on a checkpointed thread, leave a resumable + /// failure-boundary checkpoint (see [`CompiledGraph::resume`]) so the run can + /// be restarted or continued rather than lost. Without a policy (the + /// default) the first node error aborts the run immediately. + pub fn with_node_retry(mut self, policy: crate::harness::retry::RetryPolicy) -> Self { + self.node_retry = Some(policy); + self + } + /// Sets the checkpoint namespace (used by subgraph wrappers). pub fn with_namespace(mut self, namespace: Vec) -> Self { self.namespace = namespace; @@ -382,6 +451,27 @@ where .await } + /// Retries a failed run from its latest (failure-boundary) checkpoint, + /// re-running the node that failed and the not-yet-run tail of that step. + /// + /// This is the resume counterpart for the *failure* path (as opposed to a + /// human interrupt): after a node handler aborts a checkpointed run — a + /// transient outage that outlived the node-retry policy, or a hard crash — + /// the run leaves a resumable checkpoint (see + /// [`CompiledGraph::with_node_retry`]). Calling `retry` re-runs exactly what + /// did not complete, carrying no resume value. It is shorthand for + /// [`CompiledGraph::resume`] with an empty [`Command`]. + /// + /// To continue on *user feedback* instead of a bare retry, first inspect the + /// committed state with + /// [`get_state`](CompiledGraph::get_state), edit it with + /// [`update_state`](CompiledGraph::update_state), then call `retry` (or + /// `resume`) — the edited state is what the re-run sees. + pub async fn retry(&self, thread_id: impl Into) -> Result> { + self.resume_from(thread_id, ResumeTarget::Latest, Command::new()) + .await + } + /// Resumes a run from a specific checkpoint (time-travel resume). /// /// [`ResumeTarget::Latest`] behaves exactly like [`CompiledGraph::resume`]; @@ -776,7 +866,7 @@ where self.emit(GraphEvent::RunStarted { run_id: run_id.clone(), }); - self.fail_run(&run_id, &thread_id, started_at, steps, &err) + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) .await; return Err(err); } @@ -820,14 +910,14 @@ where .min(self.recursion_policy.max_total_steps); if steps >= step_limit { let err = TinyAgentsError::RecursionLimit(step_limit); - self.fail_run(&run_id, &thread_id, started_at, steps, &err) + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) .await; return Err(err); } // Node-loop recursion: enforce `max_visits_per_node` per activation. for activation in &active { if let Err(err) = recursion.record_node_visit(&mut node_visits, &activation.node) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err) + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) .await; return Err(err); } @@ -871,10 +961,11 @@ where updates, goto_map, interrupt, + failure, } = match run_result { Ok(step_run) => step_run, Err(err) => { - self.fail_run(&run_id, &thread_id, started_at, steps, &err) + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) .await; return Err(err); } @@ -893,12 +984,53 @@ where let child_runs_meta = serde_json::to_value(&step_child_runs).unwrap_or(serde_json::Value::Null); + // Node-handler failure (survived any node-retry policy): the updates + // of the branches that completed before it are already folded into + // `state` above, so persist a resumable failure-boundary checkpoint + // scheduling the failed node (and the not-yet-run tail) for a later + // `resume`/`retry`, record a `Failed` status carrying the error and + // that checkpoint, and abort. Without a checkpointer/thread the + // checkpoint is a no-op and the run aborts exactly as before. + if let Some(fail) = failure { + let StepFailure { + failed_node, + completed, + pending, + error, + } = fail; + let checkpoint_id = self + .persist_failure_checkpoint( + &thread_id, + &run_id, + &state, + &pending, + &completed, + parent_checkpoint.clone(), + steps, + &failed_node, + &error, + &recursion_meta, + &child_runs_meta, + ) + .await?; + self.fail_run( + &run_id, + &thread_id, + started_at, + steps, + &error, + checkpoint_id, + ) + .await; + return Err(error); + } + // Interrupt: persist a checkpoint whose next nodes are the // not-yet-completed members of this step (interrupted node first), // then return control to the caller. if let Some((index, emitted)) = interrupt { if let Err(err) = self.require_interrupt_durability(&thread_id) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err) + self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) .await; return Err(err); } @@ -1052,6 +1184,11 @@ where /// Emits a [`GraphEvent::RunFailed`] and records a terminal `Failed` status /// for a run that aborted with `err`. + /// + /// `checkpoint_id` is the resumable failure-boundary checkpoint when the run + /// left one (a node-handler failure on a checkpointed thread), or `None` for + /// a structural/non-resumable abort. When present it is recorded on the + /// status so an observer can locate the checkpoint to `resume`/`retry` from. async fn fail_run( &self, run_id: &RunId, @@ -1059,6 +1196,7 @@ where started_at: SystemTime, steps: usize, err: &TinyAgentsError, + checkpoint_id: Option, ) { self.emit(GraphEvent::RunFailed { run_id: run_id.clone(), @@ -1069,9 +1207,66 @@ where status.current_step = steps; status.ended_at = Some(SystemTime::now()); status.error = Some(err.to_string()); + status.checkpoint_id = checkpoint_id; self.save_status(status).await; } + /// Persists a resumable failure-boundary checkpoint for a node-handler + /// failure that survived the node-retry policy. + /// + /// Mirrors the interrupt boundary: `next_nodes` schedules the failed node + /// (and any not-yet-run members of the step) so `resume`/`retry` re-runs + /// exactly what did not complete, while `completed_tasks` records the + /// branches that already succeeded (their updates are folded into `state` + /// before this is called). The rendered error and failed node id are stamped + /// into the checkpoint metadata for diagnosis. A no-op returning `None` when + /// no checkpointer/thread is configured — the run then aborts without a + /// resumable checkpoint, exactly as before this policy existed. + #[allow(clippy::too_many_arguments)] + async fn persist_failure_checkpoint( + &self, + thread_id: &Option, + run_id: &RunId, + state: &State, + next_nodes: &[NodeId], + completed_tasks: &[NodeId], + parent: Option, + step: usize, + failed_node: &NodeId, + error: &TinyAgentsError, + recursion: &serde_json::Value, + child_runs: &serde_json::Value, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { + return Ok(None); + }; + let checkpoint = Checkpoint { + thread_id: thread.to_string(), + checkpoint_id: next_id("ckpt"), + run_id: Some(run_id.to_string()), + parent_checkpoint_id: parent, + namespace: self.namespace.clone(), + state: state.clone(), + next_nodes: next_nodes.to_vec(), + completed_tasks: completed_tasks.to_vec(), + pending_writes: Vec::new(), + interrupts: Vec::new(), + metadata: serde_json::json!({ + "source": "loop", + "step": step, + "recursion": recursion, + "child_runs": child_runs, + "failed_node": failed_node.as_str(), + "error": error.to_string(), + }), + }; + let id = checkpointer.put(checkpoint).await?; + self.emit(GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + }); + Ok(Some(id)) + } + /// Builds the per-task [`NodeContext`] for `node_id` at the given branch. /// /// `fork` carries the branch identity in a concurrent step (`None` in @@ -1123,6 +1318,51 @@ where } } + /// Runs one node handler under the graph's node-retry policy. + /// + /// Builds a fresh handler future (and re-clones the context) for each + /// attempt, so a retried node re-runs from its start — matching the durable + /// execution model, where a node is never suspended mid-flight. On a + /// [retryable][crate::harness::retry::is_retryable] error, when a + /// [`RetryPolicy`](crate::harness::retry::RetryPolicy) is configured and + /// permits another attempt, it emits + /// [`GraphEvent::NodeRetryScheduled`], sleeps the (opt-in) backoff, and + /// retries. Non-retryable errors, absence of a policy, or an exhausted + /// attempt budget return the error unchanged. The per-node timeout still + /// bounds every individual attempt via [`Self::run_node_future`]. + async fn run_node_with_retry( + &self, + node_id: &NodeId, + handler: &Arc>, + state: &State, + ctx: NodeContext, + step: usize, + ) -> Result> { + let mut attempt = 0usize; + loop { + let fut = handler(state.clone(), ctx.clone()); + match self.run_node_future(node_id, fut).await { + Ok(result) => return Ok(result), + Err(error) => { + let retry = self + .node_retry + .as_ref() + .filter(|policy| policy.should_retry(attempt) && is_retryable(&error)); + let Some(policy) = retry else { + return Err(error); + }; + attempt += 1; + self.emit(GraphEvent::NodeRetryScheduled { + node: node_id.clone(), + step, + attempt, + }); + policy.sleep_backoff(attempt).await; + } + } + } + } + /// Folds a single successful branch result into the step accumulators. /// /// Pushes the node to `visited`, records updates/goto, emits the matching @@ -1196,6 +1436,7 @@ where let mut updates: Vec = Vec::new(); let mut goto_map: HashMap> = HashMap::new(); let mut interrupt: Option<(usize, Interrupt)> = None; + let mut failure: Option = None; for (index, activation) in active.iter().enumerate() { let node_id = &activation.node; @@ -1226,7 +1467,7 @@ where child_runs, ); let result = match self - .run_node_future(node_id, (node.handler)(state.clone(), ctx)) + .run_node_with_retry(node_id, &node.handler, state, ctx, step) .await { Ok(result) => result, @@ -1236,7 +1477,16 @@ where step, error: error.to_string(), }); - return Err(error); + // Preserve the progress of the branches that already ran: + // record them as completed and schedule this node plus the + // not-yet-run tail for a resumable retry. + failure = Some(StepFailure { + failed_node: node_id.clone(), + completed: activation_nodes(&active[..index]), + pending: activation_nodes(&active[index..]), + error, + }); + break; } }; @@ -1258,6 +1508,7 @@ where updates, goto_map, interrupt, + failure, }) } @@ -1287,9 +1538,11 @@ where frames: &[RecursionFrame], child_runs: &ChildRunSink, ) -> Result> { - let timeout = self.node_timeout; // Build one forked context + future per branch. Node lookup and resume - // consumption happen up front so the futures borrow nothing local. + // consumption happen up front so the futures borrow nothing mutable; each + // branch drives its handler through the node-retry policy (which also + // applies the per-node timeout), so a transient failure in one branch is + // retried without disturbing its siblings. let mut futures = Vec::with_capacity(active.len()); for (index, activation) in active.iter().enumerate() { let node_id = &activation.node; @@ -1325,18 +1578,11 @@ where frames, child_runs, ); - let raw = (node.handler)(state.clone(), ctx); + let handler = node.handler.clone(); let owned_node = node_id.clone(); futures.push(async move { - match timeout { - Some(d) => match tokio::time::timeout(d, raw).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "node `{owned_node}` exceeded its {d:?} timeout" - ))), - }, - None => raw.await, - } + self.run_node_with_retry(&owned_node, &handler, state, ctx, step) + .await }); } @@ -1361,6 +1607,7 @@ where let mut updates: Vec = Vec::new(); let mut goto_map: HashMap> = HashMap::new(); let mut interrupt: Option<(usize, Interrupt)> = None; + let mut failure: Option = None; for (index, (activation, result)) in active.iter().zip(results).enumerate() { let node_id = &activation.node; @@ -1372,7 +1619,16 @@ where step, error: error.to_string(), }); - return Err(error); + // The lowest-index failing branch is terminal: fold the + // lower-index successes (already applied above) and schedule + // this branch plus the rest for a resumable retry. + failure = Some(StepFailure { + failed_node: node_id.clone(), + completed: activation_nodes(&active[..index]), + pending: activation_nodes(&active[index..]), + error, + }); + break; } }; @@ -1394,6 +1650,7 @@ where updates, goto_map, interrupt, + failure, }) } diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 6356e6e..27a10c8 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -9,6 +9,7 @@ use crate::graph::command::{Command, Interrupt, NodeResult, Send}; use crate::graph::reducer::ClosureStateReducer; use crate::graph::stream::{CollectingSink, GraphEvent}; use crate::harness::ids::ExecutionStatus; +use crate::harness::retry::RetryPolicy; use serde_json::json; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; @@ -1316,3 +1317,189 @@ async fn node_timeout_fails_slow_handler() { let err = graph.run(0).await.unwrap_err(); assert!(matches!(err, TinyAgentsError::Timeout(_))); } + +// ── Network resilience: node retry + resumable failures ────────────────────── + +/// A single-node graph whose handler fails (with a retryable model error) the +/// first `fail_times` invocations, then succeeds with `+1`. The shared counter +/// lets a test observe how many attempts were made. +fn flaky_graph(fail_times: usize, attempts: Arc) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("flaky", move |s, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < fail_times { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + Ok(NodeResult::Update(s + 1)) + } + } + }) + .set_entry("flaky") + .set_finish("flaky") + .compile() + .unwrap() +} + +#[tokio::test] +async fn node_retry_recovers_transient_failure() { + let attempts = Arc::new(AtomicUsize::new(0)); + let sink = Arc::new(CollectingSink::new()); + // Fail twice, succeed on the third attempt; the policy allows 1 try + 3 + // retries, so recovery is within budget. + let graph = flaky_graph(2, attempts.clone()) + .with_node_retry(RetryPolicy::default().with_max_attempts(4)) + .with_event_sink(sink.clone()); + + let run = graph.run(10).await.unwrap(); + assert_eq!(run.state, 11); + assert_eq!(run.status.status, ExecutionStatus::Completed); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 3); + + let retries = sink + .events() + .into_iter() + .filter(|e| matches!(e, GraphEvent::NodeRetryScheduled { .. })) + .count(); + assert_eq!(retries, 2, "one retry per transient failure"); +} + +#[tokio::test] +async fn exhausted_retries_leave_a_resumable_failure_checkpoint() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let attempts = Arc::new(AtomicUsize::new(0)); + // Fail the first 3 invocations. With 1 try + 1 retry the first run exhausts + // its budget (2 attempts: n=0,1) and aborts, leaving a resumable checkpoint. + let graph = flaky_graph(3, attempts.clone()) + .with_node_retry(RetryPolicy::default().with_max_attempts(2)) + .with_checkpointer(cp.clone()); + + let err = graph.run_with_thread("net", 100).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2, "1 try + 1 retry"); + + // The failure boundary is durable: the checkpoint schedules the failed node + // for re-run at the first superstep. + let list = cp.list("net").await.unwrap(); + let last = list.last().expect("a failure checkpoint was persisted"); + assert_eq!(last.next_nodes, vec![NodeId::from("flaky")]); + let snapshot = graph.get_state("net", None).await.unwrap().unwrap(); + assert_eq!( + snapshot.metadata.step, 1, + "failure boundary is at the first superstep" + ); + + // Retry: attempt n=2 fails, retry n=3 (3<3 is false) succeeds — the run + // completes without losing the earlier progress. + let resumed = graph.retry("net").await.unwrap(); + assert_eq!( + resumed.state, 101, + "resume re-runs the failed node to success" + ); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 4); +} + +#[tokio::test] +async fn node_failure_without_retry_policy_is_resumable() { + // No node-retry policy configured: the first failure aborts the run, but a + // checkpointed thread still leaves a resumable failure boundary (the + // "resumable abort" default) rather than losing the run. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp.clone()); + + let err = graph.run_with_thread("once", 7).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 1, + "no retries attempted" + ); + + // Failed status carries the resumable checkpoint id. + let status = graph.get_state("once", None).await.unwrap().unwrap(); + assert_eq!(status.next_nodes, vec![NodeId::from("flaky")]); + + // The transient condition has cleared; retry completes the run. + let resumed = graph.retry("once").await.unwrap(); + assert_eq!(resumed.state, 8); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} + +#[tokio::test] +async fn edit_state_then_retry_uses_the_edited_state() { + // User-feedback continuation: after a failure, the operator edits committed + // state via update_state, then retries; the re-run sees the edited value. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp.clone()); + + graph.run_with_thread("feedback", 0).await.unwrap_err(); + + // Operator bumps the committed state by +40, inheriting the failure + // boundary's pending nodes (`flaky`), then retries. The node adds +1 to + // whatever state it now sees. + graph.update_state("feedback", 40, None).await.unwrap(); + let resumed = graph.retry("feedback").await.unwrap(); + assert_eq!( + resumed.state, 41, + "retry runs against the edited state (40) + 1" + ); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} + +#[tokio::test] +async fn parallel_partial_progress_is_preserved_on_failure() { + // A parallel step where one branch succeeds and a lower/higher-index branch + // fails: the successful branch's update is folded into committed state and + // the failure checkpoint schedules only the failed branch for re-run. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = GraphBuilder::::new() + .with_parallel(true) + .set_reducer(ClosureStateReducer::new(|s: i32, u: i32| Ok(s + u))) + .add_node("seed", |_s, _c: NodeContext| async move { + Ok(NodeResult::Command(Command::goto(["ok", "flaky"]))) + }) + .add_node("ok", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update(1)) + }) + .add_node("flaky", move |_s, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + // Fail on the first invocation, succeed on resume. + if attempts.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + Err(TinyAgentsError::Model("branch blip".into())) + } else { + Ok(NodeResult::Update(10)) + } + } + }) + .set_entry("seed") + .mark_command_routing("seed") + .set_finish("ok") + .set_finish("flaky") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + // First run: seed fans out; the "ok" branch commits +1, "flaky" aborts. + graph.run_with_thread("fanout", 0).await.unwrap_err(); + let snapshot = graph.get_state("fanout", None).await.unwrap().unwrap(); + assert_eq!( + snapshot.values, 1, + "the successful branch's +1 is preserved" + ); + assert!( + snapshot.next_nodes.contains(&NodeId::from("flaky")), + "only the failed branch is scheduled for re-run: {:?}", + snapshot.next_nodes + ); + + // Resume: the flaky branch now succeeds (+10) without re-running "ok". + let resumed = graph.retry("fanout").await.unwrap(); + assert_eq!(resumed.state, 11, "1 (preserved) + 10 (re-run branch)"); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} diff --git a/src/graph/compiled/types.rs b/src/graph/compiled/types.rs index d6f36d4..bfb45c9 100644 --- a/src/graph/compiled/types.rs +++ b/src/graph/compiled/types.rs @@ -67,6 +67,14 @@ pub struct CompiledGraph { /// When checkpoints are persisted relative to execution (default /// [`DurabilityMode::Sync`]). pub(crate) durability: DurabilityMode, + /// Optional per-node retry policy applied around each node handler. When + /// set, a handler that fails with a + /// [retryable][crate::harness::retry::is_retryable] error is re-run from its + /// start (with opt-in backoff) up to the policy's attempt cap before the + /// failure escalates. `None` (default) preserves the original + /// abort-on-first-error behavior. Configured via + /// [`CompiledGraph::with_node_retry`](crate::graph::CompiledGraph::with_node_retry). + pub(crate) node_retry: Option, } impl std::fmt::Debug for CompiledGraph { @@ -108,6 +116,7 @@ impl Clone for CompiledGraph { max_concurrency: self.max_concurrency, node_timeout: self.node_timeout, durability: self.durability, + node_retry: self.node_retry.clone(), } } } diff --git a/src/graph/stream/types.rs b/src/graph/stream/types.rs index 1d7f16a..8836aee 100644 --- a/src/graph/stream/types.rs +++ b/src/graph/stream/types.rs @@ -84,6 +84,17 @@ pub enum GraphEvent { /// Rendered error. error: String, }, + /// A node handler failed with a retryable error and a retry was scheduled + /// under the graph's node-retry policy. Emitted before the (opt-in) backoff + /// wait and the re-run of the node from its start. + NodeRetryScheduled { + /// Node id. + node: NodeId, + /// Step number. + step: usize, + /// The 1-based retry attempt about to be made. + attempt: usize, + }, /// A node produced a state update applied at the boundary. StateUpdated { /// Node id. @@ -162,6 +173,7 @@ impl GraphEvent { GraphEvent::NodeStarted { .. } => "node.started", GraphEvent::NodeCompleted { .. } => "node.completed", GraphEvent::NodeFailed { .. } => "node.failed", + GraphEvent::NodeRetryScheduled { .. } => "node.retry_scheduled", GraphEvent::StateUpdated { .. } => "state.updated", GraphEvent::RouteSelected { .. } => "route.selected", GraphEvent::CheckpointSaved { .. } => "checkpoint.saved", @@ -185,6 +197,7 @@ impl GraphEvent { | GraphEvent::NodeStarted { step, .. } | GraphEvent::NodeCompleted { step, .. } | GraphEvent::NodeFailed { step, .. } + | GraphEvent::NodeRetryScheduled { step, .. } | GraphEvent::StateUpdated { step, .. } | GraphEvent::ContextForked { step, .. } => Some(*step), GraphEvent::RunCompleted { steps, .. } => Some(*steps), diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index dfe0b06..4e0de75 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -50,10 +50,12 @@ //! # Backoff //! //! Retry backoff durations are *computed* via -//! [`crate::harness::retry::RetryPolicy::backoff_for_attempt`] but the loop does -//! **not** sleep: keeping the loop sleep-free makes tests fast and -//! deterministic. A real provider integration may choose to sleep for the -//! computed duration before retrying. +//! [`crate::harness::retry::RetryPolicy::backoff_for_attempt`]. Whether the loop +//! actually sleeps for that duration is opt-in: it is off by default (keeping +//! tests fast and deterministic) and enabled per policy via +//! [`crate::harness::retry::RetryPolicy::with_backoff_sleep`], so a real +//! provider integration retries after a genuine, growing delay while unit tests +//! stay sleep-free. mod types; @@ -811,9 +813,10 @@ impl AgentHarness { call_id: call_id.clone(), attempt, }); - // Compute (but do not sleep on) the backoff so the - // loop stays fast and deterministic in tests. - let _ = self.policy.retry.backoff_for_attempt(attempt); + // Sleep for the backoff only when the policy opts in + // (`with_backoff_sleep`); otherwise this is a no-op so + // the loop stays fast and deterministic in tests. + self.policy.retry.sleep_backoff(attempt).await; continue; } break Err(error); diff --git a/src/harness/middleware/library/mod.rs b/src/harness/middleware/library/mod.rs index d034a6c..14265a9 100644 --- a/src/harness/middleware/library/mod.rs +++ b/src/harness/middleware/library/mod.rs @@ -19,7 +19,8 @@ //! # Testability //! //! None of these middleware sleep on the wall clock in a way that tests cannot -//! control: [`RetryMiddleware`] computes but never sleeps on backoff, +//! control: [`RetryMiddleware`] sleeps on backoff only when its policy opts in +//! via [`RetryPolicy::with_backoff_sleep`] (off by default), //! [`TimeoutMiddleware`] is exercised under `tokio::time` paused-time tests, and //! [`RateLimitMiddleware`] takes an injectable clock and a configurable poll //! interval so its wait loop can be driven deterministically. @@ -66,8 +67,9 @@ impl RetryMiddleware { /// Returns the policy-derived backoff for the given retry `attempt`. /// - /// Exposed for callers that want to sleep before a retry; the middleware - /// itself never sleeps. + /// Exposed for callers that want to inspect the backoff. The middleware + /// itself sleeps between retries only when the policy opts in via + /// [`RetryPolicy::with_backoff_sleep`]; otherwise it retries back-to-back. pub fn backoff_for_attempt(&self, attempt: usize) -> Duration { self.policy.backoff_for_attempt(attempt) } @@ -95,8 +97,9 @@ impl ModelMiddleware for Retry attempt += 1; let call_id = CallId::new(format!("{}-model", ctx.run_id())); ctx.emit(AgentEvent::RetryScheduled { call_id, attempt }); - // Compute (but do not sleep on) the backoff. - let _ = self.policy.backoff_for_attempt(attempt); + // Sleep for the backoff only when the policy opts in + // (`with_backoff_sleep`); a no-op otherwise. + self.policy.sleep_backoff(attempt).await; continue; } return Err(error); diff --git a/src/harness/retry/mod.rs b/src/harness/retry/mod.rs index d775fdd..e8906d8 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -66,6 +66,33 @@ impl RetryPolicy { self } + /// Enables or disables actually sleeping for the computed backoff between + /// retries. + /// + /// Off by default so tests stay deterministic and fast. Enable it in + /// production so a transient failure is retried after a real, growing delay + /// rather than back-to-back. See [`RetryPolicy::backoff_sleep`]. + pub fn with_backoff_sleep(mut self, sleep: bool) -> Self { + self.backoff_sleep = sleep; + self + } + + /// Sleeps for this policy's backoff before the given retry `attempt`, but + /// only when [`RetryPolicy::backoff_sleep`] is enabled. + /// + /// A single, reusable helper so every retry loop that honors a + /// [`RetryPolicy`] gets identical, opt-in backoff behavior. A no-op (returns + /// immediately) when sleeping is disabled or the computed backoff is zero. + pub async fn sleep_backoff(&self, attempt: usize) { + if !self.backoff_sleep { + return; + } + let backoff = self.backoff_for_attempt(attempt); + if backoff > Duration::ZERO { + tokio::time::sleep(backoff).await; + } + } + /// Returns `true` when another attempt should be made. /// /// `attempt` is zero-indexed: `0` is the first attempt, `1` is the first diff --git a/src/harness/retry/test.rs b/src/harness/retry/test.rs index 869ba45..c0fafcc 100644 --- a/src/harness/retry/test.rs +++ b/src/harness/retry/test.rs @@ -179,3 +179,35 @@ fn rate_limiter_refill_caps_at_capacity() { let later = start + Duration::from_secs(60); assert_eq!(limiter.available(later), 5); } + +#[test] +fn backoff_sleep_defaults_off_and_is_opt_in() { + // Default policy does not sleep, keeping retry loops deterministic in tests. + assert!(!RetryPolicy::default().backoff_sleep); + // The builder flips it on for production callers. + assert!( + RetryPolicy::default() + .with_backoff_sleep(true) + .backoff_sleep + ); +} + +#[tokio::test(start_paused = true)] +async fn sleep_backoff_waits_only_when_enabled() { + use tokio::time::Instant as TokioInstant; + + // Disabled (default): returns immediately with no virtual time elapsed. + let policy = RetryPolicy::default(); + let t0 = TokioInstant::now(); + policy.sleep_backoff(1).await; + assert_eq!(t0.elapsed(), Duration::ZERO); + + // Enabled: advances virtual time by the computed backoff (attempt 1 = + // initial_backoff_ms with the default multiplier applied at attempt^power). + let sleeping = RetryPolicy::default().with_backoff_sleep(true); + let expected = sleeping.backoff_for_attempt(1); + let t1 = TokioInstant::now(); + sleeping.sleep_backoff(1).await; + assert_eq!(t1.elapsed(), expected); + assert!(expected > Duration::ZERO); +} diff --git a/src/harness/retry/types.rs b/src/harness/retry/types.rs index 33da805..b0e373b 100644 --- a/src/harness/retry/types.rs +++ b/src/harness/retry/types.rs @@ -38,6 +38,15 @@ pub struct RetryPolicy { /// When `true`, the caller should supply a `[0, 1)` random value to /// [`RetryPolicy::backoff_for_attempt_with`] to add jitter. pub jitter: bool, + /// When `true`, retry loops that honor this policy actually + /// [`tokio::time::sleep`] for the computed backoff between attempts. + /// + /// Defaults to `false` so the harness stays fast and deterministic under + /// test: the backoff is *computed* (for observability) but not slept on. + /// Production/streaming callers opt in via + /// [`RetryPolicy::with_backoff_sleep`] so a transient provider/network + /// failure is retried after a real, growing delay instead of back-to-back. + pub backoff_sleep: bool, } impl Default for RetryPolicy { @@ -48,6 +57,7 @@ impl Default for RetryPolicy { max_backoff_ms: 30_000, multiplier: 2.0, jitter: false, + backoff_sleep: false, } } }