Skip to main content

lean_ctx/core/
litm_calibration.rs

1//! Empirical LITM placement calibration (#539, EFF-2).
2//!
3//! The static `LitmProfile` weights encode the *assumed* positional attention
4//! of each client (begin strong, middle weak, end strong). This module closes
5//! the loop with observed evidence: every wakeup injection records a manifest
6//! of placed items (begin block vs end block). If the agent later explicitly
7//! recalls an item that was already placed, that placement failed to register
8//! — a *miss* for its position. Items never re-recalled count as *hits* when
9//! the manifest rotates on the next wakeup build.
10//!
11//! The calibrated output is `begin_share`: the fraction of session items that
12//! should go to the begin block. A persistently missing begin position shifts
13//! budget toward the end block (recency wins for that client), and vice versa.
14
15use std::collections::HashMap;
16use std::sync::Mutex;
17use std::time::Instant;
18
19use serde::{Deserialize, Serialize};
20
21/// Default share of items placed at the begin position (today's layout).
22pub const DEFAULT_BEGIN_SHARE: f64 = 0.7;
23/// Calibration only activates after this many total observations per profile.
24const MIN_OBSERVATIONS: u32 = 20;
25/// Calibrated share is clamped to this range — both positions always get data.
26const 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    /// Laplace-smoothed hit rate for a position.
66    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    /// Keyed by LITM profile name ("claude" | "gpt" | "gemini").
78    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::data_dir::lean_ctx_data_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    /// Team-merge (#550): element-wise maximum per counter. Cumulative
114    /// counters only ever grow, so `max` is idempotent on re-import and never
115    /// double-counts; divergent machines converge to the strongest evidence.
116    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    /// Calibrated begin-share for a profile. Returns the default until enough
137    /// observations exist; afterwards shifts budget toward the position that
138    /// empirically holds information for this client.
139    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        // Re-center: equal hit rates map to the default layout, not to 0.5.
150        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    /// Aggregate raw counters across profiles (#549 efficacy snapshots):
155    /// `(begin_hits, begin_misses, end_hits, end_misses)`.
156    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
202/// Process-global: record a placement outcome.
203pub 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
210/// Process-global: calibrated begin-share for a profile.
211pub fn begin_share(profile: &str) -> f64 {
212    with_buffer(|c| c.begin_share(profile))
213}
214
215/// Process-global: flush to disk (shutdown paths).
216pub 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
225/// Process-global: report lines for ctx_metrics.
226pub fn report() -> Vec<String> {
227    with_buffer(|c| c.report_lines())
228}
229
230/// Process-global aggregate counters (#549).
231pub fn totals() -> (u32, u32, u32, u32) {
232    with_buffer(|c| c.totals())
233}
234
235/// Process-global: machine-readable snapshot for the dashboard (#548):
236/// `(profile, stats, calibrated begin_share)`, sorted by profile.
237pub 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
249/// Process-global: clone of the full calibration state for export (#550).
250pub fn export_state() -> LitmCalibration {
251    with_buffer(|c| c.clone())
252}
253
254/// Process-global: merge a foreign calibration in and persist (#550).
255pub fn merge_state(other: &LitmCalibration) {
256    with_buffer(|c| c.merge_from(other));
257    flush();
258}
259
260/// Loose match between a recall query and a manifest key: lowercase
261/// containment either way, or token-Jaccard >= 0.5.
262pub 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")); // identical tokens match via Jaccard
343    }
344}