aurum_core/batch/
mod.rs

1//! Bounded, resumable multi-file transcription manifests (JOE-1726).
2//!
3//! The CLI drives transcription; this module owns discovery, deterministic
4//! naming, versioned manifests, resume selection, and partial-success reports.
5
6use crate::error::{Result, UserError};
7use crate::output::OutputFormat;
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10use std::collections::BTreeSet;
11use std::fs;
12use std::io::Read;
13use std::path::{Path, PathBuf};
14use std::time::{SystemTime, UNIX_EPOCH};
15
16/// Manifest schema version.
17pub const BATCH_MANIFEST_VERSION: u32 = 1;
18
19/// Default manifest filename written into the batch output directory.
20pub const BATCH_MANIFEST_NAME: &str = "aurum-batch-manifest.json";
21
22/// Extensions treated as transcription inputs (lowercase, no dot).
23pub const AUDIO_EXTENSIONS: &[&str] = &[
24    "wav", "mp3", "m4a", "flac", "ogg", "oga", "opus", "webm", "aac", "mp4", "mpeg", "mpga",
25];
26
27/// Per-item status in the batch manifest.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
29#[serde(rename_all = "snake_case")]
30pub enum BatchItemStatus {
31    Pending,
32    Running,
33    Succeeded,
34    Failed,
35    Skipped,
36}
37
38/// One file in a batch.
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
40pub struct BatchItem {
41    /// Stable id (sha256 of normalized source path, first 16 hex chars).
42    pub id: String,
43    /// Source path as provided / discovered (UTF-8 lossy display form).
44    pub source: String,
45    /// Deterministic relative output path under `output_dir`.
46    pub output: String,
47    pub status: BatchItemStatus,
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub error: Option<String>,
50    #[serde(default)]
51    pub attempts: u32,
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    pub source_sha256: Option<String>,
54}
55
56/// Versioned batch manifest (machine-readable resume state).
57#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
58pub struct BatchManifest {
59    pub schema_version: u32,
60    pub aurum_version: String,
61    pub created_at_unix: u64,
62    pub updated_at_unix: u64,
63    pub provider: String,
64    pub model: String,
65    pub language: String,
66    pub output_format: String,
67    /// Absolute or original path string for the output directory.
68    pub output_dir: String,
69    /// Optional profile used for model selection.
70    #[serde(default, skip_serializing_if = "Option::is_none")]
71    pub profile: Option<String>,
72    pub items: Vec<BatchItem>,
73}
74
75impl BatchManifest {
76    pub fn new(
77        provider: &str,
78        model: &str,
79        language: &str,
80        format: OutputFormat,
81        output_dir: &Path,
82        profile: Option<&str>,
83    ) -> Self {
84        let now = unix_now();
85        Self {
86            schema_version: BATCH_MANIFEST_VERSION,
87            aurum_version: env!("CARGO_PKG_VERSION").into(),
88            created_at_unix: now,
89            updated_at_unix: now,
90            provider: provider.into(),
91            model: model.into(),
92            language: language.into(),
93            output_format: format.as_str().into(),
94            output_dir: output_dir.display().to_string(),
95            profile: profile.map(|s| s.to_string()),
96            items: Vec::new(),
97        }
98    }
99
100    pub fn touch(&mut self) {
101        self.updated_at_unix = unix_now();
102    }
103
104    pub fn summary(&self) -> BatchSummary {
105        let mut s = BatchSummary::default();
106        for i in &self.items {
107            s.total += 1;
108            match i.status {
109                BatchItemStatus::Pending => s.pending += 1,
110                BatchItemStatus::Running => s.running += 1,
111                BatchItemStatus::Succeeded => s.succeeded += 1,
112                BatchItemStatus::Failed => s.failed += 1,
113                BatchItemStatus::Skipped => s.skipped += 1,
114            }
115        }
116        s
117    }
118
119    pub fn to_json_pretty(&self) -> Result<String> {
120        serde_json::to_string_pretty(self).map_err(|e| {
121            UserError::Other {
122                message: format!("batch manifest json: {e}"),
123            }
124            .into()
125        })
126    }
127
128    pub fn save(&self, path: &Path) -> Result<()> {
129        if let Some(parent) = path.parent() {
130            fs::create_dir_all(parent).map_err(|e| UserError::Other {
131                message: format!("create batch output dir {}: {e}", parent.display()),
132            })?;
133        }
134        let json = self.to_json_pretty()?;
135        // Write via temp + rename for crash safety.
136        let tmp = path.with_extension("json.tmp");
137        fs::write(&tmp, json.as_bytes()).map_err(|e| UserError::Other {
138            message: format!("write batch manifest temp: {e}"),
139        })?;
140        fs::rename(&tmp, path).map_err(|e| UserError::Other {
141            message: format!("publish batch manifest: {e}"),
142        })?;
143        Ok(())
144    }
145
146    pub fn load(path: &Path) -> Result<Self> {
147        let data = fs::read_to_string(path).map_err(|e| UserError::Other {
148            message: format!("read batch manifest {}: {e}", path.display()),
149        })?;
150        let m: Self = serde_json::from_str(&data).map_err(|e| UserError::Other {
151            message: format!("parse batch manifest: {e}"),
152        })?;
153        if m.schema_version != BATCH_MANIFEST_VERSION {
154            return Err(UserError::Other {
155                message: format!(
156                    "unsupported batch manifest schema_version {} (expected {BATCH_MANIFEST_VERSION})",
157                    m.schema_version
158                ),
159            }
160            .into());
161        }
162        Ok(m)
163    }
164}
165
166/// Aggregate counters for partial-success reporting.
167#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
168pub struct BatchSummary {
169    pub total: u32,
170    pub pending: u32,
171    pub running: u32,
172    pub succeeded: u32,
173    pub failed: u32,
174    pub skipped: u32,
175}
176
177/// Discover audio files under `input` (file or directory).
178///
179/// Directories are walked non-recursively by default; set `recursive` for depth-first.
180/// Results are sorted for deterministic ordering.
181pub fn discover_inputs(input: &Path, recursive: bool) -> Result<Vec<PathBuf>> {
182    if !input.exists() {
183        return Err(UserError::FileNotFound {
184            path: input.display().to_string(),
185        }
186        .into());
187    }
188    if input.is_file() {
189        if is_audio_path(input) {
190            return Ok(vec![input.to_path_buf()]);
191        }
192        return Err(UserError::InvalidAudio {
193            reason: format!("{} is not a recognised audio extension", input.display()),
194        }
195        .into());
196    }
197    if !input.is_dir() {
198        return Err(UserError::InvalidAudio {
199            reason: format!("{} is not a file or directory", input.display()),
200        }
201        .into());
202    }
203
204    let mut out = Vec::new();
205    walk_dir(input, recursive, &mut out)?;
206    out.sort();
207    if out.is_empty() {
208        return Err(UserError::Other {
209            message: format!(
210                "no audio files found under {}\n  Hint: supported extensions: {}",
211                input.display(),
212                AUDIO_EXTENSIONS.join(", ")
213            ),
214        }
215        .into());
216    }
217    // Cap batch size for resource governance (hard ceiling).
218    const MAX_BATCH_ITEMS: usize = 10_000;
219    if out.len() > MAX_BATCH_ITEMS {
220        return Err(UserError::Other {
221            message: format!(
222                "batch has {} items (max {MAX_BATCH_ITEMS}); split the collection",
223                out.len()
224            ),
225        }
226        .into());
227    }
228    Ok(out)
229}
230
231fn walk_dir(dir: &Path, recursive: bool, out: &mut Vec<PathBuf>) -> Result<()> {
232    let entries = fs::read_dir(dir).map_err(|e| UserError::Other {
233        message: format!("read dir {}: {e}", dir.display()),
234    })?;
235    for ent in entries {
236        let ent = ent.map_err(|e| UserError::Other {
237            message: format!("read dir entry: {e}"),
238        })?;
239        let path = ent.path();
240        if path.is_dir() {
241            if recursive {
242                walk_dir(&path, true, out)?;
243            }
244        } else if is_audio_path(&path) {
245            out.push(path);
246        }
247    }
248    Ok(())
249}
250
251pub fn is_audio_path(path: &Path) -> bool {
252    path.extension()
253        .and_then(|e| e.to_str())
254        .map(|e| AUDIO_EXTENSIONS.contains(&e.to_ascii_lowercase().as_str()))
255        .unwrap_or(false)
256}
257
258/// Deterministic output file name: `<stem>.<format>` with collision disambiguation.
259pub fn output_name_for(source: &Path, format: OutputFormat, used: &mut BTreeSet<String>) -> String {
260    let stem = source
261        .file_stem()
262        .and_then(|s| s.to_str())
263        .unwrap_or("audio");
264    let ext = format.default_extension();
265    let mut name = format!("{stem}.{ext}");
266    if used.insert(name.clone()) {
267        return name;
268    }
269    // Disambiguate with a short hash of the full path.
270    let h = short_id(&source.display().to_string());
271    name = format!("{stem}-{h}.{ext}");
272    let mut n = 2u32;
273    while !used.insert(name.clone()) {
274        name = format!("{stem}-{h}-{n}.{ext}");
275        n += 1;
276    }
277    name
278}
279
280/// Build items for a fresh batch from discovered sources.
281pub fn build_items(sources: &[PathBuf], format: OutputFormat) -> Vec<BatchItem> {
282    let mut used = BTreeSet::new();
283    sources
284        .iter()
285        .map(|src| {
286            let source = src.display().to_string();
287            let id = short_id(&source);
288            let output = output_name_for(src, format, &mut used);
289            BatchItem {
290                id,
291                source,
292                output,
293                status: BatchItemStatus::Pending,
294                error: None,
295                attempts: 0,
296                source_sha256: None,
297            }
298        })
299        .collect()
300}
301
302/// Merge new sources into an existing manifest for resume (keeps terminal results).
303pub fn merge_for_resume(manifest: &mut BatchManifest, sources: &[PathBuf], format: OutputFormat) {
304    let existing: BTreeSet<String> = manifest.items.iter().map(|i| i.source.clone()).collect();
305    let mut used: BTreeSet<String> = manifest.items.iter().map(|i| i.output.clone()).collect();
306    for src in sources {
307        let source = src.display().to_string();
308        if existing.contains(&source) {
309            continue;
310        }
311        let id = short_id(&source);
312        let output = output_name_for(src, format, &mut used);
313        manifest.items.push(BatchItem {
314            id,
315            source,
316            output,
317            status: BatchItemStatus::Pending,
318            error: None,
319            attempts: 0,
320            source_sha256: None,
321        });
322    }
323    manifest.touch();
324}
325
326/// Indices that still need work.
327pub fn work_indices(manifest: &BatchManifest, retry_failed: bool) -> Vec<usize> {
328    manifest
329        .items
330        .iter()
331        .enumerate()
332        .filter(|(_, i)| match i.status {
333            BatchItemStatus::Pending | BatchItemStatus::Running => true,
334            BatchItemStatus::Failed if retry_failed => true,
335            _ => false,
336        })
337        .map(|(idx, _)| idx)
338        .collect()
339}
340
341/// Optional cheap content fingerprint (first 1 MiB + size) for resume honesty.
342pub fn fingerprint_file(path: &Path) -> Result<String> {
343    let mut f = fs::File::open(path).map_err(|e| UserError::Other {
344        message: format!("open {}: {e}", path.display()),
345    })?;
346    let meta = f.metadata().map_err(|e| UserError::Other {
347        message: format!("stat {}: {e}", path.display()),
348    })?;
349    let mut buf = vec![0u8; 1024 * 1024];
350    let n = f.read(&mut buf).map_err(|e| UserError::Other {
351        message: format!("read {}: {e}", path.display()),
352    })?;
353    let mut hasher = Sha256::new();
354    hasher.update(meta.len().to_le_bytes());
355    hasher.update(&buf[..n]);
356    Ok(hex::encode(hasher.finalize()))
357}
358
359pub fn short_id(s: &str) -> String {
360    let mut hasher = Sha256::new();
361    hasher.update(s.as_bytes());
362    let full = hex::encode(hasher.finalize());
363    full[..16].to_string()
364}
365
366fn unix_now() -> u64 {
367    SystemTime::now()
368        .duration_since(UNIX_EPOCH)
369        .map(|d| d.as_secs())
370        .unwrap_or(0)
371}
372
373/// Manifest path inside an output directory.
374pub fn manifest_path(output_dir: &Path) -> PathBuf {
375    output_dir.join(BATCH_MANIFEST_NAME)
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381    use tempfile::tempdir;
382
383    #[test]
384    fn discover_and_names_stable() {
385        let dir = tempdir().unwrap();
386        fs::write(dir.path().join("a.wav"), b"x").unwrap();
387        fs::write(dir.path().join("b.mp3"), b"y").unwrap();
388        fs::write(dir.path().join("skip.txt"), b"z").unwrap();
389        let found = discover_inputs(dir.path(), false).unwrap();
390        assert_eq!(found.len(), 2);
391        let items = build_items(&found, OutputFormat::Txt);
392        assert_eq!(items[0].output, "a.txt");
393        assert_eq!(items[1].output, "b.txt");
394        assert_eq!(items[0].status, BatchItemStatus::Pending);
395    }
396
397    #[test]
398    fn resume_keeps_succeeded() {
399        let dir = tempdir().unwrap();
400        let mut m =
401            BatchManifest::new("local", "base", "auto", OutputFormat::Txt, dir.path(), None);
402        m.items.push(BatchItem {
403            id: "1".into(),
404            source: "/x/a.wav".into(),
405            output: "a.txt".into(),
406            status: BatchItemStatus::Succeeded,
407            error: None,
408            attempts: 1,
409            source_sha256: None,
410        });
411        m.items.push(BatchItem {
412            id: "2".into(),
413            source: "/x/b.wav".into(),
414            output: "b.txt".into(),
415            status: BatchItemStatus::Failed,
416            error: Some("boom".into()),
417            attempts: 1,
418            source_sha256: None,
419        });
420        let work = work_indices(&m, false);
421        assert!(work.is_empty());
422        let retry = work_indices(&m, true);
423        assert_eq!(retry, vec![1]);
424    }
425
426    #[test]
427    fn manifest_roundtrip() {
428        let dir = tempdir().unwrap();
429        let mut m = BatchManifest::new(
430            "local",
431            "tiny-q5_1",
432            "en",
433            OutputFormat::Json,
434            dir.path(),
435            Some("speed"),
436        );
437        m.items = build_items(&[PathBuf::from("/tmp/x.wav")], OutputFormat::Json);
438        let path = manifest_path(dir.path());
439        m.save(&path).unwrap();
440        let loaded = BatchManifest::load(&path).unwrap();
441        assert_eq!(loaded.model, "tiny-q5_1");
442        assert_eq!(loaded.items.len(), 1);
443        assert_eq!(loaded.profile.as_deref(), Some("speed"));
444    }
445}
aurum_core/batch/mod.rs

aurum_core/batch/
mod.rs
