"""WebSocket连接管理器 - 管理多个并发WebSocket连接""" import asyncio import dataclasses import logging import typing import uuid from datetime import datetime import pydantic from ...api.http.context import ExecutionContext, PrincipalContext logger = logging.getLogger(__name__) _SESSION_FILTER_UNSET = object() _DEFAULT_SEND_QUEUE_SIZE = 100 @dataclasses.dataclass(frozen=True, slots=True) class WebSocketScope: """Trusted runtime placement carried by every WebSocket connection.""" instance_uuid: str workspace_uuid: str placement_generation: int def __post_init__(self) -> None: if not self.instance_uuid.strip() or not self.workspace_uuid.strip(): raise ValueError('WebSocket scope requires an instance and Workspace') if self.placement_generation <= 0: raise ValueError('WebSocket scope requires a positive placement generation') @classmethod def from_context(cls, context: typing.Any) -> 'WebSocketScope': return cls( instance_uuid=str(getattr(context, 'instance_uuid', '')), workspace_uuid=str(getattr(context, 'workspace_uuid', '')), placement_generation=int(getattr(context, 'placement_generation', 0)), ) def is_valid_session_id(value: str) -> bool: """Accept only canonical random UUIDs for client conversation identifiers.""" try: parsed = uuid.UUID(value) except (ValueError, TypeError, AttributeError): return False return parsed.version == 4 and str(parsed) == value 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""" instance_uuid: str """Owning LangBot instance.""" workspace_uuid: str """Owning Workspace.""" placement_generation: int """Workspace placement generation captured at connect time.""" pipeline_uuid: str """关联的流水线UUID""" session_type: str # 'person' or 'group' """会话类型""" session_id: str | None = None """Optional client conversation identifier used by embed widgets.""" 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=lambda: asyncio.Queue(maxsize=_DEFAULT_SEND_QUEUE_SIZE), exclude=True, ) """发送消息队列""" is_active: bool = True """连接是否活跃""" metadata: dict = pydantic.Field(default_factory=dict) """连接元数据（可存储额外信息）""" trigger_principal: PrincipalContext | None = None """Authenticated principal that opened this dashboard connection.""" @property def scope(self) -> WebSocketScope: return WebSocketScope( instance_uuid=self.instance_uuid, workspace_uuid=self.workspace_uuid, placement_generation=self.placement_generation, ) @property def execution_context(self) -> ExecutionContext: """Return the storage/runtime context captured for this connection.""" return ExecutionContext( instance_uuid=self.instance_uuid, workspace_uuid=self.workspace_uuid, placement_generation=self.placement_generation, pipeline_uuid=self.pipeline_uuid, trigger_principal=self.trigger_principal, ) class WebSocketConnectionManager: """WebSocket连接管理器 - 支持多连接并发""" def __init__(self): self.connections: dict[str, WebSocketConnection] = {} """所有活跃连接 {connection_id: connection}""" self.pipeline_connections: dict[tuple[str, str, int, str], set[str]] = {} """Scoped pipeline to connection mapping.""" self.session_connections: dict[tuple[str, str, int, str], set[str]] = {} """Scoped session-type to connection mapping.""" self._lock = asyncio.Lock() """线程锁，保护并发访问""" async def add_connection( self, websocket: typing.Any, scope: WebSocketScope, pipeline_uuid: str, session_type: str, metadata: dict | None = None, trigger_principal: PrincipalContext | None = None, session_id: str | None = None, send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE, max_connections: int = 1024, max_connections_per_workspace: int = 32, ) -> WebSocketConnection: """Register a WebSocket connection and its optional embed session.""" try: send_queue_size = max(int(send_queue_size), 1) except (TypeError, ValueError): send_queue_size = _DEFAULT_SEND_QUEUE_SIZE max_connections = max(int(max_connections), 1) max_connections_per_workspace = max( min(int(max_connections_per_workspace), max_connections), 1, ) async with self._lock: if len(self.connections) >= max_connections: raise RuntimeError(f'WebSocket connection capacity reached ({max_connections})') workspace_connection_count = sum( 1 for connection in self.connections.values() if connection.instance_uuid == scope.instance_uuid and connection.workspace_uuid == scope.workspace_uuid and connection.placement_generation == scope.placement_generation ) if workspace_connection_count >= max_connections_per_workspace: raise RuntimeError(f'Workspace WebSocket connection capacity reached ({max_connections_per_workspace})') connection = WebSocketConnection( instance_uuid=scope.instance_uuid, workspace_uuid=scope.workspace_uuid, placement_generation=scope.placement_generation, pipeline_uuid=pipeline_uuid, session_type=session_type, session_id=session_id, websocket=websocket, metadata=metadata or {}, trigger_principal=trigger_principal, send_queue=asyncio.Queue(maxsize=send_queue_size), ) self.connections[connection.connection_id] = connection # 更新流水线映射 pipeline_key = self._pipeline_key(scope, pipeline_uuid) if pipeline_key not in self.pipeline_connections: self.pipeline_connections[pipeline_key] = set() self.pipeline_connections[pipeline_key].add(connection.connection_id) # 更新会话类型映射 session_key = self._session_key(scope, session_type) if session_key not in self.session_connections: self.session_connections[session_key] = set() self.session_connections[session_key].add(connection.connection_id) logger.debug( f'WebSocket connection established: {connection.connection_id} ' f'(workspace={scope.workspace_uuid}, generation={scope.placement_generation}, ' f'pipeline={pipeline_uuid}, session_type={session_type})' ) return connection async def close_scope(self, scope: WebSocketScope) -> None: """Close and forget every live connection for one runtime placement.""" async with self._lock: connection_ids = [ connection_id for connection_id, connection in self.connections.items() if connection.scope == scope ] for connection_id in connection_ids: connection = self.connections.get(connection_id) if connection is None: continue close = getattr(connection.websocket, 'close', None) if close is not None: try: result = close() if asyncio.iscoroutine(result): await result except Exception: logger.debug( 'Failed to close WebSocket connection %s', connection_id, exc_info=True, ) await self.remove_connection(connection_id) 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 # 从流水线映射中移除 pipeline_key = self._pipeline_key(connection.scope, connection.pipeline_uuid) if pipeline_key in self.pipeline_connections: self.pipeline_connections[pipeline_key].discard(connection_id) if not self.pipeline_connections[pipeline_key]: del self.pipeline_connections[pipeline_key] # 从会话类型映射中移除 session_key = self._session_key(connection.scope, connection.session_type) if session_key in self.session_connections: self.session_connections[session_key].discard(connection_id) if not self.session_connections[session_key]: del self.session_connections[session_key] del self.connections[connection_id] logger.debug(f'WebSocket connection disconnected: {connection_id}') @staticmethod def _pipeline_key(scope: WebSocketScope, pipeline_uuid: str) -> tuple[str, str, int, str]: return ( scope.instance_uuid, scope.workspace_uuid, scope.placement_generation, pipeline_uuid, ) @staticmethod def _session_key(scope: WebSocketScope, session_type: str) -> tuple[str, str, int, str]: return ( scope.instance_uuid, scope.workspace_uuid, scope.placement_generation, session_type, ) async def get_connection( self, connection_id: str, *, scope: WebSocketScope, ) -> WebSocketConnection | None: """Get a connection only when it belongs to the expected placement.""" connection = self.connections.get(connection_id) if connection is None or connection.scope != scope: return None return connection async def get_connection_by_session_id( self, session_id: str, *, scope: WebSocketScope, pipeline_uuid: str | None = None, ) -> WebSocketConnection | None: """Get an active embed connection by its stable browser session identifier.""" candidates: typing.Iterable[WebSocketConnection] if pipeline_uuid is not None: candidates = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope) else: candidates = self.connections.values() for connection in candidates: if ( connection.session_id == session_id and connection.is_active and connection.scope == scope and (pipeline_uuid is None or connection.pipeline_uuid == pipeline_uuid) ): return connection return None async def get_connections_by_pipeline( self, pipeline_uuid: str, *, scope: WebSocketScope, ) -> list[WebSocketConnection]: """获取指定流水线的所有连接""" connection_ids = self.pipeline_connections.get(self._pipeline_key(scope, 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, *, scope: WebSocketScope, ) -> list[WebSocketConnection]: """获取指定会话类型的所有连接""" connection_ids = self.session_connections.get(self._session_key(scope, 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, *, scope: WebSocketScope, session_type: str | None = None, session_id: typing.Any = _SESSION_FILTER_UNSET, ): """Broadcast a message to matching connections for one pipeline. Args: pipeline_uuid: Pipeline identifier. message: Serialized message to enqueue. session_type: Optional session-type filter. session_id: Embed conversation filter. Omit it to broadcast across conversations; pass ``None`` to target non-embed connections. """ connections = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope) if session_type is not None: connections = [conn for conn in connections if conn.session_type == session_type] if session_id is not _SESSION_FILTER_UNSET: connections = [conn for conn in connections if conn.session_id == session_id] 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 = self.connections.get(connection_id) if not connection or not connection.is_active: logger.warning(f'Attempt to send message to invalid connection: {connection_id}') return try: try: connection.send_queue.put_nowait(message) except asyncio.QueueFull: # A slow or disconnected browser must not backpressure every # other connection or retain an unbounded response stream. try: connection.send_queue.get_nowait() except asyncio.QueueEmpty: pass connection.send_queue.put_nowait(message) logger.warning( 'WebSocket send queue full; dropped oldest message for connection %s', connection_id, ) 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 = self.connections.get(connection_id) if connection: connection.last_active = datetime.now() def get_stats(self, *, scope: WebSocketScope) -> dict: """Return connection statistics for one trusted placement.""" scoped_connections = [connection for connection in self.connections.values() if connection.scope == scope] pipelines: dict[str, int] = {} session_types: dict[str, int] = {} for connection in scoped_connections: pipelines[connection.pipeline_uuid] = pipelines.get(connection.pipeline_uuid, 0) + 1 session_types[connection.session_type] = session_types.get(connection.session_type, 0) + 1 return { 'total_connections': len(scoped_connections), 'pipelines': len(pipelines), 'connections_by_pipeline': pipelines, 'connections_by_session_type': session_types, } # 全局连接管理器实例 ws_connection_manager = WebSocketConnectionManager()