use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Mutex, OnceLock};
use serde::{Deserialize, Serialize};
const STORE_FILE: &str = "edit_metering.json";
const FLUSH_EVERY: usize = 5;
static STORE: OnceLock<Mutex<EditMeteringStore>> = OnceLock::new();
static RECORD_CALLS: AtomicUsize = AtomicUsize::new(0);
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(default)]
pub(crate) struct EditMeteringStore {
pub anchored_calls: u64,
pub anchored_ops: u64,
pub anchored_avoided_output_tokens: u64,
pub anchored_conflicts: u64,
pub str_replace_calls: u64,
pub str_replace_old_string_tokens: u64,
pub str_replace_misses: u64,
#[serde(skip)]
dirty: bool,
}
impl EditMeteringStore {
fn load_from_disk() -> Self {
let Ok(raw) = std::fs::read_to_string(store_path()) else {
return Self::default();
};
serde_json::from_str(&raw).unwrap_or_default()
}
pub(crate) fn save(&self) -> std::io::Result<()> {
let path = store_path();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let json = serde_json::to_string(self)?;
let tmp = path.with_extension("tmp");
std::fs::write(&tmp, json)?;
std::fs::rename(&tmp, &path)
}
}
fn store_path() -> std::path::PathBuf {
crate::core::data_dir::lean_ctx_data_dir()
.unwrap_or_else(|_| std::path::PathBuf::from("."))
.join(STORE_FILE)
}
fn global() -> &'static Mutex<EditMeteringStore> {
STORE.get_or_init(|| Mutex::new(EditMeteringStore::load_from_disk()))
}
fn with_store(f: impl FnOnce(&mut EditMeteringStore)) {
let Ok(mut store) = global().lock() else {
return;
};
f(&mut store);
store.dirty = true;
let n = RECORD_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
if n.is_multiple_of(FLUSH_EVERY) && store.save().is_ok() {
store.dirty = false;
}
}
pub(crate) fn record_anchored_success(ops: u64, avoided_tokens: u64) {
with_store(|s| {
s.anchored_calls = s.anchored_calls.saturating_add(1);
s.anchored_ops = s.anchored_ops.saturating_add(ops);
s.anchored_avoided_output_tokens = s
.anchored_avoided_output_tokens
.saturating_add(avoided_tokens);
});
}
pub(crate) fn record_anchored_conflict() {
with_store(|s| s.anchored_conflicts = s.anchored_conflicts.saturating_add(1));
}
pub(crate) fn record_str_replace_success(old_string_tokens: u64) {
with_store(|s| {
s.str_replace_calls = s.str_replace_calls.saturating_add(1);
s.str_replace_old_string_tokens = s
.str_replace_old_string_tokens
.saturating_add(old_string_tokens);
});
}
pub(crate) fn record_str_replace_miss() {
with_store(|s| s.str_replace_misses = s.str_replace_misses.saturating_add(1));
}
pub(crate) fn metrics_snapshot() -> serde_json::Value {
let Ok(store) = global().lock() else {
return serde_json::json!({});
};
serde_json::json!({
"anchored_calls": store.anchored_calls,
"anchored_ops": store.anchored_ops,
"anchored_avoided_output_tokens": store.anchored_avoided_output_tokens,
"anchored_conflicts": store.anchored_conflicts,
"str_replace_calls": store.str_replace_calls,
"str_replace_old_string_tokens": store.str_replace_old_string_tokens,
"str_replace_misses": store.str_replace_misses,
})
}
pub(crate) fn flush() {
if let Ok(store) = global().lock()
&& store.dirty
{
let _ = store.save();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn counters_accumulate_and_saturate() {
let mut s = EditMeteringStore::default();
s.anchored_calls = u64::MAX;
s.anchored_calls = s.anchored_calls.saturating_add(1);
assert_eq!(s.anchored_calls, u64::MAX, "saturating, never wraps");
}
#[test]
fn roundtrip_serialization() {
let s = EditMeteringStore {
anchored_calls: 3,
anchored_ops: 7,
anchored_avoided_output_tokens: 412,
anchored_conflicts: 1,
str_replace_calls: 2,
str_replace_old_string_tokens: 260,
str_replace_misses: 4,
dirty: false,
};
let json = serde_json::to_string(&s).unwrap();
let back: EditMeteringStore = serde_json::from_str(&json).unwrap();
assert_eq!(back.anchored_ops, 7);
assert_eq!(back.anchored_avoided_output_tokens, 412);
assert_eq!(back.str_replace_misses, 4);
}
#[test]
fn legacy_or_empty_store_deserializes() {
let s: EditMeteringStore = serde_json::from_str("{}").unwrap();
assert_eq!(s.anchored_calls, 0);
assert_eq!(s.str_replace_old_string_tokens, 0);
}
}