ai_agents_observability/
span.rs

1use crate::context::SpanContext;
2use crate::event::{EventStatus, EventType, ObservationError, ObservationTokenUsage};
3use crate::manager::ObservabilityManager;
4use serde_json::Value;
5use std::collections::HashMap;
6use std::sync::Arc;
7use std::time::Instant;
8
9/// Stopwatch for one observed span that records on explicit finish or Drop.
10pub struct SpanGuard {
11    manager: Arc<ObservabilityManager>,
12    context: SpanContext,
13    event_type: EventType,
14    start_time: Instant,
15    tokens: Option<ObservationTokenUsage>,
16    status: EventStatus,
17    error: Option<ObservationError>,
18    tags: HashMap<String, String>,
19    payload: Option<Value>,
20    recorded: bool,
21}
22
23impl SpanGuard {
24    /// Creates a span guard with success status and an immediate start time.
25    pub fn new(
26        manager: Arc<ObservabilityManager>,
27        context: SpanContext,
28        event_type: EventType,
29    ) -> Self {
30        Self {
31            manager,
32            context,
33            event_type,
34            start_time: Instant::now(),
35            tokens: None,
36            status: EventStatus::Success,
37            error: None,
38            tags: HashMap::new(),
39            payload: None,
40            recorded: false,
41        }
42    }
43
44    /// Attaches token usage collected from provider output or estimation.
45    pub fn set_tokens(&mut self, tokens: ObservationTokenUsage) {
46        self.tokens = Some(tokens);
47    }
48
49    /// Overrides the final status before recording.
50    pub fn set_status(&mut self, status: EventStatus) {
51        self.status = status;
52    }
53
54    /// Marks the span as failed and attaches error metadata.
55    pub fn set_error(&mut self, error: ObservationError) {
56        self.status = EventStatus::Error;
57        self.error = Some(error);
58    }
59
60    /// Attaches optional payload data that will be redacted by the manager.
61    pub fn set_payload(&mut self, payload: Value) {
62        self.payload = Some(payload);
63    }
64
65    /// Adds a safe tag to the event before recording.
66    pub fn add_tag(&mut self, key: impl Into<String>, value: impl Into<String>) {
67        self.tags.insert(key.into(), value.into());
68    }
69
70    /// Records the span immediately and prevents duplicate Drop recording.
71    pub fn record_now(&mut self) {
72        if self.recorded {
73            return;
74        }
75        self.recorded = true;
76        let event = self.manager.build_event_from_span(
77            self.context.clone(),
78            self.event_type.clone(),
79            self.start_time.elapsed(),
80            self.status.clone(),
81            self.tokens.clone(),
82            self.error.clone(),
83            self.tags.clone(),
84            self.payload.clone(),
85        );
86        let deferred = self
87            .context
88            .tags
89            .get("runtime.defer_observation")
90            .map(|value| value == "true")
91            .unwrap_or(false);
92        if deferred {
93            if let Some(branch_id) = self.context.tags.get("runtime.branch_id") {
94                self.manager.record_pending_event(branch_id.clone(), event);
95                return;
96            }
97        }
98        self.manager.record_event(event);
99    }
100}
101
102impl Drop for SpanGuard {
103    fn drop(&mut self) {
104        self.record_now();
105    }
106}
ai_agents_observability/span.rs

ai_agents_observability/
span.rs
