Skip to main content

aether_evals/agents/
docker_aether_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 aether_project::AetherSettings;
6use llm::LlmModel;
7use std::path::{Path, PathBuf};
8use testcontainers::core::Mount;
9use tokio::io::{AsyncBufRead, AsyncBufReadExt};
10use tokio::sync::mpsc::Sender;
11
12pub struct DockerAetherAgent {
13    image: DockerImage,
14    settings: Option<AetherSettings>,
15    agent: Option<String>,
16}
17
18impl DockerAetherAgent {
19    pub fn new(image: DockerImage) -> Self {
20        Self { image, settings: None, agent: None }
21    }
22
23    pub fn with_settings(mut self, settings: AetherSettings) -> Self {
24        self.settings = Some(settings);
25        self
26    }
27
28    pub fn with_agent(mut self, agent: impl Into<String>) -> Self {
29        self.agent = Some(agent.into());
30        self
31    }
32
33    fn build_command(&self, task_prompt: &str, container_cwd: &Path) -> Vec<String> {
34        let mut command = vec![
35            "aether".to_string(),
36            "headless".to_string(),
37            "--cwd".to_string(),
38            container_cwd.display().to_string(),
39        ];
40
41        if let Some(agent) = &self.agent {
42            command.push("--agent".to_string());
43            command.push(agent.clone());
44        }
45
46        command.extend(["--output".to_string(), "json".to_string()]);
47
48        if let Some(settings) = &self.settings {
49            command.push("--settings-json".to_string());
50            command.push(serde_json::to_string(settings).expect("AetherSettings should serialize to JSON"));
51        }
52        command.push(build_task_prompt(task_prompt, container_cwd));
53        command
54    }
55}
56
57impl Agent for DockerAetherAgent {
58    async fn run(&self, config: AgentConfig<'_>, tx: Sender<AgentMessage>) -> Result<(), RunError> {
59        let aether_home = tempfile::tempdir().map_err(|source| DockerError::AetherHomeTempDir { source })?;
60        let docker = {
61            let mut docker = Docker::new(self.image.clone())
62                .with_env_var("AETHER_HOME", "/root/.aether")
63                .with_mount(Mount::bind_mount(aether_home.path().display().to_string(), "/root/.aether"));
64
65            for (key, value) in std::env::vars().filter(|(key, _)| {
66                key != "AETHER_HOME"
67                    && (LlmModel::ALL_REQUIRED_ENV_VARS.contains(&key.as_str())
68                        || key == "OLLAMA_HOST"
69                        || key.starts_with("AETHER_"))
70            }) {
71                docker = docker.with_env_var(key, value);
72            }
73
74            docker
75        };
76
77        let container_cwd = config
78            .relative_cwd
79            .map_or_else(|| PathBuf::from("/workspace"), |relative_cwd| Path::new("/workspace").join(relative_cwd));
80
81        let command = self.build_command(config.task_prompt, &container_cwd);
82        let (sent_terminal, stderr) = {
83            let mut session = docker.exec(DockerExecConfig { workspace_root: config.workspace_root, command }).await?;
84            let sent_terminal = forward_stdout_messages(session.stdout(), &tx).await?;
85
86            (sent_terminal, session.stderr_to_string().await?)
87        };
88
89        if !stderr.trim().is_empty() {
90            tracing::debug!("aether headless stderr: {}", stderr.trim());
91        }
92
93        if sent_terminal { Ok(()) } else { Err(DockerError::HeadlessExit { stderr }.into()) }
94    }
95}
96
97async fn forward_stdout_messages<T: AsyncBufRead + Unpin>(
98    reader: T,
99    tx: &Sender<AgentMessage>,
100) -> Result<bool, RunError> {
101    let mut lines = reader.lines();
102    while let Some(line) = lines.next_line().await.map_err(|source| DockerError::StdoutRead { source })? {
103        if line.trim().is_empty() {
104            continue;
105        }
106        tracing::debug!("aether headless: {line}");
107        let message: AgentMessage =
108            serde_json::from_str(&line).map_err(|source| DockerError::HeadlessJsonLine { line, source })?;
109        let terminal = is_terminal(&message);
110        tx.send(message).await.map_err(|error| RunError::ChannelSendFailed(error.to_string()))?;
111        if terminal {
112            return Ok(true);
113        }
114    }
115
116    Ok(false)
117}