1use crate::Workspace;
2use std::collections::BTreeMap;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use tempfile::{TempDir, tempdir};
6use testcontainers::core::{ExecCommand, ExecResult, Mount};
7use testcontainers::runners::AsyncRunner;
8use testcontainers::{ContainerAsync, GenericImage, ImageExt};
9use tokio::io::AsyncBufRead;
10
11use super::{ContainerError, Image};
12
13#[derive(Clone)]
14pub struct Container {
15 inner: Arc<ContainerInner>,
16}
17
18struct ContainerInner {
19 container: ContainerAsync<GenericImage>,
20 workspace_root: PathBuf,
21 cwd: PathBuf,
22 _ephemeral_tempdirs: Vec<TempDir>,
23}
24
25pub struct ContainerBuilder {
26 image: Image,
27 env_vars: BTreeMap<String, String>,
28 mounts: Vec<Mount>,
29 ephemeral_mounts: Vec<String>,
30}
31
32pub struct ExecOutput {
33 pub exit_code: i64,
34 pub stdout: String,
35 pub stderr: String,
36}
37
38pub(crate) struct ExecHandle {
39 exec: ExecResult,
40}
41
42impl Container {
43 pub fn builder(image: Image) -> ContainerBuilder {
44 ContainerBuilder { image, env_vars: BTreeMap::new(), mounts: Vec::new(), ephemeral_mounts: Vec::new() }
45 }
46
47 pub fn workspace_root(&self) -> &Path {
48 &self.inner.workspace_root
49 }
50
51 pub fn cwd(&self) -> &Path {
52 &self.inner.cwd
53 }
54
55 pub async fn exec_shell(&self, script: impl Into<String>) -> Result<ExecOutput, ContainerError> {
56 let command = vec!["/bin/sh".to_string(), "-lc".to_string(), script.into()];
57 self.exec_streaming(command, BTreeMap::new()).await?.collect().await
58 }
59
60 pub(crate) async fn exec_streaming(
61 &self,
62 command: Vec<String>,
63 env_vars: BTreeMap<String, String>,
64 ) -> Result<ExecHandle, ContainerError> {
65 let exec = self
66 .inner
67 .container
68 .exec(ExecCommand::new(wrap_command(&self.inner.cwd, &command)).with_env_vars(env_vars))
69 .await?;
70 Ok(ExecHandle { exec })
71 }
72}
73
74impl ContainerBuilder {
75 pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
76 self.env_vars.insert(key.into(), value.into());
77 self
78 }
79
80 pub fn with_env_vars(mut self, env_vars: impl IntoIterator<Item = (String, String)>) -> Self {
81 self.env_vars.extend(env_vars);
82 self
83 }
84
85 pub fn with_mount(mut self, mount: Mount) -> Self {
86 self.mounts.push(mount);
87 self
88 }
89
90 pub fn with_ephemeral_mount(mut self, container_path: impl Into<String>) -> Self {
91 self.ephemeral_mounts.push(container_path.into());
92 self
93 }
94
95 pub async fn start(self, workspace: &Workspace) -> Result<Container, ContainerError> {
96 let container_workspace_root = PathBuf::from("/workspace");
97 let container_cwd = container_cwd(&container_workspace_root, workspace.relative_cwd());
98 let mut image = GenericImage::new(&self.image.name, &self.image.tag)
99 .with_entrypoint("/bin/sh")
100 .with_cmd(["-c", "sleep infinity"])
101 .with_mount(Mount::bind_mount(workspace.root_path().display().to_string(), "/workspace"))
102 .with_working_dir(container_cwd.display().to_string());
103
104 for mount in &self.mounts {
105 image = image.with_mount(mount.clone());
106 }
107
108 for (key, value) in &self.env_vars {
109 image = image.with_env_var(key.clone(), value.clone());
110 }
111
112 let mut ephemeral_tempdirs = Vec::with_capacity(self.ephemeral_mounts.len());
113 for container_path in &self.ephemeral_mounts {
114 let tempdir = tempdir().map_err(|source| ContainerError::EphemeralMountTempDir { source })?;
115 image = image.with_mount(Mount::bind_mount(tempdir.path().display().to_string(), container_path.clone()));
116 ephemeral_tempdirs.push(tempdir);
117 }
118
119 let container = image.start().await?;
120 Ok(Container {
121 inner: Arc::new(ContainerInner {
122 container,
123 workspace_root: container_workspace_root,
124 cwd: container_cwd,
125 _ephemeral_tempdirs: ephemeral_tempdirs,
126 }),
127 })
128 }
129}
130
131impl ExecHandle {
132 pub fn stdout(&mut self) -> impl AsyncBufRead + Send + '_ {
133 self.exec.stdout()
134 }
135
136 pub async fn stderr_to_string(&mut self) -> Result<String, ContainerError> {
137 let stderr = self.exec.stderr_to_vec().await?;
138 Ok(String::from_utf8_lossy(&stderr).into_owned())
139 }
140
141 pub async fn collect(mut self) -> Result<ExecOutput, ContainerError> {
142 let stdout = String::from_utf8_lossy(&self.exec.stdout_to_vec().await?).into_owned();
143 let stderr = String::from_utf8_lossy(&self.exec.stderr_to_vec().await?).into_owned();
144 let exit_code = self.exec.exit_code().await?.ok_or(ContainerError::MissingExecExitCode)?;
145 Ok(ExecOutput { exit_code, stdout, stderr })
146 }
147}
148
149fn container_cwd(container_workspace_root: &Path, relative_cwd: Option<&Path>) -> PathBuf {
150 relative_cwd.map_or_else(
151 || container_workspace_root.to_path_buf(),
152 |relative_cwd| container_workspace_root.join(relative_cwd),
153 )
154}
155
156fn wrap_command(container_cwd: &Path, command: &[String]) -> Vec<String> {
157 let mut argv = vec![
158 "/bin/sh".to_string(),
159 "-c".to_string(),
160 "cd \"$1\" && shift && exec \"$@\"".to_string(),
161 "aether-container-exec".to_string(),
162 container_cwd.display().to_string(),
163 ];
164 argv.extend(command.iter().cloned());
165 argv
166}
167
168#[cfg(test)]
169mod tests {
170 use super::*;
171
172 #[test]
173 fn container_cwd_uses_workspace_root_when_no_relative_cwd() {
174 assert_eq!(container_cwd(Path::new("/workspace"), None), Path::new("/workspace"));
175 }
176
177 #[test]
178 fn container_cwd_joins_relative_cwd() {
179 assert_eq!(container_cwd(Path::new("/workspace"), Some(Path::new("subdir"))), Path::new("/workspace/subdir"));
180 }
181
182 #[test]
183 fn wrap_command_cds_to_container_cwd_then_execs_command() {
184 let argv =
185 wrap_command(Path::new("/workspace/subdir"), &["node".to_string(), "/app/eval-agent.js".to_string()]);
186
187 assert_eq!(
188 argv,
189 vec![
190 "/bin/sh".to_string(),
191 "-c".to_string(),
192 "cd \"$1\" && shift && exec \"$@\"".to_string(),
193 "aether-container-exec".to_string(),
194 "/workspace/subdir".to_string(),
195 "node".to_string(),
196 "/app/eval-agent.js".to_string(),
197 ]
198 );
199 }
200
201 #[test]
202 fn wrap_command_preserves_command_args_after_cwd_arg() {
203 let argv = wrap_command(
204 Path::new("/workspace"),
205 &["node".to_string(), "/app/eval agent.js".to_string(), "--city".to_string(), "New York".to_string()],
206 );
207
208 assert_eq!(&argv[5..], ["node", "/app/eval agent.js", "--city", "New York"]);
209 }
210}