Improve tracing for nested tool calls and exec processes (#41950)

## Why

Code-mode callbacks can outlive their initial request and run in separate tasks,
so their trace context must be preserved explicitly.

## What changed

- Preserve the execution context for code-mode callbacks and add spans for
  nested tool invocations in the runtime and dispatch broker.
- Propagate each tool invocation span as the parent of its streamed gRPC
  callback, falling back to the outer execution context when needed.
- Add a `codex.exec_server.process` span for each process lifetime, including
  its process ID and completion result, without retaining the request span.

GitOrigin-RevId: a6059e34d895416f5517e51dad6ca0078355adca

codex-rs/Cargo.lock

@@ -2643,6 +2643,7 @@ dependencies = [
  "codex-protocol",
  "deno_core_icudata",
  "futures",
+ "opentelemetry",
  "pretty_assertions",
  "serde_json",
  "tokio",

codex-rs/code-mode-host/src/grpc/routing.rs

@@ -126,6 +126,12 @@ impl GrpcSession {
                 .collect::<Vec<_>>();
             (sequence, execution.traceparent.clone(), subscriptions)
         };
+        // The saved traceparent belongs to the outer execution. Prefer the runtime's
+        // per-tool span so the callback is nested under its invocation, with the
+        // execution context as a fallback when no current span context is available.
+        let traceparent = codex_otel::current_span_w3c_trace_context()
+            .and_then(|trace| trace.traceparent)
+            .or(traceparent);
         if subscriptions.is_empty() {
             return Err("no code-mode tool subscription matches the requested tool".to_string());
         }

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -83,7 +83,7 @@ message ToolCall {
   // Starts at one and increases independently for each execution.
   uint64 sequence = 9;
 
-  // The host execution span's W3C parent context for this streamed callback.
+  // The host tool invocation span's W3C parent context for this streamed callback.
   // gRPC metadata is fixed when the subscription stream opens, so each call
   // carries its own context in the message.
   optional string traceparent = 10;

codex-rs/code-mode-runtime/Cargo.toml

@@ -18,6 +18,7 @@ codex-code-mode-protocol = { workspace = true }
 codex-protocol = { workspace = true }
 deno_core_icudata = { workspace = true }
 futures = { workspace = true }
+opentelemetry = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
 tokio-util = { workspace = true, features = ["rt"] }

codex-rs/code-mode-runtime/src/service.rs

@@ -242,6 +242,17 @@ struct ProtocolDelegate {
 }
 
 impl runtime::SessionRuntimeDelegate for ProtocolDelegate {
+    #[tracing::instrument(
+        name = "code_mode.runtime.invoke_tool",
+        level = "info",
+        skip_all,
+        fields(
+            cell.id = %invocation.cell_id,
+            runtime_tool_call_id = invocation.runtime_tool_call_id.as_str(),
+            tool_name = invocation.tool_name.name.as_str(),
+            tool_namespace = invocation.tool_name.namespace.as_deref(),
+        )
+    )]
     async fn invoke_tool(
         &self,
         invocation: runtime::NestedToolCall,

codex-rs/code-mode-runtime/src/session_runtime/mod.rs

@@ -7,6 +7,7 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 
+use opentelemetry::context::FutureExt;
 use serde_json::Value as JsonValue;
 use tokio::sync::Mutex;
 use tokio_util::sync::CancellationToken;
@@ -163,6 +164,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
         let host = Arc::new(RuntimeCellHost {
             cell_id: cell_id.clone(),
             inner: Arc::clone(&self.inner),
+            execution_context: opentelemetry::Context::current(),
         });
         let mut cells = self.inner.cells.lock().await;
         if self.inner.shutdown_token.is_cancelled() {
@@ -235,6 +237,9 @@ impl PendingEvent {
 struct RuntimeCellHost<D: SessionRuntimeDelegate> {
     cell_id: CellId,
     inner: Arc<Inner<D>>,
+    // Callbacks outlive the initial request and run in separate tasks. Preserve
+    // their trace parent without retaining the request's tracing span.
+    execution_context: opentelemetry::Context,
 }
 
 impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
@@ -255,6 +260,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
                 },
                 cancellation_token,
             )
+            .with_context(self.execution_context.clone())
             .await
     }
 

codex-rs/code-mode-runtime/src/session_runtime/tests.rs

@@ -102,6 +102,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
     let host = RuntimeCellHost {
         cell_id: CellId::new("terminating-writer"),
         inner: Arc::clone(&runtime.inner),
+        execution_context: opentelemetry::Context::new(),
     };
     let completion = CellEvent::Completed {
         content_items: vec![OutputItem::Text {

codex-rs/core/src/tools/code_mode/delegate.rs

@@ -255,6 +255,17 @@ async fn wait_until_cell_ready_for_dispatch(
 }
 
 impl CodeModeSessionDelegate for CodeModeDispatchBroker {
+    #[tracing::instrument(
+        name = "code_mode.broker.invoke_tool",
+        level = "info",
+        skip_all,
+        fields(
+            cell.id = %invocation.cell_id,
+            runtime_tool_call_id = invocation.runtime_tool_call_id.as_str(),
+            tool_name = invocation.tool_name.name.as_str(),
+            tool_namespace = invocation.tool_name.namespace.as_deref(),
+        )
+    )]
     fn invoke_tool<'a>(
         &'a self,
         invocation: CodeModeNestedToolCall,

codex-rs/exec-server/src/local_process.rs

@@ -406,6 +406,7 @@ impl LocalProcess {
                 return Err(internal_error(err.to_string()));
             }
         };
+        let metrics = self.inner.telemetry.process_started(&process_id);
 
         let output_notify = Arc::new(Notify::new());
         let (wake_tx, _wake_rx) = watch::channel(0);
@@ -421,6 +422,7 @@ impl LocalProcess {
             ) {
                 drop(process_map);
                 spawned.session.terminate();
+                metrics.finish("terminated");
                 return Err(invalid_request(format!(
                     "process {process_id} start was cancelled"
                 )));
@@ -443,7 +445,7 @@ impl LocalProcess {
                     output_notify: Arc::clone(&output_notify),
                     open_streams: 2,
                     closed: false,
-                    metrics: Some(self.inner.telemetry.process_started()),
+                    metrics: Some(metrics),
                     termination_requested: false,
                     sandbox: prepared.sandbox,
                     sandbox_denied: false,
@@ -1893,7 +1895,7 @@ mod tests {
                 output_notify: Arc::clone(&output_notify),
                 open_streams: 2,
                 closed: false,
-                metrics: Some(backend.inner.telemetry.process_started()),
+                metrics: Some(backend.inner.telemetry.process_started(&process_id)),
                 termination_requested: false,
                 sandbox: SandboxType::None,
                 sandbox_denied: false,

codex-rs/exec-server/src/telemetry.rs

@@ -99,6 +99,7 @@ pub(crate) struct ConnectionMetricGuard {
 
 pub(crate) struct ProcessMetricGuard {
     telemetry: ExecServerTelemetry,
+    span: tracing::Span,
     started_at: Instant,
     result: &'static str,
 }
@@ -202,12 +203,28 @@ impl ExecServerTelemetry {
         });
     }
 
-    pub(crate) fn process_started(&self) -> ProcessMetricGuard {
+    pub(crate) fn process_started(&self, process_id: &str) -> ProcessMetricGuard {
         self.with_inner(|inner| {
             inner.adjust_process_count(/*delta*/ 1);
         });
+        let parent = codex_otel::current_span_w3c_trace_context();
+        // `parent:` accepts a local tracing span/ID, not a W3C context. A local
+        // parent would keep the request span alive until process exit and delay
+        // its export. Use `parent: None`, then set the W3C parent below to link
+        // the spans without retaining the request span.
+        let span = tracing::info_span!(
+            parent: None,
+            "codex.exec_server.process",
+            otel.kind = "internal",
+            process.id = process_id,
+            result = tracing::field::Empty,
+        );
+        if let Some(parent) = parent {
+            codex_otel::set_parent_from_w3c_trace_context(&span, &parent);
+        }
         ProcessMetricGuard {
             telemetry: self.clone(),
+            span,
             started_at: Instant::now(),
             result: "unknown",
         }
@@ -275,6 +292,7 @@ impl ProcessMetricGuard {
 
 impl Drop for ProcessMetricGuard {
     fn drop(&mut self) {
+        self.span.record("result", self.result);
         self.telemetry
             .process_finished(self.result, self.started_at.elapsed());
     }