diff --git a/src/pipecat/transports/smallwebrtc/connection.py b/src/pipecat/transports/smallwebrtc/connection.py index 420656f85..58537c7a1 100644 --- a/src/pipecat/transports/smallwebrtc/connection.py +++ b/src/pipecat/transports/smallwebrtc/connection.py @@ -95,15 +95,20 @@ class SmallWebRTCTrack: enable/disable control and frame discarding for audio and video streams. """ - def __init__(self, track: MediaStreamTrack): + def __init__(self, receiver): """Initialize the WebRTC track wrapper. Args: - track: The underlying MediaStreamTrack to wrap. - index: The index of the track in the transceiver (0 for mic, 1 for cam, 2 for screen) + receiver: The RemoteStreamTrack receiver instance. """ - self._track = track + self._receiver = receiver + # Configuring the receiver for not consuming the track by default to prevent memory grow + self._receiver._enabled = False + self._track = receiver.track self._enabled = True + self._last_recv_time: float = 0.0 + self._idle_task: Optional[asyncio.Task] = None + self._idle_timeout: float = 2.0 # seconds before discarding old frames def set_enabled(self, enabled: bool) -> None: """Enable or disable the track. @@ -138,13 +143,44 @@ class SmallWebRTCTrack: async def recv(self) -> Optional[Frame]: """Receive the next frame from the track. + Enables the internal receiving state and starts idle watcher. + Returns: The next frame, except for video tracks, where it returns the frame only if the track is enabled, otherwise, returns None. """ + self._receiver._enabled = True + self._last_recv_time = time.time() + + # start idle watcher if not already running + if not self._idle_task or self._idle_task.done(): + self._idle_task = asyncio.create_task(self._idle_watcher()) + if not self._enabled and self._track.kind == "video": return None return await self._track.recv() + async def _idle_watcher(self): + """Disable receiving if idle for more than _idle_timeout.""" + while self._receiver._enabled: + await asyncio.sleep(self._idle_timeout) + idle_duration = time.time() - self._last_recv_time + if idle_duration >= self._idle_timeout: + # discard old frames from the receiver queue to prevent growth + logger.debug( + f"Disabling receiver for {self._track.kind} track to prevent memory growth." + ) + await self.discard_old_frames() + self._receiver._enabled = False + + def stop(self): + """Stop receiving frames from the track.""" + self._receiver._enabled = False + if self._idle_task: + self._idle_task.cancel() + self._idle_task = None + if self._track: + self._track.stop() + def __getattr__(self, name): """Forward attribute access to the underlying track. @@ -526,8 +562,8 @@ class SmallWebRTCConnection(BaseObject): logger.warning("No audio transceiver is available") return None - track = transceivers[AUDIO_TRANSCEIVER_INDEX].receiver.track - audio_track = SmallWebRTCTrack(track) if track else None + receiver = transceivers[AUDIO_TRANSCEIVER_INDEX].receiver + audio_track = SmallWebRTCTrack(receiver) if receiver else None self._track_map[AUDIO_TRANSCEIVER_INDEX] = audio_track return audio_track @@ -548,8 +584,8 @@ class SmallWebRTCConnection(BaseObject): logger.warning("No video transceiver is available") return None - track = transceivers[VIDEO_TRANSCEIVER_INDEX].receiver.track - video_track = SmallWebRTCTrack(track) if track else None + receiver = transceivers[VIDEO_TRANSCEIVER_INDEX].receiver + video_track = SmallWebRTCTrack(receiver) if receiver else None self._track_map[VIDEO_TRANSCEIVER_INDEX] = video_track return video_track @@ -570,8 +606,8 @@ class SmallWebRTCConnection(BaseObject): logger.warning("No screen video transceiver is available") return None - track = transceivers[SCREEN_VIDEO_TRANSCEIVER_INDEX].receiver.track - video_track = SmallWebRTCTrack(track) if track else None + receiver = transceivers[SCREEN_VIDEO_TRANSCEIVER_INDEX].receiver + video_track = SmallWebRTCTrack(receiver) if receiver else None self._track_map[SCREEN_VIDEO_TRANSCEIVER_INDEX] = video_track return video_track