1use kcode_k1_audio_classification_format::{
2 AudioClassificationEventV3, ProgressUpdateV1, decode_event,
3};
4pub use kcode_k1_audio_classification_format::{ExecutedAnalysis, FragmentStageV1, SpeakerLabelV1};
5use kcode_k1_transaction::Transaction;
6use kcode_k1_txn_ordering::{K1TxnOrdering, TxId};
7use rusqlite::{Connection, Error as SqlError, ErrorCode, OptionalExtension, params};
8use serde::{Deserialize, Serialize};
9use std::ffi::OsString;
10use std::fs;
11use std::path::{Path, PathBuf};
12use std::sync::{Mutex, MutexGuard};
13
14pub type FragmentId = kcode_k1_audio_fragment_submit::FragmentId;
15
16const DATABASE_NAME: &str = "audio-classification.sqlite3";
17const SUBSYSTEM: &str = "audio-classification";
18const SCHEMA_VERSION: i64 = 1;
19const MAX_ERRORS: usize = 5_000;
20const METADATA_SQL: &str = "CREATE TABLE metadata(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), schema_version INTEGER NOT NULL, last_applied_txid BLOB)";
21const FRAGMENTS_SQL: &str = "CREATE TABLE fragments(fragment_id BLOB PRIMARY KEY, actionable_state INTEGER NOT NULL, encoded_status BLOB NOT NULL)";
22const INDEX_SQL: &str = "CREATE INDEX fragments_actionable_state ON fragments(actionable_state)";
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
25pub enum OverallState {
26 Queued,
27 Running,
28 Failed,
29 Completed,
30 Confirmed,
31 Discarded,
32}
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
34pub enum StageState {
35 Pending,
36 Running,
37 Succeeded,
38 Failed,
39}
40#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
41pub enum LlmJobState {
42 Running,
43 Succeeded,
44 Failed,
45}
46
47#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
48pub struct StageStatus {
49 pub stage: FragmentStageV1,
50 pub state: StageState,
51}
52
53#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
54pub struct LlmJobStatus {
55 pub attempt: u32,
56 pub sequence: u64,
57 pub stage: FragmentStageV1,
58 pub name: String,
59 pub state: LlmJobState,
60}
61
62#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
63pub struct FragmentStatus {
64 pub state: OverallState,
65 pub queue: StageStatus,
66 pub transcript: StageStatus,
67 pub speaker_labels: StageStatus,
68 pub speaker_features: StageStatus,
69 pub structuring: StageStatus,
70 pub label_confirmation: StageStatus,
71 pub attempt_count: u32,
72 pub jobs: Vec<LlmJobStatus>,
73 #[serde(with = "optional_fragment_id")]
74 pub interim_txid: Option<FragmentId>,
75 pub analysis: Option<ExecutedAnalysis>,
76 pub confirmed_labels: Vec<SpeakerLabelV1>,
77 pub final_transcript: Option<String>,
78 pub errors: Vec<String>,
79 pub errors_truncated: bool,
80}
81
82#[derive(Clone, Debug, PartialEq, Eq)]
83pub enum ProjectionEffect {
84 None,
85 Start,
86 Abort,
87 LabelsCommitted,
88}
89
90#[derive(Clone, Debug, PartialEq, Eq)]
91pub struct AppliedEvent {
92 pub fragment_id: FragmentId,
93 pub effect: ProjectionEffect,
94}
95
96#[derive(Clone, Debug, PartialEq, Eq)]
97pub struct InterruptedFragment {
98 pub fragment_id: FragmentId,
99 pub stage: FragmentStageV1,
100}
101
102pub struct Projection {
103 connection: Mutex<Connection>,
104}
105
106impl Projection {
107 pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, Option<TxId>), String> {
108 fs::create_dir_all(root).map_err(|error| format!("create projection root: {error}"))?;
109 let database = root.join(DATABASE_NAME);
110 let exists = database
111 .try_exists()
112 .map_err(|error| format!("inspect projection database: {error}"))?;
113 let (connection, cursor) = if exists {
114 match open_existing_database(&database) {
115 Ok(opened) => opened,
116 Err(OpenIssue::Recoverable) => recreate_database(&database)?,
117 Err(OpenIssue::Fatal(error)) => return Err(error),
118 }
119 } else {
120 remove_sidecars(&database)?;
121 create_database(&database)?
122 };
123 if let Some(cursor) = cursor {
124 let canonical = ordering
125 .get_txn(cursor)
126 .map_err(|error| format!("validate projection cursor: {error}"))?;
127 let valid = canonical.as_deref().is_some_and(|bytes| {
128 Transaction::parse(bytes)
129 .is_ok_and(|transaction| transaction.subsystem().as_str() == SUBSYSTEM)
130 });
131 if !valid {
132 drop(connection);
133 let (connection, cursor) = recreate_database(&database)?;
134 return Ok((
135 Self {
136 connection: Mutex::new(connection),
137 },
138 cursor,
139 ));
140 }
141 }
142 Ok((
143 Self {
144 connection: Mutex::new(connection),
145 },
146 cursor,
147 ))
148 }
149
150 pub fn apply(&self, callback_txid: TxId, payload: &[u8]) -> Result<AppliedEvent, String> {
151 let event = decode_event(payload).map_err(|error| format!("decode callback: {error}"))?;
152 let fragment_id = event_fragment_id(&event);
153 let mut connection = self.lock_connection()?;
154 let transaction = connection
155 .transaction()
156 .map_err(|error| format!("begin callback transaction: {error}"))?;
157 let status = load_status(&transaction, fragment_id)?;
158 let (status, effect) = reduce(status, &event, callback_txid)?;
159 write_status(&transaction, fragment_id, &status)?;
160 let updated = transaction
161 .execute(
162 "UPDATE metadata SET last_applied_txid = ?1 WHERE singleton = 1",
163 params![&callback_txid.as_bytes()[..]],
164 )
165 .map_err(|error| format!("advance projection cursor: {error}"))?;
166 if updated != 1 {
167 return Err("projection metadata row is missing".to_string());
168 }
169 transaction
170 .commit()
171 .map_err(|error| format!("commit callback transaction: {error}"))?;
172 Ok(AppliedEvent {
173 fragment_id,
174 effect,
175 })
176 }
177
178 pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
179 load_status(&self.lock_connection()?, fragment_id)
180 }
181
182 pub fn queued(&self) -> Result<Vec<FragmentId>, String> {
183 Ok(actionable_statuses(&self.lock_connection()?, 1)?
184 .into_iter()
185 .map(|(id, _)| id)
186 .collect())
187 }
188
189 pub fn running(&self) -> Result<Vec<InterruptedFragment>, String> {
190 Ok(actionable_statuses(&self.lock_connection()?, 2)?
191 .into_iter()
192 .map(|(fragment_id, status)| InterruptedFragment {
193 fragment_id,
194 stage: latest_running_stage(&status),
195 })
196 .collect())
197 }
198
199 pub fn validate_labels(
200 &self,
201 fragment_id: FragmentId,
202 labels: &[SpeakerLabelV1],
203 ) -> Result<TxId, String> {
204 let status = self
205 .status(fragment_id)?
206 .ok_or_else(|| "unknown fragment".to_string())?;
207 validate_confirmation(&status, None, labels).map(|value| value.0)
208 }
209
210 pub fn clear(&self) -> Result<(), String> {
211 let mut connection = self.lock_connection()?;
212 let transaction = connection
213 .transaction()
214 .map_err(|error| format!("begin clear transaction: {error}"))?;
215 transaction
216 .execute("DELETE FROM fragments", [])
217 .map_err(|error| format!("clear fragments: {error}"))?;
218 transaction
219 .execute(
220 "UPDATE metadata SET last_applied_txid = NULL WHERE singleton = 1",
221 [],
222 )
223 .map_err(|error| format!("clear projection cursor: {error}"))?;
224 transaction
225 .commit()
226 .map_err(|error| format!("commit clear transaction: {error}"))
227 }
228
229 #[cfg(feature = "testkit")]
230 pub fn inject_errors(
231 &self,
232 fragment_id: FragmentId,
233 errors: Vec<String>,
234 ) -> Result<(), String> {
235 let mut connection = self.lock_connection()?;
236 let transaction = connection
237 .transaction()
238 .map_err(|error| format!("begin error injection: {error}"))?;
239 let mut status = load_status(&transaction, fragment_id)?
240 .ok_or_else(|| "unknown fragment".to_string())?;
241 for error in errors {
242 append_error(&mut status, error);
243 }
244 write_status(&transaction, fragment_id, &status)?;
245 transaction
246 .commit()
247 .map_err(|error| format!("commit error injection: {error}"))
248 }
249
250 fn lock_connection(&self) -> Result<MutexGuard<'_, Connection>, String> {
251 self.connection
252 .lock()
253 .map_err(|_| "projection connection mutex is poisoned".to_string())
254 }
255}
256
257fn load_status<C: std::ops::Deref<Target = Connection>>(
258 connection: &C,
259 fragment_id: FragmentId,
260) -> Result<Option<FragmentStatus>, String> {
261 let stored = connection
262 .query_row(
263 "SELECT actionable_state, encoded_status FROM fragments WHERE fragment_id = ?1",
264 params![&fragment_id.as_bytes()[..]],
265 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
266 )
267 .optional()
268 .map_err(|error| format!("load fragment status: {error}"))?;
269 stored
270 .map(|(actionable, bytes)| decode_stored_status(&bytes, actionable))
271 .transpose()
272}
273
274fn write_status(
275 connection: &Connection,
276 fragment_id: FragmentId,
277 status: &FragmentStatus,
278) -> Result<(), String> {
279 let encoded = postcard::to_allocvec(status)
280 .map_err(|error| format!("encode fragment status: {error}"))?;
281 connection.execute(
282 "INSERT INTO fragments(fragment_id, actionable_state, encoded_status) VALUES(?1, ?2, ?3) ON CONFLICT(fragment_id) DO UPDATE SET actionable_state = excluded.actionable_state, encoded_status = excluded.encoded_status",
283 params![&fragment_id.as_bytes()[..], actionable_state(status.state), encoded],
284 ).map_err(|error| format!("write fragment status: {error}"))?;
285 Ok(())
286}
287
288fn actionable_statuses<C: std::ops::Deref<Target = Connection>>(
289 connection: &C,
290 actionable: i64,
291) -> Result<Vec<(FragmentId, FragmentStatus)>, String> {
292 let mut statement = connection
293 .prepare("SELECT fragment_id, encoded_status FROM fragments WHERE actionable_state = ?1")
294 .map_err(|error| format!("prepare actionable query: {error}"))?;
295 let mut rows = statement
296 .query(params![actionable])
297 .map_err(|error| format!("query actionable fragments: {error}"))?;
298 let mut statuses = Vec::new();
299 while let Some(row) = rows
300 .next()
301 .map_err(|error| format!("read actionable fragment: {error}"))?
302 {
303 let id: Vec<u8> = row
304 .get(0)
305 .map_err(|error| format!("read fragment ID: {error}"))?;
306 let encoded: Vec<u8> = row
307 .get(1)
308 .map_err(|error| format!("read fragment status: {error}"))?;
309 statuses.push((
310 fragment_id_from_slice(&id)?,
311 decode_stored_status(&encoded, actionable)?,
312 ));
313 }
314 Ok(statuses)
315}
316
317fn reduce(
318 current: Option<FragmentStatus>,
319 event: &AudioClassificationEventV3,
320 callback_txid: TxId,
321) -> Result<(FragmentStatus, ProjectionEffect), String> {
322 if let AudioClassificationEventV3::Queue(_) = event {
323 if current.is_some() {
324 return Err("duplicate Queue event".to_string());
325 }
326 return Ok((initial_status(), ProjectionEffect::Start));
327 }
328 let mut status = current.ok_or_else(|| "event references an unknown fragment".to_string())?;
329 if status.state == OverallState::Discarded {
330 let effect = if matches!(event, AudioClassificationEventV3::Discarded(_)) {
331 ProjectionEffect::Abort
332 } else {
333 ProjectionEffect::None
334 };
335 return Ok((status, effect));
336 }
337 match event {
338 AudioClassificationEventV3::Queue(_) => unreachable!(),
339 AudioClassificationEventV3::Progress(value) => apply_progress(&mut status, &value.update)?,
340 AudioClassificationEventV3::TranscriptionComplete(value) => {
341 if status.state != OverallState::Running {
342 return Err("TranscriptionComplete requires Running".to_string());
343 }
344 status.transcript.state = StageState::Succeeded;
345 status.speaker_labels.state = StageState::Succeeded;
346 status.speaker_features.state = StageState::Succeeded;
347 status.structuring.state = StageState::Succeeded;
348 status.label_confirmation.state = StageState::Pending;
349 status.interim_txid = Some(callback_txid);
350 status.analysis = Some(value.analysis.clone());
351 status.confirmed_labels.clear();
352 status.final_transcript = None;
353 status.state = OverallState::Completed;
354 }
355 AudioClassificationEventV3::Failed(value) => apply_terminal_failure(&mut status, value)?,
356 AudioClassificationEventV3::Discarded(_) => {
357 status.state = OverallState::Discarded;
358 return Ok((status, ProjectionEffect::Abort));
359 }
360 AudioClassificationEventV3::LabelConfirmation(value) => {
361 let (_, transcript) =
362 validate_confirmation(&status, Some(value.interim_txid), &value.speakers)?;
363 status.confirmed_labels = value.speakers.clone();
364 status.final_transcript = Some(transcript);
365 status.label_confirmation.state = StageState::Succeeded;
366 status.state = OverallState::Confirmed;
367 return Ok((status, ProjectionEffect::LabelsCommitted));
368 }
369 }
370 Ok((status, ProjectionEffect::None))
371}
372
373fn apply_progress(status: &mut FragmentStatus, update: &ProgressUpdateV1) -> Result<(), String> {
374 match update {
375 ProgressUpdateV1::LlmJobStarted {
376 sequence,
377 stage,
378 name,
379 } => start_job(status, *sequence, stage, name),
380 ProgressUpdateV1::LlmJobSucceeded { sequence } => {
381 finish_job(status, *sequence, LlmJobState::Succeeded, None)
382 }
383 ProgressUpdateV1::LlmJobFailed { sequence, error } => {
384 finish_job(status, *sequence, LlmJobState::Failed, Some(error.clone()))
385 }
386 ProgressUpdateV1::StageCompleted { stage } => complete_stage(status, stage),
387 }
388}
389
390fn start_job(
391 status: &mut FragmentStatus,
392 sequence: u64,
393 stage: &FragmentStageV1,
394 name: &str,
395) -> Result<(), String> {
396 if !is_analysis_stage(stage) {
397 return Err("LLM jobs require an analysis stage".to_string());
398 }
399 if name.trim().is_empty() {
400 return Err("LLM job name is blank".to_string());
401 }
402 if matches!(status.state, OverallState::Queued | OverallState::Failed) {
403 status.attempt_count = status
404 .attempt_count
405 .checked_add(1)
406 .ok_or_else(|| "attempt count overflow".to_string())?;
407 status.state = OverallState::Running;
408 status.transcript.state = StageState::Pending;
409 status.speaker_labels.state = StageState::Pending;
410 status.speaker_features.state = StageState::Pending;
411 status.structuring.state = StageState::Pending;
412 status.label_confirmation.state = StageState::Pending;
413 status.interim_txid = None;
414 status.analysis = None;
415 status.confirmed_labels.clear();
416 status.final_transcript = None;
417 } else if status.state != OverallState::Running {
418 return Err("LlmJobStarted requires Queued, Failed, or Running".to_string());
419 }
420 if status
421 .jobs
422 .iter()
423 .rev()
424 .find(|job| job.attempt == status.attempt_count)
425 .is_some_and(|job| sequence <= job.sequence)
426 {
427 return Err("LLM job sequence is not strictly increasing".to_string());
428 }
429 let stage_status = stage_status_mut(status, stage);
430 if !matches!(
431 stage_status.state,
432 StageState::Pending | StageState::Running
433 ) {
434 return Err("LLM job stage is not available".to_string());
435 }
436 stage_status.state = StageState::Running;
437 status.jobs.push(LlmJobStatus {
438 attempt: status.attempt_count,
439 sequence,
440 stage: *stage,
441 name: name.to_string(),
442 state: LlmJobState::Running,
443 });
444 Ok(())
445}
446
447fn finish_job(
448 status: &mut FragmentStatus,
449 sequence: u64,
450 terminal: LlmJobState,
451 error: Option<String>,
452) -> Result<(), String> {
453 if status.state != OverallState::Running {
454 return Err("LLM job terminal event requires Running".to_string());
455 }
456 let job = status
457 .jobs
458 .iter_mut()
459 .find(|job| job.attempt == status.attempt_count && job.sequence == sequence)
460 .ok_or_else(|| "unknown current-attempt LLM job".to_string())?;
461 if job.state != LlmJobState::Running {
462 return Err("LLM job is already terminal".to_string());
463 }
464 job.state = terminal;
465 if let Some(error) = error {
466 append_error(status, error);
467 }
468 Ok(())
469}
470
471fn complete_stage(status: &mut FragmentStatus, stage: &FragmentStageV1) -> Result<(), String> {
472 if status.state != OverallState::Running || !is_analysis_stage(stage) {
473 return Err("StageCompleted requires a running analysis stage".to_string());
474 }
475 let stage_status = stage_status_mut(status, stage);
476 if stage_status.state != StageState::Running {
477 return Err("StageCompleted stage is not Running".to_string());
478 }
479 stage_status.state = StageState::Succeeded;
480 Ok(())
481}
482
483fn apply_terminal_failure(
484 status: &mut FragmentStatus,
485 value: &kcode_k1_audio_classification_format::FailedV2,
486) -> Result<(), String> {
487 if !matches!(status.state, OverallState::Queued | OverallState::Running) {
488 return Err("Failed requires Queued or Running".to_string());
489 }
490 if let Some(sequence) = value.llm_job_sequence {
491 let job = status
492 .jobs
493 .iter_mut()
494 .find(|job| job.attempt == status.attempt_count && job.sequence == sequence)
495 .ok_or_else(|| "Failed references an unknown current-attempt job".to_string())?;
496 if job.stage != value.stage {
497 return Err("Failed job stage does not match".to_string());
498 }
499 match job.state {
500 LlmJobState::Running => job.state = LlmJobState::Failed,
501 LlmJobState::Failed => {}
502 LlmJobState::Succeeded => return Err("Failed references a succeeded job".to_string()),
503 }
504 }
505 stage_status_mut(status, &value.stage).state = StageState::Failed;
506 status.state = OverallState::Failed;
507 append_error(status, value.error.clone());
508 Ok(())
509}
510
511fn validate_confirmation(
512 status: &FragmentStatus,
513 supplied: Option<TxId>,
514 labels: &[SpeakerLabelV1],
515) -> Result<(TxId, String), String> {
516 if status.state != OverallState::Completed {
517 return Err("label confirmation requires Completed".to_string());
518 }
519 let interim = status
520 .interim_txid
521 .ok_or_else(|| "completed status has no interim transaction ID".to_string())?;
522 if supplied.is_some_and(|value| value != interim) {
523 return Err("label confirmation interim transaction ID does not match".to_string());
524 }
525 let analysis = status
526 .analysis
527 .as_ref()
528 .ok_or_else(|| "completed status has no analysis".to_string())?;
529 if labels.len() != analysis.envelope.analysis.speakers.len() {
530 return Err("speaker labels are not one-to-one".to_string());
531 }
532 for (label, expected) in labels.iter().zip(&analysis.envelope.analysis.speakers) {
533 if label.speaker != expected.speaker {
534 return Err("speaker labels are not in exact analysis order".to_string());
535 }
536 if invalid_person_id(&label.person_id) {
537 return Err("person ID is blank or contains a line break".to_string());
538 }
539 }
540 Ok((
541 interim,
542 replace_transcript(&analysis.envelope.analysis.transcript, labels),
543 ))
544}
545
546fn replace_transcript(transcript: &str, labels: &[SpeakerLabelV1]) -> String {
547 let mut output = String::with_capacity(transcript.len());
548 for line in transcript.split_inclusive('\n') {
549 let mut replaced = false;
550 for prefix in ["[high] ", "[medium] ", "[low] "] {
551 if let Some(rest) = line.strip_prefix(prefix) {
552 for label in labels {
553 let speaker = label.speaker.to_string();
554 if let Some(tail) = rest.strip_prefix(&speaker)
555 && (tail.starts_with(':') || tail.starts_with(" [overlap]:"))
556 {
557 output.push_str(prefix);
558 output.push_str(&label.person_id);
559 output.push_str(tail);
560 replaced = true;
561 break;
562 }
563 }
564 }
565 if replaced {
566 break;
567 }
568 }
569 if !replaced {
570 output.push_str(line);
571 }
572 }
573 output
574}
575
576fn initial_status() -> FragmentStatus {
577 FragmentStatus {
578 state: OverallState::Queued,
579 queue: stage_status(FragmentStageV1::Queue, StageState::Succeeded),
580 transcript: stage_status(FragmentStageV1::Transcript, StageState::Pending),
581 speaker_labels: stage_status(FragmentStageV1::SpeakerLabels, StageState::Pending),
582 speaker_features: stage_status(FragmentStageV1::SpeakerFeatures, StageState::Pending),
583 structuring: stage_status(FragmentStageV1::Structuring, StageState::Pending),
584 label_confirmation: stage_status(FragmentStageV1::LabelConfirmation, StageState::Pending),
585 attempt_count: 0,
586 jobs: Vec::new(),
587 interim_txid: None,
588 analysis: None,
589 confirmed_labels: Vec::new(),
590 final_transcript: None,
591 errors: Vec::new(),
592 errors_truncated: false,
593 }
594}
595
596fn stage_status(stage: FragmentStageV1, state: StageState) -> StageStatus {
597 StageStatus { stage, state }
598}
599
600fn stage_status_mut<'a>(
601 status: &'a mut FragmentStatus,
602 stage: &FragmentStageV1,
603) -> &'a mut StageStatus {
604 match stage {
605 FragmentStageV1::Queue => &mut status.queue,
606 FragmentStageV1::Transcript => &mut status.transcript,
607 FragmentStageV1::SpeakerLabels => &mut status.speaker_labels,
608 FragmentStageV1::SpeakerFeatures => &mut status.speaker_features,
609 FragmentStageV1::Structuring => &mut status.structuring,
610 FragmentStageV1::LabelConfirmation => &mut status.label_confirmation,
611 }
612}
613
614fn is_analysis_stage(stage: &FragmentStageV1) -> bool {
615 matches!(
616 stage,
617 FragmentStageV1::Transcript
618 | FragmentStageV1::SpeakerLabels
619 | FragmentStageV1::SpeakerFeatures
620 | FragmentStageV1::Structuring
621 )
622}
623
624fn latest_running_stage(status: &FragmentStatus) -> FragmentStageV1 {
625 for stage in [
626 &status.structuring,
627 &status.speaker_features,
628 &status.speaker_labels,
629 &status.transcript,
630 ] {
631 if stage.state == StageState::Running {
632 return stage.stage;
633 }
634 }
635 FragmentStageV1::Queue
636}
637
638fn append_error(status: &mut FragmentStatus, error: String) {
639 if status.errors.len() < MAX_ERRORS {
640 status.errors.push(error);
641 } else {
642 status.errors_truncated = true;
643 }
644}
645
646fn invalid_person_id(value: &str) -> bool {
647 value.trim().is_empty() || value.contains('\r') || value.contains('\n')
648}
649
650fn event_fragment_id(event: &AudioClassificationEventV3) -> FragmentId {
651 match event {
652 AudioClassificationEventV3::Queue(value) => value.audio_object_id,
653 AudioClassificationEventV3::Progress(value) => value.fragment_id,
654 AudioClassificationEventV3::TranscriptionComplete(value) => value.fragment_id,
655 AudioClassificationEventV3::Failed(value) => value.fragment_id,
656 AudioClassificationEventV3::Discarded(value) => value.fragment_id,
657 AudioClassificationEventV3::LabelConfirmation(value) => value.fragment_id,
658 }
659}
660
661fn fragment_id_from_slice(bytes: &[u8]) -> Result<FragmentId, String> {
662 let bytes: [u8; 12] = bytes
663 .try_into()
664 .map_err(|_| "stored fragment ID is not 12 bytes".to_string())?;
665 Ok(FragmentId::from_bytes(bytes))
666}
667
668fn actionable_state(state: OverallState) -> i64 {
669 match state {
670 OverallState::Queued => 1,
671 OverallState::Running => 2,
672 OverallState::Failed
673 | OverallState::Completed
674 | OverallState::Confirmed
675 | OverallState::Discarded => 0,
676 }
677}
678
679fn decode_stored_status(bytes: &[u8], actionable: i64) -> Result<FragmentStatus, String> {
680 let status: FragmentStatus = postcard::from_bytes(bytes)
681 .map_err(|error| format!("decode stored fragment status: {error}"))?;
682 let canonical = postcard::to_allocvec(&status)
683 .map_err(|error| format!("re-encode stored fragment status: {error}"))?;
684 if canonical != bytes {
685 return Err("stored fragment status is noncanonical".to_string());
686 }
687 validate_stored_status(&status, actionable)?;
688 Ok(status)
689}
690
691fn validate_stored_status(status: &FragmentStatus, actionable: i64) -> Result<(), String> {
692 let identities = [
693 (&status.queue, FragmentStageV1::Queue),
694 (&status.transcript, FragmentStageV1::Transcript),
695 (&status.speaker_labels, FragmentStageV1::SpeakerLabels),
696 (&status.speaker_features, FragmentStageV1::SpeakerFeatures),
697 (&status.structuring, FragmentStageV1::Structuring),
698 (
699 &status.label_confirmation,
700 FragmentStageV1::LabelConfirmation,
701 ),
702 ];
703 if identities
704 .iter()
705 .any(|(stored, expected)| stored.stage != *expected)
706 {
707 return Err("stored stage identity does not match its field".to_string());
708 }
709 if actionable_state(status.state) != actionable {
710 return Err("stored actionable state does not match status".to_string());
711 }
712 if status.errors.len() > MAX_ERRORS {
713 return Err("stored errors exceed the retention bound".to_string());
714 }
715 let mut previous = None;
716 for job in &status.jobs {
717 if job.attempt == 0
718 || job.attempt > status.attempt_count
719 || !is_analysis_stage(&job.stage)
720 || job.name.trim().is_empty()
721 || previous.is_some_and(|value| value >= (job.attempt, job.sequence))
722 {
723 return Err("stored LLM jobs are invalid or unordered".to_string());
724 }
725 previous = Some((job.attempt, job.sequence));
726 }
727 if !matches!(
728 status.queue.state,
729 StageState::Succeeded | StageState::Failed
730 ) {
731 return Err("stored Queue stage is neither succeeded nor failed".to_string());
732 }
733 if status.state == OverallState::Queued && status.attempt_count != 0 {
734 return Err("stored queued status has an attempt".to_string());
735 }
736 if matches!(
737 status.state,
738 OverallState::Completed | OverallState::Confirmed
739 ) && (status.interim_txid.is_none() || status.analysis.is_none())
740 {
741 return Err("stored completed status lacks its analysis".to_string());
742 }
743 if status.state == OverallState::Confirmed
744 && (status.final_transcript.is_none()
745 || status.label_confirmation.state != StageState::Succeeded)
746 {
747 return Err("stored confirmed status is incomplete".to_string());
748 }
749 if status
750 .confirmed_labels
751 .iter()
752 .any(|label| invalid_person_id(&label.person_id))
753 {
754 return Err("stored person ID is invalid".to_string());
755 }
756 Ok(())
757}
758
759fn create_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
760 let connection =
761 Connection::open(path).map_err(|error| format!("create projection database: {error}"))?;
762 configure_database(&connection)
763 .map_err(|error| format!("configure projection database: {error}"))?;
764 connection.execute_batch(&format!(
765 "{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata(singleton, schema_version, last_applied_txid) VALUES(1, {SCHEMA_VERSION}, NULL);"
766 )).map_err(|error| format!("initialize projection schema: {error}"))?;
767 Ok((connection, None))
768}
769
770fn open_existing_database(path: &Path) -> Result<(Connection, Option<TxId>), OpenIssue> {
771 let connection = Connection::open(path).map_err(classify_sql_error)?;
772 configure_database(&connection).map_err(classify_sql_error)?;
773 let quick_check: String = connection
774 .query_row("PRAGMA quick_check", [], |row| row.get(0))
775 .map_err(classify_sql_error)?;
776 if quick_check != "ok" {
777 return Err(OpenIssue::Recoverable);
778 }
779 validate_schema(&connection)?;
780 let cursor = validate_rows(&connection)?;
781 Ok((connection, cursor))
782}
783
784fn configure_database(connection: &Connection) -> Result<(), SqlError> {
785 connection.pragma_update(None, "journal_mode", "WAL")?;
786 connection.pragma_update(None, "synchronous", "FULL")
787}
788
789fn validate_schema(connection: &Connection) -> Result<(), OpenIssue> {
790 let mut statement = connection.prepare(
791 "SELECT type, name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
792 ).map_err(classify_sql_error)?;
793 let mut rows = statement.query([]).map_err(classify_sql_error)?;
794 let mut schema = Vec::new();
795 while let Some(row) = rows.next().map_err(classify_sql_error)? {
796 let kind: String = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
797 let name: String = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
798 let sql: String = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
799 schema.push((kind, name, sql));
800 }
801 let expected = vec![
802 (
803 "index".to_string(),
804 "fragments_actionable_state".to_string(),
805 INDEX_SQL.to_string(),
806 ),
807 (
808 "table".to_string(),
809 "fragments".to_string(),
810 FRAGMENTS_SQL.to_string(),
811 ),
812 (
813 "table".to_string(),
814 "metadata".to_string(),
815 METADATA_SQL.to_string(),
816 ),
817 ];
818 if schema != expected {
819 return Err(OpenIssue::Recoverable);
820 }
821 Ok(())
822}
823
824fn validate_rows(connection: &Connection) -> Result<Option<TxId>, OpenIssue> {
825 let mut metadata = connection
826 .prepare("SELECT singleton, schema_version, last_applied_txid FROM metadata")
827 .map_err(classify_sql_error)?;
828 let mut rows = metadata.query([]).map_err(classify_sql_error)?;
829 let row = rows
830 .next()
831 .map_err(classify_sql_error)?
832 .ok_or(OpenIssue::Recoverable)?;
833 let singleton: i64 = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
834 let version: i64 = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
835 let cursor: Option<Vec<u8>> = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
836 if singleton != 1
837 || version != SCHEMA_VERSION
838 || rows.next().map_err(classify_sql_error)?.is_some()
839 {
840 return Err(OpenIssue::Recoverable);
841 }
842 drop(rows);
843 drop(metadata);
844 let mut fragments = connection
845 .prepare("SELECT fragment_id, actionable_state, encoded_status FROM fragments")
846 .map_err(classify_sql_error)?;
847 let mut rows = fragments.query([]).map_err(classify_sql_error)?;
848 while let Some(row) = rows.next().map_err(classify_sql_error)? {
849 let id: Vec<u8> = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
850 let actionable: i64 = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
851 let encoded: Vec<u8> = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
852 fragment_id_from_slice(&id).map_err(|_| OpenIssue::Recoverable)?;
853 decode_stored_status(&encoded, actionable).map_err(|_| OpenIssue::Recoverable)?;
854 }
855 cursor
856 .map(|bytes| {
857 let bytes: [u8; 12] = bytes.try_into().map_err(|_| OpenIssue::Recoverable)?;
858 Ok(TxId::from_bytes(bytes))
859 })
860 .transpose()
861}
862
863fn classify_sql_error(error: SqlError) -> OpenIssue {
864 if matches!(
865 error.sqlite_error_code(),
866 Some(ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
867 ) {
868 OpenIssue::Recoverable
869 } else {
870 OpenIssue::Fatal(error.to_string())
871 }
872}
873
874fn recreate_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
875 remove_database_files(path)?;
876 create_database(path)
877}
878
879fn remove_sidecars(path: &Path) -> Result<(), String> {
880 for suffix in ["-wal", "-shm"] {
881 remove_if_present(&path_with_suffix(path, suffix))?;
882 }
883 Ok(())
884}
885
886fn remove_database_files(path: &Path) -> Result<(), String> {
887 remove_sidecars(path)?;
888 remove_if_present(path)
889}
890
891fn remove_if_present(path: &Path) -> Result<(), String> {
892 match fs::remove_file(path) {
893 Ok(()) => Ok(()),
894 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
895 Err(error) => Err(format!("remove recoverable projection state: {error}")),
896 }
897}
898
899fn path_with_suffix(path: &Path, suffix: &str) -> PathBuf {
900 let mut value = OsString::from(path.as_os_str());
901 value.push(suffix);
902 PathBuf::from(value)
903}
904
905enum OpenIssue {
906 Recoverable,
907 Fatal(String),
908}
909
910mod optional_fragment_id {
911 use super::FragmentId;
912 use serde::{Deserialize, Deserializer, Serialize, Serializer};
913
914 pub fn serialize<S: Serializer>(
915 value: &Option<FragmentId>,
916 serializer: S,
917 ) -> Result<S::Ok, S::Error> {
918 value.map(FragmentId::into_bytes).serialize(serializer)
919 }
920
921 pub fn deserialize<'de, D: Deserializer<'de>>(
922 deserializer: D,
923 ) -> Result<Option<FragmentId>, D::Error> {
924 Option::<[u8; 12]>::deserialize(deserializer).map(|value| value.map(FragmentId::from_bytes))
925 }
926}
927
928#[cfg(test)]
929mod tests {
930 use super::*;
931 use kcode_k1_audio_classification_format::{
932 DiscardedV2, FailedV2, LabelConfirmationV1, ProgressV1, QueueV2, TranscriptionCompleteV1,
933 encode_event,
934 };
935 use kcode_k1_transaction::{GENESIS_PARENT, SubsystemId, build_signed_transaction};
936 use kcode_speaker_v3_analysis::{
937 AnalysisEnvelope, FeatureVector24, GeminiCohort, LocalSpeakerLabel, OggAudioMetadata,
938 StructuredAnalysis, StructuredSpeaker, StructurerProvenance,
939 };
940 use std::sync::atomic::{AtomicU64, Ordering};
941
942 static NEXT_ROOT: AtomicU64 = AtomicU64::new(0);
943
944 struct TestRoot(PathBuf);
945
946 impl TestRoot {
947 fn new() -> Self {
948 let sequence = NEXT_ROOT.fetch_add(1, Ordering::Relaxed);
949 let path = std::env::temp_dir().join(format!(
950 "kcode-audio-projection-{}-{sequence}",
951 std::process::id()
952 ));
953 fs::create_dir_all(&path).unwrap();
954 Self(path)
955 }
956 fn join(&self, name: &str) -> PathBuf {
957 self.0.join(name)
958 }
959 }
960
961 impl Drop for TestRoot {
962 fn drop(&mut self) {
963 let _ = fs::remove_dir_all(&self.0);
964 }
965 }
966
967 fn txid(number: u32) -> TxId {
968 let mut bytes = [0; 12];
969 bytes[..4].copy_from_slice(&number.to_le_bytes());
970 TxId::from_bytes(bytes)
971 }
972
973 fn queue(fragment_id: FragmentId) -> AudioClassificationEventV3 {
974 AudioClassificationEventV3::Queue(QueueV2 {
975 audio_object_id: fragment_id,
976 })
977 }
978
979 fn apply_event(
980 projection: &Projection,
981 callback: TxId,
982 event: AudioClassificationEventV3,
983 ) -> AppliedEvent {
984 projection
985 .apply(callback, &encode_event(&event).unwrap())
986 .unwrap()
987 }
988
989 fn progress(fragment_id: FragmentId, update: ProgressUpdateV1) -> AudioClassificationEventV3 {
990 AudioClassificationEventV3::Progress(ProgressV1 {
991 fragment_id,
992 update,
993 })
994 }
995
996 fn canonical_event(
997 ordering: &K1TxnOrdering,
998 event: &AudioClassificationEventV3,
999 ) -> (TxId, Vec<u8>) {
1000 let payload = encode_event(event).unwrap();
1001 let bytes = build_signed_transaction(
1002 ordering.tip().unwrap_or(GENESIS_PARENT),
1003 1,
1004 [1; 32],
1005 SubsystemId::from_str(SUBSYSTEM).unwrap(),
1006 &payload,
1007 |_| Ok([2; 64]),
1008 )
1009 .unwrap();
1010 let id = TxId::for_transaction(&bytes);
1011 assert!(ordering.submit_txn(&bytes).is_ok());
1012 (id, payload)
1013 }
1014
1015 fn analysis(transcript: &str, speaker_numbers: &[u32]) -> ExecutedAnalysis {
1016 let mut ogg = vec![0; 47];
1017 ogg[..4].copy_from_slice(b"OggS");
1018 ogg[5] = 2;
1019 ogg[26] = 1;
1020 ogg[27] = 19;
1021 ogg[28..36].copy_from_slice(b"OpusHead");
1022 ogg[36] = 1;
1023 ogg[37] = 1;
1024 ogg[40..44].copy_from_slice(&48_000_u32.to_le_bytes());
1025 let speakers = speaker_numbers
1026 .iter()
1027 .map(|number| StructuredSpeaker {
1028 speaker: LocalSpeakerLabel::new(*number).unwrap(),
1029 language: "en".to_string(),
1030 features: FeatureVector24::default(),
1031 features_usable_for_training: true,
1032 })
1033 .collect();
1034 let provenance = StructurerProvenance {
1035 model_id: "model".to_string(),
1036 prompt_revision: "prompt".to_string(),
1037 };
1038 ExecutedAnalysis {
1039 envelope: AnalysisEnvelope {
1040 audio: OggAudioMetadata::from_bytes(&ogg, 1_000, None).unwrap(),
1041 analysis: StructuredAnalysis {
1042 transcript: transcript.to_string(),
1043 speakers,
1044 },
1045 gemini: GeminiCohort {
1046 model_id: "gemini".to_string(),
1047 transcript_prompt_revision: "transcript".to_string(),
1048 feature_prompt_revisions: [
1049 "one".to_string(),
1050 "two".to_string(),
1051 "three".to_string(),
1052 ],
1053 feature_schema_revision: "schema".to_string(),
1054 },
1055 structurer: provenance.clone(),
1056 },
1057 label_extractor: provenance,
1058 }
1059 }
1060
1061 fn label(speaker: u32, person_id: &str) -> SpeakerLabelV1 {
1062 SpeakerLabelV1 {
1063 speaker: LocalSpeakerLabel::new(speaker).unwrap(),
1064 person_id: person_id.to_string(),
1065 }
1066 }
1067
1068 #[test]
1069 fn atomic_cursor_reopen_and_clear() {
1070 let root = TestRoot::new();
1071 let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
1072 let projection_root = root.join("projection");
1073 let fragment_id = txid(10);
1074 let event = queue(fragment_id);
1075 let (callback, payload) = canonical_event(&ordering, &event);
1076 let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
1077 assert_eq!(cursor, None);
1078 projection.apply(callback, &payload).unwrap();
1079 drop(projection);
1080 let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
1081 assert_eq!(cursor, Some(callback));
1082 assert_eq!(projection.queued().unwrap(), vec![fragment_id]);
1083 projection.clear().unwrap();
1084 drop(projection);
1085 let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
1086 assert_eq!(cursor, None);
1087 assert_eq!(projection.status(fragment_id).unwrap(), None);
1088 }
1089
1090 #[test]
1091 fn queue_failure_survives_reopen() {
1092 let root = TestRoot::new();
1093 let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
1094 let projection_root = root.join("projection");
1095 let fragment_id = txid(15);
1096 let (queue_txid, queue_payload) = canonical_event(&ordering, &queue(fragment_id));
1097 let (projection, _) = Projection::open(&projection_root, &ordering).unwrap();
1098 projection.apply(queue_txid, &queue_payload).unwrap();
1099 let failure = AudioClassificationEventV3::Failed(FailedV2 {
1100 fragment_id,
1101 stage: FragmentStageV1::Queue,
1102 llm_job_sequence: None,
1103 error: "queue failed".to_string(),
1104 });
1105 let (failure_txid, failure_payload) = canonical_event(&ordering, &failure);
1106 projection.apply(failure_txid, &failure_payload).unwrap();
1107 drop(projection);
1108 let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
1109 let status = projection.status(fragment_id).unwrap().unwrap();
1110 assert_eq!(cursor, Some(failure_txid));
1111 assert_eq!(status.state, OverallState::Failed);
1112 assert_eq!(status.queue.state, StageState::Failed);
1113 }
1114
1115 #[test]
1116 fn failure_retry_completion_and_exact_confirmation() {
1117 let root = TestRoot::new();
1118 let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
1119 let (projection, _) = Projection::open(&root.join("projection"), &ordering).unwrap();
1120 let fragment_id = txid(20);
1121 apply_event(&projection, txid(21), queue(fragment_id));
1122 apply_event(
1123 &projection,
1124 txid(22),
1125 progress(
1126 fragment_id,
1127 ProgressUpdateV1::LlmJobStarted {
1128 sequence: 7,
1129 stage: FragmentStageV1::Transcript,
1130 name: "transcript".to_string(),
1131 },
1132 ),
1133 );
1134 assert_eq!(
1135 projection.running().unwrap()[0].stage,
1136 FragmentStageV1::Transcript
1137 );
1138 apply_event(
1139 &projection,
1140 txid(23),
1141 progress(
1142 fragment_id,
1143 ProgressUpdateV1::LlmJobFailed {
1144 sequence: 7,
1145 error: "provider".to_string(),
1146 },
1147 ),
1148 );
1149 apply_event(
1150 &projection,
1151 txid(24),
1152 AudioClassificationEventV3::Failed(FailedV2 {
1153 fragment_id,
1154 stage: FragmentStageV1::Transcript,
1155 llm_job_sequence: Some(7),
1156 error: "terminal".to_string(),
1157 }),
1158 );
1159 let failed = projection.status(fragment_id).unwrap().unwrap();
1160 assert_eq!(failed.state, OverallState::Failed);
1161 assert_eq!(failed.jobs[0].state, LlmJobState::Failed);
1162 apply_event(
1163 &projection,
1164 txid(25),
1165 progress(
1166 fragment_id,
1167 ProgressUpdateV1::LlmJobStarted {
1168 sequence: 1,
1169 stage: FragmentStageV1::Transcript,
1170 name: "retry".to_string(),
1171 },
1172 ),
1173 );
1174 let transcript = "[high] Speaker 2: hello\n[medium] Speaker 7 [overlap]: hi\n[low] Speaker 1: unchanged\nplain Speaker 2: unchanged\n";
1175 apply_event(
1176 &projection,
1177 txid(26),
1178 AudioClassificationEventV3::TranscriptionComplete(TranscriptionCompleteV1 {
1179 fragment_id,
1180 analysis: analysis(transcript, &[2, 7]),
1181 }),
1182 );
1183 let labels = vec![label(2, "alice"), label(7, "bob")];
1184 assert_eq!(
1185 projection.validate_labels(fragment_id, &labels).unwrap(),
1186 txid(26)
1187 );
1188 let applied = apply_event(
1189 &projection,
1190 txid(27),
1191 AudioClassificationEventV3::LabelConfirmation(LabelConfirmationV1 {
1192 fragment_id,
1193 interim_txid: txid(26),
1194 speakers: labels.clone(),
1195 }),
1196 );
1197 assert_eq!(applied.effect, ProjectionEffect::LabelsCommitted);
1198 let status = projection.status(fragment_id).unwrap().unwrap();
1199 assert_eq!(status.confirmed_labels, labels);
1200 assert_eq!(
1201 status.final_transcript.as_deref(),
1202 Some(
1203 "[high] alice: hello\n[medium] bob [overlap]: hi\n[low] Speaker 1: unchanged\nplain Speaker 2: unchanged\n"
1204 )
1205 );
1206 }
1207
1208 #[test]
1209 fn repeated_discard_suppresses_late_events() {
1210 let root = TestRoot::new();
1211 let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
1212 let (projection, _) = Projection::open(&root.join("projection"), &ordering).unwrap();
1213 let fragment_id = txid(30);
1214 apply_event(&projection, txid(31), queue(fragment_id));
1215 let discarded = || AudioClassificationEventV3::Discarded(DiscardedV2 { fragment_id });
1216 assert_eq!(
1217 apply_event(&projection, txid(32), discarded()).effect,
1218 ProjectionEffect::Abort
1219 );
1220 let status = projection.status(fragment_id).unwrap().unwrap();
1221 assert_eq!(
1222 apply_event(&projection, txid(33), discarded()).effect,
1223 ProjectionEffect::Abort
1224 );
1225 apply_event(
1226 &projection,
1227 txid(34),
1228 AudioClassificationEventV3::TranscriptionComplete(TranscriptionCompleteV1 {
1229 fragment_id,
1230 analysis: analysis("[high] Speaker 2: late", &[2]),
1231 }),
1232 );
1233 apply_event(
1234 &projection,
1235 txid(35),
1236 AudioClassificationEventV3::Failed(FailedV2 {
1237 fragment_id,
1238 stage: FragmentStageV1::Transcript,
1239 llm_job_sequence: None,
1240 error: "late".to_string(),
1241 }),
1242 );
1243 assert_eq!(projection.status(fragment_id).unwrap().unwrap(), status);
1244 }
1245
1246 #[test]
1247 fn first_five_thousand_errors_are_retained() {
1248 let mut status = initial_status();
1249 for index in 0..=MAX_ERRORS {
1250 append_error(&mut status, index.to_string());
1251 }
1252 assert_eq!(status.errors.len(), MAX_ERRORS);
1253 assert_eq!(status.errors.first().map(String::as_str), Some("0"));
1254 assert_eq!(status.errors.last().map(String::as_str), Some("4999"));
1255 assert!(status.errors_truncated);
1256 }
1257
1258 #[test]
1259 fn malformed_and_noncanonical_databases_rebuild() {
1260 let malformed = TestRoot::new();
1261 let ordering = K1TxnOrdering::open(&malformed.join("ordering")).unwrap();
1262 let projection_root = malformed.join("projection");
1263 fs::create_dir_all(&projection_root).unwrap();
1264 fs::write(projection_root.join(DATABASE_NAME), b"malformed").unwrap();
1265 assert_eq!(
1266 Projection::open(&projection_root, &ordering).unwrap().1,
1267 None
1268 );
1269
1270 let noncanonical = TestRoot::new();
1271 let ordering = K1TxnOrdering::open(&noncanonical.join("ordering")).unwrap();
1272 let projection_root = noncanonical.join("projection");
1273 let (projection, _) = Projection::open(&projection_root, &ordering).unwrap();
1274 let fragment_id = txid(40);
1275 apply_event(&projection, txid(41), queue(fragment_id));
1276 drop(projection);
1277 let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
1278 assert_eq!(cursor, None);
1279 assert_eq!(projection.status(fragment_id).unwrap(), None);
1280 }
1281}