mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-09-06 07:56:52 +08:00
fix(agent): serialize session persistence writes
This commit is contained in:
@@ -2,9 +2,11 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from typing import Any, Optional, Protocol
|
from typing import Any, Optional, Protocol
|
||||||
|
from weakref import WeakValueDictionary
|
||||||
|
|
||||||
from app.application.database import AsyncDatabaseExecutor
|
from app.application.database import AsyncDatabaseExecutor
|
||||||
from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary
|
from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary
|
||||||
@@ -350,17 +352,31 @@ class AgentChatPersistenceService:
|
|||||||
"""保存同步仓储工厂和异步执行端口。"""
|
"""保存同步仓储工厂和异步执行端口。"""
|
||||||
self._repository = repository
|
self._repository = repository
|
||||||
self._async_executor = async_executor
|
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(
|
async def _run_write(
|
||||||
self,
|
self,
|
||||||
|
session_id: str,
|
||||||
operation: Callable[[SyncAgentChatRepository], object],
|
operation: Callable[[SyncAgentChatRepository], object],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""在线程 worker 内完成同步写入并丢弃仓储对象返回值。"""
|
"""在线程 worker 内完成同步写入并丢弃仓储对象返回值。"""
|
||||||
def execute() -> None:
|
async with self._session_lock(session_id):
|
||||||
"""执行同步写入,不让 ORM 对象越过 worker 边界。"""
|
def execute() -> None:
|
||||||
operation(self._repository())
|
"""执行同步写入,不让 ORM 对象越过 worker 边界。"""
|
||||||
|
operation(self._repository())
|
||||||
|
|
||||||
await self._async_executor.run(execute)
|
await self._async_executor.run(execute)
|
||||||
|
|
||||||
async def async_append_display_messages(
|
async def async_append_display_messages(
|
||||||
self,
|
self,
|
||||||
@@ -376,6 +392,7 @@ class AgentChatPersistenceService:
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""异步追加展示消息,等待同步事务取得确定终态。"""
|
"""异步追加展示消息,等待同步事务取得确定终态。"""
|
||||||
await self._run_write(
|
await self._run_write(
|
||||||
|
session_id,
|
||||||
lambda repository: repository.append_display_messages(
|
lambda repository: repository.append_display_messages(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
@@ -403,6 +420,7 @@ class AgentChatPersistenceService:
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""异步保存展示消息快照,实际写入由有界 worker 承接。"""
|
"""异步保存展示消息快照,实际写入由有界 worker 承接。"""
|
||||||
await self._run_write(
|
await self._run_write(
|
||||||
|
session_id,
|
||||||
lambda repository: repository.save_display_messages(
|
lambda repository: repository.save_display_messages(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
@@ -425,6 +443,7 @@ class AgentChatPersistenceService:
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""异步保存可恢复的原始消息。"""
|
"""异步保存可恢复的原始消息。"""
|
||||||
await self._run_write(
|
await self._run_write(
|
||||||
|
session_id,
|
||||||
lambda repository: repository.save_agent_messages(
|
lambda repository: repository.save_agent_messages(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
@@ -446,6 +465,7 @@ class AgentChatPersistenceService:
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""异步写入首次生成的会话标题。"""
|
"""异步写入首次生成的会话标题。"""
|
||||||
await self._run_write(
|
await self._run_write(
|
||||||
|
session_id,
|
||||||
lambda repository: repository.update_title_if_empty(
|
lambda repository: repository.update_title_if_empty(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
|
|||||||
@@ -7,9 +7,12 @@ import threading
|
|||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from sqlalchemy import delete, select
|
||||||
|
|
||||||
from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService
|
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.oper.agentchat import AgentChatOper
|
||||||
|
from app.db.session import SessionFactory, async_session_scope
|
||||||
from app.db.worker import DatabaseWorker
|
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",
|
user_id="worker-user",
|
||||||
)
|
)
|
||||||
await worker.shutdown()
|
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()
|
||||||
|
|||||||
Reference in New Issue
Block a user