mirror of
https://github.com/openai/codex.git
synced 2026-05-02 18:37:01 +00:00
d473e8d56d
### 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
107 lines
3.3 KiB
Rust
107 lines
3.3 KiB
Rust
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()
|
|
);
|
|
}
|
|
}
|