fix(graph): inherit the session abort in agent-node children

A child agent spawned from a graph node now carries its parent's session
abort signal, so a user interrupt is observable from every sync wait in
the agent tree. Background agent__spawn children are unchanged.

run_child_agent_with_graph_inputs also guards the child's own supervisor
with a disarmable CancelOnDrop: if the future is dropped mid-flight or
the child fails before its normal cleanup, its spawned subagents are
cancelled instead of running on invisibly.
This commit is contained in:
Dark-Alex-17 committed 2026-09-11 10:56:27 -06:00
1 parent aa1287b902
commit 6aad82c6fd
1 file changed
+51 -1
+51 -1
View File
@@ -543,6 +543,26 @@ pub fn run_child_agent(
run_child_agent_with_graph_inputs(child_ctx, initial_input, abort_signal, None)
}
/// Stops a child agent's own spawned subagents if the future running the
/// child is dropped mid-flight (parent timeout, abort) or the child fails
/// before reaching its own cleanup. Disarmed on the success paths, which
/// already cancel explicitly.
struct CancelOnDrop(Option<Arc<RwLock<Supervisor>>>);
impl CancelOnDrop {
fn disarm(&mut self) {
self.0.take();
}
}
impl Drop for CancelOnDrop {
fn drop(&mut self) {
if let Some(supervisor) = self.0.take() {
supervisor.read().cancel_recursive();
}
}
}
/// Runs a child agent to completion. When the child is a graph agent,
/// `graph_inputs` is overlaid on its `initial_state` before the graph
/// starts; it is ignored for config-only agents.
@@ -553,14 +573,19 @@ fn run_child_agent_with_graph_inputs(
graph_inputs: Option<HashMap<String, Value>>,
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
Box::pin(async move {
let mut cancel_on_drop = CancelOnDrop(child_ctx.supervisor.clone());
if graph::active_agent_graph_name(&child_ctx).is_some() {
return graph::run_active_agent_graph_with_inputs(
let result = graph::run_active_agent_graph_with_inputs(
&mut child_ctx,
&initial_input.text(),
abort_signal,
graph_inputs,
)
.await;
if result.is_ok() {
cancel_on_drop.disarm();
}
return result;
}
let mut accumulated_output = String::new();
@@ -620,6 +645,7 @@ fn run_child_agent_with_graph_inputs(
if let Some(supervisor) = child_ctx.supervisor.clone() {
supervisor.read().cancel_recursive();
}
cancel_on_drop.disarm();
Ok(accumulated_output)
})
@@ -774,6 +800,7 @@ pub async fn run_agent_for_graph(
Arc::clone(&child_inbox),
agent_id.clone(),
);
child_ctx.session_abort = parent_ctx.session_abort.clone();
child_ctx.rag = agent.rag();
child_ctx.agent = Some(agent);
if graph_inputs.is_some() && graph::active_agent_graph_name(&child_ctx).is_none() {
@@ -2025,6 +2052,29 @@ mod tests {
);
}
#[test]
fn cancel_on_drop_guard_cancels_when_armed_and_not_when_disarmed() {
run_async(async {
let mut ctx = ctx_with_supervisor(4, 3);
let armed_abort = register_running_agent(&mut ctx, "a1", "explore");
drop(CancelOnDrop(ctx.supervisor.clone()));
assert!(
armed_abort.aborted(),
"an armed guard cancels the supervisor's agents when dropped"
);
let mut ctx = ctx_with_supervisor(4, 3);
let disarmed_abort = register_running_agent(&mut ctx, "a2", "explore");
let mut guard = CancelOnDrop(ctx.supervisor.clone());
guard.disarm();
drop(guard);
assert!(
!disarmed_abort.aborted(),
"a disarmed guard leaves the supervisor's agents running"
);
});
}
fn auto_continue_ctx() -> RequestContext {
let app = AppState {
config: Arc::new(AppConfig {