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:
1 file changed
+51
-1
+51
-1
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user