use kcode_k1_audio_classification_format::{
AudioClassificationEventV3, ProgressUpdateV1, decode_event,
};
pub use kcode_k1_audio_classification_format::{ExecutedAnalysis, FragmentStageV1, SpeakerLabelV1};
use kcode_k1_transaction::Transaction;
use kcode_k1_txn_ordering::{K1TxnOrdering, TxId};
use rusqlite::{Connection, Error as SqlError, ErrorCode, OptionalExtension, params};
use serde::{Deserialize, Serialize};
use std::ffi::OsString;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard};
pub type FragmentId = kcode_k1_audio_fragment_submit::FragmentId;
const DATABASE_NAME: &str = "audio-classification.sqlite3";
const SUBSYSTEM: &str = "audio-classification";
const SCHEMA_VERSION: i64 = 1;
const MAX_ERRORS: usize = 5_000;
const METADATA_SQL: &str = "CREATE TABLE metadata(singleton INTEGER PRIMARY KEY CHECK(singleton = 1), schema_version INTEGER NOT NULL, last_applied_txid BLOB)";
const FRAGMENTS_SQL: &str = "CREATE TABLE fragments(fragment_id BLOB PRIMARY KEY, actionable_state INTEGER NOT NULL, encoded_status BLOB NOT NULL)";
const INDEX_SQL: &str = "CREATE INDEX fragments_actionable_state ON fragments(actionable_state)";
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum OverallState {
Queued,
Running,
Failed,
Completed,
Confirmed,
Discarded,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum StageState {
Pending,
Running,
Succeeded,
Failed,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum LlmJobState {
Running,
Succeeded,
Failed,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct StageStatus {
pub stage: FragmentStageV1,
pub state: StageState,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LlmJobStatus {
pub attempt: u32,
pub sequence: u64,
pub stage: FragmentStageV1,
pub name: String,
pub state: LlmJobState,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct FragmentStatus {
pub state: OverallState,
pub queue: StageStatus,
pub transcript: StageStatus,
pub speaker_labels: StageStatus,
pub speaker_features: StageStatus,
pub structuring: StageStatus,
pub label_confirmation: StageStatus,
pub attempt_count: u32,
pub jobs: Vec<LlmJobStatus>,
#[serde(with = "optional_fragment_id")]
pub interim_txid: Option<FragmentId>,
pub analysis: Option<ExecutedAnalysis>,
pub confirmed_labels: Vec<SpeakerLabelV1>,
pub final_transcript: Option<String>,
pub errors: Vec<String>,
pub errors_truncated: bool,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ProjectionEffect {
None,
Start,
Abort,
LabelsCommitted,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AppliedEvent {
pub fragment_id: FragmentId,
pub effect: ProjectionEffect,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct InterruptedFragment {
pub fragment_id: FragmentId,
pub stage: FragmentStageV1,
}
pub struct Projection {
connection: Mutex<Connection>,
}
impl Projection {
pub fn open(root: &Path, ordering: &K1TxnOrdering) -> Result<(Self, Option<TxId>), String> {
fs::create_dir_all(root).map_err(|error| format!("create projection root: {error}"))?;
let database = root.join(DATABASE_NAME);
let exists = database
.try_exists()
.map_err(|error| format!("inspect projection database: {error}"))?;
let (connection, cursor) = if exists {
match open_existing_database(&database) {
Ok(opened) => opened,
Err(OpenIssue::Recoverable) => recreate_database(&database)?,
Err(OpenIssue::Fatal(error)) => return Err(error),
}
} else {
remove_sidecars(&database)?;
create_database(&database)?
};
if let Some(cursor) = cursor {
let canonical = ordering
.get_txn(cursor)
.map_err(|error| format!("validate projection cursor: {error}"))?;
let valid = canonical.as_deref().is_some_and(|bytes| {
Transaction::parse(bytes)
.is_ok_and(|transaction| transaction.subsystem().as_str() == SUBSYSTEM)
});
if !valid {
drop(connection);
let (connection, cursor) = recreate_database(&database)?;
return Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
));
}
}
Ok((
Self {
connection: Mutex::new(connection),
},
cursor,
))
}
pub fn apply(&self, callback_txid: TxId, payload: &[u8]) -> Result<AppliedEvent, String> {
let event = decode_event(payload).map_err(|error| format!("decode callback: {error}"))?;
let fragment_id = event_fragment_id(&event);
let mut connection = self.lock_connection()?;
let transaction = connection
.transaction()
.map_err(|error| format!("begin callback transaction: {error}"))?;
let status = load_status(&transaction, fragment_id)?;
let (status, effect) = reduce(status, &event, callback_txid)?;
write_status(&transaction, fragment_id, &status)?;
let updated = transaction
.execute(
"UPDATE metadata SET last_applied_txid = ?1 WHERE singleton = 1",
params![&callback_txid.as_bytes()[..]],
)
.map_err(|error| format!("advance projection cursor: {error}"))?;
if updated != 1 {
return Err("projection metadata row is missing".to_string());
}
transaction
.commit()
.map_err(|error| format!("commit callback transaction: {error}"))?;
Ok(AppliedEvent {
fragment_id,
effect,
})
}
pub fn status(&self, fragment_id: FragmentId) -> Result<Option<FragmentStatus>, String> {
load_status(&self.lock_connection()?, fragment_id)
}
pub fn queued(&self) -> Result<Vec<FragmentId>, String> {
Ok(actionable_statuses(&self.lock_connection()?, 1)?
.into_iter()
.map(|(id, _)| id)
.collect())
}
pub fn running(&self) -> Result<Vec<InterruptedFragment>, String> {
Ok(actionable_statuses(&self.lock_connection()?, 2)?
.into_iter()
.map(|(fragment_id, status)| InterruptedFragment {
fragment_id,
stage: latest_running_stage(&status),
})
.collect())
}
pub fn validate_labels(
&self,
fragment_id: FragmentId,
labels: &[SpeakerLabelV1],
) -> Result<TxId, String> {
let status = self
.status(fragment_id)?
.ok_or_else(|| "unknown fragment".to_string())?;
validate_confirmation(&status, None, labels).map(|value| value.0)
}
pub fn clear(&self) -> Result<(), String> {
let mut connection = self.lock_connection()?;
let transaction = connection
.transaction()
.map_err(|error| format!("begin clear transaction: {error}"))?;
transaction
.execute("DELETE FROM fragments", [])
.map_err(|error| format!("clear fragments: {error}"))?;
transaction
.execute(
"UPDATE metadata SET last_applied_txid = NULL WHERE singleton = 1",
[],
)
.map_err(|error| format!("clear projection cursor: {error}"))?;
transaction
.commit()
.map_err(|error| format!("commit clear transaction: {error}"))
}
#[cfg(feature = "testkit")]
pub fn inject_errors(
&self,
fragment_id: FragmentId,
errors: Vec<String>,
) -> Result<(), String> {
let mut connection = self.lock_connection()?;
let transaction = connection
.transaction()
.map_err(|error| format!("begin error injection: {error}"))?;
let mut status = load_status(&transaction, fragment_id)?
.ok_or_else(|| "unknown fragment".to_string())?;
for error in errors {
append_error(&mut status, error);
}
write_status(&transaction, fragment_id, &status)?;
transaction
.commit()
.map_err(|error| format!("commit error injection: {error}"))
}
fn lock_connection(&self) -> Result<MutexGuard<'_, Connection>, String> {
self.connection
.lock()
.map_err(|_| "projection connection mutex is poisoned".to_string())
}
}
fn load_status<C: std::ops::Deref<Target = Connection>>(
connection: &C,
fragment_id: FragmentId,
) -> Result<Option<FragmentStatus>, String> {
let stored = connection
.query_row(
"SELECT actionable_state, encoded_status FROM fragments WHERE fragment_id = ?1",
params![&fragment_id.as_bytes()[..]],
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
)
.optional()
.map_err(|error| format!("load fragment status: {error}"))?;
stored
.map(|(actionable, bytes)| decode_stored_status(&bytes, actionable))
.transpose()
}
fn write_status(
connection: &Connection,
fragment_id: FragmentId,
status: &FragmentStatus,
) -> Result<(), String> {
let encoded = postcard::to_allocvec(status)
.map_err(|error| format!("encode fragment status: {error}"))?;
connection.execute(
"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",
params![&fragment_id.as_bytes()[..], actionable_state(status.state), encoded],
).map_err(|error| format!("write fragment status: {error}"))?;
Ok(())
}
fn actionable_statuses<C: std::ops::Deref<Target = Connection>>(
connection: &C,
actionable: i64,
) -> Result<Vec<(FragmentId, FragmentStatus)>, String> {
let mut statement = connection
.prepare("SELECT fragment_id, encoded_status FROM fragments WHERE actionable_state = ?1")
.map_err(|error| format!("prepare actionable query: {error}"))?;
let mut rows = statement
.query(params![actionable])
.map_err(|error| format!("query actionable fragments: {error}"))?;
let mut statuses = Vec::new();
while let Some(row) = rows
.next()
.map_err(|error| format!("read actionable fragment: {error}"))?
{
let id: Vec<u8> = row
.get(0)
.map_err(|error| format!("read fragment ID: {error}"))?;
let encoded: Vec<u8> = row
.get(1)
.map_err(|error| format!("read fragment status: {error}"))?;
statuses.push((
fragment_id_from_slice(&id)?,
decode_stored_status(&encoded, actionable)?,
));
}
Ok(statuses)
}
fn reduce(
current: Option<FragmentStatus>,
event: &AudioClassificationEventV3,
callback_txid: TxId,
) -> Result<(FragmentStatus, ProjectionEffect), String> {
if let AudioClassificationEventV3::Queue(_) = event {
if current.is_some() {
return Err("duplicate Queue event".to_string());
}
return Ok((initial_status(), ProjectionEffect::Start));
}
let mut status = current.ok_or_else(|| "event references an unknown fragment".to_string())?;
if status.state == OverallState::Discarded {
let effect = if matches!(event, AudioClassificationEventV3::Discarded(_)) {
ProjectionEffect::Abort
} else {
ProjectionEffect::None
};
return Ok((status, effect));
}
match event {
AudioClassificationEventV3::Queue(_) => unreachable!(),
AudioClassificationEventV3::Progress(value) => apply_progress(&mut status, &value.update)?,
AudioClassificationEventV3::TranscriptionComplete(value) => {
if status.state != OverallState::Running {
return Err("TranscriptionComplete requires Running".to_string());
}
status.transcript.state = StageState::Succeeded;
status.speaker_labels.state = StageState::Succeeded;
status.speaker_features.state = StageState::Succeeded;
status.structuring.state = StageState::Succeeded;
status.label_confirmation.state = StageState::Pending;
status.interim_txid = Some(callback_txid);
status.analysis = Some(value.analysis.clone());
status.confirmed_labels.clear();
status.final_transcript = None;
status.state = OverallState::Completed;
}
AudioClassificationEventV3::Failed(value) => apply_terminal_failure(&mut status, value)?,
AudioClassificationEventV3::Discarded(_) => {
status.state = OverallState::Discarded;
return Ok((status, ProjectionEffect::Abort));
}
AudioClassificationEventV3::LabelConfirmation(value) => {
let (_, transcript) =
validate_confirmation(&status, Some(value.interim_txid), &value.speakers)?;
status.confirmed_labels = value.speakers.clone();
status.final_transcript = Some(transcript);
status.label_confirmation.state = StageState::Succeeded;
status.state = OverallState::Confirmed;
return Ok((status, ProjectionEffect::LabelsCommitted));
}
}
Ok((status, ProjectionEffect::None))
}
fn apply_progress(status: &mut FragmentStatus, update: &ProgressUpdateV1) -> Result<(), String> {
match update {
ProgressUpdateV1::LlmJobStarted {
sequence,
stage,
name,
} => start_job(status, *sequence, stage, name),
ProgressUpdateV1::LlmJobSucceeded { sequence } => {
finish_job(status, *sequence, LlmJobState::Succeeded, None)
}
ProgressUpdateV1::LlmJobFailed { sequence, error } => {
finish_job(status, *sequence, LlmJobState::Failed, Some(error.clone()))
}
ProgressUpdateV1::StageCompleted { stage } => complete_stage(status, stage),
}
}
fn start_job(
status: &mut FragmentStatus,
sequence: u64,
stage: &FragmentStageV1,
name: &str,
) -> Result<(), String> {
if !is_analysis_stage(stage) {
return Err("LLM jobs require an analysis stage".to_string());
}
if name.trim().is_empty() {
return Err("LLM job name is blank".to_string());
}
if matches!(status.state, OverallState::Queued | OverallState::Failed) {
status.attempt_count = status
.attempt_count
.checked_add(1)
.ok_or_else(|| "attempt count overflow".to_string())?;
status.state = OverallState::Running;
status.transcript.state = StageState::Pending;
status.speaker_labels.state = StageState::Pending;
status.speaker_features.state = StageState::Pending;
status.structuring.state = StageState::Pending;
status.label_confirmation.state = StageState::Pending;
status.interim_txid = None;
status.analysis = None;
status.confirmed_labels.clear();
status.final_transcript = None;
} else if status.state != OverallState::Running {
return Err("LlmJobStarted requires Queued, Failed, or Running".to_string());
}
if status
.jobs
.iter()
.rev()
.find(|job| job.attempt == status.attempt_count)
.is_some_and(|job| sequence <= job.sequence)
{
return Err("LLM job sequence is not strictly increasing".to_string());
}
let stage_status = stage_status_mut(status, stage);
if !matches!(
stage_status.state,
StageState::Pending | StageState::Running
) {
return Err("LLM job stage is not available".to_string());
}
stage_status.state = StageState::Running;
status.jobs.push(LlmJobStatus {
attempt: status.attempt_count,
sequence,
stage: *stage,
name: name.to_string(),
state: LlmJobState::Running,
});
Ok(())
}
fn finish_job(
status: &mut FragmentStatus,
sequence: u64,
terminal: LlmJobState,
error: Option<String>,
) -> Result<(), String> {
if status.state != OverallState::Running {
return Err("LLM job terminal event requires Running".to_string());
}
let job = status
.jobs
.iter_mut()
.find(|job| job.attempt == status.attempt_count && job.sequence == sequence)
.ok_or_else(|| "unknown current-attempt LLM job".to_string())?;
if job.state != LlmJobState::Running {
return Err("LLM job is already terminal".to_string());
}
job.state = terminal;
if let Some(error) = error {
append_error(status, error);
}
Ok(())
}
fn complete_stage(status: &mut FragmentStatus, stage: &FragmentStageV1) -> Result<(), String> {
if status.state != OverallState::Running || !is_analysis_stage(stage) {
return Err("StageCompleted requires a running analysis stage".to_string());
}
let stage_status = stage_status_mut(status, stage);
if stage_status.state != StageState::Running {
return Err("StageCompleted stage is not Running".to_string());
}
stage_status.state = StageState::Succeeded;
Ok(())
}
fn apply_terminal_failure(
status: &mut FragmentStatus,
value: &kcode_k1_audio_classification_format::FailedV2,
) -> Result<(), String> {
if !matches!(status.state, OverallState::Queued | OverallState::Running) {
return Err("Failed requires Queued or Running".to_string());
}
if let Some(sequence) = value.llm_job_sequence {
let job = status
.jobs
.iter_mut()
.find(|job| job.attempt == status.attempt_count && job.sequence == sequence)
.ok_or_else(|| "Failed references an unknown current-attempt job".to_string())?;
if job.stage != value.stage {
return Err("Failed job stage does not match".to_string());
}
match job.state {
LlmJobState::Running => job.state = LlmJobState::Failed,
LlmJobState::Failed => {}
LlmJobState::Succeeded => return Err("Failed references a succeeded job".to_string()),
}
}
stage_status_mut(status, &value.stage).state = StageState::Failed;
status.state = OverallState::Failed;
append_error(status, value.error.clone());
Ok(())
}
fn validate_confirmation(
status: &FragmentStatus,
supplied: Option<TxId>,
labels: &[SpeakerLabelV1],
) -> Result<(TxId, String), String> {
if status.state != OverallState::Completed {
return Err("label confirmation requires Completed".to_string());
}
let interim = status
.interim_txid
.ok_or_else(|| "completed status has no interim transaction ID".to_string())?;
if supplied.is_some_and(|value| value != interim) {
return Err("label confirmation interim transaction ID does not match".to_string());
}
let analysis = status
.analysis
.as_ref()
.ok_or_else(|| "completed status has no analysis".to_string())?;
if labels.len() != analysis.envelope.analysis.speakers.len() {
return Err("speaker labels are not one-to-one".to_string());
}
for (label, expected) in labels.iter().zip(&analysis.envelope.analysis.speakers) {
if label.speaker != expected.speaker {
return Err("speaker labels are not in exact analysis order".to_string());
}
if invalid_person_id(&label.person_id) {
return Err("person ID is blank or contains a line break".to_string());
}
}
Ok((
interim,
replace_transcript(&analysis.envelope.analysis.transcript, labels),
))
}
fn replace_transcript(transcript: &str, labels: &[SpeakerLabelV1]) -> String {
let mut output = String::with_capacity(transcript.len());
for line in transcript.split_inclusive('\n') {
let mut replaced = false;
for prefix in ["[high] ", "[medium] ", "[low] "] {
if let Some(rest) = line.strip_prefix(prefix) {
for label in labels {
let speaker = label.speaker.to_string();
if let Some(tail) = rest.strip_prefix(&speaker)
&& (tail.starts_with(':') || tail.starts_with(" [overlap]:"))
{
output.push_str(prefix);
output.push_str(&label.person_id);
output.push_str(tail);
replaced = true;
break;
}
}
}
if replaced {
break;
}
}
if !replaced {
output.push_str(line);
}
}
output
}
fn initial_status() -> FragmentStatus {
FragmentStatus {
state: OverallState::Queued,
queue: stage_status(FragmentStageV1::Queue, StageState::Succeeded),
transcript: stage_status(FragmentStageV1::Transcript, StageState::Pending),
speaker_labels: stage_status(FragmentStageV1::SpeakerLabels, StageState::Pending),
speaker_features: stage_status(FragmentStageV1::SpeakerFeatures, StageState::Pending),
structuring: stage_status(FragmentStageV1::Structuring, StageState::Pending),
label_confirmation: stage_status(FragmentStageV1::LabelConfirmation, StageState::Pending),
attempt_count: 0,
jobs: Vec::new(),
interim_txid: None,
analysis: None,
confirmed_labels: Vec::new(),
final_transcript: None,
errors: Vec::new(),
errors_truncated: false,
}
}
fn stage_status(stage: FragmentStageV1, state: StageState) -> StageStatus {
StageStatus { stage, state }
}
fn stage_status_mut<'a>(
status: &'a mut FragmentStatus,
stage: &FragmentStageV1,
) -> &'a mut StageStatus {
match stage {
FragmentStageV1::Queue => &mut status.queue,
FragmentStageV1::Transcript => &mut status.transcript,
FragmentStageV1::SpeakerLabels => &mut status.speaker_labels,
FragmentStageV1::SpeakerFeatures => &mut status.speaker_features,
FragmentStageV1::Structuring => &mut status.structuring,
FragmentStageV1::LabelConfirmation => &mut status.label_confirmation,
}
}
fn is_analysis_stage(stage: &FragmentStageV1) -> bool {
matches!(
stage,
FragmentStageV1::Transcript
| FragmentStageV1::SpeakerLabels
| FragmentStageV1::SpeakerFeatures
| FragmentStageV1::Structuring
)
}
fn latest_running_stage(status: &FragmentStatus) -> FragmentStageV1 {
for stage in [
&status.structuring,
&status.speaker_features,
&status.speaker_labels,
&status.transcript,
] {
if stage.state == StageState::Running {
return stage.stage;
}
}
FragmentStageV1::Queue
}
fn append_error(status: &mut FragmentStatus, error: String) {
if status.errors.len() < MAX_ERRORS {
status.errors.push(error);
} else {
status.errors_truncated = true;
}
}
fn invalid_person_id(value: &str) -> bool {
value.trim().is_empty() || value.contains('\r') || value.contains('\n')
}
fn event_fragment_id(event: &AudioClassificationEventV3) -> FragmentId {
match event {
AudioClassificationEventV3::Queue(value) => value.audio_object_id,
AudioClassificationEventV3::Progress(value) => value.fragment_id,
AudioClassificationEventV3::TranscriptionComplete(value) => value.fragment_id,
AudioClassificationEventV3::Failed(value) => value.fragment_id,
AudioClassificationEventV3::Discarded(value) => value.fragment_id,
AudioClassificationEventV3::LabelConfirmation(value) => value.fragment_id,
}
}
fn fragment_id_from_slice(bytes: &[u8]) -> Result<FragmentId, String> {
let bytes: [u8; 12] = bytes
.try_into()
.map_err(|_| "stored fragment ID is not 12 bytes".to_string())?;
Ok(FragmentId::from_bytes(bytes))
}
fn actionable_state(state: OverallState) -> i64 {
match state {
OverallState::Queued => 1,
OverallState::Running => 2,
OverallState::Failed
| OverallState::Completed
| OverallState::Confirmed
| OverallState::Discarded => 0,
}
}
fn decode_stored_status(bytes: &[u8], actionable: i64) -> Result<FragmentStatus, String> {
let status: FragmentStatus = postcard::from_bytes(bytes)
.map_err(|error| format!("decode stored fragment status: {error}"))?;
let canonical = postcard::to_allocvec(&status)
.map_err(|error| format!("re-encode stored fragment status: {error}"))?;
if canonical != bytes {
return Err("stored fragment status is noncanonical".to_string());
}
validate_stored_status(&status, actionable)?;
Ok(status)
}
fn validate_stored_status(status: &FragmentStatus, actionable: i64) -> Result<(), String> {
let identities = [
(&status.queue, FragmentStageV1::Queue),
(&status.transcript, FragmentStageV1::Transcript),
(&status.speaker_labels, FragmentStageV1::SpeakerLabels),
(&status.speaker_features, FragmentStageV1::SpeakerFeatures),
(&status.structuring, FragmentStageV1::Structuring),
(
&status.label_confirmation,
FragmentStageV1::LabelConfirmation,
),
];
if identities
.iter()
.any(|(stored, expected)| stored.stage != *expected)
{
return Err("stored stage identity does not match its field".to_string());
}
if actionable_state(status.state) != actionable {
return Err("stored actionable state does not match status".to_string());
}
if status.errors.len() > MAX_ERRORS {
return Err("stored errors exceed the retention bound".to_string());
}
let mut previous = None;
for job in &status.jobs {
if job.attempt == 0
|| job.attempt > status.attempt_count
|| !is_analysis_stage(&job.stage)
|| job.name.trim().is_empty()
|| previous.is_some_and(|value| value >= (job.attempt, job.sequence))
{
return Err("stored LLM jobs are invalid or unordered".to_string());
}
previous = Some((job.attempt, job.sequence));
}
if !matches!(
status.queue.state,
StageState::Succeeded | StageState::Failed
) {
return Err("stored Queue stage is neither succeeded nor failed".to_string());
}
if status.state == OverallState::Queued && status.attempt_count != 0 {
return Err("stored queued status has an attempt".to_string());
}
if matches!(
status.state,
OverallState::Completed | OverallState::Confirmed
) && (status.interim_txid.is_none() || status.analysis.is_none())
{
return Err("stored completed status lacks its analysis".to_string());
}
if status.state == OverallState::Confirmed
&& (status.final_transcript.is_none()
|| status.label_confirmation.state != StageState::Succeeded)
{
return Err("stored confirmed status is incomplete".to_string());
}
if status
.confirmed_labels
.iter()
.any(|label| invalid_person_id(&label.person_id))
{
return Err("stored person ID is invalid".to_string());
}
Ok(())
}
fn create_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
let connection =
Connection::open(path).map_err(|error| format!("create projection database: {error}"))?;
configure_database(&connection)
.map_err(|error| format!("configure projection database: {error}"))?;
connection.execute_batch(&format!(
"{METADATA_SQL};{FRAGMENTS_SQL};{INDEX_SQL};INSERT INTO metadata(singleton, schema_version, last_applied_txid) VALUES(1, {SCHEMA_VERSION}, NULL);"
)).map_err(|error| format!("initialize projection schema: {error}"))?;
Ok((connection, None))
}
fn open_existing_database(path: &Path) -> Result<(Connection, Option<TxId>), OpenIssue> {
let connection = Connection::open(path).map_err(classify_sql_error)?;
configure_database(&connection).map_err(classify_sql_error)?;
let quick_check: String = connection
.query_row("PRAGMA quick_check", [], |row| row.get(0))
.map_err(classify_sql_error)?;
if quick_check != "ok" {
return Err(OpenIssue::Recoverable);
}
validate_schema(&connection)?;
let cursor = validate_rows(&connection)?;
Ok((connection, cursor))
}
fn configure_database(connection: &Connection) -> Result<(), SqlError> {
connection.pragma_update(None, "journal_mode", "WAL")?;
connection.pragma_update(None, "synchronous", "FULL")
}
fn validate_schema(connection: &Connection) -> Result<(), OpenIssue> {
let mut statement = connection.prepare(
"SELECT type, name, sql FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%' ORDER BY type, name",
).map_err(classify_sql_error)?;
let mut rows = statement.query([]).map_err(classify_sql_error)?;
let mut schema = Vec::new();
while let Some(row) = rows.next().map_err(classify_sql_error)? {
let kind: String = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
let name: String = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
let sql: String = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
schema.push((kind, name, sql));
}
let expected = vec![
(
"index".to_string(),
"fragments_actionable_state".to_string(),
INDEX_SQL.to_string(),
),
(
"table".to_string(),
"fragments".to_string(),
FRAGMENTS_SQL.to_string(),
),
(
"table".to_string(),
"metadata".to_string(),
METADATA_SQL.to_string(),
),
];
if schema != expected {
return Err(OpenIssue::Recoverable);
}
Ok(())
}
fn validate_rows(connection: &Connection) -> Result<Option<TxId>, OpenIssue> {
let mut metadata = connection
.prepare("SELECT singleton, schema_version, last_applied_txid FROM metadata")
.map_err(classify_sql_error)?;
let mut rows = metadata.query([]).map_err(classify_sql_error)?;
let row = rows
.next()
.map_err(classify_sql_error)?
.ok_or(OpenIssue::Recoverable)?;
let singleton: i64 = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
let version: i64 = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
let cursor: Option<Vec<u8>> = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
if singleton != 1
|| version != SCHEMA_VERSION
|| rows.next().map_err(classify_sql_error)?.is_some()
{
return Err(OpenIssue::Recoverable);
}
drop(rows);
drop(metadata);
let mut fragments = connection
.prepare("SELECT fragment_id, actionable_state, encoded_status FROM fragments")
.map_err(classify_sql_error)?;
let mut rows = fragments.query([]).map_err(classify_sql_error)?;
while let Some(row) = rows.next().map_err(classify_sql_error)? {
let id: Vec<u8> = row.get(0).map_err(|_| OpenIssue::Recoverable)?;
let actionable: i64 = row.get(1).map_err(|_| OpenIssue::Recoverable)?;
let encoded: Vec<u8> = row.get(2).map_err(|_| OpenIssue::Recoverable)?;
fragment_id_from_slice(&id).map_err(|_| OpenIssue::Recoverable)?;
decode_stored_status(&encoded, actionable).map_err(|_| OpenIssue::Recoverable)?;
}
cursor
.map(|bytes| {
let bytes: [u8; 12] = bytes.try_into().map_err(|_| OpenIssue::Recoverable)?;
Ok(TxId::from_bytes(bytes))
})
.transpose()
}
fn classify_sql_error(error: SqlError) -> OpenIssue {
if matches!(
error.sqlite_error_code(),
Some(ErrorCode::DatabaseCorrupt | ErrorCode::NotADatabase)
) {
OpenIssue::Recoverable
} else {
OpenIssue::Fatal(error.to_string())
}
}
fn recreate_database(path: &Path) -> Result<(Connection, Option<TxId>), String> {
remove_database_files(path)?;
create_database(path)
}
fn remove_sidecars(path: &Path) -> Result<(), String> {
for suffix in ["-wal", "-shm"] {
remove_if_present(&path_with_suffix(path, suffix))?;
}
Ok(())
}
fn remove_database_files(path: &Path) -> Result<(), String> {
remove_sidecars(path)?;
remove_if_present(path)
}
fn remove_if_present(path: &Path) -> Result<(), String> {
match fs::remove_file(path) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(format!("remove recoverable projection state: {error}")),
}
}
fn path_with_suffix(path: &Path, suffix: &str) -> PathBuf {
let mut value = OsString::from(path.as_os_str());
value.push(suffix);
PathBuf::from(value)
}
enum OpenIssue {
Recoverable,
Fatal(String),
}
mod optional_fragment_id {
use super::FragmentId;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<S: Serializer>(
value: &Option<FragmentId>,
serializer: S,
) -> Result<S::Ok, S::Error> {
value.map(FragmentId::into_bytes).serialize(serializer)
}
pub fn deserialize<'de, D: Deserializer<'de>>(
deserializer: D,
) -> Result<Option<FragmentId>, D::Error> {
Option::<[u8; 12]>::deserialize(deserializer).map(|value| value.map(FragmentId::from_bytes))
}
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_k1_audio_classification_format::{
DiscardedV2, FailedV2, LabelConfirmationV1, ProgressV1, QueueV2, TranscriptionCompleteV1,
encode_event,
};
use kcode_k1_transaction::{GENESIS_PARENT, SubsystemId, build_signed_transaction};
use kcode_speaker_v3_analysis::{
AnalysisEnvelope, FeatureVector24, GeminiCohort, LocalSpeakerLabel, OggAudioMetadata,
StructuredAnalysis, StructuredSpeaker, StructurerProvenance,
};
use std::sync::atomic::{AtomicU64, Ordering};
static NEXT_ROOT: AtomicU64 = AtomicU64::new(0);
struct TestRoot(PathBuf);
impl TestRoot {
fn new() -> Self {
let sequence = NEXT_ROOT.fetch_add(1, Ordering::Relaxed);
let path = std::env::temp_dir().join(format!(
"kcode-audio-projection-{}-{sequence}",
std::process::id()
));
fs::create_dir_all(&path).unwrap();
Self(path)
}
fn join(&self, name: &str) -> PathBuf {
self.0.join(name)
}
}
impl Drop for TestRoot {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.0);
}
}
fn txid(number: u32) -> TxId {
let mut bytes = [0; 12];
bytes[..4].copy_from_slice(&number.to_le_bytes());
TxId::from_bytes(bytes)
}
fn queue(fragment_id: FragmentId) -> AudioClassificationEventV3 {
AudioClassificationEventV3::Queue(QueueV2 {
audio_object_id: fragment_id,
})
}
fn apply_event(
projection: &Projection,
callback: TxId,
event: AudioClassificationEventV3,
) -> AppliedEvent {
projection
.apply(callback, &encode_event(&event).unwrap())
.unwrap()
}
fn progress(fragment_id: FragmentId, update: ProgressUpdateV1) -> AudioClassificationEventV3 {
AudioClassificationEventV3::Progress(ProgressV1 {
fragment_id,
update,
})
}
fn canonical_event(
ordering: &K1TxnOrdering,
event: &AudioClassificationEventV3,
) -> (TxId, Vec<u8>) {
let payload = encode_event(event).unwrap();
let bytes = build_signed_transaction(
ordering.tip().unwrap_or(GENESIS_PARENT),
1,
[1; 32],
SubsystemId::from_str(SUBSYSTEM).unwrap(),
&payload,
|_| Ok([2; 64]),
)
.unwrap();
let id = TxId::for_transaction(&bytes);
assert!(ordering.submit_txn(&bytes).is_ok());
(id, payload)
}
fn analysis(transcript: &str, speaker_numbers: &[u32]) -> ExecutedAnalysis {
let mut ogg = vec![0; 47];
ogg[..4].copy_from_slice(b"OggS");
ogg[5] = 2;
ogg[26] = 1;
ogg[27] = 19;
ogg[28..36].copy_from_slice(b"OpusHead");
ogg[36] = 1;
ogg[37] = 1;
ogg[40..44].copy_from_slice(&48_000_u32.to_le_bytes());
let speakers = speaker_numbers
.iter()
.map(|number| StructuredSpeaker {
speaker: LocalSpeakerLabel::new(*number).unwrap(),
language: "en".to_string(),
features: FeatureVector24::default(),
features_usable_for_training: true,
})
.collect();
let provenance = StructurerProvenance {
model_id: "model".to_string(),
prompt_revision: "prompt".to_string(),
};
ExecutedAnalysis {
envelope: AnalysisEnvelope {
audio: OggAudioMetadata::from_bytes(&ogg, 1_000, None).unwrap(),
analysis: StructuredAnalysis {
transcript: transcript.to_string(),
speakers,
},
gemini: GeminiCohort {
model_id: "gemini".to_string(),
transcript_prompt_revision: "transcript".to_string(),
feature_prompt_revisions: [
"one".to_string(),
"two".to_string(),
"three".to_string(),
],
feature_schema_revision: "schema".to_string(),
},
structurer: provenance.clone(),
},
label_extractor: provenance,
}
}
fn label(speaker: u32, person_id: &str) -> SpeakerLabelV1 {
SpeakerLabelV1 {
speaker: LocalSpeakerLabel::new(speaker).unwrap(),
person_id: person_id.to_string(),
}
}
#[test]
fn atomic_cursor_reopen_and_clear() {
let root = TestRoot::new();
let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
let projection_root = root.join("projection");
let fragment_id = txid(10);
let event = queue(fragment_id);
let (callback, payload) = canonical_event(&ordering, &event);
let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
assert_eq!(cursor, None);
projection.apply(callback, &payload).unwrap();
drop(projection);
let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
assert_eq!(cursor, Some(callback));
assert_eq!(projection.queued().unwrap(), vec![fragment_id]);
projection.clear().unwrap();
drop(projection);
let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
assert_eq!(cursor, None);
assert_eq!(projection.status(fragment_id).unwrap(), None);
}
#[test]
fn queue_failure_survives_reopen() {
let root = TestRoot::new();
let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
let projection_root = root.join("projection");
let fragment_id = txid(15);
let (queue_txid, queue_payload) = canonical_event(&ordering, &queue(fragment_id));
let (projection, _) = Projection::open(&projection_root, &ordering).unwrap();
projection.apply(queue_txid, &queue_payload).unwrap();
let failure = AudioClassificationEventV3::Failed(FailedV2 {
fragment_id,
stage: FragmentStageV1::Queue,
llm_job_sequence: None,
error: "queue failed".to_string(),
});
let (failure_txid, failure_payload) = canonical_event(&ordering, &failure);
projection.apply(failure_txid, &failure_payload).unwrap();
drop(projection);
let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
let status = projection.status(fragment_id).unwrap().unwrap();
assert_eq!(cursor, Some(failure_txid));
assert_eq!(status.state, OverallState::Failed);
assert_eq!(status.queue.state, StageState::Failed);
}
#[test]
fn failure_retry_completion_and_exact_confirmation() {
let root = TestRoot::new();
let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
let (projection, _) = Projection::open(&root.join("projection"), &ordering).unwrap();
let fragment_id = txid(20);
apply_event(&projection, txid(21), queue(fragment_id));
apply_event(
&projection,
txid(22),
progress(
fragment_id,
ProgressUpdateV1::LlmJobStarted {
sequence: 7,
stage: FragmentStageV1::Transcript,
name: "transcript".to_string(),
},
),
);
assert_eq!(
projection.running().unwrap()[0].stage,
FragmentStageV1::Transcript
);
apply_event(
&projection,
txid(23),
progress(
fragment_id,
ProgressUpdateV1::LlmJobFailed {
sequence: 7,
error: "provider".to_string(),
},
),
);
apply_event(
&projection,
txid(24),
AudioClassificationEventV3::Failed(FailedV2 {
fragment_id,
stage: FragmentStageV1::Transcript,
llm_job_sequence: Some(7),
error: "terminal".to_string(),
}),
);
let failed = projection.status(fragment_id).unwrap().unwrap();
assert_eq!(failed.state, OverallState::Failed);
assert_eq!(failed.jobs[0].state, LlmJobState::Failed);
apply_event(
&projection,
txid(25),
progress(
fragment_id,
ProgressUpdateV1::LlmJobStarted {
sequence: 1,
stage: FragmentStageV1::Transcript,
name: "retry".to_string(),
},
),
);
let transcript = "[high] Speaker 2: hello\n[medium] Speaker 7 [overlap]: hi\n[low] Speaker 1: unchanged\nplain Speaker 2: unchanged\n";
apply_event(
&projection,
txid(26),
AudioClassificationEventV3::TranscriptionComplete(TranscriptionCompleteV1 {
fragment_id,
analysis: analysis(transcript, &[2, 7]),
}),
);
let labels = vec![label(2, "alice"), label(7, "bob")];
assert_eq!(
projection.validate_labels(fragment_id, &labels).unwrap(),
txid(26)
);
let applied = apply_event(
&projection,
txid(27),
AudioClassificationEventV3::LabelConfirmation(LabelConfirmationV1 {
fragment_id,
interim_txid: txid(26),
speakers: labels.clone(),
}),
);
assert_eq!(applied.effect, ProjectionEffect::LabelsCommitted);
let status = projection.status(fragment_id).unwrap().unwrap();
assert_eq!(status.confirmed_labels, labels);
assert_eq!(
status.final_transcript.as_deref(),
Some(
"[high] alice: hello\n[medium] bob [overlap]: hi\n[low] Speaker 1: unchanged\nplain Speaker 2: unchanged\n"
)
);
}
#[test]
fn repeated_discard_suppresses_late_events() {
let root = TestRoot::new();
let ordering = K1TxnOrdering::open(&root.join("ordering")).unwrap();
let (projection, _) = Projection::open(&root.join("projection"), &ordering).unwrap();
let fragment_id = txid(30);
apply_event(&projection, txid(31), queue(fragment_id));
let discarded = || AudioClassificationEventV3::Discarded(DiscardedV2 { fragment_id });
assert_eq!(
apply_event(&projection, txid(32), discarded()).effect,
ProjectionEffect::Abort
);
let status = projection.status(fragment_id).unwrap().unwrap();
assert_eq!(
apply_event(&projection, txid(33), discarded()).effect,
ProjectionEffect::Abort
);
apply_event(
&projection,
txid(34),
AudioClassificationEventV3::TranscriptionComplete(TranscriptionCompleteV1 {
fragment_id,
analysis: analysis("[high] Speaker 2: late", &[2]),
}),
);
apply_event(
&projection,
txid(35),
AudioClassificationEventV3::Failed(FailedV2 {
fragment_id,
stage: FragmentStageV1::Transcript,
llm_job_sequence: None,
error: "late".to_string(),
}),
);
assert_eq!(projection.status(fragment_id).unwrap().unwrap(), status);
}
#[test]
fn first_five_thousand_errors_are_retained() {
let mut status = initial_status();
for index in 0..=MAX_ERRORS {
append_error(&mut status, index.to_string());
}
assert_eq!(status.errors.len(), MAX_ERRORS);
assert_eq!(status.errors.first().map(String::as_str), Some("0"));
assert_eq!(status.errors.last().map(String::as_str), Some("4999"));
assert!(status.errors_truncated);
}
#[test]
fn malformed_and_noncanonical_databases_rebuild() {
let malformed = TestRoot::new();
let ordering = K1TxnOrdering::open(&malformed.join("ordering")).unwrap();
let projection_root = malformed.join("projection");
fs::create_dir_all(&projection_root).unwrap();
fs::write(projection_root.join(DATABASE_NAME), b"malformed").unwrap();
assert_eq!(
Projection::open(&projection_root, &ordering).unwrap().1,
None
);
let noncanonical = TestRoot::new();
let ordering = K1TxnOrdering::open(&noncanonical.join("ordering")).unwrap();
let projection_root = noncanonical.join("projection");
let (projection, _) = Projection::open(&projection_root, &ordering).unwrap();
let fragment_id = txid(40);
apply_event(&projection, txid(41), queue(fragment_id));
drop(projection);
let (projection, cursor) = Projection::open(&projection_root, &ordering).unwrap();
assert_eq!(cursor, None);
assert_eq!(projection.status(fragment_id).unwrap(), None);
}
}