cognee_http_server/observability/
span_buffer.rs

1//! Bounded in-memory ring buffer of recorded spans.
2//!
3//! Mirrors Python's `CogneeSpanExporter` storage model
4//! (`cognee/modules/observability/tracing.py`): per-trace buckets, LRU
5//! eviction once `max_traces` is exceeded, no per-span eviction beyond a
6//! safety cap.
7
8use std::collections::{HashMap, VecDeque};
9use std::env;
10use std::sync::{Arc, Mutex};
11
12use serde::{Deserialize, Serialize};
13
14/// Span status as exported via the wire shape.
15///
16/// Serialized in UPPERCASE so JSON roundtrips line up with Python's exporter
17/// (`"OK" | "ERROR" | "UNSET"`).
18#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "UPPERCASE")]
20pub enum SpanStatus {
21    #[default]
22    Unset,
23    Ok,
24    Error,
25}
26
27impl SpanStatus {
28    /// Wire string per the Python exporter (`"UNSET" | "OK" | "ERROR"`).
29    pub fn as_str(&self) -> &'static str {
30        match self {
31            Self::Unset => "UNSET",
32            Self::Ok => "OK",
33            Self::Error => "ERROR",
34        }
35    }
36}
37
38/// A single span snapshot, frozen at the moment its tracing span closed.
39///
40/// Field shape mirrors the Python `CogneeSpanExporter.export(...)` dict
41/// byte-for-byte so frontend trace viewers render identically.
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct RecordedSpan {
44    pub trace_id: String,
45    pub span_id: String,
46    pub parent_span_id: Option<String>,
47    pub name: String,
48    pub start_time_ns: u64,
49    pub end_time_ns: u64,
50    pub duration_ms: f64,
51    pub status: SpanStatus,
52    pub attributes: serde_json::Map<String, serde_json::Value>,
53}
54
55/// Configurable caps for the ring buffer.
56#[derive(Debug, Clone)]
57pub struct BufferConfig {
58    /// How many distinct traces to retain before LRU eviction kicks in.
59    pub max_traces: usize,
60    /// Per-trace span cap (defense against pathological producers).
61    pub max_spans_per_trace: usize,
62}
63
64impl Default for BufferConfig {
65    fn default() -> Self {
66        Self {
67            max_traces: 50,
68            max_spans_per_trace: 1024,
69        }
70    }
71}
72
73impl BufferConfig {
74    /// Read from `COGNEE_SPAN_BUFFER_MAX_TRACES` /
75    /// `COGNEE_SPAN_BUFFER_MAX_SPANS_PER_TRACE`. Invalid values fall back to
76    /// the defaults.
77    pub fn from_env() -> Self {
78        let mut cfg = Self::default();
79        if let Ok(v) = env::var("COGNEE_SPAN_BUFFER_MAX_TRACES")
80            && let Ok(n) = v.parse::<usize>()
81        {
82            cfg.max_traces = n;
83        }
84        if let Ok(v) = env::var("COGNEE_SPAN_BUFFER_MAX_SPANS_PER_TRACE")
85            && let Ok(n) = v.parse::<usize>()
86        {
87            cfg.max_spans_per_trace = n;
88        }
89        cfg
90    }
91}
92
93/// Per-trace summary surfaced via `/api/v1/activity/spans`.
94#[derive(Debug, Clone)]
95pub struct TraceSummary {
96    pub trace_id: String,
97    pub root_name: Option<String>,
98    pub duration_ms: f64,
99    pub span_count: usize,
100    pub status: Option<SpanStatus>,
101    pub spans: Vec<RecordedSpan>,
102}
103
104/// Lightweight stats surfaced alongside the span list.
105#[derive(Debug, Clone, Default)]
106pub struct BufferStats {
107    /// Number of spans dropped because the per-trace cap was exceeded.
108    pub dropped_overflow: u64,
109    /// Number of traces evicted via the LRU cap.
110    pub dropped_lru: u64,
111}
112
113struct BufferInner {
114    traces: HashMap<String, Vec<RecordedSpan>>,
115    /// Trace ids in insertion order; oldest at the front.
116    trace_order: VecDeque<String>,
117    stats: BufferStats,
118}
119
120/// Bounded in-memory span buffer.
121///
122/// Cheap to clone (interior `Arc<Mutex<...>>`).
123#[derive(Clone)]
124pub struct SpanBuffer {
125    inner: Arc<Mutex<BufferInner>>,
126    config: Arc<BufferConfig>,
127}
128
129impl Default for SpanBuffer {
130    fn default() -> Self {
131        Self::new(BufferConfig::default())
132    }
133}
134
135impl SpanBuffer {
136    /// Build a new buffer with the supplied config.
137    pub fn new(config: BufferConfig) -> Self {
138        Self {
139            inner: Arc::new(Mutex::new(BufferInner {
140                traces: HashMap::new(),
141                trace_order: VecDeque::new(),
142                stats: BufferStats::default(),
143            })),
144            config: Arc::new(config),
145        }
146    }
147
148    /// Record one span.
149    ///
150    /// On overflow:
151    /// - per-trace: drops the *new* span silently and bumps `dropped_overflow`.
152    /// - cross-trace: when the trace is brand-new and pushes the count past
153    ///   `max_traces`, evict the oldest trace whole and bump `dropped_lru`.
154    pub fn record(&self, span: RecordedSpan) {
155        let trace_id = span.trace_id.clone();
156        #[allow(clippy::unwrap_used, reason = "lock poison is unrecoverable")]
157        let mut inner = self.inner.lock().unwrap();
158
159        let is_new_trace = !inner.traces.contains_key(&trace_id);
160
161        if is_new_trace {
162            inner.traces.insert(trace_id.clone(), Vec::with_capacity(8));
163            inner.trace_order.push_back(trace_id.clone());
164
165            // LRU eviction once the new insertion pushed the count past the cap.
166            while inner.trace_order.len() > self.config.max_traces {
167                if let Some(oldest) = inner.trace_order.pop_front() {
168                    inner.traces.remove(&oldest);
169                    inner.stats.dropped_lru = inner.stats.dropped_lru.saturating_add(1);
170                } else {
171                    break;
172                }
173            }
174        }
175
176        let cap = self.config.max_spans_per_trace;
177        if let Some(bucket) = inner.traces.get_mut(&trace_id) {
178            if bucket.len() >= cap {
179                inner.stats.dropped_overflow = inner.stats.dropped_overflow.saturating_add(1);
180            } else {
181                bucket.push(span);
182            }
183        }
184    }
185
186    /// Snapshot every trace, most-recent first.
187    pub fn all_traces(&self) -> Vec<TraceSummary> {
188        #[allow(clippy::unwrap_used, reason = "lock poison is unrecoverable")]
189        let inner = self.inner.lock().unwrap();
190        let mut out = Vec::with_capacity(inner.trace_order.len());
191        // Iterate trace_order in reverse so the most recent trace lands first.
192        for trace_id in inner.trace_order.iter().rev() {
193            if let Some(spans) = inner.traces.get(trace_id) {
194                out.push(build_trace_summary(trace_id.clone(), spans.clone()));
195            }
196        }
197        out
198    }
199
200    /// Most recent trace, if any.
201    pub fn last_trace(&self) -> Option<TraceSummary> {
202        #[allow(clippy::unwrap_used, reason = "lock poison is unrecoverable")]
203        let inner = self.inner.lock().unwrap();
204        let trace_id = inner.trace_order.back()?.clone();
205        let spans = inner.traces.get(&trace_id)?.clone();
206        Some(build_trace_summary(trace_id, spans))
207    }
208
209    /// Drop every trace and reset stats.
210    pub fn clear(&self) {
211        #[allow(clippy::unwrap_used, reason = "lock poison is unrecoverable")]
212        let mut inner = self.inner.lock().unwrap();
213        inner.traces.clear();
214        inner.trace_order.clear();
215        inner.stats = BufferStats::default();
216    }
217
218    /// Cumulative drop counters.
219    pub fn stats(&self) -> BufferStats {
220        #[allow(clippy::unwrap_used, reason = "lock poison is unrecoverable")]
221        let inner = self.inner.lock().unwrap();
222        inner.stats.clone()
223    }
224
225    /// Configured cap on traces.
226    pub fn config(&self) -> &BufferConfig {
227        &self.config
228    }
229}
230
231/// Pick a root span from a flat list, mirroring Python's selection rules.
232///
233/// Rule: the first span whose `parent_span_id` is `None`; if none, the first
234/// span in the slice.
235fn pick_root(spans: &[RecordedSpan]) -> Option<&RecordedSpan> {
236    spans
237        .iter()
238        .find(|s| s.parent_span_id.is_none())
239        .or_else(|| spans.first())
240}
241
242/// `duration_ms = max(s.duration_ms for s in spans)`. Matches Python
243/// L86: `max((... for s in spans), default=0)`.
244fn max_duration_ms(spans: &[RecordedSpan]) -> f64 {
245    spans.iter().map(|s| s.duration_ms).fold(0.0_f64, f64::max)
246}
247
248fn build_trace_summary(trace_id: String, spans: Vec<RecordedSpan>) -> TraceSummary {
249    let span_count = spans.len();
250    let duration_ms = max_duration_ms(&spans);
251    let (root_name, status) = match pick_root(&spans) {
252        Some(root) => (Some(root.name.clone()), Some(root.status)),
253        None => (None, None),
254    };
255    TraceSummary {
256        trace_id,
257        root_name,
258        duration_ms,
259        span_count,
260        status,
261        spans,
262    }
263}
264
265#[cfg(test)]
266#[allow(
267    clippy::unwrap_used,
268    clippy::expect_used,
269    reason = "test code — panics are acceptable failures"
270)]
271mod tests {
272    use super::*;
273
274    fn span(trace: &str, span_id: &str, parent: Option<&str>, name: &str) -> RecordedSpan {
275        RecordedSpan {
276            trace_id: trace.into(),
277            span_id: span_id.into(),
278            parent_span_id: parent.map(|s| s.into()),
279            name: name.into(),
280            start_time_ns: 0,
281            end_time_ns: 1_000_000,
282            duration_ms: 1.0,
283            status: SpanStatus::Ok,
284            attributes: serde_json::Map::new(),
285        }
286    }
287
288    #[test]
289    fn record_then_snapshot_roundtrips() {
290        let buf = SpanBuffer::default();
291        buf.record(span("aa", "01", None, "root"));
292        buf.record(span("aa", "02", Some("01"), "child"));
293
294        let traces = buf.all_traces();
295        assert_eq!(traces.len(), 1);
296        assert_eq!(traces[0].trace_id, "aa");
297        assert_eq!(traces[0].span_count, 2);
298        assert_eq!(traces[0].root_name.as_deref(), Some("root"));
299        assert!((traces[0].duration_ms - 1.0).abs() < 1e-9);
300    }
301
302    #[test]
303    fn lru_evicts_oldest_trace_when_cap_exceeded() {
304        let buf = SpanBuffer::new(BufferConfig {
305            max_traces: 50,
306            max_spans_per_trace: 8,
307        });
308        for i in 0..51_u32 {
309            let trace = format!("{i:032x}");
310            buf.record(span(&trace, "0000000000000001", None, "root"));
311        }
312        let traces = buf.all_traces();
313        assert_eq!(traces.len(), 50);
314        // Oldest trace was id 0 — must be gone.
315        let zero = format!("{:032x}", 0);
316        assert!(traces.iter().all(|t| t.trace_id != zero));
317        // Most recent trace must be at the front.
318        let newest = format!("{:032x}", 50);
319        assert_eq!(traces[0].trace_id, newest);
320        assert_eq!(buf.stats().dropped_lru, 1);
321    }
322
323    #[test]
324    fn per_trace_cap_drops_silently() {
325        let buf = SpanBuffer::new(BufferConfig {
326            max_traces: 4,
327            max_spans_per_trace: 2,
328        });
329        buf.record(span("aa", "01", None, "root"));
330        buf.record(span("aa", "02", Some("01"), "c1"));
331        buf.record(span("aa", "03", Some("01"), "c2"));
332        buf.record(span("aa", "04", Some("01"), "c3"));
333
334        let traces = buf.all_traces();
335        assert_eq!(traces.len(), 1);
336        assert_eq!(traces[0].span_count, 2);
337        assert_eq!(buf.stats().dropped_overflow, 2);
338    }
339
340    #[test]
341    fn trace_summary_picks_parentless_root() {
342        let buf = SpanBuffer::default();
343        buf.record(span("aa", "02", Some("01"), "child"));
344        buf.record(span("aa", "01", None, "real_root"));
345
346        let traces = buf.all_traces();
347        assert_eq!(traces[0].root_name.as_deref(), Some("real_root"));
348    }
349
350    #[test]
351    fn trace_summary_status_unset_when_root_unset() {
352        let buf = SpanBuffer::default();
353        let mut s = span("aa", "01", None, "root");
354        s.status = SpanStatus::Unset;
355        buf.record(s);
356
357        let traces = buf.all_traces();
358        assert_eq!(traces[0].status, Some(SpanStatus::Unset));
359    }
360
361    #[test]
362    fn span_status_serializes_uppercase() {
363        let json = serde_json::to_string(&SpanStatus::Ok).expect("serialize");
364        assert_eq!(json, "\"OK\"");
365        let json = serde_json::to_string(&SpanStatus::Error).expect("serialize");
366        assert_eq!(json, "\"ERROR\"");
367    }
368
369    #[test]
370    fn last_trace_returns_most_recent() {
371        let buf = SpanBuffer::default();
372        buf.record(span("aa", "01", None, "first"));
373        buf.record(span("bb", "01", None, "second"));
374        let last = buf.last_trace().expect("last trace");
375        assert_eq!(last.trace_id, "bb");
376        assert_eq!(last.root_name.as_deref(), Some("second"));
377    }
378
379    #[test]
380    fn clear_resets_state() {
381        let buf = SpanBuffer::default();
382        buf.record(span("aa", "01", None, "root"));
383        buf.clear();
384        assert!(buf.all_traces().is_empty());
385        assert_eq!(buf.stats().dropped_lru, 0);
386    }
387}
cognee_http_server/observability/span_buffer.rs

cognee_http_server/observability/
span_buffer.rs
