cognee_http_server/observability/
span_buffer_layer.rs

1//! `tracing::Layer` that captures every span into a [`SpanBuffer`].
2//!
3//! Trace ids are synthesized: `tracing` does not have native OTEL ids, so the
4//! layer assigns a fresh 32-char lowercase-hex `trace_id` per root span and
5//! propagates it to children via the span's local `extensions_mut()` slot.
6//! `parent_span_id` is taken from the parent `TraceCtx` so the buffer's view
7//! matches Python's exporter byte-for-byte.
8
9use std::time::SystemTime;
10
11use rand::RngCore;
12use tracing::field::{Field, Visit};
13use tracing::span::{Attributes, Id, Record};
14use tracing::{Event, Level, Subscriber};
15use tracing_subscriber::Layer;
16use tracing_subscriber::layer::Context;
17use tracing_subscriber::registry::LookupSpan;
18
19use super::redaction::redact_attributes;
20use super::span_buffer::{RecordedSpan, SpanBuffer, SpanStatus};
21
22/// Per-span context attached via `extensions_mut()`.
23#[derive(Clone, Debug)]
24struct TraceCtx {
25    trace_id: String,
26    span_id: String,
27    parent_span_id: Option<String>,
28    start_time_ns: u64,
29    attributes: serde_json::Map<String, serde_json::Value>,
30    status: SpanStatus,
31}
32
33/// `tracing` layer that captures every span into a [`SpanBuffer`].
34pub struct SpanBufferLayer {
35    buffer: SpanBuffer,
36}
37
38impl SpanBufferLayer {
39    /// Build a new layer feeding `buffer`.
40    pub fn new(buffer: SpanBuffer) -> Self {
41        Self { buffer }
42    }
43}
44
45impl<S> Layer<S> for SpanBufferLayer
46where
47    S: Subscriber + for<'a> LookupSpan<'a>,
48{
49    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
50        let span_id = random_hex(8);
51        let (trace_id, parent_span_id) =
52            match attrs.parent().and_then(|pid| ctx.span(pid)).or_else(|| {
53                if attrs.is_contextual() {
54                    ctx.lookup_current()
55                } else {
56                    None
57                }
58            }) {
59                Some(parent_ref) => {
60                    let exts = parent_ref.extensions();
61                    match exts.get::<TraceCtx>() {
62                        Some(parent_ctx) => (
63                            parent_ctx.trace_id.clone(),
64                            Some(parent_ctx.span_id.clone()),
65                        ),
66                        None => (random_hex(16), None),
67                    }
68                }
69                None => (random_hex(16), None),
70            };
71
72        let mut visitor = AttrCollector::default();
73        attrs.record(&mut visitor);
74
75        let trace_ctx = TraceCtx {
76            trace_id,
77            span_id,
78            parent_span_id,
79            start_time_ns: now_ns(),
80            attributes: visitor.into_map(),
81            status: SpanStatus::Unset,
82        };
83
84        if let Some(span_ref) = ctx.span(id) {
85            span_ref.extensions_mut().insert(trace_ctx);
86        }
87    }
88
89    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
90        if let Some(span_ref) = ctx.span(id) {
91            let mut exts = span_ref.extensions_mut();
92            if let Some(trace_ctx) = exts.get_mut::<TraceCtx>() {
93                let mut visitor = AttrCollector::default();
94                values.record(&mut visitor);
95                for (k, v) in visitor.into_map() {
96                    trace_ctx.attributes.insert(k, v);
97                }
98            }
99        }
100    }
101
102    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
103        // Promote the parent span's status to ERROR when an `error!` event
104        // fires inside it. Mirrors Python's exporter behavior.
105        if *event.metadata().level() != Level::ERROR {
106            return;
107        }
108        if let Some(span_ref) = ctx.event_span(event) {
109            let mut exts = span_ref.extensions_mut();
110            if let Some(trace_ctx) = exts.get_mut::<TraceCtx>() {
111                trace_ctx.status = SpanStatus::Error;
112            }
113        }
114    }
115
116    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
117        let Some(span_ref) = ctx.span(&id) else {
118            return;
119        };
120        let metadata = span_ref.metadata();
121        let trace_ctx_opt = span_ref.extensions_mut().remove::<TraceCtx>();
122        let Some(mut trace_ctx) = trace_ctx_opt else {
123            // Foreign span (created by another layer) — drop silently.
124            return;
125        };
126
127        let end_time_ns = now_ns();
128        let duration_ns = end_time_ns.saturating_sub(trace_ctx.start_time_ns);
129        let duration_ms = duration_ns as f64 / 1_000_000.0;
130
131        redact_attributes(&mut trace_ctx.attributes);
132
133        let recorded = RecordedSpan {
134            trace_id: trace_ctx.trace_id,
135            span_id: trace_ctx.span_id,
136            parent_span_id: trace_ctx.parent_span_id,
137            name: metadata.name().to_string(),
138            start_time_ns: trace_ctx.start_time_ns,
139            end_time_ns,
140            duration_ms,
141            status: if trace_ctx.status == SpanStatus::Unset {
142                // Python's exporter normalizes "unset on close" → OK.
143                SpanStatus::Ok
144            } else {
145                trace_ctx.status
146            },
147            attributes: trace_ctx.attributes,
148        };
149        self.buffer.record(recorded);
150    }
151}
152
153// ─── helpers ──────────────────────────────────────────────────────────────────
154
155fn now_ns() -> u64 {
156    SystemTime::UNIX_EPOCH
157        .elapsed()
158        .map(|d| d.as_nanos() as u64)
159        .unwrap_or(0)
160}
161
162fn random_hex(byte_len: usize) -> String {
163    // 16-byte buf covers both 16-byte trace ids and 8-byte span ids.
164    let mut buf = [0u8; 16];
165    rand::thread_rng().fill_bytes(&mut buf[..byte_len]);
166    buf[..byte_len].iter().map(|b| format!("{b:02x}")).collect()
167}
168
169#[derive(Default)]
170struct AttrCollector {
171    map: serde_json::Map<String, serde_json::Value>,
172}
173
174impl AttrCollector {
175    fn into_map(self) -> serde_json::Map<String, serde_json::Value> {
176        self.map
177    }
178}
179
180impl Visit for AttrCollector {
181    fn record_str(&mut self, field: &Field, value: &str) {
182        self.map.insert(
183            field.name().to_string(),
184            serde_json::Value::String(value.to_string()),
185        );
186    }
187
188    fn record_i64(&mut self, field: &Field, value: i64) {
189        self.map.insert(
190            field.name().to_string(),
191            serde_json::Value::Number(value.into()),
192        );
193    }
194
195    fn record_u64(&mut self, field: &Field, value: u64) {
196        self.map.insert(
197            field.name().to_string(),
198            serde_json::Value::Number(value.into()),
199        );
200    }
201
202    fn record_i128(&mut self, field: &Field, value: i128) {
203        self.map.insert(
204            field.name().to_string(),
205            serde_json::Value::String(value.to_string()),
206        );
207    }
208
209    fn record_u128(&mut self, field: &Field, value: u128) {
210        self.map.insert(
211            field.name().to_string(),
212            serde_json::Value::String(value.to_string()),
213        );
214    }
215
216    fn record_bool(&mut self, field: &Field, value: bool) {
217        self.map
218            .insert(field.name().to_string(), serde_json::Value::Bool(value));
219    }
220
221    fn record_f64(&mut self, field: &Field, value: f64) {
222        if let Some(num) = serde_json::Number::from_f64(value) {
223            self.map
224                .insert(field.name().to_string(), serde_json::Value::Number(num));
225        } else {
226            self.map.insert(
227                field.name().to_string(),
228                serde_json::Value::String(value.to_string()),
229            );
230        }
231    }
232
233    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
234        self.map.insert(
235            field.name().to_string(),
236            serde_json::Value::String(format!("{value:?}")),
237        );
238    }
239}
240
241#[cfg(test)]
242#[allow(
243    clippy::unwrap_used,
244    clippy::expect_used,
245    reason = "test code — panics are acceptable failures"
246)]
247mod tests {
248    use super::*;
249    use tracing::Level;
250    use tracing_subscriber::Registry;
251    use tracing_subscriber::layer::SubscriberExt;
252
253    #[test]
254    fn parent_and_children_share_trace_id() {
255        let buffer = SpanBuffer::default();
256        let layer = SpanBufferLayer::new(buffer.clone());
257        let subscriber = Registry::default().with(layer);
258
259        tracing::subscriber::with_default(subscriber, || {
260            let parent = tracing::span!(Level::INFO, "parent");
261            let _g = parent.enter();
262            {
263                let child = tracing::span!(Level::INFO, "child1");
264                let _gc = child.enter();
265            }
266            {
267                let child = tracing::span!(Level::INFO, "child2");
268                let _gc = child.enter();
269            }
270        });
271
272        let traces = buffer.all_traces();
273        assert_eq!(traces.len(), 1, "all spans share one trace");
274        let summary = &traces[0];
275        assert_eq!(summary.span_count, 3);
276        let trace_id = summary.trace_id.clone();
277        for s in &summary.spans {
278            assert_eq!(s.trace_id, trace_id, "every span uses same trace_id");
279        }
280        // Find the parent span and assert children reference it.
281        let parent_span = summary
282            .spans
283            .iter()
284            .find(|s| s.parent_span_id.is_none())
285            .expect("root present");
286        for s in &summary.spans {
287            if s.span_id != parent_span.span_id {
288                assert_eq!(
289                    s.parent_span_id.as_deref(),
290                    Some(parent_span.span_id.as_str())
291                );
292            }
293        }
294    }
295
296    #[test]
297    fn recorded_attributes_are_redacted() {
298        let buffer = SpanBuffer::default();
299        let layer = SpanBufferLayer::new(buffer.clone());
300        let subscriber = Registry::default().with(layer);
301
302        tracing::subscriber::with_default(subscriber, || {
303            let span = tracing::span!(
304                Level::INFO,
305                "request",
306                auth = "Authorization: Bearer eyJabc.def.ghi-very-long-jwt-1234567890"
307            );
308            let _g = span.enter();
309        });
310
311        let traces = buffer.all_traces();
312        assert_eq!(traces.len(), 1);
313        let span = traces[0]
314            .spans
315            .iter()
316            .find(|s| s.name == "request")
317            .expect("request span recorded");
318        let auth = span
319            .attributes
320            .get("auth")
321            .and_then(|v| v.as_str())
322            .unwrap_or("");
323        assert!(auth.contains("***REDACTED***"));
324        assert!(!auth.contains("ghi-very-long-jwt"));
325    }
326
327    #[test]
328    fn error_event_marks_status_error() {
329        let buffer = SpanBuffer::default();
330        let layer = SpanBufferLayer::new(buffer.clone());
331        let subscriber = Registry::default().with(layer);
332
333        tracing::subscriber::with_default(subscriber, || {
334            let span = tracing::span!(Level::INFO, "task");
335            let _g = span.enter();
336            tracing::error!("failed to do thing");
337        });
338
339        let traces = buffer.all_traces();
340        let span = traces[0]
341            .spans
342            .iter()
343            .find(|s| s.name == "task")
344            .expect("task span");
345        assert_eq!(span.status, SpanStatus::Error);
346    }
347}
cognee_http_server/observability/span_buffer_layer.rs

cognee_http_server/observability/
span_buffer_layer.rs
