openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-14 11:57:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reduce outbound spans to semantic HTTP/MCP calls

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
nicholasclark-openai
2026-03-24 11:09:10 -07:00
parent 4322d47f74
commit dc0b2cafbc
13 changed files with 92 additions and 931 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
codex-rs/codex-client/src/default_client.rs
View File

@@ -11,7 +11,6 @@ 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)]
@@ -38,13 +37,6 @@ impl CodexHttpClient {
self.request(Method::POST, url)
}
pub fn delete<U>(&self, url: U) -> CodexRequestBuilder
where
U: IntoUrl,
{
self.request(Method::DELETE, url)
}
pub fn request<U>(&self, method: Method, url: U) -> CodexRequestBuilder
where
U: IntoUrl,
@@ -119,26 +111,10 @@ impl CodexRequestBuilder {
}
pub async fn send(self) -> Result<Response, reqwest::Error> {
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();
let headers = 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,
@@ -152,17 +128,6 @@ 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,
@@ -189,19 +154,7 @@ impl<'a> Injector for HeaderMapInjector<'a> {
}
}
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 {
fn trace_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
global::get_text_map_propagator(|prop| {
prop.inject_context(
@@ -239,7 +192,7 @@ mod tests {
let _entered = span.enter();
let span_context = span.context().span().span_context().clone();
let headers = current_trace_headers();
let headers = trace_headers();
let extractor = HeaderMapExtractor(&headers);
let extracted = TraceContextPropagator::new().extract(&extractor);

1
codex-rs/codex-client/src/lib.rs
View File

@@ -19,7 +19,6 @@ 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;

142
codex-rs/core/src/client.rs
View File

@@ -71,6 +71,7 @@ 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;
@@ -554,42 +555,30 @@ impl ModelClient {
turn_metadata_header: Option<&str>,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
model: Option<&str>,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
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 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 websocket_connect_timeout = self.state.provider.websocket_connect_timeout();
let start = Instant::now();
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 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 error_message = result.as_ref().err().map(telemetry_api_error_message);
let response_debug = result
.as_ref()
@@ -670,7 +659,6 @@ impl ModelClient {
HeaderValue::from_static("true"),
);
}
headers.extend(build_current_trace_headers());
headers
}
}
@@ -897,7 +885,6 @@ impl ModelClientSession {
/*turn_metadata_header*/ None,
auth_context,
RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
/*model*/ None,
)
.await?;
self.websocket_session.connection = Some(connection);
@@ -930,7 +917,6 @@ 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,
@@ -954,7 +940,6 @@ impl ModelClientSession {
turn_metadata_header,
auth_context,
request_route_telemetry,
model,
)
.await
{
@@ -1129,6 +1114,7 @@ impl ModelClientSession {
service_tier: Option<ServiceTier>,
turn_metadata_header: Option<&str>,
warmup: bool,
request_trace: Option<W3cTraceContext>,
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.auth_manager.clone();
@@ -1138,7 +1124,6 @@ impl ModelClientSession {
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
let api_base_url = client_setup.api_provider.base_url.clone();
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
&client_setup.api_auth,
@@ -1155,6 +1140,17 @@ 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,
@@ -1166,7 +1162,6 @@ impl ModelClientSession {
request_route_telemetry: RequestRouteTelemetry::for_endpoint(
RESPONSES_ENDPOINT,
),
model: Some(&model_info.slug),
})
.await
{
@@ -1192,44 +1187,17 @@ impl ModelClientSession {
Err(err) => return Err(map_api_error(err)),
}
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 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 (stream, last_request_rx) =
map_response_stream(stream_result, session_telemetry.clone());
self.websocket_session.last_response_rx = Some(last_request_rx);
@@ -1300,6 +1268,7 @@ impl ModelClientSession {
service_tier,
turn_metadata_header,
/*warmup*/ true,
current_span_w3c_trace_context(),
)
.await
{
@@ -1343,6 +1312,7 @@ 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,
@@ -1353,6 +1323,7 @@ impl ModelClientSession {
service_tier,
turn_metadata_header,
/*warmup*/ false,
request_trace,
)
.await?
{
@@ -1404,24 +1375,6 @@ fn parse_turn_metadata_header(turn_metadata_header: Option<&str>) -> Option<Head
turn_metadata_header.and_then(|value| HeaderValue::from_str(value).ok())
}
fn build_current_trace_headers() -> 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<HashMap<String, String>> {
let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header)?;
let turn_metadata = turn_metadata_header.to_str().ok()?.to_string();
@@ -1612,7 +1565,6 @@ struct WebsocketConnectParams<'a> {
options: &'a ApiResponsesOptions,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
model: Option<&'a str>,
}
async fn handle_unauthorized(

56
codex-rs/core/src/network_trace.rs
View File

@@ -96,62 +96,6 @@ pub(crate) fn responses_http_request_span(
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,

50
codex-rs/core/tests/responses_headers.rs
View File

@@ -9,7 +9,6 @@ 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;
@@ -19,24 +18,11 @@ 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!();
@@ -149,42 +135,6 @@ 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!();

25
codex-rs/core/tests/suite/client_websockets.rs
View File

@@ -56,17 +56,6 @@ 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()
@@ -80,10 +69,7 @@ fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTr
.as_deref()
.expect("missing expected traceparent");
assert_eq!(
trace_id(Some(actual_traceparent)),
trace_id(Some(expected_traceparent))
);
assert_eq!(actual_traceparent, expected_traceparent);
assert_eq!(
client_metadata
.get(WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY)
@@ -206,15 +192,6 @@ 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);

160
codex-rs/core/tests/suite/otel.rs
View File

@@ -23,11 +23,9 @@ 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;
@@ -721,17 +719,8 @@ async fn record_responses_sets_span_fields_for_response_events() {
}
#[tokio::test(flavor = "current_thread")]
#[traced_test]
async fn responses_request_span_records_turn_correlation_fields() {
let buffer: &'static Mutex<Vec<u8>> = 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,
@@ -754,137 +743,22 @@ async fn responses_request_span_records_turn_correlation_fields() {
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<Vec<u8>> = 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}"
);
logs_assert(|lines: &[&str]| {
lines
.iter()
.find(|line| {
line.contains("responses_http.request{")
&& line.contains("otel.kind=\"client\"")
&& line.contains("transport=\"responses_http\"")
&& line.contains("conversation.id=")
&& line.contains("session.id=")
&& line.contains("turn.id=")
})
.map(|_| Ok(()))
.unwrap_or_else(|| {
Err("missing responses_http.request span with correlation fields".to_string())
})
});
}
#[tokio::test]

349
codex-rs/core/tests/suite/rmcp_client.rs
View File

@@ -14,7 +14,6 @@ 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;
@@ -37,7 +36,6 @@ 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;
@@ -50,23 +48,13 @@ use tokio::process::Child;
use tokio::process::Command;
use tokio::time::Instant;
use tokio::time::sleep;
use tracing::Instrument;
use tracing_test::traced_test;
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)]
#[traced_test]
async fn stdio_server_round_trip() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
@@ -202,148 +190,26 @@ async fn stdio_server_round_trip() -> anyhow::Result<()> {
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_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,
logs_assert(|lines: &[&str]| {
lines
.iter()
.find(|line| {
line.contains("mcp.tools.call{")
&& line.contains("otel.kind=\"client\"")
&& line.contains("rpc.system=\"jsonrpc\"")
&& line.contains("rpc.method=\"tools/call\"")
&& line.contains("mcp.server.name=\"rmcp\"")
&& line.contains("mcp.transport=\"stdio\"")
&& line.contains("tool.name=\"echo\"")
&& line.contains("tool.call_id=\"call-123\"")
&& line.contains("turn.id=")
})
.await?;
Ok::<_, anyhow::Error>(expected_trace)
}
.instrument(tracing::info_span!("rmcp.stdio.trace_request"))
.await?;
.map(|_| Ok(()))
.unwrap_or_else(|| Err("missing mcp.tools.call span with MCP fields".to_string()))
});
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(())
}
@@ -1008,183 +874,6 @@ 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]

51
codex-rs/rmcp-client/src/bin/test_stdio_server.rs
View File

@@ -26,10 +26,6 @@ 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<Vec<Tool>>,
@@ -52,7 +48,6 @@ 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()];
@@ -168,22 +163,6 @@ 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(),
@@ -384,36 +363,6 @@ impl ServerHandler for TestToolServer {
let args = Self::parse_call_args::<ImageScenarioArgs>(&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,

86
codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs
View File

@@ -52,26 +52,22 @@ struct TestToolServer {
tools: Arc<Vec<Tool>>,
resources: Arc<Vec<Resource>>,
resource_templates: Arc<Vec<ResourceTemplate>>,
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(request_observation_state: RequestObservationState) -> Self {
let tools = vec![Self::echo_tool(), Self::trace_inspect_tool()];
fn new() -> Self {
let tools = vec![Self::echo_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,
}
}
@@ -95,22 +91,6 @@ 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(),
@@ -149,17 +129,6 @@ struct SessionFailureState {
armed_failure: Arc<Mutex<Option<ArmedFailure>>>,
}
#[derive(Clone, Debug, Default)]
struct ObservedTraceHeaders {
traceparent: Option<String>,
tracestate: Option<String>,
}
#[derive(Clone, Default)]
struct RequestObservationState {
last_mcp_post_headers: Arc<Mutex<ObservedTraceHeaders>>,
}
#[derive(Clone, Debug)]
struct ArmedFailure {
status: StatusCode,
@@ -288,24 +257,6 @@ 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,
@@ -326,7 +277,6 @@ fn parse_bind_addr() -> Result<SocketAddr, Box<dyn std::error::Error>> {
async fn main() -> Result<(), Box<dyn std::error::Error>> {
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);
@@ -377,18 +327,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.nest_service(
"/mcp",
StreamableHttpService::new(
{
let request_observation_state = request_observation_state.clone();
move || Ok(TestToolServer::new(request_observation_state.clone()))
},
|| Ok(TestToolServer::new()),
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,
@@ -426,29 +369,6 @@ async fn require_bearer(
}
}
async fn capture_mcp_request_headers(
State(state): State<RequestObservationState>,
request: Request<Body>,
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<SessionFailureState>,
Json(request): Json<ArmSessionPostFailureRequest>,

1
codex-rs/rmcp-client/src/lib.rs
View File

@@ -4,7 +4,6 @@ mod oauth;
mod perform_oauth_login;
mod program_resolver;
mod rmcp_client;
mod trace;
mod utils;
pub use auth_status::StreamableHttpOAuthDiscovery;

13
codex-rs/rmcp-client/src/rmcp_client.rs
View File

@@ -9,7 +9,6 @@ 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;
@@ -74,7 +73,6 @@ 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;
@@ -87,14 +85,12 @@ const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192;
#[derive(Clone)]
struct StreamableHttpResponseClient {
inner: CodexHttpClient,
inner: reqwest::Client,
}
impl StreamableHttpResponseClient {
fn new(inner: reqwest::Client) -> Self {
Self {
inner: CodexHttpClient::new(inner),
}
Self { inner }
}
fn reqwest_error(
@@ -726,11 +722,6 @@ 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(),

36
codex-rs/rmcp-client/src/trace.rs
View File

@@ -1,36 +0,0 @@
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<rmcp::model::Meta>,
) -> Option<rmcp::model::Meta> {
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)
}