aether_evals/agents/
docker_agent.rs

1use super::transcript::is_terminal;
2use super::{Agent, AgentRunResult, RunError};
3use crate::Task;
4use crate::containers::{Container, ContainerError};
5use aether_core::events::AgentMessage;
6use async_stream::try_stream;
7use futures::Stream;
8use std::collections::{BTreeMap, HashMap};
9use tokio::io::AsyncBufReadExt;
10
11/// Environment variable through which `DockerAgent` hands the caller-provided task prompt to the
12/// in-container command.
13pub const AETHER_EVAL_TASK_PROMPT_ENV: &str = "AETHER_EVAL_TASK_PROMPT";
14pub const AETHER_EVAL_WORKSPACE_ROOT_ENV: &str = "AETHER_EVAL_WORKSPACE_ROOT";
15pub const AETHER_EVAL_CWD_ENV: &str = "AETHER_EVAL_CWD";
16
17pub struct DockerAgent {
18    container: Container,
19    command: Vec<String>,
20    env_vars: HashMap<String, String>,
21}
22
23impl DockerAgent {
24    /// `command` is the in-container argv to exec; its stdout must be newline-delimited
25    /// `AgentMessage` JSON.
26    pub fn new(container: Container, command: Vec<String>) -> Self {
27        Self { container, command, env_vars: HashMap::new() }
28    }
29
30    pub fn with_env_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
31        self.env_vars.insert(key.into(), value.into());
32        self
33    }
34
35    pub fn with_env_vars(mut self, env_vars: impl IntoIterator<Item = (String, String)>) -> Self {
36        self.env_vars.extend(env_vars);
37        self
38    }
39}
40
41impl Agent for DockerAgent {
42    fn run(&self, task: Task) -> impl Stream<Item = AgentRunResult> + Send + '_ {
43        let env = {
44            let mut env: BTreeMap<_, _> = self.env_vars.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
45            env.insert(AETHER_EVAL_TASK_PROMPT_ENV.to_string(), task.prompt().to_string());
46            env.insert(AETHER_EVAL_CWD_ENV.to_string(), self.container.cwd().display().to_string());
47            env.insert(
48                AETHER_EVAL_WORKSPACE_ROOT_ENV.to_string(),
49                self.container.workspace_root().display().to_string(),
50            );
51            env
52        };
53
54        try_stream! {
55            let mut exec = self.container.exec_streaming(self.command.clone(), env).await?;
56            let finished = {
57                let mut lines = exec.stdout().lines();
58                let mut finished = false;
59                while let Some(line) = lines.next_line().await.map_err(|source| ContainerError::StdoutRead { source })? {
60                    if line.trim().is_empty() {
61                        continue;
62                    }
63                    tracing::debug!("docker agent stdout: {line}");
64                    let message: AgentMessage = serde_json::from_str(&line)
65                        .map_err(|source| RunError::AgentMessageJsonLine { line, source })?;
66                    let terminal = is_terminal(&message);
67                    yield message;
68                    if terminal {
69                        finished = true;
70                        break;
71                    }
72                }
73                finished
74            };
75
76            let stderr = exec.stderr_to_string().await?;
77            if !stderr.trim().is_empty() {
78                tracing::debug!("docker agent stderr: {}", stderr.trim());
79            }
80
81            if !finished {
82                Err::<AgentMessage, RunError>(RunError::CommandExitWithoutTerminal { stderr })?;
83            }
84        }
85    }
86}
aether_evals/agents/docker_agent.rs

aether_evals/agents/
docker_agent.rs
