From 51587401139bcf6be48f3d409113ade322239ad2 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 04:28:38 +0000 Subject: [PATCH 1/2] exec-server: trace terminal RPC notifications --- codex-rs/exec-server/src/client.rs | 13 ++ .../src/client/notification_tracing.rs | 37 +++++ .../src/client/notification_tracing_tests.rs | 147 ++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 codex-rs/exec-server/src/client/notification_tracing.rs create mode 100644 codex-rs/exec-server/src/client/notification_tracing_tests.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 55c7c7e0d4a1..3a08ed1ed9d5 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -105,6 +105,7 @@ use crate::rpc::RpcCallError; use crate::rpc::RpcClient; pub(crate) mod http_client; +mod notification_tracing; #[path = "client_recovery.rs"] mod recovery; @@ -1155,6 +1156,18 @@ async fn fail_all_in_flight_work(inner: &Arc, message: String) { async fn handle_server_notification( inner: &Arc, notification: JSONRPCNotification, +) -> Result<(), ExecServerError> { + let notification_span = notification_tracing::notification_span(¬ification); + let result = handle_server_notification_inner(inner, notification) + .instrument(notification_span.clone()) + .await; + notification_span.record("result", if result.is_ok() { "success" } else { "error" }); + result +} + +async fn handle_server_notification_inner( + inner: &Arc, + notification: JSONRPCNotification, ) -> Result<(), ExecServerError> { match notification.method.as_str() { EXEC_OUTPUT_DELTA_METHOD => { diff --git a/codex-rs/exec-server/src/client/notification_tracing.rs b/codex-rs/exec-server/src/client/notification_tracing.rs new file mode 100644 index 000000000000..25d20eb5bfa7 --- /dev/null +++ b/codex-rs/exec-server/src/client/notification_tracing.rs @@ -0,0 +1,37 @@ +use codex_exec_server_protocol::JSONRPCNotification; +use tracing::warn; + +use crate::rpc::should_trace_server_notification; + +pub(super) fn notification_span(notification: &JSONRPCNotification) -> tracing::Span { + let method = notification.method.as_str(); + let params = notification + .params + .as_ref() + .unwrap_or(&serde_json::Value::Null); + if !should_trace_server_notification(method, params) { + return tracing::Span::none(); + } + let span = tracing::info_span!( + "codex.exec_server.notification", + otel.kind = "server", + otel.name = method, + rpc.system = "jsonrpc", + rpc.method = method, + method, + result = tracing::field::Empty, + ); + if let Some(trace) = ¬ification.trace + && !codex_otel::set_parent_from_w3c_trace_context(&span, trace) + { + warn!( + method, + "ignoring invalid inbound exec-server notification trace carrier" + ); + } + span +} + +#[cfg(test)] +#[path = "notification_tracing_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/client/notification_tracing_tests.rs b/codex-rs/exec-server/src/client/notification_tracing_tests.rs new file mode 100644 index 000000000000..cf6e1cc7163e --- /dev/null +++ b/codex-rs/exec-server/src/client/notification_tracing_tests.rs @@ -0,0 +1,147 @@ +use codex_exec_server_protocol::JSONRPCNotification; +use codex_protocol::protocol::W3cTraceContext; +use opentelemetry::trace::SpanId; +use opentelemetry::trace::TraceId; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::InMemorySpanExporter; +use opentelemetry_sdk::trace::SdkTracerProvider; +use pretty_assertions::assert_eq; +use tracing_subscriber::filter::filter_fn; +use tracing_subscriber::prelude::*; + +use super::notification_span; +use crate::protocol::EXEC_CLOSED_METHOD; +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; +use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD; + +#[test] +#[serial_test::serial(exec_server_tracing)] +fn receive_span_uses_inbound_trace_parent() { + let span_exporter = InMemorySpanExporter::default(); + let tracer_provider = SdkTracerProvider::builder() + .with_simple_exporter(span_exporter.clone()) + .build(); + let tracer = tracer_provider.tracer("exec-server-test"); + let subscriber = tracing_subscriber::registry().with( + tracing_opentelemetry::layer() + .with_tracer(tracer) + .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), + ); + let trace_id = TraceId::from_hex("00000000000000000000000000000001").expect("trace id"); + let parent_span_id = SpanId::from_hex("0000000000000002").expect("span id"); + let notification = notification( + EXEC_EXITED_METHOD, + serde_json::json!({}), + Some(W3cTraceContext { + traceparent: Some(format!("00-{trace_id}-{parent_span_id}-01")), + tracestate: None, + }), + ); + + tracing::subscriber::with_default(subscriber, || { + tracing::callsite::rebuild_interest_cache(); + let span = notification_span(¬ification); + span.in_scope(|| {}); + drop(span); + }); + + tracer_provider.force_flush().expect("flush traces"); + let spans = span_exporter.get_finished_spans().expect("span export"); + let receive_span = spans + .iter() + .find(|span| span.name.as_ref() == EXEC_EXITED_METHOD) + .expect("process exited receive span"); + assert_eq!(receive_span.span_context.trace_id(), trace_id); + assert_eq!(receive_span.parent_span_id, parent_span_id); +} + +#[test] +#[serial_test::serial(exec_server_tracing)] +fn receive_span_policy_avoids_hot_path_volume() { + let span_exporter = InMemorySpanExporter::default(); + let tracer_provider = SdkTracerProvider::builder() + .with_simple_exporter(span_exporter.clone()) + .build(); + let tracer = tracer_provider.tracer("exec-server-test"); + let subscriber = tracing_subscriber::registry().with( + tracing_opentelemetry::layer() + .with_tracer(tracer) + .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), + ); + + tracing::subscriber::with_default(subscriber, || { + tracing::callsite::rebuild_interest_cache(); + for _ in 0..100 { + finish_span(notification_span(¬ification( + EXEC_OUTPUT_DELTA_METHOD, + serde_json::json!({}), + /*trace*/ None, + ))); + finish_span(notification_span(¬ification( + HTTP_REQUEST_BODY_DELTA_METHOD, + serde_json::json!({"done": false, "error": null}), + /*trace*/ None, + ))); + } + + for notification in [ + notification( + EXEC_EXITED_METHOD, + serde_json::json!({}), + /*trace*/ None, + ), + notification( + EXEC_CLOSED_METHOD, + serde_json::json!({}), + /*trace*/ None, + ), + notification( + HTTP_REQUEST_BODY_DELTA_METHOD, + serde_json::json!({"done": true, "error": null}), + /*trace*/ None, + ), + notification( + HTTP_REQUEST_BODY_DELTA_METHOD, + serde_json::json!({"done": false, "error": "stream failed"}), + /*trace*/ None, + ), + ] { + finish_span(notification_span(¬ification)); + } + }); + + tracer_provider.force_flush().expect("flush traces"); + let mut span_names = span_exporter + .get_finished_spans() + .expect("span export") + .into_iter() + .map(|span| span.name.into_owned()) + .collect::>(); + span_names.sort(); + let mut expected = vec![ + EXEC_CLOSED_METHOD.to_string(), + EXEC_EXITED_METHOD.to_string(), + HTTP_REQUEST_BODY_DELTA_METHOD.to_string(), + HTTP_REQUEST_BODY_DELTA_METHOD.to_string(), + ]; + expected.sort(); + assert_eq!(span_names, expected); +} + +fn notification( + method: &str, + params: serde_json::Value, + trace: Option, +) -> JSONRPCNotification { + JSONRPCNotification { + method: method.to_string(), + params: Some(params), + trace, + } +} + +fn finish_span(span: tracing::Span) { + span.in_scope(|| {}); + span.record("result", "success"); +} From a9f47886313775e97f55bf7400248d122949863e Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Fri, 10 Jul 2026 00:15:03 +0000 Subject: [PATCH 2/2] codex: fix CI failure on PR #31730 --- .../src/client/notification_tracing.rs | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/client/notification_tracing.rs b/codex-rs/exec-server/src/client/notification_tracing.rs index 25d20eb5bfa7..8cfacce91d74 100644 --- a/codex-rs/exec-server/src/client/notification_tracing.rs +++ b/codex-rs/exec-server/src/client/notification_tracing.rs @@ -1,7 +1,9 @@ use codex_exec_server_protocol::JSONRPCNotification; use tracing::warn; -use crate::rpc::should_trace_server_notification; +use crate::protocol::EXEC_CLOSED_METHOD; +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD; pub(super) fn notification_span(notification: &JSONRPCNotification) -> tracing::Span { let method = notification.method.as_str(); @@ -9,7 +11,18 @@ pub(super) fn notification_span(notification: &JSONRPCNotification) -> tracing:: .params .as_ref() .unwrap_or(&serde_json::Value::Null); - if !should_trace_server_notification(method, params) { + let should_trace = match method { + EXEC_EXITED_METHOD | EXEC_CLOSED_METHOD => true, + HTTP_REQUEST_BODY_DELTA_METHOD => { + let Some(params) = params.as_object() else { + return tracing::Span::none(); + }; + params.get("done").and_then(serde_json::Value::as_bool) == Some(true) + || params.get("error").is_some_and(|error| !error.is_null()) + } + _ => false, + }; + if !should_trace { return tracing::Span::none(); } let span = tracing::info_span!(