aether_evals/evals/

task.rs

use super::diff::GitDiff;
use super::transcript::{Transcript, format_transcript};
use super::workspace::Workspace;
use crate::EvalRunError;
use crate::agents::{Agent, AgentConfig, RunError, is_terminal};
use std::fmt::{Debug, Write as _};
use thiserror::Error;
use tokio::sync::mpsc;

pub struct Task {
    prompt: String,
    workspace: Workspace,
}

pub struct TaskRun {
    prompt: String,
    workspace: Workspace,
    transcript: Transcript,
}

#[derive(Error)]
#[error("{error}")]
pub struct TaskRunError {
    run: TaskRun,
    #[source]
    error: EvalRunError,
}

impl Task {
    pub fn new(prompt: impl Into<String>, workspace: Workspace) -> Self {
        Self { prompt: prompt.into(), workspace }
    }

    pub fn prompt(&self) -> &str {
        &self.prompt
    }

    pub fn workspace(&self) -> &Workspace {
        &self.workspace
    }

    #[tracing::instrument(skip(agent, self))]
    pub async fn run<T: Agent>(self, agent: &T) -> Result<TaskRun, TaskRunError> {
        let Task { prompt, workspace } = self;
        let (tx, mut rx) = mpsc::channel(100);
        let agent_task = agent.run(
            AgentConfig {
                workspace_root: workspace.root_path(),
                relative_cwd: workspace.relative_cwd(),
                task_prompt: &prompt,
            },
            tx,
        );

        let message_task = async {
            let mut messages = Vec::new();
            while let Some(message) = rx.recv().await {
                let is_done = is_terminal(&message);
                messages.push(message);
                if is_done {
                    break;
                }
            }
            messages
        };

        let (run_result, messages) = tokio::join!(agent_task, message_task);
        let run = TaskRun::new(prompt, workspace, Transcript::new(messages));
        match run_result {
            Ok(()) => Ok(run),
            Err(error) => Err(TaskRunError::new(run, error)),
        }
    }
}

impl TaskRun {
    pub(crate) fn new(prompt: String, workspace: Workspace, transcript: Transcript) -> Self {
        Self { prompt, workspace, transcript }
    }

    pub fn prompt(&self) -> &str {
        &self.prompt
    }

    pub fn workspace(&self) -> &Workspace {
        &self.workspace
    }

    pub fn transcript(&self) -> &Transcript {
        &self.transcript
    }

    pub fn into_workspace(self) -> Workspace {
        self.workspace
    }

    pub fn failure_context(&self) -> String {
        let mut summary = String::new();
        let _ = writeln!(summary, "Task failure context");
        let _ = writeln!(summary, "Workspace: {}", self.workspace.path().display());
        summary.push_str("Prompt:\n");
        push_indented(&mut summary, &self.prompt, 2);
        summary.push('\n');

        let (agent_diff, reference_diff) = self.workspace.capture_git_diffs();
        if let Some(diff) = &agent_diff {
            summary.push_str("Agent diff summary:\n");
            push_diff_stats(&mut summary, diff);
            summary.push('\n');
        }

        if let Some(diff) = &reference_diff {
            summary.push_str("Reference diff summary:\n");
            push_diff_stats(&mut summary, diff);
            summary.push('\n');
        }

        summary.push_str("Agent messages:\n");
        if self.transcript.messages().is_empty() {
            summary.push_str("  none\n");
        } else {
            push_indented(&mut summary, &format_transcript(self.transcript.messages()), 2);
        }

        summary
    }
}

impl TaskRunError {
    fn new(run: TaskRun, error: RunError) -> Self {
        Self { run, error: EvalRunError::from(error) }
    }

    pub fn run(&self) -> &TaskRun {
        &self.run
    }

    pub fn error(&self) -> &EvalRunError {
        &self.error
    }

    pub fn into_parts(self) -> (TaskRun, EvalRunError) {
        (self.run, self.error)
    }
}

impl Debug for TaskRunError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.debug_struct("TaskRunError").field("error", &self.error).finish_non_exhaustive()
    }
}

impl From<TaskRunError> for EvalRunError {
    fn from(error: TaskRunError) -> Self {
        error.error
    }
}

fn push_indented(output: &mut String, value: &str, spaces: usize) {
    let indentation = " ".repeat(spaces);
    for line in value.lines() {
        output.push_str(&indentation);
        output.push_str(line);
        output.push('\n');
    }
}

fn push_diff_stats(output: &mut String, diff: &GitDiff) {
    let _ = writeln!(output, "  Files changed: {}", diff.stats.files_changed);
    let _ = writeln!(output, "  Lines added: {}", diff.stats.lines_added);
    let _ = writeln!(output, "  Lines removed: {}", diff.stats.lines_removed);
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{FakeAgent, RunError};
    use aether_core::events::AgentMessage;
    use std::sync::{Arc, Mutex};
    use tokio::sync::mpsc::Sender;

    #[tokio::test]
    async fn task_run_contains_prompt_workspace_and_transcript() {
        let run = Task::new("do the thing", Workspace::empty().unwrap())
            .run(&FakeAgent::with_tool_call("bash", "success"))
            .await
            .unwrap();

        assert_eq!(run.prompt(), "do the thing");
        assert!(run.workspace().path().exists());
        assert!(run.transcript().tool_called("bash"));
        assert!(matches!(run.transcript().messages().last(), Some(AgentMessage::Done)));
    }

    #[tokio::test]
    async fn task_run_passes_raw_prompt_without_scaffolding() {
        let agent = CapturingAgent::default();
        Task::new("do the thing", Workspace::empty().unwrap()).run(&agent).await.unwrap();

        assert_eq!(agent.captured_task_prompt(), Some("do the thing".to_string()));
    }

    #[derive(Default)]
    struct CapturingAgent {
        task_prompt: Arc<Mutex<Option<String>>>,
    }

    impl CapturingAgent {
        fn captured_task_prompt(&self) -> Option<String> {
            self.task_prompt.lock().unwrap().clone()
        }
    }

    impl Agent for CapturingAgent {
        async fn run(&self, config: AgentConfig<'_>, tx: Sender<AgentMessage>) -> Result<(), RunError> {
            *self.task_prompt.lock().unwrap() = Some(config.task_prompt.to_string());
            tx.send(AgentMessage::Done).await.map_err(|e| RunError::ChannelSendFailed(e.to_string()))
        }
    }
}