From 76ab325311aa14c9d3bf55293047afbba1166403 Mon Sep 17 00:00:00 2001
From: cai <cai@nbcai.cc>
Date: Sat, 11 Jul 2026 17:59:02 +0800
Subject: [PATCH] feat: mark reply audio segment boundaries
---
src/main.rs | 210 +++++++++++++++++++++++++++++++++++++++++++++++++++-
1 files changed, 204 insertions(+), 6 deletions(-)
diff --git a/src/main.rs b/src/main.rs
index 62107a9..281351d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,8 @@
mod service;
use std::{
+ borrow::Cow,
+ collections::HashSet,
env, fs,
path::{Path, PathBuf},
sync::{
@@ -32,12 +34,18 @@
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::json;
-use tokio::time::{sleep, sleep_until};
-use tokio::{sync::mpsc::UnboundedReceiver, task::JoinHandle};
+use tokio::time::{sleep, sleep_until, timeout};
+use tokio::{
+ sync::{mpsc, mpsc::UnboundedReceiver, watch},
+ task::JoinHandle,
+};
use tracing::{info, warn};
const USER_AUDIO_SAMPLE_RATE_HZ: u32 = 48_000;
const USER_AUDIO_NUM_CHANNELS: u16 = 1;
+const INBOUND_AUDIO_QUEUE_CAPACITY: usize = 100;
+const INBOUND_AUDIO_DROP_LOG_INTERVAL: u64 = 100;
+const AUDIO_DRAIN_STOP_GRACE: Duration = Duration::from_secs(2);
const DEFAULT_BOT_AUDIO_PROFILE: &str = "pcm-16k";
const LIVEKIT_48K_SAMPLE_RATE_HZ: u32 = 48_000;
const PCM_16K_SAMPLE_RATE_HZ: u32 = 16_000;
@@ -1712,6 +1720,42 @@
.unwrap_or("pcm_s16le")
.trim()
.to_ascii_lowercase();
+ if !matches!(format.as_str(), "pcm_s16le" | "mp3" | "mpeg" | "wav") {
+ return Err(anyhow!("unsupported reply_audio_chunk format {format}"));
+ }
+ match state.reply_chunk_markers.observe(audio_chunk.segment_seq) {
+ ReplyChunkMarker::FirstReply => emit_activity(
+ call_id,
+ trace_id,
+ Some(&turn.turn_id),
+ "helper_first_reply_audio_chunk_received",
+ "ok",
+ None,
+ None,
+ json!({
+ "segmentSeq": audio_chunk.segment_seq,
+ "chunkSeq": audio_chunk.chunk_seq,
+ "format": format.as_str(),
+ "bytes": payload.len(),
+ }),
+ ),
+ ReplyChunkMarker::SegmentFirst => emit_activity(
+ call_id,
+ trace_id,
+ Some(&turn.turn_id),
+ "helper_segment_first_audio_chunk_received",
+ "ok",
+ None,
+ None,
+ json!({
+ "segmentSeq": audio_chunk.segment_seq,
+ "chunkSeq": audio_chunk.chunk_seq,
+ "format": format.as_str(),
+ "bytes": payload.len(),
+ }),
+ ),
+ ReplyChunkMarker::None => {}
+ }
let frames = if format == "pcm_s16le" {
let sample_rate = audio_chunk.sample_rate.unwrap_or(sink.sample_rate_hz);
let channels = audio_chunk.channels.unwrap_or(u32::from(sink.num_channels));
@@ -1803,7 +1847,7 @@
Err(error) => return Err(error).context("failed to decode final stream audio chunk"),
}
} else {
- return Err(anyhow!("unsupported reply_audio_chunk format {format}"));
+ unreachable!("supported encoded format checked above")
};
if frames.is_empty() {
return Ok(0);
@@ -2302,6 +2346,7 @@
encoded_audio_buffer: Vec<u8>,
pcm_stream_decoder: Option<audio::PcmS16leStreamDecoder>,
pcm_stream_network_chunk_count: u64,
+ reply_chunk_markers: ReplyChunkMarkerState,
}
impl Default for RuntimeTurnStreamState {
@@ -2317,7 +2362,37 @@
encoded_audio_buffer: Vec::new(),
pcm_stream_decoder: None,
pcm_stream_network_chunk_count: 0,
+ reply_chunk_markers: ReplyChunkMarkerState::default(),
}
+ }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum ReplyChunkMarker {
+ FirstReply,
+ SegmentFirst,
+ None,
+}
+
+#[derive(Default)]
+struct ReplyChunkMarkerState {
+ first_reply_seen: bool,
+ seen_segments: HashSet<u64>,
+}
+
+impl ReplyChunkMarkerState {
+ fn observe(&mut self, segment_seq: Option<u64>) -> ReplyChunkMarker {
+ let first_for_segment = segment_seq
+ .map(|value| self.seen_segments.insert(value))
+ .unwrap_or(false);
+ if !self.first_reply_seen {
+ self.first_reply_seen = true;
+ return ReplyChunkMarker::FirstReply;
+ }
+ if first_for_segment {
+ return ReplyChunkMarker::SegmentFirst;
+ }
+ ReplyChunkMarker::None
}
}
@@ -2372,6 +2447,8 @@
struct RuntimeTurnStreamAudioChunk {
#[serde(rename = "chunkSeq", alias = "seq")]
chunk_seq: Option<u64>,
+ #[serde(rename = "segmentSeq")]
+ segment_seq: Option<u64>,
format: Option<String>,
#[serde(rename = "sampleRate")]
sample_rate: Option<u32>,
@@ -2469,6 +2546,68 @@
i32::from(USER_AUDIO_NUM_CHANNELS),
);
let started_at = Instant::now();
+ let (frame_tx, mut frame_rx) =
+ mpsc::channel::<DrainedUserAudioFrame>(INBOUND_AUDIO_QUEUE_CAPACITY);
+ let (drain_shutdown_tx, mut drain_shutdown_rx) = watch::channel(false);
+ let drain_call_id = call_id.clone();
+ let drain_trace_id = trace_id.clone();
+ let drain_task = tokio::spawn(async move {
+ let mut received_frame_count: u64 = 0;
+ let mut dropped_frame_count: u64 = 0;
+ loop {
+ tokio::select! {
+ changed = drain_shutdown_rx.changed() => {
+ match changed {
+ Ok(()) if *drain_shutdown_rx.borrow() => break,
+ Ok(()) => {}
+ Err(_) => break,
+ }
+ }
+ maybe_frame = stream.next() => {
+ let Some(frame) = maybe_frame else {
+ break;
+ };
+ received_frame_count = received_frame_count.saturating_add(1);
+ let drained = DrainedUserAudioFrame {
+ frame_index: received_frame_count,
+ captured_elapsed_ms: started_at.elapsed().as_millis() as u64,
+ frame: AudioFrame {
+ data: Cow::Owned(frame.data.as_ref().to_vec()),
+ sample_rate: frame.sample_rate,
+ num_channels: frame.num_channels,
+ samples_per_channel: frame.samples_per_channel,
+ },
+ };
+ match frame_tx.try_send(drained) {
+ Ok(()) => {}
+ Err(mpsc::error::TrySendError::Full(_)) => {
+ dropped_frame_count = dropped_frame_count.saturating_add(1);
+ if dropped_frame_count % INBOUND_AUDIO_DROP_LOG_INTERVAL == 1 {
+ warn!(
+ call_id = %drain_call_id,
+ trace_id = %drain_trace_id,
+ dropped_frame_count,
+ received_frame_count,
+ queue_capacity = INBOUND_AUDIO_QUEUE_CAPACITY,
+ "runtime helper inbound_audio_queue_full_dropping_newest"
+ );
+ }
+ }
+ Err(mpsc::error::TrySendError::Closed(_)) => break,
+ }
+ }
+ }
+ }
+ stream.close();
+ info!(
+ call_id = %drain_call_id,
+ trace_id = %drain_trace_id,
+ received_frame_count,
+ dropped_frame_count,
+ queue_capacity = INBOUND_AUDIO_QUEUE_CAPACITY,
+ "runtime helper user_audio_drain_ended"
+ );
+ });
let mut frame_count: u64 = 0;
let mut sample_count: u64 = 0;
let mut simple_vad = if simple_vad_enabled {
@@ -2478,10 +2617,11 @@
};
let mut realtime_asr_upload: Option<RealtimeAsrUpload> = None;
- while let Some(frame) = stream.next().await {
- frame_count += 1;
+ while let Some(drained) = frame_rx.recv().await {
+ let frame = drained.frame;
+ frame_count = drained.frame_index;
sample_count += u64::from(frame.samples_per_channel) * u64::from(frame.num_channels);
- let elapsed_ms = started_at.elapsed().as_millis() as u64;
+ let elapsed_ms = drained.captured_elapsed_ms;
if frame_count == 1 {
info!(
call_id = %call_id,
@@ -2641,6 +2781,21 @@
if let Some(upload) = realtime_asr_upload.take() {
upload.cancel("stream_end").await;
}
+ let _ = drain_shutdown_tx.send(true);
+ let mut drain_task = drain_task;
+ if timeout(AUDIO_DRAIN_STOP_GRACE, &mut drain_task)
+ .await
+ .is_err()
+ {
+ drain_task.abort();
+ let _ = drain_task.await;
+ warn!(
+ call_id = %call_id,
+ trace_id = %trace_id,
+ stop_grace_ms = AUDIO_DRAIN_STOP_GRACE.as_millis() as u64,
+ "runtime helper user_audio_drain_stop_timeout"
+ );
+ }
info!(
call_id = %call_id,
@@ -2653,6 +2808,12 @@
"runtime helper user_audio_stream_ended"
);
})
+}
+
+struct DrainedUserAudioFrame {
+ frame_index: u64,
+ captured_elapsed_ms: u64,
+ frame: AudioFrame<'static>,
}
async fn finish_realtime_asr_upload(
@@ -3398,3 +3559,40 @@
Ok(())
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::{ReplyChunkMarker, ReplyChunkMarkerState, RuntimeTurnStreamEvent};
+
+ #[test]
+ fn reply_chunk_marker_state_emits_turn_first_once_and_later_segment_first_once() {
+ let mut state = ReplyChunkMarkerState::default();
+
+ assert_eq!(ReplyChunkMarker::FirstReply, state.observe(Some(1)));
+ assert_eq!(ReplyChunkMarker::None, state.observe(Some(1)));
+ assert_eq!(ReplyChunkMarker::SegmentFirst, state.observe(Some(2)));
+ assert_eq!(ReplyChunkMarker::None, state.observe(Some(2)));
+ assert_eq!(ReplyChunkMarker::SegmentFirst, state.observe(Some(3)));
+ }
+
+ #[test]
+ fn reply_chunk_marker_state_without_segment_only_emits_turn_first() {
+ let mut state = ReplyChunkMarkerState::default();
+
+ assert_eq!(ReplyChunkMarker::FirstReply, state.observe(None));
+ assert_eq!(ReplyChunkMarker::None, state.observe(None));
+ }
+
+ #[test]
+ fn runtime_turn_stream_audio_chunk_reads_segment_seq() {
+ let event: RuntimeTurnStreamEvent = serde_json::from_str(
+ r#"{"type":"reply_audio_chunk","audioChunk":{"chunkSeq":4,"segmentSeq":2,"format":"pcm_s16le","payloadBase64":"AA==","last":false}}"#,
+ )
+ .expect("turn stream event");
+
+ assert_eq!(
+ Some(2),
+ event.audio_chunk.and_then(|chunk| chunk.segment_seq)
+ );
+ }
+}
--
Gitblit v1.9.3