From fb12e331ab35935ac030d3262c2d08645d057f8d Mon Sep 17 00:00:00 2001 From: Hunter Bown Date: Tue, 28 Apr 2026 00:31:47 -0500 Subject: [PATCH] feat(snapshot): #137 add /restore command and revert_turn tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two user-facing entry points to the snapshot side-repo: - `/restore [N]` (slash command) — `/restore` with no arg lists the 10 most recent snapshots so the user can see what's available. `/restore N` restores the N-th most recent snapshot. Outside YOLO or `/trust on`, the command refuses to mutate files and tells the user how to opt in (no in-flow modal-confirm path inside slash commands today; trust mode is the explicit gate). - `revert_turn` (agent-callable tool) — `turn_offset` (default 1) counts in `pre-turn:*` snapshots, so the model can say "undo my last edit" without having to enumerate the history. Approval-gated (`ApprovalRequirement::Required`) since it mutates the workspace, and registered through `with_full_agent_surface` so children inherit it just like every other agent-mode tool. Tests for both surfaces use the process-wide env mutex (`crate::test_support::lock_test_env`) plus an RAII `HOME` guard so tempdir-based snapshot resolution stays inside the per-test sandbox even when the runner threads multiple tests in parallel. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/tui/src/commands/mod.rs | 8 + crates/tui/src/commands/restore.rs | 255 ++++++++++++++++++++++++++++ crates/tui/src/tools/mod.rs | 1 + crates/tui/src/tools/registry.rs | 11 ++ crates/tui/src/tools/revert_turn.rs | 205 ++++++++++++++++++++++ 5 files changed, 480 insertions(+) create mode 100644 crates/tui/src/commands/restore.rs create mode 100644 crates/tui/src/tools/revert_turn.rs diff --git a/crates/tui/src/commands/mod.rs b/crates/tui/src/commands/mod.rs index 8209aa26..1ad2d751 100644 --- a/crates/tui/src/commands/mod.rs +++ b/crates/tui/src/commands/mod.rs @@ -12,6 +12,7 @@ mod init; mod note; mod provider; mod queue; +mod restore; mod review; mod session; mod skills; @@ -337,6 +338,12 @@ pub const COMMANDS: &[CommandInfo] = &[ description: "Run a structured code review on a file, diff, or PR", usage: "/review ", }, + CommandInfo { + name: "restore", + aliases: &[], + description: "Roll back the workspace to a prior pre/post-turn snapshot. With no arg, lists recent snapshots.", + usage: "/restore [N]", + }, // RLM command CommandInfo { name: "rlm", @@ -412,6 +419,7 @@ pub fn execute(cmd: &str, app: &mut App) -> CommandResult { "skills" => skills::list_skills(app), "skill" => skills::run_skill(app, arg), "review" => review::review(app, arg), + "restore" => restore::restore(app, arg), // RLM command "rlm" | "recursive" => rlm(app, arg), diff --git a/crates/tui/src/commands/restore.rs b/crates/tui/src/commands/restore.rs new file mode 100644 index 00000000..2395dd43 --- /dev/null +++ b/crates/tui/src/commands/restore.rs @@ -0,0 +1,255 @@ +//! `/restore` slash command — roll back the workspace to a prior snapshot. +//! +//! `/restore` (no arg) lists the most recent snapshots so the user can +//! see what's available. `/restore ` restores the *N*th-most-recent +//! snapshot, where `N=1` is the newest. In non-YOLO mode we refuse to +//! mutate files unless the user has explicitly trusted the workspace +//! (`/trust on` or YOLO) — the user can always view the list, just not +//! one-shot revert without a safety net. + +use super::CommandResult; +use crate::snapshot::SnapshotRepo; +use crate::tui::app::App; + +const LIST_LIMIT: usize = 10; + +/// Entry point for `/restore [N]`. +pub fn restore(app: &mut App, arg: Option<&str>) -> CommandResult { + let workspace = app.workspace.clone(); + let repo = match SnapshotRepo::open_or_init(&workspace) { + Ok(r) => r, + Err(e) => { + return CommandResult::error(format!( + "Snapshot repo unavailable for {}: {e}", + workspace.display(), + )); + } + }; + + let snapshots = match repo.list(LIST_LIMIT) { + Ok(s) => s, + Err(e) => return CommandResult::error(format!("Failed to list snapshots: {e}")), + }; + + if snapshots.is_empty() { + return CommandResult::message( + "No snapshots yet. Send a message to create the first pre-turn snapshot.", + ); + } + + let Some(arg) = arg.map(str::trim).filter(|s| !s.is_empty()) else { + return CommandResult::message(format_listing(&snapshots)); + }; + + let n: usize = match arg.parse() { + Ok(n) if n >= 1 => n, + _ => { + return CommandResult::error(format!( + "Usage: /restore (N is 1-based; got '{arg}')", + )); + } + }; + + if n > snapshots.len() { + return CommandResult::error(format!( + "Only {} snapshot(s) available; asked for #{n}.", + snapshots.len(), + )); + } + + // Non-YOLO sessions get a confirmation gate. We don't have a true + // modal-confirmation path inside slash commands today, so the gate + // is "require trust mode" — `/trust on` or YOLO. Users in plain + // Agent mode get a clear message explaining how to proceed. + if !(app.yolo || app.trust_mode) { + return CommandResult::message(format!( + "Refusing to restore snapshot #{n} ('{}') outside trusted mode.\n\ + Run `/trust on` or `/yolo` first, then re-run `/restore {n}`.", + snapshots[n - 1].label, + )); + } + + let target = &snapshots[n - 1]; + if let Err(e) = repo.restore(&target.id) { + return CommandResult::error(format!("Restore failed: {e}")); + } + + CommandResult::message(format!( + "Restored snapshot #{n} ('{}', {}). Workspace files have been reverted; conversation history is unchanged.", + target.label, + short_sha(target.id.as_str()), + )) +} + +fn format_listing(snapshots: &[crate::snapshot::Snapshot]) -> String { + let mut out = String::from("Recent snapshots (newest first; pass /restore to revert):\n"); + for (i, s) in snapshots.iter().enumerate() { + out.push_str(&format!( + " #{:<2} {} {}\n", + i + 1, + short_sha(s.id.as_str()), + s.label, + )); + } + out +} + +fn short_sha(sha: &str) -> &str { + &sha[..sha.len().min(8)] +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use crate::test_support::lock_test_env; + use crate::tui::app::TuiOptions; + use std::sync::MutexGuard; + use tempfile::TempDir; + + fn make_app(tmp: &TempDir, yolo: bool) -> App { + let workspace = tmp.path().to_path_buf(); + let options = TuiOptions { + model: "deepseek-v4-pro".to_string(), + workspace, + allow_shell: false, + use_alt_screen: true, + use_mouse_capture: false, + use_bracketed_paste: true, + max_subagents: 1, + skills_dir: tmp.path().join("skills"), + memory_path: tmp.path().join("memory.md"), + notes_path: tmp.path().join("notes.txt"), + mcp_config_path: tmp.path().join("mcp.json"), + use_memory: false, + start_in_agent_mode: false, + skip_onboarding: true, + yolo, + resume_session_id: None, + }; + App::new(options, &Config::default()) + } + + /// Pins HOME to a tempdir for the duration of the test under the + /// crate-wide env mutex. + struct ScopedHome { + prev: Option, + _guard: MutexGuard<'static, ()>, + } + impl Drop for ScopedHome { + fn drop(&mut self) { + // SAFETY: process-wide lock still held. + unsafe { + match self.prev.take() { + Some(v) => std::env::set_var("HOME", v), + None => std::env::remove_var("HOME"), + } + } + } + } + fn scoped_home(tmp: &TempDir) -> ScopedHome { + let guard = lock_test_env(); + let prev = std::env::var_os("HOME"); + // SAFETY: serialised by the global env lock. + unsafe { + std::env::set_var("HOME", tmp.path()); + } + ScopedHome { + prev, + _guard: guard, + } + } + + #[test] + fn restore_with_no_snapshots_shows_empty_message() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, true); + let result = restore(&mut app, None); + let msg = result.message.expect("expected message"); + assert!(msg.contains("No snapshots")); + } + + #[test] + fn restore_lists_when_no_arg_provided() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, true); + let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); + std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + std::fs::write(app.workspace.join("a.txt"), b"v2").unwrap(); + repo.snapshot("post-turn:1").unwrap(); + + let result = restore(&mut app, None); + let msg = result.message.expect("expected message"); + assert!(msg.contains("post-turn:1")); + assert!(msg.contains("pre-turn:1")); + assert!(msg.contains("#1")); + assert!(msg.contains("#2")); + } + + #[test] + fn restore_in_yolo_reverts_workspace() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, true); + let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); + let f = app.workspace.join("a.txt"); + + std::fs::write(&f, b"original").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + std::fs::write(&f, b"clobbered").unwrap(); + repo.snapshot("post-turn:1").unwrap(); + + let result = restore(&mut app, Some("2")); + assert!(result.message.unwrap().contains("Restored")); + let after = std::fs::read_to_string(&f).unwrap(); + assert_eq!(after, "original"); + } + + #[test] + fn restore_outside_trust_mode_refuses() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, false); + let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); + std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + + let result = restore(&mut app, Some("1")); + let msg = result.message.expect("expected message"); + assert!(msg.contains("Refusing")); + assert!(msg.contains("/trust on")); + } + + #[test] + fn restore_invalid_index_returns_error() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, true); + let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); + std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + + let result = restore(&mut app, Some("99")); + let msg = result.message.expect("expected message"); + assert!(msg.contains("Only 1 snapshot")); + } + + #[test] + fn restore_zero_index_returns_error() { + let tmp = TempDir::new().unwrap(); + let _home = scoped_home(&tmp); + let mut app = make_app(&tmp, true); + // Need at least one snapshot so we exercise the parse-index + // branch instead of the "no snapshots" early return. + let repo = SnapshotRepo::open_or_init(&app.workspace).unwrap(); + std::fs::write(app.workspace.join("a.txt"), b"v1").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + + let result = restore(&mut app, Some("0")); + let msg = result.message.expect("expected message"); + assert!(msg.contains("Usage:")); + } +} diff --git a/crates/tui/src/tools/mod.rs b/crates/tui/src/tools/mod.rs index e080f0a5..075963d8 100644 --- a/crates/tui/src/tools/mod.rs +++ b/crates/tui/src/tools/mod.rs @@ -15,6 +15,7 @@ pub mod plan; pub mod project; pub mod recall_archive; pub mod registry; +pub mod revert_turn; pub mod review; pub mod rlm; pub mod search; diff --git a/crates/tui/src/tools/registry.rs b/crates/tui/src/tools/registry.rs index d6427348..5c5b3324 100644 --- a/crates/tui/src/tools/registry.rs +++ b/crates/tui/src/tools/registry.rs @@ -381,6 +381,16 @@ impl ToolRegistryBuilder { self.with_tool(Arc::new(ApplyPatchTool)) } + /// Include the `revert_turn` tool. Approval-gated since it mutates + /// the workspace; the model uses it when the user asks to "undo my + /// last edit". Backed by the per-workspace snapshot side-repo + /// (`crate::snapshot`). + #[must_use] + pub fn with_revert_turn_tool(self) -> Self { + use super::revert_turn::RevertTurnTool; + self.with_tool(Arc::new(RevertTurnTool)) + } + /// Include the RLM tool (`rlm`). Runs the full recursive language-model /// loop on a long input (file or inline content); the long input never /// enters the calling model's context window. The Python REPL exposes @@ -495,6 +505,7 @@ impl ToolRegistryBuilder { .with_review_tool(client.clone(), model.clone()) .with_rlm_tool(client, model) .with_recall_archive_tool() + .with_revert_turn_tool() .with_subagent_tools(manager, runtime) } diff --git a/crates/tui/src/tools/revert_turn.rs b/crates/tui/src/tools/revert_turn.rs new file mode 100644 index 00000000..9234d1cd --- /dev/null +++ b/crates/tui/src/tools/revert_turn.rs @@ -0,0 +1,205 @@ +//! `revert_turn` — agent-callable tool that rewinds the workspace to a +//! prior pre-turn snapshot. +//! +//! The model invokes this when the user says something like "undo the +//! last edit" or "roll back". It mirrors `/restore` but speaks JSON and +//! takes a turn-offset (default 1 = previous turn) instead of a list +//! index, so the model doesn't have to count entries. +//! +//! Approval is `Required` because this mutates the workspace. + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use super::spec::{ + ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, optional_u64, +}; +use crate::snapshot::SnapshotRepo; + +/// Default offset: revert the most-recent turn (i.e. the last `pre-turn:*` +/// snapshot in history). +const DEFAULT_OFFSET: u64 = 1; +/// Hard cap so the model can't ask to roll back to the dawn of time. +const MAX_OFFSET: u64 = 50; + +pub struct RevertTurnTool; + +#[async_trait] +impl ToolSpec for RevertTurnTool { + fn name(&self) -> &str { + "revert_turn" + } + + fn description(&self) -> &str { + "Roll back the workspace files to the snapshot taken before a recent turn. \ + Use when the user explicitly asks to undo, revert, or roll back the most recent edits. \ + `turn_offset` is 1-based: 1 reverts the most recent turn, 2 reverts the previous one, \ + and so on (max 50). Conversation history is NOT modified — only working-tree files are \ + restored from the side-git snapshot repo." + } + + fn input_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "turn_offset": { + "type": "integer", + "minimum": 1, + "maximum": MAX_OFFSET, + "description": "How many turns back to revert (default 1)." + } + }, + "additionalProperties": false + }) + } + + fn capabilities(&self) -> Vec { + vec![ + ToolCapability::WritesFiles, + ToolCapability::RequiresApproval, + ] + } + + fn approval_requirement(&self) -> ApprovalRequirement { + ApprovalRequirement::Required + } + + async fn execute(&self, input: Value, context: &ToolContext) -> Result { + let offset = optional_u64(&input, "turn_offset", DEFAULT_OFFSET); + if offset == 0 || offset > MAX_OFFSET { + return Err(ToolError::invalid_input(format!( + "turn_offset must be between 1 and {MAX_OFFSET}; got {offset}", + ))); + } + + let workspace = context.workspace.clone(); + let label = format!("revert_turn(offset={offset})"); + let result = tokio::task::spawn_blocking(move || -> Result { + let repo = SnapshotRepo::open_or_init(&workspace) + .map_err(|e| format!("Snapshot repo init failed: {e}"))?; + // Find pre-turn:* snapshots only — those mark the start of + // each turn, which is the right rollback target. We pull a + // generous list and filter so the model's `turn_offset` is + // counted in turns, not raw snapshots. + let snapshots = repo + .list((MAX_OFFSET as usize).saturating_mul(2) + 16) + .map_err(|e| format!("Snapshot list failed: {e}"))?; + let pre_turns: Vec<_> = snapshots + .into_iter() + .filter(|s| s.label.starts_with("pre-turn:")) + .collect(); + let target = pre_turns + .get((offset - 1) as usize) + .ok_or_else(|| { + format!( + "Only {} pre-turn snapshot(s) exist; turn_offset={offset} is out of range.", + pre_turns.len(), + ) + })? + .clone(); + repo.restore(&target.id) + .map_err(|e| format!("Restore failed: {e}"))?; + Ok(format!( + "{label}: restored '{}' ({}). Workspace files reverted; conversation unchanged.", + target.label, + short_sha(target.id.as_str()), + )) + }) + .await + .map_err(|e| ToolError::execution_failed(format!("revert_turn join failed: {e}")))?; + + match result { + Ok(msg) => Ok(ToolResult::success(msg)), + Err(e) => Ok(ToolResult::error(e)), + } + } +} + +fn short_sha(sha: &str) -> &str { + &sha[..sha.len().min(8)] +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::lock_test_env; + use std::sync::MutexGuard; + use tempfile::tempdir; + + /// Pins HOME to a tempdir for the duration of the test under the + /// process-wide env mutex (`crate::test_support::lock_test_env`). + struct HomeGuard { + prev: Option, + _lock: MutexGuard<'static, ()>, + } + impl Drop for HomeGuard { + fn drop(&mut self) { + // SAFETY: process-wide lock still held. + unsafe { + match self.prev.take() { + Some(v) => std::env::set_var("HOME", v), + None => std::env::remove_var("HOME"), + } + } + } + } + fn scoped_home(home: &std::path::Path) -> HomeGuard { + let lock = lock_test_env(); + let prev = std::env::var_os("HOME"); + // SAFETY: serialised by the global env lock. + unsafe { + std::env::set_var("HOME", home); + } + HomeGuard { prev, _lock: lock } + } + + #[tokio::test] + async fn revert_turn_default_offset_restores_pre_turn_one() { + let tmp = tempdir().unwrap(); + let workspace = tmp.path().join("ws"); + std::fs::create_dir_all(&workspace).unwrap(); + let _guard = scoped_home(tmp.path()); + + // Setup: create pre-turn:1, post-turn:1 with file modifications. + let repo = SnapshotRepo::open_or_init(&workspace).unwrap(); + std::fs::write(workspace.join("a.txt"), b"original").unwrap(); + repo.snapshot("pre-turn:1").unwrap(); + std::fs::write(workspace.join("a.txt"), b"modified").unwrap(); + repo.snapshot("post-turn:1").unwrap(); + + let tool = RevertTurnTool; + let ctx = ToolContext::new(workspace.clone()); + let r = tool.execute(json!({}), &ctx).await.expect("execute"); + assert!(r.success, "expected success: {r:?}"); + + let content = std::fs::read_to_string(workspace.join("a.txt")).unwrap(); + assert_eq!(content, "original"); + } + + #[tokio::test] + async fn revert_turn_invalid_offset_rejected() { + let tmp = tempdir().unwrap(); + let workspace = tmp.path().join("ws"); + std::fs::create_dir_all(&workspace).unwrap(); + let _guard = scoped_home(tmp.path()); + + let tool = RevertTurnTool; + let ctx = ToolContext::new(workspace); + let r = tool.execute(json!({"turn_offset": 0}), &ctx).await; + assert!(r.is_err()); + } + + #[tokio::test] + async fn revert_turn_no_snapshots_returns_error_result() { + let tmp = tempdir().unwrap(); + let workspace = tmp.path().join("ws"); + std::fs::create_dir_all(&workspace).unwrap(); + let _guard = scoped_home(tmp.path()); + + let tool = RevertTurnTool; + let ctx = ToolContext::new(workspace); + let r = tool.execute(json!({}), &ctx).await.expect("execute"); + assert!(!r.success); + assert!(r.content.contains("out of range")); + } +}