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
33pub 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 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 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 assert!(args.contains(&"headless".to_string()));
317 assert!(args.contains(&"-m".to_string()));
318 assert!(args.contains(&"gpt-4".to_string()));
319 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 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}