Codex Handbook
rollout-trace/src/mcp.rs 99 lines
//! Hot-path helpers for correlating concrete MCP executions with rollout traces.//!//! Core decides when an MCP request is actually going to execute. The trace//! crate owns the globally unique ID, the trace event that preserves it in the//! reduced artifact, and the bridge-private MCP request metadata key.use crate::McpCallId;use serde_json::Value as JsonValue;const MCP_CALL_ID_META_KEY: &str = "codex_bridge_mcp_call_id";/// No-op capable handle for one concrete MCP backend call.#[derive(Clone, Debug)]pub struct McpCallTraceContext {    mcp_call_id: Option<McpCallId>,}impl McpCallTraceContext {    /// Builds a context that records nothing and leaves request metadata unchanged.    pub fn disabled() -> Self {        Self { mcp_call_id: None }    }    /// Builds the trace handle for one concrete MCP execution.    pub(crate) fn enabled(mcp_call_id: McpCallId) -> Self {        Self {            mcp_call_id: Some(mcp_call_id),        }    }    /// Returns the trace-owned MCP call ID when rollout tracing is enabled.    pub(crate) fn mcp_call_id(&self) -> Option<&str> {        self.mcp_call_id.as_deref()    }    /// Adds bridge-private MCP correlation metadata to one outgoing request.    pub fn add_request_meta(&self, meta: Option<JsonValue>) -> Option<JsonValue> {        let Some(mcp_call_id) = self.mcp_call_id() else {            return meta;        };        match meta {            Some(JsonValue::Object(mut map)) => {                map.insert(                    MCP_CALL_ID_META_KEY.to_string(),                    JsonValue::String(mcp_call_id.to_string()),                );                Some(JsonValue::Object(map))            }            None => {                let mut map = serde_json::Map::new();                map.insert(                    MCP_CALL_ID_META_KEY.to_string(),                    JsonValue::String(mcp_call_id.to_string()),                );                Some(JsonValue::Object(map))            }            // This should never happen but if it does then we'll fallback to            // a noop rather than any breaking behavior. The tracing is best            // effort after all.            Some(_) => meta,        }    }}#[cfg(test)]mod tests {    use serde_json::json;    use super::MCP_CALL_ID_META_KEY;    use super::McpCallTraceContext;    #[test]    fn disabled_mcp_trace_leaves_request_meta_unchanged() {        let meta = Some(json!({"source": "test"}));        assert_eq!(            McpCallTraceContext::disabled().add_request_meta(meta.clone()),            meta        );    }    #[test]    fn enabled_mcp_trace_adds_bridge_correlation_meta() {        let trace = McpCallTraceContext::enabled("mcp-call-id".to_string());        let meta = trace            .add_request_meta(Some(json!({"source": "test"})))            .expect("enabled trace keeps request metadata");        let object = meta            .as_object()            .expect("MCP request metadata remains an object");        assert_eq!(object["source"], json!("test"));        assert_eq!(            object[MCP_CALL_ID_META_KEY],            json!(trace.mcp_call_id().expect("enabled trace has an ID"))        );    }}