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 {
|
||||
|
||||
Reference in New Issue
Block a user