diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index c48f2128..e9bb671b 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1284,6 +1284,7 @@ impl Engine { // commit-before-publish boundary. let mut terminal_call_ids = HashSet::new(); let mut pause_requested = false; + let mut pause_deadline = None; for result in synthetic_results { self.finalize_and_commit_tool_result( history, @@ -1319,12 +1320,20 @@ impl Engine { } pause = self.pause_rx.recv(), if !pause_requested => { if pause.is_some() { - // Pause is a safe-boundary request: do not cancel provider - // operations that already started. Drain all confirmed - // terminal results, then yield control to Worker. + // Pause first waits for already-started tools to reach a + // natural safe boundary. If they do not, Worker policy + // escalates to the same explicit cancel-and-confirm path. pause_requested = true; + pause_deadline = Some( + TokioInstant::now() + + self.tool_execution_policy.pause_safe_boundary_timeout, + ); } } + _ = tokio::time::sleep_until(pause_deadline.unwrap_or_else(TokioInstant::now)), if pause_deadline.is_some() => { + pause_deadline = None; + let _ = self.cancel_tx.try_send(()); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Tool execution cancellation requested"); @@ -1404,6 +1413,9 @@ impl Engine { } self.timeline.abort_current_block(); + if pause_requested { + return Ok(ToolExecutionResult::Paused); + } return Err(EngineError::Cancelled); } } diff --git a/crates/agen/src/tool.rs b/crates/agen/src/tool.rs index 9001c89c..7557d978 100644 --- a/crates/agen/src/tool.rs +++ b/crates/agen/src/tool.rs @@ -490,6 +490,9 @@ impl ToolExecutionHandle { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ToolExecutionPolicy { + /// Time a pause waits for already-started providers to reach a natural safe + /// boundary before escalating to explicit cooperative cancellation. + pub pause_safe_boundary_timeout: std::time::Duration, /// Maximum time allowed for a provider to accept one cooperative /// cancellation request. pub cancellation_request_timeout: std::time::Duration, @@ -501,6 +504,7 @@ pub struct ToolExecutionPolicy { impl Default for ToolExecutionPolicy { fn default() -> Self { Self { + pause_safe_boundary_timeout: std::time::Duration::from_millis(100), cancellation_request_timeout: std::time::Duration::from_millis(100), terminal_confirmation_timeout: std::time::Duration::from_millis(500), } diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 13108b48..61f29600 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -12,7 +12,7 @@ use agen::tool::{ Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, ToolResultDisposition, }; -use agen::{Engine, History, Item}; +use agen::{Engine, History, Item, ToolExecutionPolicy}; use async_trait::async_trait; mod common; @@ -610,7 +610,7 @@ async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() { .await .expect("tool execution starts"); pause.send(()).await.unwrap(); - tokio::time::sleep(Duration::from_millis(100)).await; + tokio::time::sleep(Duration::from_millis(50)).await; release.notify_one(); }); @@ -619,7 +619,7 @@ async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() { let output = engine.run(&mut history, "pause safely").await; control.await.unwrap(); - assert!(started_at.elapsed() >= Duration::from_millis(100)); + assert!(started_at.elapsed() >= Duration::from_millis(50)); assert_eq!(tool.calls.load(Ordering::SeqCst), 1); assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0); assert!(matches!(output.result, agen::EngineRunExit::Paused)); @@ -633,6 +633,53 @@ async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() { ))); } +#[tokio::test] +async fn pause_escalates_to_explicit_cancel_and_confirm_after_safe_boundary_deadline() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_pause_cancel", "cooperative"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + engine.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_millis(20), + cancellation_request_timeout: Duration::from_millis(50), + terminal_confirmation_timeout: Duration::from_millis(100), + }); + let tool = CooperativeCancelTool::new(); + engine.register_tool(tool.definition()); + + let pause = engine.pause_sender(); + let calls = Arc::clone(&tool.calls); + let control = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("tool execution starts"); + pause.send(()).await.unwrap(); + }); + + let mut history = History::new(); + let output = engine.run(&mut history, "pause with escalation").await; + control.await.unwrap(); + + assert!(matches!(output.result, agen::EngineRunExit::Paused)); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::Cancelled, + .. + } if call_id == "call_pause_cancel" + ))); +} + #[tokio::test] async fn cancellation_completion_race_commits_one_terminal_output() { for iteration in 0..24u64 { diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index b72dee6e..28f69623 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -2581,7 +2581,10 @@ impl Worker { result: &EngineRunExit, snapshot: &EmptyTurnRollbackSnapshot, ) -> bool { - if !matches!(result, EngineRunExit::Interrupted(StopReason::Cancelled)) { + if !matches!( + result, + EngineRunExit::Paused | EngineRunExit::Interrupted(StopReason::Cancelled) + ) { return false; } if self.ai_activity_counter.load(Ordering::SeqCst) != snapshot.ai_activity_count { @@ -5822,8 +5825,9 @@ pub fn apply_worker_manifest( // terminal result before this bounded deadline; Agen handles only the // mechanical per-call terminalization. worker.set_tool_execution_policy(ToolExecutionPolicy { + pause_safe_boundary_timeout: Duration::from_millis(100), cancellation_request_timeout: Duration::from_millis(250), - terminal_confirmation_timeout: Duration::from_secs(2), + terminal_confirmation_timeout: Duration::from_millis(500), }); worker.set_tool_output_limits(Some(ToolOutputLimits { default_max_bytes: wm.tool_output.default_max_bytes,