"""WebSocket适配器 - 支持双向通信的IM系统""" import asyncio import contextvars import logging import time import typing from datetime import datetime import pydantic import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter import langbot_plugin.api.entities.builtin.platform.message as platform_message import langbot_plugin.api.entities.builtin.platform.events as platform_events import langbot_plugin.api.entities.builtin.platform.entities as platform_entities import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger from ...core import app from ...core import entities as core_entities from .websocket_manager import WebSocketConnection, WebSocketScope, is_valid_session_id, ws_connection_manager logger = logging.getLogger(__name__) _current_pipeline_uuid: contextvars.ContextVar[str | None] = contextvars.ContextVar( 'websocket_pipeline_uuid', default=None, ) class WebSocketMessage(pydantic.BaseModel): """WebSocket消息格式""" id: int role: str # 'user' or 'assistant' content: str message_chain: list[dict] timestamp: str is_final: bool = False connection_id: str = '' """发送者连接ID""" class WebSocketSession: """WebSocket会话 - 管理单个会话的消息历史""" id: str message_lists: dict[str, list[WebSocketMessage]] = {} """消息列表 {pipeline_uuid: [messages]}""" stream_message_indexes: dict[str, dict[str, int]] = {} """流式消息索引 {pipeline_uuid: {resp_message_id: message_index}}""" def __init__( self, id: str = '', *, max_conversations: int = 200, max_messages: int = 100, idle_ttl_seconds: int = 86400, ): self.id = id self.message_lists = {} self.stream_message_indexes = {} self.message_counters: dict[str, int] = {} self.last_accessed: dict[str, float] = {} self.max_conversations = max(int(max_conversations), 1) self.max_messages = max(int(max_messages), 1) self.idle_ttl_seconds = max(int(idle_ttl_seconds), 1) def _prune(self, now: float) -> None: expired = [ key for key, last_accessed in self.last_accessed.items() if now - last_accessed >= self.idle_ttl_seconds ] for key in expired: self.reset(key) overflow = len(self.message_lists) - self.max_conversations + 1 if overflow <= 0: return oldest = sorted(self.last_accessed, key=self.last_accessed.get) for key in oldest[:overflow]: self.reset(key) def get_message_list(self, pipeline_uuid: str) -> list[WebSocketMessage]: now = time.monotonic() self._prune(now) if pipeline_uuid not in self.message_lists: self.message_lists[pipeline_uuid] = [] self.last_accessed[pipeline_uuid] = now return self.message_lists[pipeline_uuid] def get_stream_message_indexes(self, pipeline_uuid: str) -> dict[str, int]: if pipeline_uuid not in self.stream_message_indexes: self.stream_message_indexes[pipeline_uuid] = {} self.last_accessed[pipeline_uuid] = time.monotonic() return self.stream_message_indexes[pipeline_uuid] def next_message_id(self, conversation_key: str) -> int: next_id = self.message_counters.get(conversation_key, 0) + 1 self.message_counters[conversation_key] = next_id return next_id def append_message(self, conversation_key: str, message: WebSocketMessage) -> None: messages = self.get_message_list(conversation_key) messages.append(message) overflow = len(messages) - self.max_messages if overflow <= 0: return del messages[:overflow] indexes = self.stream_message_indexes.get(conversation_key, {}) adjusted_indexes = { response_id: index - overflow for response_id, index in indexes.items() if index >= overflow } indexes.clear() indexes.update(adjusted_indexes) def reset(self, conversation_key: str) -> None: self.message_lists.pop(conversation_key, None) self.stream_message_indexes.pop(conversation_key, None) self.message_counters.pop(conversation_key, None) self.last_accessed.pop(conversation_key, None) def clear(self) -> None: self.message_lists.clear() self.stream_message_indexes.clear() self.message_counters.clear() self.last_accessed.clear() class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter): """WebSocket适配器 - 支持双向实时通信""" websocket_person_session: WebSocketSession = pydantic.Field(exclude=True, default_factory=WebSocketSession) websocket_group_session: WebSocketSession = pydantic.Field(exclude=True, default_factory=WebSocketSession) listeners: dict[ typing.Type[platform_events.Event], typing.Callable[[platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], None], ] = pydantic.Field(default_factory=dict, exclude=True) ap: app.Application = pydantic.Field(exclude=True) # 主动推送消息的队列 outbound_message_queue: asyncio.Queue = pydantic.Field( default_factory=lambda: asyncio.Queue(maxsize=100), exclude=True, ) inbound_listener_tasks: set[asyncio.Task] = pydantic.Field( default_factory=set, exclude=True, ) """后端主动推送消息的队列""" # 流式输出开关 stream_enabled: bool = pydantic.Field(default=True, exclude=True) """是否启用流式输出""" def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs): super().__init__( config=config, logger=logger, **kwargs, ) application = kwargs.get('ap') instance_data = getattr(getattr(application, 'instance_config', None), 'data', {}) retention = ( instance_data.get('system', {}).get('websocket_retention', {}) if isinstance(instance_data, dict) else {} ) session_options = { 'max_conversations': retention.get('max_conversations_per_workspace', 200), 'max_messages': retention.get('max_messages_per_conversation', 100), 'idle_ttl_seconds': retention.get('conversation_idle_ttl_seconds', 86400), } self.websocket_person_session = WebSocketSession(id='websocketperson', **session_options) self.websocket_group_session = WebSocketSession(id='websocketgroup', **session_options) self.bot_account_id = 'websocketbot' try: outbound_queue_size = max(int(retention.get('send_queue_size', 100)), 1) except (TypeError, ValueError): outbound_queue_size = 100 self.outbound_message_queue = asyncio.Queue(maxsize=outbound_queue_size) self.inbound_listener_tasks = set() self.stream_enabled = True @staticmethod def _conversation_key(pipeline_uuid: str, session_id: str | None = None) -> str: """Return the history key for a pipeline/client conversation.""" return f'{pipeline_uuid}:{session_id}' if session_id else pipeline_uuid @staticmethod def _parse_embed_target(target_id: str) -> tuple[str, str] | None: """Extract pipeline and session identifiers from a stable embed launcher.""" target_value = str(target_id) for prefix in ('websocket_', 'websocketgroup_'): if target_value.startswith(prefix): target = target_value[len(prefix) :] break else: return None if ':' not in target: return None pipeline_uuid, session_id = target.rsplit(':', 1) if not pipeline_uuid or not is_valid_session_id(session_id): return None return pipeline_uuid, session_id def _scope(self) -> WebSocketScope: """Return this adapter's immutable runtime placement.""" return WebSocketScope.from_context(self.logger.execution_context) def get_pipeline_uuid_override(self) -> str | None: """Return the connection pipeline propagated into the listener task.""" return _current_pipeline_uuid.get() def _listener_task_done(self, task: asyncio.Task) -> None: listener_tasks = getattr(self, 'inbound_listener_tasks', None) if listener_tasks is not None: listener_tasks.discard(task) if not task.cancelled(): task.exception() @staticmethod def _history_message_chain(message_chain: list[dict]) -> list[dict]: """Remove large transient payloads before retaining browser history.""" history = [] for component in message_chain: copied = dict(component) if copied.get('base64'): copied['base64'] = '' history.append(copied) return history async def _get_connection_from_target(self, target_id: str): """Resolve a person or group WebSocket launcher to its connection.""" scope = self._scope() target_value = str(target_id) for prefix in ('websocket_', 'websocketgroup_'): if target_value.startswith(prefix): target = target_value[len(prefix) :] break else: return None connection = await ws_connection_manager.get_connection(target, scope=scope) if connection is not None: return connection embed_target = self._parse_embed_target(target_id) if embed_target is not None: pipeline_uuid, session_id = embed_target return await ws_connection_manager.get_connection_by_session_id( session_id, scope=scope, pipeline_uuid=pipeline_uuid, ) return await ws_connection_manager.get_connection_by_session_id(target, scope=scope) async def _get_message_context(self, message_source) -> tuple[str, str | None]: """Resolve the originating pipeline and browser session for a reply.""" sender = getattr(message_source, 'sender', None) sender_id = getattr(sender, 'id', '') connection = await self._get_connection_from_target(sender_id) if connection is not None: return connection.pipeline_uuid, connection.session_id embed_target = self._parse_embed_target(sender_id) if embed_target is not None: return embed_target raise ValueError('WebSocket reply target is not bound to this adapter scope') async def send_message( self, target_type: str, target_id: str, message: platform_message.MessageChain, ) -> dict: """发送消息 - 这里用于主动推送消息到前端 对于 WebSocket 适配器，我们需要将消息广播到正确的 pipeline 连接。 target_id 可能是 launcher_id（如 websocket_xxx）或 pipeline_uuid。 我们需要尝试两种方式来确保消息能够送达。 """ connection = await self._get_connection_from_target(target_id) if connection is not None: pipeline_uuid = connection.pipeline_uuid session_id = connection.session_id scope = connection.scope else: embed_target = self._parse_embed_target(target_id) if embed_target is not None: pipeline_uuid, session_id = embed_target else: pipeline_uuid = str(target_id).strip() if not pipeline_uuid: raise ValueError('WebSocket target pipeline is required') session_id = None scope = self._scope() session_type = 'group' if target_type == 'group' else 'person' conversation_key = self._conversation_key(pipeline_uuid, session_id) session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session msg_id = session.next_message_id(conversation_key) message_data = WebSocketMessage( id=msg_id, role='assistant', content=str(message), message_chain=[component.__dict__ for component in message], timestamp=datetime.now().isoformat(), is_final=True, ) session.append_message(conversation_key, message_data) await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { 'type': 'response', 'session_type': session_type, 'data': message_data.model_dump(), }, scope=scope, session_type=session_type, session_id=session_id, ) return message_data.model_dump() async def reply_message( self, message_source: platform_events.MessageEvent, message: platform_message.MessageChain, quote_origin: bool = False, ) -> dict: """回复消息 - 非流式""" # 获取会话和pipeline信息 session = ( self.websocket_group_session if isinstance(message_source, platform_events.GroupMessage) else self.websocket_person_session ) pipeline_uuid, session_id = await self._get_message_context(message_source) scope = self._scope() session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person' conversation_key = self._conversation_key(pipeline_uuid, session_id) msg_id = session.next_message_id(conversation_key) message_data = WebSocketMessage( id=msg_id, role='assistant', content=str(message), message_chain=[component.__dict__ for component in message], timestamp=datetime.now().isoformat(), is_final=True, ) session.append_message(conversation_key, message_data) await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { 'type': 'response', 'session_type': session_type, 'data': message_data.model_dump(), }, scope=scope, session_type=session_type, session_id=session_id, ) return message_data.model_dump() async def reply_message_chunk( self, message_source: platform_events.MessageEvent, bot_message, message: platform_message.MessageChain, quote_origin: bool = False, is_final: bool = False, ) -> dict: """回复消息块 - 流式""" # 获取会话和pipeline信息 session = ( self.websocket_group_session if isinstance(message_source, platform_events.GroupMessage) else self.websocket_person_session ) pipeline_uuid, session_id = await self._get_message_context(message_source) scope = self._scope() session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person' conversation_key = self._conversation_key(pipeline_uuid, session_id) message_list = session.get_message_list(conversation_key) stream_message_indexes = session.get_stream_message_indexes(conversation_key) # Streaming messages in LangBot have a stable resp_message_id during the same assistant reply. # Use it as the primary key to avoid overwriting an old card from a previous reply. resp_message_id = str(getattr(bot_message, 'resp_message_id', '') or '') existing_index = stream_message_indexes.get(resp_message_id) if resp_message_id else None message_is_final = is_final and bot_message.tool_calls is None if existing_index is None or existing_index >= len(message_list): # 创建新消息 msg_id = session.next_message_id(conversation_key) message_data = WebSocketMessage( id=msg_id, role='assistant', content=str(message), message_chain=[component.__dict__ for component in message], timestamp=datetime.now().isoformat(), is_final=message_is_final, ) # 立即添加到历史记录（即使is_final=False），以便后续块可以更新它 session.append_message(conversation_key, message_data) message_list = session.get_message_list(conversation_key) if resp_message_id: stream_message_indexes[resp_message_id] = len(message_list) - 1 else: # 更新同一条流式消息 old_message = message_list[existing_index] msg_id = old_message.id message_data = WebSocketMessage( id=msg_id, role='assistant', content=str(message), message_chain=[component.__dict__ for component in message], timestamp=old_message.timestamp, # 保持原始时间戳 is_final=message_is_final, ) # 更新历史记录中的对应消息 message_list[existing_index] = message_data if message_is_final and resp_message_id: stream_message_indexes.pop(resp_message_id, None) await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { 'type': 'response', 'session_type': session_type, 'data': message_data.model_dump(), }, scope=scope, session_type=session_type, session_id=session_id, ) return message_data.model_dump() async def is_stream_output_supported(self) -> bool: """根据stream_enabled标志返回是否支持流式输出""" return self.stream_enabled def register_listener( self, event_type: typing.Type[platform_events.Event], func: typing.Callable[ [platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], typing.Awaitable[None] ], ): """注册事件监听器""" self.listeners[event_type] = func def unregister_listener( self, event_type: typing.Type[platform_events.Event], func: typing.Callable[ [platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], typing.Awaitable[None] ], ): """取消注册事件监听器""" del self.listeners[event_type] async def is_muted(self, group_id: int) -> bool: return False async def run_async(self): """运行适配器""" try: while True: # 处理主动推送消息 if not self.outbound_message_queue.empty(): try: message = await asyncio.wait_for(self.outbound_message_queue.get(), timeout=0.1) # 广播到所有相关连接 target_id = message.get('target_id', '') await ws_connection_manager.broadcast_to_pipeline( target_id, message, scope=self._scope(), ) except asyncio.TimeoutError: pass await asyncio.sleep(0.1) except asyncio.CancelledError: raise async def kill(self): """停止适配器""" await ws_connection_manager.close_scope(self._scope()) listener_tasks = getattr(self, 'inbound_listener_tasks', set()) inbound_tasks = list(listener_tasks) for task in inbound_tasks: if not task.done(): task.cancel() if inbound_tasks: await asyncio.gather(*inbound_tasks, return_exceptions=True) listener_tasks.clear() self.websocket_person_session.clear() self.websocket_group_session.clear() while not self.outbound_message_queue.empty(): try: self.outbound_message_queue.get_nowait() except asyncio.QueueEmpty: break async def _process_image_components( self, connection: WebSocketConnection, message_chain_obj: list, ): """ 处理消息链中的图片、语音和文件组件，将 path 转换为 base64 Image / Voice / File components uploaded from the web client carry a storage key in ``path``. Resolve it to a base64 data URI so downstream stages (multimodal LLM input and the Box sandbox inbox) have a usable payload, then drop the now-consumed storage object. Args: message_chain_obj: 消息链对象列表 """ import base64 import mimetypes attachments = [ component for component in message_chain_obj if component.get('path') and component.get('type') in ('Image', 'Voice', 'File') ] if not attachments: return storage_mgr = self.ap.storage_mgr execution_context = connection.execution_context expected_prefix = storage_mgr.scoped_prefix(execution_context, owner_type='upload_image') for component in attachments: comp_type = component.get('type', '') comp_path = component.get('path', '') if not comp_path.startswith(expected_prefix) or not storage_mgr.is_scoped_object_key( comp_path, expected_owner_type='upload_image', ): await self.logger.warning(f'Rejected {comp_type} attachment outside the WebSocket connection scope') raise ValueError('Attachment key does not belong to this WebSocket connection') try: file_content = await storage_mgr.load_scoped_object_key( execution_context, comp_path, expected_owner_type='upload_image', ) base64_str = (await asyncio.to_thread(base64.b64encode, file_content)).decode('utf-8') lowered = comp_path.lower() if comp_type == 'Image': if lowered.endswith(('.jpg', '.jpeg')): mime_type = 'image/jpeg' elif lowered.endswith('.gif'): mime_type = 'image/gif' elif lowered.endswith('.webp'): mime_type = 'image/webp' else: mime_type = 'image/png' elif comp_type == 'Voice': mime_type = mimetypes.guess_type(comp_path)[0] or 'audio/wav' else: # File mime_type = mimetypes.guess_type(comp_path)[0] or 'application/octet-stream' component['base64'] = f'data:{mime_type};base64,{base64_str}' await storage_mgr.delete_scoped_object_key( execution_context, comp_path, expected_owner_type='upload_image', ) component['path'] = '' except Exception as e: await self.logger.error(f'Failed to load {comp_type} file {comp_path}: {e}') raise async def handle_websocket_message( self, connection: WebSocketConnection, message_data: dict, owner_bot=None, ): """ 处理从WebSocket接收的消息 这个方法只负责接收消息、保存到历史记录、并触发事件处理 不等待任何响应，响应消息会通过reply_message/reply_message_chunk直接发送 Args: connection: WebSocket连接对象 message_data: 消息数据，包含: - message: 消息链 - stream: 是否启用流式输出 (可选，默认True) owner_bot: Optional RuntimeBot that owns this pipeline (e.g. a web_page_bot). When provided, its identity is used for logging and session tracking. """ pipeline_uuid = connection.pipeline_uuid session_type = connection.session_type conversation_key = self._conversation_key(pipeline_uuid, connection.session_id) self.stream_enabled = message_data.get('stream', True) use_session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session message_chain_obj = message_data.get('message', []) await self._process_image_components(connection, message_chain_obj) message_chain = platform_message.MessageChain.model_validate(message_chain_obj) message_id = use_session.next_message_id(conversation_key) # 保存用户消息 user_message = WebSocketMessage( id=message_id, role='user', content=str(message_chain), message_chain=self._history_message_chain(message_chain_obj), timestamp=datetime.now().isoformat(), connection_id=connection.connection_id, is_final=True, # 用户消息始终是完整的，非流式 ) use_session.append_message(conversation_key, user_message) await ws_connection_manager.broadcast_to_pipeline( pipeline_uuid, { 'type': 'user_message', 'session_type': session_type, 'data': user_message.model_dump(), }, scope=connection.scope, session_type=session_type, session_id=connection.session_id, ) # 添加消息源 message_chain.insert(0, platform_message.Source(id=message_id, time=datetime.now().timestamp())) # 创建事件 launcher_id = f'{pipeline_uuid}:{connection.session_id}' if connection.session_id else connection.connection_id if session_type == 'person': sender = platform_entities.Friend(id=f'websocket_{launcher_id}', nickname='User', remark='User') event = platform_events.FriendMessage( sender=sender, message_chain=message_chain, time=datetime.now().timestamp() ) else: group = platform_entities.Group( id=f'websocketgroup_{launcher_id}' if connection.session_id else 'websocketgroup', name='Group', permission=platform_entities.Permission.Member, ) sender = platform_entities.GroupMember( id=f'websocket_{launcher_id}', member_name='User', group=group, permission=platform_entities.Permission.Member, ) event = platform_events.GroupMessage( sender=sender, message_chain=message_chain, time=datetime.now().timestamp() ) # 异步触发事件处理 # Use owner_bot's listeners if available, otherwise fall back to proxy bot listeners = ( owner_bot.adapter.listeners if (owner_bot and hasattr(owner_bot.adapter, 'listeners') and owner_bot.adapter.listeners) else self.listeners ) # Pass owner_bot's adapter so that downstream logging / dashboard # attributes the message to the correct bot adapter name. # Wire the ws adapter into the owner so replies are actually delivered. if owner_bot and hasattr(owner_bot.adapter, 'set_ws_adapter'): owner_bot.adapter.set_ws_adapter(self) callback_adapter = owner_bot.adapter if (owner_bot and hasattr(owner_bot, 'adapter')) else self if event.__class__ in listeners: listener_tasks = getattr(self, 'inbound_listener_tasks', None) if listener_tasks is None: listener_tasks = set() object.__setattr__(self, 'inbound_listener_tasks', listener_tasks) for task in tuple(listener_tasks): if task.done(): listener_tasks.discard(task) if len(listener_tasks) >= 100: await self.logger.warning('WebSocket inbound listener capacity reached; dropping message') return listener = typing.cast( typing.Callable[[typing.Any, typing.Any], typing.Awaitable[None]], listeners[event.__class__], ) async def run_listener(): token = _current_pipeline_uuid.set(pipeline_uuid) try: await listener(event, callback_adapter) finally: _current_pipeline_uuid.reset(token) listener_coro = run_listener() task_manager = getattr(self.ap, 'task_mgr', None) if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list): listener_task = asyncio.create_task(listener_coro) else: listener_task = task_manager.create_task( listener_coro, kind='websocket-message', name=f'websocket-message-{connection.connection_id}', scopes=[ core_entities.LifecycleControlScope.APPLICATION, core_entities.LifecycleControlScope.PLATFORM, ], instance_uuid=connection.instance_uuid, workspace_uuid=connection.workspace_uuid, placement_generation=connection.placement_generation, ).task listener_tasks.add(listener_task) listener_task.add_done_callback(self._listener_task_done) def get_websocket_messages( self, pipeline_uuid: str, session_type: str, session_id: str | None = None, ) -> list[dict]: """Return history for one pipeline/client conversation.""" conversation_key = self._conversation_key(pipeline_uuid, session_id) session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session return [message.model_dump() for message in session.message_lists.get(conversation_key, [])] def reset_session( self, pipeline_uuid: str, session_type: str, session_id: str | None = None, ): """Reset one pipeline/client conversation.""" conversation_key = self._conversation_key(pipeline_uuid, session_id) session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session if isinstance(session, WebSocketSession): session.reset(conversation_key) else: # Compatibility for lightweight adapter doubles. if conversation_key in session.message_lists: session.message_lists[conversation_key] = [] if conversation_key in session.stream_message_indexes: session.stream_message_indexes[conversation_key] = {} if session_id: launcher_id = ( f'websocketgroup_{pipeline_uuid}:{session_id}' if session_type == 'group' else f'websocket_{pipeline_uuid}:{session_id}' ) scope = self._scope() self.ap.sess_mgr.session_list = [ candidate_session for candidate_session in self.ap.sess_mgr.session_list if not ( getattr(candidate_session, 'instance_uuid', None) == scope.instance_uuid and getattr(candidate_session, 'workspace_uuid', None) == scope.workspace_uuid and getattr(candidate_session, 'placement_generation', None) == scope.placement_generation and str( candidate_session.launcher_type.value if hasattr(candidate_session.launcher_type, 'value') else candidate_session.launcher_type ) == session_type and str(candidate_session.launcher_id) == launcher_id ) ]