Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: preserve pause reclamation deadline during post-tool policy

Browse Source
This commit is contained in:
Hare
2026-10-12 02:54:12 +09:00
parent 2b1f68a3ea
commit e19742d509
2 changed files with 232 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+71 -53
View File
@@ -105,6 +105,29 @@ async fn await_tool_policy<T>(
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<TokioInstant>) -> Self {
deadline.map_or(Self::Interruptible, Self::UntilPauseSafeBoundary)
}
fn pause_deadline(self) -> Option<TokioInstant> {
match self {
Self::UntilPauseSafeBoundary(deadline) => Some(deadline),
_ => None,
}
}
}
/// Tool registration error
#[derive(Debug, thiserror::Error)]
pub enum ToolRegistryError {
@@ -1723,7 +1746,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
force_cancel: bool,
initial_pause_requested: bool,
) -> Result<ToolExecutionResult, EngineError> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
result,
None,
turn_id,
true,
PostToolPolicyWait::Interruptible,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
@@ -2188,7 +2211,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
result,
Some(&attempt_id),
turn_id,
false,
PostToolPolicyWait::Skip,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
@@ -2297,7 +2320,7 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
);
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
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<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
result,
Some(&attempt_id),
batch.turn_id,
true,
PostToolPolicyWait::Interruptible,
&batch.call_info_map,
&mut batch.attempt_fence,
&mut batch.terminal_call_ids,
crates/agen/tests/parallel_execution_test.rs
+161
View File
@@ -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<TurnEndAction> {
self.turn.wait().await;
Ok(TurnEndAction::Finish)
}
async fn post_tool_call(
&self,
_: &ToolResultInfo<'_>,
) -> InterceptorResult<PostToolAction> {
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<dyn Tool>,
)
})
};
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
);
}