1use std::collections::HashMap;
16use std::sync::Mutex;
17use std::time::Instant;
18
19use serde::{Deserialize, Serialize};
20
21pub const DEFAULT_BEGIN_SHARE: f64 = 0.7;
23const MIN_OBSERVATIONS: u32 = 20;
25const SHARE_CLAMP: (f64, f64) = (0.4, 0.9);
27const FLUSH_SECS: u64 = 60;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum Position {
31 Begin,
32 End,
33}
34
35impl Position {
36 pub fn as_str(self) -> &'static str {
37 match self {
38 Position::Begin => "begin",
39 Position::End => "end",
40 }
41 }
42
43 pub fn parse(s: &str) -> Option<Self> {
44 match s {
45 "begin" => Some(Position::Begin),
46 "end" => Some(Position::End),
47 _ => None,
48 }
49 }
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize, Default)]
53pub struct PlacementStats {
54 pub begin_hits: u32,
55 pub begin_misses: u32,
56 pub end_hits: u32,
57 pub end_misses: u32,
58}
59
60impl PlacementStats {
61 fn total(&self) -> u32 {
62 self.begin_hits + self.begin_misses + self.end_hits + self.end_misses
63 }
64
65 fn hit_rate(&self, pos: Position) -> f64 {
67 let (hits, misses) = match pos {
68 Position::Begin => (self.begin_hits, self.begin_misses),
69 Position::End => (self.end_hits, self.end_misses),
70 };
71 (hits as f64 + 1.0) / ((hits + misses) as f64 + 2.0)
72 }
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize, Default)]
76pub struct LitmCalibration {
77 pub per_profile: HashMap<String, PlacementStats>,
79 pub schema_version: u32,
80}
81
82static BUFFER: Mutex<Option<(LitmCalibration, Instant)>> = Mutex::new(None);
83
84fn store_path() -> std::path::PathBuf {
85 crate::core::paths::cache_dir()
86 .unwrap_or_else(|_| std::path::PathBuf::from("."))
87 .join("litm_calibration.json")
88}
89
90impl LitmCalibration {
91 fn load_from_disk() -> Self {
92 if let Ok(content) = std::fs::read_to_string(store_path()) {
93 if let Ok(c) = serde_json::from_str::<LitmCalibration>(&content) {
94 return c;
95 }
96 }
97 LitmCalibration {
98 schema_version: 1,
99 ..Default::default()
100 }
101 }
102
103 fn save_to_disk(&self) {
104 let path = store_path();
105 if let Some(parent) = path.parent() {
106 let _ = std::fs::create_dir_all(parent);
107 }
108 if let Ok(json) = serde_json::to_string_pretty(self) {
109 let _ = std::fs::write(path, json);
110 }
111 }
112
113 pub fn merge_from(&mut self, other: &Self) {
117 for (profile, theirs) in &other.per_profile {
118 let ours = self.per_profile.entry(profile.clone()).or_default();
119 ours.begin_hits = ours.begin_hits.max(theirs.begin_hits);
120 ours.begin_misses = ours.begin_misses.max(theirs.begin_misses);
121 ours.end_hits = ours.end_hits.max(theirs.end_hits);
122 ours.end_misses = ours.end_misses.max(theirs.end_misses);
123 }
124 }
125
126 pub fn record(&mut self, profile: &str, pos: Position, hit: bool) {
127 let stats = self.per_profile.entry(profile.to_string()).or_default();
128 match (pos, hit) {
129 (Position::Begin, true) => stats.begin_hits += 1,
130 (Position::Begin, false) => stats.begin_misses += 1,
131 (Position::End, true) => stats.end_hits += 1,
132 (Position::End, false) => stats.end_misses += 1,
133 }
134 }
135
136 pub fn begin_share(&self, profile: &str) -> f64 {
140 let Some(stats) = self.per_profile.get(profile) else {
141 return DEFAULT_BEGIN_SHARE;
142 };
143 if stats.total() < MIN_OBSERVATIONS {
144 return DEFAULT_BEGIN_SHARE;
145 }
146 let hb = stats.hit_rate(Position::Begin);
147 let he = stats.hit_rate(Position::End);
148 let raw = hb / (hb + he);
149 let share = DEFAULT_BEGIN_SHARE + (raw - 0.5) * 2.0 * (1.0 - DEFAULT_BEGIN_SHARE);
151 share.clamp(SHARE_CLAMP.0, SHARE_CLAMP.1)
152 }
153
154 pub fn totals(&self) -> (u32, u32, u32, u32) {
157 self.per_profile.values().fold((0, 0, 0, 0), |acc, s| {
158 (
159 acc.0 + s.begin_hits,
160 acc.1 + s.begin_misses,
161 acc.2 + s.end_hits,
162 acc.3 + s.end_misses,
163 )
164 })
165 }
166
167 pub fn report_lines(&self) -> Vec<String> {
168 let mut profiles: Vec<_> = self.per_profile.iter().collect();
169 profiles.sort_by(|a, b| a.0.cmp(b.0));
170 profiles
171 .iter()
172 .map(|(name, s)| {
173 format!(
174 " {name}: begin {}/{} hit, end {}/{} hit -> share {:.2}",
175 s.begin_hits,
176 s.begin_hits + s.begin_misses,
177 s.end_hits,
178 s.end_hits + s.end_misses,
179 self.begin_share(name)
180 )
181 })
182 .collect()
183 }
184}
185
186fn with_buffer<R>(f: impl FnOnce(&mut LitmCalibration) -> R) -> R {
187 let mut guard = BUFFER
188 .lock()
189 .unwrap_or_else(std::sync::PoisonError::into_inner);
190 if guard.is_none() {
191 *guard = Some((LitmCalibration::load_from_disk(), Instant::now()));
192 }
193 let (cal, last_flush) = guard.as_mut().expect("buffer initialized above");
194 let result = f(cal);
195 if last_flush.elapsed().as_secs() >= FLUSH_SECS {
196 cal.save_to_disk();
197 *last_flush = Instant::now();
198 }
199 result
200}
201
202pub fn record_outcome(profile: &str, pos: Position, hit: bool) {
204 if profile.is_empty() {
205 return;
206 }
207 with_buffer(|c| c.record(profile, pos, hit));
208}
209
210pub fn begin_share(profile: &str) -> f64 {
212 with_buffer(|c| c.begin_share(profile))
213}
214
215pub fn flush() {
217 let guard = BUFFER
218 .lock()
219 .unwrap_or_else(std::sync::PoisonError::into_inner);
220 if let Some((ref cal, _)) = *guard {
221 cal.save_to_disk();
222 }
223}
224
225pub fn report() -> Vec<String> {
227 with_buffer(|c| c.report_lines())
228}
229
230pub fn totals() -> (u32, u32, u32, u32) {
232 with_buffer(|c| c.totals())
233}
234
235pub fn snapshot() -> Vec<(String, PlacementStats, f64)> {
238 with_buffer(|c| {
239 let mut v: Vec<_> = c
240 .per_profile
241 .iter()
242 .map(|(p, s)| (p.clone(), s.clone(), c.begin_share(p)))
243 .collect();
244 v.sort_by(|a, b| a.0.cmp(&b.0));
245 v
246 })
247}
248
249pub fn export_state() -> LitmCalibration {
251 with_buffer(|c| c.clone())
252}
253
254pub fn merge_state(other: &LitmCalibration) {
256 with_buffer(|c| c.merge_from(other));
257 flush();
258}
259
260pub fn key_matches(manifest_key: &str, query: &str) -> bool {
263 let k = manifest_key.to_lowercase();
264 let q = query.to_lowercase();
265 if k.len() >= 6 && q.len() >= 6 && (k.contains(&q) || q.contains(&k)) {
266 return true;
267 }
268 let ks: std::collections::HashSet<&str> = k.split_whitespace().collect();
269 let qs: std::collections::HashSet<&str> = q.split_whitespace().collect();
270 if ks.is_empty() || qs.is_empty() {
271 return false;
272 }
273 let inter = ks.intersection(&qs).count() as f64;
274 let union = ks.union(&qs).count() as f64;
275 inter / union >= 0.5
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281
282 #[test]
283 fn default_share_before_min_observations() {
284 let mut c = LitmCalibration::default();
285 for _ in 0..MIN_OBSERVATIONS - 1 {
286 c.record("claude", Position::Begin, false);
287 }
288 assert!((c.begin_share("claude") - DEFAULT_BEGIN_SHARE).abs() < f64::EPSILON);
289 }
290
291 #[test]
292 fn begin_miss_series_lowers_share() {
293 let mut c = LitmCalibration::default();
294 for _ in 0..30 {
295 c.record("claude", Position::Begin, false);
296 c.record("claude", Position::End, true);
297 }
298 let share = c.begin_share("claude");
299 assert!(
300 share < DEFAULT_BEGIN_SHARE,
301 "begin misses should lower share, got {share}"
302 );
303 assert!(share >= SHARE_CLAMP.0);
304 }
305
306 #[test]
307 fn end_miss_series_raises_share() {
308 let mut c = LitmCalibration::default();
309 for _ in 0..30 {
310 c.record("gpt", Position::Begin, true);
311 c.record("gpt", Position::End, false);
312 }
313 let share = c.begin_share("gpt");
314 assert!(share > DEFAULT_BEGIN_SHARE);
315 assert!(share <= SHARE_CLAMP.1);
316 }
317
318 #[test]
319 fn balanced_hits_keep_default_layout() {
320 let mut c = LitmCalibration::default();
321 for _ in 0..50 {
322 c.record("gemini", Position::Begin, true);
323 c.record("gemini", Position::End, true);
324 }
325 assert!((c.begin_share("gemini") - DEFAULT_BEGIN_SHARE).abs() < 0.01);
326 }
327
328 #[test]
329 fn unknown_profile_uses_default() {
330 let c = LitmCalibration::default();
331 assert!((c.begin_share("nope") - DEFAULT_BEGIN_SHARE).abs() < f64::EPSILON);
332 }
333
334 #[test]
335 fn key_matching_containment_and_jaccard() {
336 assert!(key_matches("billing webhook fix", "webhook fix"));
337 assert!(key_matches(
338 "stripe cancel_at parsing",
339 "parsing stripe cancel_at"
340 ));
341 assert!(!key_matches("frontend css", "database migration"));
342 assert!(key_matches("ab", "ab")); }
344}