From 84ed5744d95418fc976114cc2381ddab401b7a23 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 10 Sep 2026 22:24:38 +0000 Subject: [PATCH] Trace tool call receipt, result readiness, and code-mode dispatch (#44661) ## Why Tool results can be ready before the sampling loop collects them, and code-mode cells can outlive a turn. Separate trace milestones make these timings visible and link nested calls to the turn that dispatches them. ## What changed - Emit `codex.tool_call_received` and `codex.tool_result_ready` events for direct and code-mode calls, using identifiers and tool names without arguments or output. Record readiness before result encoding or collection, including ordinary errors and aborted responses, but excluding fatal errors. - Assign nested call IDs at broker receipt and carry them through dispatch. Trace dispatch and interruptions caused by cancellation, cell closure, or abandonment. - Add code-mode execute and wait handler spans with conversation, turn, call, and cell identifiers and `completed`, `failed`, or `interrupted` outcomes. GitOrigin-RevId: b098a29d8ee030c823665e4b1031a0d67cc09502 --- codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/stream_events_utils.rs | 7 ++ codex-rs/core/src/tools/call_trace.rs | 88 +++++++++++++++++++ codex-rs/core/src/tools/code_mode/delegate.rs | 52 +++++++++-- .../src/tools/code_mode/delegate_tests.rs | 1 + .../src/tools/code_mode/execute_handler.rs | 18 ++++ codex-rs/core/src/tools/code_mode/mod.rs | 49 ++++++++--- .../core/src/tools/code_mode/telemetry.rs | 61 +++++++++++++ .../core/src/tools/code_mode/wait_handler.rs | 18 ++++ codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/parallel.rs | 47 ++++++++-- .../src/tools/tool_dispatch_trace_tests.rs | 1 + 13 files changed, 323 insertions(+), 23 deletions(-) create mode 100644 codex-rs/core/src/tools/call_trace.rs diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 5302eae816..047eca6d84 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1577,6 +1577,7 @@ impl Session { ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( + thread_id, Arc::clone(&code_mode_session_provider), &config.code_mode, executed_tool_calls, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 40ab7b5436..1ed7f88979 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6718,6 +6718,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( + thread_id, Arc::new(codex_code_mode::DisabledCodeModeSessionProvider), &config.code_mode, executed_tool_calls, @@ -8891,6 +8892,7 @@ where ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( + thread_id, Arc::new(codex_code_mode::DisabledCodeModeSessionProvider), &config.code_mode, executed_tool_calls, diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 285e85fa14..b3c0a0a74f 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -14,6 +14,7 @@ use crate::parse_turn_item; use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; +use crate::tools::call_trace; use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolRouter; use crate::tools::router::tool_log_payload; @@ -307,6 +308,12 @@ pub(crate) async fn handle_output_item_done( match ToolRouter::build_tool_call(item.clone()) { // The model emitted a tool call; log it, persist the item immediately, and queue the tool execution. Ok(Some(call)) => { + call_trace::received( + ctx.sess.thread_id, + &call.tool_name, + &call.call_id, + call_trace::Receipt::ModelTurn(&ctx.step_context.turn.sub_id), + ); ctx.sess .input_queue .accept_mailbox_delivery_for_current_turn( diff --git a/codex-rs/core/src/tools/call_trace.rs b/codex-rs/core/src/tools/call_trace.rs new file mode 100644 index 0000000000..b6daff424b --- /dev/null +++ b/codex-rs/core/src/tools/call_trace.rs @@ -0,0 +1,88 @@ +//! Trace milestones for every direct or code-mode tool call handled by core. +//! These events contain call identifiers and names, never arguments or output. +//! Code-mode calls reach the broker before the dispatching turn is known. + +use codex_protocol::DEFAULT_FUNCTION_NAMESPACE; +use codex_protocol::ThreadId; +use codex_tools::ToolName; + +#[derive(Clone, Copy)] +pub(crate) enum Source { + Direct, + CodeMode, +} + +impl Source { + fn as_str(self) -> &'static str { + match self { + Self::Direct => "direct", + Self::CodeMode => "code_mode", + } + } +} + +pub(crate) enum Receipt<'a> { + ModelTurn(&'a str), + CodeModeBroker { + cell_id: Option<&'a str>, + runtime_tool_call_id: Option<&'a str>, + }, +} + +pub(crate) fn received( + thread_id: ThreadId, + tool_name: &ToolName, + call_id: &str, + receipt: Receipt<'_>, +) { + let (turn_id, source, cell_id, runtime_tool_call_id) = match receipt { + Receipt::ModelTurn(turn_id) => (Some(turn_id), Source::Direct, None, None), + Receipt::CodeModeBroker { + cell_id, + runtime_tool_call_id, + } => (None, Source::CodeMode, cell_id, runtime_tool_call_id), + }; + tracing::event!( + name: "codex.tool_call_received", + target: "codex_otel.trace_safe", + tracing::Level::INFO, + event.name = "codex.tool_call_received", + conversation.id = %thread_id, + turn_id, + call_id, + tool_name = tool_name.name.as_str(), + tool_namespace = namespace(tool_name), + tool_source = source.as_str(), + cell.id = cell_id, + runtime_tool_call_id, + ); +} + +pub(crate) fn result_ready( + thread_id: ThreadId, + turn_id: &str, + tool_name: &ToolName, + call_id: &str, + source: Source, +) { + tracing::event!( + name: "codex.tool_result_ready", + target: "codex_otel.trace_safe", + tracing::Level::INFO, + event.name = "codex.tool_result_ready", + conversation.id = %thread_id, + turn_id, + call_id, + tool_name = tool_name.name.as_str(), + tool_namespace = namespace(tool_name), + tool_source = source.as_str(), + ); +} + +fn namespace(tool_name: &ToolName) -> &str { + tool_name + .namespace + .as_deref() + .filter(|namespace| !namespace.is_empty()) + .unwrap_or(DEFAULT_FUNCTION_NAMESPACE) +} diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 331622fd1c..ac8d0505ea 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -8,6 +8,7 @@ use codex_code_mode::CodeModeSessionDelegate; use codex_code_mode::NotificationFuture; use codex_code_mode::ToolInvocationFuture; use codex_protocol::ResponseItemId; +use codex_protocol::ThreadId; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; @@ -19,12 +20,17 @@ use tracing::Instrument; use super::ExecContext; use super::PUBLIC_TOOL_NAME; use super::submit_nested_tool; +use super::telemetry::DispatchInterruption; +use super::telemetry::NestedToolDispatchTrace; +use super::telemetry::trace_id; use crate::session::step_context::StepContext; use crate::tools::ExecutedToolCalls; +use crate::tools::call_trace; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::parallel::ToolCallRuntime; pub(super) struct CodeModeDispatchBroker { + thread_id: ThreadId, dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, dispatch_gates: Arc>>, @@ -38,9 +44,10 @@ struct CellDispatchGate { } impl CodeModeDispatchBroker { - pub(super) fn new(executed_tool_calls: ExecutedToolCalls) -> Self { + pub(super) fn new(thread_id: ThreadId, executed_tool_calls: ExecutedToolCalls) -> Self { let (dispatch_tx, dispatch_rx) = async_channel::unbounded(); Self { + thread_id, dispatch_tx, dispatch_rx, dispatch_gates: Arc::new(Mutex::new(HashMap::new())), @@ -141,9 +148,9 @@ impl CodeModeDispatchBroker { } DispatchMessage::InvokeTool { invocation, + mut dispatch_trace, cancellation_token, response_tx, - span, } => { let cell_id = invocation.cell_id.clone(); if !wait_until_cell_ready_for_dispatch( @@ -153,11 +160,18 @@ impl CodeModeDispatchBroker { ) .await { + let outcome = if cancellation_token.is_cancelled() { + DispatchInterruption::Cancelled + } else { + DispatchInterruption::CellClosed + }; + dispatch_trace.interruption = Some(outcome); remove_dispatch_gate(&dispatch_gates, &cell_id); continue; } let host = Arc::clone(&host); let dispatch_gates = Arc::clone(&dispatch_gates); + let span = dispatch_trace.span.clone(); tokio::spawn(async move { let invocation = { let dispatch_gate = track_completeness.then(|| { @@ -169,11 +183,22 @@ impl CodeModeDispatchBroker { cancellation_token.is_cancelled() || !gates.contains_key(&cell_id) }) { + let outcome = if cancellation_token.is_cancelled() { + DispatchInterruption::Cancelled + } else { + DispatchInterruption::CellClosed + }; + dispatch_trace.interruption = Some(outcome); return; } // Submission and cell closure share this gate. span.in_scope(|| { - host.submit_tool(invocation, cancellation_token.clone()) + dispatch_trace.interruption = None; + host.submit_tool( + invocation, + dispatch_trace.call_id.clone(), + cancellation_token.clone(), + ) }) .instrument(span) }; @@ -254,6 +279,7 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { level = "info", skip_all, fields( + conversation.id = %self.thread_id, cell.id = %invocation.cell_id, runtime_tool_call_id = invocation.runtime_tool_call_id.as_str(), tool_name = invocation.tool_name.name.as_str(), @@ -266,16 +292,30 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { cancellation_token: CancellationToken, ) -> ToolInvocationFuture<'a> { Box::pin(async move { + let call_id = format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()); + call_trace::received( + self.thread_id, + &invocation.tool_name, + &call_id, + call_trace::Receipt::CodeModeBroker { + cell_id: trace_id(invocation.cell_id.as_str()), + runtime_tool_call_id: trace_id(&invocation.runtime_tool_call_id), + }, + ); + let mut dispatch_trace = + Box::new(NestedToolDispatchTrace::new(self.thread_id, call_id)); if cancellation_token.is_cancelled() { + dispatch_trace.interruption = Some(DispatchInterruption::Cancelled); return Err("code mode nested tool call cancelled".to_string()); } let (response_tx, response_rx) = oneshot::channel(); + // Only the worker can tell whether dispatch beats cancellation once the call is queued. self.dispatch_tx .send(DispatchMessage::InvokeTool { invocation, + dispatch_trace, cancellation_token: cancellation_token.clone(), response_tx, - span: tracing::Span::current(), }) .await .map_err(|_| "code mode nested tool dispatcher is unavailable".to_string())?; @@ -329,9 +369,9 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { enum DispatchMessage { InvokeTool { invocation: CodeModeNestedToolCall, + dispatch_trace: Box, cancellation_token: CancellationToken, response_tx: oneshot::Sender>, - span: tracing::Span, }, Notify { call_id: String, @@ -363,12 +403,14 @@ impl CoreTurnHost { fn submit_tool( &self, invocation: CodeModeNestedToolCall, + call_id: String, cancellation_token: CancellationToken, ) -> impl std::future::Future> + Send + 'static { let invocation = submit_nested_tool( self.exec.clone(), self.tool_runtime.clone(), invocation, + call_id, cancellation_token, ) .map_err(|error| error.to_string()); diff --git a/codex-rs/core/src/tools/code_mode/delegate_tests.rs b/codex-rs/core/src/tools/code_mode/delegate_tests.rs index de84f82761..70340147c7 100644 --- a/codex-rs/core/src/tools/code_mode/delegate_tests.rs +++ b/codex-rs/core/src/tools/code_mode/delegate_tests.rs @@ -116,6 +116,7 @@ async fn detached_code_mode_callback_keeps_thread_id_on_dispatch_span() -> anyho )); let step = StepContext::for_test(Arc::clone(&turn)).with_tool_router_for_test(router); let broker = Arc::new(CodeModeDispatchBroker::new( + session.thread_id, /*executed_tool_calls*/ Default::default(), )); let cell_id = CellId::new("audit-cell".to_string()); diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index ac0383a9d2..f7ceb3c7c7 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -14,6 +14,7 @@ use super::PUBLIC_TOOL_NAME; use super::handle_runtime_response; use super::is_exec_tool_name; use super::telemetry::CodeModeToolCallGuard; +use super::telemetry::trace_id; type CodeModeNestedTool = (Arc, Option>); @@ -80,6 +81,7 @@ impl CodeModeExecuteHandler { .await .map_err(FunctionCallError::RespondToModel)?; let cell_id = started_cell.cell_id.clone(); + tracing::Span::current().record("cell.id", trace_id(cell_id.as_str())); telemetry.cell_id = Some(cell_id.to_string()); exec.session .services @@ -170,10 +172,25 @@ impl ToolExecutor for CodeModeExecuteHandler { } impl CodeModeExecuteHandler { + // Default to interrupted if this future is dropped; telemetry::CodeModeToolCallGuard::finish + // overwrites this handler's captured span on explicit success or failure. + #[tracing::instrument( + name = "code_mode.handler.execute", + level = "info", + skip_all, + fields( + conversation.id = %invocation.session.thread_id, + turn_id = invocation.turn.sub_id.as_str(), + call_id = trace_id(&invocation.call_id), + cell.id = tracing::field::Empty, + outcome = "interrupted", + ) + )] async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { + let handler_span = tracing::Span::current(); let originating_item_id = invocation.originating_item_id().await; let ToolInvocation { session, @@ -192,6 +209,7 @@ impl CodeModeExecuteHandler { turn.turn_metadata_state.clone(), call_id.clone(), PUBLIC_TOOL_NAME, + handler_span, ); let result = match payload { ToolPayload::Custom { input } if is_exec_tool_name(&tool_name) => self diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 7db0e8e82b..c9355ed1b8 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -17,6 +17,7 @@ use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::RuntimeResponse; +use codex_protocol::ThreadId; use codex_protocol::models::FunctionCallOutputContentItem; use futures::future::join_all; use serde_json::Value as JsonValue; @@ -31,6 +32,7 @@ use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::ExecutedToolCalls; +use crate::tools::call_trace; use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; @@ -78,11 +80,12 @@ pub(crate) struct CodeModeService { impl CodeModeService { pub(crate) fn new( + thread_id: ThreadId, session_provider: Arc, config: &CodeModeConfig, executed_tool_calls: ExecutedToolCalls, ) -> Self { - let dispatch_broker = Arc::new(CodeModeDispatchBroker::new(executed_tool_calls)); + let dispatch_broker = Arc::new(CodeModeDispatchBroker::new(thread_id, executed_tool_calls)); let availability = session_provider.availability(); Self { session: OnceCell::new(), @@ -343,6 +346,7 @@ fn submit_nested_tool( exec: ExecContext, tool_runtime: ToolCallRuntime, invocation: CodeModeNestedToolCall, + call_id: String, cancellation_token: CancellationToken, ) -> Result< impl std::future::Future> + Send + 'static, @@ -355,20 +359,43 @@ fn submit_nested_tool( tool_kind, input, } = invocation; - if is_exec_tool_name(&tool_name) { - return Err(FunctionCallError::RespondToModel(format!( - "{PUBLIC_TOOL_NAME} cannot invoke itself" - ))); - } - - let payload = match build_nested_tool_payload(tool_kind, &tool_name, input) { + let thread_id = exec.session.thread_id; + let turn_id = exec.turn.sub_id.clone(); + let tool_name = tool_name.with_default_namespace(); + // A cell can outlive a turn; the broker records arrival before a dispatching turn is known. + tracing::event!( + name: "codex.code_mode.nested_tool_dispatched", + target: "codex_otel.trace_safe", + tracing::Level::INFO, + event.name = "codex.code_mode.nested_tool_dispatched", + conversation.id = %thread_id, + turn_id = turn_id.as_str(), + cell.id = telemetry::trace_id(cell_id.as_str()), + runtime_tool_call_id = telemetry::trace_id(&runtime_tool_call_id), + call_id = call_id.as_str(), + ); + let payload = if is_exec_tool_name(&tool_name) { + Err(format!("{PUBLIC_TOOL_NAME} cannot invoke itself")) + } else { + build_nested_tool_payload(tool_kind, &tool_name, input) + }; + let payload = match payload { Ok(payload) => payload, - Err(error) => return Err(FunctionCallError::RespondToModel(error)), + Err(error) => { + call_trace::result_ready( + thread_id, + &turn_id, + &tool_name, + &call_id, + call_trace::Source::CodeMode, + ); + return Err(FunctionCallError::RespondToModel(error)); + } }; let call = ToolCall { - tool_name: tool_name.with_default_namespace(), - call_id: format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()), + tool_name, + call_id, payload, encrypted_function_args: None, }; diff --git a/codex-rs/core/src/tools/code_mode/telemetry.rs b/codex-rs/core/src/tools/code_mode/telemetry.rs index a4d122980b..ad80a0b799 100644 --- a/codex-rs/core/src/tools/code_mode/telemetry.rs +++ b/codex-rs/core/src/tools/code_mode/telemetry.rs @@ -1,9 +1,13 @@ +//! Code-mode tool analytics and tracing for handlers and queued nested dispatch. + use codex_analytics::AnalyticsEventsClient; use codex_analytics::CodeModeToolCallFact; use codex_analytics::CodeModeToolCallStatus; use codex_analytics::TurnAnalyticsMetadata; +use codex_protocol::ThreadId; use std::sync::Arc; use std::time::Duration; +use tracing::Span; pub(super) struct CodeModeToolCallGuard { analytics: AnalyticsEventsClient, @@ -15,6 +19,7 @@ pub(super) struct CodeModeToolCallGuard { tool_name: &'static str, started_at_ms: u64, status: CodeModeToolCallStatus, + handler_span: Span, } impl CodeModeToolCallGuard { @@ -25,6 +30,7 @@ impl CodeModeToolCallGuard { turn_metadata: Arc, call_id: String, tool_name: &'static str, + handler_span: Span, ) -> Self { Self { analytics, @@ -36,15 +42,19 @@ impl CodeModeToolCallGuard { tool_name, started_at_ms: codex_analytics::now_unix_millis(), status: CodeModeToolCallStatus::Interrupted, + handler_span, } } pub(super) fn finish(&mut self, success: bool) { + let mut outcome = "failed"; self.status = if success { + outcome = "completed"; CodeModeToolCallStatus::Completed } else { CodeModeToolCallStatus::Failed }; + self.handler_span.record("outcome", outcome); } pub(super) fn record_code_mode_host_duration(&self, duration: Duration) { @@ -85,3 +95,54 @@ impl Drop for CodeModeToolCallGuard { }); } } + +pub(super) struct NestedToolDispatchTrace { + thread_id: ThreadId, + pub(super) call_id: String, + pub(super) interruption: Option, + pub(super) span: Span, +} + +impl NestedToolDispatchTrace { + pub(super) fn new(thread_id: ThreadId, call_id: String) -> Self { + Self { + thread_id, + call_id, + interruption: Some(DispatchInterruption::Abandoned), + span: Span::current(), + } + } +} + +impl Drop for NestedToolDispatchTrace { + fn drop(&mut self) { + let Some(outcome) = self.interruption.take() else { + return; + }; + let outcome = match outcome { + DispatchInterruption::Cancelled => "cancelled", + DispatchInterruption::CellClosed => "cell_closed", + DispatchInterruption::Abandoned => "abandoned", + }; + tracing::event!( + name: "codex.code_mode.nested_tool_dispatch_interrupted", + target: "codex_otel.trace_safe", + parent: &self.span, + tracing::Level::INFO, + event.name = "codex.code_mode.nested_tool_dispatch_interrupted", + conversation.id = %self.thread_id, + call_id = self.call_id.as_str(), + outcome, + ); + } +} + +pub(super) enum DispatchInterruption { + Cancelled, + CellClosed, + Abandoned, +} + +pub(super) fn trace_id(id: &str) -> Option<&str> { + (!id.is_empty() && id.len() <= 256).then_some(id) +} diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 1e9e18f6d8..d4c07eaad4 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -17,6 +17,7 @@ use super::ExecContext; use super::WAIT_TOOL_NAME; use super::handle_runtime_response; use super::telemetry::CodeModeToolCallGuard; +use super::telemetry::trace_id; use super::wait_spec::create_wait_tool; pub struct CodeModeWaitHandler; @@ -63,10 +64,25 @@ impl ToolExecutor for CodeModeWaitHandler { } impl CodeModeWaitHandler { + // Default to interrupted if this future is dropped; telemetry::CodeModeToolCallGuard::finish + // overwrites this handler's captured span on explicit success or failure, including early errors. + #[tracing::instrument( + name = "code_mode.handler.wait", + level = "info", + skip_all, + fields( + conversation.id = %invocation.session.thread_id, + turn_id = invocation.turn.sub_id.as_str(), + call_id = trace_id(&invocation.call_id), + cell.id = tracing::field::Empty, + outcome = "interrupted", + ) + )] async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { + let handler_span = tracing::Span::current(); let ToolInvocation { session, turn, @@ -84,6 +100,7 @@ impl CodeModeWaitHandler { turn.turn_metadata_state.clone(), call_id.clone(), WAIT_TOOL_NAME, + handler_span, ); let result = match payload { ToolPayload::Function { arguments } @@ -123,6 +140,7 @@ impl CodeModeWaitHandler { | codex_code_mode::RuntimeResponse::Terminated { cell_id, .. } | codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id, }; + tracing::Span::current().record("cell.id", trace_id(runtime_cell_id.as_str())); telemetry.cell_id = Some(runtime_cell_id.to_string()); exec.session .services diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 7a65bc95a3..4098be8f70 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -1,4 +1,5 @@ mod approvals; +pub(crate) mod call_trace; pub(crate) mod code_mode; pub(crate) mod context; mod control_tool_analytics; diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 4b5f962bf6..e43aa5ff4d 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -17,6 +17,7 @@ use tracing::trace_span; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::step_context::StepContext; +use crate::tools::call_trace; use crate::tools::context::AbortedToolOutput; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::context::ToolPayload; @@ -88,7 +89,6 @@ impl ToolCallRuntime { )), } } - .in_current_span() } #[instrument(level = "trace", skip_all)] @@ -125,6 +125,15 @@ impl ToolCallRuntime { 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 thread_id = session.thread_id; + let trace_source = match &source { + ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { + call_trace::Source::Direct + } + ToolCallSource::CodeMode { .. } => call_trace::Source::CodeMode, + }; + let dispatch_tool_name = call.tool_name.clone(); + let dispatch_call_id = call.call_id.clone(); // Code-mode callbacks can resume outside the turn's local span ancestry. let dispatch_span = trace_span!( @@ -137,15 +146,15 @@ impl ToolCallRuntime { ); let abort_dispatch_span = dispatch_span.clone(); - let mut dispatch_handle: AbortOnDropHandle> = - AbortOnDropHandle::new(tokio::spawn(async move { + let mut dispatch_handle = AbortOnDropHandle::new(tokio::spawn( + async move { if let Some(tool_runtime) = tool_runtime && let Some(readiness) = tool_runtime.wait_until_ready(&session) { readiness.await; } - let _guard = if supports_parallel { + let guard = if supports_parallel { Either::Left(lock.read().await) } else { Either::Right(lock.write().await) @@ -156,7 +165,7 @@ impl ToolCallRuntime { let _ = execution_started_at.set(Instant::now()); } - router + let result = router .dispatch_tool_call_with_terminal_outcome( session, step_context, @@ -167,8 +176,25 @@ impl ToolCallRuntime { dispatch_terminal_outcome_reached, ) .instrument(dispatch_span.clone()) - .await - })); + .await; + drop(guard); + // The sampling loop collects results in order only after its stream ends. + // Record readiness here, before either caller encodes or collects the result. + // A fatal error still propagates to the caller instead of producing a tool + // result; unlike a normal tool failure, it has no readiness event. + if !matches!(&result, Err(FunctionCallError::Fatal(_))) { + call_trace::result_ready( + thread_id, + &turn.sub_id, + &dispatch_tool_name, + &dispatch_call_id, + trace_source, + ); + } + result + } + .in_current_span(), + )); async move { let _tool_call_timing_guard = tool_call_timing_guard; @@ -187,6 +213,13 @@ impl ToolCallRuntime { Err(err) => return Err(Self::tool_task_join_error(err)), } let response = Self::aborted_response(&call, secs); + call_trace::result_ready( + thread_id, + &abort_turn.sub_id, + &call.tool_name, + &call.call_id, + trace_source, + ); notify_tool_aborted( abort_session.as_ref(), abort_turn.as_ref(), diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index c69a3028cb..0755fbbb3f 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -282,6 +282,7 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu let temp = TempDir::new()?; let (mut session, turn) = make_session_and_context().await; session.services.code_mode_service = CodeModeService::new( + session.thread_id, Arc::new(MissingCellCodeModeSessionProvider), &turn.config.code_mode, session.services.executed_tool_calls.clone(),