lean_ctx/core/
edit_metering.rs1use std::sync::atomic::{AtomicUsize, Ordering};
26use std::sync::{Mutex, OnceLock};
27
28use serde::{Deserialize, Serialize};
29
30const STORE_FILE: &str = "edit_metering.json";
31const FLUSH_EVERY: usize = 5;
32
33static STORE: OnceLock<Mutex<EditMeteringStore>> = OnceLock::new();
34static RECORD_CALLS: AtomicUsize = AtomicUsize::new(0);
35
36#[derive(Debug, Clone, Serialize, Deserialize, Default)]
41#[serde(default)]
42pub struct EditMeteringStore {
43 pub anchored_calls: u64,
45 pub anchored_ops: u64,
47 pub anchored_avoided_output_tokens: u64,
49 pub anchored_conflicts: u64,
51 pub str_replace_calls: u64,
53 pub str_replace_old_string_tokens: u64,
55 pub str_replace_misses: u64,
57 #[serde(skip)]
58 dirty: bool,
59}
60
61impl EditMeteringStore {
62 fn load_from_disk() -> Self {
63 let Ok(raw) = std::fs::read_to_string(store_path()) else {
64 return Self::default();
65 };
66 serde_json::from_str(&raw).unwrap_or_default()
67 }
68
69 pub fn save(&self) -> std::io::Result<()> {
70 let path = store_path();
71 if let Some(parent) = path.parent() {
72 std::fs::create_dir_all(parent)?;
73 }
74 let json = serde_json::to_string(self)?;
75 let tmp = path.with_extension("tmp");
76 std::fs::write(&tmp, json)?;
77 std::fs::rename(&tmp, &path)
78 }
79}
80
81fn store_path() -> std::path::PathBuf {
82 crate::core::data_dir::lean_ctx_data_dir()
83 .unwrap_or_else(|_| std::path::PathBuf::from("."))
84 .join(STORE_FILE)
85}
86
87fn global() -> &'static Mutex<EditMeteringStore> {
88 STORE.get_or_init(|| Mutex::new(EditMeteringStore::load_from_disk()))
89}
90
91fn with_store(f: impl FnOnce(&mut EditMeteringStore)) {
92 let Ok(mut store) = global().lock() else {
93 return;
94 };
95 f(&mut store);
96 store.dirty = true;
97 let n = RECORD_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
98 if n.is_multiple_of(FLUSH_EVERY) && store.save().is_ok() {
99 store.dirty = false;
100 }
101}
102
103pub fn record_anchored_success(ops: u64, avoided_tokens: u64) {
106 with_store(|s| {
107 s.anchored_calls = s.anchored_calls.saturating_add(1);
108 s.anchored_ops = s.anchored_ops.saturating_add(ops);
109 s.anchored_avoided_output_tokens = s
110 .anchored_avoided_output_tokens
111 .saturating_add(avoided_tokens);
112 });
113}
114
115pub fn record_anchored_conflict() {
117 with_store(|s| s.anchored_conflicts = s.anchored_conflicts.saturating_add(1));
118}
119
120pub fn record_str_replace_success(old_string_tokens: u64) {
122 with_store(|s| {
123 s.str_replace_calls = s.str_replace_calls.saturating_add(1);
124 s.str_replace_old_string_tokens = s
125 .str_replace_old_string_tokens
126 .saturating_add(old_string_tokens);
127 });
128}
129
130pub fn record_str_replace_miss() {
132 with_store(|s| s.str_replace_misses = s.str_replace_misses.saturating_add(1));
133}
134
135pub fn metrics_snapshot() -> serde_json::Value {
137 let Ok(store) = global().lock() else {
138 return serde_json::json!({});
139 };
140 serde_json::json!({
141 "anchored_calls": store.anchored_calls,
142 "anchored_ops": store.anchored_ops,
143 "anchored_avoided_output_tokens": store.anchored_avoided_output_tokens,
144 "anchored_conflicts": store.anchored_conflicts,
145 "str_replace_calls": store.str_replace_calls,
146 "str_replace_old_string_tokens": store.str_replace_old_string_tokens,
147 "str_replace_misses": store.str_replace_misses,
148 })
149}
150
151pub fn flush() {
152 if let Ok(store) = global().lock()
153 && store.dirty
154 {
155 let _ = store.save();
156 }
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162
163 #[test]
164 fn counters_accumulate_and_saturate() {
165 let mut s = EditMeteringStore::default();
166 s.anchored_calls = u64::MAX;
167 s.anchored_calls = s.anchored_calls.saturating_add(1);
168 assert_eq!(s.anchored_calls, u64::MAX, "saturating, never wraps");
169 }
170
171 #[test]
172 fn roundtrip_serialization() {
173 let s = EditMeteringStore {
174 anchored_calls: 3,
175 anchored_ops: 7,
176 anchored_avoided_output_tokens: 412,
177 anchored_conflicts: 1,
178 str_replace_calls: 2,
179 str_replace_old_string_tokens: 260,
180 str_replace_misses: 4,
181 dirty: false,
182 };
183 let json = serde_json::to_string(&s).unwrap();
184 let back: EditMeteringStore = serde_json::from_str(&json).unwrap();
185 assert_eq!(back.anchored_ops, 7);
186 assert_eq!(back.anchored_avoided_output_tokens, 412);
187 assert_eq!(back.str_replace_misses, 4);
188 }
189
190 #[test]
191 fn legacy_or_empty_store_deserializes() {
192 let s: EditMeteringStore = serde_json::from_str("{}").unwrap();
193 assert_eq!(s.anchored_calls, 0);
194 assert_eq!(s.str_replace_old_string_tokens, 0);
195 }
196}