fix: preserve pause reclamation deadline during post-tool policy

crates/agen/src/engine.rs

@@ -105,6 +105,29 @@ async fn await_tool_policy<T>(
     result.map_err(|error| InterceptorFailure::new(phase, error).into())
 }
 
+/// A consumed Pause still bounds post-policy processing while owned tools drain.
+/// Terminal history and publication happen before this wait, so expiring it
+/// cannot discard the result or reopen the execution attempt.
+#[derive(Clone, Copy)]
+enum PostToolPolicyWait {
+    Skip,
+    Interruptible,
+    UntilPauseSafeBoundary(TokioInstant),
+}
+
+impl PostToolPolicyWait {
+    fn for_pause_deadline(deadline: Option<TokioInstant>) -> Self {
+        deadline.map_or(Self::Interruptible, Self::UntilPauseSafeBoundary)
+    }
+
+    fn pause_deadline(self) -> Option<TokioInstant> {
+        match self {
+            Self::UntilPauseSafeBoundary(deadline) => Some(deadline),
+            _ => None,
+        }
+    }
+}
+
 /// Tool registration error
 #[derive(Debug, thiserror::Error)]
 pub enum ToolRegistryError {
@@ -1723,7 +1746,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                 result,
                 None,
                 batch.turn_id,
-                true,
+                PostToolPolicyWait::Interruptible,
                 &batch.call_info_map,
                 &mut batch.attempt_fence,
                 &mut batch.terminal_call_ids,
@@ -1750,7 +1773,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                     result,
                     None,
                     batch.turn_id,
-                    false,
+                    PostToolPolicyWait::Skip,
                     &batch.call_info_map,
                     &mut batch.attempt_fence,
                     &mut batch.terminal_call_ids,
@@ -1798,7 +1821,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         result,
                         Some(&attempt_id),
                         batch.turn_id,
-                        false,
+                        PostToolPolicyWait::Skip,
                         &batch.call_info_map,
                         &mut batch.attempt_fence,
                         &mut batch.terminal_call_ids,
@@ -1821,7 +1844,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         ToolResult::outcome_unknown(call_id),
                         Some(attempt_id),
                         batch.turn_id,
-                        false,
+                        PostToolPolicyWait::Skip,
                         &batch.call_info_map,
                         &mut batch.attempt_fence,
                         &mut batch.terminal_call_ids,
@@ -1844,28 +1867,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         force_cancel: bool,
         initial_pause_requested: bool,
     ) -> Result<ToolExecutionResult, EngineError> {
-        while let Some(result) = batch.synthetic_results.pop_front() {
-            if let Err(error) = self
-                .finalize_and_commit_tool_result(
-                    history,
-                    annotate,
-                    result,
-                    None,
-                    batch.turn_id,
-                    true,
-                    &batch.call_info_map,
-                    &mut batch.attempt_fence,
-                    &mut batch.terminal_call_ids,
-                )
-                .await
-            {
-                let cleanup_error = self
-                    .cancel_and_terminalize_early_batch(history, annotate, &mut batch)
-                    .await;
-                return Err(cleanup_error.unwrap_or(error));
-            }
-        }
-
         if force_cancel {
             if let Some(error) = self
                 .cancel_and_terminalize_early_batch(history, annotate, &mut batch)
@@ -1879,6 +1880,28 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         let mut pause_requested = initial_pause_requested;
         let mut pause_deadline = initial_pause_requested
             .then(|| TokioInstant::now() + self.tool_execution_policy.pause_safe_boundary_timeout);
+        while let Some(result) = batch.synthetic_results.pop_front() {
+            if let Err(error) = self
+                .finalize_and_commit_tool_result(
+                    history,
+                    annotate,
+                    result,
+                    None,
+                    batch.turn_id,
+                    PostToolPolicyWait::for_pause_deadline(pause_deadline),
+                    &batch.call_info_map,
+                    &mut batch.attempt_fence,
+                    &mut batch.terminal_call_ids,
+                )
+                .await
+            {
+                let cleanup_error = self
+                    .cancel_and_terminalize_early_batch(history, annotate, &mut batch)
+                    .await;
+                return Err(cleanup_error.unwrap_or(error));
+            }
+        }
+
         while !batch.futures.is_empty() {
             tokio::select! {
                 biased;
@@ -1890,7 +1913,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         result,
                         Some(&attempt_id),
                         batch.turn_id,
-                        true,
+                        PostToolPolicyWait::for_pause_deadline(pause_deadline),
                         &batch.call_info_map,
                         &mut batch.attempt_fence,
                         &mut batch.terminal_call_ids,
@@ -2156,7 +2179,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                     result,
                     None,
                     turn_id,
-                    true,
+                    PostToolPolicyWait::Interruptible,
                     &call_info_map,
                     &mut attempt_fence,
                     &mut terminal_call_ids,
@@ -2188,7 +2211,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                         result,
                         Some(&attempt_id),
                         turn_id,
-                        true,
+                        PostToolPolicyWait::for_pause_deadline(pause_deadline),
                         &call_info_map,
                         &mut attempt_fence,
                         &mut terminal_call_ids,
@@ -2269,7 +2292,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                                     result,
                                     Some(&attempt_id),
                                     turn_id,
-                                    false,
+                                    PostToolPolicyWait::Skip,
                                     &call_info_map,
                                     &mut attempt_fence,
                                     &mut terminal_call_ids,
@@ -2297,7 +2320,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                                 ToolResult::outcome_unknown(call_id),
                                 Some(attempt_id),
                                 turn_id,
-                                false,
+                                PostToolPolicyWait::Skip,
                                 &call_info_map,
                                 &mut attempt_fence,
                                 &mut terminal_call_ids,
@@ -2347,7 +2370,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         mut tool_result: ToolResult,
         execution_attempt_id: Option<&str>,
         turn_id: usize,
-        invoke_post_hook: bool,
+        post_policy_wait: PostToolPolicyWait,
         call_info_map: &HashMap<
             String,
             (
@@ -2432,7 +2455,9 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
         );
         self.emit_tool_result(&tool_result);
 
-        if invoke_post_hook && let Some((tool_call, meta, tool, context)) = call_info {
+        if !matches!(post_policy_wait, PostToolPolicyWait::Skip)
+            && let Some((tool_call, meta, tool, context)) = call_info
+        {
             let invocation = self.interceptor_invocation(
                 InterceptorPhase::PostToolCall,
                 Some(turn_id),
@@ -2448,27 +2473,20 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                 tool: tool.clone(),
                 context: context.clone(),
             };
+            let pause_deadline = post_policy_wait.pause_deadline();
             let post_tool_action = tokio::select! {
-                result = self.interceptor.post_tool_call(&info) => {
+                biased;
-                    result.map_err(|error| {
+                // Poll the existing safe-boundary deadline even when the Pause
-                        EngineError::from(InterceptorFailure::new(
+                // that established it was consumed by an outer policy gate.
-                            InterceptorPhase::PostToolCall,
+                _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => {
-                            error,
-                        ))
-                    })?
-                }
-                pause = self.pause_rx.recv() => {
-                    if pause.is_some() {
-                        info!(call_id = %tool_call.id, "Paused while awaiting post-tool policy");
-                    }
                     return Err(EngineError::PauseRequested);
                 }
-                cancel = self.cancel_rx.recv() => {
+                result = await_tool_policy(
-                    if cancel.is_some() {
+                    self.interceptor.post_tool_call(&info),
-                        info!(call_id = %tool_call.id, "Cancelled while awaiting post-tool policy");
+                    InterceptorPhase::PostToolCall,
-                    }
+                    &mut self.cancel_rx,
-                    return Err(EngineError::Cancelled);
+                    &mut self.pause_rx,
-                }
+                ) => result?,
             };
             if let PostToolAction::Abort(reason) = post_tool_action {
                 return Err(EngineError::Aborted(reason));
@@ -2506,7 +2524,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                     ToolResult::outcome_unknown(&call_id),
                     None,
                     self.turn_count,
-                    false,
+                    PostToolPolicyWait::Skip,
                     &call_info_map,
                     &mut attempt_fence,
                     &mut terminal_call_ids,
@@ -3556,7 +3574,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
                             result,
                             Some(&attempt_id),
                             batch.turn_id,
-                            true,
+                            PostToolPolicyWait::Interruptible,
                             &batch.call_info_map,
                             &mut batch.attempt_fence,
                             &mut batch.terminal_call_ids,

crates/agen/tests/parallel_execution_test.rs

@@ -3446,3 +3446,164 @@ async fn turn_end_final_poll_signal_wins_over_finish_without_tools() {
         assert_eq!(probe.continues.load(Ordering::SeqCst), 1);
     }
 }
+
+#[tokio::test]
+async fn turn_end_pause_deadline_reclaims_started_sibling_while_post_policy_is_pending() {
+    struct TurnAndPostBarrier {
+        turn: PolicyBarrier,
+        post: PolicyBarrier,
+    }
+    #[async_trait]
+    impl Interceptor for TurnAndPostBarrier {
+        fn supports_early_tool_dispatch(&self) -> bool {
+            true
+        }
+        async fn on_assistant_turn_end(
+            &self,
+            _: AssistantTurnEndContext<'_>,
+        ) -> InterceptorResult<TurnEndAction> {
+            self.turn.wait().await;
+            Ok(TurnEndAction::Finish)
+        }
+        async fn post_tool_call(
+            &self,
+            _: &ToolResultInfo<'_>,
+        ) -> InterceptorResult<PostToolAction> {
+            self.post.wait().await;
+            Ok(PostToolAction::Continue)
+        }
+    }
+
+    let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
+    let mut engine = Engine::new(client);
+    engine.set_tool_execution_policy(ToolExecutionPolicy {
+        pause_safe_boundary_timeout: Duration::from_millis(500),
+        cancellation_request_timeout: Duration::from_secs(1),
+        terminal_confirmation_timeout: Duration::from_secs(2),
+    });
+    let turn = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early);
+    let post = PolicyBarrier::new(InterceptorPhase::PostToolCall, PolicyRoute::Early);
+    engine.set_interceptor(TurnAndPostBarrier {
+        turn: turn.clone(),
+        post: post.clone(),
+    });
+    let completed = BarrierTool::new("completed_before_pause");
+    engine.register_tool(completed.definition());
+    let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel();
+    let sibling = PolicyOwnedTool {
+        starts: Arc::new(AtomicUsize::new(0)),
+        started: Arc::new(tokio::sync::Notify::new()),
+        cancel_requested: Arc::new(tokio::sync::Notify::new()),
+        terminal: Arc::new(tokio::sync::Mutex::new(Some(terminal_rx))),
+        retired: Arc::new(tokio::sync::Notify::new()),
+    };
+    let definition: ToolDefinition = {
+        let sibling = sibling.clone();
+        Arc::new(move || {
+            (
+                ToolMeta::new("pending_sibling").input_schema(serde_json::json!({"type":"object"})),
+                Arc::new(sibling.clone()) as Arc<dyn Tool>,
+            )
+        })
+    };
+    engine.register_tool(definition.clone());
+    let committed = Arc::new(AtomicUsize::new(0));
+    let commit_probe = committed.clone();
+    engine.on_history_append(move |item| {
+        if matches!(item, Item::ToolResult { .. }) {
+            commit_probe.fetch_add(1, Ordering::SeqCst);
+        }
+        Ok(())
+    });
+    let published = Arc::new(AtomicUsize::new(0));
+    let publish_probe = published.clone();
+    let commit_probe = committed.clone();
+    engine.on_tool_result(move |_| {
+        let n = publish_probe.fetch_add(1, Ordering::SeqCst) + 1;
+        assert_eq!(commit_probe.load(Ordering::SeqCst), n);
+    });
+    let pause = engine.pause_sender();
+    let run = tokio::spawn(async move {
+        let mut history = History::new();
+        let output = engine.run(&mut history, "run").await;
+        (output, history)
+    });
+    for (index, id, name) in [
+        (0, "policy_call", "completed_before_pause"),
+        (1, "sibling_call", "pending_sibling"),
+    ] {
+        for event in [
+            Event::tool_use_start(index, id, name),
+            Event::tool_input_delta(index, "{}"),
+            Event::tool_use_stop(index),
+        ] {
+            tx.send(Ok(event)).unwrap();
+        }
+    }
+    tokio::time::timeout(Duration::from_secs(2), completed.started.notified())
+        .await
+        .expect("first provider owned");
+    tokio::time::timeout(Duration::from_secs(2), sibling.started.notified())
+        .await
+        .expect("sibling provider owned");
+    tx.send(Ok(Event::Status(StatusEvent {
+        status: ResponseStatus::Completed,
+    })))
+    .unwrap();
+    drop(tx);
+    tokio::time::timeout(Duration::from_secs(2), turn.entered.notified())
+        .await
+        .expect("turn-end pending");
+    pause.send(()).await.unwrap(); // only one interruption signal
+    tokio::time::timeout(Duration::from_secs(2), turn.retired.notified())
+        .await
+        .expect("turn-end interrupted");
+    completed.release.notify_one();
+    tokio::time::timeout(Duration::from_secs(2), post.entered.notified())
+        .await
+        .expect("completed result's post policy pending");
+    assert_eq!(committed.load(Ordering::SeqCst), 1);
+    assert_eq!(published.load(Ordering::SeqCst), 1);
+    tokio::time::timeout(Duration::from_secs(2), sibling.cancel_requested.notified()).await.expect("original Pause deadline must reach sibling cancellation without releasing either policy");
+    terminal_tx
+        .send(ToolOutput::from("confirmed sibling terminal".to_string()))
+        .unwrap();
+    let (output, history) = tokio::time::timeout(Duration::from_secs(2), run)
+        .await
+        .expect("all owned tools reclaimed")
+        .unwrap();
+    assert_policy_interruption(&output.result, true);
+    assert_eq!(committed.load(Ordering::SeqCst), 2);
+    assert_eq!(published.load(Ordering::SeqCst), 2);
+    assert_eq!(
+        history
+            .items()
+            .filter(|item| matches!(item, Item::ToolResult { .. }))
+            .count(),
+        2
+    );
+    assert_eq!(completed.starts.load(Ordering::SeqCst), 1);
+    assert_eq!(sibling.starts.load(Ordering::SeqCst), 1);
+    assert_eq!(post.invocations.load(Ordering::SeqCst), 1);
+    assert_eq!(post.continues.load(Ordering::SeqCst), 0);
+    turn.release.notify_one();
+    post.release.notify_one();
+    let mut restored_history: History =
+        serde_json::from_str(&serde_json::to_string(&history).unwrap()).unwrap();
+    let mut restored = Engine::new(MockLlmClient::new(vec![Event::Status(StatusEvent {
+        status: ResponseStatus::Completed,
+    })]));
+    restored.register_tool(completed.definition());
+    restored.register_tool(definition);
+    restored.set_max_turns(Some(1));
+    let _ = restored.resume(&mut restored_history).await;
+    assert_eq!(completed.starts.load(Ordering::SeqCst), 1);
+    assert_eq!(sibling.starts.load(Ordering::SeqCst), 1);
+    assert_eq!(
+        restored_history
+            .items()
+            .filter(|item| matches!(item, Item::ToolResult { .. }))
+            .count(),
+        2
+    );
+}