oxicode_sdk/observability/
trace.rs

1//! Distributed tracing — spans, trace IDs, and RAII guards.
2
3use parking_lot::RwLock;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::sync::Arc;
7use tokio::sync::broadcast;
8
9// ── TraceId ──────────────────────────────────────────────────────────────────
10
11/// Unique trace identifier.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub struct TraceId(u64);
14
15impl TraceId {
16    /// Generate a new random trace ID.
17    pub fn new() -> Self {
18        Self(fastrand::u64(1..))
19    }
20    /// Zero trace ID (invalid).
21    pub fn zero() -> Self {
22        Self(0)
23    }
24}
25
26impl Default for TraceId {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl std::fmt::Display for TraceId {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        write!(f, "{:016x}", self.0)
35    }
36}
37
38// ── SpanId ───────────────────────────────────────────────────────────────────
39
40/// Unique span identifier.
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
42pub struct SpanId(u64);
43
44impl SpanId {
45    /// Generate a new random span ID.
46    pub fn new() -> Self {
47        Self(fastrand::u64(1..))
48    }
49    /// Zero span ID (invalid).
50    pub fn zero() -> Self {
51        Self(0)
52    }
53}
54
55impl Default for SpanId {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl std::fmt::Display for SpanId {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        write!(f, "{:016x}", self.0)
64    }
65}
66
67// ── SpanKind ────────────────────────────────────────────────────────────────
68
69/// Role or classification of a span within a trace.
70#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
71pub enum SpanKind {
72    /// Span covering an agent's execution.
73    Agent,
74    /// Span covering a tool invocation.
75    #[default]
76    Tool,
77    /// Span covering an LLM request/response.
78    Llm,
79    /// Span covering internal SDK work.
80    Internal,
81}
82
83// ── SpanStatus ───────────────────────────────────────────────────────────────
84
85/// Outcome status of a span.
86#[derive(Debug, Clone, Default, Serialize, Deserialize)]
87pub enum SpanStatus {
88    /// The span completed successfully.
89    #[default]
90    Ok,
91    /// The span completed with an error.
92    Error {
93        /// Message describing the error.
94        message: String,
95    },
96}
97
98// ── SpanContext ──────────────────────────────────────────────────────────────
99
100/// Identifies a span within a trace and links it to its parent.
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct SpanContext {
103    /// Trace this span belongs to.
104    pub trace_id: TraceId,
105    /// Unique id of this span.
106    pub span_id: SpanId,
107    /// Id of the parent span, if any.
108    pub parent_span_id: Option<SpanId>,
109}
110
111// ── SpanEvent ─────────────────────────────────────────────────────────────────
112
113/// A timestamped event recorded on a span.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115pub struct SpanEvent {
116    /// Human-readable name of the event.
117    pub name: String,
118    /// Wall-clock time the event occurred, in milliseconds.
119    pub timestamp_ms: u64,
120    /// Arbitrary key/value attributes for the event.
121    #[serde(default)]
122    pub attributes: Vec<(String, serde_json::Value)>,
123}
124
125// ── Span ─────────────────────────────────────────────────────────────────────
126
127/// A single unit of work within a trace, with timing, status, and metadata.
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct Span {
130    /// Identifier and parent link for this span.
131    pub context: SpanContext,
132    /// Human-readable name of the operation.
133    pub name: String,
134    /// Classification of the span.
135    pub kind: SpanKind,
136    /// Wall-clock start time in milliseconds.
137    pub start_ms: u64,
138    /// Wall-clock end time in milliseconds, once the span is closed.
139    pub end_ms: Option<u64>,
140    /// Outcome of the span.
141    pub status: SpanStatus,
142    /// Arbitrary key/value metadata attached to the span.
143    #[serde(default)]
144    pub attributes: HashMap<String, serde_json::Value>,
145    /// Timestamped events recorded during the span.
146    #[serde(default)]
147    pub events: Vec<SpanEvent>,
148    /// Links to related span contexts (e.g. cross-trace references).
149    #[serde(default)]
150    pub links: Vec<SpanContext>,
151}
152
153impl Span {
154    /// Duration in milliseconds.
155    pub fn duration_ms(&self) -> Option<u64> {
156        self.end_ms.map(|end| end.saturating_sub(self.start_ms))
157    }
158    /// True if the span has an end time.
159    pub fn is_complete(&self) -> bool {
160        self.end_ms.is_some()
161    }
162}
163
164// ── Tracer ────────────────────────────────────────────────────────────────────
165
166/// Collects completed spans and broadcasts them to subscribers.
167#[derive(Debug)]
168pub struct Tracer {
169    spans: Arc<RwLock<Vec<Span>>>,
170    completed_tx: broadcast::Sender<Span>,
171}
172
173impl Tracer {
174    /// Create a new tracer.
175    pub fn new() -> Self {
176        let (tx, _) = broadcast::channel(256);
177        Self {
178            spans: Arc::new(RwLock::new(Vec::new())),
179            completed_tx: tx,
180        }
181    }
182
183    /// Start a root span.
184    pub fn start(self: &Arc<Self>, name: &str, kind: SpanKind) -> SpanGuard {
185        self.start_with_parent(name, kind, None)
186    }
187
188    /// Start a child span with optional parent context.
189    pub fn start_with_parent(
190        self: &Arc<Self>,
191        name: &str,
192        kind: SpanKind,
193        parent: Option<&SpanContext>,
194    ) -> SpanGuard {
195        let trace_id = parent.map(|c| c.trace_id).unwrap_or_default();
196        let span_id = SpanId::new();
197        let context = SpanContext {
198            trace_id,
199            span_id,
200            parent_span_id: parent.map(|c| c.span_id),
201        };
202        let span = Span {
203            context,
204            name: name.to_string(),
205            kind,
206            start_ms: now_ms(),
207            end_ms: None,
208            status: SpanStatus::Ok,
209            attributes: HashMap::new(),
210            events: Vec::new(),
211            links: Vec::new(),
212        };
213        SpanGuard {
214            tracer: Arc::clone(self),
215            span,
216        }
217    }
218
219    fn record(&self, span: Span) {
220        self.spans.write().push(span.clone());
221        let _ = self.completed_tx.send(span);
222    }
223
224    /// Retrieve all spans for a given trace ID.
225    pub fn trace(&self, trace_id: TraceId) -> Vec<Span> {
226        self.spans
227            .read()
228            .iter()
229            .filter(|s| s.context.trace_id == trace_id)
230            .cloned()
231            .collect()
232    }
233
234    /// Subscribe to completed span events.
235    pub fn subscribe(&self) -> broadcast::Receiver<Span> {
236        self.completed_tx.subscribe()
237    }
238}
239
240impl Clone for Tracer {
241    fn clone(&self) -> Self {
242        Self {
243            spans: Arc::clone(&self.spans),
244            completed_tx: self.completed_tx.clone(),
245        }
246    }
247}
248
249impl Default for Tracer {
250    fn default() -> Self {
251        Self::new()
252    }
253}
254
255// ── SpanGuard ────────────────────────────────────────────────────────────────
256
257/// RAII guard for an in-flight span; finalizes and records the span on drop.
258pub struct SpanGuard {
259    tracer: Arc<Tracer>,
260    span: Span,
261}
262
263impl SpanGuard {
264    /// Borrowed span context (trace, span, and parent ids).
265    pub fn context(&self) -> &SpanContext {
266        &self.span.context
267    }
268
269    /// Trace id this span belongs to.
270    pub fn trace_id(&self) -> TraceId {
271        self.span.context.trace_id
272    }
273
274    /// Unique id of this span.
275    pub fn span_id(&self) -> SpanId {
276        self.span.context.span_id
277    }
278
279    /// Attach a key/value attribute to the span.
280    pub fn set_attribute(&mut self, key: &str, value: serde_json::Value) {
281        self.span.attributes.insert(key.to_string(), value);
282    }
283
284    /// Record a named, timestamped event on the span.
285    pub fn add_event(&mut self, name: &str) {
286        self.span.events.push(SpanEvent {
287            name: name.to_string(),
288            timestamp_ms: now_ms(),
289            attributes: vec![],
290        });
291    }
292
293    /// Mark the span as failed with the given error message.
294    pub fn set_error(&mut self, message: &str) {
295        self.span.status = SpanStatus::Error {
296            message: message.to_string(),
297        };
298    }
299}
300
301impl Drop for SpanGuard {
302    fn drop(&mut self) {
303        let mut span = self.span.clone();
304        span.end_ms = Some(now_ms());
305        self.tracer.record(span);
306    }
307}
308
309// ── Helpers ───────────────────────────────────────────────────────────────────
310
311fn now_ms() -> u64 {
312    std::time::SystemTime::now()
313        .duration_since(std::time::UNIX_EPOCH)
314        .map(|d| d.as_millis() as u64)
315        .unwrap_or(0)
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    fn assert_send_static<T: Send + 'static>() {}
323
324    #[test]
325    fn span_guard_is_send_and_static() {
326        assert_send_static::<SpanGuard>();
327    }
328
329    #[tokio::test]
330    async fn smoke() {
331        let tracer = Arc::new(Tracer::new());
332        let guard = tracer.start("s", SpanKind::Agent);
333        let tid = guard.trace_id();
334        drop(guard);
335        let spans = tracer.trace(tid);
336        assert!(!spans.is_empty());
337        assert_eq!(spans[0].name, "s");
338        assert!(spans[0].is_complete());
339    }
340
341    #[tokio::test]
342    async fn child_span() {
343        let tracer = Arc::new(Tracer::new());
344        let parent = tracer.start("parent", SpanKind::Agent);
345        let parent_ctx = parent.context().clone();
346        drop(parent);
347        let child = tracer.start_with_parent("child", SpanKind::Tool, Some(&parent_ctx));
348        let tid = child.trace_id();
349        drop(child);
350        let spans = tracer.trace(tid);
351        assert_eq!(spans.len(), 2);
352    }
353}
oxicode_sdk/observability/trace.rs

oxicode_sdk/observability/
trace.rs
