1use crate::{AgentError, Result};
7use forge_core::Forge;
8use std::sync::Arc;
9use tokio::fs;
10
11#[derive(Clone)]
15pub struct Mutator {
16 forge: Arc<Forge>,
18 transaction: Option<Transaction>,
20}
21
22#[derive(Clone, Debug)]
24struct Transaction {
25 applied_steps: Vec<String>,
27 rollback_state: Vec<RollbackState>,
29}
30
31#[derive(Clone, Debug)]
33struct RollbackState {
34 file: String,
36 original_content: String,
38}
39
40impl Mutator {
41 pub fn new(forge: Forge) -> Self {
43 Self {
44 forge: Arc::new(forge),
45 transaction: None,
46 }
47 }
48
49 pub async fn begin_transaction(&mut self) -> Result<()> {
51 if self.transaction.is_some() {
52 return Err(AgentError::MutationFailed(
53 "Transaction already in progress".to_string(),
54 ));
55 }
56
57 self.transaction = Some(Transaction {
58 applied_steps: Vec::new(),
59 rollback_state: Vec::new(),
60 });
61
62 Ok(())
63 }
64
65 pub async fn apply_step(&mut self, step: &crate::planner::PlanStep) -> Result<()> {
67 let transaction = self
68 .transaction
69 .as_mut()
70 .ok_or_else(|| AgentError::MutationFailed("No active transaction".to_string()))?;
71
72 match &step.operation {
73 crate::planner::PlanOperation::Rename { old, new } => {
74 transaction
76 .applied_steps
77 .push(format!("Rename {} to {}", old, new));
78 }
79 crate::planner::PlanOperation::Delete { name } => {
80 transaction.applied_steps.push(format!("Delete {}", name));
81 }
82 crate::planner::PlanOperation::Create { path, content } => {
83 if let Ok(original_content) = fs::read_to_string(path).await {
85 transaction.rollback_state.push(RollbackState {
86 file: path.clone(),
87 original_content,
88 });
89 }
90
91 fs::write(path, content).await.map_err(|e| {
93 AgentError::MutationFailed(format!("Failed to write {}: {}", path, e))
94 })?;
95
96 transaction.applied_steps.push(format!("Create {}", path));
97 }
98 crate::planner::PlanOperation::Inspect { .. } => {
99 }
101 crate::planner::PlanOperation::Modify { file, .. } => {
102 if let Ok(original_content) = std::fs::read_to_string(file) {
103 transaction.rollback_state.push(RollbackState {
104 file: file.clone(),
105 original_content,
106 });
107 }
108 transaction.applied_steps.push(format!("Modify {}", file));
109 }
110 }
111
112 Ok(())
113 }
114
115 pub async fn rollback(&mut self) -> Result<()> {
117 let transaction = self
118 .transaction
119 .take()
120 .ok_or_else(|| AgentError::MutationFailed("No active transaction".to_string()))?;
121
122 for state in transaction.rollback_state.iter().rev() {
124 std::fs::write(&state.file, &state.original_content).map_err(|e| {
125 AgentError::MutationFailed(format!("Rollback failed for {}: {}", state.file, e))
126 })?;
127 }
128
129 Ok(())
130 }
131
132 pub async fn preview(&self, steps: &[crate::planner::PlanStep]) -> Result<Vec<String>> {
134 let mut previews = Vec::new();
135
136 for step in steps {
137 match &step.operation {
138 crate::planner::PlanOperation::Create { path, content } => {
139 previews.push(format!("Create {}:\n{}", path, content));
140 }
141 crate::planner::PlanOperation::Delete { name } => {
142 previews.push(format!("Delete {}", name));
143 }
144 crate::planner::PlanOperation::Rename { old, new } => {
145 previews.push(format!("Rename {} to {}", old, new));
146 }
147 _ => {}
148 }
149 }
150
151 Ok(previews)
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158 use tempfile::TempDir;
159
160 #[tokio::test]
161 async fn test_mutator_creation() {
162 let temp_dir = TempDir::new().unwrap();
163 let forge = Forge::open(temp_dir.path()).await.unwrap();
164 let mutator = Mutator::new(forge);
165
166 assert!(mutator.transaction.is_none());
167 }
168
169 #[tokio::test]
170 async fn test_begin_transaction() {
171 let temp_dir = TempDir::new().unwrap();
172 let forge = Forge::open(temp_dir.path()).await.unwrap();
173 let mut mutator = Mutator::new(forge);
174
175 mutator.begin_transaction().await.unwrap();
176 assert!(mutator.transaction.is_some());
177
178 assert!(mutator.begin_transaction().await.is_err());
180 }
181
182 #[tokio::test]
183 async fn test_rollback() {
184 let temp_dir = TempDir::new().unwrap();
185 let forge = Forge::open(temp_dir.path()).await.unwrap();
186 let mut mutator = Mutator::new(forge);
187
188 mutator.begin_transaction().await.unwrap();
189 assert!(mutator.transaction.is_some());
190
191 mutator.rollback().await.unwrap();
192 assert!(mutator.transaction.is_none());
193 }
194}