asterisk_ari_supervisor.py

"""
src/ari/asterisk_ari_supervisor.py

Manage ARI events and active call sessions.
"""

import asyncio
import json
import websockets
from urllib.parse import urlparse, quote

import aiohttp

from ari.call_session import CallSession
from env_config import EnvConfig
from utils.logger import get_logger

logger = get_logger(name=__name__)


class ARISupervisor:
    """
    Manage ARI events and active call sessions.
    """

    def __init__(self) -> None:
        """
        Initialize the ARI supervisor.
        """
        self.sessions: dict[str, CallSession] = {}
        self._http: aiohttp.ClientSession | None = None
        self._bg_tasks: dict[str, asyncio.Task] = {}
        self._running: bool = False
        self._force_reconnect: bool = False

    async def run(self) -> None:
        """
        Connect to ARI WS, reconnecting on errors, and process events.
        """
        try:
            logger.info("ARI supervisor starting")
            self._running = True

            # create persistent HTTP session
            auth = aiohttp.BasicAuth(
                login=EnvConfig.ARI_USER,
                password=EnvConfig.ARI_PASS
            )
            self._http = aiohttp.ClientSession(auth=auth)
            logger.info("Created persistent HTTP session for ARI")

            await self._run_main_loop()
            
        except Exception as e:
            logger.error("Fatal error in supervisor: %s", e, exc_info=True)
            raise
        finally:
            logger.info("ARI supervisor exiting, cleaning up HTTP session")
            if self._http:
                await self._http.close()

    async def _run_main_loop(self) -> None:
        """
        Main connection loop for ARI WebSocket.
        """
        p = urlparse(EnvConfig.ARI_BASE)
        ws_scheme = "wss" if p.scheme == "https" else "ws"
        netloc = p.netloc if p.netloc else "127.0.0.1:8088"
        app_q = quote(EnvConfig.ARI_APP, safe="")
        user_q = quote(EnvConfig.ARI_USER, safe="")
        pass_q = quote(EnvConfig.ARI_PASS, safe="")

        ws_url = f"{ws_scheme}://{netloc}/ari/events?app={app_q}&api_key={user_q}:{pass_q}"

        reconnect_count = 0
        max_reconnect_delay = 30.0

        while self._running:
            try:
                logger.info("Connecting to ARI WS: %s (attempt %d)", ws_url, reconnect_count + 1)
                async with websockets.connect(ws_url) as ws:
                    logger.info("ARI WS connected successfully")
                    reconnect_count = 0
                    self._force_reconnect = False  # Reset flag on successful connection
                    await self._event_loop(ws)
                    logger.info("ARI event loop exited, will reconnect...")

            except websockets.ConnectionClosed as e:
                logger.info("ARI WS connection closed (code: %s, reason: %s), reconnecting...", e.code, e.reason)
                reconnect_count += 1
                delay = min(1.0 * (2 ** (reconnect_count - 1)), max_reconnect_delay)
                logger.info("Waiting %0.1f seconds before reconnection attempt %d", delay, reconnect_count + 1)
                await asyncio.sleep(delay)

            except Exception as e:
                logger.error("Unexpected error in ARI connection loop: %s", e, exc_info=True)
                reconnect_count += 1
                delay = min(2.0 * (2 ** (reconnect_count - 1)), max_reconnect_delay)
                logger.info("Waiting %0.1f seconds before reconnection attempt %d", delay, reconnect_count + 1)
                await asyncio.sleep(delay)

            if not self._running:
                logger.info("Supervisor stopped, exiting reconnection loop")
                break

            # check if force reconnect was requested
            if self._force_reconnect:
                logger.info("Force reconnect requested, restarting connection loop")
                self._force_reconnect = False  # reset flag
                continue  # skip the sleep and reconnect immediately

        logger.info("Exiting supervisor run loop")

    async def _event_loop(self, ws: websockets.WebSocketClientProtocol) -> None:
        """
        Consume ARI events and manage CallSession instances.
        Args:
            ws: websockets.WebSocketClientProtocol representing the ARI WebSocket connection.
        """
        logger.info("ARI event loop started, waiting for events...")
        while self._running and not self._force_reconnect:
            try:
                raw = await ws.recv()
            except websockets.ConnectionClosed as e:
                logger.info("ARI WS connection closed during event loop (code: %s, reason: %s)", e.code, e.reason)
                break
            except Exception as e:
                logger.error("ARI WS recv error: %s", e)
                break
                
            try:
                msg = json.loads(raw)
            except Exception as e:
                logger.debug("Failed to parse ARI message: %s", e)
                continue

            try:
                t = msg.get("type")
                logger.debug("Received ARI event: %s", t)
                
                if t == "StasisStart":
                    ch = msg.get("channel", {})
                    name = ch.get("name", "")
                    if name.startswith("UnicastRTP/"):
                        logger.debug("Skipping UnicastRTP channel: %s", name)
                        continue
                    chan_id = ch.get("id")
                    if not chan_id or chan_id in self.sessions:
                        logger.debug("Skipping channel %s (no id or already exists)", chan_id)
                        continue

                    logger.info("Creating new CallSession for channel %s", chan_id)
                    session = CallSession(chan_id, self._http)
                    self.sessions[chan_id] = session
                    task = asyncio.create_task(
                        self._start_session_bg(chan_id, session),
                        name=f"start-{chan_id}",
                    )
                    self._bg_tasks[chan_id] = task
                    logger.info("[call %s] session starting", chan_id)

                elif t in ("ChannelHangupRequest", "ChannelDestroyed"):
                    ch = msg.get("channel", {})
                    chan_id = ch.get("id")
                    logger.info("Channel %s %s", chan_id, t)
                    
                    # schedule cleanup in background
                    sess = self.sessions.pop(chan_id, None)
                    tsk = self._bg_tasks.pop(chan_id, None)
                    
                    if sess or tsk:
                        cleanup_task = asyncio.create_task(
                            self._cleanup_session_bg(chan_id, sess, tsk),
                            name=f"cleanup-{chan_id}"
                        )
                        logger.info("[call %s] scheduled background cleanup", chan_id)

                else:
                    logger.debug("Unhandled ARI event type: %s", t)
                    
            except Exception as e:
                logger.error("Unexpected error processing ARI event %s: %s", msg.get("type", "unknown"), e, exc_info=True)

    async def _start_session_bg(self, chan_id: str, session: CallSession) -> None:
        """
        Start a CallSession and clean up if it fails to initialize.
        Args:
            chan_id: str representing the channel ID to start the session for.
            session: CallSession representing the session to start.
        """
        logger.info("[call %s] session starting", chan_id)
        try:
            await session.start()
        except Exception as e:
            logger.error("[call %s] session start failed: %s", chan_id, e, exc_info=True)
            try:
                await session.stop()
            except Exception:
                pass
            self.sessions.pop(chan_id, None)

    async def _cleanup_session_bg(
        self,
        chan_id: str,
        sess: CallSession | None,
        tsk: asyncio.Task | None
    ) -> None:
        """
        Clean up a call session in the background.
        Args:
            chan_id: str representing the channel ID to clean up the session for.
            sess: CallSession representing the session to clean up.
            tsk: asyncio.Task representing the background task to clean up.
        """
        logger.info("[call %s] background cleanup task starting", chan_id)
        try:
            if sess:
                try:
                    await asyncio.wait_for(sess.stop(), timeout=10.0)
                    logger.info("[call %s] session stopped successfully", chan_id)
                except asyncio.TimeoutError:
                    logger.error("[call %s] session stop timed out after 10 seconds", chan_id)
                except Exception as e:
                    logger.error("[call %s] error stopping session: %s", chan_id, e, exc_info=True)

            if tsk:
                try:
                    await asyncio.wait_for(tsk, timeout=5.0)
                except asyncio.TimeoutError:
                    logger.error("[call %s] background task cleanup timed out", chan_id)
                    tsk.cancel()
                except Exception as e:
                    logger.debug("[call %s] background task exception (expected during cleanup): %s", chan_id, e)
                    
            logger.info("[call %s] session stopped/cleaned up", chan_id)
            
        except Exception as e:
            logger.error("[call %s] unexpected error in background cleanup: %s", chan_id, e, exc_info=True)

    def get_active_channels(self) -> list[str]:
        """
        Return a snapshot list of currently active PSTN channel ids.
        Returns:
            list[str]: A list of currently active PSTN channel ids.
        """
        return list(self.sessions.keys())

    def get_status(self) -> dict:
        """
        Return the current status of the supervisor.
        Returns:
            dict[str, int | list[str] | bool | None]: The current status of the supervisor.
        """
        return {
            "running": self._running,
            "http_session_active": self._http is not None,
            "active_sessions": len(self.sessions),
            "active_channels": list(self.sessions.keys()),
            "background_tasks": len(self._bg_tasks)
        }

    async def force_reconnect(self) -> None:
        """
        Force a reconnection to ARI by stopping the current event loop.
        """
        logger.info("Force reconnection requested")
        # set a flag to break out of the event loop and trigger reconnection
        self._force_reconnect = True

    async def stop(self) -> None:
        """
        Stop the supervisor and clean up all resources.
        """
        logger.info("Stopping ARI supervisor")
        self._running = False
        
        # stop all active sessions
        for sess in list(self.sessions.values()):
            try:
                await sess.stop()
            except Exception as e:
                logger.debug("Error stopping session: %s", e)
        self.sessions.clear()
        
        # cancel background tasks
        for tsk in list(self._bg_tasks.values()):
            tsk.cancel()
        for tsk in list(self._bg_tasks.values()):
            try:
                await tsk
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.debug("Error waiting for background task: %s", e)
        self._bg_tasks.clear()
        
        logger.info("ARI supervisor stopped")