phi_telemetry/
collector.rs

1//! Metrics collector — subscribes to turn-end hooks and accumulates session metrics.
2//!
3//! Runs in an independent tokio task via channel isolation so that observer
4//! panics never affect the agent runtime.
5
6use agent_base::TurnContext;
7use serde_json::Value;
8use std::sync::{Arc, Mutex};
9use tokio::sync::RwLock;
10use tokio::sync::mpsc;
11use tokio::task::JoinHandle;
12use tracing;
13
14use crate::types::{SessionMetrics, TurnMetrics, run_outcome_to_turn_outcome};
15
16/// Message sent from the hook callback to the observer task.
17#[derive(Clone, Debug)]
18pub(crate) enum ObserverMsg {
19    /// A turn has completed — build TurnMetrics and accumulate.
20    TurnEnd(TurnContext),
21    /// Set session-level custom data.
22    SetSessionCustom(Value),
23    /// Shut down the observer and finalize metrics.
24    Shutdown,
25}
26
27/// Handle to the background observer task.
28///
29/// Call [`shutdown`](Self::shutdown) to gracefully stop the observer and wait
30/// for all pending metrics to be processed. Dropping without calling `shutdown`
31/// detaches the task — metrics from the final turns may be lost.
32pub struct ObserverHandle {
33    tx: mpsc::UnboundedSender<ObserverMsg>,
34    /// Shared session metrics, readable for external use (e.g. CLI display).
35    pub session: Arc<RwLock<SessionMetrics>>,
36    /// Handle to the observer task. None after `shutdown()` has been called.
37    task: Option<JoinHandle<()>>,
38    /// Pending turn-level custom data. Set by the consumer (e.g. phi-bard)
39    /// via `on_turn_end`, consumed by the observer task when building TurnMetrics.
40    pending_turn_custom: Arc<Mutex<Option<Value>>>,
41}
42
43impl ObserverHandle {
44    /// Send session-level custom data to the observer.
45    pub fn set_session_custom(&self, custom: Value) {
46        let _ = self.tx.send(ObserverMsg::SetSessionCustom(custom));
47    }
48
49    /// Set turn-level custom data for the NEXT turn that the observer processes.
50    ///
51    /// Call this from an `on_turn_end` callback. The custom value will be
52    /// merged into that turn's `custom` field in `TurnMetrics`.
53    ///
54    /// ```ignore
55    /// agent.runtime().on_turn_end(move |_ctx| {
56    ///     handle.set_turn_custom(json!({"check_quality_passed": true}));
57    /// });
58    /// ```
59    pub fn set_turn_custom(&self, custom: Value) {
60        if let Ok(mut pending) = self.pending_turn_custom.lock() {
61            *pending = Some(custom);
62        }
63    }
64
65    /// Shut down the observer gracefully and wait for all pending metrics
66    /// to be processed. Returns once the observer task has exited.
67    ///
68    /// After this call, [`session`](Self::session) contains the final
69    /// accumulated metrics (minus any `finalize()` call, which the caller
70    /// should perform separately).
71    pub async fn shutdown(&mut self) {
72        let _ = self.tx.send(ObserverMsg::Shutdown);
73        if let Some(task) = self.task.take() {
74            // Don't propagate panic — observer panics must not crash the caller.
75            let _ = task.await;
76        }
77    }
78}
79
80/// Initialise telemetry for an agent runtime.
81///
82/// Registers an `on_turn_end` hook that sends `TurnContext` data through a
83/// channel to an independent observer task. The observer task builds
84/// `TurnMetrics` and accumulates `SessionMetrics` — it never blocks the
85/// agent runtime hot path.
86///
87/// Returns an [`ObserverHandle`] that can be used to inject custom data or
88/// shut down the observer.
89pub fn init_telemetry(
90    runtime: &agent_base::AgentRuntime,
91    session_id: String,
92    node_id: String,
93    model: String,
94) -> ObserverHandle {
95    let (tx, mut rx) = mpsc::unbounded_channel::<ObserverMsg>();
96
97    // Hook: only sends a message — no heavy work on the hot path.
98    let hook_tx = tx.clone();
99    runtime.on_turn_end(move |ctx: &TurnContext| {
100        let _ = hook_tx.send(ObserverMsg::TurnEnd(ctx.clone()));
101    });
102
103    let session = Arc::new(RwLock::new(SessionMetrics::new(session_id, node_id, model)));
104
105    let observer_session = session.clone();
106    let pending_turn_custom = Arc::new(Mutex::new(None::<Value>));
107    let pending = pending_turn_custom.clone();
108
109    // Independent task: build metrics from TurnContext
110    let task = tokio::spawn(async move {
111        let accumulator = observer_session;
112        while let Some(msg) = rx.recv().await {
113            match msg {
114                ObserverMsg::TurnEnd(ctx) => {
115                    let turn = {
116                        let mut turn = build_turn_metrics(&ctx);
117                        // Apply any pending turn-level custom data
118                        if let Ok(mut pending) = pending.lock()
119                            && let Some(custom) = pending.take()
120                            && let Value::Object(ref mut map) = turn.custom
121                            && let Value::Object(custom_map) = custom
122                        {
123                            for (k, v) in custom_map {
124                                map.insert(k, v);
125                            }
126                        }
127                        turn
128                    };
129                    let mut session = accumulator.write().await;
130                    session.append_turn(turn);
131                    tracing::debug!(turn = session.total_turns, "metrics: turn accumulated");
132                }
133                ObserverMsg::SetSessionCustom(custom) => {
134                    let mut session = accumulator.write().await;
135                    if let Value::Object(ref mut map) = session.custom
136                        && let Value::Object(custom_map) = custom
137                    {
138                        for (k, v) in custom_map {
139                            map.insert(k, v);
140                        }
141                    }
142                }
143                ObserverMsg::Shutdown => {
144                    tracing::debug!("metrics: observer shutting down");
145                    break;
146                }
147            }
148        }
149    });
150
151    ObserverHandle {
152        tx,
153        session,
154        task: Some(task),
155        pending_turn_custom,
156    }
157}
158
159/// Build a TurnMetrics from the raw TurnContext.
160fn build_turn_metrics(ctx: &TurnContext) -> TurnMetrics {
161    let input_tokens = ctx
162        .usage
163        .as_ref()
164        .and_then(|u| u.prompt_tokens)
165        .unwrap_or(0) as u64;
166    let output_tokens = ctx
167        .usage
168        .as_ref()
169        .and_then(|u| u.completion_tokens)
170        .unwrap_or(0) as u64;
171
172    let duration_ms = ctx.duration_ms;
173
174    let turn_outcome = run_outcome_to_turn_outcome(&ctx.outcome, &ctx.tools_used);
175
176    let mut turn = TurnMetrics::new(
177        ctx.turn_number,
178        chrono::Utc::now().to_rfc3339(),
179        duration_ms,
180        ctx.model.clone(),
181        ctx.user_input.clone(),
182        turn_outcome,
183    );
184
185    turn.time_to_first_token_ms = ctx.ttft_ms;
186    turn.llm_duration_ms = ctx.llm_duration_ms;
187    turn.tool_duration_ms = ctx.tool_duration_ms;
188    turn.input_tokens = input_tokens;
189    turn.output_tokens = output_tokens;
190    turn.tool_call_count = ctx.tool_call_count;
191    turn.tools_used = ctx.tools_used.clone();
192    turn.tool_success = ctx.tool_success;
193    turn.tool_failed = ctx.tool_failed;
194    turn.text_length = ctx.full_text_len;
195    turn.has_thinking = ctx.has_thinking;
196    turn.plan_updates = ctx.plan_updates;
197    turn.approval_count = ctx.approval_count;
198    turn.llm_calls = ctx.llm_calls;
199    if let Some(ref msg) = ctx.error_message {
200        turn.error_message = Some(msg.clone());
201    }
202
203    turn
204}
phi_telemetry/collector.rs

phi_telemetry/
collector.rs
