diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index b560f568bf..c3d5bc0c42 100644 --- a/codex-rs/exec-server/README.md +++ b/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 diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 89f33d938a..37a026931a 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/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 diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs index 232dfa479e..986fd650fe 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs +++ b/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()?; diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 9969a3e130..ad2163c903 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/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, 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; } } diff --git a/codex-rs/exec-server/src/noise_relay/harness_tests.rs b/codex-rs/exec-server/src/noise_relay/harness_tests.rs index 348351e61c..16df1cbf0f 100644 --- a/codex-rs/exec-server/src/noise_relay/harness_tests.rs +++ b/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, ) -> Result { + read_outbound_frame(outbound_rx) + .await? + .into_data() + .map_err(anyhow::Error::from) +} + +async fn read_outbound_frame( + outbound_rx: &mut futures_mpsc::UnboundedReceiver, +) -> Result { 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<( diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto index 3de1bfbe99..a9d2100137 100644 --- a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto +++ b/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 { diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs index f65e1329f5..4694ed4024 100644 --- a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs +++ b/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, + #[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)] diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 3d0be695b4..82ef338c5f 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/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) -> Self { + pub(crate) fn data( + stream_id: String, + seq: u32, + payload: Vec, + trace: Option, + ) -> 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(), )) diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index 704c48d9a1..754c6e498c 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/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()))