feat(chat): 改进异步任务管理和取消处理

- 在 StreamLoopManager、MessageManager 和 DefaultReplyer 中增加子任务跟踪机制
- 统一处理 asyncio.CancelledError 异常,确保任务取消时能正确清理子任务
- 使用 child_tasks 集合管理子任务生命周期,防止任务泄漏
- 优化记忆存储等后台任务的创建方式,支持优雅取消
- 改进错误处理逻辑,确保异常情况下也能清理子任务资源
This commit is contained in:
Windpicker-owo
2025-10-23 19:31:23 +08:00
parent 5b29399d94
commit ca69b8affe
3 changed files with 92 additions and 30 deletions

View File

@@ -351,7 +351,7 @@ class StreamLoopManager:
return False
async def _process_stream_messages(self, stream_id: str, context: StreamContext) -> bool:
"""处理流消息
"""处理流消息 - 支持子任务管理
Args:
stream_id: 流ID
@@ -367,6 +367,9 @@ class StreamLoopManager:
# 设置处理状态为正在处理
self._set_stream_processing_status(stream_id, True)
# 子任务跟踪
child_tasks = set()
try:
start_time = time.time()
@@ -375,6 +378,11 @@ class StreamLoopManager:
if cached_messages:
logger.info(f"处理开始前刷新缓存消息: stream={stream_id}, 数量={len(cached_messages)}")
# 创建子任务用于刷新能量(不阻塞主流程)
energy_task = asyncio.create_task(self._refresh_focus_energy(stream_id))
child_tasks.add(energy_task)
energy_task.add_done_callback(lambda t: child_tasks.discard(t))
# 直接调用chatter_manager处理流上下文
task = asyncio.create_task(self.chatter_manager.process_stream_context(stream_id, context))
self.chatter_manager.set_processing_task(stream_id, task)
@@ -387,7 +395,6 @@ class StreamLoopManager:
if additional_messages:
logger.info(f"处理完成后刷新新消息: stream={stream_id}, 数量={len(additional_messages)}")
asyncio.create_task(self._refresh_focus_energy(stream_id))
process_time = time.time() - start_time
logger.debug(f"流处理成功: {stream_id} (耗时: {process_time:.2f}s)")
else:
@@ -395,8 +402,19 @@ class StreamLoopManager:
return success
except asyncio.CancelledError:
logger.info(f"流处理被取消: {stream_id}")
# 取消所有子任务
for child_task in child_tasks:
if not child_task.done():
child_task.cancel()
raise
except Exception as e:
logger.error(f"流处理异常: {stream_id} - {e}", exc_info=True)
# 异常时也要清理子任务
for child_task in child_tasks:
if not child_task.done():
child_task.cancel()
return False
finally:
# 无论成功或失败,都要设置处理状态为未处理