5use std::path::{Path, PathBuf};
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::{Arc, Mutex};

9use whiskers_core::*;
10use whiskers_engine::*;
11use whiskers_store::{SqliteChat, SqliteLog, SqliteMemory, SqlitePictures};
12
13const FACT: &str = "Biscuit is a purple bunny FORGOTTEN-FACT-4471";
14const FIRST: &str = "my bunny Biscuit is purple SAID-FIRST-5582";
15const AGAIN: &str = "Biscuit hopped again SAID-AGAIN-6693";
16const AFTER: &str = "what shall we play now";
17const OLD_SUMMARY: &str = "SUMMARY-SECRET-7704 she has a purple bunny";
18const OLD_TURN: &str = "OLD-TURN-SECRET-8815";

22struct Brain {
23    seen: Mutex<Vec<(String, Vec<Turn>)>>,
24    down: AtomicBool,
25    teaches: &'static str,

27    hold: Mutex<Option<Hold>>,
28}

30struct Hold {
31    kind: &'static str,
32    started: std::sync::mpsc::Sender<()>,
33    release: std::sync::mpsc::Receiver<()>,
34}
35
36impl Brain {

38    fn hold_next(&self, kind: &'static str) -> (std::sync::mpsc::Receiver<()>, std::sync::mpsc::Sender<()>) {
39        let (started_tx, started_rx) = std::sync::mpsc::channel();
40        let (release_tx, release_rx) = std::sync::mpsc::channel();
41        *self.hold.lock().unwrap() = Some(Hold { kind, started: started_tx, release: release_rx });
42        (started_rx, release_tx)
43    }

45    fn pause_if_held(&self, kind: &str) {
46        let hold = {
47            let mut g = self.hold.lock().unwrap();
48            if g.as_ref().is_some_and(|h| h.kind == kind) { g.take() } else { None }
49        };
50        if let Some(h) = hold {
51            h.started.send(()).unwrap();
52            h.release.recv().unwrap();
53        }
54    }
55}
56
57impl Model for Brain {
58    fn complete(&self, system: &str, history: &[Turn]) -> Result<String, ModelError> {
59        self.seen.lock().unwrap().push((system.to_owned(), history.to_vec()));
60        let asked = history.last().map(|t| t.text.as_str()).unwrap_or_default();
61        self.pause_if_held(if asked.starts_with("Already known:") { "remember" } else if asked.starts_with("Summary so far:") { "summary" } else { "chat" });
62        if asked.starts_with("Already known:") {
63            if self.down.load(Ordering::SeqCst) {
64                return Err(ModelError::new("down"));
65            }
66            return Ok(self.teaches.to_owned());
67        }
68        if asked.starts_with("Summary so far:") {
69            return Ok("A summary.".into());
70        }
71        Ok("What a lovely thing to say!".into())
72    }
73}
74
75struct AllowAll;
76impl Guard for AllowAll {
77    fn check(&self, _d: Direction, _a: Age, _t: &str) -> Result<Verdict, GuardError> {
78        Ok(Verdict::Allow)
79    }
80}
81struct NoEmbedding;
82impl Embedder for NoEmbedding {
83    fn embed(&self, _t: &[String]) -> Result<Vec<Vec<f32>>, EmbedError> {
84        Err(EmbedError("none in tests".into()))
85    }
86}
87struct NoRanking;
88impl Ranker for NoRanking {
89    fn rank(&self, _q: &str, _c: &[String]) -> Result<Vec<f32>, RankError> {
90        Err(RankError("none in tests".into()))
91    }
92}
93
94fn dir(name: &str) -> PathBuf {
95    let d = std::env::temp_dir().join(format!("whiskers-engine-forget-{name}-{}", std::process::id()));
96    let _ = std::fs::remove_dir_all(&d);
97    std::fs::create_dir_all(&d).unwrap();
98    d
99}
100
101fn engine(d: &Path, brain: Arc<Brain>) -> Engine {
102    let cfg = EngineConfig { data_dir: d.to_owned(), gateway_url: String::new(), guard_url: String::new(), model: String::new(), system_prompt: None };
103    let store = Store::open(&d.join(DATABASE)).unwrap();
104    let device = whiskers_engine::device_id(d).unwrap();
105    let parts = Parts {
106        model: brain,
107        guard: Arc::new(AllowAll),
108        log: SharedLog::new(Box::new(SqliteLog::new(store.clone(), &device))),
109        pictures: Box::new(SqlitePictures::new(store.clone())),
110        memory: SharedMemory::new(Box::new(SqliteMemory::open(store.clone(), &device).unwrap())),
111        chat: SharedChat::open(Box::new(SqliteChat::new(store.clone()))).unwrap(),
112        recall: Arc::new(Recall::new(Arc::new(NoEmbedding), Arc::new(NoRanking))),
113        clock: Arc::new(SystemClock),
114        profile: SharedProfile::default(),
115    };
116    Engine::with_parts(cfg, store, parts, Arc::new(NoEmbedding), Box::new(Brain { seen: Mutex::default(), down: AtomicBool::new(false), teaches: "[]", hold: Mutex::default() })).unwrap()
117}
118
119fn brain(teaches: &'static str) -> Arc<Brain> {
120    Arc::new(Brain { seen: Mutex::default(), down: AtomicBool::new(false), teaches, hold: Mutex::default() })
121}
122
123fn everything_in_the_last_prompt(b: &Brain) -> String {
124    let seen = b.seen.lock().unwrap();
125    let (system, history) = seen.iter().rev().find(|(_, h)| !h.last().is_some_and(|t| t.text.starts_with("Already known:") || t.text.starts_with("Summary so far:"))).unwrap();
126    format!("{system}\n{}", history.iter().map(|t| t.text.as_str()).collect::<Vec<_>>().join("\n"))
127}
128
129fn on_disk(d: &Path) -> Vec<u8> {
130    let mut all = std::fs::read(d.join(DATABASE)).unwrap();
131    all.extend(std::fs::read(d.join(format!("{DATABASE}-wal"))).unwrap_or_default());
132    all
133}
134
135fn contains(haystack: &[u8], needle: &str) -> bool {
136    haystack.windows(needle.len()).any(|w| w == needle.as_bytes())
137}

140fn with_an_old_chat(d: &Path) {
141    let store = Store::open(&d.join(DATABASE)).unwrap();
142    let mut chat = SqliteChat::new(store);
143    ChatStore::save(
144        &mut chat,
145        &ChatState {
146            summary: OLD_SUMMARY.into(),
147            turns: vec![ChatTurn { speaker: Speaker::Child, text: OLD_TURN.into() }, ChatTurn { speaker: Speaker::Whiskers, text: "Lovely!".into() }],
148            last_active_ms: SystemClock.now_ms(),
149            version: 3,
150            scrubbed: vec![],
151        },
152    )
153    .unwrap();
154}

156#[test]
157fn a_prompt_built_after_a_forgetting_holds_nothing_of_the_memory_or_of_what_it_was_learned_from() {
158    let d = dir("prompt");
159    with_an_old_chat(&d);
160    let b = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
161    let e = engine(&d, b.clone());
162
163    // She tells it, and it is learned.
164    assert_eq!(e.hear(FIRST, vec![]).outcome, Outcome::Answered);
165    let learned = e.reflect();
166    assert_eq!(learned.len(), 1);
167    assert_eq!(learned[0].text, FACT);
168    // She says it again in a turn that is still waiting to be thought about when it is forgotten.
169    e.hear(AGAIN, vec![]);
170    assert_eq!(e.waiting_for_memory(), 1);
171
172    // The instrument can see them: before the forgetting, what she said and the old chat are in what the cat is told, and
173    // the bytes are in the files.
174    e.hear("and now", vec![]);
175    let before = everything_in_the_last_prompt(&b);
176    for w in [FIRST, OLD_SUMMARY, OLD_TURN] {
177        assert!(before.contains(w), "{w} should be in the prompt before: {before}");
178    }
179    assert!(contains(&on_disk(&d), FIRST) && contains(&on_disk(&d), OLD_SUMMARY) && contains(&on_disk(&d), FACT));
180
181    assert!(e.forget(learned[0].id));
182
183    // The exchange that was waiting cannot teach it back (and nothing of it was lost but the learning).
184    assert_eq!(e.waiting_for_memory(), 0);
185    assert!(e.reflect().is_empty(), "it is not learned again from what was said before");
186    assert!(e.facts().is_empty());
187    // The files hold none of it: the memory, what she said in the turn it was learned in, the chat that may carry it.
188    let bytes = on_disk(&d);
189    for w in [FACT, FIRST, OLD_SUMMARY, OLD_TURN] {
190        assert!(!contains(&bytes, w), "{w} is still in the database or its write-ahead log");
191    }
192    // Whiskers is told nothing of it: not the memory, not the summary, not the recent turns that may carry it.
193    e.hear(AFTER, vec![]);
194    let after = everything_in_the_last_prompt(&b);
195    for w in [FACT, "FORGOTTEN-FACT", FIRST, "SAID-FIRST", OLD_SUMMARY, "SUMMARY-SECRET", OLD_TURN, "OLD-TURN-SECRET"] {
196        assert!(!after.contains(w), "{w} is still in what the cat is told: {after}");
197    }
198    assert!(after.contains(AFTER), "the prompt is the new turn and nothing more");
199    assert_eq!(e.waiting_for_memory(), 1, "a turn after the forgetting is thought about as ever");
200}
201
202#[test]
203fn the_parents_log_keeps_that_the_talk_happened_and_not_what_was_said() {
204    let d = dir("log");
205    let b = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
206    let e = engine(&d, b);
207    e.hear(FIRST, vec![]);
208    let learned = e.reflect();
209    e.hear("an unrelated question about the weather", vec![]);
210    let day = e.digest(0, u64::MAX).unwrap();
211    assert_eq!((day.answered, day.exchanges.len()), (2, 2));
212    assert!(e.forget(learned[0].id));
213    let day = e.digest(0, u64::MAX).unwrap();
214    assert_eq!(day.exchanges.len(), 2, "both turns are still in the log");
215    assert_eq!((day.answered, day.stopped), (2, 0), "and how each ended");
216    let heard: Vec<&str> = day.exchanges.iter().map(|x| x.heard.as_str()).collect();
217    assert_eq!(heard, ["", "an unrelated question about the weather"], "what she said in the turn it was learned in is empty; the other turn is untouched");
218    assert!(day.exchanges[0].said.is_empty());
219    assert!(day.facts_learned.is_empty(), "and what was learned from it reads as nothing");
220}
221
222#[test]
223fn what_waits_to_be_thought_about_survives_the_app_being_stopped() {
224    let d = dir("restart");
225    let down = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
226    down.down.store(true, Ordering::SeqCst);
227    let first_turn = {
228        let e = engine(&d, down);
229        assert_eq!(e.hear(FIRST, vec![]).outcome, Outcome::Answered);
230        assert!(e.reflect().is_empty(), "the model cannot be reached");
231        assert_eq!(e.waiting_for_memory(), 1, "so it waits");
232        // The process dies here.
233        SqliteLog::new(Store::open(&d.join(DATABASE)).unwrap(), &whiskers_engine::device_id(&d).unwrap())
234            .own_entries()
235            .unwrap()
236            .into_iter()
237            .find(|x| matches!(x.event, Event::Heard { .. }))
238            .unwrap()
239            .turn
240    };
241    let up = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
242    let e = engine(&d, up);
243    assert_eq!(e.waiting_for_memory(), 1, "it is still waiting when the app starts again");
244    let learned = e.reflect();
245    assert_eq!(learned.len(), 1);
246    assert_eq!(e.waiting_for_memory(), 0, "dealt with, and taken off");
247    assert_eq!(learned[0].mentions[0].turn, first_turn, "learned from the turn she said it in, which the memory names");
248    assert!(e.reflect().is_empty());
249}
250
251#[test]
252fn a_memory_forgotten_while_the_memory_work_is_thinking_about_an_exchange_is_not_filed_from_it() {
253    let d = dir("race-filing");
254    let b = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
255    let e = Arc::new(engine(&d, b.clone()));
256    // Something is learned and will be forgotten, and an exchange waits.
257    e.hear("my cat is called Tom", vec![]);
258    let learned = e.reflect();
259    assert_eq!(learned.len(), 1, "the model proposes the same fact for anything: that is the memory the parents will forget");
260    e.hear("Biscuit again", vec![]);
261    assert_eq!(e.waiting_for_memory(), 1);
262    // The memory work reads the exchange and is in the model when the parents forget.
263    let (in_the_model, release) = b.hold_next("remember");
264    let worker = {
265        let e = e.clone();
266        std::thread::spawn(move || e.reflect())
267    };
268    in_the_model.recv_timeout(std::time::Duration::from_secs(10)).expect("the memory work reached the model");
269    assert!(e.forget(learned[0].id));
270    release.send(()).unwrap();
271    let filed = worker.join().unwrap();
272    assert!(filed.is_empty(), "the exchange was dropped by the forgetting, and the model's answer to it is not filed");
273    assert!(e.facts().is_empty());
274    assert_eq!(e.waiting_for_memory(), 0);
275}
276
277#[test]
278fn a_turn_being_answered_when_a_memory_is_forgotten_is_neither_kept_in_the_chat_nor_learned_from() {
279    let d = dir("mid-turn");
280    let b = brain(r#"["Biscuit is a purple bunny FORGOTTEN-FACT-4471"]"#);
281    let e = Arc::new(engine(&d, b.clone()));
282    e.hear(FIRST, vec![]);
283    let learned = e.reflect();
284    // A turn begins and the model is thinking when the parents forget the memory.
285    let (in_the_model, release) = b.hold_next("chat");
286    let talker = {
287        let e = e.clone();
288        std::thread::spawn(move || e.hear("MID-TURN-WORDS-3326 tell me about Biscuit", vec![]))
289    };
290    in_the_model.recv_timeout(std::time::Duration::from_secs(10)).expect("the turn reached the model");
291    assert!(e.forget(learned[0].id));
292    release.send(()).unwrap();
293    assert_eq!(talker.join().unwrap().outcome, Outcome::Answered, "she is still answered");
294    assert_eq!(e.waiting_for_memory(), 0, "but what the answer may have used is not learned from");
295    // Nor is it in the chat the next prompt is made from.
296    e.hear("next", vec![]);
297    assert!(!everything_in_the_last_prompt(&b).contains("MID-TURN-WORDS"));
298}
299
300#[test]
301fn the_household_day_follows_the_local_time_the_shell_reports() {
302    let d = dir("zone");
303    let e = engine(&d, brain("[]"));
304    assert_eq!(e.household_config().utc_offset_minutes, None);
305    let now = SystemClock.now_ms();
306    let local = |ahead: i64| {
307        let m = (now / 60_000) as i64 + ahead;
308        (m.div_euclid(1440) as u32, m.rem_euclid(1440) as u16)
309    };
310    // The first thing the device sees of its own clock, with the household knowing nothing, is written.
311    let (day, minute) = local(-240);
312    e.time_status(day, minute);
313    assert_eq!(e.household_config().utc_offset_minutes, Some(-240));
314    // Again and again (the shell asks every second): nothing more is written.
315    let stamp_of = |e: &Engine| serde_json::to_string(&e.household_config()).unwrap();
316    let before = stamp_of(&e);
317    for _ in 0..3 {
318        e.time_status(day, minute);
319        e.time_sheet(day, minute);
320    }
321    assert_eq!(stamp_of(&e), before);
322    // This device's own zone changes (it travels, or the clocks change): the household follows.
323    let (day, minute) = local(-300);
324    e.time_status(day, minute);
325    assert_eq!(e.household_config().utc_offset_minutes, Some(-300));
326    // A relaunch in a zone that differs from what the household has does not take it over: another device's stands
327    // until this one's own zone changes.
328    drop(e);
329    let again = engine(&d, brain("[]"));
330    let (day, minute) = local(60);
331    again.time_status(day, minute);
332    assert_eq!(again.household_config().utc_offset_minutes, Some(-300));
333}
334
335#[test]
336fn the_log_is_cleared_by_age_only_when_the_household_says_so_and_says_it_did() {
337    let d = dir("retention");
338    let e = engine(&d, brain("[]"));
339    let (long_ago, recent) = (SystemClock.now_ms() - 40 * 86_400_000, SystemClock.now_ms() - 86_400_000);
340    {
341        let store = Store::open(&d.join(DATABASE)).unwrap();
342        let mut log = SqliteLog::new(store, &whiskers_engine::device_id(&d).unwrap());
343        log.append(&Entry::new(long_ago, Event::Heard { text: "OLD-TALK-2468".into(), pictures: vec![] })).unwrap();
344        log.append(&Entry::new(recent, Event::Heard { text: "recent talk".into(), pictures: vec![] })).unwrap();
345    }
346    assert_eq!(e.apply_retention(), 0, "the default keeps everything");
347    assert!(contains(&on_disk(&d), "OLD-TALK-2468"));
348    let mut config = e.household_config();
349    config.log_retention_days = Some(30);
350    e.set_household_config(config);
351    assert_eq!(e.apply_retention(), 1);
352    assert!(!contains(&on_disk(&d), "OLD-TALK-2468"), "cleared from the file and its log");
353    let entries = SqliteLog::new(Store::open(&d.join(DATABASE)).unwrap(), &whiskers_engine::device_id(&d).unwrap()).own_entries().unwrap();
354    assert!(entries.iter().any(|x| x.event == Event::Expired && x.at_ms == long_ago), "the place and the time stay");
355    assert!(entries.iter().any(|x| matches!(&x.event, Event::Cleared { lines: 1, .. })), "and the log says it was cleared");
356    assert!(entries.iter().any(|x| matches!(&x.event, Event::Heard { text, .. } if text == "recent talk")));
357    assert_eq!(e.apply_retention(), 0, "nothing more to clear, nothing more said");
358    let days = e.log_days(0, u64::MAX, 10).unwrap();
359    assert!(days.iter().any(|d| d.cleared == 1), "{days:?}");
360    // Zero days is not a thing the household can ask for.
361    assert!(serde_json::from_str::<HouseholdConfig>(&serde_json::to_string(&e.household_config()).unwrap()).is_ok());
362}