Skip to main content

lean_ctx/core/
feedback.rs

1use std::collections::HashMap;
2use std::sync::Mutex;
3use std::time::Instant;
4
5use serde::{Deserialize, Serialize};
6
7const FEEDBACK_FLUSH_SECS: u64 = 60;
8
9static FEEDBACK_BUFFER: Mutex<Option<(FeedbackStore, Instant)>> = Mutex::new(None);
10
11/// Feedback loop for learning optimal compression parameters.
12///
13/// Tracks compression outcomes per session and learns which
14/// threshold combinations lead to fewer turns and higher success rates.
15
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct CompressionOutcome {
18    pub session_id: String,
19    pub language: String,
20    pub entropy_threshold: f64,
21    pub jaccard_threshold: f64,
22    pub total_turns: u32,
23    pub tokens_saved: u64,
24    pub tokens_original: u64,
25    pub cache_hits: u32,
26    pub total_reads: u32,
27    pub task_completed: bool,
28    pub timestamp: String,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, Default)]
32pub struct FeedbackStore {
33    pub outcomes: Vec<CompressionOutcome>,
34    pub learned_thresholds: HashMap<String, LearnedThresholds>,
35    #[serde(skip)]
36    pub project_root: Option<String>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct LearnedThresholds {
41    pub entropy: f64,
42    pub jaccard: f64,
43    pub sample_count: u32,
44    pub avg_efficiency: f64,
45}
46
47impl FeedbackStore {
48    pub fn load() -> Self {
49        let guard = FEEDBACK_BUFFER
50            .lock()
51            .unwrap_or_else(std::sync::PoisonError::into_inner);
52        if let Some((ref store, _)) = *guard {
53            let mut s = store.clone();
54            if s.project_root.is_none() {
55                s.project_root = std::env::current_dir()
56                    .ok()
57                    .map(|p| p.to_string_lossy().to_string());
58            }
59            return s;
60        }
61        drop(guard);
62
63        let path = feedback_path();
64        if path.exists()
65            && let Ok(content) = std::fs::read_to_string(&path)
66            && let Ok(mut store) = serde_json::from_str::<FeedbackStore>(&content)
67        {
68            store.project_root = std::env::current_dir()
69                .ok()
70                .map(|p| p.to_string_lossy().to_string());
71            return store;
72        }
73        Self {
74            project_root: std::env::current_dir()
75                .ok()
76                .map(|p| p.to_string_lossy().to_string()),
77            ..Self::default()
78        }
79    }
80
81    fn save_to_disk(&self) {
82        let path = feedback_path();
83        if let Some(parent) = path.parent() {
84            let _ = std::fs::create_dir_all(parent);
85        }
86        if let Ok(json) = serde_json::to_string_pretty(self) {
87            let _ = std::fs::write(path, json);
88        }
89    }
90
91    pub fn save(&self) {
92        self.save_to_disk();
93    }
94
95    pub fn flush() {
96        let guard = FEEDBACK_BUFFER
97            .lock()
98            .unwrap_or_else(std::sync::PoisonError::into_inner);
99        if let Some((ref store, _)) = *guard {
100            store.save_to_disk();
101        }
102    }
103
104    pub fn record_outcome(&mut self, outcome: CompressionOutcome) {
105        let lang = outcome.language.clone();
106        self.update_bandit(&outcome);
107        self.outcomes.push(outcome);
108
109        if self.outcomes.len() > 200 {
110            self.outcomes.drain(0..self.outcomes.len() - 200);
111        }
112
113        self.update_learned_thresholds(&lang);
114
115        let mut guard = FEEDBACK_BUFFER
116            .lock()
117            .unwrap_or_else(std::sync::PoisonError::into_inner);
118        let should_flush = match *guard {
119            Some((_, ref last)) => last.elapsed().as_secs() >= FEEDBACK_FLUSH_SECS,
120            None => true,
121        };
122        *guard = Some((
123            self.clone(),
124            guard.as_ref().map_or_else(Instant::now, |(_, t)| *t),
125        ));
126        if should_flush {
127            self.save_to_disk();
128            if let Some((_, ref mut t)) = *guard {
129                *t = Instant::now();
130            }
131        }
132    }
133
134    fn update_bandit(&self, outcome: &CompressionOutcome) {
135        let key = format!("{}_feedback", outcome.language);
136        let project_root = self.project_root.as_deref().unwrap_or(".");
137        let mut store = crate::core::bandit::BanditStore::load(project_root);
138        let bandit = store.get_or_create(&key);
139        bandit.total_pulls = bandit.total_pulls.saturating_add(1);
140
141        let efficiency = if outcome.tokens_original > 0 {
142            outcome.tokens_saved as f64 / outcome.tokens_original as f64
143        } else {
144            0.0
145        };
146        let success = efficiency > 0.3 && outcome.task_completed;
147
148        let arm_name = if outcome.entropy_threshold >= 1.0 {
149            "conservative"
150        } else if outcome.entropy_threshold >= 0.7 {
151            "balanced"
152        } else {
153            "aggressive"
154        };
155
156        let old_mean = bandit
157            .arms
158            .iter()
159            .find(|a| a.name == arm_name)
160            .map_or(0.5, super::bandit::BanditArm::mean);
161
162        bandit.update(arm_name, success);
163
164        let new_mean = bandit
165            .arms
166            .iter()
167            .find(|a| a.name == arm_name)
168            .map_or(0.5, super::bandit::BanditArm::mean);
169
170        if (new_mean - old_mean).abs() > 0.05 {
171            crate::core::events::emit_threshold_adapted(
172                &outcome.language,
173                arm_name,
174                old_mean,
175                new_mean,
176            );
177        }
178
179        if bandit.total_pulls > 0 && bandit.total_pulls.is_multiple_of(50) {
180            bandit.decay_all(0.95);
181        }
182
183        let _ = store.save(project_root);
184    }
185
186    fn update_learned_thresholds(&mut self, language: &str) {
187        let relevant: Vec<&CompressionOutcome> = self
188            .outcomes
189            .iter()
190            .filter(|o| o.language == language && o.task_completed)
191            .collect();
192
193        if relevant.len() < 5 {
194            return; // not enough data to learn
195        }
196
197        // Find the threshold combination that maximizes efficiency
198        // Efficiency = tokens_saved / tokens_original * (1 / total_turns)
199        let mut best_entropy = 1.0;
200        let mut best_jaccard = 0.7;
201        let mut best_efficiency = 0.0;
202
203        for outcome in &relevant {
204            let compression_ratio = if outcome.tokens_original > 0 {
205                outcome.tokens_saved as f64 / outcome.tokens_original as f64
206            } else {
207                0.0
208            };
209            let turn_efficiency = 1.0 / (outcome.total_turns.max(1) as f64);
210            let efficiency = compression_ratio * 0.6 + turn_efficiency * 0.4;
211
212            if efficiency > best_efficiency {
213                best_efficiency = efficiency;
214                best_entropy = outcome.entropy_threshold;
215                best_jaccard = outcome.jaccard_threshold;
216            }
217        }
218
219        // Weighted average with current learned values for stability
220        let entry = self
221            .learned_thresholds
222            .entry(language.to_string())
223            .or_insert(LearnedThresholds {
224                entropy: best_entropy,
225                jaccard: best_jaccard,
226                sample_count: 0,
227                avg_efficiency: 0.0,
228            });
229
230        let momentum = 0.7;
231        let old_entropy = entry.entropy;
232        let old_jaccard = entry.jaccard;
233        entry.entropy = entry.entropy * momentum + best_entropy * (1.0 - momentum);
234        entry.jaccard = entry.jaccard * momentum + best_jaccard * (1.0 - momentum);
235        entry.sample_count = relevant.len() as u32;
236        entry.avg_efficiency = best_efficiency;
237
238        if (old_entropy - entry.entropy).abs() > 0.01 || (old_jaccard - entry.jaccard).abs() > 0.01
239        {
240            crate::core::events::emit(crate::core::events::EventKind::ThresholdShift {
241                language: language.to_string(),
242                old_entropy,
243                new_entropy: entry.entropy,
244                old_jaccard,
245                new_jaccard: entry.jaccard,
246            });
247        }
248    }
249
250    pub fn get_learned_entropy(&self, language: &str) -> Option<f64> {
251        self.learned_thresholds.get(language).map(|t| t.entropy)
252    }
253
254    pub fn get_learned_jaccard(&self, language: &str) -> Option<f64> {
255        self.learned_thresholds.get(language).map(|t| t.jaccard)
256    }
257
258    pub fn format_report(&self) -> String {
259        let mut lines = vec![String::from("Feedback Loop Report")];
260        lines.push(format!("Total outcomes tracked: {}", self.outcomes.len()));
261        lines.push(String::new());
262
263        if self.learned_thresholds.is_empty() {
264            lines.push(
265                "No learned thresholds yet (need 5+ completed sessions per language).".to_string(),
266            );
267        } else {
268            lines.push("Learned Thresholds:".to_string());
269            for (lang, t) in &self.learned_thresholds {
270                lines.push(format!(
271                    "  {lang}: entropy={:.2} jaccard={:.2} (n={}, eff={:.1}%)",
272                    t.entropy,
273                    t.jaccard,
274                    t.sample_count,
275                    t.avg_efficiency * 100.0
276                ));
277            }
278        }
279
280        lines.push(String::new());
281        let project_root = self.project_root.as_deref().unwrap_or(".");
282        let store = crate::core::bandit::BanditStore::load(project_root);
283        lines.push(store.format_report());
284
285        lines.join("\n")
286    }
287}
288
289fn feedback_path() -> std::path::PathBuf {
290    crate::core::paths::state_dir()
291        .unwrap_or_else(|_| std::path::PathBuf::from("."))
292        .join("feedback.json")
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    #[test]
300    fn empty_store_loads() {
301        let store = FeedbackStore::default();
302        assert!(store.outcomes.is_empty());
303        assert!(store.learned_thresholds.is_empty());
304    }
305
306    #[test]
307    fn learned_thresholds_need_minimum_samples() {
308        let mut store = FeedbackStore::default();
309        for i in 0..3 {
310            store.record_outcome(CompressionOutcome {
311                session_id: format!("s{i}"),
312                language: "rs".to_string(),
313                entropy_threshold: 0.85,
314                jaccard_threshold: 0.72,
315                total_turns: 5,
316                tokens_saved: 1000,
317                tokens_original: 2000,
318                cache_hits: 3,
319                total_reads: 10,
320                task_completed: true,
321                timestamp: String::new(),
322            });
323        }
324        assert!(store.get_learned_entropy("rs").is_none()); // only 3, need 5
325    }
326}