"""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"]