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

This commit is contained in:
InfinityPacer
2026-08-22 10:09:01 +08:00
committed by GitHub
parent 813ea77bc4
commit 1c07333909
6 changed files with 656 additions and 31 deletions
+286 -27
View File
@@ -2448,6 +2448,7 @@ class _MessageTask:
agent_factory: Optional[Callable[..., MoviePilotAgent]] = None agent_factory: Optional[Callable[..., MoviePilotAgent]] = None
agent_setup: Optional[Callable[[MoviePilotAgent], None]] = None agent_setup: Optional[Callable[[MoviePilotAgent], None]] = None
completion_future: Optional[asyncio.Future] = None completion_future: Optional[asyncio.Future] = None
enqueued_at: Optional[float] = None
class AgentManagerUnavailableError(RuntimeError): class AgentManagerUnavailableError(RuntimeError):
@@ -2456,6 +2457,23 @@ class AgentManagerUnavailableError(RuntimeError):
code = "agent_manager_unavailable" 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: class AgentManager:
""" """
AI智能体管理器 AI智能体管理器
@@ -2473,6 +2491,14 @@ class AgentManager:
self._idle_cleanup_task: Optional[asyncio.Task] = None self._idle_cleanup_task: Optional[asyncio.Task] = None
self._idle_session_ttl = timedelta(hours=24) self._idle_session_ttl = timedelta(hours=24)
self._idle_cleanup_interval = 60 * 60 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。 # 接收门禁与队列写入共用一把锁,确保关闭开始后不会再创建 worker。
self._lifecycle_lock = asyncio.Lock() self._lifecycle_lock = asyncio.Lock()
self._accepting_tasks = False self._accepting_tasks = False
@@ -2494,6 +2520,20 @@ class AgentManager:
queue = self._session_queues.get(session_id) queue = self._session_queues.get(session_id)
status["pending_messages"] = queue.qsize() if queue else 0 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"] = ( status["is_processing"] = (
session_id in self._session_workers session_id in self._session_workers
and not self._session_workers[session_id].done() and not self._session_workers[session_id].done()
@@ -2537,6 +2577,8 @@ class AgentManager:
关闭管理器 关闭管理器
""" """
async with self._lifecycle_lock: async with self._lifecycle_lock:
if self._close_finalizer_task and not self._close_finalizer_task.done():
return
# 门禁必须先关闭;锁内完成清理可阻止等待中的请求在收口期间重新入队。 # 门禁必须先关闭;锁内完成清理可阻止等待中的请求在收口期间重新入队。
self._accepting_tasks = False self._accepting_tasks = False
if self._idle_cleanup_task: if self._idle_cleanup_task:
@@ -2546,16 +2588,20 @@ class AgentManager:
except asyncio.CancelledError: except asyncio.CancelledError:
pass pass
self._idle_cleanup_task = None self._idle_cleanup_task = None
# 取消所有会话worker # 取消所有 worker,再以有限等待收口,避免关闭阶段无限挂起。
for task in list(self._session_workers.values()): workers = list(self._session_workers.items())
for session_id, task in workers:
self._session_cancel_requested.add(session_id)
task.cancel() task.cancel()
# 等待所有worker结束 timed_out_workers = []
for session_id, task in list(self._session_workers.items()): for session_id, task in workers:
try: stopped = await self._wait_for_worker_shutdown(
await task session_id,
except asyncio.CancelledError: task,
pass reason="manager_close",
self._session_workers.clear() )
if not stopped:
timed_out_workers.append((session_id, task))
for queue in list(self._session_queues.values()): for queue in list(self._session_queues.values()):
self._discard_queued_messages( self._discard_queued_messages(
queue, queue,
@@ -2563,6 +2609,32 @@ class AgentManager:
) )
self._session_queues.clear() self._session_queues.clear()
self._session_last_used.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()): for agent in list(self.active_agents.values()):
await agent.cleanup() await agent.cleanup()
self.active_agents.clear() self.active_agents.clear()
@@ -2670,16 +2742,45 @@ class AgentManager:
) )
async with self._lifecycle_lock: async with self._lifecycle_lock:
if not self._accepting_tasks: if not self._accepting_tasks:
if completion_future and not completion_future.done():
completion_future.cancel()
raise AgentManagerUnavailableError("AgentManager 未运行或已关闭") 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) self._record_session_activity(session_id, user_id)
# 获取或创建会话队列 # 获取或创建会话队列
if session_id not in self._session_queues: 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 = self._session_queues[session_id]
queue_size = queue.qsize() 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 ( if queue_size > 0 or (
session_id in self._session_workers session_id in self._session_workers
@@ -2690,8 +2791,9 @@ class AgentManager:
f"(队列中待处理: {queue_size} 条)" f"(队列中待处理: {queue_size} 条)"
) )
# 放入队列并创建 worker 与关闭门禁保持原子关系 # 非阻塞入队与 worker 创建在同一生命周期锁内完成,关闭期间不会留下悬挂入队
await queue.put(task) task.enqueued_at = asyncio.get_running_loop().time()
queue.put_nowait(task)
if ( if (
session_id not in self._session_workers session_id not in self._session_workers
or self._session_workers[session_id].done() or self._session_workers[session_id].done()
@@ -2724,10 +2826,29 @@ class AgentManager:
break break
try: 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) await self._start_task_processing_status(task)
result = await self._process_message_internal(task) result = await self._process_message_internal(task)
if task.completion_future and not task.completion_future.done(): 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: except asyncio.CancelledError:
if task.completion_future and not task.completion_future.done(): if task.completion_future and not task.completion_future.done():
if self._accepting_tasks: if self._accepting_tasks:
@@ -2744,6 +2865,8 @@ class AgentManager:
finally: finally:
await self._finish_task_processing_status(task) await self._finish_task_processing_status(task)
queue.task_done() queue.task_done()
if session_id in self._session_cancel_requested:
break
except asyncio.CancelledError: except asyncio.CancelledError:
logger.info(f"会话 {session_id} 的worker被取消") logger.info(f"会话 {session_id} 的worker被取消")
@@ -2752,6 +2875,7 @@ class AgentManager:
current_worker = asyncio.current_task() current_worker = asyncio.current_task()
if self._session_workers.get(session_id) is current_worker: if self._session_workers.get(session_id) is current_worker:
self._session_workers.pop(session_id, None) # noqa self._session_workers.pop(session_id, None) # noqa
self._session_cancel_requested.discard(session_id)
# 如果队列为空,清理队列 # 如果队列为空,清理队列
if ( if (
self._session_queues.get(session_id) is queue self._session_queues.get(session_id) is queue
@@ -2884,15 +3008,17 @@ class AgentManager:
# 先摘下旧队列再等待 worker 退出;lifecycle 锁保证清理期间不会并发建立新队列。 # 先摘下旧队列再等待 worker 退出;lifecycle 锁保证清理期间不会并发建立新队列。
if worker: if worker:
self._session_cancel_requested.add(session_id)
worker.cancel() worker.cancel()
if queue: if queue:
self._discard_queued_messages(queue) self._discard_queued_messages(queue)
if worker: if worker:
try: stopped_cleanly = await self._wait_for_worker_shutdown(
await worker session_id,
except asyncio.CancelledError: worker,
pass reason="stop_current_task",
if self._session_workers.get(session_id) is worker: )
if stopped_cleanly and self._session_workers.get(session_id) is worker:
self._session_workers.pop(session_id, None) # noqa self._session_workers.pop(session_id, None) # noqa
stopped = True stopped = True
if queue: if queue:
@@ -2905,9 +3031,10 @@ class AgentManager:
and not new_queue.empty() and not new_queue.empty()
and (not current_worker or current_worker.done()) and (not current_worker or current_worker.done())
): ):
self._session_workers[session_id] = asyncio.create_task( if session_id not in self._session_shutdown_pending:
self._session_worker(session_id) self._session_workers[session_id] = asyncio.create_task(
) self._session_worker(session_id)
)
if stopped: if stopped:
logger.info(f"会话 {session_id} 的Agent推理已应急停止") 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: async def _clear_session_locked(self, session_id: str, user_id: str) -> None:
"""在 lifecycle 互斥域内释放会话、Agent 与记忆。""" """在 lifecycle 互斥域内释放会话、Agent 与记忆。"""
if session_id in self._session_cleanup_pending:
return
self._session_last_used.pop(session_id, None) self._session_last_used.pop(session_id, None)
# 取消该会话的worker # 取消该会话的worker
if session_id in self._session_workers: if session_id in self._session_workers:
self._session_workers[session_id].cancel() worker = self._session_workers[session_id]
try: self._session_cleanup_pending.add(session_id)
await self._session_workers[session_id] self._session_cancel_requested.add(session_id)
except asyncio.CancelledError: worker.cancel()
pass stopped_cleanly = await self._wait_for_worker_shutdown(
self._session_workers.pop(session_id, None) # noqa 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 调用方永久等待。 # 清理队列时同步结束未执行请求,避免 wait_for_completion 调用方永久等待。
queue = self._session_queues.pop(session_id, None) queue = self._session_queues.pop(session_id, None)
if queue: if queue:
self._discard_queued_messages(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 # 清理agent
if session_id in self.active_agents: if session_id in self.active_agents:
@@ -2948,6 +3097,116 @@ class AgentManager:
memory_manager.clear_memory(session_id, user_id) memory_manager.clear_memory(session_id, user_id)
logger.info(f"会话 {session_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( async def run_background_prompt(
self, self,
message: str, message: str,
+7
View File
@@ -13,6 +13,7 @@ from app.schemas.openai import AnthropicTextBlock as _SchemaAnthropicTextBlock
from app.api.endpoints.openai import ( from app.api.endpoints.openai import (
MODEL_ID, MODEL_ID,
_is_manager_unavailable, _is_manager_unavailable,
_is_manager_queue_full,
_run_managed_agent, _run_managed_agent,
) )
from app.api.openai_utils import ( from app.api.openai_utils import (
@@ -246,6 +247,12 @@ async def messages(
503, 503,
error_type="api_error", 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") return _anthropic_error_response(str(exc), 500, error_type="api_error")
finally: finally:
await manager.clear_session(session_id=session_id, user_id=session_id) await manager.clear_session(session_id=session_id, user_id=session_id)
+19
View File
@@ -346,6 +346,11 @@ def _is_manager_unavailable(error: BaseException) -> bool:
return getattr(error, "code", None) == "agent_manager_unavailable" 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( async def _run_managed_agent(
*, *,
manager, manager,
@@ -552,6 +557,13 @@ async def chat_completions(
error_type="server_error", error_type="server_error",
code="ai_agent_unavailable", 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( return _error_response(
str(exc), str(exc),
500, 500,
@@ -652,6 +664,13 @@ async def responses(
error_type="server_error", error_type="server_error",
code="ai_agent_unavailable", 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( return _error_response(
str(exc), str(exc),
500, 500,
+51 -3
View File
@@ -3,6 +3,7 @@ import base64
import mimetypes import mimetypes
import re import re
import uuid import uuid
from concurrent.futures import CancelledError as FutureCancelledError
from dataclasses import dataclass from dataclasses import dataclass
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
@@ -1136,14 +1137,25 @@ class MessageChain(ChainBase):
else "" 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( 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"当前会话累计 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('model_call_count', 0)}",
f"排队消息数: {status.get('pending_messages', 0)}", f"排队消息数: {pending_text}",
f"最后更新: {status.get('last_updated_at') or '暂无'}", 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) return "\n".join(lines)
def remote_session_status( def remote_session_status(
@@ -1344,11 +1356,47 @@ class MessageChain(ChainBase):
} }
if has_audio_input: if has_audio_input:
process_kwargs["has_audio_input"] = True process_kwargs["has_audio_input"] = True
# 在事件循环中处理 # 在事件循环中处理,并消费跨线程 Future 的失败,避免队列满时静默丢消息。
asyncio.run_coroutine_threadsafe( submission_future = asyncio.run_coroutine_threadsafe(
manager.process_message(**process_kwargs), manager.process_message(**process_kwargs),
global_vars.loop, 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 return True
except Exception as e: except Exception as e:
+260 -1
View File
@@ -5,7 +5,11 @@ import pytest
import app.agent.orchestrator as agent_module import app.agent.orchestrator as agent_module
from app.agent import AgentManager 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.agent.memory import MemoryManager
from app.startup import agent_initializer, modules_initializer 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 == {} 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 @pytest.mark.anyio
async def test_agent_manager_close_serializes_racing_enqueue_and_clear( async def test_agent_manager_close_serializes_racing_enqueue_and_clear(
monkeypatch, monkeypatch,
+33
View File
@@ -1,7 +1,9 @@
import asyncio import asyncio
from concurrent.futures import Future
from unittest.mock import AsyncMock, Mock, patch from unittest.mock import AsyncMock, Mock, patch
from app.agent import MoviePilotAgent from app.agent import MoviePilotAgent
from app.agent.orchestrator import AgentManagerQueueFullError
from app.agent.tools.impl.ask_user_choice import ( from app.agent.tools.impl.ask_user_choice import (
AskUserChoiceTool, AskUserChoiceTool,
UserChoiceOptionInput, UserChoiceOptionInput,
@@ -85,6 +87,37 @@ def test_explicit_ai_message_is_not_recorded_to_message_history():
manager.process_message.assert_called_once() 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(): def test_message_chain_passes_stable_channel_admin_principal_to_agent():
"""消息链应将渠道适配器生成的管理员事实传给 Agent。""" """消息链应将渠道适配器生成的管理员事实传给 Agent。"""
chain = MessageChain() chain = MessageChain()