aether_evals/containers/
container.rs

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

aether_evals/containers/
container.rs
