Skip to main content

aether_evals/evals/
task.rs

1use super::diff::GitDiff;
2use super::transcript::{Transcript, format_transcript};
3use super::workspace::Workspace;
4use crate::EvalRunError;
5use crate::agents::{Agent, AgentConfig, RunError, is_terminal};
6use aether_core::events::AgentMessage;
7use std::fmt::{Debug, Write as _};
8use thiserror::Error;
9use tokio::sync::mpsc;
10
11pub struct Task {
12    prompt: String,
13    workspace: Workspace,
14}
15
16pub struct TaskRun {
17    prompt: String,
18    workspace: Workspace,
19    transcript: Transcript,
20}
21
22#[derive(Error)]
23#[error("{error}")]
24pub struct TaskRunError {
25    run: TaskRun,
26    #[source]
27    error: EvalRunError,
28}
29
30impl Task {
31    pub fn new(prompt: impl Into<String>, workspace: Workspace) -> Self {
32        Self { prompt: prompt.into(), workspace }
33    }
34
35    pub fn prompt(&self) -> &str {
36        &self.prompt
37    }
38
39    pub fn workspace(&self) -> &Workspace {
40        &self.workspace
41    }
42
43    /// Run the task against `agent` and collect its transcript. Use [`Task::run_observed`] to also
44    /// observe each `AgentMessage` as it streams.
45    #[tracing::instrument(skip(agent, self))]
46    pub async fn run<T: Agent>(self, agent: &T) -> Result<TaskRun, TaskRunError> {
47        self.run_observed(agent, |_| {}).await
48    }
49
50    /// Run the task against `agent`, forwarding each `AgentMessage` to `on_message` before it is
51    /// added to the transcript. Terminal messages (`Done`, `Error`, `Cancelled`) are forwarded too.
52    #[tracing::instrument(skip(agent, self, on_message))]
53    pub async fn run_observed<T: Agent, U: FnMut(&AgentMessage) + Send>(
54        self,
55        agent: &T,
56        mut on_message: U,
57    ) -> Result<TaskRun, TaskRunError> {
58        let Task { prompt, workspace } = self;
59        let (tx, mut rx) = mpsc::channel(100);
60        let agent_task = agent.run(
61            AgentConfig {
62                workspace_root: workspace.root_path(),
63                relative_cwd: workspace.relative_cwd(),
64                task_prompt: &prompt,
65            },
66            tx,
67        );
68
69        let message_task = async {
70            let mut messages = Vec::new();
71            while let Some(message) = rx.recv().await {
72                let is_done = is_terminal(&message);
73                on_message(&message);
74                messages.push(message);
75                if is_done {
76                    break;
77                }
78            }
79            messages
80        };
81
82        let (run_result, messages) = tokio::join!(agent_task, message_task);
83        let run = TaskRun::new(prompt, workspace, Transcript::new(messages));
84        match run_result {
85            Ok(()) => Ok(run),
86            Err(error) => Err(TaskRunError::new(run, error)),
87        }
88    }
89}
90
91impl TaskRun {
92    pub(crate) fn new(prompt: String, workspace: Workspace, transcript: Transcript) -> Self {
93        Self { prompt, workspace, transcript }
94    }
95
96    pub fn prompt(&self) -> &str {
97        &self.prompt
98    }
99
100    pub fn workspace(&self) -> &Workspace {
101        &self.workspace
102    }
103
104    pub fn transcript(&self) -> &Transcript {
105        &self.transcript
106    }
107
108    pub fn into_workspace(self) -> Workspace {
109        self.workspace
110    }
111
112    pub fn failure_context(&self) -> String {
113        let mut summary = String::new();
114        let _ = writeln!(summary, "Task failure context");
115        let _ = writeln!(summary, "Workspace: {}", self.workspace.path().display());
116        summary.push_str("Prompt:\n");
117        push_indented(&mut summary, &self.prompt, 2);
118        summary.push('\n');
119
120        let (agent_diff, reference_diff) = self.workspace.capture_git_diffs();
121        if let Some(diff) = &agent_diff {
122            summary.push_str("Agent diff summary:\n");
123            push_diff_stats(&mut summary, diff);
124            summary.push('\n');
125        }
126
127        if let Some(diff) = &reference_diff {
128            summary.push_str("Reference diff summary:\n");
129            push_diff_stats(&mut summary, diff);
130            summary.push('\n');
131        }
132
133        summary.push_str("Agent messages:\n");
134        if self.transcript.messages().is_empty() {
135            summary.push_str("  none\n");
136        } else {
137            push_indented(&mut summary, &format_transcript(self.transcript.messages()), 2);
138        }
139
140        summary
141    }
142}
143
144impl TaskRunError {
145    fn new(run: TaskRun, error: RunError) -> Self {
146        Self { run, error: EvalRunError::from(error) }
147    }
148
149    pub fn run(&self) -> &TaskRun {
150        &self.run
151    }
152
153    pub fn error(&self) -> &EvalRunError {
154        &self.error
155    }
156
157    pub fn into_parts(self) -> (TaskRun, EvalRunError) {
158        (self.run, self.error)
159    }
160}
161
162impl Debug for TaskRunError {
163    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164        formatter.debug_struct("TaskRunError").field("error", &self.error).finish_non_exhaustive()
165    }
166}
167
168impl From<TaskRunError> for EvalRunError {
169    fn from(error: TaskRunError) -> Self {
170        error.error
171    }
172}
173
174fn push_indented(output: &mut String, value: &str, spaces: usize) {
175    let indentation = " ".repeat(spaces);
176    for line in value.lines() {
177        output.push_str(&indentation);
178        output.push_str(line);
179        output.push('\n');
180    }
181}
182
183fn push_diff_stats(output: &mut String, diff: &GitDiff) {
184    let _ = writeln!(output, "  Files changed: {}", diff.stats.files_changed);
185    let _ = writeln!(output, "  Lines added: {}", diff.stats.lines_added);
186    let _ = writeln!(output, "  Lines removed: {}", diff.stats.lines_removed);
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::{FakeAgent, RunError};
193    use aether_core::events::AgentMessage;
194    use std::sync::{Arc, Mutex};
195    use tokio::sync::mpsc::Sender;
196
197    #[tokio::test]
198    async fn task_run_contains_prompt_workspace_and_transcript() {
199        let run = Task::new("do the thing", Workspace::empty().unwrap())
200            .run(&FakeAgent::with_tool_call("bash", "success"))
201            .await
202            .unwrap();
203
204        assert_eq!(run.prompt(), "do the thing");
205        assert!(run.workspace().path().exists());
206        assert!(run.transcript().tool_called("bash"));
207        assert!(matches!(run.transcript().messages().last(), Some(AgentMessage::Done)));
208    }
209
210    #[tokio::test]
211    async fn task_run_passes_raw_prompt_without_scaffolding() {
212        let agent = CapturingAgent::default();
213        Task::new("do the thing", Workspace::empty().unwrap()).run(&agent).await.unwrap();
214
215        assert_eq!(agent.captured_task_prompt(), Some("do the thing".to_string()));
216    }
217
218    #[tokio::test]
219    async fn task_run_with_message_observer_forwards_messages_in_order() {
220        let seen = Arc::new(Mutex::new(Vec::new()));
221        let captured = seen.clone();
222        let run = Task::new("do the thing", Workspace::empty().unwrap())
223            .run_observed(&FakeAgent::with_tool_call("bash", "success"), move |message| {
224                captured.lock().unwrap().push(message.clone());
225            })
226            .await
227            .unwrap();
228
229        let agent_messages = run.transcript().messages().to_vec();
230        let observed = seen.lock().unwrap().clone();
231
232        assert_eq!(observed.len(), agent_messages.len());
233        for (observed, transcript) in observed.iter().zip(agent_messages.iter()) {
234            assert_eq!(observed, transcript);
235        }
236        assert!(matches!(observed.last(), Some(AgentMessage::Done)));
237    }
238
239    #[derive(Default)]
240    struct CapturingAgent {
241        task_prompt: Arc<Mutex<Option<String>>>,
242    }
243
244    impl CapturingAgent {
245        fn captured_task_prompt(&self) -> Option<String> {
246            self.task_prompt.lock().unwrap().clone()
247        }
248    }
249
250    impl Agent for CapturingAgent {
251        async fn run(&self, config: AgentConfig<'_>, tx: Sender<AgentMessage>) -> Result<(), RunError> {
252            *self.task_prompt.lock().unwrap() = Some(config.task_prompt.to_string());
253            tx.send(AgentMessage::Done).await.map_err(|e| RunError::ChannelSendFailed(e.to_string()))
254        }
255    }
256}