fix: bound safe-boundary pause escalation
This commit is contained in:
@@ -1284,6 +1284,7 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
|
|||||||
// commit-before-publish boundary.
|
// commit-before-publish boundary.
|
||||||
let mut terminal_call_ids = HashSet::new();
|
let mut terminal_call_ids = HashSet::new();
|
||||||
let mut pause_requested = false;
|
let mut pause_requested = false;
|
||||||
|
let mut pause_deadline = None;
|
||||||
for result in synthetic_results {
|
for result in synthetic_results {
|
||||||
self.finalize_and_commit_tool_result(
|
self.finalize_and_commit_tool_result(
|
||||||
history,
|
history,
|
||||||
@@ -1319,12 +1320,20 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
pause = self.pause_rx.recv(), if !pause_requested => {
|
pause = self.pause_rx.recv(), if !pause_requested => {
|
||||||
if pause.is_some() {
|
if pause.is_some() {
|
||||||
// Pause is a safe-boundary request: do not cancel provider
|
// Pause first waits for already-started tools to reach a
|
||||||
// operations that already started. Drain all confirmed
|
// natural safe boundary. If they do not, Worker policy
|
||||||
// terminal results, then yield control to Worker.
|
// escalates to the same explicit cancel-and-confirm path.
|
||||||
pause_requested = true;
|
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() => {
|
cancel = self.cancel_rx.recv() => {
|
||||||
if cancel.is_some() {
|
if cancel.is_some() {
|
||||||
info!("Tool execution cancellation requested");
|
info!("Tool execution cancellation requested");
|
||||||
@@ -1404,6 +1413,9 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
self.timeline.abort_current_block();
|
self.timeline.abort_current_block();
|
||||||
|
if pause_requested {
|
||||||
|
return Ok(ToolExecutionResult::Paused);
|
||||||
|
}
|
||||||
return Err(EngineError::Cancelled);
|
return Err(EngineError::Cancelled);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -490,6 +490,9 @@ impl ToolExecutionHandle {
|
|||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub struct ToolExecutionPolicy {
|
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
|
/// Maximum time allowed for a provider to accept one cooperative
|
||||||
/// cancellation request.
|
/// cancellation request.
|
||||||
pub cancellation_request_timeout: std::time::Duration,
|
pub cancellation_request_timeout: std::time::Duration,
|
||||||
@@ -501,6 +504,7 @@ pub struct ToolExecutionPolicy {
|
|||||||
impl Default for ToolExecutionPolicy {
|
impl Default for ToolExecutionPolicy {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
|
pause_safe_boundary_timeout: std::time::Duration::from_millis(100),
|
||||||
cancellation_request_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),
|
terminal_confirmation_timeout: std::time::Duration::from_millis(500),
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ use agen::tool::{
|
|||||||
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
|
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
|
||||||
ToolResultDisposition,
|
ToolResultDisposition,
|
||||||
};
|
};
|
||||||
use agen::{Engine, History, Item};
|
use agen::{Engine, History, Item, ToolExecutionPolicy};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
mod common;
|
mod common;
|
||||||
@@ -610,7 +610,7 @@ async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() {
|
|||||||
.await
|
.await
|
||||||
.expect("tool execution starts");
|
.expect("tool execution starts");
|
||||||
pause.send(()).await.unwrap();
|
pause.send(()).await.unwrap();
|
||||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||||
release.notify_one();
|
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;
|
let output = engine.run(&mut history, "pause safely").await;
|
||||||
control.await.unwrap();
|
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.calls.load(Ordering::SeqCst), 1);
|
||||||
assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0);
|
assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0);
|
||||||
assert!(matches!(output.result, agen::EngineRunExit::Paused));
|
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]
|
#[tokio::test]
|
||||||
async fn cancellation_completion_race_commits_one_terminal_output() {
|
async fn cancellation_completion_race_commits_one_terminal_output() {
|
||||||
for iteration in 0..24u64 {
|
for iteration in 0..24u64 {
|
||||||
|
|||||||
@@ -2581,7 +2581,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
result: &EngineRunExit,
|
result: &EngineRunExit,
|
||||||
snapshot: &EmptyTurnRollbackSnapshot,
|
snapshot: &EmptyTurnRollbackSnapshot,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
if !matches!(result, EngineRunExit::Interrupted(StopReason::Cancelled)) {
|
if !matches!(
|
||||||
|
result,
|
||||||
|
EngineRunExit::Paused | EngineRunExit::Interrupted(StopReason::Cancelled)
|
||||||
|
) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
if self.ai_activity_counter.load(Ordering::SeqCst) != snapshot.ai_activity_count {
|
if self.ai_activity_counter.load(Ordering::SeqCst) != snapshot.ai_activity_count {
|
||||||
@@ -5822,8 +5825,9 @@ pub fn apply_worker_manifest<C: LlmClient + 'static, A>(
|
|||||||
// terminal result before this bounded deadline; Agen handles only the
|
// terminal result before this bounded deadline; Agen handles only the
|
||||||
// mechanical per-call terminalization.
|
// mechanical per-call terminalization.
|
||||||
worker.set_tool_execution_policy(ToolExecutionPolicy {
|
worker.set_tool_execution_policy(ToolExecutionPolicy {
|
||||||
|
pause_safe_boundary_timeout: Duration::from_millis(100),
|
||||||
cancellation_request_timeout: Duration::from_millis(250),
|
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 {
|
worker.set_tool_output_limits(Some(ToolOutputLimits {
|
||||||
default_max_bytes: wm.tool_output.default_max_bytes,
|
default_max_bytes: wm.tool_output.default_max_bytes,
|
||||||
|
|||||||
Reference in New Issue
Block a user