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