From 55e5158e1841bb4b0b392a462c24b4d9fc38d597 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 1 Sep 2026 00:20:48 +0000 Subject: [PATCH] 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 | 1 + codex-rs/code-mode-host/src/grpc/routing.rs | 6 ++++++ .../src/grpc/codex.code_mode.v1.proto | 2 +- codex-rs/code-mode-runtime/Cargo.toml | 1 + codex-rs/code-mode-runtime/src/service.rs | 11 ++++++++++ .../src/session_runtime/mod.rs | 6 ++++++ .../src/session_runtime/tests.rs | 1 + codex-rs/core/src/tools/code_mode/delegate.rs | 11 ++++++++++ codex-rs/exec-server/src/local_process.rs | 6 ++++-- codex-rs/exec-server/src/telemetry.rs | 20 ++++++++++++++++++- 10 files changed, 61 insertions(+), 4 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 07b46199c9..b59de3da7f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2643,6 +2643,7 @@ dependencies = [ "codex-protocol", "deno_core_icudata", "futures", + "opentelemetry", "pretty_assertions", "serde_json", "tokio", diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs index 4a7ddfeecd..3a16dc7792 100644 --- a/codex-rs/code-mode-host/src/grpc/routing.rs +++ b/codex-rs/code-mode-host/src/grpc/routing.rs @@ -126,6 +126,12 @@ impl GrpcSession { .collect::>(); (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()); } diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index df3b4503fc..7bd1756043 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/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; diff --git a/codex-rs/code-mode-runtime/Cargo.toml b/codex-rs/code-mode-runtime/Cargo.toml index 20c1e6cb4c..5a4d9e1d09 100644 --- a/codex-rs/code-mode-runtime/Cargo.toml +++ b/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"] } diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 65bb4095ee..8b8d71082a 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/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, diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index cc014230f9..d428112f2b 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/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 SessionRuntime { 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 { cell_id: CellId, inner: Arc>, + // 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 CellHost for RuntimeCellHost { @@ -255,6 +260,7 @@ impl CellHost for RuntimeCellHost { }, cancellation_token, ) + .with_context(self.execution_context.clone()) .await } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs index 1d90179bb3..9085e9c9bf 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs +++ b/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 { diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 0319d13d36..58904b992e 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/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, diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 27746255ce..b03cc55d19 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/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, diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index 13dbbb8a4f..725e018eeb 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/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()); }