From aaeac343d36e42399e4fd407485118c1d0d6bf77 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Tue, 28 Jul 2026 21:44:23 +0800 Subject: [PATCH] =?UTF-8?q?fix(aibox/client):=20=E4=BF=AE4=E5=A4=84?= =?UTF-8?q?=E2=80=94=E2=80=94stdin=E5=8D=95=E7=BA=BF=E7=A8=8B=E9=98=B2?= =?UTF-8?q?=E6=B3=84=E6=BC=8F=E3=80=81LRU=E5=8E=BB=E9=87=8D=E4=B8=8A?= =?UTF-8?q?=E9=99=90=E3=80=81=E8=87=AA=E5=8A=A8=E9=87=8D=E8=BF=9E=E3=80=81?= =?UTF-8?q?=E5=BC=83=E7=94=A8get=5Fevent=5Floop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- aibox/client.py | 126 ++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 110 insertions(+), 16 deletions(-) diff --git a/aibox/client.py b/aibox/client.py index 881a366..f109ed3 100644 --- a/aibox/client.py +++ b/aibox/client.py @@ -9,6 +9,7 @@ import sys import os import threading import hashlib +from collections import OrderedDict from datetime import datetime from typing import Optional, Dict, Any, List import websockets @@ -75,10 +76,15 @@ class WebSocketClient: except json.JSONDecodeError: self.callback("chunk", message) except websockets.exceptions.ConnectionClosed: - self.connected = False - self.callback("system", "❌ 连接已断开") + pass except Exception as e: self.callback("error", f"接收消息错误: {e}") + finally: + # 统一在此通知断线；手动 disconnect() 已把 connected 置 False，不会误触发自动重连 + was_connected = self.connected + self.connected = False + if was_connected: + self.callback("disconnected", "❌ 连接已断开") async def _handle_message(self, data: dict): msg_type = data.get('type', '') @@ -134,12 +140,21 @@ class SimpleCLI: self.ai_full = "" self.ai_prefix_shown = False - # 已显示消息的去重 - self.displayed_messages: set = set() - + # 已显示消息的去重（LRU 上限，防长对话内存无限增长） + self.DISPLAYED_MAX = 500 + self.displayed_messages: OrderedDict = OrderedDict() + # 锁和状态 self.lock = threading.Lock() self.current_input = "" + + # 标准输入：单守护线程阻塞读 + asyncio.Queue 中转（避免 Ctrl+C 泄漏线程） + self._input_queue: asyncio.Queue = asyncio.Queue() + self._stdin_thread: Optional[threading.Thread] = None + + # 自动重连状态 + self._manual_disconnect = False # 用户主动 /disconnect、/quit、/connect 后置 True，抑制自动重连 + self._reconnect_task: Optional[asyncio.Task] = None # 流式输出状态 self.streaming_line = False # 是否正在流式输出行 @@ -275,6 +290,13 @@ class SimpleCLI: self.ai_previous = clean self.ai_full = clean + def _remember_displayed(self, key: str): + """记录已显示消息（LRU：超 DISPLAYED_MAX 淘汰最旧，防长对话内存膨胀）。""" + self.displayed_messages[key] = None + self.displayed_messages.move_to_end(key) + if len(self.displayed_messages) > self.DISPLAYED_MAX: + self.displayed_messages.popitem(last=False) + def display(self, msg_type: str, content: str, extra: Any = None): """显示消息""" if msg_type == 'chunk': @@ -302,6 +324,11 @@ class SimpleCLI: elif msg_type == 'system': self._finish_stream() self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}") + elif msg_type == 'disconnected': + self.connected = False + self._finish_stream() + self._print_line(f"{Colors.SYSTEM}{content}{Colors.RESET}") + self._schedule_reconnect() elif msg_type == 'info': self._finish_stream() self._print_line(f"{Colors.INFO}ℹ️ {content}{Colors.RESET}") @@ -406,7 +433,7 @@ class SimpleCLI: return # 只有当内容没有被流式显示时才显示完整消息 - self.displayed_messages.add(key) + self._remember_displayed(key) # 如果之前没有流式输出，现在显示完整消息 self._print_line(f"{Colors.AI}🤖 AI助手: {clean}") @@ -466,7 +493,7 @@ class SimpleCLI: clean = content.strip() key = clean[:200] if key not in self.displayed_messages: - self.displayed_messages.add(key) + self._remember_displayed(key) self._print_line(f"{Colors.TASK}⚡ 任务 - {task_title}: {clean}") del self._task_buffers[task_title] self.streaming_line = False @@ -519,10 +546,15 @@ class SimpleCLI: self.display("error", "用法: /connect ws://host:port") return uri = parts[1] + # 切换服务器：抑制旧连接的自动重连 + self._manual_disconnect = True + if self._reconnect_task and not self._reconnect_task.done(): + self._reconnect_task.cancel() if self.client and self.client.connected: await self.client.disconnect() self.client = WebSocketClient(uri, self.display) if await self.client.connect(): + self._manual_disconnect = False self.connected = True self.display("system", f"✅ 已连接到 {uri}") asyncio.create_task(self.client.receive_messages()) @@ -530,6 +562,9 @@ class SimpleCLI: self.client = None self.connected = False elif cmd == '/disconnect': + self._manual_disconnect = True # 主动断开，不自动重连 + if self._reconnect_task and not self._reconnect_task.done(): + self._reconnect_task.cancel() if self.client: await self.client.disconnect() self.client = None @@ -556,6 +591,9 @@ class SimpleCLI: elif cmd == '/quit': self.display("system", "👋 再见！") self.running = False + self._manual_disconnect = True + if self._reconnect_task and not self._reconnect_task.done(): + self._reconnect_task.cancel() if self.client: await self.client.disconnect() return @@ -569,6 +607,35 @@ class SimpleCLI: self.display("error", f"未知命令: {cmd}") self._print_prompt(clear_input=True) + def _schedule_reconnect(self): + """启动自动重连（若未手动断开且当前无进行中的重连任务）。""" + if self._manual_disconnect or not self.running or not self.client: + return + if self._reconnect_task and not self._reconnect_task.done(): + return + try: + self._reconnect_task = asyncio.create_task(self._auto_reconnect_loop(self.client)) + except RuntimeError: + pass # 事件循环正在关闭（程序退出中），放弃重连 + + async def _auto_reconnect_loop(self, client: WebSocketClient): + """断线自动重连：指数退避 2s→30s 封顶，直到成功、/disconnect 或退出。""" + delay = 2 + while self.running and not self._manual_disconnect and self.client is client: + self.display("info", f"🔄 {delay}s 后自动重连 {client.uri} …（/disconnect 取消）") + try: + await asyncio.sleep(delay) + except asyncio.CancelledError: + return + if not self.running or self._manual_disconnect or self.client is not client: + return + if await client.connect(): + self.connected = True + self.display("system", f"✅ 已自动重连到 {client.uri}") + asyncio.create_task(client.receive_messages()) + return + delay = min(delay * 2, 30) + async def _send_interrupt(self): if self.client and self.client.connected: interrupt_msg = {"type": "interrupt", "content": "用户按下了 Ctrl+C 中断当前操作"} @@ -601,30 +668,57 @@ class SimpleCLI: await self.client.send_message(message) # 不在这里显示提示符，等待AI响应结束后再显示 + def _stdin_reader(self, loop: asyncio.AbstractEventLoop): + """独立守护线程：阻塞读 stdin，经 call_soon_threadsafe 投入队列。 + + 进程生命周期只启动一次——替代原 run_in_executor 每行一个线程的方案， + 避免 Ctrl+C 取消 await 后旧线程仍阻塞在 readline 上无限泄漏。 + """ + def _put(item): + try: + loop.call_soon_threadsafe(self._input_queue.put_nowait, item) + except RuntimeError: + pass # 事件循环已关闭（程序退出中） + + while self.running: + try: + line = sys.stdin.readline() + except Exception: + break + if not line: # EOF (Ctrl+D) + _put(None) + break + _put(line) + async def run(self): self.print_banner() self._task_buffers = {} - + + loop = asyncio.get_running_loop() + if self._stdin_thread is None: + self._stdin_thread = threading.Thread( + target=self._stdin_reader, args=(loop,), + daemon=True, name="stdin-reader") + self._stdin_thread.start() + while self.running: try: - line = await asyncio.get_event_loop().run_in_executor( - None, lambda: sys.stdin.readline() - ) - if not line: + line = await self._input_queue.get() + if line is None: # EOF (Ctrl+D) break - + self.current_input = line.rstrip('\n') line = line.strip() - + if not line: self._print_prompt(clear_input=True) continue - + if line.startswith('/'): await self.handle_command(line) else: await self.send_user_message(line) - + except KeyboardInterrupt: if self.client and self.client.connected: await self._send_interrupt()