仿生人MCP接入点

bbit_ai/app_mcp/handlers/websocket_handler.py

@@ -0,0 +1,117 @@
+"""
+WebSocket处理器
+处理工具端和小智端的WebSocket连接
+"""
+
+import json
+from core.connection_manager import connection_manager
+from utils.logger import get_logger
+from utils.jsonrpc import (
+    create_tool_not_connected_error,
+    create_forward_failed_error,
+)
+
+logger = get_logger()
+
+
+class WebSocketHandler:
+    """WebSocket处理器"""
+
+    def __init__(self):
+        pass
+
+    async def _handle_tool_message(self, agent_id: str, message: str):
+        """处理工具端消息"""
+        try:
+            # 解析消息
+            logger.debug(f"收到工具端消息: {agent_id} - {message}")
+
+            # 尝试解析JSON-RPC响应消息
+            try:
+                message_data = json.loads(message)
+                # 还原JSON-RPC ID并获取目标连接UUID
+                connection_uuid, restored_message = (
+                    connection_manager.restore_jsonrpc_message(message_data)
+                )
+
+                if connection_uuid:
+                    # 有特定的目标连接，发送给该连接
+                    success = await connection_manager.forward_to_robot_by_uuid(
+                        connection_uuid, restored_message
+                    )
+                    if not success:
+                        logger.error(f"转发消息给特定小智端连接失败: {connection_uuid}")
+                else:
+                    logger.error(f"没有特定目标，无法转发消息")
+            except json.JSONDecodeError:
+                # 如果不是JSON格式，按原来的方式处理
+                logger.error(f"由于消息不是JSON格式，已忽略: {message}")
+
+        except json.JSONDecodeError:
+            logger.error(f"工具端消息格式错误: {message}")
+        except Exception as e:
+            logger.error(f"处理工具端消息时发生错误: {e}")
+
+    async def _handle_robot_message(
+        self, agent_id: str, message: str, connection_uuid: str
+    ):
+        """处理小智端消息"""
+        try:
+            # 解析消息
+            logger.debug(
+                f"收到小智端消息: {agent_id} (UUID: {connection_uuid}) - {message}"
+            )
+
+            # 尝试解析JSON-RPC消息以获取id
+            request_id = None
+            transformed_message = message
+            try:
+                message_data = json.loads(message)
+                request_id = message_data.get("id")
+
+                # 转换JSON-RPC ID
+                transformed_message_data = connection_manager.transform_jsonrpc_message(
+                    message_data, connection_uuid
+                )
+                transformed_message = json.dumps(
+                    transformed_message_data, ensure_ascii=False
+                )
+
+                logger.debug(
+                    f"转换后的消息ID: {message_data.get('id')} -> {transformed_message_data.get('id')}"
+                )
+
+            except json.JSONDecodeError:
+                logger.warning(f"小智端消息不是有效的JSON格式: {message}")
+                # 如果消息不是JSON格式，仍然检查工具端连接状态
+
+            # 检查是否有对应的工具端连接
+            if not connection_manager.is_tool_connected(agent_id):
+                logger.warning(f"工具端未连接: {agent_id}")
+                # 发送JSON-RPC格式的错误消息给小智端
+                error_message = create_tool_not_connected_error(request_id, agent_id)
+                await connection_manager.forward_to_robot_by_uuid(
+                    connection_uuid, error_message
+                )
+                return
+
+            # 转发转换后的消息给工具端
+            success = await connection_manager.forward_to_tool(
+                agent_id, transformed_message
+            )
+            if not success:
+                logger.error(f"转发消息给工具端失败: {agent_id}")
+                # 发送JSON-RPC格式的错误消息给小智端
+                error_message = create_forward_failed_error(request_id, agent_id)
+                await connection_manager.forward_to_robot_by_uuid(
+                    connection_uuid, error_message
+                )
+
+        except json.JSONDecodeError:
+            logger.error(f"小智端消息格式错误: {message}")
+        except Exception as e:
+            logger.error(f"处理小智端消息时发生错误: {e}")
+
+
+# 全局WebSocket处理器实例
+websocket_handler = WebSocketHandler()