From 4322d47f7444f7c6e1ae5167c985ff077934022e Mon Sep 17 00:00:00 2001 From: nicholasclark-openai Date: Mon, 23 Mar 2026 13:22:38 -0700 Subject: [PATCH] Add outbound spans for responses and MCP Co-authored-by: Codex --- codex-rs/codex-client/src/default_client.rs | 53 ++- codex-rs/codex-client/src/lib.rs | 1 + codex-rs/core/src/client.rs | 276 +++++++++------ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/mcp_tool_call.rs | 114 ++++-- codex-rs/core/src/network_trace.rs | 187 ++++++++++ codex-rs/core/tests/responses_headers.rs | 50 +++ .../core/tests/suite/client_websockets.rs | 25 +- codex-rs/core/tests/suite/otel.rs | 169 +++++++++ codex-rs/core/tests/suite/rmcp_client.rs | 331 ++++++++++++++++++ .../rmcp-client/src/bin/test_stdio_server.rs | 51 +++ .../src/bin/test_streamable_http_server.rs | 86 ++++- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/rmcp_client.rs | 13 +- codex-rs/rmcp-client/src/trace.rs | 36 ++ 15 files changed, 1246 insertions(+), 148 deletions(-) create mode 100644 codex-rs/core/src/network_trace.rs create mode 100644 codex-rs/rmcp-client/src/trace.rs diff --git a/codex-rs/codex-client/src/default_client.rs b/codex-rs/codex-client/src/default_client.rs index 56b3ce4b16..5aff24720e 100644 --- a/codex-rs/codex-client/src/default_client.rs +++ b/codex-rs/codex-client/src/default_client.rs @@ -11,6 +11,7 @@ use serde::Serialize; use std::fmt::Display; use std::time::Duration; use tracing::Span; +use tracing::field::Empty; use tracing_opentelemetry::OpenTelemetrySpanExt; #[derive(Clone, Debug)] @@ -37,6 +38,13 @@ impl CodexHttpClient { self.request(Method::POST, url) } + pub fn delete(&self, url: U) -> CodexRequestBuilder + where + U: IntoUrl, + { + self.request(Method::DELETE, url) + } + pub fn request(&self, method: Method, url: U) -> CodexRequestBuilder where U: IntoUrl, @@ -111,10 +119,26 @@ impl CodexRequestBuilder { } pub async fn send(self) -> Result { - let headers = trace_headers(); + let span = tracing::info_span!( + "http.client", + otel.kind = "client", + http.request.method = %self.method, + url.full = %self.url, + server.address = Empty, + server.port = Empty, + http.response.status_code = Empty, + error.type = Empty, + ); + record_server_fields(&span, &self.url); + let _entered = span.enter(); + let headers = current_trace_headers(); match self.builder.headers(headers).send().await { Ok(response) => { + span.record( + "http.response.status_code", + response.status().as_u16() as i64, + ); tracing::debug!( method = %self.method, url = %self.url, @@ -128,6 +152,17 @@ impl CodexRequestBuilder { } Err(error) => { let status = error.status(); + if let Some(status) = status { + span.record("http.response.status_code", status.as_u16() as i64); + } + span.record( + "error.type", + if error.is_timeout() { + "timeout" + } else { + "reqwest" + }, + ); tracing::debug!( method = %self.method, url = %self.url, @@ -154,7 +189,19 @@ impl<'a> Injector for HeaderMapInjector<'a> { } } -fn trace_headers() -> HeaderMap { +fn record_server_fields(span: &Span, url: &str) { + let Ok(parsed) = reqwest::Url::parse(url) else { + return; + }; + if let Some(host) = parsed.host_str() { + span.record("server.address", host); + } + if let Some(port) = parsed.port_or_known_default() { + span.record("server.port", port as i64); + } +} + +pub fn current_trace_headers() -> HeaderMap { let mut headers = HeaderMap::new(); global::get_text_map_propagator(|prop| { prop.inject_context( @@ -192,7 +239,7 @@ mod tests { let _entered = span.enter(); let span_context = span.context().span().span_context().clone(); - let headers = trace_headers(); + let headers = current_trace_headers(); let extractor = HeaderMapExtractor(&headers); let extracted = TraceContextPropagator::new().extract(&extractor); diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs index 93dd81506f..0f0a7e11e1 100644 --- a/codex-rs/codex-client/src/lib.rs +++ b/codex-rs/codex-client/src/lib.rs @@ -19,6 +19,7 @@ pub use crate::custom_ca::build_reqwest_client_with_custom_ca; pub use crate::custom_ca::maybe_build_rustls_client_config_with_custom_ca; pub use crate::default_client::CodexHttpClient; pub use crate::default_client::CodexRequestBuilder; +pub use crate::default_client::current_trace_headers; pub use crate::error::StreamError; pub use crate::error::TransportError; pub use crate::request::Request; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index e38ff3a562..c2d8ccc468 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -71,7 +71,6 @@ use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::W3cTraceContext; use eventsource_stream::Event; use eventsource_stream::EventStreamError; use futures::StreamExt; @@ -86,6 +85,7 @@ 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; @@ -554,30 +554,42 @@ impl ModelClient { turn_metadata_header: Option<&str>, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, + model: Option<&str>, ) -> std::result::Result { - let headers = self.build_websocket_headers(turn_state.as_ref(), turn_metadata_header); - let websocket_telemetry = ModelClientSession::build_websocket_telemetry( - session_telemetry, - auth_context, - request_route_telemetry, - self.state.auth_env_telemetry.clone(), + let connect_span = crate::network_trace::responses_websocket_connect_span( + &self.state.conversation_id, + turn_metadata_header, + &self.state.provider.name, + model, + &api_provider.base_url, ); let websocket_connect_timeout = self.state.provider.websocket_connect_timeout(); let start = Instant::now(); - let result = match tokio::time::timeout( - websocket_connect_timeout, - ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( - headers, - crate::default_client::default_headers(), - turn_state, - Some(websocket_telemetry), - ), - ) - .await - { - Ok(result) => result, - Err(_) => Err(ApiError::Transport(TransportError::Timeout)), - }; + let result = async { + let headers = self.build_websocket_headers(turn_state.as_ref(), turn_metadata_header); + let websocket_telemetry = ModelClientSession::build_websocket_telemetry( + session_telemetry, + auth_context, + request_route_telemetry, + self.state.auth_env_telemetry.clone(), + ); + match tokio::time::timeout( + websocket_connect_timeout, + ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( + headers, + crate::default_client::default_headers(), + turn_state, + Some(websocket_telemetry), + ), + ) + .await + { + Ok(result) => result, + Err(_) => Err(ApiError::Transport(TransportError::Timeout)), + } + } + .instrument(connect_span) + .await; let error_message = result.as_ref().err().map(telemetry_api_error_message); let response_debug = result .as_ref() @@ -658,6 +670,7 @@ impl ModelClient { HeaderValue::from_static("true"), ); } + headers.extend(build_current_trace_headers()); headers } } @@ -884,6 +897,7 @@ impl ModelClientSession { /*turn_metadata_header*/ None, auth_context, RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT), + /*model*/ None, ) .await?; self.websocket_session.connection = Some(connection); @@ -916,6 +930,7 @@ impl ModelClientSession { options, auth_context, request_route_telemetry, + model, } = params; let needs_new = match self.websocket_session.connection.as_ref() { Some(conn) => conn.is_closed().await, @@ -939,6 +954,7 @@ impl ModelClientSession { turn_metadata_header, auth_context, request_route_telemetry, + model, ) .await { @@ -1023,6 +1039,7 @@ 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), @@ -1052,10 +1069,20 @@ impl ModelClientSession { client_setup.api_auth, ) .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); - let stream_result = client.stream_request(request, options).await; + 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; match stream_result { Ok(stream) => { + let _entered = request_span.enter(); let (stream, _) = map_response_stream(stream, session_telemetry.clone()); return Ok(stream); } @@ -1102,7 +1129,6 @@ impl ModelClientSession { service_tier: Option, turn_metadata_header: Option<&str>, warmup: bool, - request_trace: Option, ) -> Result { let auth_manager = self.client.state.auth_manager.clone(); @@ -1112,6 +1138,7 @@ 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 request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), &client_setup.api_auth, @@ -1128,17 +1155,6 @@ impl ModelClientSession { summary, service_tier, )?; - let mut ws_payload = ResponseCreateWsRequest { - client_metadata: response_create_client_metadata( - build_ws_client_metadata(turn_metadata_header), - request_trace.as_ref(), - ), - ..ResponseCreateWsRequest::from(&request) - }; - if warmup { - ws_payload.generate = Some(false); - } - match self .websocket_connection(WebsocketConnectParams { session_telemetry, @@ -1150,6 +1166,7 @@ impl ModelClientSession { request_route_telemetry: RequestRouteTelemetry::for_endpoint( RESPONSES_ENDPOINT, ), + model: Some(&model_info.slug), }) .await { @@ -1175,17 +1192,44 @@ impl ModelClientSession { Err(err) => return Err(map_api_error(err)), } - let ws_request = self.prepare_websocket_request(ws_payload, &request); - self.websocket_session.last_request = Some(request); - let stream_result = self.websocket_session.connection.as_ref().ok_or_else(|| { - map_api_error(ApiError::Stream( - "websocket connection is unavailable".to_string(), - )) - })?; - let stream_result = stream_result - .stream_request(ws_request, self.websocket_session.connection_reused()) - .await - .map_err(map_api_error)?; + let connection_reused = self.websocket_session.connection_reused(); + let request_span = crate::network_trace::responses_websocket_request_span( + &self.client.state.conversation_id, + turn_metadata_header, + &self.client.state.provider.name, + &model_info.slug, + &api_base_url, + connection_reused, + warmup, + ); + let stream_result = async { + let mut ws_payload = ResponseCreateWsRequest { + client_metadata: response_create_client_metadata( + build_ws_client_metadata(turn_metadata_header), + current_span_w3c_trace_context().as_ref(), + ), + ..ResponseCreateWsRequest::from(&request) + }; + if warmup { + ws_payload.generate = Some(false); + } + let ws_request = self.prepare_websocket_request(ws_payload, &request); + self.websocket_session.last_request = Some(request); + let stream_result = + self.websocket_session.connection.as_ref().ok_or_else(|| { + map_api_error(ApiError::Stream( + "websocket connection is unavailable".to_string(), + )) + })?; + stream_result + .stream_request(ws_request, connection_reused) + .await + .map_err(map_api_error) + } + .instrument(request_span.clone()) + .await; + let stream_result = stream_result?; + let _entered = request_span.enter(); let (stream, last_request_rx) = map_response_stream(stream_result, session_telemetry.clone()); self.websocket_session.last_response_rx = Some(last_request_rx); @@ -1256,7 +1300,6 @@ impl ModelClientSession { service_tier, turn_metadata_header, /*warmup*/ true, - current_span_w3c_trace_context(), ) .await { @@ -1300,7 +1343,6 @@ impl ModelClientSession { match wire_api { WireApi::Responses => { if self.client.responses_websocket_enabled() { - let request_trace = current_span_w3c_trace_context(); match self .stream_responses_websocket( prompt, @@ -1311,7 +1353,6 @@ impl ModelClientSession { service_tier, turn_metadata_header, /*warmup*/ false, - request_trace, ) .await? { @@ -1363,6 +1404,24 @@ fn parse_turn_metadata_header(turn_metadata_header: Option<&str>) -> Option ApiHeaderMap { + let mut headers = ApiHeaderMap::new(); + let Some(trace) = current_span_w3c_trace_context() else { + return headers; + }; + if let Some(traceparent) = trace.traceparent.as_deref() + && let Ok(header_value) = HeaderValue::from_str(traceparent) + { + headers.insert("traceparent", header_value); + } + if let Some(tracestate) = trace.tracestate.as_deref() + && let Ok(header_value) = HeaderValue::from_str(tracestate) + { + headers.insert("tracestate", header_value); + } + headers +} + fn build_ws_client_metadata(turn_metadata_header: Option<&str>) -> Option> { let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header)?; let turn_metadata = turn_metadata_header.to_str().ok()?.to_string(); @@ -1414,72 +1473,76 @@ 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, - ); + 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; + } } - 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), - }); + Ok(event) => { + if tx_event.send(Ok(event)).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; + 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) } @@ -1549,6 +1612,7 @@ struct WebsocketConnectParams<'a> { options: &'a ApiResponsesOptions, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, + model: Option<&'a str>, } async fn handle_unauthorized( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 29436a0d7f..3332dfefd7 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -52,6 +52,7 @@ mod mcp_tool_approval_templates; pub mod models_manager; mod network_policy_decision; pub mod network_proxy_loader; +mod network_trace; mod original_image_detail; mod packages; pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index f82cc73bb5..57b6a99cb1 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -50,6 +50,7 @@ use serde::Serialize; use std::path::Path; use std::sync::Arc; use toml_edit::value; +use tracing::Instrument; /// Handles the specified tool call dispatches the appropriate /// `McpToolCallBegin` and `McpToolCallEnd` events to the `Session`. @@ -121,6 +122,13 @@ pub(crate) async fn handle_mcp_tool_call( } let request_meta = build_mcp_tool_call_request_meta(turn_context.as_ref(), &server, metadata.as_ref()); + let server_origin = sess + .services + .mcp_connection_manager + .read() + .await + .server_origin(&server) + .map(str::to_string); let tool_call_begin_event = EventMsg::McpToolCallBegin(McpToolCallBeginEvent { call_id: call_id.clone(), @@ -145,25 +153,19 @@ pub(crate) async fn handle_mcp_tool_call( maybe_mark_thread_memory_mode_polluted(sess.as_ref(), turn_context.as_ref()).await; let start = Instant::now(); - let result = sess - .call_tool( - &server, - &tool_name, - arguments_value.clone(), - request_meta.clone(), - ) - .await - .map_err(|e| format!("tool call error: {e:?}")); - let result = sanitize_mcp_tool_result_for_model( - turn_context - .model_info - .input_modalities - .contains(&InputModality::Image), - result, - ); - if let Err(e) = &result { - tracing::warn!("MCP tool call error: {e:?}"); - } + let result = execute_mcp_tool_call( + &sess, + turn_context, + McpToolCallRequest { + server: &server, + tool_name: &tool_name, + call_id: &call_id, + arguments_value: arguments_value.clone(), + request_meta: request_meta.clone(), + server_origin: server_origin.as_deref(), + }, + ) + .await; let tool_call_end_event = EventMsg::McpToolCallEnd(McpToolCallEndEvent { call_id: call_id.clone(), invocation, @@ -236,20 +238,19 @@ pub(crate) async fn handle_mcp_tool_call( let start = Instant::now(); // Perform the tool call. - let result = sess - .call_tool(&server, &tool_name, arguments_value.clone(), request_meta) - .await - .map_err(|e| format!("tool call error: {e:?}")); - let result = sanitize_mcp_tool_result_for_model( - turn_context - .model_info - .input_modalities - .contains(&InputModality::Image), - result, - ); - if let Err(e) = &result { - tracing::warn!("MCP tool call error: {e:?}"); - } + let result = execute_mcp_tool_call( + &sess, + turn_context, + McpToolCallRequest { + server: &server, + tool_name: &tool_name, + call_id: &call_id, + arguments_value: arguments_value.clone(), + request_meta, + server_origin: server_origin.as_deref(), + }, + ) + .await; let tool_call_end_event = EventMsg::McpToolCallEnd(McpToolCallEndEvent { call_id: call_id.clone(), invocation, @@ -273,6 +274,53 @@ pub(crate) async fn handle_mcp_tool_call( CallToolResult::from_result(result) } +async fn execute_mcp_tool_call( + sess: &Arc, + turn_context: &Arc, + request: McpToolCallRequest<'_>, +) -> Result { + let tool_call_span = crate::network_trace::mcp_tool_call_span( + sess.as_ref(), + turn_context.as_ref(), + request.server, + request.tool_name, + request.call_id, + request.server_origin, + ); + let result = async { + sess.call_tool( + request.server, + request.tool_name, + request.arguments_value, + request.request_meta, + ) + .await + .map_err(|e| format!("tool call error: {e:?}")) + } + .instrument(tool_call_span) + .await; + let result = sanitize_mcp_tool_result_for_model( + turn_context + .model_info + .input_modalities + .contains(&InputModality::Image), + result, + ); + if let Err(error) = &result { + tracing::warn!("MCP tool call error: {error:?}"); + } + result +} + +struct McpToolCallRequest<'a> { + server: &'a str, + tool_name: &'a str, + call_id: &'a str, + arguments_value: Option, + request_meta: Option, + server_origin: Option<&'a str>, +} + async fn maybe_mark_thread_memory_mode_polluted(sess: &Session, turn_context: &TurnContext) { if !turn_context .config diff --git a/codex-rs/core/src/network_trace.rs b/codex-rs/core/src/network_trace.rs new file mode 100644 index 0000000000..73565457fc --- /dev/null +++ b/codex-rs/core/src/network_trace.rs @@ -0,0 +1,187 @@ +use codex_protocol::ThreadId; +use serde::Deserialize; +use tracing::Span; +use tracing::field::Empty; +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, + turn_id: Option, +} + +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(), + session_id: session.conversation_id.to_string(), + turn_id: Some(turn_context.sub_id.clone()), + } + } + + fn record_on(&self, span: &Span) { + span.record("conversation.id", self.conversation_id.as_str()); + span.record("session.id", self.session_id.as_str()); + if let Some(turn_id) = self.turn_id.as_deref() { + span.record("turn.id", turn_id); + } + } +} + +fn record_server_fields(span: &Span, url: Option<&str>) { + let Some(url) = url else { + return; + }; + let Ok(parsed) = Url::parse(url) else { + return; + }; + if let Some(host) = parsed.host_str() { + span.record("server.address", host); + } + if let Some(port) = parsed.port_or_known_default() { + span.record("server.port", port as i64); + } +} + +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) fn responses_websocket_connect_span( + conversation_id: &ThreadId, + turn_metadata_header: Option<&str>, + provider_name: &str, + model: Option<&str>, + base_url: &str, +) -> Span { + let span = tracing::info_span!( + "responses_websocket.connect_attempt", + otel.kind = "client", + provider = provider_name, + model = model.unwrap_or(""), + transport = "responses_websocket", + 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) fn responses_websocket_request_span( + conversation_id: &ThreadId, + turn_metadata_header: Option<&str>, + provider_name: &str, + model: &str, + base_url: &str, + connection_reused: bool, + warmup: bool, +) -> Span { + let span = tracing::info_span!( + "responses_websocket.request", + otel.kind = "client", + provider = provider_name, + model, + transport = "responses_websocket", + api.path = "responses", + connection.reused = connection_reused, + websocket.warmup = warmup, + 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) fn mcp_tool_call_span( + session: &Session, + turn_context: &TurnContext, + server_name: &str, + tool_name: &str, + call_id: &str, + server_origin: Option<&str>, +) -> Span { + let transport = match server_origin { + Some("stdio") => "stdio", + Some(_) => "streamable_http", + None => "", + }; + let span = tracing::info_span!( + "mcp.tools.call", + otel.kind = "client", + rpc.system = "jsonrpc", + rpc.method = "tools/call", + mcp.server.name = server_name, + mcp.server.origin = server_origin.unwrap_or(""), + mcp.transport = transport, + tool.name = tool_name, + tool.call_id = call_id, + conversation.id = Empty, + session.id = Empty, + turn.id = Empty, + server.address = Empty, + server.port = Empty, + ); + CorrelationFields::from_turn_context(session, turn_context).record_on(&span); + record_server_fields(&span, server_origin); + span +} diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 823057797c..9afc9aa1ba 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -9,6 +9,7 @@ use codex_core::ResponseEvent; use codex_core::WireApi; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; +use codex_otel::current_span_w3c_trace_context; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::ContentItem; @@ -18,11 +19,24 @@ use codex_protocol::protocol::SubAgentSource; use core_test_support::load_default_config_for_test; use core_test_support::responses; use core_test_support::test_codex::test_codex; +use core_test_support::tracing::install_test_tracing; use futures::StreamExt; use pretty_assertions::assert_eq; use tempfile::TempDir; +use tracing::Instrument; use wiremock::matchers::header; +fn trace_id(traceparent: Option<&str>) -> Option<&str> { + let traceparent = traceparent?; + let mut parts = traceparent.split('-'); + match (parts.next(), parts.next(), parts.next(), parts.next()) { + (Some(_version), Some(trace_id), Some(_span_id), Some(_flags)) if trace_id.len() == 32 => { + Some(trace_id) + } + _ => None, + } +} + #[tokio::test] async fn responses_stream_includes_subagent_header_on_review() { core_test_support::skip_if_no_network!(); @@ -135,6 +149,42 @@ async fn responses_stream_includes_subagent_header_on_review() { assert_eq!(request.header("x-codex-sandbox"), None); } +#[tokio::test] +async fn responses_stream_includes_traceparent_header() { + core_test_support::skip_if_no_network!(); + + let _trace_test_context = install_test_tracing("responses-headers-trace"); + let server = responses::start_mock_server().await; + let request_recorder = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + + let test = test_codex().build(&server).await.expect("build test codex"); + let expected_trace = async { + let expected_trace = + current_span_w3c_trace_context().expect("current span should have trace context"); + test.submit_turn("hello").await.expect("submit turn prompt"); + expected_trace + } + .instrument(tracing::info_span!("responses.headers.trace_request")) + .await; + + let request = request_recorder.single_request(); + assert_eq!( + trace_id(request.header("traceparent").as_deref()), + trace_id(expected_trace.traceparent.as_deref()) + ); + assert_eq!( + request.header("tracestate").as_deref(), + expected_trace.tracestate.as_deref() + ); +} + #[tokio::test] async fn responses_stream_includes_subagent_header_on_other() { core_test_support::skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index b568c6aee2..57cd678639 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -56,6 +56,17 @@ const OPENAI_BETA_HEADER: &str = "OpenAI-Beta"; const WS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; const X_CLIENT_REQUEST_ID_HEADER: &str = "x-client-request-id"; +fn trace_id(traceparent: Option<&str>) -> Option<&str> { + let traceparent = traceparent?; + let mut parts = traceparent.split('-'); + match (parts.next(), parts.next(), parts.next(), parts.next()) { + (Some(_version), Some(trace_id), Some(_span_id), Some(_flags)) if trace_id.len() == 32 => { + Some(trace_id) + } + _ => None, + } +} + fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTraceContext) { let client_metadata = body["client_metadata"] .as_object() @@ -69,7 +80,10 @@ fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTr .as_deref() .expect("missing expected traceparent"); - assert_eq!(actual_traceparent, expected_traceparent); + assert_eq!( + trace_id(Some(actual_traceparent)), + trace_id(Some(expected_traceparent)) + ); assert_eq!( client_metadata .get(WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY) @@ -192,6 +206,15 @@ async fn responses_websocket_reuses_connection_with_per_turn_trace_payloads() { }; assert_eq!(server.handshakes().len(), 1); + let handshake = server.single_handshake(); + assert_eq!( + trace_id(handshake.header("traceparent").as_deref()), + trace_id(first_trace.traceparent.as_deref()) + ); + assert_eq!( + handshake.header("tracestate").as_deref(), + first_trace.tracestate.as_deref() + ); let connection = server.single_connection(); assert_eq!(connection.len(), 2); diff --git a/codex-rs/core/tests/suite/otel.rs b/codex-rs/core/tests/suite/otel.rs index ecf6283664..39b4f60057 100644 --- a/codex-rs/core/tests/suite/otel.rs +++ b/codex-rs/core/tests/suite/otel.rs @@ -23,9 +23,11 @@ use core_test_support::responses::mount_sse_once; use core_test_support::responses::sse; use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; +use core_test_support::stdio_server_bin; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use std::collections::HashMap; use std::sync::Mutex; use tracing::Level; use tracing_test::traced_test; @@ -718,6 +720,173 @@ async fn record_responses_sets_span_fields_for_response_events() { } } +#[tokio::test(flavor = "current_thread")] +async fn responses_request_span_records_turn_correlation_fields() { + let buffer: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); + let subscriber = tracing_subscriber::fmt() + .with_level(true) + .with_ansi(false) + .with_max_level(Level::TRACE) + .with_span_events(FmtSpan::FULL) + .with_writer(MockWriter::new(buffer)) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + + 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; + + let logs = String::from_utf8(buffer.lock().unwrap().clone()).unwrap(); + let line = logs + .lines() + .find(|line| line.contains("responses_http.request{")) + .expect("missing responses_http.request span"); + assert!( + line.contains("otel.kind=\"client\"") + && line.contains("transport=\"responses_http\"") + && line.contains("conversation.id=") + && line.contains("session.id=") + && line.contains("turn.id="), + "responses_http.request span is missing expected correlation fields\nlogs:\n{logs}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn mcp_tool_call_span_records_server_and_tool_fields() { + let buffer: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); + let subscriber = tracing_subscriber::fmt() + .with_level(true) + .with_ansi(false) + .with_max_level(Level::TRACE) + .with_span_events(FmtSpan::FULL) + .with_writer(MockWriter::new(buffer)) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + + let server = start_mock_server().await; + let call_id = "mcp-span-call"; + let server_name = "rmcp_span"; + let tool_name = format!("mcp__{server_name}__echo"); + + mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, &tool_name, "{\"message\":\"ping\"}"), + ev_completed("resp-1"), + ]), + ) + .await; + mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "rmcp echo tool completed successfully."), + ev_completed("resp-2"), + ]), + ) + .await; + + let rmcp_test_server_bin = stdio_server_bin().expect("test stdio server binary"); + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex() + .with_config(move |config| { + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + server_name.to_string(), + codex_core::config::types::McpServerConfig { + transport: codex_core::config::types::McpServerTransportConfig::Stdio { + command: rmcp_test_server_bin, + args: Vec::new(), + env: Some(HashMap::from([( + "MCP_TEST_VALUE".to_string(), + "span-test".to_string(), + )])), + env_vars: Vec::new(), + cwd: None, + }, + enabled: true, + required: false, + disabled_reason: None, + startup_timeout_sec: Some(std::time::Duration::from_secs(10)), + tool_timeout_sec: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth_resource: None, + }, + ); + config + .mcp_servers + .set(servers) + .expect("test mcp servers should accept any configuration"); + }) + .build(&server) + .await + .unwrap(); + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "call the rmcp echo tool".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await + .unwrap(); + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let logs = String::from_utf8(buffer.lock().unwrap().clone()).unwrap(); + let line = logs + .lines() + .find(|line| line.contains("mcp.tools.call{")) + .expect("missing mcp.tools.call span"); + assert!( + line.contains("otel.kind=\"client\"") + && line.contains("rpc.system=\"jsonrpc\"") + && line.contains("rpc.method=\"tools/call\"") + && line.contains(&format!("mcp.server.name=\"{server_name}\"")) + && line.contains("mcp.transport=\"stdio\"") + && line.contains("tool.name=\"echo\"") + && line.contains(&format!("tool.call_id=\"{call_id}\"")) + && line.contains("turn.id="), + "mcp.tools.call span is missing expected MCP fields\nlogs:\n{logs}" + ); +} + #[tokio::test] #[traced_test] async fn handle_response_item_records_tool_result_for_custom_tool_call() { diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 772674f79a..8a01e9b531 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -14,6 +14,7 @@ use codex_core::config::types::McpServerConfig; use codex_core::config::types::McpServerTransportConfig; use codex_core::models_manager::manager::RefreshStrategy; +use codex_otel::current_span_w3c_trace_context; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::openai_models::ConfigShellToolType; use codex_protocol::openai_models::InputModality; @@ -36,6 +37,7 @@ use core_test_support::responses::mount_sse_once; use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; use core_test_support::test_codex::test_codex; +use core_test_support::tracing::install_test_tracing; use core_test_support::wait_for_event; use core_test_support::wait_for_event_with_timeout; use reqwest::Client; @@ -48,9 +50,21 @@ use tokio::process::Child; use tokio::process::Command; use tokio::time::Instant; use tokio::time::sleep; +use tracing::Instrument; static OPENAI_PNG: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAD0AAAA9CAYAAAAeYmHpAAAE6klEQVR4Aeyau44UVxCGx1fZsmRLlm3Zoe0XcGQ5cUiCCIgJeS9CHgAhMkISQnIuGQgJEkBcxLW+nqnZ6uqqc+nuWRC7q/P3qetf9e+MtOwyX25O4Nep6JPyop++0qev9HrfgZ+F6r2DuB/vHOrt/UIkqdDHYvujOW6fO7h/CNEI+a5jc+pBR8uy0jVFsziYu5HtfSUk+Io34q921hLNctFSX0gwww+S8wce8K1LfCU+cYW4888aov8NxqvQILUPPReLOrm6zyLxa4i+6VZuFbJo8d1MOHZm+7VUtB/aIvhPWc/3SWg49JcwFLlHxuXKjtyloo+YNhuW3VS+WPBuUEMvCFKjEDVgFBQHXrnazpqiSxNZCkQ1kYiozsbm9Oz7l4i2Il7vGccGNWAc3XosDrZe/9P3ZnMmzHNEQw4smf8RQ87XEAMsC7Az0Au+dgXerfH4+sHvEc0SYGic8WBBUGqFH2gN7yDrazy7m2pbRTeRmU3+MjZmr1h6LJgPbGy23SI6GlYT0brQ71IY8Us4PNQCm+zepSbaD2BY9xCaAsD9IIj/IzFmKMSdHHonwdZATbTnYREf6/VZGER98N9yCWIvXQwXDoDdhZJoT8jwLnJXDB9w4Sb3e6nK5ndzlkTLnP3JBu4LKkbrYrU69gCVceV0JvpyuW1xlsUVngzhwMetn/XamtTORF9IO5YnWNiyeF9zCAfqR3fUW+vZZKLtgP+ts8BmQRBREAdRDhH3o8QuRh/YucNFz2BEjxbRN6LGzphfKmvP6v6QhqIQyZ8XNJ0W0X83MR1PEcJBNO2KC2Z1TW/v244scp9FwRViZxIOBF0Lctk7ZVSavdLvRlV1hz/ysUi9sr8CIcB3nvWBwA93ykTz18eAYxQ6N/K2DkPA1lv3iXCwmDUT7YkjIby9siXueIJj9H+pzSqJ9oIuJWTUgSSt4WO7o/9GGg0viR4VinNRUDoIj34xoCd6pxD3aK3zfdbnx5v1J3ZNNEJsE0sBG7N27ReDrJc4sFxz7dI/ZAbOmmiKvHBitQXpAdR6+F7v+/ol/tOouUV01EeMZQF2BoQDn6dP4XNr+j9GZEtEK1/L8pFw7bd3a53tsTa7WD+054jOFmPg1XBKPQgnqFfmFcy32ZRvjmiIIQTYFvyDxQ8nH8WIwwGwlyDjDznnilYyFr6njrlZwsKkBpO59A7OwgdzPEWRm+G+oeb7IfyNuzjEEVLrOVxJsxvxwF8kmCM6I2QYmJunz4u4TrADpfl7mlbRTWQ7VmrBzh3+C9f6Grc3YoGN9dg/SXFthpRsT6vobfXRs2VBlgBHXVMLHjDNbIZv1sZ9+X3hB09cXdH1JKViyG0+W9bWZDa/r2f9zAFR71sTzGpMSWz2iI4YssWjWo3REy1MDGjdwe5e0dFSiAC1JakBvu4/CUS8Eh6dqHdU0Or0ioY3W5ClSqDXAy7/6SRfgw8vt4I+tbvvNtFT2kVDhY5+IGb1rCqYaXNF08vSALsXCPmt0kQNqJT1p5eI1mkIV/BxCY1z85lOzeFbPBQHURkkPTlwTYK9gTVE25l84IbFFN+YJDHjdpn0gq6mrHht0dkcjbM4UL9283O5p77GN+SPW/QwVB4IUYg7Or+Kp7naR6qktP98LNF2UxWo9yObPIT9KYg+hK4i56no4rfnM0qeyFf6AwAAAP//trwR3wAAAAZJREFUAwBZ0sR75itw5gAAAABJRU5ErkJggg=="; +fn trace_id(traceparent: Option<&str>) -> Option<&str> { + let traceparent = traceparent?; + let mut parts = traceparent.split('-'); + match (parts.next(), parts.next(), parts.next(), parts.next()) { + (Some(_version), Some(trace_id), Some(_span_id), Some(_flags)) if trace_id.len() == 32 => { + Some(trace_id) + } + _ => None, + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[serial(mcp_test_value)] async fn stdio_server_round_trip() -> anyhow::Result<()> { @@ -193,6 +207,146 @@ async fn stdio_server_round_trip() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(mcp_test_value)] +async fn stdio_tool_call_propagates_trace_context_in_meta() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let _trace_test_context = install_test_tracing("rmcp-stdio-trace"); + + let server = responses::start_mock_server().await; + + let call_id = "trace-call-stdio"; + let server_name = "rmcp_trace_stdio"; + let tool_name = format!("mcp__{server_name}__trace_inspect"); + + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, &tool_name, "{}"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "rmcp stdio trace inspect completed."), + responses::ev_completed("resp-2"), + ]), + ) + .await; + + let rmcp_test_server_bin = stdio_server_bin()?; + let fixture = test_codex() + .with_config(move |config| { + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + server_name.to_string(), + McpServerConfig { + transport: McpServerTransportConfig::Stdio { + command: rmcp_test_server_bin, + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + enabled: true, + required: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth_resource: None, + }, + ); + config + .mcp_servers + .set(servers) + .expect("test mcp servers should accept any configuration"); + }) + .build(&server) + .await?; + let session_model = fixture.session_configured.model.clone(); + + let expected_trace = async { + let expected_trace = + current_span_w3c_trace_context().expect("current span should have trace context"); + fixture + .codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "inspect the stdio mcp trace context".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: fixture.cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + Ok::<_, anyhow::Error>(expected_trace) + } + .instrument(tracing::info_span!("rmcp.stdio.trace_request")) + .await?; + + let end_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallEnd(_)) + }) + .await; + let EventMsg::McpToolCallEnd(end) = end_event else { + unreachable!("event guard guarantees McpToolCallEnd"); + }; + + let result = end + .result + .as_ref() + .expect("rmcp trace inspect tool should return success"); + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + let Value::Object(map) = structured else { + panic!("structured content should be an object: {structured:?}"); + }; + let expected_session_id = fixture.session_configured.session_id.to_string(); + + assert_eq!( + trace_id(map.get("traceparent").and_then(Value::as_str)), + trace_id(expected_trace.traceparent.as_deref()) + ); + assert_eq!( + map.get("tracestate").and_then(Value::as_str), + expected_trace.tracestate.as_deref() + ); + assert_eq!( + structured + .pointer("/turn_metadata/session_id") + .and_then(Value::as_str), + Some(expected_session_id.as_str()) + ); + assert!( + structured + .pointer("/turn_metadata/turn_id") + .and_then(Value::as_str) + .is_some_and(|turn_id| !turn_id.is_empty()), + "stdio trace inspect should include turn metadata turn_id: {map:?}" + ); + + wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[serial(mcp_test_value)] async fn stdio_image_responses_round_trip() -> anyhow::Result<()> { @@ -854,6 +1008,183 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tool_call_propagates_trace_headers() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let _trace_test_context = install_test_tracing("rmcp-http-trace"); + + let server = responses::start_mock_server().await; + + let call_id = "trace-call-http"; + let server_name = "rmcp_http_trace"; + let tool_name = format!("mcp__{server_name}__trace_inspect"); + + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, &tool_name, "{}"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message( + "msg-1", + "rmcp streamable http trace inspect completed.", + ), + responses::ev_completed("resp-2"), + ]), + ) + .await; + + let rmcp_http_server_bin = match cargo_bin("test_streamable_http_server") { + Ok(path) => path, + Err(err) => { + eprintln!("test_streamable_http_server binary not available, skipping test: {err}"); + return Ok(()); + } + }; + + let listener = TcpListener::bind("127.0.0.1:0")?; + let port = listener.local_addr()?.port(); + drop(listener); + let bind_addr = format!("127.0.0.1:{port}"); + let server_url = format!("http://{bind_addr}/mcp"); + + let mut http_server_child = Command::new(&rmcp_http_server_bin) + .kill_on_drop(true) + .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) + .spawn()?; + + wait_for_streamable_http_server(&mut http_server_child, &bind_addr, Duration::from_secs(5)) + .await?; + + let fixture = test_codex() + .with_config(move |config| { + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + server_name.to_string(), + McpServerConfig { + transport: McpServerTransportConfig::StreamableHttp { + url: server_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + enabled: true, + required: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth_resource: None, + }, + ); + config + .mcp_servers + .set(servers) + .expect("test mcp servers should accept any configuration"); + }) + .build(&server) + .await?; + let session_model = fixture.session_configured.model.clone(); + + let expected_trace = async { + let expected_trace = + current_span_w3c_trace_context().expect("current span should have trace context"); + fixture + .codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "inspect the streamable http mcp trace context".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: fixture.cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + }) + .await?; + Ok::<_, anyhow::Error>(expected_trace) + } + .instrument(tracing::info_span!("rmcp.http.trace_request")) + .await?; + + let end_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallEnd(_)) + }) + .await; + let EventMsg::McpToolCallEnd(end) = end_event else { + unreachable!("event guard guarantees McpToolCallEnd"); + }; + + let result = end + .result + .as_ref() + .expect("rmcp trace inspect tool should return success"); + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + let Value::Object(map) = structured else { + panic!("structured content should be an object: {structured:?}"); + }; + let expected_session_id = fixture.session_configured.session_id.to_string(); + + assert_eq!( + trace_id(map.get("traceparent").and_then(Value::as_str)), + trace_id(expected_trace.traceparent.as_deref()) + ); + assert_eq!( + map.get("tracestate").and_then(Value::as_str), + expected_trace.tracestate.as_deref() + ); + assert_eq!( + structured + .pointer("/turn_metadata/x-codex-turn-metadata/session_id") + .and_then(Value::as_str), + Some(expected_session_id.as_str()) + ); + assert!( + structured + .pointer("/turn_metadata/x-codex-turn-metadata/turn_id") + .and_then(Value::as_str) + .is_some_and(|turn_id| !turn_id.is_empty()), + "streamable http trace inspect should include turn metadata turn_id: {map:?}" + ); + + wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + server.verify().await; + + match http_server_child.try_wait() { + Ok(Some(_)) => {} + Ok(None) => { + let _ = http_server_child.kill().await; + } + Err(error) => { + eprintln!("failed to check streamable http server status: {error}"); + let _ = http_server_child.kill().await; + } + } + if let Err(error) = http_server_child.wait().await { + eprintln!("failed to await streamable http server shutdown: {error}"); + } + + Ok(()) +} + /// This test writes to a fallback credentials file in CODEX_HOME. /// Ideally, we wouldn't need to serialize the test but it's much more cumbersome to wire CODEX_HOME through the code. #[test] diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index cd08307767..9ea4e47d32 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -26,6 +26,10 @@ use serde::Deserialize; use serde_json::json; use tokio::task; +const X_CODEX_TRACEPARENT_META_KEY: &str = "x-codex-traceparent"; +const X_CODEX_TRACESTATE_META_KEY: &str = "x-codex-tracestate"; +const X_CODEX_TURN_METADATA_META_KEY: &str = "x-codex-turn-metadata"; + #[derive(Clone)] struct TestToolServer { tools: Arc>, @@ -48,6 +52,7 @@ impl TestToolServer { Self::echo_dash_tool(), Self::image_tool(), Self::image_scenario_tool(), + Self::trace_inspect_tool(), ]; let resources = vec![Self::memo_resource()]; let resource_templates = vec![Self::memo_template()]; @@ -163,6 +168,22 @@ impl TestToolServer { ) } + fn trace_inspect_tool() -> Tool { + #[expect(clippy::expect_used)] + let schema: JsonObject = serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + })) + .expect("trace inspect tool schema should deserialize"); + + Tool::new( + Cow::Borrowed("trace_inspect"), + Cow::Borrowed("Return tracing and turn metadata observed in the MCP request."), + Arc::new(schema), + ) + } + fn memo_resource() -> Resource { let raw = RawResource { uri: MEMO_URI.to_string(), @@ -363,6 +384,36 @@ impl ServerHandler for TestToolServer { let args = Self::parse_call_args::(&request, "image_scenario")?; Self::image_scenario_result(args) } + "trace_inspect" => { + let traceparent = request + .meta + .as_ref() + .and_then(|meta| meta.get(X_CODEX_TRACEPARENT_META_KEY)) + .and_then(serde_json::Value::as_str) + .map(str::to_string); + let tracestate = request + .meta + .as_ref() + .and_then(|meta| meta.get(X_CODEX_TRACESTATE_META_KEY)) + .and_then(serde_json::Value::as_str) + .map(str::to_string); + let turn_metadata = request + .meta + .as_ref() + .and_then(|meta| meta.get(X_CODEX_TURN_METADATA_META_KEY)) + .cloned(); + + Ok(CallToolResult { + content: Vec::new(), + structured_content: Some(json!({ + "traceparent": traceparent, + "tracestate": tracestate, + "turn_metadata": turn_metadata, + })), + is_error: Some(false), + meta: None, + }) + } other => Err(McpError::invalid_params( format!("unknown tool: {other}"), None, diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index 284e1194c3..f3ae370fed 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -52,22 +52,26 @@ struct TestToolServer { tools: Arc>, resources: Arc>, resource_templates: Arc>, + request_observation_state: RequestObservationState, } const MEMO_URI: &str = "memo://codex/example-note"; const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server."; const MCP_SESSION_ID_HEADER: &str = "mcp-session-id"; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; +const TRACEPARENT_HEADER: &str = "traceparent"; +const TRACESTATE_HEADER: &str = "tracestate"; impl TestToolServer { - fn new() -> Self { - let tools = vec![Self::echo_tool()]; + fn new(request_observation_state: RequestObservationState) -> Self { + let tools = vec![Self::echo_tool(), Self::trace_inspect_tool()]; let resources = vec![Self::memo_resource()]; let resource_templates = vec![Self::memo_template()]; Self { tools: Arc::new(tools), resources: Arc::new(resources), resource_templates: Arc::new(resource_templates), + request_observation_state, } } @@ -91,6 +95,22 @@ impl TestToolServer { ) } + fn trace_inspect_tool() -> Tool { + #[expect(clippy::expect_used)] + let schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + })) + .expect("trace inspect tool schema should deserialize"); + + Tool::new( + Cow::Borrowed("trace_inspect"), + Cow::Borrowed("Return tracing headers observed on the MCP HTTP request."), + Arc::new(schema), + ) + } + fn memo_resource() -> Resource { let raw = RawResource { uri: MEMO_URI.to_string(), @@ -129,6 +149,17 @@ struct SessionFailureState { armed_failure: Arc>>, } +#[derive(Clone, Debug, Default)] +struct ObservedTraceHeaders { + traceparent: Option, + tracestate: Option, +} + +#[derive(Clone, Default)] +struct RequestObservationState { + last_mcp_post_headers: Arc>, +} + #[derive(Clone, Debug)] struct ArmedFailure { status: StatusCode, @@ -257,6 +288,24 @@ impl ServerHandler for TestToolServer { meta: None, }) } + "trace_inspect" => { + let observed = self + .request_observation_state + .last_mcp_post_headers + .lock() + .await + .clone(); + Ok(CallToolResult { + content: Vec::new(), + structured_content: Some(json!({ + "traceparent": observed.traceparent, + "tracestate": observed.tracestate, + "turn_metadata": request.meta, + })), + is_error: Some(false), + meta: None, + }) + } other => Err(McpError::invalid_params( format!("unknown tool: {other}"), None, @@ -277,6 +326,7 @@ fn parse_bind_addr() -> Result> { async fn main() -> Result<(), Box> { let bind_addr = parse_bind_addr()?; let session_failure_state = SessionFailureState::default(); + let request_observation_state = RequestObservationState::default(); const MAX_BIND_RETRIES: u32 = 20; const BIND_RETRY_DELAY: Duration = Duration::from_millis(50); @@ -327,11 +377,18 @@ async fn main() -> Result<(), Box> { .nest_service( "/mcp", StreamableHttpService::new( - || Ok(TestToolServer::new()), + { + let request_observation_state = request_observation_state.clone(); + move || Ok(TestToolServer::new(request_observation_state.clone())) + }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ), ) + .layer(middleware::from_fn_with_state( + request_observation_state, + capture_mcp_request_headers, + )) .layer(middleware::from_fn_with_state( session_failure_state.clone(), fail_session_post_when_armed, @@ -369,6 +426,29 @@ async fn require_bearer( } } +async fn capture_mcp_request_headers( + State(state): State, + request: Request, + next: Next, +) -> Response { + if request.uri().path() == "/mcp" && request.method() == Method::POST { + let observed = ObservedTraceHeaders { + traceparent: request + .headers() + .get(TRACEPARENT_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + tracestate: request + .headers() + .get(TRACESTATE_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::to_string), + }; + *state.last_mcp_post_headers.lock().await = observed; + } + next.run(request).await +} + async fn arm_session_post_failure( State(state): State, Json(request): Json, diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 0edd0f1527..bce36ab8a1 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -4,6 +4,7 @@ mod oauth; mod perform_oauth_login; mod program_resolver; mod rmcp_client; +mod trace; mod utils; pub use auth_status::StreamableHttpOAuthDiscovery; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 55a3603ed7..192419d629 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -9,6 +9,7 @@ use std::time::Duration; use anyhow::Result; use anyhow::anyhow; +use codex_client::CodexHttpClient; use codex_client::build_reqwest_client_with_custom_ca; use futures::FutureExt; use futures::StreamExt; @@ -73,6 +74,7 @@ use crate::oauth::OAuthCredentialsStoreMode; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; use crate::program_resolver; +use crate::trace::inject_current_trace_into_meta; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use crate::utils::create_env_for_mcp_server; @@ -85,12 +87,14 @@ const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192; #[derive(Clone)] struct StreamableHttpResponseClient { - inner: reqwest::Client, + inner: CodexHttpClient, } impl StreamableHttpResponseClient { fn new(inner: reqwest::Client) -> Self { - Self { inner } + Self { + inner: CodexHttpClient::new(inner), + } } fn reqwest_error( @@ -722,6 +726,11 @@ impl RmcpClient { } None => None, }; + let meta = if matches!(&self.transport_recipe, TransportRecipe::Stdio { .. }) { + inject_current_trace_into_meta(meta) + } else { + meta + }; let rmcp_params = CallToolRequestParams { meta: None, name: name.into(), diff --git a/codex-rs/rmcp-client/src/trace.rs b/codex-rs/rmcp-client/src/trace.rs new file mode 100644 index 0000000000..42de1ab951 --- /dev/null +++ b/codex-rs/rmcp-client/src/trace.rs @@ -0,0 +1,36 @@ +use codex_client::current_trace_headers; + +pub(crate) const X_CODEX_TRACEPARENT_META_KEY: &str = "x-codex-traceparent"; +pub(crate) const X_CODEX_TRACESTATE_META_KEY: &str = "x-codex-tracestate"; + +pub(crate) fn inject_current_trace_into_meta( + meta: Option, +) -> Option { + let headers = current_trace_headers(); + let traceparent = headers + .get("traceparent") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + let tracestate = headers + .get("tracestate") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + if traceparent.is_none() && tracestate.is_none() { + return meta; + } + + let mut meta = meta.unwrap_or_else(|| rmcp::model::Meta(Default::default())); + if let Some(traceparent) = traceparent { + meta.insert( + X_CODEX_TRACEPARENT_META_KEY.to_string(), + serde_json::Value::String(traceparent), + ); + } + if let Some(tracestate) = tracestate { + meta.insert( + X_CODEX_TRACESTATE_META_KEY.to_string(), + serde_json::Value::String(tracestate), + ); + } + Some(meta) +}