Skip to main content

aether_telemetry/
otel_observer.rs

1use crate::content_capture::{ContentBuffer, ContentCapture};
2use crate::content_json::{messages_json, tool_definitions_json};
3use crate::gen_ai_metrics::GenAiMetrics;
4use crate::genai_constants as semconv;
5use crate::llm_call_state::LlmCallState;
6use crate::span_guard::{ErrorKind, SpanGuard};
7use aether_core::events::{AgentEvent, AgentObserver, LlmCallPurpose, MessageEvent, ToolEvent, TurnEvent, TurnOutcome};
8use llm::catalog::Provider;
9use llm::{ContentBlock, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition};
10use opentelemetry::trace::{SpanBuilder, SpanKind, TraceContextExt, Tracer as _};
11use opentelemetry::{Context, KeyValue};
12use opentelemetry_sdk::trace::SdkTracer;
13use std::collections::HashMap;
14
15/// [`AgentObserver`] that renders an agent's event stream as OpenTelemetry
16/// `GenAI` spans and metrics.
17pub struct OtelObserver {
18    turn: Option<TurnState>,
19    tool_definitions: Vec<ToolDefinition>,
20    instrumentation: OtelInstrumentation,
21}
22
23#[derive(Clone)]
24pub struct OtelInstrumentation {
25    pub tracer: SdkTracer,
26    pub metrics: GenAiMetrics,
27    pub capture_content: bool,
28}
29
30impl OtelObserver {
31    pub fn new(instrumentation: OtelInstrumentation) -> Self {
32        Self { turn: None, tool_definitions: Vec::new(), instrumentation }
33    }
34}
35
36impl AgentObserver for OtelObserver {
37    fn on_event(&mut self, message: &AgentEvent) {
38        match message {
39            AgentEvent::Turn(TurnEvent::Started { content }) => self.start_turn(content),
40            AgentEvent::Turn(TurnEvent::Ended { outcome }) => {
41                if let Some(turn) = self.turn.take() {
42                    turn.finish(outcome);
43                }
44            }
45            AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools }) => {
46                self.tool_definitions.clone_from(tools);
47            }
48            message => {
49                if let Some(turn) = &mut self.turn {
50                    turn.on_event(message, &self.instrumentation, &self.tool_definitions);
51                }
52            }
53        }
54    }
55}
56
57impl OtelObserver {
58    fn start_turn(&mut self, content: &[ContentBlock]) {
59        // Drop (and thereby cancel) any stale turn and its open spans before
60        // starting the new span, so the stale turn can't become its parent.
61        self.turn = None;
62
63        let mut input = self.instrumentation.capture().buffer();
64        input.set(&ContentBlock::join_text(content));
65        let mut attributes = vec![KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "invoke_agent")];
66        if let Some(text) = input.get() {
67            attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, messages_json("user", text)));
68        }
69        let builder = SpanBuilder::from_name("invoke_agent").with_kind(SpanKind::Internal).with_attributes(attributes);
70        let span_context = self.instrumentation.start_span(builder, None);
71        let span = SpanGuard::new(span_context, TURN_CANCEL_MESSAGE);
72        self.turn = Some(TurnState::new(span, input, self.instrumentation.capture()));
73    }
74}
75
76impl OtelInstrumentation {
77    fn capture(&self) -> ContentCapture {
78        ContentCapture::from_enabled(self.capture_content)
79    }
80
81    /// Starts a span under `parent`. Providers for disabled tracing use an
82    /// always-off sampler, so callers never need to branch.
83    fn start_span(&self, builder: SpanBuilder, parent: Option<&Context>) -> Context {
84        let parent = parent.cloned().unwrap_or_default();
85        Context::new().with_span(self.tracer.build_with_context(builder, &parent))
86    }
87}
88
89const TURN_CANCEL_MESSAGE: &str = "turn cancelled";
90const TOOL_CANCEL_MESSAGE: &str = "turn ended before the tool completed";
91
92/// All state scoped to one turn: the turn span plus any in-flight LLM-call and
93/// tool spans and captured content. A turn exists by construction inside its
94/// methods, and dropping it ends every still-open span as cancelled, so
95/// replacing the value is all a reset takes.
96struct TurnState {
97    span: SpanGuard,
98    input: ContentBuffer,
99    output: ContentBuffer,
100    chat_call: Option<LlmCallState>,
101    compaction_call: Option<LlmCallState>,
102    /// Tool-call arguments streamed before execution starts, keyed by call id.
103    streamed_arguments: HashMap<String, ContentBuffer>,
104    /// Spans of currently executing tools, keyed by call id.
105    executing_tools: HashMap<String, SpanGuard>,
106}
107
108impl TurnState {
109    fn new(span: SpanGuard, input: ContentBuffer, capture: ContentCapture) -> Self {
110        Self {
111            span,
112            input,
113            output: capture.buffer(),
114            chat_call: None,
115            compaction_call: None,
116            streamed_arguments: HashMap::new(),
117            executing_tools: HashMap::new(),
118        }
119    }
120
121    fn on_event(&mut self, message: &AgentEvent, instrumentation: &OtelInstrumentation, tools: &[ToolDefinition]) {
122        match message {
123            AgentEvent::Turn(TurnEvent::LlmCallStarted { purpose, provider, model, display_name, attempt, .. }) => {
124                self.start_llm_call(
125                    LlmCallStart {
126                        purpose: *purpose,
127                        provider: provider.as_deref(),
128                        model: model.as_deref(),
129                        display_name,
130                        attempt: *attempt,
131                    },
132                    instrumentation,
133                    tools,
134                );
135            }
136            AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
137                if let Some(call) = self.llm_call_slot(*purpose).take() {
138                    call.finish(outcome);
139                }
140            }
141            AgentEvent::Message(
142                MessageEvent::Text { message_id, chunk, is_complete: false }
143                | MessageEvent::Thought { message_id, chunk, is_complete: false },
144            ) => {
145                if let Some(chat) = &mut self.chat_call {
146                    chat.record_response_chunk(message_id, chunk);
147                }
148            }
149            AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
150                self.output.push(chunk);
151            }
152            AgentEvent::Tool(ToolEvent::Call { request, .. }) => self.on_tool_call(request, instrumentation),
153            AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id, chunk, .. }) => {
154                self.on_tool_call_update(tool_call_id, chunk);
155            }
156            AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name }) => {
157                self.on_tool_execution_started(tool_id, tool_name, instrumentation);
158            }
159            AgentEvent::Tool(ToolEvent::Result { result, .. }) => self.on_tool_result(result, instrumentation),
160            AgentEvent::Tool(ToolEvent::Error { error, .. }) => self.on_tool_error(error),
161            _ => {}
162        }
163    }
164
165    fn finish(self, outcome: &TurnOutcome) {
166        let Self { mut span, output, chat_call, compaction_call, executing_tools, .. } = self;
167        // Cancel any still-open call and tool spans before ending their parent.
168        drop(chat_call);
169        drop(compaction_call);
170        drop(executing_tools);
171        match outcome {
172            TurnOutcome::Completed => {
173                if let Some(text) = output.get() {
174                    span.set_attribute(KeyValue::new(
175                        semconv::GEN_AI_OUTPUT_MESSAGES,
176                        messages_json("assistant", text),
177                    ));
178                }
179                span.end_ok();
180            }
181            TurnOutcome::Failed { error } => span.end_error(None, error.clone()),
182            TurnOutcome::Cancelled => span.end_error(Some(ErrorKind::Cancelled), TURN_CANCEL_MESSAGE),
183        }
184    }
185
186    fn start_llm_call(
187        &mut self,
188        call: LlmCallStart<'_>,
189        instrumentation: &OtelInstrumentation,
190        tools: &[ToolDefinition],
191    ) {
192        let model_name = call.model.unwrap_or(call.display_name).to_string();
193
194        // Metrics backends key series by attribute values, so the metric set
195        // stays a small fixed vocabulary — no content, no per-call details.
196        let mut metric_attributes = vec![
197            KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "chat"),
198            KeyValue::new(semconv::GEN_AI_REQUEST_MODEL, model_name.clone()),
199        ];
200
201        if let Some(provider) = call.provider {
202            metric_attributes.push(KeyValue::new(semconv::GEN_AI_PROVIDER_NAME, genai_provider_name(provider)));
203        }
204
205        if call.purpose == LlmCallPurpose::Compaction {
206            metric_attributes.push(KeyValue::new(semconv::LLM_PURPOSE, "compaction"));
207        }
208
209        let mut attributes = metric_attributes.clone();
210        attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true));
211        attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt)));
212        // Only chat calls carry the turn's input and tool definitions; a
213        // compaction call's actual input is the internal summarization prompt.
214        if call.purpose == LlmCallPurpose::Chat {
215            if let Some(input) = self.input.get() {
216                attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, messages_json("user", input)));
217            }
218            if instrumentation.capture_content && !tools.is_empty() {
219                attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_DEFINITIONS, tool_definitions_json(tools)));
220            }
221        }
222
223        let name = if model_name.is_empty() { "chat".to_string() } else { format!("chat {model_name}") };
224        let builder = SpanBuilder::from_name(name).with_kind(SpanKind::Client).with_attributes(attributes);
225        let context = instrumentation.start_span(builder, Some(self.span.context()));
226        let state = LlmCallState::new(
227            context,
228            instrumentation.metrics.clone(),
229            call.purpose,
230            instrumentation.capture().buffer(),
231            metric_attributes,
232        );
233        *self.llm_call_slot(call.purpose) = Some(state);
234    }
235
236    fn on_tool_call(&mut self, request: &ToolCallRequest, instrumentation: &OtelInstrumentation) {
237        if let Some(chat) = &mut self.chat_call {
238            chat.record_tool_call_start(&request.id, &request.name);
239        }
240        let mut arguments = instrumentation.capture().buffer();
241        arguments.set(&request.arguments);
242        self.streamed_arguments.insert(request.id.clone(), arguments);
243    }
244
245    fn on_tool_call_update(&mut self, tool_call_id: &str, chunk: &str) {
246        if let Some(chat) = &mut self.chat_call {
247            chat.record_output_chunk();
248        }
249        if let Some(arguments) = self.streamed_arguments.get_mut(tool_call_id) {
250            arguments.push(chunk);
251        }
252    }
253
254    fn on_tool_execution_started(&mut self, tool_id: &str, tool_name: &str, instrumentation: &OtelInstrumentation) {
255        let mut attributes = vec![
256            KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "execute_tool"),
257            KeyValue::new(semconv::GEN_AI_TOOL_NAME, tool_name.to_string()),
258            KeyValue::new(semconv::GEN_AI_TOOL_CALL_ID, tool_id.to_string()),
259        ];
260        let arguments = self.streamed_arguments.remove(tool_id);
261
262        if let Some(text) = arguments.as_ref().and_then(ContentBuffer::get) {
263            attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_CALL_ARGUMENTS, text.to_string()));
264        }
265
266        let builder = SpanBuilder::from_name(format!("execute_tool {tool_name}"))
267            .with_kind(SpanKind::Internal)
268            .with_attributes(attributes);
269        let context = instrumentation.start_span(builder, Some(self.span.context()));
270        self.executing_tools.insert(tool_id.to_string(), SpanGuard::new(context, TOOL_CANCEL_MESSAGE));
271    }
272
273    fn on_tool_result(&mut self, result: &ToolCallResult, instrumentation: &OtelInstrumentation) {
274        self.streamed_arguments.remove(&result.id);
275        let Some(mut span) = self.executing_tools.remove(&result.id) else { return };
276
277        if let Some(content) = instrumentation.capture().content(&result.result) {
278            span.set_attribute(KeyValue::new(semconv::GEN_AI_TOOL_CALL_RESULT, content.to_string()));
279        }
280
281        span.end_ok();
282    }
283
284    fn on_tool_error(&mut self, error: &ToolCallError) {
285        self.streamed_arguments.remove(&error.id);
286        if let Some(mut span) = self.executing_tools.remove(&error.id) {
287            span.end_error(Some(ErrorKind::ToolError), error.error.clone());
288        }
289    }
290
291    fn llm_call_slot(&mut self, purpose: LlmCallPurpose) -> &mut Option<LlmCallState> {
292        match purpose {
293            LlmCallPurpose::Chat => &mut self.chat_call,
294            LlmCallPurpose::Compaction => &mut self.compaction_call,
295        }
296    }
297}
298
299/// Maps the agent's canonical provider name to the `GenAI` semantic-convention
300/// provider name, passing providers the catalog doesn't know through as-is.
301fn genai_provider_name(provider: &str) -> String {
302    provider.parse::<Provider>().map_or_else(|_| provider.to_string(), |p| p.genai_provider_name().to_string())
303}
304
305/// Borrowed view of [`TurnEvent::LlmCallStarted`]; named fields keep the two
306/// optional strings from being swapped at the call site.
307#[derive(Clone, Copy)]
308struct LlmCallStart<'a> {
309    purpose: LlmCallPurpose,
310    provider: Option<&'a str>,
311    model: Option<&'a str>,
312    display_name: &'a str,
313    attempt: u32,
314}