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").args(args).status().map_err(SandboxError::ExecFailed)?;
195
196 Ok(match status.code() {
197 Some(0) => ExitCode::SUCCESS,
198 _ => ExitCode::FAILURE,
199 })
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205
206 #[test]
207 fn filter_sandbox_arg_strips_separate_value() {
208 let args = vec![
209 "aether".to_string(),
210 "--sandbox-image".to_string(),
211 "my-image:latest".to_string(),
212 "headless".to_string(),
213 "-m".to_string(),
214 "gpt-4".to_string(),
215 ];
216 let filtered = filter_sandbox_arg(&args);
217 assert_eq!(filtered, vec!["aether", "headless", "-m", "gpt-4"]);
218 }
219
220 #[test]
221 fn filter_sandbox_arg_strips_equals_form() {
222 let args = vec!["aether".to_string(), "--sandbox-image=my-image:latest".to_string(), "headless".to_string()];
223 let filtered = filter_sandbox_arg(&args);
224 assert_eq!(filtered, vec!["aether", "headless"]);
225 }
226
227 #[test]
228 fn filter_sandbox_arg_noop_when_absent() {
229 let args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string()];
230 let filtered = filter_sandbox_arg(&args);
231 assert_eq!(filtered, args);
232 }
233
234 #[test]
235 fn filter_sandbox_arg_middle_position() {
236 let args = vec![
237 "aether".to_string(),
238 "headless".to_string(),
239 "--sandbox-image".to_string(),
240 "custom:v2".to_string(),
241 "-m".to_string(),
242 ];
243 let filtered = filter_sandbox_arg(&args);
244 assert_eq!(filtered, vec!["aether", "headless", "-m"]);
245 }
246
247 #[test]
248 fn select_forwarded_vars_includes_generated_provider_keys() {
249 let vars = vec![
250 ("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string()),
251 ("OPENROUTER_API_KEY".to_string(), "or-456".to_string()),
252 ("ZAI_API_KEY".to_string(), "zai-789".to_string()),
253 ("DEEPSEEK_API_KEY".to_string(), "ds-000".to_string()),
254 ("HOME".to_string(), "/root".to_string()),
255 ];
256 let forwarded = select_forwarded_vars(vars.into_iter());
257 assert_eq!(forwarded.len(), 4);
258 assert!(forwarded.iter().any(|(k, _)| k == "ANTHROPIC_API_KEY"));
259 assert!(forwarded.iter().any(|(k, _)| k == "OPENROUTER_API_KEY"));
260 assert!(forwarded.iter().any(|(k, _)| k == "ZAI_API_KEY"));
261 assert!(forwarded.iter().any(|(k, _)| k == "DEEPSEEK_API_KEY"));
262 }
263
264 #[test]
265 fn select_forwarded_vars_includes_extra_keys() {
266 let vars = vec![
267 ("OLLAMA_HOST".to_string(), "http://localhost:11434".to_string()),
268 ("HOME".to_string(), "/root".to_string()),
269 ];
270 let forwarded = select_forwarded_vars(vars.into_iter());
271 assert_eq!(forwarded.len(), 1);
272 assert!(forwarded.iter().any(|(k, _)| k == "OLLAMA_HOST"));
273 }
274
275 #[test]
276 fn select_forwarded_vars_includes_aether_prefix() {
277 let vars = vec![
278 ("AETHER_DEBUG".to_string(), "1".to_string()),
279 ("AETHER_LOG_LEVEL".to_string(), "trace".to_string()),
280 ("SOMETHING_ELSE".to_string(), "nope".to_string()),
281 ];
282 let forwarded = select_forwarded_vars(vars.into_iter());
283 assert_eq!(forwarded.len(), 2);
284 assert!(forwarded.iter().any(|(k, _)| k == "AETHER_DEBUG"));
285 assert!(forwarded.iter().any(|(k, _)| k == "AETHER_LOG_LEVEL"));
286 }
287
288 #[test]
289 fn select_forwarded_vars_excludes_unknown() {
290 let vars = vec![("HOME".to_string(), "/root".to_string()), ("EDITOR".to_string(), "vim".to_string())];
291 let forwarded = select_forwarded_vars(vars.into_iter());
292 assert!(forwarded.is_empty());
293 }
294
295 #[test]
296 fn all_required_env_vars_stays_in_sync() {
297 assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ANTHROPIC_API_KEY"));
299 assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"ZAI_API_KEY"));
300 assert!(LlmModel::ALL_REQUIRED_ENV_VARS.contains(&"DEEPSEEK_API_KEY"));
301 }
302
303 #[test]
304 fn build_docker_args_contains_expected_flags() {
305 let cwd = Path::new("/home/user/project");
306 let aether_home = Path::new("/home/user/.aether");
307 let env_vars = vec![("ANTHROPIC_API_KEY".to_string(), "sk-123".to_string())];
308 let inner_args = vec!["aether".to_string(), "headless".to_string(), "-m".to_string(), "gpt-4".to_string()];
309
310 let args = build_docker_args("test-image:latest", cwd, aether_home, &env_vars, &inner_args);
311
312 assert!(args.contains(&"run".to_string()));
313 assert!(args.contains(&"--rm".to_string()));
314 assert!(args.contains(&"-i".to_string()));
315 assert!(args.contains(&"--network".to_string()));
316 assert!(args.contains(&"host".to_string()));
317 assert!(args.contains(&"/workspace".to_string()));
318 assert!(args.contains(&format!("{}:/workspace", cwd.display())));
319 assert!(args.contains(&format!("{}:/root/.aether", aether_home.display())));
320 assert!(args.contains(&"AETHER_HOME=/root/.aether".to_string()));
321 assert!(args.contains(&"AETHER_INSIDE_SANDBOX=1".to_string()));
322 assert!(args.contains(&"ANTHROPIC_API_KEY=sk-123".to_string()));
323 assert!(args.contains(&"test-image:latest".to_string()));
324 assert!(args.contains(&"headless".to_string()));
326 assert!(args.contains(&"-m".to_string()));
327 assert!(args.contains(&"gpt-4".to_string()));
328 let image_pos = args.iter().position(|a| a == "test-image:latest").unwrap();
330 assert!(!args[image_pos..].contains(&"aether".to_string()));
331 }
332
333 #[test]
334 fn build_docker_args_uses_custom_image() {
335 let cwd = Path::new("/tmp");
336 let aether_home = Path::new("/home/user/.aether");
337 let args = build_docker_args(
338 "my-go-sandbox:v2",
339 cwd,
340 aether_home,
341 &[],
342 &["aether".to_string(), "headless".to_string()],
343 );
344
345 assert!(args.contains(&"my-go-sandbox:v2".to_string()));
346 assert!(!args.contains(&"test-image:latest".to_string()));
347 }
348
349 #[test]
350 fn build_docker_args_skips_binary_name_only() {
351 let cwd = Path::new("/tmp");
352 let aether_home = Path::new("/home/user/.aether");
353 let args = build_docker_args("test-image:latest", cwd, aether_home, &[], &["aether".to_string()]);
354
355 assert_eq!(args.last().unwrap(), "test-image:latest");
357 }
358
359 #[test]
360 fn sandbox_error_display_messages() {
361 assert_eq!(SandboxError::DockerNotFound.to_string(), "Docker is not installed or not in PATH");
362
363 assert!(SandboxError::DockerNotRunning("connection refused".into()).to_string().contains("connection refused"));
364
365 let img_err = SandboxError::ImageNotFound("aether-sandbox:latest".into());
366 assert!(img_err.to_string().contains("aether-sandbox:latest"));
367 assert!(img_err.to_string().contains("cargo build"));
368
369 assert!(SandboxError::HomeNotResolvable.to_string().contains("home directory"));
370
371 let io_err = io::Error::new(io::ErrorKind::NotFound, "not found");
372 assert!(SandboxError::ExecFailed(io_err).to_string().contains("not found"));
373 }
374}