codetether_agent/telemetry/tokens/
counter.rs

1//! Lock-light atomic counter for LLM token usage.
2//!
3//! Global and per-model counters live here. Per-model recording uses a
4//! `try_lock` fast path so that a contended counter never blocks the
5//! completion path — at worst, the stats are briefly stale.
6
7use chrono::Utc;
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use tokio::sync::Mutex;
11
12use super::snapshot::{GlobalTokenSnapshot, TokenUsageSnapshot};
13use super::totals::TokenTotals;
14
15/// Process-wide token counter. Cheap `record` on the hot path; richer
16/// per-model stats guarded by an async-compatible mutex and accessed
17/// best-effort via `try_lock`.
18///
19/// Prefer the [`super::super::TOKEN_USAGE`] singleton — constructing your
20/// own instance is only useful in tests.
21///
22/// # Examples
23///
24/// ```rust
25/// use codetether_agent::telemetry::AtomicTokenCounter;
26///
27/// let c = AtomicTokenCounter::new();
28/// c.record(100, 50);
29/// let (p, o, t) = c.get();
30/// assert_eq!((p, o, t), (100, 50, 150));
31/// ```
32#[derive(Debug)]
33pub struct AtomicTokenCounter {
34    pub(super) prompt_tokens: AtomicU64,
35    pub(super) completion_tokens: AtomicU64,
36    pub(super) total_tokens: AtomicU64,
37    pub(super) request_count: AtomicU64,
38    pub(super) model_usage: Mutex<HashMap<String, (u64, u64)>>,
39    /// Per-model prompt-cache stats: `(cache_read_tokens, cache_write_tokens)`.
40    /// Tracked separately so [`crate::provider::pricing`] can apply the
41    /// discounted / surcharged rates that Anthropic / Bedrock use.
42    pub(super) model_cache_usage: Mutex<HashMap<String, (u64, u64)>>,
43    /// Per-model **last turn** prompt token count. Used by the TUI to show
44    /// how close the current in-flight conversation is to the model's
45    /// context window, independent of cumulative lifetime usage.
46    pub(super) model_last_prompt_tokens: Mutex<HashMap<String, u64>>,
47}
48
49impl AtomicTokenCounter {
50    /// Construct a zeroed counter.
51    pub fn new() -> Self {
52        Self {
53            prompt_tokens: AtomicU64::new(0),
54            completion_tokens: AtomicU64::new(0),
55            total_tokens: AtomicU64::new(0),
56            request_count: AtomicU64::new(0),
57            model_usage: Mutex::new(HashMap::new()),
58            model_cache_usage: Mutex::new(HashMap::new()),
59            model_last_prompt_tokens: Mutex::new(HashMap::new()),
60        }
61    }
62
63    /// Record a single completion. Increments all four global atomics
64    /// relaxed-ordering — correctness only requires eventual visibility.
65    pub fn record(&self, prompt: u64, completion: u64) {
66        self.prompt_tokens.fetch_add(prompt, Ordering::Relaxed);
67        self.completion_tokens
68            .fetch_add(completion, Ordering::Relaxed);
69        self.total_tokens
70            .fetch_add(prompt + completion, Ordering::Relaxed);
71        self.request_count.fetch_add(1, Ordering::Relaxed);
72    }
73
74    /// Return the global `(prompt, completion, total)` triple.
75    pub fn get(&self) -> (u64, u64, u64) {
76        (
77            self.prompt_tokens.load(Ordering::Relaxed),
78            self.completion_tokens.load(Ordering::Relaxed),
79            self.total_tokens.load(Ordering::Relaxed),
80        )
81    }
82
83    /// Snapshot the global counters into a plain-data struct.
84    pub fn global_snapshot(&self) -> GlobalTokenSnapshot {
85        let (prompt, completion, total) = self.get();
86        let mut snapshot = GlobalTokenSnapshot::new(prompt, completion, total);
87        snapshot.request_count = self.request_count.load(Ordering::Relaxed);
88        snapshot
89    }
90
91    /// Snapshot every tracked model's usage. Returns an empty `Vec` if the
92    /// per-model map is currently contended.
93    pub fn model_snapshots(&self) -> Vec<TokenUsageSnapshot> {
94        let Ok(usage) = self.model_usage.try_lock() else {
95            return Vec::new();
96        };
97        usage
98            .iter()
99            .map(|(name, (input, output))| TokenUsageSnapshot {
100                name: name.clone(),
101                prompt_tokens: *input,
102                completion_tokens: *output,
103                total_tokens: input + output,
104                totals: TokenTotals::new(*input, *output),
105                timestamp: Utc::now(),
106                request_count: 0,
107            })
108            .collect()
109    }
110}
111
112impl Default for AtomicTokenCounter {
113    fn default() -> Self {
114        Self::new()
115    }
116}
117
118mod record_model;
codetether_agent/telemetry/tokens/counter.rs

codetether_agent/telemetry/tokens/
counter.rs
