feat(aibox): WebSocket连接用户独立AI实例，支持多用户并发对话
- 每个连接创建独立DeepSeekChat实例，拥有独立messages/工具注册表/消息锁 - 对话状态与并发控制从全局改为会话级，多用户可同时与AI对话互不干扰 - 中断/重置/重载提示词等命令改为针对当前连接的会话AI - 技能管理器懒加载共享，避免每个连接重复扫描技能
This commit is contained in:
qing committed
1 parent
69a8eee192
commit
d921f84b83
1 file changed
+92
-37
+92
-37
@@ -7669,15 +7669,47 @@ class WebSocketHandler:
|
||||
self.connected_clients = set()
|
||||
self.client_sessions = {}
|
||||
self._processing_tasks = {} # websocket id -> 当前正在处理的asyncio Task
|
||||
self._is_processing_message = False # 是否正在处理用户消息（防止并发）
|
||||
|
||||
self._shared_skill_manager = None # 共享技能管理器（懒加载，供所有连接复用）
|
||||
|
||||
def _get_shared_skill_manager(self):
|
||||
"""获取共享技能管理器（懒加载一次，供所有连接实例复用，避免重复扫描技能）"""
|
||||
if self._shared_skill_manager is None:
|
||||
self._shared_skill_manager = SkillManager(self.logger)
|
||||
return self._shared_skill_manager
|
||||
|
||||
def _create_session_ai(self):
|
||||
"""为一个WebSocket连接创建独立的AI实例，实现不同用户间上下文隔离。
|
||||
|
||||
每个连接拥有独立的 messages / 工具注册表 / 消息锁，可与其他用户并发对话，
|
||||
互不干扰。共享只读资源（记忆库、备忘录、提示词、技能管理器）复用。
|
||||
"""
|
||||
return DeepSeekChat(
|
||||
api_key=None,
|
||||
max_history=constants.DEFAULT_MAX_HISTORY,
|
||||
debug_level=self.logger.level,
|
||||
exit_callback=None, # 会话级对话结束由连接内处理，不触发全局状态
|
||||
history_manager=None, # 不加载共享历史上下文，避免多用户历史相互污染
|
||||
memory_db=self.ai_manager.memory_db,
|
||||
memo_db=self.ai_manager.memo_db,
|
||||
command_executor=self.ai_manager.command_executor,
|
||||
prompt_manager=self.ai_manager.prompt_manager,
|
||||
skill_manager=self._get_shared_skill_manager()
|
||||
)
|
||||
|
||||
def _get_session(self, websocket) -> dict:
|
||||
"""获取指定连接对应的会话状态（不存在则创建空会话）"""
|
||||
return self.client_sessions.get(websocket)
|
||||
|
||||
async def register(self, websocket):
|
||||
"""注册新客户端"""
|
||||
self.connected_clients.add(websocket)
|
||||
self.client_sessions[websocket] = {
|
||||
"connected_at": datetime.now(),
|
||||
"message_count": 0,
|
||||
"current_conversation": None
|
||||
"current_conversation": None,
|
||||
"ai": None, # 每个连接独立的AI实例
|
||||
"dialogue_in_progress": False, # 每个连接独立的对话状态
|
||||
"is_processing": False # 每个连接独立的并发处理标志
|
||||
}
|
||||
self.logger.info(i18n.get('client_connected', websocket.remote_address))
|
||||
|
||||
@@ -7695,7 +7727,14 @@ class WebSocketHandler:
|
||||
"""注销客户端"""
|
||||
self.connected_clients.remove(websocket)
|
||||
if websocket in self.client_sessions:
|
||||
del self.client_sessions[websocket]
|
||||
session = self.client_sessions.pop(websocket)
|
||||
# 释放该连接独立的AI实例资源
|
||||
ai = session.get("ai")
|
||||
if ai:
|
||||
try:
|
||||
ai.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
self.logger.info(i18n.get('client_disconnected', websocket.remote_address))
|
||||
|
||||
async def send_welcome_info(self, websocket):
|
||||
@@ -7850,16 +7889,24 @@ class WebSocketHandler:
|
||||
})
|
||||
|
||||
async def handle_interrupt(self, websocket, data):
|
||||
"""处理中断信号：取消正在进行的AI思考和工具执行"""
|
||||
"""处理中断信号：取消当前连接正在进行的AI思考和工具执行"""
|
||||
self.logger.info("收到客户端中断信号")
|
||||
# 先发送确认，让前端立即得到反馈
|
||||
await self.send_message(websocket, {
|
||||
"type": "interrupt_ack",
|
||||
"content": "已收到中断信号，正在停止..."
|
||||
})
|
||||
# 直接同步调用cancel()：cancel是线程安全的（只set Event + close连接 + kill进程），执行极快
|
||||
# 无需通过run_in_executor提交到线程池，避免额外调度延迟
|
||||
self.ai_manager.cancel()
|
||||
# 取消当前连接独立AI实例的对话（线程安全，只set Event + close连接 + kill进程）
|
||||
session = self._get_session(websocket)
|
||||
ai = session.get("ai") if session else None
|
||||
if ai:
|
||||
ai.cancel()
|
||||
# 同时确保命令执行器中的进程也被终止
|
||||
if self.ai_manager.command_executor and hasattr(self.ai_manager.command_executor, 'cancel'):
|
||||
try:
|
||||
self.ai_manager.command_executor.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def handle_task_message(self, websocket, data):
|
||||
"""处理任务相关消息 - 简化版，所有任务都是对话任务"""
|
||||
@@ -7967,11 +8014,11 @@ class WebSocketHandler:
|
||||
await websocket.close()
|
||||
|
||||
elif command == "/new":
|
||||
# 【新增】开始新对话：保存当前对话，重置AI状态，重新加载系统提示词
|
||||
self.ai_manager.history_manager.archive_current_conversation()
|
||||
if self.ai_manager.ai:
|
||||
ai = self.ai_manager.ai
|
||||
# 如果AI正在处理（可能是调度器在跑定时任务），先取消
|
||||
# 【新增】开始新对话：重置当前连接的AI状态，重新加载系统提示词
|
||||
session = self._get_session(websocket)
|
||||
ai = session.get("ai") if session else None
|
||||
if ai:
|
||||
# 如果该连接AI正在处理，先取消
|
||||
if ai.is_busy():
|
||||
self.logger.info("/new: AI正在处理中，先取消当前对话...")
|
||||
ai.cancel()
|
||||
@@ -7981,8 +8028,9 @@ class WebSocketHandler:
|
||||
break
|
||||
# 重置对话状态，重新加载系统提示词
|
||||
ai.reset_conversation(reload_prompts=True)
|
||||
# 重置对话进行中标记
|
||||
self.ai_manager.dialogue_in_progress = False
|
||||
# 重置该连接对话进行中标记
|
||||
if session:
|
||||
session["dialogue_in_progress"] = False
|
||||
# 重新初始化提示词管理器（确保最新）
|
||||
self.ai_manager.prompt_manager.reload()
|
||||
self.logger.info(i18n.get('new_conversation_reset'))
|
||||
@@ -8077,8 +8125,10 @@ class WebSocketHandler:
|
||||
})
|
||||
|
||||
elif command == "/reload_prompts":
|
||||
if self.ai_manager.ai:
|
||||
new_count = self.ai_manager.ai.reload_prompts()
|
||||
session = self._get_session(websocket)
|
||||
ai = session.get("ai") if session else None
|
||||
if ai:
|
||||
new_count = ai.reload_prompts()
|
||||
else:
|
||||
new_count = self.ai_manager.prompt_manager.reload()
|
||||
await self.send_message(websocket, {
|
||||
@@ -8163,34 +8213,40 @@ class WebSocketHandler:
|
||||
await self.handle_user_message(websocket, {"type": "message", "content": command})
|
||||
|
||||
async def handle_user_message(self, websocket, data):
|
||||
"""处理用户消息 - 实时流式输出"""
|
||||
"""处理用户消息 - 实时流式输出（每个连接使用独立的AI实例，支持多用户并发）"""
|
||||
content = data.get("content", "").strip()
|
||||
|
||||
if not content:
|
||||
return
|
||||
|
||||
# 防止并发处理用户消息：如果正在处理中，提示用户先中断
|
||||
if self._is_processing_message:
|
||||
# 获取当前连接的会话状态
|
||||
session = self._get_session(websocket)
|
||||
if session is None:
|
||||
return
|
||||
|
||||
# 防止并发处理用户消息（按连接独立）：如果正在处理中，提示用户先中断
|
||||
if session.get("is_processing"):
|
||||
await self.send_message(websocket, {
|
||||
"type": "info",
|
||||
"content": "⏳ AI正在响应中，请等待完成或发送中断信号（Ctrl+C）"
|
||||
})
|
||||
return
|
||||
|
||||
self._is_processing_message = True
|
||||
session["is_processing"] = True
|
||||
|
||||
self.logger.info(f"收到用户消息: {content[:50]}...")
|
||||
|
||||
# 确保AI已初始化
|
||||
if not self.ai_manager.ai:
|
||||
self.logger.info("初始化AI...")
|
||||
self.ai_manager.add_ai()
|
||||
# 获取该连接独立的AI实例（懒创建）
|
||||
ai = session.get("ai")
|
||||
if ai is None:
|
||||
self.logger.info("初始化独立AI实例...")
|
||||
ai = self._create_session_ai()
|
||||
session["ai"] = ai
|
||||
|
||||
# 如果还没有活跃对话，开始新对话
|
||||
if not self.ai_manager.dialogue_in_progress:
|
||||
self.logger.info("开始新对话")
|
||||
self.ai_manager.dialogue_in_progress = True
|
||||
self.ai_manager.history_manager.start_new_conversation(content[:50] + "...")
|
||||
# 如果还没有活跃对话，开始新对话（对话状态按连接独立）
|
||||
if not session.get("dialogue_in_progress"):
|
||||
self.logger.info(f"开始新对话（连接 {id(websocket)}）")
|
||||
session["dialogue_in_progress"] = True
|
||||
|
||||
# 发送对话开始消息
|
||||
await self.send_message(websocket, {
|
||||
@@ -8226,18 +8282,17 @@ class WebSocketHandler:
|
||||
# 在线程池中运行AI思考，通过回调实时发送消息
|
||||
response = await loop.run_in_executor(
|
||||
None,
|
||||
self.ai_manager.ai.think_and_respond,
|
||||
ai.think_and_respond,
|
||||
content,
|
||||
output_callback # 传递回调函数，AI会在生成每个chunk时调用
|
||||
)
|
||||
|
||||
self.logger.info(f"AI思考完成，完整响应长度: {len(response) if response else 0}")
|
||||
|
||||
# 检查对话是否结束
|
||||
if not self.ai_manager.ai.conversation_active:
|
||||
# 检查对话是否结束（按连接独立）
|
||||
if not ai.conversation_active:
|
||||
self.logger.info("对话结束")
|
||||
self.ai_manager.dialogue_in_progress = False
|
||||
self.ai_manager.history_manager.end_current_conversation()
|
||||
session["dialogue_in_progress"] = False
|
||||
await self.send_message(websocket, {
|
||||
"type": "info",
|
||||
"content": i18n.get('conversation_ended')
|
||||
@@ -8260,8 +8315,8 @@ class WebSocketHandler:
|
||||
except:
|
||||
pass
|
||||
finally:
|
||||
# 无论成功、失败还是被中断，都清除处理标志
|
||||
self._is_processing_message = False
|
||||
# 无论成功、失败还是被中断，都清除该连接的处理标志
|
||||
session["is_processing"] = False
|
||||
|
||||
|
||||
# ==================== 对话管理器（修改版：添加重置对话功能） ====================
|
||||
|
||||
Reference in new issue
Block a user