diff --git a/app/agent/orchestrator.py b/app/agent/orchestrator.py index 13a43b80d..7053d9acd 100644 --- a/app/agent/orchestrator.py +++ b/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, diff --git a/app/api/endpoints/anthropic.py b/app/api/endpoints/anthropic.py index b9b72df5d..43bd1e06e 100644 --- a/app/api/endpoints/anthropic.py +++ b/app/api/endpoints/anthropic.py @@ -13,6 +13,7 @@ from app.schemas.openai import AnthropicTextBlock as _SchemaAnthropicTextBlock from app.api.endpoints.openai import ( MODEL_ID, _is_manager_unavailable, + _is_manager_queue_full, _run_managed_agent, ) from app.api.openai_utils import ( @@ -246,6 +247,12 @@ async def messages( 503, error_type="api_error", ) + if _is_manager_queue_full(exc): + return _anthropic_error_response( + str(exc), + 429, + error_type="rate_limit_error", + ) return _anthropic_error_response(str(exc), 500, error_type="api_error") finally: await manager.clear_session(session_id=session_id, user_id=session_id) diff --git a/app/api/endpoints/openai.py b/app/api/endpoints/openai.py index f9cdc03c2..5368badd9 100644 --- a/app/api/endpoints/openai.py +++ b/app/api/endpoints/openai.py @@ -346,6 +346,11 @@ def _is_manager_unavailable(error: BaseException) -> bool: return getattr(error, "code", None) == "agent_manager_unavailable" +def _is_manager_queue_full(error: BaseException) -> bool: + """识别 Agent 会话排队已满，供兼容 API 返回可重试状态。""" + return getattr(error, "code", None) == "agent_manager_queue_full" + + async def _run_managed_agent( *, manager, @@ -552,6 +557,13 @@ async def chat_completions( error_type="server_error", code="ai_agent_unavailable", ) + if _is_manager_queue_full(exc): + return _error_response( + str(exc), + 429, + error_type="rate_limit_error", + code="ai_agent_queue_full", + ) return _error_response( str(exc), 500, @@ -652,6 +664,13 @@ async def responses( error_type="server_error", code="ai_agent_unavailable", ) + if _is_manager_queue_full(exc): + return _error_response( + str(exc), + 429, + error_type="rate_limit_error", + code="ai_agent_queue_full", + ) return _error_response( str(exc), 500, diff --git a/app/chain/message.py b/app/chain/message.py index 135669742..6d01790c2 100644 --- a/app/chain/message.py +++ b/app/chain/message.py @@ -3,6 +3,7 @@ import base64 import mimetypes import re import uuid +from concurrent.futures import CancelledError as FutureCancelledError from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -1136,14 +1137,25 @@ class MessageChain(ChainBase): else "" ), ) + pending_messages = status.get("pending_messages", 0) + queue_capacity = status.get("queue_capacity") + pending_text = ( + f"{pending_messages} / {queue_capacity}" + if queue_capacity + else str(pending_messages) + ) lines.extend( [ f"当前会话累计 tokens: 输入 {cls._format_token_count(status.get('total_input_tokens'))} / 输出 {cls._format_token_count(status.get('total_output_tokens'))} / 总计 {cls._format_token_count(status.get('total_tokens'))}", f"模型调用次数: {status.get('model_call_count', 0)}", - f"排队消息数: {status.get('pending_messages', 0)}", + f"排队消息数: {pending_text}", f"最后更新: {status.get('last_updated_at') or '暂无'}", ] ) + if status.get("queue_rejections"): + lines.append(f"排队拒绝次数: {status['queue_rejections']}") + if status.get("shutdown_pending"): + lines.append("会话状态: 正在停止") return "\n".join(lines) def remote_session_status( @@ -1344,11 +1356,47 @@ class MessageChain(ChainBase): } if has_audio_input: process_kwargs["has_audio_input"] = True - # 在事件循环中处理 - asyncio.run_coroutine_threadsafe( + # 在事件循环中处理，并消费跨线程 Future 的失败，避免队列满时静默丢消息。 + submission_future = asyncio.run_coroutine_threadsafe( manager.process_message(**process_kwargs), global_vars.loop, ) + + def _report_agent_submission_failure(completed) -> None: + try: + completed.result() + except BaseException as error: + if isinstance( + error, + (asyncio.CancelledError, FutureCancelledError), + ): + return + error_code = getattr(error, "code", None) + if error_code == "agent_manager_queue_full": + title = "智能助手当前排队已满，请稍后重试" + elif error_code == "agent_manager_unavailable": + title = "智能助手服务暂不可用，请稍后重试" + else: + title = "智能助手处理失败，请查看日志" + logger.warning(f"Agent 消息提交失败: {error}") + try: + self.post_message( + Message( + channel=channel, + source=source, + userid=userid, + username=username, + title=title, + original_message_id=original_message_id, + original_chat_id=original_chat_id, + save_history=False, + ) + ) + except Exception as report_error: + logger.error(f"发送 Agent 提交失败提示失败: {report_error}") + + if submission_future is not None: + submission_future.add_done_callback(_report_agent_submission_failure) return True except Exception as e: diff --git a/tests/test_agent_lifecycle.py b/tests/test_agent_lifecycle.py index 39d61077e..2ec843909 100644 --- a/tests/test_agent_lifecycle.py +++ b/tests/test_agent_lifecycle.py @@ -5,7 +5,11 @@ import pytest import app.agent.orchestrator as agent_module from app.agent import AgentManager -from app.agent.orchestrator import AgentManagerUnavailableError +from app.agent.orchestrator import ( + AGENT_SESSION_QUEUE_MAX_SIZE, + AgentManagerQueueFullError, + AgentManagerUnavailableError, +) from app.agent.memory import MemoryManager from app.startup import agent_initializer, modules_initializer @@ -186,6 +190,261 @@ async def test_agent_manager_acceptance_gate_rejects_stale_references( assert manager.active_agents == {} +@pytest.mark.anyio +async def test_agent_manager_rejects_messages_when_session_queue_is_full( + monkeypatch, +) -> None: + """会话达到待处理容量后应立即拒绝，不得在生命周期锁内无限等待。""" + manager = AgentManager() + memory_manager = MemoryManager() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + started = asyncio.Event() + + async def block_current(_task): + started.set() + await asyncio.Event().wait() + + manager._process_message_internal = block_current + await manager.initialize() + current = asyncio.create_task( + manager.process_message( + "bounded", + "1", + "current", + wait_for_completion=True, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + queued = [ + asyncio.create_task( + manager.process_message( + "bounded", + "1", + f"queued-{index}", + wait_for_completion=True, + ) + ) + for index in range(AGENT_SESSION_QUEUE_MAX_SIZE) + ] + await asyncio.sleep(0) + + with pytest.raises(AgentManagerQueueFullError) as error_info: + await manager.process_message( + "bounded", + "1", + "rejected", + wait_for_completion=True, + ) + assert error_info.value.code == "agent_manager_queue_full" + + status = manager.get_session_status("bounded") + assert status["pending_messages"] == AGENT_SESSION_QUEUE_MAX_SIZE + assert status["queue_capacity"] == AGENT_SESSION_QUEUE_MAX_SIZE + assert status["queue_saturated"] is True + assert status["queue_rejections"] == 1 + + await manager.close() + results = await asyncio.gather(current, *queued, return_exceptions=True) + assert all(isinstance(result, AgentManagerUnavailableError) for result in results) + + +@pytest.mark.anyio +async def test_agent_manager_records_queue_wait_time(monkeypatch) -> None: + """任务开始执行后应保留最近一次排队等待的可观测值。""" + manager = AgentManager() + memory_manager = MemoryManager() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + started = asyncio.Event() + + async def process(_task): + started.set() + return "done" + + manager._process_message_internal = process + await manager.initialize() + assert await manager.process_message( + "queue-observe", + "1", + "message", + wait_for_completion=True, + ) == "done" + await asyncio.wait_for(started.wait(), timeout=1) + + status = manager.get_session_status("queue-observe") + assert status["last_queue_wait_ms"] >= 0 + await manager.close() + + +@pytest.mark.anyio +async def test_agent_manager_rejects_new_messages_while_worker_shutdown_is_pending( + monkeypatch, +) -> None: + """worker 未在关停上限内收敛时，同一会话必须保持停止态。""" + manager = AgentManager() + manager._shutdown_timeout = 0.01 + memory_manager = MemoryManager() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + started = asyncio.Event() + release = asyncio.Event() + + async def ignore_cancellation(_task): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + + manager._process_message_internal = ignore_cancellation + await manager.initialize() + execution = asyncio.create_task( + manager.process_message( + "shutdown-boundary", + "1", + "current", + wait_for_completion=True, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + assert await manager.stop_current_task("shutdown-boundary") is True + status = manager.get_session_status("shutdown-boundary") + assert status["shutdown_pending"] is True + with pytest.raises(AgentManagerUnavailableError): + await manager.process_message( + "shutdown-boundary", + "1", + "late", + ) + + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(execution, timeout=1) + for _ in range(20): + if not manager.get_session_status("shutdown-boundary")["shutdown_pending"]: + break + await asyncio.sleep(0) + assert manager.get_session_status("shutdown-boundary")["shutdown_pending"] is False + await manager.close() + + +@pytest.mark.anyio +async def test_clear_session_defers_agent_cleanup_until_worker_finishes( + monkeypatch, +) -> None: + """clear_session 超时期间不得清理仍被 worker 使用的 Agent 和记忆。""" + manager = AgentManager() + manager._shutdown_timeout = 0.01 + memory_manager = MemoryManager() + memory_manager.clear_memory = MagicMock() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + started = asyncio.Event() + release = asyncio.Event() + cleanup_called = asyncio.Event() + + class BlockingAgent: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + async def process(self, _message, **_kwargs): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + + async def cleanup(self): + cleanup_called.set() + + def get_session_status(self): + return {} + + await manager.initialize() + execution = asyncio.create_task( + manager.process_message( + "clear-timeout", + "1", + "current", + agent_factory=BlockingAgent, + wait_for_completion=True, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + await manager.clear_session("clear-timeout", "1") + assert not cleanup_called.is_set() + assert memory_manager.clear_memory.call_count == 0 + assert "clear-timeout" in manager.active_agents + assert manager.get_session_status("clear-timeout")["shutdown_pending"] is True + + release.set() + await asyncio.wait_for(cleanup_called.wait(), timeout=1) + with pytest.raises(asyncio.CancelledError): + await execution + for _ in range(20): + if "clear-timeout" not in manager.active_agents: + break + await asyncio.sleep(0) + assert "clear-timeout" not in manager.active_agents + assert memory_manager.clear_memory.call_count == 1 + await manager.close() + + +@pytest.mark.anyio +async def test_close_defers_shared_agent_teardown_after_worker_timeout( + monkeypatch, +) -> None: + """管理器关闭超时后，旧 worker 收敛前不得拆除共享 Agent 资源。""" + manager = AgentManager() + manager._shutdown_timeout = 0.01 + memory_manager = MemoryManager() + monkeypatch.setattr(agent_module, "memory_manager", memory_manager) + started = asyncio.Event() + release = asyncio.Event() + cleanup_called = asyncio.Event() + + class BlockingAgent: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + async def process(self, _message, **_kwargs): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + + async def cleanup(self): + cleanup_called.set() + + await manager.initialize() + execution = asyncio.create_task( + manager.process_message( + "close-timeout", + "1", + "current", + agent_factory=BlockingAgent, + wait_for_completion=True, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + await manager.close() + assert not cleanup_called.is_set() + assert "close-timeout" in manager.active_agents + assert manager._close_finalizer_task is not None + + release.set() + await asyncio.wait_for(cleanup_called.wait(), timeout=1) + with pytest.raises(asyncio.CancelledError): + await execution + for _ in range(20): + if manager._close_finalizer_task is None: + break + await asyncio.sleep(0) + assert manager.active_agents == {} + + @pytest.mark.anyio async def test_agent_manager_close_serializes_racing_enqueue_and_clear( monkeypatch, diff --git a/tests/test_agent_message_routing.py b/tests/test_agent_message_routing.py index f850bd5c7..a57a165bf 100644 --- a/tests/test_agent_message_routing.py +++ b/tests/test_agent_message_routing.py @@ -1,7 +1,9 @@ import asyncio +from concurrent.futures import Future from unittest.mock import AsyncMock, Mock, patch from app.agent import MoviePilotAgent +from app.agent.orchestrator import AgentManagerQueueFullError from app.agent.tools.impl.ask_user_choice import ( AskUserChoiceTool, UserChoiceOptionInput, @@ -85,6 +87,37 @@ def test_explicit_ai_message_is_not_recorded_to_message_history(): manager.process_message.assert_called_once() +def test_agent_queue_full_is_reported_to_the_originating_channel(): + """消息队列满时应消费 Future 异常并向原渠道返回可重试提示。""" + chain = MessageChain() + manager = Mock(process_message=AsyncMock()) + failed = Future() + failed.set_exception(AgentManagerQueueFullError("session-1", 8)) + + def submit(coro, _loop): + coro.close() + return failed + + with patch.object(settings, "AI_AGENT_ENABLE", True), patch( + "app.chain.message.get_running_agent_manager", return_value=manager + ), patch( + "app.chain.message.asyncio.run_coroutine_threadsafe", + side_effect=submit, + ), patch.object(chain, "post_message") as post_message: + assert chain._handle_ai_message( + text="/ai 检查状态", + channel=NotificationChannel.Telegram, + source="telegram-test", + userid="10001", + username="tester", + ) is True + + notification = post_message.call_args.args[0] + assert notification.title == "智能助手当前排队已满，请稍后重试" + assert notification.userid == "10001" + assert notification.save_history is False + + def test_message_chain_passes_stable_channel_admin_principal_to_agent(): """消息链应将渠道适配器生成的管理员事实传给 Agent。""" chain = MessageChain()