Skip to main content

aether_evals/evals/
workspace.rs

1use super::report::{DiffStats, GitDiff};
2use crate::WorkspaceError;
3use crate::git_repo::GitRepo;
4use std::path::{Path, PathBuf};
5use tempfile::TempDir;
6
7pub struct Workspace {
8    path: PathBuf,
9    root_path: PathBuf,
10    relative_cwd: Option<PathBuf>,
11    source: WorkspaceSource,
12    _drop_guard: TempDir,
13}
14
15#[derive(Debug, Clone)]
16pub enum WorkspaceSource {
17    Local,
18    GitRepo { url: String, start_commit: String, gold_commit: String },
19}
20
21#[derive(Debug, Clone, serde::Deserialize)]
22#[serde(rename_all = "camelCase", deny_unknown_fields)]
23pub struct GitRepoSpec {
24    pub url: String,
25    pub start_commit: String,
26    pub gold_commit: String,
27    #[serde(default)]
28    pub subdir: Option<PathBuf>,
29}
30
31impl Workspace {
32    pub fn empty() -> Result<Self, WorkspaceError> {
33        let temp_dir = new_temp_dir()?;
34        let path = temp_dir.path().to_path_buf();
35        Ok(Self {
36            path: path.clone(),
37            root_path: path,
38            relative_cwd: None,
39            source: WorkspaceSource::Local,
40            _drop_guard: temp_dir,
41        })
42    }
43
44    pub fn from_dir(src_path: impl Into<PathBuf>) -> Result<Self, WorkspaceError> {
45        let src_path = src_path.into();
46        let temp_dir = new_temp_dir()?;
47        let path = temp_dir.path().to_path_buf();
48
49        copy_dir_contents(&src_path, &path).map_err(|source| WorkspaceError::CopyFixture {
50            from: src_path.clone(),
51            to: path.clone(),
52            source,
53        })?;
54
55        Ok(Self {
56            path: path.clone(),
57            root_path: path,
58            relative_cwd: None,
59            source: WorkspaceSource::Local,
60            _drop_guard: temp_dir,
61        })
62    }
63
64    pub fn from_files<T, U>(files: impl IntoIterator<Item = (T, U)>) -> Result<Self, WorkspaceError>
65    where
66        T: AsRef<Path>,
67        U: AsRef<str>,
68    {
69        let workspace = Self::empty()?;
70        for (relative_path, contents) in files {
71            let path = workspace.path().join(relative_path.as_ref());
72            if let Some(parent) = path.parent() {
73                std::fs::create_dir_all(parent)
74                    .map_err(|source| WorkspaceError::WriteFile { path: parent.to_path_buf(), source })?;
75            }
76            std::fs::write(&path, contents.as_ref()).map_err(|source| WorkspaceError::WriteFile { path, source })?;
77        }
78        Ok(workspace)
79    }
80
81    pub fn from_git_repo(spec: GitRepoSpec) -> Result<Self, WorkspaceError> {
82        let GitRepoSpec { url, start_commit, gold_commit, subdir } = spec;
83        let temp_dir = new_temp_dir()?;
84
85        tracing::debug!("Cloning git repo {} at commit {}", url, start_commit);
86        let repo = GitRepo::clone(&url, temp_dir.path())?;
87        repo.checkout(&start_commit)?;
88
89        let root_path = temp_dir.path().to_path_buf();
90        let (path, relative_cwd) = match subdir {
91            None => (root_path.clone(), None),
92            Some(relative_cwd) => {
93                let working_path = root_path.join(&relative_cwd);
94                if !working_path.exists() {
95                    return Err(WorkspaceError::MissingSubdir { path: working_path });
96                }
97                (working_path, Some(relative_cwd))
98            }
99        };
100
101        Ok(Self {
102            path,
103            root_path,
104            relative_cwd,
105            source: WorkspaceSource::GitRepo { url, start_commit, gold_commit },
106            _drop_guard: temp_dir,
107        })
108    }
109
110    pub fn path(&self) -> &Path {
111        &self.path
112    }
113
114    pub fn root_path(&self) -> &Path {
115        &self.root_path
116    }
117
118    pub fn relative_cwd(&self) -> Option<&Path> {
119        self.relative_cwd.as_deref()
120    }
121
122    pub fn source(&self) -> &WorkspaceSource {
123        &self.source
124    }
125
126    pub fn capture_git_diffs(&self) -> (Option<GitDiff>, Option<GitDiff>) {
127        let WorkspaceSource::GitRepo { start_commit, gold_commit, .. } = self.source() else {
128            return (None, None);
129        };
130
131        let repo = GitRepo::from_path(self.path());
132        let agent_diff = repo.diff_unstaged().ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
133        let reference_diff =
134            repo.diff(start_commit, gold_commit).ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
135
136        (agent_diff, reference_diff)
137    }
138}
139
140fn new_temp_dir() -> Result<tempfile::TempDir, WorkspaceError> {
141    tempfile::tempdir().map_err(WorkspaceError::CreateTempDir)
142}
143
144fn copy_dir_contents(src: &Path, dst: &Path) -> std::io::Result<()> {
145    for entry in std::fs::read_dir(src)? {
146        let entry = entry?;
147        let source_path = entry.path();
148        let dest_path = dst.join(entry.file_name());
149        let file_type = entry.file_type()?;
150
151        if file_type.is_dir() {
152            std::fs::create_dir_all(&dest_path)?;
153            copy_dir_contents(&source_path, &dest_path)?;
154        } else if file_type.is_file() {
155            std::fs::copy(&source_path, &dest_path)?;
156        }
157    }
158    Ok(())
159}