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

fix: terminalize interrupted tool executions

Browse Source
This commit is contained in:
Hare
2026-08-27 20:54:13 +09:00
parent ccabea59c9
commit 58cc94d4b7
16 changed files with 914 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/src/engine.rs
+251 -23
View File
@@ -1,9 +1,10 @@
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::{marker::PhantomData, sync::Arc, time::Instant};
use futures::StreamExt;
use serde_json::{Value, json};
use tokio::sync::mpsc;
use tokio::time::{Duration, Instant as TokioInstant};
use tracing::{debug, info, trace, warn};
use crate::{
@@ -27,11 +28,14 @@ use crate::{
timeline::{TextBlockCollector, ThinkingBlockCollector, Timeline, ToolCallCollector},
tool::{
ToolCall, ToolDefinition as EngineToolDefinition, ToolError, ToolExecutionContext,
ToolOutputLimits, ToolResult, truncate_content,
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 {
@@ -53,6 +57,9 @@ pub enum EngineError {
/// A durable-history observer rejected an item before it entered history.
#[error("History append failed: {0}")]
HistoryAppend(String),
/// Tool terminalization lost its execution-attempt compare-and-set fence.
#[error("Tool execution attempt fence failed: {0}")]
ToolAttemptFence(String),
}
/// Tool registration error
@@ -183,6 +190,64 @@ enum ToolExecutionResult {
Paused,
}
#[derive(Debug, Clone)]
struct ToolExecutionAttempt {
attempt_id: String,
terminal: bool,
}
/// Per-batch compare-and-set fence for terminal ToolResult commits.
///
/// A completion may commit only when its attempt id still matches the active
/// execution for that call and no prior terminal output has won the fence.
#[derive(Debug, Default)]
struct ToolExecutionAttemptFence {
attempts: HashMap<String, ToolExecutionAttempt>,
}
impl ToolExecutionAttemptFence {
fn register(&mut self, call_id: String, attempt_id: String) {
self.attempts.insert(
call_id,
ToolExecutionAttempt {
attempt_id,
terminal: false,
},
);
}
fn can_commit(&self, call_id: &str, attempt_id: &str) -> bool {
matches!(
self.attempts.get(call_id),
Some(attempt) if attempt.attempt_id == attempt_id && !attempt.terminal
)
}
fn commit_terminal(&mut self, call_id: &str, attempt_id: &str) -> bool {
let Some(attempt) = self.attempts.get_mut(call_id) else {
return false;
};
if attempt.attempt_id != attempt_id || attempt.terminal {
return false;
}
attempt.terminal = true;
true
}
fn is_terminal(&self, call_id: &str) -> bool {
self.attempts
.get(call_id)
.is_some_and(|attempt| attempt.terminal)
}
#[cfg(test)]
fn attempt_id(&self, call_id: &str) -> Option<&str> {
self.attempts
.get(call_id)
.map(|attempt| attempt.attempt_id.as_str())
}
}
const MAX_STREAM_CONTINUATIONS: u32 = 3;
/// Central component for managing LLM interactions
@@ -1100,40 +1165,73 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
context.clone(),
),
);
approved_calls.push((tool_call, context));
approved_calls.push((tool_call, context, Some(info.tool)));
} else {
// Unknown tools go into approved list as-is (will error at execution)
let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index);
approved_calls.push((tool_call, context));
approved_calls.push((tool_call, context, None));
}
}
// 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.
let started_calls: Vec<_> = approved_calls
.iter()
.map(|(tool_call, context, _)| (tool_call.id.clone(), context.batch_id.clone()))
.collect();
let mut attempt_fence = ToolExecutionAttemptFence::default();
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)| {
let tool_server = self.tool_server.clone();
async move {
let input_json = serde_json::to_string(&tool_call.input).unwrap_or_default();
match tool_server
.call_tool(&tool_call.name, &input_json, context)
.await
{
.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(e) => ToolResult::error(&tool_call.id, e.to_string()),
}
}
Err(ToolError::Cancelled(output)) => {
ToolResult::from_output_with_disposition(
&tool_call.id,
output,
ToolResultDisposition::Cancelled,
)
}
Err(ToolError::Interrupted(output)) => {
ToolResult::from_output_with_disposition(
&tool_call.id,
output,
ToolResultDisposition::Interrupted,
)
}
Err(error) => ToolResult::error(&tool_call.id, error.to_string()),
},
};
(attempt_id, result)
})
.collect();
// 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();
for result in synthetic_results {
self.finalize_and_commit_tool_result(history, annotate, result, &call_info_map)
.await?;
self.finalize_and_commit_tool_result(
history,
annotate,
result,
None,
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
)
.await?;
}
let mut futures = futures;
@@ -1144,18 +1242,94 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
// output observed before the cancellation boundary.
biased;
result = futures.next() => {
let result = result.expect("non-empty FuturesUnordered returns a result");
let (attempt_id, result) =
result.expect("non-empty FuturesUnordered returns a result");
self.finalize_and_commit_tool_result(
history,
annotate,
result,
Some(&attempt_id),
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
).await?;
}
cancel = self.cancel_rx.recv() => {
if cancel.is_some() {
info!("Tool execution cancelled");
info!("Tool execution cancellation requested");
}
let cancellation_requests = call_info_map
.iter()
.filter(|(call_id, _)| !terminal_call_ids.contains(*call_id))
.map(|(call_id, (_, _, tool, _))| {
let call_id = call_id.clone();
let tool = tool.clone();
async move { (call_id.clone(), tool.cancel(&call_id).await) }
});
let cancellation_requests: FuturesUnordered<_> =
cancellation_requests.collect();
match tokio::time::timeout(
TOOL_CANCEL_SIGNAL_TIMEOUT,
cancellation_requests.collect::<Vec<_>>(),
)
.await
{
Ok(results) => {
for (call_id, result) in results {
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;
while !futures.is_empty() {
tokio::select! {
biased;
result = futures.next() => {
let (attempt_id, result) =
result.expect("non-empty FuturesUnordered returns a result");
self.finalize_and_commit_tool_result(
history,
annotate,
result,
Some(&attempt_id),
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
).await?;
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
// Calls that did not confirm a terminal outcome inside the
// grace period are durably closed as OutcomeUnknown before
// Engine/Worker final status becomes observable.
for (call_id, attempt_id) in &started_calls {
if !attempt_fence.is_terminal(call_id) {
self.finalize_and_commit_tool_result(
history,
annotate,
ToolResult::outcome_unknown(call_id),
Some(attempt_id),
&call_info_map,
&mut attempt_fence,
&mut terminal_call_ids,
).await?;
}
}
self.timeline.abort_current_block();
return Err(EngineError::Cancelled);
}
@@ -1172,6 +1346,7 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
history: &mut History<A>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
mut tool_result: ToolResult,
execution_attempt_id: Option<&str>,
call_info_map: &HashMap<
String,
(
@@ -1181,7 +1356,24 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
ToolExecutionContext,
),
>,
) -> Result<(), EngineError> {
attempt_fence: &mut ToolExecutionAttemptFence,
terminal_call_ids: &mut HashSet<String>,
) -> Result<bool, EngineError> {
let call_id = tool_result.tool_use_id.as_str();
let may_commit = match execution_attempt_id {
Some(attempt_id) => attempt_fence.can_commit(call_id, attempt_id),
None => !terminal_call_ids.contains(call_id),
};
if !may_commit {
warn!(
call_id,
execution_attempt_id,
disposition = ?tool_result.disposition,
"Ignoring stale or duplicate tool result after terminal output commit"
);
return Ok(false);
}
let call_info = call_info_map.get(&tool_result.tool_use_id);
if let Some((tool_call, meta, tool, context)) = call_info {
let mut info = ToolResultInfo {
@@ -1200,6 +1392,10 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
}
tool_result = info.result;
}
if tool_result.is_error && tool_result.disposition.is_success() {
tool_result.disposition = ToolResultDisposition::Error;
}
tool_result.is_error = !tool_result.disposition.is_success();
// Cap content only after post_tool_call so interceptors still observe
// the full payload and any content they inject is bounded too.
@@ -1229,16 +1425,33 @@ impl<C: LlmClient, S: EngineState, A> Engine<C, S, A> {
}
}
let item = Item::tool_result_item_with_attachments(
let item = Item::tool_result_item_with_disposition_and_attachments(
&tool_result.tool_use_id,
&tool_result.summary,
tool_result.content.clone(),
tool_result.is_error,
tool_result.disposition,
tool_result.attachments.clone(),
);
self.append_history_items(history, std::iter::once(item), annotate)?;
if let Some(attempt_id) = execution_attempt_id
&& !attempt_fence.commit_terminal(&tool_result.tool_use_id, attempt_id)
{
return Err(EngineError::ToolAttemptFence(
"tool execution attempt fence changed during terminal commit".to_string(),
));
}
terminal_call_ids.insert(tool_result.tool_use_id.clone());
debug!(
tool = call_info
.map(|(call, _, _, _)| call.name.as_str())
.unwrap_or("unknown"),
call_id = %tool_result.tool_use_id,
execution_attempt_id,
disposition = ?tool_result.disposition,
"Tool execution terminalized"
);
self.emit_tool_result(&tool_result);
Ok(())
Ok(true)
}
/// Internal turn execution logic
@@ -2378,6 +2591,21 @@ mod tests {
use crate::tool::{Attachment, ImageAttachment};
use std::time::Duration;
#[test]
fn tool_execution_attempt_fence_rejects_duplicate_and_stale_results() {
let mut fence = ToolExecutionAttemptFence::default();
fence.register("call".to_string(), "attempt-1".to_string());
assert!(fence.can_commit("call", "attempt-1"));
assert!(fence.commit_terminal("call", "attempt-1"));
assert!(!fence.commit_terminal("call", "attempt-1"));
fence.register("call".to_string(), "attempt-2".to_string());
assert_eq!(fence.attempt_id("call"), Some("attempt-2"));
assert!(!fence.can_commit("call", "attempt-1"));
assert!(!fence.commit_terminal("call", "attempt-1"));
assert!(fence.commit_terminal("call", "attempt-2"));
}
#[test]
fn provider_projection_reorders_results_and_remaps_cache_anchor() {
let items = vec![
crates/agen/src/lib.rs
+3 -1
View File
@@ -28,7 +28,9 @@ pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
pub use interceptor::Interceptor;
pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult};
pub use tool::{
ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult, ToolResultDisposition,
};
pub use usage_record::UsageRecord;
/// Implementation dependencies used by code generated from `agen` macros.
crates/agen/src/llm_client/types.rs
+37 -2
View File
@@ -9,7 +9,7 @@
use std::{fmt, sync::Arc};
use crate::tool::Attachment;
use crate::tool::{Attachment, ToolResultDisposition};
use base64::Engine as _;
use serde::{Deserialize, Serialize};
@@ -121,6 +121,9 @@ pub enum Item {
/// Detailed output (removed by pruning when old enough)
#[serde(default, skip_serializing_if = "Option::is_none")]
content: Option<String>,
/// Typed terminal state used for replay and recovery.
#[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")]
disposition: ToolResultDisposition,
/// Whether the tool result represents an execution error.
#[serde(default, skip_serializing_if = "is_false")]
is_error: bool,
@@ -261,7 +264,17 @@ impl Item {
content: Option<String>,
is_error: bool,
) -> Self {
Self::tool_result_item_with_attachments(call_id, summary, content, is_error, Vec::new())
Self::tool_result_item_with_disposition_and_attachments(
call_id,
summary,
content,
if is_error {
ToolResultDisposition::Error
} else {
ToolResultDisposition::Success
},
Vec::new(),
)
}
/// Create a tool result item with durable, prunable structured attachments.
@@ -272,11 +285,33 @@ impl Item {
is_error: bool,
attachments: Vec<Attachment>,
) -> Self {
Self::tool_result_item_with_disposition_and_attachments(
call_id,
summary,
content,
if is_error {
ToolResultDisposition::Error
} else {
ToolResultDisposition::Success
},
attachments,
)
}
pub fn tool_result_item_with_disposition_and_attachments(
call_id: impl Into<String>,
summary: impl Into<String>,
content: Option<String>,
disposition: ToolResultDisposition,
attachments: Vec<Attachment>,
) -> Self {
let is_error = !disposition.is_success();
Self::ToolResult {
id: None,
call_id: call_id.into(),
summary: summary.into(),
content,
disposition,
is_error,
attachments,
}
crates/agen/src/tool.rs
+68 -1
View File
@@ -23,6 +23,12 @@ pub enum ToolError {
/// Internal error
#[error("Internal error: {0}")]
Internal(String),
/// Cooperative cancellation completed with bounded terminal output.
#[error("Tool execution cancelled")]
Cancelled(ToolOutput),
/// Execution was interrupted with a confirmed bounded terminal output.
#[error("Tool execution interrupted")]
Interrupted(ToolOutput),
}
// =============================================================================
@@ -158,6 +164,28 @@ pub enum Attachment {
Image(ImageAttachment),
}
/// Terminal disposition of one started tool call.
///
/// `Cancelled` means the tool confirmed cancellation. `OutcomeUnknown` means
/// execution stopped without confirmation, so neither completion nor side
/// effects may be inferred.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum ToolResultDisposition {
#[default]
Success,
Error,
Interrupted,
Cancelled,
OutcomeUnknown,
}
impl ToolResultDisposition {
pub const fn is_success(&self) -> bool {
matches!(self, Self::Success)
}
}
/// Tool execution result.
///
/// Every output has a mandatory `summary` (1-2 lines) that persists in
@@ -402,6 +430,17 @@ pub trait Tool: Send + Sync {
input_json: &str,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError>;
/// 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.
async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> {
Ok(())
}
}
// =============================================================================
@@ -429,6 +468,9 @@ pub struct ToolCall {
pub struct ToolResult {
/// Corresponding tool call ID
pub tool_use_id: String,
/// Typed terminal state.
#[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")]
pub disposition: ToolResultDisposition,
/// Short summary (always kept in history)
pub summary: String,
/// Detailed output (prunable)
@@ -445,11 +487,20 @@ pub struct ToolResult {
impl ToolResult {
/// Create a success result from a [`ToolOutput`].
pub fn from_output(tool_use_id: impl Into<String>, output: ToolOutput) -> Self {
Self::from_output_with_disposition(tool_use_id, output, ToolResultDisposition::Success)
}
pub fn from_output_with_disposition(
tool_use_id: impl Into<String>,
output: ToolOutput,
disposition: ToolResultDisposition,
) -> Self {
Self {
tool_use_id: tool_use_id.into(),
disposition,
summary: output.summary,
content: output.content,
is_error: false,
is_error: !disposition.is_success(),
attachments: output.attachments,
}
}
@@ -458,12 +509,28 @@ impl ToolResult {
pub fn error(tool_use_id: impl Into<String>, message: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
disposition: ToolResultDisposition::Error,
summary: message.into(),
content: None,
is_error: true,
attachments: Vec::new(),
}
}
/// Close an execution whose completion and side effects cannot be confirmed.
pub fn outcome_unknown(tool_use_id: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
disposition: ToolResultDisposition::OutcomeUnknown,
summary: "Tool execution outcome unknown".to_string(),
content: Some(
"Execution was interrupted before completion could be confirmed. Completion and side effects are unknown."
.to_string(),
),
is_error: true,
attachments: Vec::new(),
}
}
}
#[cfg(test)]
crates/agen/tests/parallel_execution_test.rs
+212 -11
View File
@@ -10,6 +10,7 @@ use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
ToolResultDisposition,
};
use agen::{Engine, History, Item};
use async_trait::async_trait;
@@ -112,6 +113,53 @@ impl Tool for FirstAttemptHangsTool {
}
}
#[derive(Clone)]
struct CooperativeCancelTool {
calls: Arc<AtomicUsize>,
cancelled: Arc<tokio::sync::Notify>,
}
impl CooperativeCancelTool {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
cancelled: Arc::new(tokio::sync::Notify::new()),
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new("cooperative")
.description("Returns bounded progress after cancellation")
.input_schema(serde_json::json!({"type": "object"}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
}
#[async_trait]
impl Tool for CooperativeCancelTool {
async fn execute(
&self,
_input_json: &str,
_ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.cancelled.notified().await;
Err(ToolError::Cancelled(ToolOutput {
summary: "cooperative command cancelled".to_string(),
content: Some("stdout before cancellation\nstderr before cancellation".to_string()),
attachments: Vec::new(),
}))
}
async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> {
self.cancelled.notify_one();
Ok(())
}
}
#[derive(Clone)]
struct ContextRecordingTool {
name: String,
@@ -269,6 +317,7 @@ async fn completed_results_commit_before_publish_without_waiting_for_siblings()
let _ = engine
.run_with_annotation(&mut history, "run both", &mut annotate)
.await;
observed.lock().unwrap().push("run-returned".to_string());
assert_eq!(
observed.lock().unwrap().as_slice(),
@@ -277,6 +326,7 @@ async fn completed_results_commit_before_publish_without_waiting_for_siblings()
"publish:call_fast",
"commit:call_slow",
"publish:call_slow",
"run-returned",
]
);
@@ -308,9 +358,12 @@ async fn cancellation_preserves_completed_results_and_resume_skips_them() {
Event::tool_use_start(0, "call_hang", "hang_once"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_fast", "fast"),
Event::tool_use_start(1, "call_fast_a", "fast_a"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_use_start(2, "call_fast_b", "fast_b"),
Event::tool_input_delta(2, r#"{}"#),
Event::tool_use_stop(2),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
@@ -326,9 +379,11 @@ async fn cancellation_preserves_completed_results_and_resume_skips_them() {
]);
let mut engine = Engine::new(client);
let hanging = FirstAttemptHangsTool::new();
let fast = SlowTool::new("fast", 1);
let fast_a = SlowTool::new("fast_a", 1);
let fast_b = SlowTool::new("fast_b", 2);
engine.register_tool(hanging.definition());
engine.register_tool(fast.definition());
engine.register_tool(fast_a.definition());
engine.register_tool(fast_b.definition());
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
@@ -345,36 +400,182 @@ async fn cancellation_preserves_completed_results_and_resume_skips_them() {
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. } if call_id == "call_fast"
Item::ToolResult { call_id, .. }
if call_id == "call_fast_a" || call_id == "call_fast_b"
)
})
.count();
assert_eq!(completed_before_resume, 1);
assert_eq!(fast.call_count(), 1);
let unknown_before_resume = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_hang"
)
})
.count();
assert_eq!(completed_before_resume, 2);
assert_eq!(unknown_before_resume, 1);
assert_eq!(fast_a.call_count(), 1);
assert_eq!(fast_b.call_count(), 1);
assert_eq!(hanging.call_count(), 1);
let _ = engine.resume(&mut history).await;
assert_eq!(
fast.call_count(),
fast_a.call_count(),
1,
"completed call must not be re-executed"
);
assert_eq!(
fast_b.call_count(),
1,
"completed call must not be re-executed"
);
assert_eq!(
hanging.call_count(),
2,
"only the unresolved call is retried"
1,
"OutcomeUnknown is terminal and must not be re-executed"
);
let completed_after_resume = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. } if call_id == "call_fast"
Item::ToolResult { call_id, .. }
if call_id == "call_fast_a" || call_id == "call_fast_b"
)
})
.count();
assert_eq!(completed_after_resume, 1);
assert_eq!(completed_after_resume, 2);
assert_eq!(
history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_hang"
)
})
.count(),
1
);
}
#[tokio::test]
async fn cooperative_cancellation_commits_bounded_terminal_output() {
let client = MockLlmClient::with_responses(vec![vec![
Event::tool_use_start(0, "call_cooperative", "cooperative"),
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 = CooperativeCancelTool::new();
engine.register_tool(tool.definition());
let observed = Arc::new(Mutex::new(Vec::<&'static str>::new()));
let published = observed.clone();
engine.on_tool_result(move |_| published.lock().unwrap().push("published"));
let committed = observed.clone();
let mut annotate = move |item: &Item| {
if matches!(item, Item::ToolResult { .. }) {
committed.lock().unwrap().push("committed");
}
Ok(())
};
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(30)).await;
cancel.send(()).await.unwrap();
});
let mut history = History::new();
let output = engine
.run_with_annotation(&mut history, "start", &mut annotate)
.await;
observed.lock().unwrap().push("run-returned");
cancel_task.await.unwrap();
assert_eq!(
observed.lock().unwrap().as_slice(),
["committed", "published", "run-returned"]
);
assert_eq!(tool.calls.load(Ordering::SeqCst), 1);
let terminal: Vec<_> = history
.iter()
.filter_map(|entry| match &entry.item {
Item::ToolResult {
call_id,
disposition,
content,
..
} if call_id == "call_cooperative" => Some((*disposition, content.as_deref())),
_ => None,
})
.collect();
assert_eq!(terminal.len(), 1);
assert_eq!(terminal[0].0, ToolResultDisposition::Cancelled);
assert_eq!(
terminal[0].1,
Some("stdout before cancellation\nstderr before cancellation")
);
assert!(matches!(
output.result,
agen::EngineRunExit::Interrupted(agen::StopReason::Cancelled)
));
}
#[tokio::test]
async fn cancellation_completion_race_commits_one_terminal_output() {
for iteration in 0..24u64 {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_racy", "racy"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})],
]);
let mut engine = Engine::new(client);
let delay = 2 + iteration % 3;
let tool = SlowTool::new("racy", delay);
engine.register_tool(tool.definition());
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(delay)).await;
let _ = cancel.send(()).await;
});
let mut history = History::new();
let _ = engine.run(&mut history, "race").await;
cancel_task.await.unwrap();
let terminal_count = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. } if call_id == "call_racy"
)
})
.count();
assert_eq!(terminal_count, 1, "iteration {iteration}");
assert_eq!(tool.call_count(), 1, "iteration {iteration}");
}
}
#[tokio::test]