diff --git a/docs/modules/graph/builder.md b/docs/modules/graph/builder.md index e0bb8e8..d140620 100644 --- a/docs/modules/graph/builder.md +++ b/docs/modules/graph/builder.md @@ -43,6 +43,15 @@ parallel superstep (the active set runs in chunks of at most `n`). `with_node_timeout(d)` fails the run with `TinyAgentsError::Timeout` if any node handler does not resolve within `d`. +`CompiledGraph::with_run_deadline(d)` bounds the *whole run* by a wall-clock +`d`, checked at every super-step boundary: when the elapsed run time first +reaches `d` the run stops *between* super-steps with `TinyAgentsError::Timeout`, +leaving the last committed boundary checkpoint intact and resumable. Prefer this +over wrapping `run` in an external `tokio::time::timeout`, which aborts +mid-super-step and cannot leave a clean checkpoint. It bounds scheduling, not a +single in-flight node — pair it with `with_node_timeout` to also bound +individual handlers. + Compile-time options: ```rust diff --git a/src/graph/compiled/executor.rs b/src/graph/compiled/executor.rs index 9105fc6..7b3330e 100644 --- a/src/graph/compiled/executor.rs +++ b/src/graph/compiled/executor.rs @@ -418,6 +418,23 @@ where .await; return Err(err); } + // Whole-run wall-clock deadline: stop *between* super-steps once the + // elapsed run time reaches it, leaving the last committed boundary + // checkpoint intact (unlike an external `tokio::time::timeout`, which + // aborts mid-super-step and cannot). The already-completed super-steps + // and their checkpoints are preserved; the run fails with `Timeout`. + if let Some(deadline) = self.run_deadline { + let elapsed = started_at.elapsed().unwrap_or_default(); + if elapsed >= deadline { + let err = TinyAgentsError::Timeout(format!( + "graph run exceeded its {deadline:?} deadline after {steps} super-step(s) \ + ({elapsed:?} elapsed)" + )); + 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) { diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 922adf8..9bfe022 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -277,6 +277,7 @@ impl CompiledGraph { parallel, max_concurrency, node_timeout, + run_deadline: None, durability: crate::graph::checkpoint::DurabilityMode::default(), node_retry: None, } @@ -352,6 +353,26 @@ impl CompiledGraph { self } + /// Bounds the whole run by a wall-clock `deadline`, checked at every + /// super-step boundary. + /// + /// Unlike wrapping [`run`](Self::run) in an external + /// [`tokio::time::timeout`] — which aborts mid-super-step and cannot leave a + /// clean checkpoint — this stops the run *between* super-steps: when the + /// elapsed run time first reaches `deadline`, the run fails with + /// [`TinyAgentsError::Timeout`] and (on a checkpointed thread) the last + /// committed boundary checkpoint stays intact, so the run can be resumed or + /// inspected rather than lost. + /// + /// The deadline bounds *scheduling*, not a single in-flight node: a + /// long-running node still runs to completion within its super-step (bound + /// it independently with a per-node timeout). `None` (the default) imposes + /// no deadline. + pub fn with_run_deadline(mut self, deadline: std::time::Duration) -> Self { + self.run_deadline = Some(deadline); + self + } + /// Sets the checkpoint namespace (used by subgraph wrappers). pub fn with_namespace(mut self, namespace: Vec) -> Self { self.namespace = namespace; diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index a24a8d4..705a203 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -1861,6 +1861,103 @@ async fn node_timeout_fails_slow_handler() { assert!(matches!(err, TinyAgentsError::Timeout(_))); } +// ── Whole-run wall-clock deadline ──────────────────────────────────────────── + +/// A per-run deadline stops the run *between* super-steps once the elapsed run +/// time reaches it, surfacing [`TinyAgentsError::Timeout`]. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn run_deadline_stops_between_supersteps() { + let graph = GraphBuilder::::overwrite() + .add_node("a", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_millis(40)).await; + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", |s: i32, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .set_finish("b") + .compile() + .unwrap() + .with_run_deadline(Duration::from_millis(20)); + + // The first boundary (elapsed ~0) admits node `a`; the next boundary + // (elapsed ~40ms ≥ 20ms) trips the deadline before `b` ever runs. + let err = graph.run(0).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); +} + +/// A run that finishes within its deadline is unaffected — no false trip. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn run_deadline_allows_a_run_that_finishes_in_time() { + let graph = GraphBuilder::::overwrite() + .add_node("a", |s: i32, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", |s: i32, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .set_finish("b") + .compile() + .unwrap() + .with_run_deadline(Duration::from_secs(30)); + + let run = graph.run(0).await.unwrap(); + assert_eq!(run.state, 2); +} + +/// On a checkpointed thread, a deadline trip leaves the last committed boundary +/// checkpoint intact — so the run can be resumed to completion rather than lost +/// (the durability win over an external `tokio::time::timeout`). +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn run_deadline_leaves_last_checkpoint_resumable() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let topology = || { + GraphBuilder::::overwrite() + .add_node("a", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_millis(40)).await; + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", |s: i32, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("c", |s: i32, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("b", "c") + .set_finish("c") + .compile() + .unwrap() + }; + + // Trips after `a`'s boundary (state=1, next=[b]) but before `b` runs. + let deadlined = topology() + .with_checkpointer(cp.clone()) + .with_run_deadline(Duration::from_millis(20)); + let err = deadlined.run_with_thread("t", 0).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + + // The boundary checkpoint from the completed super-step survived intact. + let list = cp.list("t").await.unwrap(); + assert!( + !list.is_empty(), + "the pre-deadline boundary checkpoint is intact" + ); + + // Resuming (no deadline) continues from that checkpoint to completion. + let resumed = topology().with_checkpointer(cp.clone()); + let run = resumed + .resume("t", Command::resume(json!(null))) + .await + .unwrap(); + assert_eq!(run.state, 3, "resume ran the remaining super-steps b and c"); +} + // ── Network resilience: node retry + resumable failures ────────────────────── /// A single-node graph whose handler fails (with a retryable model error) the diff --git a/src/graph/compiled/types.rs b/src/graph/compiled/types.rs index a595798..e220827 100644 --- a/src/graph/compiled/types.rs +++ b/src/graph/compiled/types.rs @@ -64,6 +64,12 @@ pub struct CompiledGraph { pub(crate) max_concurrency: Option, /// Default per-node handler timeout (`None` = no timeout). pub(crate) node_timeout: Option, + /// Optional whole-run wall-clock deadline (`None` = no deadline). Checked at + /// every super-step boundary: when the elapsed run time reaches it the run + /// stops *between* super-steps with a [`TinyAgentsError::Timeout`], leaving + /// the last committed boundary checkpoint intact and resumable. Configured + /// via [`CompiledGraph::with_run_deadline`](crate::graph::CompiledGraph::with_run_deadline). + pub(crate) run_deadline: Option, /// When checkpoints are persisted relative to execution (default /// [`DurabilityMode::Sync`]). pub(crate) durability: DurabilityMode, @@ -115,6 +121,7 @@ impl Clone for CompiledGraph { parallel: self.parallel, max_concurrency: self.max_concurrency, node_timeout: self.node_timeout, + run_deadline: self.run_deadline, durability: self.durability, node_retry: self.node_retry.clone(), }