147 lines
5.5 KiB
Python
147 lines
5.5 KiB
Python
"""BotOps WebSocket 协议层（提炼自 aibox/client.py，增加状态回调与自动重连）。
|
|
|
|
协议：JSON 消息。
|
|
发送：{"type": "message", "content": ...}
|
|
{"type": "command", "command": "/..."}
|
|
{"type": "interrupt", "content": ...}
|
|
接收：chunk / complete / dialogue_ended / task_process / task_status / task_error /
|
|
interrupt_ack / error / info / command_result / connection_ack / notification
|
|
|
|
说明：服务端 chunk 为**累计全文**（与 aibox/client.py 前缀匹配逻辑一致）。
|
|
"""
|
|
import asyncio
|
|
import json
|
|
from typing import Awaitable, Callable, Optional
|
|
|
|
import websockets
|
|
|
|
MessageHandler = Callable[[dict], Awaitable[None]]
|
|
StateHandler = Callable[[str, str], Awaitable[None]] # (state, detail)
|
|
|
|
|
|
class BotWSClient:
|
|
"""aibox 协议 WebSocket 客户端：断线指数退避自动重连，手动断开则抑制。"""
|
|
|
|
def __init__(self, uri: str):
|
|
self.uri = uri
|
|
self.ws = None
|
|
self.connected = False
|
|
self._manual_disconnect = False
|
|
self._listen_task: Optional[asyncio.Task] = None
|
|
self._reconnect_task: Optional[asyncio.Task] = None
|
|
self.on_message: Optional[MessageHandler] = None
|
|
self.on_state: Optional[StateHandler] = None
|
|
|
|
async def _emit_state(self, state: str, detail: str = ""):
|
|
if self.on_state:
|
|
await self.on_state(state, detail)
|
|
|
|
async def connect(self) -> bool:
|
|
"""连接并开始监听；成功返回 True。"""
|
|
await self._close_ws()
|
|
try:
|
|
self.ws = await websockets.connect(self.uri)
|
|
except Exception as e:
|
|
await self._emit_state("error", f"连接失败: {e}")
|
|
return False
|
|
self.connected = True
|
|
self._manual_disconnect = False
|
|
self._listen_task = asyncio.create_task(self._listen())
|
|
await self._emit_state("connected", self.uri)
|
|
return True
|
|
|
|
async def disconnect(self):
|
|
"""手动断开（抑制自动重连）。"""
|
|
self._manual_disconnect = True
|
|
if self._reconnect_task and not self._reconnect_task.done():
|
|
self._reconnect_task.cancel()
|
|
self.connected = False
|
|
await self._close_ws()
|
|
if self._listen_task and not self._listen_task.done():
|
|
self._listen_task.cancel()
|
|
await self._emit_state("disconnected", "已手动断开")
|
|
|
|
async def _close_ws(self):
|
|
if self.ws is not None:
|
|
try:
|
|
await self.ws.close()
|
|
except Exception:
|
|
pass
|
|
self.ws = None
|
|
|
|
async def send_message(self, text: str) -> bool:
|
|
"""发送用户消息（/ 开头按 command 发送，与 aibox 协议一致）。"""
|
|
if not self.connected or not self.ws:
|
|
return False
|
|
data = ({"type": "command", "command": text} if text.startswith("/")
|
|
else {"type": "message", "content": text})
|
|
try:
|
|
await self.ws.send(json.dumps(data, ensure_ascii=False))
|
|
return True
|
|
except Exception as e:
|
|
await self._emit_state("error", f"发送失败: {e}")
|
|
return False
|
|
|
|
async def send_interrupt(self) -> bool:
|
|
"""发送中断信号（对应 client.py 的 Ctrl+C 行为）。"""
|
|
if not self.connected or not self.ws:
|
|
return False
|
|
try:
|
|
await self.ws.send(json.dumps(
|
|
{"type": "interrupt", "content": "用户按下了 Ctrl+C 中断当前操作"},
|
|
ensure_ascii=False))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
async def _listen(self):
|
|
"""接收循环：断线时在 finally 中触发自动重连。"""
|
|
try:
|
|
async for raw in self.ws:
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
data = {"type": "chunk", "content": raw}
|
|
if self.on_message:
|
|
await self.on_message(data)
|
|
except websockets.exceptions.ConnectionClosed:
|
|
pass
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
await self._emit_state("error", f"接收错误: {e}")
|
|
finally:
|
|
was_connected = self.connected
|
|
self.connected = False
|
|
if was_connected:
|
|
await self._emit_state("disconnected", "连接已断开")
|
|
self._schedule_reconnect()
|
|
|
|
def _schedule_reconnect(self):
|
|
if self._manual_disconnect:
|
|
return
|
|
if self._reconnect_task and not self._reconnect_task.done():
|
|
return
|
|
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
|
|
|
|
async def _reconnect_loop(self):
|
|
"""指数退避 2s→30s 封顶，直到成功或手动断开。"""
|
|
delay = 2
|
|
while not self._manual_disconnect and not self.connected:
|
|
await self._emit_state("reconnecting", f"{delay}s 后自动重连 {self.uri} …")
|
|
try:
|
|
await asyncio.sleep(delay)
|
|
except asyncio.CancelledError:
|
|
return
|
|
if self._manual_disconnect:
|
|
return
|
|
try:
|
|
self.ws = await websockets.connect(self.uri)
|
|
except Exception:
|
|
delay = min(delay * 2, 30)
|
|
continue
|
|
self.connected = True
|
|
self._listen_task = asyncio.create_task(self._listen())
|
|
await self._emit_state("connected", self.uri)
|
|
return
|