dci_tool/
telemetry.rs

1//! Telemetry helpers for DCI.
2//!
3//! Provides shims around `rig-tap` to capture latency and costs.
4
5#[cfg(feature = "telemetry")]
6use rig_tap::{EventKind, emit_kind};
7
8#[cfg(feature = "telemetry")]
9tokio::task_local! {
10    /// Task-local tracking of the current active session/conversation for telemetry.
11    pub static CURRENT_SESSION: String;
12}
13
14/// Run an asynchronous closure within a telemetry session scope.
15pub async fn with_session<F, Fut>(session_id: String, f: F) -> Fut::Output
16where
17    F: FnOnce() -> Fut,
18    Fut: std::future::Future,
19{
20    #[cfg(feature = "telemetry")]
21    {
22        CURRENT_SESSION.scope(session_id, f()).await
23    }
24    #[cfg(not(feature = "telemetry"))]
25    {
26        let _ = session_id;
27        f().await
28    }
29}
30
31/// Retrieve the active session ID if one is set.
32pub fn current_session() -> String {
33    #[cfg(feature = "telemetry")]
34    {
35        CURRENT_SESSION
36            .try_with(|s| s.clone())
37            .unwrap_or_else(|_| "unknown".to_string())
38    }
39    #[cfg(not(feature = "telemetry"))]
40    {
41        "unknown".to_string()
42    }
43}
44
45/// Emit a prompt-completed telemetry event carrying token usage.
46///
47/// This is the version-safe alternative to attaching `rig-tap`'s typed
48/// `TelemetryHook` directly: that hook is generic over `rig-tap`'s pinned
49/// rig-core (0.37) `CompletionModel`, which does not unify with this crate's
50/// rig-core (0.38). `EventKind::PromptCompleted` is a plain event envelope, so
51/// emitting it ourselves crosses no rig-core type and works across the skew.
52///
53/// `cost_usd` is intentionally left unset: the crate ships no pricing table, so
54/// token counts are reported and cost is left to downstream consumers.
55#[cfg_attr(not(feature = "telemetry"), allow(unused_variables))]
56pub fn record_prompt(
57    model: &str,
58    tokens_in: u64,
59    tokens_out: u64,
60    cached_in: u64,
61    reasoning: u64,
62    duration_ms: u64,
63) {
64    #[cfg(feature = "telemetry")]
65    emit_kind(
66        current_session(),
67        EventKind::PromptCompleted {
68            model: model.to_string(),
69            tokens_in: Some(tokens_in),
70            tokens_out: Some(tokens_out),
71            cached_tokens_in: Some(cached_in),
72            reasoning_tokens: Some(reasoning),
73            cost_usd: None,
74            finish_reason: None,
75            response_id: None,
76            previous_response_id: None,
77            time_to_first_token_ms: None,
78            duration_ms: Some(duration_ms),
79        },
80    );
81}
82
83/// Record a tool invocation and completion with latency.
84pub async fn record_tool_call<F, Fut, T, E>(
85    #[allow(unused_variables)] tool_name: &'static str,
86    #[allow(unused_variables)] args: &str,
87    f: F,
88) -> Result<T, E>
89where
90    F: FnOnce() -> Fut,
91    Fut: std::future::Future<Output = Result<T, E>>,
92    T: serde::Serialize,
93    E: std::fmt::Display,
94{
95    #[cfg(feature = "telemetry")]
96    let start = std::time::Instant::now();
97    #[cfg(feature = "telemetry")]
98    let call_id = format!("{tool_name}-{}", uuid::Uuid::new_v4());
99    #[cfg(feature = "telemetry")]
100    let session = current_session();
101
102    #[cfg(feature = "telemetry")]
103    emit_kind(
104        &session,
105        EventKind::ToolInvoked {
106            tool_name: tool_name.to_string(),
107            provider_call_id: None,
108            call_id: call_id.clone(),
109            args_json: args.to_string(),
110            truncated: false,
111        },
112    );
113
114    let res = f().await;
115
116    #[cfg(feature = "telemetry")]
117    match &res {
118        Ok(val) => {
119            let result_str = serde_json::to_string(val).unwrap_or_else(|_| "{}".to_string());
120            emit_kind(
121                &session,
122                EventKind::ToolCompleted {
123                    tool_name: tool_name.to_string(),
124                    provider_call_id: None,
125                    call_id,
126                    result: result_str,
127                    truncated: false,
128                    duration_ms: Some(start.elapsed().as_millis() as u64),
129                },
130            );
131        }
132        Err(err) => {
133            emit_kind(
134                &session,
135                EventKind::ToolFailed {
136                    tool_name: tool_name.to_string(),
137                    call_id,
138                    error_class: rig_tap::ErrorClass::Unknown,
139                    message: err.to_string(),
140                },
141            );
142        }
143    }
144
145    res
146}
dci_tool/telemetry.rs

dci_tool/
telemetry.rs
