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

fix: integrate signal-aware Agen policies with Codex refresh bounds

Browse Source
This commit is contained in:
Hare
2026-10-12 03:14:25 +09:00
parent 9e1de93a04 e19742d509
commit ae26cdfbef
2 changed files with 899 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+165 -121
View File
@@ -71,6 +71,63 @@ pub enum EngineError {
ToolAttemptFence(String),
}
/// Policy completion is not permission to start work when an interruption was
/// queued during its final poll. Cancel wins when both requests are observed.
fn check_policy_interruption(
cancel_rx: &mut mpsc::Receiver<()>,
pause_rx: &mut mpsc::Receiver<()>,
) -> Result<(), EngineError> {
if cancel_rx.try_recv().is_ok() {
return Err(EngineError::Cancelled);
}
if pause_rx.try_recv().is_ok() {
return Err(EngineError::PauseRequested);
}
Ok(())
}
/// Poll one policy invocation, never a second copy of its side effects. Dropping
/// a pending policy stops admission, not any already-owned tool execution; the
/// caller must drain that execution through its original attempt fence.
async fn await_tool_policy<T>(
policy: impl Future<Output = Result<T, InterceptorError>>,
phase: InterceptorPhase,
cancel_rx: &mut mpsc::Receiver<()>,
pause_rx: &mut mpsc::Receiver<()>,
) -> Result<T, EngineError> {
let result = tokio::select! {
biased;
Some(()) = cancel_rx.recv() => return Err(EngineError::Cancelled),
Some(()) = pause_rx.recv() => return Err(EngineError::PauseRequested),
result = policy => result,
};
check_policy_interruption(cancel_rx, pause_rx)?;
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 {
@@ -748,14 +805,6 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
}
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
@@ -1531,43 +1580,24 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
tool,
context,
};
let action = tokio::select! {
result = self.interceptor.pre_tool_call(&mut info) => {
result.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
error,
))
})?
}
pause = self.pause_rx.recv() => {
if pause.is_some() {
info!(call_id = %expected_tool_use_id, "Paused while awaiting pre-tool policy");
}
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
None,
)?;
let action = await_tool_policy(
self.interceptor.pre_tool_call(&mut info),
InterceptorPhase::PreToolCall,
&mut self.cancel_rx,
&mut self.pause_rx,
)
.await;
let action = match action {
Ok(action) => action,
Err(error @ (EngineError::PauseRequested | EngineError::Cancelled)) => {
self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?;
batch.dispatched_any = true;
return Ok(EarlyToolAdmission::Pause);
}
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!(call_id = %expected_tool_use_id, "Cancelled while awaiting pre-tool policy");
}
self.append_early_tool_call(
history,
annotate,
&tool_call,
call_index,
None,
)?;
batch.dispatched_any = true;
return Err(EngineError::Cancelled);
return match error {
EngineError::PauseRequested => Ok(EarlyToolAdmission::Pause),
_ => Err(error),
};
}
Err(error) => return Err(error),
};
if info.call.id != expected_tool_use_id {
return Err(InterceptorFailure::new(
@@ -1716,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,
@@ -1743,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,
@@ -1791,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,
@@ -1814,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,
@@ -1837,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)
@@ -1872,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;
@@ -1883,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,
@@ -1945,8 +1975,10 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
// 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);
match check_policy_interruption(&mut self.cancel_rx, &mut self.pause_rx) {
Ok(()) => {}
Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused),
Err(error) => return Err(error),
}
// Map from tool call ID to (ToolCall, Meta, Tool, Context)
@@ -1985,16 +2017,18 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
context,
};
let pre_tool_action =
self.interceptor
.pre_tool_call(&mut info)
.await
.map_err(|error| {
EngineError::from(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
error,
))
})?;
let pre_tool_action = match await_tool_policy(
self.interceptor.pre_tool_call(&mut info),
InterceptorPhase::PreToolCall,
&mut self.cancel_rx,
&mut self.pause_rx,
)
.await
{
Ok(action) => action,
Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused),
Err(error) => return Err(error),
};
if info.call.id != expected_tool_use_id {
return Err(InterceptorFailure::new(
InterceptorPhase::PreToolCall,
@@ -2061,6 +2095,14 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
}
}
// Recheck after all policies and immediately before provider ownership.
// No tool task exists yet: these calls remain proven NotStarted.
match check_policy_interruption(&mut self.cancel_rx, &mut self.pause_rx) {
Ok(()) => {}
Err(EngineError::PauseRequested) => return Ok(ToolExecutionResult::Paused),
Err(error) => return Err(error),
}
// Phase 2: Execute approved tools in parallel. FuturesUnordered yields
// each terminal result as soon as that call completes instead of
// holding fast siblings behind the slowest call in the batch.
@@ -2137,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,
@@ -2169,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,
@@ -2250,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,
@@ -2278,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,
@@ -2328,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,
(
@@ -2413,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),
@@ -2429,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));
@@ -2487,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,
@@ -2949,22 +2986,29 @@ impl<C: LlmClient, S: EngineState, A: Send + Sync> Engine<C, S, A> {
)),
0,
);
let assistant_turn_action = self
.interceptor
.on_assistant_turn_end(AssistantTurnEndContext {
invocation: assistant_invocation,
assistant_entries: &history.entries()[assistant_start..],
history: history.entries(),
tool_calls: &tool_calls,
})
.await;
let assistant_turn_action = await_tool_policy(
self.interceptor
.on_assistant_turn_end(AssistantTurnEndContext {
invocation: assistant_invocation,
assistant_entries: &history.entries()[assistant_start..],
history: history.entries(),
tool_calls: &tool_calls,
}),
InterceptorPhase::AssistantTurnEnd,
&mut self.cancel_rx,
&mut self.pause_rx,
)
.await;
let assistant_turn_action = match assistant_turn_action {
Ok(action) => action,
Err(EngineError::PauseRequested) => {
if let Some(batch) = early_tools.take() {
self.drain_early_tool_batch(history, annotate, batch, false, true)
.await?;
}
return Ok(EngineResult::Paused);
}
Err(error) => {
let error = EngineError::from(InterceptorFailure::new(
InterceptorPhase::AssistantTurnEnd,
error,
));
if let Some(batch) = early_tools.take()
&& let Err(cleanup_error) = self
.drain_early_tool_batch(history, annotate, batch, true, false)
@@ -3530,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
+734
View File
@@ -2873,3 +2873,737 @@ async fn post_tool_stop_terminalizes_started_parallel_siblings_before_returning(
assert!(matches!(next, EngineRunExit::Finished));
}
}
// The same policy barrier is exercised through the real response-complete,
// provider fallback, interceptor fallback, and early-dispatch engine paths.
#[derive(Clone, Copy, Debug)]
enum PolicyRoute {
AfterResponse,
ProviderFallback,
InterceptorFallback,
Early,
}
impl PolicyRoute {
fn completion_support(self) -> ToolCallCompletionSupport {
match self {
Self::ProviderFallback => ToolCallCompletionSupport::ResponseComplete,
_ => ToolCallCompletionSupport::PerBlock,
}
}
fn configure(self, engine: &mut Engine<ControlledStreamClient>) {
engine.set_max_turns(Some(1));
if matches!(self, Self::AfterResponse) {
engine.set_tool_call_dispatch_mode(ToolCallDispatchMode::AfterResponse);
}
}
}
struct NotifyOnDrop(Arc<tokio::sync::Notify>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
self.0.notify_one();
}
}
#[derive(Clone)]
struct PolicyBarrier {
phase: InterceptorPhase,
early_compatible: bool,
entered: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
invocations: Arc<AtomicUsize>,
continues: Arc<AtomicUsize>,
retired: Arc<tokio::sync::Notify>,
// Enqueue from the policy's final poll, after select has polled signals.
// This deterministically exercises the policy-completion/ownership race.
signal_on_continue: Option<tokio::sync::mpsc::Sender<()>>,
}
impl PolicyBarrier {
fn new(phase: InterceptorPhase, route: PolicyRoute) -> Self {
Self {
phase,
early_compatible: !matches!(route, PolicyRoute::InterceptorFallback),
entered: Arc::new(tokio::sync::Notify::new()),
release: Arc::new(tokio::sync::Notify::new()),
invocations: Arc::new(AtomicUsize::new(0)),
continues: Arc::new(AtomicUsize::new(0)),
retired: Arc::new(tokio::sync::Notify::new()),
signal_on_continue: None,
}
}
async fn wait(&self) {
let _retired = NotifyOnDrop(self.retired.clone());
self.invocations.fetch_add(1, Ordering::SeqCst);
self.entered.notify_one();
self.release.notified().await;
if let Some(signal) = &self.signal_on_continue {
signal.try_send(()).unwrap();
}
self.continues.fetch_add(1, Ordering::SeqCst);
}
}
#[async_trait]
impl Interceptor for PolicyBarrier {
fn supports_early_tool_dispatch(&self) -> bool {
self.early_compatible
}
async fn pre_tool_call(
&self,
_info: &mut ToolCallInfo<'_>,
) -> InterceptorResult<PreToolAction> {
if self.phase == InterceptorPhase::PreToolCall {
self.wait().await;
}
Ok(PreToolAction::Continue)
}
async fn on_assistant_turn_end(
&self,
_context: AssistantTurnEndContext<'_>,
) -> InterceptorResult<TurnEndAction> {
if self.phase == InterceptorPhase::AssistantTurnEnd {
self.wait().await;
}
Ok(TurnEndAction::Finish)
}
}
fn send_policy_response(tx: &tokio::sync::mpsc::UnboundedSender<Result<Event, ClientError>>) {
for event in [
Event::tool_use_start(0, "policy_call", "policy_tool"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::usage(10, 3),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
] {
tx.send(Ok(event)).unwrap();
}
}
fn assert_policy_interruption(exit: &EngineRunExit, pause: bool) {
if pause {
assert!(matches!(exit, EngineRunExit::Paused), "{exit:?}");
} else {
assert!(
matches!(
exit,
EngineRunExit::Interrupted(RunInterruptionReason::Cancelled)
),
"{exit:?}"
);
}
}
fn call_and_result_count(history: &History) -> (usize, usize) {
(
history
.items()
.filter(
|item| matches!(item, Item::ToolCall { call_id, .. } if call_id == "policy_call"),
)
.count(),
history
.items()
.filter(
|item| matches!(item, Item::ToolResult { call_id, .. } if call_id == "policy_call"),
)
.count(),
)
}
#[tokio::test]
async fn response_complete_pending_policies_interrupt_before_release_and_preserve_not_started() {
for route in [
PolicyRoute::AfterResponse,
PolicyRoute::ProviderFallback,
PolicyRoute::InterceptorFallback,
] {
for phase in [
InterceptorPhase::PreToolCall,
InterceptorPhase::AssistantTurnEnd,
] {
for pause in [false, true] {
let (client, tx) = ControlledStreamClient::new(route.completion_support());
let mut engine = Engine::new(client);
route.configure(&mut engine);
let policy = PolicyBarrier::new(phase, route);
let probe = policy.clone();
engine.set_interceptor(policy);
let tool = SlowTool::new("policy_tool", 0);
engine.register_tool(tool.definition());
let usage = Arc::new(AtomicUsize::new(0));
let usage_probe = usage.clone();
engine.on_usage(move |_| {
usage_probe.fetch_add(1, Ordering::SeqCst);
});
let signal = if pause {
engine.pause_sender()
} else {
engine.cancel_sender()
};
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
send_policy_response(&tx);
drop(tx);
tokio::time::timeout(Duration::from_secs(2), probe.entered.notified())
.await
.expect("policy entered");
signal.send(()).await.unwrap();
let (output, history) = tokio::time::timeout(Duration::from_secs(2), run).await
.unwrap_or_else(|_| panic!("{route:?}/{phase:?}/pause={pause}: interruption waited for policy release"))
.unwrap();
assert_policy_interruption(&output.result, pause);
assert_eq!(tool.call_count(), 0);
assert_eq!(call_and_result_count(&history), (1, 0));
assert!(history.items().any(|item| matches!(
item,
Item::ToolCall {
execution_id: None,
status: None | Some(agen::llm_client::ItemStatus::Completed),
..
}
)));
assert_eq!(usage.load(Ordering::SeqCst), 1);
assert_eq!(probe.invocations.load(Ordering::SeqCst), 1);
probe.release.notify_one(); // late permission cannot create a task
assert_eq!(probe.continues.load(Ordering::SeqCst), 0);
assert_eq!(tool.call_count(), 0);
// Recreate both history and engine as a restore boundary. Only
// proven NotStarted is allowed to run, with normal pre-tool policy.
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.set_max_turns(Some(1));
restored.register_tool(tool.definition());
let _ = restored.resume(&mut restored_history).await;
assert_eq!(tool.call_count(), 1);
assert_eq!(call_and_result_count(&restored_history), (1, 1));
assert_eq!(usage.load(Ordering::SeqCst), 1);
}
}
}
}
#[tokio::test]
async fn policy_final_poll_signal_prevents_provider_ownership_even_with_continue() {
for route in [
PolicyRoute::AfterResponse,
PolicyRoute::ProviderFallback,
PolicyRoute::InterceptorFallback,
PolicyRoute::Early,
] {
for pause in [false, true] {
let (client, tx) = ControlledStreamClient::new(route.completion_support());
let mut engine = Engine::new(client);
route.configure(&mut engine);
let mut policy = PolicyBarrier::new(InterceptorPhase::PreToolCall, route);
policy.signal_on_continue = Some(if pause {
engine.pause_sender()
} else {
engine.cancel_sender()
});
let probe = policy.clone();
engine.set_interceptor(policy);
let tool = SlowTool::new("policy_tool", 0);
engine.register_tool(tool.definition());
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
send_policy_response(&tx);
drop(tx);
tokio::time::timeout(Duration::from_secs(2), probe.entered.notified())
.await
.expect("policy entered");
probe.release.notify_one();
let (output, history) = tokio::time::timeout(Duration::from_secs(2), run)
.await
.expect("final-poll signal processed")
.unwrap();
assert_policy_interruption(&output.result, pause);
assert_eq!(probe.continues.load(Ordering::SeqCst), 1);
assert_eq!(probe.invocations.load(Ordering::SeqCst), 1);
assert_eq!(
tool.call_count(),
0,
"{route:?}/pause={pause}: Continue bypassed queued interruption"
);
assert_eq!(call_and_result_count(&history), (1, 0));
}
}
}
// Provider owns the terminal channel. Cancel acknowledges receipt but does not
// manufacture a result, so tests can release a real terminal output or deliver
// it after the engine's confirmation deadline (without a live side effect).
#[derive(Clone)]
struct PolicyOwnedTool {
starts: Arc<AtomicUsize>,
started: Arc<tokio::sync::Notify>,
cancel_requested: Arc<tokio::sync::Notify>,
terminal: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Receiver<ToolOutput>>>>,
retired: Arc<tokio::sync::Notify>,
}
#[async_trait]
impl Tool for PolicyOwnedTool {
async fn execute(&self, _: &str, _: ToolExecutionContext) -> Result<ToolOutput, ToolError> {
let _retired = NotifyOnDrop(self.retired.clone());
self.starts.fetch_add(1, Ordering::SeqCst);
self.started.notify_one();
Ok(self.terminal.lock().await.take().unwrap().await.unwrap())
}
async fn cancel(&self, _: &str) -> Result<(), ToolError> {
self.cancel_requested.notify_one();
Ok(())
}
}
#[tokio::test]
async fn pending_turn_end_interrupts_drain_owned_early_attempt_and_fence_late_result_on_restore() {
for pause in [false, true] {
for late in [false, true] {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let mut engine = Engine::new(client);
PolicyRoute::Early.configure(&mut engine);
engine.set_tool_execution_policy(ToolExecutionPolicy {
pause_safe_boundary_timeout: Duration::ZERO,
cancellation_request_timeout: Duration::from_secs(1),
terminal_confirmation_timeout: if late {
Duration::ZERO
} else {
Duration::from_secs(2)
},
});
let policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early);
let probe = policy.clone();
engine.set_interceptor(policy);
let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel();
let tool = 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 tool = tool.clone();
Arc::new(move || {
(
ToolMeta::new("policy_tool")
.input_schema(serde_json::json!({"type":"object"})),
Arc::new(tool.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 |_| {
assert_eq!(
commit_probe.load(Ordering::SeqCst),
1,
"commit before publication"
);
publish_probe.fetch_add(1, Ordering::SeqCst);
});
let usage = Arc::new(AtomicUsize::new(0));
let usage_probe = usage.clone();
engine.on_usage(move |_| {
usage_probe.fetch_add(1, Ordering::SeqCst);
});
let signal = if pause {
engine.pause_sender()
} else {
engine.cancel_sender()
};
let run = tokio::spawn(async move {
let mut history = History::new();
let output = engine.run(&mut history, "run").await;
(output, history)
});
// Wait for ownership before completing the assistant response.
for event in [
Event::tool_use_start(0, "policy_call", "policy_tool"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
] {
tx.send(Ok(event)).unwrap();
}
tokio::time::timeout(Duration::from_secs(2), tool.started.notified())
.await
.expect("early provider started");
tx.send(Ok(Event::usage(10, 3))).unwrap();
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
tokio::time::timeout(Duration::from_secs(2), probe.entered.notified())
.await
.expect("turn-end policy pending");
signal.send(()).await.unwrap();
tokio::time::timeout(Duration::from_secs(2), tool.cancel_requested.notified())
.await
.expect("owned provider cancellation progresses before policy release");
let terminal_tx = if late {
Some(terminal_tx)
} else {
terminal_tx
.send(ToolOutput::from("confirmed after interruption".to_string()))
.unwrap();
None
};
let (output, history) = tokio::time::timeout(Duration::from_secs(2), run)
.await
.expect("terminalization does not await turn-end policy")
.unwrap();
assert_policy_interruption(&output.result, pause);
assert_eq!(call_and_result_count(&history), (1, 1));
assert_eq!(committed.load(Ordering::SeqCst), 1);
assert_eq!(published.load(Ordering::SeqCst), 1);
assert_eq!(usage.load(Ordering::SeqCst), 1);
assert!(history.items().any(|item| matches!(
item,
Item::ToolCall {
execution_id: Some(_),
status: Some(agen::llm_client::ItemStatus::InProgress),
..
}
)));
assert!(history.items().any(|item| matches!(item, Item::ToolResult { disposition, .. } if *disposition == if late { ToolResultDisposition::OutcomeUnknown } else { ToolResultDisposition::Success })));
if let Some(terminal_tx) = terminal_tx {
tokio::time::timeout(Duration::from_secs(2), tool.retired.notified())
.await
.expect("forced-close retired the owned receiver");
// Force-close retires the receiver; a provider's late completion
// cannot reenter history or replace OutcomeUnknown.
assert!(
terminal_tx
.send(ToolOutput::from("late completion".to_string()))
.is_err()
);
}
probe.release.notify_one();
assert_eq!(probe.continues.load(Ordering::SeqCst), 0);
let original_items = history.items().cloned().collect::<Vec<_>>();
let mut resumed_history = history.clone();
let mut resumed_engine = output.engine;
let _ = resumed_engine.resume(&mut resumed_history).await;
assert_eq!(
tool.starts.load(Ordering::SeqCst),
1,
"same-engine resume cannot rerun a terminal attempt"
);
assert_eq!(call_and_result_count(&resumed_history), (1, 1));
let mut restored_history: History =
serde_json::from_str(&serde_json::to_string(&history).unwrap()).unwrap();
let restored_client = MockLlmClient::new(vec![Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})]);
let request_probe = restored_client.clone();
let mut restored = Engine::new(restored_client);
restored.register_tool(definition);
restored.set_max_turns(Some(1));
let _ = restored.resume(&mut restored_history).await;
let requests = request_probe.requests();
assert_eq!(requests.len(), 1);
let replay = requests[0]
.items
.iter()
.filter_map(|item| match item {
Item::ToolCall { call_id, .. } => Some(format!("call:{call_id}")),
Item::ToolResult { call_id, .. } => Some(format!("result:{call_id}")),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(replay, ["call:policy_call", "result:policy_call"]);
assert_eq!(tool.starts.load(Ordering::SeqCst), 1);
assert_eq!(call_and_result_count(&restored_history), (1, 1));
assert_eq!(
restored_history
.items()
.take(original_items.len())
.cloned()
.collect::<Vec<_>>(),
original_items,
"restore must not rewrite original history"
);
assert_eq!(published.load(Ordering::SeqCst), 1);
}
}
}
#[tokio::test]
async fn pending_turn_end_pause_preserves_natural_safe_boundary_without_cancelling_started_tool() {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let mut engine = Engine::new(client);
PolicyRoute::Early.configure(&mut engine);
engine.set_tool_execution_policy(ToolExecutionPolicy {
pause_safe_boundary_timeout: Duration::from_secs(2),
..ToolExecutionPolicy::default()
});
let policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early);
let probe = policy.clone();
engine.set_interceptor(policy);
let tool = SafePauseTool::new();
// Use this tool's actual public name in the stream fixture.
engine.register_tool(tool.definition());
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 event in [
Event::tool_use_start(0, "policy_call", "safe_pause"),
Event::tool_input_delta(0, "{}"),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
] {
tx.send(Ok(event)).unwrap();
}
drop(tx);
tokio::time::timeout(Duration::from_secs(2), probe.entered.notified())
.await
.expect("turn-end entered");
pause.send(()).await.unwrap();
tokio::time::timeout(Duration::from_secs(2), probe.retired.notified())
.await
.expect("pending policy interrupted before tool release");
tool.release.notify_one();
let (output, history) = tokio::time::timeout(Duration::from_secs(2), run)
.await
.expect("safe boundary collected")
.unwrap();
assert_policy_interruption(&output.result, true);
assert_eq!(tool.calls.load(Ordering::SeqCst), 1);
assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0);
assert_eq!(call_and_result_count(&history), (1, 1));
probe.release.notify_one();
assert_eq!(probe.continues.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn turn_end_final_poll_signal_wins_over_finish_without_tools() {
for pause in [false, true] {
let (client, tx) = ControlledStreamClient::new(ToolCallCompletionSupport::PerBlock);
let mut engine = Engine::new(client);
let mut policy = PolicyBarrier::new(InterceptorPhase::AssistantTurnEnd, PolicyRoute::Early);
policy.signal_on_continue = Some(if pause {
engine.pause_sender()
} else {
engine.cancel_sender()
});
let probe = policy.clone();
engine.set_interceptor(policy);
let run = tokio::spawn(async move {
let mut history = History::new();
engine.run(&mut history, "run").await
});
tx.send(Ok(Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})))
.unwrap();
drop(tx);
tokio::time::timeout(Duration::from_secs(2), probe.entered.notified())
.await
.expect("turn-end entered");
probe.release.notify_one();
let output = tokio::time::timeout(Duration::from_secs(2), run)
.await
.expect("final-poll signal processed")
.unwrap();
assert_policy_interruption(&output.result, pause);
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
);
}