provider-core 0.1.1

// SPDX-License-Identifier: AGPL-3.0-or-later

//! Per-credential usage accounting (PRD-001s AC-08).
//!
//! A provider call is billed at most once per (request id, attempt id). The
//! router may retry an attempt (budget-gated 429 recovery), and a dispatch
//! that fails ambiguously may be re-driven; without an idempotency boundary a
//! single logical provider interaction could be metered twice. This ledger
//! keys usage records by `(request_id, attempt_id)` so that the billing
//! authority sees exactly one metered record per attempt, regardless of how
//! many times the daemon re-presents the attempt result to the ledger.
//!
//! The ledger is append-only per key. `insert` returns whether the key was new
//! (a duplicate is a no-op, so the caller can skip a second deduction). Rows
//! carry a monotonic timestamp and expire after a configurable TTL so a
//! long-lived daemon does not accumulate unbounded state; eviction is lazy on
//! access and explicit via `sweep_expired`.

use router_protocol::{AttemptId, RequestId};
use std::collections::{BTreeMap, HashMap};
use std::hash::{Hash, Hasher};
use std::time::{Duration, SystemTime};

pub const DEFAULT_USAGE_LEDGER_TTL: Duration = Duration::from_secs(24 * 60 * 60);

/// Default cap on distinct usage records retained by a ledger. When the cap is
/// exceeded the oldest recorded attempts (by recorded time) are evicted.
pub const DEFAULT_USAGE_LEDGER_MAX_RECORDS: usize = 100_000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UsageRecord {
    pub request_id: RequestId,
    pub attempt_id: AttemptId,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
    pub recorded_at_unix_seconds: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InsertOutcome {
    /// True when this (request, attempt) pair was not previously recorded and
    /// the record was stored. False means a record already existed and this
    /// insert was a no-op (caller must not bill a second time).
    pub inserted: bool,
    /// True when an existing record was found and its usage exactly matched
    /// the incoming usage (idempotent re-presentation of the same attempt).
    pub deduplicated: bool,
    /// True when an existing record was found but with different token counts
    /// than the incoming usage. The caller can distinguish this from a plain
    /// dedup and treat the conflicting re-presentation as an anomaly.
    pub conflict: bool,
}

impl InsertOutcome {
    /// Convenience for tests and the adapter wiring: a record is only billable
    /// when it was actually inserted.
    pub fn is_new(&self) -> bool {
        self.inserted
    }
}

/// Stable hashable key for `(request_id, attempt_id)`.
#[derive(Clone, Debug, PartialEq, Eq)]
struct UsageKey {
    request_id: RequestId,
    attempt_id: AttemptId,
}

impl Hash for UsageKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.request_id.as_str().hash(state);
        self.attempt_id.as_str().hash(state);
    }
}

#[derive(Debug)]
pub struct UsageLedger {
    records: HashMap<UsageKey, UsageRecord>,
    /// Chronological index of `recorded_at_unix_seconds` to keys for lazy
    /// sweep. Buckets of equal timestamps are collected under one key.
    by_time: BTreeMap<u64, Vec<UsageKey>>,
    ttl: Duration,
    clock: fn() -> u64,
    /// Monotonic floor: the largest timestamp ever emitted by this ledger. A
    /// backwards wall-clock jump never produces decreasing stored timestamps.
    last_emitted_seconds: u64,
    /// Maximum distinct records retained. When exceeded, the oldest recorded
    /// attempts (by `by_time` order) are evicted deterministically.
    max_records: usize,
}

impl Default for UsageLedger {
    fn default() -> Self {
        Self::with_ttl(DEFAULT_USAGE_LEDGER_TTL)
    }
}

impl UsageLedger {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_ttl(ttl: Duration) -> Self {
        Self::with_limits(ttl, unix_now, DEFAULT_USAGE_LEDGER_MAX_RECORDS)
    }

    #[cfg(test)]
    fn with_clock(ttl: Duration, clock: fn() -> u64) -> Self {
        Self::with_limits(ttl, clock, DEFAULT_USAGE_LEDGER_MAX_RECORDS)
    }

    #[cfg(test)]
    fn with_clock_and_cap(ttl: Duration, clock: fn() -> u64, max_records: usize) -> Self {
        Self::with_limits(ttl, clock, max_records)
    }

    fn with_limits(ttl: Duration, clock: fn() -> u64, max_records: usize) -> Self {
        Self {
            records: HashMap::new(),
            by_time: BTreeMap::new(),
            ttl,
            clock,
            last_emitted_seconds: 0,
            max_records,
        }
    }

    /// Returns the current wall clock floored to the largest timestamp this
    /// ledger has ever emitted, so stored timestamps never decrease.
    fn now(&mut self) -> u64 {
        let now = (self.clock)().max(self.last_emitted_seconds);
        self.last_emitted_seconds = now;
        now
    }

    /// Record a completed attempt's usage. Returns `InsertOutcome::is_new()`
    /// false when the same (request, attempt) was already recorded, which the
    /// caller must treat as a no-op billing event.
    pub fn record(
        &mut self,
        request_id: RequestId,
        attempt_id: AttemptId,
        prompt_tokens: u64,
        completion_tokens: u64,
        total_tokens: u64,
    ) -> InsertOutcome {
        self.sweep_expired();
        let key = UsageKey {
            request_id: request_id.clone(),
            attempt_id: attempt_id.clone(),
        };
        if let Some(existing) = self.records.get(&key) {
            let same_usage = existing.prompt_tokens == prompt_tokens
                && existing.completion_tokens == completion_tokens
                && existing.total_tokens == total_tokens;
            return InsertOutcome {
                inserted: false,
                deduplicated: same_usage,
                conflict: !same_usage,
            };
        }
        let recorded_at_unix_seconds = self.now();
        self.by_time
            .entry(recorded_at_unix_seconds)
            .or_default()
            .push(key.clone());
        self.records.insert(
            key,
            UsageRecord {
                request_id,
                attempt_id,
                prompt_tokens,
                completion_tokens,
                total_tokens,
                recorded_at_unix_seconds,
            },
        );
        self.enforce_record_cap();
        InsertOutcome {
            inserted: true,
            deduplicated: false,
            conflict: false,
        }
    }

    /// Enforces the record-count cap by evicting the oldest recorded attempts
    /// in `by_time` order (deterministic: smallest timestamp first, then
    /// insertion order within a timestamp bucket).
    fn enforce_record_cap(&mut self) {
        while self.records.len() > self.max_records {
            let Some((&bucket, _)) = self.by_time.iter().next() else {
                break;
            };
            let key = {
                let keys = self.by_time.get_mut(&bucket).expect("bucket exists");
                let key = keys.remove(0);
                if keys.is_empty() {
                    self.by_time.remove(&bucket);
                }
                key
            };
            self.records.remove(&key);
        }
    }

    /// The number of distinct attempts currently recorded.
    pub fn len(&self) -> usize {
        self.records.len()
    }

    pub fn is_empty(&self) -> bool {
        self.records.is_empty()
    }

    /// True when a usage record exists for this (request, attempt) pair.
    pub fn contains(&self, request_id: &RequestId, attempt_id: &AttemptId) -> bool {
        self.records.contains_key(&UsageKey {
            request_id: request_id.clone(),
            attempt_id: attempt_id.clone(),
        })
    }

    /// Total tokens recorded across all distinct attempts.
    pub fn total_tokens(&self) -> u64 {
        self.records
            .values()
            .map(|record| record.total_tokens)
            .sum()
    }

    /// Drops records older than the TTL. Called lazily on every `record` and
    /// available explicitly for maintenance. The clock read is floored to the
    /// ledger's monotonic timestamp so a backwards jump cannot expire freshly
    /// recorded rows.
    pub fn sweep_expired(&mut self) {
        let now = self.now();
        let ttl_seconds = self.ttl.as_secs();
        let cutoff = now.saturating_sub(ttl_seconds);
        let expired: Vec<u64> = self
            .by_time
            .range(..=cutoff)
            .map(|(bucket, _)| *bucket)
            .collect();
        for bucket in expired {
            if let Some(keys) = self.by_time.remove(&bucket) {
                for key in keys {
                    self.records.remove(&key);
                }
            }
        }
    }
}

fn unix_now() -> u64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

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

    fn request_id(value: &str) -> RequestId {
        RequestId::parse(value).expect("fixture request id")
    }

    fn attempt_id(value: &str) -> AttemptId {
        AttemptId::parse(value).expect("fixture attempt id")
    }

    #[test]
    fn same_attempt_is_recorded_at_most_once() {
        let mut ledger = UsageLedger::new();
        let request = request_id("request-dedupe");
        let attempt = attempt_id("attempt-dedupe");
        let first = ledger.record(request.clone(), attempt.clone(), 10, 20, 30);
        assert!(first.is_new());
        assert!(!first.deduplicated);
        assert!(!first.conflict);
        let second = ledger.record(request.clone(), attempt.clone(), 10, 20, 30);
        assert!(!second.is_new());
        assert!(second.deduplicated);
        assert!(!second.conflict);
        let third = ledger.record(request, attempt, 99, 99, 99);
        assert!(!third.is_new());
        assert!(!third.deduplicated);
        assert!(third.conflict);
        assert_eq!(ledger.len(), 1);
        assert_eq!(ledger.total_tokens(), 30);
    }

    #[test]
    fn backwards_clock_jump_keeps_timestamps_monotonic() {
        use std::sync::atomic::{AtomicU64, Ordering};
        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
        fn clock() -> u64 {
            TICK.load(Ordering::Relaxed)
        }
        TICK.store(1_000_000, Ordering::Relaxed);
        let mut ledger = UsageLedger::with_clock(Duration::from_secs(1000), clock);
        let request = request_id("request-mono");
        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
        // Jump the wall clock backwards; stored timestamps must not decrease.
        TICK.store(999_500, Ordering::Relaxed);
        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
        TICK.store(999_000, Ordering::Relaxed);
        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
        assert_eq!(ledger.len(), 3);
        let stamps: Vec<u64> = ledger
            .records
            .values()
            .map(|record| record.recorded_at_unix_seconds)
            .collect();
        assert!(
            stamps.iter().all(|&stamp| stamp == 1_000_000),
            "all records pinned to the monotonic floor, got {stamps:?}"
        );
        // A sweep from a clock earlier than the floor must not drop the
        // just-recorded rows.
        ledger.sweep_expired();
        assert_eq!(ledger.len(), 3);
    }

    #[test]
    fn record_cap_evicts_oldest_deterministically() {
        use std::sync::atomic::{AtomicU64, Ordering};
        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
        fn clock() -> u64 {
            TICK.load(Ordering::Relaxed)
        }
        TICK.store(1_000_000, Ordering::Relaxed);
        let mut ledger = UsageLedger::with_clock_and_cap(Duration::from_secs(1000), clock, 2);
        let request = request_id("request-cap");
        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
        TICK.store(1_000_001, Ordering::Relaxed);
        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
        assert_eq!(ledger.len(), 2);
        TICK.store(1_000_002, Ordering::Relaxed);
        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
        // The oldest record (attempt-1) is evicted, the two newest survive.
        assert_eq!(ledger.len(), 2);
        assert!(!ledger.contains(&request_id("request-cap"), &attempt_id("attempt-1")));
        assert!(ledger.contains(&request_id("request-cap"), &attempt_id("attempt-2")));
        assert!(ledger.contains(&request_id("request-cap"), &attempt_id("attempt-3")));
        assert_eq!(ledger.total_tokens(), 10);
    }

    #[test]
    fn distinct_attempts_are_separate_billing_events() {
        let mut ledger = UsageLedger::new();
        let request = request_id("request-multi");
        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
        assert_eq!(ledger.len(), 3);
        assert_eq!(ledger.total_tokens(), 12);
    }

    #[test]
    fn expired_records_are_swept() {
        use std::sync::atomic::{AtomicU64, Ordering};
        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
        fn clock() -> u64 {
            TICK.load(Ordering::Relaxed)
        }
        TICK.store(1_000_000, Ordering::Relaxed);
        let mut ledger = UsageLedger::with_clock(Duration::from_secs(100), clock);
        let request = request_id("request-ttl");
        ledger.record(request.clone(), attempt_id("attempt-old"), 1, 1, 2);
        assert_eq!(ledger.len(), 1);
        TICK.store(1_000_200, Ordering::Relaxed);
        ledger.record(request.clone(), attempt_id("attempt-new"), 1, 1, 2);
        assert_eq!(ledger.len(), 1);
        assert!(ledger.contains(&request, &attempt_id("attempt-new")));
        assert!(!ledger.contains(&request, &attempt_id("attempt-old")));
        ledger.sweep_expired();
        assert_eq!(ledger.len(), 1);
    }
}