dv-pipecat-ai 0.0.82.dev857__py3-none-any.whl → 0.0.85.dev837__py3-none-any.whl

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 and monitor queue size."""
+        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 to prevent memory growth
+                logger.debug(
+                    f"Disabling receiver for {self._track.kind} track after {idle_duration:.2f}s idle"
+                )
+                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.
 
@@ -170,11 +206,16 @@ class SmallWebRTCConnection(BaseObject):
     for real-time audio/video communication.
     """
 
-    def __init__(self, ice_servers: Optional[Union[List[str], List[IceServer]]] = None):
+    def __init__(
+        self,
+        ice_servers: Optional[Union[List[str], List[IceServer]]] = None,
+        connection_timeout_secs: int = 60,
+    ):
         """Initialize the WebRTC connection.
 
         Args:
             ice_servers: List of ICE servers as URLs or IceServer objects.
+            connection_timeout_secs: Timeout in seconds for connecting to the peer.
 
         Raises:
             TypeError: If ice_servers contains mixed types or unsupported types.
@@ -195,6 +236,7 @@ class SmallWebRTCConnection(BaseObject):
             VIDEO_TRANSCEIVER_INDEX: self.video_input_track,
             SCREEN_VIDEO_TRANSCEIVER_INDEX: self.screen_video_input_track,
         }
+        self.connection_timeout_secs = connection_timeout_secs
 
         self._initialize()
 
@@ -241,8 +283,8 @@ class SmallWebRTCConnection(BaseObject):
         self._data_channel = None
         self._renegotiation_in_progress = False
         self._last_received_time = None
-        self._message_queue = []
         self._pending_app_messages = []
+        self._connecting_timeout_task = None
 
     def _setup_listeners(self):
         """Set up event listeners for the peer connection."""
@@ -254,10 +296,7 @@ class SmallWebRTCConnection(BaseObject):
             # Flush queued messages once the data channel is open
             @channel.on("open")
             async def on_open():
-                logger.debug("Data channel is open, flushing queued messages")
-                while self._message_queue:
-                    message = self._message_queue.pop(0)
-                    self._data_channel.send(message)
+                logger.debug("Data channel is open!")
 
             @channel.on("message")
             async def on_message(message):
@@ -454,11 +493,15 @@ class SmallWebRTCConnection(BaseObject):
 
     async def _close(self):
         """Close the peer connection and cleanup resources."""
+        for track in self._track_map.values():
+            if track:
+                track.stop()
+        self._track_map.clear()
         if self._pc:
             await self._pc.close()
-        self._message_queue.clear()
         self._pending_app_messages.clear()
         self._track_map = {}
+        self._cancel_monitoring_connecting_state()
 
     def get_answer(self):
         """Get the SDP answer for the current connection.
@@ -476,9 +519,45 @@ class SmallWebRTCConnection(BaseObject):
             "pc_id": self._pc_id,
         }
 
+    def _monitoring_connecting_state(self) -> None:
+        """Start monitoring the peer connection while it is in the *connecting* state.
+
+        This method schedules a timeout task that will automatically close the
+        connection if it remains in the connecting state for more than the specified
+        timeout, default to 60 seconds.
+        """
+        logger.debug("Monitoring connecting state")
+
+        async def timeout_handler():
+            # We will close the connection in case we have remained in the connecting state for over 1 minute
+            await asyncio.sleep(self.connection_timeout_secs)
+            logger.warning("Timeout establishing the connection to the remote peer. Closing.")
+
+            await self._close()
+
+        # Create and store the timeout task
+        self._connecting_timeout_task = asyncio.create_task(timeout_handler())
+
+    def _cancel_monitoring_connecting_state(self) -> None:
+        """Cancel the ongoing connecting-state timeout task, if any.
+
+        This method should be called once the connection has either succeeded or
+        transitioned out of the connecting state. If the timeout task is still
+        pending, it will be canceled and the reference cleared.
+        """
+        if self._connecting_timeout_task and not self._connecting_timeout_task.done():
+            logger.debug("Cancelling the connecting timeout task")
+            self._connecting_timeout_task.cancel()
+        self._connecting_timeout_task = None
+
     async def _handle_new_connection_state(self):
         """Handle changes in the peer connection state."""
         state = self._pc.connectionState
+        if state == "connecting":
+            self._monitoring_connecting_state()
+        else:
+            self._cancel_monitoring_connecting_state()
+
         if state == "connected" and not self._connect_invoked:
             # We are going to wait until the pipeline is ready before triggering the event
             return
@@ -526,8 +605,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 +627,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 +649,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
 
@@ -585,8 +664,8 @@ class SmallWebRTCConnection(BaseObject):
         if self._data_channel and self._data_channel.readyState == "open":
             self._data_channel.send(json_message)
         else:
-            logger.debug("Data channel not ready, queuing message")
-            self._message_queue.append(json_message)
+            # The client might choose never to create a data channel.
+            logger.trace("Data channel not ready, discarding message!")
 
     def ask_to_renegotiate(self):
         """Request renegotiation of the WebRTC connection."""
@@ -610,3 +689,8 @@ class SmallWebRTCConnection(BaseObject):
                 )()
                 if track:
                     track.set_enabled(signalling_message.enabled)
+
+    async def add_ice_candidate(self, candidate):
+        """Handle incoming ICE candidates."""
+        logger.debug(f"Adding remote candidate: {candidate}")
+        await self.pc.addIceCandidate(candidate)

pipecat/transports/smallwebrtc/request_handler.py

@@ -0,0 +1,246 @@
+#
+# Copyright (c) 2024–2025, Daily
+#
+# SPDX-License-Identifier: BSD 2-Clause License
+#
+
+"""SmallWebRTC request handler for managing peer connections.
+
+This module provides a client for handling web requests and managing WebRTC connections.
+"""
+
+import asyncio
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from aiortc.sdp import candidate_from_sdp
+from fastapi import HTTPException
+from loguru import logger
+
+from pipecat.transports.smallwebrtc.connection import IceServer, SmallWebRTCConnection
+
+
+@dataclass
+class SmallWebRTCRequest:
+    """Small WebRTC transport session arguments for the runner.
+
+    Parameters:
+        sdp: The SDP string (Session Description Protocol).
+        type: The type of the SDP, either "offer" or "answer".
+        pc_id: Optional identifier for the peer connection.
+        restart_pc: Optional whether to restart the peer connection.
+        request_data: Optional custom data sent by the customer.
+    """
+
+    sdp: str
+    type: str
+    pc_id: Optional[str] = None
+    restart_pc: Optional[bool] = None
+    request_data: Optional[Any] = None
+
+
+@dataclass
+class IceCandidate:
+    """The remote ice candidate object received from the peer connection.
+
+    Parameters:
+        candidate: The ice candidate patch SDP string (Session Description Protocol).
+        sdp_mid: The SDP mid for the candidate patch.
+        sdp_mline_index: The SDP mline index for the candidate patch.
+    """
+
+    candidate: str
+    sdp_mid: str
+    sdp_mline_index: int
+
+
+@dataclass
+class SmallWebRTCPatchRequest:
+    """Small WebRTC transport session arguments for the runner.
+
+    Parameters:
+        pc_id: Identifier for the peer connection.
+        candidates: A list of ICE candidate patches.
+    """
+
+    pc_id: str
+    candidates: List[IceCandidate]
+
+
+class ConnectionMode(Enum):
+    """Enum defining the connection handling modes."""
+
+    SINGLE = "single"  # Only one active connection allowed
+    MULTIPLE = "multiple"  # Multiple simultaneous connections allowed
+
+
+class SmallWebRTCRequestHandler:
+    """SmallWebRTC request handler for managing peer connections.
+
+    This class is responsible for:
+      - Handling incoming SmallWebRTC requests.
+      - Creating and managing WebRTC peer connections.
+      - Supporting ESP32-specific SDP munging if enabled.
+      - Invoking callbacks for newly initialized connections.
+      - Supporting both single and multiple connection modes.
+    """
+
+    def __init__(
+        self,
+        ice_servers: Optional[List[IceServer]] = None,
+        esp32_mode: bool = False,
+        host: Optional[str] = None,
+        connection_mode: ConnectionMode = ConnectionMode.MULTIPLE,
+    ) -> None:
+        """Initialize a SmallWebRTC request handler.
+
+        Args:
+            ice_servers (Optional[List[IceServer]]): List of ICE servers to use for WebRTC
+                connections.
+            esp32_mode (bool): If True, enables ESP32-specific SDP munging.
+            host (Optional[str]): Host address used for SDP munging in ESP32 mode.
+                Ignored if `esp32_mode` is False.
+            connection_mode (ConnectionMode): Mode of operation for handling connections.
+                SINGLE allows only one active connection, MULTIPLE allows several.
+        """
+        self._ice_servers = ice_servers
+        self._esp32_mode = esp32_mode
+        self._host = host
+        self._connection_mode = connection_mode
+
+        # Store connections by pc_id
+        self._pcs_map: Dict[str, SmallWebRTCConnection] = {}
+
+    def _check_single_connection_constraints(self, pc_id: Optional[str]) -> None:
+        """Check if the connection request satisfies single connection mode constraints.
+
+        Args:
+            pc_id: The peer connection ID from the request
+
+        Raises:
+            HTTPException: If constraints are violated in single connection mode
+        """
+        if self._connection_mode != ConnectionMode.SINGLE:
+            return
+
+        if not self._pcs_map:  # No existing connections
+            return
+
+        # Get the existing connection (should be only one in single mode)
+        existing_connection = next(iter(self._pcs_map.values()))
+
+        if existing_connection.pc_id != pc_id and pc_id:
+            logger.warning(
+                f"Connection pc_id mismatch: existing={existing_connection.pc_id}, received={pc_id}"
+            )
+            raise HTTPException(status_code=400, detail="PC ID mismatch with existing connection")
+
+        if not pc_id:
+            logger.warning(
+                "Cannot create new connection: existing connection found but no pc_id received"
+            )
+            raise HTTPException(
+                status_code=400,
+                detail="Cannot create new connection with existing connection active",
+            )
+
+    def update_ice_servers(self, ice_servers: Optional[List[IceServer]] = None):
+        """Update the list of ICE servers used for WebRTC connections."""
+        self._ice_servers = ice_servers
+
+    async def handle_web_request(
+        self,
+        request: SmallWebRTCRequest,
+        webrtc_connection_callback: Callable[[Any], Awaitable[None]],
+    ) -> None:
+        """Handle a SmallWebRTC request and resolve the pending answer.
+
+        This method will:
+          - Reuse an existing WebRTC connection if `pc_id` exists.
+          - Otherwise, create a new `SmallWebRTCConnection`.
+          - Invoke the provided callback with the connection.
+          - Manage ESP32-specific munging if enabled.
+          - Enforce single/multiple connection mode constraints.
+
+        Args:
+            request (SmallWebRTCRequest): The incoming WebRTC request, containing
+                SDP, type, and optionally a `pc_id`.
+            webrtc_connection_callback (Callable[[Any], Awaitable[None]]): An
+                asynchronous callback function that is invoked with the WebRTC connection.
+
+        Raises:
+            HTTPException: If connection mode constraints are violated
+            Exception: Any exception raised during request handling or callback execution
+                will be logged and propagated.
+        """
+        try:
+            pc_id = request.pc_id
+
+            # Check connection mode constraints first
+            self._check_single_connection_constraints(pc_id)
+
+            # After constraints are satisfied, get the existing connection if any
+            existing_connection = self._pcs_map.get(pc_id) if pc_id else None
+
+            if existing_connection:
+                pipecat_connection = existing_connection
+                logger.info(f"Reusing existing connection for pc_id: {pc_id}")
+                await pipecat_connection.renegotiate(
+                    sdp=request.sdp,
+                    type=request.type,
+                    restart_pc=request.restart_pc or False,
+                )
+            else:
+                pipecat_connection = SmallWebRTCConnection(ice_servers=self._ice_servers)
+                await pipecat_connection.initialize(sdp=request.sdp, type=request.type)
+
+                @pipecat_connection.event_handler("closed")
+                async def handle_disconnected(webrtc_connection: SmallWebRTCConnection):
+                    logger.info(f"Discarding peer connection for pc_id: {webrtc_connection.pc_id}")
+                    self._pcs_map.pop(webrtc_connection.pc_id, None)
+
+                # Invoke callback provided in runner arguments
+                try:
+                    await webrtc_connection_callback(pipecat_connection)
+                    logger.debug(
+                        f"webrtc_connection_callback executed successfully for peer: {pipecat_connection.pc_id}"
+                    )
+                except Exception as callback_error:
+                    logger.error(
+                        f"webrtc_connection_callback failed for peer {pipecat_connection.pc_id}: {callback_error}"
+                    )
+
+            answer = pipecat_connection.get_answer()
+
+            if self._esp32_mode:
+                from pipecat.runner.utils import smallwebrtc_sdp_munging
+
+                answer["sdp"] = smallwebrtc_sdp_munging(answer["sdp"], self._host)
+
+            self._pcs_map[answer["pc_id"]] = pipecat_connection
+
+            return answer
+        except Exception as e:
+            logger.error(f"Error processing SmallWebRTC request: {e}")
+            logger.debug(f"SmallWebRTC request details: {request}")
+            raise
+
+    async def handle_patch_request(self, request: SmallWebRTCPatchRequest):
+        """Handle a SmallWebRTC patch candidate request."""
+        peer_connection = self._pcs_map.get(request.pc_id)
+
+        if not peer_connection:
+            raise HTTPException(status_code=404, detail="Peer connection not found")
+
+        for c in request.candidates:
+            candidate = candidate_from_sdp(c.candidate)
+            candidate.sdpMid = c.sdp_mid
+            candidate.sdpMLineIndex = c.sdp_mline_index
+            await peer_connection.add_ice_candidate(candidate)
+
+    async def close(self):
+        """Clear the connection map."""
+        coros = [pc.disconnect() for pc in self._pcs_map.values()]
+        await asyncio.gather(*coros)
+        self._pcs_map.clear()