Skip to main content

aether_telemetry/
otel_observer.rs

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