From 598c004efeaba201b9ea093226c71e5c7727a153 Mon Sep 17 00:00:00 2001 From: InfinityPacer Date: Sun, 23 Aug 2026 03:21:34 +0800 Subject: [PATCH] fix(agent): bound session persistence admission --- app/application/messaging/chat.py | 36 +++++++++++++++--- app/runtime/observability/__init__.py | 2 + app/startup/modules_initializer.py | 1 + tests/test_agent_chat_persistence.py | 55 +++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 6 deletions(-) diff --git a/app/application/messaging/chat.py b/app/application/messaging/chat.py index ce0e9e2bc..24c012611 100644 --- a/app/application/messaging/chat.py +++ b/app/application/messaging/chat.py @@ -8,8 +8,15 @@ from collections.abc import Callable from typing import Any, Optional, Protocol from weakref import WeakValueDictionary -from app.application.database import AsyncDatabaseExecutor +from app.application.database import ( + AsyncDatabaseExecutor, + DatabaseWorkerOverloadedError, +) from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary +from app.runtime.observability import record_metric + + +DEFAULT_AGENT_CHAT_WRITE_CAPACITY = 32 def has_custom_agent_chat_title(value: Optional[str]) -> bool: @@ -348,10 +355,15 @@ class AgentChatPersistenceService: self, repository: SyncAgentChatRepositoryFactory, async_executor: AsyncDatabaseExecutor, + capacity: int = DEFAULT_AGENT_CHAT_WRITE_CAPACITY, ) -> None: """保存同步仓储工厂和异步执行端口。""" + if capacity < 1: + raise ValueError("AgentChat 写入容量必须大于 0") self._repository = repository self._async_executor = async_executor + self._capacity = capacity + self._pending_writes = 0 # append_display_messages 属于读取旧快照后整列写回的复合操作；按会话串行化， # 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行 # 中为一次性会话永久保留锁对象，不限制不同会话之间的 worker 并行度。 @@ -371,12 +383,24 @@ class AgentChatPersistenceService: operation: Callable[[SyncAgentChatRepository], object], ) -> None: """在线程 worker 内完成同步写入并丢弃仓储对象返回值。""" - async with self._session_lock(session_id): - def execute() -> None: - """执行同步写入，不让 ORM 对象越过 worker 边界。""" - operation(self._repository()) + # 会话锁前的等待也纳入固定总量，避免公开展示保存入口形成无界应用层队列。 + if self._pending_writes >= self._capacity: + record_metric("agent.chat.persistence.rejected") + raise DatabaseWorkerOverloadedError( + f"AgentChat 写入容量已用尽（上限 {self._capacity}）" + ) + self._pending_writes += 1 + record_metric("agent.chat.persistence.pending", self._pending_writes) + try: + 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) + finally: + self._pending_writes -= 1 + record_metric("agent.chat.persistence.pending", self._pending_writes) async def async_append_display_messages( self, diff --git a/app/runtime/observability/__init__.py b/app/runtime/observability/__init__.py index 1da6e1d10..bbebbca15 100644 --- a/app/runtime/observability/__init__.py +++ b/app/runtime/observability/__init__.py @@ -60,6 +60,8 @@ METRIC_SPECS = { MetricSpec("scheduler.job.dead_letter", MetricKind.COUNTER, frozenset({"owner"})), MetricSpec("plugin.lifecycle.duration", MetricKind.HISTOGRAM, frozenset({"operation", "outcome"})), MetricSpec("agent.active_tasks", MetricKind.GAUGE, frozenset({"task_type"})), + MetricSpec("agent.chat.persistence.pending", MetricKind.GAUGE, frozenset()), + MetricSpec("agent.chat.persistence.rejected", MetricKind.COUNTER, frozenset()), MetricSpec("agent.cancel", MetricKind.COUNTER, frozenset({"task_type", "outcome"})), MetricSpec("agent.provider.duration", MetricKind.HISTOGRAM, frozenset({"provider_type", "outcome"})), MetricSpec("agent.token_usage", MetricKind.COUNTER, frozenset({"provider_type", "direction"})), diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index f14c415a8..b889cc9e1 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -718,6 +718,7 @@ async def init_modules() -> HostRuntime: AgentChatPersistenceService( repository=AgentChatOper, async_executor=database_worker, + capacity=database_worker.snapshot().capacity, ) ) configure_user_lookups( diff --git a/tests/test_agent_chat_persistence.py b/tests/test_agent_chat_persistence.py index 225125a96..b894d0c0b 100644 --- a/tests/test_agent_chat_persistence.py +++ b/tests/test_agent_chat_persistence.py @@ -9,6 +9,7 @@ from uuid import uuid4 import pytest from sqlalchemy import delete, select +from app.application.database import DatabaseWorkerOverloadedError from app.application.messaging.chat import AgentChatPersistenceService, AgentChatService from app.db.models.agentchat import AgentChat from app.db.oper.agentchat import AgentChatOper @@ -122,6 +123,60 @@ async def test_agent_chat_persistence_propagates_worker_failure() -> None: ) +@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=_Repository, + async_executor=executor, + 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(DatabaseWorkerOverloadedError): + 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_uses_real_worker_and_sqlite_transaction() -> None: """真实 AgentChat Oper 经 worker 写入后可被 native async 查询恢复。"""