kcode-intelligence-router 0.2.2

use std::{
    collections::BTreeMap,
    fs::{self, OpenOptions},
    io::Write,
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use chrono::{DateTime, NaiveDate, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::{CostEstimate, Error, Result, estimate_cost};

const RECEIPT_VERSION: u32 = 3;

/// The four token categories Kennedy accounts independently.
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TokenUsage {
    pub input_tokens: u64,
    pub cached_input_tokens: u64,
    pub thinking_tokens: u64,
    pub output_tokens: u64,
}

impl TokenUsage {
    fn add_assign(&mut self, other: Self) {
        self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
        self.cached_input_tokens = self
            .cached_input_tokens
            .saturating_add(other.cached_input_tokens);
        self.thinking_tokens = self.thinking_tokens.saturating_add(other.thinking_tokens);
        self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
    }
}

/// Provider metering returned by one call.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Metering {
    Tokens(TokenUsage),
    DurationSeconds { seconds: f64 },
    Unavailable,
}

/// Durable evidence for one completed provider call.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageReceipt {
    pub version: u32,
    pub id: Uuid,
    pub recorded_at: DateTime<Utc>,
    pub user_id: String,
    /// Stable identity assigned by the application to this provider call.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub operation_id: Option<Uuid>,
    /// Causal parent operation, when this call is descendant inference.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub parent_operation_id: Option<Uuid>,
    pub operation: String,
    pub requested_model: String,
    pub actual_model: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub provider_request_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub provider_thread_id: Option<String>,
    pub metering: Metering,
    /// Locally calculated standard-tier provider cost, when the model and
    /// provider usage can be priced.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost: Option<CostEstimate>,
}

impl UsageReceipt {
    pub(crate) fn new(
        user_id: impl Into<String>,
        operation_id: Uuid,
        parent_operation_id: Option<Uuid>,
        operation: impl Into<String>,
        requested_model: impl Into<String>,
        actual_model: impl Into<String>,
        metering: Metering,
    ) -> Self {
        let actual_model = actual_model.into();
        let cost = estimate_cost(&actual_model, &metering);
        Self {
            version: RECEIPT_VERSION,
            id: Uuid::new_v4(),
            recorded_at: Utc::now(),
            user_id: user_id.into(),
            operation_id: Some(operation_id),
            parent_operation_id,
            operation: operation.into(),
            requested_model: requested_model.into(),
            actual_model,
            provider_request_id: None,
            provider_thread_id: None,
            metering,
            cost,
        }
    }
}

/// One grouping key in a linearly calculated daily report.
#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DailyUsageKey {
    pub user_id: String,
    pub model: String,
}

/// Totals calculated by opening every receipt for the requested UTC day.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DailyUsage {
    pub calls: u64,
    pub token_calls: u64,
    pub duration_calls: u64,
    pub unmetered_calls: u64,
    pub tokens: TokenUsage,
    pub duration_seconds: f64,
    pub estimated_cost_usd_nanos: u64,
    pub unpriced_calls: u64,
}

/// Intentionally simple one-file-per-call receipt storage.
#[derive(Clone, Debug)]
pub(crate) struct ReceiptStore {
    root: Arc<PathBuf>,
    writer: Arc<Mutex<()>>,
}

impl ReceiptStore {
    pub(crate) fn open(root: impl Into<PathBuf>) -> Result<Self> {
        let root = root.into();
        fs::create_dir_all(&root).map_err(|error| receipt_io("create", &root, error))?;
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            fs::set_permissions(&root, fs::Permissions::from_mode(0o700))
                .map_err(|error| receipt_io("protect", &root, error))?;
        }
        Ok(Self {
            root: Arc::new(root),
            writer: Arc::new(Mutex::new(())),
        })
    }

    /// Atomically writes exactly one new JSON receipt.
    pub(crate) fn record(&self, receipt: &UsageReceipt) -> Result<()> {
        let _guard = self.writer.lock().map_err(|_| {
            Error::internal(
                "receipt_store_unavailable",
                "The usage receipt writer is unavailable.",
            )
        })?;
        let bytes = serde_json::to_vec_pretty(receipt).map_err(|error| {
            Error::internal(
                "receipt_serialization_failed",
                format!("Could not serialize usage receipt: {error}"),
            )
        })?;
        let final_path = self.root.join(format!("{}.json", receipt.id));
        let temporary_path = self.root.join(format!(".{}.tmp", receipt.id));
        let result = (|| -> std::io::Result<()> {
            let mut options = OpenOptions::new();
            options.write(true).create_new(true);
            #[cfg(unix)]
            {
                use std::os::unix::fs::OpenOptionsExt;
                options.mode(0o600);
            }
            let mut file = options.open(&temporary_path)?;
            file.write_all(&bytes)?;
            file.sync_all()?;
            fs::rename(&temporary_path, &final_path)?;
            sync_directory(&self.root)?;
            Ok(())
        })();
        if result.is_err() {
            let _ = fs::remove_file(&temporary_path);
        }
        result.map_err(|error| receipt_io("write", &final_path, error))
    }

    /// Opens and decodes every receipt file. There are deliberately no indexes yet.
    pub(crate) fn receipts(&self) -> Result<Vec<UsageReceipt>> {
        let mut receipts = Vec::new();
        let entries =
            fs::read_dir(&*self.root).map_err(|error| receipt_io("list", &self.root, error))?;
        for entry in entries {
            let entry = entry.map_err(|error| receipt_io("list", &self.root, error))?;
            let path = entry.path();
            if path.extension().and_then(|value| value.to_str()) != Some("json") {
                continue;
            }
            let bytes = fs::read(&path).map_err(|error| receipt_io("read", &path, error))?;
            let receipt = serde_json::from_slice::<UsageReceipt>(&bytes).map_err(|error| {
                Error::internal(
                    "invalid_usage_receipt",
                    format!("Invalid usage receipt {}: {error}", path.display()),
                )
            })?;
            receipts.push(receipt);
        }
        receipts.sort_by_key(|receipt| (receipt.recorded_at, receipt.id));
        Ok(receipts)
    }

    /// Calculates per-user, per-model totals for one UTC day by scanning all receipts.
    pub(crate) fn daily_usage(
        &self,
        day: NaiveDate,
    ) -> Result<BTreeMap<DailyUsageKey, DailyUsage>> {
        let mut groups = BTreeMap::new();
        for receipt in self.receipts()? {
            if receipt.recorded_at.date_naive() != day {
                continue;
            }
            let totals = groups
                .entry(DailyUsageKey {
                    user_id: receipt.user_id,
                    model: receipt.actual_model,
                })
                .or_insert_with(DailyUsage::default);
            totals.calls = totals.calls.saturating_add(1);
            if let Some(cost) = receipt.cost {
                totals.estimated_cost_usd_nanos = totals
                    .estimated_cost_usd_nanos
                    .saturating_add(cost.usd_nanos);
            } else {
                totals.unpriced_calls = totals.unpriced_calls.saturating_add(1);
            }
            match receipt.metering {
                Metering::Tokens(tokens) => {
                    totals.token_calls = totals.token_calls.saturating_add(1);
                    totals.tokens.add_assign(tokens);
                }
                Metering::DurationSeconds { seconds } => {
                    totals.duration_calls = totals.duration_calls.saturating_add(1);
                    totals.duration_seconds += seconds;
                }
                Metering::Unavailable => {
                    totals.unmetered_calls = totals.unmetered_calls.saturating_add(1);
                }
            }
        }
        Ok(groups)
    }
}

fn receipt_io(operation: &str, path: &Path, error: std::io::Error) -> Error {
    Error::internal(
        "receipt_io_failed",
        format!(
            "Could not {operation} usage receipts at {}: {error}",
            path.display()
        ),
    )
}

#[cfg(unix)]
fn sync_directory(path: &Path) -> std::io::Result<()> {
    std::fs::File::open(path)?.sync_all()
}

#[cfg(not(unix))]
fn sync_directory(_path: &Path) -> std::io::Result<()> {
    Ok(())
}

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

    fn temporary_store() -> (PathBuf, ReceiptStore) {
        let root = std::env::temp_dir().join(format!("kennedy-usage-{}", Uuid::new_v4()));
        let store = ReceiptStore::open(&root).unwrap();
        (root, store)
    }

    #[test]
    fn one_file_is_written_for_each_call_and_daily_totals_scan_them() {
        let (root, store) = temporary_store();
        let mut first = UsageReceipt::new(
            "user-a",
            Uuid::new_v4(),
            None,
            "search",
            "gemini-2.5-flash",
            "gemini-2.5-flash",
            Metering::Tokens(TokenUsage {
                input_tokens: 10,
                cached_input_tokens: 4,
                thinking_tokens: 3,
                output_tokens: 2,
            }),
        );
        first.recorded_at = DateTime::parse_from_rfc3339("2026-07-26T01:00:00Z")
            .unwrap()
            .with_timezone(&Utc);
        assert!(first.operation_id.is_some());
        let mut second = first.clone();
        second.id = Uuid::new_v4();
        second.user_id = "user-b".into();
        store.record(&first).unwrap();
        store.record(&second).unwrap();

        assert_eq!(
            fs::read_dir(&root).unwrap().flatten().count(),
            2,
            "there should be exactly one durable file per call"
        );
        let totals = store
            .daily_usage(NaiveDate::from_ymd_opt(2026, 7, 26).unwrap())
            .unwrap();
        assert_eq!(totals.len(), 2);
        assert_eq!(
            totals[&DailyUsageKey {
                user_id: "user-a".into(),
                model: "gemini-2.5-flash".into(),
            }]
                .tokens
                .cached_input_tokens,
            4
        );
        fs::remove_dir_all(root).unwrap();
    }
}