update:规范websocket连接

Author: openrz

src/handlers/websocket_handler.py

@@ -11,8 +11,6 @@
 from ..core.connection_manager import connection_manager
 from ..utils.logger import get_logger
 from ..utils.jsonrpc import (
-    JSONRPCProtocol,
-    create_connection_established_message,
     create_tool_not_connected_error,
     create_forward_failed_error,
 )
@@ -26,78 +24,6 @@ class WebSocketHandler:
     def __init__(self):
         pass
 
-    async def handle_tool_connection(
-        self, websocket: WebSocketServerProtocol, path: str
-    ):
-        """处理工具端连接"""
-        try:
-            # 解析URL参数获取agentId
-            agent_id = self._extract_agent_id(path)
-            if not agent_id:
-                await websocket.close(1008, "缺少agentId参数")
-                return
-
-            # 注册连接
-            await connection_manager.register_tool_connection(agent_id, websocket)
-
-            # 发送连接确认消息
-            connection_message = create_connection_established_message(
-                agent_id, "工具端连接已建立"
-            )
-            await websocket.send(connection_message)
-
-            logger.info(f"工具端连接已建立: {agent_id}")
-
-            # 处理消息
-            async for message in websocket:
-                await self._handle_tool_message(agent_id, message)
-
-        except ConnectionClosed:
-            logger.info(
-                f"工具端连接已关闭: {agent_id if 'agent_id' in locals() else 'unknown'}"
-            )
-        except Exception as e:
-            logger.error(f"处理工具端连接时发生错误: {e}")
-        finally:
-            if "agent_id" in locals():
-                await connection_manager.unregister_tool_connection(agent_id)
-
-    async def handle_robot_connection(
-        self, websocket: WebSocketServerProtocol, path: str
-    ):
-        """处理小智端连接"""
-        try:
-            # 解析URL参数获取agentId
-            agent_id = self._extract_agent_id(path)
-            if not agent_id:
-                await websocket.close(1008, "缺少agentId参数")
-                return
-
-            # 注册连接
-            await connection_manager.register_robot_connection(agent_id, websocket)
-
-            # 发送连接确认消息
-            connection_message = create_connection_established_message(
-                agent_id, "小智端连接已建立"
-            )
-            await websocket.send(connection_message)
-
-            logger.info(f"小智端连接已建立: {agent_id}")
-
-            # 处理消息
-            async for message in websocket:
-                await self._handle_robot_message(agent_id, message)
-
-        except ConnectionClosed:
-            logger.info(
-                f"小智端连接已关闭: {agent_id if 'agent_id' in locals() else 'unknown'}"
-            )
-        except Exception as e:
-            logger.error(f"处理小智端连接时发生错误: {e}")
-        finally:
-            if "agent_id" in locals():
-                await connection_manager.unregister_robot_connection(agent_id)
-
     async def _handle_tool_message(self, agent_id: str, message: str):
         """处理工具端消息"""
         try:

src/server.py

@@ -12,8 +12,6 @@
 from .utils.aes_utils import decrypt
 from .utils.jsonrpc import (
     JSONRPCProtocol,
-    create_connection_established_message,
-    create_authentication_error,
 )
 from src.utils.util import get_local_ip
 from contextlib import asynccontextmanager
@@ -149,13 +147,6 @@ async def websocket_tool_endpoint(websocket: WebSocket):
     try:
         # 注册连接
         await connection_manager.register_tool_connection(agent_id, websocket)
-
-        # 发送连接确认消息
-        connection_message = create_connection_established_message(
-            agent_id, "工具端连接已建立"
-        )
-        await websocket.send_text(connection_message)
-
         logger.info(f"工具端连接已建立: {agent_id}")
 
         # 处理消息
@@ -190,12 +181,6 @@ async def websocket_robot_endpoint(websocket: WebSocket):
         # 注册连接
         await connection_manager.register_robot_connection(agent_id, websocket)
 
-        # 发送连接确认消息
-        connection_message = create_connection_established_message(
-            agent_id, "小智端连接已建立"
-        )
-        await websocket.send_text(connection_message)
-
         logger.info(f"小智端连接已建立: {agent_id}")
 
         # 处理消息

src/utils/jsonrpc.py

@@ -178,21 +178,6 @@ def is_notification(json_str: str) -> bool:
         return request is not None and request.id is None
 
 
-# 便捷函数
-def create_connection_established_message(
-    agent_id: str, message: str = "连接已建立"
-) -> str:
-    """创建连接建立成功的消息"""
-    response = JSONRPCProtocol.create_success_response(
-        result={
-            "type": "connection_established",
-            "message": message,
-            "agent_id": agent_id,
-        }
-    )
-    return JSONRPCProtocol.to_json(response, ensure_ascii=False)
-
-
 def create_tool_not_connected_error(
     request_id: Optional[Union[str, int]] = None, agent_id: Optional[str] = None
 ) -> str: