aether_evals/agents/
image_builder.rs1use futures::{StreamExt, TryStreamExt};
2use std::collections::BTreeMap;
3use std::num::NonZeroUsize;
4use std::path::PathBuf;
5use thiserror::Error;
6use tokio::process::Command;
7
8#[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
31pub 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
60async 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}