dci-tool 0.1.0

//! Telemetry helpers for DCI.
//!
//! Provides shims around `rig-tap` to capture latency and costs.

#[cfg(feature = "telemetry")]
use rig_tap::{EventKind, emit_kind};

#[cfg(feature = "telemetry")]
tokio::task_local! {
    /// Task-local tracking of the current active session/conversation for telemetry.
    pub static CURRENT_SESSION: String;
}

/// Run an asynchronous closure within a telemetry session scope.
pub async fn with_session<F, Fut>(session_id: String, f: F) -> Fut::Output
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future,
{
    #[cfg(feature = "telemetry")]
    {
        CURRENT_SESSION.scope(session_id, f()).await
    }
    #[cfg(not(feature = "telemetry"))]
    {
        let _ = session_id;
        f().await
    }
}

/// Retrieve the active session ID if one is set.
pub fn current_session() -> String {
    #[cfg(feature = "telemetry")]
    {
        CURRENT_SESSION
            .try_with(|s| s.clone())
            .unwrap_or_else(|_| "unknown".to_string())
    }
    #[cfg(not(feature = "telemetry"))]
    {
        "unknown".to_string()
    }
}

/// Emit a prompt-completed telemetry event carrying token usage.
///
/// This is the version-safe alternative to attaching `rig-tap`'s typed
/// `TelemetryHook` directly: that hook is generic over `rig-tap`'s pinned
/// rig-core (0.37) `CompletionModel`, which does not unify with this crate's
/// rig-core (0.38). `EventKind::PromptCompleted` is a plain event envelope, so
/// emitting it ourselves crosses no rig-core type and works across the skew.
///
/// `cost_usd` is intentionally left unset: the crate ships no pricing table, so
/// token counts are reported and cost is left to downstream consumers.
#[cfg_attr(not(feature = "telemetry"), allow(unused_variables))]
pub fn record_prompt(
    model: &str,
    tokens_in: u64,
    tokens_out: u64,
    cached_in: u64,
    reasoning: u64,
    duration_ms: u64,
) {
    #[cfg(feature = "telemetry")]
    emit_kind(
        current_session(),
        EventKind::PromptCompleted {
            model: model.to_string(),
            tokens_in: Some(tokens_in),
            tokens_out: Some(tokens_out),
            cached_tokens_in: Some(cached_in),
            reasoning_tokens: Some(reasoning),
            cost_usd: None,
            finish_reason: None,
            response_id: None,
            previous_response_id: None,
            time_to_first_token_ms: None,
            duration_ms: Some(duration_ms),
        },
    );
}

/// Record a tool invocation and completion with latency.
pub async fn record_tool_call<F, Fut, T, E>(
    #[allow(unused_variables)] tool_name: &'static str,
    #[allow(unused_variables)] args: &str,
    f: F,
) -> Result<T, E>
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = Result<T, E>>,
    T: serde::Serialize,
    E: std::fmt::Display,
{
    #[cfg(feature = "telemetry")]
    let start = std::time::Instant::now();
    #[cfg(feature = "telemetry")]
    let call_id = format!("{tool_name}-{}", uuid::Uuid::new_v4());
    #[cfg(feature = "telemetry")]
    let session = current_session();

    #[cfg(feature = "telemetry")]
    emit_kind(
        &session,
        EventKind::ToolInvoked {
            tool_name: tool_name.to_string(),
            provider_call_id: None,
            call_id: call_id.clone(),
            args_json: args.to_string(),
            truncated: false,
        },
    );

    let res = f().await;

    #[cfg(feature = "telemetry")]
    match &res {
        Ok(val) => {
            let result_str = serde_json::to_string(val).unwrap_or_else(|_| "{}".to_string());
            emit_kind(
                &session,
                EventKind::ToolCompleted {
                    tool_name: tool_name.to_string(),
                    provider_call_id: None,
                    call_id,
                    result: result_str,
                    truncated: false,
                    duration_ms: Some(start.elapsed().as_millis() as u64),
                },
            );
        }
        Err(err) => {
            emit_kind(
                &session,
                EventKind::ToolFailed {
                    tool_name: tool_name.to_string(),
                    call_id,
                    error_class: rig_tap::ErrorClass::Unknown,
                    message: err.to_string(),
                },
            );
        }
    }

    res
}