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

aether_evals/containers/
container.rs
