diff --git a/codex-rs/voice-host/src/audio_track.rs b/codex-rs/voice-host/src/audio_track.rs index 2dde7d7599..1d3de94890 100644 --- a/codex-rs/voice-host/src/audio_track.rs +++ b/codex-rs/voice-host/src/audio_track.rs @@ -14,6 +14,7 @@ use webrtc::media_stream::MediaStreamTrack; use webrtc::media_stream::track_local::static_sample::TrackLocalStaticSample; use webrtc::peer_connection::MediaEngine; +pub(crate) const OPUS_PAYLOAD_TYPE: u8 = 111; pub(crate) const SEND_TIMEOUT: Duration = Duration::from_millis(/*millis*/ 100); pub(crate) struct EncodedAudio { @@ -42,7 +43,7 @@ impl AudioTrack { .register_codec( RTCRtpCodecParameters { rtp_codec: codec.clone(), - payload_type: 111, + payload_type: OPUS_PAYLOAD_TYPE, }, RtpCodecKind::Audio, ) @@ -84,7 +85,7 @@ impl AudioTrack { self.track .write_sample( self.ssrc, - /*payload_type*/ 111, + /*payload_type*/ OPUS_PAYLOAD_TYPE, &Sample { duration, ..Default::default() @@ -99,7 +100,7 @@ impl AudioTrack { self.track .write_sample( self.ssrc, - /*payload_type*/ 111, + /*payload_type*/ OPUS_PAYLOAD_TYPE, &Sample { data: frame.data.into(), duration, diff --git a/codex-rs/voice-host/src/incoming.rs b/codex-rs/voice-host/src/incoming.rs new file mode 100644 index 0000000000..b230d0eda5 --- /dev/null +++ b/codex-rs/voice-host/src/incoming.rs @@ -0,0 +1,169 @@ +//! Takes incoming Opus RTP before the upstream track queue, preserving network arrival time. +//! Packet and byte permits follow owned data through the consumer, not merely through this queue. + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; + +use rtc::interceptor::Interceptor; +use rtc::interceptor::NoopInterceptor; +use rtc::interceptor::Packet; +use rtc::interceptor::StreamInfo; +use rtc::interceptor::TaggedPacket; +use rtc::interceptor::interceptor; +use rtc::sansio; +use rtc::shared::error::Error; +use rtc::shared::marshal::Marshal; +use rtc::shared::marshal::MarshalSize; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; + +const PACKETS: usize = 64; +const BYTES: usize = 2 * 1024 * 1024; +const PACKET_BYTES: usize = 64 * 1024; + +struct State { + epoch: AtomicU64, + failed: AtomicBool, + packets: Arc, + bytes: Arc, +} + +pub(crate) struct ReceivedRtp { + data: Vec, + pub(crate) at: Instant, + epoch: u64, + _packet: OwnedSemaphorePermit, + _bytes: OwnedSemaphorePermit, +} + +impl AsRef<[u8]> for ReceivedRtp { + fn as_ref(&self) -> &[u8] { + &self.data + } +} + +pub(crate) struct Incoming { + state: Arc, + receiver: mpsc::Receiver, +} + +#[derive(Interceptor)] +pub(crate) struct Ingress { + #[next] + next: NoopInterceptor, + state: Arc, + sender: mpsc::Sender, + stream: Option, +} + +impl Incoming { + pub(crate) fn new() -> (Self, Ingress) { + let state = Arc::new(State { + epoch: AtomicU64::new(/*v*/ 1), + failed: AtomicBool::new(false), + packets: Arc::new(Semaphore::new(/*permits*/ PACKETS)), + bytes: Arc::new(Semaphore::new(/*permits*/ BYTES)), + }); + let (sender, receiver) = mpsc::channel(/*buffer*/ PACKETS); + ( + Self { + state: state.clone(), + receiver, + }, + Ingress { + next: NoopInterceptor::new(), + state, + sender, + stream: None, + }, + ) + } + + pub(crate) fn set_suppressed(&self, suppressed: bool) -> Result<(), &'static str> { + let epoch = self.state.epoch.load(Ordering::Acquire); + if (epoch % 2 == 1) != suppressed { + self.state.epoch.store( + epoch.checked_add(1).ok_or("audio epoch exhausted")?, + Ordering::Release, + ); + } + Ok(()) + } + + pub(crate) fn take(&mut self) -> Result, &'static str> { + if self.state.failed.load(Ordering::Acquire) { + return Err("incoming audio failed"); + } + let epoch = self.state.epoch.load(Ordering::Acquire); + for _ in 0..PACKETS { + let Ok(packet) = self.receiver.try_recv() else { + return Ok(None); + }; + if epoch.is_multiple_of(2) && packet.epoch == epoch { + if packet.at.elapsed() > Duration::from_secs(/*secs*/ 1) { + return Err("incoming audio fell behind"); + } + return Ok(Some(packet)); + } + } + Ok(None) + } +} + +#[interceptor] +impl Ingress { + #[overrides] + fn handle_read(&mut self, message: TaggedPacket) -> Result<(), Self::Error> { + let Packet::Rtp(packet) = message.message else { + return Ok(()); + }; + let epoch = self.state.epoch.load(Ordering::Acquire); + if epoch % 2 == 1 || self.state.failed.load(Ordering::Acquire) { + return Ok(()); + } + let size = packet.marshal_size(); + let accepted = (|| { + if size > PACKET_BYTES + || packet.header.payload_type != crate::audio_track::OPUS_PAYLOAD_TYPE + || self.stream.is_some_and(|ssrc| ssrc != packet.header.ssrc) + { + return None; + } + self.stream = Some(packet.header.ssrc); + let packet_permit = self.state.packets.clone().try_acquire_owned().ok()?; + let bytes_permit = self + .state + .bytes + .clone() + .try_acquire_many_owned(size as u32) + .ok()?; + let mut data = vec![0; size]; + if packet.marshal_to(&mut data).ok()? != size { + return None; + } + self.sender + .try_send(ReceivedRtp { + data, + at: message.now, + epoch, + _packet: packet_permit, + _bytes: bytes_permit, + }) + .ok() + })(); + if accepted.is_none() { + self.state.failed.store(true, Ordering::Release); + } + // Consume here: no second copy is retained by the upstream track-event queue. + Ok(()) + } +} + +#[cfg(test)] +#[path = "incoming_tests.rs"] +mod tests; diff --git a/codex-rs/voice-host/src/incoming_tests.rs b/codex-rs/voice-host/src/incoming_tests.rs new file mode 100644 index 0000000000..bc32d83845 --- /dev/null +++ b/codex-rs/voice-host/src/incoming_tests.rs @@ -0,0 +1,100 @@ +use super::*; +use pretty_assertions::assert_eq; +use rtc::rtp::header::Header; +use rtc::sansio::Protocol; + +fn pair() -> (Incoming, Ingress) { + let (incoming, ingress) = Incoming::new(); + incoming.set_suppressed(/*suppressed*/ false).unwrap(); + (incoming, ingress) +} + +fn packet(size: usize, at: Instant) -> TaggedPacket { + TaggedPacket { + now: at, + transport: Default::default(), + message: Packet::Rtp(rtc::rtp::Packet { + header: Header { + version: 2, + payload_type: 111, + ssrc: 7, + ..Default::default() + }, + payload: vec![5; size - 12].into(), + }), + } +} + +#[test] +fn arrival_and_rtp_bytes_survive_the_adapter_without_an_upstream_copy() { + let (mut incoming, mut ingress) = pair(); + let at = Instant::now() - Duration::from_millis(/*millis*/ 20); + let input = packet(/*size*/ 32, at); + let Packet::Rtp(expected) = &input.message else { + panic!() + }; + let expected = expected.marshal().unwrap(); + ingress.handle_read(input).unwrap(); + let received = incoming.take().unwrap().unwrap(); + assert_eq!((received.at, received.as_ref()), (at, expected.as_ref())); + assert!(ingress.poll_read().is_none()); +} + +#[test] +fn consumed_packets_keep_their_budget_until_the_last_owner_drops_them() { + for size in [32, PACKET_BYTES] { + let (mut incoming, mut ingress) = pair(); + let count = PACKETS.min(BYTES / size); + let mut held = Vec::new(); + for _ in 0..count { + ingress.handle_read(packet(size, Instant::now())).unwrap(); + held.push(incoming.take().unwrap().unwrap()); + } + drop(held.pop()); + ingress.handle_read(packet(size, Instant::now())).unwrap(); + held.push(incoming.take().unwrap().unwrap()); + ingress.handle_read(packet(size, Instant::now())).unwrap(); + assert_eq!(incoming.take().err(), Some("incoming audio failed")); + drop(held); + assert_eq!( + ( + incoming.state.packets.available_permits(), + incoming.state.bytes.available_permits() + ), + (PACKETS, BYTES) + ); + } +} + +#[test] +fn suppression_discards_queued_and_in_flight_old_epochs() { + let (mut incoming, mut ingress) = pair(); + ingress + .handle_read(packet(/*size*/ 32, Instant::now())) + .unwrap(); + incoming.set_suppressed(/*suppressed*/ true).unwrap(); + ingress + .handle_read(packet(/*size*/ 32, Instant::now())) + .unwrap(); + incoming.set_suppressed(/*suppressed*/ false).unwrap(); + assert!(incoming.take().unwrap().is_none()); + ingress + .handle_read(packet(/*size*/ 32, Instant::now())) + .unwrap(); + assert!(incoming.take().unwrap().is_some()); +} + +#[test] +fn invalid_stream_and_oversize_packets_fail_without_returning_media() { + for size in [32, PACKET_BYTES + 1] { + let (mut incoming, mut ingress) = pair(); + let mut input = packet(size, Instant::now()); + if size == 32 + && let Packet::Rtp(packet) = &mut input.message + { + packet.header.payload_type = 96; + } + ingress.handle_read(input).unwrap(); + assert_eq!(incoming.take().err(), Some("incoming audio failed")); + } +} diff --git a/codex-rs/voice-host/src/main.rs b/codex-rs/voice-host/src/main.rs index 58fbda71be..580de49418 100644 --- a/codex-rs/voice-host/src/main.rs +++ b/codex-rs/voice-host/src/main.rs @@ -20,6 +20,7 @@ mod audio_track; path = "devices_unavailable.rs" )] mod devices; +mod incoming; mod runtime; mod transport; mod transport_runtime; @@ -104,6 +105,10 @@ fn run( let mut answered = false; let mut devices: Option = None; loop { + if let Some(peer) = &mut transport { + // The next media stage feeds these owned packets into the jitter/decode pipeline. + while peer.incoming.take().map_err(io::Error::other)?.is_some() {} + } let message = if let Some(devices) = &mut devices { let peer = transport .as_mut() @@ -136,6 +141,12 @@ fn run( Message::DevicesOpened {} } Ok(Message::SetAudioControls { controls }) => { + transport + .as_ref() + .ok_or_else(|| io::Error::other("voice peer not started"))? + .incoming + .set_suppressed(controls.speaker_suppressed) + .map_err(io::Error::other)?; devices .as_mut() .ok_or_else(|| io::Error::other("audio devices not open"))? diff --git a/codex-rs/voice-host/src/transport.rs b/codex-rs/voice-host/src/transport.rs index f120a70477..3b0646576d 100644 --- a/codex-rs/voice-host/src/transport.rs +++ b/codex-rs/voice-host/src/transport.rs @@ -60,6 +60,7 @@ impl PeerConnectionEventHandler for Events { } pub(crate) struct Transport { + pub(crate) incoming: crate::incoming::Incoming, pub(crate) audio: crate::audio_track::AudioTrack, connection: Arc, gathered: Arc, @@ -74,6 +75,7 @@ impl Transport { async fn with_runtime(runtime: Arc) -> Result { let (media, audio) = crate::audio_track::AudioTrack::new()?; + let (incoming, ingress) = crate::incoming::Incoming::new(); let gathered = Arc::new(Notify::new()); // Upstream defaults exhaust checks after 1.4s, including checks sent before // a TCP connection exists. Keep probing throughout our negotiation deadline. @@ -84,6 +86,7 @@ impl Transport { settings.set_ice_connection_attempts(Some(check_interval), Some(attempts)); let connection: Arc = Arc::new( PeerConnectionBuilder::new() + .with_interceptor_registry(rtc::interceptor::Registry::from(ingress)) .with_media_engine(media) .with_runtime(runtime) .with_setting_engine(settings) @@ -128,6 +131,7 @@ impl Transport { sender.send_replace(false); }); Ok(Self { + incoming, audio, connection, gathered, diff --git a/codex-rs/voice-host/src/transport_tests.rs b/codex-rs/voice-host/src/transport_tests.rs index 9c9f338b71..20eda40ad9 100644 --- a/codex-rs/voice-host/src/transport_tests.rs +++ b/codex-rs/voice-host/src/transport_tests.rs @@ -55,7 +55,7 @@ async fn check_negotiation(runtime: Arc) { let gathered = Arc::new(Notify::new()); let mut settings = webrtc::peer_connection::SettingEngine::default(); settings.set_lite(/*lite*/ true); - let (media, _) = crate::audio_track::AudioTrack::new().unwrap(); + let (media, mut remote_audio) = crate::audio_track::AudioTrack::new().unwrap(); let builder = PeerConnectionBuilder::new() .with_media_engine(media) .with_setting_engine(settings) @@ -68,6 +68,7 @@ async fn check_negotiation(runtime: Arc) { .build() .await .unwrap(); + remote.add_track(remote_audio.track.clone()).await.unwrap(); let mut local = Transport::with_runtime(runtime.clone()).await.unwrap(); let offer = RTCSessionDescription::offer(local.offer().await.unwrap()).unwrap(); remote.set_remote_description(offer).await.unwrap(); @@ -118,6 +119,26 @@ async fn check_negotiation(runtime: Arc) { (RTCDataChannelState::Open, true) ); channel.send_text("synthetic sideband event").await.unwrap(); + local.incoming.set_suppressed(/*suppressed*/ false).unwrap(); + let sent = std::time::Instant::now(); + remote_audio + .send(crate::audio_track::EncodedAudio { + data: vec![0xf8, 0xff, 0xfe], + at: sent, + }) + .await + .unwrap(); + let received = loop { + if let Some(packet) = local.incoming.take().unwrap() { + break packet; + } + tokio::time::sleep(Duration::from_millis(/*millis*/ 5)).await; + }; + assert!(received.at >= sent); + assert_eq!( + &received.as_ref()[received.as_ref().len() - 3..], + &[0xf8, 0xff, 0xfe] + ); local.close().await.unwrap(); assert!(local.ready.has_changed().is_err()); remote.close().await.unwrap();