From 171bbc836442a8d5cbec56994ceded37dc85a8e3 Mon Sep 17 00:00:00 2001 From: nicholasclark-openai Date: Wed, 25 Mar 2026 10:50:59 -0700 Subject: [PATCH] Narrow tracing branch to MCP spans Drop the Responses-specific tracing changes from this branch so the draft PR focuses on MCP spans and MCP transport timing. Co-authored-by: Codex --- codex-rs/codex-client/src/default_client.rs | 79 ++---------- codex-rs/core/src/client.rs | 134 +++++++++----------- codex-rs/core/src/network_trace.rs | 52 -------- codex-rs/core/src/tasks/mod.rs | 5 +- codex-rs/core/tests/suite/otel.rs | 51 -------- 5 files changed, 69 insertions(+), 252 deletions(-) diff --git a/codex-rs/codex-client/src/default_client.rs b/codex-rs/codex-client/src/default_client.rs index 31ea3458e3..56b3ce4b16 100644 --- a/codex-rs/codex-client/src/default_client.rs +++ b/codex-rs/codex-client/src/default_client.rs @@ -10,9 +10,7 @@ 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)] @@ -113,39 +111,10 @@ impl CodexRequestBuilder { } pub async fn send(self) -> Result { - 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); + let headers = trace_headers(); - match async { self.builder.headers(headers).send().await } - .instrument(request_span.clone()) - .await - { + match self.builder.headers(headers).send().await { Ok(response) => { - request_span.record( - "http.response.status_code", - response.status().as_u16() as i64, - ); tracing::debug!( method = %self.method, url = %self.url, @@ -158,14 +127,11 @@ impl CodexRequestBuilder { Ok(response) } Err(error) => { - let status = error.status().map(|status| status.as_u16() as i64); - if let Some(status) = status { - request_span.record("http.response.status_code", status); - } + let status = error.status(); tracing::debug!( method = %self.method, url = %self.url, - status, + status = status.map(|s| s.as_u16()), error = %error, "Request failed" ); @@ -188,15 +154,13 @@ 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.context(), &mut HeaderMapInjector(&mut headers)); + prop.inject_context( + &Span::current().context(), + &mut HeaderMapInjector(&mut headers), + ); }); headers } @@ -240,33 +204,6 @@ 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/core/src/client.rs b/codex-rs/core/src/client.rs index d879cd97cb..e38ff3a562 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -86,7 +86,6 @@ use tokio::sync::oneshot; use tokio::sync::oneshot::error::TryRecvError; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::Message; -use tracing::Instrument; use tracing::instrument; use tracing::trace; use tracing::warn; @@ -1024,7 +1023,6 @@ impl ModelClientSession { let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; - let api_base_url = client_setup.api_provider.base_url.clone(); let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), @@ -1054,20 +1052,10 @@ impl ModelClientSession { client_setup.api_auth, ) .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); - let request_span = crate::network_trace::responses_http_request_span( - &self.client.state.conversation_id, - turn_metadata_header, - &self.client.state.provider.name, - &model_info.slug, - &api_base_url, - ); - let stream_result = async { client.stream_request(request, options).await } - .instrument(request_span.clone()) - .await; + let stream_result = client.stream_request(request, options).await; match stream_result { Ok(stream) => { - let _entered = request_span.enter(); let (stream, _) = map_response_stream(stream, session_telemetry.clone()); return Ok(stream); } @@ -1426,76 +1414,72 @@ where { let (tx_event, rx_event) = mpsc::channel::>(1600); let (tx_last_response, rx_last_response) = oneshot::channel::(); - let current_span = tracing::Span::current(); - tokio::spawn( - async move { - let mut logged_error = false; - let mut tx_last_response = Some(tx_last_response); - let mut items_added: Vec = Vec::new(); - let mut api_stream = api_stream; - while let Some(event) = api_stream.next().await { - match event { - Ok(ResponseEvent::OutputItemDone(item)) => { - items_added.push(item.clone()); - if tx_event - .send(Ok(ResponseEvent::OutputItemDone(item))) - .await - .is_err() - { - return; - } + tokio::spawn(async move { + let mut logged_error = false; + let mut tx_last_response = Some(tx_last_response); + let mut items_added: Vec = Vec::new(); + let mut api_stream = api_stream; + while let Some(event) = api_stream.next().await { + match event { + Ok(ResponseEvent::OutputItemDone(item)) => { + items_added.push(item.clone()); + if tx_event + .send(Ok(ResponseEvent::OutputItemDone(item))) + .await + .is_err() + { + return; } - Ok(ResponseEvent::Completed { - response_id, - token_usage, - }) => { - if let Some(usage) = &token_usage { - session_telemetry.sse_event_completed( - usage.input_tokens, - usage.output_tokens, - Some(usage.cached_input_tokens), - Some(usage.reasoning_output_tokens), - usage.total_tokens, - ); - } - if let Some(sender) = tx_last_response.take() { - let _ = sender.send(LastResponse { - response_id: response_id.clone(), - items_added: std::mem::take(&mut items_added), - }); - } - if tx_event - .send(Ok(ResponseEvent::Completed { - response_id, - token_usage, - })) - .await - .is_err() - { - return; - } + } + Ok(ResponseEvent::Completed { + response_id, + token_usage, + }) => { + if let Some(usage) = &token_usage { + session_telemetry.sse_event_completed( + usage.input_tokens, + usage.output_tokens, + Some(usage.cached_input_tokens), + Some(usage.reasoning_output_tokens), + usage.total_tokens, + ); } - Ok(event) => { - if tx_event.send(Ok(event)).await.is_err() { - return; - } + if let Some(sender) = tx_last_response.take() { + let _ = sender.send(LastResponse { + response_id: response_id.clone(), + items_added: std::mem::take(&mut items_added), + }); } - Err(err) => { - let mapped = map_api_error(err); - if !logged_error { - session_telemetry.see_event_completed_failed(&mapped); - logged_error = true; - } - if tx_event.send(Err(mapped)).await.is_err() { - return; - } + if tx_event + .send(Ok(ResponseEvent::Completed { + response_id, + token_usage, + })) + .await + .is_err() + { + return; + } + } + Ok(event) => { + if tx_event.send(Ok(event)).await.is_err() { + return; + } + } + Err(err) => { + let mapped = map_api_error(err); + if !logged_error { + session_telemetry.see_event_completed_failed(&mapped); + logged_error = true; + } + if tx_event.send(Err(mapped)).await.is_err() { + return; } } } } - .instrument(current_span), - ); + }); (ResponseStream { rx_event }, rx_last_response) } diff --git a/codex-rs/core/src/network_trace.rs b/codex-rs/core/src/network_trace.rs index e243d1192c..cdfa15075c 100644 --- a/codex-rs/core/src/network_trace.rs +++ b/codex-rs/core/src/network_trace.rs @@ -1,5 +1,3 @@ -use codex_protocol::ThreadId; -use serde::Deserialize; use tracing::Span; use tracing::field::Empty; use url::Url; @@ -7,12 +5,6 @@ use url::Url; use crate::codex::Session; use crate::codex::TurnContext; -#[derive(Debug, Default, Deserialize)] -struct TurnTraceCorrelation { - session_id: Option, - turn_id: Option, -} - struct CorrelationFields { conversation_id: String, session_id: String, @@ -20,24 +12,6 @@ struct CorrelationFields { } impl CorrelationFields { - fn from_turn_metadata_header( - conversation_id: &ThreadId, - turn_metadata_header: Option<&str>, - ) -> Self { - let conversation_id = conversation_id.to_string(); - let correlation = turn_metadata_header - .and_then(|header| serde_json::from_str::(header).ok()) - .unwrap_or_default(); - let session_id = correlation - .session_id - .unwrap_or_else(|| conversation_id.clone()); - Self { - conversation_id, - session_id, - turn_id: correlation.turn_id, - } - } - fn from_turn_context(session: &Session, turn_context: &TurnContext) -> Self { Self { conversation_id: session.conversation_id.to_string(), @@ -70,32 +44,6 @@ fn record_server_fields(span: &Span, url: Option<&str>) { } } -pub(crate) fn responses_http_request_span( - conversation_id: &ThreadId, - turn_metadata_header: Option<&str>, - provider_name: &str, - model: &str, - base_url: &str, -) -> Span { - let span = tracing::info_span!( - "responses_http.request", - otel.kind = "client", - provider = provider_name, - model, - transport = "responses_http", - api.path = "responses", - conversation.id = Empty, - session.id = Empty, - turn.id = Empty, - server.address = Empty, - server.port = Empty, - ); - CorrelationFields::from_turn_metadata_header(conversation_id, turn_metadata_header) - .record_on(&span); - record_server_fields(&span, Some(base_url)); - span -} - pub(crate) struct McpToolCallTrace<'a> { pub server_name: &'a str, pub tool_name: &'a str, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 58b0752dd6..601f5cb27f 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -197,12 +197,11 @@ impl Session { let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); let task_cancellation_token = cancellation_token.child_token(); - // Task-owned turn spans keep a core-owned parent span open for the - // full turn lifecycle after the submission dispatch span ends. + // Task-owned turn spans keep a core-owned span open for the + // full task lifecycle after the submission dispatch span ends. let task_span = info_span!( "turn", otel.name = span_name, - conversation.id = %self.conversation_id, thread.id = %self.conversation_id, turn.id = %turn_context.sub_id, model = %turn_context.model_info.slug, diff --git a/codex-rs/core/tests/suite/otel.rs b/codex-rs/core/tests/suite/otel.rs index b66d34f8b9..96df3fd2f1 100644 --- a/codex-rs/core/tests/suite/otel.rs +++ b/codex-rs/core/tests/suite/otel.rs @@ -718,57 +718,6 @@ async fn record_responses_sets_span_fields_for_response_events() { } } -#[tokio::test(flavor = "current_thread")] -#[traced_test] -async fn responses_request_span_records_turn_correlation_fields() { - let server = start_mock_server().await; - mount_sse_once( - &server, - sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), - ) - .await; - - let TestCodex { codex, .. } = test_codex().build(&server).await.unwrap(); - - codex - .submit(Op::UserInput { - items: vec![UserInput::Text { - text: "hello".into(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - }) - .await - .unwrap(); - - wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; - - logs_assert(|lines: &[&str]| { - lines - .iter() - .find(|line| { - line.contains("turn{otel.name=\"session_task.turn\"") - && line.contains("conversation.id=") - && line.contains("thread.id=") - && line.contains("turn.id=") - && line.contains("model=") - && line.contains("responses_http.request{") - && line.contains("otel.kind=\"client\"") - && line.contains("transport=\"responses_http\"") - && line.contains("conversation.id=") - && line.contains("session.id=") - && line.contains("turn.id=") - }) - .map(|_| Ok(())) - .unwrap_or_else(|| { - Err( - "missing responses_http.request span nested under session_task.turn" - .to_string(), - ) - }) - }); -} - #[tokio::test] #[traced_test] async fn handle_response_item_records_tool_result_for_custom_tool_call() {