Skip to main content

aether_evals/agents/
docker.rs

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