From df302e2d8d48e86f0ad60ecc557e657423c0cffc Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Mon, 29 Jun 2026 22:10:27 -0700 Subject: [PATCH] telemetry(core): trace tool dispatch stages --- .../handlers/unified_exec/exec_command.rs | 83 +++++++++++++------ codex-rs/core/src/tools/orchestrator.rs | 41 ++++++++- codex-rs/core/src/tools/parallel.rs | 64 ++++++++------ codex-rs/core/src/tools/registry.rs | 25 +++++- 4 files changed, 156 insertions(+), 57 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index 546ad459e1..fde4803b97 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -1,5 +1,6 @@ use std::path::Path; use std::sync::Arc; +use tracing::Instrument; use crate::function_tool::FunctionCallError; use crate::maybe_emit_implicit_skill_invocation; @@ -98,7 +99,12 @@ impl ToolExecutor for ExecCommandHandler { } fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { - Box::pin(self.handle_call(invocation)) + let span = tracing::info_span!( + "unified_exec.handle_call", + call_id = %invocation.call_id, + ); + let future: codex_tools::ToolExecutorFuture<'_> = Box::pin(self.handle_call(invocation)); + Box::pin(future.instrument(span)) } } @@ -128,11 +134,16 @@ impl ExecCommandHandler { let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone()); - let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?; - let Some(turn_environment) = resolve_tool_environment( - &step_context.environments, - environment_args.environment_id.as_deref(), - )? + let environment_args: ExecCommandEnvironmentArgs = + tracing::info_span!("unified_exec.handle_call.parse_environment_args") + .in_scope(|| parse_arguments(&arguments))?; + let Some(turn_environment) = + tracing::info_span!("unified_exec.handle_call.resolve_environment").in_scope(|| { + resolve_tool_environment( + &step_context.environments, + environment_args.environment_id.as_deref(), + ) + })? else { return Err(FunctionCallError::RespondToModel( "unified exec is unavailable in this session".to_string(), @@ -177,18 +188,21 @@ impl ExecCommandHandler { ))); } }; - let mut args: ExecCommandArgs = match native_cwd.as_ref() { - Some(native_cwd) => { - // The base path only resolves paths nested in the permissions config types. - parse_arguments_with_base_path(&arguments, native_cwd)? - } - None => { - // Parsing without a base only skips relative-path resolution inside the - // permissions config. That is safe only for a truly unsandboxed attempt; - // sandboxed attempts fall through and return the conversion error below. - parse_arguments(&arguments)? - } - }; + let mut args: ExecCommandArgs = + tracing::info_span!("unified_exec.handle_call.parse_exec_args").in_scope(|| { + match native_cwd.as_ref() { + Some(native_cwd) => { + // The base path only resolves paths nested in the permissions config types. + parse_arguments_with_base_path(&arguments, native_cwd) + } + None => { + // Parsing without a base only skips relative-path resolution inside the + // permissions config. That is safe only for a truly unsandboxed attempt; + // sandboxed attempts fall through and return the conversion error below. + parse_arguments(&arguments) + } + } + })?; let hook_command = args.cmd.clone(); // TODO(anp) wire PathUri through implicit skills instead of skipping on foreign paths if let Some(native_cwd) = native_cwd.as_ref() { @@ -198,6 +212,9 @@ impl ExecCommandHandler { &hook_command, native_cwd, ) + .instrument(tracing::info_span!( + "unified_exec.handle_call.implicit_skill_detection" + )) .await; } let shell_mode = @@ -228,14 +245,22 @@ impl ExecCommandHandler { ))); } } - let process_id = manager.allocate_process_id().await; - let resolved_command = get_command( - &args, - shell, - &shell_mode, - turn.config.permissions.allow_login_shell, - ) - .map_err(FunctionCallError::RespondToModel)?; + let process_id = manager + .allocate_process_id() + .instrument(tracing::info_span!( + "unified_exec.handle_call.allocate_process_id" + )) + .await; + let resolved_command = tracing::info_span!("unified_exec.handle_call.resolve_command") + .in_scope(|| { + get_command( + &args, + shell, + &shell_mode, + turn.config.permissions.allow_login_shell, + ) + }) + .map_err(FunctionCallError::RespondToModel)?; let command = resolved_command.command; let shell_type = resolved_command.shell_type; let command_for_display = codex_shell_command::parse_command::shlex_join(&command); @@ -263,6 +288,9 @@ impl ExecCommandHandler { sandbox_permissions, additional_permissions, ) + .instrument(tracing::info_span!( + "unified_exec.handle_call.apply_granted_permissions" + )) .await; let additional_permissions_allowed = exec_permission_approvals_enabled || (session.features().enabled(Feature::RequestPermissionsTool) @@ -322,6 +350,9 @@ impl ExecCommandHandler { &context.call_id, "exec_command", ) + .instrument(tracing::info_span!( + "unified_exec.handle_call.intercept_apply_patch" + )) .await? { manager.release_process_id(process_id).await; diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index f82ad5bb70..fcf289636e 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -42,6 +42,7 @@ use codex_sandboxing::SandboxManager; use codex_sandboxing::SandboxType; use codex_utils_path_uri::PathUri; use std::time::Instant; +use tracing::Instrument; pub(crate) struct ToolOrchestrator { sandbox: SandboxManager, @@ -59,6 +60,16 @@ impl ToolOrchestrator { } } + #[tracing::instrument( + name = "tool_orchestrator.run_attempt", + level = "info", + skip_all, + fields( + sandbox = ?attempt.sandbox, + sandbox_requested = attempt.sandbox_requested, + managed_network_active, + ) + )] async fn run_attempt( tool: &mut T, req: &Rq, @@ -76,6 +87,9 @@ impl ToolOrchestrator { attempt.sandbox, tool.network_approval_spec(req, tool_ctx), ) + .instrument(tracing::info_span!( + "tool_orchestrator.begin_network_approval" + )) .await { Ok(network_approval) => network_approval, @@ -110,6 +124,7 @@ impl ToolOrchestrator { }; let run_result = tool .run(req, &attempt_with_network_approval, &attempt_tool_ctx) + .instrument(tracing::info_span!("tool_orchestrator.runtime_run")) .await; let Some(network_approval) = network_approval else { @@ -119,7 +134,11 @@ impl ToolOrchestrator { match network_approval.mode() { NetworkApprovalMode::Immediate => { let finalize_result = - finish_immediate_network_approval(&tool_ctx.session, network_approval).await; + finish_immediate_network_approval(&tool_ctx.session, network_approval) + .instrument(tracing::info_span!( + "tool_orchestrator.finish_immediate_network_approval" + )) + .await; if let Err(err) = finalize_result { return (Err(err), None); } @@ -129,7 +148,11 @@ impl ToolOrchestrator { let deferred = network_approval.into_deferred(); if run_result.is_err() { let finalize_result = - finish_deferred_network_approval(&tool_ctx.session, deferred).await; + finish_deferred_network_approval(&tool_ctx.session, deferred) + .instrument(tracing::info_span!( + "tool_orchestrator.finish_deferred_network_approval" + )) + .await; if let Err(err) = finalize_result { return (Err(err), None); } @@ -154,7 +177,13 @@ impl ToolOrchestrator { let otel = turn_ctx.session_telemetry.clone(); let otel_tn = flat_tool_name(&tool_ctx.tool_name).into_owned(); let otel_ci = &tool_ctx.call_id; - let strict_auto_review = tool_ctx.session.strict_auto_review_enabled_for_turn().await; + let strict_auto_review = tool_ctx + .session + .strict_auto_review_enabled_for_turn() + .instrument(tracing::info_span!( + "tool_orchestrator.resolve_auto_review_mode" + )) + .await; let use_guardian = routes_approval_to_guardian(turn_ctx) || strict_auto_review; // 1) Approval @@ -521,6 +550,12 @@ impl ToolOrchestrator { // PermissionRequest hooks take top precedence for answering approval // prompts. If no matching hook returns a decision, fall back to the // normal guardian or user approval path. + #[tracing::instrument( + name = "tool_orchestrator.request_approval", + level = "info", + skip_all, + fields(evaluate_permission_request_hooks) + )] async fn request_approval( tool: &mut T, req: &Rq, diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index e201f0ffce..de8a33dc36 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -129,31 +129,47 @@ impl ToolCallRuntime { let abort_dispatch_span = dispatch_span.clone(); let mut dispatch_handle: AbortOnDropHandle> = - AbortOnDropHandle::new(tokio::spawn(async move { - let _guard = if supports_parallel { - Either::Left(lock.read().await) - } else { - Either::Right(lock.write().await) - }; - // Admission through the parallel-execution gate marks the end - // of dispatch waiting and the start of handler execution. - if let Some(execution_started_at) = execution_started_at { - let _ = execution_started_at.set(Instant::now()); - } + AbortOnDropHandle::new(tokio::spawn( + async move { + let _guard = if supports_parallel { + Either::Left( + lock.read() + .instrument(tracing::info_span!( + "tool_dispatch.parallel_execution_lock", + mode = "shared", + )) + .await, + ) + } else { + Either::Right( + lock.write() + .instrument(tracing::info_span!( + "tool_dispatch.parallel_execution_lock", + mode = "exclusive", + )) + .await, + ) + }; + // Admission through the parallel-execution gate marks the end + // of dispatch waiting and the start of handler execution. + 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 - })); + router + .dispatch_tool_call_with_terminal_outcome( + session, + step_context, + invocation_cancellation_token, + tracker, + dispatch_call, + source, + dispatch_terminal_outcome_reached, + ) + .await + } + .instrument(dispatch_span), + )); async move { let _tool_call_timing_guard = tool_call_timing_guard; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 2220d4244a..f88501e587 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -34,6 +34,7 @@ use codex_tools::ToolSearchInfo; use codex_tools::ToolSpec; use futures::future::BoxFuture; use serde_json::Value; +use tracing::Instrument; use tracing::instrument; pub(crate) type ToolTelemetryTags = Vec<(&'static str, String)>; @@ -431,9 +432,18 @@ impl ToolRegistry { ]; { - let mut active = invocation.session.active_turn.lock().await; + let mut active = invocation + .session + .active_turn + .lock() + .instrument(tracing::info_span!("tool_dispatch.active_turn_lock")) + .await; if let Some(active_turn) = active.as_mut() { - let mut turn_state = active_turn.turn_state.lock().await; + let mut turn_state = active_turn + .turn_state + .lock() + .instrument(tracing::info_span!("tool_dispatch.turn_state_lock")) + .await; turn_state.tool_calls = turn_state.tool_calls.saturating_add(1); } } @@ -460,7 +470,10 @@ impl ToolRegistry { } }; - let telemetry_tags = tool.telemetry_tags(&invocation).await; + let telemetry_tags = tool + .telemetry_tags(&invocation) + .instrument(tracing::info_span!("tool_dispatch.telemetry_tags")) + .await; let mut tool_result_tags = Vec::with_capacity(base_tool_result_tags.len() + telemetry_tags.len()); let mut extra_trace_fields = Vec::new(); @@ -490,7 +503,9 @@ impl ToolRegistry { return Err(err); } - notify_tool_start(&invocation).await; + notify_tool_start(&invocation) + .instrument(tracing::info_span!("tool_dispatch.notify_tool_start")) + .await; if let Some(pre_tool_use_payload) = tool.pre_tool_use_payload(&invocation) { match run_pre_tool_use_hooks( @@ -500,6 +515,7 @@ impl ToolRegistry { &pre_tool_use_payload.tool_name, &pre_tool_use_payload.tool_input, ) + .instrument(tracing::info_span!("tool_dispatch.pre_tool_use_hooks")) .await { PreToolUseHookResult::Blocked(message) => { @@ -683,6 +699,7 @@ async fn notify_tool_finish_if_unclaimed( true } +#[instrument(name = "tool_dispatch.invoke_handler", level = "info", skip_all)] async fn handle_any_tool( tool: &dyn CoreToolRuntime, invocation: ToolInvocation,