Vapora/crates/vapora-llm-router/src/cost_metrics.rs

use prometheus::{GaugeVec, IntCounterVec, Registry};
use std::sync::Arc;

/// Prometheus metrics for cost tracking and budget enforcement.
/// Exposes budget utilization, spending, and fallback events.
pub struct CostMetrics {
    /// Remaining budget per role in cents (gauge)
    pub budget_remaining_cents: GaugeVec,
    /// Budget utilization per role (0.0-1.0) (gauge)
    pub budget_utilization: GaugeVec,
    /// Cost per provider in cents (counter)
    pub cost_per_provider_cents: IntCounterVec,
    /// Fallback triggered events with reason (counter)
    pub fallback_triggered_total: IntCounterVec,
    /// Total tokens used per provider (counter)
    pub tokens_per_provider: IntCounterVec,
}

impl CostMetrics {
    /// Create new cost metrics collection (registers with default global registry)
    pub fn new() -> Result<Arc<Self>, prometheus::Error> {
        let registry = prometheus::default_registry();
        Self::with_registry(registry)
    }

    /// Create metrics with existing registry
    pub fn with_registry(registry: &Registry) -> Result<Arc<Self>, prometheus::Error> {
        let budget_remaining_cents = GaugeVec::new(
            prometheus::Opts::new(
                "vapora_llm_budget_remaining_cents",
                "Remaining budget for agent role in cents",
            ),
            &["role"],
        )?;
        registry.register(Box::new(budget_remaining_cents.clone()))?;

        let budget_utilization = GaugeVec::new(
            prometheus::Opts::new(
                "vapora_llm_budget_utilization",
                "Budget utilization percentage for agent role (0.0-1.0)",
            ),
            &["role"],
        )?;
        registry.register(Box::new(budget_utilization.clone()))?;

        let cost_per_provider_cents = IntCounterVec::new(
            prometheus::Opts::new(
                "vapora_llm_cost_per_provider_cents",
                "Total cost per provider in cents",
            ),
            &["provider"],
        )?;
        registry.register(Box::new(cost_per_provider_cents.clone()))?;

        let fallback_triggered_total = IntCounterVec::new(
            prometheus::Opts::new(
                "vapora_llm_fallback_triggered_total",
                "Total times fallback provider was triggered",
            ),
            &["role", "reason"],
        )?;
        registry.register(Box::new(fallback_triggered_total.clone()))?;

        let tokens_per_provider = IntCounterVec::new(
            prometheus::Opts::new(
                "vapora_llm_tokens_per_provider",
                "Total tokens processed per provider",
            ),
            &["provider", "token_type"],
        )?;
        registry.register(Box::new(tokens_per_provider.clone()))?;

        Ok(Arc::new(Self {
            budget_remaining_cents,
            budget_utilization,
            cost_per_provider_cents,
            fallback_triggered_total,
            tokens_per_provider,
        }))
    }

    /// Record budget update for role
    pub fn record_budget_update(&self, role: &str, remaining_cents: u32, utilization: f64) {
        self.budget_remaining_cents
            .with_label_values(&[role])
            .set(remaining_cents as f64);
        self.budget_utilization
            .with_label_values(&[role])
            .set(utilization);
    }

    /// Record cost for provider
    pub fn record_provider_cost(&self, provider: &str, cost_cents: u32) {
        self.cost_per_provider_cents
            .with_label_values(&[provider])
            .inc_by(cost_cents as u64);
    }

    /// Record fallback provider activation
    pub fn record_fallback_triggered(&self, role: &str, reason: &str) {
        self.fallback_triggered_total
            .with_label_values(&[role, reason])
            .inc();
    }

    /// Record tokens used per provider
    pub fn record_tokens(&self, provider: &str, input_tokens: u64, output_tokens: u64) {
        self.tokens_per_provider
            .with_label_values(&[provider, "input"])
            .inc_by(input_tokens);
        self.tokens_per_provider
            .with_label_values(&[provider, "output"])
            .inc_by(output_tokens);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_metrics() -> Arc<CostMetrics> {
        let registry = Registry::new();
        CostMetrics::with_registry(&registry).expect("Failed to create test metrics")
    }

    #[test]
    fn test_cost_metrics_creation() {
        let registry = Registry::new();
        let metrics = CostMetrics::with_registry(&registry);
        assert!(metrics.is_ok());
    }

    #[test]
    fn test_record_budget_update() {
        let metrics = create_test_metrics();
        metrics.record_budget_update("developer", 25000, 0.167);
        // Metric recorded (would verify via Prometheus gather in integration test)
    }

    #[test]
    fn test_record_provider_cost() {
        let metrics = create_test_metrics();
        metrics.record_provider_cost("claude", 500);
        metrics.record_provider_cost("claude", 300);
        // Counter incremented by 800 total
    }

    #[test]
    fn test_record_fallback_triggered() {
        let metrics = create_test_metrics();
        metrics.record_fallback_triggered("developer", "budget_exceeded");
        metrics.record_fallback_triggered("architect", "budget_exceeded");
        metrics.record_fallback_triggered("developer", "budget_near_threshold");
        // Multiple fallback events recorded
    }

    #[test]
    fn test_record_tokens() {
        let metrics = create_test_metrics();
        metrics.record_tokens("claude", 5000, 1000);
        metrics.record_tokens("gpt4", 3000, 500);
        // Token counts recorded per provider
    }
}