Skip to main content

aether_cli/
sandbox.rs

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