1use serde::{Deserialize, Serialize};
2use std::fmt::{Display, Formatter};
3use std::path::{Component, Path, PathBuf};
4
5pub use kcode_k1_transaction_id::TxId;
6pub use kcode_speaker_v3_analysis::{ExecutedAnalysis, FeatureVector24, LocalSpeakerLabel};
7
8const EVENT_VERSION: u8 = 4;
9const QUEUE_TAG: u8 = 1;
10const TRANSCRIPTION_COMPLETE_TAG: u8 = 2;
11const FAILED_TAG: u8 = 3;
12const DISCARDED_TAG: u8 = 4;
13const LABEL_CONFIRMATION_TAG: u8 = 5;
14const PROGRESS_TAG: u8 = 6;
15
16const FRAGMENT_VERSION: u8 = 1;
17const STAGED_KIND: u8 = 1;
18const FINAL_KIND: u8 = 2;
19const HEADER_LEN: usize = 16;
20const ANALYSIS_SLOT_OFFSET: usize = 16;
21const CONFIRMATION_SLOT_OFFSET: usize = 32;
22const BODY_OFFSET: usize = 48;
23
24const BASE64_ALPHABET: &[u8; 64] =
25 b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
28pub enum FragmentStageV1 {
29 Queue,
30 Transcript,
31 SpeakerLabels,
32 SpeakerFeatures,
33 Structuring,
34 LabelConfirmation,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
38pub struct QueueV2 {
39 #[serde(with = "txid_serde")]
40 pub audio_object_id: TxId,
41}
42
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44pub struct ProgressV1 {
45 #[serde(with = "txid_serde")]
46 pub fragment_id: TxId,
47 pub update: ProgressUpdateV1,
48}
49
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51pub enum ProgressUpdateV1 {
52 LlmJobStarted {
53 sequence: u64,
54 stage: FragmentStageV1,
55 name: String,
56 },
57 LlmJobSucceeded {
58 sequence: u64,
59 },
60 LlmJobFailed {
61 sequence: u64,
62 error: String,
63 },
64 StageCompleted {
65 stage: FragmentStageV1,
66 },
67}
68
69#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
70pub struct TranscriptionCompleteV1 {
71 #[serde(with = "txid_serde")]
72 pub fragment_id: TxId,
73 pub analysis: ExecutedAnalysis,
74}
75
76#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
77pub struct FailedV2 {
78 #[serde(with = "txid_serde")]
79 pub fragment_id: TxId,
80 pub stage: FragmentStageV1,
81 pub llm_job_sequence: Option<u64>,
82 pub error: String,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct DiscardedV2 {
87 #[serde(with = "txid_serde")]
88 pub fragment_id: TxId,
89}
90
91#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
92pub struct SpeakerLabelV1 {
93 pub speaker: LocalSpeakerLabel,
94 pub person_id: String,
95}
96
97#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
98pub struct LabelConfirmationV1 {
99 #[serde(with = "txid_serde")]
100 pub fragment_id: TxId,
101 #[serde(with = "txid_serde")]
102 pub interim_txid: TxId,
103 pub speakers: Vec<SpeakerLabelV1>,
104}
105
106#[allow(clippy::large_enum_variant)]
107#[derive(Debug, Clone, PartialEq)]
108pub enum AudioClassificationEventV3 {
109 Queue(QueueV2),
110 Progress(ProgressV1),
111 TranscriptionComplete(TranscriptionCompleteV1),
112 Failed(FailedV2),
113 Discarded(DiscardedV2),
114 LabelConfirmation(LabelConfirmationV1),
115}
116
117#[derive(Debug, Clone, PartialEq)]
118pub struct StagedSpeakerV1 {
119 pub speaker: LocalSpeakerLabel,
120 pub language: String,
121 pub features: FeatureVector24,
122 pub usable_for_training: bool,
123}
124
125#[derive(Debug, Clone, PartialEq)]
126pub struct StagedFragmentV1 {
127 pub analysis_txid: TxId,
128 pub transcript: String,
129 pub speakers: Vec<StagedSpeakerV1>,
130}
131
132#[derive(Debug, Clone, PartialEq)]
133pub struct FinalSpeakerV1 {
134 pub speaker: LocalSpeakerLabel,
135 pub person_id: String,
136 pub language: String,
137 pub features: FeatureVector24,
138 pub usable_for_training: bool,
139}
140
141#[derive(Debug, Clone, PartialEq)]
142pub struct FinalFragmentV1 {
143 pub analysis_txid: TxId,
144 pub confirmation_txid: TxId,
145 pub transcript: String,
146 pub speakers: Vec<FinalSpeakerV1>,
147}
148
149#[derive(Debug, Clone, PartialEq, Eq)]
150pub enum FormatError {
151 Truncated,
152 LengthOverflow,
153 UnsupportedVersion(u8),
154 UnknownEventTag(u8),
155 InvalidEventBody,
156 InvalidFragmentKind(u8),
157 NonZeroReserved,
158 NonZeroPadding,
159 NonZeroStagedConfirmation,
160 InvalidTxIdSlotLength(usize),
161 InvalidUtf8,
162 InvalidSpeakerLabel,
163 InvalidFeatureBody,
164 InvalidBoolean(u8),
165 TrailingBytes,
166 InvalidPath,
167}
168
169impl Display for FormatError {
170 fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
171 match self {
172 Self::Truncated => formatter.write_str("truncated input"),
173 Self::LengthOverflow => formatter.write_str("encoded length overflow"),
174 Self::UnsupportedVersion(version) => {
175 write!(formatter, "unsupported format version {version}")
176 }
177 Self::UnknownEventTag(tag) => write!(formatter, "unknown event tag {tag}"),
178 Self::InvalidEventBody => formatter.write_str("invalid event body"),
179 Self::InvalidFragmentKind(kind) => {
180 write!(formatter, "invalid fragment kind {kind}")
181 }
182 Self::NonZeroReserved => formatter.write_str("nonzero reserved bytes"),
183 Self::NonZeroPadding => formatter.write_str("nonzero transaction ID slot padding"),
184 Self::NonZeroStagedConfirmation => {
185 formatter.write_str("nonzero staged confirmation slot")
186 }
187 Self::InvalidTxIdSlotLength(length) => {
188 write!(formatter, "invalid transaction ID slot length {length}")
189 }
190 Self::InvalidUtf8 => formatter.write_str("invalid UTF-8"),
191 Self::InvalidSpeakerLabel => formatter.write_str("invalid speaker label"),
192 Self::InvalidFeatureBody => formatter.write_str("invalid feature body"),
193 Self::InvalidBoolean(value) => write!(formatter, "invalid boolean byte {value}"),
194 Self::TrailingBytes => formatter.write_str("trailing bytes"),
195 Self::InvalidPath => formatter.write_str("invalid transaction ID path"),
196 }
197 }
198}
199
200impl std::error::Error for FormatError {}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub struct TxIdSlot {
204 txid: TxId,
205}
206
207impl TxIdSlot {
208 pub const LEN: usize = 16;
209 pub const PADDING_LEN: usize = 4;
210
211 pub const fn new(txid: TxId) -> Self {
212 Self { txid }
213 }
214
215 pub const fn txid(self) -> TxId {
216 self.txid
217 }
218
219 pub fn encode(self) -> [u8; Self::LEN] {
220 let mut output = [0; Self::LEN];
221 output[..12].copy_from_slice(self.txid.as_bytes());
222 output
223 }
224
225 pub fn decode(bytes: &[u8]) -> Result<Self, FormatError> {
226 if bytes.len() != Self::LEN {
227 return Err(FormatError::InvalidTxIdSlotLength(bytes.len()));
228 }
229 if bytes[12..].iter().any(|byte| *byte != 0) {
230 return Err(FormatError::NonZeroPadding);
231 }
232 let mut txid = [0; 12];
233 txid.copy_from_slice(&bytes[..12]);
234 Ok(Self::new(TxId::from_bytes(txid)))
235 }
236}
237
238pub fn encode_event(event: &AudioClassificationEventV3) -> Result<Vec<u8>, FormatError> {
239 let (tag, body) = match event {
240 AudioClassificationEventV3::Queue(value) => (QUEUE_TAG, postcard::to_allocvec(value)),
241 AudioClassificationEventV3::TranscriptionComplete(value) => {
242 (TRANSCRIPTION_COMPLETE_TAG, postcard::to_allocvec(value))
243 }
244 AudioClassificationEventV3::Failed(value) => (FAILED_TAG, postcard::to_allocvec(value)),
245 AudioClassificationEventV3::Discarded(value) => {
246 (DISCARDED_TAG, postcard::to_allocvec(value))
247 }
248 AudioClassificationEventV3::LabelConfirmation(value) => {
249 (LABEL_CONFIRMATION_TAG, postcard::to_allocvec(value))
250 }
251 AudioClassificationEventV3::Progress(value) => (PROGRESS_TAG, postcard::to_allocvec(value)),
252 };
253 let body = body.map_err(|_| FormatError::InvalidEventBody)?;
254 let mut output = Vec::with_capacity(2 + body.len());
255 output.extend_from_slice(&[EVENT_VERSION, tag]);
256 output.extend_from_slice(&body);
257 Ok(output)
258}
259
260pub fn decode_event(bytes: &[u8]) -> Result<AudioClassificationEventV3, FormatError> {
261 if bytes.len() < 2 {
262 return Err(FormatError::Truncated);
263 }
264 if bytes[0] != EVENT_VERSION {
265 return Err(FormatError::UnsupportedVersion(bytes[0]));
266 }
267 match bytes[1] {
268 QUEUE_TAG => decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::Queue),
269 TRANSCRIPTION_COMPLETE_TAG => {
270 decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::TranscriptionComplete)
271 }
272 FAILED_TAG => decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::Failed),
273 DISCARDED_TAG => decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::Discarded),
274 LABEL_CONFIRMATION_TAG => {
275 decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::LabelConfirmation)
276 }
277 PROGRESS_TAG => decode_event_body(&bytes[2..]).map(AudioClassificationEventV3::Progress),
278 tag => Err(FormatError::UnknownEventTag(tag)),
279 }
280}
281
282pub fn encode_staged_fragment(value: &StagedFragmentV1) -> Result<Vec<u8>, FormatError> {
283 let mut output = encode_fragment_prefix(STAGED_KIND, value.analysis_txid, None);
284 append_string(&mut output, &value.transcript)?;
285 append_length(&mut output, value.speakers.len())?;
286 for speaker in &value.speakers {
287 append_string(&mut output, &speaker.speaker.to_string())?;
288 append_string(&mut output, &speaker.language)?;
289 append_features(&mut output, &speaker.features)?;
290 output.push(u8::from(speaker.usable_for_training));
291 }
292 Ok(output)
293}
294
295pub fn decode_staged_fragment(bytes: &[u8]) -> Result<StagedFragmentV1, FormatError> {
296 let (analysis_txid, _) = decode_fragment_header(bytes, STAGED_KIND, false)?;
297 let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
298 let transcript = decoder.read_string()?;
299 let speaker_count = decoder.read_count()?;
300 let mut speakers = Vec::new();
301 for _ in 0..speaker_count {
302 speakers.push(StagedSpeakerV1 {
303 speaker: decoder.read_speaker_label()?,
304 language: decoder.read_string()?,
305 features: decoder.read_features()?,
306 usable_for_training: decoder.read_boolean()?,
307 });
308 }
309 decoder.finish()?;
310 Ok(StagedFragmentV1 {
311 analysis_txid,
312 transcript,
313 speakers,
314 })
315}
316
317pub fn encode_final_fragment(value: &FinalFragmentV1) -> Result<Vec<u8>, FormatError> {
318 let mut output = encode_fragment_prefix(
319 FINAL_KIND,
320 value.analysis_txid,
321 Some(value.confirmation_txid),
322 );
323 append_string(&mut output, &value.transcript)?;
324 append_length(&mut output, value.speakers.len())?;
325 for speaker in &value.speakers {
326 append_string(&mut output, &speaker.speaker.to_string())?;
327 append_string(&mut output, &speaker.person_id)?;
328 append_string(&mut output, &speaker.language)?;
329 append_features(&mut output, &speaker.features)?;
330 output.push(u8::from(speaker.usable_for_training));
331 }
332 Ok(output)
333}
334
335pub fn decode_final_fragment(bytes: &[u8]) -> Result<FinalFragmentV1, FormatError> {
336 let (analysis_txid, confirmation_txid) = decode_fragment_header(bytes, FINAL_KIND, true)?;
337 let mut decoder = BodyDecoder::new(&bytes[BODY_OFFSET..]);
338 let transcript = decoder.read_string()?;
339 let speaker_count = decoder.read_count()?;
340 let mut speakers = Vec::new();
341 for _ in 0..speaker_count {
342 speakers.push(FinalSpeakerV1 {
343 speaker: decoder.read_speaker_label()?,
344 person_id: decoder.read_string()?,
345 language: decoder.read_string()?,
346 features: decoder.read_features()?,
347 usable_for_training: decoder.read_boolean()?,
348 });
349 }
350 decoder.finish()?;
351 Ok(FinalFragmentV1 {
352 analysis_txid,
353 confirmation_txid: confirmation_txid
354 .expect("final fragment header always contains a confirmation slot"),
355 transcript,
356 speakers,
357 })
358}
359
360pub fn txid_path(txid: TxId) -> PathBuf {
361 let encoded = encode_txid_base64(txid);
362 let shard = String::from_utf8(encoded[..1].to_vec()).expect("base64 alphabet is ASCII");
363 let name = String::from_utf8(encoded[1..].to_vec()).expect("base64 alphabet is ASCII");
364 PathBuf::from(shard).join(format!("{name}.dat"))
365}
366
367pub fn txid_from_path(path: impl AsRef<Path>) -> Result<TxId, FormatError> {
368 let mut components = path.as_ref().components();
369 let shard = normal_utf8_component(components.next())?;
370 let filename = normal_utf8_component(components.next())?;
371 if components.next().is_some() || shard.len() != 1 {
372 return Err(FormatError::InvalidPath);
373 }
374 let name = filename
375 .strip_suffix(".dat")
376 .ok_or(FormatError::InvalidPath)?;
377 if name.len() != 15 {
378 return Err(FormatError::InvalidPath);
379 }
380 let mut encoded = [0; 16];
381 encoded[0] = shard.as_bytes()[0];
382 encoded[1..].copy_from_slice(name.as_bytes());
383 decode_txid_base64(encoded)
384}
385
386fn decode_event_body<T: for<'de> Deserialize<'de>>(bytes: &[u8]) -> Result<T, FormatError> {
387 let (value, remaining) =
388 postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidEventBody)?;
389 if !remaining.is_empty() {
390 return Err(FormatError::TrailingBytes);
391 }
392 Ok(value)
393}
394
395fn encode_fragment_prefix(
396 kind: u8,
397 analysis_txid: TxId,
398 confirmation_txid: Option<TxId>,
399) -> Vec<u8> {
400 let mut output = vec![0; BODY_OFFSET];
401 output[0] = FRAGMENT_VERSION;
402 output[1] = kind;
403 output[ANALYSIS_SLOT_OFFSET..ANALYSIS_SLOT_OFFSET + TxIdSlot::LEN]
404 .copy_from_slice(&TxIdSlot::new(analysis_txid).encode());
405 if let Some(txid) = confirmation_txid {
406 output[CONFIRMATION_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET + TxIdSlot::LEN]
407 .copy_from_slice(&TxIdSlot::new(txid).encode());
408 }
409 output
410}
411
412fn decode_fragment_header(
413 bytes: &[u8],
414 expected_kind: u8,
415 has_confirmation: bool,
416) -> Result<(TxId, Option<TxId>), FormatError> {
417 if bytes.len() < BODY_OFFSET {
418 return Err(FormatError::Truncated);
419 }
420 if bytes[0] != FRAGMENT_VERSION {
421 return Err(FormatError::UnsupportedVersion(bytes[0]));
422 }
423 if bytes[1] != expected_kind {
424 return Err(FormatError::InvalidFragmentKind(bytes[1]));
425 }
426 if bytes[2..HEADER_LEN].iter().any(|byte| *byte != 0) {
427 return Err(FormatError::NonZeroReserved);
428 }
429
430 let analysis_txid =
431 TxIdSlot::decode(&bytes[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET])?.txid();
432 let confirmation_slot = &bytes[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET];
433 let confirmation_txid = if has_confirmation {
434 Some(TxIdSlot::decode(confirmation_slot)?.txid())
435 } else {
436 if confirmation_slot.iter().any(|byte| *byte != 0) {
437 return Err(FormatError::NonZeroStagedConfirmation);
438 }
439 None
440 };
441 Ok((analysis_txid, confirmation_txid))
442}
443
444fn append_length(output: &mut Vec<u8>, length: usize) -> Result<(), FormatError> {
445 output.extend_from_slice(
446 &u64::try_from(length)
447 .map_err(|_| FormatError::LengthOverflow)?
448 .to_le_bytes(),
449 );
450 Ok(())
451}
452
453fn append_bytes(output: &mut Vec<u8>, bytes: &[u8]) -> Result<(), FormatError> {
454 append_length(output, bytes.len())?;
455 output.extend_from_slice(bytes);
456 Ok(())
457}
458
459fn append_string(output: &mut Vec<u8>, value: &str) -> Result<(), FormatError> {
460 append_bytes(output, value.as_bytes())
461}
462
463fn append_features(output: &mut Vec<u8>, features: &FeatureVector24) -> Result<(), FormatError> {
464 let bytes = postcard::to_allocvec(features).map_err(|_| FormatError::InvalidFeatureBody)?;
465 append_bytes(output, &bytes)
466}
467
468struct BodyDecoder<'a> {
469 bytes: &'a [u8],
470 position: usize,
471}
472
473impl<'a> BodyDecoder<'a> {
474 fn new(bytes: &'a [u8]) -> Self {
475 Self { bytes, position: 0 }
476 }
477
478 fn take(&mut self, length: usize) -> Result<&'a [u8], FormatError> {
479 let end = self
480 .position
481 .checked_add(length)
482 .ok_or(FormatError::LengthOverflow)?;
483 let value = self
484 .bytes
485 .get(self.position..end)
486 .ok_or(FormatError::Truncated)?;
487 self.position = end;
488 Ok(value)
489 }
490
491 fn read_u64(&mut self) -> Result<u64, FormatError> {
492 let mut bytes = [0; 8];
493 bytes.copy_from_slice(self.take(8)?);
494 Ok(u64::from_le_bytes(bytes))
495 }
496
497 fn read_length(&mut self) -> Result<usize, FormatError> {
498 usize::try_from(self.read_u64()?).map_err(|_| FormatError::LengthOverflow)
499 }
500
501 fn read_count(&mut self) -> Result<usize, FormatError> {
502 self.read_length()
503 }
504
505 fn read_bytes(&mut self) -> Result<&'a [u8], FormatError> {
506 let length = self.read_length()?;
507 self.take(length)
508 }
509
510 fn read_string(&mut self) -> Result<String, FormatError> {
511 Ok(std::str::from_utf8(self.read_bytes()?)
512 .map_err(|_| FormatError::InvalidUtf8)?
513 .to_owned())
514 }
515
516 fn read_speaker_label(&mut self) -> Result<LocalSpeakerLabel, FormatError> {
517 self.read_string()?
518 .parse()
519 .map_err(|_| FormatError::InvalidSpeakerLabel)
520 }
521
522 fn read_features(&mut self) -> Result<FeatureVector24, FormatError> {
523 let bytes = self.read_bytes()?;
524 let (value, remaining) =
525 postcard::take_from_bytes(bytes).map_err(|_| FormatError::InvalidFeatureBody)?;
526 if !remaining.is_empty() {
527 return Err(FormatError::InvalidFeatureBody);
528 }
529 Ok(value)
530 }
531
532 fn read_boolean(&mut self) -> Result<bool, FormatError> {
533 match self.take(1)?[0] {
534 0 => Ok(false),
535 1 => Ok(true),
536 value => Err(FormatError::InvalidBoolean(value)),
537 }
538 }
539
540 fn finish(self) -> Result<(), FormatError> {
541 if self.position == self.bytes.len() {
542 Ok(())
543 } else {
544 Err(FormatError::TrailingBytes)
545 }
546 }
547}
548
549fn normal_utf8_component(component: Option<Component<'_>>) -> Result<&str, FormatError> {
550 match component {
551 Some(Component::Normal(value)) => value.to_str().ok_or(FormatError::InvalidPath),
552 _ => Err(FormatError::InvalidPath),
553 }
554}
555
556fn encode_txid_base64(txid: TxId) -> [u8; 16] {
557 let bytes = txid.into_bytes();
558 let mut encoded = [0; 16];
559 for group in 0..4 {
560 let input = group * 3;
561 let output = group * 4;
562 encoded[output] = BASE64_ALPHABET[(bytes[input] >> 2) as usize];
563 encoded[output + 1] =
564 BASE64_ALPHABET[(((bytes[input] & 3) << 4) | (bytes[input + 1] >> 4)) as usize];
565 encoded[output + 2] =
566 BASE64_ALPHABET[(((bytes[input + 1] & 15) << 2) | (bytes[input + 2] >> 6)) as usize];
567 encoded[output + 3] = BASE64_ALPHABET[(bytes[input + 2] & 63) as usize];
568 }
569 encoded
570}
571
572fn decode_txid_base64(encoded: [u8; 16]) -> Result<TxId, FormatError> {
573 let mut bytes = [0; 12];
574 for group in 0..4 {
575 let input = group * 4;
576 let output = group * 3;
577 let first = decode_base64_character(encoded[input])?;
578 let second = decode_base64_character(encoded[input + 1])?;
579 let third = decode_base64_character(encoded[input + 2])?;
580 let fourth = decode_base64_character(encoded[input + 3])?;
581 bytes[output] = (first << 2) | (second >> 4);
582 bytes[output + 1] = (second << 4) | (third >> 2);
583 bytes[output + 2] = (third << 6) | fourth;
584 }
585 Ok(TxId::from_bytes(bytes))
586}
587
588fn decode_base64_character(value: u8) -> Result<u8, FormatError> {
589 match value {
590 b'A'..=b'Z' => Ok(value - b'A'),
591 b'a'..=b'z' => Ok(value - b'a' + 26),
592 b'0'..=b'9' => Ok(value - b'0' + 52),
593 b'-' => Ok(62),
594 b'_' => Ok(63),
595 _ => Err(FormatError::InvalidPath),
596 }
597}
598
599mod txid_serde {
600 use super::TxId;
601 use serde::{Deserialize, Deserializer, Serialize, Serializer};
602
603 pub fn serialize<S: Serializer>(value: &TxId, serializer: S) -> Result<S::Ok, S::Error> {
604 value.into_bytes().serialize(serializer)
605 }
606
607 pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<TxId, D::Error> {
608 Ok(TxId::from_bytes(<[u8; 12]>::deserialize(deserializer)?))
609 }
610}
611
612#[cfg(test)]
613mod tests {
614 use super::*;
615 use kcode_speaker_v3_analysis::{
616 AnalysisEnvelope, GeminiCohort, OggAudioMetadata, StructuredAnalysis, StructuredSpeaker,
617 StructurerProvenance,
618 };
619
620 fn txid(seed: u8) -> TxId {
621 let mut bytes = [0; 12];
622 for (index, value) in bytes.iter_mut().enumerate() {
623 *value = seed.wrapping_add(index as u8);
624 }
625 TxId::from_bytes(bytes)
626 }
627
628 fn features(seed: f64) -> FeatureVector24 {
629 FeatureVector24 {
630 median_f0_hz: Some(seed),
631 ..FeatureVector24::default()
632 }
633 }
634
635 fn analysis() -> ExecutedAnalysis {
636 let mut ogg = vec![0; 29];
637 ogg[..4].copy_from_slice(b"OggS");
638 ogg[26] = 1;
639 ogg[27] = 1;
640 ExecutedAnalysis {
641 envelope: AnalysisEnvelope {
642 audio: OggAudioMetadata::from_bytes(&ogg, 1_250, Some("sample.ogg".into()))
643 .expect("Ogg metadata"),
644 analysis: StructuredAnalysis {
645 transcript: "Speaker 1".into(),
646 speakers: vec![StructuredSpeaker {
647 speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
648 language: "English".into(),
649 features: features(182.5),
650 features_usable_for_training: true,
651 }],
652 },
653 gemini: GeminiCohort::new("gemini-3.1-pro"),
654 structurer: StructurerProvenance::new("terra-structurer"),
655 },
656 label_extractor: StructurerProvenance::new("terra-labeler"),
657 }
658 }
659
660 fn staged_fragment() -> StagedFragmentV1 {
661 StagedFragmentV1 {
662 analysis_txid: txid(20),
663 transcript: "hello".into(),
664 speakers: vec![StagedSpeakerV1 {
665 speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
666 language: "English".into(),
667 features: features(201.25),
668 usable_for_training: true,
669 }],
670 }
671 }
672
673 fn final_fragment() -> FinalFragmentV1 {
674 FinalFragmentV1 {
675 analysis_txid: txid(30),
676 confirmation_txid: txid(40),
677 transcript: "hello".into(),
678 speakers: vec![FinalSpeakerV1 {
679 speaker: LocalSpeakerLabel::new(1).expect("speaker label"),
680 person_id: "person-1".into(),
681 language: "English".into(),
682 features: features(201.25),
683 usable_for_training: false,
684 }],
685 }
686 }
687
688 fn assert_event_roundtrip(event: AudioClassificationEventV3, expected_tag: u8) {
689 let bytes = encode_event(&event).expect("encode event");
690 assert_eq!(bytes[0], EVENT_VERSION);
691 assert_eq!(bytes[1], expected_tag);
692 assert_eq!(decode_event(&bytes), Ok(event));
693 }
694
695 #[test]
696 fn slot_roundtrips_and_rejects_padding() {
697 let id = txid(3);
698 let encoded = TxIdSlot::new(id).encode();
699 assert_eq!(TxIdSlot::decode(&encoded).expect("decode slot").txid(), id);
700 let mut bad_padding = encoded;
701 bad_padding[12] = 1;
702 assert_eq!(
703 TxIdSlot::decode(&bad_padding),
704 Err(FormatError::NonZeroPadding)
705 );
706 }
707
708 #[test]
709 fn all_v4_events_roundtrip_with_stable_tags() {
710 let speaker = LocalSpeakerLabel::new(1).expect("speaker label");
711 assert_event_roundtrip(
712 AudioClassificationEventV3::Queue(QueueV2 {
713 audio_object_id: txid(1),
714 }),
715 QUEUE_TAG,
716 );
717 assert_event_roundtrip(
718 AudioClassificationEventV3::TranscriptionComplete(TranscriptionCompleteV1 {
719 fragment_id: txid(2),
720 analysis: analysis(),
721 }),
722 TRANSCRIPTION_COMPLETE_TAG,
723 );
724 assert_event_roundtrip(
725 AudioClassificationEventV3::Failed(FailedV2 {
726 fragment_id: txid(3),
727 stage: FragmentStageV1::SpeakerFeatures,
728 llm_job_sequence: Some(7),
729 error: "failure".into(),
730 }),
731 FAILED_TAG,
732 );
733 assert_event_roundtrip(
734 AudioClassificationEventV3::Discarded(DiscardedV2 {
735 fragment_id: txid(4),
736 }),
737 DISCARDED_TAG,
738 );
739 assert_event_roundtrip(
740 AudioClassificationEventV3::LabelConfirmation(LabelConfirmationV1 {
741 fragment_id: txid(5),
742 interim_txid: txid(6),
743 speakers: vec![SpeakerLabelV1 {
744 speaker,
745 person_id: "person-1".into(),
746 }],
747 }),
748 LABEL_CONFIRMATION_TAG,
749 );
750 assert_event_roundtrip(
751 AudioClassificationEventV3::Progress(ProgressV1 {
752 fragment_id: txid(7),
753 update: ProgressUpdateV1::LlmJobStarted {
754 sequence: 9,
755 stage: FragmentStageV1::Transcript,
756 name: "transcribe".into(),
757 },
758 }),
759 PROGRESS_TAG,
760 );
761 }
762
763 #[test]
764 fn every_progress_update_roundtrips() {
765 let updates = [
766 ProgressUpdateV1::LlmJobStarted {
767 sequence: 1,
768 stage: FragmentStageV1::Transcript,
769 name: "transcript".into(),
770 },
771 ProgressUpdateV1::LlmJobSucceeded { sequence: 1 },
772 ProgressUpdateV1::LlmJobFailed {
773 sequence: 2,
774 error: "provider failed".into(),
775 },
776 ProgressUpdateV1::StageCompleted {
777 stage: FragmentStageV1::Transcript,
778 },
779 ];
780 for update in updates {
781 assert_event_roundtrip(
782 AudioClassificationEventV3::Progress(ProgressV1 {
783 fragment_id: txid(8),
784 update,
785 }),
786 PROGRESS_TAG,
787 );
788 }
789 }
790
791 #[test]
792 fn events_reject_v3_unknown_tags_malformed_bodies_and_trailing_bytes() {
793 let event = AudioClassificationEventV3::Queue(QueueV2 {
794 audio_object_id: txid(1),
795 });
796 let bytes = encode_event(&event).expect("encode event");
797
798 let mut version_three = bytes.clone();
799 version_three[0] = 3;
800 assert_eq!(
801 decode_event(&version_three),
802 Err(FormatError::UnsupportedVersion(3))
803 );
804
805 let mut unknown_tag = bytes.clone();
806 unknown_tag[1] = 7;
807 assert_eq!(
808 decode_event(&unknown_tag),
809 Err(FormatError::UnknownEventTag(7))
810 );
811
812 assert_eq!(
813 decode_event(&bytes[..2]),
814 Err(FormatError::InvalidEventBody)
815 );
816
817 let mut trailing = bytes;
818 trailing.push(0);
819 assert_eq!(decode_event(&trailing), Err(FormatError::TrailingBytes));
820 }
821
822 #[test]
823 fn fragments_roundtrip_and_corruption_rejects() {
824 let staged = staged_fragment();
825 let staged_bytes = encode_staged_fragment(&staged).expect("encode staged fragment");
826 assert_eq!(
827 &staged_bytes[ANALYSIS_SLOT_OFFSET..CONFIRMATION_SLOT_OFFSET],
828 &TxIdSlot::new(staged.analysis_txid).encode()
829 );
830 assert_eq!(
831 &staged_bytes[CONFIRMATION_SLOT_OFFSET..BODY_OFFSET],
832 &[0; TxIdSlot::LEN]
833 );
834 assert_eq!(decode_staged_fragment(&staged_bytes), Ok(staged));
835
836 let final_value = final_fragment();
837 assert_eq!(
838 decode_final_fragment(
839 &encode_final_fragment(&final_value).expect("encode final fragment")
840 ),
841 Ok(final_value)
842 );
843
844 let mut bad_reserved = staged_bytes;
845 bad_reserved[2] = 1;
846 assert_eq!(
847 decode_staged_fragment(&bad_reserved),
848 Err(FormatError::NonZeroReserved)
849 );
850 }
851
852 #[test]
853 fn paths_roundtrip() {
854 let id = txid(99);
855 assert_eq!(txid_from_path(txid_path(id)), Ok(id));
856 assert_eq!(
857 txid_from_path("A/invalid.dat"),
858 Err(FormatError::InvalidPath)
859 );
860 }
861}