fix(agent): bound session persistence admission

This commit is contained in:
InfinityPacer
2026-08-23 11:43:07 +08:00
parent 8d05c6d8bf
commit 598c004efe
4 changed files with 88 additions and 6 deletions
+30 -6
View File
@@ -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,
+2
View File
@@ -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"})),
+1
View File
@@ -718,6 +718,7 @@ async def init_modules() -> HostRuntime:
AgentChatPersistenceService(
repository=AgentChatOper,
async_executor=database_worker,
capacity=database_worker.snapshot().capacity,
)
)
configure_user_lookups(
+55
View File
@@ -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 查询恢复。"""