MoviePilot/tests/test_agent_memory.py

import asyncio
import json
from datetime import datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage

from app.agent.middleware.memory import (
    SEARCH_MEMORY_TOOL_NAME,
    MemoryMiddleware,
    _summarize_with_llm,
    query_memory_files,
)
from app.agent.tools.factory import MoviePilotToolFactory
from app.agent.tools.tags import ToolTag
from app.runtime.tasks import TaskRegistry


def _write_activity_log(activity_dir, date_str: str, lines: list[str]) -> None:
    """写入测试用活动记忆文件。"""
    activity_dir.mkdir(parents=True, exist_ok=True)
    body = "\n".join(lines)
    (activity_dir / f"{date_str}.md").write_text(
        f"# {date_str} 活动记忆\n\n{body}\n",
        encoding="utf-8",
    )


async def _wait_memory_tasks(middleware: MemoryMiddleware) -> None:
    """等待活动记忆后台任务完成，避免测试与后台写入竞态。"""
    tasks = list(middleware._background_tasks)
    if tasks:
        await asyncio.gather(*tasks)


def test_memory_loads_only_primary_file(tmp_path):
    """每轮默认上下文只能装载 MEMORY.md，不应装载主题或活动正文。"""
    primary = tmp_path / "MEMORY.md"
    primary.write_text("用户偏好：简洁回复。", encoding="utf-8")
    (tmp_path / "MEDIA_RULES.md").write_text("主题记忆：优先 Remux。", encoding="utf-8")
    _write_activity_log(
        tmp_path / "activity",
        datetime.now().strftime("%Y-%m-%d"),
        ["- **10:00** 活动正文不应自动进入上下文"],
    )
    middleware = MemoryMiddleware(
        memory_dir=str(tmp_path),
        activity_dir=str(tmp_path / "activity"),
    )

    state_update = asyncio.run(middleware.abefore_agent({}, runtime=None, config=None))
    request = SimpleNamespace(
        state=state_update,
        system_message=SystemMessage(content="SYSTEM"),
        override=lambda **kwargs: SimpleNamespace(
            state=state_update,
            system_message=kwargs["system_message"],
        ),
    )
    modified = middleware.modify_request(request)
    system_text = str(modified.system_message.content)

    assert state_update["memory_contents"] == {str(primary): "用户偏好：简洁回复。"}
    assert "用户偏好：简洁回复" in system_text
    assert "主题记忆：优先 Remux" not in system_text
    assert "活动正文不应自动进入上下文" not in system_text
    assert "search_memory" in system_text
    assert "first tool call" in system_text


def test_memory_loads_global_and_current_user_primary_files_only(tmp_path):
    """默认上下文应按全局再当前用户顺序加载，且不触及其他用户记忆。"""
    global_memory = tmp_path / "MEMORY.md"
    global_memory.write_text("公共规则：使用简洁风格。", encoding="utf-8")
    user_memory_dir = tmp_path / "users" / "user-a"
    user_memory = user_memory_dir / "MEMORY.md"
    user_memory.parent.mkdir(parents=True)
    user_memory.write_text("用户偏好：优先中文回答。", encoding="utf-8")
    other_memory = tmp_path / "users" / "user-b" / "MEMORY.md"
    other_memory.parent.mkdir(parents=True)
    other_memory.write_text("其他用户的私密偏好。", encoding="utf-8")

    middleware = MemoryMiddleware(
        memory_dir=str(tmp_path),
        user_memory_dir=str(user_memory_dir),
    )
    state_update = asyncio.run(middleware.abefore_agent({}, runtime=None, config=None))
    prompt = middleware._format_agent_memory(
        state_update["memory_contents"],
        memory_empty=state_update["memory_empty"],
    )

    assert state_update["memory_contents"] == {
        str(global_memory): "公共规则：使用简洁风格。",
        str(user_memory): "用户偏好：优先中文回答。",
    }
    assert prompt.index("公共规则：使用简洁风格") < prompt.index("用户偏好：优先中文回答")
    assert "其他用户的私密偏好" not in prompt
    assert str(user_memory) in prompt


def test_memory_onboarding_still_requires_search_before_task(tmp_path):
    """主记忆为空时也必须要求 Agent 在执行任务前检索其它记忆。"""
    (tmp_path / "MEDIA_RULES.md").write_text("主题记忆：偏好 HEVC。", encoding="utf-8")
    middleware = MemoryMiddleware(memory_dir=str(tmp_path))
    state_update = asyncio.run(middleware.abefore_agent({}, runtime=None, config=None))
    prompt = middleware._format_agent_memory(
        state_update["memory_contents"],
        memory_empty=state_update["memory_empty"],
    )

    assert "primary memory file is empty" in prompt
    assert "first tool call" in prompt
    assert "search_memory" in prompt
    assert "偏好 HEVC" not in prompt


def test_query_memory_files_searches_topic_and_activity_categories(tmp_path):
    """统一检索工具应能按分类、关键词和日期读取主题与活动记忆。"""
    (tmp_path / "MEDIA_RULES.md").write_text(
        "# 媒体规则\n优先 Remux，字幕使用简体中文。\n",
        encoding="utf-8",
    )
    activity_dir = tmp_path / "activity"
    _write_activity_log(
        activity_dir,
        "2026-06-18",
        [
            "- **10:00** 帮用户整理了电影 A",
            "- **10:30** 查询了站点状态",
        ],
    )

    topic_payload = query_memory_files(
        str(tmp_path),
        activity_dir=str(activity_dir),
        query="Remux",
        category="topic",
        limit=10,
    )
    activity_payload = query_memory_files(
        str(tmp_path),
        activity_dir=str(activity_dir),
        query="整理",
        category="activity",
        date="2026-06-18",
        limit=10,
    )

    assert topic_payload["success"] is True
    assert topic_payload["entries"][0]["category"] == "topic"
    assert topic_payload["entries"][0]["text"] == "优先 Remux，字幕使用简体中文。"
    assert activity_payload["success"] is True
    assert activity_payload["entries"][0]["category"] == "activity"
    assert activity_payload["entries"][0]["summary"] == "帮用户整理了电影 A"
    assert activity_payload["entries"][0]["date"] == "2026-06-18"


def test_query_memory_files_scopes_global_and_current_user_memory(tmp_path):
    """检索只能看到公共记忆与当前用户记忆，不能跨用户读取。"""
    (tmp_path / "PUBLIC_RULES.md").write_text("公共主题：优先 Remux。\n", encoding="utf-8")
    user_memory_dir = tmp_path / "users" / "user-a"
    (user_memory_dir / "USER_RULES.md").parent.mkdir(parents=True)
    (user_memory_dir / "USER_RULES.md").write_text("当前用户主题：偏好 HEVC。\n", encoding="utf-8")
    other_memory = tmp_path / "users" / "user-b" / "USER_RULES.md"
    other_memory.parent.mkdir(parents=True)
    other_memory.write_text("其他用户主题：偏好 AV1。\n", encoding="utf-8")

    payload = query_memory_files(
        str(tmp_path),
        user_memory_dir=str(user_memory_dir),
        query="偏好",
        category="topic",
        limit=10,
    )
    other_file = query_memory_files(
        str(tmp_path),
        user_memory_dir=str(user_memory_dir),
        file_path=str(other_memory),
    )

    assert payload["success"] is True
    assert {entry["text"] for entry in payload["entries"]} == {
        "当前用户主题：偏好 HEVC。",
    }
    assert all(entry["scope"] == "user" for entry in payload["entries"])
    assert other_file["success"] is False


def test_query_memory_files_supports_regex_and_bounds_results(tmp_path):
    """记忆检索应支持显式正则，并对返回条数做有界处理。"""
    (tmp_path / "one.md").write_text("整理电影 A\n查询站点\n", encoding="utf-8")
    (tmp_path / "two.md").write_text("整理电影 B\n", encoding="utf-8")

    payload = query_memory_files(
        str(tmp_path),
        query="整理|站点",
        category="topic",
        use_regex=True,
        limit=1,
    )
    invalid = query_memory_files(
        str(tmp_path),
        query="[",
        category="topic",
        use_regex=True,
    )

    assert payload["success"] is True
    assert payload["total_count"] == 3
    assert payload["returned_count"] == 1
    assert payload["truncated"] is True
    assert invalid["success"] is False
    assert "无效的记忆检索正则表达式" in invalid["message"]


def test_memory_middleware_exposes_search_tool(tmp_path):
    """统一记忆中间件应通过一个只读系统工具提供按需检索。"""
    middleware = MemoryMiddleware(memory_dir=str(tmp_path))

    assert [tool.name for tool in middleware.tools] == [SEARCH_MEMORY_TOOL_NAME]
    assert ToolTag.Read in middleware.tools[0].tags
    assert ToolTag.System in middleware.tools[0].tags
    assert middleware.tools[0]._agent_tool_source == "middleware:memory"
    assert "topic memory" in middleware.tools[0].description


def test_memory_search_tool_returns_json_payload(tmp_path):
    """search_memory 工具应返回结构化 JSON，而不是将文件正文隐式注入上下文。"""
    (tmp_path / "MEDIA_RULES.md").write_text("优先 Remux。\n", encoding="utf-8")
    middleware = MemoryMiddleware(memory_dir=str(tmp_path))

    result = asyncio.run(
        middleware.tools[0].ainvoke({"query": "Remux", "category": "topic", "limit": 5})
    )
    payload = json.loads(result)

    assert payload["success"] is True
    assert payload["returned_count"] == 1
    assert payload["entries"][0]["text"] == "优先 Remux。"


def test_memory_search_tool_reports_streaming_execution(tmp_path):
    """search_memory 工具执行时应复用统一的工具显示策略。"""

    async def _run_test():
        calls = []
        stream_handler = SimpleNamespace(
            is_streaming=True,
            report_tool_call=lambda **kwargs: calls.append(kwargs) or "tool-1",
            tool_call_finished=MagicMock(),
        )
        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            stream_handler=stream_handler,
        )
        request = SimpleNamespace(
            tool=SimpleNamespace(name=SEARCH_MEMORY_TOOL_NAME),
            tool_call={"args": {"query": "整理", "category": "activity"}},
        )

        async def _fake_handler(_request):
            """返回模拟工具结果。"""
            return "ok"

        result = await middleware.awrap_tool_call(request, _fake_handler)
        return result, calls, stream_handler.tool_call_finished

    result, calls, finished = asyncio.run(_run_test())

    assert result == "ok"
    assert calls == [
        {
            "tool_name": SEARCH_MEMORY_TOOL_NAME,
            "tool_message": '检索记忆，主要参数：{"query": "整理", "category": "activity"}',
            "tool_kwargs": {"query": "整理", "category": "activity"},
        }
    ]
    finished.assert_called_once_with("tool-1", "done")


def test_memory_middleware_sanitizes_its_own_logs(tmp_path):
    """记忆中间件读取参数和异常写日志时必须脱敏。"""

    async def _run_test():
        secret_marker = "memory-secret-marker-6825"
        stream_handler = SimpleNamespace(is_streaming=True, report_tool_call=MagicMock())
        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            stream_handler=stream_handler,
        )
        request = SimpleNamespace(
            tool=SimpleNamespace(name=SEARCH_MEMORY_TOOL_NAME),
            tool_call={"args": {"query": f"token={secret_marker}"}},
        )
        mock_logger = MagicMock()

        async def _failing_handler(_request):
            raise RuntimeError(f"Authorization: Bearer {secret_marker}")

        with patch("app.agent.middleware.memory.logger", mock_logger):
            try:
                await middleware.awrap_tool_call(request, _failing_handler)
            except RuntimeError:
                pass
            else:
                raise AssertionError("middleware should re-raise handler errors")
        return secret_marker, mock_logger

    secret_marker, mock_logger = asyncio.run(_run_test())

    assert secret_marker not in str(mock_logger.method_calls)
    assert "***" in str(mock_logger.method_calls)


def test_memory_provider_error_does_not_echo_secret(tmp_path):
    """记忆 provider 内部异常不能进入日志或模型错误结果。"""
    secret_marker = "memory-provider-secret-3584"
    middleware = MemoryMiddleware(memory_dir=str(tmp_path))
    mock_logger = MagicMock()

    with (
        patch(
            "app.agent.middleware.memory.query_memory_files",
            side_effect=RuntimeError(f"OPENAI_API_KEY={secret_marker}"),
        ),
        patch("app.agent.middleware.memory.logger", mock_logger),
    ):
        result = asyncio.run(
            middleware._tool_provider.search_memory(query="visible")
        )

    assert secret_marker not in result
    assert secret_marker not in str(mock_logger.method_calls)
    assert "***" in result
    assert "***" in str(mock_logger.method_calls)


def test_activity_memory_records_under_unified_memory_directory(tmp_path):
    """活动摘要应写入 memory/activity，而不是旧的 agent/activity 目录。"""
    summary = "用户要求整理电影文件，助手调用 transfer_file 完成处理，结果成功。"
    activity_dir = tmp_path / "activity"

    async def _run_test():
        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            activity_dir=str(activity_dir),
        )
        with patch(
            "app.agent.middleware.memory._summarize_with_llm",
            new=AsyncMock(return_value=summary),
        ):
            await middleware.aafter_agent(
                {
                    "messages": [
                        HumanMessage(content="帮我整理电影"),
                        AIMessage(
                            content="",
                            tool_calls=[
                                {"name": "transfer_file", "args": {}, "id": "call_1"}
                            ],
                        ),
                        ToolMessage(content='{"success": true}', tool_call_id="call_1"),
                    ]
                },
                runtime=None,
            )
            await _wait_memory_tasks(middleware)

    asyncio.run(_run_test())

    log_files = list(activity_dir.glob("*.md"))
    assert len(log_files) == 1
    assert summary in log_files[0].read_text(encoding="utf-8")


def test_activity_memory_records_under_current_user_directory(tmp_path):
    """带用户上下文的活动摘要只能写入当前用户的活动目录。"""
    summary = "当前用户请求整理电影，助手完成了文件处理。"
    user_memory_dir = tmp_path / "users" / "user-a"
    user_activity_dir = user_memory_dir / "activity"

    async def _run_test():
        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            user_memory_dir=str(user_memory_dir),
        )
        with patch(
            "app.agent.middleware.memory._summarize_with_llm",
            new=AsyncMock(return_value=summary),
        ):
            await middleware.aafter_agent(
                {
                    "messages": [
                        HumanMessage(content="帮我整理电影"),
                        AIMessage(
                            content="",
                            tool_calls=[
                                {"name": "transfer_file", "args": {}, "id": "call_1"}
                            ],
                        ),
                        ToolMessage(content='{"success": true}', tool_call_id="call_1"),
                    ]
                },
                runtime=None,
            )
            await _wait_memory_tasks(middleware)

    asyncio.run(_run_test())

    log_files = list(user_activity_dir.glob("*.md"))
    assert len(log_files) == 1
    assert summary in log_files[0].read_text(encoding="utf-8")
    assert not list((tmp_path / "activity").glob("*.md"))


def test_activity_memory_skips_trivial_greeting_without_llm(tmp_path):
    """无实际任务的寒暄不应调用 LLM，也不应写入活动记忆。"""

    async def _run_test():
        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            activity_dir=str(tmp_path / "activity"),
        )
        summarize_mock = AsyncMock(return_value="不应写入")
        with patch("app.agent.middleware.memory._summarize_with_llm", new=summarize_mock):
            await middleware.aafter_agent(
                {
                    "messages": [
                        HumanMessage(content="你好"),
                        AIMessage(content="你好，有什么可以帮你？"),
                    ]
                },
                runtime=None,
            )
            await _wait_memory_tasks(middleware)
        return summarize_mock

    summarize_mock = asyncio.run(_run_test())

    summarize_mock.assert_not_awaited()
    assert not list((tmp_path / "activity").glob("*.md"))


def test_activity_summary_background_task_follows_host_shutdown(tmp_path):
    """活动摘要任务必须登记统一 owner，并随宿主关停取消和收敛。"""

    async def _run_test():
        registry = TaskRegistry()
        started = asyncio.Event()
        cancelled = asyncio.Event()

        async def _blocked_record(_messages: list) -> None:
            """保持记录任务运行，直到宿主关停发出取消。"""
            started.set()
            try:
                await asyncio.Event().wait()
            except asyncio.CancelledError:
                cancelled.set()
                raise

        middleware = MemoryMiddleware(
            memory_dir=str(tmp_path),
            activity_dir=str(tmp_path / "activity"),
            task_registry=registry,
        )
        with patch.object(middleware, "_record_activity", side_effect=_blocked_record):
            middleware._schedule_activity_recording([])
            await started.wait()
            owners = tuple(record.owner for record in registry.records)
            converged = await registry.shutdown(timeout_seconds=1.0)
            await asyncio.sleep(0)
            return owners, converged, cancelled.is_set(), middleware._background_tasks

    owners, converged, cancelled, background_tasks = asyncio.run(_run_test())

    assert owners == ("agent.memory.activity_record",)
    assert converged is True
    assert cancelled is True
    assert background_tasks == set()


def test_summarize_with_llm_ignores_skip_marker():
    """LLM 返回 SKIP 时应视为无需记录活动记忆。"""
    llm = SimpleNamespace(ainvoke=AsyncMock(return_value=SimpleNamespace(content="SKIP")))

    with patch(
        "app.agent.llm.LLMHelper.get_llm",
        new=AsyncMock(return_value=llm),
    ):
        summary = asyncio.run(_summarize_with_llm("用户: 你好"))

    assert summary is None
    llm.ainvoke.assert_awaited_once()


def test_activity_summary_hides_image_payload():
    """活动摘要输入只能保留图片占位符，不能把 Base64 写入活动记忆。"""
    from app.agent.middleware.memory import _format_conversation_for_summary

    content = [
        {"type": "text", "text": "请看看图片"},
        {"type": "image_url", "image_url": {"url": "data:image/png;base64,secret"}},
    ]

    formatted = _format_conversation_for_summary([HumanMessage(content=content)])

    assert "[图片]" in formatted
    assert "secret" not in formatted


def test_factory_does_not_register_memory_tool():
    """记忆检索工具由统一中间件注册，不应进入全局工具工厂。"""
    with patch(
        "app.agent.tools.factory._get_plugin_agent_tools",
        return_value=[],
    ):
        tools = MoviePilotToolFactory.create_tools(
            session_id="memory-session",
            user_id="10001",
        )

    assert SEARCH_MEMORY_TOOL_NAME not in {tool.name for tool in tools}