lean_ctx/core/
metering.rs

1//! Local, append-only metering for the free product.
2//!
3//! Metering is deliberately observational: failures are logged and never affect
4//! an MCP tool response.  The JSONL file contains token counts only.
5
6use chrono::{DateTime, NaiveDate, Utc};
7use serde::{Deserialize, Serialize};
8use std::fs::{OpenOptions, create_dir_all};
9use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
10use std::path::{Path, PathBuf};
11
12use fs2::FileExt;
13
14const METERING_FILE: &str = "metering.jsonl";
15
16/// `tool_name` of a native shell call the hook left to the agent's own tool
17/// (#1285): observed, but not routed through lean-ctx and not tokenized.
18pub const NATIVE_PASSTHROUGH_TOOL: &str = "native_shell_passthrough";
19
20/// Block size for reading the metering file backwards.
21const TAIL_BLOCK_BYTES: usize = 64 * 1024;
22
23/// Totals for one UTC day.
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
25pub struct DayTotals {
26    pub input_tokens: u64,
27    pub savings_tokens: u64,
28    /// Calls that went through lean-ctx (every metered `ctx_*` call).
29    pub routed_calls: u64,
30    pub native_passthrough_calls: u64,
31}
32
33/// One completed tool call, recorded locally as a JSONL line.
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35pub struct MeterEntry {
36    pub timestamp: DateTime<Utc>,
37    pub tool_name: String,
38    pub input_tokens: u64,
39    pub output_tokens: u64,
40    pub savings_tokens: u64,
41    /// Project key of the work (`savings_ledger::project_scope::project_key_of`);
42    /// absent on entries written before projects were kept, and outside a project.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub project: Option<String>,
45}
46
47impl MeterEntry {
48    #[must_use]
49    pub fn new(
50        tool_name: impl Into<String>,
51        input_tokens: u64,
52        output_tokens: u64,
53        savings_tokens: u64,
54    ) -> Self {
55        Self {
56            timestamp: Utc::now(),
57            tool_name: tool_name.into(),
58            input_tokens,
59            output_tokens,
60            savings_tokens,
61            project: crate::core::savings_ledger::project_scope::current_project_key(),
62        }
63    }
64}
65
66/// Aggregate used by local value displays such as "Pro would have saved you $X".
67#[derive(Debug, Clone, Copy, Default, PartialEq)]
68pub struct MeterAggregate {
69    pub total_savings_tokens: u64,
70    pub total_calls: u64,
71    /// Weighted output/input ratio. `0.0` means no metered input exists yet.
72    pub avg_compression_ratio: f64,
73}
74
75/// Append-only local JSONL store.
76#[derive(Debug, Clone)]
77pub struct MeterStore {
78    path: PathBuf,
79}
80
81impl MeterStore {
82    #[must_use]
83    pub fn new(data_dir: impl AsRef<Path>) -> Self {
84        Self {
85            path: data_dir.as_ref().join(METERING_FILE),
86        }
87    }
88
89    pub fn from_data_dir() -> Result<Self, String> {
90        Ok(Self::new(crate::core::data_dir::lean_ctx_data_dir()?))
91    }
92
93    #[must_use]
94    pub fn path(&self) -> &Path {
95        &self.path
96    }
97
98    /// Appends one complete JSON object and newline. Call this only off the hot path.
99    pub fn append(&self, entry: &MeterEntry) -> Result<(), String> {
100        let Some(parent) = self.path.parent() else {
101            return Err("metering path has no parent".to_string());
102        };
103        create_dir_all(parent).map_err(|error| error.to_string())?;
104        let line = serde_json::to_string(entry).map_err(|error| error.to_string())?;
105        let mut file = OpenOptions::new()
106            .create(true)
107            .read(true)
108            .append(true)
109            .open(&self.path)
110            .map_err(|error| error.to_string())?;
111        file.lock_exclusive().map_err(|error| error.to_string())?;
112        let write_result = writeln!(file, "{line}").map_err(|error| error.to_string());
113        let unlock_result = FileExt::unlock(&file).map_err(|error| error.to_string());
114        write_result.and(unlock_result)
115    }
116
117    /// Schedules persistence and returns immediately; metering cannot delay a tool call.
118    pub fn append_best_effort(entry: MeterEntry) {
119        tokio::task::spawn_blocking(move || {
120            let result = Self::from_data_dir().and_then(|store| store.append(&entry));
121            if let Err(error) = result {
122                tracing::warn!(%error, "lean-ctx: failed to append local metering entry");
123            }
124        });
125    }
126
127    #[must_use]
128    pub fn aggregate(&self) -> MeterAggregate {
129        let Ok(file) = std::fs::File::open(&self.path) else {
130            return MeterAggregate::default();
131        };
132        let mut aggregate = MeterAggregate::default();
133        let mut total_input = 0_u64;
134        let mut total_output = 0_u64;
135        for line in BufReader::new(file).lines().map_while(Result::ok) {
136            let Ok(entry) = serde_json::from_str::<MeterEntry>(&line) else {
137                continue;
138            };
139            aggregate.total_calls = aggregate.total_calls.saturating_add(1);
140            aggregate.total_savings_tokens = aggregate
141                .total_savings_tokens
142                .saturating_add(entry.savings_tokens);
143            total_input = total_input.saturating_add(entry.input_tokens);
144            total_output = total_output.saturating_add(entry.output_tokens);
145        }
146        if total_input > 0 {
147            aggregate.avg_compression_ratio = total_output as f64 / total_input as f64;
148        }
149        aggregate
150    }
151
152    /// Totals for the UTC day `day`. The file is append-only and chronological and
153    /// grows to tens of MB, so it is read backwards and parsing stops at the first
154    /// earlier entry: the cost follows the day's volume, not the history.
155    #[must_use]
156    pub fn day_totals(&self, day: NaiveDate) -> DayTotals {
157        self.day_totals_in_blocks(day, TAIL_BLOCK_BYTES, |_| true)
158    }
159
160    /// [`Self::day_totals`] over the entries whose project key `keep` accepts.
161    #[must_use]
162    pub fn day_totals_where(
163        &self,
164        day: NaiveDate,
165        keep: impl Fn(Option<&str>) -> bool,
166    ) -> DayTotals {
167        self.day_totals_in_blocks(day, TAIL_BLOCK_BYTES, keep)
168    }
169
170    fn day_totals_in_blocks(
171        &self,
172        day: NaiveDate,
173        block: usize,
174        keep: impl Fn(Option<&str>) -> bool,
175    ) -> DayTotals {
176        let mut totals = DayTotals::default();
177        let Ok(mut file) = std::fs::File::open(&self.path) else {
178            return totals;
179        };
180        let Ok(mut pos) = file.seek(SeekFrom::End(0)) else {
181            return totals;
182        };
183        // Bytes of a line that started in an earlier block, carried forward.
184        let mut carry = Vec::new();
185        while pos > 0 {
186            let len = usize::try_from(pos).map_or(block, |p| p.min(block));
187            pos -= len as u64;
188            let mut buf = vec![0; len];
189            if file.seek(SeekFrom::Start(pos)).is_err() || file.read_exact(&mut buf).is_err() {
190                return totals;
191            }
192            buf.extend_from_slice(&carry);
193            let mut lines: Vec<&[u8]> = buf.split(|b| *b == b'\n').collect();
194            // The first segment may be cut off unless the block starts the file.
195            carry = if pos > 0 {
196                lines.remove(0).to_vec()
197            } else {
198                Vec::new()
199            };
200            for line in lines.into_iter().rev() {
201                let Ok(entry) = serde_json::from_slice::<MeterEntry>(line) else {
202                    continue;
203                };
204                let date = entry.timestamp.date_naive();
205                if date < day {
206                    return totals;
207                }
208                if date > day || !keep(entry.project.as_deref()) {
209                    continue;
210                }
211                totals.input_tokens = totals.input_tokens.saturating_add(entry.input_tokens);
212                totals.savings_tokens = totals.savings_tokens.saturating_add(entry.savings_tokens);
213                if entry.tool_name == NATIVE_PASSTHROUGH_TOOL {
214                    totals.native_passthrough_calls += 1;
215                } else {
216                    totals.routed_calls += 1;
217                }
218            }
219        }
220        totals
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    #[test]
229    fn aggregates_jsonl_entries() {
230        let dir = tempfile::tempdir().unwrap();
231        let store = MeterStore::new(dir.path());
232        store
233            .append(&MeterEntry::new("ctx_read", 100, 20, 80))
234            .unwrap();
235        store
236            .append(&MeterEntry::new("ctx_search", 200, 100, 100))
237            .unwrap();
238
239        let aggregate = store.aggregate();
240        assert_eq!(aggregate.total_calls, 2);
241        assert_eq!(aggregate.total_savings_tokens, 180);
242        assert!((aggregate.avg_compression_ratio - 0.4).abs() < f64::EPSILON);
243    }
244
245    #[test]
246    fn day_totals_reads_only_the_requested_day_across_block_boundaries() {
247        let dir = tempfile::tempdir().unwrap();
248        let store = MeterStore::new(dir.path());
249        let at = |ts: &str, tool: &str, input, saved| MeterEntry {
250            timestamp: ts.parse().unwrap(),
251            tool_name: tool.into(),
252            input_tokens: input,
253            output_tokens: input - saved,
254            savings_tokens: saved,
255            project: (tool == "ctx_search").then(|| "0123456789abcdef".to_string()),
256        };
257        store
258            .append(&at("2026-10-01T23:59:59Z", "ctx_read", 1_000, 900))
259            .unwrap();
260        store
261            .append(&at("2026-10-02T00:00:01Z", "ctx_read", 100, 80))
262            .unwrap();
263        std::fs::OpenOptions::new()
264            .append(true)
265            .open(store.path())
266            .unwrap()
267            .write_all(b"{not json\n")
268            .unwrap();
269        store
270            .append(&at("2026-10-02T09:00:00Z", NATIVE_PASSTHROUGH_TOOL, 0, 0))
271            .unwrap();
272        store
273            .append(&at("2026-10-02T10:00:00Z", "ctx_search", 50, 10))
274            .unwrap();
275
276        let day = NaiveDate::from_ymd_opt(2026, 10, 2).unwrap();
277        let expected = DayTotals {
278            input_tokens: 150,
279            savings_tokens: 90,
280            routed_calls: 2,
281            native_passthrough_calls: 1,
282        };
283        assert_eq!(store.day_totals(day), expected);
284        // Blocks smaller than a line: every line is split across reads.
285        assert_eq!(store.day_totals_in_blocks(day, 7, |_| true), expected);
286        // One project's share: only the entries carrying its key.
287        let project = store.day_totals_where(day, |key| key == Some("0123456789abcdef"));
288        assert_eq!(
289            project,
290            DayTotals {
291                input_tokens: 50,
292                savings_tokens: 10,
293                routed_calls: 1,
294                native_passthrough_calls: 0,
295            }
296        );
297        assert_eq!(
298            MeterStore::new(dir.path().join("missing")).day_totals(day),
299            DayTotals::default()
300        );
301    }
302
303    #[tokio::test]
304    async fn best_effort_append_does_not_block_calling_task() {
305        let entry = MeterEntry::new("ctx_read", 10, 2, 8);
306        let task = tokio::spawn(async move {
307            MeterStore::append_best_effort(entry);
308            42_u8
309        });
310        assert_eq!(task.await.unwrap(), 42);
311    }
312}
lean_ctx/core/metering.rs

lean_ctx/core/
metering.rs
