Skip to main content

apm_core/wrapper/
mod.rs

1pub mod builtin;
2pub mod custom;
3pub mod path_guard;
4pub mod hook_config;
5pub use builtin::ClaudeWrapper;
6pub use custom::{WrapperKind, Manifest};
7pub use path_guard::PathGuard;
8
9use std::collections::HashMap;
10use std::path::{Path, PathBuf};
11
12pub const CONTRACT_VERSION: u32 = 1;
13
14pub struct WrapperContext {
15    /// Unique per-run identifier — e.g. "pi-0514-0628-7348". Stable for the
16    /// lifetime of one worker process; exposed to the wrapper via
17    /// `APM_AGENT_NAME` for log lines and worker-registry lookups.
18    pub worker_name: String,
19    /// Agent type — e.g. "pi", "claude". Exposed via `APM_AGENT_TYPE` and
20    /// used as the actor in ticket history rows when the worker calls
21    /// `apm state` (cleaner than the unique worker_name).
22    pub agent_type: String,
23    pub ticket_id: String,
24    pub ticket_branch: String,
25    pub worktree_path: PathBuf,
26    pub system_prompt_file: PathBuf,
27    pub user_message_file: PathBuf,
28    pub skip_permissions: bool,
29    pub profile: String,
30    pub role_prefix: Option<String>,
31    pub options: HashMap<String, String>,
32    pub model: Option<String>,
33    pub log_path: PathBuf,
34    pub container: Option<String>,
35    pub extra_env: HashMap<String, String>,
36    pub root: PathBuf,
37    pub keychain: HashMap<String, String>,
38    pub current_state: String,
39    /// Override for the wrapper-specific binary (e.g. for ClaudeWrapper, the
40    /// claude binary path). Honoured by built-ins that shell out to a fixed
41    /// binary; legacy `[workers].command` flows in here.
42    pub command: Option<String>,
43}
44
45pub trait Wrapper {
46    fn spawn(&self, ctx: &WrapperContext) -> anyhow::Result<std::process::Child>;
47}
48
49pub fn resolve_builtin(name: &str) -> Option<Box<dyn Wrapper>> {
50    match name {
51        "claude" => Some(Box::new(builtin::ClaudeWrapper)),
52        "mock-happy" => Some(Box::new(builtin::MockHappyWrapper)),
53        "mock-sad" => Some(Box::new(builtin::MockSadWrapper)),
54        "mock-random" => Some(Box::new(builtin::MockRandomWrapper)),
55        "debug" => Some(Box::new(builtin::DebugWrapper)),
56        _ => None,
57    }
58}
59
60pub fn list_builtin_names() -> &'static [&'static str] {
61    &["claude", "mock-happy", "mock-sad", "mock-random", "debug"]
62}
63
64pub fn resolve_wrapper(root: &Path, name: &str) -> anyhow::Result<Option<WrapperKind>> {
65    if let Some(script_path) = custom::find_script(root, name) {
66        let manifest = custom::parse_manifest(root, name)?;
67        return Ok(Some(WrapperKind::Custom { script_path, manifest }));
68    }
69    if resolve_builtin(name).is_some() {
70        return Ok(Some(WrapperKind::Builtin(name.to_owned())));
71    }
72    Ok(None)
73}
74
75pub fn write_temp_file(prefix: &str, content: &str) -> anyhow::Result<PathBuf> {
76    let path = std::env::temp_dir().join(format!("apm-{prefix}-{:04x}.txt", rand_u16()));
77    std::fs::write(&path, content)?;
78    Ok(path)
79}
80
81pub(crate) fn rand_u16() -> u16 {
82    use std::time::{SystemTime, UNIX_EPOCH};
83    SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().subsec_nanos() as u16
84}
85
86pub(crate) fn resolve_apm_cli_bin() -> String {
87    std::env::current_exe()
88        .and_then(|p| p.canonicalize())
89        .ok()
90        .map(|exe| resolve_cli_bin_from_exe(&exe))
91        .map(|p| p.to_string_lossy().into_owned())
92        .unwrap_or_default()
93}
94
95fn resolve_cli_bin_from_exe(exe: &std::path::Path) -> std::path::PathBuf {
96    // Check sibling first (production layout: apm-server and apm share the same directory).
97    let candidate = exe
98        .parent()
99        .map(|dir| dir.join("apm"))
100        .filter(|p| p.is_file() && *p != exe);
101    if let Some(p) = candidate {
102        return p;
103    }
104    // Check grandparent directory (test layout: test binary lives in deps/, apm lives
105    // one level up in the profile output directory, e.g. target/debug/).
106    let grandparent_candidate = exe
107        .parent()
108        .and_then(|parent| parent.parent())
109        .map(|grandparent| grandparent.join("apm"))
110        .filter(|p| p.is_file() && *p != exe);
111    if let Some(p) = grandparent_candidate {
112        return p;
113    }
114    exe.to_path_buf()
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    // --- resolve_cli_bin_from_exe ---
122
123    #[test]
124    fn resolve_cli_bin_from_exe_uses_sibling_apm_when_running_as_apm_server() {
125        use std::os::unix::fs::PermissionsExt;
126        let dir = tempfile::tempdir().unwrap();
127        let apm_server = dir.path().join("apm-server");
128        std::fs::write(&apm_server, "#!/bin/sh").unwrap();
129        std::fs::set_permissions(&apm_server, std::fs::Permissions::from_mode(0o755)).unwrap();
130        let apm = dir.path().join("apm");
131        std::fs::write(&apm, "#!/bin/sh").unwrap();
132        std::fs::set_permissions(&apm, std::fs::Permissions::from_mode(0o755)).unwrap();
133        let result = resolve_cli_bin_from_exe(&apm_server);
134        assert_eq!(
135            result.file_stem().and_then(|s| s.to_str()),
136            Some("apm"),
137            "APM_BIN must point to the apm CLI binary, not apm-server: {result:?}"
138        );
139    }
140
141    #[test]
142    fn resolve_cli_bin_from_exe_no_change_when_already_apm() {
143        use std::os::unix::fs::PermissionsExt;
144        let dir = tempfile::tempdir().unwrap();
145        let apm = dir.path().join("apm");
146        std::fs::write(&apm, "#!/bin/sh").unwrap();
147        std::fs::set_permissions(&apm, std::fs::Permissions::from_mode(0o755)).unwrap();
148        let result = resolve_cli_bin_from_exe(&apm);
149        assert_eq!(result, apm);
150    }
151
152    #[test]
153    fn resolve_cli_bin_from_exe_falls_back_when_no_sibling_apm() {
154        let dir = tempfile::tempdir().unwrap();
155        let apm_server = dir.path().join("apm-server");
156        std::fs::write(&apm_server, "#!/bin/sh").unwrap();
157        // No sibling apm file — must fall back to the exe itself
158        let result = resolve_cli_bin_from_exe(&apm_server);
159        assert_eq!(result, apm_server);
160    }
161
162    #[test]
163    fn resolve_builtin_claude_returns_some() {
164        assert!(resolve_builtin("claude").is_some());
165    }
166
167    #[test]
168    fn resolve_builtin_unknown_returns_none() {
169        assert!(resolve_builtin("bogus").is_none());
170        assert!(resolve_builtin("").is_none());
171    }
172
173    #[test]
174    fn resolve_builtin_mock_happy_returns_some() {
175        assert!(resolve_builtin("mock-happy").is_some());
176    }
177
178    #[test]
179    fn resolve_builtin_mock_sad_returns_some() {
180        assert!(resolve_builtin("mock-sad").is_some());
181    }
182
183    #[test]
184    fn resolve_builtin_mock_random_returns_some() {
185        assert!(resolve_builtin("mock-random").is_some());
186    }
187
188    #[test]
189    fn resolve_builtin_debug_returns_some() {
190        assert!(resolve_builtin("debug").is_some());
191    }
192}