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    let candidate = exe
97        .parent()
98        .map(|dir| dir.join("apm"))
99        .filter(|p| p.is_file() && *p != exe);
100    candidate.unwrap_or_else(|| exe.to_path_buf())
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    // --- resolve_cli_bin_from_exe ---
108
109    #[test]
110    fn resolve_cli_bin_from_exe_uses_sibling_apm_when_running_as_apm_server() {
111        use std::os::unix::fs::PermissionsExt;
112        let dir = tempfile::tempdir().unwrap();
113        let apm_server = dir.path().join("apm-server");
114        std::fs::write(&apm_server, "#!/bin/sh").unwrap();
115        std::fs::set_permissions(&apm_server, std::fs::Permissions::from_mode(0o755)).unwrap();
116        let apm = dir.path().join("apm");
117        std::fs::write(&apm, "#!/bin/sh").unwrap();
118        std::fs::set_permissions(&apm, std::fs::Permissions::from_mode(0o755)).unwrap();
119        let result = resolve_cli_bin_from_exe(&apm_server);
120        assert_eq!(
121            result.file_stem().and_then(|s| s.to_str()),
122            Some("apm"),
123            "APM_BIN must point to the apm CLI binary, not apm-server: {result:?}"
124        );
125    }
126
127    #[test]
128    fn resolve_cli_bin_from_exe_no_change_when_already_apm() {
129        use std::os::unix::fs::PermissionsExt;
130        let dir = tempfile::tempdir().unwrap();
131        let apm = dir.path().join("apm");
132        std::fs::write(&apm, "#!/bin/sh").unwrap();
133        std::fs::set_permissions(&apm, std::fs::Permissions::from_mode(0o755)).unwrap();
134        let result = resolve_cli_bin_from_exe(&apm);
135        assert_eq!(result, apm);
136    }
137
138    #[test]
139    fn resolve_cli_bin_from_exe_falls_back_when_no_sibling_apm() {
140        let dir = tempfile::tempdir().unwrap();
141        let apm_server = dir.path().join("apm-server");
142        std::fs::write(&apm_server, "#!/bin/sh").unwrap();
143        // No sibling apm file — must fall back to the exe itself
144        let result = resolve_cli_bin_from_exe(&apm_server);
145        assert_eq!(result, apm_server);
146    }
147
148    #[test]
149    fn resolve_builtin_claude_returns_some() {
150        assert!(resolve_builtin("claude").is_some());
151    }
152
153    #[test]
154    fn resolve_builtin_unknown_returns_none() {
155        assert!(resolve_builtin("bogus").is_none());
156        assert!(resolve_builtin("").is_none());
157    }
158
159    #[test]
160    fn resolve_builtin_mock_happy_returns_some() {
161        assert!(resolve_builtin("mock-happy").is_some());
162    }
163
164    #[test]
165    fn resolve_builtin_mock_sad_returns_some() {
166        assert!(resolve_builtin("mock-sad").is_some());
167    }
168
169    #[test]
170    fn resolve_builtin_mock_random_returns_some() {
171        assert!(resolve_builtin("mock-random").is_some());
172    }
173
174    #[test]
175    fn resolve_builtin_debug_returns_some() {
176        assert!(resolve_builtin("debug").is_some());
177    }
178}