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)
|
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,
|
/// Runs a child agent to completion. When the child is a graph agent,
|
||||||
/// `graph_inputs` is overlaid on its `initial_state` before the graph
|
/// `graph_inputs` is overlaid on its `initial_state` before the graph
|
||||||
/// starts; it is ignored for config-only agents.
|
/// 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>>,
|
graph_inputs: Option<HashMap<String, Value>>,
|
||||||
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
) -> Pin<Box<dyn Future<Output = Result<String>> + Send>> {
|
||||||
Box::pin(async move {
|
Box::pin(async move {
|
||||||
|
let mut cancel_on_drop = CancelOnDrop(child_ctx.supervisor.clone());
|
||||||
if graph::active_agent_graph_name(&child_ctx).is_some() {
|
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,
|
&mut child_ctx,
|
||||||
&initial_input.text(),
|
&initial_input.text(),
|
||||||
abort_signal,
|
abort_signal,
|
||||||
graph_inputs,
|
graph_inputs,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
if result.is_ok() {
|
||||||
|
cancel_on_drop.disarm();
|
||||||
|
}
|
||||||
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut accumulated_output = String::new();
|
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() {
|
if let Some(supervisor) = child_ctx.supervisor.clone() {
|
||||||
supervisor.read().cancel_recursive();
|
supervisor.read().cancel_recursive();
|
||||||
}
|
}
|
||||||
|
cancel_on_drop.disarm();
|
||||||
|
|
||||||
Ok(accumulated_output)
|
Ok(accumulated_output)
|
||||||
})
|
})
|
||||||
@@ -774,6 +800,7 @@ pub async fn run_agent_for_graph(
|
|||||||
Arc::clone(&child_inbox),
|
Arc::clone(&child_inbox),
|
||||||
agent_id.clone(),
|
agent_id.clone(),
|
||||||
);
|
);
|
||||||
|
child_ctx.session_abort = parent_ctx.session_abort.clone();
|
||||||
child_ctx.rag = agent.rag();
|
child_ctx.rag = agent.rag();
|
||||||
child_ctx.agent = Some(agent);
|
child_ctx.agent = Some(agent);
|
||||||
if graph_inputs.is_some() && graph::active_agent_graph_name(&child_ctx).is_none() {
|
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 {
|
fn auto_continue_ctx() -> RequestContext {
|
||||||
let app = AppState {
|
let app = AppState {
|
||||||
config: Arc::new(AppConfig {
|
config: Arc::new(AppConfig {
|
||||||
|
|||||||
Reference in new issue
Block a user