claude_agent/observability/
spans.rs

1//! Structured span definitions for tracing.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5use tracing::{Level, Span, field, span};
6
7/// Tracing configuration.
8#[derive(Clone, Default)]
9pub struct TracingConfig {
10    pub service_name: Option<String>,
11    pub enabled: bool,
12    pub level: TracingLevel,
13}
14
15#[derive(Clone, Copy, Default, PartialEq, Eq)]
16pub enum TracingLevel {
17    #[default]
18    Info,
19    Debug,
20    Trace,
21}
22
23impl TracingConfig {
24    pub fn new() -> Self {
25        Self {
26            enabled: true,
27            ..Default::default()
28        }
29    }
30
31    pub fn disabled() -> Self {
32        Self {
33            enabled: false,
34            ..Default::default()
35        }
36    }
37}
38
39/// Context for creating structured spans.
40pub struct SpanContext {
41    session_id: String,
42    request_id: AtomicU64,
43}
44
45impl SpanContext {
46    pub fn new(session_id: impl Into<String>) -> Self {
47        Self {
48            session_id: session_id.into(),
49            request_id: AtomicU64::new(0),
50        }
51    }
52
53    pub fn next_request_id(&self) -> u64 {
54        self.request_id.fetch_add(1, Ordering::Relaxed)
55    }
56
57    pub fn agent_execute_span(&self, model: &str) -> Span {
58        let request_id = self.next_request_id();
59        span!(
60            Level::INFO,
61            "agent.execute",
62            session_id = %self.session_id,
63            request_id = request_id,
64            model = model,
65            otel.name = "agent.execute",
66        )
67    }
68
69    pub fn api_call_span(&self, model: &str) -> ApiCallSpan {
70        ApiCallSpan::new(model)
71    }
72
73    pub fn tool_execute_span(&self, tool_name: &str, tool_use_id: &str) -> Span {
74        span!(
75            Level::INFO,
76            "tool.execute",
77            tool_name = tool_name,
78            tool_use_id = tool_use_id,
79            session_id = %self.session_id,
80            otel.name = format!("tool.{}", tool_name),
81            is_error = field::Empty,
82            duration_ms = field::Empty,
83        )
84    }
85}
86
87/// Helper for tracking API call metrics within a span.
88pub struct ApiCallSpan {
89    span: Span,
90    start: Instant,
91}
92
93impl ApiCallSpan {
94    pub fn new(model: &str) -> Self {
95        let span = span!(
96            Level::INFO,
97            "api.call",
98            model = model,
99            otel.name = "api.call",
100            input_tokens = field::Empty,
101            output_tokens = field::Empty,
102            latency_ms = field::Empty,
103            cache_read_tokens = field::Empty,
104            cache_creation_tokens = field::Empty,
105        );
106        Self {
107            span,
108            start: Instant::now(),
109        }
110    }
111
112    pub fn record_usage(&self, input_tokens: u32, output_tokens: u32) {
113        self.span.record("input_tokens", input_tokens);
114        self.span.record("output_tokens", output_tokens);
115    }
116
117    pub fn record_cache(&self, read_tokens: u32, creation_tokens: u32) {
118        self.span.record("cache_read_tokens", read_tokens);
119        self.span.record("cache_creation_tokens", creation_tokens);
120    }
121
122    pub fn finish(self) {
123        let latency_ms = self.start.elapsed().as_millis() as u64;
124        self.span.record("latency_ms", latency_ms);
125    }
126
127    pub fn span(&self) -> &Span {
128        &self.span
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn test_span_context() {
138        let span_context = SpanContext::new("test-session");
139        assert_eq!(span_context.next_request_id(), 0);
140        assert_eq!(span_context.next_request_id(), 1);
141    }
142
143    #[test]
144    fn test_api_call_span() {
145        let span = ApiCallSpan::new("claude-sonnet-4-5");
146        span.record_usage(100, 50);
147        span.record_cache(20, 10);
148        span.finish();
149    }
150}