agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(app-server): add tracing to all app-server APIs (#13285)

Browse Source 
### 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
This commit is contained in:
Owen Lin
2026-03-02 16:01:41 -08:00
committed by GitHub
Unverified
parent 14fcb6645c
commit d473e8d56d
17 changed files with 464 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -1428,6 +1428,7 @@ dependencies = [
"codex-feedback",
"codex-file-search",
"codex-login",
"codex-otel",
"codex-protocol",
"codex-rmcp-client",
"codex-shell-command",
codex-rs/app-server-protocol/schema/json/JSONRPCMessage.json
+29 -1
View File
@@ -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.",
codex-rs/app-server-protocol/schema/json/JSONRPCRequest.json
+29 -1
View File
@@ -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",
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json
+29 -1
View File
@@ -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.",
codex-rs/app-server-protocol/src/jsonrpc_lite.rs
+5
View File
@@ -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<serde_json::Value>,
/// Optional W3C Trace Context for distributed tracing.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub trace: Option<W3cTraceContext>,
}
/// A notification which does not expect a response.
codex-rs/app-server/Cargo.toml
+1
View File
@@ -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 }
codex-rs/app-server/src/app_server_tracing.rs
+101
View File
@@ -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<InitializeParams> {
if request.method != "initialize" {
return None;
}
let params = request.params.clone()?;
serde_json::from_value(params).ok()
}
codex-rs/app-server/src/lib.rs
+3 -1
View File
@@ -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,
)
codex-rs/app-server/src/message_processor.rs
+131 -117
View File
@@ -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<String>,
pub(crate) app_server_client_name: Option<String>,
pub(crate) client_version: Option<String>,
}
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::<ClientRequest>(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::<ClientRequest>(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) {
codex-rs/app-server/src/transport.rs
+2
View File
@@ -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(
codex-rs/app-server/tests/common/mcp_process.rs
+1
View File
@@ -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)
codex-rs/app-server/tests/suite/v2/analytics.rs
+2 -2
View File
@@ -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()))?;
codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs
+1
View File
@@ -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
}
codex-rs/otel/src/lib.rs
+6
View File
@@ -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")]
codex-rs/otel/src/otel_provider.rs
+7 -24
View File
@@ -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<Context> {
pub fn traceparent_context_from_env() -> Option<Context> {
TRACEPARENT_CONTEXT
.get_or_init(load_traceparent_context)
.clone()
@@ -194,7 +192,7 @@ fn load_traceparent_context() -> Option<Context> {
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<Context> {
}
}
fn extract_traceparent_context(traceparent: String, tracestate: Option<String>) -> Option<Context> {
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]
codex-rs/otel/src/trace_context.rs
+106
View File
@@ -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<W3cTraceContext> {
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> {
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<Context> {
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()
);
}
}
codex-rs/protocol/src/protocol.rs
+10
View File
@@ -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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub tracestate: Option<String>,
}
/// Config payload for refreshing MCP servers.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
pub struct McpServerRefreshConfig {