From d3b39dec027579f10b941faab662fa95fe08ce89 Mon Sep 17 00:00:00 2001
From: Jackkp0t <184381499+Jackkp0t@users.noreply.github.com>
Date: Sat, 26 Sep 2026 16:48:11 +0800
Subject: [PATCH] fix(snapshot): remove work products added after rollback
 target

---
 ms_agent/utils/snapshot.py         | 10 +++++-
 tests/utils/test_snapshot_smoke.py | 52 ++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/ms_agent/utils/snapshot.py b/ms_agent/utils/snapshot.py
index f6cf854d7..9afdf6474 100644
--- a/ms_agent/utils/snapshot.py
+++ b/ms_agent/utils/snapshot.py
@@ -220,7 +220,7 @@ def list_snapshots(output_dir: str) -> list[dict]:
 
 def restore_snapshot(output_dir: str, commit_hash: str) -> tuple[bool, int]:
     """
-    Restore output_dir to the state at commit_hash.
+    Restore tracked work products and cached state at commit_hash.
 
     Returns (success, message_count) where message_count is the number of
     messages in history at snapshot time (0 if unknown).
@@ -233,6 +233,14 @@ def restore_snapshot(output_dir: str, commit_hash: str) -> tuple[bool, int]:
         _git(['checkout', commit_hash, '--', '.'],
              work_tree=output_dir,
              git_dir=git_dir)
+        # Remove later work products, but preserve later history caches so
+        # LLMAgent.rollback can still read and truncate conversation history.
+        paths = [
+            '.', ':(exclude,literal).ms_agent/memory',
+            ':(exclude,literal).memory'
+        ]
+        command = ['checkout', '--no-overlay', commit_hash, '--'] + paths
+        _git(command, work_tree=output_dir, git_dir=git_dir)
         logger.info(f'[snapshot] Restored to {commit_hash}')
         meta = _load_meta(output_dir)
         message_count = meta.get(commit_hash, {}).get('message_count', 0)
diff --git a/tests/utils/test_snapshot_smoke.py b/tests/utils/test_snapshot_smoke.py
index 7a02ec899..1c2a19322 100644
--- a/tests/utils/test_snapshot_smoke.py
+++ b/tests/utils/test_snapshot_smoke.py
@@ -168,6 +168,58 @@ def test_restore_returns_message_count(self):
             self.assertTrue(ok)
             self.assertEqual(mc, 7)
 
+    def test_restore_removes_files_added_in_later_snapshots(self):
+        with tempfile.TemporaryDirectory() as td:
+            original = os.path.join(td, 'original.txt')
+            added = os.path.join(td, 'nested', 'added.txt')
+            untracked = os.path.join(td, 'untracked.txt')
+            _write(original, 'v1')
+            h1 = take_snapshot(td, 'first state', message_count=2)
+            self.assertIsNotNone(h1)
+
+            _write(original, 'v2')
+            _write(added, 'added later')
+            h2 = take_snapshot(td, 'second state', message_count=4)
+            self.assertIsNotNone(h2)
+            _write(untracked, 'not in either snapshot')
+
+            ok, mc = restore_snapshot(td, h1)
+            self.assertTrue(ok)
+            self.assertEqual(mc, 2)
+            self.assertEqual(_read(original), 'v1')
+            self.assertFalse(os.path.exists(added))
+            self.assertEqual(_read(untracked), 'not in either snapshot')
+            self.assertEqual([s['hash'] for s in list_snapshots(td)], [h2, h1])
+
+            ok, mc = restore_snapshot(td, h2)
+            self.assertTrue(ok)
+            self.assertEqual(mc, 4)
+            self.assertEqual(_read(original), 'v2')
+            self.assertEqual(_read(added), 'added later')
+            self.assertEqual(_read(untracked), 'not in either snapshot')
+
+    def test_restore_preserves_later_history_and_restores_existing_history(self):
+        for cache_dir in ('.ms_agent/memory', '.memory'):
+            with self.subTest(cache_dir=cache_dir), tempfile.TemporaryDirectory() as td:
+                artifact = os.path.join(td, 'work.txt')
+                existing = os.path.join(td, cache_dir, 'existing.json')
+                later = os.path.join(td, cache_dir, 'later.json')
+                _write(artifact, 'v1')
+                _write(existing, 'original history')
+                h1 = take_snapshot(td, 'first state')
+                self.assertIsNotNone(h1)
+
+                _write(artifact, 'v2')
+                _write(existing, 'modified history')
+                _write(later, 'later history')
+                self.assertIsNotNone(take_snapshot(td, 'second state'))
+
+                ok, _ = restore_snapshot(td, h1)
+                self.assertTrue(ok)
+                self.assertEqual(_read(artifact), 'v1')
+                self.assertEqual(_read(existing), 'original history')
+                self.assertEqual(_read(later), 'later history')
+
     def test_restore_deleted_file(self):
         """A file deleted after snapshot is recreated on restore."""
         with tempfile.TemporaryDirectory() as td: