use std::sync::Arc;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Instant;
use tokio::sync::RwLock;
use tokio::task::JoinError;
use tokio_util::either::Either;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use tracing::Instrument;
use tracing::info;
use tracing::instrument;
use tracing::trace_span;
use crate::function_tool::FunctionCallError;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::tools::context::AbortedToolOutput;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::context::ToolPayload;
use crate::tools::lifecycle::notify_tool_aborted;
use crate::tools::registry::AnyToolResult;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::router::ToolCall;
use crate::tools::router::ToolCallSource;
use crate::tools::router::ToolRouter;
use codex_protocol::error::CodexErr;
use codex_protocol::models::ResponseInputItem;
struct ToolCallTimingGuard {
started_at: Instant,
execution_started_at: Arc<OnceLock<Instant>>,
conversation_id: String,
turn_id: String,
call_id: String,
tool_name: codex_tools::ToolName,
}
#[derive(Clone)]
pub(crate) struct ToolCallRuntime {
router: Arc<ToolRouter>,
session: Arc<Session>,
step_context: Arc<StepContext>,
tracker: SharedTurnDiffTracker,
parallel_execution: Arc<RwLock<()>>,
}
impl ToolCallRuntime {
pub(crate) fn new(
router: Arc<ToolRouter>,
session: Arc<Session>,
step_context: Arc<StepContext>,
tracker: SharedTurnDiffTracker,
) -> Self {
Self {
router,
session,
step_context,
tracker,
parallel_execution: Arc::new(RwLock::new(())),
}
}
pub(crate) fn create_diff_consumer(
&self,
tool_name: &codex_tools::ToolName,
) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
self.router.create_diff_consumer(tool_name)
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn handle_tool_call(
self,
call: ToolCall,
cancellation_token: CancellationToken,
) -> impl std::future::Future<Output = Result<ResponseInputItem, CodexErr>> {
let error_call = call.clone();
let future =
self.handle_tool_call_with_source(call, ToolCallSource::Direct, cancellation_token);
async move {
match future.await {
Ok(response) => Ok(response.into_response()),
Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)),
Err(other) => Ok(Self::failure_response(error_call, other)),
}
}
.in_current_span()
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn handle_tool_call_with_source(
self,
call: ToolCall,
source: ToolCallSource,
cancellation_token: CancellationToken,
) -> impl std::future::Future<Output = Result<AnyToolResult, FunctionCallError>> {
let supports_parallel = self.router.tool_supports_parallel(&call);
let router = Arc::clone(&self.router);
let session = Arc::clone(&self.session);
let step_context = Arc::clone(&self.step_context);
let turn = Arc::clone(&step_context.turn);
let tracker = Arc::clone(&self.tracker);
let lock = Arc::clone(&self.parallel_execution);
let invocation_cancellation_token = cancellation_token.clone();
let wait_for_runtime_cancellation = self.router.tool_waits_for_runtime_cancellation(&call);
let started = Instant::now();
let tool_call_timing_guard =
ToolCallTimingGuard::capture(started, &session.thread_id, &turn.sub_id, &call, &source);
let execution_started_at = tool_call_timing_guard
.as_ref()
.map(|timing| Arc::clone(&timing.execution_started_at));
let abort_session = Arc::clone(&session);
let abort_source = source.clone();
let abort_turn = Arc::clone(&turn);
let terminal_outcome_reached = Arc::new(AtomicBool::new(false));
let dispatch_terminal_outcome_reached = Arc::clone(&terminal_outcome_reached);
let dispatch_call = call.clone();
let dispatch_span = trace_span!(
"dispatch_tool_call_with_code_mode_result",
otel.name = %call.tool_name,
tool_name = %call.tool_name,
call_id = call.call_id.as_str(),
aborted = false,
);
let abort_dispatch_span = dispatch_span.clone();
let mut dispatch_handle: AbortOnDropHandle<Result<AnyToolResult, FunctionCallError>> =
AbortOnDropHandle::new(tokio::spawn(async move {
let _guard = if supports_parallel {
Either::Left(lock.read().await)
} else {
Either::Right(lock.write().await)
};
if let Some(execution_started_at) = execution_started_at {
let _ = execution_started_at.set(Instant::now());
}
router
.dispatch_tool_call_with_terminal_outcome(
session,
step_context,
invocation_cancellation_token,
tracker,
dispatch_call,
source,
dispatch_terminal_outcome_reached,
)
.instrument(dispatch_span.clone())
.await
}));
async move {
let _tool_call_timing_guard = tool_call_timing_guard;
tokio::select! {
res = &mut dispatch_handle => res.map_err(Self::tool_task_join_error)?,
_ = cancellation_token.cancelled() => {
if terminal_outcome_reached.load(Ordering::Acquire) || dispatch_handle.is_finished() {
dispatch_handle.await.map_err(Self::tool_task_join_error)?
} else {
let secs = started.elapsed().as_secs_f32().max(0.1);
abort_dispatch_span.record("aborted", true);
if wait_for_runtime_cancellation {
if terminal_outcome_reached.swap(true, Ordering::AcqRel) {
return dispatch_handle.await.map_err(Self::tool_task_join_error)?;
}
match dispatch_handle.await {
Ok(_) => {}
Err(err) if err.is_cancelled() => {}
Err(err) => return Err(Self::tool_task_join_error(err)),
}
} else {
dispatch_handle.abort();
match dispatch_handle.await {
Ok(result) => return result,
Err(err) if err.is_cancelled() => {}
Err(err) => return Err(Self::tool_task_join_error(err)),
}
}
let response = Self::aborted_response(&call, secs);
notify_tool_aborted(
abort_session.as_ref(),
abort_turn.as_ref(),
call.call_id.as_str(),
&call.tool_name,
abort_source,
)
.await;
Ok(response)
}
},
}
}
.in_current_span()
}
}
impl ToolCallRuntime {
fn tool_task_join_error(err: JoinError) -> FunctionCallError {
FunctionCallError::Fatal(format!("tool task failed to receive: {err:?}"))
}
fn failure_response(call: ToolCall, err: FunctionCallError) -> ResponseInputItem {
let message = err.to_string();
match call.payload {
ToolPayload::ToolSearch { .. } => ResponseInputItem::ToolSearchOutput {
call_id: call.call_id,
status: "completed".to_string(),
execution: "client".to_string(),
tools: Vec::new(),
},
ToolPayload::Custom { .. } => ResponseInputItem::CustomToolCallOutput {
call_id: call.call_id,
name: None,
output: codex_protocol::models::FunctionCallOutputPayload {
body: codex_protocol::models::FunctionCallOutputBody::Text(message),
success: Some(false),
},
},
_ => ResponseInputItem::FunctionCallOutput {
call_id: call.call_id,
output: codex_protocol::models::FunctionCallOutputPayload {
body: codex_protocol::models::FunctionCallOutputBody::Text(message),
success: Some(false),
},
},
}
}
fn aborted_response(call: &ToolCall, secs: f32) -> AnyToolResult {
AnyToolResult {
call_id: call.call_id.clone(),
payload: call.payload.clone(),
result: Box::new(AbortedToolOutput {
message: Self::abort_message(call, secs),
}),
post_tool_use_payload: None,
}
}
fn abort_message(call: &ToolCall, secs: f32) -> String {
if call.tool_name.namespace.is_none()
&& matches!(
call.tool_name.name.as_str(),
"shell_command" | "unified_exec"
)
{
format!("Wall time: {secs:.1} seconds\naborted by user")
} else {
format!("aborted by user after {secs:.1}s")
}
}
}
impl ToolCallTimingGuard {
fn capture(
started_at: Instant,
conversation_id: &impl std::fmt::Display,
turn_id: &str,
call: &ToolCall,
source: &ToolCallSource,
) -> Option<Self> {
if !matches!(source, ToolCallSource::Direct) || !tracing::enabled!(tracing::Level::INFO) {
return None;
}
Some(Self {
started_at,
execution_started_at: Arc::new(OnceLock::new()),
conversation_id: conversation_id.to_string(),
turn_id: turn_id.to_string(),
call_id: call.call_id.clone(),
tool_name: call.tool_name.clone(),
})
}
}
impl Drop for ToolCallTimingGuard {
fn drop(&mut self) {
let completed_at = Instant::now();
let execution_started_at = self
.execution_started_at
.get()
.copied()
.filter(|execution_started_at| *execution_started_at <= completed_at);
let duration_ms = |duration: std::time::Duration| u64::try_from(duration.as_millis()).ok();
let total_duration_ms = duration_ms(completed_at.duration_since(self.started_at));
let dispatch_duration_ms = execution_started_at.map_or_else(
|| total_duration_ms,
|execution_started_at| {
duration_ms(execution_started_at.duration_since(self.started_at))
},
);
let handler_duration_ms = execution_started_at.map_or(Some(0), |execution_started_at| {
duration_ms(completed_at.duration_since(execution_started_at))
});
macro_rules! log_tool_call {
($dispatch_duration_ms:expr, $handler_duration_ms:expr, $total_duration_ms:expr) => {
info!(
event.name = "codex.tool_call",
trace_id = %codex_otel::current_span_trace_id().unwrap_or_default(),
conversation.id = %self.conversation_id,
turn_id = %self.turn_id,
tool_name = %self.tool_name,
call_id = %self.call_id,
tool_source = "direct",
execution_started = execution_started_at.is_some(),
dispatch_duration_ms = $dispatch_duration_ms,
handler_duration_ms = $handler_duration_ms,
total_duration_ms = $total_duration_ms,
"tool call completed"
);
};
}
match (dispatch_duration_ms, handler_duration_ms, total_duration_ms) {
(Some(dispatch_duration_ms), Some(handler_duration_ms), Some(total_duration_ms)) => {
log_tool_call!(dispatch_duration_ms, handler_duration_ms, total_duration_ms);
}
_ => {
log_tool_call!(
tracing::field::Empty,
tracing::field::Empty,
tracing::field::Empty
);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use crate::session::step_context::StepContext;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolInvocation;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::ToolExecutor;
use crate::tools::registry::ToolRegistry;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_extension_api::ToolCallOutcome;
use codex_protocol::models::FunctionCallOutputBody;
use codex_protocol::models::FunctionCallOutputPayload;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
use tokio::sync::oneshot;
use tracing_test::internal::MockWriter;
#[test]
fn tool_call_timing_guard_ignores_code_mode_source() {
let subscriber = tracing_subscriber::fmt()
.with_max_level(tracing::Level::INFO)
.finish();
tracing::subscriber::with_default(subscriber, || {
let call = ToolCall {
tool_name: codex_tools::ToolName::plain("test_tool"),
call_id: "call-1".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
let direct_guard = ToolCallTimingGuard::capture(
Instant::now(),
&"conversation-id",
"turn-id",
&call,
&ToolCallSource::Direct,
);
assert!(
direct_guard.is_some(),
"direct tool calls should create a timing guard"
);
drop(direct_guard);
let code_mode_guard = ToolCallTimingGuard::capture(
Instant::now(),
&"conversation-id",
"turn-id",
&call,
&ToolCallSource::CodeMode {
cell_id: "cell-1".to_string(),
runtime_tool_call_id: "runtime-call-1".to_string(),
},
);
assert!(
code_mode_guard.is_none(),
"nested code-mode calls should not create overlapping timing events"
);
});
}
#[tokio::test]
async fn cancellation_before_dispatch_admission_logs_dispatch_only_timing() -> anyhow::Result<()>
{
let (session, turn_context) = crate::session::tests::make_session_and_context().await;
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let tool_name = codex_tools::ToolName::plain("test_tool");
let handler = Arc::new(ImmediateHandler {
tool_name: tool_name.clone(),
}) as Arc<dyn CoreToolRuntime>;
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, step_context, tracker);
let execution_gate = Arc::clone(&runtime.parallel_execution);
let execution_gate_guard = execution_gate
.try_write_owned()
.expect("execution gate should be available before dispatch starts");
let (release_execution_gate_tx, release_execution_gate_rx) = std::sync::mpsc::channel();
let execution_gate_task = tokio::task::spawn_blocking(move || {
let _execution_gate_guard = execution_gate_guard;
release_execution_gate_rx
.recv()
.expect("test should release the execution gate");
});
let buffer: &'static std::sync::Mutex<Vec<u8>> =
Box::leak(Box::new(std::sync::Mutex::new(Vec::new())));
let subscriber = tracing_subscriber::fmt()
.with_ansi(false)
.with_max_level(tracing::Level::INFO)
.with_writer(MockWriter::new(buffer))
.finish();
let _subscriber_guard = tracing::subscriber::set_default(subscriber);
let cancellation_token = CancellationToken::new();
let call = ToolCall {
tool_name,
call_id: "call-1".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
let response_task =
tokio::spawn(runtime.handle_tool_call(call, cancellation_token.clone()));
cancellation_token.cancel();
tokio::time::timeout(Duration::from_secs(1), response_task)
.await
.expect("timed out waiting for cancelled tool response")
.expect("cancelled tool response task should join")
.expect("cancelled tool call should produce a response");
let logs = String::from_utf8(
buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone(),
)?;
let timing_events = logs
.lines()
.filter(|line| line.contains("event.name=\"codex.tool_call\""))
.collect::<Vec<_>>();
assert_eq!(
timing_events.len(),
1,
"cancelled tool call should emit exactly one timing event; logs:\n{logs}"
);
let timing_event = timing_events[0];
assert!(
timing_event.contains("execution_started=false"),
"tool cancelled before admission should not report execution started: {timing_event}"
);
assert!(
timing_event.contains("handler_duration_ms=0"),
"tool cancelled before admission should report zero handler duration: {timing_event}"
);
let duration_field = |name: &str| {
timing_event.split_whitespace().find_map(|field| {
field
.strip_prefix(&format!("{name}="))
.and_then(|value| value.parse::<u64>().ok())
})
};
let dispatch_duration_ms = duration_field("dispatch_duration_ms")
.expect("timing event should include dispatch_duration_ms");
let total_duration_ms = duration_field("total_duration_ms")
.expect("timing event should include total_duration_ms");
assert_eq!(
dispatch_duration_ms, total_duration_ms,
"tool cancelled before admission should attribute all elapsed time to dispatch: {timing_event}"
);
release_execution_gate_tx
.send(())
.expect("execution gate task should remain available");
execution_gate_task
.await
.expect("execution gate task should join");
Ok(())
}
struct ImmediateHandler {
tool_name: codex_tools::ToolName,
}
impl ToolExecutor<ToolInvocation> for ImmediateHandler {
fn tool_name(&self) -> codex_tools::ToolName {
self.tool_name.clone()
}
fn spec(&self) -> codex_tools::ToolSpec {
codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {
name: self.tool_name.name.clone(),
description: "Immediate test tool.".to_string(),
strict: false,
defer_loading: None,
parameters: codex_tools::JsonSchema::default(),
output_schema: None,
})
}
fn handle(&self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
Box::pin(async {
Ok(
Box::new(FunctionToolOutput::from_text("ok".to_string(), Some(true)))
as Box<dyn crate::tools::context::ToolOutput>,
)
})
}
}
impl CoreToolRuntime for ImmediateHandler {}
struct CancellationCleanupHandler {
tool_name: codex_tools::ToolName,
started: std::sync::Mutex<Option<oneshot::Sender<()>>>,
cleanup_started: std::sync::Mutex<Option<oneshot::Sender<()>>>,
allow_cleanup: Arc<Notify>,
}
impl ToolExecutor<ToolInvocation> for CancellationCleanupHandler {
fn tool_name(&self) -> codex_tools::ToolName {
self.tool_name.clone()
}
fn spec(&self) -> codex_tools::ToolSpec {
codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {
name: self.tool_name.name.clone(),
description: "Cancellation cleanup test tool.".to_string(),
strict: false,
defer_loading: None,
parameters: codex_tools::JsonSchema::default(),
output_schema: None,
})
}
fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
Box::pin(self.handle_call(invocation))
}
}
impl CancellationCleanupHandler {
async fn handle_call(
&self,
invocation: ToolInvocation,
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let started = self
.started
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
if let Some(started) = started {
let _ = started.send(());
}
invocation.cancellation_token.cancelled().await;
let cleanup_started = self
.cleanup_started
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
if let Some(cleanup_started) = cleanup_started {
let _ = cleanup_started.send(());
}
self.allow_cleanup.notified().await;
Ok(Box::new(FunctionToolOutput::from_text(
"cleanup complete".to_string(),
Some(false),
)) as Box<dyn crate::tools::context::ToolOutput>)
}
}
impl CoreToolRuntime for CancellationCleanupHandler {
fn waits_for_runtime_cancellation(&self) -> bool {
true
}
}
struct FinishRecorder {
records: Arc<std::sync::Mutex<Vec<ToolCallOutcome>>>,
}
impl codex_extension_api::ToolLifecycleContributor for FinishRecorder {
fn on_tool_finish<'a>(
&'a self,
input: codex_extension_api::ToolFinishInput<'a>,
) -> codex_extension_api::ToolLifecycleFuture<'a> {
let records = Arc::clone(&self.records);
let outcome = input.outcome;
Box::pin(async move {
records
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(outcome);
})
}
}
struct BlockingFinishContributor {
records: Arc<std::sync::Mutex<Vec<ToolCallOutcome>>>,
finish_started: std::sync::Mutex<Option<oneshot::Sender<()>>>,
allow_finish: Arc<Notify>,
}
impl codex_extension_api::ToolLifecycleContributor for BlockingFinishContributor {
fn on_tool_finish<'a>(
&'a self,
input: codex_extension_api::ToolFinishInput<'a>,
) -> codex_extension_api::ToolLifecycleFuture<'a> {
let records = Arc::clone(&self.records);
let allow_finish = Arc::clone(&self.allow_finish);
let finish_started = self
.finish_started
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
let outcome = input.outcome;
Box::pin(async move {
if let Some(finish_started) = finish_started {
let _ = finish_started.send(());
}
allow_finish.notified().await;
records
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(outcome);
})
}
}
#[tokio::test]
async fn cancellation_after_handler_finishes_preserves_completed_lifecycle()
-> anyhow::Result<()> {
let (mut session, turn_context) = crate::session::tests::make_session_and_context().await;
let records = Arc::new(std::sync::Mutex::new(Vec::new()));
let (finish_started_tx, finish_started_rx) = oneshot::channel();
let allow_finish = Arc::new(Notify::new());
let mut builder =
codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
builder.tool_lifecycle_contributor(Arc::new(BlockingFinishContributor {
records: Arc::clone(&records),
finish_started: std::sync::Mutex::new(Some(finish_started_tx)),
allow_finish: Arc::clone(&allow_finish),
}));
session.services.extensions = Arc::new(builder.build());
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let tool_name = codex_tools::ToolName::plain("test_tool");
let handler = Arc::new(ImmediateHandler {
tool_name: tool_name.clone(),
}) as Arc<dyn CoreToolRuntime>;
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, step_context, tracker);
let cancellation_token = CancellationToken::new();
let call = ToolCall {
tool_name,
call_id: "call-1".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
let response_task =
tokio::spawn(runtime.handle_tool_call(call, cancellation_token.clone()));
tokio::time::timeout(Duration::from_secs(1), finish_started_rx)
.await
.expect("timed out waiting for lifecycle notification to start")
.expect("lifecycle notification should start");
cancellation_token.cancel();
tokio::time::sleep(Duration::from_millis(10)).await;
allow_finish.notify_waiters();
let response = tokio::time::timeout(Duration::from_secs(1), response_task)
.await
.expect("timed out waiting for tool response")
.expect("tool response task should join")?;
let expected_response = ResponseInputItem::FunctionCallOutput {
call_id: "call-1".to_string(),
output: FunctionCallOutputPayload {
body: FunctionCallOutputBody::Text("ok".to_string()),
success: Some(true),
},
};
assert_eq!(expected_response, response);
let actual = records
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.drain(..)
.collect::<Vec<_>>();
assert_eq!(vec![ToolCallOutcome::Completed { success: true }], actual);
Ok(())
}
#[tokio::test]
async fn cancellation_waiting_for_runtime_cleanup_emits_only_aborted_lifecycle()
-> anyhow::Result<()> {
let (mut session, turn_context) = crate::session::tests::make_session_and_context().await;
let records = Arc::new(std::sync::Mutex::new(Vec::new()));
let mut builder =
codex_extension_api::ExtensionRegistryBuilder::<crate::config::Config>::new();
builder.tool_lifecycle_contributor(Arc::new(FinishRecorder {
records: Arc::clone(&records),
}));
session.services.extensions = Arc::new(builder.build());
let session = Arc::new(session);
let turn_context = Arc::new(turn_context);
let tool_name = codex_tools::ToolName::plain("cleanup_tool");
let (started_tx, started_rx) = oneshot::channel();
let (cleanup_started_tx, cleanup_started_rx) = oneshot::channel();
let allow_cleanup = Arc::new(Notify::new());
let handler = Arc::new(CancellationCleanupHandler {
tool_name: tool_name.clone(),
started: std::sync::Mutex::new(Some(started_tx)),
cleanup_started: std::sync::Mutex::new(Some(cleanup_started_tx)),
allow_cleanup: Arc::clone(&allow_cleanup),
}) as Arc<dyn CoreToolRuntime>;
let step_context = StepContext::for_test(Arc::clone(&turn_context));
let router = Arc::new(ToolRouter::from_parts(
ToolRegistry::from_tools([handler]),
Vec::new(),
));
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let runtime = ToolCallRuntime::new(router, session, step_context, tracker);
let cancellation_token = CancellationToken::new();
let call = ToolCall {
tool_name,
call_id: "call-1".to_string(),
payload: ToolPayload::Function {
arguments: "{}".to_string(),
},
};
let response_task =
tokio::spawn(runtime.handle_tool_call(call, cancellation_token.clone()));
started_rx.await.expect("handler should start");
cancellation_token.cancel();
cleanup_started_rx
.await
.expect("handler should start cleanup");
tokio::time::sleep(Duration::from_millis(10)).await;
allow_cleanup.notify_one();
let response = tokio::time::timeout(Duration::from_secs(1), response_task)
.await
.expect("timed out waiting for tool response")
.expect("tool response task should join")?;
let ResponseInputItem::FunctionCallOutput { output, .. } = response else {
anyhow::bail!("cancelled tool should return function output");
};
let FunctionCallOutputBody::Text(text) = output.body else {
anyhow::bail!("cancelled tool output should be text");
};
assert!(text.contains("aborted by user"));
let actual = records
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.drain(..)
.collect::<Vec<_>>();
assert_eq!(vec![ToolCallOutcome::Aborted], actual);
Ok(())
}
}