Skip to main content

aether_core/events/
trace_context.rs

1use rmcp::model::Meta;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4
5/// Metadata key carrying the W3C traceparent, in MCP request metadata and HTTP headers alike.
6pub const TRACEPARENT_KEY: &str = "traceparent";
7/// Metadata key carrying the W3C tracestate.
8pub const TRACESTATE_KEY: &str = "tracestate";
9
10/// W3C trace context linking telemetry across process boundaries.
11#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
12#[serde(rename_all = "camelCase", deny_unknown_fields)]
13pub struct TraceContext {
14    /// W3C traceparent header: version, trace ID, parent span ID, and trace flags.
15    pub traceparent: String,
16    /// Optional W3C tracestate header carrying vendor-specific trace state.
17    #[serde(default, skip_serializing_if = "Option::is_none")]
18    pub tracestate: Option<String>,
19}
20
21impl TraceContext {
22    /// Reads the trace context a caller attached to MCP request metadata via
23    /// [`to_meta`](Self::to_meta), if any.
24    pub fn from_meta(meta: &Meta) -> Option<Self> {
25        let traceparent = meta.0.get(TRACEPARENT_KEY)?.as_str()?.to_string();
26        let tracestate = meta.0.get(TRACESTATE_KEY).and_then(|value| value.as_str()).map(str::to_string);
27        Some(Self { traceparent, tracestate })
28    }
29
30    /// Attaches the trace context to MCP request metadata.
31    pub fn to_meta(&self) -> Meta {
32        let mut meta = Meta::new();
33        meta.0.insert(TRACEPARENT_KEY.to_string(), self.traceparent.clone().into());
34        if let Some(tracestate) = &self.tracestate {
35            meta.0.insert(TRACESTATE_KEY.to_string(), tracestate.clone().into());
36        }
37        meta
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn meta_round_trips_a_trace_context() {
47        let trace_context = TraceContext {
48            traceparent: "00-00112233445566778899aabbccddeeff-0123456789abcdef-01".to_string(),
49            tracestate: Some("vendor=value".to_string()),
50        };
51
52        assert_eq!(TraceContext::from_meta(&trace_context.to_meta()), Some(trace_context));
53    }
54
55    #[test]
56    fn meta_round_trips_without_tracestate() {
57        let trace_context = TraceContext {
58            traceparent: "00-00112233445566778899aabbccddeeff-0123456789abcdef-01".to_string(),
59            tracestate: None,
60        };
61
62        let meta = trace_context.to_meta();
63
64        assert_eq!(meta.0.len(), 1);
65        assert_eq!(TraceContext::from_meta(&meta), Some(trace_context));
66    }
67
68    #[test]
69    fn from_meta_returns_none_without_a_traceparent() {
70        assert_eq!(TraceContext::from_meta(&Meta::new()), None);
71    }
72}