LangBot/src/langbot/pkg/platform/sources/websocket_manager.py

"""WebSocket连接管理器 - 管理多个并发WebSocket连接"""

import asyncio
import logging
import typing
import uuid
from datetime import datetime

import pydantic

logger = logging.getLogger(__name__)


class WebSocketConnection(pydantic.BaseModel):
    """单个WebSocket连接"""

    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)

    connection_id: str = pydantic.Field(default_factory=lambda: str(uuid.uuid4()))
    """连接唯一ID"""

    pipeline_uuid: str
    """关联的流水线UUID"""

    session_type: str  # 'person' or 'group'
    """会话类型"""

    websocket: typing.Any = pydantic.Field(exclude=True)
    """WebSocket连接对象 (quart.websocket)"""

    created_at: datetime = pydantic.Field(default_factory=datetime.now)
    """连接创建时间"""

    last_active: datetime = pydantic.Field(default_factory=datetime.now)
    """最后活跃时间"""

    send_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
    """发送消息队列"""

    is_active: bool = True
    """连接是否活跃"""

    metadata: dict = pydantic.Field(default_factory=dict)
    """连接元数据（可存储额外信息）"""


class WebSocketConnectionManager:
    """WebSocket连接管理器 - 支持多连接并发"""

    def __init__(self):
        self.connections: dict[str, WebSocketConnection] = {}
        """所有活跃连接 {connection_id: connection}"""

        self.pipeline_connections: dict[str, set[str]] = {}
        """流水线到连接的映射 {pipeline_uuid: {connection_id, ...}}"""

        self.session_connections: dict[str, set[str]] = {}
        """会话类型到连接的映射 {session_type: {connection_id, ...}}"""

        self._lock = asyncio.Lock()
        """线程锁，保护并发访问"""

    async def add_connection(
        self,
        websocket: typing.Any,
        pipeline_uuid: str,
        session_type: str,
        metadata: dict = None,
    ) -> WebSocketConnection:
        """添加新的WebSocket连接"""
        async with self._lock:
            connection = WebSocketConnection(
                pipeline_uuid=pipeline_uuid,
                session_type=session_type,
                websocket=websocket,
                metadata=metadata or {},
            )

            self.connections[connection.connection_id] = connection

            # 更新流水线映射
            if pipeline_uuid not in self.pipeline_connections:
                self.pipeline_connections[pipeline_uuid] = set()
            self.pipeline_connections[pipeline_uuid].add(connection.connection_id)

            # 更新会话类型映射
            if session_type not in self.session_connections:
                self.session_connections[session_type] = set()
            self.session_connections[session_type].add(connection.connection_id)

            logger.debug(
                f'WebSocket connection established: {connection.connection_id} '
                f'(pipeline={pipeline_uuid}, session_type={session_type})'
            )

            return connection

    async def remove_connection(self, connection_id: str):
        """移除WebSocket连接"""
        async with self._lock:
            if connection_id not in self.connections:
                return

            connection = self.connections[connection_id]
            connection.is_active = False

            # 从流水线映射中移除
            if connection.pipeline_uuid in self.pipeline_connections:
                self.pipeline_connections[connection.pipeline_uuid].discard(connection_id)
                if not self.pipeline_connections[connection.pipeline_uuid]:
                    del self.pipeline_connections[connection.pipeline_uuid]

            # 从会话类型映射中移除
            if connection.session_type in self.session_connections:
                self.session_connections[connection.session_type].discard(connection_id)
                if not self.session_connections[connection.session_type]:
                    del self.session_connections[connection.session_type]

            del self.connections[connection_id]

            logger.debug(f'WebSocket connection disconnected: {connection_id}')

    async def get_connection(self, connection_id: str) -> typing.Optional[WebSocketConnection]:
        """获取指定连接"""
        return self.connections.get(connection_id)

    async def get_connections_by_pipeline(self, pipeline_uuid: str) -> list[WebSocketConnection]:
        """获取指定流水线的所有连接"""
        connection_ids = self.pipeline_connections.get(pipeline_uuid, set())
        return [self.connections[cid] for cid in connection_ids if cid in self.connections]

    async def get_connections_by_session_type(self, session_type: str) -> list[WebSocketConnection]:
        """获取指定会话类型的所有连接"""
        connection_ids = self.session_connections.get(session_type, set())
        return [self.connections[cid] for cid in connection_ids if cid in self.connections]

    async def broadcast_to_pipeline(self, pipeline_uuid: str, message: dict, session_type: str = None):
        """向指定流水线的所有连接广播消息

        Args:
            pipeline_uuid: 流水线UUID
            message: 要广播的消息
            session_type: 可选的会话类型过滤器，如果提供则只向匹配的session_type连接广播
        """
        connections = await self.get_connections_by_pipeline(pipeline_uuid)

        # 如果指定了session_type，只向匹配的连接广播
        if session_type is not None:
            connections = [conn for conn in connections if conn.session_type == session_type]

        tasks = []
        for conn in connections:
            tasks.append(self.send_to_connection(conn.connection_id, message))
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def send_to_connection(self, connection_id: str, message: dict):
        """向指定连接发送消息"""
        connection = await self.get_connection(connection_id)
        if not connection or not connection.is_active:
            logger.warning(f'Attempt to send message to invalid connection: {connection_id}')
            return

        try:
            await connection.send_queue.put(message)
            connection.last_active = datetime.now()
        except Exception as e:
            logger.error(f'Failed to send message to connection {connection_id}: {e}')
            await self.remove_connection(connection_id)

    async def update_activity(self, connection_id: str):
        """更新连接活跃时间"""
        connection = await self.get_connection(connection_id)
        if connection:
            connection.last_active = datetime.now()

    def get_stats(self) -> dict:
        """获取连接统计信息"""
        return {
            'total_connections': len(self.connections),
            'pipelines': len(self.pipeline_connections),
            'connections_by_pipeline': {k: len(v) for k, v in self.pipeline_connections.items()},
            'connections_by_session_type': {k: len(v) for k, v in self.session_connections.items()},
        }


# 全局连接管理器实例
ws_connection_manager = WebSocketConnectionManager()