fix: 收敛 Agent 会话队列与 worker 关闭 (#6386)

app/agent/orchestrator.py

@@ -2448,6 +2448,7 @@ class _MessageTask:
     agent_factory: Optional[Callable[..., MoviePilotAgent]] = None
     agent_setup: Optional[Callable[[MoviePilotAgent], None]] = None
     completion_future: Optional[asyncio.Future] = None
+    enqueued_at: Optional[float] = None
 
 
 class AgentManagerUnavailableError(RuntimeError):
@@ -2456,6 +2457,23 @@ class AgentManagerUnavailableError(RuntimeError):
     code = "agent_manager_unavailable"
 
 
+class AgentManagerQueueFullError(RuntimeError):
+    """Agent 会话的待处理消息达到容量上限。"""
+
+    code = "agent_manager_queue_full"
+
+    def __init__(self, session_id: str, limit: int):
+        self.session_id = session_id
+        self.limit = limit
+        super().__init__(
+            f"Agent 会话当前排队消息已达上限（{limit} 条），请稍后重试"
+        )
+
+
+AGENT_SESSION_QUEUE_MAX_SIZE = 8
+AGENT_MANAGER_SHUTDOWN_TIMEOUT = 10.0
+
+
 class AgentManager:
     """
     AI智能体管理器
@@ -2473,6 +2491,14 @@ class AgentManager:
         self._idle_cleanup_task: Optional[asyncio.Task] = None
         self._idle_session_ttl = timedelta(hours=24)
         self._idle_cleanup_interval = 60 * 60
+        self._session_queue_rejections: Dict[str, int] = {}
+        self._session_last_queue_wait_ms: Dict[str, float] = {}
+        self._session_shutdown_pending: Dict[str, asyncio.Task] = {}
+        self._session_cleanup_pending: set[str] = set()
+        self._session_deferred_cleanup_tasks: Dict[str, asyncio.Task] = {}
+        self._session_cancel_requested: set[str] = set()
+        self._close_finalizer_task: Optional[asyncio.Task] = None
+        self._shutdown_timeout = AGENT_MANAGER_SHUTDOWN_TIMEOUT
         # 接收门禁与队列写入共用一把锁，确保关闭开始后不会再创建 worker。
         self._lifecycle_lock = asyncio.Lock()
         self._accepting_tasks = False
@@ -2494,6 +2520,20 @@ class AgentManager:
 
         queue = self._session_queues.get(session_id)
         status["pending_messages"] = queue.qsize() if queue else 0
+        status["queue_capacity"] = AGENT_SESSION_QUEUE_MAX_SIZE
+        status["queue_saturated"] = bool(queue and queue.full())
+        status["queue_rejections"] = self._session_queue_rejections.get(
+            session_id,
+            0,
+        )
+        status["last_queue_wait_ms"] = self._session_last_queue_wait_ms.get(
+            session_id,
+            0.0,
+        )
+        pending_shutdown = self._session_shutdown_pending.get(session_id)
+        status["shutdown_pending"] = bool(
+            pending_shutdown and not pending_shutdown.done()
+        )
         status["is_processing"] = (
                 session_id in self._session_workers
                 and not self._session_workers[session_id].done()
@@ -2537,6 +2577,8 @@ class AgentManager:
         关闭管理器
         """
         async with self._lifecycle_lock:
+            if self._close_finalizer_task and not self._close_finalizer_task.done():
+                return
             # 门禁必须先关闭；锁内完成清理可阻止等待中的请求在收口期间重新入队。
             self._accepting_tasks = False
             if self._idle_cleanup_task:
@@ -2546,16 +2588,20 @@ class AgentManager:
                 except asyncio.CancelledError:
                     pass
                 self._idle_cleanup_task = None
-            # 取消所有会话worker
-            for task in list(self._session_workers.values()):
+            # 先取消所有 worker，再以有限等待收口，避免关闭阶段无限挂起。
+            workers = list(self._session_workers.items())
+            for session_id, task in workers:
+                self._session_cancel_requested.add(session_id)
                 task.cancel()
-            # 等待所有worker结束
-            for session_id, task in list(self._session_workers.items()):
-                try:
-                    await task
-                except asyncio.CancelledError:
-                    pass
-            self._session_workers.clear()
+            timed_out_workers = []
+            for session_id, task in workers:
+                stopped = await self._wait_for_worker_shutdown(
+                    session_id,
+                    task,
+                    reason="manager_close",
+                )
+                if not stopped:
+                    timed_out_workers.append((session_id, task))
             for queue in list(self._session_queues.values()):
                 self._discard_queued_messages(
                     queue,
@@ -2563,6 +2609,32 @@ class AgentManager:
                 )
             self._session_queues.clear()
             self._session_last_used.clear()
+            self._session_queue_rejections.clear()
+            self._session_last_queue_wait_ms.clear()
+
+            if timed_out_workers:
+                timed_out_session_ids = {
+                    session_id for session_id, _ in timed_out_workers
+                }
+                for session_id, task in workers:
+                    if session_id in timed_out_session_ids:
+                        continue
+                    if self._session_workers.get(session_id) is task:
+                        self._session_workers.pop(session_id, None)
+                for session_id, agent in list(self.active_agents.items()):
+                    if session_id not in timed_out_session_ids:
+                        await agent.cleanup()
+                        self.active_agents.pop(session_id, None)
+                logger.error(
+                    "AgentManager 关闭时仍有 worker 未收敛，"
+                    f"保留 {len(timed_out_workers)} 个会话资源直到 worker 结束"
+                )
+                self._close_finalizer_task = asyncio.create_task(
+                    self._finish_deferred_close(timed_out_workers)
+                )
+                return
+
+            self._session_workers.clear()
             for agent in list(self.active_agents.values()):
                 await agent.cleanup()
             self.active_agents.clear()
@@ -2670,16 +2742,45 @@ class AgentManager:
         )
         async with self._lifecycle_lock:
             if not self._accepting_tasks:
+                if completion_future and not completion_future.done():
+                    completion_future.cancel()
                 raise AgentManagerUnavailableError("AgentManager 未运行或已关闭")
+            pending_shutdown = self._session_shutdown_pending.get(session_id)
+            if pending_shutdown:
+                if pending_shutdown.done():
+                    self._session_shutdown_pending.pop(session_id, None)
+                else:
+                    if completion_future and not completion_future.done():
+                        completion_future.cancel()
+                    raise AgentManagerUnavailableError(
+                        f"Agent 会话 {session_id} 仍在停止，暂时不能接收新任务"
+                    )
             self._record_session_activity(session_id, user_id)
 
             # 获取或创建会话队列
             if session_id not in self._session_queues:
-                self._session_queues[session_id] = asyncio.Queue()
+                self._session_queues[session_id] = asyncio.Queue(
+                    maxsize=AGENT_SESSION_QUEUE_MAX_SIZE
+                )
 
             queue = self._session_queues[session_id]
             queue_size = queue.qsize()
 
+            if queue.full():
+                self._session_queue_rejections[session_id] = (
+                    self._session_queue_rejections.get(session_id, 0) + 1
+                )
+                logger.warning(
+                    f"会话 {session_id} 的 Agent 排队已满，拒绝新消息 "
+                    f"(上限: {AGENT_SESSION_QUEUE_MAX_SIZE})"
+                )
+                if completion_future and not completion_future.done():
+                    completion_future.cancel()
+                raise AgentManagerQueueFullError(
+                    session_id=session_id,
+                    limit=AGENT_SESSION_QUEUE_MAX_SIZE,
+                )
+
             # 如果队列中已有等待的消息，通知用户消息已排队
             if queue_size > 0 or (
                     session_id in self._session_workers
@@ -2690,8 +2791,9 @@ class AgentManager:
                     f"(队列中待处理: {queue_size} 条)"
                 )
 
-            # 放入队列并创建 worker 与关闭门禁保持原子关系。
-            await queue.put(task)
+            # 非阻塞入队与 worker 创建在同一生命周期锁内完成，关闭期间不会留下悬挂入队。
+            task.enqueued_at = asyncio.get_running_loop().time()
+            queue.put_nowait(task)
             if (
                     session_id not in self._session_workers
                     or self._session_workers[session_id].done()
@@ -2724,10 +2826,29 @@ class AgentManager:
                     break
 
                 try:
+                    if task.enqueued_at is not None:
+                        queue_wait_ms = max(
+                            0.0,
+                            (
+                                asyncio.get_running_loop().time()
+                                - task.enqueued_at
+                            )
+                            * 1000,
+                        )
+                        self._session_last_queue_wait_ms[session_id] = round(
+                            queue_wait_ms,
+                            3,
+                        )
                     await self._start_task_processing_status(task)
                     result = await self._process_message_internal(task)
                     if task.completion_future and not task.completion_future.done():
-                        task.completion_future.set_result(result)
+                        if (
+                                not self._accepting_tasks
+                                or session_id in self._session_cancel_requested
+                        ):
+                            task.completion_future.cancel()
+                        else:
+                            task.completion_future.set_result(result)
                 except asyncio.CancelledError:
                     if task.completion_future and not task.completion_future.done():
                         if self._accepting_tasks:
@@ -2744,6 +2865,8 @@ class AgentManager:
                 finally:
                     await self._finish_task_processing_status(task)
                     queue.task_done()
+                if session_id in self._session_cancel_requested:
+                    break
 
         except asyncio.CancelledError:
             logger.info(f"会话 {session_id} 的worker被取消")
@@ -2752,6 +2875,7 @@ class AgentManager:
             current_worker = asyncio.current_task()
             if self._session_workers.get(session_id) is current_worker:
                 self._session_workers.pop(session_id, None)  # noqa
+            self._session_cancel_requested.discard(session_id)
             # 如果队列为空，清理队列
             if (
                     self._session_queues.get(session_id) is queue
@@ -2884,15 +3008,17 @@ class AgentManager:
 
         # 先摘下旧队列再等待 worker 退出；lifecycle 锁保证清理期间不会并发建立新队列。
         if worker:
+            self._session_cancel_requested.add(session_id)
             worker.cancel()
         if queue:
             self._discard_queued_messages(queue)
         if worker:
-            try:
-                await worker
-            except asyncio.CancelledError:
-                pass
-            if self._session_workers.get(session_id) is worker:
+            stopped_cleanly = await self._wait_for_worker_shutdown(
+                session_id,
+                worker,
+                reason="stop_current_task",
+            )
+            if stopped_cleanly and self._session_workers.get(session_id) is worker:
                 self._session_workers.pop(session_id, None)  # noqa
             stopped = True
         if queue:
@@ -2905,9 +3031,10 @@ class AgentManager:
                 and not new_queue.empty()
                 and (not current_worker or current_worker.done())
         ):
-            self._session_workers[session_id] = asyncio.create_task(
-                self._session_worker(session_id)
-            )
+            if session_id not in self._session_shutdown_pending:
+                self._session_workers[session_id] = asyncio.create_task(
+                    self._session_worker(session_id)
+                )
 
         if stopped:
             logger.info(f"会话 {session_id} 的Agent推理已应急停止")
@@ -2925,20 +3052,42 @@ class AgentManager:
 
     async def _clear_session_locked(self, session_id: str, user_id: str) -> None:
         """在 lifecycle 互斥域内释放会话、Agent 与记忆。"""
+        if session_id in self._session_cleanup_pending:
+            return
         self._session_last_used.pop(session_id, None)
         # 取消该会话的worker
         if session_id in self._session_workers:
-            self._session_workers[session_id].cancel()
-            try:
-                await self._session_workers[session_id]
-            except asyncio.CancelledError:
-                pass
-            self._session_workers.pop(session_id, None)  # noqa
+            worker = self._session_workers[session_id]
+            self._session_cleanup_pending.add(session_id)
+            self._session_cancel_requested.add(session_id)
+            worker.cancel()
+            stopped_cleanly = await self._wait_for_worker_shutdown(
+                session_id,
+                worker,
+                reason="clear_session",
+            )
+            if not stopped_cleanly:
+                queue = self._session_queues.pop(session_id, None)
+                if queue:
+                    self._discard_queued_messages(queue)
+                worker.add_done_callback(
+                    lambda done: self._schedule_deferred_session_cleanup(
+                        session_id,
+                        user_id,
+                        worker,
+                    )
+                )
+                return
+            if self._session_workers.get(session_id) is worker:
+                self._session_workers.pop(session_id, None)  # noqa
+            self._session_cleanup_pending.discard(session_id)
 
         # 清理队列时同步结束未执行请求，避免 wait_for_completion 调用方永久等待。
         queue = self._session_queues.pop(session_id, None)
         if queue:
             self._discard_queued_messages(queue)
+        self._session_queue_rejections.pop(session_id, None)
+        self._session_last_queue_wait_ms.pop(session_id, None)
 
         # 清理agent
         if session_id in self.active_agents:
@@ -2948,6 +3097,116 @@ class AgentManager:
             memory_manager.clear_memory(session_id, user_id)
             logger.info(f"会话 {session_id} 的记忆已清空")
 
+    def _schedule_deferred_session_cleanup(
+            self,
+            session_id: str,
+            user_id: str,
+            worker: asyncio.Task,
+    ) -> None:
+        """worker 超时后延迟释放会话资源，避免与仍在运行的 Agent 竞态。"""
+        if session_id in self._session_shutdown_pending:
+            cleanup_task = asyncio.create_task(
+                self._finish_deferred_session_cleanup(
+                    session_id=session_id,
+                    user_id=user_id,
+                    worker=worker,
+                )
+            )
+            self._session_deferred_cleanup_tasks[session_id] = cleanup_task
+
+    async def _finish_deferred_session_cleanup(
+            self,
+            session_id: str,
+            user_id: str,
+            worker: asyncio.Task,
+    ) -> None:
+        """等待超时 worker 真正结束后，再完成 clear_session 的资源释放。"""
+        try:
+            await worker
+        except BaseException:
+            pass
+        async with self._lifecycle_lock:
+            if self._session_workers.get(session_id) is worker:
+                self._session_workers.pop(session_id, None)
+            self._session_shutdown_pending.pop(session_id, None)
+            self._session_cleanup_pending.discard(session_id)
+            self._session_deferred_cleanup_tasks.pop(session_id, None)
+            self._session_queue_rejections.pop(session_id, None)
+            self._session_last_queue_wait_ms.pop(session_id, None)
+            agent = self.active_agents.pop(session_id, None)
+            if agent:
+                await agent.cleanup()
+                memory_manager.clear_memory(session_id, user_id)
+                logger.info(f"会话 {session_id} 的记忆已清空")
+
+    async def _finish_deferred_close(
+            self,
+            workers: list[tuple[str, asyncio.Task]],
+    ) -> None:
+        """关闭超时后等待遗留 worker，再释放共享 Agent 资源。"""
+        try:
+            await asyncio.gather(
+                *(worker for _, worker in workers),
+                return_exceptions=True,
+            )
+            async with self._lifecycle_lock:
+                for session_id, worker in workers:
+                    if self._session_workers.get(session_id) is worker:
+                        self._session_workers.pop(session_id, None)
+                    agent = self.active_agents.pop(session_id, None)
+                    if agent:
+                        await agent.cleanup()
+                for session_id, agent in list(self.active_agents.items()):
+                    await agent.cleanup()
+                    self.active_agents.pop(session_id, None)
+                self._session_shutdown_pending.clear()
+                self._session_cancel_requested.clear()
+                await memory_manager.close()
+        finally:
+            self._close_finalizer_task = None
+
+    async def _wait_for_worker_shutdown(
+            self,
+            session_id: str,
+            worker: asyncio.Task,
+            *,
+            reason: str,
+    ) -> bool:
+        """有限等待 worker 结束，超时会话保持停止态直到旧 worker 收敛。"""
+        try:
+            await asyncio.wait_for(
+                asyncio.shield(worker),
+                timeout=self._shutdown_timeout,
+            )
+            return True
+        except asyncio.CancelledError:
+            raise
+        except asyncio.TimeoutError:
+            self._session_shutdown_pending[session_id] = worker
+
+            def _clear_pending(done: asyncio.Task) -> None:
+                if (
+                        self._session_shutdown_pending.get(session_id) is done
+                        and session_id not in self._session_cleanup_pending
+                ):
+                    self._session_shutdown_pending.pop(session_id, None)
+                    logger.info(
+                        f"会话 {session_id} 的 Agent worker 已在超时后收敛"
+                    )
+
+            worker.add_done_callback(_clear_pending)
+            logger.error(
+                f"会话 {session_id} 的 Agent worker 关闭超时，"
+                f"已阻止新任务进入，reason={reason}, timeout={self._shutdown_timeout:g}s"
+            )
+            return False
+        except Exception as error:
+            logger.error(
+                f"等待会话 {session_id} 的 Agent worker 关闭失败，"
+                f"reason={reason}: {error}"
+            )
+            return True
+
     async def run_background_prompt(
             self,
             message: str,