Skip to main content

aether_cli/
sandbox.rs

1use std::env;
2use std::fmt;
3use std::io;
4use std::io::IsTerminal;
5use std::path::{Path, PathBuf};
6use std::process::{Command, ExitCode};
7
8use llm::LlmModel;
9
10const EXTRA_FORWARDED_KEYS: &[&str] = &["OLLAMA_HOST"];
11
12const AETHER_ENV_PREFIX: &str = "AETHER_";
13
14#[derive(Debug)]
15pub enum SandboxError {
16    DockerNotFound,
17    DockerNotRunning(String),
18    ImageNotFound(String),
19    ExecFailed(io::Error),
20    HomeNotResolvable,
21}
22
23impl fmt::Display for SandboxError {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            SandboxError::DockerNotFound => {
27                write!(f, "Docker is not installed or not in PATH")
28            }
29            SandboxError::DockerNotRunning(msg) => {
30                write!(f, "Docker daemon is not running: {msg}")
31            }
32            SandboxError::ImageNotFound(image) => {
33                write!(
34                    f,
35                    "Sandbox image '{image}' not found. Build it with:\n\
36                     cargo build --release -p aether-agent-cli\n\
37                     cp target/release/aether docker/\n\
38                     docker build -t {image} -f docker/Dockerfile.sandbox docker/"
39                )
40            }
41            SandboxError::ExecFailed(err) => write!(f, "Failed to exec docker: {err}"),
42            SandboxError::HomeNotResolvable => {
43                write!(f, "Could not determine home directory")
44            }
45        }
46    }
47}
48
49impl std::error::Error for SandboxError {}
50
51/// Entry point called from `main()` when `--sandbox-image` is present.
52pub fn exec_in_container(image: &str) -> ExitCode {
53    match try_exec_in_container(image) {
54        Ok(code) => code,
55        Err(err) => {
56            eprintln!("Sandbox error: {err}");
57            ExitCode::FAILURE
58        }
59    }
60}
61
62fn try_exec_in_container(image: &str) -> Result<ExitCode, SandboxError> {
63    check_docker()?;
64    check_image(image)?;
65
66    let cwd = env::current_dir().map_err(SandboxError::ExecFailed)?;
67    let aether_home = resolve_aether_home()?;
68    let args: Vec<String> = env::args().collect();
69    let inner_args = filter_sandbox_arg(&args);
70    let env_vars = select_forwarded_vars(env::vars());
71
72    let tty = io::stdin().is_terminal();
73    let docker_args = build_docker_args(image, &cwd, &aether_home, &env_vars, &inner_args, tty);
74
75    exec_docker(&docker_args)
76}
77
78fn check_docker() -> Result<(), SandboxError> {
79    let output = Command::new("docker")
80        .arg("info")
81        .stdout(std::process::Stdio::null())
82        .stderr(std::process::Stdio::piped())
83        .output()
84        .map_err(|_| SandboxError::DockerNotFound)?;
85
86    if !output.status.success() {
87        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
88        return Err(SandboxError::DockerNotRunning(stderr));
89    }
90
91    Ok(())
92}
93
94fn check_image(image: &str) -> Result<(), SandboxError> {
95    let output = Command::new("docker")
96        .args(["image", "inspect", image])
97        .stdout(std::process::Stdio::null())
98        .stderr(std::process::Stdio::null())
99        .output()
100        .map_err(|_| SandboxError::DockerNotFound)?;
101
102    if !output.status.success() {
103        return Err(SandboxError::ImageNotFound(image.to_string()));
104    }
105
106    Ok(())
107}
108
109fn resolve_aether_home() -> Result<PathBuf, SandboxError> {
110    if let Ok(val) = env::var("AETHER_HOME") {
111        return Ok(PathBuf::from(val));
112    }
113    let home = dirs::home_dir().ok_or(SandboxError::HomeNotResolvable)?;
114    Ok(home.join(".aether"))
115}
116
117fn filter_sandbox_arg(args: &[String]) -> Vec<String> {
118    let mut result = Vec::new();
119    let mut skip_next = false;
120    for arg in args {
121        if skip_next {
122            skip_next = false;
123            continue;
124        }
125        if arg == "--sandbox-image" {
126            skip_next = true;
127            continue;
128        }
129        if arg.starts_with("--sandbox-image=") {
130            continue;
131        }
132        result.push(arg.clone());
133    }
134    result
135}
136
137fn select_forwarded_vars(vars: impl Iterator<Item = (String, String)>) -> Vec<(String, String)> {
138    vars.filter(|(key, _)| {
139        LlmModel::ALL_REQUIRED_ENV_VARS.contains(&key.as_str())
140            || EXTRA_FORWARDED_KEYS.contains(&key.as_str())
141            || key.starts_with(AETHER_ENV_PREFIX)
142    })
143    .collect()
144}
145
146fn build_docker_args(
147    image: &str,
148    cwd: &Path,
149    aether_home: &Path,
150    env_vars: &[(String, String)],
151    inner_args: &[String],
152    tty: bool,
153) -> Vec<String> {
154    let mut args = vec!["run".to_string(), "--rm".to_string(), "-i".to_string()];
155    if tty {
156        args.push("-t".to_string());
157    }
158    args.extend(
159        [
160            "--network",
161            "host",
162            "-w",
163            "/workspace",
164            "-v",
165            &format!("{}:/workspace", cwd.display()),
166            "-v",
167            &format!("{}:/root/.aether", aether_home.display()),
168            "-e",
169            "AETHER_HOME=/root/.aether",
170            "-e",
171            "AETHER_INSIDE_SANDBOX=1",
172        ]
173        .iter()
174        .map(ToString::to_string),
175    );
176
177    for (key, value) in env_vars {
178        args.push("-e".to_string());
179        args.push(format!("{key}={value}"));
180    }
181
182    args.push(image.to_string());
183
184    // Skip the binary name (first element) — the ENTRYPOINT already provides it
185    if inner_args.len() > 1 {
186        args.extend(inner_args[1..].iter().cloned());
187    }
188
189    args
190}
191
192#[cfg(unix)]
193fn exec_docker(args: &[String]) -> Result<ExitCode, SandboxError> {
194    use std::os::unix::process::CommandExt;
195
196    let err = Command::new("docker").args(args).exec();
197    Err(SandboxError::ExecFailed(err))
198}
199
200#[cfg(not(unix))]
201fn exec_docker(args: &[String]) -> Result<ExitCode, SandboxError> {
202    let status = Command::new("docker").args(args).status().map_err(SandboxError::ExecFailed)?;
203
204    Ok(match status.code() {
205        Some(0) => ExitCode::SUCCESS,
206        _ => ExitCode::FAILURE,
207    })
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn filter_sandbox_arg_strips_separate_value() {
216        let args = vec![
217            "aether".to_string(),
218            "--sandbox-image".to_string(),
219            "my-image:latest".to_string(),
220            "headless".to_string(),
221            "-m".to_string(),
222            "gpt-4".to_string(),
223        ];
224        let filtered = filter_sandbox_arg(&args);
225        assert_eq!(filtered, vec!["aether", "headless", "-m", "gpt-4"]);
226    }
227
228    #[test]
229    fn filter_sandbox_arg_strips_equals_form() {
230        let args = vec!["aether".to_string(), "--sandbox-image=my-image:latest".to_string(), "headless".to_string()];
231        let filtered = filter_sandbox_arg(&args);
232        assert_eq!(filtered, vec!["aether", "headless"]);
233    }
234
235    #[test]
236    fn filter_sandbox_arg_noop_when_absent() {
237        let args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string()];
238        let filtered = filter_sandbox_arg(&args);
239        assert_eq!(filtered, args);
240    }
241
242    #[test]
243    fn filter_sandbox_arg_middle_position() {
244        let args = vec![
245            "aether".to_string(),
246            "headless".to_string(),
247            "--sandbox-image".to_string(),
248            "custom:v2".to_string(),
249            "-m".to_string(),
250        ];
251        let filtered = filter_sandbox_arg(&args);
252        assert_eq!(filtered, vec!["aether", "headless", "-m"]);
253    }
254
255    #[test]
256    fn select_forwarded_vars_includes_generated_provider_keys() {
257        let vars = vec![
258            ("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string()),
259            ("OPENROUTER_API_KEY".to_string(), "or-456".to_string()),
260            ("ZAI_API_KEY".to_string(), "zai-789".to_string()),
261            ("DEEPSEEK_API_KEY".to_string(), "ds-000".to_string()),
262            ("HOME".to_string(), "/root".to_string()),
263        ];
264        let forwarded = select_forwarded_vars(vars.into_iter());
265        assert_eq!(forwarded.len(), 4);
266        assert!(forwarded.iter().any(|(k, _)| k == "ANTHROPIC_API_KEY"));
267        assert!(forwarded.iter().any(|(k, _)| k == "OPENROUTER_API_KEY"));
268        assert!(forwarded.iter().any(|(k, _)| k == "ZAI_API_KEY"));
269        assert!(forwarded.iter().any(|(k, _)| k == "DEEPSEEK_API_KEY"));
270    }
271
272    #[test]
273    fn select_forwarded_vars_includes_extra_keys() {
274        let vars = vec![
275            ("OLLAMA_HOST".to_string(), "http://localhost:11434".to_string()),
276            ("HOME".to_string(), "/root".to_string()),
277        ];
278        let forwarded = select_forwarded_vars(vars.into_iter());
279        assert_eq!(forwarded.len(), 1);
280        assert!(forwarded.iter().any(|(k, _)| k == "OLLAMA_HOST"));
281    }
282
283    #[test]
284    fn select_forwarded_vars_includes_aether_prefix() {
285        let vars = vec![
286            ("AETHER_DEBUG".to_string(), "1".to_string()),
287            ("AETHER_LOG_LEVEL".to_string(), "trace".to_string()),
288            ("SOMETHING_ELSE".to_string(), "nope".to_string()),
289        ];
290        let forwarded = select_forwarded_vars(vars.into_iter());
291        assert_eq!(forwarded.len(), 2);
292        assert!(forwarded.iter().any(|(k, _)| k == "AETHER_DEBUG"));
293        assert!(forwarded.iter().any(|(k, _)| k == "AETHER_LOG_LEVEL"));
294    }
295
296    #[test]
297    fn select_forwarded_vars_excludes_unknown() {
298        let vars = vec![("HOME".to_string(), "/root".to_string()), ("EDITOR".to_string(), "vim".to_string())];
299        let forwarded = select_forwarded_vars(vars.into_iter());
300        assert!(forwarded.is_empty());
301    }
302
303    #[test]
304    fn all_required_env_vars_stays_in_sync() {
305        // If a new provider is added to codegen, this test reminds us it's auto-forwarded
306        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ANTHROPIC_API_KEY"));
307        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ZAI_API_KEY"));
308        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"DEEPSEEK_API_KEY"));
309    }
310
311    #[test]
312    fn build_docker_args_contains_expected_flags() {
313        let cwd = Path::new("/home/user/project");
314        let aether_home = Path::new("/home/user/.aether");
315        let env_vars = vec![("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string())];
316        let inner_args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string(), "gpt-4".to_string()];
317
318        let args = build_docker_args("test-image:latest", cwd, aether_home, &env_vars, &inner_args, false);
319
320        assert!(args.contains(&"run".to_string()));
321        assert!(args.contains(&"--rm".to_string()));
322        assert!(args.contains(&"-i".to_string()));
323        assert!(!args.contains(&"-t".to_string()));
324        assert!(args.contains(&"--network".to_string()));
325        assert!(args.contains(&"host".to_string()));
326        assert!(args.contains(&"/workspace".to_string()));
327        assert!(args.contains(&format!("{}:/workspace", cwd.display())));
328        assert!(args.contains(&format!("{}:/root/.aether", aether_home.display())));
329        assert!(args.contains(&"AETHER_HOME=/root/.aether".to_string()));
330        assert!(args.contains(&"AETHER_INSIDE_SANDBOX=1".to_string()));
331        assert!(args.contains(&"ANTHROPIC_API_KEY=sk-123".to_string()));
332        assert!(args.contains(&"test-image:latest".to_string()));
333        // Inner args skip the binary name
334        assert!(args.contains(&"headless".to_string()));
335        assert!(args.contains(&"-m".to_string()));
336        assert!(args.contains(&"gpt-4".to_string()));
337        // Binary name must NOT appear after the image
338        let image_pos = args.iter().position(|a| a == "test-image:latest").unwrap();
339        assert!(!args[image_pos..].contains(&"aether".to_string()));
340    }
341
342    #[test]
343    fn build_docker_args_uses_custom_image() {
344        let cwd = Path::new("/tmp");
345        let aether_home = Path::new("/home/user/.aether");
346        let args = build_docker_args(
347            "my-go-sandbox:v2",
348            cwd,
349            aether_home,
350            &[],
351            &["aether".to_string(), "headless".to_string()],
352            false,
353        );
354
355        assert!(args.contains(&"my-go-sandbox:v2".to_string()));
356        assert!(!args.contains(&"test-image:latest".to_string()));
357    }
358
359    #[test]
360    fn build_docker_args_adds_tty_flag_when_requested() {
361        let cwd = Path::new("/tmp");
362        let aether_home = Path::new("/home/user/.aether");
363        let args = build_docker_args("test-image", cwd, aether_home, &[], &["aether".to_string()], true);
364
365        assert!(args.contains(&"-t".to_string()));
366        assert!(args.contains(&"-i".to_string()));
367    }
368
369    #[test]
370    fn build_docker_args_skips_binary_name_only() {
371        let cwd = Path::new("/tmp");
372        let aether_home = Path::new("/home/user/.aether");
373        let args = build_docker_args("test-image:latest", cwd, aether_home, &[], &["aether".to_string()], false);
374
375        // Only the binary name — nothing after image
376        assert_eq!(args.last().unwrap(), "test-image:latest");
377    }
378
379    #[test]
380    fn sandbox_error_display_messages() {
381        assert_eq!(SandboxError::DockerNotFound.to_string(), "Docker is not installed or not in PATH");
382
383        assert!(SandboxError::DockerNotRunning("connection refused".into()).to_string().contains("connection refused"));
384
385        let img_err = SandboxError::ImageNotFound("aether-sandbox:latest".into());
386        assert!(img_err.to_string().contains("aether-sandbox:latest"));
387        assert!(img_err.to_string().contains("cargo build"));
388
389        assert!(SandboxError::HomeNotResolvable.to_string().contains("home directory"));
390
391        let io_err = io::Error::new(io::ErrorKind::NotFound, "not found");
392        assert!(SandboxError::ExecFailed(io_err).to_string().contains("not found"));
393    }
394}