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 = ctx.resolve_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 invocation_provenance(
250        &self,
251        args: &ToolArgs,
252        ctx: &ToolCtx,
253    ) -> Result<crate::permission::ResourceProvenance, RuntimeError> {
254        let proposal = extract_proposal(args)?;
255        Ok(crate::permission::ResourceProvenance::for_ctx(ctx)
256            .with_path(ctx, &proposal.path)?
257            .with_risk(crate::trust::RiskKind::FilesystemWrite))
258    }
259
260    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
261        Box::pin(async move {
262            let proposal = extract_proposal(&args)?;
263            crate::fs_access::authorize_write(ctx, &proposal.path, self.name(), true).await?;
264            let selection = resolve_selection(&args, &proposal)?;
265            let applied = proposal
266                .apply_selected(&selection)
267                .map_err(|e| RuntimeError::ToolFailed(format!("hunk.apply: {e}")))?;
268            tokio::fs::write(&proposal.path, applied.as_bytes())
269                .await
270                .map_err(|e| {
271                    RuntimeError::ToolFailed(format!(
272                        "hunk.apply write {}: {e}",
273                        proposal.path.display()
274                    ))
275                })?;
276            let all_ids: Vec<u32> = proposal.hunks.iter().map(|h| h.id).collect();
277            let skipped: Vec<Value> = all_ids
278                .iter()
279                .filter(|id| !selection.contains(id))
280                .map(|id| Value::Int(*id as i64))
281                .collect();
282            let applied_ids: Vec<Value> =
283                selection.iter().map(|id| Value::Int(*id as i64)).collect();
284            Ok(Value::Struct(vec![
285                ("status".into(), Value::Str("applied".into())),
286                ("path".into(), Value::Path(proposal.path.clone())),
287                ("applied_hunks".into(), Value::List(applied_ids)),
288                ("skipped_hunks".into(), Value::List(skipped)),
289                ("total_hunks".into(), Value::Int(all_ids.len() as i64)),
290            ]))
291        })
292    }
293}
294
295fn resolve_selection(args: &ToolArgs, proposal: &EditProposal) -> Result<Vec<u32>, RuntimeError> {
296    let value = args
297        .named("hunks")
298        .cloned()
299        .or_else(|| args.positional(1).ok().cloned())
300        .unwrap_or(Value::Str("all".into()));
301    match value {
302        Value::Str(s) => match s.as_str() {
303            "all" => Ok(proposal.hunks.iter().map(|h| h.id).collect()),
304            "none" => Ok(Vec::new()),
305            other => Err(RuntimeError::ToolFailed(format!(
306                "hunk.apply: unknown selection string `{other}` (want `all` | `none` | [1,3,...])"
307            ))),
308        },
309        Value::List(items) => {
310            let mut out = Vec::with_capacity(items.len());
311            for item in items {
312                match item {
313                    Value::Int(n) if n > 0 => out.push(n as u32),
314                    other => {
315                        return Err(RuntimeError::TypeMismatch {
316                            expected: "positive int (hunk id)".into(),
317                            actual: other.kind_name().into(),
318                        });
319                    }
320                }
321            }
322            Ok(out)
323        }
324        other => Err(RuntimeError::TypeMismatch {
325            expected: "`all` | `none` | list of int (hunk ids)".into(),
326            actual: other.kind_name().into(),
327        }),
328    }
329}
330
331fn extract_proposal(args: &ToolArgs) -> Result<EditProposal, RuntimeError> {
332    let value = match args.named("proposal") {
333        Some(v) => v,
334        None => args.positional(0)?,
335    };
336    match value {
337        Value::EditProposal(p) => Ok((**p).clone()),
338        other => Err(RuntimeError::TypeMismatch {
339            expected: "edit_proposal".into(),
340            actual: other.kind_name().into(),
341        }),
342    }
343}
344
345fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
346    let value = match args.named(name) {
347        Some(v) => v,
348        None => args.positional(pos)?,
349    };
350    match value {
351        Value::Str(s) => Ok(s.clone()),
352        other => Err(RuntimeError::TypeMismatch {
353            expected: "string".into(),
354            actual: other.kind_name().into(),
355        }),
356    }
357}
358
359fn extract_path(args: &ToolArgs, name: &str, pos: usize) -> Result<PathBuf, RuntimeError> {
360    let value = match args.named(name) {
361        Some(v) => v,
362        None => args.positional(pos)?,
363    };
364    match value {
365        Value::Path(p) => Ok(p.clone()),
366        Value::Str(s) => Ok(PathBuf::from(s)),
367        other => Err(RuntimeError::TypeMismatch {
368            expected: "path".into(),
369            actual: other.kind_name().into(),
370        }),
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    #[tokio::test]
379    async fn fs_edit_returns_edit_proposal_with_hunks() {
380        let dir = tempfile::tempdir().unwrap();
381        let path = dir.path().join("f.txt");
382        std::fs::write(&path, "a\nb\nc\n").unwrap();
383        let ctx = ToolCtx::new();
384        let args = ToolArgs {
385            positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
386            named: vec![],
387        };
388        let v = FsEdit.call(args, &ctx).await.unwrap();
389        let Value::EditProposal(p) = v else {
390            panic!("expected EditProposal");
391        };
392        assert_eq!(p.hunks.len(), 1);
393        assert_eq!(p.original, "a\nb\nc\n");
394        assert_eq!(p.proposed, "a\nB\nc\n");
395    }
396
397    #[tokio::test]
398    async fn hunk_apply_all_writes_full_proposed() {
399        let dir = tempfile::tempdir().unwrap();
400        let path = dir.path().join("f.txt");
401        std::fs::write(&path, "a\nb\nc\n").unwrap();
402        let ctx = ToolCtx::new();
403        let proposal = FsEdit
404            .call(
405                ToolArgs {
406                    positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
407                    named: vec![],
408                },
409                &ctx,
410            )
411            .await
412            .unwrap();
413        let apply_args = ToolArgs {
414            positional: vec![proposal, Value::Str("all".into())],
415            named: vec![],
416        };
417        let out = HunkApply.call(apply_args, &ctx).await.unwrap();
418        let Value::Struct(fields) = out else {
419            panic!("expected struct");
420        };
421        assert!(matches!(
422            fields.iter().find(|(k, _)| k == "status").unwrap().1,
423            Value::Str(ref s) if s == "applied"
424        ));
425        let on_disk = std::fs::read_to_string(&path).unwrap();
426        assert_eq!(on_disk, "a\nB\nc\n");
427    }
428
429    #[tokio::test]
430    async fn hunk_apply_none_leaves_file_untouched() {
431        let dir = tempfile::tempdir().unwrap();
432        let path = dir.path().join("f.txt");
433        std::fs::write(&path, "a\nb\nc\n").unwrap();
434        let ctx = ToolCtx::new();
435        let proposal = FsEdit
436            .call(
437                ToolArgs {
438                    positional: vec![Value::Path(path.clone()), Value::Str("a\nB\nc\n".into())],
439                    named: vec![],
440                },
441                &ctx,
442            )
443            .await
444            .unwrap();
445        let apply_args = ToolArgs {
446            positional: vec![proposal],
447            named: vec![("hunks".into(), Value::Str("none".into()))],
448        };
449        HunkApply.call(apply_args, &ctx).await.unwrap();
450        let on_disk = std::fs::read_to_string(&path).unwrap();
451        assert_eq!(on_disk, "a\nb\nc\n");
452    }
453
454    #[tokio::test]
455    async fn hunk_apply_with_id_list_writes_only_selected() {
456        let dir = tempfile::tempdir().unwrap();
457        let path = dir.path().join("f.txt");
458        let original: String = (0..20).map(|i| format!("l{i}\n")).collect();
459        std::fs::write(&path, &original).unwrap();
460        let mut proposed = original.clone();
461        proposed = proposed.replace("l3\n", "L3\n");
462        proposed = proposed.replace("l15\n", "L15\n");
463        let ctx = ToolCtx::new();
464        let proposal_v = FsEdit
465            .call(
466                ToolArgs {
467                    positional: vec![Value::Path(path.clone()), Value::Str(proposed.clone())],
468                    named: vec![],
469                },
470                &ctx,
471            )
472            .await
473            .unwrap();
474        let apply_args = ToolArgs {
475            positional: vec![proposal_v],
476            named: vec![("hunks".into(), Value::List(vec![Value::Int(1)]))],
477        };
478        let out = HunkApply.call(apply_args, &ctx).await.unwrap();
479        let Value::Struct(fields) = out else {
480            panic!("expected struct");
481        };
482        let f = |k: &str| fields.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone());
483        assert!(matches!(f("total_hunks"), Some(Value::Int(2))));
484        let on_disk = std::fs::read_to_string(&path).unwrap();
485        assert!(on_disk.contains("L3\n"));
486        assert!(!on_disk.contains("L15\n"));
487        assert!(on_disk.contains("l15\n"));
488    }
489
490    #[tokio::test]
491    async fn hunk_apply_rejects_external_path_without_changing_file() {
492        let workspace = tempfile::tempdir().unwrap();
493        let fixture = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
494            .join("target")
495            .join(format!("r4-hunk-{}.txt", uuid::Uuid::now_v7()));
496        std::fs::create_dir_all(fixture.parent().unwrap()).unwrap();
497        std::fs::write(&fixture, "original\n").unwrap();
498        let proposal = Value::EditProposal(Box::new(EditProposal::compute(
499            fixture.clone(),
500            "original\n".into(),
501            "changed\n".into(),
502        )));
503        let ctx = ToolCtx::new()
504            .with_fs_access(crate::fs_access::FsAccessPolicy::workspace_write(
505                workspace.path().into(),
506            ))
507            .with_workspace(crate::git_workspace::WorkspaceBinding {
508                workspace_id: "test".into(),
509                repository_root: workspace.path().into(),
510                path: workspace.path().into(),
511                branch: None,
512            });
513        let error = HunkApply
514            .call(
515                ToolArgs {
516                    positional: vec![proposal],
517                    named: vec![],
518                },
519                &ctx,
520            )
521            .await
522            .unwrap_err();
523        assert!(error.to_string().contains("outside workspace"));
524        assert_eq!(std::fs::read_to_string(&fixture).unwrap(), "original\n");
525        std::fs::remove_file(fixture).unwrap();
526    }
527
528    #[tokio::test]
529    async fn hunk_apply_rejects_unknown_selection_string() {
530        let dir = tempfile::tempdir().unwrap();
531        let path = dir.path().join("f.txt");
532        std::fs::write(&path, "a\n").unwrap();
533        let ctx = ToolCtx::new();
534        let proposal = FsEdit
535            .call(
536                ToolArgs {
537                    positional: vec![Value::Path(path), Value::Str("b\n".into())],
538                    named: vec![],
539                },
540                &ctx,
541            )
542            .await
543            .unwrap();
544        let args = ToolArgs {
545            positional: vec![proposal, Value::Str("some".into())],
546            named: vec![],
547        };
548        let err = HunkApply.call(args, &ctx).await.unwrap_err();
549        assert!(format!("{err}").contains("unknown selection"));
550    }
551}