starweaver_runtime/trace/
memory.rs

1use std::sync::{Arc, Mutex};
2
3use starweaver_core::{Metadata, TraceContext};
4use uuid::Uuid;
5
6use super::{export_otel_gen_ai_spans, OtelGenAiSpan};
7use super::{RecordedSpan, SpanEvent, SpanHandle, SpanSpec, SpanStatus, TraceRecorder};
8
9/// Deterministic in-memory trace recorder for tests and CLI inspection.
10#[derive(Clone, Debug, Default)]
11pub struct InMemoryTraceRecorder {
12    spans: Arc<Mutex<Vec<RecordedSpan>>>,
13}
14
15impl InMemoryTraceRecorder {
16    /// Create an empty recorder.
17    #[must_use]
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    /// Return recorded spans.
23    #[must_use]
24    pub fn spans(&self) -> Vec<RecordedSpan> {
25        self.spans
26            .lock()
27            .map_or_else(|_| Vec::new(), |spans| spans.clone())
28    }
29
30    fn update_span(&self, span_id: &str, f: impl FnOnce(&mut RecordedSpan)) {
31        if let Ok(mut spans) = self.spans.lock() {
32            if let Some(span) = spans.iter_mut().find(|span| span.span_id == span_id) {
33                f(span);
34            }
35        }
36    }
37}
38
39impl TraceRecorder for InMemoryTraceRecorder {
40    fn start_span(&self, spec: SpanSpec, parent: &TraceContext) -> SpanHandle {
41        let trace_id = parent
42            .trace_id
43            .clone()
44            .unwrap_or_else(|| format!("trace_{}", Uuid::new_v4()));
45        let parent_span_id = parent
46            .span_id
47            .clone()
48            .or_else(|| parent.parent_span_id.clone());
49        let span_id = format!("span_{}", Uuid::new_v4());
50        let mut metadata = Metadata::default();
51        metadata.insert("span_name".to_string(), serde_json::json!(spec.name));
52        let context = TraceContext {
53            trace_id: Some(trace_id.clone()),
54            span_id: Some(span_id.clone()),
55            parent_span_id: parent_span_id.clone(),
56            trace_state: parent.trace_state.clone(),
57            metadata,
58        };
59        if let Ok(mut spans) = self.spans.lock() {
60            spans.push(RecordedSpan {
61                span_id: span_id.clone(),
62                trace_id,
63                parent_span_id,
64                name: spec.name,
65                kind: spec.kind,
66                level: spec.level,
67                attributes: spec.attributes,
68                events: Vec::new(),
69                status: SpanStatus::Open,
70            });
71        }
72        SpanHandle::new(context, span_id)
73    }
74
75    fn record_event(&self, span: &SpanHandle, event: SpanEvent) {
76        self.update_span(span.span_id(), |record| record.events.push(event));
77    }
78
79    fn close_span(&self, span: &SpanHandle, status: SpanStatus) {
80        self.update_span(span.span_id(), |record| record.status = status);
81    }
82}
83
84/// Adapter seam for feature-gated tracing/OpenTelemetry exporters.
85#[derive(Clone, Debug, Default)]
86pub struct AdapterTraceRecorder {
87    inner: InMemoryTraceRecorder,
88}
89
90impl AdapterTraceRecorder {
91    /// Create an adapter seam backed by deterministic in-memory storage.
92    #[must_use]
93    pub fn new() -> Self {
94        Self::default()
95    }
96
97    /// Return the spans captured before export.
98    #[must_use]
99    pub fn spans(&self) -> Vec<RecordedSpan> {
100        self.inner.spans()
101    }
102
103    /// Return spans projected into deterministic OpenTelemetry `GenAI` fields.
104    #[must_use]
105    pub fn otel_gen_ai_spans(&self) -> Vec<OtelGenAiSpan> {
106        export_otel_gen_ai_spans(&self.spans())
107    }
108}
109
110impl TraceRecorder for AdapterTraceRecorder {
111    fn start_span(&self, spec: SpanSpec, parent: &TraceContext) -> SpanHandle {
112        self.inner.start_span(spec, parent)
113    }
114
115    fn record_event(&self, span: &SpanHandle, event: SpanEvent) {
116        self.inner.record_event(span, event);
117    }
118
119    fn close_span(&self, span: &SpanHandle, status: SpanStatus) {
120        self.inner.close_span(span, status);
121    }
122}
starweaver_runtime/trace/memory.rs

starweaver_runtime/trace/
memory.rs
