Skip to main content

a_agent/tools/
patch.rs

1use std::collections::BTreeSet;
2use std::fs::{self, OpenOptions};
3use std::io::Write;
4use std::path::{Path, PathBuf};
5
6use anyhow::{Context, Result};
7
8use super::path::unrestricted_path;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct PatchSummary {
12    pub files: Vec<PatchFileSummary>,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct PatchFileSummary {
17    pub path: String,
18    pub added: usize,
19    pub removed: usize,
20}
21
22#[derive(Debug)]
23enum Operation {
24    Add { path: String, content: String },
25    Delete { path: String },
26    Update { path: String, hunks: Vec<Hunk> },
27}
28
29#[derive(Debug)]
30struct Hunk {
31    anchor: Option<String>,
32    lines: Vec<HunkLine>,
33}
34
35#[derive(Debug)]
36enum HunkLine {
37    Context(String),
38    Add(String),
39    Remove(String),
40}
41
42enum Prepared {
43    Add {
44        path: PathBuf,
45        content: String,
46        summary: PatchFileSummary,
47    },
48    Update {
49        path: PathBuf,
50        content: String,
51        summary: PatchFileSummary,
52    },
53    Delete {
54        path: PathBuf,
55        summary: PatchFileSummary,
56    },
57}
58
59pub fn affected_paths(patch: &str) -> Result<Vec<String>> {
60    let lines = patch.lines().collect::<Vec<_>>();
61    if lines.first() != Some(&"*** Begin Patch") || lines.last() != Some(&"*** End Patch") {
62        anyhow::bail!("patch must start with '*** Begin Patch' and end with '*** End Patch'");
63    }
64    let mut paths = BTreeSet::new();
65    for line in lines {
66        let path = line
67            .strip_prefix("*** Add File: ")
68            .or_else(|| line.strip_prefix("*** Delete File: "))
69            .or_else(|| line.strip_prefix("*** Update File: "));
70        if let Some(path) = path {
71            validate_patch_path(path)?;
72            paths.insert(path.to_owned());
73        }
74    }
75    Ok(paths.into_iter().collect())
76}
77
78pub async fn apply_patch(root: &Path, patch: &str) -> Result<PatchSummary> {
79    let operations = parse_patch(patch)?;
80    let mut prepared = Vec::with_capacity(operations.len());
81    let mut seen = BTreeSet::new();
82    for operation in operations {
83        let requested = match &operation {
84            Operation::Add { path, .. }
85            | Operation::Delete { path }
86            | Operation::Update { path, .. } => path,
87        };
88        if !seen.insert(requested.clone()) {
89            anyhow::bail!("patch contains multiple operations for {requested}");
90        }
91        match operation {
92            Operation::Add { path, content } => {
93                let resolved = unrestricted_path(root, &path, false)?;
94                if resolved.exists() {
95                    anyhow::bail!("cannot add existing file: {path}");
96                }
97                let added = content.lines().count();
98                prepared.push(Prepared::Add {
99                    path: resolved,
100                    content,
101                    summary: PatchFileSummary {
102                        path,
103                        added,
104                        removed: 0,
105                    },
106                });
107            }
108            Operation::Delete { path } => {
109                let resolved = unrestricted_path(root, &path, true)?;
110                let source = fs::read_to_string(&resolved)
111                    .with_context(|| format!("read file before delete: {path}"))?;
112                prepared.push(Prepared::Delete {
113                    path: resolved,
114                    summary: PatchFileSummary {
115                        path,
116                        added: 0,
117                        removed: source.lines().count(),
118                    },
119                });
120            }
121            Operation::Update { path, hunks } => {
122                let resolved = unrestricted_path(root, &path, true)?;
123                let source = fs::read_to_string(&resolved)
124                    .with_context(|| format!("read file before update: {path}"))?;
125                let (content, added, removed) = apply_hunks(&source, &hunks, &path)?;
126                prepared.push(Prepared::Update {
127                    path: resolved,
128                    content,
129                    summary: PatchFileSummary {
130                        path,
131                        added,
132                        removed,
133                    },
134                });
135            }
136        }
137    }
138
139    let mut summaries = Vec::new();
140    for operation in prepared {
141        match operation {
142            Prepared::Add {
143                path,
144                content,
145                summary,
146            } => {
147                let parent = path.parent().context("patch target has no parent")?;
148                fs::create_dir_all(parent)?;
149                let mut file = OpenOptions::new()
150                    .write(true)
151                    .create_new(true)
152                    .open(&path)
153                    .with_context(|| format!("create {}", path.display()))?;
154                let result = file
155                    .write_all(content.as_bytes())
156                    .and_then(|()| file.sync_all());
157                if let Err(error) = result {
158                    drop(file);
159                    let _ = fs::remove_file(&path);
160                    return Err(error)
161                        .with_context(|| format!("write new file {}", path.display()));
162                }
163                summaries.push(summary);
164            }
165            Prepared::Update {
166                path,
167                content,
168                summary,
169            } => {
170                let mut file = OpenOptions::new()
171                    .write(true)
172                    .truncate(true)
173                    .open(&path)
174                    .with_context(|| format!("open {} for update", path.display()))?;
175                file.write_all(content.as_bytes())?;
176                file.sync_all()?;
177                summaries.push(summary);
178            }
179            Prepared::Delete { path, summary } => {
180                fs::remove_file(&path).with_context(|| format!("delete {}", path.display()))?;
181                summaries.push(summary);
182            }
183        }
184    }
185    Ok(PatchSummary { files: summaries })
186}
187
188fn parse_patch(patch: &str) -> Result<Vec<Operation>> {
189    let lines = patch.lines().collect::<Vec<_>>();
190    if lines.first() != Some(&"*** Begin Patch") || lines.last() != Some(&"*** End Patch") {
191        anyhow::bail!("patch must start with '*** Begin Patch' and end with '*** End Patch'");
192    }
193    let mut operations = Vec::new();
194    let mut index = 1;
195    while index + 1 < lines.len() {
196        let line = lines[index];
197        if let Some(path) = line.strip_prefix("*** Add File: ") {
198            validate_patch_path(path)?;
199            index += 1;
200            let mut content = Vec::new();
201            while index + 1 < lines.len() && !lines[index].starts_with("*** ") {
202                let value = lines[index].strip_prefix('+').with_context(|| {
203                    format!("add-file line must start with '+': {}", lines[index])
204                })?;
205                content.push(value);
206                index += 1;
207            }
208            let mut content = content.join("\n");
209            if !content.is_empty() {
210                content.push('\n');
211            }
212            operations.push(Operation::Add {
213                path: path.into(),
214                content,
215            });
216        } else if let Some(path) = line.strip_prefix("*** Delete File: ") {
217            validate_patch_path(path)?;
218            operations.push(Operation::Delete { path: path.into() });
219            index += 1;
220        } else if let Some(path) = line.strip_prefix("*** Update File: ") {
221            validate_patch_path(path)?;
222            index += 1;
223            let mut hunks = Vec::new();
224            while index + 1 < lines.len() && !lines[index].starts_with("*** ") {
225                let header = lines[index]
226                    .strip_prefix("@@")
227                    .with_context(|| format!("expected hunk header, found: {}", lines[index]))?;
228                let anchor = (!header.trim().is_empty()).then(|| header.trim().to_owned());
229                index += 1;
230                let mut hunk_lines = Vec::new();
231                while index + 1 < lines.len()
232                    && !lines[index].starts_with("@@")
233                    && !lines[index].starts_with("*** ")
234                {
235                    let value = lines[index];
236                    let (prefix, content) = value.split_at(1);
237                    hunk_lines.push(match prefix {
238                        " " => HunkLine::Context(content.into()),
239                        "+" => HunkLine::Add(content.into()),
240                        "-" => HunkLine::Remove(content.into()),
241                        _ => anyhow::bail!("hunk line must start with ' ', '+' or '-': {value}"),
242                    });
243                    index += 1;
244                }
245                if hunk_lines.is_empty() {
246                    anyhow::bail!("empty update hunk for {path}");
247                }
248                hunks.push(Hunk {
249                    anchor,
250                    lines: hunk_lines,
251                });
252            }
253            if hunks.is_empty() {
254                anyhow::bail!("update operation has no hunks: {path}");
255            }
256            operations.push(Operation::Update {
257                path: path.into(),
258                hunks,
259            });
260        } else {
261            anyhow::bail!("unknown patch operation: {line}");
262        }
263    }
264    if operations.is_empty() {
265        anyhow::bail!("patch contains no operations");
266    }
267    Ok(operations)
268}
269
270fn validate_patch_path(path: &str) -> Result<()> {
271    if path.trim().is_empty() {
272        anyhow::bail!("patch path must not be empty");
273    }
274    Ok(())
275}
276
277fn apply_hunks(source: &str, hunks: &[Hunk], path: &str) -> Result<(String, usize, usize)> {
278    let had_newline = source.ends_with('\n');
279    let mut lines = source.lines().map(ToOwned::to_owned).collect::<Vec<_>>();
280    let mut cursor = 0;
281    let mut total_added = 0;
282    let mut total_removed = 0;
283
284    for hunk in hunks {
285        if let Some(anchor) = &hunk.anchor
286            && let Some(position) = lines[cursor..]
287                .iter()
288                .position(|line| line.contains(anchor))
289        {
290            cursor += position;
291        }
292        let old = hunk
293            .lines
294            .iter()
295            .filter_map(|line| match line {
296                HunkLine::Context(value) | HunkLine::Remove(value) => Some(value.as_str()),
297                HunkLine::Add(_) => None,
298            })
299            .collect::<Vec<_>>();
300        let new = hunk
301            .lines
302            .iter()
303            .filter_map(|line| match line {
304                HunkLine::Context(value) | HunkLine::Add(value) => Some(value.clone()),
305                HunkLine::Remove(_) => None,
306            })
307            .collect::<Vec<_>>();
308        let position = if old.is_empty() {
309            cursor
310        } else {
311            find_sequence(&lines, &old, cursor)
312                .with_context(|| format!("context not found in {path}"))?
313        };
314        total_added += hunk
315            .lines
316            .iter()
317            .filter(|line| matches!(line, HunkLine::Add(_)))
318            .count();
319        total_removed += hunk
320            .lines
321            .iter()
322            .filter(|line| matches!(line, HunkLine::Remove(_)))
323            .count();
324        lines.splice(position..position + old.len(), new.iter().cloned());
325        cursor = position + new.len();
326    }
327    let mut result = lines.join("\n");
328    if had_newline {
329        result.push('\n');
330    }
331    Ok((result, total_added, total_removed))
332}
333
334fn find_sequence(lines: &[String], needle: &[&str], start: usize) -> Option<usize> {
335    (start..=lines.len().saturating_sub(needle.len())).find(|&position| {
336        lines[position..position + needle.len()]
337            .iter()
338            .map(String::as_str)
339            .eq(needle.iter().copied())
340    })
341}