diff --git a/Cargo.lock b/Cargo.lock index 3965e16..1fa8618 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1603,7 +1603,7 @@ dependencies = [ [[package]] name = "tinyflows" -version = "0.6.0" +version = "0.6.1" dependencies = [ "async-trait", "axum", diff --git a/Cargo.toml b/Cargo.toml index 408cb02..a2dc850 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tinyflows" -version = "0.6.0" +version = "0.6.1" edition = "2024" rust-version = "1.85" license = "GPL-3.0-or-later" diff --git a/src/engine.rs b/src/engine.rs index 487058a..224f47c 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -808,6 +808,15 @@ pub const MAX_SUB_WORKFLOW_DEPTH: u64 = 8; /// child can read it back from `ctx.run` and enforce [`MAX_SUB_WORKFLOW_DEPTH`]. /// Used only by the `sub_workflow` node's recursive execution. /// +/// `token` is the **parent run's** cancellation token, forwarded so cancelling +/// the parent winds the whole subtree down: the child observes the same flipped +/// flag at its next node boundary and returns a cancelled [`RunOutcome`] instead +/// of running to completion orphaned from the parent. The child in turn hands +/// this token to its own node contexts, so a deeper `sub_workflow` propagates it +/// on — the whole nesting chain shares one signal. Historically this seeded a +/// fresh [`CancellationToken`], which severed cancellation at every sub-workflow +/// boundary. +/// /// # Errors /// Same as [`run`]. pub(crate) async fn run_sub_workflow( @@ -816,6 +825,7 @@ pub(crate) async fn run_sub_workflow( capabilities: &Capabilities, depth: u64, max_depth: u64, + token: CancellationToken, ) -> Result { let checkpointer: Arc> = Arc::new(InMemoryCheckpointer::::default()); @@ -833,7 +843,7 @@ pub(crate) async fn run_sub_workflow( "sub_workflow_depth": depth, "max_sub_workflow_depth": max_depth, })), - CancellationToken::new(), + token, ) .await?; Ok(outcome) @@ -1266,6 +1276,10 @@ fn build_graph( run: &run_meta, nodes: &nodes_state, caps: &caps, + // Handed to the executor so a nested engine call (today the + // `sub_workflow` node) can thread this run's cancellation + // into its child; a plain executor never reads it. + token: token.clone(), }; // BUG-8: bound THIS attempt (not the whole retry loop) to // `node_timeout`. Race the attempt future against a @@ -1386,6 +1400,7 @@ fn build_graph( run: &run_meta, nodes: &nodes_state, caps: &caps, + token: token.clone(), }; let scope = crate::nodes::expr_scope(&ctx); crate::expr::resolve_traced(&node.config, &scope).1 diff --git a/src/nodes/control_flow/condition.rs b/src/nodes/control_flow/condition.rs index 04d7007..c91c959 100644 --- a/src/nodes/control_flow/condition.rs +++ b/src/nodes/control_flow/condition.rs @@ -89,6 +89,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ConditionNode.execute(ctx).await.expect("execute"); ( diff --git a/src/nodes/control_flow/dedup.rs b/src/nodes/control_flow/dedup.rs index 6203da4..95bafe6 100644 --- a/src/nodes/control_flow/dedup.rs +++ b/src/nodes/control_flow/dedup.rs @@ -325,6 +325,7 @@ mod tests { run: &run, nodes: &Value::Null, caps, + token: crate::engine::CancellationToken::new(), }; DedupNode.execute(ctx).await.expect("execute") } diff --git a/src/nodes/control_flow/loop_node.rs b/src/nodes/control_flow/loop_node.rs index 10547b2..9ab4d35 100644 --- a/src/nodes/control_flow/loop_node.rs +++ b/src/nodes/control_flow/loop_node.rs @@ -180,6 +180,7 @@ mod tests { run: &Value::Null, nodes: &nodes, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await } diff --git a/src/nodes/control_flow/merge.rs b/src/nodes/control_flow/merge.rs index d443636..ae49dd3 100644 --- a/src/nodes/control_flow/merge.rs +++ b/src/nodes/control_flow/merge.rs @@ -53,6 +53,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let output = MergeNode.execute(ctx).await.expect("execute"); @@ -71,6 +72,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; MergeNode.execute(ctx).await.expect("execute").items } diff --git a/src/nodes/control_flow/split_out.rs b/src/nodes/control_flow/split_out.rs index e50d01d..8b0ad9a 100644 --- a/src/nodes/control_flow/split_out.rs +++ b/src/nodes/control_flow/split_out.rs @@ -88,6 +88,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let output = SplitOutNode.execute(ctx).await.expect("execute"); @@ -110,6 +111,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let output = SplitOutNode.execute(ctx).await.expect("execute"); @@ -131,6 +133,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let output = SplitOutNode.execute(ctx).await.expect("execute"); @@ -149,6 +152,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; SplitOutNode.execute(ctx).await.expect("execute").items } diff --git a/src/nodes/control_flow/switch.rs b/src/nodes/control_flow/switch.rs index 6418960..b9edd88 100644 --- a/src/nodes/control_flow/switch.rs +++ b/src/nodes/control_flow/switch.rs @@ -91,6 +91,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = SwitchNode.execute(ctx).await.expect("execute"); (out.port.expect("switch always routes to a port"), out.items) @@ -112,6 +113,7 @@ mod tests { run: &run, nodes: &nodes, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = SwitchNode.execute(ctx).await.expect("execute"); assert_eq!(out.port.as_deref(), Some("urgent")); diff --git a/src/nodes/control_flow/transform.rs b/src/nodes/control_flow/transform.rs index a843c20..df60ea7 100644 --- a/src/nodes/control_flow/transform.rs +++ b/src/nodes/control_flow/transform.rs @@ -86,6 +86,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; TransformNode.execute(ctx).await.expect("execute").items } @@ -110,6 +111,7 @@ mod tests { run: &run, nodes: &nodes, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = TransformNode.execute(ctx).await.expect("execute").items; assert_eq!(out[0].json["who"], json!("a@b.com")); diff --git a/src/nodes/integration/agent.rs b/src/nodes/integration/agent.rs index 860207e..015ce32 100644 --- a/src/nodes/integration/agent.rs +++ b/src/nodes/integration/agent.rs @@ -291,6 +291,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await .expect("execute"); @@ -306,6 +307,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await .expect("execute"); @@ -327,6 +329,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = AgentNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1); @@ -352,6 +355,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = AgentNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["completion"]["prompt"], "X"); @@ -369,6 +373,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = AgentNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["connection"], Value::Null); @@ -392,6 +397,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = AgentNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1); @@ -419,6 +425,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps, + token: crate::engine::CancellationToken::new(), }; AgentNode .execute(ctx) @@ -579,6 +586,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = AgentNode .execute(ctx) diff --git a/src/nodes/integration/code.rs b/src/nodes/integration/code.rs index 5f2a9c1..dcc1385 100644 --- a/src/nodes/integration/code.rs +++ b/src/nodes/integration/code.rs @@ -109,6 +109,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; CodeNode.execute(ctx).await.expect("execute").items } diff --git a/src/nodes/integration/http_request.rs b/src/nodes/integration/http_request.rs index 2096a17..97782dc 100644 --- a/src/nodes/integration/http_request.rs +++ b/src/nodes/integration/http_request.rs @@ -139,6 +139,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = HttpRequestNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1); @@ -175,6 +176,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = HttpRequestNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["request"]["url"], "https://a"); @@ -202,6 +204,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = HttpRequestNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["connection"], Value::Null); diff --git a/src/nodes/integration/memory.rs b/src/nodes/integration/memory.rs index 63f35e7..62a51c7 100644 --- a/src/nodes/integration/memory.rs +++ b/src/nodes/integration/memory.rs @@ -404,6 +404,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = MemoryNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 2, "per_item default maps over input"); @@ -458,6 +459,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = MemoryNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["opts"]["operation"], "search"); @@ -476,6 +478,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = MemoryNode .execute(ctx) @@ -501,6 +504,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = MemoryNode .execute(ctx) @@ -543,6 +547,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = MemoryNode .execute(ctx) @@ -567,6 +572,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = MemoryNode .execute(ctx) @@ -591,6 +597,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = MemoryNode .execute(ctx) @@ -619,6 +626,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = MemoryNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1, "once mode emits a single item"); diff --git a/src/nodes/integration/output_parser.rs b/src/nodes/integration/output_parser.rs index 772d1fa..49bcdc1 100644 --- a/src/nodes/integration/output_parser.rs +++ b/src/nodes/integration/output_parser.rs @@ -82,6 +82,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = OutputParserNode.execute(ctx).await.expect("execute"); assert_eq!(out.items, input); @@ -109,6 +110,7 @@ mod tests { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; OutputParserNode.execute(ctx).await.expect("execute").items } @@ -171,6 +173,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps, + token: crate::engine::CancellationToken::new(), }; OutputParserNode.execute(ctx).await.map(|o| o.items) } diff --git a/src/nodes/integration/shell_tests.rs b/src/nodes/integration/shell_tests.rs index b69b5a3..a1b8cae 100644 --- a/src/nodes/integration/shell_tests.rs +++ b/src/nodes/integration/shell_tests.rs @@ -36,6 +36,7 @@ async fn execute_with(caps: Capabilities, config: Value) -> Result { run: &Value::Null, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await } diff --git a/src/nodes/integration/sub_workflow.rs b/src/nodes/integration/sub_workflow.rs index 50a26cd..8205dbe 100644 --- a/src/nodes/integration/sub_workflow.rs +++ b/src/nodes/integration/sub_workflow.rs @@ -184,16 +184,36 @@ impl NodeExecutor for SubWorkflowNode { // so `=item.x` addresses the element this run is for, and // receives that single item as its input. let scope = crate::nodes::expr_scope_for(ctx, item.json.clone()); - let child = run_child(ctx, &scope, std::slice::from_ref(item)).await?; + // `run_child` yields `None` only when the parent cancelled + // this run mid-child (`ctx.token` is then set). The map slots + // exactly one output per input index, so stand in with an + // empty item — the whole node's output is discarded by the + // token check below, so this placeholder never surfaces. + let child = run_child(ctx, &scope, std::slice::from_ref(item)) + .await? + .unwrap_or_else(|| crate::data::Item::new(Value::Null)); Ok((child, vec![])) }) .await?; + // Parent-initiated cancel: wind down with no output, mirroring the + // top-level cancelled-node contract. `ctx.token` is a one-way flag, + // so if any child wound down (returned `None`) it is set here; the + // parent's next boundary check sees the same flip and settles + // `cancelled = true`. + if ctx.token.is_cancelled() { + return Ok(NodeOutput::empty()); + } return Ok(NodeOutput::main(items)); } let scope = crate::nodes::expr_scope(&ctx); - let item = run_child(&ctx, &scope, ctx.input).await?; - Ok(NodeOutput::main(vec![item])) + match run_child(&ctx, &scope, ctx.input).await? { + Some(item) => Ok(NodeOutput::main(vec![item])), + // Parent-initiated cancel wound the child down: emit nothing, the + // same clean wind-down a top-level cancelled node performs. The + // parent's next boundary check settles `cancelled = true`. + None => Ok(NodeOutput::empty()), + } } } @@ -203,11 +223,17 @@ impl NodeExecutor for SubWorkflowNode { /// `scope` is the expression scope `workflow_id` is resolved against (the whole /// input for `once`, the current element for `per_item`), and `child_input` is /// the item array seeded into the child run. +/// +/// Returns `Ok(None)` when the parent run cancelled this child mid-flight +/// (`ctx.token` is set): the child is a clean cooperative wind-down, not a +/// failure, so it emits no item and lets the parent settle as cancelled. A child +/// that stops for any *other* reason (a `requires_approval` pause, or a cancel +/// arriving through a channel independent of the parent's token) still errors. async fn run_child( ctx: &NodeContext<'_>, scope: &Value, child_input: &[crate::data::Item], -) -> Result { +) -> Result> { // The inline `workflow` graph carries its *own* `=`-expressions, scoped // to the CHILD run — it must pass through untouched. Only the fields the // sub_workflow node itself reads (here `workflow_id`) are resolved @@ -271,12 +297,17 @@ async fn run_child( // once at the call site. let child_inputs = child_inputs(&ctx.node.config, scope)?; // Box the recursive engine call so the async future type stays sized. + // Forward the parent run's cancellation token: cancelling the parent must + // wind down this child too, rather than letting it run on orphaned behind a + // fresh token. The child threads it into its own node contexts, so the whole + // nesting chain shares one cancellation signal. let outcome = Box::pin(crate::engine::run_sub_workflow( &compiled, crate::engine::RunInput::new(trigger).with_inputs(child_inputs), ctx.caps, child_depth, depth_cap, + ctx.token.clone(), )) .await?; @@ -313,6 +344,28 @@ async fn run_child( ))); } if outcome.cancelled { + // Two cancellations look the same on the child's `RunOutcome` but mean + // opposite things to the parent, so split on *who* cancelled: + // + // - The parent's own token is set: this is a cooperative wind-down of + // the whole run (the parent is being cancelled and forwarded the same + // token in, per `run_sub_workflow`). Halting with an error here would + // turn a clean cancel into a spurious failure. Emit nothing and let + // the parent settle: its next node-boundary check sees the same + // flipped token and reports `cancelled = true`, exactly as a + // top-level cancelled node does. + // - The parent's token is NOT set, yet the child still reports + // cancelled: the child was cancelled through some channel independent + // of this run (none exists today — the only token a child receives is + // the parent's clone — but keep the arm so a future independent-cancel + // path can never be silently treated as a completed child). + if ctx.token.is_cancelled() { + tracing::debug!( + node = %ctx.node.id, + "sub_workflow: child wound down under the parent's cancellation; emitting no output" + ); + return Ok(None); + } return Err(EngineError::Capability(format!( "sub_workflow node {:?}: child run was cancelled before completing; the parent \ run is halted rather than falsely completed", @@ -320,7 +373,7 @@ async fn run_child( ))); } - Ok(crate::data::Item::new(outcome.output)) + Ok(Some(crate::data::Item::new(outcome.output))) } #[cfg(test)] @@ -362,6 +415,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; SubWorkflowNode .execute(ctx) @@ -384,6 +438,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps, + token: crate::engine::CancellationToken::new(), }; SubWorkflowNode.execute(ctx).await.expect("execute") } @@ -647,6 +702,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps, + token: crate::engine::CancellationToken::new(), }; SubWorkflowNode.execute(ctx).await } @@ -815,3 +871,334 @@ mod tests { ); } } + +/// Cross-boundary cancellation: a parent run's [`CancellationToken`] must reach +/// its `sub_workflow` children so a parent cancel winds the whole subtree down +/// instead of orphaning it behind a fresh token. These pin the propagation +/// end-to-end through a real `run_cancellable` drive (T1–T5). +/// +/// The mid-flight cancel is made **deterministic under parallel test load** by +/// having the `slow` node hold the run at its boundary until the token actually +/// flips (a bounded spin), rather than racing a wall-clock sleep against the +/// scheduler — so the boundary check before `marker` is guaranteed to observe +/// the cancellation. +#[cfg(test)] +mod cancellation_propagation_tests { + use std::sync::{Arc, Mutex}; + use std::time::{Duration, Instant}; + + use async_trait::async_trait; + use serde_json::{Value, json}; + use tokio::sync::Notify; + use tokio::time::sleep; + + use crate::caps::mock::mock_capabilities; + use crate::caps::{Capabilities, ToolInvoker}; + use crate::compiler::compile; + use crate::engine::{CancellationToken, RunOutcome, run_cancellable}; + use crate::error::Result; + use crate::model::{Edge, Node, NodeKind, WorkflowGraph}; + + /// The bounded spin `slow` uses to wait for cancellation, and the cap on + /// `marker`'s sleep — both generous enough to never trip under load, small + /// enough that a broken build (where `marker` runs) is still obvious. + const SPIN_CAP_MS: u64 = 5_000; + + /// A [`ToolInvoker`] that records every slug it runs and lets a test suspend a + /// run *inside* a node. `slow` fires `slow_started` and then — when + /// `block_until_cancel` is set — holds the run at that node until `run_token` + /// flips (bounded by [`SPIN_CAP_MS`]); this pins the cancel to land while + /// `slow` is mid-flight, with no dependency on scheduler timing. `marker` is a + /// node that must never run once cancellation has propagated: its appearance + /// in `invoked` (and its long sleep in the elapsed time) is what a broken + /// propagation reveals. + #[derive(Clone)] + struct ProbeTools { + invoked: Arc>>, + slow_started: Arc, + run_token: CancellationToken, + block_until_cancel: bool, + marker_ms: u64, + } + + #[async_trait] + impl ToolInvoker for ProbeTools { + async fn invoke(&self, slug: &str, _args: Value, _conn: Option<&str>) -> Result { + self.invoked + .lock() + .expect("invoked mutex") + .push(slug.to_string()); + match slug { + "slow" => { + self.slow_started.notify_one(); + // Hold the child at this node until the run is cancelled, so + // the boundary check before `marker` deterministically sees + // the flip. Bounded so a broken build cannot hang CI. + if self.block_until_cancel { + let mut waited = 0; + while !self.run_token.is_cancelled() && waited < SPIN_CAP_MS { + sleep(Duration::from_millis(1)).await; + waited += 1; + } + } + } + // When cancellation propagated, `marker` never runs. A cancel test + // sets `marker_ms` long so a *broken* propagation (marker runs) + // blows the elapsed bound too; the uncancelled control sets it to 0 + // so the run stays fast when `marker` legitimately runs. + "marker" => sleep(Duration::from_millis(self.marker_ms)).await, + _ => {} + } + Ok(json!({ "tool": slug })) + } + } + + fn node(id: &str, kind: NodeKind) -> Node { + Node { + id: id.to_string(), + kind, + type_version: 1, + name: id.to_string(), + config: Value::Null, + ports: Vec::new(), + position: None, + } + } + + fn edge(from: &str, to: &str) -> Edge { + Edge { + from_node: from.to_string(), + from_port: "main".to_string(), + to_node: to.to_string(), + to_port: "main".to_string(), + } + } + + /// A single-invocation `tool_call` node bound to `slug`. + fn tool(id: &str, slug: &str) -> Node { + let mut n = node(id, NodeKind::ToolCall); + n.config = json!({ "slug": slug, "execution": "once" }); + n + } + + /// `trigger -> slow -> marker`: the innermost chain every test cancels within. + fn slow_then_marker() -> WorkflowGraph { + WorkflowGraph { + nodes: vec![ + node("ct", NodeKind::Trigger), + tool("slow", "slow"), + tool("marker", "marker"), + ], + edges: vec![edge("ct", "slow"), edge("slow", "marker")], + ..Default::default() + } + } + + /// `trigger -> sub_workflow(child)`, with `child` embedded inline. + fn wrap_inline(child: &WorkflowGraph) -> WorkflowGraph { + let inline = serde_json::to_value(child).expect("serialize child graph"); + let mut sw = node("sw", NodeKind::SubWorkflow); + sw.config = json!({ "workflow": inline, "execution": "once" }); + WorkflowGraph { + nodes: vec![node("pt", NodeKind::Trigger), sw], + edges: vec![edge("pt", "sw")], + ..Default::default() + } + } + + fn probe( + token: &CancellationToken, + block_until_cancel: bool, + marker_ms: u64, + ) -> (Capabilities, Arc>>, Arc) { + let invoked = Arc::new(Mutex::new(Vec::new())); + let slow_started = Arc::new(Notify::new()); + let tools = ProbeTools { + invoked: invoked.clone(), + slow_started: slow_started.clone(), + run_token: token.clone(), + block_until_cancel, + marker_ms, + }; + let caps = Capabilities { + tools: Arc::new(tools), + ..mock_capabilities() + }; + (caps, invoked, slow_started) + } + + /// Drives `run_cancellable` to completion while **actively polling** the run + /// future, cancelling `token` the moment the innermost `slow` node starts. + /// Polling matters: an un-awaited run future is the documented trap that makes + /// a cancellation test hollow, so the future is raced against the start signal + /// rather than cancelled blind. + async fn run_cancelling_on_slow( + graph: &WorkflowGraph, + caps: &Capabilities, + token: CancellationToken, + slow_started: Arc, + ) -> (RunOutcome, Duration) { + let compiled = compile(graph).expect("compile"); + let fut = run_cancellable(&compiled, json!({}), caps, token.clone()); + tokio::pin!(fut); + let notified = slow_started.notified(); + tokio::pin!(notified); + let mut cancelled = false; + let started = Instant::now(); + let outcome = loop { + tokio::select! { + out = &mut fut => break out.expect("cancelled run still returns Ok"), + // Guarded so the one-shot start signal is not polled after it fires. + () = &mut notified, if !cancelled => { + token.cancel(); + cancelled = true; + } + } + }; + assert!( + cancelled, + "the `slow` node must have started so the cancel landed mid-flight" + ); + (outcome, started.elapsed()) + } + + // T1 — repro→green. Parent -> sub_workflow(child), child is + // trigger -> slow -> marker. Cancelling mid-`slow` must wind the child down: + // `slow` (already running) completes, `marker` never runs, and the run + // settles cancelled. Before the fix the child ran behind a fresh token, so + // the cancel never crossed the boundary and `marker` executed. + #[tokio::test] + async fn t1_parent_cancel_stops_child_before_marker() { + let token = CancellationToken::new(); + let (caps, invoked, slow_started) = probe(&token, true, SPIN_CAP_MS); + let graph = wrap_inline(&slow_then_marker()); + + let (outcome, elapsed) = run_cancelling_on_slow(&graph, &caps, token, slow_started).await; + + assert!(outcome.cancelled, "the parent run should report cancelled"); + let slugs = invoked.lock().expect("invoked mutex").clone(); + assert!( + slugs.contains(&"slow".to_string()), + "the in-flight `slow` node ran: {slugs:?}" + ); + assert!( + !slugs.contains(&"marker".to_string()), + "cancellation must reach the child: `marker` should never run, got {slugs:?}" + ); + // Elapsed is bounded by the in-flight `slow` node's remainder, not by + // `marker`'s long sleep — the run did not wait on the skipped node. + assert!( + elapsed < Duration::from_millis(2_000), + "wind-down should be bounded by `slow`, not `marker`; took {elapsed:?}" + ); + } + + // T2 — transitive inheritance at depth ≥ 2. Parent -> sub -> sub, the + // innermost being trigger -> slow -> marker. Cancelling mid-innermost-`slow` + // must still skip the innermost `marker`: each level forwards the same token + // down through its own node contexts. + #[tokio::test] + async fn t2_cancel_propagates_through_two_levels() { + let token = CancellationToken::new(); + let (caps, invoked, slow_started) = probe(&token, true, SPIN_CAP_MS); + let inner = wrap_inline(&slow_then_marker()); + let graph = wrap_inline(&inner); + + let (outcome, elapsed) = run_cancelling_on_slow(&graph, &caps, token, slow_started).await; + + assert!(outcome.cancelled, "the top run should report cancelled"); + let slugs = invoked.lock().expect("invoked mutex").clone(); + assert!( + slugs.contains(&"slow".to_string()), + "innermost `slow` ran: {slugs:?}" + ); + assert!( + !slugs.contains(&"marker".to_string()), + "the token must reach depth 2: innermost `marker` should never run, got {slugs:?}" + ); + assert!( + elapsed < Duration::from_millis(2_000), + "two-level wind-down should still be bounded by `slow`; took {elapsed:?}" + ); + } + + // T3 — the guardrail: an *uncancelled* run of the identical graph must be + // untouched, so the fix cannot be "cancel everything". Both child slugs run + // and the outcome is not cancelled. `block_until_cancel` is off so `slow` + // returns immediately (nothing will ever cancel it). + #[tokio::test] + async fn t3_uncancelled_child_runs_to_completion() { + let token = CancellationToken::new(); + let (caps, invoked, _slow_started) = probe(&token, false, 0); + let graph = wrap_inline(&slow_then_marker()); + let compiled = compile(&graph).expect("compile"); + + let outcome = run_cancellable(&compiled, json!({}), &caps, token) + .await + .expect("run"); + + assert!( + !outcome.cancelled, + "an uncancelled run must not report cancelled" + ); + let slugs = invoked.lock().expect("invoked mutex").clone(); + assert!( + slugs.contains(&"slow".to_string()), + "`slow` should run: {slugs:?}" + ); + assert!( + slugs.contains(&"marker".to_string()), + "`marker` must still run when nothing cancels: {slugs:?}" + ); + } + + // T4 — a token already cancelled before the run starts: the parent's + // `sub_workflow` node short-circuits at its own boundary, so the child never + // starts and neither tool is invoked. + #[tokio::test] + async fn t4_pre_cancelled_token_never_starts_child() { + let token = CancellationToken::new(); + let (caps, invoked, _slow_started) = probe(&token, true, 0); + let graph = wrap_inline(&slow_then_marker()); + let compiled = compile(&graph).expect("compile"); + + token.cancel(); + let outcome = run_cancellable(&compiled, json!({}), &caps, token) + .await + .expect("pre-cancelled run still returns Ok"); + + assert!(outcome.cancelled, "a pre-cancelled run reports cancelled"); + let slugs = invoked.lock().expect("invoked mutex").clone(); + assert!( + slugs.is_empty(), + "the sub_workflow node short-circuits, so no child tool runs: {slugs:?}" + ); + } + + // T5 — the defensive arm. When a child reports cancelled but the parent's own + // token is *not* set, `run_child` still errors rather than silently treating + // the child as completed. That state is unreachable through `run_sub_workflow` + // today (a child only ever receives the parent's own token clone, so a + // cancelled child implies a cancelled parent token — see T1), so this pins the + // arm at the source level: a future refactor cannot delete it and let an + // independently-cancelled child fall through as a false completion. + #[test] + fn t5_defensive_independent_cancel_arm_is_present() { + // Scope the check to the production region — everything before the test + // module. `include_str!` pulls the whole file, so an unscoped `contains` + // would match the assertion strings *in this test itself* and pass even if + // the production arm were deleted. Slicing at `mod tests` (the sole such + // marker) makes deleting the arm from `run_child` actually fail this test. + let src = include_str!("sub_workflow.rs"); + let production = src + .split("mod tests") + .next() + .expect("source file has a body before its test module"); + assert!( + production.contains("if ctx.token.is_cancelled() {") + && production.contains("run is halted rather than falsely completed"), + "run_child must keep BOTH the parent-cancel wind-down (Ok(None)) and the \ + defensive independent-cancel error arm" + ); + } +} diff --git a/src/nodes/integration/tool_call.rs b/src/nodes/integration/tool_call.rs index a6d4258..a5832b5 100644 --- a/src/nodes/integration/tool_call.rs +++ b/src/nodes/integration/tool_call.rs @@ -174,6 +174,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let err = ToolCallNode .execute(ctx) @@ -199,6 +200,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ToolCallNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["tool"], "x.y"); @@ -223,6 +225,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ToolCallNode.execute(ctx).await.expect("execute"); assert_eq!(out.items[0].json["json"]["args"]["to"], Value::Null); @@ -243,6 +246,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ToolCallNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1); @@ -269,6 +273,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ToolCallNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 3, "one output per input item"); @@ -296,6 +301,7 @@ mod tests { run: &run_meta, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let out = ToolCallNode.execute(ctx).await.expect("execute"); assert_eq!(out.items.len(), 1, "once mode emits a single item"); diff --git a/src/nodes/mod.rs b/src/nodes/mod.rs index 9ca43c7..1174fca 100644 --- a/src/nodes/mod.rs +++ b/src/nodes/mod.rs @@ -15,6 +15,7 @@ use serde_json::Value; use crate::caps::Capabilities; use crate::data::Item; +use crate::engine::CancellationToken; use crate::error::Result; use crate::model::{Node, NodeKind}; @@ -38,6 +39,15 @@ pub struct NodeContext<'a> { pub nodes: &'a Value, /// Host-provided capabilities. pub caps: &'a Capabilities, + /// The run's cooperative-cancellation token (see + /// [`crate::engine::CancellationToken`]). An **owned clone** of the run + /// token, not a borrow — an executor that spawns nested engine work (today + /// only [`sub_workflow`](crate::nodes::integration)) must thread a clone + /// into that child run, so a parent cancel winds the whole subtree down at + /// the next node boundary instead of orphaning it. Executors that touch the + /// outside world within a single node need not consult it; the engine + /// already checks it at the node boundary before this node runs. + pub token: CancellationToken, } /// Builds the expression scope for a node from its runtime [`NodeContext`]. @@ -430,6 +440,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await; assert!( @@ -453,6 +464,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }) .await .expect("execute"); @@ -481,6 +493,7 @@ mod tests { run: &run, nodes: &nodes_state, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let scope = expr_scope(&ctx); // Existing keys unchanged (back-compat). @@ -510,6 +523,7 @@ mod tests { run: &run, nodes: &Value::Null, caps: &caps, + token: crate::engine::CancellationToken::new(), }; let scope = expr_scope(&ctx); assert_eq!(scope["nodes"], json!({}));