MoviePilot/tests/test_agent_chat_persistence.py

"""
AgentChat 同步短事务经有界 worker 委托的应用端口测试。
"""

from __future__ import annotations

import asyncio
import threading
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, call, patch
from uuid import uuid4

import pytest
from sqlalchemy import delete, select

from app.api.endpoints.agent import save_agent_chat_display
from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService
from app.db.models.agentchat import AgentChat
from app.db.oper.agentchat import AgentChatOper
from app.db.session import SessionFactory, async_session_scope
from app.db.uow import run_sync_transaction
from app.db.worker import DatabaseWorker
from app.schemas.agent import AgentChatDisplaySaveRequest
from app.schemas.exception import (
    AgentChatPersistenceUnavailableError,
    DatabaseWorkerOverloadedError,
)


class _Executor:
    """
    用独立线程模拟 G2B worker，验证调用方不会直接执行同步仓储。
    """

    def __init__(self) -> None:
        self.calls = 0
        self.worker_thread_id: int | None = None
        self.results: list[object] = []

    async def run(self, operation):
        """
        在线程中执行一个完整的同步操作。
        """
        self.calls += 1

        def invoke():
            self.worker_thread_id = threading.get_ident()
            result = operation()
            self.results.append(result)
            return result

        return await asyncio.to_thread(invoke)


class _Repository:
    """
    记录 AgentChat 端口调用的同步仓储替身。
    """

    def __init__(self) -> None:
        self.calls: list[tuple[str, dict]] = []

    def append_display_messages(self, **kwargs):
        """
        记录展示消息的追加参数，不执行真实写入。
        """
        self.calls.append(("append_display_messages", kwargs))
        return None

    def save_display_messages(self, **kwargs):
        """
        记录展示消息快照的保存参数，不执行真实写入。
        """
        self.calls.append(("save_display_messages", kwargs))
        return None

    def save_agent_messages(self, **kwargs):
        """
        记录 Agent 历史消息的保存参数，不执行真实写入。
        """
        self.calls.append(("save_agent_messages", kwargs))

    def update_title_if_empty(self, **kwargs):
        """
        记录空标题更新参数，供调用顺序断言使用。
        """
        self.calls.append(("update_title_if_empty", kwargs))


@pytest.mark.asyncio
async def test_agent_chat_persistence_runs_sync_repository_inside_worker() -> None:
    """
    同步 AgentChat 写入必须经过一次 worker admission。
    """
    executor = _Executor()
    repository = _Repository()
    service = AgentChatPersistenceService(
        repository=lambda _session: repository,
        async_executor=executor,
        sync_transaction=lambda operation: operation(object()),
    )
    caller_thread_id = threading.get_ident()

    await service.async_append_display_messages(
        session_id="session-1",
        user_id="1",
        messages=[{"role": "user", "content": "hello"}],
    )
    await service.async_save_display_messages(
        session_id="session-1",
        user_id="1",
        messages=[],
    )
    await service.async_save_agent_messages(
        session_id="session-1",
        user_id="1",
        messages=[],
    )
    await service.async_update_title_if_empty(
        session_id="session-1",
        user_id="1",
        title="标题",
    )

    assert executor.calls == 4
    assert executor.results == [None, None, None, None]
    assert executor.worker_thread_id != caller_thread_id
    assert [name for name, _kwargs in repository.calls] == [
        "append_display_messages",
        "save_display_messages",
        "save_agent_messages",
        "update_title_if_empty",
    ]


@pytest.mark.asyncio
async def test_agent_chat_persistence_propagates_worker_failure() -> None:
    """
    worker admission 或事务异常必须原样返回给 async 应用调用方。
    """

    class FailingExecutor:
        """
        模拟无法接收同步操作的后台执行器。
        """
        async def run(self, _operation):
            """
            模拟后台执行器失败，原样抛出错误。
            """
            raise RuntimeError("worker failed")

    service = AgentChatPersistenceService(
        repository=lambda _session: _Repository(),
        async_executor=FailingExecutor(),
        sync_transaction=lambda operation: operation(object()),
    )

    with pytest.raises(RuntimeError, match="worker failed"):
        await service.async_save_agent_messages(
            session_id="session-1",
            user_id="1",
            messages=[],
        )


@pytest.mark.asyncio
async def test_agent_chat_persistence_pending_metric_uses_deltas() -> None:
    """
    pending 是 UpDownCounter，准入和释放必须分别记录增减量。
    """
    service = AgentChatPersistenceService(
        repository=lambda _session: _Repository(),
        async_executor=_Executor(),
        sync_transaction=lambda operation: operation(object()),
    )
    with patch("app.application.messaging.chat.record_metric") as record_metric:
        await service.async_save_agent_messages(
            session_id="metric-session",
            user_id="1",
            messages=[],
        )
    record_metric.assert_has_calls(
        [
            call("agent.chat.persistence.pending", 1),
            call("agent.chat.persistence.pending", -1),
        ]
    )


@pytest.mark.asyncio
async def test_authoritative_display_save_propagates_worker_overload() -> None:
    """
    权威 PUT 保存不能把 worker 背压吞成成功或普通业务失败。
    """
    repository = AsyncMock()
    repository.async_get.return_value = None
    service = AgentChatService(repository=repository)

    class OverloadedPersistence:
        """
        模拟写入容量耗尽的聊天持久化端口。
        """
        async def async_save_display_messages(self, **_kwargs):
            """
            模拟展示消息写入过载，向调用方传播数据库背压错误。
            """
            raise DatabaseWorkerOverloadedError("busy")

    with pytest.raises(DatabaseWorkerOverloadedError, match="busy"):
        await save_agent_chat_display(
            session_id="overloaded-session",
            payload=AgentChatDisplaySaveRequest(messages=[]),
            current_user=SimpleNamespace(id=1, name="admin", is_superuser=True),
            service=service,
            persistence=OverloadedPersistence(),
        )


@pytest.mark.asyncio
async def test_authoritative_display_save_reads_fresh_projection_after_worker_write(
    monkeypatch,
) -> None:
    """
    权威展示保存的响应必须读取 worker 提交后的最新投影。
    """
    existing_chat = SimpleNamespace(
        user_id="1",
        username="admin",
        channel="WebAgent",
        source="web-agent",
        original_chat_id=None,
        client_session_id="client-1",
    )
    updated_chat = SimpleNamespace(
        session_id="fresh-session",
        message_count=2,
    )
    request_service = SimpleNamespace(
        get_accessible=AsyncMock(return_value=existing_chat),
        get=AsyncMock(return_value=existing_chat),
    )
    canonical_service = SimpleNamespace(
        get_accessible=AsyncMock(return_value=updated_chat),
        to_summary=MagicMock(return_value="fresh-summary"),
    )
    persistence = SimpleNamespace(async_save_display_messages=AsyncMock())
    current_user = SimpleNamespace(id=1, name="admin", is_superuser=True)
    monkeypatch.setattr(
        "app.api.endpoints.agent.get_configured_agent_chat_service",
        MagicMock(return_value=canonical_service),
    )

    response = await save_agent_chat_display(
        session_id="fresh-session",
        payload=AgentChatDisplaySaveRequest(messages=[]),
        current_user=current_user,
        service=request_service,
        persistence=persistence,
    )

    assert response.success is True
    assert response.data == "fresh-summary"
    canonical_service.get_accessible.assert_awaited_once_with("fresh-session", current_user)


@pytest.mark.asyncio
async def test_agent_chat_persistence_rolls_back_compound_write(monkeypatch) -> None:
    """
    复合写入中途失败时，创建或更新不能留下半成品。
    """
    worker = DatabaseWorker(max_workers=1, capacity=4)
    await worker.start()
    session_id = f"worker-rollback-{uuid4().hex}"
    persistence = AgentChatPersistenceService(
        repository=lambda session: AgentChatOper(session),
        async_executor=worker,
        sync_transaction=run_sync_transaction,
    )
    original = AgentChatOper.save_display_messages

    def fail_after_stage(self, *args, **kwargs):
        original(self, *args, **kwargs)
        raise RuntimeError("display snapshot failed")

    monkeypatch.setattr(AgentChatOper, "save_display_messages", fail_after_stage)
    try:
        with pytest.raises(RuntimeError, match="display snapshot failed"):
            await persistence.async_append_display_messages(
                session_id=session_id,
                user_id="rollback-user",
                messages=[{"role": "user", "content": "not committed"}],
            )
        async with async_session_scope() as session:
            result = await session.execute(select(AgentChat).where(AgentChat.session_id == session_id))
            assert result.scalars().first() is None
    finally:
        await worker.shutdown()


@pytest.mark.asyncio
async def test_agent_chat_persistence_bounds_session_waiters_and_releases_cancelled() -> None:
    """
    同会话锁等待受总量限制，取消等待不会遗留 admission。
    """

    class BlockingExecutor:
        """
        通过事件控制执行时机，放大会话准入和关闭竞态。
        """
        def __init__(self) -> None:
            self.started = asyncio.Event()
            self.release = asyncio.Event()

        async def run(self, operation):
            """
            通知测试任务已开始，等待释放事件后执行同步操作。
            """
            self.started.set()
            await self.release.wait()
            return operation()

    executor = BlockingExecutor()
    service = AgentChatPersistenceService(
        repository=lambda _session: _Repository(),
        async_executor=executor,
        sync_transaction=lambda operation: operation(object()),
        capacity=2,
        session_capacity=2,
    )
    first = asyncio.create_task(
        service.async_save_agent_messages(
            session_id="session-admission",
            user_id="1",
            messages=[],
        )
    )
    await executor.started.wait()
    second = asyncio.create_task(
        service.async_save_agent_messages(
            session_id="session-admission",
            user_id="1",
            messages=[],
        )
    )
    await asyncio.sleep(0)
    third = asyncio.create_task(
        service.async_save_agent_messages(
            session_id="session-admission",
            user_id="1",
            messages=[],
        )
    )
    with pytest.raises(AgentChatPersistenceUnavailableError):
        await third
    second.cancel()
    with pytest.raises(asyncio.CancelledError):
        await second
    assert service._pending_writes == 1
    executor.release.set()
    await first
    assert service._pending_writes == 0


@pytest.mark.asyncio
async def test_agent_chat_persistence_session_admission_is_fair() -> None:
    """
    热点会话的锁等待不能占满全局容量并拒绝其他会话。
    """

    class BlockingExecutor:
        """
        通过事件控制执行时机，放大会话准入和关闭竞态。
        """
        def __init__(self) -> None:
            self.started = asyncio.Event()
            self.release = asyncio.Event()

        async def run(self, operation):
            """
            通知测试任务已开始，等待释放事件后执行同步操作。
            """
            self.started.set()
            await self.release.wait()
            return operation()

    executor = BlockingExecutor()
    service = AgentChatPersistenceService(
        repository=lambda _session: _Repository(),
        async_executor=executor,
        sync_transaction=lambda operation: operation(object()),
        capacity=4,
        session_capacity=2,
    )
    first = asyncio.create_task(service.async_save_agent_messages(session_id="hot-session", user_id="1", messages=[]))
    await executor.started.wait()
    second = asyncio.create_task(service.async_save_agent_messages(session_id="hot-session", user_id="1", messages=[]))
    await asyncio.sleep(0)
    with pytest.raises(AgentChatPersistenceUnavailableError):
        await service.async_save_agent_messages(session_id="hot-session", user_id="1", messages=[])
    other = asyncio.create_task(service.async_save_agent_messages(session_id="other-session", user_id="1", messages=[]))
    await asyncio.sleep(0)
    assert not other.done()
    executor.release.set()
    await first
    await second
    await other


@pytest.mark.asyncio
async def test_agent_chat_persistence_shutdown_drains_active_writes() -> None:
    """
    关闭持久化端口时拒绝新写入并等待现有会话写入收口。
    """

    class BlockingExecutor:
        """
        通过事件控制执行时机，放大会话准入和关闭竞态。
        """
        def __init__(self) -> None:
            self.started = asyncio.Event()
            self.release = asyncio.Event()

        async def run(self, operation):
            """
            通知测试任务已开始，等待释放事件后执行同步操作。
            """
            self.started.set()
            await self.release.wait()
            return operation()

    executor = BlockingExecutor()
    service = AgentChatPersistenceService(
        repository=lambda _session: _Repository(),
        async_executor=executor,
        sync_transaction=lambda operation: operation(object()),
    )
    write = asyncio.create_task(
        service.async_save_agent_messages(session_id="shutdown-session", user_id="1", messages=[])
    )
    await executor.started.wait()
    shutdown = asyncio.create_task(service.shutdown())
    await asyncio.sleep(0)
    assert not shutdown.done()
    with pytest.raises(AgentChatPersistenceUnavailableError):
        await service.async_save_agent_messages(session_id="new-session", user_id="1", messages=[])
    executor.release.set()
    await write
    await shutdown


@pytest.mark.asyncio
async def test_agent_chat_shutdown_timeout_keeps_worker_owner_until_write_finishes() -> None:
    """
    持久化关闭超时时保留运行中的写入和数据库 worker owner。
    """
    started = threading.Event()
    release = threading.Event()

    class BlockingRepository(_Repository):
        """
        通过线程事件延迟写入，验证关闭超时后的资源所有权。
        """
        def save_agent_messages(self, **kwargs):
            """
            通知写入已开始，等待释放事件或一秒超时后记录保存参数。
            """
            started.set()
            release.wait(1)
            super().save_agent_messages(**kwargs)

    worker = DatabaseWorker(max_workers=1, capacity=1)
    await worker.start()
    service = AgentChatPersistenceService(
        repository=lambda _session: BlockingRepository(),
        async_executor=worker,
        sync_transaction=lambda operation: operation(object()),
    )
    write = asyncio.create_task(
        service.async_save_agent_messages(
            session_id="shutdown-timeout-session",
            user_id="1",
            messages=[],
        )
    )
    assert await asyncio.to_thread(started.wait, 1)
    shutdown = asyncio.create_task(service.shutdown())
    try:
        with pytest.raises(asyncio.TimeoutError):
            await asyncio.wait_for(shutdown, timeout=0.01)
        assert service._closing is True
        assert write.done() is False
        assert worker._executor is not None
    finally:
        release.set()
        await write
        await worker.shutdown()

    assert worker._executor is None


@pytest.mark.asyncio
async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() -> None:
    """
    真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。
    """
    worker = DatabaseWorker(max_workers=1, capacity=4)
    await worker.start()
    session_id = f"worker-{uuid4().hex}"
    persistence = AgentChatPersistenceService(
        repository=lambda session: AgentChatOper(session),
        async_executor=worker,
        sync_transaction=run_sync_transaction,
    )
    query = AgentChatService(repository=AgentChatOper())

    try:
        await persistence.async_save_display_messages(
            session_id=session_id,
            user_id="worker-user",
            username="worker-user",
            channel="WebAgent",
            source="worker-test",
            messages=[{"role": "user", "content": "worker"}],
        )
        chat = await query.get(
            session_id,
            user_id="worker-user",
        )
        assert chat is not None
        assert chat.message_count == 1
        assert chat.messages[0]["content"] == "worker"
    finally:
        await AgentChatOper().async_delete(
            session_id=session_id,
            user_id="worker-user",
        )
        await worker.shutdown()


@pytest.mark.asyncio
async def test_agent_chat_persistence_serializes_same_session_writes() -> None:
    """
    同一会话的首次创建和既有快照追加都必须串行。
    """
    worker = DatabaseWorker(max_workers=4, capacity=16)
    await worker.start()
    session_id = f"worker-race-{uuid4().hex}"
    existing_session_id = f"worker-race-existing-{uuid4().hex}"
    persistence = AgentChatPersistenceService(
        repository=lambda session: AgentChatOper(session),
        async_executor=worker,
        sync_transaction=run_sync_transaction,
    )

    async def append(content: str) -> None:
        await persistence.async_append_display_messages(
            session_id=session_id,
            user_id="worker-race-user",
            messages=[{"role": "user", "content": content}],
        )

    async def append_existing(content: str) -> None:
        await persistence.async_append_display_messages(
            session_id=existing_session_id,
            user_id="worker-race-user",
            messages=[{"role": "user", "content": content}],
        )

    try:
        await asyncio.gather(*(append(f"message-{index}") for index in range(4)))
        await persistence.async_save_display_messages(
            session_id=existing_session_id,
            user_id="worker-race-user",
            messages=[{"role": "user", "content": "seed"}],
        )
        await asyncio.gather(*(append_existing(f"existing-{index}") for index in range(4)))
        async with async_session_scope() as session:
            result = await session.execute(
                select(AgentChat).where(AgentChat.session_id.in_((session_id, existing_session_id)))
            )
            rows = list(result.scalars().all())
        assert len(rows) == 2
        row_by_session = {row.session_id: row for row in rows}
        assert {message["content"] for message in row_by_session[session_id].display_messages} == {
            f"message-{index}" for index in range(4)
        }
        assert {message["content"] for message in row_by_session[existing_session_id].display_messages} == {"seed"} | {
            f"existing-{index}" for index in range(4)
        }
    finally:
        with SessionFactory() as session:
            session.execute(delete(AgentChat).where(AgentChat.session_id.in_((session_id, existing_session_id))))
            session.commit()
        await worker.shutdown()