From 5d61da481bc0f3a70c99a114656e8683230c7dd6 Mon Sep 17 00:00:00 2001
From: Hare <kei.hiracchi.0928@gmail.com>
Date: Fri, 4 Sep 2026 13:01:37 +0900
Subject: [PATCH] fix: consume internal tool batch cancellation

---
 crates/agen/src/engine.rs                    | 10 ++++++++
 crates/agen/tests/parallel_execution_test.rs | 26 +++++++++++++++-----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs
index 6c23e0cd..da3cead9 100644
--- a/crates/agen/src/engine.rs
+++ b/crates/agen/src/engine.rs
@@ -1390,6 +1390,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         let mut pause_requested = false;
         let mut pause_deadline = None;
         let mut batch_error = None;
+        let mut locally_enqueued_cancel = false;
         for result in synthetic_results {
             if let Err(error) = self
                 .finalize_and_commit_tool_result(
@@ -1411,6 +1412,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         let mut futures = futures;
         if batch_error.is_some() && !futures.is_empty() {
             let _ = self.cancel_tx.try_send(());
+            locally_enqueued_cancel = true;
         }
         while !futures.is_empty() {
             tokio::select! {
@@ -1435,6 +1437,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         }
                         if !futures.is_empty() {
                             let _ = self.cancel_tx.try_send(());
+                            locally_enqueued_cancel = true;
                         }
                     }
                 }
@@ -1453,6 +1456,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                 _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => {
                     pause_deadline = None;
                     let _ = self.cancel_tx.try_send(());
+                    locally_enqueued_cancel = true;
                 }
                 cancel = self.cancel_rx.recv() => {
                     if cancel.is_some() {
@@ -1552,6 +1556,12 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
             }
         }
 
+        // A result-biased ready sibling can empty the batch before the local
+        // cancel signal is selected. Never let that current-batch signal leak
+        // into the next run or resume call.
+        if locally_enqueued_cancel {
+            let _ = self.cancel_rx.try_recv();
+        }
         if let Some(error) = batch_error {
             self.timeline.abort_current_block();
             return Err(error);
diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs
index 3b7c9c18..5df9f386 100644
--- a/crates/agen/tests/parallel_execution_test.rs
+++ b/crates/agen/tests/parallel_execution_test.rs
@@ -1367,6 +1367,7 @@ impl Interceptor for StopFirstParallelResult {
         if info.call.id != "call_fast" {
             return Ok(PostToolAction::Continue);
         }
+        tokio::time::sleep(Duration::from_millis(5)).await;
         match self.0 {
             PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())),
             PostToolStopMode::Failure => Err(InterceptorError::new(
@@ -1380,20 +1381,29 @@ impl Interceptor for StopFirstParallelResult {
 #[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![
+        let first_response = 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_use_start(1, "call_ready", "ready"),
             Event::tool_input_delta(1, r#"{}"#),
             Event::tool_use_stop(1),
             Event::Status(StatusEvent {
                 status: ResponseStatus::Completed,
             }),
-        ]);
+        ];
+        let second_response = vec![
+            Event::text_block_start(0),
+            Event::text_delta(0, "next run completed"),
+            Event::text_block_stop(0, None),
+            Event::Status(StatusEvent {
+                status: ResponseStatus::Completed,
+            }),
+        ];
+        let client = MockLlmClient::with_responses(vec![first_response, second_response]);
         let mut engine = Engine::new(client);
-        engine.register_tool(SlowTool::new("fast", 1).definition());
-        engine.register_tool(SlowTool::new("slow", 10_000).definition());
+        engine.register_tool(SlowTool::new("fast", 0).definition());
+        engine.register_tool(SlowTool::new("ready", 1).definition());
         engine.set_interceptor(StopFirstParallelResult(mode));
         let mut history = History::new();
 
@@ -1422,6 +1432,10 @@ async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning(
             .collect();
         assert_eq!(terminal_ids.len(), 2);
         assert!(terminal_ids.contains(&"call_fast"));
-        assert!(terminal_ids.contains(&"call_slow"));
+        assert!(terminal_ids.contains(&"call_ready"));
+
+        let mut engine = output.engine;
+        let next = engine.run(&mut history, "next run").await;
+        assert!(matches!(next, EngineRunExit::Finished));
     }
 }