From 3601d9387a4b3ccbd8284fec48e7e02faa4a417f Mon Sep 17 00:00:00 2001 From: nicholasclark-openai Date: Tue, 24 Mar 2026 18:42:55 -0700 Subject: [PATCH] Add lower-level HTTP and MCP child spans Co-authored-by: Codex --- codex-rs/codex-client/src/default_client.rs | 79 ++++++++++++++++++-- codex-rs/rmcp-client/src/rmcp_client.rs | 82 +++++++++++++++++---- 2 files changed, 139 insertions(+), 22 deletions(-) diff --git a/codex-rs/codex-client/src/default_client.rs b/codex-rs/codex-client/src/default_client.rs index 56b3ce4b16..31ea3458e3 100644 --- a/codex-rs/codex-client/src/default_client.rs +++ b/codex-rs/codex-client/src/default_client.rs @@ -10,7 +10,9 @@ use reqwest::Response; use serde::Serialize; use std::fmt::Display; use std::time::Duration; +use tracing::Instrument; use tracing::Span; +use tracing::field::Empty; use tracing_opentelemetry::OpenTelemetrySpanExt; #[derive(Clone, Debug)] @@ -111,10 +113,39 @@ impl CodexRequestBuilder { } pub async fn send(self) -> Result { - let headers = trace_headers(); + let parsed_url = reqwest::Url::parse(&self.url).ok(); + let path = parsed_url + .as_ref() + .map(|url| url.path().to_string()) + .unwrap_or_else(|| self.url.clone()); + let request_span = tracing::info_span!( + "http.client", + otel.kind = "client", + http.request.method = %self.method, + http.response.status_code = Empty, + url.path = %path, + server.address = Empty, + server.port = Empty, + ); + if let Some(url) = parsed_url.as_ref() { + if let Some(host) = url.host_str() { + request_span.record("server.address", host); + } + if let Some(port) = url.port_or_known_default() { + request_span.record("server.port", port as i64); + } + } + let headers = trace_headers_for_span(&request_span); - match self.builder.headers(headers).send().await { + match async { self.builder.headers(headers).send().await } + .instrument(request_span.clone()) + .await + { Ok(response) => { + request_span.record( + "http.response.status_code", + response.status().as_u16() as i64, + ); tracing::debug!( method = %self.method, url = %self.url, @@ -127,11 +158,14 @@ impl CodexRequestBuilder { Ok(response) } Err(error) => { - let status = error.status(); + let status = error.status().map(|status| status.as_u16() as i64); + if let Some(status) = status { + request_span.record("http.response.status_code", status); + } tracing::debug!( method = %self.method, url = %self.url, - status = status.map(|s| s.as_u16()), + status, error = %error, "Request failed" ); @@ -154,13 +188,15 @@ impl<'a> Injector for HeaderMapInjector<'a> { } } +#[cfg(test)] fn trace_headers() -> HeaderMap { + trace_headers_for_span(&Span::current()) +} + +fn trace_headers_for_span(span: &Span) -> HeaderMap { let mut headers = HeaderMap::new(); global::get_text_map_propagator(|prop| { - prop.inject_context( - &Span::current().context(), - &mut HeaderMapInjector(&mut headers), - ); + prop.inject_context(&span.context(), &mut HeaderMapInjector(&mut headers)); }); headers } @@ -204,6 +240,33 @@ mod tests { assert_eq!(extracted_context.span_id(), span_context.span_id()); } + #[test] + fn inject_trace_headers_for_span_uses_explicit_span_context() { + global::set_text_map_propagator(TraceContextPropagator::new()); + + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("test-tracer"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + + let parent = trace_span!("parent"); + let _parent_entered = parent.enter(); + let child = trace_span!("child"); + let child_context = child.context().span().span_context().clone(); + + let headers = trace_headers_for_span(&child); + + let extractor = HeaderMapExtractor(&headers); + let extracted = TraceContextPropagator::new().extract(&extractor); + let extracted_span = extracted.span(); + let extracted_context = extracted_span.span_context(); + + assert!(extracted_context.is_valid()); + assert_eq!(extracted_context.trace_id(), child_context.trace_id()); + assert_eq!(extracted_context.span_id(), child_context.span_id()); + } + struct HeaderMapExtractor<'a>(&'a HeaderMap); impl<'a> Extractor for HeaderMapExtractor<'a> { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 55a3603ed7..f13472a21b 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -64,6 +64,8 @@ use tokio::io::BufReader; use tokio::process::Command; use tokio::sync::Mutex; use tokio::time; +use tracing::Instrument; +use tracing::field::Empty; use tracing::info; use tracing::warn; @@ -1052,41 +1054,93 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; - match Self::run_service_operation_once(Arc::clone(&service), label, timeout, &operation) - .await + match Self::run_service_operation_once( + Arc::clone(&service), + label, + timeout, + self.service_operation_span(label), + &operation, + ) + .await { Ok(result) => Ok(result), Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; - Self::run_service_operation_once(recovered_service, label, timeout, &operation) - .await - .map_err(Into::into) + Self::run_service_operation_once( + recovered_service, + label, + timeout, + self.service_operation_span(label), + &operation, + ) + .await + .map_err(Into::into) } Err(error) => Err(error.into()), } } + fn service_operation_span(&self, label: &str) -> tracing::Span { + let span = tracing::info_span!( + "mcp.client.operation", + otel.kind = "client", + rpc.system = "jsonrpc", + rpc.method = label, + mcp.transport = Empty, + mcp.server.name = Empty, + server.address = Empty, + server.port = Empty, + ); + + match &self.transport_recipe { + TransportRecipe::Stdio { .. } => { + span.record("mcp.transport", "stdio"); + } + TransportRecipe::StreamableHttp { + server_name, url, .. + } => { + span.record("mcp.transport", "streamable_http"); + span.record("mcp.server.name", server_name.as_str()); + if let Ok(parsed_url) = reqwest::Url::parse(url) { + if let Some(host) = parsed_url.host_str() { + span.record("server.address", host); + } + if let Some(port) = parsed_url.port_or_known_default() { + span.record("server.port", port as i64); + } + } + } + } + + span + } + async fn run_service_operation_once( service: Arc>, label: &str, timeout: Option, + operation_span: tracing::Span, operation: &F, ) -> std::result::Result where F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - match timeout { - Some(duration) => time::timeout(duration, operation(service)) - .await - .map_err(|_| ClientOperationError::Timeout { - label: label.to_string(), - duration, - })? - .map_err(ClientOperationError::from), - None => operation(service).await.map_err(ClientOperationError::from), + async move { + match timeout { + Some(duration) => time::timeout(duration, operation(service)) + .await + .map_err(|_| ClientOperationError::Timeout { + label: label.to_string(), + duration, + })? + .map_err(ClientOperationError::from), + None => operation(service).await.map_err(ClientOperationError::from), + } } + .instrument(operation_span) + .await } fn is_session_expired_404(error: &ClientOperationError) -> bool {