From 9b48b1ff5d43f57ff38e932aa5ab269415f00c30 Mon Sep 17 00:00:00 2001 From: Hare Date: Fri, 4 Sep 2026 12:27:59 +0900 Subject: [PATCH] fix: terminalize parallel tool siblings on interceptor stop --- crates/agen/src/engine.rs | 63 +++++++++++---- crates/agen/tests/parallel_execution_test.rs | 80 +++++++++++++++++++- 2 files changed, 125 insertions(+), 18 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 644cf161..6c23e0cd 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1389,20 +1389,29 @@ impl Engine { let mut terminal_call_ids = HashSet::new(); let mut pause_requested = false; let mut pause_deadline = None; + let mut batch_error = None; for result in synthetic_results { - self.finalize_and_commit_tool_result( - history, - annotate, - result, - None, - &call_info_map, - &mut attempt_fence, - &mut terminal_call_ids, - ) - .await?; + if let Err(error) = self + .finalize_and_commit_tool_result( + history, + annotate, + result, + None, + &call_info_map, + &mut attempt_fence, + &mut terminal_call_ids, + ) + .await + && batch_error.is_none() + { + batch_error = Some(error); + } } let mut futures = futures; + if batch_error.is_some() && !futures.is_empty() { + let _ = self.cancel_tx.try_send(()); + } while !futures.is_empty() { tokio::select! { // If cancellation and a completed result are both ready, drain @@ -1412,7 +1421,7 @@ impl Engine { result = futures.next() => { let (attempt_id, result) = result.expect("non-empty FuturesUnordered returns a result"); - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, result, @@ -1420,7 +1429,14 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await { + if batch_error.is_none() { + batch_error = Some(error); + } + if !futures.is_empty() { + let _ = self.cancel_tx.try_send(()); + } + } } pause = self.pause_rx.recv(), if !pause_requested => { if pause.is_some() { @@ -1482,7 +1498,7 @@ impl Engine { result = futures.next() => { let (attempt_id, result) = result.expect("non-empty FuturesUnordered returns a result"); - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, result, @@ -1490,7 +1506,11 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await + && batch_error.is_none() + { + batch_error = Some(error); + } } _ = tokio::time::sleep_until(deadline) => break, } @@ -1504,7 +1524,7 @@ impl Engine { if let Some(handle) = execution_handles.get(call_id) { handle.force_close(); } - self.finalize_and_commit_tool_result( + if let Err(error) = self.finalize_and_commit_tool_result( history, annotate, ToolResult::outcome_unknown(call_id), @@ -1512,11 +1532,18 @@ impl Engine { &call_info_map, &mut attempt_fence, &mut terminal_call_ids, - ).await?; + ).await + && batch_error.is_none() + { + batch_error = Some(error); + } } } self.timeline.abort_current_block(); + if let Some(error) = batch_error.take() { + return Err(error); + } if pause_requested { return Ok(ToolExecutionResult::Paused); } @@ -1525,6 +1552,10 @@ impl Engine { } } + if let Some(error) = batch_error { + self.timeline.abort_current_block(); + return Err(error); + } Ok(if pause_requested { ToolExecutionResult::Paused } else { diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index c79c0a2e..3b7c9c18 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -7,8 +7,8 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use agen::interceptor::{ - Interceptor, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, PostToolAction, - PreToolAction, ToolCallInfo, ToolResultInfo, + Interceptor, InterceptorError, InterceptorErrorCategory, InterceptorPhase, InterceptorResult, + PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo, }; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::tool::{ @@ -1349,3 +1349,79 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() { } if call_id == "call_confirmed" ))); } + +#[derive(Clone, Copy)] +enum PostToolStopMode { + Abort, + Failure, +} + +struct StopFirstParallelResult(PostToolStopMode); + +#[async_trait] +impl Interceptor for StopFirstParallelResult { + async fn post_tool_call( + &self, + info: &ToolResultInfo<'_, ()>, + ) -> InterceptorResult { + if info.call.id != "call_fast" { + return Ok(PostToolAction::Continue); + } + match self.0 { + PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())), + PostToolStopMode::Failure => Err(InterceptorError::new( + InterceptorErrorCategory::Policy, + "reject parallel batch", + )), + } + } +} + +#[tokio::test] +async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning() { + for mode in [PostToolStopMode::Abort, PostToolStopMode::Failure] { + let client = MockLlmClient::new(vec![ + Event::tool_use_start(0, "call_fast", "fast"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::tool_use_start(1, "call_slow", "slow"), + Event::tool_input_delta(1, r#"{}"#), + Event::tool_use_stop(1), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]); + let mut engine = Engine::new(client); + engine.register_tool(SlowTool::new("fast", 1).definition()); + engine.register_tool(SlowTool::new("slow", 10_000).definition()); + engine.set_interceptor(StopFirstParallelResult(mode)); + let mut history = History::new(); + + let output = engine.run(&mut history, "parallel stop").await; + match mode { + PostToolStopMode::Abort => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Aborted(ref reason) + )) if reason == "stop parallel batch" + )), + PostToolStopMode::Failure => assert!(matches!( + output.result, + EngineRunExit::Interrupted(RunInterruptionReason::Unexpected( + EngineError::Interceptor(ref failure) + )) if failure.phase() == InterceptorPhase::PostToolCall + )), + } + + let terminal_ids: Vec<_> = history + .iter() + .filter_map(|entry| match &entry.item { + Item::ToolResult { call_id, .. } => Some(call_id.as_str()), + _ => None, + }) + .collect(); + assert_eq!(terminal_ids.len(), 2); + assert!(terminal_ids.contains(&"call_fast")); + assert!(terminal_ids.contains(&"call_slow")); + } +}