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:
3 files changed
+82
-32
No files matched your search
+56
-13
@@ -247,20 +247,21 @@ impl GraphExecutor {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
let observer = frontier_observer.clone();
|
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 {
|
let task = tokio::spawn(async move {
|
||||||
// Retires the teammate identity however this task ends:
|
let retire = retire;
|
||||||
// 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 _permit = sem_clone
|
let _permit = sem_clone
|
||||||
.acquire()
|
.acquire()
|
||||||
.await
|
.await
|
||||||
@@ -1722,6 +1723,48 @@ nodes:
|
|||||||
assert!(!registry.is_finished(&id_b), "the sibling is untouched");
|
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(®istry),
|
||||||
|
id: id_a.clone(),
|
||||||
|
observer: None,
|
||||||
|
};
|
||||||
|
let inside_registry = Arc::clone(®istry);
|
||||||
|
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
|
/// A session already aborted when the run starts trips the synchronous
|
||||||
/// pre-check: no bridge task, no super-step, the start script never runs.
|
/// pre-check: no bridge task, no super-step, the start script never runs.
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
|
|||||||
+13
-11
@@ -92,18 +92,20 @@ async fn run_map(
|
|||||||
let sem = semaphore.clone();
|
let sem = semaphore.clone();
|
||||||
let abort = step_ctx.abort_signal.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 {
|
let task = tokio::spawn(async move {
|
||||||
// Cancellation net behind the chain runner's own retirement: a
|
let _retire = retire;
|
||||||
// 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 _permit = sem
|
let _permit = sem
|
||||||
.acquire()
|
.acquire()
|
||||||
.await
|
.await
|
||||||
|
|||||||
+13
-8
@@ -628,8 +628,8 @@ impl GraphValidator {
|
|||||||
if graph.settings.timeout == Some(0) {
|
if graph.settings.timeout == Some(0) {
|
||||||
result.warning(ValidationError::new(
|
result.warning(ValidationError::new(
|
||||||
"settings.timeout: 0 disables the graph wall-clock bound; the run \
|
"settings.timeout: 0 disables the graph wall-clock bound; the run \
|
||||||
ends only when the graph completes, is aborted, or hits \
|
ends only when the graph completes, is aborted, or another \
|
||||||
max_loop_iterations",
|
limit stops it",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
for (node_id, node) in &graph.nodes {
|
for (node_id, node) in &graph.nodes {
|
||||||
@@ -655,7 +655,8 @@ impl GraphValidator {
|
|||||||
node_id,
|
node_id,
|
||||||
format!(
|
format!(
|
||||||
"timeout: 0 disables the wall-clock bound for {kind} node \
|
"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!(
|
assert_eq!(
|
||||||
w[0].message,
|
w[0].message,
|
||||||
"settings.timeout: 0 disables the graph wall-clock bound; the run ends only when \
|
"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].node_id.as_deref(), Some("a"));
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
w[0].message,
|
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].node_id.as_deref(), Some("s"));
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
w[0].message,
|
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].node_id.as_deref(), Some("l"));
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
w[0].message,
|
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].node_id.as_deref(), Some("r"));
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
w[0].message,
|
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"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user