fix(aibox): 修复中断信号延迟生效问题,确保Ctrl+C能立即中断AI响应
根因:WebSocket消息接收循环使用await串行处理,导致interrupt消息在 AI整段对话完成后才被读取和处理。 修复内容: 1. 消息接收与处理解耦:用create_task异步处理消息,不再await阻塞接收循环 2. cancel()直接同步调用而非通过run_in_executor,消除线程池调度延迟 3. 强制关闭response.raw底层socket连接,中断阻塞的iter_lines() recv() 4. 添加_is_processing_message标志防止多消息并发污染对话历史 5. 连接关闭时自动取消正在进行的处理任务
This commit is contained in:
1 parent
f319686571
commit
f701730016
1 file changed
+44
-8
+44
-8
@@ -6792,10 +6792,19 @@ class DeepSeekChat:
|
||||
"""中断当前正在进行的AI思考/工具执行。线程安全,可从任意线程调用。"""
|
||||
self.logger.info("收到中断信号,正在取消当前操作...")
|
||||
self._cancel_event.set()
|
||||
# 关闭当前HTTP响应,中断 iter_lines() 阻塞
|
||||
# 强制关闭当前HTTP响应,立即中断 iter_lines() 的阻塞读取
|
||||
# 仅调用 response.close() 可能无法中断底层socket recv(),需要先关闭raw连接
|
||||
if self._current_response is not None:
|
||||
try:
|
||||
self._current_response.close()
|
||||
resp = self._current_response
|
||||
self._current_response = None
|
||||
# 先关闭底层raw socket连接,强制阻塞的recv()立即抛出异常
|
||||
if hasattr(resp, 'raw') and resp.raw is not None:
|
||||
try:
|
||||
resp.raw.close()
|
||||
except Exception:
|
||||
pass
|
||||
resp.close()
|
||||
except Exception:
|
||||
pass
|
||||
# 终止正在运行的子进程
|
||||
@@ -7584,6 +7593,8 @@ class WebSocketHandler:
|
||||
self.logger = Logger("WebSocket")
|
||||
self.connected_clients = set()
|
||||
self.client_sessions = {}
|
||||
self._processing_tasks = {} # websocket id -> 当前正在处理的asyncio Task
|
||||
self._is_processing_message = False # 是否正在处理用户消息(防止并发)
|
||||
|
||||
async def register(self, websocket):
|
||||
"""注册新客户端"""
|
||||
@@ -7771,9 +7782,9 @@ class WebSocketHandler:
|
||||
"type": "interrupt_ack",
|
||||
"content": "已收到中断信号,正在停止..."
|
||||
})
|
||||
# 在事件循环线程外执行cancel(cancel是线程安全的,会设置Event并关闭response/kill进程)
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self.ai_manager.cancel)
|
||||
# 直接同步调用cancel():cancel是线程安全的(只set Event + close连接 + kill进程),执行极快
|
||||
# 无需通过run_in_executor提交到线程池,避免额外调度延迟
|
||||
self.ai_manager.cancel()
|
||||
|
||||
async def handle_task_message(self, websocket, data):
|
||||
"""处理任务相关消息 - 简化版,所有任务都是对话任务"""
|
||||
@@ -8074,6 +8085,16 @@ class WebSocketHandler:
|
||||
if not content:
|
||||
return
|
||||
|
||||
# 防止并发处理用户消息:如果正在处理中,提示用户先中断
|
||||
if self._is_processing_message:
|
||||
await self.send_message(websocket, {
|
||||
"type": "info",
|
||||
"content": "⏳ AI正在响应中,请等待完成或发送中断信号(Ctrl+C)"
|
||||
})
|
||||
return
|
||||
|
||||
self._is_processing_message = True
|
||||
|
||||
self.logger.info(f"收到用户消息: {content[:50]}...")
|
||||
|
||||
# 确保AI已初始化
|
||||
@@ -8115,7 +8136,6 @@ class WebSocketHandler:
|
||||
# 提交到事件循环
|
||||
asyncio.run_coroutine_threadsafe(send_msg(), loop)
|
||||
|
||||
# 让AI思考并回应 - 不等待完整响应,通过回调实时发送
|
||||
try:
|
||||
self.logger.info("开始AI思考...")
|
||||
|
||||
@@ -8155,6 +8175,9 @@ class WebSocketHandler:
|
||||
}, ensure_ascii=False))
|
||||
except:
|
||||
pass
|
||||
finally:
|
||||
# 无论成功、失败还是被中断,都清除处理标志
|
||||
self._is_processing_message = False
|
||||
|
||||
|
||||
# ==================== 对话管理器(修改版:添加重置对话功能) ====================
|
||||
@@ -8319,17 +8342,30 @@ class EnhancedAIManager:
|
||||
print(i18n.get('websocket_started', host, port))
|
||||
|
||||
async def handle_connection(self, websocket, path):
|
||||
"""处理WebSocket连接"""
|
||||
"""处理WebSocket连接 - 消息接收与处理解耦,确保中断信号能被即时处理"""
|
||||
self.logger.info(f"新连接: {websocket.remote_address}, path: {path}")
|
||||
await self.ws_handler.register(websocket)
|
||||
try:
|
||||
async for message in websocket:
|
||||
await self.ws_handler.handle_message(websocket, message)
|
||||
# 【关键修复】用create_task异步处理消息,不await阻塞接收循环
|
||||
# 这样当AI在后台生成响应时,interrupt消息仍能被立即接收和处理
|
||||
task = asyncio.create_task(
|
||||
self.ws_handler.handle_message(websocket, message)
|
||||
)
|
||||
self.ws_handler._processing_tasks[id(websocket)] = task
|
||||
|
||||
def _cleanup_task(t, ws_id=id(websocket)):
|
||||
self.ws_handler._processing_tasks.pop(ws_id, None)
|
||||
task.add_done_callback(_cleanup_task)
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
self.logger.info(f"连接关闭: {e}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"处理错误: {e}")
|
||||
finally:
|
||||
# 连接关闭时取消正在进行的处理任务
|
||||
task = self.ws_handler._processing_tasks.pop(id(websocket), None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
await self.ws_handler.unregister(websocket)
|
||||
|
||||
async def websocket_server(self):
|
||||
|
||||
Reference in new issue
Block a user