| 1 | //! `/restore` slash command — roll back the workspace to a prior snapshot. |
| 2 | //! |
| 3 | //! `/restore` (no arg) lists the most recent snapshots so the user can |
| 4 | //! see what's available. `/restore <N>` restores the *N*th-most-recent |
| 5 | //! snapshot, where `N=1` is the newest. In non-YOLO mode we refuse to |
| 6 | //! mutate files unless the user has explicitly trusted the workspace |
| 7 | //! (`/trust on` or YOLO) — the user can always view the list, just not |
| 8 | //! one-shot revert without a safety net. |
| 9 | |
| 10 | use super::CommandResult; |
| 11 | use crate::snapshot::SnapshotRepo; |
| 12 | use crate::tui::app::App; |
| 13 | |
| 14 | const LIST_LIMIT: usize = 10; |
| 15 | |
| 16 | /// Entry point for `/restore [N]`. |
| 17 | pub fn restore(app: &mut App, arg: Option<&str>) -> CommandResult { |
| 18 | let workspace = app.workspace.clone(); |
| 19 | let repo = match SnapshotRepo::open_or_init(&workspace) { |
| 20 | Ok(r) => r, |
| 21 | Err(e) => { |
| 22 | return CommandResult::error(format!( |
| 23 | "Snapshot repo unavailable for {}: {e}", |
| 24 | workspace.display(), |
| 25 | )); |
| 26 | } |
| 27 | }; |
| 28 | |
| 29 | let snapshots = match repo.list(LIST_LIMIT) { |
| 30 | Ok(s) => s, |
| 31 | Err(e) => return CommandResult::error(format!("Failed to list snapshots: {e}")), |
| 32 | }; |
| 33 | |
| 34 | if snapshots.is_empty() { |
| 35 | return CommandResult::message( |
| 36 | "No snapshots yet. Send a message to create the first pre-turn snapshot.", |
| 37 | ); |
| 38 | } |
| 39 | |
| 40 | let Some(arg) = arg.map(str::trim).filter(|s| !s.is_empty()) else { |
| 41 | return CommandResult::message(format_listing(&snapshots)); |
| 42 | }; |
| 43 | |
| 44 | let n: usize = match arg.parse() { |
| 45 | Ok(n) if n >= 1 => n, |
| 46 | _ => { |
| 47 | return CommandResult::error(format!( |
| 48 | "Usage: /restore <N> (N is 1-based; got '{arg}')", |
| 49 | )); |
| 50 | } |
| 51 | }; |
| 52 | |
| 53 | if n > snapshots.len() { |
| 54 | return CommandResult::error(format!( |
| 55 | "Only {} snapshot(s) available; asked for #{n}.", |
| 56 | snapshots.len(), |
| 57 | )); |
| 58 | } |
| 59 | |
| 60 | // Non-YOLO sessions get a confirmation gate. We don't have a true |
| 61 | // modal-confirmation path inside slash commands today, so the gate |
| 62 | // is "require trust mode" — `/trust on` or YOLO. Users in plain |
| 63 | // Agent mode get a clear message explaining how to proceed. |
| 64 | if !(app.yolo || app.trust_mode) { |
| 65 | return CommandResult::message(format!( |
| 66 | "Refusing to restore snapshot #{n} ('{}') outside trusted mode.\n\ |
| 67 | Run `/trust on` or `/yolo` first, then re-run `/restore {n}`.", |
| 68 | snapshots[n - 1].label, |
| 69 | )); |
| 70 | } |
| 71 | |
| 72 | let target = &snapshots[n - 1]; |
| 73 | if let Err(e) = repo.restore(&target.id) { |
| 74 | return CommandResult::error(format!("Restore failed: {e}")); |
| 75 | } |
| 76 | |
| 77 | CommandResult::message(format!( |
| 78 | "Restored snapshot #{n} ('{}', {}). Workspace files have been reverted; conversation history is unchanged.", |
| 79 | target.label, |
| 80 | short_sha(target.id.as_str()), |
| 81 | )) |
| 82 | } |
| 83 | |
| 84 | fn format_listing(snapshots: &[crate::snapshot::Snapshot]) -> String { |
| 85 | let mut out = String::from("Recent snapshots (newest first; pass /restore <N> to revert):\n"); |
| 86 | for (i, s) in snapshots.iter().enumerate() { |
| 87 | out.push_str(&format!( |
| 88 | " #{:<2} {} {}\n", |
| 89 | i + 1, |
| 90 | short_sha(s.id.as_str()), |
| 91 | s.label, |
| 92 | )); |
| 93 | } |
| 94 | out |
| 95 | } |
| 96 | |
| 97 | fn short_sha(sha: &str) -> &str { |
| 98 | &sha[..sha.len().min(8)] |
| 99 | } |
| 100 | |
| 101 | #[cfg(test)] |
| 102 | mod tests { |
| 103 | use super::*; |
| 104 | use crate::config::Config; |
| 105 | use crate::test_support::lock_test_env; |
| 106 | use crate::tui::app::TuiOptions; |
| 107 | use std::sync::MutexGuard; |
| 108 | use tempfile::TempDir; |
| 109 | |
| 110 | fn make_app(tmp: &TempDir, yolo: bool) -> App { |
| 111 | let workspace = tmp.path().to_path_buf(); |
| 112 | let options = TuiOptions { |
| 113 | model: "deepseek-v4-pro".to_string(), |
| 114 | workspace, |
| 115 | config_path: None, |
| 116 | config_profile: None, |
| 117 | allow_shell: false, |
| 118 | use_alt_screen: true, |
| 119 | use_mouse_capture: false, |
| 120 | use_bracketed_paste: true, |
| 121 | max_subagents: 1, |
| 122 | skills_dir: tmp.path().join("skills"), |
| 123 | memory_path: tmp.path().join("memory.md"), |
| 124 | notes_path: tmp.path().join("notes.txt"), |
| 125 | mcp_config_path: tmp.path().join("mcp.json"), |
| 126 | use_memory: false, |
| 127 | start_in_agent_mode: false, |
| 128 | skip_onboarding: true, |
| 129 | yolo, |
| 130 | resume_session_id: None, |
| 131 | initial_input: None, |
| 132 | }; |
| 133 | App::new(options, &Config::default()) |
| 134 | } |
| 135 | |
| 136 | /// Pins HOME to a tempdir for the duration of the test under the |
| 137 | /// crate-wide env mutex. |
| 138 | struct ScopedHome { |
| 139 | prev: Option<std::ffi::OsString>, |
| 140 | _home: TempDir, |
| 141 | _guard: MutexGuard<'static, ()>, |
| 142 | } |
| 143 | impl Drop for ScopedHome { |
| 144 | fn drop(&mut self) { |
| 145 | // SAFETY: process-wide lock still held. |
| 146 | unsafe { |
| 147 | match self.prev.take() { |
| 148 | Some(v) => std::env::set_var("HOME", v), |
| 149 | None => std::env::remove_var("HOME"), |
| 150 | } |
| 151 | } |
| 152 | } |
| 153 | } |
| 154 | fn scoped_home(_workspace: &TempDir) -> ScopedHome { |
| 155 | let guard = lock_test_env(); |
| 156 | let prev = std::env::var_os("HOME"); |
| 157 | let home = TempDir::new().expect("home tempdir"); |
| 158 | // SAFETY: serialised by the global env lock. |
| 159 | unsafe { |
| 160 | std::env::set_var("HOME", home.path()); |
| 161 | } |
| 162 | ScopedHome { |
| 163 | prev, |
| 164 | _home: home, |
| 165 | _guard: guard, |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | #[test] |
| 170 | fn restore_with_no_snapshots_shows_empty_message() { |
| 171 | let tmp = TempDir::new().unwrap(); |
| 172 | let _home = scoped_home(&tmp); |
| 173 | let mut app = make_app(&tmp, true); |
| 174 | let result = restore(&mut app, None); |
| 175 | let msg = result.message.expect("expected message"); |
| 176 | assert!(msg.contains("No snapshots")); |
| 177 | } |
| 178 | |
| 179 | #[test] |
| 180 | fn restore_lists_when_no_arg_provided() { |
| 181 | let tmp = TempDir::new().unwrap(); |
| 182 | let _home = scoped_home(&tmp); |
| 183 | let mut app = make_app(&tmp, true); |
| 184 | let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); |
| 185 | std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); |
| 186 | repo.snapshot("pre-turn:1").unwrap(); |
| 187 | std::fs::write(app.workspace.join("a.txt"), b"v2").unwrap(); |
| 188 | repo.snapshot("post-turn:1").unwrap(); |
| 189 | |
| 190 | let result = restore(&mut app, None); |
| 191 | let msg = result.message.expect("expected message"); |
| 192 | assert!(msg.contains("post-turn:1")); |
| 193 | assert!(msg.contains("pre-turn:1")); |
| 194 | assert!(msg.contains("#1")); |
| 195 | assert!(msg.contains("#2")); |
| 196 | } |
| 197 | |
| 198 | #[test] |
| 199 | fn restore_in_yolo_reverts_workspace() { |
| 200 | let tmp = TempDir::new().unwrap(); |
| 201 | let _home = scoped_home(&tmp); |
| 202 | let mut app = make_app(&tmp, true); |
| 203 | let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); |
| 204 | let f = app.workspace.join("a.txt"); |
| 205 | |
| 206 | std::fs::write(&f, b"original").unwrap(); |
| 207 | repo.snapshot("pre-turn:1").unwrap(); |
| 208 | std::fs::write(&f, b"clobbered").unwrap(); |
| 209 | repo.snapshot("post-turn:1").unwrap(); |
| 210 | |
| 211 | let result = restore(&mut app, Some("2")); |
| 212 | assert!(result.message.unwrap().contains("Restored")); |
| 213 | let after = std::fs::read_to_string(&f).unwrap(); |
| 214 | assert_eq!(after, "original"); |
| 215 | } |
| 216 | |
| 217 | #[test] |
| 218 | fn restore_outside_trust_mode_refuses() { |
| 219 | let tmp = TempDir::new().unwrap(); |
| 220 | let _home = scoped_home(&tmp); |
| 221 | let mut app = make_app(&tmp, false); |
| 222 | let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); |
| 223 | std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); |
| 224 | repo.snapshot("pre-turn:1").unwrap(); |
| 225 | |
| 226 | let result = restore(&mut app, Some("1")); |
| 227 | let msg = result.message.expect("expected message"); |
| 228 | assert!(msg.contains("Refusing")); |
| 229 | assert!(msg.contains("/trust on")); |
| 230 | } |
| 231 | |
| 232 | #[test] |
| 233 | fn restore_invalid_index_returns_error() { |
| 234 | let tmp = TempDir::new().unwrap(); |
| 235 | let _home = scoped_home(&tmp); |
| 236 | let mut app = make_app(&tmp, true); |
| 237 | let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); |
| 238 | std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); |
| 239 | repo.snapshot("pre-turn:1").unwrap(); |
| 240 | |
| 241 | let result = restore(&mut app, Some("99")); |
| 242 | let msg = result.message.expect("expected message"); |
| 243 | assert!(msg.contains("Only 1 snapshot")); |
| 244 | } |
| 245 | |
| 246 | #[test] |
| 247 | fn restore_zero_index_returns_error() { |
| 248 | let tmp = TempDir::new().unwrap(); |
| 249 | let _home = scoped_home(&tmp); |
| 250 | let mut app = make_app(&tmp, true); |
| 251 | // Need at least one snapshot so we exercise the parse-index |
| 252 | // branch instead of the "no snapshots" early return. |
| 253 | let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); |
| 254 | std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); |
| 255 | repo.snapshot("pre-turn:1").unwrap(); |
| 256 | |
| 257 | let result = restore(&mut app, Some("0")); |
| 258 | let msg = result.message.expect("expected message"); |
| 259 | assert!(msg.contains("Usage:")); |
| 260 | } |
| 261 | } |
| 262 |