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
10/// Credentials and settings not covered by `LlmModel::ALL_REQUIRED_ENV_VARS`.
11///
12/// Bedrock resolves credentials through the AWS chain rather than a single
13/// required env var, so its variables have to be listed explicitly or a
14/// sandboxed run cannot authenticate at all.
15const EXTRA_FORWARDED_KEYS: &[&str] = &[
16    "OLLAMA_HOST",
17    "AWS_ACCESS_KEY_ID",
18    "AWS_BEARER_TOKEN_BEDROCK",
19    "AWS_DEFAULT_REGION",
20    "AWS_PROFILE",
21    "AWS_REGION",
22    "AWS_SECRET_ACCESS_KEY",
23    "AWS_SESSION_TOKEN",
24];
25
26const AWS_PATH_KEYS: &[&str] = &["AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "AWS_WEB_IDENTITY_TOKEN_FILE"];
27const AWS_SHARED_PATHS: &[&str] = &["config", "credentials", "sso/cache", "login/cache"];
28const AETHER_ENV_PREFIX: &str = "AETHER_";
29
30/// Host paths made available inside the container, and the env values that point at them.
31struct AwsBindings {
32    mounts: Vec<BindMount>,
33    env_overrides: Vec<(String, String)>,
34}
35
36struct BindMount {
37    host: PathBuf,
38    container: PathBuf,
39}
40
41/// Whether the sandboxed process runs attached to an interactive terminal.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43enum TerminalMode {
44    Interactive,
45    NonInteractive,
46}
47
48#[derive(Debug, Error)]
49pub enum SandboxError {
50    #[error("Docker is not installed or not in PATH")]
51    DockerNotFound,
52    #[error("Docker daemon is not running: {0}")]
53    DockerNotRunning(String),
54    #[error(
55        "Sandbox image '{0}' not found. Build it with:\n\
56             cargo build -p aether-agent-cli --bin aether\n\
57             docker build -t {0} -f crates/internal-evals/examples/Dockerfile ."
58    )]
59    ImageNotFound(String),
60    #[error("Failed to exec docker: {0}")]
61    ExecFailed(#[from] io::Error),
62    #[error("Could not determine home directory")]
63    HomeNotResolvable,
64}
65
66/// Entry point called from `main()` when `--sandbox-image` is present.
67pub fn exec_in_container(image: &str) -> ExitCode {
68    match try_exec_in_container(image) {
69        Ok(code) => code,
70        Err(err) => {
71            eprintln!("Sandbox error: {err}");
72            ExitCode::FAILURE
73        }
74    }
75}
76
77fn try_exec_in_container(image: &str) -> Result<ExitCode, SandboxError> {
78    check_docker()?;
79    check_image(image)?;
80
81    let cwd = env::current_dir().map_err(SandboxError::ExecFailed)?;
82    let home = dirs::home_dir().ok_or(SandboxError::HomeNotResolvable)?;
83    let aether_home = resolve_aether_home(&home);
84    let args: Vec<String> = env::args().collect();
85    let inner_args = filter_sandbox_arg(&args);
86    let mut env_vars = select_forwarded_vars(env::vars());
87    let aws = aws_bindings(&home, &env_vars);
88    apply_overrides(&mut env_vars, aws.env_overrides);
89
90    let terminal = if io::stdin().is_terminal() { TerminalMode::Interactive } else { TerminalMode::NonInteractive };
91    let docker_args = build_docker_args(image, &cwd, &aether_home, &aws.mounts, &env_vars, &inner_args, terminal);
92
93    exec_docker(&docker_args)
94}
95
96fn check_docker() -> Result<(), SandboxError> {
97    let output = Command::new("docker")
98        .arg("info")
99        .stdout(std::process::Stdio::null())
100        .stderr(std::process::Stdio::piped())
101        .output()
102        .map_err(|_| SandboxError::DockerNotFound)?;
103
104    if !output.status.success() {
105        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
106        return Err(SandboxError::DockerNotRunning(stderr));
107    }
108
109    Ok(())
110}
111
112fn check_image(image: &str) -> Result<(), SandboxError> {
113    let output = Command::new("docker")
114        .args(["image", "inspect", image])
115        .stdout(std::process::Stdio::null())
116        .stderr(std::process::Stdio::null())
117        .output()
118        .map_err(|_| SandboxError::DockerNotFound)?;
119
120    if !output.status.success() {
121        return Err(SandboxError::ImageNotFound(image.to_string()));
122    }
123
124    Ok(())
125}
126
127fn resolve_aether_home(home: &Path) -> PathBuf {
128    env::var("AETHER_HOME").map_or_else(|_| home.join(".aether"), PathBuf::from)
129}
130
131fn aws_bindings(home: &Path, env_vars: &[(String, String)]) -> AwsBindings {
132    let mut mounts = Vec::new();
133    let mut env_overrides = Vec::new();
134
135    let aws_home = home.join(".aws");
136    for relative in AWS_SHARED_PATHS {
137        let host = aws_home.join(relative);
138        if host.exists() {
139            mounts.push(BindMount { host, container: Path::new("/root/.aws").join(relative) });
140        }
141    }
142
143    for (key, value) in env_vars.iter().filter(|(key, _)| AWS_PATH_KEYS.contains(&key.as_str())) {
144        let host = PathBuf::from(value);
145        if !host.is_file() {
146            continue;
147        }
148        let container = PathBuf::from(format!("/run/aether-aws/{key}"));
149        env_overrides.push((key.clone(), container.to_string_lossy().into_owned()));
150        mounts.push(BindMount { host, container });
151    }
152
153    AwsBindings { mounts, env_overrides }
154}
155
156fn apply_overrides(env_vars: &mut [(String, String)], overrides: Vec<(String, String)>) {
157    for (key, replacement) in overrides {
158        if let Some((_, value)) = env_vars.iter_mut().find(|(existing, _)| *existing == key) {
159            *value = replacement;
160        }
161    }
162}
163
164fn filter_sandbox_arg(args: &[String]) -> Vec<String> {
165    let mut result = Vec::new();
166    let mut skip_next = false;
167    for arg in args {
168        if skip_next {
169            skip_next = false;
170            continue;
171        }
172        if arg == "--sandbox-image" {
173            skip_next = true;
174            continue;
175        }
176        if arg.starts_with("--sandbox-image=") {
177            continue;
178        }
179        result.push(arg.clone());
180    }
181    result
182}
183
184fn select_forwarded_vars(vars: impl Iterator<Item = (String, String)>) -> Vec<(String, String)> {
185    vars.filter(|(key, _)| {
186        LlmModel::ALL_REQUIRED_ENV_VARS.contains(&key.as_str())
187            || EXTRA_FORWARDED_KEYS.contains(&key.as_str())
188            || AWS_PATH_KEYS.contains(&key.as_str())
189            || key.starts_with(AETHER_ENV_PREFIX)
190    })
191    .collect()
192}
193
194fn build_docker_args(
195    image: &str,
196    cwd: &Path,
197    aether_home: &Path,
198    mounts: &[BindMount],
199    env_vars: &[(String, String)],
200    inner_args: &[String],
201    terminal: TerminalMode,
202) -> Vec<String> {
203    let mut args = vec!["run".to_string(), "--rm".to_string(), "-i".to_string()];
204    if terminal == TerminalMode::Interactive {
205        args.push("-t".to_string());
206    }
207    args.extend(
208        [
209            "--network",
210            "host",
211            "-w",
212            "/workspace",
213            "-v",
214            &format!("{}:/workspace", cwd.display()),
215            "-v",
216            &format!("{}:/root/.aether", aether_home.display()),
217            "-e",
218            "AETHER_HOME=/root/.aether",
219            "-e",
220            "AETHER_INSIDE_SANDBOX=1",
221        ]
222        .iter()
223        .map(ToString::to_string),
224    );
225
226    for mount in mounts {
227        args.push("-v".to_string());
228        args.push(format!("{}:{}:ro", mount.host.display(), mount.container.display()));
229    }
230
231    for (key, value) in env_vars {
232        args.push("-e".to_string());
233        args.push(format!("{key}={value}"));
234    }
235
236    args.push(image.to_string());
237
238    // Skip the binary name (first element) — the ENTRYPOINT already provides it
239    if inner_args.len() > 1 {
240        args.extend(inner_args[1..].iter().cloned());
241    }
242
243    args
244}
245
246#[cfg(unix)]
247fn exec_docker(args: &[String]) -> Result<ExitCode, SandboxError> {
248    use std::os::unix::process::CommandExt;
249
250    let err = Command::new("docker").args(args).exec();
251    Err(SandboxError::ExecFailed(err))
252}
253
254#[cfg(not(unix))]
255fn exec_docker(args: &[String]) -> Result<ExitCode, SandboxError> {
256    let status = Command::new("docker").args(args).status().map_err(SandboxError::ExecFailed)?;
257
258    Ok(match status.code() {
259        Some(0) => ExitCode::SUCCESS,
260        _ => ExitCode::FAILURE,
261    })
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    #[test]
269    fn filter_sandbox_arg_strips_separate_value() {
270        let args = vec![
271            "aether".to_string(),
272            "--sandbox-image".to_string(),
273            "my-image:latest".to_string(),
274            "headless".to_string(),
275            "-m".to_string(),
276            "gpt-4".to_string(),
277        ];
278        let filtered = filter_sandbox_arg(&args);
279        assert_eq!(filtered, vec!["aether", "headless", "-m", "gpt-4"]);
280    }
281
282    #[test]
283    fn filter_sandbox_arg_strips_equals_form() {
284        let args = vec!["aether".to_string(), "--sandbox-image=my-image:latest".to_string(), "headless".to_string()];
285        let filtered = filter_sandbox_arg(&args);
286        assert_eq!(filtered, vec!["aether", "headless"]);
287    }
288
289    #[test]
290    fn filter_sandbox_arg_noop_when_absent() {
291        let args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string()];
292        let filtered = filter_sandbox_arg(&args);
293        assert_eq!(filtered, args);
294    }
295
296    #[test]
297    fn filter_sandbox_arg_middle_position() {
298        let args = vec![
299            "aether".to_string(),
300            "headless".to_string(),
301            "--sandbox-image".to_string(),
302            "custom:v2".to_string(),
303            "-m".to_string(),
304        ];
305        let filtered = filter_sandbox_arg(&args);
306        assert_eq!(filtered, vec!["aether", "headless", "-m"]);
307    }
308
309    #[test]
310    fn select_forwarded_vars_includes_generated_provider_keys() {
311        let vars = vec![
312            ("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string()),
313            ("OPENROUTER_API_KEY".to_string(), "or-456".to_string()),
314            ("ZAI_API_KEY".to_string(), "zai-789".to_string()),
315            ("DEEPSEEK_API_KEY".to_string(), "ds-000".to_string()),
316            ("HOME".to_string(), "/root".to_string()),
317        ];
318        let forwarded = select_forwarded_vars(vars.into_iter());
319        assert_eq!(forwarded.len(), 4);
320        assert!(forwarded.iter().any(|(k, _)| k == "ANTHROPIC_API_KEY"));
321        assert!(forwarded.iter().any(|(k, _)| k == "OPENROUTER_API_KEY"));
322        assert!(forwarded.iter().any(|(k, _)| k == "ZAI_API_KEY"));
323        assert!(forwarded.iter().any(|(k, _)| k == "DEEPSEEK_API_KEY"));
324    }
325
326    #[test]
327    fn select_forwarded_vars_includes_extra_keys() {
328        let vars = vec![
329            ("OLLAMA_HOST".to_string(), "http://localhost:11434".to_string()),
330            ("HOME".to_string(), "/root".to_string()),
331        ];
332        let forwarded = select_forwarded_vars(vars.into_iter());
333        assert_eq!(forwarded.len(), 1);
334        assert!(forwarded.iter().any(|(k, _)| k == "OLLAMA_HOST"));
335    }
336
337    #[test]
338    fn select_forwarded_vars_includes_aether_prefix() {
339        let vars = vec![
340            ("AETHER_DEBUG".to_string(), "1".to_string()),
341            ("AETHER_LOG_LEVEL".to_string(), "trace".to_string()),
342            ("SOMETHING_ELSE".to_string(), "nope".to_string()),
343        ];
344        let forwarded = select_forwarded_vars(vars.into_iter());
345        assert_eq!(forwarded.len(), 2);
346        assert!(forwarded.iter().any(|(k, _)| k == "AETHER_DEBUG"));
347        assert!(forwarded.iter().any(|(k, _)| k == "AETHER_LOG_LEVEL"));
348    }
349
350    #[test]
351    fn select_forwarded_vars_excludes_unknown() {
352        let vars = vec![("HOME".to_string(), "/root".to_string()), ("EDITOR".to_string(), "vim".to_string())];
353        let forwarded = select_forwarded_vars(vars.into_iter());
354        assert_eq!(forwarded.len(), 0);
355    }
356
357    #[test]
358    fn all_required_env_vars_stays_in_sync() {
359        // If a new provider is added to codegen, this test reminds us it's auto-forwarded
360        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ANTHROPIC_API_KEY"));
361        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ZAI_API_KEY"));
362        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"DEEPSEEK_API_KEY"));
363        assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"XIAOMI_API_KEY"));
364    }
365
366    #[test]
367    fn build_docker_args_contains_expected_flags() {
368        let cwd = Path::new("/home/user/project");
369        let aether_home = Path::new("/home/user/.aether");
370        let env_vars = vec![("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string())];
371        let inner_args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string(), "gpt-4".to_string()];
372
373        let args = build_docker_args(
374            "test-image:latest",
375            cwd,
376            aether_home,
377            &[],
378            &env_vars,
379            &inner_args,
380            TerminalMode::NonInteractive,
381        );
382
383        assert!(args.contains(&"run".to_string()));
384        assert!(args.contains(&"--rm".to_string()));
385        assert!(args.contains(&"-i".to_string()));
386        assert!(!args.contains(&"-t".to_string()));
387        assert!(args.contains(&"--network".to_string()));
388        assert!(args.contains(&"host".to_string()));
389        assert!(args.contains(&"/workspace".to_string()));
390        assert!(args.contains(&format!("{}:/workspace", cwd.display())));
391        assert!(args.contains(&format!("{}:/root/.aether", aether_home.display())));
392        assert!(args.contains(&"AETHER_HOME=/root/.aether".to_string()));
393        assert!(args.contains(&"AETHER_INSIDE_SANDBOX=1".to_string()));
394        assert!(args.contains(&"ANTHROPIC_API_KEY=sk-123".to_string()));
395        assert!(args.contains(&"test-image:latest".to_string()));
396        // Inner args skip the binary name
397        assert!(args.contains(&"headless".to_string()));
398        assert!(args.contains(&"-m".to_string()));
399        assert!(args.contains(&"gpt-4".to_string()));
400        // Binary name must NOT appear after the image
401        let image_pos = args.iter().position(|a| a == "test-image:latest").unwrap();
402        assert!(!args[image_pos..].contains(&"aether".to_string()));
403    }
404
405    #[test]
406    fn build_docker_args_uses_custom_image() {
407        let cwd = Path::new("/tmp");
408        let aether_home = Path::new("/home/user/.aether");
409        let args = build_docker_args(
410            "my-go-sandbox:v2",
411            cwd,
412            aether_home,
413            &[],
414            &[],
415            &["aether".to_string(), "headless".to_string()],
416            TerminalMode::NonInteractive,
417        );
418
419        assert!(args.contains(&"my-go-sandbox:v2".to_string()));
420        assert!(!args.contains(&"test-image:latest".to_string()));
421    }
422
423    #[test]
424    fn build_docker_args_adds_tty_flag_when_requested() {
425        let cwd = Path::new("/tmp");
426        let aether_home = Path::new("/home/user/.aether");
427        let args = build_docker_args(
428            "test-image",
429            cwd,
430            aether_home,
431            &[],
432            &[],
433            &["aether".to_string()],
434            TerminalMode::Interactive,
435        );
436
437        assert!(args.contains(&"-t".to_string()));
438        assert!(args.contains(&"-i".to_string()));
439    }
440
441    #[test]
442    fn build_docker_args_skips_binary_name_only() {
443        let cwd = Path::new("/tmp");
444        let aether_home = Path::new("/home/user/.aether");
445        let args = build_docker_args(
446            "test-image:latest",
447            cwd,
448            aether_home,
449            &[],
450            &[],
451            &["aether".to_string()],
452            TerminalMode::NonInteractive,
453        );
454
455        // Only the binary name — nothing after image
456        assert_eq!(args.last().unwrap(), "test-image:latest");
457    }
458
459    /// A `~/.aws` holding one of everything: what the SDK reads, plus the AWS
460    /// CLI's own credential cache.
461    fn fake_aws_home(home: &Path) {
462        let aws = home.join(".aws");
463        std::fs::create_dir_all(aws.join("sso/cache")).unwrap();
464        std::fs::create_dir_all(aws.join("cli/cache")).unwrap();
465        std::fs::write(aws.join("config"), "[default]\n").unwrap();
466        std::fs::write(aws.join("credentials"), "[default]\n").unwrap();
467        std::fs::write(aws.join("sso/cache/token.json"), "{}").unwrap();
468        std::fs::write(aws.join("cli/cache/assumed-role.json"), "{}").unwrap();
469    }
470
471    #[test]
472    fn aws_bindings_mount_the_profile_and_rewrite_custom_files() {
473        let home = tempfile::tempdir().unwrap();
474        fake_aws_home(home.path());
475        let token = home.path().join("web-identity-token");
476        std::fs::write(&token, "token").unwrap();
477        let mut env_vars = vec![("AWS_WEB_IDENTITY_TOKEN_FILE".to_string(), token.to_string_lossy().into_owned())];
478
479        let aws = aws_bindings(home.path(), &env_vars);
480        apply_overrides(&mut env_vars, aws.env_overrides);
481
482        assert!(aws.mounts.iter().any(|mount| mount.container == Path::new("/root/.aws/config")));
483        assert!(aws.mounts.iter().any(|mount| mount.container == Path::new("/root/.aws/sso/cache")));
484        assert!(aws.mounts.iter().any(|mount| mount.host == token));
485        assert_eq!(env_vars[0].1, "/run/aether-aws/AWS_WEB_IDENTITY_TOKEN_FILE");
486
487        let args = build_docker_args(
488            "test-image",
489            Path::new("/workspace"),
490            &home.path().join(".aether"),
491            &aws.mounts,
492            &env_vars,
493            &["aether".to_string()],
494            TerminalMode::NonInteractive,
495        );
496        assert!(args.contains(&format!("{}:/root/.aws/config:ro", home.path().join(".aws/config").display())));
497        assert!(args.contains(&"AWS_WEB_IDENTITY_TOKEN_FILE=/run/aether-aws/AWS_WEB_IDENTITY_TOKEN_FILE".to_string()));
498    }
499
500    #[test]
501    fn aws_bindings_never_share_the_aws_cli_credential_cache() {
502        let home = tempfile::tempdir().unwrap();
503        fake_aws_home(home.path());
504
505        let aws = aws_bindings(home.path(), &[]);
506
507        assert!(
508            !aws.mounts.iter().any(|mount| mount.host.to_string_lossy().contains("/.aws/cli")),
509            "the AWS CLI's cache holds credentials the SDK never reads: {:?}",
510            aws.mounts.iter().map(|mount| mount.host.display().to_string()).collect::<Vec<_>>()
511        );
512        assert!(
513            !aws.mounts.iter().any(|mount| mount.container == Path::new("/root/.aws")),
514            "sharing the whole directory would expose everything under it"
515        );
516    }
517
518    #[test]
519    fn aws_bindings_mount_only_the_paths_that_exist() {
520        let home = tempfile::tempdir().unwrap();
521        std::fs::create_dir(home.path().join(".aws")).unwrap();
522        std::fs::write(home.path().join(".aws/config"), "[default]\n").unwrap();
523
524        let aws = aws_bindings(home.path(), &[]);
525
526        assert_eq!(aws.mounts.len(), 1);
527        assert_eq!(aws.mounts[0].container, Path::new("/root/.aws/config"));
528    }
529
530    #[test]
531    fn aws_bindings_skip_env_paths_that_do_not_exist_on_the_host() {
532        let home = tempfile::tempdir().unwrap();
533        let env_vars = vec![("AWS_CONFIG_FILE".to_string(), "/nowhere/config".to_string())];
534
535        let aws = aws_bindings(home.path(), &env_vars);
536
537        assert!(aws.mounts.is_empty(), "nothing to mount when neither ~/.aws nor the named file exists");
538        assert!(aws.env_overrides.is_empty(), "a missing host file must keep its original value");
539    }
540
541    #[test]
542    fn forwarded_vars_include_the_aws_path_variables() {
543        let vars = [
544            ("AWS_SHARED_CREDENTIALS_FILE".to_string(), "/home/user/creds".to_string()),
545            ("UNRELATED".to_string(), "x".to_string()),
546        ];
547
548        let forwarded = select_forwarded_vars(vars.into_iter());
549
550        assert_eq!(forwarded.len(), 1);
551        assert_eq!(forwarded[0].0, "AWS_SHARED_CREDENTIALS_FILE");
552    }
553
554    #[test]
555    fn sandbox_error_display_messages() {
556        assert_eq!(SandboxError::DockerNotFound.to_string(), "Docker is not installed or not in PATH");
557
558        assert!(SandboxError::DockerNotRunning("connection refused".into()).to_string().contains("connection refused"));
559
560        let img_err = SandboxError::ImageNotFound("aether-sandbox:latest".into());
561        assert!(img_err.to_string().contains("aether-sandbox:latest"));
562        assert!(img_err.to_string().contains("cargo build"));
563
564        assert!(SandboxError::HomeNotResolvable.to_string().contains("home directory"));
565
566        let io_err = io::Error::new(io::ErrorKind::NotFound, "not found");
567        assert!(SandboxError::ExecFailed(io_err).to_string().contains("not found"));
568    }
569}