embacle 0.32.0

// ABOUTME: Decorator wrapping any LlmProvider to measure latency, token usage, cost, and call counts
// ABOUTME: Provides MetricsReport snapshots for cost and performance normalization, optional OTel export
//
// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 dravr.ai

//! # Cost/Latency Normalization
//!
//! [`MetricsProvider`] is a decorator that wraps any `Box<dyn LlmProvider>`,
//! measuring per-call latency and token usage. Callers can retrieve a
//! [`MetricsReport`] snapshot at any time via [`MetricsProvider::report()`].
//!
//! ## Cost Tracking
//!
//! Every call is priced through [`crate::pricing`], keyed on the wrapped
//! provider's `name()` and the response's model, with the cache and
//! reasoning counts the provider reported. A pair the table does not price
//! costs `0.0` (see [`crate::pricing::calculate_cost_for`] for how that is
//! logged). Costs accumulate in the report.
//!
//! ## OpenTelemetry (feature: `otel`)
//!
//! When built with `--features otel`, instruments are created via the
//! global meter `embacle` and recorded on each `complete()` call.
//!
//! Token estimation: when `TokenUsage` is not provided by the inner provider,
//! tokens are estimated at ~4 characters per token.

use std::sync::{Arc, Mutex};
use std::time::Instant;

#[cfg(feature = "otel")]
use opentelemetry::global;
#[cfg(feature = "otel")]
use opentelemetry::metrics::{Counter, Histogram};

use async_trait::async_trait;
use tracing::info;

use crate::pricing::{calculate_cost_for, TokenCounts};
use crate::types::{
    ChatRequest, ChatResponse, ChatStream, LlmCapabilities, LlmProvider, RunnerError, TokenUsage,
};

/// Characters-per-token estimate used when the provider does not report usage
const CHARS_PER_TOKEN_ESTIMATE: u32 = 4;

/// Accumulated metrics state protected by a mutex
#[derive(Debug, Default)]
struct MetricsState {
    call_count: u64,
    total_latency_ms: u64,
    total_prompt_tokens: u64,
    total_completion_tokens: u64,
    total_tokens: u64,
    errors_count: u64,
    total_cost: f64,
}

/// Snapshot of accumulated metrics for a provider
#[derive(Debug, Clone)]
pub struct MetricsReport {
    /// Name of the wrapped provider
    pub provider_name: String,
    /// Total number of `complete()` calls
    pub call_count: u64,
    /// Total latency across all calls (milliseconds)
    pub total_latency_ms: u64,
    /// Average latency per call (milliseconds)
    pub avg_latency_ms: u64,
    /// Total prompt tokens consumed
    pub total_prompt_tokens: u64,
    /// Total completion tokens generated
    pub total_completion_tokens: u64,
    /// Total tokens (prompt + completion)
    pub total_tokens: u64,
    /// Number of calls that returned an error
    pub errors_count: u64,
    /// Accumulated cost in USD (0.0 for every call the price table does not cover)
    pub total_cost: f64,
}

/// OpenTelemetry instruments for metrics export
#[cfg(feature = "otel")]
struct OtelInstruments {
    requests_total: Counter<u64>,
    requests_duration_ms: Histogram<f64>,
    tokens_prompt: Counter<u64>,
    tokens_completion: Counter<u64>,
    errors_total: Counter<u64>,
    cost_total: Counter<f64>,
}

#[cfg(feature = "otel")]
impl OtelInstruments {
    fn new() -> Self {
        let meter = global::meter("embacle");
        Self {
            requests_total: meter
                .u64_counter("embacle.requests.total")
                .with_description("Total LLM requests")
                .build(),
            requests_duration_ms: meter
                .f64_histogram("embacle.requests.duration_ms")
                .with_description("Request duration in milliseconds")
                .build(),
            tokens_prompt: meter
                .u64_counter("embacle.tokens.prompt")
                .with_description("Total prompt tokens consumed")
                .build(),
            tokens_completion: meter
                .u64_counter("embacle.tokens.completion")
                .with_description("Total completion tokens generated")
                .build(),
            errors_total: meter
                .u64_counter("embacle.errors.total")
                .with_description("Total error count")
                .build(),
            cost_total: meter
                .f64_counter("embacle.cost.total")
                .with_description("Total cost in USD")
                .build(),
        }
    }
}

/// Decorator wrapping any `Box<dyn LlmProvider>` to collect latency, token, and cost metrics.
///
/// # Usage
///
/// ```rust,no_run
/// # use embacle::metrics::MetricsProvider;
/// # use embacle::types::LlmProvider;
/// # fn example(provider: Box<dyn LlmProvider>) {
/// let metered = MetricsProvider::new(provider);
/// // ... use metered as LlmProvider ...
/// let report = metered.report().unwrap(); // Safe: test assertion
/// println!("calls={} avg_latency={}ms", report.call_count, report.avg_latency_ms);
/// # }
/// ```
pub struct MetricsProvider {
    inner: Box<dyn LlmProvider>,
    state: Arc<Mutex<MetricsState>>,
    #[cfg(feature = "otel")]
    otel: OtelInstruments,
}

impl MetricsProvider {
    /// Wrap a provider with metrics collection
    pub fn new(inner: Box<dyn LlmProvider>) -> Self {
        Self {
            inner,
            state: Arc::new(Mutex::new(MetricsState::default())),
            #[cfg(feature = "otel")]
            otel: OtelInstruments::new(),
        }
    }

    /// Return a snapshot of the current metrics
    ///
    /// # Errors
    ///
    /// Returns [`RunnerError`] if the internal mutex is poisoned.
    pub fn report(&self) -> Result<MetricsReport, RunnerError> {
        let state = self
            .state
            .lock()
            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
        let divisor = state.call_count.max(1);
        Ok(MetricsReport {
            provider_name: self.inner.name().to_owned(),
            call_count: state.call_count,
            total_latency_ms: state.total_latency_ms,
            avg_latency_ms: state.total_latency_ms / divisor,
            total_prompt_tokens: state.total_prompt_tokens,
            total_completion_tokens: state.total_completion_tokens,
            total_tokens: state.total_tokens,
            errors_count: state.errors_count,
            total_cost: state.total_cost,
        })
    }

    /// Reset all counters to zero
    ///
    /// # Errors
    ///
    /// Returns [`RunnerError`] if the internal mutex is poisoned.
    pub fn reset(&self) -> Result<(), RunnerError> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
        *state = MetricsState::default();
        Ok(())
    }

    /// Price one call: the wrapped provider's `name()` is the table key, and
    /// the cache and reasoning counts the provider reported ride along.
    /// `prompt_tokens`/`completion_tokens` are the (possibly estimated)
    /// headline counts; `usage` supplies the optional categories.
    fn compute_cost(
        &self,
        model: &str,
        prompt_tokens: u64,
        completion_tokens: u64,
        usage: Option<&TokenUsage>,
    ) -> f64 {
        let count = |n: u64| i64::try_from(n).unwrap_or(i64::MAX);
        let optional = |n: Option<u32>| i64::from(n.unwrap_or(0));
        let counts = TokenCounts::new(count(prompt_tokens), count(completion_tokens))
            .with_cache(
                optional(usage.and_then(|u| u.cached_read_tokens)),
                optional(usage.and_then(|u| u.cached_write_tokens)),
            )
            .with_reasoning(optional(usage.and_then(|u| u.reasoning_tokens)));
        calculate_cost_for(self.inner.name(), model, &counts)
    }
}

/// Estimate token count from character length (~4 chars per token)
fn estimate_tokens(text: &str) -> u32 {
    #[allow(clippy::cast_possible_truncation)]
    let len = text.len() as u32;
    len / CHARS_PER_TOKEN_ESTIMATE.max(1)
}

#[async_trait]
impl LlmProvider for MetricsProvider {
    fn name(&self) -> &'static str {
        self.inner.name()
    }

    fn display_name(&self) -> &str {
        self.inner.display_name()
    }

    fn capabilities(&self) -> LlmCapabilities {
        self.inner.capabilities()
    }

    fn default_model(&self) -> &str {
        self.inner.default_model()
    }

    fn available_models(&self) -> &[String] {
        self.inner.available_models()
    }

    async fn complete(&self, request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
        let start = Instant::now();
        let result = self.inner.complete(request).await;
        #[allow(clippy::cast_possible_truncation)]
        let elapsed_ms = start.elapsed().as_millis() as u64;

        let mut state = self
            .state
            .lock()
            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
        state.call_count += 1;
        state.total_latency_ms += elapsed_ms;

        #[cfg(feature = "otel")]
        let otel_attrs: Vec<opentelemetry::KeyValue> = {
            let mut attrs = vec![opentelemetry::KeyValue::new("provider", self.inner.name())];
            if let Some(turn_id) = request.turn_id {
                attrs.push(opentelemetry::KeyValue::new("turn_id", turn_id.to_string()));
            }
            attrs
        };

        if let Ok(response) = &result {
            let usage = response.usage.as_ref();
            let prompt_tokens = u64::from(
                usage.map_or_else(|| estimate_prompt_tokens(request), |u| u.prompt_tokens),
            );
            let completion_tokens = u64::from(usage.map_or_else(
                || estimate_tokens(&response.content),
                |u| u.completion_tokens,
            ));
            let total = prompt_tokens + completion_tokens;

            state.total_prompt_tokens += prompt_tokens;
            state.total_completion_tokens += completion_tokens;
            state.total_tokens += total;

            let cost = self.compute_cost(&response.model, prompt_tokens, completion_tokens, usage);
            state.total_cost += cost;

            info!(
                provider = self.inner.name(),
                elapsed_ms, prompt_tokens, completion_tokens, cost, "metrics: complete() succeeded"
            );

            #[cfg(feature = "otel")]
            {
                let attrs = otel_attrs.as_slice();
                self.otel.requests_total.add(1, attrs);
                #[allow(clippy::cast_precision_loss)]
                self.otel
                    .requests_duration_ms
                    .record(elapsed_ms as f64, attrs);
                self.otel.tokens_prompt.add(prompt_tokens, attrs);
                self.otel.tokens_completion.add(completion_tokens, attrs);
                if cost > 0.0 {
                    self.otel.cost_total.add(cost, attrs);
                }
            }
        } else {
            state.errors_count += 1;
            info!(
                provider = self.inner.name(),
                elapsed_ms, "metrics: complete() failed"
            );

            #[cfg(feature = "otel")]
            {
                self.otel.errors_total.add(1, otel_attrs.as_slice());
            }
        }

        drop(state);

        result
    }

    /// Delegate streaming directly; only measures stream setup time (documented limitation)
    async fn complete_stream(&self, request: &ChatRequest) -> Result<ChatStream, RunnerError> {
        self.inner.complete_stream(request).await
    }

    async fn health_check(&self) -> Result<bool, RunnerError> {
        self.inner.health_check().await
    }
}

/// Estimate prompt tokens from request messages
fn estimate_prompt_tokens(request: &ChatRequest) -> u32 {
    let total_chars: usize = request.messages.iter().map(|m| m.content.len()).sum();
    #[allow(clippy::cast_possible_truncation)]
    let len = total_chars as u32;
    len / CHARS_PER_TOKEN_ESTIMATE.max(1)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{
        ChatMessage, ChatRequest, ChatResponse, ChatStream, LlmCapabilities, LlmProvider,
        RunnerError, TokenUsage,
    };
    use async_trait::async_trait;
    use std::sync::atomic::{AtomicU32, Ordering};

    struct TestProvider {
        provider_name: &'static str,
        responses: Mutex<Vec<Result<ChatResponse, RunnerError>>>,
        call_count: AtomicU32,
    }

    impl TestProvider {
        fn new(responses: Vec<Result<ChatResponse, RunnerError>>) -> Self {
            Self::named("test", responses)
        }

        /// A provider reporting `name`, so the price table keys on it.
        fn named(name: &'static str, responses: Vec<Result<ChatResponse, RunnerError>>) -> Self {
            Self {
                provider_name: name,
                responses: Mutex::new(responses),
                call_count: AtomicU32::new(0),
            }
        }
    }

    #[async_trait]
    impl LlmProvider for TestProvider {
        fn name(&self) -> &'static str {
            self.provider_name
        }
        fn display_name(&self) -> &str {
            "Test Provider"
        }
        fn capabilities(&self) -> LlmCapabilities {
            LlmCapabilities::text_only()
        }
        fn default_model(&self) -> &'static str {
            "test-model"
        }
        fn available_models(&self) -> &[String] {
            &[]
        }

        async fn complete(&self, _request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
            self.call_count.fetch_add(1, Ordering::SeqCst);
            let mut responses = self.responses.lock().expect("test lock"); // Safe: test assertion
            if responses.is_empty() {
                Ok(ChatResponse {
                    content: "default".to_owned(),
                    model: "test-model".to_owned(),
                    usage: None,
                    finish_reason: Some("stop".to_owned()),
                    warnings: None,
                    tool_calls: None,
                })
            } else {
                responses.remove(0)
            }
        }

        async fn complete_stream(&self, _request: &ChatRequest) -> Result<ChatStream, RunnerError> {
            Err(RunnerError::internal("streaming not supported in test"))
        }

        async fn health_check(&self) -> Result<bool, RunnerError> {
            Ok(true)
        }
    }

    #[test]
    fn fresh_report_is_zeroed() {
        let provider = TestProvider::new(vec![]);
        let metered = MetricsProvider::new(Box::new(provider));
        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.call_count, 0);
        assert_eq!(report.total_latency_ms, 0);
        assert_eq!(report.avg_latency_ms, 0);
        assert_eq!(report.total_prompt_tokens, 0);
        assert_eq!(report.total_completion_tokens, 0);
        assert_eq!(report.total_tokens, 0);
        assert_eq!(report.errors_count, 0);
        assert!(
            report.total_cost.abs() < f64::EPSILON,
            "{}",
            report.total_cost
        );
        assert_eq!(report.provider_name, "test");
    }

    #[tokio::test]
    async fn call_count_increments() {
        let provider = TestProvider::new(vec![
            Ok(ChatResponse {
                content: "hello world".to_owned(),
                model: "test-model".to_owned(),
                usage: Some(TokenUsage::new(10, 5, 15)),
                finish_reason: Some("stop".to_owned()),
                warnings: None,
                tool_calls: None,
            }),
            Ok(ChatResponse {
                content: "second".to_owned(),
                model: "test-model".to_owned(),
                usage: Some(TokenUsage::new(8, 3, 11)),
                finish_reason: Some("stop".to_owned()),
                warnings: None,
                tool_calls: None,
            }),
        ]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);

        metered.complete(&request).await.expect("first call"); // Safe: test assertion
        metered.complete(&request).await.expect("second call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.call_count, 2);
        assert_eq!(report.total_prompt_tokens, 18);
        assert_eq!(report.total_completion_tokens, 8);
        assert_eq!(report.total_tokens, 26);
        assert_eq!(report.errors_count, 0);
    }

    #[tokio::test]
    async fn errors_count_on_failure() {
        let provider = TestProvider::new(vec![Err(RunnerError::external_service("test", "boom"))]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);

        let result = metered.complete(&request).await;
        assert!(result.is_err());

        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.call_count, 1);
        assert_eq!(report.errors_count, 1);
    }

    #[tokio::test]
    async fn token_estimation_when_no_usage() {
        let provider = TestProvider::new(vec![Ok(ChatResponse {
            content: "abcdefghijklmnop".to_owned(), // 16 chars => 4 tokens
            model: "test-model".to_owned(),
            usage: None,
            finish_reason: Some("stop".to_owned()),
            warnings: None,
            tool_calls: None,
        })]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("12345678")]); // 8 chars => 2 tokens

        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.total_prompt_tokens, 2);
        assert_eq!(report.total_completion_tokens, 4);
        assert_eq!(report.total_tokens, 6);
    }

    #[test]
    fn div_by_zero_guard_on_avg_latency() {
        let provider = TestProvider::new(vec![]);
        let metered = MetricsProvider::new(Box::new(provider));
        // No calls made — avg should be 0/max(0,1) = 0, not panic
        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.avg_latency_ms, 0);
    }

    #[tokio::test]
    async fn reset_zeroes_counters() {
        let provider = TestProvider::new(vec![Ok(ChatResponse {
            content: "hello".to_owned(),
            model: "test-model".to_owned(),
            usage: Some(TokenUsage::new(5, 2, 7)),
            finish_reason: Some("stop".to_owned()),
            warnings: None,
            tool_calls: None,
        })]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);

        metered.complete(&request).await.expect("call"); // Safe: test assertion
        assert_eq!(metered.report().unwrap().call_count, 1); // Safe: test assertion

        metered.reset().unwrap(); // Safe: test assertion
        let report = metered.report().unwrap(); // Safe: test assertion
        assert_eq!(report.call_count, 0);
        assert_eq!(report.total_tokens, 0);
        assert_eq!(report.errors_count, 0);
        assert!(
            report.total_cost.abs() < f64::EPSILON,
            "{}",
            report.total_cost
        );
    }

    // ========================================================================
    // Cost tracking tests — priced through `crate::pricing`, keyed on name()
    // ========================================================================

    /// gemini-2.5-flash: $0.15/M input, $0.60/M output.
    const FLASH_IN: f64 = 0.15 / 1_000_000.0;
    const FLASH_OUT: f64 = 0.60 / 1_000_000.0;

    fn flash_response(content: &str, usage: Option<TokenUsage>) -> ChatResponse {
        ChatResponse {
            content: content.to_owned(),
            model: "gemini-2.5-flash".to_owned(),
            usage,
            finish_reason: Some("stop".to_owned()),
            warnings: None,
            tool_calls: None,
        }
    }

    #[tokio::test]
    async fn cost_with_known_model() {
        let provider = TestProvider::named(
            "gemini",
            vec![Ok(flash_response(
                "response",
                Some(TokenUsage::new(1000, 500, 1500)),
            ))],
        );
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        let expected = 1000.0f64.mul_add(FLASH_IN, 500.0 * FLASH_OUT);
        assert!(
            (report.total_cost - expected).abs() < 1e-12,
            "{}",
            report.total_cost
        );
    }

    /// The table is keyed on the provider's `name()`: the same model under a
    /// provider the table does not price costs nothing.
    #[tokio::test]
    async fn cost_with_unknown_provider_is_zero() {
        let provider = TestProvider::new(vec![Ok(flash_response(
            "response",
            Some(TokenUsage::new(1000, 500, 1500)),
        ))]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        assert!(
            report.total_cost.abs() < f64::EPSILON,
            "{}",
            report.total_cost
        );
    }

    #[tokio::test]
    async fn cost_with_unknown_model() {
        let provider = TestProvider::named(
            "gemini",
            vec![Ok(ChatResponse {
                content: "response".to_owned(),
                model: "some-unknown-model".to_owned(),
                usage: Some(TokenUsage::new(1000, 500, 1500)),
                finish_reason: Some("stop".to_owned()),
                warnings: None,
                tool_calls: None,
            })],
        );
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        assert!(
            report.total_cost.abs() < f64::EPSILON,
            "{}",
            report.total_cost
        );
    }

    #[tokio::test]
    async fn cost_accumulates() {
        let provider = TestProvider::named(
            "gemini",
            vec![
                Ok(flash_response("r1", Some(TokenUsage::new(1000, 500, 1500)))),
                Ok(flash_response(
                    "r2",
                    Some(TokenUsage::new(2000, 1000, 3000)),
                )),
            ],
        );
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call 1"); // Safe: test assertion
        metered.complete(&request).await.expect("call 2"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        let expected = 3000.0f64.mul_add(FLASH_IN, 1500.0 * FLASH_OUT);
        assert!(
            (report.total_cost - expected).abs() < 1e-12,
            "{}",
            report.total_cost
        );
    }

    /// Cache reads the provider reports bill at the model's read multiplier
    /// (Gemini 0.25×), not as fresh input.
    #[tokio::test]
    async fn cost_credits_reported_cache_reads() {
        let provider = TestProvider::named(
            "gemini",
            vec![Ok(flash_response(
                "r",
                Some(TokenUsage::new(1000, 0, 1000).with_cache(Some(1000), None)),
            ))],
        );
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        let expected = 1000.0 * FLASH_IN * 0.25;
        assert!(
            (report.total_cost - expected).abs() < 1e-12,
            "{}",
            report.total_cost
        );
    }

    #[tokio::test]
    async fn cost_with_estimated_tokens() {
        // 16 chars => 4 completion tokens; the request's 8 chars => 2 prompt tokens
        let provider =
            TestProvider::named("gemini", vec![Ok(flash_response("abcdefghijklmnop", None))]);
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("12345678")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion

        let report = metered.report().unwrap(); // Safe: test assertion
        let expected = 2.0f64.mul_add(FLASH_IN, 4.0 * FLASH_OUT);
        assert!(report.total_cost > 0.0);
        assert!(
            (report.total_cost - expected).abs() < 1e-12,
            "{}",
            report.total_cost
        );
    }

    #[tokio::test]
    async fn reset_zeroes_cost() {
        let provider = TestProvider::named(
            "gemini",
            vec![Ok(flash_response(
                "response",
                Some(TokenUsage::new(1000, 500, 1500)),
            ))],
        );
        let metered = MetricsProvider::new(Box::new(provider));
        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
        metered.complete(&request).await.expect("call"); // Safe: test assertion
        assert!(metered.report().unwrap().total_cost > 0.0); // Safe: test assertion

        metered.reset().unwrap(); // Safe: test assertion
        let cost = metered.report().unwrap().total_cost; // Safe: test assertion
        assert!(cost.abs() < f64::EPSILON, "{cost}");
    }
}