Close abandoned tool dispatch traces as aborted (#61434)

## Why

Interrupting a turn while a tool awaits approval can leave its rollout trace execution marked as running without an end timestamp.

## What changed

Record an `Aborted` end event when a tool dispatch trace is dropped before producing a result. Make completion and failure methods consume the trace handle so each dispatch is closed once.

## Testing

Add coverage for dropping a pending dispatch and for interrupting a turn with a pending command approval, verifying that replay reports an aborted execution with an end timestamp.

GitOrigin-RevId: 63cf68312fcbff8c6b85a01a7713cf4460c9d795
This commit is contained in:
Max Johnson
2026-11-14 01:32:55 +00:00
committed by copyberry
parent de83e7674a
commit 9ac8c3f8d7
6 changed files with 125 additions and 56 deletions
@@ -256,6 +256,7 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<()
];
let tmp = TempDir::new()?;
let traces = TempDir::new()?;
let codex_home = tmp.path().join("codex_home");
std::fs::create_dir(&codex_home)?;
let working_directory = tmp.path().join("workdir");
@@ -276,6 +277,7 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<()
let mut mcp = TestAppServer::builder()
.with_codex_home(&codex_home)
.with_env_overrides(&[("CODEX_ROLLOUT_TRACE_ROOT", traces.path().to_str())])
.build_initialized()
.await?;
@@ -341,5 +343,18 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<()
assert_eq!(completed.thread_id, thread.id);
assert_eq!(completed.turn.status, TurnStatus::Interrupted);
assert!(mcp.shutdown_gracefully().await?.success());
let bundles = std::fs::read_dir(traces.path())?.collect::<std::io::Result<Vec<_>>>()?;
let [bundle] = bundles.as_slice() else {
anyhow::bail!("expected one trace bundle, got {}", bundles.len());
};
let trace = codex_rollout_trace::replay_bundle(bundle.path())?;
let call = &trace.tool_calls["call_sleep_approval"];
assert_eq!(
call.execution.status,
codex_rollout_trace::ExecutionStatus::Aborted
);
assert!(call.execution.ended_at_unix_ms >= Some(call.execution.started_at_unix_ms));
Ok(())
}
+33 -30
View File
@@ -853,40 +853,43 @@ impl ToolRegistry {
return Err(err);
}
let fail_before_execution = async |invocation: &ToolInvocation, err: FunctionCallError| {
let log_payload = tool_log_payload(&invocation.payload, &invocation.source);
otel.tool_result_with_tags(
&tool_name,
&call_id_owned,
log_payload.as_ref(),
Duration::ZERO,
/*success*/ false,
&err.to_string(),
&tool_result_tags,
&extra_trace_fields,
);
if tool.is_builtin_control_tool() {
let mut analytics = ControlToolCallGuard::new(invocation);
analytics.finish(ControlToolCallStatus::Failed);
}
dispatch_trace.record_failed(&err);
notify_tool_finish_if_unclaimed(
invocation,
call_state.as_deref(),
ToolCallOutcome::Failed {
handler_executed: false,
},
)
.await;
err
};
let fail_before_execution =
async |dispatch_trace: ToolDispatchTrace,
invocation: &ToolInvocation,
err: FunctionCallError| {
let log_payload = tool_log_payload(&invocation.payload, &invocation.source);
otel.tool_result_with_tags(
&tool_name,
&call_id_owned,
log_payload.as_ref(),
Duration::ZERO,
/*success*/ false,
&err.to_string(),
&tool_result_tags,
&extra_trace_fields,
);
if tool.is_builtin_control_tool() {
let mut analytics = ControlToolCallGuard::new(invocation);
analytics.finish(ControlToolCallStatus::Failed);
}
dispatch_trace.record_failed(&err);
notify_tool_finish_if_unclaimed(
invocation,
call_state.as_deref(),
ToolCallOutcome::Failed {
handler_executed: false,
},
)
.await;
err
};
let environments = match tool.live_environment_resolver() {
Some(resolve) => resolve(&invocation).await,
None => tool.environments_for_input(&invocation),
};
invocation.environments = match environments {
Ok(environments) => environments,
Err(err) => return Err(fail_before_execution(&invocation, err).await),
Err(err) => return Err(fail_before_execution(dispatch_trace, &invocation, err).await),
};
match apply_pre_tool_use(tool.as_ref(), &mut invocation).await {
Ok(()) => {}
@@ -906,7 +909,7 @@ impl ToolRegistry {
return Err(err);
}
Err(PreToolUseError::Failed(err)) => {
return Err(fail_before_execution(&invocation, err).await);
return Err(fail_before_execution(dispatch_trace, &invocation, err).await);
}
Err(PreToolUseError::Cancelled(err)) => return Err(err),
}
@@ -924,7 +927,7 @@ impl ToolRegistry {
if registered_tool.native_review == NativeReview::Required
&& let Err(err) = super::review_native_tool_call(&invocation).await
{
return Err(fail_before_execution(&invocation, err).await);
return Err(fail_before_execution(dispatch_trace, &invocation, err).await);
}
let invocation_guard = reservation.activate(&invocation);
@@ -16,7 +16,7 @@ use codex_rollout_trace::ToolDispatchRequester;
use codex_rollout_trace::ToolDispatchResult;
use codex_rollout_trace::ToolDispatchTraceContext;
/// Keeps registry early-return paths paired with trace end events.
/// Keeps registry results and abandoned dispatches paired with trace end events.
pub(crate) struct ToolDispatchTrace {
context: ToolDispatchTraceContext,
}
@@ -32,7 +32,7 @@ impl ToolDispatchTrace {
}
pub(crate) fn record_completed(
&self,
self,
invocation: &ToolInvocation,
call_id: &str,
payload: &ToolPayload,
@@ -54,7 +54,7 @@ impl ToolDispatchTrace {
self.context.record_completed(status, result_payload);
}
pub(crate) fn record_failed(&self, error: &FunctionCallError) {
pub(crate) fn record_failed(self, error: &FunctionCallError) {
self.context.record_failed(error);
}
}
@@ -410,6 +410,7 @@ fn cancelled_turn_terminates_unfinished_code_cell() -> anyhow::Result<()> {
fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let writer = Arc::new(create_started_writer(&temp)?);
let mut pending_dispatches = Vec::new();
writer.append(RawTraceEventPayload::ThreadStarted {
thread_id: "thread-child".to_string(),
agent_path: "/root/child".to_string(),
@@ -482,8 +483,11 @@ fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow:
);
if wait_id == "rejected" {
dispatch.record_failed("invalid wait arguments");
} else if wait_id == "confirmed" {
// The source item is still pending. No registry return is required for this edge.
continue;
}
// Replay these dispatches while they are still pending.
pending_dispatches.push(dispatch);
if wait_id == "confirmed" {
crate::CodeCellTraceContext::enabled(Arc::clone(&writer), thread_id, turn_id, "1")
.record_wait_target(wait_id);
}
@@ -368,6 +368,46 @@ fn protocol_wrapper_records_selected_events_as_raw_payloads() -> anyhow::Result<
Ok(())
}
#[test]
fn dropping_tool_dispatch_closes_the_execution_window() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let thread_id = ThreadId::new();
let trace =
ThreadTraceContext::start_root_in_root_for_test(temp.path(), minimal_metadata(thread_id))?;
trace.record_codex_turn_started("turn-1");
let dispatch = trace.start_tool_dispatch_trace(|| {
Some(ToolDispatchInvocation {
thread_id: thread_id.to_string(),
codex_turn_id: "turn-1".into(),
tool_call_id: "call".into(),
tool_name: "test_tool".into(),
tool_namespace: None,
requester: ToolDispatchRequester::Model {
model_visible_call_id: "call".into(),
},
payload: ToolDispatchPayload::Function {
arguments: "{}".into(),
},
})
});
trace.flush()?;
let bundle = single_bundle_dir(temp.path())?;
let before = replay_bundle(&bundle)?.tool_calls["call"].execution.clone();
assert_eq!(before.status, ExecutionStatus::Running);
assert_eq!(before.ended_at_unix_ms, None);
drop(dispatch);
trace.flush()?;
let rollout = replay_bundle(&bundle)?;
let call = &rollout.tool_calls["call"];
assert_eq!(call.execution.status, ExecutionStatus::Aborted);
assert_eq!(call.execution.started_at_unix_ms, before.started_at_unix_ms);
assert!(call.execution.ended_at_unix_ms >= Some(before.started_at_unix_ms));
assert!(call.execution.ended_seq > Some(before.started_seq));
assert_eq!(call.raw_result_payload_id, None);
Ok(())
}
#[test]
fn intercepted_exec_records_patch_without_a_terminal() -> anyhow::Result<()> {
let temp = TempDir::new()?;
+28 -21
View File
@@ -29,19 +29,13 @@ use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
use crate::writer::TraceWriter;
/// No-op capable trace handle for one resolved tool dispatch.
#[derive(Clone, Debug)]
/// Records one tool dispatch, closing it as aborted if dropped before a result.
#[derive(Debug)]
pub struct ToolDispatchTraceContext {
state: ToolDispatchTraceContextState,
active: Option<EnabledToolDispatchTraceContext>,
}
#[derive(Clone, Debug)]
enum ToolDispatchTraceContextState {
Disabled,
Enabled(EnabledToolDispatchTraceContext),
}
#[derive(Clone, Debug)]
#[derive(Debug)]
struct EnabledToolDispatchTraceContext {
writer: Arc<TraceWriter>,
thread_id: AgentThreadId,
@@ -127,9 +121,7 @@ enum DispatchedToolTraceResponse<'a> {
impl ToolDispatchTraceContext {
/// Builds a context that accepts trace calls and records nothing.
pub(crate) fn disabled() -> Self {
Self {
state: ToolDispatchTraceContextState::Disabled,
}
Self { active: None }
}
/// Returns whether caller-side result conversion would be recorded.
@@ -137,7 +129,7 @@ impl ToolDispatchTraceContext {
/// Core uses this to avoid formatting or cloning tool outputs when the
/// dispatch lifecycle is disabled.
pub fn is_enabled(&self) -> bool {
matches!(self.state, ToolDispatchTraceContextState::Enabled(_))
self.active.is_some()
}
/// Starts one dispatch-level lifecycle and returns the handle for its result.
@@ -152,13 +144,13 @@ impl ToolDispatchTraceContext {
};
record_started(&context, invocation);
Self {
state: ToolDispatchTraceContextState::Enabled(context),
active: Some(context),
}
}
/// Records the caller-facing successful or failed tool result.
pub fn record_completed(&self, status: ExecutionStatus, result: ToolDispatchResult) {
let ToolDispatchTraceContextState::Enabled(context) = &self.state else {
pub fn record_completed(mut self, status: ExecutionStatus, result: ToolDispatchResult) {
let Some(context) = self.active.take() else {
return;
};
let response = match &result {
@@ -169,16 +161,16 @@ impl ToolDispatchTraceContext {
DispatchedToolTraceResponse::CodeModeResponse { value }
}
};
append_tool_call_ended(context, status, &response);
append_tool_call_ended(&context, status, &response);
}
/// Records a dispatch failure before the tool produced a normal result payload.
pub fn record_failed(&self, error: impl Display) {
let ToolDispatchTraceContextState::Enabled(context) = &self.state else {
pub fn record_failed(mut self, error: impl Display) {
let Some(context) = self.active.take() else {
return;
};
append_tool_call_ended(
context,
&context,
ExecutionStatus::Failed,
&DispatchedToolTraceResponse::Error {
error: error.to_string(),
@@ -187,6 +179,21 @@ impl ToolDispatchTraceContext {
}
}
impl Drop for ToolDispatchTraceContext {
fn drop(&mut self) {
if let Some(context) = self.active.take() {
let _ = context.writer.enqueue_event(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawTraceEventPayload::ToolCallEnded {
tool_call_id: context.tool_call_id,
status: ExecutionStatus::Aborted,
result_payload: None,
},
);
}
}
}
fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDispatchInvocation) {
let tool_name = invocation.tool_name;
let tool_namespace = invocation.tool_namespace;