Skip to main content

aether_evals/evals/
workspace.rs

1use super::diff::{DiffStats, GitDiff};
2use crate::WorkspaceError;
3use crate::git_repo::GitRepo;
4use schemars::JsonSchema;
5use serde::Serialize;
6use std::{
7    fs::{create_dir_all, write},
8    path::{Path, PathBuf},
9};
10use tempfile::TempDir;
11
12pub struct Workspace {
13    path: PathBuf,
14    root_path: PathBuf,
15    relative_cwd: Option<PathBuf>,
16    source: WorkspaceSource,
17    temp_dir: TempDir,
18}
19
20#[derive(Debug, Clone)]
21pub enum WorkspaceSource {
22    Local,
23    GitRepo { url: String, start_commit: String, gold_commit: String },
24    Bundle { start_commit: String, gold_commit: String },
25}
26
27#[derive(Debug, Clone, serde::Deserialize, JsonSchema)]
28#[serde(rename_all = "camelCase", deny_unknown_fields)]
29pub struct GitRepoSpec {
30    pub url: String,
31    pub start_commit: String,
32    pub gold_commit: String,
33    #[serde(default)]
34    pub subdir: Option<PathBuf>,
35}
36
37#[derive(Debug, Clone, serde::Deserialize, JsonSchema)]
38#[serde(rename_all = "camelCase", deny_unknown_fields)]
39pub struct GitBundleSpec {
40    pub bundle_path: PathBuf,
41    pub start_commit: String,
42    pub gold_commit: String,
43    #[serde(default)]
44    pub subdir: Option<PathBuf>,
45}
46
47#[derive(Debug, Clone, Serialize, JsonSchema)]
48#[serde(rename_all = "camelCase", deny_unknown_fields)]
49pub struct RetainedWorkspaceInfo {
50    pub root_path: PathBuf,
51    pub path: PathBuf,
52}
53
54impl Workspace {
55    pub fn empty() -> Result<Self, WorkspaceError> {
56        let temp_dir = new_temp_dir()?;
57        let path = temp_dir.path().to_path_buf();
58        Ok(Self { path: path.clone(), root_path: path, relative_cwd: None, source: WorkspaceSource::Local, temp_dir })
59    }
60
61    pub fn from_dir(src_path: impl Into<PathBuf>) -> Result<Self, WorkspaceError> {
62        let src_path = src_path.into();
63        let temp_dir = new_temp_dir()?;
64        let path = temp_dir.path().to_path_buf();
65
66        copy_dir_contents(&src_path, &path).map_err(|source| WorkspaceError::CopyFixture {
67            from: src_path.clone(),
68            to: path.clone(),
69            source,
70        })?;
71
72        Ok(Self { path: path.clone(), root_path: path, relative_cwd: None, source: WorkspaceSource::Local, temp_dir })
73    }
74
75    pub fn from_files<T: AsRef<Path>, U: AsRef<str>>(
76        files: impl IntoIterator<Item = (T, U)>,
77    ) -> Result<Self, WorkspaceError> {
78        let workspace = Self::empty()?;
79        for (relative_path, contents) in files {
80            let path = workspace.path().join(relative_path.as_ref());
81            if let Some(parent) = path.parent() {
82                create_dir_all(parent)
83                    .map_err(|source| WorkspaceError::WriteFile { path: parent.to_path_buf(), source })?;
84            }
85            write(&path, contents.as_ref()).map_err(|source| WorkspaceError::WriteFile { path, source })?;
86        }
87        Ok(workspace)
88    }
89
90    pub fn from_git_repo(spec: GitRepoSpec) -> Result<Self, WorkspaceError> {
91        let GitRepoSpec { url, start_commit, gold_commit, subdir } = spec;
92        let temp_dir = new_temp_dir()?;
93
94        tracing::debug!("Cloning git repo {} at commit {}", url, start_commit);
95        let repo = GitRepo::clone(&url, temp_dir.path(), true)?;
96
97        let source = WorkspaceSource::GitRepo { url, start_commit: start_commit.clone(), gold_commit };
98        Self::from_cloned(temp_dir, &repo, &start_commit, subdir, source)
99    }
100
101    /// Instantiate a workspace from a local git bundle file
102    pub fn from_git_bundle(spec: GitBundleSpec) -> Result<Self, WorkspaceError> {
103        let GitBundleSpec { bundle_path, start_commit, gold_commit, subdir } = spec;
104        if !bundle_path.exists() {
105            return Err(WorkspaceError::MissingBundle { path: bundle_path });
106        }
107
108        tracing::debug!("Cloning git bundle {} at commit {}", bundle_path.display(), start_commit);
109        let temp_dir = new_temp_dir()?;
110        let repo = GitRepo::clone(&bundle_path, temp_dir.path(), false)?;
111        let source = WorkspaceSource::Bundle { start_commit: start_commit.clone(), gold_commit };
112        Self::from_cloned(temp_dir, &repo, &start_commit, subdir, source)
113    }
114
115    fn from_cloned(
116        temp_dir: TempDir,
117        repo: &GitRepo,
118        start_commit: &str,
119        subdir: Option<PathBuf>,
120        source: WorkspaceSource,
121    ) -> Result<Self, WorkspaceError> {
122        repo.checkout(start_commit)?;
123        let root_path = temp_dir.path().to_path_buf();
124        let (path, relative_cwd) = resolve_subdir(&root_path, subdir)?;
125        Ok(Self { path, root_path, relative_cwd, source, temp_dir })
126    }
127
128    pub fn path(&self) -> &Path {
129        &self.path
130    }
131
132    pub fn join(&self, relative_path: impl AsRef<Path>) -> PathBuf {
133        self.path.join(relative_path)
134    }
135
136    pub fn root_path(&self) -> &Path {
137        &self.root_path
138    }
139
140    pub fn relative_cwd(&self) -> Option<&Path> {
141        self.relative_cwd.as_deref()
142    }
143
144    pub fn source(&self) -> &WorkspaceSource {
145        &self.source
146    }
147
148    /// Prevents the workspace from getting automatically removed and returns its retained root and effective cwd. The caller is responsible for cleanup.
149    pub fn persist(self) -> RetainedWorkspaceInfo {
150        let root_path = self.temp_dir.keep();
151        let path = self.relative_cwd.map_or_else(|| root_path.clone(), |relative_cwd| root_path.join(relative_cwd));
152        RetainedWorkspaceInfo { root_path, path }
153    }
154
155    pub fn capture_git_diffs(&self) -> (Option<GitDiff>, Option<GitDiff>) {
156        let Some((start_commit, gold_commit)) = self.diff_commits() else {
157            return (None, None);
158        };
159
160        let repo = GitRepo::from_path(self.path());
161        let agent_diff = repo.diff_unstaged().ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
162        let reference_diff =
163            repo.diff(start_commit, gold_commit).ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
164
165        (agent_diff, reference_diff)
166    }
167
168    fn diff_commits(&self) -> Option<(&str, &str)> {
169        match &self.source {
170            WorkspaceSource::Local => None,
171            WorkspaceSource::GitRepo { start_commit, gold_commit, .. }
172            | WorkspaceSource::Bundle { start_commit, gold_commit } => Some((start_commit, gold_commit)),
173        }
174    }
175}
176
177const EVAL_START_REF: &str = "eval-start";
178const EVAL_GOLD_REF: &str = "eval-gold";
179
180/// Create a self-contained git bundle at `out` containing `spec`'s start and gold commits.
181pub fn create_git_bundle(spec: &GitRepoSpec, out: &Path) -> Result<(), WorkspaceError> {
182    let temp_dir = new_temp_dir()?;
183    let repo = GitRepo::init(temp_dir.path())?;
184    repo.fetch(&spec.url, &[&spec.start_commit, &spec.gold_commit])?;
185    repo.update_ref(&format!("refs/heads/{EVAL_START_REF}"), &spec.start_commit)?;
186    repo.update_ref(&format!("refs/heads/{EVAL_GOLD_REF}"), &spec.gold_commit)?;
187    repo.bundle(&[EVAL_START_REF, EVAL_GOLD_REF], out)?;
188    Ok(())
189}
190
191fn new_temp_dir() -> Result<tempfile::TempDir, WorkspaceError> {
192    tempfile::tempdir().map_err(WorkspaceError::CreateTempDir)
193}
194
195fn resolve_subdir(root_path: &Path, subdir: Option<PathBuf>) -> Result<(PathBuf, Option<PathBuf>), WorkspaceError> {
196    match subdir {
197        None => Ok((root_path.to_path_buf(), None)),
198        Some(relative_cwd) => {
199            let working_path = root_path.join(&relative_cwd);
200            if !working_path.exists() {
201                return Err(WorkspaceError::MissingSubdir { path: working_path });
202            }
203            Ok((working_path, Some(relative_cwd)))
204        }
205    }
206}
207
208fn copy_dir_contents(src: &Path, dst: &Path) -> std::io::Result<()> {
209    for entry in std::fs::read_dir(src)? {
210        let entry = entry?;
211        let source_path = entry.path();
212        let dest_path = dst.join(entry.file_name());
213        let file_type = entry.file_type()?;
214
215        if file_type.is_dir() {
216            std::fs::create_dir_all(&dest_path)?;
217            copy_dir_contents(&source_path, &dest_path)?;
218        } else if file_type.is_file() {
219            std::fs::copy(&source_path, &dest_path)?;
220        }
221    }
222    Ok(())
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use std::fs::{read_to_string, remove_dir_all};
229    use std::process::Command;
230    use tempfile::TempDir;
231
232    #[test]
233    fn persist_reports_root_and_path_for_local_workspace() {
234        let workspace = Workspace::from_files([("notes.txt", "hi\n")]).unwrap();
235
236        let retained = workspace.persist();
237
238        assert_eq!(retained.root_path, retained.path);
239        assert_eq!(read_to_string(retained.path.join("notes.txt")).unwrap(), "hi\n");
240        remove_dir_all(retained.root_path).unwrap();
241    }
242
243    #[test]
244    fn from_git_bundle_round_trips_checkout_subdir_and_diffs() {
245        let (repo, start, gold) = init_repo();
246        let bundle_dir = tempfile::tempdir().unwrap();
247        let bundle_path = bundle_dir.path().join("repo.bundle");
248        create_git_bundle(
249            &GitRepoSpec {
250                url: format!("file://{}", repo.path().display()),
251                start_commit: start.clone(),
252                gold_commit: gold.clone(),
253                subdir: None,
254            },
255            &bundle_path,
256        )
257        .unwrap();
258
259        let workspace = Workspace::from_git_bundle(GitBundleSpec {
260            bundle_path,
261            start_commit: start,
262            gold_commit: gold,
263            subdir: Some(PathBuf::from("pkg")),
264        })
265        .unwrap();
266
267        let (agent_diff, reference_diff) = workspace.capture_git_diffs();
268
269        assert_eq!(read_to_string(workspace.root_path().join("root.txt")).unwrap(), "root v1\n");
270        assert_eq!(workspace.relative_cwd(), Some(Path::new("pkg")));
271        assert_eq!(workspace.path(), workspace.root_path().join("pkg"));
272        assert!(reference_diff.unwrap().diff.contains("root v2"), "reference diff should span start..gold");
273        assert!(agent_diff.unwrap().diff.is_empty(), "no agent edits yet");
274
275        write(workspace.root_path().join("root.txt"), "root edited\n").unwrap();
276        let (agent_diff, _) = workspace.capture_git_diffs();
277
278        assert!(agent_diff.unwrap().diff.contains("root edited"), "agent diff should capture the edit");
279    }
280
281    #[test]
282    fn from_git_bundle_missing_file_errors() {
283        let result = Workspace::from_git_bundle(GitBundleSpec {
284            bundle_path: PathBuf::from("/nonexistent/repo.bundle"),
285            start_commit: "abc".into(),
286            gold_commit: "def".into(),
287            subdir: None,
288        });
289
290        assert!(matches!(result, Err(WorkspaceError::MissingBundle { .. })));
291    }
292
293    fn init_repo() -> (TempDir, String, String) {
294        let dir = tempfile::tempdir().unwrap();
295        let path = dir.path();
296        git(path, &["init", "--initial-branch", "main"]);
297        git(path, &["config", "user.email", "eval@example.com"]);
298        git(path, &["config", "user.name", "Eval"]);
299
300        write(path.join("root.txt"), "root v1\n").unwrap();
301        create_dir_all(path.join("pkg")).unwrap();
302        write(path.join("pkg").join("inner.txt"), "inner v1\n").unwrap();
303        git(path, &["add", "."]);
304        git(path, &["commit", "-m", "start"]);
305        let start = git(path, &["rev-parse", "HEAD"]);
306
307        write(path.join("root.txt"), "root v2\n").unwrap();
308        git(path, &["add", "."]);
309        git(path, &["commit", "-m", "gold"]);
310        let gold = git(path, &["rev-parse", "HEAD"]);
311
312        (dir, start, gold)
313    }
314
315    fn git(repo: &Path, args: &[&str]) -> String {
316        let output = Command::new("git").arg("-C").arg(repo).args(args).output().unwrap();
317        assert!(output.status.success(), "git {args:?} failed: {}", String::from_utf8_lossy(&output.stderr));
318        String::from_utf8(output.stdout).unwrap().trim().to_string()
319    }
320}