1use std::collections::VecDeque;
4use std::path::{Path, PathBuf};
5use std::process::Command;
6use std::sync::{Arc, Mutex};
7use std::thread::JoinHandle;
8
9use whycodes_core::types::Message;
10
11#[derive(Debug, Clone)]
18pub struct StashSlot(Arc<Mutex<Stash>>);
19
20#[derive(Debug)]
21enum Stash {
22 Running(JoinHandle<Option<String>>),
23 Done(Option<String>),
24}
25
26impl StashSlot {
27 fn spawn(project_path: &Path) -> Self {
28 let dir = project_path.to_path_buf();
29 Self::with(Stash::Running(std::thread::spawn(move || {
30 git_stash_create(&dir)
31 })))
32 }
33
34 fn done(stash: Option<String>) -> Self {
35 Self::with(Stash::Done(stash))
36 }
37
38 fn with(stash: Stash) -> Self {
39 Self(Arc::new(Mutex::new(stash)))
40 }
41
42 pub fn wait(&self) -> Option<String> {
44 let mut state = self.0.lock().unwrap_or_else(|e| e.into_inner());
45 let stash = match std::mem::replace(&mut *state, Stash::Done(None)) {
46 Stash::Running(handle) => handle.join().unwrap_or_default(),
48 Stash::Done(stash) => stash,
49 };
50 *state = Stash::Done(stash.clone());
51 stash
52 }
53}
54
55pub const DEFAULT_UNDO_DEPTH: usize = 50;
57
58#[derive(Debug, Clone)]
64struct HistorySnapshot {
65 store: u64,
66 len: usize,
67 stash: StashSlot,
69 new_files: Vec<PathBuf>,
71}
72
73#[derive(Debug, Clone)]
78struct Store {
79 id: u64,
80 messages: Vec<Message>,
81}
82
83#[derive(Debug, Clone)]
85pub struct SessionHistory {
86 undo_stack: VecDeque<HistorySnapshot>,
87 redo_stack: Vec<HistorySnapshot>,
88 stores: Vec<Store>,
90 next_store: u64,
91 depth: usize,
93 dropped: bool,
95}
96
97impl Default for SessionHistory {
98 fn default() -> Self {
99 Self::new()
100 }
101}
102
103impl SessionHistory {
104 pub fn new() -> Self {
106 Self::with_depth(DEFAULT_UNDO_DEPTH)
107 }
108
109 pub fn with_depth(depth: usize) -> Self {
112 Self {
113 undo_stack: VecDeque::new(),
114 redo_stack: Vec::new(),
115 stores: Vec::new(),
116 next_store: 0,
117 depth,
118 dropped: false,
119 }
120 }
121
122 pub fn clear(&mut self) {
124 *self = Self::with_depth(self.depth);
125 }
126
127 pub fn can_undo(&self) -> bool {
128 !self.undo_stack.is_empty()
129 }
130
131 pub fn can_redo(&self) -> bool {
132 !self.redo_stack.is_empty()
133 }
134
135 pub fn exhausted_depth(&self) -> Option<usize> {
138 (self.dropped && self.undo_stack.is_empty()).then_some(self.depth)
139 }
140
141 pub fn push_before_turn(&mut self, messages: &[Message], project_path: &Path) -> StashSlot {
145 let stash = StashSlot::spawn(project_path);
146 self.record_turn(messages, stash.clone());
147 stash
148 }
149
150 fn record_turn(&mut self, messages: &[Message], stash: StashSlot) {
151 self.redo_stack.clear();
152 let (store, len) = self.capture(messages);
153 self.push_undo(HistorySnapshot {
154 store,
155 len,
156 stash,
157 new_files: Vec::new(), });
159 self.gc();
160 }
161
162 pub fn mark_new_files(&mut self, files: Vec<PathBuf>) {
164 if let Some(snap) = self.undo_stack.back_mut() {
165 snap.new_files = files;
166 }
167 }
168
169 pub fn undo(
171 &mut self,
172 current_messages: &[Message],
173 project_path: &Path,
174 ) -> Option<Vec<Message>> {
175 let snap = self.undo_stack.pop_back()?;
176 let messages = self.messages_of(&snap);
178 let (store, len) = self.capture(current_messages);
180 self.redo_stack.push(HistorySnapshot {
181 store,
182 len,
183 stash: StashSlot::done(git_stash_create(project_path)),
184 new_files: Vec::new(),
185 });
186 self.gc();
187
188 if let Some(ref stash) = snap.stash.wait() {
190 let _ = Command::new("git")
191 .args(["checkout", stash, "--", "."])
192 .current_dir(project_path)
193 .output();
194 } else {
195 let _ = Command::new("git")
197 .args(["checkout", "--", "."])
198 .current_dir(project_path)
199 .output();
200 }
201
202 for f in &snap.new_files {
203 let _ = std::fs::remove_file(f);
204 }
205
206 Some(messages)
207 }
208
209 pub fn redo(
211 &mut self,
212 current_messages: &[Message],
213 project_path: &Path,
214 ) -> Option<Vec<Message>> {
215 let snap = self.redo_stack.pop()?;
216 let messages = self.messages_of(&snap);
217 let (store, len) = self.capture(current_messages);
218 self.push_undo(HistorySnapshot {
219 store,
220 len,
221 stash: StashSlot::done(git_stash_create(project_path)),
222 new_files: Vec::new(),
223 });
224 self.gc();
225
226 if let Some(ref stash) = snap.stash.wait() {
227 let _ = Command::new("git")
228 .args(["checkout", stash, "--", "."])
229 .current_dir(project_path)
230 .output();
231 }
232
233 Some(messages)
234 }
235
236 fn capture(&mut self, messages: &[Message]) -> (u64, usize) {
244 if let Some(live) = self.stores.last_mut() {
245 let needed = self
246 .undo_stack
247 .iter()
248 .chain(&self.redo_stack)
249 .filter(|snap| snap.store == live.id)
250 .map(|snap| snap.len)
251 .max()
252 .unwrap_or(0);
253 let shared = live
254 .messages
255 .iter()
256 .zip(messages)
257 .take_while(|(kept, now)| kept == now)
258 .count();
259 if shared >= needed {
260 live.messages.truncate(shared);
261 live.messages.extend_from_slice(&messages[shared..]);
262 return (live.id, messages.len());
263 }
264 }
265 self.next_store += 1;
266 self.stores.push(Store {
267 id: self.next_store,
268 messages: messages.to_vec(),
269 });
270 (self.next_store, messages.len())
271 }
272
273 fn messages_of(&self, snap: &HistorySnapshot) -> Vec<Message> {
274 self.stores
275 .iter()
276 .find(|store| store.id == snap.store)
277 .and_then(|store| store.messages.get(..snap.len))
278 .map(<[Message]>::to_vec)
279 .unwrap_or_default()
280 }
281
282 fn push_undo(&mut self, snap: HistorySnapshot) {
283 self.undo_stack.push_back(snap);
284 while self.depth > 0 && self.undo_stack.len() > self.depth {
285 self.undo_stack.pop_front();
286 self.dropped = true;
287 }
288 }
289
290 fn gc(&mut self) {
292 let live = self.stores.last().map(|store| store.id);
293 let (undo, redo) = (&self.undo_stack, &self.redo_stack);
294 self.stores.retain(|store| {
295 Some(store.id) == live || undo.iter().chain(redo).any(|snap| snap.store == store.id)
296 });
297 }
298}
299
300fn git_stash_create(project_path: &Path) -> Option<String> {
306 let output = Command::new("git")
307 .args(["stash", "create"])
308 .current_dir(project_path)
309 .output()
310 .ok()?;
311 if !output.status.success() {
312 return None;
313 }
314 let hash = String::from_utf8_lossy(&output.stdout).trim().to_string();
315 if hash.is_empty() { None } else { Some(hash) }
316}
317
318pub fn git_changed_files(project_path: &Path) -> Vec<String> {
320 let output = Command::new("git")
321 .args(["status", "--porcelain"])
322 .current_dir(project_path)
323 .output();
324 match output {
325 Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
326 .lines()
327 .filter(|l| !l.is_empty())
328 .map(|l| l.to_string())
329 .collect(),
330 _ => Vec::new(),
331 }
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337 use whycodes_core::types::{Message, MessageContent, Role};
338
339 fn msg(content: &str) -> Message {
340 Message {
341 role: Role::User,
342 content: MessageContent::text(content),
343 tool_call_id: None,
344 name: None,
345 created_at: None,
346 }
347 }
348
349 fn git_available() -> bool {
350 Command::new("git")
351 .arg("--version")
352 .output()
353 .map(|o| o.status.success())
354 .unwrap_or(false)
355 }
356
357 fn run_git(dir: &Path, args: &[&str]) -> bool {
358 Command::new("git")
359 .args(args)
360 .current_dir(dir)
361 .output()
362 .map(|o| o.status.success())
363 .unwrap_or(false)
364 }
365
366 fn texts(messages: &[Message]) -> Vec<&str> {
367 messages
368 .iter()
369 .map(|m| m.content.as_text().unwrap_or_default())
370 .collect()
371 }
372
373 fn retained_bytes(h: &SessionHistory) -> usize {
375 h.stores
376 .iter()
377 .flat_map(|store| &store.messages)
378 .map(|m| m.content.as_text().unwrap_or_default().len())
379 .sum()
380 }
381
382 fn record(h: &mut SessionHistory, messages: &[Message]) {
383 h.record_turn(messages, StashSlot::done(None));
384 }
385
386 #[test]
387 fn two_hundred_turns_retain_linear_bytes() {
388 let mut h = SessionHistory::with_depth(0);
389 let mut transcript = Vec::new();
390 for turn in 0..200 {
391 record(&mut h, &transcript);
392 transcript.push(msg(&format!("{turn:04}{}", "x".repeat(2044))));
393 }
394 let transcript_bytes: usize = transcript
395 .iter()
396 .map(|m| m.content.as_text().unwrap_or_default().len())
397 .sum();
398 assert_eq!(transcript_bytes, 200 * 2048);
399 let retained = retained_bytes(&h);
401 assert!(retained < 3 * transcript_bytes, "{retained} retained");
402 assert_eq!(h.stores.len(), 1, "appended turns share one log");
403 assert_eq!(h.undo_stack.len(), 200);
404
405 let dir = tempfile::tempdir().expect("tempdir");
406 let back = h.undo(&transcript, dir.path()).expect("undo");
407 assert_eq!(back.as_slice(), &transcript[..199]);
408 let back = h.undo(&back, dir.path()).expect("undo again");
409 assert_eq!(back.as_slice(), &transcript[..198]);
410 let forward = h.redo(&back, dir.path()).expect("redo");
411 assert_eq!(forward.as_slice(), &transcript[..199]);
412 assert!(retained_bytes(&h) < 3 * transcript_bytes);
413 }
414
415 #[test]
416 fn rewritten_prefix_keeps_both_logs_until_unreferenced() {
417 let dir = tempfile::tempdir().expect("tempdir");
418 let mut h = SessionHistory::with_depth(0);
419 record(&mut h, &[msg("a")]);
420 record(&mut h, &[msg("a"), msg("b")]);
421 record(&mut h, &[msg("summary")]);
423 assert_eq!(h.stores.len(), 2);
424
425 let now = [msg("summary"), msg("c")];
426 let back = h.undo(&now, dir.path()).expect("undo 1");
427 assert_eq!(texts(&back), ["summary"]);
428 let back = h.undo(&back, dir.path()).expect("undo 2");
429 assert_eq!(texts(&back), ["a", "b"]);
430 let back = h.undo(&back, dir.path()).expect("undo 3");
431 assert_eq!(texts(&back), ["a"]);
432 assert!(!h.can_undo());
433
434 record(&mut h, &back);
436 assert_eq!(h.stores.len(), 1);
437 h.clear();
438 assert!(h.stores.is_empty());
439 assert!(!h.can_undo() && !h.can_redo());
440 }
441
442 #[test]
443 fn new_turn_after_undo_reuses_the_log() {
444 let dir = tempfile::tempdir().expect("tempdir");
445 let mut h = SessionHistory::new();
446 record(&mut h, &[msg("a")]);
447 let back = h.undo(&[msg("a"), msg("b")], dir.path()).expect("undo");
448 assert_eq!(texts(&back), ["a"]);
449
450 record(&mut h, &back);
452 assert_eq!(h.stores.len(), 1);
453 let back = h.undo(&[msg("a"), msg("c")], dir.path()).expect("undo");
454 assert_eq!(texts(&back), ["a"]);
455 let forward = h.redo(&back, dir.path()).expect("redo");
456 assert_eq!(texts(&forward), ["a", "c"]);
457 assert_eq!(h.stores.len(), 1);
458 }
459
460 #[test]
461 fn depth_cap_drops_oldest_and_reports_exhaustion() {
462 let dir = tempfile::tempdir().expect("tempdir");
463 let mut h = SessionHistory::with_depth(2);
464 record(&mut h, &[msg("1")]);
465 record(&mut h, &[msg("1"), msg("2")]);
466 assert_eq!(h.exhausted_depth(), None);
467 record(&mut h, &[msg("1"), msg("2"), msg("3")]);
468 assert_eq!(h.undo_stack.len(), 2, "oldest entry dropped");
469 assert_eq!(h.exhausted_depth(), None, "entries remain");
470
471 let now = [msg("1"), msg("2"), msg("3"), msg("4")];
472 let back = h.undo(&now, dir.path()).expect("undo 1");
473 assert_eq!(texts(&back), ["1", "2", "3"]);
474 let back = h.undo(&back, dir.path()).expect("undo 2");
475 assert_eq!(texts(&back), ["1", "2"]);
476 assert!(h.undo(&back, dir.path()).is_none());
477 assert_eq!(h.exhausted_depth(), Some(2));
478
479 h.clear();
480 assert_eq!(h.exhausted_depth(), None, "clear forgets the drop");
481 assert_eq!(h.depth, 2, "clear keeps the cap");
482 }
483
484 #[test]
485 fn default_depth_is_fifty_and_zero_is_unlimited() {
486 assert_eq!(SessionHistory::new().depth, DEFAULT_UNDO_DEPTH);
487 assert_eq!(SessionHistory::default().depth, 50);
488 let mut h = SessionHistory::with_depth(0);
489 for n in 0..60 {
490 record(&mut h, &vec![msg("m"); n]);
491 }
492 assert_eq!(h.undo_stack.len(), 60);
493 assert_eq!(h.exhausted_depth(), None);
494 }
495
496 #[test]
497 fn new_history_has_no_undo_or_redo() {
498 let h = SessionHistory::new();
499 assert!(!h.can_undo());
500 assert!(!h.can_redo());
501 }
502
503 #[test]
504 fn push_then_undo_then_redo_round_trips_messages() {
505 let dir = tempfile::tempdir().expect("tempdir");
506 let mut h = SessionHistory::new();
507
508 h.push_before_turn(&[msg("m1")], dir.path());
510 assert!(h.can_undo());
511 assert!(!h.can_redo());
512
513 let restored = h.undo(&[msg("m1"), msg("m2")], dir.path());
515 let restored = restored.expect("undo returns a snapshot");
516 assert_eq!(restored.len(), 1);
517 assert_eq!(restored[0].content.as_text(), Some("m1"));
518 assert!(!h.can_undo(), "undo popped the only snapshot");
519 assert!(h.can_redo());
520
521 let forward = h.redo(&[msg("m1")], dir.path());
523 let forward = forward.expect("redo returns a snapshot");
524 assert_eq!(forward.len(), 2);
525 assert!(!h.can_redo());
526 assert!(h.can_undo());
527 }
528
529 #[test]
530 fn push_clears_the_redo_stack() {
531 let dir = tempfile::tempdir().expect("tempdir");
532 let mut h = SessionHistory::new();
533 h.push_before_turn(&[msg("a")], dir.path());
534 let _ = h.undo(&[msg("a"), msg("b")], dir.path());
535 assert!(h.can_redo());
536
537 h.push_before_turn(&[msg("a")], dir.path());
539 assert!(!h.can_redo(), "new turn clears redo");
540 }
541
542 #[test]
543 fn undo_and_redo_on_empty_stacks_return_none() {
544 let dir = tempfile::tempdir().expect("tempdir");
545 let mut h = SessionHistory::new();
546 assert!(h.undo(&[], dir.path()).is_none());
547 assert!(h.redo(&[], dir.path()).is_none());
548 }
549
550 #[test]
551 fn mark_new_files_removes_created_files_on_undo() {
552 let dir = tempfile::tempdir().expect("tempdir");
553 let new_file = dir.path().join("scratch.txt");
554 std::fs::write(&new_file, "data").expect("write");
555 assert!(new_file.exists());
556
557 let mut h = SessionHistory::new();
558 h.push_before_turn(&[msg("m")], dir.path());
559 h.mark_new_files(vec![new_file.clone()]);
560
561 let _ = h.undo(&[msg("m"), msg("m2")], dir.path());
562 assert!(!new_file.exists(), "undo must remove marked new files");
563 }
564
565 #[test]
566 fn git_changed_files_empty_outside_repo() {
567 let dir = tempfile::tempdir().expect("tempdir");
568 assert!(git_changed_files(dir.path()).is_empty());
570 }
571
572 #[test]
573 fn undo_with_stash_restores_tree_when_git_available() {
574 assert!(git_available(), "git is required by the session test suite");
575 let dir = tempfile::tempdir().expect("tempdir");
576 assert!(run_git(dir.path(), &["init", "-q"]), "git init");
577 let f = dir.path().join("a.txt");
578 std::fs::write(&f, "v1").expect("write v1");
579 assert!(run_git(dir.path(), &["add", "a.txt"]), "git add");
580 assert!(
581 run_git(
582 dir.path(),
583 &[
584 "-c",
585 "user.name=t",
586 "-c",
587 "user.email=t@example.com",
588 "commit",
589 "-q",
590 "-m",
591 "init",
592 ],
593 ),
594 "git commit"
595 );
596
597 std::fs::write(&f, "v2").expect("write v2");
599 let mut h = SessionHistory::new();
600 let stash = h.push_before_turn(&[msg("m")], dir.path()).wait();
602 assert!(stash.is_some(), "dirty tree yields a stash ref");
603 std::fs::write(&f, "v3").expect("write v3");
604
605 let restored = h.undo(&[msg("m"), msg("m2")], dir.path());
606 assert!(restored.is_some(), "undo proceeds even with git restore");
607 let content = std::fs::read_to_string(&f).unwrap_or_default();
608 assert_eq!(content, "v2", "undo checks out the pre-turn tree");
609 }
610
611 #[test]
612 fn stash_slot_wait_survives_a_poisoned_lock() {
613 let slot = StashSlot::done(Some("kept".into()));
614 let held = slot.clone();
615 let poisoned = std::thread::spawn(move || {
616 let _guard = held.0.lock().expect("first lock");
617 panic!("poison the slot");
618 })
619 .join();
620 assert!(poisoned.is_err());
621 assert_eq!(slot.wait().as_deref(), Some("kept"));
622 }
623
624 #[test]
625 fn stash_slot_waits_once_and_keeps_the_result() {
626 let dir = tempfile::tempdir().expect("tempdir");
627 let mut h = SessionHistory::new();
628 let slot = h.push_before_turn(&[msg("m")], dir.path());
629 assert_eq!(slot.wait(), None);
632 assert_eq!(slot.wait(), None);
633 let done = StashSlot::done(Some("abc".into()));
634 assert_eq!(done.clone().wait().as_deref(), Some("abc"));
635 assert_eq!(done.wait().as_deref(), Some("abc"));
636 }
637
638 #[test]
639 fn mark_new_files_without_snapshot_is_a_noop() {
640 let mut history = SessionHistory::new();
641 history.mark_new_files(vec![PathBuf::from("unused")]);
642 assert!(!history.can_undo());
643 }
644
645 #[test]
646 fn changed_files_reports_porcelain_lines_in_repo() {
647 assert!(git_available(), "git is required by the session test suite");
648 let dir = tempfile::tempdir().expect("tempdir");
649 assert!(run_git(dir.path(), &["init", "-q"]));
650 std::fs::write(dir.path().join("new.txt"), "new").unwrap();
651 let changed = git_changed_files(dir.path());
652 assert_eq!(changed, vec!["?? new.txt"]);
653 }
654
655 #[test]
656 fn redo_restores_its_stashed_tree_when_available() {
657 assert!(git_available(), "git is required by the session test suite");
658 let dir = tempfile::tempdir().expect("tempdir");
659 assert!(run_git(dir.path(), &["init", "-q"]));
660 let file = dir.path().join("a.txt");
661 std::fs::write(&file, "v1").unwrap();
662 assert!(run_git(dir.path(), &["add", "a.txt"]));
663 assert!(run_git(
664 dir.path(),
665 &[
666 "-c",
667 "user.name=t",
668 "-c",
669 "user.email=t@example.com",
670 "commit",
671 "-q",
672 "-m",
673 "init"
674 ]
675 ));
676
677 let mut history = SessionHistory::new();
678 history
679 .push_before_turn(&[msg("before")], dir.path())
680 .wait();
681 std::fs::write(&file, "after").unwrap();
682 let _ = history.undo(&[msg("after")], dir.path());
683 assert_eq!(std::fs::read_to_string(&file).unwrap(), "v1");
684 let restored = history.redo(&[msg("before")], dir.path()).unwrap();
685 assert_eq!(restored[0].content.as_text(), Some("after"));
686 assert_eq!(std::fs::read_to_string(&file).unwrap(), "after");
687 }
688
689 #[test]
690 fn git_helpers_handle_disappearing_working_directory() {
691 let dir = tempfile::tempdir().unwrap();
692 let path = dir.path().to_path_buf();
693 drop(dir);
694 assert!(git_stash_create(&path).is_none());
695 assert!(git_changed_files(&path).is_empty());
696
697 let mut history = SessionHistory::new();
698 history.push_before_turn(&[msg("before")], &path);
699 assert_eq!(history.undo(&[msg("after")], &path).unwrap().len(), 1);
700 assert_eq!(history.redo(&[msg("before")], &path).unwrap().len(), 1);
701 }
702
703 #[test]
704 fn stash_creation_failure_in_damaged_repo_returns_none() {
705 assert!(git_available(), "git is required by the session test suite");
706 let dir = tempfile::tempdir().unwrap();
707 assert!(run_git(dir.path(), &["init", "-q"]));
708 std::fs::write(dir.path().join("tracked"), "data").unwrap();
709 assert!(run_git(dir.path(), &["add", "tracked"]));
710 assert!(run_git(
711 dir.path(),
712 &[
713 "-c",
714 "user.name=t",
715 "-c",
716 "user.email=t@example.com",
717 "commit",
718 "-q",
719 "-m",
720 "init"
721 ]
722 ));
723 std::fs::write(dir.path().join("tracked"), "changed").unwrap();
724 std::fs::remove_file(dir.path().join(".git/index")).unwrap();
725 std::fs::create_dir(dir.path().join(".git/index")).unwrap();
726 assert!(git_stash_create(dir.path()).is_none());
727 }
728}