返回 CodeWhale
restore.rs
根目录 / crates / tui / src / commands / groups / skills / restore.rs
1 //! `/restore` slash command — roll back the workspace to a prior snapshot.
2 //!
3 //! `/restore` (no arg) lists the 20 most recent snapshots so the user can
4 //! see what's available. `/restore list [N]` lists more snapshots, capped
5 //! at 100. `/restore <N>` restores the *N*th-most-recent snapshot, where
6 //! `N=1` is the newest. Without trusted/full access we refuse to mutate files unless
7 //! the user has explicitly trusted the workspace (`/trust on` or Full Access) —
8 //! the user can always view the list, just not one-shot revert without a
9 //! safety net.
10
11 use crate::commands::CommandResult;
12 use crate::snapshot::{Snapshot, SnapshotRepo};
13 use crate::tui::app::App;
14 use chrono::TimeZone;
15
16 const DEFAULT_LIST_LIMIT: usize = 20;
17 const MAX_LIST_LIMIT: usize = 100;
18 const MAX_RESTORE_INDEX: usize = 1000;
19
20 /// Entry point for `/restore [N|list [N]]`.
21 fn restore(app: &mut App, arg: Option<&str>) -> CommandResult {
22 let workspace = app.workspace.clone();
23 let repo = match SnapshotRepo::open_or_init(&workspace) {
24 Ok(r) => r,
25 Err(e) => {
26 return CommandResult::error(format!(
27 "Snapshot repo unavailable for {}: {e}",
28 workspace.display(),
29 ));
30 }
31 };
32
33 let Some(arg) = arg.map(str::trim).filter(|s| !s.is_empty()) else {
34 let snapshots = match repo.list(DEFAULT_LIST_LIMIT) {
35 Ok(s) => s,
36 Err(e) => return CommandResult::error(format!("Failed to list snapshots: {e}")),
37 };
38 if snapshots.is_empty() {
39 return no_snapshots_message();
40 }
41 return CommandResult::message(format_listing(&snapshots));
42 };
43
44 if let Some(limit) = match parse_list_arg(arg) {
45 Ok(limit) => limit,
46 Err(message) => return CommandResult::error(message),
47 } {
48 let snapshots = match repo.list(limit) {
49 Ok(s) => s,
50 Err(e) => return CommandResult::error(format!("Failed to list snapshots: {e}")),
51 };
52 if snapshots.is_empty() {
53 return no_snapshots_message();
54 }
55 return CommandResult::message(format_listing(&snapshots));
56 }
57
58 let n: usize = match arg.parse() {
59 Ok(n) if (1..=MAX_RESTORE_INDEX).contains(&n) => n,
60 Ok(n) if n > MAX_RESTORE_INDEX => {
61 return CommandResult::error(format!(
62 "Restore index must be <= {MAX_RESTORE_INDEX}; got {n}. Use /restore list [N] to inspect snapshots first.",
63 ));
64 }
65 _ => {
66 return CommandResult::error(format!(
67 "Usage: /restore <N> or /restore list [N] (N is 1-based; got '{arg}')",
68 ));
69 }
70 };
71 let snapshots = match repo.list(n.max(DEFAULT_LIST_LIMIT)) {
72 Ok(s) => s,
73 Err(e) => return CommandResult::error(format!("Failed to list snapshots: {e}")),
74 };
75 if snapshots.is_empty() {
76 return no_snapshots_message();
77 }
78
79 if n > snapshots.len() {
80 return CommandResult::error(format!(
81 "Only {} snapshot(s) available; asked for #{n}.",
82 snapshots.len(),
83 ));
84 }
85
86 // Sessions without trusted/full access get a confirmation gate. We don't have a true
87 // modal-confirmation path inside slash commands today, so the gate
88 // is "require trust mode" — `/trust on` or Full Access. Users in plain
89 // Agent mode get a clear message explaining how to proceed.
90 if !(app.yolo || app.trust_mode) {
91 return CommandResult::message(format!(
92 "Refusing to restore snapshot #{n} ('{}') outside trusted mode.\n\
93 Run `/trust on` or select Full Access with Shift+Tab, then re-run `/restore {n}`.",
94 snapshots[n - 1].label,
95 ));
96 }
97
98 let target = &snapshots[n - 1];
99 if let Err(e) = repo.restore(&target.id) {
100 return CommandResult::error(format!("Restore failed: {e}"));
101 }
102
103 CommandResult::message(format!(
104 "Restored snapshot #{n} ('{}', {}). Workspace files have been reverted; conversation history is unchanged.",
105 target.label,
106 short_sha(target.id.as_str()),
107 ))
108 }
109
110 fn parse_list_arg(arg: &str) -> Result<Option<usize>, String> {
111 let mut parts = arg.split_whitespace();
112 let action = match parts.next() {
113 Some(action) => action,
114 None => return Ok(None),
115 };
116 if action != "list" {
117 return Ok(None);
118 }
119 let Some(value) = parts.next() else {
120 return Ok(Some(DEFAULT_LIST_LIMIT));
121 };
122 if parts.next().is_some() {
123 return Err(format!(
124 "Usage: /restore list [N] (got extra arguments in '{arg}')",
125 ));
126 }
127 match value.parse::<usize>() {
128 Ok(limit @ 1..=MAX_LIST_LIMIT) => Ok(Some(limit)),
129 Ok(limit) if limit > MAX_LIST_LIMIT => Err(format!(
130 "Restore list limit must be <= {MAX_LIST_LIMIT}; got {limit}.",
131 )),
132 _ => Err(format!(
133 "Usage: /restore list [N] (N must be >= 1; got '{value}')",
134 )),
135 }
136 }
137
138 fn no_snapshots_message() -> CommandResult {
139 CommandResult::message(
140 "No snapshots yet. Send a message to create the first pre-turn snapshot.",
141 )
142 }
143
144 fn format_listing(snapshots: &[Snapshot]) -> String {
145 let mut out = String::from(
146 "Recent snapshots (newest first; pass /restore <N> to revert; /restore list 50 shows more):\n",
147 );
148 for (i, s) in snapshots.iter().enumerate() {
149 out.push_str(&format!(
150 " #{:<2} {} {} {}\n",
151 i + 1,
152 format_snapshot_time(s.timestamp),
153 short_sha(s.id.as_str()),
154 s.label,
155 ));
156 }
157 out
158 }
159
160 fn format_snapshot_time(timestamp: i64) -> String {
161 match chrono::Utc.timestamp_opt(timestamp, 0).single() {
162 Some(dt) => dt.format("%Y-%m-%d %H:%M UTC").to_string(),
163 None => "unknown time".to_string(),
164 }
165 }
166
167 fn short_sha(sha: &str) -> &str {
168 &sha[..sha.len().min(8)]
169 }
170
171 pub(in crate::commands) const COMMAND_INFO: crate::commands::traits::CommandInfo =
172 crate::commands::traits::CommandInfo {
173 name: "restore",
174 aliases: &[],
175 usage: "/restore [N|list [N]]",
176 description_id: crate::localization::MessageId::CmdRestoreDescription,
177 };
178
179 pub(in crate::commands) struct RestoreCmd;
180
181 impl crate::commands::traits::RegisterCommand for RestoreCmd {
182 fn info() -> &'static crate::commands::traits::CommandInfo {
183 &COMMAND_INFO
184 }
185
186 fn execute(
187 app: &mut crate::tui::app::App,
188 arg: Option<&str>,
189 ) -> crate::commands::CommandResult {
190 restore(app, arg)
191 }
192 }
193
194 #[cfg(test)]
195 mod tests {
196 use super::*;
197 use crate::config::Config;
198 use crate::test_support::lock_test_env;
199 use crate::tui::app::TuiOptions;
200 use tempfile::TempDir;
201
202 fn make_app(tmp: &TempDir, yolo: bool) -> App {
203 let workspace = tmp.path().to_path_buf();
204 let options = TuiOptions {
205 skills_dir: tmp.path().join("skills"),
206 memory_path: tmp.path().join("memory.md"),
207 notes_path: tmp.path().join("notes.txt"),
208 mcp_config_path: tmp.path().join("mcp.json"),
209 yolo,
210 ..crate::test_support::test_tui_options(workspace)
211 };
212 App::new(options, &Config::default())
213 }
214
215 /// Pins HOME to a tempdir for the duration of the test under the
216 /// crate-wide env mutex.
217 struct ScopedHome {
218 prev: Option<std::ffi::OsString>,
219 _home: TempDir,
220 _guard: crate::test_support::TestEnvLock,
221 }
222 impl Drop for ScopedHome {
223 fn drop(&mut self) {
224 // SAFETY: process-wide lock still held.
225 unsafe {
226 match self.prev.take() {
227 Some(v) => std::env::set_var("HOME", v),
228 None => std::env::remove_var("HOME"),
229 }
230 }
231 }
232 }
233 fn scoped_home(_workspace: &TempDir) -> ScopedHome {
234 let guard = lock_test_env();
235 let prev = std::env::var_os("HOME");
236 let home = TempDir::new().expect("home tempdir");
237 // SAFETY: serialised by the global env lock.
238 unsafe {
239 std::env::set_var("HOME", home.path());
240 }
241 ScopedHome {
242 prev,
243 _home: home,
244 _guard: guard,
245 }
246 }
247
248 #[test]
249 fn restore_with_no_snapshots_shows_empty_message() {
250 let tmp = TempDir::new().unwrap();
251 let _home = scoped_home(&tmp);
252 let mut app = make_app(&tmp, true);
253 let result = restore(&mut app, None);
254 let msg = result.message.expect("expected message");
255 assert!(msg.contains("No snapshots"));
256 }
257
258 #[test]
259 fn restore_lists_when_no_arg_provided() {
260 let tmp = TempDir::new().unwrap();
261 let _home = scoped_home(&tmp);
262 let mut app = make_app(&tmp, true);
263 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
264 std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap();
265 repo.snapshot("pre-turn:1").unwrap();
266 std::fs::write(app.workspace.join("a.txt"), b"v2").unwrap();
267 repo.snapshot("post-turn:1").unwrap();
268
269 let result = restore(&mut app, None);
270 let msg = result.message.expect("expected message");
271 assert!(msg.contains("post-turn:1"));
272 assert!(msg.contains("pre-turn:1"));
273 assert!(msg.contains("#1"));
274 assert!(msg.contains("#2"));
275 }
276
277 #[test]
278 fn restore_lists_more_than_ten_snapshots_by_default() {
279 let tmp = TempDir::new().unwrap();
280 let _home = scoped_home(&tmp);
281 let mut app = make_app(&tmp, true);
282 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
283 for i in 0..12 {
284 std::fs::write(app.workspace.join("a.txt"), format!("v{i}")).unwrap();
285 repo.snapshot(&format!("turn:{i}")).unwrap();
286 }
287
288 let result = restore(&mut app, None);
289 let msg = result.message.expect("expected message");
290 assert!(msg.contains("#12"), "{msg}");
291 assert!(msg.contains("turn:0"), "{msg}");
292 }
293
294 #[test]
295 fn restore_listing_includes_snapshot_utc_time() {
296 let snapshots = [Snapshot {
297 id: crate::snapshot::SnapshotId("abcdef123456".to_string()),
298 label: "turn:demo".to_string(),
299 timestamp: 1_700_000_000,
300 session_id: None,
301 }];
302
303 let msg = format_listing(&snapshots);
304
305 assert!(msg.contains("2023-11-14 22:13 UTC"), "{msg}");
306 assert!(msg.contains("abcdef12"), "{msg}");
307 assert!(msg.contains("turn:demo"), "{msg}");
308 }
309
310 #[test]
311 fn restore_list_subcommand_accepts_explicit_limit() {
312 let tmp = TempDir::new().unwrap();
313 let _home = scoped_home(&tmp);
314 let mut app = make_app(&tmp, true);
315 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
316 for i in 0..15 {
317 std::fs::write(app.workspace.join("a.txt"), format!("v{i}")).unwrap();
318 repo.snapshot(&format!("turn:{i}")).unwrap();
319 }
320
321 let result = restore(&mut app, Some("list 12"));
322 let msg = result.message.expect("expected message");
323 assert!(msg.contains("#12"), "{msg}");
324 assert!(!msg.contains("#13"), "{msg}");
325 }
326
327 #[test]
328 fn restore_list_subcommand_rejects_invalid_limit() {
329 let tmp = TempDir::new().unwrap();
330 let _home = scoped_home(&tmp);
331 let mut app = make_app(&tmp, true);
332
333 let result = restore(&mut app, Some("list nope"));
334 assert!(result.is_error);
335 assert!(result.message.unwrap().contains("Usage: /restore list [N]"));
336 }
337
338 #[test]
339 fn restore_list_subcommand_rejects_limit_above_cap() {
340 let tmp = TempDir::new().unwrap();
341 let _home = scoped_home(&tmp);
342 let mut app = make_app(&tmp, true);
343
344 let result = restore(&mut app, Some("list 101"));
345 assert!(result.is_error);
346 assert!(
347 result
348 .message
349 .unwrap()
350 .contains("Restore list limit must be <= 100")
351 );
352 }
353
354 #[test]
355 fn restore_numeric_index_can_target_beyond_default_listing() {
356 let tmp = TempDir::new().unwrap();
357 let _home = scoped_home(&tmp);
358 let mut app = make_app(&tmp, true);
359 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
360 let f = app.workspace.join("a.txt");
361 for i in 0..12 {
362 std::fs::write(&f, format!("v{i}")).unwrap();
363 repo.snapshot(&format!("turn:{i}")).unwrap();
364 }
365 std::fs::write(&f, "changed").unwrap();
366
367 let result = restore(&mut app, Some("12"));
368 assert!(result.message.unwrap().contains("Restored"));
369 assert_eq!(std::fs::read_to_string(&f).unwrap(), "v0");
370 }
371
372 #[test]
373 fn restore_numeric_index_rejects_unbounded_query() {
374 let tmp = TempDir::new().unwrap();
375 let _home = scoped_home(&tmp);
376 let mut app = make_app(&tmp, true);
377
378 let result = restore(&mut app, Some("1001"));
379
380 assert!(result.is_error);
381 assert!(
382 result
383 .message
384 .unwrap()
385 .contains("Restore index must be <= 1000")
386 );
387 }
388
389 #[test]
390 fn restore_in_yolo_reverts_workspace() {
391 let tmp = TempDir::new().unwrap();
392 let _home = scoped_home(&tmp);
393 let mut app = make_app(&tmp, true);
394 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
395 let f = app.workspace.join("a.txt");
396
397 std::fs::write(&f, b"original").unwrap();
398 repo.snapshot("pre-turn:1").unwrap();
399 std::fs::write(&f, b"clobbered").unwrap();
400 repo.snapshot("post-turn:1").unwrap();
401
402 let result = restore(&mut app, Some("2"));
403 assert!(result.message.unwrap().contains("Restored"));
404 let after = std::fs::read_to_string(&f).unwrap();
405 assert_eq!(after, "original");
406 }
407
408 #[test]
409 fn restore_outside_trust_mode_refuses() {
410 let tmp = TempDir::new().unwrap();
411 let _home = scoped_home(&tmp);
412 let mut app = make_app(&tmp, false);
413 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
414 std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap();
415 repo.snapshot("pre-turn:1").unwrap();
416
417 let result = restore(&mut app, Some("1"));
418 let msg = result.message.expect("expected message");
419 assert!(msg.contains("Refusing"));
420 assert!(msg.contains("/trust on"));
421 }
422
423 #[test]
424 fn restore_invalid_index_returns_error() {
425 let tmp = TempDir::new().unwrap();
426 let _home = scoped_home(&tmp);
427 let mut app = make_app(&tmp, true);
428 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
429 std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap();
430 repo.snapshot("pre-turn:1").unwrap();
431
432 let result = restore(&mut app, Some("99"));
433 let msg = result.message.expect("expected message");
434 assert!(msg.contains("Only 1 snapshot"));
435 }
436
437 #[test]
438 fn restore_zero_index_returns_error() {
439 let tmp = TempDir::new().unwrap();
440 let _home = scoped_home(&tmp);
441 let mut app = make_app(&tmp, true);
442 // Need at least one snapshot so we exercise the parse-index
443 // branch instead of the "no snapshots" early return.
444 let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap();
445 std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap();
446 repo.snapshot("pre-turn:1").unwrap();
447
448 let result = restore(&mut app, Some("0"));
449 let msg = result.message.expect("expected message");
450 assert!(msg.contains("Usage:"));
451 }
452 }
453
453 lines RUST