openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add outbound spans for responses and MCP

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