From 74d0af93e3b80b795a66d0647226810f8fdc4055 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:38:38 +0900 Subject: [PATCH 1/3] 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) From 2b1f68a3ead74876314737326b8d2deae1548a46 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:38:38 +0900 Subject: [PATCH 2/3] test: fence policy interruptions across tool ownership and restore --- crates/agen/tests/parallel_execution_test.rs | 573 +++++++++++++++++++ 1 file changed, 573 insertions(+) diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index de3156607..e31493751 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -2873,3 +2873,576 @@ async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning( assert!(matches!(next, EngineRunExit::Finished)); } } + +// The same policy barrier is exercised through the real response-complete, +// provider fallback, interceptor fallback, and early-dispatch engine paths. +#[derive(Clone, Copy, Debug)] +enum PolicyRoute { + AfterResponse, + ProviderFallback, + InterceptorFallback, + Early, +} + +impl PolicyRoute { + fn completion_support(self) -> ToolCallCompletionSupport { + match self { + Self::ProviderFallback => ToolCallCompletionSupport::ResponseComplete, + _ => ToolCallCompletionSupport::PerBlock, + } + } + + fn configure(self, engine: &mut Engine) { + engine.set_max_turns(Some(1)); + if matches!(self, Self::AfterResponse) { + engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::AfterResponse); + } + } +} + +struct NotifyOnDrop(Arc); + +impl Drop for NotifyOnDrop { + fn drop(&mut self) { + self.0.notify_one(); + } +} + +#[derive(Clone)] +struct PolicyBarrier { + phase: InterceptorPhase, + early_compatible: bool, + entered: Arc, + release: Arc, + invocations: Arc, + continues: Arc, + retired: Arc, + // Enqueue from the policy's final poll, after select has polled signals. + // This deterministically exercises the policy-completion/ownership race. + signal_on_continue: Option>, +} + +impl PolicyBarrier { + fn new(phase: InterceptorPhase, route: PolicyRoute) -> Self { + Self { + phase, + early_compatible: !matches!(route, PolicyRoute::InterceptorFallback), + entered: Arc::new(tokio::sync::Notify::new()), + release: Arc::new(tokio::sync::Notify::new()), + invocations: Arc::new(AtomicUsize::new(0)), + continues: Arc::new(AtomicUsize::new(0)), + retired: Arc::new(tokio::sync::Notify::new()), + signal_on_continue: None, + } + } + + async fn wait(&self) { + let _retired = NotifyOnDrop(self.retired.clone()); + self.invocations.fetch_add(1, Ordering::SeqCst); + self.entered.notify_one(); + self.release.notified().await; + if let Some(signal) = &self.signal_on_continue { + signal.try_send(()).unwrap(); + } + self.continues.fetch_add(1, Ordering::SeqCst); + } +} + +#[async_trait] +impl Interceptor for PolicyBarrier { + fn supports_early_tool_dispatch(&self) -> bool { + self.early_compatible + } + + async fn pre_tool_call( + &self, + _info: &mut ToolCallInfo<'_>, + ) -> InterceptorResult { + if self.phase == InterceptorPhase::PreToolCall { + self.wait().await; + } + Ok(PreToolAction::Continue) + } + + async fn on_assistant_turn_end( + &self, + _context: AssistantTurnEndContext<'_>, + ) -> InterceptorResult { + if self.phase == InterceptorPhase::AssistantTurnEnd { + self.wait().await; + } + Ok(TurnEndAction::Finish) + } +} + +fn send_policy_response(tx: &tokio::sync::mpsc::UnboundedSender>) { + for event in [ + Event::tool_use_start(0, "policy_call", "policy_tool"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::usage(10, 3), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] { + tx.send(Ok(event)).unwrap(); + } +} + +fn assert_policy_interruption(exit: &EngineRunExit, pause: bool) { + if pause { + assert!(matches!(exit, EngineRunExit::Paused), "{exit:?}"); + } else { + assert!( + matches!( + exit, + EngineRunExit::Interrupted(RunInterruptionReason::Cancelled) + ), + "{exit:?}" + ); + } +} + +fn call_and_result_count(history: &History) -> (usize, usize) { + ( + history + .items() + .filter( + |item| matches!(item, Item::ToolCall { call_id, .. } if call_id == "policy_call"), + ) + .count(), + history + .items() + .filter( + |item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "policy_call"), + ) + .count(), + ) +} + +#[tokio::test] +async fn response_complete_pending_policies_interrupt_before_release_and_preserve_not_started() { + for route in [ + PolicyRoute::AfterResponse, + PolicyRoute::ProviderFallback, + PolicyRoute::InterceptorFallback, + ] { + for phase in [ + InterceptorPhase::PreToolCall, + InterceptorPhase::AssistantTurnEnd, + ] { + for pause in [false, true] { + let (client, tx) = ControlledStreamClient::new(route.completion_support()); + let mut engine = Engine::new(client); + route.configure(&mut engine); + let policy = PolicyBarrier::new(phase, route); + let probe = policy.clone(); + engine.set_interceptor(policy); + let tool = SlowTool::new("policy_tool", 0); + engine.register_tool(tool.definition()); + let usage = Arc::new(AtomicUsize::new(0)); + let usage_probe = usage.clone(); + engine.on_usage(move |_| { + usage_probe.fetch_add(1, Ordering::SeqCst); + }); + let signal = if pause { + engine.pause_sender() + } else { + engine.cancel_sender() + }; + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + send_policy_response(&tx); + drop(tx); + tokio::time::timeout(Duration::from_secs(2), probe.entered.notified()) + .await + .expect("policy entered"); + signal.send(()).await.unwrap(); + let (output, history) = tokio::time::timeout(Duration::from_secs(2), run).await + .unwrap_or_else(|_| panic!("{route:?}/{phase:?}/pause={pause}: interruption waited for policy release")) + .unwrap(); + assert_policy_interruption(&output.result, pause); + assert_eq!(tool.call_count(), 0); + assert_eq!(call_and_result_count(&history), (1, 0)); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + execution_id: None, + status: None | Some(agen::llm_client::ItemStatus::Completed), + .. + } + ))); + assert_eq!(usage.load(Ordering::SeqCst), 1); + assert_eq!(probe.invocations.load(Ordering::SeqCst), 1); + probe.release.notify_one(); // late permission cannot create a task + assert_eq!(probe.continues.load(Ordering::SeqCst), 0); + assert_eq!(tool.call_count(), 0); + + // Recreate both history and engine as a restore boundary. Only + // proven NotStarted is allowed to run, with normal pre-tool policy. + 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.set_max_turns(Some(1)); + restored.register_tool(tool.definition()); + let _ = restored.resume(&mut restored_history).await; + assert_eq!(tool.call_count(), 1); + assert_eq!(call_and_result_count(&restored_history), (1, 1)); + assert_eq!(usage.load(Ordering::SeqCst), 1); + } + } + } +} + +#[tokio::test] +async fn policy_final_poll_signal_prevents_provider_ownership_even_with_continue() { + for route in [ + PolicyRoute::AfterResponse, + PolicyRoute::ProviderFallback, + PolicyRoute::InterceptorFallback, + PolicyRoute::Early, + ] { + for pause in [false, true] { + let (client, tx) = ControlledStreamClient::new(route.completion_support()); + let mut engine = Engine::new(client); + route.configure(&mut engine); + let mut policy = PolicyBarrier::new(InterceptorPhase::PreToolCall, route); + policy.signal_on_continue = Some(if pause { + engine.pause_sender() + } else { + engine.cancel_sender() + }); + let probe = policy.clone(); + engine.set_interceptor(policy); + let tool = SlowTool::new("policy_tool", 0); + engine.register_tool(tool.definition()); + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + send_policy_response(&tx); + drop(tx); + tokio::time::timeout(Duration::from_secs(2), probe.entered.notified()) + .await + .expect("policy entered"); + probe.release.notify_one(); + let (output, history) = tokio::time::timeout(Duration::from_secs(2), run) + .await + .expect("final-poll signal processed") + .unwrap(); + assert_policy_interruption(&output.result, pause); + assert_eq!(probe.continues.load(Ordering::SeqCst), 1); + assert_eq!(probe.invocations.load(Ordering::SeqCst), 1); + assert_eq!( + tool.call_count(), + 0, + "{route:?}/pause={pause}: Continue bypassed queued interruption" + ); + assert_eq!(call_and_result_count(&history), (1, 0)); + } + } +} + +// Provider owns the terminal channel. Cancel acknowledges receipt but does not +// manufacture a result, so tests can release a real terminal output or deliver +// it after the engine's confirmation deadline (without a live side effect). +#[derive(Clone)] +struct PolicyOwnedTool { + starts: Arc, + started: Arc, + cancel_requested: Arc, + terminal: Arc>>>, + retired: Arc, +} + +#[async_trait] +impl Tool for PolicyOwnedTool { + async fn execute(&self, _: &str, _: ToolExecutionContext) -> Result { + let _retired = NotifyOnDrop(self.retired.clone()); + self.starts.fetch_add(1, Ordering::SeqCst); + self.started.notify_one(); + Ok(self.terminal.lock().await.take().unwrap().await.unwrap()) + } + async fn cancel(&self, _: &str) -> Result<(), ToolError> { + self.cancel_requested.notify_one(); + Ok(()) + } +} + +#[tokio::test] +async fn pending_turn_end_interrupts_drain_owned_early_attempt_and_fence_late_result_on_restore() { + for pause in [false, true] { + for late in [false, true] { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let mut engine = Engine::new(client); + PolicyRoute::Early.configure(&mut engine); + engine.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::ZERO, + cancellation_request_timeout: Duration::from_secs(1), + terminal_confirmation_timeout: if late { + Duration::ZERO + } else { + Duration::from_secs(2) + }, + }); + let policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early); + let probe = policy.clone(); + engine.set_interceptor(policy); + let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel(); + let tool = 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 tool = tool.clone(); + Arc::new(move || { + ( + ToolMeta::new("policy_tool") + .input_schema(serde_json::json!({"type":"object"})), + Arc::new(tool.clone()) as Arc, + ) + }) + }; + 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 |_| { + assert_eq!( + commit_probe.load(Ordering::SeqCst), + 1, + "commit before publication" + ); + publish_probe.fetch_add(1, Ordering::SeqCst); + }); + let usage = Arc::new(AtomicUsize::new(0)); + let usage_probe = usage.clone(); + engine.on_usage(move |_| { + usage_probe.fetch_add(1, Ordering::SeqCst); + }); + let signal = if pause { + engine.pause_sender() + } else { + engine.cancel_sender() + }; + let run = tokio::spawn(async move { + let mut history = History::new(); + let output = engine.run(&mut history, "run").await; + (output, history) + }); + // Wait for ownership before completing the assistant response. + for event in [ + Event::tool_use_start(0, "policy_call", "policy_tool"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + ] { + tx.send(Ok(event)).unwrap(); + } + tokio::time::timeout(Duration::from_secs(2), tool.started.notified()) + .await + .expect("early provider started"); + tx.send(Ok(Event::usage(10, 3))).unwrap(); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + tokio::time::timeout(Duration::from_secs(2), probe.entered.notified()) + .await + .expect("turn-end policy pending"); + signal.send(()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(2), tool.cancel_requested.notified()) + .await + .expect("owned provider cancellation progresses before policy release"); + let terminal_tx = if late { + Some(terminal_tx) + } else { + terminal_tx + .send(ToolOutput::from("confirmed after interruption".to_string())) + .unwrap(); + None + }; + let (output, history) = tokio::time::timeout(Duration::from_secs(2), run) + .await + .expect("terminalization does not await turn-end policy") + .unwrap(); + assert_policy_interruption(&output.result, pause); + assert_eq!(call_and_result_count(&history), (1, 1)); + assert_eq!(committed.load(Ordering::SeqCst), 1); + assert_eq!(published.load(Ordering::SeqCst), 1); + assert_eq!(usage.load(Ordering::SeqCst), 1); + assert!(history.items().any(|item| matches!( + item, + Item::ToolCall { + execution_id: Some(_), + status: Some(agen::llm_client::ItemStatus::InProgress), + .. + } + ))); + assert!(history.items().any(|item| matches!(item, Item::ToolResult { disposition, .. } if *disposition == if late { ToolResultDisposition::OutcomeUnknown } else { ToolResultDisposition::Success }))); + if let Some(terminal_tx) = terminal_tx { + tokio::time::timeout(Duration::from_secs(2), tool.retired.notified()) + .await + .expect("forced-close retired the owned receiver"); + // Force-close retires the receiver; a provider's late completion + // cannot reenter history or replace OutcomeUnknown. + assert!( + terminal_tx + .send(ToolOutput::from("late completion".to_string())) + .is_err() + ); + } + probe.release.notify_one(); + assert_eq!(probe.continues.load(Ordering::SeqCst), 0); + let original_items = history.items().cloned().collect::>(); + let mut resumed_history = history.clone(); + let mut resumed_engine = output.engine; + let _ = resumed_engine.resume(&mut resumed_history).await; + assert_eq!( + tool.starts.load(Ordering::SeqCst), + 1, + "same-engine resume cannot rerun a terminal attempt" + ); + assert_eq!(call_and_result_count(&resumed_history), (1, 1)); + let mut restored_history: History = + serde_json::from_str(&serde_json::to_string(&history).unwrap()).unwrap(); + let restored_client = MockLlmClient::new(vec![Event::Status(StatusEvent { + status: ResponseStatus::Completed, + })]); + let request_probe = restored_client.clone(); + let mut restored = Engine::new(restored_client); + restored.register_tool(definition); + restored.set_max_turns(Some(1)); + let _ = restored.resume(&mut restored_history).await; + let requests = request_probe.requests(); + assert_eq!(requests.len(), 1); + let replay = requests[0] + .items + .iter() + .filter_map(|item| match item { + Item::ToolCall { call_id, .. } => Some(format!("call:{call_id}")), + Item::ToolResult { call_id, .. } => Some(format!("result:{call_id}")), + _ => None, + }) + .collect::>(); + assert_eq!(replay, ["call:policy_call", "result:policy_call"]); + assert_eq!(tool.starts.load(Ordering::SeqCst), 1); + assert_eq!(call_and_result_count(&restored_history), (1, 1)); + assert_eq!( + restored_history + .items() + .take(original_items.len()) + .cloned() + .collect::>(), + original_items, + "restore must not rewrite original history" + ); + assert_eq!(published.load(Ordering::SeqCst), 1); + } + } +} + +#[tokio::test] +async fn pending_turn_end_pause_preserves_natural_safe_boundary_without_cancelling_started_tool() { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let mut engine = Engine::new(client); + PolicyRoute::Early.configure(&mut engine); + engine.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_secs(2), + ..ToolExecutionPolicy::default() + }); + let policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early); + let probe = policy.clone(); + engine.set_interceptor(policy); + let tool = SafePauseTool::new(); + // Use this tool's actual public name in the stream fixture. + engine.register_tool(tool.definition()); + 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 event in [ + Event::tool_use_start(0, "policy_call", "safe_pause"), + Event::tool_input_delta(0, "{}"), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] { + tx.send(Ok(event)).unwrap(); + } + drop(tx); + tokio::time::timeout(Duration::from_secs(2), probe.entered.notified()) + .await + .expect("turn-end entered"); + pause.send(()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(2), probe.retired.notified()) + .await + .expect("pending policy interrupted before tool release"); + tool.release.notify_one(); + let (output, history) = tokio::time::timeout(Duration::from_secs(2), run) + .await + .expect("safe boundary collected") + .unwrap(); + assert_policy_interruption(&output.result, true); + assert_eq!(tool.calls.load(Ordering::SeqCst), 1); + assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0); + assert_eq!(call_and_result_count(&history), (1, 1)); + probe.release.notify_one(); + assert_eq!(probe.continues.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn turn_end_final_poll_signal_wins_over_finish_without_tools() { + for pause in [false, true] { + let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock); + let mut engine = Engine::new(client); + let mut policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early); + policy.signal_on_continue = Some(if pause { + engine.pause_sender() + } else { + engine.cancel_sender() + }); + let probe = policy.clone(); + engine.set_interceptor(policy); + let run = tokio::spawn(async move { + let mut history = History::new(); + engine.run(&mut history, "run").await + }); + tx.send(Ok(Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }))) + .unwrap(); + drop(tx); + tokio::time::timeout(Duration::from_secs(2), probe.entered.notified()) + .await + .expect("turn-end entered"); + probe.release.notify_one(); + let output = tokio::time::timeout(Duration::from_secs(2), run) + .await + .expect("final-poll signal processed") + .unwrap(); + assert_policy_interruption(&output.result, pause); + assert_eq!(probe.continues.load(Ordering::SeqCst), 1); + } +} From e19742d509e7e64dd623319594df2ae5f7a7f93c Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 02:54:12 +0900 Subject: [PATCH 3/3] fix: preserve pause reclamation deadline during post-tool policy --- crates/agen/src/engine.rs | 124 ++++++++------ crates/agen/tests/parallel_execution_test.rs | 161 +++++++++++++++++++ 2 files changed, 232 insertions(+), 53 deletions(-) diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index beef6816a..d741d0015 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -105,6 +105,29 @@ async fn await_tool_policy( 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) -> Self { + deadline.map_or(Self::Interruptible, Self::UntilPauseSafeBoundary) + } + + fn pause_deadline(self) -> Option { + match self { + Self::UntilPauseSafeBoundary(deadline) => Some(deadline), + _ => None, + } + } +} + /// Tool registration error #[derive(Debug, thiserror::Error)] pub enum ToolRegistryError { @@ -1723,7 +1746,7 @@ impl Engine { 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 Engine { 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 Engine { 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 Engine { 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 Engine { force_cancel: bool, initial_pause_requested: bool, ) -> Result { - 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 Engine { 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 Engine { 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 Engine { result, None, turn_id, - true, + PostToolPolicyWait::Interruptible, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2188,7 +2211,7 @@ impl Engine { 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 Engine { result, Some(&attempt_id), turn_id, - false, + PostToolPolicyWait::Skip, &call_info_map, &mut attempt_fence, &mut terminal_call_ids, @@ -2297,7 +2320,7 @@ impl Engine { 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 Engine { 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 Engine { ); 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 Engine { 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) => { - result.map_err(|error| { - EngineError::from(InterceptorFailure::new( - InterceptorPhase::PostToolCall, - error, - )) - })? - } - pause = self.pause_rx.recv() => { - if pause.is_some() { - info!(call_id = %tool_call.id, "Paused while awaiting post-tool policy"); - } + biased; + // Poll the existing safe-boundary deadline even when the Pause + // that established it was consumed by an outer policy gate. + _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => { return Err(EngineError::PauseRequested); } - cancel = self.cancel_rx.recv() => { - if cancel.is_some() { - info!(call_id = %tool_call.id, "Cancelled while awaiting post-tool policy"); - } - return Err(EngineError::Cancelled); - } + result = await_tool_policy( + self.interceptor.post_tool_call(&info), + InterceptorPhase::PostToolCall, + &mut self.cancel_rx, + &mut self.pause_rx, + ) => result?, }; if let PostToolAction::Abort(reason) = post_tool_action { return Err(EngineError::Aborted(reason)); @@ -2506,7 +2524,7 @@ impl Engine { 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 Engine { result, Some(&attempt_id), batch.turn_id, - true, + PostToolPolicyWait::Interruptible, &batch.call_info_map, &mut batch.attempt_fence, &mut batch.terminal_call_ids, diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index e31493751..ae3d63e78 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/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 { + self.turn.wait().await; + Ok(TurnEndAction::Finish) + } + async fn post_tool_call( + &self, + _: &ToolResultInfo<'_>, + ) -> InterceptorResult { + 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, + ) + }) + }; + 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 + ); +}