use crate::McpCallId;
use serde_json::Value as JsonValue;
const MCP_CALL_ID_META_KEY: &str = "codex_bridge_mcp_call_id";
#[derive(Clone, Debug)]
pub struct McpCallTraceContext {
mcp_call_id: Option<McpCallId>,
}
impl McpCallTraceContext {
pub fn disabled() -> Self {
Self { mcp_call_id: None }
}
pub(crate) fn enabled(mcp_call_id: McpCallId) -> Self {
Self {
mcp_call_id: Some(mcp_call_id),
}
}
pub(crate) fn mcp_call_id(&self) -> Option<&str> {
self.mcp_call_id.as_deref()
}
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))
}
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"))
);
}
}