| 1 | //! `revert_turn` — agent-callable tool that rewinds the workspace to a |
| 2 | //! prior pre-turn snapshot. |
| 3 | //! |
| 4 | //! The model invokes this when the user says something like "undo the |
| 5 | //! last edit" or "roll back". It mirrors `/restore` but speaks JSON and |
| 6 | //! takes a turn-offset (default 1 = previous turn) instead of a list |
| 7 | //! index, so the model doesn't have to count entries. |
| 8 | //! |
| 9 | //! Approval is `Required` because this mutates the workspace. |
| 10 | |
| 11 | use async_trait::async_trait; |
| 12 | use serde_json::{Value, json}; |
| 13 | |
| 14 | use super::spec::{ |
| 15 | ApprovalRequirement, ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, optional_u64, |
| 16 | }; |
| 17 | use crate::snapshot::SnapshotRepo; |
| 18 | |
| 19 | /// Default offset: revert the most-recent turn (i.e. the last `pre-turn:*` |
| 20 | /// snapshot in history). |
| 21 | const DEFAULT_OFFSET: u64 = 1; |
| 22 | /// Hard cap so the model can't ask to roll back to the dawn of time. |
| 23 | const MAX_OFFSET: u64 = 50; |
| 24 | |
| 25 | pub struct RevertTurnTool; |
| 26 | |
| 27 | #[async_trait] |
| 28 | impl ToolSpec for RevertTurnTool { |
| 29 | fn name(&self) -> &str { |
| 30 | "revert_turn" |
| 31 | } |
| 32 | |
| 33 | fn description(&self) -> &str { |
| 34 | "Roll back the workspace files to the snapshot taken before a recent turn. \ |
| 35 | Use when the user explicitly asks to undo, revert, or roll back the most recent edits. \ |
| 36 | `turn_offset` is 1-based: 1 reverts the most recent turn, 2 reverts the previous one, \ |
| 37 | and so on (max 50). Conversation history is NOT modified — only working-tree files are \ |
| 38 | restored from the side-git snapshot repo." |
| 39 | } |
| 40 | |
| 41 | fn input_schema(&self) -> Value { |
| 42 | json!({ |
| 43 | "type": "object", |
| 44 | "properties": { |
| 45 | "turn_offset": { |
| 46 | "type": "integer", |
| 47 | "minimum": 1, |
| 48 | "maximum": MAX_OFFSET, |
| 49 | "description": "How many turns back to revert (default 1)." |
| 50 | } |
| 51 | }, |
| 52 | "additionalProperties": false |
| 53 | }) |
| 54 | } |
| 55 | |
| 56 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 57 | vec![ |
| 58 | ToolCapability::WritesFiles, |
| 59 | ToolCapability::RequiresApproval, |
| 60 | ] |
| 61 | } |
| 62 | |
| 63 | fn approval_requirement(&self) -> ApprovalRequirement { |
| 64 | ApprovalRequirement::Required |
| 65 | } |
| 66 | |
| 67 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 68 | let offset = optional_u64(&input, "turn_offset", DEFAULT_OFFSET); |
| 69 | if offset == 0 || offset > MAX_OFFSET { |
| 70 | return Err(ToolError::invalid_input(format!( |
| 71 | "turn_offset must be between 1 and {MAX_OFFSET}; got {offset}", |
| 72 | ))); |
| 73 | } |
| 74 | |
| 75 | let workspace = context.workspace.clone(); |
| 76 | let label = format!("revert_turn(offset={offset})"); |
| 77 | let result = tokio::task::spawn_blocking(move || -> Result<String, String> { |
| 78 | let repo = SnapshotRepo::open_or_init(&workspace) |
| 79 | .map_err(|e| format!("Snapshot repo init failed: {e}"))?; |
| 80 | // Find pre-turn:* snapshots only — those mark the start of |
| 81 | // each turn, which is the right rollback target. We pull a |
| 82 | // generous list and filter so the model's `turn_offset` is |
| 83 | // counted in turns, not raw snapshots. |
| 84 | let snapshots = repo |
| 85 | .list((MAX_OFFSET as usize).saturating_mul(2) + 16) |
| 86 | .map_err(|e| format!("Snapshot list failed: {e}"))?; |
| 87 | let pre_turns: Vec<_> = snapshots |
| 88 | .into_iter() |
| 89 | .filter(|s| s.label.starts_with("pre-turn:")) |
| 90 | .collect(); |
| 91 | let target = pre_turns |
| 92 | .get((offset - 1) as usize) |
| 93 | .ok_or_else(|| { |
| 94 | format!( |
| 95 | "Only {} pre-turn snapshot(s) exist; turn_offset={offset} is out of range.", |
| 96 | pre_turns.len(), |
| 97 | ) |
| 98 | })? |
| 99 | .clone(); |
| 100 | repo.restore(&target.id) |
| 101 | .map_err(|e| format!("Restore failed: {e}"))?; |
| 102 | Ok(format!( |
| 103 | "{label}: restored '{}' ({}). Workspace files reverted; conversation unchanged.", |
| 104 | target.label, |
| 105 | short_sha(target.id.as_str()), |
| 106 | )) |
| 107 | }) |
| 108 | .await |
| 109 | .map_err(|e| ToolError::execution_failed(format!("revert_turn join failed: {e}")))?; |
| 110 | |
| 111 | match result { |
| 112 | Ok(msg) => Ok(ToolResult::success(msg)), |
| 113 | Err(e) => Ok(ToolResult::error(e)), |
| 114 | } |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | fn short_sha(sha: &str) -> &str { |
| 119 | &sha[..sha.len().min(8)] |
| 120 | } |
| 121 | |
| 122 | #[cfg(test)] |
| 123 | mod tests { |
| 124 | use super::*; |
| 125 | use crate::test_support::lock_test_env; |
| 126 | use std::sync::MutexGuard; |
| 127 | use tempfile::tempdir; |
| 128 | |
| 129 | /// Pins HOME to a tempdir for the duration of the test under the |
| 130 | /// process-wide env mutex (`crate::test_support::lock_test_env`). |
| 131 | struct HomeGuard { |
| 132 | prev: Option<std::ffi::OsString>, |
| 133 | _lock: MutexGuard<'static, ()>, |
| 134 | } |
| 135 | impl Drop for HomeGuard { |
| 136 | fn drop(&mut self) { |
| 137 | // SAFETY: process-wide lock still held. |
| 138 | unsafe { |
| 139 | match self.prev.take() { |
| 140 | Some(v) => std::env::set_var("HOME", v), |
| 141 | None => std::env::remove_var("HOME"), |
| 142 | } |
| 143 | } |
| 144 | } |
| 145 | } |
| 146 | fn scoped_home(home: &std::path::Path) -> HomeGuard { |
| 147 | let lock = lock_test_env(); |
| 148 | let prev = std::env::var_os("HOME"); |
| 149 | // SAFETY: serialised by the global env lock. |
| 150 | unsafe { |
| 151 | std::env::set_var("HOME", home); |
| 152 | } |
| 153 | HomeGuard { prev, _lock: lock } |
| 154 | } |
| 155 | |
| 156 | #[tokio::test] |
| 157 | async fn revert_turn_default_offset_restores_pre_turn_one() { |
| 158 | let tmp = tempdir().unwrap(); |
| 159 | let workspace = tmp.path().join("ws"); |
| 160 | std::fs::create_dir_all(&workspace).unwrap(); |
| 161 | let _guard = scoped_home(tmp.path()); |
| 162 | |
| 163 | // Setup: create pre-turn:1, post-turn:1 with file modifications. |
| 164 | let repo = SnapshotRepo::open_or_init(&workspace).unwrap(); |
| 165 | std::fs::write(workspace.join("a.txt"), b"original").unwrap(); |
| 166 | repo.snapshot("pre-turn:1").unwrap(); |
| 167 | std::fs::write(workspace.join("a.txt"), b"modified").unwrap(); |
| 168 | repo.snapshot("post-turn:1").unwrap(); |
| 169 | |
| 170 | let tool = RevertTurnTool; |
| 171 | let ctx = ToolContext::new(workspace.clone()); |
| 172 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 173 | assert!(r.success, "expected success: {r:?}"); |
| 174 | |
| 175 | let content = std::fs::read_to_string(workspace.join("a.txt")).unwrap(); |
| 176 | assert_eq!(content, "original"); |
| 177 | } |
| 178 | |
| 179 | #[tokio::test] |
| 180 | async fn revert_turn_invalid_offset_rejected() { |
| 181 | let tmp = tempdir().unwrap(); |
| 182 | let workspace = tmp.path().join("ws"); |
| 183 | std::fs::create_dir_all(&workspace).unwrap(); |
| 184 | let _guard = scoped_home(tmp.path()); |
| 185 | |
| 186 | let tool = RevertTurnTool; |
| 187 | let ctx = ToolContext::new(workspace); |
| 188 | let r = tool.execute(json!({"turn_offset": 0}), &ctx).await; |
| 189 | assert!(r.is_err()); |
| 190 | } |
| 191 | |
| 192 | #[tokio::test] |
| 193 | async fn revert_turn_no_snapshots_returns_error_result() { |
| 194 | let tmp = tempdir().unwrap(); |
| 195 | let workspace = tmp.path().join("ws"); |
| 196 | std::fs::create_dir_all(&workspace).unwrap(); |
| 197 | let _guard = scoped_home(tmp.path()); |
| 198 | |
| 199 | let tool = RevertTurnTool; |
| 200 | let ctx = ToolContext::new(workspace); |
| 201 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 202 | assert!(!r.success); |
| 203 | assert!(r.content.contains("out of range")); |
| 204 | } |
| 205 | } |
| 206 |