| 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 session = context.state_namespace.clone(); |
| 78 | let result = tokio::task::spawn_blocking(move || -> Result<String, String> { |
| 79 | let repo = SnapshotRepo::open_or_init(&workspace) |
| 80 | .map_err(|e| format!("Snapshot repo init failed: {e}"))?; |
| 81 | // Find pre-turn:* snapshots only — those mark the start of |
| 82 | // each turn, which is the right rollback target. We pull a |
| 83 | // generous list and filter so the model's `turn_offset` is |
| 84 | // counted in turns, not raw snapshots. |
| 85 | let snapshots = repo |
| 86 | .list((MAX_OFFSET as usize).saturating_mul(2) + 16) |
| 87 | .map_err(|e| format!("Snapshot list failed: {e}"))?; |
| 88 | let pre_turns: Vec<_> = snapshots |
| 89 | .into_iter() |
| 90 | .filter(|s| s.label.starts_with("pre-turn:")) |
| 91 | // Session ownership must be exact. Legacy snapshots are not |
| 92 | // safe automatic targets because they may belong to another |
| 93 | // conversation that used this workspace. |
| 94 | .filter(|s| s.session_id.as_deref() == Some(session.as_str())) |
| 95 | .collect(); |
| 96 | let target = pre_turns |
| 97 | .get((offset - 1) as usize) |
| 98 | .ok_or_else(|| { |
| 99 | format!( |
| 100 | "Only {} current-session pre-turn snapshot(s) exist; turn_offset={offset} is out of range.", |
| 101 | pre_turns.len(), |
| 102 | ) |
| 103 | })? |
| 104 | .clone(); |
| 105 | if repo |
| 106 | .work_tree_matches_snapshot(&target.id) |
| 107 | .map_err(|e| format!("Snapshot comparison failed: {e}"))? |
| 108 | { |
| 109 | return Err(format!( |
| 110 | "NoSnapshotForTurn: target '{}' ({}) already matches the current workspace. \ |
| 111 | Revert operates at completed turn boundaries; there is no distinct later snapshot to restore.", |
| 112 | target.label, |
| 113 | short_sha(target.id.as_str()), |
| 114 | )); |
| 115 | } |
| 116 | repo.restore(&target.id) |
| 117 | .map_err(|e| format!("Restore failed: {e}"))?; |
| 118 | Ok(format!( |
| 119 | "{label}: restored '{}' ({}). Workspace files reverted; conversation unchanged.", |
| 120 | target.label, |
| 121 | short_sha(target.id.as_str()), |
| 122 | )) |
| 123 | }) |
| 124 | .await |
| 125 | .map_err(|e| ToolError::execution_failed(format!("revert_turn join failed: {e}")))?; |
| 126 | |
| 127 | match result { |
| 128 | Ok(msg) => Ok(ToolResult::success(msg)), |
| 129 | Err(e) => Ok(ToolResult::error(e)), |
| 130 | } |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | fn short_sha(sha: &str) -> &str { |
| 135 | &sha[..sha.len().min(8)] |
| 136 | } |
| 137 | |
| 138 | #[cfg(test)] |
| 139 | mod tests { |
| 140 | use super::*; |
| 141 | use crate::test_support::lock_test_env; |
| 142 | use tempfile::tempdir; |
| 143 | |
| 144 | /// Pins HOME to a tempdir for the duration of the test under the |
| 145 | /// process-wide env mutex (`crate::test_support::lock_test_env`). |
| 146 | struct HomeGuard { |
| 147 | prev: Option<std::ffi::OsString>, |
| 148 | _lock: crate::test_support::TestEnvLock, |
| 149 | } |
| 150 | impl Drop for HomeGuard { |
| 151 | fn drop(&mut self) { |
| 152 | // SAFETY: process-wide lock still held. |
| 153 | unsafe { |
| 154 | match self.prev.take() { |
| 155 | Some(v) => std::env::set_var("HOME", v), |
| 156 | None => std::env::remove_var("HOME"), |
| 157 | } |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | fn scoped_home(home: &std::path::Path) -> HomeGuard { |
| 162 | let lock = lock_test_env(); |
| 163 | let prev = std::env::var_os("HOME"); |
| 164 | // SAFETY: serialised by the global env lock. |
| 165 | unsafe { |
| 166 | std::env::set_var("HOME", home); |
| 167 | } |
| 168 | HomeGuard { prev, _lock: lock } |
| 169 | } |
| 170 | |
| 171 | #[tokio::test] |
| 172 | async fn revert_turn_default_offset_restores_pre_turn_one() { |
| 173 | let tmp = tempdir().unwrap(); |
| 174 | let workspace = tmp.path().join("ws"); |
| 175 | std::fs::create_dir_all(&workspace).unwrap(); |
| 176 | let _guard = scoped_home(tmp.path()); |
| 177 | |
| 178 | // Setup: create pre-turn:1, post-turn:1 with file modifications. |
| 179 | let repo = SnapshotRepo::open_or_init(&workspace).unwrap(); |
| 180 | std::fs::write(workspace.join("a.txt"), b"original").unwrap(); |
| 181 | repo.snapshot_with_session("pre-turn:1", Some("workspace")) |
| 182 | .unwrap(); |
| 183 | std::fs::write(workspace.join("a.txt"), b"modified").unwrap(); |
| 184 | repo.snapshot_with_session("post-turn:1", Some("workspace")) |
| 185 | .unwrap(); |
| 186 | |
| 187 | let tool = RevertTurnTool; |
| 188 | let ctx = ToolContext::new(workspace.clone()); |
| 189 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 190 | assert!(r.success, "expected success: {r:?}"); |
| 191 | |
| 192 | let content = std::fs::read_to_string(workspace.join("a.txt")).unwrap(); |
| 193 | assert_eq!(content, "original"); |
| 194 | } |
| 195 | |
| 196 | #[tokio::test] |
| 197 | async fn revert_turn_invalid_offset_rejected() { |
| 198 | let tmp = tempdir().unwrap(); |
| 199 | let workspace = tmp.path().join("ws"); |
| 200 | std::fs::create_dir_all(&workspace).unwrap(); |
| 201 | let _guard = scoped_home(tmp.path()); |
| 202 | |
| 203 | let tool = RevertTurnTool; |
| 204 | let ctx = ToolContext::new(workspace); |
| 205 | let r = tool.execute(json!({"turn_offset": 0}), &ctx).await; |
| 206 | assert!(r.is_err()); |
| 207 | } |
| 208 | |
| 209 | #[tokio::test] |
| 210 | async fn revert_turn_rejects_snapshot_matching_current_workspace() { |
| 211 | let tmp = tempdir().unwrap(); |
| 212 | let workspace = tmp.path().join("ws"); |
| 213 | std::fs::create_dir_all(&workspace).unwrap(); |
| 214 | let _guard = scoped_home(tmp.path()); |
| 215 | |
| 216 | let repo = SnapshotRepo::open_or_init(&workspace).unwrap(); |
| 217 | std::fs::write(workspace.join("a.txt"), b"unchanged").unwrap(); |
| 218 | repo.snapshot_with_session("pre-turn:1", Some("workspace")) |
| 219 | .unwrap(); |
| 220 | |
| 221 | let tool = RevertTurnTool; |
| 222 | let ctx = ToolContext::new(workspace); |
| 223 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 224 | assert!(!r.success); |
| 225 | assert!(r.content.contains("NoSnapshotForTurn"), "{}", r.content); |
| 226 | } |
| 227 | |
| 228 | #[tokio::test] |
| 229 | async fn revert_turn_no_snapshots_returns_error_result() { |
| 230 | let tmp = tempdir().unwrap(); |
| 231 | let workspace = tmp.path().join("ws"); |
| 232 | std::fs::create_dir_all(&workspace).unwrap(); |
| 233 | let _guard = scoped_home(tmp.path()); |
| 234 | |
| 235 | let tool = RevertTurnTool; |
| 236 | let ctx = ToolContext::new(workspace); |
| 237 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 238 | assert!(!r.success); |
| 239 | assert!(r.content.contains("out of range")); |
| 240 | } |
| 241 | |
| 242 | #[tokio::test] |
| 243 | async fn revert_turn_rejects_legacy_and_foreign_session_snapshots() { |
| 244 | let tmp = tempdir().unwrap(); |
| 245 | let workspace = tmp.path().join("ws"); |
| 246 | std::fs::create_dir_all(&workspace).unwrap(); |
| 247 | let _guard = scoped_home(tmp.path()); |
| 248 | |
| 249 | let repo = SnapshotRepo::open_or_init(&workspace).unwrap(); |
| 250 | std::fs::write(workspace.join("a.txt"), b"legacy").unwrap(); |
| 251 | repo.snapshot("pre-turn:legacy").unwrap(); |
| 252 | std::fs::write(workspace.join("a.txt"), b"foreign").unwrap(); |
| 253 | repo.snapshot_with_session("pre-turn:foreign", Some("other-session")) |
| 254 | .unwrap(); |
| 255 | std::fs::write(workspace.join("a.txt"), b"current").unwrap(); |
| 256 | |
| 257 | let tool = RevertTurnTool; |
| 258 | let ctx = ToolContext::new(workspace.clone()); |
| 259 | let r = tool.execute(json!({}), &ctx).await.expect("execute"); |
| 260 | assert!(!r.success); |
| 261 | assert!( |
| 262 | r.content.contains("Only 0 current-session"), |
| 263 | "{}", |
| 264 | r.content |
| 265 | ); |
| 266 | assert_eq!( |
| 267 | std::fs::read_to_string(workspace.join("a.txt")).unwrap(), |
| 268 | "current" |
| 269 | ); |
| 270 | } |
| 271 | } |
| 272 |