1use std::path::{Path, PathBuf};
18use std::time::{SystemTime, UNIX_EPOCH};
19
20use serde_json::json;
21use transcribe_cpp::CancelToken;
22
23use crate::audio::SAMPLE_RATE;
24use crate::config::ModelSpec;
25use crate::engine::Stages;
26
27pub(crate) struct WindowRecord {
30 pub(crate) start_ms: u64,
31 pub(crate) end_ms: u64,
32 pub(crate) raw: String,
33}
34
35pub(crate) struct TakeModels {
38 pub(crate) asr: ModelSpec,
39 pub(crate) normalizer: Option<ModelSpec>,
41 pub(crate) language: Option<String>,
43}
44
45impl TakeModels {
46 pub(crate) fn from_config(cfg: &crate::Config) -> Self {
47 Self {
48 asr: cfg.asr_model.clone(),
49 normalizer: cfg.normalizer.clone(),
50 language: cfg.language.clone(),
51 }
52 }
53}
54
55pub(crate) struct TakeRecord<'a> {
57 pub(crate) audio: &'a [f32],
58 pub(crate) windows: &'a [WindowRecord],
59 pub(crate) joined: &'a str,
61 pub(crate) text: &'a str,
62 pub(crate) stages: &'a Stages,
63 pub(crate) models: &'a TakeModels,
65 pub(crate) start_lag_ms: u64,
70 pub(crate) processing_ms: u64,
71 pub(crate) truncated: bool,
72 pub(crate) outcome: &'a str,
74 pub(crate) recognition_error: Option<(usize, String)>,
77}
78
79const RETAINED_TAKES: usize = 100;
89
90pub(crate) fn take_stamp() -> u128 {
94 SystemTime::now().duration_since(UNIX_EPOCH).map_or(0, |d| d.as_millis())
95}
96
97pub(crate) fn claim_take_dir(dir: &Path, stamp: u128) -> std::io::Result<(String, PathBuf)> {
115 let mut stamp = stamp;
116 loop {
117 let name = format!("take-{stamp:013}");
118 let take_dir = dir.join(&name);
119 match std::fs::create_dir(&take_dir) {
120 Ok(()) => return Ok((name, take_dir)),
121 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => stamp += 1,
122 Err(source) => return Err(source),
123 }
124 }
125}
126
127pub(crate) fn capture_take(dir: &Path, take_id: u128, take: &TakeRecord<'_>, cancel: &CancelToken) {
137 if cancel.is_cancelled() {
140 return;
141 }
142 if let Err(error) = std::fs::create_dir_all(dir) {
144 tracing::warn!(
145 event_name = "diagnostics_store_uncreatable",
146 %error,
147 dir = %dir.display(),
148 "diagnostics: store directory could not be made"
149 );
150 return;
151 }
152 let take_dir = match claim_take_dir(dir, take_id) {
153 Ok((_, take_dir)) => take_dir,
154 Err(error) => {
155 tracing::warn!(
156 event_name = "diagnostics_take_dir_unclaimed",
157 %error,
158 dir = %dir.display(),
159 "diagnostics: store directory not claimed"
160 );
161 return;
162 }
163 };
164 let wrote = write_take(&take_dir, take);
165 if cancel.is_cancelled() {
166 if let Err(error) = std::fs::remove_dir_all(&take_dir) {
167 tracing::warn!(
168 event_name = "diagnostics_cancelled_take_not_removed",
169 %error,
170 dir = %take_dir.display(),
171 "diagnostics: cancelled take's directory could not be removed"
172 );
173 }
174 return;
175 }
176 if wrote {
179 prune(dir);
180 }
181}
182
183fn write_take(take_dir: &Path, take: &TakeRecord<'_>) -> bool {
191 if let Err(error) = std::fs::create_dir_all(take_dir.join("raw")) {
196 tracing::warn!(
197 event_name = "diagnostics_store_dir_unwritable",
198 %error,
199 dir = %take_dir.display(),
200 "diagnostics: store directory not writable"
201 );
202 return false;
203 }
204
205 if let Err(error) = write_wav(&take_dir.join("output.wav"), take.audio) {
206 tracing::warn!(
207 event_name = "diagnostics_capture_not_written",
208 %error,
209 dir = %take_dir.display(),
210 "diagnostics: capture not written"
211 );
212 return false;
213 }
214 for (k, window) in take.windows.iter().enumerate() {
215 if let Err(error) = std::fs::write(take_dir.join(format!("raw/{k}.txt")), &window.raw) {
216 tracing::warn!(
217 event_name = "diagnostics_window_transcript_not_written",
218 %error,
219 dir = %take_dir.display(),
220 window = k,
221 "diagnostics: window transcript not written"
222 );
223 return false;
224 }
225 }
226 if let Err(error) = std::fs::write(take_dir.join("joined.txt"), take.joined) {
227 tracing::warn!(
228 event_name = "diagnostics_joined_transcript_not_written",
229 %error,
230 dir = %take_dir.display(),
231 "diagnostics: joined transcript not written"
232 );
233 return false;
234 }
235 if let Err(error) = std::fs::write(take_dir.join("text.txt"), take.text) {
236 tracing::warn!(
237 event_name = "diagnostics_normalized_transcript_not_written",
238 %error,
239 dir = %take_dir.display(),
240 "diagnostics: normalized transcript not written"
241 );
242 return false;
243 }
244
245 let ms = |d: std::time::Duration| u64::try_from(d.as_millis()).unwrap_or(u64::MAX);
246 let mut meta = json!({
247 "duration_ms": ms(take.stages.audio),
248 "start_lag_ms": take.start_lag_ms,
251 "processing_ms": take.processing_ms,
252 "truncated": take.truncated,
253 "outcome": take.outcome,
254 "asr": model_identity(&take.models.asr),
255 "normalizer": take.models.normalizer.as_ref().map(model_identity),
256 "normalizer_ran": take.stages.normalize.is_some(),
261 "language": take.models.language,
264 "forge_version": env!("CARGO_PKG_VERSION"),
267 "stages_ms": {
268 "mel": ms(take.stages.mel),
269 "encode": ms(take.stages.encode),
270 "decode": ms(take.stages.decode),
271 },
272 "windows": take.windows.iter().enumerate().map(|(k, window)| json!({
273 "index": k,
274 "file": format!("raw/{k}.txt"),
275 "start_ms": window.start_ms,
276 "end_ms": window.end_ms,
277 })).collect::<Vec<_>>(),
278 });
279 if let Some((window, error)) = &take.recognition_error {
280 meta["recognition_error"] = json!({ "window": window, "error": error });
281 }
282 match serde_json::to_vec_pretty(&meta) {
283 Ok(bytes) => {
284 if let Err(error) = std::fs::write(take_dir.join("meta.json"), bytes) {
285 tracing::warn!(
286 event_name = "diagnostics_metadata_not_written",
287 %error,
288 dir = %take_dir.display(),
289 "diagnostics: metadata not written"
290 );
291 return false;
292 }
293 }
294 Err(error) => {
295 tracing::warn!(
296 event_name = "diagnostics_metadata_not_serializable",
297 %error,
298 dir = %take_dir.display(),
299 "diagnostics: metadata not serializable"
300 );
301 return false;
302 }
303 }
304 true
305}
306
307fn model_identity(spec: &ModelSpec) -> serde_json::Value {
310 json!({ "file": spec.file, "sha256": spec.sha256, "size": spec.size })
311}
312
313pub(crate) fn sha256_file(path: &Path) -> String {
315 use sha2::Digest as _;
316 let Ok(bytes) = std::fs::read(path) else {
317 return String::new();
318 };
319 hex::encode(sha2::Sha256::digest(&bytes))
320}
321
322pub(crate) fn write_wav(path: &Path, audio: &[f32]) -> Result<(), String> {
325 let spec = hound::WavSpec {
326 channels: 1,
327 sample_rate: SAMPLE_RATE,
328 bits_per_sample: 16,
329 sample_format: hound::SampleFormat::Int,
330 };
331 let mut writer = hound::WavWriter::create(path, spec).map_err(|error| error.to_string())?;
332 for &sample in audio {
333 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
335 let value = (sample * 32767.0).clamp(-32768.0, 32767.0).round() as i16;
336 writer.write_sample(value).map_err(|error| error.to_string())?;
337 }
338 writer.finalize().map_err(|error| error.to_string())
339}
340
341fn prune(dir: &Path) {
347 let Ok(entries) = std::fs::read_dir(dir) else { return };
348 let mut takes: Vec<PathBuf> = entries
349 .flatten()
350 .map(|entry| entry.path())
351 .filter(|path| path.is_dir() && path.file_name().is_some_and(starts_take))
352 .collect();
353 takes.sort();
354 for stale in takes.iter().rev().skip(RETAINED_TAKES) {
355 if let Err(error) = std::fs::remove_dir_all(stale) {
356 tracing::warn!(
357 event_name = "diagnostics_stale_take_not_pruned",
358 %error,
359 dir = %stale.display(),
360 "diagnostics: stale take not pruned"
361 );
362 }
363 }
364}
365
366fn starts_take(name: &std::ffi::OsStr) -> bool {
370 name.to_str().is_some_and(|s| {
371 let id = s.strip_prefix("take-").unwrap_or_default();
372 !id.is_empty() && id.len() >= 13 && id.bytes().all(|b| b.is_ascii_digit())
373 })
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use std::time::Duration;
380
381 use crate::{ConfigBuilder, ModelFacts, ModelSpec};
382
383 fn stub_normalizer() -> ModelSpec {
386 ModelSpec {
387 file: "cleanup-stub.gguf".into(),
388 url: String::new(),
389 size: 7,
390 sha256: None,
391 facts: ModelFacts::default(),
392 }
393 }
394
395 fn stub_models(normalizer: Option<ModelSpec>, language: Option<&str>) -> TakeModels {
398 let mut builder = ConfigBuilder::new()
399 .asr_model(ModelSpec {
400 file: "asr-stub.gguf".into(),
401 url: String::new(),
402 size: 4242,
403 sha256: Some("a1b2c3".into()),
404 facts: ModelFacts::default(),
405 })
406 .normalizer(normalizer);
407 if let Some(language) = language {
408 builder = builder.language(language);
409 }
410 TakeModels::from_config(&builder.build())
411 }
412
413 fn take_record<'a>(
414 models: &'a TakeModels,
415 audio: &'a [f32],
416 windows: &'a [WindowRecord],
417 joined: &'a str,
418 text: &'a str,
419 stages: &'a Stages,
420 ) -> TakeRecord<'a> {
421 TakeRecord {
422 audio,
423 windows,
424 joined,
425 text,
426 stages,
427 models,
428 start_lag_ms: 3,
429 processing_ms: 7,
430 truncated: false,
431 outcome: "transcript",
432 recognition_error: None,
433 }
434 }
435
436 #[test]
441 fn a_take_cancelled_before_its_write_leaves_no_directory() {
442 let dir = tempfile::tempdir().unwrap();
443 let audio = vec![0.5; 16];
444 let stages = Stages::default();
445 let models = stub_models(None, None);
446 let cancel = CancelToken::new();
447 cancel.cancel();
448 capture_take(
449 dir.path(),
450 41,
451 &take_record(&models, &audio, &[], "early", "early", &stages),
452 &cancel,
453 );
454 let entries: Vec<_> = std::fs::read_dir(dir.path()).unwrap().flatten().collect();
455 assert!(
456 entries.is_empty(),
457 "a cancelled take persists nothing, found {} entries",
458 entries.len()
459 );
460 }
461
462 #[test]
466 fn a_take_writes_the_whole_store() {
467 let dir = tempfile::tempdir().unwrap();
468 let audio = vec![0.5; SAMPLE_RATE as usize];
469 let stages = Stages { audio: Duration::from_millis(1000), ..Stages::default() };
470 let windows = vec![
471 WindowRecord { start_ms: 0, end_ms: 500, raw: "first take".into() },
472 WindowRecord { start_ms: 500, end_ms: 1000, raw: "second take".into() },
473 ];
474 let models = stub_models(None, None);
475 capture_take(
476 dir.path(),
477 42,
478 &take_record(
479 &models,
480 &audio,
481 &windows,
482 "first take second take",
483 "First take, second take.",
484 &stages,
485 ),
486 &CancelToken::new(),
487 );
488
489 let take = dir.path().join("take-0000000000042");
490 let mut reader = hound::WavReader::open(take.join("output.wav")).unwrap();
491 assert_eq!(
492 reader.spec().sample_rate,
493 SAMPLE_RATE,
494 "the capture is stored at the model rate"
495 );
496 assert_eq!(
497 reader.samples::<i16>().count(),
498 SAMPLE_RATE as usize,
499 "the whole capture is stored"
500 );
501 assert_eq!(
502 std::fs::read_to_string(take.join("raw/0.txt")).unwrap(),
503 "first take",
504 "window 0's pre-normalization text"
505 );
506 assert_eq!(
507 std::fs::read_to_string(take.join("raw/1.txt")).unwrap(),
508 "second take",
509 "window 1's pre-normalization text"
510 );
511 assert_eq!(
512 std::fs::read_to_string(take.join("joined.txt")).unwrap(),
513 "first take second take",
514 "joined.txt is the exact normalizer input"
515 );
516 assert_eq!(
517 std::fs::read_to_string(take.join("text.txt")).unwrap(),
518 "First take, second take.",
519 "text.txt is the final normalized text"
520 );
521 let meta: serde_json::Value =
522 serde_json::from_str(&std::fs::read_to_string(take.join("meta.json")).unwrap())
523 .unwrap();
524 assert_eq!(meta["duration_ms"], 1000);
525 assert_eq!(meta["start_lag_ms"], 3, "the pre-transcribe lag reaches the metadata");
526 assert_eq!(meta["processing_ms"], 7);
527 assert_eq!(meta["truncated"], false);
528 assert_eq!(meta["outcome"], "transcript");
529 assert_eq!(meta["stages_ms"]["mel"], 0);
530 assert_eq!(meta["windows"][0]["file"], "raw/0.txt");
531 assert_eq!(meta["windows"][1]["end_ms"], 1000);
532 assert!(
533 meta.get("recognition_error").is_none(),
534 "a successful take names no failed window"
535 );
536 }
537
538 #[test]
542 fn a_recognition_error_take_records_the_failed_window() {
543 let dir = tempfile::tempdir().unwrap();
544 let audio = vec![0.5; 16];
545 let stages = Stages::default();
546 let windows = vec![WindowRecord { start_ms: 0, end_ms: 500, raw: "first".into() }];
547 let models = stub_models(None, None);
548 let mut record = take_record(&models, &audio, &windows, "first", "", &stages);
549 record.outcome = "recognition_error";
550 record.recognition_error = Some((1, "decode failed: bad input".into()));
551 capture_take(dir.path(), 7, &record, &CancelToken::new());
552
553 let meta: serde_json::Value = serde_json::from_str(
554 &std::fs::read_to_string(dir.path().join("take-0000000000007/meta.json")).unwrap(),
555 )
556 .unwrap();
557 assert_eq!(meta["outcome"], "recognition_error");
558 assert_eq!(meta["recognition_error"]["window"], 1, "the failed window index");
559 assert_eq!(
560 meta["recognition_error"]["error"], "decode failed: bad input",
561 "the recognition failure text"
562 );
563 }
564
565 #[test]
570 fn a_take_names_the_models_that_made_it() {
571 let dir = tempfile::tempdir().unwrap();
572 let audio = vec![0.5; 16];
573 let stages = Stages {
574 audio: Duration::from_millis(1000),
575 normalize: Some(Duration::from_millis(11)),
576 ..Stages::default()
577 };
578 let windows = vec![];
579 let models = stub_models(Some(stub_normalizer()), Some("en"));
580 let record = take_record(&models, &audio, &windows, "", "first take", &stages);
581 capture_take(dir.path(), 9, &record, &CancelToken::new());
582
583 let meta: serde_json::Value = serde_json::from_str(
584 &std::fs::read_to_string(dir.path().join("take-0000000000009/meta.json")).unwrap(),
585 )
586 .unwrap();
587 assert_eq!(meta["asr"]["file"], "asr-stub.gguf", "the transcribing spec's file");
588 assert_eq!(meta["asr"]["sha256"], "a1b2c3", "the digest the transcribing spec declares");
589 assert_eq!(meta["asr"]["size"], 4242, "the size the transcribing spec declares");
590 assert_eq!(meta["normalizer"]["file"], "cleanup-stub.gguf", "the cleanup spec's file");
591 assert!(
592 meta["normalizer"]["sha256"].is_null(),
593 "a spec whose upstream publishes no digest records none, got {}",
594 meta["normalizer"]["sha256"]
595 );
596 assert_eq!(meta["normalizer"]["size"], 7, "the size the cleanup spec declares");
597 assert_eq!(
598 meta["normalizer_ran"], true,
599 "a take whose stages carry a normalize pass ran the normalizer"
600 );
601 assert_eq!(meta["language"], "en", "the language the engine was configured with");
602 assert_eq!(
603 meta["forge_version"],
604 env!("CARGO_PKG_VERSION"),
605 "the release the take was made under"
606 );
607 }
608
609 #[test]
613 fn a_take_with_the_normalizer_off_names_no_cleanup_model() {
614 let dir = tempfile::tempdir().unwrap();
615 let audio = vec![0.5; 16];
616 let stages = Stages::default();
617 let windows = vec![];
618 let models = stub_models(None, None);
619 let record = take_record(&models, &audio, &windows, "", "first take", &stages);
620 capture_take(dir.path(), 8, &record, &CancelToken::new());
621
622 let meta: serde_json::Value = serde_json::from_str(
623 &std::fs::read_to_string(dir.path().join("take-0000000000008/meta.json")).unwrap(),
624 )
625 .unwrap();
626 assert!(meta["normalizer"].is_null(), "no configured cleanup names no model");
627 assert_eq!(meta["normalizer_ran"], false, "no normalize stage means it did not run");
628 assert!(meta["language"].is_null(), "autodetect records no language");
629 assert_eq!(
630 meta["asr"]["file"], "asr-stub.gguf",
631 "the transcribing spec is named whatever the cleanup role is"
632 );
633 }
634
635 #[test]
639 fn prune_never_touches_a_foreign_directory() {
640 let dir = tempfile::tempdir().unwrap();
641 let foreign = dir.path().join("take-2026-photos");
642 std::fs::create_dir_all(&foreign).unwrap();
643 std::fs::write(foreign.join("holiday.jpg"), "not ours").unwrap();
644 let stages = Stages::default();
645 let audio = vec![0.5; 16];
646 let windows = vec![];
647 let models = stub_models(None, None);
648 for id in 1..=RETAINED_TAKES as u128 + 1 {
649 capture_take(
650 dir.path(),
651 id,
652 &take_record(&models, &audio, &windows, "", "", &stages),
653 &CancelToken::new(),
654 );
655 }
656 assert!(foreign.is_dir(), "a user directory sharing the take- prefix must survive pruning");
657 let takes: Vec<_> = std::fs::read_dir(dir.path())
658 .unwrap()
659 .flatten()
660 .map(|e| e.file_name().to_string_lossy().into_owned())
661 .filter(|n| n.starts_with("take-0"))
662 .collect();
663 assert_eq!(
664 takes.len(),
665 RETAINED_TAKES,
666 "the store kept a different number than its own cap, got {takes:?}"
667 );
668 }
669
670 #[test]
675 fn the_store_keeps_the_last_of_its_retention() {
676 let dir = tempfile::tempdir().unwrap();
677 let audio = vec![0.5; 16];
678 let stages = Stages::default();
679 let windows = vec![];
680 let models = stub_models(None, None);
681 let newest = RETAINED_TAKES as u128 + 2;
682 for id in 1..=newest {
683 capture_take(
684 dir.path(),
685 id,
686 &take_record(&models, &audio, &windows, "", "", &stages),
687 &CancelToken::new(),
688 );
689 }
690 let takes: Vec<_> = std::fs::read_dir(dir.path())
691 .unwrap()
692 .flatten()
693 .map(|e| e.file_name().to_string_lossy().into_owned())
694 .collect();
695 assert_eq!(
696 takes.len(),
697 RETAINED_TAKES,
698 "{newest} takes leave the cap's number, got {takes:?}"
699 );
700 assert!(
701 !takes.contains(&"take-0000000000001".to_owned()),
702 "the oldest take is pruned, got {takes:?}"
703 );
704 assert!(
705 takes.contains(&format!("take-{newest:013}")),
706 "the newest take survives, got {takes:?}"
707 );
708 }
709
710 #[test]
713 fn a_diagnostics_failure_is_swallowed() {
714 let dir = tempfile::tempdir().unwrap();
715 let blocker = dir.path().join("blocker");
716 std::fs::write(&blocker, "a regular file, so subdirectories cannot exist").unwrap();
717 let audio = vec![0.5; 16];
718 let stages = Stages::default();
719 let windows = vec![];
720 let models = stub_models(None, None);
721 capture_take(
722 &blocker.join("under"),
723 1,
724 &take_record(&models, &audio, &windows, "", "", &stages),
725 &CancelToken::new(),
726 );
727 assert!(!blocker.is_dir(), "the unwritable location must not have been turned into one");
728 }
729}