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