Preventing memory growth if we are not consuming the track.

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.
+            receiver: The RemoteStreamTrack receiver instance.
-            index: The index of the track in the transceiver (0 for mic, 1 for cam, 2 for screen)
         """
-        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
+        receiver = transceivers[AUDIO_TRANSCEIVER_INDEX].receiver
-        audio_track = SmallWebRTCTrack(track) if track else None
+        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
+        receiver = transceivers[VIDEO_TRANSCEIVER_INDEX].receiver
-        video_track = SmallWebRTCTrack(track) if track else None
+        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
+        receiver = transceivers[SCREEN_VIDEO_TRANSCEIVER_INDEX].receiver
-        video_track = SmallWebRTCTrack(track) if track else None
+        video_track = SmallWebRTCTrack(receiver) if receiver else None
         self._track_map[SCREEN_VIDEO_TRANSCEIVER_INDEX] = video_track
         return video_track