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