praxis/server/assist/webrtc.py

"""WarmWebRTCManager — shift-bounded warm WebRTC connection (D-067, REQ-IDEATE-08).

The connection opens at shift start, stays warm (keepalive only between turns),
and closes at shift-end. 30s app-level heartbeat (in addition to the
SmallWebRTCTransport's ICE keepalive) prevents NAT timeouts.

Reconnect state machine (REQ-IDEATE-08):
  - connected → (disconnect) → reconnecting (wait 30s for a new offer)
  - reconnecting + new offer within 30s → connected (pipeline rebuilt)
  - reconnecting + no offer within 30s → disconnected
  - The shift is NOT auto-ended on disconnect (the learner can reconnect or
    end explicitly). The 8h auto-end (D-069) still fires on disconnected shifts.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Any

from loguru import logger

log = logging.getLogger(__name__)

_HEARTBEAT_INTERVAL_S = 30
_RECONNECT_WAIT_S = 30


@dataclass
class WarmConnection:
    """One active warm WebRTC connection for an assist shift."""

    connection: Any  # SmallWebRTCConnection
    task: Any  # PipelineTask
    runner: Any  # PipelineRunner
    heartbeat_task: asyncio.Task | None = None
    shift_id: str = ""
    reconnect_state: str = "connected"  # 'connected' | 'reconnecting' | 'disconnected'


class WarmWebRTCManager:
    """Manages warm WebRTC connections for assist shifts (D-067, REQ-IDEATE-08)."""

    def __init__(self) -> None:
        self._connections: dict[str, WarmConnection] = {}

    async def open(
        self, shift_id: str, webrtc_offer: dict, *, context: Any, session: Any | None = None
    ) -> dict:
        """Accept a WebRTC offer, build the assist pipeline, start the heartbeat.

        Returns the WebRTC answer dict ({sdp, type}).
        """
        from pipecat.transports.smallwebrtc.connection import SmallWebRTCConnection

        from server.assist.pipeline import build_assist_pipeline

        connection = SmallWebRTCConnection(
            ice_servers=[{"urls": "stun:stun.l.google.com:19302"}],
        )
        await connection.receive_offer(webrtc_offer)
        await connection.accept()
        answer = connection.get_answer()

        pipeline, task, runner, transport = build_assist_pipeline(
            connection, context=context, session=session
        )
        # Run the pipeline task in the background.
        runner_task = asyncio.create_task(runner.run(task))

        heartbeat = asyncio.create_task(self._heartbeat(shift_id))

        warm = WarmConnection(
            connection=connection,
            task=task,
            runner=runner,
            heartbeat_task=heartbeat,
            shift_id=shift_id,
            reconnect_state="connected",
        )
        self._connections[shift_id] = warm
        logger.info("warm WebRTC opened for shift %s", shift_id)
        return answer

    async def close(self, shift_id: str) -> None:
        """Close the warm connection + cancel the heartbeat."""
        warm = self._connections.pop(shift_id, None)
        if warm is None:
            return
        if warm.heartbeat_task is not None:
            warm.heartbeat_task.cancel()
            try:
                await warm.heartbeat_task
            except asyncio.CancelledError:
                pass
        # The pipeline task is cancelled when the connection closes.
        try:
            await warm.connection.close()
        except Exception:
            pass
        logger.info("warm WebRTC closed for shift %s", shift_id)

    def get(self, shift_id: str) -> WarmConnection | None:
        return self._connections.get(shift_id)

    def get_reconnect_state(self, shift_id: str) -> str:
        """Return 'connected' | 'reconnecting' | 'disconnected' (REQ-IDEATE-08)."""
        warm = self._connections.get(shift_id)
        if warm is None:
            return "disconnected"
        return warm.reconnect_state

    async def _heartbeat(self, shift_id: str) -> None:
        """App-level heartbeat every 30s (D-067 — prevents NAT timeouts)."""
        try:
            while True:
                await asyncio.sleep(_HEARTBEAT_INTERVAL_S)
                warm = self._connections.get(shift_id)
                if warm is None:
                    return
                # The SmallWebRTCTransport's ICE keepalive (15-30s) is the
                # transport-level keepalive; this app-level heartbeat is an
                # additional safety. We send a no-op ping (in a real impl this
                # would be a Pipecat frame; here we just check the connection).
                if not _connection_alive(warm.connection):
                    await self._on_disconnect(shift_id)
                    return
        except asyncio.CancelledError:
            return

    async def _on_disconnect(self, shift_id: str) -> None:
        """Reconnect state machine (REQ-IDEATE-08).

        1. Log the disconnection (timestamp + shift_id + turn count).
        2. Mark the shift 'reconnecting' + wait up to 30s for a new offer.
        3. New offer within 30s → rebuild the pipeline + resume.
        4. No offer within 30s → mark 'disconnected'. The shift is NOT auto-ended
           (the learner can reconnect or end explicitly; the 8h auto-end still fires).
        """
        warm = self._connections.get(shift_id)
        if warm is None:
            return
        warm.reconnect_state = "reconnecting"
        logger.warning(
            "WebRTC disconnect for shift %s — reconnecting (waiting %ds for a new offer)",
            shift_id, _RECONNECT_WAIT_S,
        )
        # Wait for a new offer. In a real impl this would be an event the
        # /api/assist/webrtc endpoint sets when a new offer arrives. For the
        # pilot we wait then transition to 'disconnected' if no offer came.
        await asyncio.sleep(_RECONNECT_WAIT_S)
        warm = self._connections.get(shift_id)
        if warm is None:
            return
        if warm.reconnect_state == "reconnecting":
            # No new offer arrived within 30s → disconnected.
            warm.reconnect_state = "disconnected"
            logger.warning(
                "WebRTC reconnect timed out for shift %s — disconnected (shift NOT auto-ended; 8h auto-end still fires)",
                shift_id,
            )

    async def reconnect(self, shift_id: str, webrtc_offer: dict, *, context: Any, session: Any | None = None) -> dict:
        """Handle a reconnect offer (REQ-IDEATE-08). Rebuilds the pipeline + resumes."""
        warm = self._connections.get(shift_id)
        if warm is None:
            # Shift not in the map — treat as a fresh open.
            return await self.open(shift_id, webrtc_offer, context=context, session=session)
        # Close the old connection + rebuild.
        if warm.heartbeat_task is not None:
            warm.heartbeat_task.cancel()
            try:
                await warm.heartbeat_task
            except asyncio.CancelledError:
                pass
        try:
            await warm.connection.close()
        except Exception:
            pass
        # Rebuild with the new offer.
        answer = await self.open(shift_id, webrtc_offer, context=context, session=session)
        logger.info("WebRTC reconnected for shift %s", shift_id)
        return answer


def _connection_alive(connection: Any) -> bool:
    """Best-effort check that a SmallWebRTCConnection is still alive."""
    try:
        # The SmallWebRTCConnection has a closed/ready state; this is a heuristic.
        return not getattr(connection, "_closed", False)
    except Exception:
        return True


__all__ = ["WarmWebRTCManager", "WarmConnection"]