Skip to main content

pma_core/
agent.rs

1//! Runs a coding agent, and then the project's verify command, in a worktree.
2//!
3//! Both run without push credentials, in their own process group, under a
4//! timeout. Which agent runs, and how it is invoked, is a record in the
5//! store: see `worker.rs`.
6
7use std::ffi::OsStr;
8use std::fs::File;
9use std::path::Path;
10use std::process::{Child, ChildStdin, Command, Stdio};
11use std::time::{Duration, Instant};
12
13use sanduk_sandbox::Policy;
14
15/// Removes what lets a child push, using `dir`, a directory `pma` owns outside
16/// every worktree, for an empty `gh` configuration and a `pre-push` hook that
17/// refuses. It stops an accidental push, not a process set on pushing:
18/// `git push --no-verify` to a remote with an explicit `pushurl` still gets
19/// through, and so does anything that reads a credential from disk.
20pub fn restrict(cmd: &mut Command, dir: &Path) -> std::io::Result<()> {
21    let gh = dir.join("gh");
22    std::fs::create_dir_all(&gh)?;
23    let hooks = dir.join("hooks");
24    write_hook(&hooks)?;
25    for var in [
26        "GH_TOKEN",
27        "GITHUB_TOKEN",
28        "GH_ENTERPRISE_TOKEN",
29        "GITHUB_ENTERPRISE_TOKEN",
30        "SSH_AUTH_SOCK",
31        "GIT_ASKPASS",
32        "SSH_ASKPASS",
33    ] {
34        cmd.env_remove(var);
35    }
36    cmd.env("GH_CONFIG_DIR", &gh)
37        .env("GIT_TERMINAL_PROMPT", "0");
38    let settings: [(&str, &OsStr); 3] = [
39        // An empty helper clears the user's credential helpers.
40        ("credential.helper", OsStr::new("")),
41        // An empty prefix matches every URL, so a push to any remote or URL
42        // is rewritten to one that cannot be reached. Git ignores this for
43        // a remote with an explicit `pushurl`, which the hook covers.
44        ("url.pma-agents-cannot-push:.pushInsteadOf", OsStr::new("")),
45        // Also keeps the project's own hooks off the agent's git commands;
46        // it is told not to commit.
47        ("core.hooksPath", hooks.as_os_str()),
48    ];
49    // Appended after any the user set, which a fixed count would drop.
50    let first = cmd
51        .get_envs()
52        .find(|(k, _)| *k == "GIT_CONFIG_COUNT")
53        .map(|(_, v)| v.map(OsStr::to_os_string))
54        .unwrap_or_else(|| std::env::var_os("GIT_CONFIG_COUNT"))
55        .and_then(|v| v.to_str()?.parse::<usize>().ok())
56        .unwrap_or(0);
57    for (i, (key, value)) in settings.iter().enumerate() {
58        cmd.env(format!("GIT_CONFIG_KEY_{}", first + i), key)
59            .env(format!("GIT_CONFIG_VALUE_{}", first + i), value);
60    }
61    cmd.env("GIT_CONFIG_COUNT", (first + settings.len()).to_string());
62    Ok(())
63}
64
65const PRE_PUSH: &str = "#!/bin/sh\necho 'pma: agents cannot push' >&2\nexit 1\n";
66
67/// Writes the refusing `pre-push` hook unless it is already there. Through a
68/// rename, because agents running in parallel share the file.
69fn write_hook(hooks: &Path) -> std::io::Result<()> {
70    use std::os::unix::fs::PermissionsExt;
71
72    let path = hooks.join("pre-push");
73    let executable =
74        |p: &Path| std::fs::metadata(p).is_ok_and(|m| m.permissions().mode() & 0o111 == 0o111);
75    if std::fs::read_to_string(&path).is_ok_and(|t| t == PRE_PUSH) && executable(&path) {
76        return Ok(());
77    }
78    std::fs::create_dir_all(hooks)?;
79    let tmp = hooks.join(format!(
80        ".pre-push.{}.{:?}",
81        std::process::id(),
82        std::thread::current().id()
83    ));
84    std::fs::write(&tmp, PRE_PUSH)?;
85    std::fs::set_permissions(&tmp, std::fs::Permissions::from_mode(0o755))?;
86    std::fs::rename(&tmp, &path)
87}
88
89#[derive(Debug)]
90pub struct Finished {
91    /// `None` when killed at the timeout.
92    pub success: Option<bool>,
93    pub seconds: i64,
94}
95
96/// Runs `cmd` with stdout and stderr in `log`, then kills its process group:
97/// at `timeout`, and also after a normal exit, so nothing it started in the
98/// background outlives it. With a `sandbox`, the whole group is confined.
99pub fn run_limited(
100    cmd: Command,
101    sandbox: Option<&Policy>,
102    log: &Path,
103    timeout: Duration,
104) -> std::io::Result<Finished> {
105    let started = Instant::now();
106    // Held until the group is killed. Dropped any earlier, it ends the run.
107    let (mut child, _lifeline) = spawn_guarded(cmd, sandbox, log)?;
108    let success = loop {
109        if let Some(status) = child.try_wait()? {
110            break Some(status.success());
111        }
112        if started.elapsed() >= timeout {
113            let _ = child.kill();
114            break None;
115        }
116        std::thread::sleep(Duration::from_millis(100));
117    };
118    // The watchdog is still waiting on its pipe, so the group exists and its
119    // id cannot have been reused.
120    let group = format!("-{}", child.id());
121    let _ = Command::new("kill").args(["-KILL", "--", &group]).status();
122    child.wait()?;
123    Ok(Finished {
124        success,
125        seconds: started.elapsed().as_secs() as i64,
126    })
127}
128
129/// Runs the program with its arguments, and kills the whole process group
130/// once its standard input, a pipe `pma` holds, reaches end of file. `pma`
131/// exiting for any reason closes the pipe, so a child cannot outlive the
132/// session whose lock it ran under. `kill -KILL -$$` rather than `kill --`,
133/// which `dash` refuses.
134const WATCHDOG: &str = r#"exec 3<&0 0</dev/null; "$@" 3<&- & pid=$!; { read _ <&3; kill -KILL -$$; } & exec 3<&-; wait "$pid""#;
135
136/// Starts `cmd` under the watchdog, as the leader of a new process group,
137/// with stdout and stderr in `log`. Dropping the returned pipe kills the
138/// group.
139///
140/// The sandbox wraps the watchdog's own `sh`: on Linux it is a `pre_exec`
141/// hook on that `Command`, which rebuilding `cmd` below would drop.
142pub fn spawn_guarded(
143    cmd: Command,
144    sandbox: Option<&Policy>,
145    log: &Path,
146) -> std::io::Result<(Child, ChildStdin)> {
147    use std::os::unix::process::CommandExt;
148
149    let mut guarded = match sandbox {
150        Some(policy) => policy.command("sh")?,
151        None => Command::new("sh"),
152    };
153    guarded
154        .arg("-c")
155        .arg(WATCHDOG)
156        .arg("pma-watchdog")
157        .arg(cmd.get_program())
158        .args(cmd.get_args());
159    if let Some(dir) = cmd.get_current_dir() {
160        guarded.current_dir(dir);
161    }
162    for (key, value) in cmd.get_envs() {
163        match value {
164            Some(v) => guarded.env(key, v),
165            None => guarded.env_remove(key),
166        };
167    }
168    let out = File::create(log)?;
169    guarded
170        .stdin(Stdio::piped())
171        .stdout(out.try_clone()?)
172        .stderr(out)
173        .process_group(0);
174    let mut child = guarded.spawn()?;
175    let lifeline = child
176        .stdin
177        .take()
178        .ok_or_else(|| std::io::Error::other("the watchdog's pipe was not created"))?;
179    Ok((child, lifeline))
180}
181
182/// The subcommands of a compound command. Claude Code checks each part of
183/// `a && b` against its rules separately, so an allowlist needs one rule per
184/// part. Quoting is not parsed: a rule that fails to match only denies the
185/// agent that command.
186pub fn bash_parts(command: &str) -> Vec<String> {
187    // `>&` and `&>` are redirections, not separators.
188    let command = command.replace(">&", "\u{0}").replace("&>", "\u{1}");
189    command
190        .replace("|&", "\n")
191        .split(['\n', ';', '|', '&'])
192        .map(|part| part.trim().replace('\u{0}', ">&").replace('\u{1}', "&>"))
193        .filter(|part| !part.is_empty())
194        .collect()
195}
196
197#[derive(Debug, PartialEq)]
198pub struct Report {
199    pub ok: bool,
200    pub summary: String,
201    pub cost_usd: Option<f64>,
202}
203
204/// Reads `claude -p --output-format json` output. Anything else is reported as
205/// a failure with the output's tail as the summary.
206pub fn parse_claude(output: &str) -> Report {
207    let value = output
208        .lines()
209        .rev()
210        .find_map(|l| serde_json::from_str::<serde_json::Value>(l.trim()).ok())
211        .filter(|v| v.get("type").and_then(|t| t.as_str()) == Some("result"));
212    match value {
213        Some(v) => Report {
214            ok: !v["is_error"].as_bool().unwrap_or(true)
215                && v["subtype"].as_str() == Some("success"),
216            summary: v["result"]
217                .as_str()
218                .map(String::from)
219                .or_else(|| v["subtype"].as_str().map(String::from))
220                .unwrap_or_default(),
221            cost_usd: v["total_cost_usd"].as_f64(),
222        },
223        None => Report {
224            ok: false,
225            summary: tail(output, 20),
226            cost_usd: None,
227        },
228    }
229}
230
231pub fn tail(text: &str, lines: usize) -> String {
232    let all: Vec<&str> = text.lines().collect();
233    all[all.len().saturating_sub(lines)..].join("\n")
234}
235
236/// The first of `make test`, `cargo test`, `go test ./...`, `npm test` and
237/// pytest that the project's files support.
238pub fn detect_verify(dir: &Path) -> Option<String> {
239    let read = |name: &str| std::fs::read_to_string(dir.join(name)).ok();
240    let has_test_target = ["GNUmakefile", "makefile", "Makefile"]
241        .iter()
242        .filter_map(|m| read(m))
243        .any(|text| text.lines().any(|l| l.starts_with("test:")));
244    if has_test_target {
245        return Some("make test".into());
246    }
247    if dir.join("Cargo.toml").is_file() {
248        return Some("cargo test".into());
249    }
250    if dir.join("go.mod").is_file() {
251        return Some("go test ./...".into());
252    }
253    if read("package.json").is_some_and(|t| t.contains("\"test\"")) {
254        return Some("npm test".into());
255    }
256    if dir.join("pyproject.toml").is_file() {
257        return Some(if dir.join("uv.lock").is_file() {
258            "uv run pytest".into()
259        } else {
260            "python3 -m pytest".into()
261        });
262    }
263    None
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    fn scratch(name: &str) -> std::path::PathBuf {
271        let dir = std::env::temp_dir().join(format!("pma-agent-{}-{name}", std::process::id()));
272        let _ = std::fs::remove_dir_all(&dir);
273        std::fs::create_dir_all(&dir).unwrap();
274        dir
275    }
276
277    #[test]
278    fn claude_results_are_read() {
279        let ok = r#"{"type":"result","subtype":"success","is_error":false,"result":"Fixed it.","total_cost_usd":0.42}"#;
280        assert_eq!(
281            parse_claude(ok),
282            Report {
283                ok: true,
284                summary: "Fixed it.".into(),
285                cost_usd: Some(0.42)
286            }
287        );
288        let over = r#"{"type":"result","subtype":"error_max_budget_usd","is_error":true,"total_cost_usd":1.01}"#;
289        let r = parse_claude(&format!("noise\n{over}\n"));
290        assert!(!r.ok);
291        assert_eq!(
292            (r.summary.as_str(), r.cost_usd),
293            ("error_max_budget_usd", Some(1.01))
294        );
295        let junk = parse_claude("Error: not logged in\n");
296        assert!(!junk.ok);
297        assert_eq!(junk.summary, "Error: not logged in");
298    }
299
300    #[test]
301    fn a_compound_verify_command_splits_into_its_parts() {
302        assert_eq!(
303            bash_parts("make check && uv run pytest -q 2>&1 | tail -5; cargo test &> log"),
304            [
305                "make check",
306                "uv run pytest -q 2>&1",
307                "tail -5",
308                "cargo test &> log"
309            ]
310        );
311        assert_eq!(bash_parts("a || b |& c"), ["a", "b", "c"]);
312    }
313
314    #[test]
315    fn verify_is_detected_from_project_files() {
316        let dir = scratch("verify");
317        assert_eq!(detect_verify(&dir), None);
318        std::fs::write(dir.join("pyproject.toml"), "").unwrap();
319        assert_eq!(detect_verify(&dir).as_deref(), Some("python3 -m pytest"));
320        std::fs::write(dir.join("Cargo.toml"), "").unwrap();
321        assert_eq!(detect_verify(&dir).as_deref(), Some("cargo test"));
322        std::fs::write(dir.join("Makefile"), "build:\n\tcc\n").unwrap();
323        assert_eq!(
324            detect_verify(&dir).as_deref(),
325            Some("cargo test"),
326            "a Makefile without a test target"
327        );
328        std::fs::write(dir.join("Makefile"), "test: build\n\tcc\n").unwrap();
329        assert_eq!(detect_verify(&dir).as_deref(), Some("make test"));
330        let _ = std::fs::remove_dir_all(dir);
331    }
332
333    #[test]
334    fn runs_are_logged_and_killed_at_the_timeout() {
335        let dir = scratch("run");
336        let log = dir.join("log");
337        let mut cmd = Command::new("sh");
338        cmd.args(["-c", "echo out; echo err >&2; exit 3"]);
339        let done = run_limited(cmd, None, &log, Duration::from_secs(10)).unwrap();
340        assert_eq!(done.success, Some(false));
341        assert_eq!(std::fs::read_to_string(&log).unwrap(), "out\nerr\n");
342
343        // The grandchild `sleep` must die with its group.
344        let mut cmd = Command::new("sh");
345        cmd.args(["-c", "sleep 30 & wait"]);
346        let started = Instant::now();
347        let done = run_limited(cmd, None, &log, Duration::from_millis(300)).unwrap();
348        assert_eq!(done.success, None);
349        assert!(started.elapsed() < Duration::from_secs(5));
350        let _ = std::fs::remove_dir_all(dir);
351    }
352
353    /// Whether `pid` is gone within 5 s. A killed process can linger as a
354    /// zombie until its new parent reaps it, and `kill -0` still finds one.
355    fn gone(pid: &str) -> bool {
356        let started = Instant::now();
357        while started.elapsed() < Duration::from_secs(5) {
358            let alive = Command::new("kill")
359                .args(["-0", pid.trim()])
360                .stderr(Stdio::null())
361                .status()
362                .unwrap()
363                .success();
364            if !alive {
365                return true;
366            }
367            std::thread::sleep(Duration::from_millis(50));
368        }
369        false
370    }
371
372    /// Reads the pid a child wrote to `file`, waiting for it to appear.
373    fn pid_in(file: &Path) -> String {
374        for _ in 0..50 {
375            if let Ok(pid) = std::fs::read_to_string(file)
376                && !pid.trim().is_empty()
377            {
378                return pid;
379            }
380            std::thread::sleep(Duration::from_millis(100));
381        }
382        panic!("{} never appeared", file.display());
383    }
384
385    #[test]
386    fn a_background_process_does_not_outlive_a_normal_exit() {
387        let dir = scratch("orphan");
388        let pid = dir.join("pid");
389        let mut cmd = Command::new("sh");
390        cmd.args(["-c", "sleep 30 & echo $! > pid; exit 0"])
391            .current_dir(&dir);
392        let done = run_limited(cmd, None, &dir.join("log"), Duration::from_secs(10)).unwrap();
393        assert_eq!(done.success, Some(true), "the child's own status");
394        assert!(gone(&pid_in(&pid)), "the background sleep was left running");
395        let _ = std::fs::remove_dir_all(dir);
396    }
397
398    /// What happens when `pma` dies mid-run: its end of the pipe closes.
399    #[test]
400    fn a_child_dies_with_the_process_that_started_it() {
401        let dir = scratch("lifeline");
402        let pid = dir.join("pid");
403        let mut cmd = Command::new("sh");
404        cmd.args(["-c", "sleep 30 & echo $! > pid; wait"])
405            .current_dir(&dir);
406        let (mut child, lifeline) = spawn_guarded(cmd, None, &dir.join("log")).unwrap();
407        let grandchild = pid_in(&pid);
408        drop(lifeline);
409        let started = Instant::now();
410        child.wait().unwrap();
411        assert!(started.elapsed() < Duration::from_secs(5));
412        assert!(gone(&grandchild));
413        let _ = std::fs::remove_dir_all(dir);
414    }
415
416    #[test]
417    fn restricted_children_cannot_see_tokens() {
418        let dir = scratch("env");
419        let mut cmd = Command::new("sh");
420        cmd.args([
421            "-c",
422            "echo \"${GH_TOKEN-unset} ${GH_ENTERPRISE_TOKEN-unset} ${GIT_ASKPASS-unset} $GIT_TERMINAL_PROMPT\"",
423        ])
424        .env("GH_TOKEN", "secret")
425        .env("GH_ENTERPRISE_TOKEN", "secret")
426        .env("GIT_ASKPASS", "/bin/echo")
427        .current_dir(&dir);
428        restrict(&mut cmd, &dir).unwrap();
429        let out = cmd.output().unwrap();
430        assert_eq!(
431            String::from_utf8_lossy(&out.stdout),
432            "unset unset unset 0\n"
433        );
434        let _ = std::fs::remove_dir_all(dir);
435    }
436
437    /// No remote, URL or explicit `pushurl` takes a push, and the settings a
438    /// user passed through the environment survive.
439    #[test]
440    fn restricted_children_cannot_push_anywhere() {
441        let dir = scratch("push");
442        let git = |args: &[&str]| {
443            let out = Command::new("git")
444                .arg("-C")
445                .arg(&dir)
446                .args(["-c", "user.name=t", "-c", "user.email=t@example.com"])
447                .args(args)
448                .env("GIT_CONFIG_GLOBAL", "/dev/null")
449                .output()
450                .unwrap();
451            assert!(out.status.success(), "git {args:?}");
452        };
453        git(&["init", "-q", "--bare", "real.git"]);
454        git(&["init", "-q", "--bare", "other.git"]);
455        git(&["init", "-q", "wt"]);
456        let wt = dir.join("wt");
457        let real = dir.join("real.git");
458        let url = real.to_str().unwrap();
459        git(&["-C", "wt", "commit", "-q", "--allow-empty", "-m", "a"]);
460        git(&["-C", "wt", "remote", "add", "origin", url]);
461        git(&["-C", "wt", "remote", "add", "upstream", "../other.git"]);
462        git(&["-C", "wt", "config", "remote.upstream.pushurl", url]);
463
464        let env = dir.join("env");
465        let pushes = |args: &[&str]| {
466            let mut cmd = Command::new("git");
467            cmd.arg("-C")
468                .arg(&wt)
469                .arg("push")
470                .args(args)
471                .env("GIT_CONFIG_GLOBAL", "/dev/null")
472                .stderr(Stdio::null());
473            restrict(&mut cmd, &env).unwrap();
474            cmd.status().unwrap().success()
475        };
476        assert!(!pushes(&["origin", "HEAD:refs/heads/a"]));
477        assert!(!pushes(&[url, "HEAD:refs/heads/b"]));
478        assert!(
479            !pushes(&["upstream", "HEAD:refs/heads/c"]),
480            "explicit pushurl"
481        );
482        let refs = Command::new("git")
483            .args(["-C", url, "for-each-ref"])
484            .output()
485            .unwrap();
486        assert!(
487            refs.stdout.is_empty(),
488            "{}",
489            String::from_utf8_lossy(&refs.stdout)
490        );
491
492        let mut cmd = Command::new("git");
493        cmd.args(["config", "user.name"])
494            .current_dir(&dir)
495            .env("GIT_CONFIG_GLOBAL", "/dev/null")
496            .env("GIT_CONFIG_COUNT", "1")
497            .env("GIT_CONFIG_KEY_0", "user.name")
498            .env("GIT_CONFIG_VALUE_0", "kept");
499        restrict(&mut cmd, &env).unwrap();
500        assert_eq!(
501            String::from_utf8_lossy(&cmd.output().unwrap().stdout),
502            "kept\n"
503        );
504        let _ = std::fs::remove_dir_all(dir);
505    }
506}