From 74d0af93e3b80b795a66d0647226810f8fdc4055 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:38:38 +0900 Subject: [PATCH] fix: interrupt pending tool policies before provider admission --- crates/agen/src/engine.rs | 162 ++++++++++++++++++++++---------------- 1 file changed, 94 insertions(+), 68 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 16f76d790..beef6816a 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -71,6 +71,40 @@ pub enum EngineError { ToolAttemptFence(String), } +/// Policy completion is not permission to start work when an interruption was +/// queued during its final poll. Cancel wins when both requests are observed. +fn check_policy_interruption( + cancel_rx: &mut mpsc::Receiver<()>, + pause_rx: &mut mpsc::Receiver<()>, +) -> Result<(), EngineError> { + if cancel_rx.try_recv().is_ok() { + return Err(EngineError::Cancelled); + } + if pause_rx.try_recv().is_ok() { + return Err(EngineError::PauseRequested); + } + Ok(()) +} + +/// Poll one policy invocation, never a second copy of its side effects. Dropping +/// a pending policy stops admission, not any already-owned tool execution; the +/// caller must drain that execution through its original attempt fence. +async fn await_tool_policy( + policy: impl Future>, + phase: InterceptorPhase, + cancel_rx: &mut mpsc::Receiver<()>, + pause_rx: &mut mpsc::Receiver<()>, +) -> Result { + let result = tokio::select! { + biased; + Some(()) = cancel_rx.recv() => return Err(EngineError::Cancelled), + Some(()) = pause_rx.recv() => return Err(EngineError::PauseRequested), + result = policy => result, + }; + check_policy_interruption(cancel_rx, pause_rx)?; + result.map_err(|error| InterceptorFailure::new(phase, error).into()) +} + /// Tool registration error #[derive(Debug, thiserror::Error)] pub enum ToolRegistryError { @@ -748,14 +782,6 @@ impl Engine { } } - fn try_paused(&mut self) -> bool { - use tokio::sync::mpsc::error::TryRecvError; - match self.pause_rx.try_recv() { - Ok(()) => true, - Err(TryRecvError::Empty | TryRecvError::Disconnected) => false, - } - } - /// Register a text block observer with scoped callbacks. /// /// The setup closure is called once per text block. Inside it, register @@ -1531,43 +1557,24 @@ impl Engine { tool, context, }; - let action = tokio::select! { - result = self.interceptor.pre_tool_call(&mut info) => { - result.map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPhase::PreToolCall, - error, - )) - })? - } - pause = self.pause_rx.recv() => { - if pause.is_some() { - info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy"); - } - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - None, - )?; + let action = await_tool_policy( + self.interceptor.pre_tool_call(&mut info), + InterceptorPhase::PreToolCall, + &mut self.cancel_rx, + &mut self.pause_rx, + ) + .await; + let action = match action { + Ok(action) => action, + Err(error @ (EngineError::PauseRequested | EngineError::Cancelled)) => { + self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; batch.dispatched_any = true; - return Ok(EarlyToolAdmission::Pause); - } - cancel = self.cancel_rx.recv() => { - if cancel.is_some() { - info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy"); - } - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - None, - )?; - batch.dispatched_any = true; - return Err(EngineError::Cancelled); + return match error { + EngineError::PauseRequested => Ok(EarlyToolAdmission::Pause), + _ => Err(error), + }; } + Err(error) => return Err(error), }; if info.call.id != expected_tool_use_id { return Err(InterceptorFailure::new( @@ -1945,8 +1952,10 @@ impl Engine { // A pause observed before provider ownership starts leaves every call // NotStarted and therefore eligible for an explicit later retry. - if self.try_paused() { - return Ok(ToolExecutionResult::Paused); + match check_policy_interruption(&mut self.cancel_rx, &mut self.pause_rx) { + Ok(()) => {} + Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused), + Err(error) => return Err(error), } // Map from tool call ID to (ToolCall, Meta, Tool, Context) @@ -1985,16 +1994,18 @@ impl Engine { context, }; - let pre_tool_action = - self.interceptor - .pre_tool_call(&mut info) - .await - .map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPhase::PreToolCall, - error, - )) - })?; + let pre_tool_action = match await_tool_policy( + self.interceptor.pre_tool_call(&mut info), + InterceptorPhase::PreToolCall, + &mut self.cancel_rx, + &mut self.pause_rx, + ) + .await + { + Ok(action) => action, + Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused), + Err(error) => return Err(error), + }; if info.call.id != expected_tool_use_id { return Err(InterceptorFailure::new( InterceptorPhase::PreToolCall, @@ -2061,6 +2072,14 @@ impl Engine { } } + // Recheck after all policies and immediately before provider ownership. + // No tool task exists yet: these calls remain proven NotStarted. + match check_policy_interruption(&mut self.cancel_rx, &mut self.pause_rx) { + Ok(()) => {} + Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused), + Err(error) => return Err(error), + } + // Phase 2: Execute approved tools in parallel. FuturesUnordered yields // each terminal result as soon as that call completes instead of // holding fast siblings behind the slowest call in the batch. @@ -2949,22 +2968,29 @@ impl Engine { )), 0, ); - let assistant_turn_action = self - .interceptor - .on_assistant_turn_end(AssistantTurnEndContext { - invocation: assistant_invocation, - assistant_entries: &history.entries()[assistant_start..], - history: history.entries(), - tool_calls: &tool_calls, - }) - .await; + let assistant_turn_action = await_tool_policy( + self.interceptor + .on_assistant_turn_end(AssistantTurnEndContext { + invocation: assistant_invocation, + assistant_entries: &history.entries()[assistant_start..], + history: history.entries(), + tool_calls: &tool_calls, + }), + InterceptorPhase::AssistantTurnEnd, + &mut self.cancel_rx, + &mut self.pause_rx, + ) + .await; let assistant_turn_action = match assistant_turn_action { Ok(action) => action, + Err(EngineError::PauseRequested) => { + if let Some(batch) = early_tools.take() { + self.drain_early_tool_batch(history, annotate, batch, false, true) + .await?; + } + return Ok(EngineResult::Paused); + } Err(error) => { - let error = EngineError::from(InterceptorFailure::new( - InterceptorPhase::AssistantTurnEnd, - error, - )); if let Some(batch) = early_tools.take() && let Err(cleanup_error) = self .drain_early_tool_batch(history, annotate, batch, true, false)