1#![forbid(unsafe_code)]
2
3pub mod admission;
11pub mod audio;
12pub mod health;
13
14use std::{
15 env, fmt,
16 ops::Range,
17 sync::{
18 Arc, OnceLock,
19 atomic::{AtomicBool, Ordering},
20 mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError},
21 },
22 thread,
23 time::{Duration, Instant},
24};
25
26use asupersync::runtime::{Runtime, RuntimeBuilder};
27
28pub const SCAFFOLD_REVISION: u8 = 2;
30
31const DEFAULT_QUEUE_CAPACITY: usize = 8;
32const DEFAULT_SYNTHESIS_BUDGET: Duration = Duration::from_secs(30);
33const DEFAULT_ENROLL_BUDGET: Duration = Duration::from_secs(30);
34const BACKPRESSURE_POLL: Duration = Duration::from_millis(1);
35
36const DEFAULT_SYNTHESIS_FRAME_BUDGET: Duration = Duration::from_secs(8);
49
50const DEBUG_BUILD_SLOWDOWN: u32 = 32;
68
69pub fn process_engine_config() -> EngineConfig {
75 static CONFIG: OnceLock<EngineConfig> = OnceLock::new();
76 CONFIG.get_or_init(EngineConfig::from_environment).clone()
77}
78
79#[derive(Clone, Debug, PartialEq, Eq)]
81pub struct EngineConfig {
82 pub stream_queue_capacity: usize,
84 pub synthesis_stage_budget: Duration,
91 pub synthesis_frame_budget: Duration,
97 pub enroll_stage_budget: Duration,
99 pub admission: admission::AdmissionPolicy,
101}
102
103impl Default for EngineConfig {
104 fn default() -> Self {
105 let slowdown = build_profile_slowdown();
108 Self {
109 stream_queue_capacity: DEFAULT_QUEUE_CAPACITY,
110 synthesis_stage_budget: DEFAULT_SYNTHESIS_BUDGET * slowdown,
111 synthesis_frame_budget: DEFAULT_SYNTHESIS_FRAME_BUDGET * slowdown,
112 enroll_stage_budget: DEFAULT_ENROLL_BUDGET,
113 admission: admission::AdmissionPolicy::default(),
114 }
115 }
116}
117
118const fn build_profile_slowdown() -> u32 {
125 if cfg!(debug_assertions) {
126 DEBUG_BUILD_SLOWDOWN
127 } else {
128 1
129 }
130}
131
132impl EngineConfig {
133 fn from_environment() -> Self {
134 let mut config = Self::default();
135 config.synthesis_stage_budget = stage_budget_from_environment(
136 "FTTS_STAGE_BUDGET_SYNTHESIS_MS",
137 config.synthesis_stage_budget,
138 );
139 config.synthesis_frame_budget = stage_budget_from_environment(
140 "FTTS_STAGE_BUDGET_FRAME_MS",
141 config.synthesis_frame_budget,
142 );
143 config.enroll_stage_budget = stage_budget_from_environment(
144 "FTTS_STAGE_BUDGET_ENROLL_MS",
145 config.enroll_stage_budget,
146 );
147 config.admission.budget_bytes = positive_u64_from_environment("FTTS_MEMORY_BUDGET_MB")
151 .and_then(|megabytes| megabytes.checked_mul(1024 * 1024))
152 .unwrap_or(config.admission.budget_bytes);
153 if let Some(max_frames) = positive_u64_from_environment("FTTS_MAX_FRAMES") {
154 config.admission.max_new_tokens = max_frames;
157 config.admission.heuristic_eos_backstop = false;
158 }
159 config
160 }
161
162 fn validate(&self) -> Result<(), EngineError> {
163 if self.stream_queue_capacity == 0 {
164 return Err(EngineError::InvalidConfiguration(
165 "stream queue capacity must be greater than zero",
166 ));
167 }
168 if self.synthesis_stage_budget.is_zero()
169 || self.synthesis_frame_budget.is_zero()
170 || self.enroll_stage_budget.is_zero()
171 {
172 return Err(EngineError::InvalidConfiguration(
173 "stage budgets must be greater than zero",
174 ));
175 }
176 Ok(())
177 }
178}
179
180fn positive_u64_from_environment(name: &str) -> Option<u64> {
182 env::var(name)
183 .ok()
184 .and_then(|value| value.parse::<u64>().ok())
185 .filter(|value| *value > 0)
186}
187
188fn stage_budget_from_environment(name: &str, fallback: Duration) -> Duration {
189 env::var(name)
190 .ok()
191 .and_then(|value| value.parse::<u64>().ok())
192 .filter(|milliseconds| *milliseconds > 0)
193 .map(Duration::from_millis)
194 .unwrap_or(fallback)
195}
196
197#[derive(Clone, Debug, Default)]
203pub struct CancellationToken {
204 cancelled: Arc<AtomicBool>,
205}
206
207impl CancellationToken {
208 #[must_use]
210 pub fn new() -> Self {
211 Self::default()
212 }
213
214 pub fn cancel(&self) {
216 self.cancelled.store(true, Ordering::Release);
217 }
218
219 #[must_use]
221 pub fn is_cancelled(&self) -> bool {
222 self.cancelled.load(Ordering::Acquire)
223 }
224
225 pub fn checkpoint(&self) -> Result<(), EngineError> {
227 if self.is_cancelled() {
228 Err(EngineError::Cancelled)
229 } else {
230 Ok(())
231 }
232 }
233}
234
235#[derive(Clone, Copy, Debug, PartialEq, Eq)]
237pub enum StreamKind {
238 Pcm,
240 Events,
242}
243
244#[derive(Clone, Debug, PartialEq, Eq)]
246pub struct PcmPacket {
247 pub frame_count: u8,
249 pub samples: Vec<i16>,
251}
252
253pub struct StreamQueues {
259 pub pcm: BoundedSender<PcmPacket>,
261 pub pcm_receiver: BoundedReceiver<PcmPacket>,
263 pub events: BoundedSender<SynthesisEvent>,
265 pub event_receiver: BoundedReceiver<SynthesisEvent>,
267}
268
269impl StreamQueues {
270 pub fn new(capacity: usize) -> Result<Self, EngineError> {
272 if capacity == 0 {
273 return Err(EngineError::InvalidConfiguration(
274 "stream queue capacity must be greater than zero",
275 ));
276 }
277 let (pcm, pcm_receiver) = bounded_queue(capacity, StreamKind::Pcm);
278 let (events, event_receiver) = bounded_queue(capacity, StreamKind::Events);
279 Ok(Self {
280 pcm,
281 pcm_receiver,
282 events,
283 event_receiver,
284 })
285 }
286}
287
288#[derive(Clone)]
290pub struct BoundedSender<T> {
291 kind: StreamKind,
292 sender: SyncSender<T>,
293}
294
295impl<T> BoundedSender<T> {
296 pub fn send(&self, mut item: T, cancellation: &CancellationToken) -> Result<(), EngineError> {
298 loop {
299 cancellation.checkpoint()?;
300 match self.sender.try_send(item) {
301 Ok(()) => return Ok(()),
302 Err(TrySendError::Full(returned)) => {
303 item = returned;
304 thread::sleep(BACKPRESSURE_POLL);
305 }
306 Err(TrySendError::Disconnected(_)) => {
307 return Err(EngineError::StreamDisconnected(self.kind));
308 }
309 }
310 }
311 }
312}
313
314pub struct BoundedReceiver<T> {
316 kind: StreamKind,
317 receiver: Receiver<T>,
318}
319
320impl<T> BoundedReceiver<T> {
321 pub fn recv_timeout(&self, timeout: Duration) -> Result<T, EngineError> {
323 match self.receiver.recv_timeout(timeout) {
324 Ok(item) => Ok(item),
325 Err(RecvTimeoutError::Timeout) => Err(EngineError::QueueTimeout),
326 Err(RecvTimeoutError::Disconnected) => Err(EngineError::StreamDisconnected(self.kind)),
327 }
328 }
329}
330
331fn bounded_queue<T>(capacity: usize, kind: StreamKind) -> (BoundedSender<T>, BoundedReceiver<T>) {
332 let (sender, receiver) = mpsc::sync_channel(capacity);
333 (
334 BoundedSender { kind, sender },
335 BoundedReceiver { kind, receiver },
336 )
337}
338
339#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
341pub enum NormalizationMode {
342 #[default]
344 Verbatim,
345 Conservative,
347 LocaleAware,
349}
350
351#[derive(Clone, Debug, Eq, PartialEq)]
353pub struct LanguageSpan {
354 pub range: Range<usize>,
356 pub language: String,
358}
359
360#[derive(Clone, Debug, Eq, PartialEq)]
365pub struct PronunciationEntry {
366 pub language: String,
368 pub surface: String,
370 pub spoken: String,
372}
373
374#[derive(Clone, Debug, Default, Eq, PartialEq)]
376pub struct NormalizationOptions {
377 pub mode: NormalizationMode,
379 pub language_spans: Vec<LanguageSpan>,
381 pub pronunciation_lexicon: Vec<PronunciationEntry>,
383}
384
385#[derive(Clone, Debug, Eq, PartialEq)]
391pub struct NormalizationChange {
392 pub rule: &'static str,
394 pub before: String,
396 pub after: String,
398}
399
400#[derive(Clone, Debug, Eq, PartialEq)]
402pub struct NormalizationTrace {
403 pub mode: NormalizationMode,
405 pub unicode_version: String,
407 pub changes: Vec<NormalizationChange>,
409}
410
411impl NormalizationTrace {
412 #[must_use]
414 pub fn summary(&self) -> NormalizationTraceSummary {
415 let mut rules = self
416 .changes
417 .iter()
418 .map(|change| change.rule.to_owned())
419 .collect::<Vec<_>>();
420 rules.sort_unstable();
421 rules.dedup();
422 NormalizationTraceSummary {
423 mode: self.mode,
424 unicode_version: self.unicode_version.clone(),
425 rules,
426 change_count: self.changes.len(),
427 }
428 }
429}
430
431#[derive(Clone, Debug, Eq, PartialEq)]
433pub struct NormalizationTraceSummary {
434 pub mode: NormalizationMode,
436 pub unicode_version: String,
438 pub rules: Vec<String>,
440 pub change_count: usize,
442}
443
444#[derive(Clone, Debug, Eq, PartialEq)]
446pub struct PreparedText {
447 pub token_ids: Vec<u32>,
449 pub normalization_trace: NormalizationTrace,
451}
452
453impl PreparedText {
454 #[must_use]
456 pub fn new(token_ids: Vec<u32>, normalization_trace: NormalizationTrace) -> Self {
457 Self {
458 token_ids,
459 normalization_trace,
460 }
461 }
462}
463
464#[derive(Clone, Debug, Eq, PartialEq)]
466pub struct TextPreparationError {
467 message: String,
468}
469
470impl TextPreparationError {
471 #[must_use]
473 pub fn new(message: impl Into<String>) -> Self {
474 Self {
475 message: message.into(),
476 }
477 }
478}
479
480impl fmt::Display for TextPreparationError {
481 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
482 formatter.write_str(&self.message)
483 }
484}
485
486impl std::error::Error for TextPreparationError {}
487
488pub trait TextPreparer: Send + Sync {
494 fn prepare(
496 &self,
497 text: &str,
498 options: &NormalizationOptions,
499 ) -> Result<PreparedText, TextPreparationError>;
500}
501
502#[derive(Clone, Debug, PartialEq, Eq)]
504pub struct CodeFrame {
505 pub codes: Vec<u32>,
507}
508
509#[derive(Clone, Debug, PartialEq, Eq)]
511pub struct GenerationError {
512 message: String,
513}
514
515impl GenerationError {
516 #[must_use]
518 pub fn new(message: impl Into<String>) -> Self {
519 Self {
520 message: message.into(),
521 }
522 }
523}
524
525impl fmt::Display for GenerationError {
526 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
527 formatter.write_str(&self.message)
528 }
529}
530
531impl std::error::Error for GenerationError {}
532
533pub trait FrameGenerator {
539 fn begin_utterance(&mut self, prepared: &PreparedText) -> Result<(), GenerationError>;
541
542 fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError>;
544}
545
546#[derive(Clone, Debug, PartialEq, Eq)]
548pub struct SynthesisRequest {
549 pub text: String,
551 pub normalization_options: NormalizationOptions,
553 pub trace_normalization: bool,
555}
556
557impl SynthesisRequest {
558 #[must_use]
560 pub fn new(text: impl Into<String>) -> Self {
561 Self {
562 text: text.into(),
563 normalization_options: NormalizationOptions::default(),
564 trace_normalization: false,
565 }
566 }
567
568 #[must_use]
570 pub fn with_normalization_options(
571 mut self,
572 normalization_options: NormalizationOptions,
573 ) -> Self {
574 self.normalization_options = normalization_options;
575 self
576 }
577
578 #[must_use]
580 pub const fn with_normalization_trace(mut self, trace_normalization: bool) -> Self {
581 self.trace_normalization = trace_normalization;
582 self
583 }
584}
585
586#[derive(Clone, Debug, PartialEq, Eq)]
588pub struct EnrollmentRequest {
589 pub reference_audio: Vec<u8>,
591}
592
593#[derive(Clone, Debug, PartialEq, Eq)]
595pub struct SynthesisResult {
596 pub generated_frames: u64,
598 pub code_frames: Vec<CodeFrame>,
600 pub prepared_token_count: usize,
602}
603
604#[derive(Clone, Debug, PartialEq, Eq)]
606pub struct EnrollmentResult {
607 pub accepted_reference_bytes: usize,
609}
610
611#[derive(Clone, Copy, Debug, PartialEq, Eq)]
613pub enum EngineStage {
614 Synthesis,
616 Enrollment,
618}
619
620#[derive(Clone, Copy, Debug, PartialEq, Eq)]
622pub enum HealthEvent {
623 BudgetExceeded,
625 Cancelled,
627 Violation(health::HealthViolation),
634}
635
636impl HealthEvent {
637 #[must_use]
639 pub const fn invalidates_output(self) -> bool {
640 match self {
641 Self::BudgetExceeded | Self::Cancelled => true,
642 Self::Violation(violation) => violation.invalidates_output(),
643 }
644 }
645
646 #[must_use]
648 pub const fn as_str(self) -> &'static str {
649 match self {
650 Self::BudgetExceeded => "budget_exceeded",
651 Self::Cancelled => "cancelled",
652 Self::Violation(violation) => violation.as_str(),
653 }
654 }
655}
656
657#[derive(Clone, Debug, PartialEq, Eq)]
659pub enum SynthesisEvent {
660 Admission { accepted: bool },
662 ResourceAdmission {
668 admitted: bool,
670 predicted_max_frames: u64,
672 predicted_peak_bytes: u64,
674 budget_bytes: u64,
676 },
677 StageStarted { stage: EngineStage },
679 StageFinished {
681 stage: EngineStage,
682 elapsed: Duration,
683 },
684 FrameProgress { frame: u64 },
686 TextPrepared {
688 token_count: usize,
690 normalization: NormalizationTraceSummary,
692 },
693 PacketEmitted {
695 frame_count: u8,
696 sample_count: usize,
697 },
698 Health { event: HealthEvent },
700}
701
702pub trait SynthesisObserver: Send + Sync {
707 fn on_event(&self, event: SynthesisEvent);
709}
710
711impl<F> SynthesisObserver for F
712where
713 F: Fn(SynthesisEvent) + Send + Sync,
714{
715 fn on_event(&self, event: SynthesisEvent) {
716 self(event);
717 }
718}
719
720#[derive(Clone, Debug, PartialEq, Eq)]
722pub enum EngineError {
723 Busy,
725 Cancelled,
727 BudgetExceeded(EngineStage),
729 StreamDisconnected(StreamKind),
731 QueueTimeout,
733 TextPreparation(TextPreparationError),
735 Generation(GenerationError),
737 ResourceAdmission(admission::AdmissionRejection),
742 InvalidConfiguration(&'static str),
744 Runtime(String),
746}
747
748impl fmt::Display for EngineError {
749 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
750 match self {
751 Self::Busy => formatter.write_str("another synthesis is already active"),
752 Self::Cancelled => formatter.write_str("synthesis cancelled"),
753 Self::BudgetExceeded(stage) => write!(formatter, "{stage:?} stage budget exceeded"),
754 Self::StreamDisconnected(kind) => write!(formatter, "{kind:?} stream disconnected"),
755 Self::QueueTimeout => formatter.write_str("bounded queue receive timed out"),
756 Self::TextPreparation(error) => write!(formatter, "text preparation failed: {error}"),
757 Self::Generation(error) => write!(formatter, "frame generation failed: {error}"),
758 Self::ResourceAdmission(rejection) => {
759 write!(
760 formatter,
761 "resource admission refused the request: {rejection}"
762 )
763 }
764 Self::InvalidConfiguration(message) => formatter.write_str(message),
765 Self::Runtime(message) => write!(formatter, "runtime initialization failed: {message}"),
766 }
767 }
768}
769
770impl std::error::Error for EngineError {}
771
772pub struct TtsEngine {
778 runtime: Runtime,
779 config: EngineConfig,
780 synthesis_active: AtomicBool,
781}
782
783impl TtsEngine {
784 pub fn new(config: EngineConfig) -> Result<Self, EngineError> {
786 config.validate()?;
787 let runtime = RuntimeBuilder::current_thread()
788 .blocking_threads(1, 1)
789 .build()
790 .map_err(|error| EngineError::Runtime(error.to_string()))?;
791 Ok(Self {
792 runtime,
793 config,
794 synthesis_active: AtomicBool::new(false),
795 })
796 }
797
798 pub fn from_process_environment() -> Result<Self, EngineError> {
800 Self::new(process_engine_config())
801 }
802
803 pub fn synthesize<P: TextPreparer + ?Sized>(
809 &self,
810 request: SynthesisRequest,
811 text_preparer: &P,
812 frame_generator: &mut dyn FrameGenerator,
813 cancellation: &CancellationToken,
814 observer: &dyn SynthesisObserver,
815 ) -> Result<SynthesisResult, EngineError> {
816 let _admission = self.acquire_synthesis_admission(observer)?;
817 cancellation.checkpoint().inspect_err(|_| {
818 observer.on_event(SynthesisEvent::Health {
819 event: HealthEvent::Cancelled,
820 });
821 })?;
822 let prepared = text_preparer
823 .prepare(&request.text, &request.normalization_options)
824 .map_err(EngineError::TextPreparation)?;
825 if request.trace_normalization {
826 observer.on_event(SynthesisEvent::TextPrepared {
827 token_count: prepared.token_ids.len(),
828 normalization: prepared.normalization_trace.summary(),
829 });
830 }
831
832 let prompt_tokens = prepared.token_ids.len() as u64;
837 let plan = match self.config.admission.admit(prompt_tokens) {
838 Ok(plan) => {
839 observer.on_event(SynthesisEvent::ResourceAdmission {
840 admitted: true,
841 predicted_max_frames: plan.predicted_max_frames,
842 predicted_peak_bytes: plan.predicted_peak_bytes,
843 budget_bytes: plan.budget_bytes,
844 });
845 plan
846 }
847 Err(rejection) => {
848 if let admission::AdmissionRejection::BudgetExceeded { plan } = rejection {
849 observer.on_event(SynthesisEvent::ResourceAdmission {
850 admitted: false,
851 predicted_max_frames: plan.predicted_max_frames,
852 predicted_peak_bytes: plan.predicted_peak_bytes,
853 budget_bytes: plan.budget_bytes,
854 });
855 }
856 return Err(EngineError::ResourceAdmission(rejection));
857 }
858 };
859
860 observer.on_event(SynthesisEvent::StageStarted {
861 stage: EngineStage::Synthesis,
862 });
863 let started = Instant::now();
864 let startup_budget = self.config.synthesis_stage_budget;
871 let frame_budget = self.config.synthesis_frame_budget;
872 let mut code_frames: Vec<CodeFrame> = Vec::new();
873 frame_generator
874 .begin_utterance(&prepared)
875 .map_err(EngineError::Generation)?;
876 while (code_frames.len() as u64) < plan.predicted_max_frames {
877 cancellation.checkpoint().inspect_err(|_| {
878 observer.on_event(SynthesisEvent::Health {
879 event: HealthEvent::Cancelled,
880 });
881 })?;
882 let deadline = frame_budget
885 .checked_mul(u32::try_from(code_frames.len()).unwrap_or(u32::MAX))
886 .and_then(|earned| earned.checked_add(startup_budget))
887 .unwrap_or(Duration::MAX);
888 if started.elapsed() > deadline {
889 observer.on_event(SynthesisEvent::Health {
890 event: HealthEvent::BudgetExceeded,
891 });
892 return Err(EngineError::BudgetExceeded(EngineStage::Synthesis));
893 }
894 match frame_generator
895 .next_frame()
896 .map_err(EngineError::Generation)?
897 {
898 Some(frame) => {
899 observer.on_event(SynthesisEvent::FrameProgress {
900 frame: code_frames.len() as u64,
901 });
902 code_frames.push(frame);
903 }
904 None => break,
905 }
906 }
907 observer.on_event(SynthesisEvent::StageFinished {
908 stage: EngineStage::Synthesis,
909 elapsed: started.elapsed(),
910 });
911 Ok(SynthesisResult {
912 generated_frames: code_frames.len() as u64,
913 code_frames,
914 prepared_token_count: prepared.token_ids.len(),
915 })
916 }
917
918 pub fn enroll(
920 &self,
921 request: EnrollmentRequest,
922 cancellation: &CancellationToken,
923 observer: &dyn SynthesisObserver,
924 ) -> Result<EnrollmentResult, EngineError> {
925 observer.on_event(SynthesisEvent::Admission { accepted: true });
926 self.run_stage(
927 EngineStage::Enrollment,
928 self.config.enroll_stage_budget,
929 cancellation,
930 observer,
931 |_| Ok(()),
932 )?;
933 Ok(EnrollmentResult {
934 accepted_reference_bytes: request.reference_audio.len(),
935 })
936 }
937
938 fn acquire_synthesis_admission(
939 &self,
940 observer: &dyn SynthesisObserver,
941 ) -> Result<SynthesisAdmission<'_>, EngineError> {
942 match self.synthesis_active.compare_exchange(
943 false,
944 true,
945 Ordering::AcqRel,
946 Ordering::Acquire,
947 ) {
948 Ok(_) => {
949 observer.on_event(SynthesisEvent::Admission { accepted: true });
950 Ok(SynthesisAdmission { engine: self })
951 }
952 Err(_) => {
953 observer.on_event(SynthesisEvent::Admission { accepted: false });
954 Err(EngineError::Busy)
955 }
956 }
957 }
958
959 fn run_stage<R, F>(
960 &self,
961 stage: EngineStage,
962 budget: Duration,
963 cancellation: &CancellationToken,
964 observer: &dyn SynthesisObserver,
965 work: F,
966 ) -> Result<R, EngineError>
967 where
968 R: Send + 'static,
969 F: FnOnce(CancellationToken) -> Result<R, EngineError> + Send + 'static,
970 {
971 cancellation.checkpoint().inspect_err(|_| {
972 observer.on_event(SynthesisEvent::Health {
973 event: HealthEvent::Cancelled,
974 });
975 })?;
976 observer.on_event(SynthesisEvent::StageStarted { stage });
977 let started = Instant::now();
978 let (sender, receiver) = mpsc::sync_channel(1);
979 let stage_cancellation = cancellation.clone();
980 let task_cancellation = cancellation.clone();
981 let task = self
982 .runtime
983 .spawn_blocking(move || {
984 let result = task_cancellation
985 .checkpoint()
986 .and_then(|()| work(task_cancellation));
987 let _ignored_if_timed_out = sender.send(result);
988 })
989 .ok_or_else(|| EngineError::Runtime("blocking pool was not configured".to_owned()))?;
990
991 match receiver.recv_timeout(budget) {
992 Ok(result) => {
993 let result = result?;
994 observer.on_event(SynthesisEvent::StageFinished {
995 stage,
996 elapsed: started.elapsed(),
997 });
998 Ok(result)
999 }
1000 Err(RecvTimeoutError::Timeout) => {
1001 stage_cancellation.cancel();
1002 task.cancel();
1003 observer.on_event(SynthesisEvent::Health {
1004 event: HealthEvent::BudgetExceeded,
1005 });
1006 Err(EngineError::BudgetExceeded(stage))
1007 }
1008 Err(RecvTimeoutError::Disconnected) => Err(EngineError::Runtime(
1009 "blocking stage disconnected before producing a result".to_owned(),
1010 )),
1011 }
1012 }
1013}
1014
1015struct SynthesisAdmission<'a> {
1016 engine: &'a TtsEngine,
1017}
1018
1019impl Drop for SynthesisAdmission<'_> {
1020 fn drop(&mut self) {
1021 self.engine.synthesis_active.store(false, Ordering::Release);
1022 }
1023}
1024
1025#[cfg(test)]
1026mod tests {
1027 use super::*;
1028 use std::sync::Mutex;
1029
1030 #[derive(Default)]
1031 struct RecordingObserver {
1032 events: Mutex<Vec<SynthesisEvent>>,
1033 }
1034
1035 impl RecordingObserver {
1036 fn events(&self) -> Vec<SynthesisEvent> {
1037 self.events
1038 .lock()
1039 .expect("test observer lock poisoned")
1040 .clone()
1041 }
1042 }
1043
1044 impl SynthesisObserver for RecordingObserver {
1045 fn on_event(&self, event: SynthesisEvent) {
1046 self.events
1047 .lock()
1048 .expect("test observer lock poisoned")
1049 .push(event);
1050 }
1051 }
1052
1053 fn engine_with_budget(budget: Duration) -> TtsEngine {
1054 TtsEngine::new(EngineConfig {
1055 synthesis_stage_budget: budget,
1056 ..EngineConfig::default()
1057 })
1058 .expect("test engine builds")
1059 }
1060
1061 fn engine_with_frame_budget(startup: Duration, per_frame: Duration) -> TtsEngine {
1063 TtsEngine::new(EngineConfig {
1064 synthesis_stage_budget: startup,
1065 synthesis_frame_budget: per_frame,
1066 ..EngineConfig::default()
1067 })
1068 .expect("test engine builds")
1069 }
1070
1071 struct PacedFrameGenerator {
1079 remaining: usize,
1080 per_frame: Duration,
1081 stall: Option<Duration>,
1082 began: bool,
1083 }
1084
1085 impl FrameGenerator for PacedFrameGenerator {
1086 fn begin_utterance(&mut self, _prepared: &PreparedText) -> Result<(), GenerationError> {
1087 self.began = true;
1088 Ok(())
1089 }
1090
1091 fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError> {
1092 assert!(self.began, "next_frame before begin_utterance");
1093 if self.remaining == 0 {
1094 let Some(stall) = self.stall else {
1095 return Ok(None);
1096 };
1097 thread::sleep(stall);
1098 return Ok(Some(CodeFrame { codes: vec![0; 16] }));
1099 }
1100 self.remaining -= 1;
1101 thread::sleep(self.per_frame);
1102 Ok(Some(CodeFrame { codes: vec![0; 16] }))
1103 }
1104 }
1105
1106 #[test]
1113 fn steady_progress_past_the_startup_grace_is_not_refused_for_being_long() {
1114 let engine = engine_with_frame_budget(Duration::from_millis(50), Duration::from_millis(30));
1115 let observer = RecordingObserver::default();
1116 let mut generator = PacedFrameGenerator {
1117 remaining: 10,
1118 per_frame: Duration::from_millis(20),
1119 stall: None,
1120 began: false,
1121 };
1122
1123 let result = engine
1124 .synthesize(
1125 SynthesisRequest::new(""),
1126 &TestTextPreparer,
1127 &mut generator,
1128 &CancellationToken::new(),
1129 &observer,
1130 )
1131 .expect("a steadily-progressing run must not be refused");
1132
1133 assert_eq!(
1134 result.generated_frames, 10,
1135 "all ten frames must survive; a flat 50 ms ceiling would have cut this at ~2"
1136 );
1137 }
1138
1139 #[test]
1143 fn a_generator_that_stops_progressing_is_still_caught_within_its_earned_deadline() {
1144 let engine = engine_with_frame_budget(Duration::from_millis(50), Duration::from_millis(30));
1145 let observer = RecordingObserver::default();
1146 let mut generator = PacedFrameGenerator {
1147 remaining: 3,
1148 per_frame: Duration::from_millis(1),
1149 stall: Some(Duration::from_millis(400)),
1150 began: false,
1151 };
1152
1153 let started = Instant::now();
1154 let error = engine
1155 .synthesize(
1156 SynthesisRequest::new(""),
1157 &TestTextPreparer,
1158 &mut generator,
1159 &CancellationToken::new(),
1160 &observer,
1161 )
1162 .expect_err("a stalled generator must still be refused");
1163 let elapsed = started.elapsed();
1164
1165 assert_eq!(error, EngineError::BudgetExceeded(EngineStage::Synthesis));
1166 assert!(
1167 observer.events().contains(&SynthesisEvent::Health {
1168 event: HealthEvent::BudgetExceeded,
1169 }),
1170 "the stall must be reported on the health channel, not only as a return value"
1171 );
1172 assert!(
1175 elapsed < Duration::from_millis(2000),
1176 "stall detection took {elapsed:?}; the deadline is not supposed to keep growing while \
1177 no frames are produced"
1178 );
1179 }
1180
1181 #[test]
1183 fn a_zero_frame_budget_is_rejected_rather_than_collapsing_to_a_flat_deadline() {
1184 let built = TtsEngine::new(EngineConfig {
1185 synthesis_frame_budget: Duration::ZERO,
1186 ..EngineConfig::default()
1187 });
1188 assert!(
1189 matches!(built, Err(EngineError::InvalidConfiguration(_))),
1190 "a zero per-frame budget must be refused, not accepted as a flat deadline"
1191 );
1192 }
1193
1194 #[test]
1198 fn an_unoptimized_build_is_granted_a_larger_synthesis_budget() {
1199 let config = EngineConfig::default();
1200 let expected = u32::from(cfg!(debug_assertions)) * (DEBUG_BUILD_SLOWDOWN - 1) + 1;
1201 assert_eq!(
1202 config.synthesis_frame_budget,
1203 DEFAULT_SYNTHESIS_FRAME_BUDGET * expected
1204 );
1205 assert_eq!(
1206 config.synthesis_stage_budget,
1207 DEFAULT_SYNTHESIS_BUDGET * expected
1208 );
1209 assert_eq!(config.enroll_stage_budget, DEFAULT_ENROLL_BUDGET);
1211 }
1212
1213 struct ScriptedFrameGenerator {
1219 remaining: usize,
1220 began: bool,
1221 endless: bool,
1222 polls: usize,
1223 }
1224
1225 impl ScriptedFrameGenerator {
1226 fn emitting(frames: usize) -> Self {
1227 Self {
1228 remaining: frames,
1229 began: false,
1230 endless: false,
1231 polls: 0,
1232 }
1233 }
1234
1235 fn endless() -> Self {
1237 Self {
1238 remaining: 0,
1239 began: false,
1240 endless: true,
1241 polls: 0,
1242 }
1243 }
1244 }
1245
1246 impl FrameGenerator for ScriptedFrameGenerator {
1247 fn begin_utterance(&mut self, _prepared: &PreparedText) -> Result<(), GenerationError> {
1248 self.began = true;
1249 Ok(())
1250 }
1251
1252 fn next_frame(&mut self) -> Result<Option<CodeFrame>, GenerationError> {
1253 assert!(self.began, "next_frame before begin_utterance");
1254 self.polls += 1;
1255 if self.endless {
1256 return Ok(Some(CodeFrame { codes: vec![0; 16] }));
1257 }
1258 if self.remaining == 0 {
1259 return Ok(None);
1260 }
1261 self.remaining -= 1;
1262 Ok(Some(CodeFrame { codes: vec![0; 16] }))
1263 }
1264 }
1265
1266 fn engine_with_frame_ceiling(max_new_tokens: u64) -> TtsEngine {
1268 TtsEngine::new(EngineConfig {
1269 synthesis_stage_budget: Duration::from_secs(5),
1270 admission: admission::AdmissionPolicy {
1271 max_new_tokens,
1272 ..admission::AdmissionPolicy::default()
1273 },
1274 ..EngineConfig::default()
1275 })
1276 .expect("test engine builds")
1277 }
1278
1279 fn admitted_ceiling(observer: &RecordingObserver) -> u64 {
1281 observer
1282 .events()
1283 .into_iter()
1284 .find_map(|event| match event {
1285 SynthesisEvent::ResourceAdmission {
1286 admitted: true,
1287 predicted_max_frames,
1288 ..
1289 } => Some(predicted_max_frames),
1290 _ => None,
1291 })
1292 .expect("an admitted request reports its frame ceiling")
1293 }
1294
1295 struct TestTextPreparer;
1296
1297 impl TextPreparer for TestTextPreparer {
1298 fn prepare(
1299 &self,
1300 _text: &str,
1301 options: &NormalizationOptions,
1302 ) -> Result<PreparedText, TextPreparationError> {
1303 Ok(PreparedText::new(
1304 vec![7, 11],
1305 NormalizationTrace {
1306 mode: options.mode,
1307 unicode_version: "15.1.0".to_owned(),
1308 changes: vec![NormalizationChange {
1309 rule: "unicode_nfc",
1310 before: "caller-owned secret".to_owned(),
1311 after: "caller-owned secret".to_owned(),
1312 }],
1313 },
1314 ))
1315 }
1316 }
1317
1318 #[test]
1324 fn the_decode_loop_stops_on_the_generators_eos_and_polls_exactly_once_past_it() {
1325 let engine = engine_with_frame_ceiling(64);
1326 let observer = RecordingObserver::default();
1327 let mut generator = ScriptedFrameGenerator::emitting(3);
1328
1329 let result = engine
1330 .synthesize(
1331 SynthesisRequest::new(""),
1332 &TestTextPreparer,
1333 &mut generator,
1334 &CancellationToken::new(),
1335 &observer,
1336 )
1337 .expect("scripted pipeline succeeds");
1338
1339 let ceiling = admitted_ceiling(&observer);
1340 assert!(
1341 ceiling > 3,
1342 "ceiling {ceiling} must exceed the 3 emitted frames, or the stop is ambiguous"
1343 );
1344 assert_eq!(result.generated_frames, 3, "EOS bounds the utterance");
1345 assert_eq!(result.code_frames.len(), 3);
1346 assert_eq!(
1347 generator.polls, 4,
1348 "the loop must poll once past the last frame to observe EOS, and then stop"
1349 );
1350 }
1351
1352 #[test]
1357 fn a_generator_that_never_stops_is_truncated_exactly_at_the_admitted_ceiling() {
1358 let engine = engine_with_frame_ceiling(5);
1359 let observer = RecordingObserver::default();
1360 let mut generator = ScriptedFrameGenerator::endless();
1361
1362 let result = engine
1363 .synthesize(
1364 SynthesisRequest::new(""),
1365 &TestTextPreparer,
1366 &mut generator,
1367 &CancellationToken::new(),
1368 &observer,
1369 )
1370 .expect("a ceiling-bound utterance still completes");
1371
1372 let ceiling = admitted_ceiling(&observer);
1373 assert_eq!(
1374 ceiling, 5,
1375 "the policy's max_new_tokens is the ceiling here"
1376 );
1377 assert_eq!(
1378 result.generated_frames, ceiling,
1379 "an endless generator must be cut at the ceiling, not one frame either side"
1380 );
1381 assert_eq!(result.code_frames.len() as u64, ceiling);
1382 assert_eq!(
1383 generator.polls as u64, ceiling,
1384 "once the ceiling is reached the loop must stop asking, not poll a discarded frame"
1385 );
1386 }
1387
1388 #[test]
1390 fn eos_landing_exactly_on_the_ceiling_yields_the_ceiling_frames() {
1391 let engine = engine_with_frame_ceiling(4);
1392 let observer = RecordingObserver::default();
1393 let mut generator = ScriptedFrameGenerator::emitting(4);
1394
1395 let result = engine
1396 .synthesize(
1397 SynthesisRequest::new(""),
1398 &TestTextPreparer,
1399 &mut generator,
1400 &CancellationToken::new(),
1401 &observer,
1402 )
1403 .expect("scripted pipeline succeeds");
1404
1405 assert_eq!(admitted_ceiling(&observer), 4);
1406 assert_eq!(result.generated_frames, 4);
1407 assert_eq!(
1408 generator.polls, 4,
1409 "the ceiling is reached first, so the generator is never asked for a fifth frame"
1410 );
1411 }
1412
1413 #[test]
1414 fn the_decode_loop_drives_the_generator_and_reports_every_frame() {
1415 let engine = engine_with_budget(Duration::from_secs(1));
1416 let cancellation = CancellationToken::new();
1417 let observer = RecordingObserver::default();
1418 let mut generator = ScriptedFrameGenerator::emitting(2);
1419
1420 let result = engine
1421 .synthesize(
1422 SynthesisRequest::new(""),
1423 &TestTextPreparer,
1424 &mut generator,
1425 &cancellation,
1426 &observer,
1427 )
1428 .expect("scripted pipeline succeeds");
1429
1430 assert_eq!(result.generated_frames, 2);
1431 assert_eq!(result.code_frames.len(), 2);
1432 assert!(
1433 result
1434 .code_frames
1435 .iter()
1436 .all(|frame| frame.codes.len() == 16)
1437 );
1438 assert_eq!(result.prepared_token_count, 2);
1439 let events = observer.events();
1440 assert!(
1441 matches!(
1442 events.as_slice(),
1443 [
1444 SynthesisEvent::Admission { accepted: true },
1445 SynthesisEvent::ResourceAdmission { admitted: true, .. },
1447 SynthesisEvent::StageStarted {
1448 stage: EngineStage::Synthesis,
1449 },
1450 SynthesisEvent::FrameProgress { frame: 0 },
1451 SynthesisEvent::FrameProgress { frame: 1 },
1452 SynthesisEvent::StageFinished {
1453 stage: EngineStage::Synthesis,
1454 ..
1455 },
1456 ]
1457 ),
1458 "unexpected event sequence: {events:?}"
1459 );
1460 }
1461
1462 #[test]
1467 fn an_unaffordable_request_is_refused_before_any_stage_runs() {
1468 let mut config = EngineConfig {
1469 synthesis_stage_budget: Duration::from_secs(1),
1470 ..EngineConfig::default()
1471 };
1472 config.admission.budget_bytes = 1;
1474 let engine = TtsEngine::new(config).expect("engine builds");
1475 let cancellation = CancellationToken::new();
1476 let observer = RecordingObserver::default();
1477
1478 let error = engine
1479 .synthesize(
1480 SynthesisRequest::new(""),
1481 &TestTextPreparer,
1482 &mut ScriptedFrameGenerator::emitting(0),
1483 &cancellation,
1484 &observer,
1485 )
1486 .expect_err("an unaffordable request must be refused");
1487
1488 assert!(
1489 matches!(error, EngineError::ResourceAdmission(_)),
1490 "got {error}"
1491 );
1492
1493 let events = observer.events();
1494 assert!(
1495 !events.iter().any(|event| matches!(
1496 event,
1497 SynthesisEvent::StageStarted { .. }
1498 | SynthesisEvent::StageFinished { .. }
1499 | SynthesisEvent::FrameProgress { .. }
1500 )),
1501 "a refused request must not start any stage; got {events:?}"
1502 );
1503 assert!(
1505 events.iter().any(|event| matches!(
1506 event,
1507 SynthesisEvent::ResourceAdmission {
1508 admitted: false,
1509 ..
1510 }
1511 )),
1512 "a capacity refusal must appear in the event stream: {events:?}"
1513 );
1514 }
1515
1516 #[test]
1517 fn the_admission_policy_is_configurable_and_defaults_are_documented() {
1518 let config = EngineConfig::default();
1519 assert_eq!(
1520 config.admission.budget_bytes,
1521 admission::DEFAULT_BUDGET_BYTES
1522 );
1523 assert_eq!(
1524 config.admission.max_new_tokens,
1525 admission::DEFAULT_MAX_NEW_TOKENS
1526 );
1527 let plan = config
1532 .admission
1533 .admit(512)
1534 .expect("the documented default must admit its own worked case");
1535 assert_eq!(
1536 plan.predicted_max_frames,
1537 512 * admission::HEURISTIC_FRAMES_PER_PROMPT_TOKEN
1538 + admission::HEURISTIC_FRAME_HEADROOM
1539 );
1540 assert!(plan.fits());
1541
1542 let mut explicit = config.admission;
1544 explicit.heuristic_eos_backstop = false;
1545 let plan = explicit
1546 .admit(512)
1547 .expect("the documented explicit-cap case must admit");
1548 assert_eq!(plan.predicted_max_frames, admission::DEFAULT_MAX_NEW_TOKENS);
1549 assert!(plan.fits());
1550 }
1551
1552 #[test]
1553 fn cancellation_is_observed_before_the_cpu_stage_starts() {
1554 let engine = engine_with_budget(Duration::from_secs(1));
1555 let cancellation = CancellationToken::new();
1556 cancellation.cancel();
1557 let observer = RecordingObserver::default();
1558
1559 let error = engine
1560 .synthesize(
1561 SynthesisRequest::new("cancelled"),
1562 &TestTextPreparer,
1563 &mut ScriptedFrameGenerator::emitting(0),
1564 &cancellation,
1565 &observer,
1566 )
1567 .expect_err("cancelled request must not run");
1568
1569 assert_eq!(error, EngineError::Cancelled);
1570 assert_eq!(
1571 observer.events(),
1572 vec![
1573 SynthesisEvent::Admission { accepted: true },
1574 SynthesisEvent::Health {
1575 event: HealthEvent::Cancelled,
1576 },
1577 ]
1578 );
1579 }
1580
1581 #[test]
1582 fn stage_budget_cancels_cooperative_cpu_work() {
1583 let engine = engine_with_budget(Duration::from_millis(5));
1584 let cancellation = CancellationToken::new();
1585 let observer = RecordingObserver::default();
1586
1587 let error = engine
1588 .run_stage(
1589 EngineStage::Synthesis,
1590 Duration::from_millis(5),
1591 &cancellation,
1592 &observer,
1593 |token| -> Result<(), EngineError> {
1594 loop {
1595 token.checkpoint()?;
1596 thread::sleep(Duration::from_millis(1));
1597 }
1598 },
1599 )
1600 .expect_err("long stage must time out");
1601
1602 assert_eq!(error, EngineError::BudgetExceeded(EngineStage::Synthesis));
1603 assert!(cancellation.is_cancelled());
1604 assert!(observer.events().contains(&SynthesisEvent::Health {
1605 event: HealthEvent::BudgetExceeded,
1606 }));
1607 }
1608
1609 #[test]
1610 fn pcm_and_events_have_independent_bounded_queues() {
1611 let queues = StreamQueues::new(1).expect("queue config is valid");
1612 let cancellation = CancellationToken::new();
1613 queues
1614 .events
1615 .send(SynthesisEvent::Admission { accepted: true }, &cancellation)
1616 .expect("event queue accepts first event");
1617 queues
1618 .pcm
1619 .send(
1620 PcmPacket {
1621 frame_count: 1,
1622 samples: vec![1, -1],
1623 },
1624 &cancellation,
1625 )
1626 .expect("full event queue cannot block PCM queue");
1627
1628 assert_eq!(
1629 queues
1630 .pcm_receiver
1631 .recv_timeout(Duration::from_millis(10))
1632 .expect("PCM arrives"),
1633 PcmPacket {
1634 frame_count: 1,
1635 samples: vec![1, -1],
1636 }
1637 );
1638 }
1639
1640 #[test]
1641 fn explicit_normalization_trace_is_text_free() {
1642 let engine = engine_with_budget(Duration::from_secs(1));
1643 let observer = RecordingObserver::default();
1644 let request = SynthesisRequest::new("caller-owned secret")
1645 .with_normalization_options(NormalizationOptions {
1646 mode: NormalizationMode::LocaleAware,
1647 ..NormalizationOptions::default()
1648 })
1649 .with_normalization_trace(true);
1650
1651 engine
1652 .synthesize(
1653 request,
1654 &TestTextPreparer,
1655 &mut ScriptedFrameGenerator::emitting(0),
1656 &CancellationToken::new(),
1657 &observer,
1658 )
1659 .expect("explicit trace request succeeds");
1660
1661 let trace = observer
1662 .events()
1663 .into_iter()
1664 .find_map(|event| match event {
1665 SynthesisEvent::TextPrepared {
1666 token_count,
1667 normalization,
1668 } => Some((token_count, normalization)),
1669 _ => None,
1670 })
1671 .expect("explicit request emits a trace summary");
1672 assert_eq!(trace.0, 2);
1673 assert_eq!(trace.1.mode, NormalizationMode::LocaleAware);
1674 assert_eq!(trace.1.unicode_version, "15.1.0");
1675 assert_eq!(trace.1.rules, vec!["unicode_nfc"]);
1676 assert_eq!(trace.1.change_count, 1);
1677 assert!(
1678 !format!("{:?}", trace.1).contains("caller-owned secret"),
1679 "observer trace summaries must never contain sensitive before/after text"
1680 );
1681 }
1682
1683 #[test]
1684 fn a_health_violation_reaches_the_caller_through_the_observer() {
1685 let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
1689 let sink = Arc::clone(&seen);
1690 let observer = move |event: SynthesisEvent| {
1691 if let SynthesisEvent::Health { event } = event {
1692 sink.lock().expect("observer lock").push(event);
1693 }
1694 };
1695
1696 let violation = health::HealthViolation::OutputSilent {
1697 silent_millis: 1_500,
1698 };
1699 observer(SynthesisEvent::Health {
1700 event: HealthEvent::Violation(violation),
1701 });
1702 let demotion = health::HealthViolation::KernelDemoted {
1703 from: health::KernelTier::Optimized("i8mm"),
1704 to: health::KernelTier::Scalar,
1705 };
1706 observer(SynthesisEvent::Health {
1707 event: HealthEvent::Violation(demotion),
1708 });
1709
1710 let events = seen.lock().expect("observer lock").clone();
1711 assert_eq!(events.len(), 2);
1712 assert_eq!(events[0], HealthEvent::Violation(violation));
1713 assert_eq!(events[0].as_str(), "output_silent");
1714 assert!(events[0].invalidates_output());
1716 assert!(!events[1].invalidates_output());
1717 assert_eq!(events[1].as_str(), "kernel_demoted");
1718 }
1719
1720 #[test]
1721 fn many_utterances_without_deadlock_watchdog() {
1722 let (done_sender, done_receiver) = mpsc::sync_channel(1);
1723 let worker = thread::spawn(move || {
1724 let engine = engine_with_budget(Duration::from_secs(1));
1725 let observer = RecordingObserver::default();
1726 for _ in 0..64 {
1727 engine
1728 .synthesize(
1729 SynthesisRequest::new("watchdog"),
1730 &TestTextPreparer,
1731 &mut ScriptedFrameGenerator::emitting(1),
1732 &CancellationToken::new(),
1733 &observer,
1734 )
1735 .expect("empty utterance succeeds");
1736 }
1737 done_sender
1738 .send(())
1739 .expect("watchdog completion receiver lives");
1740 });
1741
1742 done_receiver
1743 .recv_timeout(Duration::from_secs(2))
1744 .expect("many utterances watchdog expired");
1745 worker.join().expect("watchdog worker does not panic");
1746 }
1747}