Skip to main content

forge_agent/
mutate.rs

1//! Mutation engine - Transaction-based code mutation.
2//!
3//! This module implements the mutation phase, applying changes through
4//! the edit module with transaction support.
5
6use crate::{AgentError, Result};
7use forge_core::Forge;
8use std::sync::Arc;
9use tokio::fs;
10
11/// Mutator for transaction-based code changes.
12///
13/// The Mutator uses the EditModule to apply changes atomically.
14#[derive(Clone)]
15pub struct Mutator {
16    /// Forge SDK for graph queries
17    forge: Arc<Forge>,
18    /// Current transaction state
19    transaction: Option<Transaction>,
20}
21
22/// Active transaction state.
23#[derive(Clone, Debug)]
24struct Transaction {
25    /// Steps applied in this transaction
26    applied_steps: Vec<String>,
27    /// Original state for rollback
28    rollback_state: Vec<RollbackState>,
29}
30
31/// Rollback state for a transaction.
32#[derive(Clone, Debug)]
33struct RollbackState {
34    /// File path
35    file: String,
36    /// Original content
37    original_content: String,
38}
39
40impl Mutator {
41    /// Creates a new mutator.
42    pub fn new(forge: Forge) -> Self {
43        Self {
44            forge: Arc::new(forge),
45            transaction: None,
46        }
47    }
48
49    /// Begins a new transaction.
50    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    /// Applies a single step in the current transaction.
66    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                // Record for rollback
75                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                // Save original for rollback
84                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                // Write new content
92                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                // Inspect doesn't modify files
100            }
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    /// Rolls back the current transaction.
116    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        // Rollback in reverse order
123        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    /// Returns preview of changes without applying.
133    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        // Second begin should fail
179        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}