gemini_adk_rs/telemetry/
mod.rs

1//! Agent-level observability — OpenTelemetry tracing, structured logging, Prometheus metrics.
2//!
3//! Mirrors the wire crate's telemetry layer but for agent lifecycle operations.
4//! All components are feature-gated for zero overhead when disabled:
5//! - `tracing-support`: OTel spans and structured logging
6//! - `metrics`: Prometheus metric definitions
7
8use async_trait::async_trait;
9
10use gemini_genai_rs::prelude::FunctionCall;
11
12use crate::context::InvocationContext;
13use crate::error::{AgentError, ToolError};
14use crate::llm::{LlmRequest, LlmResponse};
15use crate::middleware::Middleware;
16
17pub mod logging;
18pub mod metrics;
19pub mod setup;
20pub mod spans;
21
22pub use setup::TelemetrySetup;
23
24/// Auto-registered middleware that calls telemetry functions.
25/// Zero-overhead when `tracing-support` and `metrics` features are disabled
26/// (all telemetry functions compile to no-ops).
27///
28/// Automatically prepended to every LlmAgent's middleware chain at build time,
29/// so all agents get observability by default.
30pub struct TelemetryMiddleware {
31    agent_name: String,
32}
33
34impl TelemetryMiddleware {
35    /// Create a new telemetry middleware for the given agent.
36    pub fn new(agent_name: impl Into<String>) -> Self {
37        Self {
38            agent_name: agent_name.into(),
39        }
40    }
41
42    /// Returns the agent name this middleware is tracking.
43    pub fn agent_name(&self) -> &str {
44        &self.agent_name
45    }
46}
47
48#[async_trait]
49impl Middleware for TelemetryMiddleware {
50    fn name(&self) -> &str {
51        "telemetry"
52    }
53
54    async fn before_agent(&self, _ctx: &InvocationContext) -> Result<(), AgentError> {
55        metrics::record_agent_started(&self.agent_name);
56        logging::log_agent_started(&self.agent_name, 0);
57        Ok(())
58    }
59
60    async fn after_agent(&self, _ctx: &InvocationContext) -> Result<(), AgentError> {
61        // Duration tracked in LlmAgent::run_live() directly
62        Ok(())
63    }
64
65    async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
66        metrics::record_agent_tool_dispatched(&self.agent_name, &call.name);
67        logging::log_tool_dispatch(&self.agent_name, &call.name, "function");
68        Ok(())
69    }
70
71    async fn after_tool(
72        &self,
73        call: &FunctionCall,
74        _result: &serde_json::Value,
75    ) -> Result<(), AgentError> {
76        logging::log_tool_result(&self.agent_name, &call.name, true, 0.0);
77        Ok(())
78    }
79
80    async fn on_tool_error(&self, call: &FunctionCall, _err: &ToolError) -> Result<(), AgentError> {
81        logging::log_tool_result(&self.agent_name, &call.name, false, 0.0);
82        Ok(())
83    }
84
85    async fn on_error(&self, err: &AgentError) -> Result<(), AgentError> {
86        metrics::record_agent_error(&self.agent_name, &err.to_string());
87        logging::log_agent_error(&self.agent_name, &err.to_string());
88        Ok(())
89    }
90
91    async fn before_model(&self, _request: &LlmRequest) -> Result<Option<LlmResponse>, AgentError> {
92        logging::log_tool_dispatch(&self.agent_name, "llm", "model_call");
93        Ok(None)
94    }
95
96    async fn after_model(
97        &self,
98        _request: &LlmRequest,
99        _response: &LlmResponse,
100    ) -> Result<Option<LlmResponse>, AgentError> {
101        if let Some(_usage) = &_response.usage {
102            metrics::record_agent_completed(&self.agent_name, 0.0); // Duration tracked elsewhere
103        }
104        Ok(None)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[tokio::test]
113    async fn telemetry_middleware_hooks_dont_panic() {
114        let mw = TelemetryMiddleware::new("test_agent");
115        assert_eq!(mw.name(), "telemetry");
116        assert_eq!(mw.agent_name(), "test_agent");
117
118        let call = FunctionCall {
119            name: "my_tool".to_string(),
120            args: serde_json::json!({}),
121            id: None,
122        };
123        let result = serde_json::json!({"ok": true});
124        let tool_err = ToolError::ExecutionFailed("boom".to_string());
125        let agent_err = AgentError::Other("oops".to_string());
126
127        // All hooks should complete without panic
128        assert!(mw.before_tool(&call).await.is_ok());
129        assert!(mw.after_tool(&call, &result).await.is_ok());
130        assert!(mw.on_tool_error(&call, &tool_err).await.is_ok());
131        assert!(mw.on_error(&agent_err).await.is_ok());
132    }
133}