fix: consume internal tool batch cancellation

This commit is contained in:
2026-09-04 13:01:37 +09:00
parent 42d109cae3
commit 5d61da481b
2 changed files with 30 additions and 6 deletions
+10
View File
@@ -1390,6 +1390,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
let mut pause_requested = false; let mut pause_requested = false;
let mut pause_deadline = None; let mut pause_deadline = None;
let mut batch_error = None; let mut batch_error = None;
let mut locally_enqueued_cancel = false;
for result in synthetic_results { for result in synthetic_results {
if let Err(error) = self if let Err(error) = self
.finalize_and_commit_tool_result( .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; let mut futures = futures;
if batch_error.is_some() && !futures.is_empty() { if batch_error.is_some() && !futures.is_empty() {
let _ = self.cancel_tx.try_send(()); let _ = self.cancel_tx.try_send(());
locally_enqueued_cancel = true;
} }
while !futures.is_empty() { while !futures.is_empty() {
tokio::select! { tokio::select! {
@@ -1435,6 +1437,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
} }
if !futures.is_empty() { if !futures.is_empty() {
let _ = self.cancel_tx.try_send(()); 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() => { _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => {
pause_deadline = None; pause_deadline = None;
let _ = self.cancel_tx.try_send(()); let _ = self.cancel_tx.try_send(());
locally_enqueued_cancel = true;
} }
cancel = self.cancel_rx.recv() => { cancel = self.cancel_rx.recv() => {
if cancel.is_some() { 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 { if let Some(error) = batch_error {
self.timeline.abort_current_block(); self.timeline.abort_current_block();
return Err(error); return Err(error);
+20 -6
View File
@@ -1367,6 +1367,7 @@ impl Interceptor for StopFirstParallelResult {
if info.call.id != "call_fast" { if info.call.id != "call_fast" {
return Ok(PostToolAction::Continue); return Ok(PostToolAction::Continue);
} }
tokio::time::sleep(Duration::from_millis(5)).await;
match self.0 { match self.0 {
PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())), PostToolStopMode::Abort => Ok(PostToolAction::Abort("stop parallel batch".to_string())),
PostToolStopMode::Failure => Err(InterceptorError::new( PostToolStopMode::Failure => Err(InterceptorError::new(
@@ -1380,20 +1381,29 @@ impl Interceptor for StopFirstParallelResult {
#[tokio::test] #[tokio::test]
async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning() { async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning() {
for mode in [PostToolStopMode::Abort, PostToolStopMode::Failure] { 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_use_start(0, "call_fast", "fast"),
Event::tool_input_delta(0, r#"{}"#), Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0), 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_input_delta(1, r#"{}"#),
Event::tool_use_stop(1), Event::tool_use_stop(1),
Event::Status(StatusEvent { Event::Status(StatusEvent {
status: ResponseStatus::Completed, 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); let mut engine = Engine::new(client);
engine.register_tool(SlowTool::new("fast", 1).definition()); engine.register_tool(SlowTool::new("fast", 0).definition());
engine.register_tool(SlowTool::new("slow", 10_000).definition()); engine.register_tool(SlowTool::new("ready", 1).definition());
engine.set_interceptor(StopFirstParallelResult(mode)); engine.set_interceptor(StopFirstParallelResult(mode));
let mut history = History::new(); let mut history = History::new();
@@ -1422,6 +1432,10 @@ async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning(
.collect(); .collect();
assert_eq!(terminal_ids.len(), 2); assert_eq!(terminal_ids.len(), 2);
assert!(terminal_ids.contains(&"call_fast")); 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));
} }
} }