From d473e8d56d40b990e5f1e90130529df1978041f8 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Mon, 2 Mar 2026 16:01:41 -0800 Subject: [PATCH] feat(app-server): add tracing to all app-server APIs (#13285) ### Overview This PR adds the first piece of tracing for app-server JSON-RPC requests. There are two main changes: - JSON-RPC requests can now take an optional W3C trace context at the top level via a `trace` field (`traceparent` / `tracestate`). - app-server now creates a dedicated request span for every inbound JSON-RPC request in `MessageProcessor`, and uses the request-level trace context as the parent when present. For compatibility with existing flows, app-server still falls back to the TRACEPARENT env var when there is no request-level traceparent. This PR is intentionally scoped to the app-server boundary. In a followup, we'll actually propagate trace context through the async handoff into core execution spans like run_turn, which will make app-server traces much more useful. ### Spans A few details on the app-server span shape: - each inbound request gets its own server span - span/resource names are based on the JSON-RPC method (`initialize`, `thread/start`, `turn/start`, etc.) - spans record transport (stdio vs websocket), request id, connection id, and client name/version when available - `initialize` stores client metadata in session state so later requests on the same connection can reuse it --- codex-rs/Cargo.lock | 1 + .../schema/json/JSONRPCMessage.json | 30 ++- .../schema/json/JSONRPCRequest.json | 30 ++- .../codex_app_server_protocol.schemas.json | 30 ++- .../app-server-protocol/src/jsonrpc_lite.rs | 5 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/app_server_tracing.rs | 101 +++++++ codex-rs/app-server/src/lib.rs | 4 +- codex-rs/app-server/src/message_processor.rs | 248 +++++++++--------- codex-rs/app-server/src/transport.rs | 2 + .../app-server/tests/common/mcp_process.rs | 1 + .../app-server/tests/suite/v2/analytics.rs | 4 +- .../suite/v2/connection_handling_websocket.rs | 1 + codex-rs/otel/src/lib.rs | 6 + codex-rs/otel/src/otel_provider.rs | 31 +-- codex-rs/otel/src/trace_context.rs | 106 ++++++++ codex-rs/protocol/src/protocol.rs | 10 + 17 files changed, 464 insertions(+), 147 deletions(-) create mode 100644 codex-rs/app-server/src/app_server_tracing.rs create mode 100644 codex-rs/otel/src/trace_context.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2ab1d5943..6e13edc7b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1428,6 +1428,7 @@ dependencies = [ "codex-feedback", "codex-file-search", "codex-login", + "codex-otel", "codex-protocol", "codex-rmcp-client", "codex-shell-command", diff --git a/codex-rs/app-server-protocol/schema/json/JSONRPCMessage.json b/codex-rs/app-server-protocol/schema/json/JSONRPCMessage.json index b2f6c3101..27b78b90c 100644 --- a/codex-rs/app-server-protocol/schema/json/JSONRPCMessage.json +++ b/codex-rs/app-server-protocol/schema/json/JSONRPCMessage.json @@ -70,7 +70,18 @@ "method": { "type": "string" }, - "params": true + "params": true, + "trace": { + "anyOf": [ + { + "$ref": "#/definitions/W3cTraceContext" + }, + { + "type": "null" + } + ], + "description": "Optional W3C Trace Context for distributed tracing." + } }, "required": [ "id", @@ -102,6 +113,23 @@ "type": "integer" } ] + }, + "W3cTraceContext": { + "properties": { + "traceparent": { + "type": [ + "string", + "null" + ] + }, + "tracestate": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" } }, "description": "Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent.", diff --git a/codex-rs/app-server-protocol/schema/json/JSONRPCRequest.json b/codex-rs/app-server-protocol/schema/json/JSONRPCRequest.json index cc7a16f1b..e4ea7c206 100644 --- a/codex-rs/app-server-protocol/schema/json/JSONRPCRequest.json +++ b/codex-rs/app-server-protocol/schema/json/JSONRPCRequest.json @@ -11,6 +11,23 @@ "type": "integer" } ] + }, + "W3cTraceContext": { + "properties": { + "traceparent": { + "type": [ + "string", + "null" + ] + }, + "tracestate": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" } }, "description": "A request that expects a response.", @@ -21,7 +38,18 @@ "method": { "type": "string" }, - "params": true + "params": true, + "trace": { + "anyOf": [ + { + "$ref": "#/definitions/W3cTraceContext" + }, + { + "type": "null" + } + ], + "description": "Optional W3C Trace Context for distributed tracing." + } }, "required": [ "id", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index edee815e0..275e0dd94 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -5016,7 +5016,18 @@ "method": { "type": "string" }, - "params": true + "params": true, + "trace": { + "anyOf": [ + { + "$ref": "#/definitions/W3cTraceContext" + }, + { + "type": "null" + } + ], + "description": "Optional W3C Trace Context for distributed tracing." + } }, "required": [ "id", @@ -7220,6 +7231,23 @@ } ] }, + "W3cTraceContext": { + "properties": { + "traceparent": { + "type": [ + "string", + "null" + ] + }, + "tracestate": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, "v2": { "AbsolutePathBuf": { "description": "A path that is guaranteed to be absolute and normalized (though it is not guaranteed to be canonicalized or exist on the filesystem).\n\nIMPORTANT: When deserializing an `AbsolutePathBuf`, a base path must be set using [AbsolutePathBufGuard::new]. If no base path is set, the deserialization will fail unless the path being deserialized is already absolute.", diff --git a/codex-rs/app-server-protocol/src/jsonrpc_lite.rs b/codex-rs/app-server-protocol/src/jsonrpc_lite.rs index 372bdbe40..13d3e0cc9 100644 --- a/codex-rs/app-server-protocol/src/jsonrpc_lite.rs +++ b/codex-rs/app-server-protocol/src/jsonrpc_lite.rs @@ -1,6 +1,7 @@ //! We do not do true JSON-RPC 2.0, as we neither send nor expect the //! "jsonrpc": "2.0" field. +use codex_protocol::protocol::W3cTraceContext; use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; @@ -38,6 +39,10 @@ pub struct JSONRPCRequest { #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] pub params: Option, + /// Optional W3C Trace Context for distributed tracing. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub trace: Option, } /// A notification which does not expect a response. diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index eebaaba10..8a4ec05fd 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -21,6 +21,7 @@ async-trait = { workspace = true } codex-arg0 = { workspace = true } codex-cloud-requirements = { workspace = true } codex-core = { workspace = true } +codex-otel = { workspace = true } codex-shell-command = { workspace = true } codex-utils-cli = { workspace = true } codex-backend-client = { workspace = true } diff --git a/codex-rs/app-server/src/app_server_tracing.rs b/codex-rs/app-server/src/app_server_tracing.rs new file mode 100644 index 000000000..779cef967 --- /dev/null +++ b/codex-rs/app-server/src/app_server_tracing.rs @@ -0,0 +1,101 @@ +use crate::message_processor::ConnectionSessionState; +use crate::outgoing_message::ConnectionId; +use crate::transport::AppServerTransport; +use codex_app_server_protocol::InitializeParams; +use codex_app_server_protocol::JSONRPCRequest; +use codex_otel::set_parent_from_context; +use codex_otel::set_parent_from_w3c_trace_context; +use codex_otel::traceparent_context_from_env; +use codex_protocol::protocol::W3cTraceContext; +use tracing::Span; +use tracing::field; +use tracing::info_span; + +pub(crate) fn request_span( + request: &JSONRPCRequest, + transport: AppServerTransport, + connection_id: ConnectionId, + session: &ConnectionSessionState, +) -> Span { + let span = info_span!( + "app_server.request", + otel.kind = "server", + otel.name = request.method.as_str(), + rpc.system = "jsonrpc", + rpc.method = request.method.as_str(), + rpc.transport = transport_name(transport), + rpc.request_id = ?request.id, + app_server.connection_id = ?connection_id, + app_server.api_version = "v2", + app_server.client_name = field::Empty, + app_server.client_version = field::Empty, + ); + + let initialize_client_info = initialize_client_info(request); + if let Some(client_name) = client_name(initialize_client_info.as_ref(), session) { + span.record("app_server.client_name", client_name); + } + if let Some(client_version) = client_version(initialize_client_info.as_ref(), session) { + span.record("app_server.client_version", client_version); + } + + if let Some(traceparent) = request + .trace + .as_ref() + .and_then(|trace| trace.traceparent.as_deref()) + { + let trace = W3cTraceContext { + traceparent: Some(traceparent.to_string()), + tracestate: request + .trace + .as_ref() + .and_then(|value| value.tracestate.clone()), + }; + if !set_parent_from_w3c_trace_context(&span, &trace) { + tracing::warn!( + rpc_method = request.method.as_str(), + rpc_request_id = ?request.id, + "ignoring invalid inbound request trace carrier" + ); + } + } else if let Some(context) = traceparent_context_from_env() { + set_parent_from_context(&span, context); + } + + span +} + +fn transport_name(transport: AppServerTransport) -> &'static str { + match transport { + AppServerTransport::Stdio => "stdio", + AppServerTransport::WebSocket { .. } => "websocket", + } +} + +fn client_name<'a>( + initialize_client_info: Option<&'a InitializeParams>, + session: &'a ConnectionSessionState, +) -> Option<&'a str> { + if let Some(params) = initialize_client_info { + return Some(params.client_info.name.as_str()); + } + session.app_server_client_name.as_deref() +} + +fn client_version<'a>( + initialize_client_info: Option<&'a InitializeParams>, + session: &'a ConnectionSessionState, +) -> Option<&'a str> { + if let Some(params) = initialize_client_info { + return Some(params.client_info.version.as_str()); + } + session.client_version.as_deref() +} + +fn initialize_client_info(request: &JSONRPCRequest) -> Option { + if request.method != "initialize" { + return None; + } + let params = request.params.clone()?; + serde_json::from_value(params).ok() +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index dcc438e29..fec663b48 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -52,6 +52,7 @@ use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::registry::Registry; use tracing_subscriber::util::SubscriberInitExt; +mod app_server_tracing; mod bespoke_event_handling; mod codex_message_processor; mod config_api; @@ -447,7 +448,7 @@ pub async fn run_main_with_transport( let otel = codex_core::otel_init::build_provider( &config, env!("CARGO_PKG_VERSION"), - Some("codex_app_server"), + Some("codex-app-server"), default_analytics_enabled, ) .map_err(|e| { @@ -675,6 +676,7 @@ pub async fn run_main_with_transport( .process_request( connection_id, request, + transport, &mut connection_state.session, &connection_state.outbound_initialized, ) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 2c3034aa1..9c22e56ed 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -12,6 +12,7 @@ use crate::external_agent_config_api::ExternalAgentConfigApi; use crate::outgoing_message::ConnectionId; use crate::outgoing_message::ConnectionRequestId; use crate::outgoing_message::OutgoingMessageSender; +use crate::transport::AppServerTransport; use async_trait::async_trait; use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; @@ -59,6 +60,7 @@ use tokio::sync::watch; use tokio::time::Duration; use tokio::time::timeout; use toml::Value as TomlValue; +use tracing::Instrument; const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); @@ -141,6 +143,7 @@ pub(crate) struct ConnectionSessionState { pub(crate) experimental_api_enabled: bool, pub(crate) opted_out_notification_methods: HashSet, pub(crate) app_server_client_name: Option, + pub(crate) client_version: Option, } pub(crate) struct MessageProcessorArgs { @@ -224,46 +227,50 @@ impl MessageProcessor { &mut self, connection_id: ConnectionId, request: JSONRPCRequest, + transport: AppServerTransport, session: &mut ConnectionSessionState, outbound_initialized: &AtomicBool, ) { - let request_method = request.method.as_str(); - tracing::trace!( - ?connection_id, - request_id = ?request.id, - "app-server request: {request_method}" - ); - let request_id = ConnectionRequestId { - connection_id, - request_id: request.id.clone(), - }; - let request_json = match serde_json::to_value(&request) { - Ok(request_json) => request_json, - Err(err) => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("Invalid request: {err}"), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - }; + let request_span = + crate::app_server_tracing::request_span(&request, transport, connection_id, session); + async { + let request_method = request.method.as_str(); + tracing::trace!( + ?connection_id, + request_id = ?request.id, + "app-server request: {request_method}" + ); + let request_id = ConnectionRequestId { + connection_id, + request_id: request.id.clone(), + }; + let request_json = match serde_json::to_value(&request) { + Ok(request_json) => request_json, + Err(err) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("Invalid request: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + }; - let codex_request = match serde_json::from_value::(request_json) { - Ok(codex_request) => codex_request, - Err(err) => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("Invalid request: {err}"), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - }; + let codex_request = match serde_json::from_value::(request_json) { + Ok(codex_request) => codex_request, + Err(err) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("Invalid request: {err}"), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } + }; - match codex_request { + match codex_request { // Handle Initialize internally so CodexMessageProcessor does not have to concern // itself with the `initialized` bool. ClientRequest::Initialize { request_id, params } => { @@ -304,6 +311,8 @@ impl MessageProcessor { title: _title, version, } = params.client_info; + session.app_server_client_name = Some(name.clone()); + session.client_version = Some(version.clone()); if let Err(error) = set_default_originator(name.clone()) { match error { SetOriginatorError::InvalidHeaderValue => { @@ -330,7 +339,6 @@ impl MessageProcessor { if let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() { *suffix = Some(user_agent_suffix); } - session.app_server_client_name = Some(name.clone()); let user_agent = get_codex_user_agent(); let response = InitializeResponse { user_agent }; @@ -355,91 +363,97 @@ impl MessageProcessor { return; } } - } + } + if let Some(reason) = codex_request.experimental_reason() + && !session.experimental_api_enabled + { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: experimental_required_message(reason), + data: None, + }; + self.outgoing.send_error(request_id, error).await; + return; + } - if let Some(reason) = codex_request.experimental_reason() - && !session.experimental_api_enabled - { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: experimental_required_message(reason), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - - match codex_request { - ClientRequest::ConfigRead { request_id, params } => { - self.handle_config_read( - ConnectionRequestId { - connection_id, - request_id, - }, - params, - ) - .await; - } - ClientRequest::ExternalAgentConfigDetect { request_id, params } => { - self.handle_external_agent_config_detect( - ConnectionRequestId { - connection_id, - request_id, - }, - params, - ) - .await; - } - ClientRequest::ExternalAgentConfigImport { request_id, params } => { - self.handle_external_agent_config_import( - ConnectionRequestId { - connection_id, - request_id, - }, - params, - ) - .await; - } - ClientRequest::ConfigValueWrite { request_id, params } => { - self.handle_config_value_write( - ConnectionRequestId { - connection_id, - request_id, - }, - params, - ) - .await; - } - ClientRequest::ConfigBatchWrite { request_id, params } => { - self.handle_config_batch_write( - ConnectionRequestId { - connection_id, - request_id, - }, - params, - ) - .await; - } - ClientRequest::ConfigRequirementsRead { - request_id, - params: _, - } => { - self.handle_config_requirements_read(ConnectionRequestId { - connection_id, - request_id, - }) - .await; - } - other => { - // Box the delegated future so this wrapper's async state machine does not - // inline the full `CodexMessageProcessor::process_request` future, which - // can otherwise push worker-thread stack usage over the edge. - self.codex_message_processor - .process_request(connection_id, other, session.app_server_client_name.clone()) - .boxed() + match codex_request { + ClientRequest::ConfigRead { request_id, params } => { + self.handle_config_read( + ConnectionRequestId { + connection_id, + request_id, + }, + params, + ) .await; + } + ClientRequest::ExternalAgentConfigDetect { request_id, params } => { + self.handle_external_agent_config_detect( + ConnectionRequestId { + connection_id, + request_id, + }, + params, + ) + .await; + } + ClientRequest::ExternalAgentConfigImport { request_id, params } => { + self.handle_external_agent_config_import( + ConnectionRequestId { + connection_id, + request_id, + }, + params, + ) + .await; + } + ClientRequest::ConfigValueWrite { request_id, params } => { + self.handle_config_value_write( + ConnectionRequestId { + connection_id, + request_id, + }, + params, + ) + .await; + } + ClientRequest::ConfigBatchWrite { request_id, params } => { + self.handle_config_batch_write( + ConnectionRequestId { + connection_id, + request_id, + }, + params, + ) + .await; + } + ClientRequest::ConfigRequirementsRead { + request_id, + params: _, + } => { + self.handle_config_requirements_read(ConnectionRequestId { + connection_id, + request_id, + }) + .await; + } + other => { + // Box the delegated future so this wrapper's async state machine does not + // inline the full `CodexMessageProcessor::process_request` future, which + // can otherwise push worker-thread stack usage over the edge. + self.codex_message_processor + .process_request( + connection_id, + other, + session.app_server_client_name.clone(), + ) + .boxed() + .await; + } } } + .instrument(request_span) + .await; } pub(crate) async fn process_notification(&self, notification: JSONRPCNotification) { diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 2bae5b44e..019d00e45 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -744,6 +744,7 @@ mod tests { id: codex_app_server_protocol::RequestId::Integer(7), method: "config/read".to_string(), params: Some(json!({ "includeLayers": false })), + trace: None, }); assert!( enqueue_incoming_message(&transport_event_tx, &writer_tx, connection_id, request).await @@ -885,6 +886,7 @@ mod tests { id: codex_app_server_protocol::RequestId::Integer(7), method: "config/read".to_string(), params: Some(json!({ "includeLayers": false })), + trace: None, }); let enqueue_result = tokio::time::timeout( diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index cb0c98806..9f0fb31cb 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -891,6 +891,7 @@ impl McpProcess { id: RequestId::Integer(request_id), method: method.to_string(), params, + trace: None, }); self.send_jsonrpc_message(message).await?; Ok(request_id) diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index 0d05d6446..e64287cb2 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -30,7 +30,7 @@ async fn app_server_default_analytics_disabled_without_flag() -> Result<()> { let provider = codex_core::otel_init::build_provider( &config, SERVICE_VERSION, - Some("codex_app_server"), + Some("codex-app-server"), false, ) .map_err(|err| anyhow::anyhow!(err.to_string()))?; @@ -55,7 +55,7 @@ async fn app_server_default_analytics_enabled_with_flag() -> Result<()> { let provider = codex_core::otel_init::build_provider( &config, SERVICE_VERSION, - Some("codex_app_server"), + Some("codex-app-server"), true, ) .map_err(|err| anyhow::anyhow!(err.to_string()))?; diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index 0bbbc70b5..979994538 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -174,6 +174,7 @@ pub(super) async fn send_request( id: RequestId::Integer(id), method: method.to_string(), params, + trace: None, }); send_jsonrpc(stream, message).await } diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index 5aaee5f5a..217099524 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -1,6 +1,7 @@ pub mod config; pub mod metrics; pub mod otel_provider; +pub mod trace_context; pub mod traces; mod otlp; @@ -23,6 +24,11 @@ use tracing::debug; pub use crate::metrics::runtime_metrics::RuntimeMetricTotals; pub use crate::metrics::runtime_metrics::RuntimeMetricsSummary; +pub use crate::otel_provider::traceparent_context_from_env; +pub use crate::trace_context::context_from_w3c_trace_context; +pub use crate::trace_context::current_span_w3c_trace_context; +pub use crate::trace_context::set_parent_from_context; +pub use crate::trace_context::set_parent_from_w3c_trace_context; #[derive(Debug, Clone, Serialize, Display)] #[serde(rename_all = "snake_case")] diff --git a/codex-rs/otel/src/otel_provider.rs b/codex-rs/otel/src/otel_provider.rs index 6f2dc5a09..c9233f5d6 100644 --- a/codex-rs/otel/src/otel_provider.rs +++ b/codex-rs/otel/src/otel_provider.rs @@ -3,13 +3,12 @@ use crate::config::OtelHttpProtocol; use crate::config::OtelSettings; use crate::metrics::MetricsClient; use crate::metrics::MetricsConfig; +use crate::trace_context::context_from_trace_headers; use gethostname::gethostname; use opentelemetry::Context; use opentelemetry::KeyValue; use opentelemetry::context::ContextGuard; use opentelemetry::global; -use opentelemetry::propagation::TextMapPropagator; -use opentelemetry::trace::TraceContextExt; use opentelemetry::trace::TracerProvider as _; use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge; use opentelemetry_otlp::LogExporter; @@ -30,7 +29,6 @@ use opentelemetry_sdk::trace::SdkTracerProvider; use opentelemetry_sdk::trace::Tracer; use opentelemetry_semantic_conventions as semconv; use std::cell::RefCell; -use std::collections::HashMap; use std::env; use std::error::Error; use std::sync::OnceLock; @@ -172,7 +170,7 @@ impl Drop for OtelProvider { } } -pub(crate) fn traceparent_context_from_env() -> Option { +pub fn traceparent_context_from_env() -> Option { TRACEPARENT_CONTEXT .get_or_init(load_traceparent_context) .clone() @@ -194,7 +192,7 @@ fn load_traceparent_context() -> Option { let traceparent = env::var(TRACEPARENT_ENV_VAR).ok()?; let tracestate = env::var(TRACESTATE_ENV_VAR).ok(); - match extract_traceparent_context(traceparent, tracestate) { + match context_from_trace_headers(Some(&traceparent), tracestate.as_deref()) { Some(context) => { debug!("TRACEPARENT detected; continuing trace from parent context"); Some(context) @@ -206,22 +204,6 @@ fn load_traceparent_context() -> Option { } } -fn extract_traceparent_context(traceparent: String, tracestate: Option) -> Option { - let mut headers = HashMap::new(); - headers.insert("traceparent".to_string(), traceparent); - if let Some(tracestate) = tracestate { - headers.insert("tracestate".to_string(), tracestate); - } - - let context = TraceContextPropagator::new().extract(&headers); - let span = context.span(); - let span_context = span.span_context(); - if !span_context.is_valid() { - return None; - } - Some(context) -} - fn make_resource(settings: &OtelSettings) -> Resource { Resource::builder() .with_service_name(settings.service_name.clone()) @@ -407,8 +389,9 @@ mod tests { fn parses_valid_traceparent() { let trace_id = "00000000000000000000000000000001"; let span_id = "0000000000000002"; - let context = extract_traceparent_context(format!("00-{trace_id}-{span_id}-01"), None) - .expect("trace context"); + let context = + context_from_trace_headers(Some(&format!("00-{trace_id}-{span_id}-01")), None) + .expect("trace context"); let span = context.span(); let span_context = span.span_context(); assert_eq!( @@ -421,7 +404,7 @@ mod tests { #[test] fn invalid_traceparent_returns_none() { - assert!(extract_traceparent_context("not-a-traceparent".to_string(), None).is_none()); + assert!(context_from_trace_headers(Some("not-a-traceparent"), None).is_none()); } #[test] diff --git a/codex-rs/otel/src/trace_context.rs b/codex-rs/otel/src/trace_context.rs new file mode 100644 index 000000000..3ef38062d --- /dev/null +++ b/codex-rs/otel/src/trace_context.rs @@ -0,0 +1,106 @@ +use std::collections::HashMap; + +use codex_protocol::protocol::W3cTraceContext; +use opentelemetry::Context; +use opentelemetry::propagation::TextMapPropagator; +use opentelemetry::trace::TraceContextExt; +use opentelemetry_sdk::propagation::TraceContextPropagator; +use tracing::Span; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +pub fn current_span_w3c_trace_context() -> Option { + let context = Span::current().context(); + if !context.span().span_context().is_valid() { + return None; + } + + let mut headers = HashMap::new(); + TraceContextPropagator::new().inject_context(&context, &mut headers); + + Some(W3cTraceContext { + traceparent: headers.remove("traceparent"), + tracestate: headers.remove("tracestate"), + }) +} + +pub fn context_from_w3c_trace_context(trace: &W3cTraceContext) -> Option { + context_from_trace_headers(trace.traceparent.as_deref(), trace.tracestate.as_deref()) +} + +pub fn set_parent_from_w3c_trace_context(span: &Span, trace: &W3cTraceContext) -> bool { + if let Some(context) = context_from_w3c_trace_context(trace) { + set_parent_from_context(span, context); + true + } else { + false + } +} + +pub fn set_parent_from_context(span: &Span, context: Context) { + let _ = span.set_parent(context); +} + +pub(crate) fn context_from_trace_headers( + traceparent: Option<&str>, + tracestate: Option<&str>, +) -> Option { + let traceparent = traceparent?; + let mut headers = HashMap::new(); + headers.insert("traceparent".to_string(), traceparent.to_string()); + if let Some(tracestate) = tracestate { + headers.insert("tracestate".to_string(), tracestate.to_string()); + } + + let context = TraceContextPropagator::new().extract(&headers); + if !context.span().span_context().is_valid() { + return None; + } + Some(context) +} + +#[cfg(test)] +mod tests { + use super::context_from_trace_headers; + use super::context_from_w3c_trace_context; + use codex_protocol::protocol::W3cTraceContext; + use opentelemetry::trace::SpanId; + use opentelemetry::trace::TraceContextExt; + use opentelemetry::trace::TraceId; + use pretty_assertions::assert_eq; + + #[test] + fn parses_valid_w3c_trace_context() { + let trace_id = "00000000000000000000000000000001"; + let span_id = "0000000000000002"; + let context = context_from_w3c_trace_context(&W3cTraceContext { + traceparent: Some(format!("00-{trace_id}-{span_id}-01")), + tracestate: None, + }) + .expect("trace context"); + + let span = context.span(); + let span_context = span.span_context(); + assert_eq!( + span_context.trace_id(), + TraceId::from_hex(trace_id).unwrap() + ); + assert_eq!(span_context.span_id(), SpanId::from_hex(span_id).unwrap()); + assert!(span_context.is_remote()); + } + + #[test] + fn invalid_traceparent_returns_none() { + assert!(context_from_trace_headers(Some("not-a-traceparent"), None).is_none()); + } + + #[test] + fn missing_traceparent_returns_none() { + assert!( + context_from_w3c_trace_context(&W3cTraceContext { + traceparent: None, + tracestate: Some("vendor=value".to_string()), + }) + .is_none() + ); + } +} diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 3e9e9fe9a..7e10951cd 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -83,6 +83,16 @@ pub struct Submission { pub op: Op, } +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +pub struct W3cTraceContext { + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub traceparent: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub tracestate: Option, +} + /// Config payload for refreshing MCP servers. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)] pub struct McpServerRefreshConfig {