aether_telemetry/
trace_context.rs1use aether_core::events::{TRACEPARENT_KEY, TRACESTATE_KEY, TraceContext};
2use opentelemetry::Context;
3use opentelemetry::propagation::TextMapPropagator;
4use opentelemetry::trace::TraceContextExt;
5use opentelemetry_sdk::propagation::TraceContextPropagator;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
11#[serde(untagged, rename_all = "camelCase", deny_unknown_fields)]
12pub enum AgentTraceContext {
13 Parent(TraceContext),
15 Root {
17 #[serde(rename = "traceId")]
19 #[schemars(rename = "traceId")]
20 trace_id: String,
21 },
22}
23
24pub(crate) fn inject_trace_context(context: &Context) -> Option<TraceContext> {
27 if !context.span().span_context().is_valid() {
28 return None;
29 }
30
31 let mut carrier = HashMap::new();
32 TraceContextPropagator::new().inject_context(context, &mut carrier);
33 Some(TraceContext { traceparent: carrier.remove(TRACEPARENT_KEY)?, tracestate: carrier.remove(TRACESTATE_KEY) })
34}
35
36pub(crate) fn extract_trace_context(trace_context: &TraceContext) -> Option<Context> {
39 let mut carrier = HashMap::from([(TRACEPARENT_KEY.to_string(), trace_context.traceparent.clone())]);
40 carrier.extend(trace_context.tracestate.clone().map(|tracestate| (TRACESTATE_KEY.to_string(), tracestate)));
41 let context = TraceContextPropagator::new().extract(&carrier);
42 context.span().span_context().is_valid().then_some(context)
43}