diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index 984a69e8a..ce0e9e2bc 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -2,9 +2,11 @@ from __future__ import annotations +import asyncio from dataclasses import dataclass from collections.abc import Callable from typing import Any, Optional, Protocol +from weakref import WeakValueDictionary from app.application.database import AsyncDatabaseExecutor from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary @@ -350,17 +352,31 @@ class AgentChatPersistenceService: """保存同步仓储工厂和异步执行端口。""" self._repository = repository self._async_executor = async_executor + # append_display_messages 属于读取旧快照后整列写回的复合操作；按会话串行化， + # 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行 + # 中为一次性会话永久保留锁对象，不限制不同会话之间的 worker 并行度。 + self._session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() + + def _session_lock(self, session_id: str) -> asyncio.Lock: + """返回当前进程内指定会话的写锁。""" + lock = self._session_locks.get(session_id) + if lock is None: + lock = asyncio.Lock() + self._session_locks[session_id] = lock + return lock async def _run_write( self, + session_id: str, operation: Callable[[SyncAgentChatRepository], object], ) -> None: """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" - def execute() -> None: - """执行同步写入，不让 ORM 对象越过 worker 边界。""" - operation(self._repository()) + async with self._session_lock(session_id): + def execute() -> None: + """执行同步写入，不让 ORM 对象越过 worker 边界。""" + operation(self._repository()) - await self._async_executor.run(execute) + await self._async_executor.run(execute) async def async_append_display_messages( self, @@ -376,6 +392,7 @@ class AgentChatPersistenceService: ) -> None: """异步追加展示消息，等待同步事务取得确定终态。""" await self._run_write( + session_id, lambda repository: repository.append_display_messages( session_id=session_id, user_id=user_id, @@ -403,6 +420,7 @@ class AgentChatPersistenceService: ) -> None: """异步保存展示消息快照，实际写入由有界 worker 承接。""" await self._run_write( + session_id, lambda repository: repository.save_display_messages( session_id=session_id, user_id=user_id, @@ -425,6 +443,7 @@ class AgentChatPersistenceService: ) -> None: """异步保存可恢复的原始消息。""" await self._run_write( + session_id, lambda repository: repository.save_agent_messages( session_id=session_id, user_id=user_id, @@ -446,6 +465,7 @@ class AgentChatPersistenceService: ) -> None: """异步写入首次生成的会话标题。""" await self._run_write( + session_id, lambda repository: repository.update_title_if_empty( session_id=session_id, user_id=user_id, diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index ff5f7520f..225125a96 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -7,9 +7,12 @@ import threading from uuid import uuid4 import pytest +from sqlalchemy import delete, select 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.worker import DatabaseWorker @@ -153,3 +156,67 @@ async def test_agent_chat_persistence_uses_real_worker_and_sqlite_transaction() 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=AgentChatOper, + async_executor=worker, + ) + + 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()