aether_evals/agents/
docker.rs1use 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}