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