Skip to main content

aether_evals/agents/
docker.rs

1use std::collections::BTreeMap;
2use std::path::Path;
3use testcontainers::core::{ExecCommand, ExecResult, Mount};
4use testcontainers::runners::AsyncRunner;
5use testcontainers::{ContainerAsync, GenericImage, ImageExt};
6use tokio::io::AsyncBufRead;
7
8pub(crate) struct Docker {
9    image: DockerImage,
10    env_vars: BTreeMap<String, String>,
11    mounts: Vec<Mount>,
12}
13
14pub(crate) struct DockerExecConfig<'a> {
15    pub workspace_root: &'a Path,
16    pub command: Vec<String>,
17}
18
19pub(crate) struct DockerExecSession {
20    _container: ContainerAsync<GenericImage>,
21    exec: ExecResult,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct DockerImage {
26    pub name: String,
27    pub tag: String,
28}
29
30impl Docker {
31    pub fn new(image: DockerImage) -> Self {
32        Self { image, env_vars: BTreeMap::new(), mounts: Vec::new() }
33    }
34
35    pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
36        self.env_vars.insert(key.into(), value.into());
37        self
38    }
39
40    pub fn with_mount(mut self, mount: Mount) -> Self {
41        self.mounts.push(mount);
42        self
43    }
44
45    pub async fn exec(&self, config: DockerExecConfig<'_>) -> Result<DockerExecSession, DockerError> {
46        let mut image = GenericImage::new(&self.image.name, &self.image.tag)
47            .with_entrypoint("/bin/sh")
48            .with_cmd(["-c", "sleep infinity"])
49            .with_mount(Mount::bind_mount(config.workspace_root.display().to_string(), "/workspace"))
50            .with_working_dir("/workspace");
51
52        for mount in &self.mounts {
53            image = image.with_mount(mount.clone());
54        }
55
56        for (key, value) in &self.env_vars {
57            image = image.with_env_var(key.clone(), value.clone());
58        }
59
60        let container = image.start().await?;
61        let exec = container.exec(ExecCommand::new(config.command)).await?;
62        Ok(DockerExecSession { _container: container, exec })
63    }
64}
65
66impl DockerExecSession {
67    pub fn stdout(&mut self) -> impl AsyncBufRead + Send + '_ {
68        self.exec.stdout()
69    }
70
71    pub async fn stderr_to_string(&mut self) -> Result<String, DockerError> {
72        let stderr = self.exec.stderr_to_vec().await?;
73        Ok(String::from_utf8_lossy(&stderr).into_owned())
74    }
75}
76
77impl DockerImage {
78    pub fn new(name: impl Into<String>, tag: impl Into<String>) -> Self {
79        Self { name: name.into(), tag: tag.into() }
80    }
81
82    pub fn parse(reference: &str) -> Result<Self, DockerImageParseError> {
83        parse_image_reference(reference).ok_or_else(|| DockerImageParseError { reference: reference.to_string() })
84    }
85}
86
87#[derive(Debug, thiserror::Error)]
88pub enum DockerError {
89    #[error("Docker container operation failed: {0}")]
90    Testcontainers(#[from] testcontainers::TestcontainersError),
91
92    #[error("Failed to create ephemeral mount temporary directory: {source}")]
93    EphemeralMountTempDir {
94        #[source]
95        source: std::io::Error,
96    },
97
98    #[error("Failed to read Docker stdout: {source}")]
99    StdoutRead {
100        #[source]
101        source: std::io::Error,
102    },
103
104    #[error("Docker command exited without emitting a terminal AgentMessage: {stderr}")]
105    CommandExitWithoutTerminal { stderr: String },
106
107    #[error("Docker command emitted invalid AgentMessage JSON line: {line}")]
108    AgentMessageJsonLine {
109        line: String,
110        #[source]
111        source: serde_json::Error,
112    },
113}
114
115#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
116#[error("invalid Docker image reference '{reference}'")]
117pub struct DockerImageParseError {
118    pub(crate) reference: String,
119}
120
121fn parse_image_reference(reference: &str) -> Option<DockerImage> {
122    if reference.is_empty() || reference.starts_with(':') || reference.ends_with(':') {
123        return None;
124    }
125
126    let last_slash = reference.rfind('/');
127    let last_colon = reference.rfind(':');
128    match last_colon.filter(|colon| last_slash.is_none_or(|slash| *colon > slash)) {
129        Some(colon) => Some(DockerImage::new(&reference[..colon], &reference[colon + 1..])),
130        None => Some(DockerImage::new(reference, "latest")),
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    #[test]
139    fn docker_image_parse_handles_name_tag() {
140        assert_eq!(DockerImage::parse("aether-sandbox:dev").unwrap(), DockerImage::new("aether-sandbox", "dev"));
141    }
142
143    #[test]
144    fn docker_image_parse_handles_registry_path() {
145        assert_eq!(
146            DockerImage::parse("ghcr.io/org/aether:sha").unwrap(),
147            DockerImage::new("ghcr.io/org/aether", "sha")
148        );
149    }
150
151    #[test]
152    fn docker_image_parse_defaults_latest() {
153        assert_eq!(DockerImage::parse("aether-sandbox").unwrap(), DockerImage::new("aether-sandbox", "latest"));
154    }
155
156    #[test]
157    fn docker_image_parse_rejects_invalid_reference() {
158        assert!(DockerImage::parse(":latest").is_err());
159        assert!(DockerImage::parse("aether:").is_err());
160    }
161}