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}
25
26#[derive(Debug, Clone, serde::Deserialize, JsonSchema)]
27#[serde(rename_all = "camelCase", deny_unknown_fields)]
28pub struct GitRepoSpec {
29    pub url: String,
30    pub start_commit: String,
31    pub gold_commit: String,
32    #[serde(default)]
33    pub subdir: Option<PathBuf>,
34}
35
36#[derive(Debug, Clone, Serialize, JsonSchema)]
37#[serde(rename_all = "camelCase", deny_unknown_fields)]
38pub struct RetainedWorkspaceInfo {
39    pub root_path: PathBuf,
40    pub path: PathBuf,
41}
42
43impl Workspace {
44    pub fn empty() -> Result<Self, WorkspaceError> {
45        let temp_dir = new_temp_dir()?;
46        let path = temp_dir.path().to_path_buf();
47        Ok(Self { path: path.clone(), root_path: path, relative_cwd: None, source: WorkspaceSource::Local, temp_dir })
48    }
49
50    pub fn from_dir(src_path: impl Into<PathBuf>) -> Result<Self, WorkspaceError> {
51        let src_path = src_path.into();
52        let temp_dir = new_temp_dir()?;
53        let path = temp_dir.path().to_path_buf();
54
55        copy_dir_contents(&src_path, &path).map_err(|source| WorkspaceError::CopyFixture {
56            from: src_path.clone(),
57            to: path.clone(),
58            source,
59        })?;
60
61        Ok(Self { path: path.clone(), root_path: path, relative_cwd: None, source: WorkspaceSource::Local, temp_dir })
62    }
63
64    pub fn from_files<T: AsRef<Path>, U: AsRef<str>>(
65        files: impl IntoIterator<Item = (T, U)>,
66    ) -> Result<Self, WorkspaceError> {
67        let workspace = Self::empty()?;
68        for (relative_path, contents) in files {
69            let path = workspace.path().join(relative_path.as_ref());
70            if let Some(parent) = path.parent() {
71                create_dir_all(parent)
72                    .map_err(|source| WorkspaceError::WriteFile { path: parent.to_path_buf(), source })?;
73            }
74            write(&path, contents.as_ref()).map_err(|source| WorkspaceError::WriteFile { path, source })?;
75        }
76        Ok(workspace)
77    }
78
79    pub fn from_git_repo(spec: GitRepoSpec) -> Result<Self, WorkspaceError> {
80        let GitRepoSpec { url, start_commit, gold_commit, subdir } = spec;
81        let temp_dir = new_temp_dir()?;
82
83        tracing::debug!("Cloning git repo {} at commit {}", url, start_commit);
84        let repo = GitRepo::clone(&url, temp_dir.path())?;
85        repo.checkout(&start_commit)?;
86
87        let root_path = temp_dir.path().to_path_buf();
88        let (path, relative_cwd) = match subdir {
89            None => (root_path.clone(), None),
90            Some(relative_cwd) => {
91                let working_path = root_path.join(&relative_cwd);
92                if !working_path.exists() {
93                    return Err(WorkspaceError::MissingSubdir { path: working_path });
94                }
95                (working_path, Some(relative_cwd))
96            }
97        };
98
99        Ok(Self {
100            path,
101            root_path,
102            relative_cwd,
103            source: WorkspaceSource::GitRepo { url, start_commit, gold_commit },
104            temp_dir,
105        })
106    }
107
108    pub fn path(&self) -> &Path {
109        &self.path
110    }
111
112    pub fn join(&self, relative_path: impl AsRef<Path>) -> PathBuf {
113        self.path.join(relative_path)
114    }
115
116    pub fn root_path(&self) -> &Path {
117        &self.root_path
118    }
119
120    pub fn relative_cwd(&self) -> Option<&Path> {
121        self.relative_cwd.as_deref()
122    }
123
124    pub fn source(&self) -> &WorkspaceSource {
125        &self.source
126    }
127
128    /// Prevents the workspace from getting automatically removed and returns its retained root and effective cwd. The caller is responsible for cleanup.
129    pub fn persist(self) -> RetainedWorkspaceInfo {
130        let root_path = self.temp_dir.keep();
131        let path = self.relative_cwd.map_or_else(|| root_path.clone(), |relative_cwd| root_path.join(relative_cwd));
132        RetainedWorkspaceInfo { root_path, path }
133    }
134
135    pub fn capture_git_diffs(&self) -> (Option<GitDiff>, Option<GitDiff>) {
136        let WorkspaceSource::GitRepo { start_commit, gold_commit, .. } = self.source() else {
137            return (None, None);
138        };
139
140        let repo = GitRepo::from_path(self.path());
141        let agent_diff = repo.diff_unstaged().ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
142        let reference_diff =
143            repo.diff(start_commit, gold_commit).ok().map(|diff| GitDiff { stats: DiffStats::from_diff(&diff), diff });
144
145        (agent_diff, reference_diff)
146    }
147}
148
149fn new_temp_dir() -> Result<tempfile::TempDir, WorkspaceError> {
150    tempfile::tempdir().map_err(WorkspaceError::CreateTempDir)
151}
152
153fn copy_dir_contents(src: &Path, dst: &Path) -> std::io::Result<()> {
154    for entry in std::fs::read_dir(src)? {
155        let entry = entry?;
156        let source_path = entry.path();
157        let dest_path = dst.join(entry.file_name());
158        let file_type = entry.file_type()?;
159
160        if file_type.is_dir() {
161            std::fs::create_dir_all(&dest_path)?;
162            copy_dir_contents(&source_path, &dest_path)?;
163        } else if file_type.is_file() {
164            std::fs::copy(&source_path, &dest_path)?;
165        }
166    }
167    Ok(())
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use std::fs::{read_to_string, remove_dir_all};
174
175    #[test]
176    fn persist_reports_root_and_path_for_local_workspace() {
177        let workspace = Workspace::from_files([("notes.txt", "hi\n")]).unwrap();
178
179        let retained = workspace.persist();
180
181        assert_eq!(retained.root_path, retained.path);
182        assert_eq!(read_to_string(retained.path.join("notes.txt")).unwrap(), "hi\n");
183        remove_dir_all(retained.root_path).unwrap();
184    }
185}
aether_evals/evals/workspace.rs

aether_evals/evals/
workspace.rs
