AI/ontology_platform/ont_platform/realtime/websocket.py

"""WebSocket 연결 관리 (Phase 8).

기능:
- 클라이언트 연결 관리
- 조직별 격리
- 메시지 브로드캐스트
"""

import logging
from typing import Dict, Set, Optional
from fastapi import WebSocket

logger = logging.getLogger(__name__)


class ConnectionManager:
    """WebSocket 연결 관리."""

    def __init__(self):
        """초기화."""
        # org_id → {WebSocket 객체들}
        self.active_connections: Dict[str, Set[WebSocket]] = {}
        # WebSocket → org_id (역 매핑)
        self.connection_to_org: Dict[WebSocket, str] = {}

    async def connect(self, org_id: str, websocket: WebSocket) -> None:
        """클라이언트 연결.

        Args:
            org_id: 조직 ID
            websocket: WebSocket 연결
        """
        await websocket.accept()

        if org_id not in self.active_connections:
            self.active_connections[org_id] = set()

        self.active_connections[org_id].add(websocket)
        self.connection_to_org[websocket] = org_id

        logger.info(f"WebSocket connected for org {org_id}")

    async def disconnect(self, websocket: WebSocket) -> None:
        """클라이언트 연결 해제.

        Args:
            websocket: WebSocket 연결
        """
        org_id = self.connection_to_org.get(websocket)

        if org_id:
            if org_id in self.active_connections:
                self.active_connections[org_id].discard(websocket)

                if not self.active_connections[org_id]:
                    del self.active_connections[org_id]

            del self.connection_to_org[websocket]

            logger.info(f"WebSocket disconnected for org {org_id}")

    async def broadcast(self, org_id: str, message: dict) -> int:
        """조직의 모든 클라이언트에게 메시지 브로드캐스트.

        Args:
            org_id: 조직 ID
            message: 전송할 메시지

        Returns:
            전송 성공 수
        """
        if org_id not in self.active_connections:
            return 0

        disconnected = set()
        sent_count = 0

        for connection in self.active_connections[org_id]:
            try:
                await connection.send_json(message)
                sent_count += 1
            except Exception as e:
                logger.warning(f"Failed to send message: {e}")
                disconnected.add(connection)

        # 연결 끊긴 클라이언트 제거
        for connection in disconnected:
            await self.disconnect(connection)

        return sent_count

    async def broadcast_all(self, message: dict) -> int:
        """모든 클라이언트에게 메시지 브로드캐스트.

        Args:
            message: 전송할 메시지

        Returns:
            전송 성공 수
        """
        total_sent = 0

        for org_id in list(self.active_connections.keys()):
            sent = await self.broadcast(org_id, message)
            total_sent += sent

        return total_sent

    def get_connection_count(self, org_id: Optional[str] = None) -> int:
        """연결 수 조회.

        Args:
            org_id: 조직 ID (None이면 전체)

        Returns:
            연결 수
        """
        if org_id is None:
            return sum(len(connections) for connections in self.active_connections.values())

        return len(self.active_connections.get(org_id, set()))

    def get_org_ids(self) -> list:
        """활성 조직 ID 리스트 조회.

        Returns:
            조직 ID 리스트
        """
        return list(self.active_connections.keys())