Skip to main content

aether_cli/
sandbox.rs

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