claude_agent/observability/
spans.rs1use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5use tracing::{Level, Span, field, span};
6
7#[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
39pub 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
87pub 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}