codex_rollout_trace/
mcp.rs

1//! Hot-path helpers for correlating concrete MCP executions with rollout traces.
2//!
3//! Core decides when an MCP request is actually going to execute. The trace
4//! crate owns the globally unique ID, the trace event that preserves it in the
5//! reduced artifact, and the bridge-private MCP request metadata key.
6
7use crate::McpCallId;
8use serde_json::Value as JsonValue;
9
10const MCP_CALL_ID_META_KEY: &str = "codex_bridge_mcp_call_id";
11
12/// No-op capable handle for one concrete MCP backend call.
13#[derive(Clone, Debug)]
14pub struct McpCallTraceContext {
15    mcp_call_id: Option<McpCallId>,
16}
17
18impl McpCallTraceContext {
19    /// Builds a context that records nothing and leaves request metadata unchanged.
20    pub fn disabled() -> Self {
21        Self { mcp_call_id: None }
22    }
23
24    /// Builds the trace handle for one concrete MCP execution.
25    pub(crate) fn enabled(mcp_call_id: McpCallId) -> Self {
26        Self {
27            mcp_call_id: Some(mcp_call_id),
28        }
29    }
30
31    /// Returns the trace-owned MCP call ID when rollout tracing is enabled.
32    pub(crate) fn mcp_call_id(&self) -> Option<&str> {
33        self.mcp_call_id.as_deref()
34    }
35
36    /// Adds bridge-private MCP correlation metadata to one outgoing request.
37    pub fn add_request_meta(&self, meta: Option<JsonValue>) -> Option<JsonValue> {
38        let Some(mcp_call_id) = self.mcp_call_id() else {
39            return meta;
40        };
41
42        match meta {
43            Some(JsonValue::Object(mut map)) => {
44                map.insert(
45                    MCP_CALL_ID_META_KEY.to_string(),
46                    JsonValue::String(mcp_call_id.to_string()),
47                );
48                Some(JsonValue::Object(map))
49            }
50            None => {
51                let mut map = serde_json::Map::new();
52                map.insert(
53                    MCP_CALL_ID_META_KEY.to_string(),
54                    JsonValue::String(mcp_call_id.to_string()),
55                );
56                Some(JsonValue::Object(map))
57            }
58            // This should never happen but if it does then we'll fallback to
59            // a noop rather than any breaking behavior. The tracing is best
60            // effort after all.
61            Some(_) => meta,
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use serde_json::json;
69
70    use super::MCP_CALL_ID_META_KEY;
71    use super::McpCallTraceContext;
72
73    #[test]
74    fn disabled_mcp_trace_leaves_request_meta_unchanged() {
75        let meta = Some(json!({"source": "test"}));
76
77        assert_eq!(
78            McpCallTraceContext::disabled().add_request_meta(meta.clone()),
79            meta
80        );
81    }
82
83    #[test]
84    fn enabled_mcp_trace_adds_bridge_correlation_meta() {
85        let trace = McpCallTraceContext::enabled("mcp-call-id".to_string());
86        let meta = trace
87            .add_request_meta(Some(json!({"source": "test"})))
88            .expect("enabled trace keeps request metadata");
89        let object = meta
90            .as_object()
91            .expect("MCP request metadata remains an object");
92
93        assert_eq!(object["source"], json!("test"));
94        assert_eq!(
95            object[MCP_CALL_ID_META_KEY],
96            json!(trace.mcp_call_id().expect("enabled trace has an ID"))
97        );
98    }
99}
codex_rollout_trace/mcp.rs

codex_rollout_trace/
mcp.rs
