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
19pub struct OtelObserver {
22 turn: Option<TurnState>,
23 tool_definitions: Vec<ToolDefinition>,
24 otel: OtelInstrumentation,
25}
26
27#[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 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 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
104struct TurnState {
109 span: SpanGuard,
110 input: ContentBuffer,
111 output: ContentBuffer,
112 chat_call: Option<LlmCallState>,
113 compaction_call: Option<LlmCallState>,
114 streamed_arguments: HashMap<String, ContentBuffer>,
116 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 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 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 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 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
324fn 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#[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;