Propagate trace context through gRPC code mode (#41017)

## What changed

- Inject W3C `traceparent` metadata into code-mode session and execution requests.
- Carry each execution's trace context on streamed tool callbacks so callback and nested-tool spans remain connected across the gRPC boundary.
- Preserve the relevant spans when execution, callback, and dispatch work moves into spawned tasks.

GitOrigin-RevId: 6482963cb9b2e2f0ea27f1345bba208b70a80d0b

codex-rs/code-mode/src/grpc_session/callbacks.rs

@@ -6,8 +6,10 @@ use std::time::Duration;
 
 use codex_code_mode_protocol::CellId;
 use codex_code_mode_protocol::grpc;
+use codex_protocol::protocol::W3cTraceContext;
 use futures::FutureExt;
 use tokio_util::sync::CancellationToken;
+use tracing::Instrument;
 use tracing::warn;
 
 use super::SessionInner;
@@ -112,6 +114,23 @@ impl SessionInner {
             .lock()
             .unwrap_or_else(PoisonError::into_inner)
             .admit_invocation(&call)?;
+        let callback_span = tracing::info_span!(
+            "code_mode.grpc.callback",
+            otel.name = "code_mode.grpc.callback",
+            session.id = %call.session_id,
+            execution.id = %call.execution_id,
+            cell.id = %call.cell_id,
+            invocation.id = %call.invocation_id,
+        );
+        if let Some(traceparent) = call.traceparent.as_ref() {
+            codex_otel::set_parent_from_w3c_trace_context(
+                &callback_span,
+                &W3cTraceContext {
+                    traceparent: Some(traceparent.clone()),
+                    tracestate: None,
+                },
+            );
+        }
         let invocation_id = call.invocation_id.clone();
         let cancellation = match admission {
             CallbackAdmission::Active(cancellation) => Ok(cancellation),
@@ -133,31 +152,34 @@ impl SessionInner {
         };
         let invocation = conversion::tool_call(call);
         let inner = Arc::clone(self);
-        tokio::spawn(async move {
-            let result = match invocation {
-                Ok(invocation) => {
-                    let callback = AssertUnwindSafe(async {
-                        inner
-                            .delegate
-                            .invoke_tool(invocation, cancellation.child_token())
-                            .await
-                    })
-                    .catch_unwind();
-                    tokio::select! {
-                        biased;
-                        _ = cancellation.cancelled() => return,
-                        result = callback => match result {
-                            Ok(result) => result,
-                            Err(_) => Err("code-mode tool delegate panicked".to_string()),
-                        },
+        tokio::spawn(
+            async move {
+                let result = match invocation {
+                    Ok(invocation) => {
+                        let callback = AssertUnwindSafe(async {
+                            inner
+                                .delegate
+                                .invoke_tool(invocation, cancellation.child_token())
+                                .await
+                        })
+                        .catch_unwind();
+                        tokio::select! {
+                            biased;
+                            _ = cancellation.cancelled() => return,
+                            result = callback => match result {
+                                Ok(result) => result,
+                                Err(_) => Err("code-mode tool delegate panicked".to_string()),
+                            },
+                        }
                     }
-                }
-                Err(error) => Err(error),
-            };
-            inner
-                .complete_tool_call(invocation_id, cancellation, result)
-                .await;
-        });
+                    Err(error) => Err(error),
+                };
+                inner
+                    .complete_tool_call(invocation_id, cancellation, result)
+                    .await;
+            }
+            .instrument(callback_span),
+        );
         Ok(())
     }
 

codex-rs/code-mode/src/grpc_session/conversion_tests.rs

@@ -73,6 +73,7 @@ fn tool_call_decodes_structured_input_and_namespace() {
         tool_kind: grpc::ToolKind::Function as i32,
         input_json: Some(br#"{"query":"hello"}"#.to_vec()),
         sequence: 1,
+        traceparent: None,
     };
 
     assert_eq!(

codex-rs/code-mode/src/grpc_session/mod.rs

@@ -27,6 +27,7 @@ use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 use tonic::transport::Channel;
+use tracing::Instrument;
 
 use self::operations::WaitSlot;
 use self::state::SessionState;
@@ -49,6 +50,15 @@ type GrpcClient = CodeModeHostClient<GrpcTransport>;
 
 const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
 
+fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::Span) {
+    if let Some(traceparent) =
+        codex_otel::span_w3c_trace_context(span).and_then(|trace| trace.traceparent)
+        && let Ok(traceparent) = traceparent.parse()
+    {
+        request.metadata_mut().insert("traceparent", traceparent);
+    }
+}
+
 /// Creates code-mode sessions over an HTTP/2 gRPC connection.
 #[derive(Clone)]
 pub struct GrpcCodeModeSessionProvider {
@@ -83,6 +93,7 @@ impl GrpcCodeModeSessionProvider {
         }
     }
 
+    #[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)]
     async fn open_binding(
         &self,
         delegate: Arc<dyn CodeModeSessionDelegate>,
@@ -100,17 +111,23 @@ impl GrpcCodeModeSessionProvider {
         let cell_execution_limits = (limits.max_yield_time_ms.is_some()
             || limits.max_heap_size_bytes.is_some())
         .then_some(limits);
-        let mut lease = deadline::startup(
-            "session opening",
-            client.open_session(grpc::OpenSessionRequest {
-                cell_execution_limits,
-            }),
-        )
-        .await?
-        .into_inner();
-        let first = deadline::startup("session lease opening", lease.message())
-            .await?
-            .ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
+        let open_session_span = tracing::info_span!("code_mode.grpc.open_session");
+        let mut open_session_request = tonic::Request::new(grpc::OpenSessionRequest {
+            cell_execution_limits,
+        });
+        inject_span_traceparent(&mut open_session_request, &open_session_span);
+        let (lease, first) = async {
+            let mut lease =
+                deadline::startup("session opening", client.open_session(open_session_request))
+                    .await?
+                    .into_inner();
+            let first = deadline::startup("session lease opening", lease.message())
+                .await?
+                .ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
+            Ok::<_, String>((lease, first))
+        }
+        .instrument(open_session_span)
+        .await?;
         let Some(grpc::session_event::Event::Opened(opened)) = first.event else {
             return Err("gRPC code-mode session lease omitted its opening event".to_string());
         };
@@ -134,13 +151,19 @@ impl GrpcCodeModeSessionProvider {
         };
         inner.spawn_session_events(lease);
 
-        let request = grpc::SubscribeToToolCallsRequest {
+        let subscribe_span = tracing::info_span!(
+            "code_mode.grpc.subscribe_to_tool_calls",
+            session.id = %inner.id,
+        );
+        let mut request = tonic::Request::new(grpc::SubscribeToToolCallsRequest {
             session_id: inner.id.clone(),
             tool_names: Vec::new(),
-        };
+        });
+        inject_span_traceparent(&mut request, &subscribe_span);
         let mut client = inner.client();
         let response =
             match deadline::startup("tool subscription", client.subscribe_to_tool_calls(request))
+                .instrument(subscribe_span)
                 .await
             {
                 Ok(response) => response,

codex-rs/code-mode/src/grpc_session/operations.rs

@@ -12,8 +12,10 @@ use codex_code_mode_protocol::StartedCell;
 use codex_code_mode_protocol::WaitOutcome;
 use codex_code_mode_protocol::WaitRequest;
 use codex_code_mode_protocol::grpc;
+use codex_protocol::protocol::W3cTraceContext;
 use tokio::sync::OwnedMutexGuard;
 use tokio::sync::oneshot;
+use tracing::Instrument;
 use tracing::debug;
 use uuid::Uuid;
 
@@ -68,6 +70,13 @@ impl SessionInner {
     ) -> Result<StartedCell, String> {
         self.require_open()?;
         let execution_id = Uuid::new_v4().to_string();
+        let execute_span = tracing::info_span!(
+            "code_mode.grpc.execute",
+            otel.name = "code_mode.grpc.execute",
+            execution.id = %execution_id,
+            call_id = %request.tool_call_id,
+        );
+        let trace = codex_otel::span_w3c_trace_context(&execute_span);
         let request = conversion::execute_request(&self.id, execution_id.clone(), request)?;
         self.state
             .lock()
@@ -81,7 +90,10 @@ impl SessionInner {
         let (started_tx, started_rx) = oneshot::channel();
         let inner = Arc::clone(self);
         self.stream_tasks.spawn(async move {
-            inner.drive_execution(request, ownership, started_tx).await;
+            inner
+                .drive_execution(request, ownership, started_tx, trace)
+                .instrument(execute_span)
+                .await;
         });
         started_rx
             .await
@@ -93,12 +105,19 @@ impl SessionInner {
         request: grpc::ExecuteRequest,
         ownership: ExecutionOwnership,
         started_tx: oneshot::Sender<Result<StartedCell, String>>,
+        trace: Option<W3cTraceContext>,
     ) {
         let runtime_timeout =
             Duration::from_millis(request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS))
                 .saturating_add(Duration::from_secs(1));
         let opening = async {
             let mut client = self.client();
+            let mut request = tonic::Request::new(request);
+            if let Some(traceparent) = trace.and_then(|trace| trace.traceparent)
+                && let Ok(traceparent) = traceparent.parse()
+            {
+                request.metadata_mut().insert("traceparent", traceparent);
+            }
             let mut stream =
                 deadline::request(&self, "execution", Duration::ZERO, client.execute(request))
                     .await?

codex-rs/code-mode/src/grpc_session/state_tests.rs

@@ -44,6 +44,7 @@ fn tool_call(execution_id: &str, invocation_id: u128) -> grpc::ToolCall {
         tool_kind: grpc::ToolKind::Function as i32,
         input_json: None,
         sequence: 1,
+        traceparent: None,
     }
 }