Skip to main content

agent_framework_core/
observability.rs

1//! Lightweight OpenTelemetry GenAI-style span instrumentation, built on the
2//! [`tracing`] crate, plus optional GenAI metrics behind the `otel-metrics`
3//! feature.
4//!
5//! This is a dependency-light port of the Python `observability.py`
6//! instrumentation. It emits `tracing` spans that follow the OpenTelemetry
7//! GenAI semantic conventions, so that an OTel bridge (e.g.
8//! `tracing-opentelemetry`) can export them without any additional glue:
9//!
10//! * span names: `chat {model}`, `invoke_agent {agent}`, `execute_tool {tool}`
11//!   — the human-readable name is carried in the `otel.name` field (the static
12//!   `tracing` metadata name is the bare operation, since `tracing` requires a
13//!   literal span name).
14//! * chat-span attributes: `gen_ai.operation.name`, `gen_ai.system` /
15//!   `gen_ai.provider.name` (dual-emitted — see [`attr::PROVIDER_NAME`]),
16//!   `gen_ai.request.model`, `gen_ai.response.model`, `gen_ai.response.id`,
17//!   `gen_ai.response.finish_reasons`, `gen_ai.usage.{input,output}_tokens`,
18//!   the request parameters (`gen_ai.request.{temperature,top_p,max_tokens,
19//!   seed,frequency_penalty,presence_penalty,stop_sequences}`,
20//!   `gen_ai.conversation.id`), `error.type` plus the `tracing-opentelemetry`
21//!   "special fields" `otel.status_code` / `otel.status_message` — and, only
22//!   when content capture is explicitly enabled, `gen_ai.input.messages` /
23//!   `gen_ai.output.messages`, `gen_ai.system_instructions`, and
24//!   `gen_ai.tool.definitions`.
25//! * tool-span attributes: `gen_ai.tool.name`, `gen_ai.tool.call.id`,
26//!   `gen_ai.tool.description`, `gen_ai.tool.type`, and (content-capture-gated)
27//!   `gen_ai.tool.call.arguments` / `gen_ai.tool.call.result`.
28//!
29//! The main entry point is [`ObservableChatClient`], a [`ChatClient`] decorator.
30//! Tool execution inside [`FunctionInvokingChatClient`] and the
31//! [`Agent`](crate::agent::Agent) run paths are instrumented directly by
32//! those types using the span constructors here.
33//!
34//! ## Metrics (`otel-metrics` feature)
35//!
36//! With the `otel-metrics` feature enabled, [`ObservableChatClient`] also
37//! records two histograms through the `opentelemetry` **API** crate only
38//! (mirroring `observability.py:788-803`, bucket boundaries at `:65-96`):
39//! [`metrics::TOKEN_USAGE_METRIC`] (`gen_ai.client.token.usage`, unit
40//! `"tokens"`) and [`metrics::OPERATION_DURATION_METRIC`]
41//! (`gen_ai.client.operation.duration`, unit `"s"`). A third histogram,
42//! [`metrics::FUNCTION_INVOCATION_DURATION_METRIC`]
43//! (`agent_framework.function.invocation.duration`), is defined for tool-call
44//! timing; see [`metrics::record_function_invocation_duration`] for why it
45//! isn't wired to a call site yet. This crate never depends on an OTel SDK:
46//! without an application-installed `MeterProvider` (via
47//! [`opentelemetry::global::set_meter_provider`]) the instruments are no-ops,
48//! so the feature is safe to enable unconditionally.
49//!
50//! ## Wiring to a real OTel backend
51//!
52//! Neither the `tracing` spans nor the `otel-metrics` histograms are exported
53//! anywhere by this crate — that stays the application's job. A minimal
54//! bridge, using `tracing-opentelemetry` for spans and `opentelemetry_sdk` for
55//! metrics:
56//!
57//! ```ignore
58//! use opentelemetry_sdk::trace::SdkTracerProvider;
59//! use opentelemetry_sdk::metrics::SdkMeterProvider;
60//! use tracing_subscriber::layer::SubscriberExt;
61//!
62//! let tracer_provider = SdkTracerProvider::builder()
63//!     // .with_batch_exporter(otlp_span_exporter) / .with_simple_exporter(...) …
64//!     .build();
65//! let meter_provider = SdkMeterProvider::builder()
66//!     // .with_reader(periodic_reader_wrapping_your_metric_exporter) …
67//!     .build();
68//! opentelemetry::global::set_meter_provider(meter_provider); // powers `otel-metrics`
69//!
70//! let tracer = tracer_provider.tracer("agent_framework");
71//! let subscriber = tracing_subscriber::registry()
72//!     .with(tracing_opentelemetry::layer().with_tracer(tracer));
73//! tracing::subscriber::set_global_default(subscriber).unwrap();
74//! ```
75//!
76//! Without any of this, spans are still emitted to whatever plain `tracing`
77//! subscriber you do have (e.g. for structured logging), and `otel-metrics`
78//! instruments silently drop their measurements — zero required setup either
79//! way.
80//!
81//! ## Environment configuration
82//!
83//! [`ObservabilityConfig::from_env`] reads `ENABLE_SENSITIVE_DATA` (mirrors
84//! Python's `enable_sensitive_data`, `observability.py:347-394`) for use with
85//! [`ObservableChatClient::from_env`]. Unlike Python, there is no
86//! `ENABLE_OTEL` equivalent to read: this crate's spans are plain `tracing`
87//! spans, already effectively free without a subscriber attached, so there is
88//! no separate "enable observability" switch.
89//!
90//! [`FunctionInvokingChatClient`]: crate::client::FunctionInvokingChatClient
91
92use async_trait::async_trait;
93use futures::StreamExt;
94use tracing::field::Empty;
95use tracing::{Instrument, Span};
96
97use crate::client::{ChatClient, ChatStream};
98use crate::error::{Error, Result};
99use crate::tools::ToolDefinition;
100use crate::types::{ChatOptions, ChatResponse, Message};
101
102/// OpenTelemetry GenAI semantic-convention attribute keys.
103pub mod attr {
104    pub const OPERATION: &str = "gen_ai.operation.name";
105    /// The provider/system tag, e.g. `"openai"`. Supplied by the client.
106    pub const SYSTEM: &str = "gen_ai.system";
107    /// Current semantic-convention replacement for [`SYSTEM`]; both are set
108    /// from the same value so older and newer consumers each find what they
109    /// expect on the span.
110    pub const PROVIDER_NAME: &str = "gen_ai.provider.name";
111    pub const REQUEST_MODEL: &str = "gen_ai.request.model";
112    pub const RESPONSE_MODEL: &str = "gen_ai.response.model";
113    pub const RESPONSE_ID: &str = "gen_ai.response.id";
114    pub const FINISH_REASONS: &str = "gen_ai.response.finish_reasons";
115    pub const INPUT_TOKENS: &str = "gen_ai.usage.input_tokens";
116    pub const OUTPUT_TOKENS: &str = "gen_ai.usage.output_tokens";
117    /// Input tokens written to a provider-managed cache.
118    pub const CACHE_CREATION_INPUT_TOKENS: &str = "gen_ai.usage.cache_creation.input_tokens";
119    /// Input tokens served from a provider-managed cache.
120    pub const CACHE_READ_INPUT_TOKENS: &str = "gen_ai.usage.cache_read.input_tokens";
121    /// Output tokens spent on reasoning.
122    pub const REASONING_OUTPUT_TOKENS: &str = "gen_ai.usage.reasoning.output_tokens";
123    /// Low-cardinality prompt name (e.g. for a named/templated prompt).
124    pub const PROMPT_NAME: &str = "gen_ai.prompt.name";
125    pub const REQUEST_TEMPERATURE: &str = "gen_ai.request.temperature";
126    pub const REQUEST_TOP_P: &str = "gen_ai.request.top_p";
127    pub const REQUEST_MAX_TOKENS: &str = "gen_ai.request.max_tokens";
128    pub const REQUEST_SEED: &str = "gen_ai.request.seed";
129    pub const REQUEST_FREQUENCY_PENALTY: &str = "gen_ai.request.frequency_penalty";
130    pub const REQUEST_PRESENCE_PENALTY: &str = "gen_ai.request.presence_penalty";
131    pub const REQUEST_STOP_SEQUENCES: &str = "gen_ai.request.stop_sequences";
132    pub const CONVERSATION_ID: &str = "gen_ai.conversation.id";
133    /// Content-capture-gated: the system/instructions prompt, JSON-encoded as
134    /// `[{"type":"text","content":...}]` (mirrors
135    /// `observability.py:1444-1448`).
136    pub const SYSTEM_INSTRUCTIONS: &str = "gen_ai.system_instructions";
137    /// Content-capture-gated: the request's tool list, JSON-encoded (mirrors
138    /// `_tools_to_dict`, `observability.py:1388-1391`).
139    pub const TOOL_DEFINITIONS: &str = "gen_ai.tool.definitions";
140    pub const ERROR_TYPE: &str = "error.type";
141    pub const TOOL_NAME: &str = "gen_ai.tool.name";
142    pub const TOOL_CALL_ID: &str = "gen_ai.tool.call.id";
143    pub const TOOL_DESCRIPTION: &str = "gen_ai.tool.description";
144    pub const TOOL_TYPE: &str = "gen_ai.tool.type";
145    /// Content-capture-gated tool-call arguments (JSON-encoded).
146    pub const TOOL_CALL_ARGUMENTS: &str = "gen_ai.tool.call.arguments";
147    /// Content-capture-gated tool-call result (JSON-encoded).
148    pub const TOOL_CALL_RESULT: &str = "gen_ai.tool.call.result";
149    pub const AGENT_NAME: &str = "gen_ai.agent.name";
150    pub const AGENT_ID: &str = "gen_ai.agent.id";
151    pub const INPUT_MESSAGES: &str = "gen_ai.input.messages";
152    pub const OUTPUT_MESSAGES: &str = "gen_ai.output.messages";
153    /// The human-readable span name override consumed by OTel bridges.
154    pub const OTEL_NAME: &str = "otel.name";
155    /// `tracing-opentelemetry` special field: sets the OTel span status code
156    /// (`"ERROR"` or `"OK"`) when a bridge is attached.
157    pub const OTEL_STATUS_CODE: &str = "otel.status_code";
158    /// `tracing-opentelemetry` special field: sets the OTel span status
159    /// description. Used here to carry the exception message, mirroring
160    /// Python's `capture_exception` (`observability.py:1407-1411`, which
161    /// calls `set_status(description=repr(exception))`).
162    pub const OTEL_STATUS_MESSAGE: &str = "otel.status_message";
163}
164
165/// OpenTelemetry GenAI operation names.
166pub mod op {
167    pub const CHAT: &str = "chat";
168    pub const INVOKE_AGENT: &str = "invoke_agent";
169    pub const EXECUTE_TOOL: &str = "execute_tool";
170    pub const EMBEDDINGS: &str = "embeddings";
171}
172
173/// The `error.type` value for a framework [`Error`]: its variant discriminant.
174pub fn error_type(err: &Error) -> String {
175    // The Display of the thiserror variants is prefixed with a stable label,
176    // e.g. "tool error: …"; take the label as a compact type tag.
177    match err {
178        Error::AgentInitialization(_) => "agent_initialization",
179        Error::AgentExecution(_) => "agent_execution",
180        Error::Serialization(_) => "serialization",
181        Error::Content(_) => "content",
182        Error::Tool(_) => "tool",
183        Error::Service(_) => "service",
184        Error::ServiceStatus { .. } => "service",
185        Error::ServiceInvalidAuth { .. } => "service_invalid_auth",
186        Error::ServiceInvalidRequest { .. } => "service_invalid_request",
187        Error::ServiceContentFilter { .. } => "service_content_filter",
188        Error::Workflow(_) => "workflow",
189        Error::AdditionItemMismatch(_) => "addition_item_mismatch",
190        Error::Configuration(_) => "configuration",
191        Error::Json(_) => "json",
192        Error::Other(_) => "other",
193    }
194    .to_string()
195}
196
197/// Build a `chat {model}` span for a chat-completion request.
198///
199/// Only the fields that are cheap/always-known at call time are set here
200/// (`gen_ai.system` / `gen_ai.provider.name`, `gen_ai.request.model`); the
201/// rest of the request/response/error attribute set is filled in afterward
202/// via [`record_request`], [`record_response`], and [`record_error`] — mirrors
203/// `_get_span_attributes` (`observability.py:1345-1404`).
204pub fn chat_span(system: &str, model: &str) -> Span {
205    let span = tracing::info_span!(
206        "chat",
207        otel.name = Empty,
208        gen_ai.operation.name = op::CHAT,
209        gen_ai.system = system,
210        gen_ai.provider.name = system,
211        gen_ai.request.model = model,
212        gen_ai.response.model = Empty,
213        gen_ai.response.id = Empty,
214        gen_ai.response.finish_reasons = Empty,
215        gen_ai.usage.input_tokens = Empty,
216        gen_ai.usage.output_tokens = Empty,
217        gen_ai.usage.cache_creation.input_tokens = Empty,
218        gen_ai.usage.cache_read.input_tokens = Empty,
219        gen_ai.usage.reasoning.output_tokens = Empty,
220        gen_ai.request.temperature = Empty,
221        gen_ai.request.top_p = Empty,
222        gen_ai.request.max_tokens = Empty,
223        gen_ai.request.seed = Empty,
224        gen_ai.request.frequency_penalty = Empty,
225        gen_ai.request.presence_penalty = Empty,
226        gen_ai.request.stop_sequences = Empty,
227        gen_ai.conversation.id = Empty,
228        gen_ai.system_instructions = Empty,
229        gen_ai.tool.definitions = Empty,
230        error.type = Empty,
231        otel.status_code = Empty,
232        otel.status_message = Empty,
233        gen_ai.input.messages = Empty,
234        gen_ai.output.messages = Empty,
235    );
236    span.record(attr::OTEL_NAME, format!("{} {}", op::CHAT, model).as_str());
237    span
238}
239
240/// Build an `invoke_agent {agent}` span for an agent run.
241pub fn agent_span(agent_name: &str, agent_id: &str) -> Span {
242    let span = tracing::info_span!(
243        "invoke_agent",
244        otel.name = Empty,
245        gen_ai.operation.name = op::INVOKE_AGENT,
246        gen_ai.agent.name = agent_name,
247        gen_ai.agent.id = agent_id,
248        gen_ai.usage.input_tokens = Empty,
249        gen_ai.usage.output_tokens = Empty,
250        error.type = Empty,
251        // Declared for forward-compatibility with `record_error`; the
252        // current `agent.rs` error path only records `error.type` directly,
253        // so these stay unset until that call site migrates.
254        otel.status_code = Empty,
255        otel.status_message = Empty,
256    );
257    span.record(
258        attr::OTEL_NAME,
259        format!("{} {}", op::INVOKE_AGENT, agent_name).as_str(),
260    );
261    span
262}
263
264/// Build an `execute_tool {tool}` span with the full OTel GenAI tool
265/// attribute set: name, call id, description, a fixed `"function"` tool type
266/// (the only kind executed through this in-process loop — mirrors Python's
267/// `get_function_span_attributes`, `observability.py:1284-1302`), and
268/// placeholders for the content-capture-gated call arguments/result (fill
269/// with [`record_tool_arguments`] / [`record_tool_result`]) and the
270/// error/status fields (fill with [`record_error`]).
271///
272/// [`tool_span`] is the source-compatible two-argument form used today by
273/// `client.rs`'s `FunctionInvokingChatClient`; it delegates here with
274/// `description = None`. New call sites — and that eventual migration —
275/// should call this directly to get a real tool description and light up the
276/// content-capture-gated attributes.
277pub fn tool_span_ex(tool_name: &str, call_id: &str, description: Option<&str>) -> Span {
278    let span = tracing::info_span!(
279        "execute_tool",
280        otel.name = Empty,
281        gen_ai.operation.name = op::EXECUTE_TOOL,
282        gen_ai.tool.name = tool_name,
283        gen_ai.tool.call.id = call_id,
284        gen_ai.tool.description = Empty,
285        gen_ai.tool.type = "function",
286        gen_ai.tool.call.arguments = Empty,
287        gen_ai.tool.call.result = Empty,
288        error.type = Empty,
289        otel.status_code = Empty,
290        otel.status_message = Empty,
291    );
292    if let Some(description) = description {
293        if !description.is_empty() {
294            span.record(attr::TOOL_DESCRIPTION, description);
295        }
296    }
297    span.record(
298        attr::OTEL_NAME,
299        format!("{} {}", op::EXECUTE_TOOL, tool_name).as_str(),
300    );
301    span
302}
303
304/// Build an `execute_tool {tool}` span (source-compatible two-argument form).
305///
306/// Delegates to [`tool_span_ex`] with no description. Prefer `tool_span_ex`
307/// for new call sites.
308pub fn tool_span(tool_name: &str, call_id: &str) -> Span {
309    tool_span_ex(tool_name, call_id, None)
310}
311
312/// Record tool-call arguments onto a tool span, gated by content capture
313/// (mirrors the `SENSITIVE_DATA_ENABLED`-gated `gen_ai.tool.call.arguments`
314/// capture in Python's `AIFunction.invoke`, `_tools.py:751-759`).
315pub fn record_tool_arguments(span: &Span, arguments: &serde_json::Value, capture_content: bool) {
316    if !capture_content {
317        return;
318    }
319    span.record(attr::TOOL_CALL_ARGUMENTS, arguments.to_string().as_str());
320}
321
322/// Record a tool-call result onto a tool span, gated by content capture
323/// (mirrors the `gen_ai.tool.call.result` capture in Python's
324/// `AIFunction.invoke`, `_tools.py:779-787`).
325pub fn record_tool_result(span: &Span, result: &serde_json::Value, capture_content: bool) {
326    if !capture_content {
327        return;
328    }
329    span.record(attr::TOOL_CALL_RESULT, result.to_string().as_str());
330}
331
332/// Record an error onto a span following the existing `error.type` pattern:
333/// `error.type` (the framework [`Error`] variant tag), plus the
334/// `tracing-opentelemetry` "special fields" `otel.status_code` (`"ERROR"`)
335/// and `otel.status_message` (the exception's `Display` text) so that a
336/// bridge sets real OTel span status. This mirrors Python's
337/// `capture_exception` (`record_exception` + `set_status`,
338/// `observability.py:1407-1411`) as far as bare `tracing` fields allow —
339/// there is no span-events API here without also taking on an SDK
340/// dependency.
341pub fn record_error(span: &Span, err: &Error) {
342    span.record(attr::ERROR_TYPE, error_type(err).as_str());
343    span.record(attr::OTEL_STATUS_CODE, "ERROR");
344    span.record(attr::OTEL_STATUS_MESSAGE, err.to_string().as_str());
345}
346
347/// Record the response-side attributes (finish reason, usage, id, model) onto
348/// `span`, mirroring `_get_response_attributes` (`observability.py:1488-1512`).
349pub fn record_response(span: &Span, response: &ChatResponse, capture_content: bool) {
350    if let Some(reason) = &response.finish_reason {
351        span.record(attr::FINISH_REASONS, reason.as_str());
352    }
353    if let Some(id) = &response.response_id {
354        span.record(attr::RESPONSE_ID, id.as_str());
355    }
356    if let Some(model) = &response.model {
357        span.record(attr::RESPONSE_MODEL, model.as_str());
358    }
359    if let Some(usage) = &response.usage_details {
360        if let Some(input) = usage.input_token_count {
361            span.record(attr::INPUT_TOKENS, input);
362        }
363        if let Some(output) = usage.output_token_count {
364            span.record(attr::OUTPUT_TOKENS, output);
365        }
366        if let Some(v) = usage.cache_creation_input_token_count {
367            span.record(attr::CACHE_CREATION_INPUT_TOKENS, v);
368        }
369        if let Some(v) = usage.cache_read_input_token_count {
370            span.record(attr::CACHE_READ_INPUT_TOKENS, v);
371        }
372        if let Some(v) = usage.reasoning_output_token_count {
373            span.record(attr::REASONING_OUTPUT_TOKENS, v);
374        }
375    }
376    if capture_content {
377        span.record(
378            attr::OUTPUT_MESSAGES,
379            messages_json(&response.messages).as_str(),
380        );
381    }
382}
383
384/// Record request-side attributes onto a `chat` span from [`ChatOptions`],
385/// mirroring `_get_span_attributes` (`observability.py:1345-1404`). System
386/// instructions and the serialized tool list are additionally gated by
387/// `capture_content` (mirrors Python's `SENSITIVE_DATA_ENABLED` gate).
388pub fn record_request(span: &Span, options: &ChatOptions, capture_content: bool) {
389    if let Some(v) = options.temperature {
390        span.record(attr::REQUEST_TEMPERATURE, f64::from(v));
391    }
392    if let Some(v) = options.top_p {
393        span.record(attr::REQUEST_TOP_P, f64::from(v));
394    }
395    if let Some(v) = options.max_tokens {
396        span.record(attr::REQUEST_MAX_TOKENS, u64::from(v));
397    }
398    if let Some(v) = options.seed {
399        span.record(attr::REQUEST_SEED, v);
400    }
401    if let Some(v) = options.frequency_penalty {
402        span.record(attr::REQUEST_FREQUENCY_PENALTY, f64::from(v));
403    }
404    if let Some(v) = options.presence_penalty {
405        span.record(attr::REQUEST_PRESENCE_PENALTY, f64::from(v));
406    }
407    if let Some(stop) = &options.stop {
408        if !stop.is_empty() {
409            span.record(
410                attr::REQUEST_STOP_SEQUENCES,
411                serde_json::to_string(stop).unwrap_or_default().as_str(),
412            );
413        }
414    }
415    if let Some(id) = &options.conversation_id {
416        span.record(attr::CONVERSATION_ID, id.as_str());
417    }
418    if capture_content {
419        if let Some(instructions) = &options.instructions {
420            if !instructions.is_empty() {
421                span.record(
422                    attr::SYSTEM_INSTRUCTIONS,
423                    system_instructions_json(instructions).as_str(),
424                );
425            }
426        }
427        if !options.tools.is_empty() {
428            span.record(
429                attr::TOOL_DEFINITIONS,
430                tool_definitions_json(&options.tools).as_str(),
431            );
432        }
433    }
434}
435
436/// JSON-encode a system/instructions prompt as
437/// `[{"type":"text","content":...}]`, mirroring
438/// `observability.py:1444-1448`.
439fn system_instructions_json(instructions: &str) -> String {
440    serde_json::json!([{ "type": "text", "content": instructions }]).to_string()
441}
442
443/// JSON-encode a tool list for `gen_ai.tool.definitions`, mirroring the
444/// shape produced by `_tools_to_dict` (`_tools.py:827-857`) closely enough to
445/// be useful without depending on any provider-specific wire format.
446fn tool_definitions_json(tools: &[ToolDefinition]) -> String {
447    let list: Vec<serde_json::Value> = tools
448        .iter()
449        .map(|t| {
450            serde_json::json!({
451                "name": t.name,
452                "description": t.description,
453                "parameters": t.parameters,
454            })
455        })
456        .collect();
457    serde_json::to_string(&list).unwrap_or_default()
458}
459
460/// Serialize messages to a compact JSON string for content-capture attributes.
461fn messages_json(messages: &[Message]) -> String {
462    serde_json::to_string(messages).unwrap_or_default()
463}
464
465/// Per-stream bookkeeping threaded through [`ObservableChatClient`]'s
466/// `get_streaming_response` `unfold` state, so the finalization arm can
467/// record response attributes and, with `otel-metrics`, the completion
468/// histograms. `system` / `request_model` / `start` are only needed for the
469/// latter, so they (and the work to populate them) are compiled out
470/// entirely when the feature is off.
471struct StreamTelemetry {
472    capture: bool,
473    #[cfg(feature = "otel-metrics")]
474    system: String,
475    #[cfg(feature = "otel-metrics")]
476    request_model: String,
477    #[cfg(feature = "otel-metrics")]
478    start: std::time::Instant,
479}
480
481/// A [`ChatClient`] decorator that emits a `chat` span per request following the
482/// OpenTelemetry GenAI semantic conventions.
483///
484/// Message-content capture (`gen_ai.input.messages` / `gen_ai.output.messages`,
485/// plus `gen_ai.system_instructions` / `gen_ai.tool.definitions`) is **off by
486/// default**; enable it with [`ObservableChatClient::with_content_capture`],
487/// mirroring Python's `enable_sensitive_data` flag — or construct via
488/// [`ObservableChatClient::from_env`] to read that flag from
489/// `ENABLE_SENSITIVE_DATA`.
490pub struct ObservableChatClient<C: ChatClient> {
491    inner: C,
492    system: String,
493    capture_content: bool,
494}
495
496impl<C: ChatClient> ObservableChatClient<C> {
497    /// Wrap `inner`, tagging spans with the given provider/system name (the
498    /// `gen_ai.system` / `gen_ai.provider.name` attributes), e.g. `"openai"`
499    /// or `"anthropic"`.
500    pub fn new(inner: C, system: impl Into<String>) -> Self {
501        Self {
502            inner,
503            system: system.into(),
504            capture_content: false,
505        }
506    }
507
508    /// Wrap `inner` using [`ObservabilityConfig::from_env`] to decide whether
509    /// content capture is enabled (`ENABLE_SENSITIVE_DATA`). Equivalent to:
510    ///
511    /// ```ignore
512    /// ObservableChatClient::new(inner, system)
513    ///     .with_content_capture(ObservabilityConfig::from_env().enable_sensitive_data)
514    /// ```
515    pub fn from_env(inner: C, system: impl Into<String>) -> Self {
516        let config = ObservabilityConfig::from_env();
517        Self::new(inner, system).with_content_capture(config.enable_sensitive_data)
518    }
519
520    /// Enable or disable capturing message content on spans (default: off).
521    pub fn with_content_capture(mut self, capture: bool) -> Self {
522        self.capture_content = capture;
523        self
524    }
525
526    /// A reference to the wrapped client.
527    pub fn inner(&self) -> &C {
528        &self.inner
529    }
530
531    fn model_for(&self, options: &ChatOptions) -> String {
532        options
533            .model
534            .clone()
535            .or_else(|| self.inner.model().map(str::to_string))
536            .unwrap_or_default()
537    }
538}
539
540#[async_trait]
541impl<C: ChatClient> ChatClient for ObservableChatClient<C> {
542    async fn get_response(
543        &self,
544        messages: Vec<Message>,
545        options: ChatOptions,
546    ) -> Result<ChatResponse> {
547        let request_model = self.model_for(&options);
548        let span = chat_span(&self.system, &request_model);
549        record_request(&span, &options, self.capture_content);
550        if self.capture_content {
551            span.record(attr::INPUT_MESSAGES, messages_json(&messages).as_str());
552        }
553        let capture = self.capture_content;
554        #[cfg(feature = "otel-metrics")]
555        let start = std::time::Instant::now();
556        async move {
557            let result = self.inner.get_response(messages, options).await;
558            let span = Span::current();
559            match &result {
560                Ok(response) => {
561                    record_response(&span, response, capture);
562                    #[cfg(feature = "otel-metrics")]
563                    {
564                        let (input_tokens, output_tokens) = response
565                            .usage_details
566                            .as_ref()
567                            .map(|u| (u.input_token_count, u.output_token_count))
568                            .unwrap_or((None, None));
569                        metrics::record_chat_completion(
570                            &self.system,
571                            &request_model,
572                            response.model.as_deref(),
573                            input_tokens,
574                            output_tokens,
575                            start.elapsed(),
576                        );
577                    }
578                }
579                Err(err) => {
580                    record_error(&span, err);
581                }
582            }
583            result
584        }
585        .instrument(span)
586        .await
587    }
588
589    async fn get_streaming_response(
590        &self,
591        messages: Vec<Message>,
592        options: ChatOptions,
593    ) -> Result<ChatStream> {
594        let request_model = self.model_for(&options);
595        let span = chat_span(&self.system, &request_model);
596        record_request(&span, &options, self.capture_content);
597        if self.capture_content {
598            span.record(attr::INPUT_MESSAGES, messages_json(&messages).as_str());
599        }
600        let capture = self.capture_content;
601        // Instrument the initiation future with the span (never hold an
602        // `enter()` guard across an await); attribute recording happens as
603        // the stream drains and completes.
604        let inner = self
605            .inner
606            .get_streaming_response(messages, options)
607            .instrument(span.clone())
608            .await;
609
610        let inner = match inner {
611            Ok(s) => s,
612            Err(err) => {
613                record_error(&span, &err);
614                return Err(err);
615            }
616        };
617
618        // Aggregate updates to recover finish reason / usage from the final
619        // chunks, recording them onto the span (and, with `otel-metrics`, the
620        // completion histograms) when the stream ends.
621        let telemetry = StreamTelemetry {
622            capture,
623            #[cfg(feature = "otel-metrics")]
624            system: self.system.clone(),
625            #[cfg(feature = "otel-metrics")]
626            request_model: request_model.clone(),
627            #[cfg(feature = "otel-metrics")]
628            start: std::time::Instant::now(),
629        };
630        let state = (inner, ChatResponse::default(), Some(span), false, telemetry);
631        let stream = futures::stream::unfold(
632            state,
633            |(mut inner, mut agg, mut span, done, telemetry)| async move {
634                if done {
635                    return None;
636                }
637                match inner.next().await {
638                    Some(Ok(update)) => {
639                        agg.absorb_update(update.clone());
640                        Some((Ok(update), (inner, agg, span, false, telemetry)))
641                    }
642                    Some(Err(err)) => {
643                        if let Some(span) = &span {
644                            record_error(span, &err);
645                        }
646                        Some((Err(err), (inner, agg, span.take(), true, telemetry)))
647                    }
648                    None => {
649                        if let Some(span) = span.take() {
650                            agg.finalize();
651                            record_response(&span, &agg, telemetry.capture);
652                            #[cfg(feature = "otel-metrics")]
653                            {
654                                let (input_tokens, output_tokens) = agg
655                                    .usage_details
656                                    .as_ref()
657                                    .map(|u| (u.input_token_count, u.output_token_count))
658                                    .unwrap_or((None, None));
659                                metrics::record_chat_completion(
660                                    &telemetry.system,
661                                    &telemetry.request_model,
662                                    agg.model.as_deref(),
663                                    input_tokens,
664                                    output_tokens,
665                                    telemetry.start.elapsed(),
666                                );
667                            }
668                        }
669                        None
670                    }
671                }
672            },
673        );
674        Ok(stream.boxed())
675    }
676
677    fn model(&self) -> Option<&str> {
678        self.inner.model()
679    }
680}
681
682/// Observability configuration read from the process environment, mirroring
683/// (a subset of) Python's `ObservabilitySettings` (`observability.py:347-394`).
684///
685/// This port intentionally does **not** read `ENABLE_OTEL`: Python uses that
686/// flag to skip constructing spans entirely when no heavier OTel SDK is
687/// configured. This crate's spans are plain [`tracing`] spans, which are
688/// already effectively free when no subscriber is attached, so there is no
689/// separate "enable" switch — attach a subscriber (optionally bridged to
690/// OTel; see the [module docs](self)) to turn observability on.
691///
692/// This crate does **not** build an OTel `Resource`, exporter, or
693/// `MeterProvider` — see the [module docs](self) for wiring those yourself.
694#[derive(Debug, Clone, Default)]
695pub struct ObservabilityConfig {
696    /// Whether message / system-instructions / tool-definition / tool-call
697    /// content capture is enabled. Reads `ENABLE_SENSITIVE_DATA` (mirrors
698    /// Python's `enable_sensitive_data` — "Warning: Sensitive events should
699    /// only be enabled on test and development environments.").
700    pub enable_sensitive_data: bool,
701}
702
703impl ObservabilityConfig {
704    /// Read configuration from the process environment. Unset or
705    /// unrecognized values default to `false` (matching Python's default).
706    /// Recognized truthy values (case-insensitive, surrounding whitespace
707    /// ignored): `"1"`, `"true"`, `"yes"`, `"on"`.
708    pub fn from_env() -> Self {
709        Self {
710            enable_sensitive_data: env_flag("ENABLE_SENSITIVE_DATA"),
711        }
712    }
713
714    /// The `OTEL_SERVICE_NAME` passthrough, defaulting to `"agent_framework"`
715    /// (mirrors Python's `_create_resource`, `observability.py:336-344`).
716    /// This crate does not build a `Resource` from it — it's exposed so an
717    /// app wiring its own OTel `Resource` (see the [module docs](self)) can
718    /// stay consistent with whatever this reads. Reads the environment fresh
719    /// on every call rather than caching it at [`from_env`](Self::from_env)
720    /// time.
721    pub fn otel_service_name(&self) -> String {
722        std::env::var("OTEL_SERVICE_NAME").unwrap_or_else(|_| "agent_framework".to_string())
723    }
724}
725
726/// Parse a boolean-ish environment variable, matching Python's
727/// pydantic-settings-style truthy strings. Missing or unrecognized values are
728/// `false`.
729fn env_flag(key: &str) -> bool {
730    std::env::var(key).is_ok_and(|v| {
731        matches!(
732            v.trim().to_ascii_lowercase().as_str(),
733            "1" | "true" | "yes" | "on"
734        )
735    })
736}
737
738/// Metrics instruments behind the `otel-metrics` feature (see the crate's
739/// `[features]` table). Uses the `opentelemetry` **API** crate only
740/// (`"metrics"` feature, no default features): instrument creation and
741/// `.record()` calls are inert until an application installs a real
742/// `MeterProvider` via [`opentelemetry::global::set_meter_provider`] — the
743/// API's own default global provider is a no-op — so enabling this feature
744/// never *requires* an application to also wire up an OTel SDK. See the
745/// [module docs](super) for how to do that when you do want real metrics out
746/// the other end.
747///
748/// Mirrors upstream's two chat-client histograms (`observability.py:788-803`,
749/// bucket boundaries at `:65-96`) and the function-invocation-duration
750/// histogram (`_tools.py`'s `_default_histogram`).
751#[cfg(feature = "otel-metrics")]
752pub mod metrics {
753    use std::sync::OnceLock;
754    use std::time::Duration;
755
756    use opentelemetry::metrics::Histogram;
757    use opentelemetry::KeyValue;
758
759    use super::{attr, op};
760
761    /// Bucket boundaries for `gen_ai.client.token.usage`, matching upstream's
762    /// `TOKEN_USAGE_BUCKET_BOUNDARIES` (`observability.py:65-80`).
763    pub const TOKEN_USAGE_BUCKET_BOUNDARIES: &[f64] = &[
764        1.0,
765        4.0,
766        16.0,
767        64.0,
768        256.0,
769        1024.0,
770        4096.0,
771        16384.0,
772        65536.0,
773        262_144.0,
774        1_048_576.0,
775        4_194_304.0,
776        16_777_216.0,
777        67_108_864.0,
778    ];
779
780    /// Bucket boundaries for `gen_ai.client.operation.duration` and
781    /// `agent_framework.function.invocation.duration`, matching upstream's
782    /// `OPERATION_DURATION_BUCKET_BOUNDARIES` (`observability.py:81-96`).
783    pub const OPERATION_DURATION_BUCKET_BOUNDARIES: &[f64] = &[
784        0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92,
785    ];
786
787    /// `Meters.LLM_TOKEN_USAGE`: the token-usage histogram name.
788    pub const TOKEN_USAGE_METRIC: &str = "gen_ai.client.token.usage";
789    /// `Meters.LLM_OPERATION_DURATION`: the operation-duration histogram name.
790    pub const OPERATION_DURATION_METRIC: &str = "gen_ai.client.operation.duration";
791    /// `OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION`: the
792    /// function-invocation-duration histogram name.
793    pub const FUNCTION_INVOCATION_DURATION_METRIC: &str =
794        "agent_framework.function.invocation.duration";
795
796    const TOKEN_TYPE: &str = "gen_ai.token.type";
797    const TOKEN_TYPE_INPUT: &str = "input";
798    const TOKEN_TYPE_OUTPUT: &str = "output";
799    const FUNCTION_NAME: &str = "agent_framework.function.name";
800
801    struct Metrics {
802        token_usage: Histogram<u64>,
803        operation_duration: Histogram<f64>,
804        function_invocation_duration: Histogram<f64>,
805    }
806
807    static METRICS: OnceLock<Metrics> = OnceLock::new();
808
809    /// The lazily-initialized instrument set, bound to whichever
810    /// [`opentelemetry::global`] meter provider is installed the first time
811    /// this is called in this process — matches `global::meter`'s own
812    /// "bound at call time" semantics, so install your provider before
813    /// running any instrumented code.
814    fn instruments() -> &'static Metrics {
815        METRICS.get_or_init(|| {
816            let meter = opentelemetry::global::meter("agent_framework");
817            Metrics {
818                token_usage: meter
819                    .u64_histogram(TOKEN_USAGE_METRIC)
820                    .with_unit("tokens")
821                    .with_description("Captures the token usage of chat clients")
822                    .with_boundaries(TOKEN_USAGE_BUCKET_BOUNDARIES.to_vec())
823                    .build(),
824                operation_duration: meter
825                    .f64_histogram(OPERATION_DURATION_METRIC)
826                    .with_unit("s")
827                    .with_description("Captures the duration of chat client operations")
828                    .with_boundaries(OPERATION_DURATION_BUCKET_BOUNDARIES.to_vec())
829                    .build(),
830                function_invocation_duration: meter
831                    .f64_histogram(FUNCTION_INVOCATION_DURATION_METRIC)
832                    .with_unit("s")
833                    .with_description("Measures the duration of a function's execution")
834                    .with_boundaries(OPERATION_DURATION_BUCKET_BOUNDARIES.to_vec())
835                    .build(),
836            }
837        })
838    }
839
840    /// Record one completed chat call's histograms, mirroring
841    /// `_capture_response` (`observability.py:1525-1543`): the token-usage
842    /// histogram records once per token type present on the response; the
843    /// operation-duration histogram always records. Only called from the
844    /// success path of [`super::ObservableChatClient`] — like upstream, a
845    /// failed call records neither histogram.
846    pub(super) fn record_chat_completion(
847        provider: &str,
848        request_model: &str,
849        response_model: Option<&str>,
850        input_tokens: Option<u64>,
851        output_tokens: Option<u64>,
852        duration: Duration,
853    ) {
854        let m = instruments();
855        let mut base = vec![
856            KeyValue::new(attr::OPERATION, op::CHAT),
857            KeyValue::new(attr::PROVIDER_NAME, provider.to_string()),
858            KeyValue::new(attr::REQUEST_MODEL, request_model.to_string()),
859        ];
860        if let Some(model) = response_model {
861            base.push(KeyValue::new(attr::RESPONSE_MODEL, model.to_string()));
862        }
863        if let Some(input) = input_tokens {
864            let mut attrs = base.clone();
865            attrs.push(KeyValue::new(TOKEN_TYPE, TOKEN_TYPE_INPUT));
866            m.token_usage.record(input, &attrs);
867        }
868        if let Some(output) = output_tokens {
869            let mut attrs = base.clone();
870            attrs.push(KeyValue::new(TOKEN_TYPE, TOKEN_TYPE_OUTPUT));
871            m.token_usage.record(output, &attrs);
872        }
873        m.operation_duration.record(duration.as_secs_f64(), &base);
874    }
875
876    /// Record the function-invocation-duration histogram for one tool call.
877    ///
878    /// Not yet called anywhere in this crate: the timing measurement belongs
879    /// around `exec.invoke(...)` in `client.rs`'s
880    /// `FunctionInvokingChatClient::execute_tool_call`, which is out of scope
881    /// here (see the observability task's final report for the exact
882    /// follow-up). The instrument and its recording logic are complete and
883    /// tested on their own so that call site only needs to wrap its
884    /// invocation with a timer and call this — plus, ideally, switch
885    /// `tool_span` to [`super::tool_span_ex`] and add
886    /// [`super::record_tool_arguments`] / [`super::record_tool_result`] calls
887    /// at the same time.
888    pub fn record_function_invocation_duration(
889        tool_name: &str,
890        duration: Duration,
891        error_type: Option<&str>,
892    ) {
893        let m = instruments();
894        let mut attrs = vec![KeyValue::new(FUNCTION_NAME, tool_name.to_string())];
895        if let Some(err) = error_type {
896            attrs.push(KeyValue::new(attr::ERROR_TYPE, err.to_string()));
897        }
898        m.function_invocation_duration
899            .record(duration.as_secs_f64(), &attrs);
900    }
901}
902
903#[cfg(test)]
904mod tests {
905    use super::*;
906
907    // -- Attribute string values (cross-language wire contract) ----------
908
909    #[test]
910    fn cache_reasoning_and_embedding_attrs_match_upstream() {
911        // These strings are the OTel GenAI attribute keys upstream emits;
912        // they must match exactly for cross-tool/cross-language consistency.
913        assert_eq!(
914            attr::CACHE_CREATION_INPUT_TOKENS,
915            "gen_ai.usage.cache_creation.input_tokens"
916        );
917        assert_eq!(
918            attr::CACHE_READ_INPUT_TOKENS,
919            "gen_ai.usage.cache_read.input_tokens"
920        );
921        assert_eq!(
922            attr::REASONING_OUTPUT_TOKENS,
923            "gen_ai.usage.reasoning.output_tokens"
924        );
925        assert_eq!(attr::PROMPT_NAME, "gen_ai.prompt.name");
926        assert_eq!(op::EMBEDDINGS, "embeddings");
927    }
928
929    // -- ObservabilityConfig::from_env -----------------------------------
930
931    /// Guards env var mutation: tests run on multiple threads within a
932    /// crate, and env vars are process-global (same pattern as
933    /// `agent-framework-copilotstudio`'s / `agent-framework-mem0`'s
934    /// `ENV_MUTEX`).
935    static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
936
937    fn clear_env() {
938        // SAFETY: serialized by ENV_MUTEX; no other test in this crate
939        // touches these variables (confirmed via a repo-wide grep before
940        // adding this).
941        unsafe {
942            std::env::remove_var("ENABLE_SENSITIVE_DATA");
943            std::env::remove_var("OTEL_SERVICE_NAME");
944        }
945    }
946
947    #[test]
948    fn from_env_defaults_to_disabled_when_unset() {
949        let _guard = ENV_MUTEX.lock().unwrap();
950        clear_env();
951        let config = ObservabilityConfig::from_env();
952        assert!(!config.enable_sensitive_data);
953        assert_eq!(config.otel_service_name(), "agent_framework");
954    }
955
956    #[test]
957    fn from_env_reads_truthy_enable_sensitive_data() {
958        let _guard = ENV_MUTEX.lock().unwrap();
959        for value in ["1", "true", "TRUE", "  yes  ", "on"] {
960            clear_env();
961            // SAFETY: serialized by ENV_MUTEX.
962            unsafe { std::env::set_var("ENABLE_SENSITIVE_DATA", value) };
963            let config = ObservabilityConfig::from_env();
964            assert!(
965                config.enable_sensitive_data,
966                "expected {value:?} to be truthy"
967            );
968        }
969        clear_env();
970    }
971
972    #[test]
973    fn from_env_rejects_unrecognized_values() {
974        let _guard = ENV_MUTEX.lock().unwrap();
975        clear_env();
976        // SAFETY: serialized by ENV_MUTEX.
977        unsafe { std::env::set_var("ENABLE_SENSITIVE_DATA", "nope") };
978        let config = ObservabilityConfig::from_env();
979        clear_env();
980        assert!(!config.enable_sensitive_data);
981    }
982
983    #[test]
984    fn otel_service_name_reads_env_override() {
985        let _guard = ENV_MUTEX.lock().unwrap();
986        clear_env();
987        // SAFETY: serialized by ENV_MUTEX.
988        unsafe { std::env::set_var("OTEL_SERVICE_NAME", "my-service") };
989        let config = ObservabilityConfig::from_env();
990        let name = config.otel_service_name();
991        clear_env();
992        assert_eq!(name, "my-service");
993    }
994
995    // -- error_type --------------------------------------------------------
996
997    #[test]
998    fn error_type_gives_the_granular_service_errors_distinct_tags() {
999        // The three newer variants get their own `error.type` tags, distinct
1000        // from both each other and the generic "service"/"service" tags that
1001        // `Error::Service`/`Error::ServiceStatus` share — that granularity is
1002        // the whole point of adding them.
1003        assert_eq!(
1004            error_type(&Error::service_invalid_auth("x")),
1005            "service_invalid_auth"
1006        );
1007        assert_eq!(
1008            error_type(&Error::service_invalid_request("x")),
1009            "service_invalid_request"
1010        );
1011        assert_eq!(
1012            error_type(&Error::service_content_filter("x")),
1013            "service_content_filter"
1014        );
1015        assert_eq!(error_type(&Error::service("x")), "service");
1016        assert_eq!(
1017            error_type(&Error::service_status(500, "x", None)),
1018            "service"
1019        );
1020    }
1021}