Skip to main content

atman_runtime/tools/
hunk.rs

1use std::path::PathBuf;
2
3use crate::error::RuntimeError;
4use crate::hunk::EditProposal;
5use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
6use crate::value::Value;
7
8pub struct FsEdit;
9
10impl Tool for FsEdit {
11    fn name(&self) -> &str {
12        "hunk.plan_edit"
13    }
14
15    fn tier(&self) -> Tier {
16        Tier::Zero
17    }
18
19    fn description(&self) -> Option<&str> {
20        Some(
21            "Compute a hunk-level EditProposal for replacing a file with new content. \
22             Nothing is written; feed the proposal into hunk.review or hunk.apply. \
23             For straightforward str_replace edits, prefer fs.edit instead.",
24        )
25    }
26
27    fn input_schema(&self) -> serde_json::Value {
28        serde_json::json!({
29            "type": "object",
30            "properties": {
31                "path": {"type": "string", "description": "File to edit."},
32                "new_content": {"type": "string", "description": "Proposed replacement content."}
33            },
34            "required": ["path", "new_content"]
35        })
36    }
37
38    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
39        Box::pin(async move {
40            let path = extract_path(&args, "path", 0)?;
41            let new_content = extract_string(&args, "new_content", 1)?;
42            let original = tokio::fs::read_to_string(&path).await.map_err(|e| {
43                RuntimeError::ToolFailed(format!("fs.edit({}): {e}", path.display()))
44            })?;
45            let proposal = EditProposal::compute(path, original, new_content);
46            Ok(Value::EditProposal(Box::new(proposal)))
47        })
48    }
49}
50
51pub struct HunkReview;
52
53impl Tool for HunkReview {
54    fn name(&self) -> &str {
55        "hunk.review"
56    }
57
58    fn tier(&self) -> Tier {
59        Tier::One
60    }
61
62    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> crate::tool::ApprovalLevel {
63        crate::tool::ApprovalLevel::Auto
64    }
65
66    fn description(&self) -> Option<&str> {
67        Some(
68            "Present an EditProposal to a human reviewer (or auto-approve if no resolver is \
69             configured). Returns a struct with mode = auto|resolved and a hunks id list \
70             the caller should pass to hunk.apply.",
71        )
72    }
73
74    fn input_schema(&self) -> serde_json::Value {
75        serde_json::json!({
76            "type": "object",
77            "properties": {
78                "proposal": {"description": "EditProposal value from hunk.plan_edit."},
79                "timeout_secs": {"type": "integer", "description": "Seconds to wait for a reviewer answer (default 300)."}
80            },
81            "required": ["proposal"]
82        })
83    }
84
85    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
86        Box::pin(async move {
87            let proposal = extract_proposal(&args)?;
88            let timeout_secs = match args.named("timeout_secs") {
89                Some(Value::Int(n)) if *n > 0 => *n as u64,
90                _ => 300,
91            };
92            let default_selection: Vec<u32> = proposal.hunks.iter().map(|h| h.id).collect();
93            let Some(resolver) = ctx.prompt_resolver.clone() else {
94                return Ok(Value::Struct(vec![
95                    ("mode".into(), Value::Str("auto".into())),
96                    (
97                        "hunks".into(),
98                        Value::List(
99                            default_selection
100                                .into_iter()
101                                .map(|id| Value::Int(id as i64))
102                                .collect(),
103                        ),
104                    ),
105                ]));
106            };
107            let id = crate::rendezvous::PromptId::now();
108            let payload = hunk_review_payload(&proposal);
109            let answer = crate::rendezvous::await_prompt_with_payload(
110                &resolver,
111                id,
112                "hunk_selection",
113                payload,
114                std::time::Duration::from_secs(timeout_secs),
115            )
116            .await?;
117            let selection = parse_answer_hunk_ids(&answer, &default_selection)?;
118            Ok(Value::Struct(vec![
119                ("mode".into(), Value::Str("resolved".into())),
120                ("prompt_id".into(), Value::Str(id.to_string())),
121                (
122                    "hunks".into(),
123                    Value::List(
124                        selection
125                            .into_iter()
126                            .map(|id| Value::Int(id as i64))
127                            .collect(),
128                    ),
129                ),
130            ]))
131        })
132    }
133}
134
135fn hunk_review_payload(proposal: &EditProposal) -> serde_json::Value {
136    let hunks: Vec<serde_json::Value> = proposal
137        .hunks
138        .iter()
139        .map(|h| {
140            let mut diff = String::new();
141            for line in &h.lines {
142                match line {
143                    crate::hunk::HunkLine::Add { text } => {
144                        diff.push('+');
145                        diff.push_str(text);
146                        if !text.ends_with('\n') {
147                            diff.push('\n');
148                        }
149                    }
150                    crate::hunk::HunkLine::Delete { text } => {
151                        diff.push('-');
152                        diff.push_str(text);
153                        if !text.ends_with('\n') {
154                            diff.push('\n');
155                        }
156                    }
157                    crate::hunk::HunkLine::Context { text } => {
158                        diff.push(' ');
159                        diff.push_str(text);
160                        if !text.ends_with('\n') {
161                            diff.push('\n');
162                        }
163                    }
164                }
165            }
166            serde_json::json!({
167                "id": h.id,
168                "old_start": h.old_start,
169                "old_len": h.old_len,
170                "new_start": h.new_start,
171                "new_len": h.new_len,
172                "unified_diff": diff,
173            })
174        })
175        .collect();
176    serde_json::json!({
177        "path": proposal.path.display().to_string(),
178        "hunks": hunks,
179        "options": ["all", "none", "select"],
180    })
181}
182
183fn parse_answer_hunk_ids(
184    answer: &serde_json::Value,
185    default: &[u32],
186) -> Result<Vec<u32>, RuntimeError> {
187    if answer.is_null() {
188        return Ok(default.to_vec());
189    }
190    if let Some(s) = answer.as_str() {
191        match s {
192            "all" => return Ok(default.to_vec()),
193            "none" => return Ok(Vec::new()),
194            other => {
195                return Err(RuntimeError::ToolFailed(format!(
196                    "hunk.review answer: unknown string `{other}`"
197                )));
198            }
199        }
200    }
201    if let Some(hunks) = answer.get("hunks").and_then(|v| v.as_array()) {
202        let mut ids = Vec::with_capacity(hunks.len());
203        for h in hunks {
204            let n = h.as_u64().ok_or_else(|| {
205                RuntimeError::ToolFailed(format!("hunk.review answer: hunk id not u64: {h:?}"))
206            })?;
207            ids.push(n as u32);
208        }
209        return Ok(ids);
210    }
211    Err(RuntimeError::ToolFailed(format!(
212        "hunk.review answer: unrecognized shape: {answer:?}"
213    )))
214}
215
216pub struct HunkApply;
217
218impl Tool for HunkApply {
219    fn name(&self) -> &str {
220        "hunk.apply"
221    }
222
223    fn tier(&self) -> Tier {
224        Tier::Two
225    }
226
227    fn description(&self) -> Option<&str> {
228        Some(
229            "Apply selected hunks from an EditProposal to disk. `hunks` is a list of hunk ids \
230             (usually from a hunk.review result). Returns which ids were applied vs skipped.",
231        )
232    }
233
234    fn input_schema(&self) -> serde_json::Value {
235        serde_json::json!({
236            "type": "object",
237            "properties": {
238                "proposal": {"description": "EditProposal value from hunk.plan_edit."},
239                "hunks": {
240                    "type": "array",
241                    "items": {"type": "integer"},
242                    "description": "Hunk ids to apply. Omit or pass \"all\" to apply everything."
243                }
244            },
245            "required": ["proposal"]
246        })
247    }
248
249    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
250        Box::pin(async move {
251            let proposal = extract_proposal(&args)?;
252            let selection = resolve_selection(&args, &proposal)?;
253            let applied = proposal
254                .apply_selected(&selection)
255                .map_err(|e| RuntimeError::ToolFailed(format!("hunk.apply: {e}")))?;
256            tokio::fs::write(&proposal.path, applied.as_bytes())
257                .await
258                .map_err(|e| {
259                    RuntimeError::ToolFailed(format!(
260                        "hunk.apply write {}: {e}",
261                        proposal.path.display()
262                    ))
263                })?;
264            let all_ids: Vec<u32> = proposal.hunks.iter().map(|h| h.id).collect();
265            let skipped: Vec<Value> = all_ids
266                .iter()
267                .filter(|id| !selection.contains(id))
268                .map(|id| Value::Int(*id as i64))
269                .collect();
270            let applied_ids: Vec<Value> =
271                selection.iter().map(|id| Value::Int(*id as i64)).collect();
272            Ok(Value::Struct(vec![
273                ("status".into(), Value::Str("applied".into())),
274                ("path".into(), Value::Path(proposal.path.clone())),
275                ("applied_hunks".into(), Value::List(applied_ids)),
276                ("skipped_hunks".into(), Value::List(skipped)),
277                ("total_hunks".into(), Value::Int(all_ids.len() as i64)),
278            ]))
279        })
280    }
281}
282
283fn resolve_selection(args: &ToolArgs, proposal: &EditProposal) -> Result<Vec<u32>, RuntimeError> {
284    let value = args
285        .named("hunks")
286        .cloned()
287        .or_else(|| args.positional(1).ok().cloned())
288        .unwrap_or(Value::Str("all".into()));
289    match value {
290        Value::Str(s) => match s.as_str() {
291            "all" => Ok(proposal.hunks.iter().map(|h| h.id).collect()),
292            "none" => Ok(Vec::new()),
293            other => Err(RuntimeError::ToolFailed(format!(
294                "hunk.apply: unknown selection string `{other}` (want `all` | `none` | [1,3,...])"
295            ))),
296        },
297        Value::List(items) => {
298            let mut out = Vec::with_capacity(items.len());
299            for item in items {
300                match item {
301                    Value::Int(n) if n > 0 => out.push(n as u32),
302                    other => {
303                        return Err(RuntimeError::TypeMismatch {
304                            expected: "positive int (hunk id)".into(),
305                            actual: other.kind_name().into(),
306                        });
307                    }
308                }
309            }
310            Ok(out)
311        }
312        other => Err(RuntimeError::TypeMismatch {
313            expected: "`all` | `none` | list of int (hunk ids)".into(),
314            actual: other.kind_name().into(),
315        }),
316    }
317}
318
319fn extract_proposal(args: &ToolArgs) -> Result<EditProposal, RuntimeError> {
320    let value = match args.named("proposal") {
321        Some(v) => v,
322        None => args.positional(0)?,
323    };
324    match value {
325        Value::EditProposal(p) => Ok((**p).clone()),
326        other => Err(RuntimeError::TypeMismatch {
327            expected: "edit_proposal".into(),
328            actual: other.kind_name().into(),
329        }),
330    }
331}
332
333fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
334    let value = match args.named(name) {
335        Some(v) => v,
336        None => args.positional(pos)?,
337    };
338    match value {
339        Value::Str(s) => Ok(s.clone()),
340        other => Err(RuntimeError::TypeMismatch {
341            expected: "string".into(),
342            actual: other.kind_name().into(),
343        }),
344    }
345}
346
347fn extract_path(args: &ToolArgs, name: &str, pos: usize) -> Result<PathBuf, RuntimeError> {
348    let value = match args.named(name) {
349        Some(v) => v,
350        None => args.positional(pos)?,
351    };
352    match value {
353        Value::Path(p) => Ok(p.clone()),
354        Value::Str(s) => Ok(PathBuf::from(s)),
355        other => Err(RuntimeError::TypeMismatch {
356            expected: "path".into(),
357            actual: other.kind_name().into(),
358        }),
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    #[tokio::test]
367    async fn fs_edit_returns_edit_proposal_with_hunks() {
368        let dir = tempfile::tempdir().unwrap();
369        let path = dir.path().join("f.txt");
370        std::fs::write(&path, "a\nb\nc\n").unwrap();
371        let ctx = ToolCtx::new();
372        let args = ToolArgs {
373            positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
374            named: vec![],
375        };
376        let v = FsEdit.call(args, &ctx).await.unwrap();
377        let Value::EditProposal(p) = v else {
378            panic!("expected EditProposal");
379        };
380        assert_eq!(p.hunks.len(), 1);
381        assert_eq!(p.original, "a\nb\nc\n");
382        assert_eq!(p.proposed, "a\nB\nc\n");
383    }
384
385    #[tokio::test]
386    async fn hunk_apply_all_writes_full_proposed() {
387        let dir = tempfile::tempdir().unwrap();
388        let path = dir.path().join("f.txt");
389        std::fs::write(&path, "a\nb\nc\n").unwrap();
390        let ctx = ToolCtx::new();
391        let proposal = FsEdit
392            .call(
393                ToolArgs {
394                    positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
395                    named: vec![],
396                },
397                &ctx,
398            )
399            .await
400            .unwrap();
401        let apply_args = ToolArgs {
402            positional: vec![proposal, Value::Str("all".into())],
403            named: vec![],
404        };
405        let out = HunkApply.call(apply_args, &ctx).await.unwrap();
406        let Value::Struct(fields) = out else {
407            panic!("expected struct");
408        };
409        assert!(matches!(
410            fields.iter().find(|(k, _)| k == "status").unwrap().1,
411            Value::Str(ref s) if s == "applied"
412        ));
413        let on_disk = std::fs::read_to_string(&path).unwrap();
414        assert_eq!(on_disk, "a\nB\nc\n");
415    }
416
417    #[tokio::test]
418    async fn hunk_apply_none_leaves_file_untouched() {
419        let dir = tempfile::tempdir().unwrap();
420        let path = dir.path().join("f.txt");
421        std::fs::write(&path, "a\nb\nc\n").unwrap();
422        let ctx = ToolCtx::new();
423        let proposal = FsEdit
424            .call(
425                ToolArgs {
426                    positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
427                    named: vec![],
428                },
429                &ctx,
430            )
431            .await
432            .unwrap();
433        let apply_args = ToolArgs {
434            positional: vec![proposal],
435            named: vec![("hunks".into(), Value::Str("none".into()))],
436        };
437        HunkApply.call(apply_args, &ctx).await.unwrap();
438        let on_disk = std::fs::read_to_string(&path).unwrap();
439        assert_eq!(on_disk, "a\nb\nc\n");
440    }
441
442    #[tokio::test]
443    async fn hunk_apply_with_id_list_writes_only_selected() {
444        let dir = tempfile::tempdir().unwrap();
445        let path = dir.path().join("f.txt");
446        let original: String = (0..20).map(|i| format!("l{i}\n")).collect();
447        std::fs::write(&path, &original).unwrap();
448        let mut proposed = original.clone();
449        proposed = proposed.replace("l3\n", "L3\n");
450        proposed = proposed.replace("l15\n", "L15\n");
451        let ctx = ToolCtx::new();
452        let proposal_v = FsEdit
453            .call(
454                ToolArgs {
455                    positional: vec![Value::Path(path.clone()), Value::Str(proposed.clone())],
456                    named: vec![],
457                },
458                &ctx,
459            )
460            .await
461            .unwrap();
462        let apply_args = ToolArgs {
463            positional: vec![proposal_v],
464            named: vec![("hunks".into(), Value::List(vec![Value::Int(1)]))],
465        };
466        let out = HunkApply.call(apply_args, &ctx).await.unwrap();
467        let Value::Struct(fields) = out else {
468            panic!("expected struct");
469        };
470        let f = |k: &str| fields.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone());
471        assert!(matches!(f("total_hunks"), Some(Value::Int(2))));
472        let on_disk = std::fs::read_to_string(&path).unwrap();
473        assert!(on_disk.contains("L3\n"));
474        assert!(!on_disk.contains("L15\n"));
475        assert!(on_disk.contains("l15\n"));
476    }
477
478    #[tokio::test]
479    async fn hunk_apply_rejects_unknown_selection_string() {
480        let dir = tempfile::tempdir().unwrap();
481        let path = dir.path().join("f.txt");
482        std::fs::write(&path, "a\n").unwrap();
483        let ctx = ToolCtx::new();
484        let proposal = FsEdit
485            .call(
486                ToolArgs {
487                    positional: vec![Value::Path(path), Value::Str("b\n".into())],
488                    named: vec![],
489                },
490                &ctx,
491            )
492            .await
493            .unwrap();
494        let args = ToolArgs {
495            positional: vec![proposal, Value::Str("some".into())],
496            named: vec![],
497        };
498        let err = HunkApply.call(args, &ctx).await.unwrap_err();
499        assert!(format!("{err}").contains("unknown selection"));
500    }
501}