aether_evals/agents/
docker_agent.rs

1use super::docker::{Docker, DockerExecConfig};
2use super::transcript::is_terminal;
3use super::{Agent, AgentConfig, DockerError, DockerImage, RunError, build_task_prompt};
4use aether_core::events::AgentMessage;
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7use tempfile::tempdir;
8use testcontainers::core::Mount;
9use tokio::io::{AsyncBufRead, AsyncBufReadExt};
10use tokio::sync::mpsc::Sender;
11
12/// Environment variable through which `DockerAgent` hands the wrapped task prompt to the
13/// in-container command.
14pub const AETHER_EVAL_WRAPPED_TASK_PROMPT_ENV: &str = "AETHER_EVAL_WRAPPED_TASK_PROMPT";
15pub const AETHER_EVAL_WORKSPACE_ROOT_ENV: &str = "AETHER_EVAL_WORKSPACE_ROOT";
16pub const AETHER_EVAL_CWD_ENV: &str = "AETHER_EVAL_CWD";
17
18pub struct DockerAgent {
19    image: DockerImage,
20    command: Vec<String>,
21    env_vars: HashMap<String, String>,
22    mounts: Vec<Mount>,
23    ephemeral_mounts: Vec<String>,
24}
25
26impl DockerAgent {
27    /// `command` is the in-container argv to exec; its stdout must be newline-delimited
28    /// `AgentMessage` JSON.
29    pub fn new(image: DockerImage, command: Vec<String>) -> Self {
30        Self { image, command, env_vars: HashMap::new(), mounts: Vec::new(), ephemeral_mounts: Vec::new() }
31    }
32
33    pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
34        self.env_vars.insert(key.into(), value.into());
35        self
36    }
37
38    pub fn with_env_vars(mut self, env_vars: impl IntoIterator<Item = (String, String)>) -> Self {
39        self.env_vars.extend(env_vars);
40        self
41    }
42
43    pub fn with_mount(mut self, mount: Mount) -> Self {
44        self.mounts.push(mount);
45        self
46    }
47
48    /// Register a container path to be backed by a fresh host tempdir, created at run time.
49    pub fn with_ephemeral_mount(mut self, container_path: impl Into<String>) -> Self {
50        self.ephemeral_mounts.push(container_path.into());
51        self
52    }
53}
54
55impl Agent for DockerAgent {
56    async fn run(&self, config: AgentConfig<'_>, tx: Sender<AgentMessage>) -> Result<(), RunError> {
57        let container_workspace_root = Path::new("/workspace");
58        let container_cwd = container_cwd(container_workspace_root, config.relative_cwd);
59        let wrapped_task_prompt = build_task_prompt(config.task_prompt, &container_cwd);
60        let command = wrap_command(&container_cwd, &self.command);
61
62        let mut docker = Docker::new(self.image.clone());
63
64        for mount in &self.mounts {
65            docker = docker.with_mount(mount.clone());
66        }
67
68        for (key, value) in &self.env_vars {
69            docker = docker.with_env_var(key.clone(), value.clone());
70        }
71
72        docker = docker
73            .with_env_var(AETHER_EVAL_WRAPPED_TASK_PROMPT_ENV, wrapped_task_prompt)
74            .with_env_var(AETHER_EVAL_WORKSPACE_ROOT_ENV, container_workspace_root.display().to_string())
75            .with_env_var(AETHER_EVAL_CWD_ENV, container_cwd.display().to_string());
76
77        let mut ephemeral_tempdirs = Vec::with_capacity(self.ephemeral_mounts.len());
78        for container_path in &self.ephemeral_mounts {
79            let tempdir = tempdir().map_err(|source| DockerError::EphemeralMountTempDir { source })?;
80            docker = docker.with_mount(Mount::bind_mount(tempdir.path().display().to_string(), container_path.clone()));
81            ephemeral_tempdirs.push(tempdir);
82        }
83
84        let (sent_terminal, stderr) = {
85            let mut session = docker.exec(DockerExecConfig { workspace_root: config.workspace_root, command }).await?;
86            let sent_terminal = forward_stdout_messages(session.stdout(), &tx).await?;
87            (sent_terminal, session.stderr_to_string().await?)
88        };
89
90        if !stderr.trim().is_empty() {
91            tracing::debug!("docker agent stderr: {}", stderr.trim());
92        }
93
94        if sent_terminal { Ok(()) } else { Err(DockerError::CommandExitWithoutTerminal { stderr }.into()) }
95    }
96}
97
98async fn forward_stdout_messages<T: AsyncBufRead + Unpin>(
99    reader: T,
100    tx: &Sender<AgentMessage>,
101) -> Result<bool, RunError> {
102    let mut lines = reader.lines();
103    while let Some(line) = lines.next_line().await.map_err(|source| DockerError::StdoutRead { source })? {
104        if line.trim().is_empty() {
105            continue;
106        }
107        tracing::debug!("docker agent stdout: {line}");
108        let message: AgentMessage =
109            serde_json::from_str(&line).map_err(|source| DockerError::AgentMessageJsonLine { line, source })?;
110        let terminal = is_terminal(&message);
111        tx.send(message).await.map_err(|error| RunError::ChannelSendFailed(error.to_string()))?;
112        if terminal {
113            return Ok(true);
114        }
115    }
116
117    Ok(false)
118}
119
120fn container_cwd(container_workspace_root: &Path, relative_cwd: Option<&Path>) -> PathBuf {
121    relative_cwd.map_or_else(
122        || container_workspace_root.to_path_buf(),
123        |relative_cwd| container_workspace_root.join(relative_cwd),
124    )
125}
126
127/// Wrap the agent argv in a shell that `cd`s into the eval cwd first. `testcontainers` execs inherit
128/// the image working dir (`/workspace`) and expose no per-exec cwd, so the cwd is passed as the
129/// shell's first positional (`$1`) and dropped with `shift` before exec-ing the real command. `$0`
130/// is a throwaway name.
131fn wrap_command(container_cwd: &Path, command: &[String]) -> Vec<String> {
132    let mut argv = vec![
133        "/bin/sh".to_string(),
134        "-c".to_string(),
135        "cd \"$1\" && shift && exec \"$@\"".to_string(),
136        "aether-eval-command".to_string(),
137        container_cwd.display().to_string(),
138    ];
139    argv.extend(command.iter().cloned());
140    argv
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn with_env_vars_merges_explicit_container_env() {
149        let agent = DockerAgent::new(DockerImage::new("sandbox", "latest"), vec!["agent".to_string()])
150            .with_env_var("ONE", "1")
151            .with_env_vars(HashMap::from([("TWO".to_string(), "2".to_string())]));
152
153        assert_eq!(agent.env_vars.get("ONE"), Some(&"1".to_string()));
154        assert_eq!(agent.env_vars.get("TWO"), Some(&"2".to_string()));
155    }
156
157    #[test]
158    fn container_cwd_uses_workspace_root_when_no_relative_cwd() {
159        assert_eq!(container_cwd(Path::new("/workspace"), None), Path::new("/workspace"));
160    }
161
162    #[test]
163    fn container_cwd_joins_relative_cwd() {
164        assert_eq!(container_cwd(Path::new("/workspace"), Some(Path::new("subdir"))), Path::new("/workspace/subdir"));
165    }
166
167    #[test]
168    fn wrap_command_cds_to_container_cwd_then_execs_command() {
169        let argv =
170            wrap_command(Path::new("/workspace/subdir"), &["node".to_string(), "/app/eval-agent.js".to_string()]);
171
172        assert_eq!(
173            argv,
174            vec![
175                "/bin/sh".to_string(),
176                "-c".to_string(),
177                "cd \"$1\" && shift && exec \"$@\"".to_string(),
178                "aether-eval-command".to_string(),
179                "/workspace/subdir".to_string(),
180                "node".to_string(),
181                "/app/eval-agent.js".to_string(),
182            ]
183        );
184    }
185
186    #[test]
187    fn wrap_command_preserves_command_args_after_cwd_arg() {
188        let argv = wrap_command(
189            Path::new("/workspace"),
190            &["node".to_string(), "/app/eval agent.js".to_string(), "--city".to_string(), "New York".to_string()],
191        );
192
193        assert_eq!(&argv[5..], ["node", "/app/eval agent.js", "--city", "New York"]);
194    }
195}
aether_evals/agents/docker_agent.rs

aether_evals/agents/
docker_agent.rs
