fix(graph): build peer retirement guards before spawning their tasks

A task aborted before its first poll never runs its body, so a guard
constructed inside the future did not exist yet and the pre-provisioned
teammate identity stayed messageable. Built outside and moved in, the
guard drops with the unpolled future on both the frontier and map items.

Also rewords the two timeout: 0 validator warnings: the graph-level one
no longer claims max_loop_iterations always applies (it may be 0 too),
and the node-level one names the enclosing bounds that still end the
node instead of saying it runs until it returns.
This commit is contained in:
Dark-Alex-17 committed 2026-09-11 10:56:28 -06:00
1 parent 02e2227aa0
commit 24b5504e21
3 files changed
+82 -32

No files matched your search

+56 -13
View File
@@ -247,20 +247,21 @@ impl GraphExecutor {
#[cfg(test)]
let observer = frontier_observer.clone();
// Retires the teammate identity however this task ends:
// dropped explicitly right after step() on the normal path,
// by unwinding on the abort return, task abort, or panic.
// Built before the spawn so a task aborted before its first
// poll still drops it with the future.
let retire = registry_for_task
.zip(peer_id)
.map(|(registry, id)| PeerRetireGuard {
registry,
id,
#[cfg(test)]
observer,
});
let task = tokio::spawn(async move {
// Retires the teammate identity however this task ends:
// dropped explicitly right after step() on the normal
// path, by unwinding on the abort return, task abort, or
// panic.
let retire =
registry_for_task
.zip(peer_id)
.map(|(registry, id)| PeerRetireGuard {
registry,
id,
#[cfg(test)]
observer,
});
let retire = retire;
let _permit = sem_clone
.acquire()
.await
@@ -1722,6 +1723,48 @@ nodes:
assert!(!registry.is_finished(&id_b), "the sibling is untouched");
}
/// A task aborted before its first poll never runs its body, so a guard
/// built inside the future would never exist. Built outside and moved in,
/// it drops with the unpolled future and still retires the identity.
#[tokio::test]
async fn peer_retire_guard_fires_when_task_is_aborted_before_first_poll() {
let (registry, assignments) = provision_frontier_peers(vec![
("a".to_string(), "worker".to_string()),
("b".to_string(), "worker".to_string()),
]);
let registry = registry.expect("two flagged nodes should get a registry");
let id_a = assignments["a"].0.clone();
let id_b = assignments["b"].0.clone();
let outside = PeerRetireGuard {
registry: Arc::clone(&registry),
id: id_a.clone(),
observer: None,
};
let inside_registry = Arc::clone(&registry);
let inside_id = id_b.clone();
let task = tokio::spawn(async move {
let _outside = outside;
let _inside = PeerRetireGuard {
registry: inside_registry,
id: inside_id,
observer: None,
};
std::future::pending::<()>().await;
});
// No await between spawn and abort: on the current-thread test
// runtime the task has not been polled yet.
task.abort();
let err = task.await.expect_err("the task was aborted");
assert!(err.is_cancelled(), "{err}");
assert!(registry.is_finished(&id_a), "the moved-in guard retires");
assert!(
!registry.is_finished(&id_b),
"a guard built inside never existed"
);
}
/// A session already aborted when the run starts trips the synchronous
/// pre-check: no bridge task, no super-step, the start script never runs.
#[tokio::test]
+13 -11
View File
@@ -92,18 +92,20 @@ async fn run_map(
let sem = semaphore.clone();
let abort = step_ctx.abort_signal.clone();
// Cancellation net behind the chain runner's own retirement: a task
// aborted before or during its chain still leaves the item's identity
// finished. Built before the spawn so a task aborted before its first
// poll drops it with the future.
let retire = item_peers
.as_ref()
.map(|(registry, assignment)| PeerRetireGuard {
registry: Arc::clone(registry),
id: assignment.0.clone(),
#[cfg(test)]
observer: None,
});
let task = tokio::spawn(async move {
// Cancellation net behind the chain runner's own retirement: a
// task aborted before or during its chain still leaves the item's
// identity finished.
let _retire = item_peers
.as_ref()
.map(|(registry, assignment)| PeerRetireGuard {
registry: Arc::clone(registry),
id: assignment.0.clone(),
#[cfg(test)]
observer: None,
});
let _retire = retire;
let _permit = sem
.acquire()
.await
+13 -8
View File
@@ -628,8 +628,8 @@ impl GraphValidator {
if graph.settings.timeout == Some(0) {
result.warning(ValidationError::new(
"settings.timeout: 0 disables the graph wall-clock bound; the run \
ends only when the graph completes, is aborted, or hits \
max_loop_iterations",
ends only when the graph completes, is aborted, or another \
limit stops it",
));
}
for (node_id, node) in &graph.nodes {
@@ -655,7 +655,8 @@ impl GraphValidator {
node_id,
format!(
"timeout: 0 disables the wall-clock bound for {kind} node \
'{node_id}'; it runs until it returns"
'{node_id}'; it runs until it returns or an enclosing bound \
(settings.timeout, an outer agent node's timeout, or abort) ends it"
),
));
}
@@ -3293,7 +3294,7 @@ mod tests {
assert_eq!(
w[0].message,
"settings.timeout: 0 disables the graph wall-clock bound; the run ends only when \
the graph completes, is aborted, or hits max_loop_iterations"
the graph completes, is aborted, or another limit stops it"
);
}
@@ -3312,7 +3313,8 @@ mod tests {
assert_eq!(w[0].node_id.as_deref(), Some("a"));
assert_eq!(
w[0].message,
"timeout: 0 disables the wall-clock bound for agent node 'a'; it runs until it returns"
"timeout: 0 disables the wall-clock bound for agent node 'a'; it runs until it returns \
or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it"
);
}
@@ -3332,7 +3334,8 @@ mod tests {
assert_eq!(w[0].node_id.as_deref(), Some("s"));
assert_eq!(
w[0].message,
"timeout: 0 disables the wall-clock bound for script node 's'; it runs until it returns"
"timeout: 0 disables the wall-clock bound for script node 's'; it runs until it returns \
or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it"
);
}
@@ -3352,7 +3355,8 @@ mod tests {
assert_eq!(w[0].node_id.as_deref(), Some("l"));
assert_eq!(
w[0].message,
"timeout: 0 disables the wall-clock bound for llm node 'l'; it runs until it returns"
"timeout: 0 disables the wall-clock bound for llm node 'l'; it runs until it returns \
or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it"
);
}
@@ -3372,7 +3376,8 @@ mod tests {
assert_eq!(w[0].node_id.as_deref(), Some("r"));
assert_eq!(
w[0].message,
"timeout: 0 disables the wall-clock bound for rag node 'r'; it runs until it returns"
"timeout: 0 disables the wall-clock bound for rag node 'r'; it runs until it returns \
or an enclosing bound (settings.timeout, an outer agent node's timeout, or abort) ends it"
);
}