Propagate request trace context through exec-server relays (#38690)

## What changed

- Add optional W3C `traceparent` and `tracestate` fields to relay frames.
- Copy trace context from JSON-RPC requests onto relay data frames.
- For encrypted requests split across multiple Noise records, attach the context only to the first record while keeping the request payload encrypted.

## Testing

- Cover trace propagation for both Noise relay paths, including fragmented encrypted requests.

GitOrigin-RevId: a61bbbefef31e2e7e93a43f439c5f296700feb7b

codex-rs/exec-server/README.md

@@ -59,6 +59,8 @@ identity plus endpoint-owned reliability metadata:
 ```text
 version
 stream_id
+traceparent       // optional W3C parent on the first frame of a traced request
+tracestate        // optional W3C vendor state paired with traceparent
 body              // handshake | data | ack_frame | resume | reset | heartbeat
 ack               // highest contiguous peer segment seq received
 ack_bits          // bitset for peer segment seqs after ack

codex-rs/exec-server/src/noise_relay/executor_stream.rs

@@ -7,6 +7,7 @@
 use std::sync::Arc;
 use std::sync::Mutex;
 
+use codex_exec_server_protocol::JSONRPCMessage;
 use tokio::sync::mpsc;
 use tokio::sync::watch;
 use tracing::warn;
@@ -119,6 +120,12 @@ pub(crate) fn spawn_noise_virtual_stream(
                     break;
                 }
             };
+            let mut trace = match message {
+                JSONRPCMessage::Request(request) => request.trace,
+                JSONRPCMessage::Notification(_)
+                | JSONRPCMessage::Response(_)
+                | JSONRPCMessage::Error(_) => None,
+            };
             for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {
                 let seq = match take_next_sequence(&mut next_seq) {
                     Ok(seq) => seq,
@@ -140,7 +147,12 @@ pub(crate) fn spawn_noise_virtual_stream(
                         break 'writer;
                     }
                 };
-                let frame = RelayMessageFrame::data(writer_stream_id.clone(), seq, ciphertext);
+                let frame = RelayMessageFrame::data(
+                    writer_stream_id.clone(),
+                    seq,
+                    ciphertext,
+                    trace.take(),
+                );
                 if physical_outgoing_tx
                     .send(encode_relay_message_frame(&frame))
                     .await

codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs

@@ -2,8 +2,10 @@ use std::time::Duration;
 
 use anyhow::Result;
 use codex_exec_server_protocol::JSONRPCMessage;
+use codex_exec_server_protocol::JSONRPCRequest;
 use codex_exec_server_protocol::JSONRPCResponse;
 use codex_exec_server_protocol::RequestId;
+use codex_protocol::protocol::W3cTraceContext;
 use tokio::sync::mpsc;
 use tokio::time::timeout;
 
@@ -16,8 +18,60 @@ use crate::noise_channel::NoiseChannelIdentity;
 use crate::noise_channel::PendingResponderHandshake;
 use crate::noise_relay::message_framing::frame_jsonrpc_message;
 use crate::relay_proto::RelayData;
+use crate::relay_proto::RelayMessageFrame;
 use crate::server::ConnectionProcessor;
 
+#[test]
+fn executor_requests_attach_trace_context_only_to_the_first_noise_record() {
+    let traceparent = "00-00000000000000000000000000000001-0000000000000002-01";
+    let tracestate = "dd=s:1";
+    let owned_traceparent = traceparent.to_string();
+    let owned_tracestate = tracestate.to_string();
+    let traceparent_ptr = owned_traceparent.as_ptr();
+    let tracestate_ptr = owned_tracestate.as_ptr();
+    let mut request = JSONRPCRequest {
+        id: RequestId::Integer(1),
+        method: "approval/request".to_string(),
+        params: None,
+        trace: Some(W3cTraceContext {
+            traceparent: Some(owned_traceparent),
+            tracestate: Some(owned_tracestate),
+        }),
+    };
+
+    let first = RelayMessageFrame::data(
+        "stream-1".to_string(),
+        /*seq*/ 0,
+        vec![1],
+        request.trace.take(),
+    );
+    assert_eq!(first.traceparent.as_deref(), Some(traceparent));
+    assert_eq!(first.tracestate.as_deref(), Some(tracestate));
+    assert_eq!(
+        first.traceparent.as_ref().unwrap().as_ptr(),
+        traceparent_ptr
+    );
+    assert_eq!(first.tracestate.as_ref().unwrap().as_ptr(), tracestate_ptr);
+
+    let second = RelayMessageFrame::data(
+        "stream-1".to_string(),
+        /*seq*/ 1,
+        vec![2],
+        request.trace.take(),
+    );
+    assert!(second.traceparent.is_none());
+    assert!(second.tracestate.is_none());
+
+    let response = RelayMessageFrame::data(
+        "stream-1".to_string(),
+        /*seq*/ 2,
+        vec![3],
+        /*trace*/ None,
+    );
+    assert!(response.traceparent.is_none());
+    assert!(response.tracestate.is_none());
+}
+
 #[tokio::test]
 async fn processor_exit_reports_closed_virtual_stream() -> Result<()> {
     let executor_identity = NoiseChannelIdentity::generate()?;

codex-rs/exec-server/src/noise_relay/harness.rs

@@ -275,7 +275,7 @@ where
         // Keep one framed message as a cursor. Sending one Noise record per loop
         // creates a scheduling point for keepalive and inbound control frames
         // without splitting the WebSocket reader and writer.
-        let mut pending_outbound: Option<(Vec<u8>, usize)> = None;
+        let mut pending_outbound = None;
         let mut force_incoming = false;
         let mut frames_drained_after_pong_deadline = 0usize;
         'relay: loop {
@@ -344,13 +344,20 @@ where
                     let Some(message) = maybe_message else {
                         break;
                     };
-                    pending_outbound = Some(match frame_jsonrpc_message(&message) {
-                        Ok(framed) => (framed, 0),
+                    let framed = match frame_jsonrpc_message(&message) {
+                        Ok(framed) => framed,
                         Err(error) => {
                             warn!("failed to frame JSON-RPC payload for Noise relay: {error}");
                             break;
                         }
-                    });
+                    };
+                    let request_trace = match message {
+                        codex_exec_server_protocol::JSONRPCMessage::Request(request) => request.trace,
+                        codex_exec_server_protocol::JSONRPCMessage::Notification(_)
+                        | codex_exec_server_protocol::JSONRPCMessage::Response(_)
+                        | codex_exec_server_protocol::JSONRPCMessage::Error(_) => None,
+                    };
+                    pending_outbound = Some((framed, 0, request_trace));
                 }
                 _ = std::future::ready(()), if pending_outbound.is_some() && !force_incoming && !pong_deadline_expired => {
                     let seq = match take_next_sequence(&mut next_outbound_seq) {
@@ -360,8 +367,8 @@ where
                             break 'relay;
                         }
                     };
-                    let (ciphertext, next_offset, message_complete) = {
-                        let Some((framed, offset)) = pending_outbound.as_ref() else {
+                    let (ciphertext, next_offset, message_complete, request_trace) = {
+                        let Some((framed, offset, request_trace)) = pending_outbound.as_mut() else {
                             continue;
                         };
                         let next_offset = (*offset + NOISE_RECORD_PLAINTEXT_LEN).min(framed.len());
@@ -372,9 +379,15 @@ where
                                 break 'relay;
                             }
                         };
-                        (ciphertext, next_offset, next_offset == framed.len())
+                        (
+                            ciphertext,
+                            next_offset,
+                            next_offset == framed.len(),
+                            request_trace.take(),
+                        )
                     };
-                    let frame = RelayMessageFrame::data(stream_id.clone(), seq, ciphertext);
+                    let frame =
+                        RelayMessageFrame::data(stream_id.clone(), seq, ciphertext, request_trace);
                     // A Pong can arrive after the readiness check while this write owns the
                     // combined sink and stream. A single bounded record can therefore hit the
                     // deadline and disconnect with that Pong queued. Treat that as write
@@ -391,7 +404,7 @@ where
                     }
                     if message_complete {
                         pending_outbound = None;
-                    } else if let Some((_framed, offset)) = pending_outbound.as_mut() {
+                    } else if let Some((_framed, offset, _request_trace)) = pending_outbound.as_mut() {
                         *offset = next_offset;
                     }
                 }

codex-rs/exec-server/src/noise_relay/harness_tests.rs

@@ -11,6 +11,7 @@ use anyhow::Result;
 use codex_exec_server_protocol::JSONRPCMessage;
 use codex_exec_server_protocol::JSONRPCRequest;
 use codex_exec_server_protocol::RequestId;
+use codex_protocol::protocol::W3cTraceContext;
 use futures::Sink;
 use futures::SinkExt;
 use futures::StreamExt;
@@ -31,6 +32,52 @@ use crate::noise_channel::PendingResponderHandshake;
 const ENVIRONMENT_ID: &str = "environment-1";
 const EXECUTOR_REGISTRATION_ID: &str = "registration-1";
 
+#[tokio::test(start_paused = true)]
+async fn first_encrypted_request_frame_exposes_only_its_trace_context() -> Result<()> {
+    let (connection, mut control, mut outbound_rx) = connected_controlled_harness().await?;
+    let traceparent = "00-00000000000000000000000000000001-0000000000000002-01";
+    let tracestate = "dd=s:1";
+
+    connection
+        .outgoing_tx
+        .send(JSONRPCMessage::Request(JSONRPCRequest {
+            id: RequestId::Integer(1),
+            method: "fs/getMetadata".to_string(),
+            params: Some(serde_json::json!({
+                "payload": "x".repeat(NOISE_RECORD_PLAINTEXT_LEN * 2),
+            })),
+            trace: Some(W3cTraceContext {
+                traceparent: Some(traceparent.to_string()),
+                tracestate: Some(tracestate.to_string()),
+            }),
+        }))
+        .await?;
+
+    control.wait_for_blocked_write(/*expected*/ 1).await?;
+    control.grant_writes(/*count*/ 1);
+    let first = read_outbound_frame(&mut outbound_rx).await?;
+    assert_eq!(
+        (first.traceparent.as_deref(), first.tracestate.as_deref()),
+        (Some(traceparent), Some(tracestate))
+    );
+    let first_payload = first.into_data()?.payload;
+    assert!(
+        !first_payload
+            .windows("fs/getMetadata".len())
+            .any(|window| window == b"fs/getMetadata")
+    );
+
+    control.wait_for_blocked_write(/*expected*/ 2).await?;
+    control.grant_writes(/*count*/ 1);
+    let second = read_outbound_frame(&mut outbound_rx).await?;
+    assert_eq!((second.traceparent, second.tracestate), (None, None));
+
+    for task in &connection.task_handles {
+        task.abort();
+    }
+    Ok(())
+}
+
 #[tokio::test(start_paused = true)]
 async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> {
     let (connection, mut control, mut outbound_rx) = connected_controlled_harness().await?;
@@ -240,6 +287,15 @@ async fn application_event_delivery_is_bounded() -> Result<()> {
 async fn read_outbound_data(
     outbound_rx: &mut futures_mpsc::UnboundedReceiver<Message>,
 ) -> Result<RelayData> {
+    read_outbound_frame(outbound_rx)
+        .await?
+        .into_data()
+        .map_err(anyhow::Error::from)
+}
+
+async fn read_outbound_frame(
+    outbound_rx: &mut futures_mpsc::UnboundedReceiver<Message>,
+) -> Result<RelayMessageFrame> {
     let Message::Binary(payload) = timeout(Duration::from_secs(1), outbound_rx.next())
         .await?
         .context("harness closed before sending data")?
@@ -248,7 +304,7 @@ async fn read_outbound_data(
     };
     let frame = decode_relay_message_frame(payload.as_ref())?;
     assert_eq!(frame.validate()?, RelayFrameBodyKind::Data);
-    frame.into_data().map_err(anyhow::Error::from)
+    Ok(frame)
 }
 
 async fn connected_controlled_harness() -> Result<(

codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto

@@ -16,6 +16,9 @@ message RelayMessageFrame {
     RelayHeartbeat heartbeat = 9;
     RelayHandshake handshake = 10;
   }
+
+  optional string traceparent = 11;
+  optional string tracestate = 12;
 }
 
 message RelayData {

codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs

@@ -11,6 +11,10 @@ pub struct RelayMessageFrame {
     pub ack_bits: u32,
     #[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9, 10")]
     pub body: ::core::option::Option<relay_message_frame::Body>,
+    #[prost(string, optional, tag = "11")]
+    pub traceparent: ::core::option::Option<::prost::alloc::string::String>,
+    #[prost(string, optional, tag = "12")]
+    pub tracestate: ::core::option::Option<::prost::alloc::string::String>,
 }
 pub mod relay_message_frame {
     #[derive(Clone, PartialEq, ::prost::Oneof)]

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

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use codex_exec_server_protocol::JSONRPCMessage;
+use codex_protocol::protocol::W3cTraceContext;
 use futures::Sink;
 use futures::SinkExt;
 use futures::Stream;
@@ -81,18 +82,27 @@ pub(crate) enum RelayFrameBodyKind {
 }
 
 impl RelayMessageFrame {
-    pub(crate) fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
+    pub(crate) fn data(
+        stream_id: String,
+        seq: u32,
+        payload: Vec<u8>,
+        trace: Option<W3cTraceContext>,
+    ) -> Self {
+        let (traceparent, tracestate) = trace
+            .map(|trace| (trace.traceparent, trace.tracestate))
+            .unwrap_or_default();
         Self {
             version: RELAY_MESSAGE_FRAME_VERSION,
             stream_id,
-            ack: 0,
-            ack_bits: 0,
+            traceparent,
+            tracestate,
             body: Some(relay_message_frame::Body::Data(RelayData {
                 seq,
                 segment_index: 0,
                 segment_count: 1,
                 payload,
             })),
+            ..Self::default()
         }
     }
 
@@ -100,11 +110,10 @@ impl RelayMessageFrame {
         Self {
             version: RELAY_MESSAGE_FRAME_VERSION,
             stream_id,
-            ack: 0,
-            ack_bits: 0,
             body: Some(relay_message_frame::Body::Resume(RelayResume {
                 next_seq: 0,
             })),
+            ..Self::default()
         }
     }
 
@@ -112,11 +121,10 @@ impl RelayMessageFrame {
         Self {
             version: RELAY_MESSAGE_FRAME_VERSION,
             stream_id,
-            ack: 0,
-            ack_bits: 0,
             body: Some(relay_message_frame::Body::Handshake(RelayHandshake {
                 payload,
             })),
+            ..Self::default()
         }
     }
 
@@ -124,9 +132,8 @@ impl RelayMessageFrame {
         Self {
             version: RELAY_MESSAGE_FRAME_VERSION,
             stream_id,
-            ack: 0,
-            ack_bits: 0,
             body: Some(relay_message_frame::Body::Reset(RelayReset { reason })),
+            ..Self::default()
         }
     }
 
@@ -313,7 +320,13 @@ where
                             break;
                         }
                     };
-                    let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload);
+                    let trace = match message {
+                        JSONRPCMessage::Request(request) => request.trace,
+                        JSONRPCMessage::Notification(_)
+                        | JSONRPCMessage::Response(_)
+                        | JSONRPCMessage::Error(_) => None,
+                    };
+                    let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload, trace);
                     next_seq = next_seq.wrapping_add(1);
                     if websocket
                         .send(Message::Binary(encode_relay_message_frame(&frame).into()))
@@ -961,6 +974,7 @@ mod tests {
                     stream_id,
                     /*seq*/ 0,
                     jsonrpc_payload(&message)?,
+                    /*trace*/ None,
                 ))
                 .into(),
             ))

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

@@ -426,7 +426,7 @@ async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Re
         )?;
         for frame in [
             RelayMessageFrame::handshake(stream_id.clone(), request),
-            RelayMessageFrame::data(stream_id, /*seq*/ 0, vec![0]),
+            RelayMessageFrame::data(stream_id, /*seq*/ 0, vec![0], /*trace*/ None),
         ] {
             harness_websocket
                 .send(Message::Binary(encode_relay_message_frame(&frame).into()))