Skip to main content

qualia_core_db/inference/runtime/artifacts/
run_dir.rs

1use std::fs;
2use std::path::{Path, PathBuf};
3
4use tempfile::{Builder, TempDir};
5
6use super::budget::{validate_label, validate_relative_artifact_path, ArtifactError};
7
8pub const RUN_MARKER_FILE: &str = ".qualia-inference-run";
9
10#[derive(Debug, Clone)]
11pub enum ArtifactRetention {
12    Ephemeral,
13    RetainTo(PathBuf),
14}
15
16#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
17pub struct ArtifactStats {
18    pub created_bytes: u64,
19    pub removed_bytes: u64,
20    pub retained_bytes: u64,
21    pub cleanup_failures: u64,
22}
23
24#[derive(Debug)]
25pub struct ArtifactFinish {
26    pub retained_path: Option<PathBuf>,
27    pub stats: ArtifactStats,
28}
29
30/// One bounded, marker-owned scratch directory.
31///
32/// Dropping an unfinished run removes it through [`TempDir`]. Call [`Self::finish`] to obtain
33/// cleanup counters or atomically promote retained evidence.
34#[derive(Debug)]
35pub struct RunArtifactDir {
36    dir: Option<TempDir>,
37    retention: ArtifactRetention,
38    byte_budget: u64,
39    stats: ArtifactStats,
40}
41
42impl RunArtifactDir {
43    pub fn new_in(
44        scratch_parent: &Path,
45        label: &str,
46        byte_budget: u64,
47        retention: ArtifactRetention,
48    ) -> Result<Self, ArtifactError> {
49        validate_label(label)?;
50        fs::create_dir_all(scratch_parent)?;
51        let dir = Builder::new()
52            .prefix(&format!("qualia-inference-{label}-"))
53            .tempdir_in(scratch_parent)?;
54        fs::write(
55            dir.path().join(RUN_MARKER_FILE),
56            b"qualia-inference-run-v1\n",
57        )?;
58        Ok(Self {
59            dir: Some(dir),
60            retention,
61            byte_budget,
62            stats: ArtifactStats::default(),
63        })
64    }
65
66    pub fn path(&self) -> &Path {
67        self.dir
68            .as_ref()
69            .expect("artifact directory is unavailable after finish")
70            .path()
71    }
72
73    pub fn remaining_bytes(&self) -> u64 {
74        self.byte_budget.saturating_sub(self.stats.created_bytes)
75    }
76
77    pub fn stats(&self) -> ArtifactStats {
78        self.stats
79    }
80
81    pub fn write_bounded(
82        &mut self,
83        relative_path: impl AsRef<Path>,
84        bytes: &[u8],
85    ) -> Result<PathBuf, ArtifactError> {
86        let relative_path = relative_path.as_ref();
87        validate_relative_artifact_path(relative_path)?;
88        let attempted_bytes = self.stats.created_bytes.saturating_add(bytes.len() as u64);
89        if attempted_bytes > self.byte_budget {
90            return Err(ArtifactError::BudgetExceeded {
91                budget_bytes: self.byte_budget,
92                attempted_bytes,
93            });
94        }
95        let output = self.path().join(relative_path);
96        if let Some(parent) = output.parent() {
97            fs::create_dir_all(parent)?;
98        }
99        fs::write(&output, bytes)?;
100        self.stats.created_bytes = attempted_bytes;
101        Ok(output)
102    }
103
104    pub fn finish(mut self) -> Result<ArtifactFinish, ArtifactError> {
105        let dir = self
106            .dir
107            .take()
108            .expect("artifact directory can only be finished once");
109        match self.retention {
110            ArtifactRetention::Ephemeral => {
111                let path = dir.path().to_path_buf();
112                match dir.close() {
113                    Ok(()) => {
114                        self.stats.removed_bytes = self.stats.created_bytes;
115                        Ok(ArtifactFinish {
116                            retained_path: None,
117                            stats: self.stats,
118                        })
119                    }
120                    Err(source) => {
121                        self.stats.cleanup_failures = self.stats.cleanup_failures.saturating_add(1);
122                        Err(ArtifactError::Cleanup { path, source })
123                    }
124                }
125            }
126            ArtifactRetention::RetainTo(target) => {
127                if target.exists() {
128                    return Err(ArtifactError::TargetExists);
129                }
130                if let Some(parent) = target.parent() {
131                    fs::create_dir_all(parent)?;
132                }
133                let staging = dir.keep();
134                if let Err(source) = fs::rename(&staging, &target) {
135                    return Err(ArtifactError::Promotion {
136                        staging,
137                        target,
138                        source,
139                    });
140                }
141                self.stats.retained_bytes = self.stats.created_bytes;
142                Ok(ArtifactFinish {
143                    retained_path: Some(target),
144                    stats: self.stats,
145                })
146            }
147        }
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use std::panic::{catch_unwind, AssertUnwindSafe};
154
155    use super::*;
156
157    #[test]
158    fn ephemeral_finish_removes_run_directory() {
159        let parent = tempfile::tempdir().unwrap();
160        let mut run =
161            RunArtifactDir::new_in(parent.path(), "cleanup", 16, ArtifactRetention::Ephemeral)
162                .unwrap();
163        let run_path = run.path().to_path_buf();
164        run.write_bounded("nested/out.bin", b"1234").unwrap();
165        let finish = run.finish().unwrap();
166        assert!(!run_path.exists());
167        assert_eq!(finish.stats.created_bytes, 4);
168        assert_eq!(finish.stats.removed_bytes, 4);
169        assert_eq!(finish.stats.retained_bytes, 0);
170    }
171
172    #[test]
173    fn drop_removes_directory_during_unwind() {
174        let parent = tempfile::tempdir().unwrap();
175        let mut run_path = PathBuf::new();
176        let result = catch_unwind(AssertUnwindSafe(|| {
177            let mut run =
178                RunArtifactDir::new_in(parent.path(), "panic", 16, ArtifactRetention::Ephemeral)
179                    .unwrap();
180            run_path = run.path().to_path_buf();
181            run.write_bounded("out.bin", b"1234").unwrap();
182            panic!("test unwind");
183        }));
184        assert!(result.is_err());
185        assert!(!run_path.exists());
186    }
187
188    #[test]
189    fn budget_fails_closed_without_partial_file() {
190        let parent = tempfile::tempdir().unwrap();
191        let mut run =
192            RunArtifactDir::new_in(parent.path(), "budget", 3, ArtifactRetention::Ephemeral)
193                .unwrap();
194        let err = run.write_bounded("too-large.bin", b"1234").unwrap_err();
195        assert!(matches!(err, ArtifactError::BudgetExceeded { .. }));
196        assert!(!run.path().join("too-large.bin").exists());
197    }
198
199    #[test]
200    fn traversal_is_rejected() {
201        let parent = tempfile::tempdir().unwrap();
202        let mut run =
203            RunArtifactDir::new_in(parent.path(), "paths", 16, ArtifactRetention::Ephemeral)
204                .unwrap();
205        assert!(matches!(
206            run.write_bounded("../escape.bin", b"x"),
207            Err(ArtifactError::InvalidRelativePath)
208        ));
209    }
210
211    #[test]
212    fn retained_run_is_promoted_out_of_staging() {
213        let parent = tempfile::tempdir().unwrap();
214        let target = parent.path().join("evidence").join("run-1");
215        let mut run = RunArtifactDir::new_in(
216            parent.path(),
217            "retain",
218            16,
219            ArtifactRetention::RetainTo(target.clone()),
220        )
221        .unwrap();
222        let staging = run.path().to_path_buf();
223        run.write_bounded("receipt.json", b"{}").unwrap();
224        let finish = run.finish().unwrap();
225        assert!(!staging.exists());
226        assert_eq!(finish.retained_path.as_deref(), Some(target.as_path()));
227        assert!(target.join(RUN_MARKER_FILE).is_file());
228        assert_eq!(finish.stats.retained_bytes, 2);
229    }
230}