kcode_speaker_store/

lib.rs

use kcode_speaker_types::{FeatureVector, Key, LabeledSample, ObjectId, RecordingKind, SegmentRef};
use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
    collections::{BTreeMap, BTreeSet},
    error::Error,
    fmt,
    path::Path,
    sync::Mutex,
};

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct EventEnvelope {
    pub event_id: Key,
    pub event: Event,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Event {
    RegisterSegment(SegmentRegistration),
    RecordAttempt(AttemptRecord),
    SelectAttempt(AttemptSelection),
    SetSampleState(SampleStateChange),
    PutArtifact(ArtifactRecord),
    SetActiveModel(ActiveModelChange),
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SegmentRegistration {
    pub segment: SegmentRef,
    pub group_id: Key,
    pub recording_kind: RecordingKind,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptRecord {
    pub attempt_id: Key,
    pub clip_object: ObjectId,
    pub extraction_key: Key,
    pub provider_result_object: Option<ObjectId>,
    pub outcome: StoredAttemptOutcome,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum StoredAttemptOutcome {
    Scored {
        recording_quality: u8,
        speakers: Vec<StoredSpeaker>,
    },
    Unscorable {
        reason: String,
    },
    InvalidResponse {
        error: String,
    },
    ProviderFailure {
        code: Key,
        detail: String,
    },
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct StoredSpeaker {
    pub speaker_ordinal: u16,
    pub primary_language: Key,
    pub closest_dialect: String,
    pub usable_speech_ms: u32,
    pub features: FeatureVector,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptSelection {
    pub clip_object: ObjectId,
    pub extraction_key: Key,
    pub attempt_id: Key,
    pub reason: String,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SampleStateChange {
    pub sample_id: Key,
    pub state: SampleState,
    pub reason: String,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum SampleState {
    Unlabeled,
    Confirmed { speaker_id: Key },
    Retracted,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ArtifactRecord {
    pub artifact_id: Key,
    pub cohort_id: Key,
    pub kind: ArtifactKind,
    pub sha256: [u8; 32],
    pub bytes: Vec<u8>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum ArtifactKind {
    Model,
    Evaluation,
    ReplayManifest,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ActiveModelChange {
    pub cohort_id: Key,
    pub model_artifact_id: Option<Key>,
    pub reason: String,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Request {
    Commit(Vec<EventEnvelope>),
    Read(Query),
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Query {
    Attempt {
        attempt_id: Key,
    },
    Attempts {
        extraction_key: Option<Key>,
        source_object: Option<ObjectId>,
    },
    TrainingRows {
        cohort_id: Key,
    },
    RepeatabilityRows {
        cohort_id: Key,
    },
    Artifact {
        artifact_id: Key,
    },
    ActiveModel {
        cohort_id: Key,
    },
    Inventory {
        extraction_key: Option<Key>,
    },
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct Response {
    pub revision: u64,
    pub result: ResponseKind,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum ResponseKind {
    Commit { applied: u64 },
    Attempt(Option<AttemptSnapshot>),
    Attempts(Vec<AttemptSnapshot>),
    TrainingRows(Vec<LabeledSample>),
    RepeatabilityRows(Vec<LabeledSample>),
    Artifact(Option<ArtifactRecord>),
    ActiveModel(Option<Key>),
    Inventory(InventorySnapshot),
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptSnapshot {
    pub attempt: AttemptRecord,
    pub sample_ids: Vec<Key>,
    pub selected: bool,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SelectionSnapshot {
    pub selection: AttemptSelection,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SampleStateSnapshot {
    pub change: SampleStateChange,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ArtifactMetadata {
    pub artifact_id: Key,
    pub cohort_id: Key,
    pub kind: ArtifactKind,
    pub sha256: [u8; 32],
    pub byte_len: u64,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ActiveModelSnapshot {
    pub change: ActiveModelChange,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct MissingSegmentOrdinals {
    pub source_object: ObjectId,
    pub ordinals: Vec<u16>,
}

#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
pub struct OutcomeCounts {
    pub scored: u64,
    pub unscorable: u64,
    pub invalid_response: u64,
    pub provider_failure: u64,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct InventorySnapshot {
    pub registrations: Vec<SegmentRegistration>,
    pub attempts: Vec<AttemptSnapshot>,
    pub selections: Vec<SelectionSnapshot>,
    pub sample_states: Vec<SampleStateSnapshot>,
    pub artifacts: Vec<ArtifactMetadata>,
    pub active_models: Vec<ActiveModelSnapshot>,
    pub missing_segment_ordinals: Vec<MissingSegmentOrdinals>,
    pub outcome_counts: OutcomeCounts,
}

#[derive(Debug)]
pub enum StoreError {
    Validation(String),
    Conflict(String),
    SchemaVersion(i64),
    Storage(String),
}

impl fmt::Display for StoreError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Validation(v) => write!(f, "validation: {v}"),
            Self::Conflict(v) => write!(f, "conflict: {v}"),
            Self::SchemaVersion(v) => write!(f, "unsupported schema version {v}"),
            Self::Storage(v) => write!(f, "storage: {v}"),
        }
    }
}

impl Error for StoreError {}

impl From<rusqlite::Error> for StoreError {
    fn from(value: rusqlite::Error) -> Self {
        Self::Storage(value.to_string())
    }
}

impl From<serde_json::Error> for StoreError {
    fn from(value: serde_json::Error) -> Self {
        Self::Storage(value.to_string())
    }
}

pub struct Store {
    connection: Mutex<Connection>,
}

pub fn open(path: impl AsRef<Path>) -> Result<Store, StoreError> {
    let connection = Connection::open(path)?;
    let version: i64 = connection.query_row("PRAGMA user_version", [], |row| row.get(0))?;
    if version == 0 {
        let objects: i64 = connection.query_row(
            "SELECT count(*) FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%'",
            [],
            |row| row.get(0),
        )?;
        if objects != 0 {
            return Err(StoreError::SchemaVersion(0));
        }
        connection.execute_batch(
            "CREATE TABLE events(
                 event_id TEXT PRIMARY KEY NOT NULL,
                 event_json BLOB NOT NULL
             );
             PRAGMA user_version = 1;",
        )?;
    } else if version != 1 {
        return Err(StoreError::SchemaVersion(version));
    }
    connection.prepare("SELECT event_id,event_json FROM events LIMIT 0")?;
    Ok(Store {
        connection: Mutex::new(connection),
    })
}

impl Store {
    pub fn execute(&self, request: Request) -> Result<Response, StoreError> {
        let mut connection = self
            .connection
            .lock()
            .map_err(|_| StoreError::Storage("connection lock poisoned".into()))?;
        let (mut event_ids, mut projection) = load(&connection)?;
        match request {
            Request::Read(query) => Ok(Response {
                revision: event_ids.len() as u64,
                result: projection.read(query)?,
            }),
            Request::Commit(events) => {
                let mut additions = Vec::new();
                for envelope in events {
                    let id = envelope.event_id.as_ref().to_owned();
                    let json = serde_json::to_vec(&envelope.event)?;
                    if let Some(old) = event_ids.get(&id) {
                        if old != &json {
                            return Err(StoreError::Conflict(format!(
                                "event ID {id} has different content"
                            )));
                        }
                        continue;
                    }
                    projection.apply(&envelope.event)?;
                    event_ids.insert(id.clone(), json.clone());
                    additions.push((id, json));
                }
                let transaction = connection.transaction()?;
                for (id, json) in &additions {
                    transaction.execute(
                        "INSERT INTO events(event_id,event_json) VALUES(?1,?2)",
                        params![id, json],
                    )?;
                }
                transaction.commit()?;
                Ok(Response {
                    revision: event_ids.len() as u64,
                    result: ResponseKind::Commit {
                        applied: additions.len() as u64,
                    },
                })
            }
        }
    }
}

#[derive(Default)]
struct Projection {
    registrations: BTreeMap<ObjectId, SegmentRegistration>,
    attempts: BTreeMap<Key, AttemptRecord>,
    selections: BTreeMap<(ObjectId, Key), AttemptSelection>,
    states: BTreeMap<Key, SampleStateChange>,
    artifacts: BTreeMap<Key, ArtifactRecord>,
    active_models: BTreeMap<Key, ActiveModelChange>,
}

fn load(connection: &Connection) -> Result<(BTreeMap<String, Vec<u8>>, Projection), StoreError> {
    let mut statement =
        connection.prepare("SELECT event_id,event_json FROM events ORDER BY rowid")?;
    let rows = statement.query_map([], |row| {
        Ok((row.get::<_, String>(0)?, row.get::<_, Vec<u8>>(1)?))
    })?;
    let mut ids = BTreeMap::new();
    let mut projection = Projection::default();
    for row in rows {
        let (id, json) = row?;
        Key::parse(&id)
            .map_err(|e| StoreError::Storage(format!("invalid stored event ID: {e}")))?;
        let event: Event = serde_json::from_slice(&json)?;
        projection.apply(&event)?;
        ids.insert(id, json);
    }
    Ok((ids, projection))
}

impl Projection {
    fn apply(&mut self, event: &Event) -> Result<(), StoreError> {
        match event {
            Event::RegisterSegment(registration) => self.register(registration),
            Event::RecordAttempt(attempt) => self.record_attempt(attempt),
            Event::SelectAttempt(selection) => self.select(selection),
            Event::SetSampleState(change) => self.set_sample_state(change),
            Event::PutArtifact(artifact) => self.put_artifact(artifact),
            Event::SetActiveModel(change) => self.set_active_model(change),
        }
    }

    fn register(&mut self, registration: &SegmentRegistration) -> Result<(), StoreError> {
        registration
            .segment
            .validate()
            .map_err(|e| StoreError::Validation(e.to_string()))?;
        let segment = &registration.segment;
        let duration = segment.end_ms - segment.start_ms;
        if duration == 0 || duration >= 240_000 {
            return validation("segment duration must be 1..239999 ms");
        }
        if segment.ordinal == 0 && segment.start_ms != 0 {
            return validation("source ordinal zero must begin at zero");
        }
        if self.registrations.contains_key(&segment.clip_object) {
            return conflict("clip object is already registered");
        }
        for old in self
            .registrations
            .values()
            .filter(|old| old.segment.source_object == segment.source_object)
        {
            if old.segment.policy != segment.policy
                || old.segment.segment_count != segment.segment_count
                || old.group_id != registration.group_id
                || old.recording_kind != registration.recording_kind
            {
                return conflict("source plan, group, or recording kind differs");
            }
            if old.segment.ordinal == segment.ordinal {
                return conflict("source ordinal is already registered");
            }
            let old_duration = old.segment.end_ms - old.segment.start_ms;
            if old_duration.abs_diff(duration) > 1 {
                return validation("source segment lengths differ by more than one millisecond");
            }
            let ordered = if old.segment.ordinal < segment.ordinal {
                old.segment.start_ms < segment.start_ms && old.segment.end_ms < segment.end_ms
            } else {
                segment.start_ms < old.segment.start_ms && segment.end_ms < old.segment.end_ms
            };
            if !ordered {
                return validation("source segment ranges do not follow ordinal order");
            }
            if old.segment.ordinal.checked_add(1) == Some(segment.ordinal)
                && old.segment.end_ms.checked_sub(segment.start_ms) != Some(5_000)
            {
                return validation("adjacent source segments must overlap by 5000 ms");
            }
            if segment.ordinal.checked_add(1) == Some(old.segment.ordinal)
                && segment.end_ms.checked_sub(old.segment.start_ms) != Some(5_000)
            {
                return validation("adjacent source segments must overlap by 5000 ms");
            }
        }
        self.registrations
            .insert(segment.clip_object.clone(), registration.clone());
        Ok(())
    }

    fn record_attempt(&mut self, attempt: &AttemptRecord) -> Result<(), StoreError> {
        if self.attempts.contains_key(&attempt.attempt_id) {
            return conflict("attempt ID is already recorded");
        }
        let registration = self
            .registrations
            .get(&attempt.clip_object)
            .ok_or_else(|| StoreError::Conflict("clip must be registered first".into()))?;
        if !matches!(
            attempt.outcome,
            StoredAttemptOutcome::ProviderFailure { .. }
        ) && attempt.provider_result_object.is_none()
        {
            return validation("response-bearing attempt needs a result object");
        }
        match &attempt.outcome {
            StoredAttemptOutcome::Scored {
                recording_quality,
                speakers,
            } => {
                if *recording_quality > 100 || speakers.is_empty() {
                    return validation("scored attempt needs quality <=100 and speakers");
                }
                let duration = registration.segment.end_ms - registration.segment.start_ms;
                for (index, speaker) in speakers.iter().enumerate() {
                    if usize::from(speaker.speaker_ordinal) != index {
                        return validation("speaker ordinals must be contiguous from zero");
                    }
                    bounded(&speaker.closest_dialect, 128, "dialect")?;
                    if speaker.usable_speech_ms == 0
                        || u64::from(speaker.usable_speech_ms) > duration
                    {
                        return validation("usable speech is outside segment duration");
                    }
                    sample_id(&attempt.attempt_id, speaker.speaker_ordinal)?;
                }
            }
            StoredAttemptOutcome::Unscorable { reason } => bounded(reason, 240, "reason")?,
            StoredAttemptOutcome::InvalidResponse { error } => bounded(error, 240, "error")?,
            StoredAttemptOutcome::ProviderFailure { detail, .. } => bounded(detail, 240, "detail")?,
        }
        self.attempts
            .insert(attempt.attempt_id.clone(), attempt.clone());
        Ok(())
    }

    fn select(&mut self, selection: &AttemptSelection) -> Result<(), StoreError> {
        bounded(&selection.reason, 240, "selection reason")?;
        let attempt = self
            .attempts
            .get(&selection.attempt_id)
            .ok_or_else(|| StoreError::Conflict("selected attempt does not exist".into()))?;
        if attempt.clip_object != selection.clip_object
            || attempt.extraction_key != selection.extraction_key
            || !matches!(attempt.outcome, StoredAttemptOutcome::Scored { .. })
        {
            return conflict("selection does not match a scored clip attempt");
        }
        self.selections.insert(
            (
                selection.clip_object.clone(),
                selection.extraction_key.clone(),
            ),
            selection.clone(),
        );
        Ok(())
    }

    fn set_sample_state(&mut self, change: &SampleStateChange) -> Result<(), StoreError> {
        bounded(&change.reason, 240, "sample-state reason")?;
        if self.sample_owner(&change.sample_id).is_none() {
            return conflict("sample ID does not exist");
        }
        self.states.insert(change.sample_id.clone(), change.clone());
        Ok(())
    }

    fn put_artifact(&mut self, artifact: &ArtifactRecord) -> Result<(), StoreError> {
        let actual: [u8; 32] = Sha256::digest(&artifact.bytes).into();
        if actual != artifact.sha256 {
            return validation("artifact SHA-256 does not match bytes");
        }
        if let Some(old) = self.artifacts.get(&artifact.artifact_id) {
            if old != artifact {
                return conflict("artifact ID is immutable");
            }
            return Ok(());
        }
        self.artifacts
            .insert(artifact.artifact_id.clone(), artifact.clone());
        Ok(())
    }

    fn set_active_model(&mut self, change: &ActiveModelChange) -> Result<(), StoreError> {
        bounded(&change.reason, 240, "activation reason")?;
        if let Some(id) = &change.model_artifact_id {
            let artifact = self
                .artifacts
                .get(id)
                .ok_or_else(|| StoreError::Conflict("model artifact does not exist".into()))?;
            if artifact.kind != ArtifactKind::Model || artifact.cohort_id != change.cohort_id {
                return conflict("active artifact must be a Model in the same cohort");
            }
        }
        self.active_models
            .insert(change.cohort_id.clone(), change.clone());
        Ok(())
    }

    fn sample_owner(&self, wanted: &Key) -> Option<(&AttemptRecord, &StoredSpeaker, u8)> {
        for attempt in self.attempts.values() {
            if let StoredAttemptOutcome::Scored {
                recording_quality,
                speakers,
            } = &attempt.outcome
            {
                for speaker in speakers {
                    if sample_id(&attempt.attempt_id, speaker.speaker_ordinal)
                        .ok()
                        .as_ref()
                        == Some(wanted)
                    {
                        return Some((attempt, speaker, *recording_quality));
                    }
                }
            }
        }
        None
    }

    fn snapshot(&self, attempt: &AttemptRecord) -> Result<AttemptSnapshot, StoreError> {
        let sample_ids = match &attempt.outcome {
            StoredAttemptOutcome::Scored { speakers, .. } => speakers
                .iter()
                .map(|speaker| sample_id(&attempt.attempt_id, speaker.speaker_ordinal))
                .collect::<Result<_, _>>()?,
            _ => Vec::new(),
        };
        let selected = self
            .selections
            .get(&(attempt.clip_object.clone(), attempt.extraction_key.clone()))
            .is_some_and(|selection| selection.attempt_id == attempt.attempt_id);
        Ok(AttemptSnapshot {
            attempt: attempt.clone(),
            sample_ids,
            selected,
        })
    }

    fn rows(&self, cohort_id: &Key, selected_only: bool) -> Result<Vec<LabeledSample>, StoreError> {
        let mut rows = Vec::new();
        for attempt in self
            .attempts
            .values()
            .filter(|attempt| &attempt.extraction_key == cohort_id)
        {
            if selected_only
                && !self
                    .selections
                    .get(&(attempt.clip_object.clone(), attempt.extraction_key.clone()))
                    .is_some_and(|selection| selection.attempt_id == attempt.attempt_id)
            {
                continue;
            }
            let StoredAttemptOutcome::Scored {
                recording_quality,
                speakers,
            } = &attempt.outcome
            else {
                continue;
            };
            let registration = self
                .registrations
                .get(&attempt.clip_object)
                .ok_or_else(|| StoreError::Storage("attempt lost its registration".into()))?;
            for speaker in speakers {
                let id = sample_id(&attempt.attempt_id, speaker.speaker_ordinal)?;
                let Some(SampleStateChange {
                    state: SampleState::Confirmed { speaker_id },
                    ..
                }) = self.states.get(&id)
                else {
                    continue;
                };
                rows.push(LabeledSample {
                    sample_id: id,
                    attempt_id: attempt.attempt_id.clone(),
                    speaker_id: speaker_id.clone(),
                    cohort_id: attempt.extraction_key.clone(),
                    group_id: registration.group_id.clone(),
                    clip_object: attempt.clip_object.clone(),
                    primary_language: speaker.primary_language.clone(),
                    recording_kind: registration.recording_kind,
                    usable_speech_ms: speaker.usable_speech_ms,
                    recording_quality: *recording_quality,
                    features: speaker.features.clone(),
                });
            }
        }
        Ok(rows)
    }

    fn read(&self, query: Query) -> Result<ResponseKind, StoreError> {
        match query {
            Query::Attempt { attempt_id } => Ok(ResponseKind::Attempt(
                self.attempts
                    .get(&attempt_id)
                    .map(|attempt| self.snapshot(attempt))
                    .transpose()?,
            )),
            Query::Attempts {
                extraction_key,
                source_object,
            } => Ok(ResponseKind::Attempts(
                self.attempts
                    .values()
                    .filter(|attempt| {
                        extraction_key
                            .as_ref()
                            .is_none_or(|key| key == &attempt.extraction_key)
                            && source_object.as_ref().is_none_or(|source| {
                                self.registrations.get(&attempt.clip_object).is_some_and(
                                    |registration| &registration.segment.source_object == source,
                                )
                            })
                    })
                    .map(|attempt| self.snapshot(attempt))
                    .collect::<Result<_, _>>()?,
            )),
            Query::TrainingRows { cohort_id } => {
                Ok(ResponseKind::TrainingRows(self.rows(&cohort_id, true)?))
            }
            Query::RepeatabilityRows { cohort_id } => Ok(ResponseKind::RepeatabilityRows(
                self.rows(&cohort_id, false)?,
            )),
            Query::Artifact { artifact_id } => Ok(ResponseKind::Artifact(
                self.artifacts.get(&artifact_id).cloned(),
            )),
            Query::ActiveModel { cohort_id } => Ok(ResponseKind::ActiveModel(
                self.active_models
                    .get(&cohort_id)
                    .and_then(|change| change.model_artifact_id.clone()),
            )),
            Query::Inventory { extraction_key } => {
                Ok(ResponseKind::Inventory(self.inventory(extraction_key)?))
            }
        }
    }

    fn inventory(&self, filter: Option<Key>) -> Result<InventorySnapshot, StoreError> {
        let included = |attempt: &&AttemptRecord| {
            filter
                .as_ref()
                .is_none_or(|key| key == &attempt.extraction_key)
        };
        let attempts = self
            .attempts
            .values()
            .filter(included)
            .map(|attempt| self.snapshot(attempt))
            .collect::<Result<Vec<_>, _>>()?;
        let included_samples: BTreeSet<Key> = attempts
            .iter()
            .flat_map(|attempt| attempt.sample_ids.iter().cloned())
            .collect();
        let selections = self
            .selections
            .values()
            .filter(|selection| {
                filter
                    .as_ref()
                    .is_none_or(|key| key == &selection.extraction_key)
            })
            .cloned()
            .map(|selection| SelectionSnapshot { selection })
            .collect();
        let sample_states = self
            .states
            .values()
            .filter(|change| filter.is_none() || included_samples.contains(&change.sample_id))
            .cloned()
            .map(|change| SampleStateSnapshot { change })
            .collect();
        let mut registrations: Vec<_> = self.registrations.values().cloned().collect();
        registrations.sort_by(|a, b| {
            (&a.segment.source_object, a.segment.ordinal)
                .cmp(&(&b.segment.source_object, b.segment.ordinal))
        });
        let mut sources: BTreeMap<ObjectId, (u16, BTreeSet<u16>)> = BTreeMap::new();
        for registration in &registrations {
            let entry = sources
                .entry(registration.segment.source_object.clone())
                .or_insert((registration.segment.segment_count, BTreeSet::new()));
            entry.1.insert(registration.segment.ordinal);
        }
        let missing_segment_ordinals = sources
            .into_iter()
            .filter_map(|(source_object, (count, present))| {
                let ordinals: Vec<_> = (0..count)
                    .filter(|ordinal| !present.contains(ordinal))
                    .collect();
                (!ordinals.is_empty()).then_some(MissingSegmentOrdinals {
                    source_object,
                    ordinals,
                })
            })
            .collect();
        let mut outcome_counts = OutcomeCounts::default();
        for attempt in self.attempts.values().filter(included) {
            match attempt.outcome {
                StoredAttemptOutcome::Scored { .. } => outcome_counts.scored += 1,
                StoredAttemptOutcome::Unscorable { .. } => outcome_counts.unscorable += 1,
                StoredAttemptOutcome::InvalidResponse { .. } => {
                    outcome_counts.invalid_response += 1
                }
                StoredAttemptOutcome::ProviderFailure { .. } => {
                    outcome_counts.provider_failure += 1
                }
            }
        }
        Ok(InventorySnapshot {
            registrations,
            attempts,
            selections,
            sample_states,
            artifacts: self
                .artifacts
                .values()
                .map(|artifact| ArtifactMetadata {
                    artifact_id: artifact.artifact_id.clone(),
                    cohort_id: artifact.cohort_id.clone(),
                    kind: artifact.kind,
                    sha256: artifact.sha256,
                    byte_len: artifact.bytes.len() as u64,
                })
                .collect(),
            active_models: self
                .active_models
                .values()
                .cloned()
                .map(|change| ActiveModelSnapshot { change })
                .collect(),
            missing_segment_ordinals,
            outcome_counts,
        })
    }
}

fn sample_id(attempt_id: &Key, ordinal: u16) -> Result<Key, StoreError> {
    let mut digest = Sha256::new();
    digest.update(attempt_id.as_ref().as_bytes());
    digest.update([0]);
    digest.update(ordinal.to_be_bytes());
    let hex = digest
        .finalize()
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect::<String>();
    Key::parse(&format!("sample:sha256:{hex}"))
        .map_err(|e| StoreError::Storage(format!("generated invalid sample ID: {e}")))
}

fn bounded(value: &str, maximum: usize, name: &str) -> Result<(), StoreError> {
    if value.trim().is_empty() || value.len() > maximum {
        return validation(format!(
            "{name} must be nonblank and at most {maximum} bytes"
        ));
    }
    Ok(())
}

fn validation<T>(message: impl Into<String>) -> Result<T, StoreError> {
    Err(StoreError::Validation(message.into()))
}

fn conflict<T>(message: impl Into<String>) -> Result<T, StoreError> {
    Err(StoreError::Conflict(message.into()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{
        fs,
        path::PathBuf,
        sync::{
            Arc,
            atomic::{AtomicU64, Ordering},
        },
        thread,
    };

    static NEXT_PATH: AtomicU64 = AtomicU64::new(0);

    fn path() -> PathBuf {
        std::env::temp_dir().join(format!(
            "kcode-speaker-store-{}-{}.db",
            std::process::id(),
            NEXT_PATH.fetch_add(1, Ordering::Relaxed)
        ))
    }

    fn key(value: &str) -> Key {
        Key::parse(value).unwrap()
    }

    fn object(value: &str) -> ObjectId {
        ObjectId::parse(value).unwrap()
    }

    fn features(seed: u8) -> FeatureVector {
        FeatureVector::new([seed; 35]).unwrap()
    }

    fn registration(source: &str, clip: &str, ordinal: u16, count: u16) -> SegmentRegistration {
        SegmentRegistration {
            segment: SegmentRef {
                source_object: object(source),
                clip_object: object(clip),
                ordinal,
                segment_count: count,
                start_ms: u64::from(ordinal) * 95_000,
                end_ms: u64::from(ordinal) * 95_000 + 100_000,
                policy: key("speaker-segments/1"),
            },
            group_id: key("group/1"),
            recording_kind: RecordingKind::VoiceNote,
        }
    }

    fn scored(id: &str, clip: &str, extraction: &str, speakers: u16) -> AttemptRecord {
        AttemptRecord {
            attempt_id: key(id),
            clip_object: object(clip),
            extraction_key: key(extraction),
            provider_result_object: Some(object("RESULT01")),
            outcome: StoredAttemptOutcome::Scored {
                recording_quality: 80,
                speakers: (0..speakers)
                    .map(|speaker_ordinal| StoredSpeaker {
                        speaker_ordinal,
                        primary_language: key("en"),
                        closest_dialect: "General American".into(),
                        usable_speech_ms: 20_000,
                        features: features(50),
                    })
                    .collect(),
            },
        }
    }

    fn failure(id: &str, clip: &str, extraction: &str) -> AttemptRecord {
        AttemptRecord {
            attempt_id: key(id),
            clip_object: object(clip),
            extraction_key: key(extraction),
            provider_result_object: None,
            outcome: StoredAttemptOutcome::ProviderFailure {
                code: key("timeout"),
                detail: "provider timed out".into(),
            },
        }
    }

    fn event(id: &str, event: Event) -> EventEnvelope {
        EventEnvelope {
            event_id: key(id),
            event,
        }
    }

    fn commit(store: &Store, events: Vec<EventEnvelope>) -> Response {
        store.execute(Request::Commit(events)).unwrap()
    }

    fn attempt(store: &Store, id: &str) -> AttemptSnapshot {
        let response = store
            .execute(Request::Read(Query::Attempt {
                attempt_id: key(id),
            }))
            .unwrap();
        let ResponseKind::Attempt(Some(attempt)) = response.result else {
            panic!("missing attempt")
        };
        attempt
    }

    fn selection(id: &str, clip: &str, extraction: &str) -> AttemptSelection {
        AttemptSelection {
            clip_object: object(clip),
            extraction_key: key(extraction),
            attempt_id: key(id),
            reason: "reviewed".into(),
        }
    }

    fn state(sample_id: Key, state: SampleState) -> SampleStateChange {
        SampleStateChange {
            sample_id,
            state,
            reason: "human review".into(),
        }
    }

    #[test]
    fn fresh_database_restarts() {
        let path = path();
        {
            let store = open(&path).unwrap();
            let response = commit(
                &store,
                vec![event(
                    "event/1",
                    Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
                )],
            );
            assert_eq!(response.revision, 1);
        }
        let store = open(&path).unwrap();
        let response = store
            .execute(Request::Read(Query::Inventory {
                extraction_key: None,
            }))
            .unwrap();
        assert_eq!(response.revision, 1);
        let ResponseKind::Inventory(inventory) = response.result else {
            panic!()
        };
        assert_eq!(inventory.registrations.len(), 1);
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn commit_is_atomic_and_event_ids_are_idempotent() {
        let path = path();
        let store = open(&path).unwrap();
        let registration_event = event(
            "event/register",
            Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
        );
        let bad = event(
            "event/bad",
            Event::RecordAttempt(scored("attempt/1", "UNKNOWN1", "extract/1", 1)),
        );
        assert!(
            store
                .execute(Request::Commit(vec![registration_event.clone(), bad]))
                .is_err()
        );
        let response = store
            .execute(Request::Read(Query::Inventory {
                extraction_key: None,
            }))
            .unwrap();
        assert_eq!(response.revision, 0);

        assert!(matches!(
            commit(&store, vec![registration_event.clone()]).result,
            ResponseKind::Commit { applied: 1 }
        ));
        assert!(matches!(
            commit(
                &store,
                vec![registration_event.clone(), registration_event.clone()]
            )
            .result,
            ResponseKind::Commit { applied: 0 }
        ));
        let mut changed = registration_event;
        changed.event = Event::RegisterSegment(registration("SOURCE02", "CLIP0002", 0, 1));
        assert!(matches!(
            store.execute(Request::Commit(vec![changed])),
            Err(StoreError::Conflict(_))
        ));
        assert_eq!(
            store
                .execute(Request::Read(Query::Inventory {
                    extraction_key: None
                }))
                .unwrap()
                .revision,
            1
        );
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn registration_order_and_source_plan_conflicts_fail() {
        let path = path();
        let store = open(&path).unwrap();
        assert!(matches!(
            store.execute(Request::Commit(vec![event(
                "attempt/early",
                Event::RecordAttempt(scored("attempt/1", "CLIP0001", "extract/1", 1))
            )])),
            Err(StoreError::Conflict(_))
        ));
        commit(
            &store,
            vec![event(
                "register/0",
                Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 2)),
            )],
        );
        let mut wrong_count = registration("SOURCE01", "CLIP0002", 1, 3);
        assert!(
            store
                .execute(Request::Commit(vec![event(
                    "register/wrong-count",
                    Event::RegisterSegment(wrong_count.clone())
                )]))
                .is_err()
        );
        wrong_count.segment.segment_count = 2;
        wrong_count.segment.policy = key("different-plan");
        assert!(
            store
                .execute(Request::Commit(vec![event(
                    "register/wrong-plan",
                    Event::RegisterSegment(wrong_count)
                )]))
                .is_err()
        );
        let duplicate_ordinal = registration("SOURCE01", "CLIP0003", 0, 2);
        assert!(
            store
                .execute(Request::Commit(vec![event(
                    "register/duplicate",
                    Event::RegisterSegment(duplicate_ordinal)
                )]))
                .is_err()
        );
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn retries_replacement_ids_and_label_projections() {
        let path = path();
        let store = open(&path).unwrap();
        commit(
            &store,
            vec![
                event(
                    "register/1",
                    Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
                ),
                event(
                    "attempt/1",
                    Event::RecordAttempt(scored("try/one", "CLIP0001", "extract/1", 2)),
                ),
                event(
                    "attempt/2",
                    Event::RecordAttempt(scored("try/two", "CLIP0001", "extract/1", 1)),
                ),
                event(
                    "select/1",
                    Event::SelectAttempt(selection("try/one", "CLIP0001", "extract/1")),
                ),
            ],
        );
        let first = attempt(&store, "try/one");
        assert_eq!(first.sample_ids.len(), 2);
        assert_ne!(first.sample_ids[0], first.sample_ids[1]);
        let mut digest = Sha256::new();
        digest.update(b"try/one");
        digest.update([0]);
        digest.update(0u16.to_be_bytes());
        let expected = key(&format!("sample:sha256:{:x}", digest.finalize()));
        assert_eq!(first.sample_ids[0], expected);

        let second = attempt(&store, "try/two");
        commit(
            &store,
            vec![
                event(
                    "state/1",
                    Event::SetSampleState(state(
                        first.sample_ids[0].clone(),
                        SampleState::Confirmed {
                            speaker_id: key("speaker/alice"),
                        },
                    )),
                ),
                event(
                    "state/2",
                    Event::SetSampleState(state(
                        second.sample_ids[0].clone(),
                        SampleState::Confirmed {
                            speaker_id: key("speaker/alice"),
                        },
                    )),
                ),
            ],
        );
        let training = store
            .execute(Request::Read(Query::TrainingRows {
                cohort_id: key("extract/1"),
            }))
            .unwrap();
        let ResponseKind::TrainingRows(rows) = training.result else {
            panic!()
        };
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0].attempt_id, key("try/one"));
        let repeat = store
            .execute(Request::Read(Query::RepeatabilityRows {
                cohort_id: key("extract/1"),
            }))
            .unwrap();
        let ResponseKind::RepeatabilityRows(rows) = repeat.result else {
            panic!()
        };
        assert_eq!(rows.len(), 2);

        commit(
            &store,
            vec![event(
                "select/2",
                Event::SelectAttempt(selection("try/two", "CLIP0001", "extract/1")),
            )],
        );
        let ResponseKind::TrainingRows(rows) = store
            .execute(Request::Read(Query::TrainingRows {
                cohort_id: key("extract/1"),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert_eq!(rows[0].attempt_id, key("try/two"));

        commit(
            &store,
            vec![event(
                "state/correct",
                Event::SetSampleState(state(
                    second.sample_ids[0].clone(),
                    SampleState::Confirmed {
                        speaker_id: key("speaker/bob"),
                    },
                )),
            )],
        );
        let ResponseKind::TrainingRows(rows) = store
            .execute(Request::Read(Query::TrainingRows {
                cohort_id: key("extract/1"),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert_eq!(rows[0].speaker_id, key("speaker/bob"));
        commit(
            &store,
            vec![event(
                "state/retract",
                Event::SetSampleState(state(second.sample_ids[0].clone(), SampleState::Retracted)),
            )],
        );
        let ResponseKind::TrainingRows(rows) = store
            .execute(Request::Read(Query::TrainingRows {
                cohort_id: key("extract/1"),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert!(rows.is_empty());
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn artifacts_are_hashed_immutable_and_cohort_checked() {
        let path = path();
        let store = open(&path).unwrap();
        let bytes = b"model bytes".to_vec();
        let hash: [u8; 32] = Sha256::digest(&bytes).into();
        let model = ArtifactRecord {
            artifact_id: key("artifact/model"),
            cohort_id: key("cohort/1"),
            kind: ArtifactKind::Model,
            sha256: hash,
            bytes,
        };
        let evaluation = ArtifactRecord {
            artifact_id: key("artifact/eval"),
            cohort_id: key("cohort/1"),
            kind: ArtifactKind::Evaluation,
            sha256: Sha256::digest(b"evaluation").into(),
            bytes: b"evaluation".to_vec(),
        };
        commit(
            &store,
            vec![
                event("artifact/model", Event::PutArtifact(model.clone())),
                event("artifact/eval", Event::PutArtifact(evaluation)),
            ],
        );
        let activate = |id: &str, cohort: &str, artifact: &str| {
            event(
                id,
                Event::SetActiveModel(ActiveModelChange {
                    cohort_id: key(cohort),
                    model_artifact_id: Some(key(artifact)),
                    reason: "approved".into(),
                }),
            )
        };
        assert!(
            store
                .execute(Request::Commit(vec![activate(
                    "activate/wrong-kind",
                    "cohort/1",
                    "artifact/eval"
                )]))
                .is_err()
        );
        assert!(
            store
                .execute(Request::Commit(vec![activate(
                    "activate/wrong-cohort",
                    "cohort/2",
                    "artifact/model"
                )]))
                .is_err()
        );

        let mut changed = model.clone();
        changed.bytes.push(1);
        changed.sha256 = Sha256::digest(&changed.bytes).into();
        assert!(matches!(
            store.execute(Request::Commit(vec![event(
                "artifact/change",
                Event::PutArtifact(changed)
            )])),
            Err(StoreError::Conflict(_))
        ));
        let mut bad_hash = model;
        bad_hash.artifact_id = key("artifact/bad");
        bad_hash.sha256 = [0; 32];
        assert!(matches!(
            store.execute(Request::Commit(vec![event(
                "artifact/bad",
                Event::PutArtifact(bad_hash)
            )])),
            Err(StoreError::Validation(_))
        ));
        commit(
            &store,
            vec![activate("activate/good", "cohort/1", "artifact/model")],
        );
        let ResponseKind::ActiveModel(Some(active)) = store
            .execute(Request::Read(Query::ActiveModel {
                cohort_id: key("cohort/1"),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert_eq!(active, key("artifact/model"));
        commit(
            &store,
            vec![event(
                "activate/none",
                Event::SetActiveModel(ActiveModelChange {
                    cohort_id: key("cohort/1"),
                    model_artifact_id: None,
                    reason: "retired".into(),
                }),
            )],
        );
        assert!(matches!(
            store
                .execute(Request::Read(Query::ActiveModel {
                    cohort_id: key("cohort/1")
                }))
                .unwrap()
                .result,
            ResponseKind::ActiveModel(None)
        ));
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn filtering_inventory_missing_and_counts_are_typed() {
        let path = path();
        let store = open(&path).unwrap();
        let unscorable = AttemptRecord {
            attempt_id: key("unscorable"),
            clip_object: object("CLIP0001"),
            extraction_key: key("extract/1"),
            provider_result_object: Some(object("RESULT02")),
            outcome: StoredAttemptOutcome::Unscorable {
                reason: "too faint".into(),
            },
        };
        let invalid = AttemptRecord {
            attempt_id: key("invalid"),
            clip_object: object("CLIP0001"),
            extraction_key: key("extract/1"),
            provider_result_object: Some(object("RESULT03")),
            outcome: StoredAttemptOutcome::InvalidResponse {
                error: "extra field".into(),
            },
        };
        commit(
            &store,
            vec![
                event(
                    "register",
                    Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 2)),
                ),
                event(
                    "attempt/scored",
                    Event::RecordAttempt(scored("scored", "CLIP0001", "extract/1", 1)),
                ),
                event("attempt/unscorable", Event::RecordAttempt(unscorable)),
                event("attempt/invalid", Event::RecordAttempt(invalid)),
                event(
                    "attempt/failure",
                    Event::RecordAttempt(failure("failure", "CLIP0001", "extract/1")),
                ),
                event(
                    "attempt/other",
                    Event::RecordAttempt(failure("other", "CLIP0001", "extract/2")),
                ),
            ],
        );
        let ResponseKind::Attempts(filtered) = store
            .execute(Request::Read(Query::Attempts {
                extraction_key: Some(key("extract/1")),
                source_object: Some(object("SOURCE01")),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert_eq!(filtered.len(), 4);
        let ResponseKind::Inventory(inventory) = store
            .execute(Request::Read(Query::Inventory {
                extraction_key: Some(key("extract/1")),
            }))
            .unwrap()
            .result
        else {
            panic!()
        };
        assert_eq!(inventory.attempts.len(), 4);
        assert_eq!(
            inventory.outcome_counts,
            OutcomeCounts {
                scored: 1,
                unscorable: 1,
                invalid_response: 1,
                provider_failure: 1,
            }
        );
        assert_eq!(inventory.missing_segment_ordinals.len(), 1);
        assert_eq!(inventory.missing_segment_ordinals[0].ordinals, vec![1]);
        drop(store);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn future_schema_is_rejected() {
        let path = path();
        let connection = Connection::open(&path).unwrap();
        connection
            .execute_batch("PRAGMA user_version = 2;")
            .unwrap();
        drop(connection);
        assert!(matches!(open(&path), Err(StoreError::SchemaVersion(2))));
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn one_store_serializes_concurrent_calls() {
        let path = path();
        let store = Arc::new(open(&path).unwrap());
        commit(
            &store,
            vec![event(
                "register",
                Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
            )],
        );
        let threads: Vec<_> = (0..8)
            .map(|index| {
                let store = Arc::clone(&store);
                thread::spawn(move || {
                    commit(
                        &store,
                        vec![event(
                            &format!("event/{index}"),
                            Event::RecordAttempt(failure(
                                &format!("attempt/{index}"),
                                "CLIP0001",
                                "extract/1",
                            )),
                        )],
                    );
                })
            })
            .collect();
        for thread in threads {
            thread.join().unwrap();
        }
        let response = store
            .execute(Request::Read(Query::Attempts {
                extraction_key: None,
                source_object: None,
            }))
            .unwrap();
        assert_eq!(response.revision, 9);
        let ResponseKind::Attempts(attempts) = response.result else {
            panic!()
        };
        assert_eq!(attempts.len(), 8);
        drop(store);
        fs::remove_file(path).unwrap();
    }
}