fix: interrupt pending tool policies before provider admission

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<T>(
+    policy: impl Future<Output = Result<T, InterceptorError>>,
+    phase: InterceptorPhase,
+    cancel_rx: &mut mpsc::Receiver<()>,
+    pause_rx: &mut mpsc::Receiver<()>,
+) -> Result<T, EngineError> {
+    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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         }
     }
 
-    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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                 tool,
                 context,
             };
-            let action = tokio::select! {
+            let action = await_tool_policy(
-                result = self.interceptor.pre_tool_call(&mut info) => {
+                self.interceptor.pre_tool_call(&mut info),
-                    result.map_err(|error| {
+                InterceptorPhase::PreToolCall,
-                        EngineError::from(InterceptorFailure::new(
+                &mut self.cancel_rx,
-                            InterceptorPhase::PreToolCall,
+                &mut self.pause_rx,
-                            error,
+            )
-                        ))
+            .await;
-                    })?
+            let action = match action {
-                }
+                Ok(action) => action,
-                pause = self.pause_rx.recv() => {
+                Err(error @ (EngineError::PauseRequested | EngineError::Cancelled)) => {
-                    if pause.is_some() {
+                    self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
-                        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,
-                    )?;
                     batch.dispatched_any = true;
-                    return Ok(EarlyToolAdmission::Pause);
+                    return match error {
-                }
+                        EngineError::PauseRequested => Ok(EarlyToolAdmission::Pause),
-                cancel = self.cancel_rx.recv() => {
+                        _ => Err(error),
-                    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);
                 }
+                Err(error) => return Err(error),
             };
             if info.call.id != expected_tool_use_id {
                 return Err(InterceptorFailure::new(
@@ -1945,8 +1952,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
 
         // A pause observed before provider ownership starts leaves every call
         // NotStarted and therefore eligible for an explicit later retry.
-        if self.try_paused() {
+        match check_policy_interruption(&mut self.cancel_rx, &mut self.pause_rx) {
-            return Ok(ToolExecutionResult::Paused);
+            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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                     context,
                 };
 
-                let pre_tool_action =
+                let pre_tool_action = match await_tool_policy(
-                    self.interceptor
+                    self.interceptor.pre_tool_call(&mut info),
-                        .pre_tool_call(&mut info)
+                    InterceptorPhase::PreToolCall,
-                        .await
+                    &mut self.cancel_rx,
-                        .map_err(|error| {
+                    &mut self.pause_rx,
-                            EngineError::from(InterceptorFailure::new(
+                )
-                                InterceptorPhase::PreToolCall,
+                .await
-                                error,
+                {
-                            ))
+                    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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
             }
         }
 
+        // 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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                 )),
                 0,
             );
-            let assistant_turn_action = self
+            let assistant_turn_action = await_tool_policy(
-                .interceptor
+                self.interceptor
-                .on_assistant_turn_end(AssistantTurnEndContext {
+                    .on_assistant_turn_end(AssistantTurnEndContext {
-                    invocation: assistant_invocation,
+                        invocation: assistant_invocation,
-                    assistant_entries: &history.entries()[assistant_start..],
+                        assistant_entries: &history.entries()[assistant_start..],
-                    history: history.entries(),
+                        history: history.entries(),
-                    tool_calls: &tool_calls,
+                        tool_calls: &tool_calls,
-                })
+                    }),
-                .await;
+                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)