mirror of
https://github.com/openai/codex.git
synced 2026-11-10 21:00:47 +00:00
Preserve MCP trace context when tracing is disabled (#60506)
## Why MCP requests need to retain inbound W3C trace context even when no OpenTelemetry tracing layer is installed. ## What changed - Carry validated request trace context through app-server MCP tasks, connection startup, and resource listing tasks, while preferring active span context. - Forward fallback `traceparent` and `tracestate` headers for MCP HTTP requests, including discovery and initialization. - Use HTTP trace headers as a fallback for exec-server RPC calls when active span context is unavailable. ## Testing Add an app-server integration test that disables trace exporting and verifies that `mcpServer/resource/read` forwards the inbound `traceparent` and `tracestate` to the MCP server. GitOrigin-RevId: c886a55ce0a3992f1b2e04614efec80a20cb4b49
This commit is contained in:
@@ -10,6 +10,7 @@ use codex_mcp::ema_auth_scope;
|
||||
use codex_mcp::resolve_oauth_callback;
|
||||
use codex_otel::auth_storage::AuthStorageOriginator;
|
||||
use codex_rmcp_client::EnterpriseOAuthLoginRequest;
|
||||
use codex_rmcp_client::McpRequestTrace;
|
||||
use std::future::Future;
|
||||
|
||||
use crate::thread_state::ThreadStateManager;
|
||||
@@ -394,7 +395,7 @@ impl McpRequestProcessor {
|
||||
.await;
|
||||
outgoing.send_result(request, result).await;
|
||||
};
|
||||
tokio::spawn(AuthStorageOriginator::current().scope(task));
|
||||
self.spawn_mcp_request(request_id, task).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -683,7 +684,7 @@ impl McpRequestProcessor {
|
||||
|
||||
if let Some(thread_id) = thread_id {
|
||||
let (_, thread) = self.load_thread(&thread_id).await?;
|
||||
let request_id = request_id.clone();
|
||||
let response_request_id = request_id.clone();
|
||||
|
||||
let task = async move {
|
||||
let origin_call_id = origin_call_id.filter(|_| is_codex_apps);
|
||||
@@ -695,10 +696,15 @@ impl McpRequestProcessor {
|
||||
}
|
||||
None => thread.read_mcp_resource(&server, resource_params).await,
|
||||
};
|
||||
Self::send_mcp_resource_read_response(outgoing, request_id, result, origin_call_id)
|
||||
.await;
|
||||
Self::send_mcp_resource_read_response(
|
||||
outgoing,
|
||||
response_request_id,
|
||||
result,
|
||||
origin_call_id,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
tokio::spawn(AuthStorageOriginator::current().scope(task));
|
||||
self.spawn_mcp_request(request_id, task).await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -722,7 +728,7 @@ impl McpRequestProcessor {
|
||||
// environment stdio MCPs must declare their own absolute cwd.
|
||||
let runtime_context =
|
||||
McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
|
||||
let request_id = request_id.clone();
|
||||
let response_request_id = request_id.clone();
|
||||
|
||||
let task = async move {
|
||||
let result = read_mcp_resource_without_thread(
|
||||
@@ -738,11 +744,14 @@ impl McpRequestProcessor {
|
||||
.await
|
||||
.and_then(|result| serde_json::to_value(result).map_err(anyhow::Error::from));
|
||||
Self::send_mcp_resource_read_response(
|
||||
outgoing, request_id, result, /*origin_call_id*/ None,
|
||||
outgoing,
|
||||
response_request_id,
|
||||
result,
|
||||
/*origin_call_id*/ None,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
tokio::spawn(AuthStorageOriginator::current().scope(task));
|
||||
self.spawn_mcp_request(request_id, task).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -778,7 +787,7 @@ impl McpRequestProcessor {
|
||||
let (_, thread) = self.load_thread(&thread_id).await?;
|
||||
ensure_direct_input_allowed(thread.as_ref()).await?;
|
||||
let meta = with_mcp_tool_call_thread_id_meta(params.meta, &thread_id);
|
||||
let request_id = request_id.clone();
|
||||
let response_request_id = request_id.clone();
|
||||
|
||||
let task = async move {
|
||||
let result = thread
|
||||
@@ -786,11 +795,25 @@ impl McpRequestProcessor {
|
||||
.await
|
||||
.map(McpServerToolCallResponse::from)
|
||||
.map_err(mcp_operation_error);
|
||||
outgoing.send_result(request_id, result).await;
|
||||
outgoing.send_result(response_request_id, result).await;
|
||||
};
|
||||
tokio::spawn(AuthStorageOriginator::current().scope(task));
|
||||
self.spawn_mcp_request(request_id, task).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_mcp_request<'a>(
|
||||
&'a self,
|
||||
request_id: &'a ConnectionRequestId,
|
||||
task: impl Future<Output = ()> + Send + 'static,
|
||||
) -> impl Future<Output = ()> + 'a {
|
||||
// Bound the dispatcher future before awaiting the request trace lookup.
|
||||
let task = Box::pin(task);
|
||||
async move {
|
||||
let trace =
|
||||
McpRequestTrace::from_w3c(self.outgoing.request_trace_context(request_id).await);
|
||||
tokio::spawn(AuthStorageOriginator::current().scope(trace.scope(task)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn mcp_operation_error(error: anyhow::Error) -> JSONRPCErrorError {
|
||||
|
||||
@@ -1651,7 +1651,7 @@ impl TestAppServer {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_jsonrpc_message(&mut self, message: JSONRPCMessage) -> anyhow::Result<()> {
|
||||
pub async fn send_jsonrpc_message(&mut self, message: JSONRPCMessage) -> anyhow::Result<()> {
|
||||
eprintln!("writing message to stdin: {message:?}");
|
||||
let Some(stdin) = self.stdin.as_mut() else {
|
||||
anyhow::bail!("mcp stdin closed");
|
||||
|
||||
@@ -1,9 +1,50 @@
|
||||
//! Public JSON-RPC coverage for caller metadata on a thread's stdio resource read.
|
||||
//! Public JSON-RPC coverage for caller metadata on MCP resource reads.
|
||||
|
||||
use super::*;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use core_test_support::stdio_server_bin;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn mcp_resource_read_forwards_rpc_trace_context() -> Result<()> {
|
||||
let responses = responses::start_mock_server().await;
|
||||
let (url, calls, server) = start_resource_apps_mcp_server().await?;
|
||||
let (_home, mut mcp) = start_resource_test_app_server_with_extra_config(
|
||||
&url,
|
||||
&responses.uri(),
|
||||
"[otel]\ntrace_exporter = \"none\"",
|
||||
ResourceTestEnvironment::Auto,
|
||||
)
|
||||
.await?;
|
||||
let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
|
||||
let tracestate = "vendor=origin";
|
||||
let id = RequestId::String("traced-read".to_string());
|
||||
mcp.send_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
|
||||
id: id.clone(),
|
||||
method: "mcpServer/resource/read".to_string(),
|
||||
params: Some(json!({"server": "codex_apps", "uri": TEST_RESOURCE_URI})),
|
||||
trace: Some(W3cTraceContext {
|
||||
traceparent: Some(traceparent.to_string()),
|
||||
tracestate: Some(tracestate.to_string()),
|
||||
}),
|
||||
}))
|
||||
.await?;
|
||||
timeout(
|
||||
DEFAULT_READ_TIMEOUT,
|
||||
mcp.read_stream_until_response_message(id),
|
||||
)
|
||||
.await??;
|
||||
server.abort();
|
||||
|
||||
let reads = calls.read_metadata.lock().expect("resource read metadata");
|
||||
let read = reads.last().expect("resource read reached the MCP server");
|
||||
assert_eq!(read.get("traceparent"), Some(&json!(traceparent)));
|
||||
assert_eq!(read.get("tracestate"), Some(&json!(tracestate)));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stdio_resource_receives_caller_metadata() -> Result<()> {
|
||||
let responses_server = responses::start_mock_server().await;
|
||||
|
||||
@@ -1097,7 +1097,8 @@ impl McpConnectionSet {
|
||||
// Dormant servers must not hold the initial startup summary open.
|
||||
deferred_workers.push(startup);
|
||||
} else {
|
||||
eager_workers.push(startup.boxed());
|
||||
let trace = codex_rmcp_client::McpRequestTrace::current();
|
||||
eager_workers.push(trace.scope(startup).boxed());
|
||||
}
|
||||
}
|
||||
let manager = Self {
|
||||
|
||||
@@ -12,6 +12,7 @@ use std::time::Duration;
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use codex_rmcp_client::McpRequestTrace;
|
||||
use rmcp::model::ListResourceTemplatesResult;
|
||||
use rmcp::model::ListResourcesResult;
|
||||
use rmcp::model::PaginatedRequestParams;
|
||||
@@ -57,7 +58,7 @@ impl McpConnectionSet {
|
||||
continue;
|
||||
};
|
||||
let timeout = view.tool_timeout;
|
||||
join_set.spawn(async move {
|
||||
let task = async move {
|
||||
let resources = collect_paginated("resources/list", timeout, |params| {
|
||||
let client = client.clone();
|
||||
async move {
|
||||
@@ -67,7 +68,8 @@ impl McpConnectionSet {
|
||||
})
|
||||
.await;
|
||||
(server_name, resources)
|
||||
});
|
||||
};
|
||||
join_set.spawn(McpRequestTrace::current().scope(task));
|
||||
}
|
||||
|
||||
let mut resources = HashMap::new();
|
||||
@@ -101,7 +103,7 @@ impl McpConnectionSet {
|
||||
continue;
|
||||
};
|
||||
let timeout = view.tool_timeout;
|
||||
join_set.spawn(async move {
|
||||
let task = async move {
|
||||
let templates = collect_paginated("resources/templates/list", timeout, |params| {
|
||||
let client = client.clone();
|
||||
async move {
|
||||
@@ -111,7 +113,8 @@ impl McpConnectionSet {
|
||||
})
|
||||
.await;
|
||||
(server_name, templates)
|
||||
});
|
||||
};
|
||||
join_set.spawn(McpRequestTrace::current().scope(task));
|
||||
}
|
||||
|
||||
let mut templates = HashMap::new();
|
||||
|
||||
@@ -103,7 +103,6 @@ use rmcp::model::Tool as RmcpTool;
|
||||
use tokio::sync::watch;
|
||||
use tokio::time::Instant as TokioInstant;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::Instrument;
|
||||
use tracing::instrument;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -447,8 +446,7 @@ impl ManagedClientStartup {
|
||||
outcome
|
||||
};
|
||||
originator
|
||||
.scope(Box::pin(startup))
|
||||
.in_current_span()
|
||||
.scope(codex_rmcp_client::McpRequestTrace::current().scope(Box::pin(startup)))
|
||||
.boxed()
|
||||
.shared()
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use futures::FutureExt;
|
||||
use futures::future::BoxFuture;
|
||||
use tokio::sync::mpsc;
|
||||
@@ -17,8 +18,10 @@ use crate::HttpClient;
|
||||
use crate::client::ExecServerClient;
|
||||
use crate::client::ExecServerError;
|
||||
use crate::protocol::HTTP_REQUEST_METHOD;
|
||||
use crate::protocol::HttpHeader;
|
||||
use crate::protocol::HttpRequestParams;
|
||||
use crate::protocol::HttpRequestResponse;
|
||||
use crate::rpc::RpcClient;
|
||||
|
||||
/// Maximum queued body frames per streamed HTTP response.
|
||||
const HTTP_BODY_DELTA_CHANNEL_CAPACITY: usize = 256;
|
||||
@@ -30,7 +33,8 @@ impl ExecServerClient {
|
||||
mut params: HttpRequestParams,
|
||||
) -> Result<HttpRequestResponse, ExecServerError> {
|
||||
params.stream_response = false;
|
||||
self.call(HTTP_REQUEST_METHOD, ¶ms).await
|
||||
let rpc_client = self.rpc_client().await?;
|
||||
self.call_http_request(&rpc_client, ¶ms).await
|
||||
}
|
||||
|
||||
/// Performs an HTTP request and returns a body stream.
|
||||
@@ -52,10 +56,7 @@ impl ExecServerClient {
|
||||
.await?;
|
||||
let mut registration =
|
||||
HttpBodyStreamRegistration::new(Arc::clone(&self.inner), request_id.clone());
|
||||
let response = match self
|
||||
.call_rpc(&rpc_client, HTTP_REQUEST_METHOD, ¶ms)
|
||||
.await
|
||||
{
|
||||
let response = match self.call_http_request(&rpc_client, ¶ms).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
self.inner.remove_http_body_stream(&request_id).await;
|
||||
@@ -69,6 +70,34 @@ impl ExecServerClient {
|
||||
HttpResponseBodyStream::remote(Arc::clone(&self.inner), request_id, rx),
|
||||
))
|
||||
}
|
||||
|
||||
async fn call_http_request(
|
||||
&self,
|
||||
rpc_client: &RpcClient,
|
||||
params: &HttpRequestParams,
|
||||
) -> Result<HttpRequestResponse, ExecServerError> {
|
||||
let trace = trace_from_http_headers(¶ms.headers);
|
||||
self.map_rpc_call_result(
|
||||
rpc_client
|
||||
.call_with_trace_fallback(HTTP_REQUEST_METHOD, params, trace)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn trace_from_http_headers(headers: &[HttpHeader]) -> Option<W3cTraceContext> {
|
||||
let header = |name: &str| {
|
||||
headers
|
||||
.iter()
|
||||
.find(|header| header.name.eq_ignore_ascii_case(name) && header.value_env_var.is_none())
|
||||
.map(|header| header.value.clone())
|
||||
};
|
||||
let trace = W3cTraceContext {
|
||||
traceparent: Some(header("traceparent")?),
|
||||
tracestate: header("tracestate"),
|
||||
};
|
||||
codex_otel::context_from_w3c_trace_context(&trace)
|
||||
.and_then(|context| codex_otel::w3c_trace_context_from_context(&context))
|
||||
}
|
||||
|
||||
impl HttpClient for ExecServerClient {
|
||||
|
||||
@@ -608,6 +608,16 @@ impl RpcClient {
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn call<P, T>(&self, method: &str, params: &P) -> Result<T, RpcCallError>
|
||||
where
|
||||
P: Serialize,
|
||||
T: DeserializeOwned,
|
||||
{
|
||||
self.call_with_trace_fallback(method, params, /*trace_fallback*/ None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Uses the supplied carrier only when this process cannot create an active W3C span.
|
||||
#[tracing::instrument(
|
||||
name = "codex.exec_server.request",
|
||||
level = "info",
|
||||
@@ -618,12 +628,20 @@ impl RpcClient {
|
||||
method,
|
||||
)
|
||||
)]
|
||||
pub(crate) async fn call<P, T>(&self, method: &str, params: &P) -> Result<T, RpcCallError>
|
||||
pub(crate) async fn call_with_trace_fallback<P, T>(
|
||||
&self,
|
||||
method: &str,
|
||||
params: &P,
|
||||
trace_fallback: Option<W3cTraceContext>,
|
||||
) -> Result<T, RpcCallError>
|
||||
where
|
||||
P: Serialize,
|
||||
T: DeserializeOwned,
|
||||
{
|
||||
self.call_untraced(method, params).await
|
||||
record_client_request(self.metrics.as_ref(), method);
|
||||
let _call_slot = self.acquire_regular_call_slot()?;
|
||||
self.call_inner(method, params, RpcCallTimeout::None, trace_fallback)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Send one request without creating the standard request span.
|
||||
@@ -641,7 +659,13 @@ impl RpcClient {
|
||||
{
|
||||
record_client_request(self.metrics.as_ref(), method);
|
||||
let _call_slot = self.acquire_regular_call_slot()?;
|
||||
self.call_inner(method, params, RpcCallTimeout::None).await
|
||||
self.call_inner(
|
||||
method,
|
||||
params,
|
||||
RpcCallTimeout::None,
|
||||
/*trace_fallback*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn call_with_timeout<P, T>(
|
||||
@@ -656,8 +680,13 @@ impl RpcClient {
|
||||
{
|
||||
record_client_request(self.metrics.as_ref(), method);
|
||||
let _call_slot = self.acquire_regular_call_slot()?;
|
||||
self.call_inner(method, params, RpcCallTimeout::After(call_timeout))
|
||||
.await
|
||||
self.call_inner(
|
||||
method,
|
||||
params,
|
||||
RpcCallTimeout::After(call_timeout),
|
||||
/*trace_fallback*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
@@ -690,7 +719,13 @@ impl RpcClient {
|
||||
}
|
||||
},
|
||||
};
|
||||
self.call_inner(method, params, RpcCallTimeout::None).await
|
||||
self.call_inner(
|
||||
method,
|
||||
params,
|
||||
RpcCallTimeout::None,
|
||||
/*trace_fallback*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn call_inner<P, T>(
|
||||
@@ -698,6 +733,7 @@ impl RpcClient {
|
||||
method: &str,
|
||||
params: &P,
|
||||
call_timeout: RpcCallTimeout,
|
||||
trace_fallback: Option<W3cTraceContext>,
|
||||
) -> Result<T, RpcCallError>
|
||||
where
|
||||
P: Serialize,
|
||||
@@ -725,7 +761,14 @@ impl RpcClient {
|
||||
};
|
||||
let params = serde_json::to_value(params).map_err(RpcCallError::Json)?;
|
||||
let response = self
|
||||
.send_request_and_wait(method, request_id, params, response_rx, call_timeout)
|
||||
.send_request_and_wait(
|
||||
method,
|
||||
request_id,
|
||||
params,
|
||||
response_rx,
|
||||
call_timeout,
|
||||
trace_fallback,
|
||||
)
|
||||
.await?;
|
||||
serde_json::from_value(response).map_err(RpcCallError::Json)
|
||||
}
|
||||
@@ -740,6 +783,7 @@ impl RpcClient {
|
||||
params: Value,
|
||||
response_rx: oneshot::Receiver<RpcCompletion>,
|
||||
call_timeout: RpcCallTimeout,
|
||||
trace_fallback: Option<W3cTraceContext>,
|
||||
) -> Result<Value, RpcCallError> {
|
||||
let enqueue_started = Instant::now();
|
||||
if self
|
||||
@@ -748,7 +792,7 @@ impl RpcClient {
|
||||
id: request_id.clone(),
|
||||
method: method.to_string(),
|
||||
params: Some(params),
|
||||
trace: codex_otel::current_span_w3c_trace_context(),
|
||||
trace: codex_otel::current_span_w3c_trace_context().or(trace_fallback),
|
||||
}))
|
||||
.await
|
||||
.is_err()
|
||||
|
||||
@@ -47,6 +47,7 @@ pub use crate::trace_context::span_w3c_trace_context;
|
||||
pub use crate::trace_context::traceparent_context_from_env;
|
||||
pub use crate::trace_context::validate_tracestate_entries;
|
||||
pub use crate::trace_context::validate_tracestate_member;
|
||||
pub use crate::trace_context::w3c_trace_context_from_context;
|
||||
pub use codex_utils_string::sanitize_metric_tag_value;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Display)]
|
||||
|
||||
@@ -31,13 +31,17 @@ pub fn current_span_w3c_trace_context() -> Option<W3cTraceContext> {
|
||||
}
|
||||
|
||||
pub fn span_w3c_trace_context(span: &Span) -> Option<W3cTraceContext> {
|
||||
let context = span.context();
|
||||
w3c_trace_context_from_context(&span.context())
|
||||
}
|
||||
|
||||
/// Serializes a valid OpenTelemetry context into a W3C trace carrier.
|
||||
pub fn w3c_trace_context_from_context(context: &Context) -> Option<W3cTraceContext> {
|
||||
if !context.span().span_context().is_valid() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut headers = HashMap::new();
|
||||
TraceContextPropagator::new().inject_context(&context, &mut headers);
|
||||
TraceContextPropagator::new().inject_context(context, &mut headers);
|
||||
let tracestate = headers.remove("tracestate");
|
||||
let configured_tracestate_guard = tracestate_entries()
|
||||
.read()
|
||||
|
||||
@@ -23,6 +23,7 @@ use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRedirectPolicy;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use codex_exec_server::HttpResponseBodyStream;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use futures::StreamExt;
|
||||
use futures::stream;
|
||||
use futures::stream::BoxStream;
|
||||
@@ -84,6 +85,8 @@ pub(crate) struct StreamableHttpClientAdapter {
|
||||
has_configured_headers: bool,
|
||||
redirect_mode: StreamableHttpRedirectMode,
|
||||
initialize_deadline: Arc<Mutex<Option<Instant>>>,
|
||||
// The SDK sends handshakes on its worker without request metadata.
|
||||
handshake_trace: Option<W3cTraceContext>,
|
||||
}
|
||||
|
||||
struct EventStreamCancellation {
|
||||
@@ -131,6 +134,7 @@ impl StreamableHttpClientAdapter {
|
||||
has_configured_headers,
|
||||
redirect_mode,
|
||||
initialize_deadline,
|
||||
handshake_trace: crate::trace_context::current_w3c_trace_context(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -235,7 +239,17 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
};
|
||||
|
||||
let body = serde_json::to_value(&message).map_err(StreamableHttpError::Deserialize)?;
|
||||
let request_span = crate::trace_context::request_span(&body);
|
||||
let handshake_trace = if is_discovery_request || mcp_method.as_deref() == Some("initialize")
|
||||
{
|
||||
self.handshake_trace.as_ref()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let request_span = crate::trace_context::request_span_and_trace_headers(
|
||||
&body,
|
||||
&mut headers,
|
||||
handshake_trace,
|
||||
);
|
||||
let body = serde_json::to_vec(&body).map_err(StreamableHttpError::Deserialize)?;
|
||||
let has_authorization_header = headers.contains_key(AUTHORIZATION);
|
||||
|
||||
|
||||
@@ -118,4 +118,5 @@ pub use stdio_server_launcher::ExecutorStdioServerLauncher;
|
||||
pub use stdio_server_launcher::LocalStdioServerLauncher;
|
||||
pub use stdio_server_launcher::StdioServerLauncher;
|
||||
pub use tool_list_changes::ToolListRevision;
|
||||
pub use trace_context::McpRequestTrace;
|
||||
pub use user_verification::UserVerificationReason;
|
||||
|
||||
@@ -1,13 +1,60 @@
|
||||
//! 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 std::future::Future;
|
||||
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use http::HeaderMap;
|
||||
use http::HeaderValue;
|
||||
use rmcp::model::CustomRequest;
|
||||
use rmcp::model::PaginatedRequestParams;
|
||||
use rmcp::model::RequestMetaObject;
|
||||
use serde_json::Value;
|
||||
use tracing::Span;
|
||||
|
||||
tokio::task_local! {
|
||||
static CURRENT_REQUEST_TRACE: McpRequestTrace;
|
||||
}
|
||||
|
||||
/// A request-scoped W3C fallback for MCP work when no tracing layer is installed.
|
||||
#[derive(Clone)]
|
||||
pub struct McpRequestTrace {
|
||||
trace: Option<W3cTraceContext>,
|
||||
parent: Span,
|
||||
}
|
||||
|
||||
impl McpRequestTrace {
|
||||
/// Retains only a valid inbound carrier; active spans take precedence over this fallback.
|
||||
pub fn from_w3c(trace: Option<W3cTraceContext>) -> Self {
|
||||
Self {
|
||||
trace: trace
|
||||
.as_ref()
|
||||
.and_then(codex_otel::context_from_w3c_trace_context)
|
||||
.and_then(|context| codex_otel::w3c_trace_context_from_context(&context)),
|
||||
parent: Span::current(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Captures the fallback for an explicitly spawned child task.
|
||||
pub fn current() -> Self {
|
||||
Self {
|
||||
trace: CURRENT_REQUEST_TRACE
|
||||
.try_with(|current| current.trace.clone())
|
||||
.unwrap_or_default(),
|
||||
parent: Span::current(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "mcp.request", parent = &self.parent, skip_all)]
|
||||
pub async fn scope<F: Future>(self, future: F) -> F::Output {
|
||||
CURRENT_REQUEST_TRACE.scope(self, future).await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn current_w3c_trace_context() -> Option<W3cTraceContext> {
|
||||
codex_otel::current_span_w3c_trace_context().or_else(|| McpRequestTrace::current().trace)
|
||||
}
|
||||
|
||||
/// 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.
|
||||
@@ -16,7 +63,7 @@ fn merged_request_meta(message: &Value) -> Option<&serde_json::Map<String, Value
|
||||
}
|
||||
|
||||
pub(crate) fn with_current_trace(mut meta: Option<RequestMetaObject>) -> Option<RequestMetaObject> {
|
||||
if let Some(trace) = codex_otel::current_span_w3c_trace_context() {
|
||||
if let Some(trace) = current_w3c_trace_context() {
|
||||
let meta = meta.get_or_insert_default();
|
||||
if let Some(traceparent) = trace.traceparent {
|
||||
meta.set_traceparent(traceparent);
|
||||
@@ -59,11 +106,12 @@ pub(crate) fn traced_custom_request(method: &str, params: Option<Value>) -> Cust
|
||||
request
|
||||
}
|
||||
|
||||
pub(crate) fn request_span(message: &Value) -> Span {
|
||||
let Some(meta) = merged_request_meta(message) else {
|
||||
return Span::none();
|
||||
};
|
||||
let trace = W3cTraceContext {
|
||||
pub(crate) fn request_span_and_trace_headers(
|
||||
message: &Value,
|
||||
headers: &mut HeaderMap,
|
||||
handshake_trace: Option<&W3cTraceContext>,
|
||||
) -> Span {
|
||||
let trace = merged_request_meta(message).map(|meta| W3cTraceContext {
|
||||
traceparent: meta
|
||||
.get("traceparent")
|
||||
.and_then(Value::as_str)
|
||||
@@ -72,8 +120,12 @@ pub(crate) fn request_span(message: &Value) -> Span {
|
||||
.get("tracestate")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned),
|
||||
};
|
||||
let Some(context) = codex_otel::context_from_w3c_trace_context(&trace) else {
|
||||
});
|
||||
let Some(context) = trace
|
||||
.as_ref()
|
||||
.and_then(codex_otel::context_from_w3c_trace_context)
|
||||
.or_else(|| handshake_trace.and_then(codex_otel::context_from_w3c_trace_context))
|
||||
else {
|
||||
return Span::none();
|
||||
};
|
||||
let method = message
|
||||
@@ -82,7 +134,22 @@ pub(crate) fn request_span(message: &Value) -> Span {
|
||||
.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);
|
||||
codex_otel::set_parent_from_context(&span, context.clone());
|
||||
if codex_otel::span_w3c_trace_context(&span).is_none()
|
||||
&& let Some(trace) = codex_otel::w3c_trace_context_from_context(&context)
|
||||
{
|
||||
// Without a tracing layer, the HTTP client cannot derive headers from the bridge span.
|
||||
for (name, value) in [
|
||||
("traceparent", trace.traceparent),
|
||||
("tracestate", trace.tracestate),
|
||||
] {
|
||||
if let Some(value) = value.and_then(|value| HeaderValue::from_str(&value).ok()) {
|
||||
headers.insert(name, value);
|
||||
} else {
|
||||
headers.remove(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
span
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user