diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 3af11903..c48f2128 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1,10 +1,10 @@ use std::collections::{HashMap, HashSet}; -use std::{marker::PhantomData, sync::Arc, time::Instant}; +use std::{future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Instant}; use futures::StreamExt; use serde_json::{Value, json}; use tokio::sync::mpsc; -use tokio::time::{Duration, Instant as TokioInstant}; +use tokio::time::Instant as TokioInstant; use tracing::{debug, info, trace, warn}; use crate::{ @@ -28,14 +28,12 @@ use crate::{ timeline::{TextBlockCollector, ThinkingBlockCollector, Timeline, ToolCallCollector}, tool::{ ToolCall, ToolDefinition as EngineToolDefinition, ToolError, ToolExecutionContext, - ToolOutputLimits, ToolResult, ToolResultDisposition, truncate_content, + ToolExecutionHandle, ToolExecutionPolicy, ToolExecutionTerminal, ToolOutputLimits, + ToolResult, ToolResultDisposition, truncate_content, }, tool_server::{ToolServer, ToolServerHandle}, }; -const TOOL_CANCEL_SIGNAL_TIMEOUT: Duration = Duration::from_millis(100); -const TOOL_CANCEL_GRACE_PERIOD: Duration = Duration::from_millis(500); - /// Engine errors #[derive(Debug, thiserror::Error)] pub enum EngineError { @@ -51,6 +49,9 @@ pub enum EngineError { /// Cancelled by CancellationToken #[error("Cancelled")] Cancelled, + /// Paused by the caller at the next safe boundary. + #[error("Paused")] + PauseRequested, /// Config warnings (unsupported options) #[error("Config warnings: {}", .0.iter().map(|w| w.to_string()).collect::>().join(", "))] ConfigWarnings(Vec), @@ -169,6 +170,7 @@ impl From> for EngineRunExit { Self::Interrupted(StopReason::ContextWindowExceeded) } Err(EngineError::Cancelled) => Self::Interrupted(StopReason::Cancelled), + Err(EngineError::PauseRequested) => Self::Paused, Err(error) => Self::Interrupted(StopReason::Unexpected(error)), } } @@ -346,6 +348,8 @@ pub struct Engine { tool_execution_batch_count: usize, /// Maximum number of AgentTurns (None = unlimited) max_turns: Option, + /// Caller-selected policy for interrupting started provider operations. + tool_execution_policy: ToolExecutionPolicy, /// AgentTurn-start callbacks (1:1 with LlmCall today) turn_start_cbs: Vec>, /// AgentTurn-end callbacks (1:1 with LlmCall today) @@ -384,6 +388,10 @@ pub struct Engine { /// Cancel notification channel (for interrupting execution) cancel_tx: mpsc::Sender<()>, cancel_rx: mpsc::Receiver<()>, + /// Pause notification channel. Unlike cancellation, this waits for already + /// started tools to reach provider-confirmed terminal results. + pause_tx: mpsc::Sender<()>, + pause_rx: mpsc::Receiver<()>, /// Byte-size caps applied to tool `content` before it reaches history. /// `None` disables truncation (tests and minimal setups). tool_output_limits: Option, @@ -421,7 +429,10 @@ impl Engine { } fn finish_logical_run(&mut self, result: &Result) { - if !matches!(result, Ok(EngineResult::Paused) | Ok(EngineResult::Yielded)) { + if !matches!( + result, + Ok(EngineResult::Paused | EngineResult::Yielded) | Err(EngineError::PauseRequested) + ) { self.active_run_turn_count = None; } } @@ -430,14 +441,19 @@ impl Engine { while self.cancel_rx.try_recv().is_ok() {} } - /// Discard pending cancellation notifications while the engine is idle. + fn drain_pause_queue(&mut self) { + while self.pause_rx.try_recv().is_ok() {} + } + + /// Discard pending interruption notifications while the engine is idle. /// - /// Cancellation is a running-turn control signal. Callers that own a higher - /// level run state can use this before starting a new turn so an old idle - /// signal does not poison the next request, while cancellation queued after - /// the run has been accepted remains observable by the turn loop. + /// Cancellation and pause are running-turn control signals. Callers that own + /// a higher level run state can use this before starting a new turn so an old + /// idle signal does not poison the next request, while interruption queued + /// after the run has been accepted remains observable by the turn loop. pub fn clear_pending_cancel(&mut self) { self.drain_cancel_queue(); + self.drain_pause_queue(); } fn try_cancelled(&mut self) -> bool { @@ -449,6 +465,14 @@ 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 @@ -912,6 +936,22 @@ impl Engine { self.cancel_tx.clone() } + /// Get the safe-boundary pause notification sender. + pub fn pause_sender(&self) -> mpsc::Sender<()> { + self.pause_tx.clone() + } + + /// Select the deadline policy applied to already-started provider operations. + /// Worker/controller layers own this lifecycle policy; Agen owns only the + /// mechanical terminalization of each call. + pub fn set_tool_execution_policy(&mut self, policy: ToolExecutionPolicy) { + self.tool_execution_policy = policy; + } + + pub fn tool_execution_policy(&self) -> ToolExecutionPolicy { + self.tool_execution_policy + } + /// Set request configuration at once pub fn set_request_config(&mut self, config: RequestConfig) { self.request_config = config; @@ -1108,6 +1148,12 @@ impl Engine { ) -> Result { use futures::stream::{FuturesUnordered, StreamExt}; + // 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); + } + // Map from tool call ID to (ToolCall, Meta, Tool, Context) // Retained because it's needed for PostToolCall hooks let mut call_info_map = HashMap::new(); @@ -1184,43 +1230,60 @@ impl Engine { for (call_id, attempt_id) in &started_calls { attempt_fence.register(call_id.clone(), attempt_id.clone()); } - let futures: FuturesUnordered<_> = approved_calls - .into_iter() - .map(|(tool_call, context, tool)| async move { - let attempt_id = context.batch_id.clone(); - let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); - let result = match tool { - None => ToolResult::error( - &tool_call.id, - format!("Tool not found: {}", tool_call.name), - ), - Some(tool) => match tool.execute(&input_json, context).await { - Ok(output) => ToolResult::from_output(&tool_call.id, output), - Err(ToolError::Cancelled(output)) => { - ToolResult::from_output_with_disposition( - &tool_call.id, + let futures: FuturesUnordered + Send>>> = + FuturesUnordered::new(); + let mut execution_handles = HashMap::new(); + for (tool_call, context, tool) in approved_calls { + let attempt_id = context.batch_id.clone(); + let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default(); + let call_id = tool_call.id.clone(); + let future: Pin + Send>> = match tool { + None => { + let result = + ToolResult::error(&call_id, format!("Tool not found: {}", tool_call.name)); + Box::pin(async move { (attempt_id, result) }) + } + Some(tool) => { + let (handle, terminal) = ToolExecutionHandle::start(tool, input_json, context); + execution_handles.insert(call_id.clone(), handle); + Box::pin(async move { + let result = match terminal.await { + ToolExecutionTerminal::Confirmed(Ok(output)) => { + ToolResult::from_output(&call_id, output) + } + ToolExecutionTerminal::Confirmed(Err(ToolError::Cancelled(output))) => { + ToolResult::from_output_with_disposition( + &call_id, + output, + ToolResultDisposition::Cancelled, + ) + } + ToolExecutionTerminal::Confirmed(Err(ToolError::Interrupted( output, - ToolResultDisposition::Cancelled, - ) - } - Err(ToolError::Interrupted(output)) => { - ToolResult::from_output_with_disposition( - &tool_call.id, + ))) => ToolResult::from_output_with_disposition( + &call_id, output, ToolResultDisposition::Interrupted, - ) - } - Err(error) => ToolResult::error(&tool_call.id, error.to_string()), - }, - }; - (attempt_id, result) - }) - .collect(); + ), + ToolExecutionTerminal::Confirmed(Err(error)) => { + ToolResult::error(&call_id, error.to_string()) + } + ToolExecutionTerminal::OutcomeUnknown => { + ToolResult::outcome_unknown(&call_id) + } + }; + (attempt_id, result) + }) + } + }; + futures.push(future); + } // Synthetic results are already terminal and need no execution wait. // Commit them before polling ordinary calls so they obey the same // commit-before-publish boundary. let mut terminal_call_ids = HashSet::new(); + let mut pause_requested = false; for result in synthetic_results { self.finalize_and_commit_tool_result( history, @@ -1254,45 +1317,52 @@ impl Engine { &mut terminal_call_ids, ).await?; } + 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_requested = true; + } + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Tool execution cancellation requested"); } - let cancellation_requests = call_info_map + let cancellation_request_deadline = TokioInstant::now() + + self.tool_execution_policy.cancellation_request_timeout; + let cancellation_requests = execution_handles .iter() .filter(|(call_id, _)| !terminal_call_ids.contains(*call_id)) - .map(|(call_id, (_, _, tool, _))| { + .map(|(call_id, handle)| { let call_id = call_id.clone(); - let tool = tool.clone(); - async move { (call_id.clone(), tool.cancel(&call_id).await) } + let handle = handle.clone(); + async move { + ( + call_id, + handle.cancel_before(cancellation_request_deadline).await, + ) + } }); let cancellation_requests: FuturesUnordered<_> = cancellation_requests.collect(); - match tokio::time::timeout( - TOOL_CANCEL_SIGNAL_TIMEOUT, - cancellation_requests.collect::>(), - ) - .await - { - Ok(results) => { - for (call_id, result) in results { - if let Err(error) = result { - warn!( - %call_id, - error = %error, - "Tool cooperative cancellation request failed" - ); - } - } + for (call_id, result) in cancellation_requests.collect::>().await { + if let Err(error) = result { + warn!( + %call_id, + error = %error, + "Tool cooperative cancellation request failed" + ); } - Err(_) => warn!("Tool cooperative cancellation request timed out"), } - // Keep polling the original execution futures for a bounded - // grace period so cooperative providers can return their - // confirmed terminal output, including bounded progress. - let deadline = TokioInstant::now() + TOOL_CANCEL_GRACE_PERIOD; + // Keep polling the original execution handles to their + // provider-confirmed terminal result until the caller-selected + // deadline. The execution remains owned even if this polling + // future is later dropped. + let deadline = TokioInstant::now() + + self.tool_execution_policy.terminal_confirmation_timeout; while !futures.is_empty() { tokio::select! { biased; @@ -1318,6 +1388,9 @@ impl Engine { // Engine/Worker final status becomes observable. for (call_id, attempt_id) in &started_calls { if !attempt_fence.is_terminal(call_id) { + if let Some(handle) = execution_handles.get(call_id) { + handle.force_close(); + } self.finalize_and_commit_tool_result( history, annotate, @@ -1336,7 +1409,11 @@ impl Engine { } } - Ok(ToolExecutionResult::Completed) + Ok(if pause_requested { + ToolExecutionResult::Paused + } else { + ToolExecutionResult::Completed + }) } /// Apply post-execution policy, bound the model-visible payload, durably @@ -1751,6 +1828,13 @@ impl Engine { let stream_started = Instant::now(); let stream_result = tokio::select! { stream_result = self.client.stream(request.clone()) => stream_result, + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused before stream started"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled before stream started"); @@ -1782,6 +1866,13 @@ impl Engine { ); let first_event_result = tokio::select! { first_event = wait_for_first_stream_event(stream, DEFAULT_FIRST_STREAM_EVENT_TIMEOUT) => first_event, + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused before first stream event"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled before first stream event"); @@ -1868,6 +1959,13 @@ impl Engine { tokio::select! { _ = tokio::time::sleep(wait) => {} + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused during LLM retry backoff"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Cancelled during LLM retry backoff"); @@ -1946,6 +2044,13 @@ impl Engine { None => break, } } + pause = self.pause_rx.recv() => { + if pause.is_some() { + info!("Paused during response stream"); + } + self.timeline.abort_current_block(); + return Err(EngineError::PauseRequested); + } cancel = self.cancel_rx.recv() => { if cancel.is_some() { info!("Stream cancelled"); @@ -1987,6 +2092,7 @@ impl Engine { let thinking_block_collector = ThinkingBlockCollector::new(); let mut timeline = Timeline::new(); let (cancel_tx, cancel_rx) = mpsc::channel(1); + let (pause_tx, pause_rx) = mpsc::channel(1); // Register collectors with Timeline timeline.on_text_block(text_block_collector.clone()); @@ -2009,6 +2115,7 @@ impl Engine { llm_call_count: 0, tool_execution_batch_count: 0, max_turns: None, + tool_execution_policy: ToolExecutionPolicy::default(), turn_start_cbs: Vec::new(), turn_end_cbs: Vec::new(), llm_call_start_cbs: Vec::new(), @@ -2023,6 +2130,8 @@ impl Engine { request_config: RequestConfig::default(), cancel_tx, cancel_rx, + pause_tx, + pause_rx, tool_output_limits: None, prune_config: None, token_estimator: None, @@ -2281,6 +2390,7 @@ impl Engine { llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, llm_call_start_cbs: self.llm_call_start_cbs, @@ -2296,6 +2406,8 @@ impl Engine { cancel_tx: self.cancel_tx, cancel_rx: self.cancel_rx, + pause_tx: self.pause_tx, + pause_rx: self.pause_rx, tool_output_limits: self.tool_output_limits, prune_config: self.prune_config, token_estimator: self.token_estimator, @@ -2390,7 +2502,10 @@ impl Engine { self.append_history_items(history, extras, annotate)?; } self.start_logical_run(); - let result = self.run_turn_loop(history, annotate).await; + let result = match self.run_turn_loop(history, annotate).await { + Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), + other => other, + }; let result = self.finalize_interruption(result).await; self.finish_logical_run(&result); result @@ -2413,7 +2528,10 @@ impl Engine { annotate: &mut impl FnMut(&Item) -> Result, ) -> Result { self.ensure_logical_run(); - let result = self.run_turn_loop(history, annotate).await; + let result = match self.run_turn_loop(history, annotate).await { + Err(EngineError::PauseRequested) => Ok(EngineResult::Paused), + other => other, + }; let result = self.finalize_interruption(result).await; self.finish_logical_run(&result); result @@ -2445,6 +2563,7 @@ impl Engine { llm_call_count: self.llm_call_count, tool_execution_batch_count: self.tool_execution_batch_count, max_turns: self.max_turns, + tool_execution_policy: self.tool_execution_policy, turn_start_cbs: self.turn_start_cbs, turn_end_cbs: self.turn_end_cbs, llm_call_start_cbs: self.llm_call_start_cbs, @@ -2460,6 +2579,8 @@ impl Engine { cancel_tx: self.cancel_tx, cancel_rx: self.cancel_rx, + pause_tx: self.pause_tx, + pause_rx: self.pause_rx, tool_output_limits: self.tool_output_limits, prune_config: self.prune_config, token_estimator: self.token_estimator, diff --git a/crates/agen/src/lib.rs b/crates/agen/src/lib.rs index aac1727a..e6ea9953 100644 --- a/crates/agen/src/lib.rs +++ b/crates/agen/src/lib.rs @@ -29,7 +29,9 @@ pub use history::{History, HistoryEntry}; pub use interceptor::Interceptor; pub use message::{ContentPart, Item, Message, Role}; pub use tool::{ - ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult, ToolResultDisposition, + ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy, + ToolExecutionTerminal, ToolExecutionTerminalFuture, ToolOutputLimits, ToolResult, + ToolResultDisposition, }; pub use usage_record::UsageRecord; diff --git a/crates/agen/src/tool.rs b/crates/agen/src/tool.rs index 0ae5e96d..9001c89c 100644 --- a/crates/agen/src/tool.rs +++ b/crates/agen/src/tool.rs @@ -3,7 +3,14 @@ //! Traits for defining tools callable by LLM. //! Usually auto-implemented using the `#[tool]` macro. -use std::{collections::HashMap, fmt, sync::Arc}; +use std::{ + collections::HashMap, + fmt, + future::Future, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose::STANDARD}; @@ -350,6 +357,12 @@ impl ToolExecutionContext { } } + /// Identifies one live execution attempt without making the batch id a durable + /// replay or idempotency authority. + pub fn execution_id(&self) -> String { + format!("{}:{}", self.batch_id, self.call_id) + } + /// Context for direct, non-engine calls in unit tests and low-level callers. pub fn direct() -> Self { Self::new("direct", "direct", 0) @@ -362,6 +375,138 @@ impl Default for ToolExecutionContext { } } +/// The provider-confirmed terminal result of one started tool execution. +/// +/// `OutcomeUnknown` is reserved for an execution task that had to be force-closed +/// or failed before the provider could confirm its terminal result. +#[derive(Debug)] +pub enum ToolExecutionTerminal { + Confirmed(Result), + OutcomeUnknown, +} + +/// The completion future paired with a [`ToolExecutionHandle`]. Dropping this +/// future does not drop the provider execution: the spawned execution remains +/// owned by its handle until it completes or is explicitly force-closed. +pub struct ToolExecutionTerminalFuture { + task: tokio::task::JoinHandle>, +} + +impl Future for ToolExecutionTerminalFuture { + type Output = ToolExecutionTerminal; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match Pin::new(&mut self.task).poll(cx) { + Poll::Ready(Ok(result)) => Poll::Ready(ToolExecutionTerminal::Confirmed(result)), + Poll::Ready(Err(_)) => Poll::Ready(ToolExecutionTerminal::OutcomeUnknown), + Poll::Pending => Poll::Pending, + } + } +} + +/// Live ownership and control for one started tool execution. +/// +/// Execution, cancellation, and terminal confirmation remain provider-owned: +/// this handle starts `Tool::execute`, delegates cooperative cancellation to +/// `Tool::cancel_execution`, and treats execution-future completion as the +/// provider's terminal confirmation. Agen may force-close only after its caller's +/// deadline expires, at which point the outcome is necessarily unknown. +#[derive(Clone)] +pub struct ToolExecutionHandle { + inner: Arc, +} + +struct ToolExecutionHandleInner { + tool: Arc, + context: ToolExecutionContext, + abort: tokio::task::AbortHandle, +} + +impl Drop for ToolExecutionHandleInner { + fn drop(&mut self) { + // Losing the final live owner is an explicit forced close, never a + // best-effort detached provider future. + self.abort.abort(); + } +} + +impl fmt::Debug for ToolExecutionHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ToolExecutionHandle") + .field("call_id", &self.inner.context.call_id) + .field("batch_id", &self.inner.context.batch_id) + .finish_non_exhaustive() + } +} + +impl ToolExecutionHandle { + pub fn start( + tool: Arc, + input_json: String, + context: ToolExecutionContext, + ) -> (Self, ToolExecutionTerminalFuture) { + let execution_tool = Arc::clone(&tool); + let execution_context = context.clone(); + let task = + tokio::spawn( + async move { execution_tool.execute(&input_json, execution_context).await }, + ); + let abort = task.abort_handle(); + ( + Self { + inner: Arc::new(ToolExecutionHandleInner { + tool, + context, + abort, + }), + }, + ToolExecutionTerminalFuture { task }, + ) + } + + pub fn context(&self) -> &ToolExecutionContext { + &self.inner.context + } + + pub async fn cancel_before(&self, deadline: tokio::time::Instant) -> Result<(), ToolError> { + match tokio::time::timeout_at( + deadline, + self.inner.tool.cancel_execution(&self.inner.context), + ) + .await + { + Ok(result) => result, + Err(_) => Err(ToolError::Internal(format!( + "tool cancellation request exceeded its deadline for call {}", + self.inner.context.call_id + ))), + } + } + + pub fn force_close(&self) { + self.inner.abort.abort(); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ToolExecutionPolicy { + /// Maximum time allowed for a provider to accept one cooperative + /// cancellation request. + pub cancellation_request_timeout: std::time::Duration, + /// Maximum time allowed for all providers to confirm terminal results after + /// cancellation has been requested. + pub terminal_confirmation_timeout: std::time::Duration, +} + +impl Default for ToolExecutionPolicy { + fn default() -> Self { + Self { + cancellation_request_timeout: std::time::Duration::from_millis(100), + terminal_confirmation_timeout: std::time::Duration::from_millis(500), + } + } +} + // ============================================================================= // Tool trait // ============================================================================= @@ -434,13 +579,22 @@ pub trait Tool: Send + Sync { /// Request cooperative cancellation for one started call. /// /// Implementations that own cancellable provider operations should signal - /// the exact execution identified by `call_id`, then let `execute` return - /// the confirmed bounded terminal output. The Engine applies a bounded - /// grace period and falls back to `OutcomeUnknown` when confirmation never - /// arrives. + /// every live execution identified by `call_id`, then let `execute` return + /// the confirmed bounded terminal output. Direct callers may use this + /// compatibility surface; Agen uses [`Tool::cancel_execution`] so providers + /// can bind cancellation to one exact live attempt. async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> { Ok(()) } + + /// Request cooperative cancellation for one exact started execution. + /// + /// The default preserves existing tools by delegating to `cancel(call_id)`. + /// Providers with their own execution registry should override this method + /// and key cancellation by [`ToolExecutionContext::execution_id`]. + async fn cancel_execution(&self, ctx: &ToolExecutionContext) -> Result<(), ToolError> { + self.cancel(&ctx.call_id).await + } } // ============================================================================= diff --git a/crates/agen/tests/parallel_execution_test.rs b/crates/agen/tests/parallel_execution_test.rs index 295ff3d8..13108b48 100644 --- a/crates/agen/tests/parallel_execution_test.rs +++ b/crates/agen/tests/parallel_execution_test.rs @@ -160,6 +160,55 @@ impl Tool for CooperativeCancelTool { } } +#[derive(Clone)] +struct SafePauseTool { + calls: Arc, + cancellations: Arc, + release: Arc, +} + +impl SafePauseTool { + fn new() -> Self { + Self { + calls: Arc::new(AtomicUsize::new(0)), + cancellations: Arc::new(AtomicUsize::new(0)), + release: Arc::new(tokio::sync::Notify::new()), + } + } + + fn definition(&self) -> ToolDefinition { + let tool = self.clone(); + Arc::new(move || { + let meta = ToolMeta::new("safe_pause") + .description("Waits for a safe-boundary release") + .input_schema(serde_json::json!({"type": "object"})); + (meta, Arc::new(tool.clone()) as Arc) + }) + } +} + +#[async_trait] +impl Tool for SafePauseTool { + async fn execute( + &self, + _input_json: &str, + _ctx: ToolExecutionContext, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.release.notified().await; + Ok(ToolOutput { + summary: "safe-boundary complete".to_string(), + content: Some("safe-boundary complete".to_string()), + attachments: Vec::new(), + }) + } + + async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> { + self.cancellations.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + #[derive(Clone)] struct ContextRecordingTool { name: String, @@ -535,6 +584,55 @@ async fn cooperative_cancellation_commits_bounded_terminal_output() { )); } +#[tokio::test] +async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() { + let client = MockLlmClient::with_responses(vec![vec![ + Event::tool_use_start(0, "call_safe_pause", "safe_pause"), + Event::tool_input_delta(0, r#"{}"#), + Event::tool_use_stop(0), + Event::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ]]); + let mut engine = Engine::new(client); + let tool = SafePauseTool::new(); + engine.register_tool(tool.definition()); + + let pause = engine.pause_sender(); + let calls = Arc::clone(&tool.calls); + let release = Arc::clone(&tool.release); + 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(); + tokio::time::sleep(Duration::from_millis(100)).await; + release.notify_one(); + }); + + let started_at = std::time::Instant::now(); + let mut history = History::new(); + let output = engine.run(&mut history, "pause safely").await; + control.await.unwrap(); + + assert!(started_at.elapsed() >= Duration::from_millis(100)); + assert_eq!(tool.calls.load(Ordering::SeqCst), 1); + assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0); + assert!(matches!(output.result, agen::EngineRunExit::Paused)); + assert!(history.iter().any(|entry| matches!( + &entry.item, + Item::ToolResult { + call_id, + disposition: ToolResultDisposition::Success, + .. + } if call_id == "call_safe_pause" + ))); +} + #[tokio::test] async fn cancellation_completion_race_commits_one_terminal_output() { for iteration in 0..24u64 { diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index c0e48ead..237721b2 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -485,6 +485,7 @@ impl WorkerController { // into the controller task so the in-flight turn can be reached // via these handles while worker itself is borrowed by drive_turn. let cancel_tx = worker.engine_mut().cancel_sender(); + let pause_tx = worker.engine_mut().pause_sender(); let notify_buffer = worker.notify_buffer_handle(); tokio::spawn(controller_loop( @@ -494,6 +495,7 @@ impl WorkerController { shared_state, runtime_dir, cancel_tx, + pause_tx, notify_buffer, self_parent_socket, spawner_name, @@ -1136,6 +1138,7 @@ async fn controller_loop( shared_state: Arc, runtime_dir: Arc, cancel_tx: mpsc::Sender<()>, + pause_tx: mpsc::Sender<()>, notify_buffer: NotifyBuffer, self_parent_socket: Option, spawner_name: String, @@ -1207,6 +1210,7 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, &runtime_dir, Some(input_commit_rx), @@ -1231,6 +1235,7 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, &runtime_dir, Some(input_commit_rx), @@ -1248,6 +1253,7 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, &runtime_dir, None, @@ -1265,6 +1271,7 @@ async fn controller_loop( &mut method_rx, &event_tx, &cancel_tx, + &pause_tx, &shared_state, &runtime_dir, None, @@ -1664,6 +1671,7 @@ async fn drive_turn( method_rx: &mut mpsc::Receiver, event_tx: &broadcast::Sender, cancel_tx: &mpsc::Sender<()>, + pause_tx: &mpsc::Sender<()>, shared_state: &Arc, runtime_dir: &RuntimeDir, mut input_commit_rx: Option>, @@ -1707,6 +1715,9 @@ where return match result { Ok(r) => { let (status, run_result) = match r { + WorkerRunResult::Finished if pause_requested => { + (WorkerStatus::Paused, RunResult::Paused) + } WorkerRunResult::Finished => (WorkerStatus::Idle, RunResult::Finished), WorkerRunResult::Paused => (WorkerStatus::Paused, RunResult::Paused), WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached), @@ -1779,7 +1790,7 @@ where } Some(Method::Pause) => { pause_requested = true; - let _ = cancel_tx.try_send(()); + let _ = pause_tx.try_send(()); } Some(Method::Shutdown) => { shutdown_requested = true; @@ -2031,6 +2042,8 @@ mod tests { event_tx: broadcast::Sender, cancel_tx: mpsc::Sender<()>, _cancel_rx: mpsc::Receiver<()>, + pause_tx: mpsc::Sender<()>, + _pause_rx: mpsc::Receiver<()>, shared_state: Arc, notify_buffer: NotifyBuffer, spawned_registry: Arc, @@ -2049,6 +2062,7 @@ mod tests { let (method_tx, method_rx) = mpsc::channel::(16); let (event_tx, _) = broadcast::channel::(16); let (cancel_tx, cancel_rx) = mpsc::channel::<()>(1); + let (pause_tx, pause_rx) = mpsc::channel::<()>(1); let shared_state = Arc::new(WorkerSharedState::new( "child-worker".to_string(), session_store::new_segment_id(), @@ -2074,6 +2088,8 @@ mod tests { event_tx, cancel_tx, _cancel_rx: cancel_rx, + pause_tx, + _pause_rx: pause_rx, shared_state, notify_buffer, spawned_registry, @@ -2131,6 +2147,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2154,6 +2171,44 @@ mod tests { } } + #[tokio::test] + async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() { + let mut env = make_env().await; + let method_tx = env._method_tx.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + method_tx.send(Method::Pause).await.expect("send pause"); + }); + + let worker_future = async { + tokio::time::sleep(Duration::from_millis(100)).await; + Ok::<_, WorkerError>(WorkerRunResult::Finished) + }; + let started_at = std::time::Instant::now(); + let (status, shutdown) = drive_turn( + worker_future, + &mut env.method_rx, + &env.event_tx, + &env.cancel_tx, + &env.pause_tx, + &env.shared_state, + &env.runtime_dir, + None, + &env.notify_buffer, + None, + "child-worker", + &env.spawned_registry, + true, + ) + .await; + + assert_eq!(status, WorkerStatus::Paused); + assert!(!shutdown); + assert!(started_at.elapsed() >= Duration::from_millis(100)); + assert!(env._pause_rx.try_recv().is_ok()); + assert!(env._cancel_rx.try_recv().is_err()); + } + #[tokio::test] async fn non_parent_originated_finished_stays_silent() { let mut env = make_env().await; @@ -2165,6 +2220,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2202,6 +2258,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2245,6 +2302,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2286,6 +2344,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2324,6 +2383,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2360,6 +2420,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, @@ -2395,6 +2456,7 @@ mod tests { &mut env.method_rx, &env.event_tx, &env.cancel_tx, + &env.pause_tx, &env.shared_state, &env.runtime_dir, None, diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 1991bc34..b72dee6e 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -11,7 +11,7 @@ use agen::llm_client::types::Role; use agen::state::Mutable; use agen::{ Engine, EngineError, EngineResult, EngineRunExit, History, HistoryEntry, Item, StopReason, - ToolOutputLimits, UsageRecord, + ToolExecutionPolicy, ToolOutputLimits, UsageRecord, }; use arc_swap::ArcSwap; use session_store::{ @@ -5817,6 +5817,14 @@ pub fn apply_worker_manifest( ) { worker.set_request_config(request_config_from_engine_manifest(wm)); worker.set_max_turns(wm.max_turns.map(|n| n.get())); + // Worker owns the lifecycle strategy for already-started tool operations. + // The provider must first accept cooperative cancellation, then confirm a + // terminal result before this bounded deadline; Agen handles only the + // mechanical per-call terminalization. + worker.set_tool_execution_policy(ToolExecutionPolicy { + cancellation_request_timeout: Duration::from_millis(250), + terminal_confirmation_timeout: Duration::from_secs(2), + }); worker.set_tool_output_limits(Some(ToolOutputLimits { default_max_bytes: wm.tool_output.default_max_bytes, per_tool: wm.tool_output.per_tool.clone(), @@ -5932,6 +5940,7 @@ fn stop_reason_error_code(reason: &StopReason) -> ErrorCode { | StopReason::Unexpected( EngineError::Aborted(_) | EngineError::Cancelled + | EngineError::PauseRequested | EngineError::ConfigWarnings(_) | EngineError::HistoryAppend(_) | EngineError::ToolAttemptFence(_),