Fully configure frame processors when they are used directly on an audio stream (#679)

Author: 1egoman

livekit-rtc/livekit/rtc/audio_stream.py

@@ -65,6 +65,7 @@ def __init__(
         num_channels: int = 1,
         frame_size_ms: int | None = None,
         noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = True,
         **kwargs: Any,
     ) -> None:
         """Initialize an `AudioStream` instance.
@@ -81,6 +82,9 @@ def __init__(
             noise_cancellation (Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]], optional):
                 If noise cancellation is used, pass a `NoiseCancellationOptions` or `FrameProcessor[AudioFrame]` instance
                 created by the noise cancellation module.
+            auto_close_noise_cancellation (bool):
+                When the audio stream closes, should the FrameProcessor's close method be run? If
+                False, then the frame processor can be reused with another AudioStream.
 
         Example:
             ```python
@@ -113,11 +117,13 @@ def __init__(
         self._audio_filter_module: str | None = None
         self._audio_filter_options: dict[str, Any] | None = None
         self._processor: FrameProcessor[AudioFrame] | None = None
+        self._processor_auto_close = True
         if isinstance(noise_cancellation, NoiseCancellationOptions):
             self._audio_filter_module = noise_cancellation.module_id
             self._audio_filter_options = noise_cancellation.options
         elif isinstance(noise_cancellation, FrameProcessor):
             self._processor = noise_cancellation
+            self._processor_auto_close = auto_close_noise_cancellation
 
         self._task = self._loop.create_task(self._run())
         self._task.add_done_callback(task_done_logger)
@@ -132,6 +138,9 @@ def __init__(
         self._ffi_handle = FfiHandle(stream.handle.id)
         self._info = stream.info
 
+        if self._track is not None:
+            self._track._register_audio_stream(self)
+
     @classmethod
     def from_participant(
         cls,
@@ -144,6 +153,7 @@ def from_participant(
         num_channels: int = 1,
         frame_size_ms: int | None = None,
         noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = True,
     ) -> AudioStream:
         """Create an `AudioStream` from a participant's audio track.
 
@@ -179,8 +189,9 @@ def from_participant(
             track=None,  # type: ignore
             sample_rate=sample_rate,
             num_channels=num_channels,
-            noise_cancellation=noise_cancellation,
             frame_size_ms=frame_size_ms,
+            noise_cancellation=noise_cancellation,
+            auto_close_noise_cancellation=auto_close_noise_cancellation,
         )
 
     @classmethod
@@ -194,6 +205,7 @@ def from_track(
         num_channels: int = 1,
         frame_size_ms: int | None = None,
         noise_cancellation: Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]] = None,
+        auto_close_noise_cancellation: bool = False,
     ) -> AudioStream:
         """Create an `AudioStream` from an existing audio track.
 
@@ -203,9 +215,12 @@ def from_track(
             capacity (int, optional): The capacity of the internal frame queue. Defaults to 0 (unbounded).
             sample_rate (int, optional): The sample rate for the audio stream in Hz. Defaults to 48000.
             num_channels (int, optional): The number of audio channels. Defaults to 1.
-            noise_cancellation (Optional[NoiseCancellationOptions], optional):
-                If noise cancellation is used, pass a `NoiseCancellationOptions` instance
+            noise_cancellation (Optional[NoiseCancellationOptions | FrameProcessor[AudioFrame]], optional):
+                If noise cancellation is used, pass a `NoiseCancellationOptions` or `FrameProcessor[AudioFrame]` instance
                 created by the noise cancellation module.
+            auto_close_noise_cancellation (bool):
+                When the audio stream closes, leaves the FrameProcessor in an unclosed state so it
+                can be used with another AudioStream.
 
         Returns:
             AudioStream: An instance of `AudioStream` that can be used to receive audio frames.
@@ -225,8 +240,9 @@ def from_track(
             capacity=capacity,
             sample_rate=sample_rate,
             num_channels=num_channels,
-            noise_cancellation=noise_cancellation,
             frame_size_ms=frame_size_ms,
+            noise_cancellation=noise_cancellation,
+            auto_close_noise_cancellation=auto_close_noise_cancellation,
         )
 
     def __del__(self) -> None:
@@ -303,8 +319,12 @@ async def aclose(self) -> None:
         This method cleans up resources associated with the audio stream and waits for
         any pending operations to complete.
         """
+        if self._track is not None:
+            self._track._unregister_audio_stream(self)
         self._ffi_handle.dispose()
         await self._task
+        if self._processor is not None and self._processor_auto_close:
+            self._processor._close()
 
     def _is_event(self, e: proto_ffi.FfiEvent) -> bool:
         return e.audio_stream_event.stream_handle == self._ffi_handle.handle

livekit-rtc/livekit/rtc/frame_processor.py

@@ -24,8 +24,12 @@ def _on_stream_info_updated(
         publication_sid: str,
     ) -> None: ...
 
+    def _on_stream_info_cleared(self) -> None: ...
+
     def _on_credentials_updated(self, *, token: str, url: str) -> None: ...
 
+    def _on_credentials_cleared(self) -> None: ...
+
     @abstractmethod
     def _process(self, frame: T) -> T: ...
 

livekit-rtc/livekit/rtc/participant.py

@@ -852,6 +852,15 @@ async def unpublish_track(self, track_sid: str) -> None:
             # so when it processed local_track_unpublished first.
             self._track_publications.pop(track_sid, None)
             if publication is not None:
+                # Clear the processor's room context here too: this path races
+                # the local_track_unpublished room event, and whichever loses
+                # the race finds the publication already gone and skips its own
+                # _set_room(None). Calling it from both paths guarantees the
+                # processor is cleared (and the token_refreshed listener
+                # detached) exactly once it matters; _set_room(None) is
+                # idempotent, so a double-clear when this path wins is safe.
+                if publication._track is not None:
+                    publication._track._set_room(None)
                 publication._track = None
             queue.task_done()
         finally:

livekit-rtc/livekit/rtc/room.py

@@ -754,6 +754,8 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None:
             sid = event.local_track_published.track_sid
             lpublication = self.local_participant.track_publications[sid]
             ltrack = lpublication.track
+            if ltrack is not None:
+                ltrack._set_room(self)
             self.emit("local_track_published", lpublication, ltrack)
         elif which == "local_track_unpublished":
             # During teardown the publication may already have been removed
@@ -767,7 +769,21 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None:
             unpublished = self.local_participant._track_publications.get(sid)
             if unpublished is not None:
                 del self.local_participant._track_publications[sid]
+                track = unpublished.track
+                if track is not None:
+                    track._set_room(None)
+                # Emit while `publication.track` is still set, preserving the
+                # pre-existing payload for callbacks. This handler is synchronous
+                # and emit() invokes listeners synchronously, so nulling the track
+                # right after still completes before any other coroutine (e.g.
+                # unpublish_track) can interleave.
                 self.emit("local_track_unpublished", unpublished)
+                # Mirror track_unsubscribed: drop the publication's track
+                # reference. This also makes unpublish_track's own _set_room(None)
+                # a no-op when it loses the race (its `publication._track is not
+                # None` guard short-circuits), avoiding a redundant clear.
+                if track is not None:
+                    unpublished._track = None
             else:
                 logging.debug("local_track_unpublished for untracked publication sid %s", sid)
         elif which == "local_track_republished":
@@ -784,6 +800,15 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None:
                 del self.local_participant._track_publications[previous_sid]
                 republished._info = event.local_track_republished.info
                 self.local_participant._track_publications[republished.sid] = republished
+                if republished.track is not None:
+                    # Keep the local-track invariant (track.sid == publication.sid,
+                    # set at publish_track) intact across republish, then re-push
+                    # metadata so any attached FrameProcessor learns the new
+                    # publication SID / credentials. _set_room with the same room
+                    # is a no-op for the token_refreshed listener but re-fans the
+                    # metadata to every registered AudioStream.
+                    republished.track._info.sid = republished.sid
+                    republished.track._set_room(self)
                 self.emit("local_track_republished", republished, previous_sid)
         elif which == "local_track_subscribed":
             sid = event.local_track_subscribed.track_sid
@@ -809,17 +834,21 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None:
             rpublication._subscribed = True
             if track_info.kind == TrackKind.KIND_VIDEO:
                 remote_video_track = RemoteVideoTrack(owned_track_info)
+                remote_video_track._set_room(self)
                 rpublication._track = remote_video_track
                 self.emit("track_subscribed", remote_video_track, rpublication, rparticipant)
             elif track_info.kind == TrackKind.KIND_AUDIO:
                 remote_audio_track = RemoteAudioTrack(owned_track_info)
+                remote_audio_track._set_room(self)
                 rpublication._track = remote_audio_track
                 self.emit("track_subscribed", remote_audio_track, rpublication, rparticipant)
         elif which == "track_unsubscribed":
             identity = event.track_unsubscribed.participant_identity
             rparticipant = self._remote_participants[identity]
             rpublication = rparticipant.track_publications[event.track_unsubscribed.track_sid]
             rtrack = rpublication.track
+            if rtrack is not None:
+                rtrack._set_room(None)
             rpublication._track = None
             rpublication._subscribed = False
             self.emit("track_unsubscribed", rtrack, rpublication, rparticipant)

livekit-rtc/livekit/rtc/track.py

@@ -12,14 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, List, Union
+from __future__ import annotations
+
+import weakref
+from typing import TYPE_CHECKING, List, Optional, Union
 from ._ffi_client import FfiHandle, FfiClient
 from ._proto import ffi_pb2 as proto_ffi
 from ._proto import track_pb2 as proto_track
 from ._proto import stats_pb2 as proto_stats
 
 if TYPE_CHECKING:
     from .audio_source import AudioSource
+    from .audio_stream import AudioStream
+    from .room import Room
     from .video_source import VideoSource
     from .platform_audio import PlatformAudioSource
 
@@ -28,6 +33,83 @@ class Track:
     def __init__(self, owned_info: proto_track.OwnedTrack):
         self._info = owned_info.info
         self._ffi_handle = FfiHandle(owned_info.handle.id)
+        self._room_ref: Optional[weakref.ref[Room]] = None
+        self._audio_streams: weakref.WeakSet[AudioStream] = weakref.WeakSet()
+
+    def _resolve_room(self) -> Optional[Room]:
+        return self._room_ref() if self._room_ref is not None else None
+
+    def _set_room(self, room: Optional[Room]) -> None:
+        old_room = self._resolve_room()
+        if old_room is None and room is None:
+            # Already roomless — nothing to detach and nothing to re-clear.
+            # Without this guard a second _set_room(None) (e.g. the unpublish /
+            # local_track_unpublished race calling it from both paths) would
+            # re-fire _on_*_cleared on every registered processor.
+            return
+        if old_room is not room:
+            if old_room is not None:
+                old_room.off("token_refreshed", self._on_room_token_refreshed)
+            if room is not None:
+                room.on("token_refreshed", self._on_room_token_refreshed)
+
+        self._room_ref = weakref.ref(room) if room is not None else None
+
+        for stream in self._audio_streams:
+            self._push_processor_metadata_to_stream(stream, room)
+
+    def _on_room_token_refreshed(self) -> None:
+        room = self._resolve_room()
+        if room is None or room._token is None or room._server_url is None:
+            return
+        for stream in self._audio_streams:
+            if not stream._processor:
+                continue
+            stream._processor._on_credentials_updated(token=room._token, url=room._server_url)
+
+    def _push_processor_metadata_to_stream(self, stream: AudioStream, room: Optional[Room]) -> None:
+        if not stream._processor:
+            return
+
+        if room is None:
+            # track left a room - clear processor's room context
+            stream._processor._on_stream_info_cleared()
+            stream._processor._on_credentials_cleared()
+            return
+
+        identity = ""
+        pub_sid = ""
+        track_sid = self.sid
+        if track_sid:
+            for participant in room.remote_participants.values():
+                publication = participant.track_publications.get(track_sid)
+                if publication is not None:
+                    identity, pub_sid = participant.identity, publication.sid
+                    break
+            else:
+                local = room._local_participant
+                if local is not None:
+                    for local_publication in local.track_publications.values():
+                        if local_publication.sid == track_sid:
+                            identity, pub_sid = local.identity, local_publication.sid
+                            break
+
+        stream._processor._on_stream_info_updated(
+            room_name=room.name,
+            participant_identity=identity,
+            publication_sid=pub_sid,
+        )
+        if room._token is not None and room._server_url is not None:
+            stream._processor._on_credentials_updated(token=room._token, url=room._server_url)
+
+    def _register_audio_stream(self, stream: AudioStream) -> None:
+        self._audio_streams.add(stream)
+        room = self._resolve_room()
+        if room is not None:
+            self._push_processor_metadata_to_stream(stream, room)
+
+    def _unregister_audio_stream(self, stream: AudioStream) -> None:
+        self._audio_streams.discard(stream)
 
     @property
     def sid(self) -> str: