aether_evals/evals/
workspace.rs

1use crate::WorkspaceError;
2use schemars::JsonSchema;
3use serde::Serialize;
4use std::{
5    fs::{create_dir_all, write},
6    path::{Path, PathBuf},
7};
8use tempfile::TempDir;
9
10pub struct Workspace {
11    path: PathBuf,
12    root_path: PathBuf,
13    relative_cwd: Option<PathBuf>,
14    temp_dir: TempDir,
15}
16
17#[derive(Debug, Clone, Serialize, JsonSchema)]
18#[serde(rename_all = "camelCase", deny_unknown_fields)]
19pub struct RetainedWorkspaceInfo {
20    pub root_path: PathBuf,
21    pub path: PathBuf,
22}
23
24impl Workspace {
25    pub fn empty() -> Result<Self, WorkspaceError> {
26        Self::from_temp_dir(new_temp_dir()?, None)
27    }
28
29    pub fn from_dir(src_path: impl Into<PathBuf>) -> Result<Self, WorkspaceError> {
30        let src_path = src_path.into();
31        let temp_dir = new_temp_dir()?;
32
33        copy_dir_contents(&src_path, temp_dir.path()).map_err(|source| WorkspaceError::CopyFixture {
34            from: src_path.clone(),
35            to: temp_dir.path().to_path_buf(),
36            source,
37        })?;
38
39        Self::from_temp_dir(temp_dir, None)
40    }
41
42    pub fn from_files<T: AsRef<Path>, U: AsRef<str>>(
43        files: impl IntoIterator<Item = (T, U)>,
44    ) -> Result<Self, WorkspaceError> {
45        let workspace = Self::empty()?;
46        for (relative_path, contents) in files {
47            let path = workspace.path().join(relative_path.as_ref());
48            if let Some(parent) = path.parent() {
49                create_dir_all(parent)
50                    .map_err(|source| WorkspaceError::WriteFile { path: parent.to_path_buf(), source })?;
51            }
52            write(&path, contents.as_ref()).map_err(|source| WorkspaceError::WriteFile { path, source })?;
53        }
54        Ok(workspace)
55    }
56
57    pub fn path(&self) -> &Path {
58        &self.path
59    }
60
61    pub fn join(&self, relative_path: impl AsRef<Path>) -> PathBuf {
62        self.path.join(relative_path)
63    }
64
65    pub fn root_path(&self) -> &Path {
66        &self.root_path
67    }
68
69    pub fn relative_cwd(&self) -> Option<&Path> {
70        self.relative_cwd.as_deref()
71    }
72
73    /// Prevents the workspace from getting automatically removed and returns its retained root and effective cwd. The caller is responsible for cleanup.
74    pub fn persist(self) -> RetainedWorkspaceInfo {
75        let root_path = self.temp_dir.keep();
76        let path = self.relative_cwd.map_or_else(|| root_path.clone(), |relative_cwd| root_path.join(relative_cwd));
77        RetainedWorkspaceInfo { root_path, path }
78    }
79
80    pub(super) fn from_temp_dir(temp_dir: TempDir, subdir: Option<PathBuf>) -> Result<Self, WorkspaceError> {
81        let root_path = temp_dir.path().to_path_buf();
82        let (path, relative_cwd) = resolve_subdir(&root_path, subdir)?;
83        Ok(Self { path, root_path, relative_cwd, temp_dir })
84    }
85}
86
87pub(super) fn new_temp_dir() -> Result<TempDir, WorkspaceError> {
88    tempfile::tempdir().map_err(WorkspaceError::CreateTempDir)
89}
90
91fn resolve_subdir(root_path: &Path, subdir: Option<PathBuf>) -> Result<(PathBuf, Option<PathBuf>), WorkspaceError> {
92    match subdir {
93        None => Ok((root_path.to_path_buf(), None)),
94        Some(relative_cwd) => {
95            let working_path = root_path.join(&relative_cwd);
96            if !working_path.exists() {
97                return Err(WorkspaceError::MissingSubdir { path: working_path });
98            }
99            Ok((working_path, Some(relative_cwd)))
100        }
101    }
102}
103
104fn copy_dir_contents(src: &Path, dst: &Path) -> std::io::Result<()> {
105    for entry in std::fs::read_dir(src)? {
106        let entry = entry?;
107        let source_path = entry.path();
108        let dest_path = dst.join(entry.file_name());
109        let file_type = entry.file_type()?;
110
111        if file_type.is_dir() {
112            std::fs::create_dir_all(&dest_path)?;
113            copy_dir_contents(&source_path, &dest_path)?;
114        } else if file_type.is_file() {
115            std::fs::copy(&source_path, &dest_path)?;
116        }
117    }
118    Ok(())
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use std::fs::{read_to_string, remove_dir_all};
125
126    #[test]
127    fn persist_reports_root_and_path_for_local_workspace() {
128        let workspace = Workspace::from_files([("notes.txt", "hi\n")]).unwrap();
129
130        let retained = workspace.persist();
131
132        assert_eq!(retained.root_path, retained.path);
133        assert_eq!(read_to_string(retained.path.join("notes.txt")).unwrap(), "hi\n");
134        remove_dir_all(retained.root_path).unwrap();
135    }
136}
aether_evals/evals/workspace.rs

aether_evals/evals/
workspace.rs
