Skip to main content

livekit_datatrack/remote/
manager.rs

1// Copyright 2025 LiveKit, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::{
16    events::*,
17    pipeline::{Pipeline, PipelineOptions},
18    RemoteDataTrack, RemoteTrackInner,
19};
20use crate::{
21    api::{
22        DataTrackFrame, DataTrackInfo, DataTrackSid, DataTrackSubscribeError, InternalError,
23        RemoteDataTrackPipelineOptions,
24    },
25    e2ee::DecryptionProvider,
26    packet::{Handle, Packet},
27};
28use anyhow::{anyhow, Context};
29use bytes::Bytes;
30use std::{
31    collections::{HashMap, HashSet},
32    mem,
33    pin::Pin,
34    sync::{
35        atomic::{AtomicUsize, Ordering},
36        Arc,
37    },
38    task::{Context as TaskContext, Poll},
39};
40use tokio::sync::{broadcast, mpsc, oneshot, watch};
41use tokio_stream::{wrappers::ReceiverStream, Stream};
42
43/// Options for creating a [`Manager`].
44#[derive(Debug)]
45pub struct ManagerOptions {
46    /// Provider to use for decrypting incoming frame payloads.
47    ///
48    /// If none, remote tracks using end-to-end encryption will not be available
49    /// for subscription.
50    ///
51    pub decryption_provider: Option<Arc<dyn DecryptionProvider>>,
52}
53
54/// An actor that manages the state of data tracks published by remote participants.
55pub struct Manager {
56    decryption_provider: Option<Arc<dyn DecryptionProvider>>,
57    event_in_tx: mpsc::Sender<InputEvent>,
58    event_in_rx: mpsc::Receiver<InputEvent>,
59    event_out_tx: mpsc::Sender<OutputEvent>,
60
61    /// Mapping between track SID and descriptor.
62    descriptors: HashMap<DataTrackSid, Descriptor>,
63
64    /// Mapping between subscriber handle and track SID.
65    ///
66    /// This is an index that allows track descriptors to be looked up
67    /// by subscriber handle in O(1) time—necessary for routing incoming packets.
68    ///
69    sub_handles: HashMap<Handle, DataTrackSid>,
70}
71
72impl Manager {
73    /// Creates a new manager.
74    ///
75    /// Returns a tuple containing the following:
76    ///
77    /// - The manager itself to be spawned by the caller (see [`Manager::run`]).
78    /// - Channel for sending [`InputEvent`]s to be processed by the manager.
79    /// - Stream for receiving [`OutputEvent`]s produced by the manager.
80    ///
81    pub fn new(options: ManagerOptions) -> (Self, ManagerInput, ManagerOutput) {
82        let (event_in_tx, event_in_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT);
83        let (event_out_tx, event_out_rx) = mpsc::channel(Self::EVENT_BUFFER_COUNT);
84
85        let event_in = ManagerInput::new(event_in_tx.clone());
86        let manager = Manager {
87            decryption_provider: options.decryption_provider,
88            event_in_tx,
89            event_in_rx,
90            event_out_tx,
91            descriptors: HashMap::default(),
92            sub_handles: HashMap::default(),
93        };
94
95        let event_out = ManagerOutput(ReceiverStream::new(event_out_rx));
96        (manager, event_in, event_out)
97    }
98
99    /// Run the manager task, consuming self.
100    ///
101    /// The manager will continue running until receiving [`InputEvent::Shutdown`].
102    ///
103    pub async fn run(mut self) {
104        log::debug!("Task started");
105        while let Some(event) = self.event_in_rx.recv().await {
106            match event {
107                InputEvent::SubscribeRequest(event) => self.on_subscribe_request(event).await,
108                InputEvent::UnsubscribeRequest(event) => self.on_unsubscribe_request(event).await,
109                InputEvent::SfuPublicationUpdates(event) => {
110                    self.on_sfu_publication_updates(event).await
111                }
112                InputEvent::SfuSubscriberHandles(event) => self.on_sfu_subscriber_handles(event),
113                InputEvent::SetPipelineOptions(event) => self.on_set_pipeline_options(event),
114                InputEvent::PacketReceived(bytes) => self.on_packet_received(bytes),
115                InputEvent::ResendSubscriptionUpdates => {
116                    self.on_resend_subscription_updates().await
117                }
118                InputEvent::Shutdown => break,
119            }
120        }
121        self.shutdown().await;
122        log::debug!("Task ended");
123    }
124
125    async fn on_subscribe_request(&mut self, event: SubscribeRequest) {
126        let Some(descriptor) = self.descriptors.get_mut(&event.sid) else {
127            let error = DataTrackSubscribeError::Internal(
128                anyhow!("Cannot subscribe to unknown track").into(),
129            );
130            _ = event.result_tx.send(Err(error));
131            return;
132        };
133        match &mut descriptor.subscription {
134            SubscriptionState::None => {
135                let update_event = SfuUpdateSubscription { sid: event.sid, subscribe: true };
136                _ = self.event_out_tx.send(update_event.into()).await;
137                descriptor.subscription = SubscriptionState::Pending {
138                    result_txs: vec![event.result_tx],
139                    buffer_size: event.options.buffer_size,
140                };
141                // TODO: schedule timeout internally
142            }
143            SubscriptionState::Pending { result_txs, .. } => {
144                result_txs.push(event.result_tx);
145            }
146            SubscriptionState::Active { frame_tx, .. } => {
147                let frame_rx = frame_tx.subscribe();
148                _ = event.result_tx.send(Ok(frame_rx))
149            }
150        }
151    }
152
153    async fn on_unsubscribe_request(&mut self, event: UnsubscribeRequest) {
154        let Some(descriptor) = self.descriptors.get_mut(&event.sid) else {
155            return;
156        };
157
158        let SubscriptionState::Active { sub_handle, .. } = descriptor.subscription else {
159            log::warn!("Unexpected state");
160            return;
161        };
162        descriptor.subscription = SubscriptionState::None;
163        self.sub_handles.remove(&sub_handle);
164
165        let event = SfuUpdateSubscription { sid: event.sid, subscribe: false };
166        _ = self.event_out_tx.send(event.into()).await;
167    }
168
169    async fn on_sfu_publication_updates(&mut self, event: SfuPublicationUpdates) {
170        if event.updates.is_empty() {
171            return;
172        }
173        let mut participant_to_sids: HashMap<String, HashSet<DataTrackSid>> = HashMap::new();
174
175        // Detect published tracks
176        for (publisher_identity, tracks) in event.updates {
177            let sids_in_update = participant_to_sids.entry(publisher_identity.clone()).or_default();
178            for info in tracks {
179                let sid = info.sid();
180                sids_in_update.insert(sid.clone());
181                if self.descriptors.contains_key(&sid) {
182                    continue;
183                }
184                if self.handle_sid_reassigned(&publisher_identity, &info).await {
185                    continue;
186                }
187                self.handle_track_published(publisher_identity.clone(), info).await;
188            }
189        }
190
191        // Detect unpublished tracks (scoped per publisher in the update)
192        for (publisher_identity, sids_in_update) in &participant_to_sids {
193            let unpublished_sids: Vec<_> = self
194                .descriptors
195                .iter()
196                .filter(|(_, desc)| desc.publisher_identity.as_ref() == publisher_identity)
197                .filter(|(sid, _)| !sids_in_update.contains(*sid))
198                .map(|(sid, _)| sid.clone())
199                .collect();
200            for sid in unpublished_sids {
201                self.handle_track_unpublished(sid).await;
202            }
203        }
204    }
205
206    async fn handle_track_published(&mut self, publisher_identity: String, info: DataTrackInfo) {
207        let sid = info.sid();
208        if self.descriptors.contains_key(&sid) {
209            log::error!("Existing descriptor for track {}", sid);
210            return;
211        }
212        let info = Arc::new(info);
213        let publisher_identity: Arc<str> = publisher_identity.into();
214
215        let (published_tx, published_rx) = watch::channel(true);
216
217        let descriptor = Descriptor {
218            info: info.clone(),
219            publisher_identity: publisher_identity.clone(),
220            published_tx,
221            subscription: SubscriptionState::None,
222            max_partial_frames: Arc::new(AtomicUsize::new(
223                RemoteDataTrackPipelineOptions::default().max_partial_frames(),
224            )),
225        };
226        self.descriptors.insert(sid, descriptor);
227
228        let inner = RemoteTrackInner {
229            published_rx,
230            event_in_tx: self.event_in_tx.downgrade(), // TODO: wrap
231            publisher_identity,
232        };
233        let track = RemoteDataTrack::new(info, inner);
234        _ = self.event_out_tx.send(TrackPublished { track }.into()).await;
235    }
236
237    /// Detects and handles SID reassignment, which occurs when the publisher
238    /// republishes its tracks after a full reconnect.
239    ///
240    /// Returns `true` if an SID reassignment occurred, `false` otherwise.
241    ///
242    async fn handle_sid_reassigned(
243        &mut self,
244        publisher_identity: &str,
245        info: &DataTrackInfo,
246    ) -> bool {
247        // Publisher identity and pub handle are stable across republications.
248        let Some((old_sid, descriptor)) = self.descriptors.iter().find(|(_, desc)| {
249            desc.publisher_identity.as_ref() == publisher_identity
250                && desc.info.pub_handle == info.pub_handle
251        }) else {
252            return false;
253        };
254
255        // Invariant: other than SID, info should not have changed.
256        // TODO: consider refactoring to move SID out of info to allow for direct comparison.
257        let DataTrackInfo { sid: _, pub_handle: _, name, uses_e2ee, schema, frame_encoding } =
258            &*descriptor.info;
259        if *name != info.name
260            || *uses_e2ee != info.uses_e2ee
261            || *schema != info.schema
262            || *frame_encoding != info.frame_encoding
263        {
264            log::warn!("Info mismatch for {}, treating as new publication", old_sid);
265            return false;
266        }
267        let old_sid = old_sid.clone();
268
269        let new_sid = info.sid();
270        log::debug!("SID reassigned: {} -> {}", old_sid, new_sid);
271
272        let Some(descriptor) = self.descriptors.remove(&old_sid) else {
273            return false;
274        };
275        *descriptor.info.sid.write().unwrap() = new_sid.clone();
276
277        match &descriptor.subscription {
278            SubscriptionState::None => {}
279            SubscriptionState::Pending { .. } | SubscriptionState::Active { .. } => {
280                // The SFU does not carry subscriptions across a publisher's full
281                // reconnect; re-request the subscription under the new SID.
282                let event = SfuUpdateSubscription { sid: new_sid.clone(), subscribe: true };
283                _ = self.event_out_tx.send(event.into()).await;
284            }
285        }
286        if let SubscriptionState::Active { sub_handle, .. } = &descriptor.subscription {
287            // Keep the routing index consistent until the SFU assigns a new handle
288            // (see `register_subscriber_handle`).
289            self.sub_handles.insert(*sub_handle, new_sid.clone());
290        }
291        self.descriptors.insert(new_sid, descriptor);
292        true
293    }
294
295    fn on_set_pipeline_options(&mut self, event: SetPipelineOptions) {
296        let Some(descriptor) = self.descriptors.get(&event.sid) else {
297            log::warn!("Unknown track {}, cannot set pipeline options", event.sid);
298            return;
299        };
300        descriptor.max_partial_frames.store(event.options.max_partial_frames(), Ordering::Relaxed);
301    }
302
303    async fn handle_track_unpublished(&mut self, sid: DataTrackSid) {
304        let Some(descriptor) = self.descriptors.remove(&sid) else {
305            log::error!("Unknown track {}", sid);
306            return;
307        };
308        if let SubscriptionState::Active { sub_handle, .. } = descriptor.subscription {
309            self.sub_handles.remove(&sub_handle);
310        };
311        _ = descriptor.published_tx.send(false);
312        _ = self.event_out_tx.send(TrackUnpublished { sid }.into()).await;
313    }
314
315    fn on_sfu_subscriber_handles(&mut self, event: SfuSubscriberHandles) {
316        for (handle, sid) in event.mapping {
317            self.register_subscriber_handle(handle, sid);
318        }
319    }
320
321    fn register_subscriber_handle(&mut self, assigned_handle: Handle, sid: DataTrackSid) {
322        let Some(descriptor) = self.descriptors.get_mut(&sid) else {
323            log::warn!("Unknown track: {}", sid);
324            return;
325        };
326        let (result_txs, buffer_size) = match &mut descriptor.subscription {
327            SubscriptionState::None => {
328                // Handle assigned when there is no pending or active subscription is unexpected.
329                log::warn!("No subscription for {}", sid);
330                return;
331            }
332            SubscriptionState::Active { sub_handle, .. } => {
333                // Update handle for an active subscription. This can occur following a full reconnect.
334                self.sub_handles.remove(sub_handle);
335                *sub_handle = assigned_handle;
336                self.sub_handles.insert(assigned_handle, sid);
337                return;
338            }
339            SubscriptionState::Pending { result_txs, buffer_size } => {
340                // Handle assigned for pending subscription, transition to active.
341                (mem::take(result_txs), *buffer_size)
342            }
343        };
344
345        let (packet_tx, packet_rx) = mpsc::channel(Self::PACKET_BUFFER_COUNT);
346        let (frame_tx, frame_rx) = broadcast::channel(buffer_size);
347
348        let decryption_provider = if descriptor.info.uses_e2ee() {
349            self.decryption_provider.as_ref().map(Arc::clone)
350        } else {
351            None
352        };
353
354        let pipeline_opts = PipelineOptions {
355            info: descriptor.info.clone(),
356            publisher_identity: descriptor.publisher_identity.clone(),
357            decryption_provider,
358            max_partial_frames: descriptor.max_partial_frames.clone(),
359        };
360        let pipeline = Pipeline::new(pipeline_opts);
361
362        let track_task = TrackTask {
363            info: descriptor.info.clone(),
364            pipeline,
365            published_rx: descriptor.published_tx.subscribe(),
366            packet_rx,
367            frame_tx: frame_tx.clone(),
368            event_in_tx: self.event_in_tx.clone(),
369        };
370        let task_handle = livekit_runtime::spawn(track_task.run());
371
372        descriptor.subscription = SubscriptionState::Active {
373            sub_handle: assigned_handle,
374            packet_tx,
375            frame_tx,
376            task_handle,
377        };
378        self.sub_handles.insert(assigned_handle, sid);
379
380        for result_tx in result_txs {
381            _ = result_tx.send(Ok(frame_rx.resubscribe()));
382        }
383    }
384
385    fn on_packet_received(&mut self, bytes: Bytes) {
386        let packet = match Packet::deserialize(bytes) {
387            Ok(packet) => packet,
388            Err(err) => {
389                log::error!("Failed to deserialize packet: {}", err);
390                return;
391            }
392        };
393        let Some(sid) = self.sub_handles.get(&packet.header.track_handle) else {
394            log::warn!("Unknown subscriber handle {}", packet.header.track_handle);
395            return;
396        };
397        let Some(descriptor) = self.descriptors.get(sid) else {
398            log::warn!("Missing descriptor for track {}", sid);
399            return;
400        };
401        let SubscriptionState::Active { packet_tx, .. } = &descriptor.subscription else {
402            log::warn!("Received packet for track {} without subscription", sid);
403            return;
404        };
405        _ = packet_tx
406            .try_send(packet)
407            .inspect_err(|err| log::debug!("Cannot send packet to track pipeline: {}", err));
408    }
409
410    async fn on_resend_subscription_updates(&self) {
411        let update_events =
412            self.descriptors.iter().filter_map(|(sid, descriptor)| match descriptor.subscription {
413                SubscriptionState::None => None,
414                SubscriptionState::Pending { .. } | SubscriptionState::Active { .. } => {
415                    Some(SfuUpdateSubscription { sid: sid.clone(), subscribe: true })
416                }
417            });
418        for event in update_events {
419            _ = self.event_out_tx.send(event.into()).await;
420        }
421    }
422
423    /// Performs cleanup before the task ends.
424    async fn shutdown(self) {
425        for (_, descriptor) in self.descriptors {
426            _ = descriptor.published_tx.send(false);
427            match descriptor.subscription {
428                SubscriptionState::None => {}
429                SubscriptionState::Pending { result_txs, .. } => {
430                    for result_tx in result_txs {
431                        _ = result_tx.send(Err(DataTrackSubscribeError::Disconnected));
432                    }
433                }
434                SubscriptionState::Active { task_handle, .. } => task_handle.await,
435            }
436        }
437    }
438
439    /// Maximum number of incoming packets to buffer per track to be sent
440    /// to the track's pipeline.
441    const PACKET_BUFFER_COUNT: usize = 16;
442
443    /// Maximum number of input and output events to buffer.
444    const EVENT_BUFFER_COUNT: usize = 16;
445}
446
447/// Information and state for a remote data track.
448#[derive(Debug)]
449struct Descriptor {
450    info: Arc<DataTrackInfo>,
451    publisher_identity: Arc<str>,
452    published_tx: watch::Sender<bool>,
453    subscription: SubscriptionState,
454    max_partial_frames: Arc<AtomicUsize>,
455}
456
457#[derive(Debug)]
458enum SubscriptionState {
459    /// Track is not subscribed to.
460    None,
461    /// Track is being subscribed to, waiting for subscriber handle.
462    Pending {
463        /// All currently pending requests to subscribe to the track.
464        result_txs: Vec<oneshot::Sender<SubscribeResult>>,
465        /// Internal frame buffer size to use once active.
466        buffer_size: usize,
467    },
468    /// Track has an active subscription.
469    Active {
470        sub_handle: Handle,
471        packet_tx: mpsc::Sender<Packet>,
472        frame_tx: broadcast::Sender<DataTrackFrame>,
473        task_handle: livekit_runtime::JoinHandle<()>,
474    },
475}
476
477/// Task for an individual data track with an active subscription.
478struct TrackTask {
479    info: Arc<DataTrackInfo>,
480    pipeline: Pipeline,
481    published_rx: watch::Receiver<bool>,
482    packet_rx: mpsc::Receiver<Packet>,
483    frame_tx: broadcast::Sender<DataTrackFrame>,
484    event_in_tx: mpsc::Sender<InputEvent>,
485}
486
487impl TrackTask {
488    async fn run(mut self) {
489        log::debug!("Track task started: name={}", self.info.name);
490
491        let mut is_published = *self.published_rx.borrow();
492        while is_published {
493            tokio::select! {
494                biased;  // State updates take priority
495                _ = self.published_rx.changed() => {
496                    is_published = *self.published_rx.borrow();
497                },
498                _ = self.frame_tx.closed() => {
499                    let event = UnsubscribeRequest { sid: self.info.sid() };
500                    _ = self.event_in_tx.send(event.into()).await;
501                    break;  // No more subscribers
502                },
503                Some(packet) = self.packet_rx.recv() => {
504                    self.receive(packet);
505                },
506                else => break
507            }
508        }
509
510        log::debug!("Track task ended: name={}", self.info.name);
511    }
512
513    fn receive(&mut self, packet: Packet) {
514        let Some(frame) = self.pipeline.process_packet(packet) else { return };
515        _ = self
516            .frame_tx
517            .send(frame)
518            .inspect_err(|err| log::debug!("Cannot send frame to subscribers: {}", err));
519    }
520}
521
522/// Channel for sending [`InputEvent`]s to [`Manager`].
523#[derive(Debug, Clone)]
524pub struct ManagerInput {
525    event_in_tx: mpsc::Sender<InputEvent>,
526    _drop_guard: Arc<DropGuard>,
527}
528
529/// Stream of [`OutputEvent`]s produced by [`Manager`].
530#[derive(Debug)]
531pub struct ManagerOutput(ReceiverStream<OutputEvent>);
532
533impl Stream for ManagerOutput {
534    type Item = OutputEvent;
535
536    fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
537        Pin::new(&mut self.0).poll_next(cx)
538    }
539}
540
541/// Guard that sends shutdown event when the last reference is dropped.
542#[derive(Debug)]
543struct DropGuard {
544    event_in_tx: mpsc::Sender<InputEvent>,
545}
546
547impl Drop for DropGuard {
548    fn drop(&mut self) {
549        _ = self.event_in_tx.try_send(InputEvent::Shutdown);
550    }
551}
552
553impl ManagerInput {
554    fn new(event_in_tx: mpsc::Sender<InputEvent>) -> Self {
555        Self { event_in_tx: event_in_tx.clone(), _drop_guard: DropGuard { event_in_tx }.into() }
556    }
557
558    /// Sends an input event to the manager's task to be processed.
559    pub fn send(&self, event: InputEvent) -> Result<(), InternalError> {
560        Ok(self.event_in_tx.try_send(event).context("Failed to send input event")?)
561    }
562}
563
564#[cfg(test)]
565mod tests {
566    use super::*;
567    use crate::{
568        api::DataTrackSubscribeOptions,
569        e2ee::{DecryptionError, DecryptionProvider, EncryptedPayload},
570        packet::{E2eeExt, Extensions, FrameMarker, Header, Timestamp},
571        utils::testing::expect_event,
572    };
573    use fake::{Fake, Faker};
574    use futures_util::{future::join, StreamExt};
575    use std::{collections::HashMap, sync::RwLock, time::Duration};
576    use test_case::test_case;
577    use tokio::time;
578
579    #[derive(Debug)]
580    struct PrefixStrippingDecryptor;
581
582    impl DecryptionProvider for PrefixStrippingDecryptor {
583        fn decrypt(
584            &self,
585            payload: EncryptedPayload,
586            _sender_identity: String,
587        ) -> Result<Bytes, DecryptionError> {
588            Ok(payload.payload.slice(4..))
589        }
590    }
591
592    #[tokio::test]
593    async fn test_manager_task_shutdown() {
594        let options = ManagerOptions { decryption_provider: None };
595        let (manager, input, _) = Manager::new(options);
596
597        let join_handle = livekit_runtime::spawn(manager.run());
598        _ = input.send(InputEvent::Shutdown);
599
600        time::timeout(Duration::from_secs(1), join_handle).await.unwrap();
601    }
602
603    #[test_case(true; "via_unpublish")]
604    #[test_case(false; "via_unsubscribe")]
605    #[tokio::test]
606    async fn test_track_task_shutdown(via_unpublish: bool) {
607        let mut info: DataTrackInfo = Faker.fake();
608        info.uses_e2ee = false;
609
610        let info = Arc::new(info);
611        let sid = info.sid();
612        let publisher_identity: Arc<str> = Faker.fake::<String>().into();
613
614        let pipeline_opts = PipelineOptions {
615            info: info.clone(),
616            publisher_identity,
617            decryption_provider: None,
618            max_partial_frames: Arc::new(AtomicUsize::new(
619                RemoteDataTrackPipelineOptions::default().max_partial_frames(),
620            )),
621        };
622        let pipeline = Pipeline::new(pipeline_opts);
623
624        let (published_tx, published_rx) = watch::channel(true);
625        let (_packet_tx, packet_rx) = mpsc::channel(4);
626        let (frame_tx, frame_rx) = broadcast::channel(4);
627        let (event_in_tx, mut event_in_rx) = mpsc::channel(4);
628
629        let task =
630            TrackTask { info: info, pipeline, published_rx, packet_rx, frame_tx, event_in_tx };
631        let task_handle = livekit_runtime::spawn(task.run());
632
633        let trigger_shutdown = async {
634            if via_unpublish {
635                // Simulates SFU publication update
636                published_tx.send(false).unwrap();
637                return;
638            }
639            // Simulates all subscribers dropped
640            mem::drop(frame_rx);
641
642            while let Some(event) = event_in_rx.recv().await {
643                let InputEvent::UnsubscribeRequest(event) = event else {
644                    panic!("Unexpected event type");
645                };
646                assert_eq!(event.sid, sid);
647                return;
648            }
649            panic!("Did not receive unsubscribe");
650        };
651        time::timeout(Duration::from_secs(1), join(task_handle, trigger_shutdown)).await.unwrap();
652    }
653
654    #[tokio::test]
655    async fn test_subscribe() {
656        let publisher_identity: String = Faker.fake();
657        let track_name: String = Faker.fake();
658        let track_sid: DataTrackSid = Faker.fake();
659        let sub_handle: Handle = Faker.fake();
660
661        let options = ManagerOptions { decryption_provider: None };
662        let (manager, input, mut output) = Manager::new(options);
663        livekit_runtime::spawn(manager.run());
664
665        // Simulate track published
666        let event = SfuPublicationUpdates {
667            updates: HashMap::from([(
668                publisher_identity.clone(),
669                vec![DataTrackInfo {
670                    sid: RwLock::new(track_sid.clone()).into(),
671                    pub_handle: Faker.fake(), // Pub handle
672                    name: track_name.clone(),
673                    uses_e2ee: false,
674                    schema: None,
675                    frame_encoding: None,
676                }],
677            )]),
678        };
679        _ = input.send(event.into());
680
681        let wait_for_track = async {
682            while let Some(event) = output.next().await {
683                match event {
684                    OutputEvent::TrackPublished(track) => return track,
685                    _ => continue,
686                }
687            }
688            panic!("No track received");
689        };
690
691        let track = wait_for_track.await.track;
692        assert!(track.is_published());
693        assert_eq!(track.info().name, track_name);
694        assert_eq!(track.info().sid(), track_sid);
695        assert_eq!(track.publisher_identity(), publisher_identity);
696
697        let simulate_subscriber_handles = async {
698            while let Some(event) = output.next().await {
699                match event {
700                    OutputEvent::SfuUpdateSubscription(event) => {
701                        assert!(event.subscribe);
702                        assert_eq!(event.sid, track_sid);
703                        time::sleep(Duration::from_millis(20)).await;
704
705                        // Simulate SFU reply
706                        let event = SfuSubscriberHandles {
707                            mapping: HashMap::from([(sub_handle, track_sid.clone())]),
708                        };
709                        _ = input.send(event.into());
710                    }
711                    _ => {}
712                }
713            }
714        };
715
716        time::timeout(Duration::from_secs(1), async {
717            tokio::select! {
718                _ = simulate_subscriber_handles => {}
719                _ = track.subscribe() => {}
720            }
721        })
722        .await
723        .unwrap();
724    }
725
726    #[tokio::test]
727    async fn test_track_publication_add_and_remove() {
728        let options = ManagerOptions { decryption_provider: None };
729        let (manager, input, mut output) = Manager::new(options);
730        livekit_runtime::spawn(manager.run());
731
732        let track_sid: DataTrackSid = Faker.fake();
733        let info = DataTrackInfo {
734            sid: RwLock::new(track_sid.clone()).into(),
735            pub_handle: Faker.fake(),
736            name: "test".into(),
737            uses_e2ee: false,
738            schema: None,
739            frame_encoding: None,
740        };
741
742        // Simulate track published
743        let event =
744            SfuPublicationUpdates { updates: HashMap::from([("identity1".into(), vec![info])]) };
745        input.send(event.into()).unwrap();
746
747        let track = expect_event!(output, OutputEvent::TrackPublished).track;
748        assert_eq!(track.info().sid(), track_sid);
749        assert_eq!(track.info().name, "test");
750        assert!(track.is_published());
751
752        // Simulate track unpublished
753        let event =
754            SfuPublicationUpdates { updates: HashMap::from([("identity1".into(), vec![])]) };
755        input.send(event.into()).unwrap();
756
757        time::timeout(Duration::from_secs(1), track.wait_for_unpublish()).await.unwrap();
758        assert!(!track.is_published());
759
760        let event = expect_event!(output, OutputEvent::TrackUnpublished);
761        assert_eq!(event.sid, track_sid);
762    }
763
764    #[tokio::test]
765    async fn test_sfu_publication_updates_idempotent() {
766        let options = ManagerOptions { decryption_provider: None };
767        let (manager, input, mut output) = Manager::new(options);
768        livekit_runtime::spawn(manager.run());
769
770        let track_sid: DataTrackSid = Faker.fake();
771        let info = DataTrackInfo {
772            sid: RwLock::new(track_sid.clone()).into(),
773            pub_handle: Faker.fake(),
774            name: "test".into(),
775            uses_e2ee: false,
776            schema: None,
777            frame_encoding: None,
778        };
779
780        // Simulate three identical publication updates
781        for _ in 0..3 {
782            let event = SfuPublicationUpdates {
783                updates: HashMap::from([("identity1".into(), vec![info.clone()])]),
784            };
785            input.send(event.into()).unwrap();
786        }
787
788        expect_event!(output, OutputEvent::TrackPublished);
789
790        // Drain remaining events; no second TrackAvailable should appear
791        input.send(InputEvent::Shutdown).unwrap();
792        while let Some(event) = output.next().await {
793            assert!(!matches!(event, OutputEvent::TrackPublished(_)));
794        }
795    }
796
797    #[tokio::test]
798    async fn test_sid_reassignment_does_not_republish() {
799        let options = ManagerOptions { decryption_provider: None };
800        let (manager, input, mut output) = Manager::new(options);
801        livekit_runtime::spawn(manager.run());
802
803        let pub_handle: Handle = Faker.fake();
804        let old_sid: DataTrackSid = Faker.fake();
805        let new_sid: DataTrackSid = Faker.fake();
806
807        // Simulate track published
808        let info = DataTrackInfo {
809            sid: RwLock::new(old_sid.clone()).into(),
810            pub_handle,
811            name: "test".into(),
812            uses_e2ee: false,
813            schema: None,
814            frame_encoding: None,
815        };
816        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
817        input.send(event.into()).unwrap();
818
819        let track = expect_event!(output, OutputEvent::TrackPublished).track;
820        assert_eq!(track.info().sid(), old_sid);
821
822        // Simulate publisher full reconnect: same track, new SID
823        let info = DataTrackInfo {
824            sid: RwLock::new(new_sid.clone()).into(),
825            pub_handle,
826            name: "test".into(),
827            uses_e2ee: false,
828            schema: None,
829            frame_encoding: None,
830        };
831        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
832        input.send(event.into()).unwrap();
833
834        // Drain remaining events; no publish/unpublish should appear
835        input.send(InputEvent::Shutdown).unwrap();
836        while let Some(event) = output.next().await {
837            assert!(!matches!(
838                event,
839                OutputEvent::TrackPublished(_) | OutputEvent::TrackUnpublished(_)
840            ));
841        }
842        assert_eq!(track.info().sid(), new_sid);
843    }
844
845    #[tokio::test]
846    async fn test_sid_reassignment_resubscribes_active_subscription() {
847        let options = ManagerOptions { decryption_provider: None };
848        let (manager, input, mut output) = Manager::new(options);
849        livekit_runtime::spawn(manager.run());
850
851        let pub_handle: Handle = Faker.fake();
852        let old_sid: DataTrackSid = Faker.fake();
853        let new_sid: DataTrackSid = Faker.fake();
854        let old_sub_handle: Handle = Faker.fake();
855        let new_sub_handle: Handle = Faker.fake();
856
857        // Simulate track published
858        let info = DataTrackInfo {
859            sid: RwLock::new(old_sid.clone()).into(),
860            pub_handle,
861            name: "test".into(),
862            uses_e2ee: false,
863            schema: None,
864            frame_encoding: None,
865        };
866        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
867        input.send(event.into()).unwrap();
868        let track = expect_event!(output, OutputEvent::TrackPublished).track;
869
870        // Subscribe to the track
871        let (result_tx, result_rx) = oneshot::channel();
872        let event = SubscribeRequest {
873            sid: old_sid.clone(),
874            options: DataTrackSubscribeOptions::default(),
875            result_tx,
876        };
877        input.send(event.into()).unwrap();
878        expect_event!(output, OutputEvent::SfuUpdateSubscription);
879
880        // Simulate SFU assigning subscriber handle
881        let event = SfuSubscriberHandles { mapping: HashMap::from([(old_sub_handle, old_sid)]) };
882        input.send(event.into()).unwrap();
883
884        let mut frame_rx =
885            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
886
887        // Simulate publisher full reconnect: same track, new SID
888        let info = DataTrackInfo {
889            sid: RwLock::new(new_sid.clone()).into(),
890            pub_handle,
891            name: "test".into(),
892            uses_e2ee: false,
893            schema: None,
894            frame_encoding: None,
895        };
896        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
897        input.send(event.into()).unwrap();
898
899        // Manager should re-subscribe under the new SID
900        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
901        assert!(event.subscribe);
902        assert_eq!(event.sid, new_sid);
903        assert_eq!(track.info().sid(), new_sid);
904        assert!(track.is_published());
905
906        // Simulate SFU assigning a new subscriber handle
907        let event = SfuSubscriberHandles { mapping: HashMap::from([(new_sub_handle, new_sid)]) };
908        input.send(event.into()).unwrap();
909
910        // Frames received on the new handle reach the existing subscriber
911        let packet = Packet {
912            header: Header {
913                marker: FrameMarker::Single,
914                track_handle: new_sub_handle,
915                extensions: Extensions::default(),
916                ..Faker.fake()
917            },
918            payload: Bytes::from_static(&[1, 2, 3, 4, 5]),
919        };
920        input.send(InputEvent::PacketReceived(packet.serialize())).unwrap();
921
922        let frame = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
923        assert_eq!(frame.payload.as_ref(), &[1, 2, 3, 4, 5]);
924    }
925
926    #[tokio::test]
927    async fn test_subscribe_receives_frame() {
928        let options = ManagerOptions { decryption_provider: None };
929        let (manager, input, mut output) = Manager::new(options);
930        livekit_runtime::spawn(manager.run());
931
932        let track_sid: DataTrackSid = Faker.fake();
933        let sub_handle: Handle = Faker.fake();
934        let info = DataTrackInfo {
935            sid: RwLock::new(track_sid.clone()).into(),
936            pub_handle: Faker.fake(),
937            name: "test".into(),
938            uses_e2ee: false,
939            schema: None,
940            frame_encoding: None,
941        };
942
943        // Simulate track published
944        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
945        input.send(event.into()).unwrap();
946        expect_event!(output, OutputEvent::TrackPublished);
947
948        // Subscribe to the track
949        let (result_tx, result_rx) = oneshot::channel();
950        let event = SubscribeRequest {
951            sid: track_sid.clone(),
952            options: DataTrackSubscribeOptions::default(),
953            result_tx,
954        };
955        input.send(event.into()).unwrap();
956
957        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
958        assert!(event.subscribe);
959        assert_eq!(event.sid, track_sid);
960
961        // Simulate SFU assigning subscriber handle
962        let event = SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid)]) };
963        input.send(event.into()).unwrap();
964
965        let mut frame_rx =
966            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
967
968        // Simulate receiving a single-frame packet
969        let packet = Packet {
970            header: Header {
971                marker: FrameMarker::Single,
972                track_handle: sub_handle,
973                sequence: 0,
974                frame_number: 0,
975                timestamp: Timestamp::from_ticks(0),
976                extensions: Extensions::default(),
977            },
978            payload: Bytes::from_static(&[1, 2, 3, 4, 5]),
979        };
980        input.send(InputEvent::PacketReceived(packet.serialize())).unwrap();
981
982        let frame = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
983        assert_eq!(frame.payload.as_ref(), &[1, 2, 3, 4, 5]);
984    }
985
986    #[tokio::test]
987    async fn test_subscribe_with_e2ee() {
988        let options =
989            ManagerOptions { decryption_provider: Some(Arc::new(PrefixStrippingDecryptor)) };
990        let (manager, input, mut output) = Manager::new(options);
991        livekit_runtime::spawn(manager.run());
992
993        let track_sid: DataTrackSid = Faker.fake();
994        let sub_handle: Handle = Faker.fake();
995        let info = DataTrackInfo {
996            sid: RwLock::new(track_sid.clone()).into(),
997            pub_handle: Faker.fake(),
998            name: "test".into(),
999            uses_e2ee: true,
1000            schema: None,
1001            frame_encoding: None,
1002        };
1003
1004        // Simulate track published (with e2ee)
1005        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1006        input.send(event.into()).unwrap();
1007        expect_event!(output, OutputEvent::TrackPublished);
1008
1009        // Subscribe to the track
1010        let (result_tx, result_rx) = oneshot::channel();
1011        let event = SubscribeRequest {
1012            sid: track_sid.clone(),
1013            options: DataTrackSubscribeOptions::default(),
1014            result_tx,
1015        };
1016        input.send(event.into()).unwrap();
1017
1018        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1019        assert!(event.subscribe);
1020
1021        // Simulate SFU assigning subscriber handle
1022        let event = SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid)]) };
1023        input.send(event.into()).unwrap();
1024
1025        let mut frame_rx =
1026            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1027
1028        // Simulate receiving an encrypted single-frame packet
1029        let packet = Packet {
1030            header: Header {
1031                marker: FrameMarker::Single,
1032                track_handle: sub_handle,
1033                sequence: 0,
1034                frame_number: 0,
1035                timestamp: Timestamp::from_ticks(0),
1036                extensions: Extensions {
1037                    e2ee: Some(E2eeExt { key_index: 0, iv: [0; 12] }),
1038                    ..Default::default()
1039                },
1040            },
1041            payload: Bytes::from_static(&[0xDE, 0xAD, 0xBE, 0xEF, 1, 2, 3, 4, 5]),
1042        };
1043        input.send(InputEvent::PacketReceived(packet.serialize())).unwrap();
1044
1045        // Payload should have fake encryption prefix stripped by decryptor
1046        let frame = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1047        assert_eq!(frame.payload.as_ref(), &[1, 2, 3, 4, 5]);
1048    }
1049
1050    #[tokio::test]
1051    async fn test_subscribe_fan_out_to_multiple_subscribers() {
1052        let options = ManagerOptions { decryption_provider: None };
1053        let (manager, input, mut output) = Manager::new(options);
1054        livekit_runtime::spawn(manager.run());
1055
1056        let track_sid: DataTrackSid = Faker.fake();
1057        let sub_handle: Handle = Faker.fake();
1058        let info = DataTrackInfo {
1059            sid: RwLock::new(track_sid.clone()).into(),
1060            pub_handle: Faker.fake(),
1061            name: "test".into(),
1062            uses_e2ee: false,
1063            schema: None,
1064            frame_encoding: None,
1065        };
1066
1067        // Simulate track published
1068        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1069        input.send(event.into()).unwrap();
1070        expect_event!(output, OutputEvent::TrackPublished);
1071
1072        // First subscriber triggers SFU interaction
1073        let (result_tx1, result_rx1) = oneshot::channel();
1074        let event = SubscribeRequest {
1075            sid: track_sid.clone(),
1076            options: DataTrackSubscribeOptions::default(),
1077            result_tx: result_tx1,
1078        };
1079        input.send(event.into()).unwrap();
1080
1081        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1082        assert!(event.subscribe);
1083
1084        // Simulate SFU assigning subscriber handle
1085        let event =
1086            SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid.clone())]) };
1087        input.send(event.into()).unwrap();
1088
1089        let mut rx1 =
1090            time::timeout(Duration::from_secs(1), result_rx1).await.unwrap().unwrap().unwrap();
1091
1092        // Additional subscribers attach directly (no further SFU interaction)
1093        let (result_tx2, result_rx2) = oneshot::channel();
1094        let event = SubscribeRequest {
1095            sid: track_sid.clone(),
1096            options: DataTrackSubscribeOptions::default(),
1097            result_tx: result_tx2,
1098        };
1099        input.send(event.into()).unwrap();
1100        let mut rx2 = result_rx2.await.unwrap().unwrap();
1101
1102        let (result_tx3, result_rx3) = oneshot::channel();
1103        let event = SubscribeRequest {
1104            sid: track_sid.clone(),
1105            options: DataTrackSubscribeOptions::default(),
1106            result_tx: result_tx3,
1107        };
1108        input.send(event.into()).unwrap();
1109        let mut rx3 = result_rx3.await.unwrap().unwrap();
1110
1111        // Simulate receiving a single-frame packet
1112        let packet = Packet {
1113            header: Header {
1114                marker: FrameMarker::Single,
1115                track_handle: sub_handle,
1116                sequence: 0,
1117                frame_number: 0,
1118                timestamp: Timestamp::from_ticks(0),
1119                extensions: Extensions::default(),
1120            },
1121            payload: Bytes::from_static(&[1, 2, 3, 4, 5]),
1122        };
1123        input.send(InputEvent::PacketReceived(packet.serialize())).unwrap();
1124
1125        // All subscribers should receive the same frame
1126        for rx in [&mut rx1, &mut rx2, &mut rx3] {
1127            let frame = time::timeout(Duration::from_secs(1), rx.recv()).await.unwrap().unwrap();
1128            assert_eq!(frame.payload.as_ref(), &[1, 2, 3, 4, 5]);
1129        }
1130    }
1131
1132    #[tokio::test]
1133    async fn test_subscribe_unknown_track_fails() {
1134        let options = ManagerOptions { decryption_provider: None };
1135        let (manager, input, _) = Manager::new(options);
1136        livekit_runtime::spawn(manager.run());
1137
1138        // Subscribe to a track that was never published
1139        let (result_tx, result_rx) = oneshot::channel();
1140        let event = SubscribeRequest {
1141            sid: Faker.fake(),
1142            options: DataTrackSubscribeOptions::default(),
1143            result_tx,
1144        };
1145        input.send(event.into()).unwrap();
1146
1147        let result = result_rx.await.unwrap();
1148        assert!(result.is_err());
1149    }
1150
1151    #[tokio::test]
1152    async fn test_unpublish_terminates_pending_subscription() {
1153        let options = ManagerOptions { decryption_provider: None };
1154        let (manager, input, mut output) = Manager::new(options);
1155        livekit_runtime::spawn(manager.run());
1156
1157        let track_sid: DataTrackSid = Faker.fake();
1158        let info = DataTrackInfo {
1159            sid: RwLock::new(track_sid.clone()).into(),
1160            pub_handle: Faker.fake(),
1161            name: "test".into(),
1162            uses_e2ee: false,
1163            schema: None,
1164            frame_encoding: None,
1165        };
1166
1167        // Simulate track published
1168        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1169        input.send(event.into()).unwrap();
1170        expect_event!(output, OutputEvent::TrackPublished);
1171
1172        // Subscribe (enters Pending state)
1173        let (result_tx, result_rx) = oneshot::channel();
1174        let event = SubscribeRequest {
1175            sid: track_sid.clone(),
1176            options: DataTrackSubscribeOptions::default(),
1177            result_tx,
1178        };
1179        input.send(event.into()).unwrap();
1180
1181        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1182        assert!(event.subscribe);
1183
1184        // Simulate track unpublished before SFU assigns a handle
1185        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![])]) };
1186        input.send(event.into()).unwrap();
1187
1188        let result = time::timeout(Duration::from_secs(1), result_rx).await.unwrap();
1189        assert!(result.is_err());
1190
1191        let event = expect_event!(output, OutputEvent::TrackUnpublished);
1192        assert_eq!(event.sid, track_sid);
1193    }
1194
1195    #[tokio::test]
1196    async fn test_unpublish_terminates_active_subscription() {
1197        let options = ManagerOptions { decryption_provider: None };
1198        let (manager, input, mut output) = Manager::new(options);
1199        livekit_runtime::spawn(manager.run());
1200
1201        let track_sid: DataTrackSid = Faker.fake();
1202        let sub_handle: Handle = Faker.fake();
1203        let info = DataTrackInfo {
1204            sid: RwLock::new(track_sid.clone()).into(),
1205            pub_handle: Faker.fake(),
1206            name: "test".into(),
1207            uses_e2ee: false,
1208            schema: None,
1209            frame_encoding: None,
1210        };
1211
1212        // Simulate track published
1213        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1214        input.send(event.into()).unwrap();
1215        expect_event!(output, OutputEvent::TrackPublished);
1216
1217        // Subscribe to the track
1218        let (result_tx, result_rx) = oneshot::channel();
1219        let event = SubscribeRequest {
1220            sid: track_sid.clone(),
1221            options: DataTrackSubscribeOptions::default(),
1222            result_tx,
1223        };
1224        input.send(event.into()).unwrap();
1225
1226        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1227        assert!(event.subscribe);
1228
1229        // Simulate SFU assigning subscriber handle
1230        let event =
1231            SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid.clone())]) };
1232        input.send(event.into()).unwrap();
1233
1234        let mut frame_rx =
1235            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1236
1237        // Simulate track unpublished while subscription is active
1238        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![])]) };
1239        input.send(event.into()).unwrap();
1240
1241        let result = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap();
1242        assert!(result.is_err());
1243
1244        let event = expect_event!(output, OutputEvent::TrackUnpublished);
1245        assert_eq!(event.sid, track_sid);
1246    }
1247
1248    #[tokio::test]
1249    async fn test_all_subscribers_dropped_terminates_sfu_subscription() {
1250        let options = ManagerOptions { decryption_provider: None };
1251        let (manager, input, mut output) = Manager::new(options);
1252        livekit_runtime::spawn(manager.run());
1253
1254        let track_sid: DataTrackSid = Faker.fake();
1255        let sub_handle: Handle = Faker.fake();
1256        let info = DataTrackInfo {
1257            sid: RwLock::new(track_sid.clone()).into(),
1258            pub_handle: Faker.fake(),
1259            name: "test".into(),
1260            uses_e2ee: false,
1261            schema: None,
1262            frame_encoding: None,
1263        };
1264
1265        // Simulate track published
1266        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1267        input.send(event.into()).unwrap();
1268        expect_event!(output, OutputEvent::TrackPublished);
1269
1270        // Subscribe to the track
1271        let (result_tx, result_rx) = oneshot::channel();
1272        let event = SubscribeRequest {
1273            sid: track_sid.clone(),
1274            options: DataTrackSubscribeOptions::default(),
1275            result_tx,
1276        };
1277        input.send(event.into()).unwrap();
1278
1279        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1280        assert!(event.subscribe);
1281
1282        // Simulate SFU assigning subscriber handle
1283        let event =
1284            SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid.clone())]) };
1285        input.send(event.into()).unwrap();
1286
1287        let frame_rx =
1288            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1289
1290        // Drop the only subscriber
1291        drop(frame_rx);
1292
1293        // Manager should request SFU to unsubscribe
1294        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1295        assert!(!event.subscribe);
1296        assert_eq!(event.sid, track_sid);
1297    }
1298
1299    /// Should depacketize multiple interleaved partial frames when
1300    /// `max_partial_frames` is set before subscribe.
1301    #[tokio::test]
1302    async fn test_max_partial_frames_set_before_subscribe() {
1303        let options = ManagerOptions { decryption_provider: None };
1304        let (manager, input, mut output) = Manager::new(options);
1305        livekit_runtime::spawn(manager.run());
1306
1307        let mut info: DataTrackInfo = Faker.fake();
1308        info.uses_e2ee = false;
1309        let track_sid = info.sid();
1310        let sub_handle: Handle = Faker.fake();
1311
1312        // Simulate track published
1313        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1314        input.send(event.into()).unwrap();
1315        let track = expect_event!(output, OutputEvent::TrackPublished).track;
1316
1317        // Configure the track BEFORE any subscribe.
1318        track.set_pipeline_options(
1319            RemoteDataTrackPipelineOptions::default().with_max_partial_frames(3),
1320        );
1321
1322        // Subscribe to the track
1323        let (result_tx, result_rx) = oneshot::channel();
1324        let event = SubscribeRequest {
1325            sid: track_sid.clone(),
1326            options: DataTrackSubscribeOptions::default(),
1327            result_tx,
1328        };
1329        input.send(event.into()).unwrap();
1330
1331        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1332        assert!(event.subscribe);
1333
1334        // Simulate SFU assigning subscriber handle
1335        let event = SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid)]) };
1336        input.send(event.into()).unwrap();
1337
1338        let mut frame_rx =
1339            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1340
1341        // Two interleaved partial frames: Start(1), Start(2), Final(1), Final(2). With the default
1342        // max_partial_frames=1 frame 1 would be evicted by frame 2; with max_partial_frames=3 both
1343        // frames coexist and emerge.
1344        push_interleaved_two_frame_pair(
1345            &input,
1346            sub_handle,
1347            1,
1348            0,
1349            [&[0xA1], &[0xA2]],
1350            2,
1351            100,
1352            [&[0xB1], &[0xB2]],
1353        );
1354
1355        let first = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1356        assert_eq!(first.payload.as_ref(), &[0xA1, 0xA2]);
1357
1358        let second = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1359        assert_eq!(second.payload.as_ref(), &[0xB1, 0xB2]);
1360    }
1361
1362    /// Should pick up `max_partial_frames` live on an already-active subscription.
1363    #[tokio::test]
1364    async fn test_max_partial_frames_set_live() {
1365        let options = ManagerOptions { decryption_provider: None };
1366        let (manager, input, mut output) = Manager::new(options);
1367        livekit_runtime::spawn(manager.run());
1368
1369        let mut info: DataTrackInfo = Faker.fake();
1370        info.uses_e2ee = false;
1371        let track_sid = info.sid();
1372        let sub_handle: Handle = Faker.fake();
1373
1374        // Simulate track published
1375        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1376        input.send(event.into()).unwrap();
1377        let track = expect_event!(output, OutputEvent::TrackPublished).track;
1378
1379        // Subscribe to the track
1380        let (result_tx, result_rx) = oneshot::channel();
1381        let event = SubscribeRequest {
1382            sid: track_sid.clone(),
1383            options: DataTrackSubscribeOptions::default(),
1384            result_tx,
1385        };
1386        input.send(event.into()).unwrap();
1387
1388        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1389        assert!(event.subscribe);
1390
1391        // Simulate SFU assigning subscriber handle
1392        let event = SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid)]) };
1393        input.send(event.into()).unwrap();
1394
1395        let mut frame_rx =
1396            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1397
1398        // Subscription is now active; flip the cap on the live pipeline.
1399        track.set_pipeline_options(
1400            RemoteDataTrackPipelineOptions::default().with_max_partial_frames(3),
1401        );
1402
1403        push_interleaved_two_frame_pair(
1404            &input,
1405            sub_handle,
1406            1,
1407            0,
1408            [&[0xA1], &[0xA2]],
1409            2,
1410            100,
1411            [&[0xB1], &[0xB2]],
1412        );
1413
1414        let first = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1415        assert_eq!(first.payload.as_ref(), &[0xA1, 0xA2]);
1416
1417        let second = time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1418        assert_eq!(second.payload.as_ref(), &[0xB1, 0xB2]);
1419    }
1420
1421    /// Should drop the older partial frame by default (no `max_partial_frames` set).
1422    #[tokio::test]
1423    async fn test_default_drops_older_partial_frame() {
1424        let options = ManagerOptions { decryption_provider: None };
1425        let (manager, input, mut output) = Manager::new(options);
1426        livekit_runtime::spawn(manager.run());
1427
1428        let mut info: DataTrackInfo = Faker.fake();
1429        info.uses_e2ee = false;
1430        let track_sid = info.sid();
1431        let sub_handle: Handle = Faker.fake();
1432
1433        // Simulate track published
1434        let event = SfuPublicationUpdates { updates: HashMap::from([("id".into(), vec![info])]) };
1435        input.send(event.into()).unwrap();
1436        expect_event!(output, OutputEvent::TrackPublished);
1437
1438        // Subscribe to the track
1439        let (result_tx, result_rx) = oneshot::channel();
1440        let event = SubscribeRequest {
1441            sid: track_sid.clone(),
1442            options: DataTrackSubscribeOptions::default(),
1443            result_tx,
1444        };
1445        input.send(event.into()).unwrap();
1446
1447        let event = expect_event!(output, OutputEvent::SfuUpdateSubscription);
1448        assert!(event.subscribe);
1449
1450        // Simulate SFU assigning subscriber handle
1451        let event = SfuSubscriberHandles { mapping: HashMap::from([(sub_handle, track_sid)]) };
1452        input.send(event.into()).unwrap();
1453
1454        let mut frame_rx =
1455            time::timeout(Duration::from_secs(1), result_rx).await.unwrap().unwrap().unwrap();
1456
1457        // Default cap of 1: Start(2) evicts Start(1), so Final(1) is unknown and only frame 2
1458        // makes it through.
1459        push_interleaved_two_frame_pair(
1460            &input,
1461            sub_handle,
1462            1,
1463            0,
1464            [&[0xA1], &[0xA2]],
1465            2,
1466            100,
1467            [&[0xB1], &[0xB2]],
1468        );
1469
1470        let only_frame =
1471            time::timeout(Duration::from_secs(1), frame_rx.recv()).await.unwrap().unwrap();
1472        assert_eq!(only_frame.payload.as_ref(), &[0xB1, 0xB2]);
1473    }
1474
1475    /// Pushes Start(frame1), Start(frame2), Final(frame1), Final(frame2) packets through the
1476    /// manager to exercise the depacketizer's concurrent-partial-frame handling.
1477    fn push_interleaved_two_frame_pair(
1478        input: &ManagerInput,
1479        handle: Handle,
1480        frame_one_number: u16,
1481        frame_one_start_sequence: u16,
1482        frame_one_payloads: [&[u8]; 2],
1483        frame_two_number: u16,
1484        frame_two_start_sequence: u16,
1485        frame_two_payloads: [&[u8]; 2],
1486    ) {
1487        let push = |frame_number: u16, sequence: u16, marker: FrameMarker, payload: &[u8]| {
1488            let mut packet: Packet = Faker.fake();
1489            packet.header.marker = marker;
1490            packet.header.track_handle = handle;
1491            packet.header.frame_number = frame_number;
1492            packet.header.sequence = sequence;
1493            packet.header.extensions.e2ee = None;
1494            packet.payload = Bytes::copy_from_slice(payload);
1495            input.send(InputEvent::PacketReceived(packet.serialize())).unwrap();
1496        };
1497        push(frame_one_number, frame_one_start_sequence, FrameMarker::Start, frame_one_payloads[0]);
1498        push(frame_two_number, frame_two_start_sequence, FrameMarker::Start, frame_two_payloads[0]);
1499        push(
1500            frame_one_number,
1501            frame_one_start_sequence + 1,
1502            FrameMarker::Final,
1503            frame_one_payloads[1],
1504        );
1505        push(
1506            frame_two_number,
1507            frame_two_start_sequence + 1,
1508            FrameMarker::Final,
1509            frame_two_payloads[1],
1510        );
1511    }
1512}