diff --git a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs index 35459f29e6..8dbb8d37c0 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs @@ -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::>>()?; + 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(()) } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 59ff5bbda9..e9d6e86ee1 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -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); diff --git a/codex-rs/core/src/tools/tool_dispatch_trace.rs b/codex-rs/core/src/tools/tool_dispatch_trace.rs index deb5fa02bc..a16e70055d 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace.rs @@ -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); } } diff --git a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs index c80f0511a8..ac6ff63dd4 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs @@ -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); } diff --git a/codex-rs/rollout-trace/src/thread_tests.rs b/codex-rs/rollout-trace/src/thread_tests.rs index 285cfc3216..00d38746f6 100644 --- a/codex-rs/rollout-trace/src/thread_tests.rs +++ b/codex-rs/rollout-trace/src/thread_tests.rs @@ -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()?; diff --git a/codex-rs/rollout-trace/src/tool_dispatch.rs b/codex-rs/rollout-trace/src/tool_dispatch.rs index 34d57308e6..2fbf87205f 100644 --- a/codex-rs/rollout-trace/src/tool_dispatch.rs +++ b/codex-rs/rollout-trace/src/tool_dispatch.rs @@ -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, } -#[derive(Clone, Debug)] -enum ToolDispatchTraceContextState { - Disabled, - Enabled(EnabledToolDispatchTraceContext), -} - -#[derive(Clone, Debug)] +#[derive(Debug)] struct EnabledToolDispatchTraceContext { writer: Arc, 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;