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