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-host/src/grpc/mod.rs

@@ -14,6 +14,7 @@ use codex_code_mode_protocol::CellId;
 use codex_code_mode_protocol::WaitRequest;
 use codex_code_mode_protocol::grpc as proto;
 use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
+use codex_protocol::protocol::W3cTraceContext;
 use futures::Stream;
 use futures::StreamExt;
 use tokio::sync::mpsc;
@@ -21,6 +22,7 @@ use tokio_stream::wrappers::ReceiverStream;
 use tonic::Request;
 use tonic::Response;
 use tonic::Status;
+use tracing::Instrument;
 
 use self::session::GrpcHostState;
 use self::session::GrpcSession;
@@ -29,6 +31,17 @@ use self::waits::WaitRegistration;
 type GrpcStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
 type GrpcFuture<'a, T> = Pin<Box<dyn Future<Output = Result<Response<T>, Status>> + Send + 'a>>;
 
+fn trace_context_from_request<T>(request: &Request<T>) -> Option<W3cTraceContext> {
+    request
+        .metadata()
+        .get("traceparent")
+        .and_then(|value| value.to_str().ok())
+        .map(|traceparent| W3cTraceContext {
+            traceparent: Some(traceparent.to_string()),
+            tracestate: None,
+        })
+}
+
 /// Serves transport-independent, leased code-mode sessions over gRPC.
 #[derive(Clone)]
 pub struct GrpcCodeModeHost {
@@ -43,12 +56,6 @@ impl GrpcCodeModeHost {
         }
     }
 
-    #[tracing::instrument(
-        name = "code_mode_host.grpc.open_session",
-        level = "info",
-        skip_all,
-        fields(otel.name = "code_mode_host.grpc.open_session")
-    )]
     async fn open_session_request(
         &self,
         request: proto::OpenSessionRequest,
@@ -119,20 +126,10 @@ impl GrpcCodeModeHost {
         Ok(Response::new(proto::AcknowledgeNotificationResponse {}))
     }
 
-    #[tracing::instrument(
-        name = "code_mode_host.grpc.execute",
-        level = "info",
-        skip_all,
-        fields(
-            otel.name = "code_mode_host.grpc.execute",
-            session.id = %request.session_id,
-            execution.id = %request.execution_id,
-            call_id = %request.tool_call_id,
-        )
-    )]
     async fn execute_request(
         &self,
         request: proto::ExecuteRequest,
+        callback_traceparent: Option<String>,
     ) -> Result<Response<GrpcStream<proto::ExecuteEvent>>, Status> {
         let session = self.state.session(&request.session_id)?;
         validation::identifier(&request.execution_id, "execution ID")?;
@@ -154,7 +151,12 @@ impl GrpcCodeModeHost {
             }
         };
         let cell_id = started.cell_id.clone();
-        session.admit_execution(execution_id.clone(), cell_id.to_string(), cell_permit)?;
+        session.admit_execution(
+            execution_id.clone(),
+            cell_id.to_string(),
+            cell_permit,
+            callback_traceparent,
+        )?;
 
         let (sender, receiver) = mpsc::channel(/*buffer*/ 2);
         sender
@@ -167,22 +169,26 @@ impl GrpcCodeModeHost {
                 )),
             }))
             .map_err(|_| Status::internal("failed to publish code-mode execution admission"))?;
-        tokio::spawn(async move {
-            let _request_permit = request_permit;
-            tokio::select! {
-                biased;
-                _ = sender.closed() => {}
-                response = started.initial_response() => {
-                    let event = response.map(|response| proto::ExecuteEvent {
-                        event: Some(proto::execute_event::Event::Outcome(
-                            conversions::execution_outcome(response),
-                        )),
-                    }).map_err(Status::internal);
-                    let _ = sender.send(event).await;
+        let outcome_span = tracing::Span::current();
+        tokio::spawn(
+            async move {
+                let _request_permit = request_permit;
+                tokio::select! {
+                    biased;
+                    _ = sender.closed() => {}
+                    response = started.initial_response() => {
+                        let event = response.map(|response| proto::ExecuteEvent {
+                            event: Some(proto::execute_event::Event::Outcome(
+                                conversions::execution_outcome(response),
+                            )),
+                        }).map_err(Status::internal);
+                        let _ = sender.send(event).await;
+                    }
+                    _ = session.closed.cancelled() => {}
                 }
-                _ = session.closed.cancelled() => {}
             }
-        });
+            .instrument(outcome_span),
+        );
 
         let stream = ReceiverStream::new(receiver).inspect(move |event| {
             if matches!(
@@ -278,7 +284,16 @@ impl CodeModeHost for GrpcCodeModeHost {
         'a: 'async_trait,
         Self: 'async_trait,
     {
-        Box::pin(self.open_session_request(request.into_inner()))
+        let trace = trace_context_from_request(&request);
+        let request = request.into_inner();
+        let open_session_span = tracing::info_span!("code_mode_host.grpc.open_session");
+        if let Some(trace) = trace.as_ref() {
+            codex_otel::set_parent_from_w3c_trace_context(&open_session_span, trace);
+        }
+        Box::pin(
+            self.open_session_request(request)
+                .instrument(open_session_span),
+        )
     }
 
     fn close_session<'a, 'async_trait>(
@@ -333,7 +348,24 @@ impl CodeModeHost for GrpcCodeModeHost {
         'a: 'async_trait,
         Self: 'async_trait,
     {
-        Box::pin(self.execute_request(request.into_inner()))
+        let trace = trace_context_from_request(&request);
+        let request = request.into_inner();
+        let execute_span = tracing::info_span!(
+            "code_mode_host.grpc.execute",
+            otel.name = "code_mode_host.grpc.execute",
+            session.id = %request.session_id,
+            execution.id = %request.execution_id,
+            call_id = %request.tool_call_id,
+        );
+        if let Some(trace) = trace.as_ref() {
+            codex_otel::set_parent_from_w3c_trace_context(&execute_span, trace);
+        }
+        let callback_traceparent =
+            codex_otel::span_w3c_trace_context(&execute_span).and_then(|trace| trace.traceparent);
+        Box::pin(
+            self.execute_request(request, callback_traceparent)
+                .instrument(execute_span),
+        )
     }
 
     fn wait<'a, 'async_trait>(

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

@@ -139,6 +139,7 @@ async fn dropping_execution_before_admission_releases_its_reservation() {
             execution_id,
             "cell".to_string(),
             host.state.cell_permit().expect("reserve cell permit"),
+            /*traceparent*/ None,
         )
         .expect_err("abandoned execution must not admit a runtime cell");
     assert_eq!(error.code(), Code::Cancelled);

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

@@ -103,7 +103,7 @@ impl GrpcSession {
         let canonical_tool_name =
             ToolName::new(tool_name.namespace.clone(), tool_name.name.clone())
                 .with_default_namespace();
-        let (sequence, subscriptions) = {
+        let (sequence, traceparent, subscriptions) = {
             let state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
             let Some(execution) = state.cells.get(&cell_id) else {
                 return Err("code-mode cell closed before dispatching its tool call".to_string());
@@ -124,7 +124,7 @@ impl GrpcSession {
                 })
                 .map(|subscription| (subscription.id, subscription.sender.clone()))
                 .collect::<Vec<_>>();
-            (sequence, subscriptions)
+            (sequence, execution.traceparent.clone(), subscriptions)
         };
         if subscriptions.is_empty() {
             return Err("no code-mode tool subscription matches the requested tool".to_string());
@@ -139,6 +139,7 @@ impl GrpcSession {
             tool_kind: conversions::tool_kind(invocation.tool_kind),
             input_json,
             sequence,
+            traceparent,
         };
         if message.encoded_len() > MAX_FRAME_BYTES {
             return Err(format!(

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

@@ -71,6 +71,7 @@ pub(super) struct SessionState {
 
 pub(super) struct ExecutionState {
     pub(super) execution_id: String,
+    pub(super) traceparent: Option<String>,
     pub(super) tool_call_sequence: u64,
     permit: OwnedSemaphorePermit,
 }
@@ -296,6 +297,7 @@ impl GrpcSession {
         execution_id: String,
         cell_id: String,
         permit: OwnedSemaphorePermit,
+        traceparent: Option<String>,
     ) -> Result<(), Status> {
         let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
         if !state.pending_executions.remove(&execution_id) {
@@ -308,6 +310,7 @@ impl GrpcSession {
         };
         entry.insert(ExecutionState {
             execution_id,
+            traceparent,
             tool_call_sequence: 0,
             permit,
         });