provider_core/
usage_ledger.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3//! Per-credential usage accounting (PRD-001s AC-08).
4//!
5//! A provider call is billed at most once per (request id, attempt id). The
6//! router may retry an attempt (budget-gated 429 recovery), and a dispatch
7//! that fails ambiguously may be re-driven; without an idempotency boundary a
8//! single logical provider interaction could be metered twice. This ledger
9//! keys usage records by `(request_id, attempt_id)` so that the billing
10//! authority sees exactly one metered record per attempt, regardless of how
11//! many times the daemon re-presents the attempt result to the ledger.
12//!
13//! The ledger is append-only per key. `insert` returns whether the key was new
14//! (a duplicate is a no-op, so the caller can skip a second deduction). Rows
15//! carry a monotonic timestamp and expire after a configurable TTL so a
16//! long-lived daemon does not accumulate unbounded state; eviction is lazy on
17//! access and explicit via `sweep_expired`.
18
19use router_protocol::{AttemptId, RequestId};
20use std::collections::{BTreeMap, HashMap};
21use std::hash::{Hash, Hasher};
22use std::time::{Duration, SystemTime};
23
24pub const DEFAULT_USAGE_LEDGER_TTL: Duration = Duration::from_secs(24 * 60 * 60);
25
26/// Default cap on distinct usage records retained by a ledger. When the cap is
27/// exceeded the oldest recorded attempts (by recorded time) are evicted.
28pub const DEFAULT_USAGE_LEDGER_MAX_RECORDS: usize = 100_000;
29
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub struct UsageRecord {
32    pub request_id: RequestId,
33    pub attempt_id: AttemptId,
34    pub prompt_tokens: u64,
35    pub completion_tokens: u64,
36    pub total_tokens: u64,
37    pub recorded_at_unix_seconds: u64,
38}
39
40#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub struct InsertOutcome {
42    /// True when this (request, attempt) pair was not previously recorded and
43    /// the record was stored. False means a record already existed and this
44    /// insert was a no-op (caller must not bill a second time).
45    pub inserted: bool,
46    /// True when an existing record was found and its usage exactly matched
47    /// the incoming usage (idempotent re-presentation of the same attempt).
48    pub deduplicated: bool,
49    /// True when an existing record was found but with different token counts
50    /// than the incoming usage. The caller can distinguish this from a plain
51    /// dedup and treat the conflicting re-presentation as an anomaly.
52    pub conflict: bool,
53}
54
55impl InsertOutcome {
56    /// Convenience for tests and the adapter wiring: a record is only billable
57    /// when it was actually inserted.
58    pub fn is_new(&self) -> bool {
59        self.inserted
60    }
61}
62
63/// Stable hashable key for `(request_id, attempt_id)`.
64#[derive(Clone, Debug, PartialEq, Eq)]
65struct UsageKey {
66    request_id: RequestId,
67    attempt_id: AttemptId,
68}
69
70impl Hash for UsageKey {
71    fn hash<H: Hasher>(&self, state: &mut H) {
72        self.request_id.as_str().hash(state);
73        self.attempt_id.as_str().hash(state);
74    }
75}
76
77#[derive(Debug)]
78pub struct UsageLedger {
79    records: HashMap<UsageKey, UsageRecord>,
80    /// Chronological index of `recorded_at_unix_seconds` to keys for lazy
81    /// sweep. Buckets of equal timestamps are collected under one key.
82    by_time: BTreeMap<u64, Vec<UsageKey>>,
83    ttl: Duration,
84    clock: fn() -> u64,
85    /// Monotonic floor: the largest timestamp ever emitted by this ledger. A
86    /// backwards wall-clock jump never produces decreasing stored timestamps.
87    last_emitted_seconds: u64,
88    /// Maximum distinct records retained. When exceeded, the oldest recorded
89    /// attempts (by `by_time` order) are evicted deterministically.
90    max_records: usize,
91}
92
93impl Default for UsageLedger {
94    fn default() -> Self {
95        Self::with_ttl(DEFAULT_USAGE_LEDGER_TTL)
96    }
97}
98
99impl UsageLedger {
100    pub fn new() -> Self {
101        Self::default()
102    }
103
104    pub fn with_ttl(ttl: Duration) -> Self {
105        Self::with_limits(ttl, unix_now, DEFAULT_USAGE_LEDGER_MAX_RECORDS)
106    }
107
108    #[cfg(test)]
109    fn with_clock(ttl: Duration, clock: fn() -> u64) -> Self {
110        Self::with_limits(ttl, clock, DEFAULT_USAGE_LEDGER_MAX_RECORDS)
111    }
112
113    #[cfg(test)]
114    fn with_clock_and_cap(ttl: Duration, clock: fn() -> u64, max_records: usize) -> Self {
115        Self::with_limits(ttl, clock, max_records)
116    }
117
118    fn with_limits(ttl: Duration, clock: fn() -> u64, max_records: usize) -> Self {
119        Self {
120            records: HashMap::new(),
121            by_time: BTreeMap::new(),
122            ttl,
123            clock,
124            last_emitted_seconds: 0,
125            max_records,
126        }
127    }
128
129    /// Returns the current wall clock floored to the largest timestamp this
130    /// ledger has ever emitted, so stored timestamps never decrease.
131    fn now(&mut self) -> u64 {
132        let now = (self.clock)().max(self.last_emitted_seconds);
133        self.last_emitted_seconds = now;
134        now
135    }
136
137    /// Record a completed attempt's usage. Returns `InsertOutcome::is_new()`
138    /// false when the same (request, attempt) was already recorded, which the
139    /// caller must treat as a no-op billing event.
140    pub fn record(
141        &mut self,
142        request_id: RequestId,
143        attempt_id: AttemptId,
144        prompt_tokens: u64,
145        completion_tokens: u64,
146        total_tokens: u64,
147    ) -> InsertOutcome {
148        self.sweep_expired();
149        let key = UsageKey {
150            request_id: request_id.clone(),
151            attempt_id: attempt_id.clone(),
152        };
153        if let Some(existing) = self.records.get(&key) {
154            let same_usage = existing.prompt_tokens == prompt_tokens
155                && existing.completion_tokens == completion_tokens
156                && existing.total_tokens == total_tokens;
157            return InsertOutcome {
158                inserted: false,
159                deduplicated: same_usage,
160                conflict: !same_usage,
161            };
162        }
163        let recorded_at_unix_seconds = self.now();
164        self.by_time
165            .entry(recorded_at_unix_seconds)
166            .or_default()
167            .push(key.clone());
168        self.records.insert(
169            key,
170            UsageRecord {
171                request_id,
172                attempt_id,
173                prompt_tokens,
174                completion_tokens,
175                total_tokens,
176                recorded_at_unix_seconds,
177            },
178        );
179        self.enforce_record_cap();
180        InsertOutcome {
181            inserted: true,
182            deduplicated: false,
183            conflict: false,
184        }
185    }
186
187    /// Enforces the record-count cap by evicting the oldest recorded attempts
188    /// in `by_time` order (deterministic: smallest timestamp first, then
189    /// insertion order within a timestamp bucket).
190    fn enforce_record_cap(&mut self) {
191        while self.records.len() > self.max_records {
192            let Some((&bucket, _)) = self.by_time.iter().next() else {
193                break;
194            };
195            let key = {
196                let keys = self.by_time.get_mut(&bucket).expect("bucket exists");
197                let key = keys.remove(0);
198                if keys.is_empty() {
199                    self.by_time.remove(&bucket);
200                }
201                key
202            };
203            self.records.remove(&key);
204        }
205    }
206
207    /// The number of distinct attempts currently recorded.
208    pub fn len(&self) -> usize {
209        self.records.len()
210    }
211
212    pub fn is_empty(&self) -> bool {
213        self.records.is_empty()
214    }
215
216    /// True when a usage record exists for this (request, attempt) pair.
217    pub fn contains(&self, request_id: &RequestId, attempt_id: &AttemptId) -> bool {
218        self.records.contains_key(&UsageKey {
219            request_id: request_id.clone(),
220            attempt_id: attempt_id.clone(),
221        })
222    }
223
224    /// Total tokens recorded across all distinct attempts.
225    pub fn total_tokens(&self) -> u64 {
226        self.records
227            .values()
228            .map(|record| record.total_tokens)
229            .sum()
230    }
231
232    /// Drops records older than the TTL. Called lazily on every `record` and
233    /// available explicitly for maintenance. The clock read is floored to the
234    /// ledger's monotonic timestamp so a backwards jump cannot expire freshly
235    /// recorded rows.
236    pub fn sweep_expired(&mut self) {
237        let now = self.now();
238        let ttl_seconds = self.ttl.as_secs();
239        let cutoff = now.saturating_sub(ttl_seconds);
240        let expired: Vec<u64> = self
241            .by_time
242            .range(..=cutoff)
243            .map(|(bucket, _)| *bucket)
244            .collect();
245        for bucket in expired {
246            if let Some(keys) = self.by_time.remove(&bucket) {
247                for key in keys {
248                    self.records.remove(&key);
249                }
250            }
251        }
252    }
253}
254
255fn unix_now() -> u64 {
256    SystemTime::now()
257        .duration_since(SystemTime::UNIX_EPOCH)
258        .unwrap_or_default()
259        .as_secs()
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265
266    fn request_id(value: &str) -> RequestId {
267        RequestId::parse(value).expect("fixture request id")
268    }
269
270    fn attempt_id(value: &str) -> AttemptId {
271        AttemptId::parse(value).expect("fixture attempt id")
272    }
273
274    #[test]
275    fn same_attempt_is_recorded_at_most_once() {
276        let mut ledger = UsageLedger::new();
277        let request = request_id("request-dedupe");
278        let attempt = attempt_id("attempt-dedupe");
279        let first = ledger.record(request.clone(), attempt.clone(), 10, 20, 30);
280        assert!(first.is_new());
281        assert!(!first.deduplicated);
282        assert!(!first.conflict);
283        let second = ledger.record(request.clone(), attempt.clone(), 10, 20, 30);
284        assert!(!second.is_new());
285        assert!(second.deduplicated);
286        assert!(!second.conflict);
287        let third = ledger.record(request, attempt, 99, 99, 99);
288        assert!(!third.is_new());
289        assert!(!third.deduplicated);
290        assert!(third.conflict);
291        assert_eq!(ledger.len(), 1);
292        assert_eq!(ledger.total_tokens(), 30);
293    }
294
295    #[test]
296    fn backwards_clock_jump_keeps_timestamps_monotonic() {
297        use std::sync::atomic::{AtomicU64, Ordering};
298        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
299        fn clock() -> u64 {
300            TICK.load(Ordering::Relaxed)
301        }
302        TICK.store(1_000_000, Ordering::Relaxed);
303        let mut ledger = UsageLedger::with_clock(Duration::from_secs(1000), clock);
304        let request = request_id("request-mono");
305        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
306        // Jump the wall clock backwards; stored timestamps must not decrease.
307        TICK.store(999_500, Ordering::Relaxed);
308        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
309        TICK.store(999_000, Ordering::Relaxed);
310        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
311        assert_eq!(ledger.len(), 3);
312        let stamps: Vec<u64> = ledger
313            .records
314            .values()
315            .map(|record| record.recorded_at_unix_seconds)
316            .collect();
317        assert!(
318            stamps.iter().all(|&stamp| stamp == 1_000_000),
319            "all records pinned to the monotonic floor, got {stamps:?}"
320        );
321        // A sweep from a clock earlier than the floor must not drop the
322        // just-recorded rows.
323        ledger.sweep_expired();
324        assert_eq!(ledger.len(), 3);
325    }
326
327    #[test]
328    fn record_cap_evicts_oldest_deterministically() {
329        use std::sync::atomic::{AtomicU64, Ordering};
330        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
331        fn clock() -> u64 {
332            TICK.load(Ordering::Relaxed)
333        }
334        TICK.store(1_000_000, Ordering::Relaxed);
335        let mut ledger = UsageLedger::with_clock_and_cap(Duration::from_secs(1000), clock, 2);
336        let request = request_id("request-cap");
337        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
338        TICK.store(1_000_001, Ordering::Relaxed);
339        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
340        assert_eq!(ledger.len(), 2);
341        TICK.store(1_000_002, Ordering::Relaxed);
342        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
343        // The oldest record (attempt-1) is evicted, the two newest survive.
344        assert_eq!(ledger.len(), 2);
345        assert!(!ledger.contains(&request_id("request-cap"), &attempt_id("attempt-1")));
346        assert!(ledger.contains(&request_id("request-cap"), &attempt_id("attempt-2")));
347        assert!(ledger.contains(&request_id("request-cap"), &attempt_id("attempt-3")));
348        assert_eq!(ledger.total_tokens(), 10);
349    }
350
351    #[test]
352    fn distinct_attempts_are_separate_billing_events() {
353        let mut ledger = UsageLedger::new();
354        let request = request_id("request-multi");
355        ledger.record(request.clone(), attempt_id("attempt-1"), 1, 1, 2);
356        ledger.record(request.clone(), attempt_id("attempt-2"), 2, 2, 4);
357        ledger.record(request, attempt_id("attempt-3"), 3, 3, 6);
358        assert_eq!(ledger.len(), 3);
359        assert_eq!(ledger.total_tokens(), 12);
360    }
361
362    #[test]
363    fn expired_records_are_swept() {
364        use std::sync::atomic::{AtomicU64, Ordering};
365        static TICK: AtomicU64 = AtomicU64::new(1_000_000);
366        fn clock() -> u64 {
367            TICK.load(Ordering::Relaxed)
368        }
369        TICK.store(1_000_000, Ordering::Relaxed);
370        let mut ledger = UsageLedger::with_clock(Duration::from_secs(100), clock);
371        let request = request_id("request-ttl");
372        ledger.record(request.clone(), attempt_id("attempt-old"), 1, 1, 2);
373        assert_eq!(ledger.len(), 1);
374        TICK.store(1_000_200, Ordering::Relaxed);
375        ledger.record(request.clone(), attempt_id("attempt-new"), 1, 1, 2);
376        assert_eq!(ledger.len(), 1);
377        assert!(ledger.contains(&request, &attempt_id("attempt-new")));
378        assert!(!ledger.contains(&request, &attempt_id("attempt-old")));
379        ledger.sweep_expired();
380        assert_eq!(ledger.len(), 1);
381    }
382}
provider_core/usage_ledger.rs

provider_core/
usage_ledger.rs
