diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index aabaad812c..ae86e1ac43 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4283,6 +4283,7 @@ dependencies = [ "codex-http-client", "codex-keyring-store", "codex-network-proxy", + "codex-otel", "codex-protocol", "codex-secrets", "codex-utils-cargo-bin", diff --git a/codex-rs/core/tests/mcp_trace_propagation.rs b/codex-rs/core/tests/mcp_trace_propagation.rs new file mode 100644 index 0000000000..1024b90027 --- /dev/null +++ b/codex-rs/core/tests/mcp_trace_propagation.rs @@ -0,0 +1,205 @@ +//! Verifies request-local trace ancestry through real RMCP queues and HTTP headers. + +use std::time::Duration; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::McpProtocolMode; +use codex_rmcp_client::RmcpClient; +use futures::FutureExt; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::InMemorySpanExporter; +use opentelemetry_sdk::trace::SdkTracerProvider; +use pretty_assertions::assert_eq; +use rmcp::model::ClientCapabilities; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::RequestMetaObject; +use serde_json::Value; +use serde_json::json; +use tracing::Instrument; +use tracing_subscriber::prelude::*; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; + +#[tokio::test(flavor = "current_thread")] +async fn mcp_requests_preserve_trace_context_across_workers_and_continuations() -> anyhow::Result<()> +{ + let exported = InMemorySpanExporter::default(); + let telemetry = SdkTracerProvider::builder() + .with_simple_exporter(exported.clone()) + .build(); + let _subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(telemetry.tracer("mcp-trace-test"))) + .set_default(); + + for mode in [McpProtocolMode::Legacy, McpProtocolMode::V20260728] { + let modern = mode == McpProtocolMode::V20260728; + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(move |request: &Request| { + let body: Value = request.body_json().unwrap(); + let result = match body["method"].as_str().unwrap() { + "initialize" => json!({"protocolVersion":"2025-06-18", "capabilities":{"tools":{}}, "serverInfo":{"name":"test","version":"1"}}), + "notifications/initialized" => return ResponseTemplate::new(202), + "server/discover" => json!({"resultType":"complete", "supportedVersions":["2026-07-28"], "capabilities":{"tools":{}}, "_meta":{"io.modelcontextprotocol/serverInfo":{"name":"test","version":"1"}}, "ttlMs":0, "cacheScope":"private"}), + "tools/list" => json!({"resultType":"complete", "tools":[], "ttlMs":60000, "cacheScope":"private"}), + "tools/call" if modern && body["params"].get("requestState").is_none() => json!({"resultType":"input_required", "requestState":"continue"}), + "tools/call" => json!({"resultType":"complete", "content":[]}), + "test/custom" => json!({"resultType":"complete", "ok":true}), + other => panic!("unexpected method: {other}"), + }; + ResponseTemplate::new(200).set_body_json(json!({"jsonrpc":"2.0", "id":body["id"], "result":result})) + }).mount(&server).await; + let client = RmcpClient::new_streamable_http_client_with_protocol_mode( + "trace-test", + &server.uri(), + Some("test-bearer".to_string()), + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + mode, + ) + .await? + .with_read_only_tools(/*requires_read_only_tools*/ true); + client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("test", "1"), + ), + Some(Duration::from_secs(5)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: None, + meta: None, + }) + } + .boxed() + }), + ) + .await?; + + let first = tracing::info_span!(parent: None, "first_page"); + let second = tracing::info_span!(parent: None, "second_page"); + let first_trace = codex_otel::span_w3c_trace_context(&first).unwrap(); + let second_trace = codex_otel::span_w3c_trace_context(&second).unwrap(); + let mut first_params = PaginatedRequestParams::default(); + first_params.meta = Some(RequestMetaObject::from( + json!({"application":"preserved", "traceparent":"invalid", "tracestate":"stale=value"}) + .as_object() + .unwrap() + .clone(), + )); + let second_params = + PaginatedRequestParams::default().with_cursor(Some("page2".to_string())); + let (first_result, second_result) = tokio::join!( + client + .list_tools_with_connector_ids(Some(first_params), Some(Duration::from_secs(5))) + .instrument(first), + client + .list_tools(Some(second_params), Some(Duration::from_secs(5))) + .instrument(second), + ); + first_result?; + second_result?; + + let tool = tracing::info_span!(parent: None, "tool"); + let tool_trace = codex_otel::span_w3c_trace_context(&tool).unwrap(); + client + .call_tool( + "test".to_string(), + /*arguments*/ None, + Some(json!({"application":"tool-meta"})), + Some(Duration::from_secs(5)), + ) + .instrument(tool) + .await?; + let custom = tracing::info_span!(parent: None, "custom"); + let custom_trace = codex_otel::span_w3c_trace_context(&custom).unwrap(); + client.send_custom_request("test/custom", Some(json!({"argument":"preserved", "_meta":{"application":"custom-meta", "traceparent":"invalid", "tracestate":"stale=value"}}))) + .instrument(custom).await?; + client.shutdown().await; + telemetry.force_flush()?; + let spans = exported.get_finished_spans()?; + let requests = server.received_requests().await.unwrap(); + let requests = requests + .iter() + .filter_map(|request| { + let body: Value = request.body_json().unwrap(); + matches!( + body["method"].as_str(), + Some("tools/list" | "tools/call" | "test/custom") + ) + .then_some((request, body)) + }) + .collect::>(); + assert_eq!(requests.len(), if modern { 5 } else { 4 }); + for (request, body) in requests { + if matches!(body["method"].as_str(), Some("tools/list" | "tools/call")) { + assert_eq!(body["params"]["_meta"]["openai/readOnly"], true); + } + let expected = if body["method"] == "test/custom" { + assert_eq!(body["params"]["_meta"]["application"], "custom-meta"); + assert_eq!(body["params"]["argument"], "preserved"); + assert!(body["params"]["_meta"].get("tracestate").is_none()); + &custom_trace + } else if body["method"] == "tools/call" { + assert_eq!(body["params"]["_meta"]["application"], "tool-meta"); + &tool_trace + } else if body["params"]["cursor"] == "page2" { + &second_trace + } else { + assert_eq!(body["params"]["_meta"]["application"], "preserved"); + assert!(body["params"]["_meta"].get("tracestate").is_none()); + &first_trace + }; + let parent = expected + .traceparent + .as_deref() + .unwrap() + .split('-') + .collect::>(); + let metadata_parent = body["params"]["_meta"]["traceparent"] + .as_str() + .unwrap() + .split('-') + .collect::>(); + assert_eq!(metadata_parent[1], parent[1]); + if metadata_parent[2] != parent[2] { + let caller = spans + .iter() + .find(|span| span.span_context.span_id().to_string() == metadata_parent[2]) + .unwrap(); + assert_eq!(caller.name, "list_tools_with_connector_ids"); + assert_eq!(caller.parent_span_id.to_string(), parent[2]); + } + let header = request.headers.get("traceparent").unwrap().to_str()?; + let header = header.split('-').collect::>(); + assert_eq!(header[1], parent[1]); + let http = spans + .iter() + .find(|span| span.span_context.span_id().to_string() == header[2]) + .unwrap(); + assert_eq!(http.name, "codex.exec_server.http_request"); + let bridge = spans + .iter() + .find(|span| span.span_context.span_id() == http.parent_span_id) + .unwrap(); + assert_eq!(bridge.name, "mcp.http.request"); + assert_eq!(bridge.parent_span_id.to_string(), metadata_parent[2]); + } + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index c3c0f07db3..c661aa0ed8 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -21,6 +21,7 @@ codex-exec-server = { workspace = true } codex-http-client = { workspace = true } codex-keyring-store = { workspace = true } codex-network-proxy = { workspace = true } +codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-secrets = { workspace = true } codex-utils-path-uri = { workspace = true } diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index fee686286b..bc49bd06c5 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -54,6 +54,7 @@ use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; use sse_stream::Sse; use sse_stream::SseStream; use tokio::sync::oneshot; +use tracing::Instrument; use crate::bounded_stdio_transport::MAX_MCP_STDIO_LINE_BYTES; use crate::event_notification_transport::MAX_EVENT_NOTIFICATION_BYTES; @@ -218,7 +219,9 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { None }; - let body = serde_json::to_vec(&message).map_err(StreamableHttpError::Deserialize)?; + let body = serde_json::to_value(&message).map_err(StreamableHttpError::Deserialize)?; + let request_span = crate::trace_context::request_span(&body); + let body = serde_json::to_vec(&body).map_err(StreamableHttpError::Deserialize)?; let has_authorization_header = headers.contains_key(AUTHORIZATION); if let JsonRpcMessage::Notification(notification) = &message && let ClientNotification::CancelledNotification(cancelled) = ¬ification.notification @@ -233,16 +236,19 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { return Ok(StreamableHttpPostResponse::Accepted); } - let request = self.http_client.http_request_stream(HttpRequestParams { - method: "POST".to_string(), - url: uri.to_string(), - headers: protocol_headers(&headers), - body: Some(body.into()), - timeout_ms, - redirect_policy, - request_id: "buffered-request".to_string(), - stream_response: true, - }); + let request = self + .http_client + .http_request_stream(HttpRequestParams { + method: "POST".to_string(), + url: uri.to_string(), + headers: protocol_headers(&headers), + body: Some(body.into()), + timeout_ms, + redirect_policy, + request_id: "buffered-request".to_string(), + stream_response: true, + }) + .instrument(request_span); let response = if is_event_stream_request { tokio::time::timeout(EVENT_STREAM_RESPONSE_TIMEOUT, request) .await diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 51105cf535..ab270cc95f 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -27,6 +27,7 @@ mod service_error; mod startup_error; mod stdio_server_launcher; mod tool_input; +mod trace_context; mod user_verification; mod utils; mod www_authenticate; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 1bb4970bb4..408d1b53ff 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -31,7 +31,6 @@ use rmcp::model::ClientNotification; use rmcp::model::ClientRequest; use rmcp::model::ContentBlock; use rmcp::model::CustomNotification; -use rmcp::model::CustomRequest; use rmcp::model::ElicitRequestParams; use rmcp::model::ElicitResult; use rmcp::model::ElicitationAction; @@ -687,9 +686,10 @@ impl RmcpClient { pub async fn list_tools( &self, - mut params: Option, + params: Option, timeout: Option, ) -> Result { + let mut params = crate::trace_context::traced_pagination(params); self.refresh_oauth_if_needed().await?; if self.requires_read_only_tools { self.apply_read_only_tools_meta(&mut params.get_or_insert_default().meta); @@ -748,6 +748,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { + let params = crate::trace_context::traced_pagination(params); self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/list", timeout, move |service| { @@ -764,6 +765,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { + let params = crate::trace_context::traced_pagination(params); self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { @@ -780,6 +782,8 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { + let mut params = params; + params.meta = crate::trace_context::with_current_trace(params.meta); self.refresh_oauth_if_needed().await?; let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; let result = self @@ -850,6 +854,7 @@ impl RmcpClient { None => None, }; self.apply_read_only_tools_meta(&mut meta); + let meta = crate::trace_context::with_current_trace(meta); let mut rmcp_params = CallToolRequestParams::new(name); rmcp_params.arguments = arguments; let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; @@ -968,10 +973,9 @@ impl RmcpClient { .run_service_operation("requests/custom", timeout, move |service| { let params = params.clone(); async move { + let request = crate::trace_context::traced_custom_request(method, params); service - .send_request(ClientRequest::CustomRequest(CustomRequest::new( - method, params, - ))) + .send_request(ClientRequest::CustomRequest(request)) .await } .boxed() @@ -988,7 +992,7 @@ impl RmcpClient { ) -> Result { let service = self.service().await?; let (sender, notifications) = event_notification_channel(); - let mut request = CustomRequest::new("events/stream", params); + let mut request = crate::trace_context::traced_custom_request("events/stream", params); request.extensions.insert(sender); let handle = service .peer() diff --git a/codex-rs/rmcp-client/src/trace_context.rs b/codex-rs/rmcp-client/src/trace_context.rs new file mode 100644 index 0000000000..6ba1b6f835 --- /dev/null +++ b/codex-rs/rmcp-client/src/trace_context.rs @@ -0,0 +1,91 @@ +//! Carries request-local W3C snapshots through RMCP's background transport queues. +//! Metadata does not retain caller spans or change the SDK's cursor-based cache keys. + +use codex_protocol::protocol::W3cTraceContext; +use rmcp::model::CustomRequest; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::RequestMetaObject; +use serde_json::Value; +use tracing::Span; + +/// Reads metadata from an SDK-serialized message, where envelope extensions have +/// already been merged with params._meta and take precedence on key conflicts. +/// Unlike RMCP's GetMeta, this also includes metadata supplied in typed params. +fn merged_request_meta(message: &Value) -> Option<&serde_json::Map> { + message.pointer("/params/_meta").and_then(Value::as_object) +} + +pub(crate) fn with_current_trace(mut meta: Option) -> Option { + if let Some(trace) = codex_otel::current_span_w3c_trace_context() { + let meta = meta.get_or_insert_default(); + if let Some(traceparent) = trace.traceparent { + meta.set_traceparent(traceparent); + } + // A supplied tracestate must not be paired with a different caller's parent. + meta.remove("tracestate"); + if let Some(tracestate) = trace.tracestate { + meta.set_tracestate(tracestate); + } + } + meta +} + +pub(crate) fn traced_pagination( + mut params: Option, +) -> Option { + let meta = with_current_trace(params.as_mut().and_then(|params| params.meta.take())); + if params.is_none() && meta.is_none() { + return None; + } + let mut params = params.unwrap_or_default(); + params.meta = meta; + Some(params) +} + +pub(crate) fn traced_custom_request(method: &str, params: Option) -> CustomRequest { + let mut request = CustomRequest::new(method, params); + let existing_meta = request + .params + .as_mut() + .and_then(|params| params.get_mut("_meta")) + .and_then(Value::as_object_mut) + .map(std::mem::take) + .map(RequestMetaObject::from); + if let Some(meta) = with_current_trace(existing_meta) { + // Move valid metadata into the canonical SDK extension slot. Otherwise + // its merge could retain an old params.tracestate beside the new parent. + request.extensions.insert(meta); + } + request +} + +pub(crate) fn request_span(message: &Value) -> Span { + let Some(meta) = merged_request_meta(message) else { + return Span::none(); + }; + let trace = W3cTraceContext { + traceparent: meta + .get("traceparent") + .and_then(Value::as_str) + .map(str::to_owned), + tracestate: meta + .get("tracestate") + .and_then(Value::as_str) + .map(str::to_owned), + }; + let Some(context) = codex_otel::context_from_w3c_trace_context(&trace) else { + return Span::none(); + }; + let method = message + .get("method") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let span = tracing::info_span!("mcp.http.request", mcp.method = method); + // Set the parent before the span is entered; never reparent the shared worker. + codex_otel::set_parent_from_context(&span, context); + span +} + +#[cfg(test)] +#[path = "trace_context_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/trace_context_tests.rs b/codex-rs/rmcp-client/src/trace_context_tests.rs new file mode 100644 index 0000000000..0c15848882 --- /dev/null +++ b/codex-rs/rmcp-client/src/trace_context_tests.rs @@ -0,0 +1,45 @@ +//! Checks that trace metadata uses the SDK's envelope/params merge precedence. + +use super::merged_request_meta; +use pretty_assertions::assert_eq; +use rmcp::model::ClientRequest; +use rmcp::model::CustomRequest; +use rmcp::model::GetMeta; +use rmcp::model::ListToolsRequest; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::RequestMetaObject; +use serde_json::json; + +#[test] +fn merged_metadata_preserves_params_and_prefers_extensions() -> serde_json::Result<()> { + let params_meta = json!({"application": "preserved", "traceparent": "params-parent"}); + let mut params = PaginatedRequestParams::default(); + params.meta = Some(serde_json::from_value::( + params_meta.clone(), + )?); + let requests = [ + ClientRequest::ListToolsRequest(ListToolsRequest::with_param(params)), + ClientRequest::CustomRequest(CustomRequest::new( + "test/custom", + Some(json!({"_meta": params_meta})), + )), + ]; + for mut request in requests { + let params_only = serde_json::to_value(&request)?; + assert_eq!(merged_request_meta(¶ms_only), params_meta.as_object()); + + request.get_meta_mut().set_traceparent("extension-parent"); + request.get_meta_mut().set_tracestate("vendor=value"); + let merged = serde_json::to_value(&request)?; + assert_eq!( + merged_request_meta(&merged), + json!({ + "application": "preserved", + "traceparent": "extension-parent", + "tracestate": "vendor=value", + }) + .as_object(), + ); + } + Ok(()) +}