Skip to main content

aether_evals/evals/
workspace.rs

1use super::diff::GitDiff;
2use crate::WorkspaceError;
3use crate::git_repo::GitRepo;
4use schemars::JsonSchema;
5use serde::Serialize;
6use std::{
7    ffi::OsStr,
8    fs::{create_dir_all, write},
9    path::{Path, PathBuf},
10};
11use tempfile::TempDir;
12
13pub struct Workspace {
14    path: PathBuf,
15    root_path: PathBuf,
16    relative_cwd: Option<PathBuf>,
17    git: Option<GitBaseline>,
18    temp_dir: TempDir,
19}
20
21#[derive(Debug, Clone, serde::Deserialize, JsonSchema)]
22#[serde(rename_all = "camelCase", deny_unknown_fields)]
23pub struct GitSpec {
24    pub source: GitSource,
25    pub start_commit: String,
26    pub gold_commit: String,
27    #[serde(default)]
28    pub subdir: Option<PathBuf>,
29    #[serde(default)]
30    pub branch: Option<String>,
31    #[serde(default)]
32    pub diff_exclude: Vec<String>,
33}
34
35#[derive(Debug, Clone, serde::Deserialize, JsonSchema)]
36#[serde(rename_all = "camelCase")]
37pub enum GitSource {
38    Url(String),
39    Bundle(PathBuf),
40}
41
42#[derive(Debug, Clone, Serialize, JsonSchema)]
43#[serde(rename_all = "camelCase", deny_unknown_fields)]
44pub struct RetainedWorkspaceInfo {
45    pub root_path: PathBuf,
46    pub path: PathBuf,
47}
48
49impl Workspace {
50    pub fn empty() -> Result<Self, WorkspaceError> {
51        Ok(Self::local(new_temp_dir()?))
52    }
53
54    pub fn from_dir(src_path: impl Into<PathBuf>) -> Result<Self, WorkspaceError> {
55        let src_path = src_path.into();
56        let temp_dir = new_temp_dir()?;
57
58        copy_dir_contents(&src_path, temp_dir.path()).map_err(|source| WorkspaceError::CopyFixture {
59            from: src_path.clone(),
60            to: temp_dir.path().to_path_buf(),
61            source,
62        })?;
63
64        Ok(Self::local(temp_dir))
65    }
66
67    pub fn from_files<T: AsRef<Path>, U: AsRef<str>>(
68        files: impl IntoIterator<Item = (T, U)>,
69    ) -> Result<Self, WorkspaceError> {
70        let workspace = Self::empty()?;
71        for (relative_path, contents) in files {
72            let path = workspace.path().join(relative_path.as_ref());
73            if let Some(parent) = path.parent() {
74                create_dir_all(parent)
75                    .map_err(|source| WorkspaceError::WriteFile { path: parent.to_path_buf(), source })?;
76            }
77            write(&path, contents.as_ref()).map_err(|source| WorkspaceError::WriteFile { path, source })?;
78        }
79        Ok(workspace)
80    }
81
82    pub fn from_git(spec: GitSpec) -> Result<Self, WorkspaceError> {
83        let GitSpec { source, start_commit, gold_commit, subdir, branch, diff_exclude } = spec;
84        if let GitSource::Bundle(path) = &source
85            && !path.exists()
86        {
87            return Err(WorkspaceError::MissingBundle { path: path.clone() });
88        }
89        tracing::debug!("Fetching {:?} at commit {}", source, start_commit);
90
91        let scratch_dir = new_temp_dir()?;
92        let scratch = GitRepo::init(scratch_dir.path())?;
93        scratch.fetch(source.remote(), &[&start_commit, &gold_commit])?;
94        let reference_diff = GitDiff::new(scratch.diff(&start_commit, &gold_commit, &diff_exclude)?);
95
96        let temp_dir = new_temp_dir()?;
97        let repo = GitRepo::init(temp_dir.path())?;
98        repo.fetch(scratch_dir.path(), &[&start_commit])?;
99        repo.checkout_branch(branch.as_deref().unwrap_or(DEFAULT_BRANCH), &start_commit)?;
100
101        let root_path = temp_dir.path().to_path_buf();
102        let (path, relative_cwd) = resolve_subdir(&root_path, subdir)?;
103        let git = GitBaseline { start_commit, reference_diff, diff_exclude };
104        Ok(Self { path, root_path, relative_cwd, git: Some(git), temp_dir })
105    }
106
107    pub fn path(&self) -> &Path {
108        &self.path
109    }
110
111    pub fn join(&self, relative_path: impl AsRef<Path>) -> PathBuf {
112        self.path.join(relative_path)
113    }
114
115    pub fn root_path(&self) -> &Path {
116        &self.root_path
117    }
118
119    pub fn relative_cwd(&self) -> Option<&Path> {
120        self.relative_cwd.as_deref()
121    }
122
123    /// Prevents the workspace from getting automatically removed and returns its retained root and effective cwd. The caller is responsible for cleanup.
124    pub fn persist(self) -> RetainedWorkspaceInfo {
125        let root_path = self.temp_dir.keep();
126        let path = self.relative_cwd.map_or_else(|| root_path.clone(), |relative_cwd| root_path.join(relative_cwd));
127        RetainedWorkspaceInfo { root_path, path }
128    }
129
130    pub fn reference_diff(&self) -> Option<&GitDiff> {
131        self.git.as_ref().map(|git| &git.reference_diff)
132    }
133
134    pub fn capture_agent_diff(&self) -> Result<Option<GitDiff>, WorkspaceError> {
135        let Some(git) = &self.git else {
136            return Ok(None);
137        };
138        let diff = GitRepo::from_path(&self.root_path).diff_working_tree(&git.start_commit, &git.diff_exclude)?;
139        Ok(Some(GitDiff::new(diff)))
140    }
141
142    fn local(temp_dir: TempDir) -> Self {
143        let path = temp_dir.path().to_path_buf();
144        Self { path: path.clone(), root_path: path, relative_cwd: None, git: None, temp_dir }
145    }
146}
147
148const EVAL_START_REF: &str = "eval-start";
149const EVAL_GOLD_REF: &str = "eval-gold";
150const DEFAULT_BRANCH: &str = "main";
151
152/// Create a self-contained git bundle at `out` containing `spec`'s start and gold commits.
153pub fn create_git_bundle(spec: &GitSpec, out: &Path) -> Result<(), WorkspaceError> {
154    let temp_dir = new_temp_dir()?;
155    let repo = GitRepo::init(temp_dir.path())?;
156    repo.fetch(spec.source.remote(), &[&spec.start_commit, &spec.gold_commit])?;
157    repo.update_ref(&format!("refs/heads/{EVAL_START_REF}"), &spec.start_commit)?;
158    repo.update_ref(&format!("refs/heads/{EVAL_GOLD_REF}"), &spec.gold_commit)?;
159    repo.bundle(&[EVAL_START_REF, EVAL_GOLD_REF], out)?;
160    Ok(())
161}
162
163struct GitBaseline {
164    start_commit: String,
165    reference_diff: GitDiff,
166    diff_exclude: Vec<String>,
167}
168
169impl GitSource {
170    fn remote(&self) -> &OsStr {
171        match self {
172            Self::Url(url) => url.as_ref(),
173            Self::Bundle(path) => path.as_os_str(),
174        }
175    }
176}
177
178fn new_temp_dir() -> Result<tempfile::TempDir, WorkspaceError> {
179    tempfile::tempdir().map_err(WorkspaceError::CreateTempDir)
180}
181
182fn resolve_subdir(root_path: &Path, subdir: Option<PathBuf>) -> Result<(PathBuf, Option<PathBuf>), WorkspaceError> {
183    match subdir {
184        None => Ok((root_path.to_path_buf(), None)),
185        Some(relative_cwd) => {
186            let working_path = root_path.join(&relative_cwd);
187            if !working_path.exists() {
188                return Err(WorkspaceError::MissingSubdir { path: working_path });
189            }
190            Ok((working_path, Some(relative_cwd)))
191        }
192    }
193}
194
195fn copy_dir_contents(src: &Path, dst: &Path) -> std::io::Result<()> {
196    for entry in std::fs::read_dir(src)? {
197        let entry = entry?;
198        let source_path = entry.path();
199        let dest_path = dst.join(entry.file_name());
200        let file_type = entry.file_type()?;
201
202        if file_type.is_dir() {
203            std::fs::create_dir_all(&dest_path)?;
204            copy_dir_contents(&source_path, &dest_path)?;
205        } else if file_type.is_file() {
206            std::fs::copy(&source_path, &dest_path)?;
207        }
208    }
209    Ok(())
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use std::fs::{read_to_string, remove_dir_all};
216    use utils::testing::{TestWorkspace, git};
217
218    #[test]
219    fn persist_reports_root_and_path_for_local_workspace() {
220        let workspace = Workspace::from_files([("notes.txt", "hi\n")]).unwrap();
221
222        let retained = workspace.persist();
223
224        assert_eq!(retained.root_path, retained.path);
225        assert_eq!(read_to_string(retained.path.join("notes.txt")).unwrap(), "hi\n");
226        remove_dir_all(retained.root_path).unwrap();
227    }
228
229    #[test]
230    fn local_workspaces_have_no_diffs() {
231        let workspace = Workspace::empty().unwrap();
232
233        assert!(workspace.reference_diff().is_none());
234        assert!(workspace.capture_agent_diff().unwrap().is_none());
235    }
236
237    #[test]
238    fn git_spec_deserializes_from_camel_case_json() {
239        let spec: GitSpec = serde_json::from_value(serde_json::json!({
240            "source": { "bundle": "repo.bundle" },
241            "startCommit": "abc",
242            "goldCommit": "def",
243            "diffExclude": ["*.lock"],
244        }))
245        .unwrap();
246
247        assert!(matches!(spec.source, GitSource::Bundle(path) if path == Path::new("repo.bundle")));
248        assert_eq!(spec.diff_exclude, ["*.lock"]);
249    }
250
251    #[test]
252    fn from_git_with_a_bundle_round_trips_checkout_subdir_and_diffs() {
253        let (_repo, spec) = init_repo();
254        let bundle_dir = tempfile::tempdir().unwrap();
255        let spec = GitSpec { subdir: Some("pkg".into()), ..bundle_spec(&spec, bundle_dir.path()) };
256        let workspace = Workspace::from_git(spec).unwrap();
257
258        assert_eq!(read_to_string(workspace.root_path().join("root.txt")).unwrap(), "root v1\n");
259        assert_eq!(workspace.relative_cwd(), Some(Path::new("pkg")));
260        assert_eq!(workspace.path(), workspace.root_path().join("pkg"));
261        assert!(workspace.reference_diff().unwrap().diff.contains("root v2"), "reference diff should span start..gold");
262        assert!(agent_diff(&workspace).diff.is_empty(), "no agent edits yet");
263
264        write(workspace.root_path().join("root.txt"), "root edited\n").unwrap();
265
266        assert!(agent_diff(&workspace).diff.contains("root edited"), "agent diff should capture the edit");
267    }
268
269    #[test]
270    fn git_workspaces_hold_only_the_start_commit() {
271        let (_repo, spec) = init_repo();
272        let bundle_dir = tempfile::tempdir().unwrap();
273        let workspaces = [
274            Workspace::from_git(spec.clone()).unwrap(),
275            Workspace::from_git(bundle_spec(&spec, bundle_dir.path())).unwrap(),
276        ];
277
278        for workspace in workspaces {
279            let root = workspace.root_path();
280            assert_eq!(git(root, &["rev-parse", "HEAD"]), spec.start_commit);
281            assert_eq!(git(root, &["for-each-ref", "--format=%(refname)"]), "refs/heads/main");
282            assert_eq!(git(root, &["remote"]), "");
283            assert!(
284                utils::git::git(Some(root), &["cat-file", "-e", &spec.gold_commit]).is_err(),
285                "the gold commit must not be in the agent's repository"
286            );
287            assert!(workspace.reference_diff().unwrap().diff.contains("+root v2"));
288        }
289    }
290
291    #[test]
292    fn from_git_checks_out_the_requested_branch() {
293        let (_repo, spec) = init_repo();
294
295        let workspace = Workspace::from_git(GitSpec { branch: Some("trunk".into()), ..spec }).unwrap();
296
297        assert_eq!(git(workspace.root_path(), &["branch", "--show-current"]), "trunk");
298    }
299
300    #[test]
301    fn diff_exclude_leaves_paths_out_of_the_reference_and_agent_diffs() {
302        let (_repo, spec) = init_repo();
303        let workspace = Workspace::from_git(GitSpec { diff_exclude: vec!["*root.txt".into()], ..spec }).unwrap();
304        write(workspace.join("root.txt"), "agent\n").unwrap();
305        write(workspace.join("pkg/inner.txt"), "agent inner\n").unwrap();
306
307        let agent_diff = agent_diff(&workspace);
308
309        assert_eq!(workspace.reference_diff().unwrap().diff, "");
310        assert!(agent_diff.diff.contains("+agent inner"));
311        assert!(!agent_diff.diff.contains("root.txt"));
312    }
313
314    #[test]
315    fn capture_agent_diff_includes_committed_changes() {
316        let (_repo, spec) = init_repo();
317        let workspace = Workspace::from_git(spec).unwrap();
318
319        write(workspace.join("root.txt"), "agent committed\n").unwrap();
320        git(workspace.root_path(), &["add", "."]);
321        git(
322            workspace.root_path(),
323            &["-c", "user.email=agent@example.com", "-c", "user.name=Agent", "commit", "-m", "agent change"],
324        );
325
326        let agent_diff = agent_diff(&workspace);
327
328        assert!(agent_diff.diff.contains("+agent committed"));
329        assert_eq!(agent_diff.stats.files_changed, 1);
330        assert_eq!(agent_diff.stats.lines_added, 1);
331        assert_eq!(agent_diff.stats.lines_removed, 1);
332    }
333
334    #[test]
335    fn capture_agent_diff_includes_staged_unstaged_added_deleted_and_untracked_changes() {
336        let (_repo, spec) = init_repo();
337        let workspace = Workspace::from_git(spec).unwrap();
338
339        write(workspace.join("root.txt"), "staged change\n").unwrap();
340        write(workspace.join("added.txt"), "added\n").unwrap();
341        git(workspace.root_path(), &["add", "root.txt", "added.txt"]);
342        write(workspace.join("pkg/inner.txt"), "unstaged change\n").unwrap();
343        std::fs::remove_file(workspace.join("deleted.txt")).unwrap();
344        create_dir_all(workspace.join("dir")).unwrap();
345        write(workspace.join("dir/untracked.txt"), "untracked\n").unwrap();
346
347        let agent_diff = agent_diff(&workspace);
348
349        assert!(agent_diff.diff.contains("root.txt"));
350        assert!(agent_diff.diff.contains("+staged change"));
351        assert!(agent_diff.diff.contains("pkg/inner.txt"));
352        assert!(agent_diff.diff.contains("+unstaged change"));
353        assert!(agent_diff.diff.contains("added.txt"));
354        assert!(agent_diff.diff.contains("+added"));
355        assert!(agent_diff.diff.contains("deleted.txt"));
356        assert!(agent_diff.diff.contains("-before deleted"));
357        assert!(agent_diff.diff.contains("+untracked"));
358        assert_eq!(agent_diff.stats.files_changed, 5);
359        assert_eq!(agent_diff.stats.lines_added, 4);
360        assert_eq!(agent_diff.stats.lines_removed, 3);
361        assert!(git(workspace.root_path(), &["status", "--porcelain"]).contains("?? dir/"), "the index is untouched");
362    }
363
364    #[test]
365    fn capture_agent_diff_errors_when_git_fails() {
366        let (_repo, spec) = init_repo();
367        let workspace = Workspace::from_git(spec).unwrap();
368        write(workspace.root_path().join(".git/index"), "corrupt").unwrap();
369
370        assert!(matches!(workspace.capture_agent_diff(), Err(WorkspaceError::Git(_))));
371    }
372
373    #[test]
374    fn from_git_missing_bundle_errors() {
375        let result = Workspace::from_git(GitSpec {
376            source: GitSource::Bundle(PathBuf::from("/nonexistent/repo.bundle")),
377            start_commit: "abc".into(),
378            gold_commit: "def".into(),
379            subdir: None,
380            branch: None,
381            diff_exclude: Vec::new(),
382        });
383
384        assert!(matches!(result, Err(WorkspaceError::MissingBundle { .. })));
385    }
386
387    fn agent_diff(workspace: &Workspace) -> GitDiff {
388        workspace.capture_agent_diff().unwrap().expect("git workspaces have an agent diff")
389    }
390
391    fn bundle_spec(spec: &GitSpec, dir: &Path) -> GitSpec {
392        let bundle_path = dir.join("repo.bundle");
393        create_git_bundle(spec, &bundle_path).unwrap();
394        GitSpec { source: GitSource::Bundle(bundle_path), ..spec.clone() }
395    }
396
397    fn init_repo() -> (TestWorkspace, GitSpec) {
398        let repo = TestWorkspace::new()
399            .file("root.txt", "root v1\n")
400            .file("deleted.txt", "before deleted\n")
401            .file("pkg/inner.txt", "inner v1\n")
402            .git_init();
403        let start_commit = repo.commit("start");
404        repo.write("root.txt", "root v2\n");
405        let gold_commit = repo.commit("gold");
406        let source = GitSource::Url(format!("file://{}", repo.root().display()));
407        (repo, GitSpec { source, start_commit, gold_commit, subdir: None, branch: None, diff_exclude: Vec::new() })
408    }
409}