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 #[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 #[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}