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 + ); +}