BBIT-Kai
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
"""
|
|
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()
|