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 aether_core::events::{AgentEvent, AgentObserver, LlmCallPurpose, MessageEvent, ToolEvent, TurnEvent, TurnOutcome};
8use llm::catalog::Provider;
9use llm::{ContentBlock, ModelPricing, 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, input_messages_json(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 {
125                purpose,
126                provider,
127                model,
128                display_name,
129                pricing,
130                attempt,
131                ..
132            }) => {
133                self.start_llm_call(
134                    LlmCallStart {
135                        purpose: *purpose,
136                        provider: provider.as_deref(),
137                        model: model.as_deref(),
138                        display_name,
139                        pricing: *pricing,
140                        attempt: *attempt,
141                    },
142                    instrumentation,
143                    tools,
144                );
145            }
146            AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
147                if let Some(call) = self.llm_call_slot(*purpose).take() {
148                    call.finish(outcome);
149                }
150            }
151            AgentEvent::Message(
152                MessageEvent::Text { message_id, chunk, is_complete: false }
153                | MessageEvent::Thought { message_id, chunk, is_complete: false },
154            ) => {
155                if let Some(chat) = &mut self.chat_call {
156                    chat.record_response_chunk(message_id, chunk);
157                }
158            }
159            AgentEvent::Message(MessageEvent::Text { chunk, is_complete: true, .. }) => {
160                self.output.push(chunk);
161            }
162            AgentEvent::Tool(ToolEvent::Call { request, .. }) => self.on_tool_call(request, instrumentation),
163            AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id, chunk, .. }) => {
164                self.on_tool_call_update(tool_call_id, chunk);
165            }
166            AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name }) => {
167                self.on_tool_execution_started(tool_id, tool_name, instrumentation);
168            }
169            AgentEvent::Tool(ToolEvent::Result { result, .. }) => self.on_tool_result(result, instrumentation),
170            AgentEvent::Tool(ToolEvent::Error { error, .. }) => self.on_tool_error(error),
171            _ => {}
172        }
173    }
174
175    fn finish(self, outcome: &TurnOutcome) {
176        let Self { mut span, output, chat_call, compaction_call, executing_tools, .. } = self;
177        // Cancel any still-open call and tool spans before ending their parent.
178        drop(chat_call);
179        drop(compaction_call);
180        drop(executing_tools);
181        match outcome {
182            TurnOutcome::Completed => {
183                if let Some(text) = output_messages_json(output.get(), &[], None) {
184                    span.set_attribute(KeyValue::new(semconv::GEN_AI_OUTPUT_MESSAGES, text));
185                }
186                span.end_ok();
187            }
188            TurnOutcome::Failed { error } => span.end_error(None, error.clone()),
189            TurnOutcome::Cancelled => span.end_error(Some(ErrorKind::Cancelled), TURN_CANCEL_MESSAGE),
190        }
191    }
192
193    fn start_llm_call(
194        &mut self,
195        call: LlmCallStart<'_>,
196        instrumentation: &OtelInstrumentation,
197        tools: &[ToolDefinition],
198    ) {
199        let model_name = call.model.unwrap_or(call.display_name).to_string();
200
201        // Metrics backends key series by attribute values, so the metric set
202        // stays a small fixed vocabulary — no content, no per-call details.
203        let mut metric_attributes = vec![
204            KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "chat"),
205            KeyValue::new(semconv::GEN_AI_REQUEST_MODEL, model_name.clone()),
206        ];
207
208        if let Some(provider) = call.provider {
209            metric_attributes.push(KeyValue::new(semconv::GEN_AI_PROVIDER_NAME, genai_provider_name(provider)));
210        }
211
212        if call.purpose == LlmCallPurpose::Compaction {
213            metric_attributes.push(KeyValue::new(semconv::LLM_PURPOSE, "compaction"));
214        }
215
216        let mut attributes = metric_attributes.clone();
217        attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true));
218        attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt)));
219        if let Some(pricing) = call.pricing {
220            attributes.extend(pricing_attributes(pricing));
221        }
222        // Only chat calls carry the turn's input and tool definitions; a
223        // compaction call's actual input is the internal summarization prompt.
224        if call.purpose == LlmCallPurpose::Chat {
225            if let Some(input) = self.input.get() {
226                attributes.push(KeyValue::new(semconv::GEN_AI_INPUT_MESSAGES, input_messages_json(input)));
227            }
228            if instrumentation.capture_content && !tools.is_empty() {
229                attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_DEFINITIONS, tool_definitions_json(tools)));
230            }
231        }
232
233        let name = if model_name.is_empty() { "chat".to_string() } else { format!("chat {model_name}") };
234        let builder = SpanBuilder::from_name(name).with_kind(SpanKind::Client).with_attributes(attributes);
235        let context = instrumentation.start_span(builder, Some(self.span.context()));
236        let state = LlmCallState::new(
237            context,
238            instrumentation.metrics.clone(),
239            call.purpose,
240            instrumentation.capture(),
241            metric_attributes,
242        );
243        *self.llm_call_slot(call.purpose) = Some(state);
244    }
245
246    fn on_tool_call(&mut self, request: &ToolCallRequest, instrumentation: &OtelInstrumentation) {
247        if let Some(chat) = &mut self.chat_call {
248            chat.record_tool_call_start(request);
249        }
250        let mut arguments = instrumentation.capture().buffer();
251        arguments.set(&request.arguments);
252        self.streamed_arguments.insert(request.id.clone(), arguments);
253    }
254
255    fn on_tool_call_update(&mut self, tool_call_id: &str, chunk: &str) {
256        if let Some(chat) = &mut self.chat_call {
257            chat.record_tool_call_update(tool_call_id, chunk);
258        }
259        if let Some(arguments) = self.streamed_arguments.get_mut(tool_call_id) {
260            arguments.push(chunk);
261        }
262    }
263
264    fn on_tool_execution_started(&mut self, tool_id: &str, tool_name: &str, instrumentation: &OtelInstrumentation) {
265        let mut attributes = vec![
266            KeyValue::new(semconv::GEN_AI_OPERATION_NAME, "execute_tool"),
267            KeyValue::new(semconv::GEN_AI_TOOL_NAME, tool_name.to_string()),
268            KeyValue::new(semconv::GEN_AI_TOOL_CALL_ID, tool_id.to_string()),
269        ];
270        let arguments = self.streamed_arguments.remove(tool_id);
271
272        if let Some(text) = arguments.as_ref().and_then(ContentBuffer::get) {
273            attributes.push(KeyValue::new(semconv::GEN_AI_TOOL_CALL_ARGUMENTS, text.to_string()));
274        }
275
276        let builder = SpanBuilder::from_name(format!("execute_tool {tool_name}"))
277            .with_kind(SpanKind::Internal)
278            .with_attributes(attributes);
279        let context = instrumentation.start_span(builder, Some(self.span.context()));
280        self.executing_tools.insert(tool_id.to_string(), SpanGuard::new(context, TOOL_CANCEL_MESSAGE));
281    }
282
283    fn on_tool_result(&mut self, result: &ToolCallResult, instrumentation: &OtelInstrumentation) {
284        self.streamed_arguments.remove(&result.id);
285        let Some(mut span) = self.executing_tools.remove(&result.id) else { return };
286
287        if let Some(content) = instrumentation.capture().content(&result.result) {
288            span.set_attribute(KeyValue::new(semconv::GEN_AI_TOOL_CALL_RESULT, content.to_string()));
289        }
290
291        span.end_ok();
292    }
293
294    fn on_tool_error(&mut self, error: &ToolCallError) {
295        self.streamed_arguments.remove(&error.id);
296        if let Some(mut span) = self.executing_tools.remove(&error.id) {
297            span.end_error(Some(ErrorKind::ToolError), error.error.clone());
298        }
299    }
300
301    fn llm_call_slot(&mut self, purpose: LlmCallPurpose) -> &mut Option<LlmCallState> {
302        match purpose {
303            LlmCallPurpose::Chat => &mut self.chat_call,
304            LlmCallPurpose::Compaction => &mut self.compaction_call,
305        }
306    }
307}
308
309/// Maps the agent's canonical provider name to the `GenAI` semantic-convention
310/// provider name, passing providers the catalog doesn't know through as-is.
311fn genai_provider_name(provider: &str) -> String {
312    provider.parse::<Provider>().map_or_else(|_| provider.to_string(), |p| p.genai_provider_name().to_string())
313}
314
315/// Borrowed view of [`TurnEvent::LlmCallStarted`]; named fields keep the two
316/// optional strings from being swapped at the call site.
317#[derive(Clone, Copy)]
318struct LlmCallStart<'a> {
319    purpose: LlmCallPurpose,
320    provider: Option<&'a str>,
321    model: Option<&'a str>,
322    display_name: &'a str,
323    pricing: Option<ModelPricing>,
324    attempt: u32,
325}
326
327fn pricing_attributes(pricing: ModelPricing) -> Vec<KeyValue> {
328    let mut attributes = vec![
329        KeyValue::new(semconv::AI_INPUT_TOKEN_PRICE, pricing.input_per_million / TOKENS_PER_MILLION),
330        KeyValue::new(semconv::AI_OUTPUT_TOKEN_PRICE, pricing.output_per_million / TOKENS_PER_MILLION),
331    ];
332
333    if let Some(price) = pricing.cache_read_per_million {
334        attributes.push(KeyValue::new(semconv::AI_CACHE_READ_TOKEN_PRICE, price / TOKENS_PER_MILLION));
335    }
336
337    if let Some(price) = pricing.cache_write_per_million {
338        attributes.push(KeyValue::new(semconv::AI_CACHE_WRITE_TOKEN_PRICE, price / TOKENS_PER_MILLION));
339    }
340
341    attributes
342}
343
344const TOKENS_PER_MILLION: f64 = 1_000_000.0;