Skip to main content

aether_evals/agents/
image_builder.rs

1use futures::{StreamExt, TryStreamExt};
2use std::collections::BTreeMap;
3use std::num::NonZeroUsize;
4use std::path::PathBuf;
5use thiserror::Error;
6use tokio::process::Command;
7
8/// A request to build a Docker image from a Dockerfile.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ImageBuildRequest {
11    pub dockerfile: PathBuf,
12    pub context: PathBuf,
13    pub tag: String,
14}
15
16#[derive(Debug, Error)]
17pub enum ImageBuildError {
18    #[error("failed to spawn `docker build`: {0}")]
19    Spawn(#[source] std::io::Error),
20
21    #[error("`docker build -t {tag}` failed{}:\n{output_tail}", match code {
22        Some(code) => format!(" with exit code {code}"),
23        None => String::new(),
24    })]
25    Failed { tag: String, code: Option<i32>, output_tail: String },
26
27    #[error("multiple builds target the image tag `{tag}` with different Dockerfiles or contexts")]
28    ConflictingTag { tag: String },
29}
30
31/// Build the requested images concurrently (bounded by `max_concurrency`), deduplicating
32/// identical requests. Two requests with the same tag but a different Dockerfile or context are
33/// rejected before anything builds, since the second build would silently replace the first
34/// image.
35pub async fn build_images(
36    requests: impl IntoIterator<Item = ImageBuildRequest>,
37    max_concurrency: NonZeroUsize,
38) -> Result<(), ImageBuildError> {
39    let mut requests_by_tag: BTreeMap<String, ImageBuildRequest> = BTreeMap::new();
40    for request in requests {
41        match requests_by_tag.get(&request.tag) {
42            Some(existing) if existing != &request => {
43                return Err(ImageBuildError::ConflictingTag { tag: request.tag });
44            }
45            _ => {
46                requests_by_tag.insert(request.tag.clone(), request);
47            }
48        }
49    }
50
51    futures::stream::iter(requests_by_tag.into_values().map(build_image))
52        .buffer_unordered(max_concurrency.get())
53        .try_collect::<Vec<()>>()
54        .await
55        .map(|_| ())
56}
57
58const BUILD_OUTPUT_TAIL_CHARS: usize = 8_000;
59
60/// Build one image. Output is captured rather than streamed so concurrent builds do not
61/// interleave; on failure the tail of the build output is attached to the error.
62async fn build_image(request: ImageBuildRequest) -> Result<(), ImageBuildError> {
63    let output = Command::new("docker")
64        .args(["build", "-t", &request.tag, "-f"])
65        .arg(&request.dockerfile)
66        .arg(&request.context)
67        .kill_on_drop(true)
68        .output()
69        .await
70        .map_err(ImageBuildError::Spawn)?;
71
72    if output.status.success() {
73        return Ok(());
74    }
75
76    let mut combined = String::from_utf8_lossy(&output.stdout).into_owned();
77    combined.push_str(&String::from_utf8_lossy(&output.stderr));
78    Err(ImageBuildError::Failed {
79        tag: request.tag,
80        code: output.status.code(),
81        output_tail: tail_chars(&combined, BUILD_OUTPUT_TAIL_CHARS),
82    })
83}
84
85fn tail_chars(value: &str, max_chars: usize) -> String {
86    let total = value.chars().count();
87    if total <= max_chars {
88        return value.to_string();
89    }
90
91    let tail: String = value.chars().skip(total - max_chars).collect();
92    format!("[truncated]...{tail}")
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[tokio::test]
100    async fn rejects_conflicting_builds_for_the_same_tag() {
101        let requests = [
102            ImageBuildRequest { dockerfile: "a/Dockerfile".into(), context: "a".into(), tag: "same:tag".into() },
103            ImageBuildRequest { dockerfile: "b/Dockerfile".into(), context: "b".into(), tag: "same:tag".into() },
104        ];
105
106        let error = build_images(requests, NonZeroUsize::new(2).unwrap()).await.unwrap_err();
107
108        assert!(matches!(error, ImageBuildError::ConflictingTag { tag } if tag == "same:tag"));
109    }
110
111    #[test]
112    fn tail_chars_keeps_the_end_of_long_output() {
113        assert_eq!(tail_chars("short", 10), "short");
114        assert_eq!(tail_chars("abcdefghij", 3), "[truncated]...hij");
115    }
116}