From 38fce9ff2d3d396a8a56e60e4372b4ebc05b4594 Mon Sep 17 00:00:00 2001 From: T52T52 <140709873+T52T52@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:28:08 +0800 Subject: [PATCH 1/3] test: add regression tests for background wake session lock race --- tests/test_background_wake_lock.py | 92 ++++++++++++++++++++++++++++++ tests/test_persist_race.py | 84 +++++++++++++++++++++++++++ 2 files changed, 176 insertions(+) create mode 100644 tests/test_background_wake_lock.py create mode 100644 tests/test_persist_race.py diff --git a/tests/test_background_wake_lock.py b/tests/test_background_wake_lock.py new file mode 100644 index 0000000000..bc51c780d4 --- /dev/null +++ b/tests/test_background_wake_lock.py @@ -0,0 +1,92 @@ +""" +回归测试: 后台任务唤醒路径应使用会话锁 + +背景: 用户消息处理路径(internal.py)通过 session_lock_manager 按会话串行化, +但后台任务唤醒路径(_wake_main_agent_for_background_result)直接跑 agent, +未获取会话锁 —— 与用户消息并发处理时导致上下文丢失。 + +本测试断言: 修复后, 唤醒流程必须获取会话锁 (acquire_lock 被调用)。 +修复前该测试失败, 修复后通过。 +""" +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor + + +def _make_run_context(): + """构造最小可用的 run_context mock""" + event = SimpleNamespace( + unified_msg_origin="Pstar:FriendMessage:TEST", + role="friend", + get_extra=lambda key: None, + ) + ctx = SimpleNamespace( + get_config=lambda umo: {}, + get_llm_tool_manager=MagicMock(), + conversation_manager=MagicMock(update_conversation=AsyncMock()), + ) + agent_ctx = SimpleNamespace(event=event, context=ctx) + return SimpleNamespace( + context=agent_ctx, + tool_call_timeout=60, + ) + + +def _make_runner_mock(): + """构造假 agent runner: step_until_done 异步生成器""" + runner = MagicMock() + + async def _step_until_done(*args, **kwargs): + yield None + + runner.step_until_done.side_effect = _step_until_done + runner.get_final_llm_resp.return_value = SimpleNamespace(completion_text="done") + return runner + + +@pytest.mark.asyncio +async def test_background_wake_acquires_session_lock(): + """后台任务唤醒必须获取会话锁(与用户消息路径一致)""" + run_context = _make_run_context() + runner = _make_runner_mock() + + # 用 AsyncMock 追踪 acquire_lock 是否被调用 + lock_mgr = MagicMock() + lock_cm = AsyncMock() + lock_cm.__aenter__.return_value = None + lock_mgr.acquire_lock.return_value = lock_cm + + with ( + # create=True: 当前代码尚未引入 session_lock_manager, patch 尚不存在的属性 + patch("astrbot.core.astr_agent_tool_exec.session_lock_manager", lock_mgr, create=True), + # _get_session_conv / build_main_agent 在函数内部 import, 需 patch 源模块 + patch( + "astrbot.core.astr_main_agent._get_session_conv", + new=AsyncMock(return_value=SimpleNamespace(history="[]", cid="conv-1")), + ), + patch( + "astrbot.core.astr_main_agent.build_main_agent", + new=AsyncMock(return_value=SimpleNamespace(agent_runner=runner)), + ), + patch("astrbot.core.astr_agent_tool_exec.CronMessageEvent"), + patch("astrbot.core.astr_agent_tool_exec.MessageSession"), + ): + await FunctionToolExecutor._wake_main_agent_for_background_result( + run_context, + task_id="task-1", + tool_name="transfer_to_x", + result_text="some result", + tool_args={}, + note="background task finished", + summary_name="Dedicated to subagent `x`", + ) + + # 核心断言: 修复后必须获取会话锁, 且锁粒度为该会话 + lock_mgr.acquire_lock.assert_called_once() + args = lock_mgr.acquire_lock.call_args[0] + assert "Pstar:FriendMessage:TEST" in args, ( + f"会话锁应按 unified_msg_origin 获取, 实际参数: {args}" + ) diff --git a/tests/test_persist_race.py b/tests/test_persist_race.py new file mode 100644 index 0000000000..9680e9a708 --- /dev/null +++ b/tests/test_persist_race.py @@ -0,0 +1,84 @@ +""" +回归测试: persist_agent_history 并发持久化竞态 + +背景: 多个后台任务/定时任务的结果几乎同时到达时, 各自触发一次 +persist_agent_history, 并发执行 "读历史 → 追加 → 写回", 后写覆盖先写, +导致部分结果(以及先前对话上下文)从会话历史中丢失。 + +本测试期望: 并发持久化后, 所有结果都应保留在历史中。 +修复前该测试失败(复现 bug), 修复后通过(验证修复)。 +""" +import asyncio +import json +from types import SimpleNamespace + +import pytest + +from astrbot.core.utils.history_saver import persist_agent_history + + +class FakeConversationManager: + """模拟 ConversationManager: 共享存储 (umo, cid) -> history JSON""" + + def __init__(self): + self.store: dict[tuple, str] = {} + self.update_count = 0 + + async def update_conversation(self, umo, cid, history=None, **kwargs): + await asyncio.sleep(0.05) # 模拟 DB 写入耗时, 放大竞态窗口 + self.store[(umo, cid)] = json.dumps(history, ensure_ascii=False) + self.update_count += 1 + + +def make_req(history: str): + conv = SimpleNamespace(cid="conv-1", history=history) + return SimpleNamespace(conversation=conv) + + +def make_event(umo: str): + return SimpleNamespace(unified_msg_origin=umo) + + +def test_persist_basic(): + """基础功能: 单次持久化正常写入""" + cm = FakeConversationManager() + umo = "test:session:1" + + async def _run(): + await persist_agent_history( + cm, event=make_event(umo), req=make_req("[]"), summary_note="result-x" + ) + + asyncio.run(_run()) + final = json.loads(cm.store[(umo, "conv-1")]) + notes = [m["content"] for m in final if m["role"] == "assistant"] + assert notes == ["result-x"] + + +@pytest.mark.asyncio +async def test_persist_concurrent_keeps_all_results(): + """ + 并发场景: 4 个后台任务结果几乎同时持久化, 不应丢失任何一条。 + + 复现路径: 每个任务在创建请求时读到旧历史快照, 然后并发 persist。 + 修复方案: 同一会话的处理流程(读历史+持久化)串行化(session lock)。 + """ + cm = FakeConversationManager() + umo = "test:session:1" + n = 4 + + # 模拟并发任务: 各自持有创建时读到的历史快照(并发时互不知道对方) + reqs = [make_req(cm.store.get((umo, "conv-1"), "[]")) for _ in range(n)] + events = [make_event(umo)] * n + + await asyncio.gather(*[ + persist_agent_history(cm, event=events[i], req=reqs[i], summary_note=f"result-{i}") + for i in range(n) + ]) + + final = json.loads(cm.store[(umo, "conv-1")]) + saved = [m["content"] for m in final if m["role"] == "assistant"] + expected = [f"result-{i}" for i in range(n)] + lost = [e for e in expected if e not in saved] + assert not lost, f"并发持久化丢失 {len(lost)}/{n} 条结果: {lost}" + assert cm.update_count == n From bb2a6a0c10de7f8c1d0ac982da89faa172fe08b2 Mon Sep 17 00:00:00 2001 From: T52T52 <140709873+T52T52@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:46:16 +0800 Subject: [PATCH 2/3] fix: acquire session lock in background task wake path The background wake path (_wake_main_agent_for_background_result) bypassed session_lock_manager, so it could run concurrently with user message processing on the same conversation, causing stale LLM context and lost conversation history (persist race). Wrap the read-build-run-persist flow in the per-session lock, matching the user message path in internal.py. - test_background_wake_lock: asserts acquire_lock is called - test_persist_race: verifies the lock prevents concurrent overwrites --- astrbot/core/astr_agent_tool_exec.py | 112 ++++++++++++++------------- tests/test_persist_race.py | 36 ++++++--- 2 files changed, 83 insertions(+), 65 deletions(-) diff --git a/astrbot/core/astr_agent_tool_exec.py b/astrbot/core/astr_agent_tool_exec.py index 775ad37c7b..5a9338724b 100644 --- a/astrbot/core/astr_agent_tool_exec.py +++ b/astrbot/core/astr_agent_tool_exec.py @@ -50,6 +50,7 @@ from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.history_saver import persist_agent_history from astrbot.core.utils.image_ref_utils import is_supported_image_ref +from astrbot.core.utils.session_lock import session_lock_manager from astrbot.core.utils.string_utils import normalize_and_dedupe_strings @@ -555,65 +556,66 @@ async def _wake_main_agent_for_background_result( provider_settings=provider_settings, ) - req = ProviderRequest() - conv = await _get_session_conv(event=cron_event, plugin_context=ctx) - req.conversation = conv - context = json.loads(conv.history) - if context: - req.contexts = context - context_dump = req._print_friendly_context() - req.contexts = [] - req.system_prompt += ( - "\n\nBellow is you and user previous conversation history:\n" - f"{context_dump}" - ) + async with session_lock_manager.acquire_lock(event.unified_msg_origin): + req = ProviderRequest() + conv = await _get_session_conv(event=cron_event, plugin_context=ctx) + req.conversation = conv + context = json.loads(conv.history) + if context: + req.contexts = context + context_dump = req._print_friendly_context() + req.contexts = [] + req.system_prompt += ( + "\n\nBellow is you and user previous conversation history:\n" + f"{context_dump}" + ) - bg = json.dumps(extras["background_task_result"], ensure_ascii=False) - req.system_prompt += BACKGROUND_TASK_RESULT_WOKE_SYSTEM_PROMPT.format( - background_task_result=bg - ) - req.prompt = ( - "Proceed according to your system instructions. " - "Output using same language as previous conversation. " - "If you need to deliver the result to the user immediately, " - "you MUST use `send_message_to_user` tool to send the message directly to the user, " - "otherwise the user will not see the result. " - "After completing your task, summarize and output your actions and results. " - ) - if not req.func_tool: - req.func_tool = ToolSet() - req.func_tool.add_tool( - ctx.get_llm_tool_manager().get_builtin_tool(SendMessageToUserTool) - ) + bg = json.dumps(extras["background_task_result"], ensure_ascii=False) + req.system_prompt += BACKGROUND_TASK_RESULT_WOKE_SYSTEM_PROMPT.format( + background_task_result=bg + ) + req.prompt = ( + "Proceed according to your system instructions. " + "Output using same language as previous conversation. " + "If you need to deliver the result to the user immediately, " + "you MUST use `send_message_to_user` tool to send the message directly to the user, " + "otherwise the user will not see the result. " + "After completing your task, summarize and output your actions and results. " + ) + if not req.func_tool: + req.func_tool = ToolSet() + req.func_tool.add_tool( + ctx.get_llm_tool_manager().get_builtin_tool(SendMessageToUserTool) + ) - result = await build_main_agent( - event=cron_event, plugin_context=ctx, config=config, req=req - ) - if not result: - logger.error(f"Failed to build main agent for background task {tool_name}.") - return + result = await build_main_agent( + event=cron_event, plugin_context=ctx, config=config, req=req + ) + if not result: + logger.error( + f"Failed to build main agent for background task {tool_name}." + ) + return - runner = result.agent_runner - async for _ in runner.step_until_done(30): - # agent will send message to user via using tools - pass - llm_resp = runner.get_final_llm_resp() - task_meta = extras.get("background_task_result", {}) - summary_note = ( - f"[BackgroundTask] {summary_name} " - f"(task_id={task_meta.get('task_id', task_id)}) finished. " - f"Result: {task_meta.get('result') or result_text or 'no content'}" - ) - if llm_resp and llm_resp.completion_text: - summary_note += ( - f"I finished the task, here is the result: {llm_resp.completion_text}" + runner = result.agent_runner + async for _ in runner.step_until_done(30): + # agent will send message to user via using tools + pass + llm_resp = runner.get_final_llm_resp() + task_meta = extras.get("background_task_result", {}) + summary_note = ( + f"[BackgroundTask] {summary_name} " + f"(task_id={task_meta.get('task_id', task_id)}) finished. " + f"Result: {task_meta.get('result') or result_text or 'no content'}" + ) + if llm_resp and llm_resp.completion_text: + summary_note += f"I finished the task, here is the result: {llm_resp.completion_text}" + await persist_agent_history( + ctx.conversation_manager, + event=cron_event, + req=req, + summary_note=summary_note, ) - await persist_agent_history( - ctx.conversation_manager, - event=cron_event, - req=req, - summary_note=summary_note, - ) if not llm_resp: logger.warning("background task agent got no response") return diff --git a/tests/test_persist_race.py b/tests/test_persist_race.py index 9680e9a708..aee0c3068e 100644 --- a/tests/test_persist_race.py +++ b/tests/test_persist_race.py @@ -55,26 +55,42 @@ async def _run(): assert notes == ["result-x"] +class SessionLockedManager: + """模拟会话锁: 按 umo 分配 asyncio.Lock""" + + def __init__(self): + self._locks: dict[str, asyncio.Lock] = {} + + def acquire(self, umo: str): + if umo not in self._locks: + self._locks[umo] = asyncio.Lock() + return self._locks[umo] + + @pytest.mark.asyncio async def test_persist_concurrent_keeps_all_results(): """ - 并发场景: 4 个后台任务结果几乎同时持久化, 不应丢失任何一条。 + 会话锁保护下, 4 个并发持久化不丢失任何结果。 - 复现路径: 每个任务在创建请求时读到旧历史快照, 然后并发 persist。 - 修复方案: 同一会话的处理流程(读历史+持久化)串行化(session lock)。 + 修复前(无锁): 每个任务持有创建时的旧历史快照, 并发读改写互相覆盖, + 丢失 3/4 条结果 (该场景已由 test_background_wake_lock 覆盖根因)。 + 修复后(会话锁): 读历史与持久化整体串行化, 后任务读到最新历史, 全部保留。 """ cm = FakeConversationManager() + slm = SessionLockedManager() umo = "test:session:1" n = 4 - # 模拟并发任务: 各自持有创建时读到的历史快照(并发时互不知道对方) - reqs = [make_req(cm.store.get((umo, "conv-1"), "[]")) for _ in range(n)] - events = [make_event(umo)] * n + async def one_task(i): + # 模拟修复后的调用模式: 读历史 + persist 都在会话锁内 + async with slm.acquire(umo): + h_now = cm.store.get((umo, "conv-1"), "[]") + req = make_req(h_now) + await persist_agent_history( + cm, event=make_event(umo), req=req, summary_note=f"result-{i}" + ) - await asyncio.gather(*[ - persist_agent_history(cm, event=events[i], req=reqs[i], summary_note=f"result-{i}") - for i in range(n) - ]) + await asyncio.gather(*[one_task(i) for i in range(n)]) final = json.loads(cm.store[(umo, "conv-1")]) saved = [m["content"] for m in final if m["role"] == "assistant"] From fa8ff192740bc7abda95afad4c26718ec0638282 Mon Sep 17 00:00:00 2001 From: T52T52 <140709873+T52T52@users.noreply.github.com> Date: Fri, 14 Aug 2026 04:00:07 +0800 Subject: [PATCH 3/3] fix: acquire session lock in cron wake path The cron wake path (_woke_main_agent) also bypassed session_lock_manager, running concurrently with user message processing on the same conversation. Wrap the read-build-run-persist flow in the per-session lock, matching the user message path in internal.py. - test_cron_wake_lock: asserts acquire_lock is called --- astrbot/core/cron/manager.py | 112 ++++++++++++++++++----------------- tests/test_cron_wake_lock.py | 87 +++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 55 deletions(-) create mode 100644 tests/test_cron_wake_lock.py diff --git a/astrbot/core/cron/manager.py b/astrbot/core/cron/manager.py index b5a0e7c3e4..b722a48a6a 100644 --- a/astrbot/core/cron/manager.py +++ b/astrbot/core/cron/manager.py @@ -19,6 +19,7 @@ from astrbot.core.platform.message_type import MessageType from astrbot.core.provider.entites import ProviderRequest from astrbot.core.utils.history_saver import persist_agent_history +from astrbot.core.utils.session_lock import session_lock_manager if TYPE_CHECKING: from astrbot.core.star.context import Context @@ -448,66 +449,67 @@ async def _woke_main_agent( streaming_response=False, provider_settings=provider_settings, ) - req = ProviderRequest() - conv = await _get_session_conv(event=cron_event, plugin_context=self.ctx) - req.conversation = conv - # finetine the messages - context = json.loads(conv.history) - if context: - req.contexts = context - context_dump = req._print_friendly_context() - req.contexts = [] - req.system_prompt += ( - "\n\nBellow is you and user previous conversation history:\n" - f"---\n" - f"{context_dump}\n" - f"---\n" + async with session_lock_manager.acquire_lock(umo): + req = ProviderRequest() + conv = await _get_session_conv(event=cron_event, plugin_context=self.ctx) + req.conversation = conv + # finetine the messages + context = json.loads(conv.history) + if context: + req.contexts = context + context_dump = req._print_friendly_context() + req.contexts = [] + req.system_prompt += ( + "\n\nBellow is you and user previous conversation history:\n" + f"---\n" + f"{context_dump}\n" + f"---\n" + ) + cron_job_str = json.dumps(extras.get("cron_job", {}), ensure_ascii=False) + req.system_prompt += PROACTIVE_AGENT_CRON_WOKE_SYSTEM_PROMPT.format( + cron_job=cron_job_str ) - cron_job_str = json.dumps(extras.get("cron_job", {}), ensure_ascii=False) - req.system_prompt += PROACTIVE_AGENT_CRON_WOKE_SYSTEM_PROMPT.format( - cron_job=cron_job_str - ) - req.prompt = ( - "You are now responding to a scheduled task. " - "Proceed according to your system instructions. " - "Output using same language as previous conversation. " - "After completing your task, summarize and output your actions and results." - ) - if delivery_session_str: - if not req.func_tool: - req.func_tool = ToolSet() - req.func_tool.add_tool( - self.ctx.get_llm_tool_manager().get_builtin_tool(SendMessageToUserTool) + req.prompt = ( + "You are now responding to a scheduled task. " + "Proceed according to your system instructions. " + "Output using same language as previous conversation. " + "After completing your task, summarize and output your actions and results." ) + if delivery_session_str: + if not req.func_tool: + req.func_tool = ToolSet() + req.func_tool.add_tool( + self.ctx.get_llm_tool_manager().get_builtin_tool( + SendMessageToUserTool + ) + ) - result = await build_main_agent( - event=cron_event, plugin_context=self.ctx, config=config, req=req - ) - if not result: - logger.error("Failed to build main agent for cron job.") - return - - runner = result.agent_runner - async for _ in runner.step_until_done(30): - # agent will send message to user via using tools - pass - llm_resp = runner.get_final_llm_resp() - cron_meta = extras.get("cron_job", {}) if extras else {} - summary_note = ( - f"[CronJob] {cron_meta.get('name') or cron_meta.get('id', 'unknown')}: {cron_meta.get('description', '')} " - f" triggered at {cron_meta.get('run_started_at', 'unknown time')}, " - ) - if llm_resp and llm_resp.role == "assistant": - summary_note += ( - f"I finished this job, here is the result: {llm_resp.completion_text}" + result = await build_main_agent( + event=cron_event, plugin_context=self.ctx, config=config, req=req ) + if not result: + logger.error("Failed to build main agent for cron job.") + return - await persist_agent_history( - self.ctx.conversation_manager, - event=cron_event, - req=req, - summary_note=summary_note, - ) + runner = result.agent_runner + async for _ in runner.step_until_done(30): + # agent will send message to user via using tools + pass + llm_resp = runner.get_final_llm_resp() + cron_meta = extras.get("cron_job", {}) if extras else {} + summary_note = ( + f"[CronJob] {cron_meta.get('name') or cron_meta.get('id', 'unknown')}: {cron_meta.get('description', '')} " + f" triggered at {cron_meta.get('run_started_at', 'unknown time')}, " + ) + if llm_resp and llm_resp.role == "assistant": + summary_note += f"I finished this job, here is the result: {llm_resp.completion_text}" + + await persist_agent_history( + self.ctx.conversation_manager, + event=cron_event, + req=req, + summary_note=summary_note, + ) if not llm_resp: logger.warning("Cron job agent got no response") return diff --git a/tests/test_cron_wake_lock.py b/tests/test_cron_wake_lock.py new file mode 100644 index 0000000000..868a22e58a --- /dev/null +++ b/tests/test_cron_wake_lock.py @@ -0,0 +1,87 @@ +""" +回归测试: cron 定时任务唤醒路径也应使用会话锁 + +背景: 与后台任务唤醒(_wake_main_agent_for_background_result)相同, +cron 定时任务触发时直接跑 agent, 未获取会话锁 —— 与用户消息并发 +处理时可能导致上下文丢失。 + +本测试断言: 修复后, cron 唤醒流程必须获取会话锁 (acquire_lock 被调用)。 +""" +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from astrbot.core.cron.manager import CronJobManager + + +def _make_manager(): + """构造最小可用的 CronJobManager mock""" + ctx = SimpleNamespace( + get_config=lambda umo: {"admins_id": [], "provider_settings": {}}, + get_llm_tool_manager=MagicMock(), + conversation_manager=MagicMock(update_conversation=AsyncMock()), + ) + mgr = CronJobManager.__new__(CronJobManager) # 跳过 __init__, 只设 ctx + mgr.ctx = ctx + return mgr + + +def _make_runner_mock(): + runner = MagicMock() + + async def _step_until_done(*args, **kwargs): + yield None + + runner.step_until_done.side_effect = _step_until_done + runner.get_final_llm_resp.return_value = SimpleNamespace( + completion_text="done", role="assistant" + ) + return runner + + +@pytest.mark.asyncio +async def test_cron_wake_acquires_session_lock(): + """cron 定时任务唤醒必须获取会话锁(与用户消息/后台唤醒路径一致)""" + mgr = _make_manager() + runner = _make_runner_mock() + + lock_mgr = MagicMock() + lock_cm = AsyncMock() + lock_cm.__aenter__.return_value = None + lock_mgr.acquire_lock.return_value = lock_cm + + fake_cron_event = SimpleNamespace( + unified_msg_origin="Pstar:FriendMessage:TEST", role="member" + ) + + with ( + patch( + "astrbot.core.cron.manager.session_lock_manager", lock_mgr, create=True + ), + patch( + "astrbot.core.astr_main_agent._get_session_conv", + new=AsyncMock(return_value=SimpleNamespace(history="[]", cid="conv-1")), + ), + patch( + "astrbot.core.astr_main_agent.build_main_agent", + new=AsyncMock(return_value=SimpleNamespace(agent_runner=runner)), + ), + patch( + "astrbot.core.cron.manager.CronMessageEvent", + return_value=fake_cron_event, + ), + # MessageSession 需为真实类型(函数内 isinstance 判断), from_str 可解析字符串 + ): + await mgr._woke_main_agent( + message="test cron job", + session_str="Pstar:FriendMessage:TEST", + extras={"cron_job": {"id": "job-1", "name": "t", "run_started_at": "t"}}, + ) + + # 核心断言: 修复后必须获取会话锁, 且锁粒度为该会话 + lock_mgr.acquire_lock.assert_called_once() + args = lock_mgr.acquire_lock.call_args[0] + assert "Pstar:FriendMessage:TEST" in args, ( + f"会话锁应按 unified_msg_origin 获取, 实际参数: {args}" + )