fix(agent): close async chat persistence lifecycle

This commit is contained in:
InfinityPacer
2026-08-23 11:44:18 +08:00
parent 6f5ee96152
commit 0ba4a7e5e3
14 changed files with 448 additions and 72 deletions
+9
View File
@@ -4,6 +4,7 @@ from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession
from app.api.context import (
get_agent_chat_runtime,
get_agent_chat_repository,
get_agent_chat_transaction,
get_async_session,
@@ -11,6 +12,7 @@ from app.api.context import (
)
from app.application.messaging.chat import (
AgentChatService,
AgentChatPersistenceService,
AsyncAgentChatRepository,
AsyncUnitOfWork,
)
@@ -26,6 +28,13 @@ def get_agent_chat_service(
return AgentChatService(chat_repository, unit_of_work)
def get_agent_chat_persistence(
runtime: HostRuntime = Depends(get_agent_chat_runtime),
) -> AgentChatPersistenceService:
"""从类型化 Agent 运行时获取有界会话写入端口。"""
return runtime.persistence
def get_message_query_service(
db: AsyncSession = Depends(get_async_session),
runtime: HostRuntime = Depends(get_host_runtime),
+2
View File
@@ -5,6 +5,7 @@
"""
from app.api.dependencies.agent import (
get_agent_chat_persistence,
get_agent_chat_service,
get_message_query_service,
)
@@ -52,6 +53,7 @@ from app.api.dependencies.workflow import (
# 兼容聚合入口只显式列出既有 FastAPI 依赖,不向插件制造新的动态导出规则。
__all__ = [
"get_agent_chat_persistence",
"get_agent_chat_service",
"get_auth_service",
"get_current_active_manage_user",
+91 -45
View File
@@ -47,17 +47,25 @@ from app.command import Command
from app.runtime.config import global_vars
from app.runtime.events import Event, EventManager
from app.api.principal import ApiPrincipal
from app.api.dependencies.agent import get_agent_chat_service
from app.api.dependencies.agent import (
get_agent_chat_persistence,
get_agent_chat_service,
)
from app.api.dependencies.auth import get_current_active_user
from app.application.messaging.chat import (
AgentChatRecord,
AgentChatPersistenceService,
AgentChatService,
get_configured_agent_chat_service,
get_configured_agent_chat_persistence,
)
from app.application.security.user import get_configured_user_id_lookup
from app.application.configuration import get_api_runtime_config_snapshot
from app.application.messaging.agent import attach_web_agent_edit_queue, detach_web_agent_edit_queue
from app.application.messaging.agent import (
attach_web_agent_edit_queue,
create_web_agent_background_task,
detach_web_agent_edit_queue,
)
from app.application.messaging.agent import agent_interaction_manager
from app.application.messaging.agent import (
build_agent_choice_button_rows,
@@ -88,7 +96,6 @@ _WEB_AGENT_FILE_REGISTRY: dict[str, dict[str, Any]] = {}
_WEB_AGENT_MESSAGE_QUEUES: dict[str, list[Queue[_SchemaMessage]]] = {}
_WEB_AGENT_MESSAGE_LOCK = Lock()
_WEB_AGENT_MESSAGE_LISTENER_REGISTERED = False
_WEB_AGENT_BACKGROUND_TASKS: set[asyncio.Task] = set()
class _WebAgentEventPublisher:
@@ -517,13 +524,16 @@ def _build_web_agent_session_id(user: ApiPrincipal, session_id: Optional[str]) -
async def _build_web_agent_session_id_async(
user: ApiPrincipal,
session_id: Optional[str],
service: Optional[AgentChatService] = None,
) -> str:
"""异步解析 Web Agent 会话 ID,并复用异步会话查询端口。"""
seed = str(session_id or "").strip() or uuid.uuid4().hex
if seed.startswith(WEB_AGENT_SESSION_PREFIX):
return seed
try:
existing_chat = await get_configured_agent_chat_service().get(seed)
if service is None:
service = get_configured_agent_chat_service()
existing_chat = await service.get(seed)
if existing_chat and AgentChatService.can_access(existing_chat, user):
return seed
except Exception as e:
@@ -658,36 +668,41 @@ async def _save_web_agent_display_snapshot(
current_user: ApiPrincipal,
messages: list[dict],
client_session_id: Optional[str] = None,
service: Optional[AgentChatService] = None,
persistence: Optional[AgentChatPersistenceService] = None,
) -> None:
"""
保存 WebAgent 当前展示消息快照。
"""
try:
existing_chat = await get_configured_agent_chat_service().get(session_id)
await get_configured_agent_chat_persistence().async_save_display_messages(
session_id=session_id,
user_id=(existing_chat.user_id if existing_chat else str(current_user.id)),
username=(existing_chat.username if existing_chat else current_user.name),
channel=(
existing_chat.channel
if existing_chat and existing_chat.channel
else NotificationChannel.WebAgent
),
source=(
existing_chat.source
if existing_chat and existing_chat.source
else WEB_AGENT_SOURCE
),
original_chat_id=existing_chat.original_chat_id if existing_chat else None,
client_session_id=(
existing_chat.client_session_id
if existing_chat and existing_chat.client_session_id
else client_session_id
),
messages=messages,
)
except Exception as e:
logger.debug(f"保存WebAgent展示历史失败: {e}")
if service is None:
# 直接调用该内部 helper 时没有 FastAPI 依赖注入上下文。
service = get_configured_agent_chat_service()
existing_chat = await service.get(session_id)
if persistence is None:
# 直接调用该内部 helper 时没有 FastAPI 依赖注入上下文。
persistence = get_configured_agent_chat_persistence()
await persistence.async_save_display_messages(
session_id=session_id,
user_id=(existing_chat.user_id if existing_chat else str(current_user.id)),
username=(existing_chat.username if existing_chat else current_user.name),
channel=(
existing_chat.channel
if existing_chat and existing_chat.channel
else NotificationChannel.WebAgent
),
source=(
existing_chat.source
if existing_chat and existing_chat.source
else WEB_AGENT_SOURCE
),
original_chat_id=existing_chat.original_chat_id if existing_chat else None,
client_session_id=(
existing_chat.client_session_id
if existing_chat and existing_chat.client_session_id
else client_session_id
),
messages=messages,
)
def _build_web_agent_sse(
@@ -757,6 +772,7 @@ def _sanitize_web_agent_upload_name(
async def _get_web_agent_upload_dir(
user: ApiPrincipal,
session_id: Optional[str],
service: Optional[AgentChatService] = None,
) -> Path:
"""
计算当前 Web Agent 会话的临时附件目录。
@@ -765,7 +781,11 @@ async def _get_web_agent_upload_dir(
:param session_id: 前端会话标识
:return: 已创建的临时附件目录
"""
server_session_id = await _build_web_agent_session_id_async(user, session_id)
server_session_id = await _build_web_agent_session_id_async(
user,
session_id,
service,
)
safe_session_id = server_session_id.replace(":", "_")
upload_dir = get_api_runtime_config_snapshot().temp_path / "agent_uploads" / safe_session_id
upload_dir.mkdir(parents=True, exist_ok=True)
@@ -1701,6 +1721,7 @@ async def upload_web_agent_file(
file: UploadFile = File(...),
session_id: Optional[str] = Form(None),
current_user: ApiPrincipal = Depends(get_current_active_user),
service: AgentChatService = Depends(get_agent_chat_service),
) -> _SchemaResponse:
"""
上传 Web 智能助手对话附件。
@@ -1712,7 +1733,7 @@ async def upload_web_agent_file(
"""
mime_type = file.content_type or mimetypes.guess_type(file.filename or "")[0]
safe_name = _sanitize_web_agent_upload_name(file.filename, mime_type)
upload_dir = await _get_web_agent_upload_dir(current_user, session_id)
upload_dir = await _get_web_agent_upload_dir(current_user, session_id, service)
target_path = upload_dir / f"{uuid.uuid4().hex[:8]}_{safe_name}"
size = await _save_web_agent_upload(file, target_path)
attachment = _register_web_agent_file(
@@ -1846,6 +1867,7 @@ async def get_agent_chat_session(
server_session_id = await _build_web_agent_session_id_async(
current_user,
session_id,
service,
)
if server_session_id != session_id:
chat = await _get_accessible_agent_chat(
@@ -1884,6 +1906,7 @@ async def save_agent_chat_display(
payload: _SchemaAgentChatDisplaySaveRequest,
current_user: ApiPrincipal = Depends(get_current_active_user),
service: AgentChatService = Depends(get_agent_chat_service),
persistence: AgentChatPersistenceService = Depends(get_agent_chat_persistence),
) -> _SchemaResponse:
"""
保存前端聚合后的 Agent 展示消息。
@@ -1911,6 +1934,8 @@ async def save_agent_chat_display(
current_user=current_user,
messages=messages,
client_session_id=existing_chat.client_session_id if existing_chat else session_id,
service=service,
persistence=persistence,
)
chat = await service.get_accessible(session_id, current_user)
if not chat:
@@ -1964,6 +1989,7 @@ async def stop_web_agent_session_task(
server_session_id = await _build_web_agent_session_id_async(
current_user,
session_id,
service,
)
chat = await _get_accessible_agent_chat(
service,
@@ -1988,6 +2014,8 @@ async def _web_agent_stream_impl(
payload: _SchemaAgentWebChatRequest,
request: Request,
current_user: ApiPrincipal = Depends(get_current_active_user),
service: Optional[AgentChatService] = None,
persistence: Optional[AgentChatPersistenceService] = None,
) -> StreamingResponse:
"""
Web 智能助手流式对话。
@@ -1998,11 +2026,18 @@ async def _web_agent_stream_impl(
:return: SSE 流式响应
"""
prompt = payload.text.strip()
if not isinstance(service, AgentChatService):
# 直接调用公开函数时不经过 FastAPI 依赖解析;生产路由总是传入运行时服务。
service = get_configured_agent_chat_service()
if not isinstance(persistence, AgentChatPersistenceService):
# 直接调用公开函数时不经过 FastAPI 依赖解析;生产路由总是传入运行时端口。
persistence = get_configured_agent_chat_persistence()
locale = LocaleHelper.get_locale_from_request(request)
display_prompt = (payload.display_text or payload.text).strip()
session_id = await _build_web_agent_session_id_async(
current_user,
payload.session_id,
service,
)
is_secret_confirmation_candidate = (
prompt in {"确认", "取消"}
@@ -2110,13 +2145,13 @@ async def _web_agent_stream_impl(
current_user=current_user,
messages=display_messages,
client_session_id=payload.session_id or session_id,
service=service,
persistence=persistence,
)
except Exception as err:
logger.error(f"保存WebAgent传统消息快照失败: {str(err)}")
snapshot_task = asyncio.create_task(save_display_snapshot())
_WEB_AGENT_BACKGROUND_TASKS.add(snapshot_task)
snapshot_task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard)
snapshot_task = create_web_agent_background_task(save_display_snapshot())
await asyncio.sleep(0)
for event in events:
event_payload = copy.deepcopy(event)
@@ -2268,16 +2303,19 @@ async def _web_agent_stream_impl(
# 终态先进入 SSE 队列,避免展示快照落库延迟前端结束动画。
event_publisher.publish(done_event)
if not is_secret_confirmation_control:
await _save_web_agent_display_snapshot(
session_id=session_id,
current_user=current_user,
messages=display_messages,
client_session_id=payload.session_id or session_id,
)
try:
await _save_web_agent_display_snapshot(
session_id=session_id,
current_user=current_user,
messages=display_messages,
client_session_id=payload.session_id or session_id,
service=service,
persistence=persistence,
)
except Exception as err:
logger.error(f"保存WebAgent展示历史失败:{err}")
task = asyncio.create_task(run_agent())
_WEB_AGENT_BACKGROUND_TASKS.add(task)
task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard)
task = create_web_agent_background_task(run_agent())
disconnected = False
terminal_sent = False
try:
@@ -2349,6 +2387,14 @@ async def web_agent_stream(
payload: _SchemaAgentWebChatRequest,
request: Request,
current_user: ApiPrincipal = Depends(get_current_active_user),
service: AgentChatService = Depends(get_agent_chat_service),
persistence: AgentChatPersistenceService = Depends(get_agent_chat_persistence),
) -> StreamingResponse:
"""Web 智能助手流式对话的稳定公开路由入口。"""
return await _web_agent_stream_impl(payload, request, current_user)
return await _web_agent_stream_impl(
payload,
request,
current_user,
service,
persistence,
)
+22 -1
View File
@@ -1,9 +1,10 @@
import asyncio
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from queue import Queue
from threading import Lock
from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Tuple, Union
from app.schemas.types import NotificationChannel
@@ -177,6 +178,26 @@ _WEB_AGENT_EDIT_QUEUES: dict[str, list[Queue[dict]]] = {}
_WEB_AGENT_EDIT_LOCK = Lock()
_ChannelAdminResolver = Callable[[Optional[dict]], Iterable[Union[str, int]]]
_CHANNEL_ADMIN_RESOLVERS: dict[str, _ChannelAdminResolver] = {}
_WEB_AGENT_BACKGROUND_TASKS: set[asyncio.Task[object]] = set()
def create_web_agent_background_task(
coroutine: Awaitable[object],
) -> asyncio.Task[object]:
"""登记 Web Agent 后台任务,使应用关闭时可以统一收口。"""
task = asyncio.create_task(coroutine)
_WEB_AGENT_BACKGROUND_TASKS.add(task)
task.add_done_callback(_WEB_AGENT_BACKGROUND_TASKS.discard)
return task
async def shutdown_web_agent_background_tasks() -> None:
"""取消并等待 Web Agent 后台任务,避免关闭数据库后仍提交快照。"""
tasks = tuple(_WEB_AGENT_BACKGROUND_TASKS)
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
def register_channel_admin_resolver(
+49 -9
View File
@@ -10,6 +10,7 @@ from weakref import WeakValueDictionary
from app.application.database import (
AsyncDatabaseExecutor,
DatabaseWorkerClosedError,
DatabaseWorkerOverloadedError,
)
from app.schemas.agent import AgentChatSessionDetail, AgentChatSessionSummary
@@ -17,6 +18,7 @@ from app.runtime.observability import record_metric
DEFAULT_AGENT_CHAT_WRITE_CAPACITY = 32
DEFAULT_AGENT_CHAT_SESSION_CAPACITY = 4
def has_custom_agent_chat_title(value: Optional[str]) -> bool:
@@ -143,7 +145,8 @@ class SyncAgentChatRepository(Protocol):
...
SyncAgentChatRepositoryFactory = Callable[[], SyncAgentChatRepository]
SyncAgentChatRepositoryFactory = Callable[[object], SyncAgentChatRepository]
SyncAgentChatTransaction = Callable[[Callable[[object], object]], object]
@dataclass(frozen=True, slots=True)
@@ -355,15 +358,24 @@ class AgentChatPersistenceService:
self,
repository: SyncAgentChatRepositoryFactory,
async_executor: AsyncDatabaseExecutor,
sync_transaction: SyncAgentChatTransaction,
capacity: int = DEFAULT_AGENT_CHAT_WRITE_CAPACITY,
session_capacity: int = DEFAULT_AGENT_CHAT_SESSION_CAPACITY,
) -> None:
"""保存同步仓储工厂和异步执行端口"""
"""保存同步仓储工厂、事务端口和两级写入容量"""
if capacity < 1:
raise ValueError("AgentChat 写入容量必须大于 0")
if session_capacity < 1:
raise ValueError("AgentChat 单会话写入容量必须大于 0")
self._repository = repository
self._async_executor = async_executor
self._sync_transaction = sync_transaction
self._capacity = capacity
self._session_capacity = session_capacity
self._pending_writes = 0
self._pending_by_session: dict[str, int] = {}
self._active_tasks: set[asyncio.Task[object]] = set()
self._closing = False
# append_display_messages 属于读取旧快照后整列写回的复合操作;按会话串行化,
# 才能在 worker 并发下保持首次建行和既有会话追加的完整性。弱引用避免长期运行
# 中为一次性会话永久保留锁对象,不限制不同会话之间的 worker 并行度。
@@ -383,24 +395,52 @@ class AgentChatPersistenceService:
operation: Callable[[SyncAgentChatRepository], object],
) -> None:
"""在线程 worker 内完成同步写入并丢弃仓储对象返回值。"""
# 会话锁前的等待也纳入固定总量,避免公开展示保存入口形成无界应用层队列
if self._pending_writes >= self._capacity:
# 同时限制全局和单会话等待量,避免一个热点会话占满总 admission 后饿死其他会话
if self._closing:
raise DatabaseWorkerClosedError("AgentChat 持久化服务当前不可接收任务")
session_pending = self._pending_by_session.get(session_id, 0)
if (
self._pending_writes >= self._capacity
or session_pending >= self._session_capacity
):
record_metric("agent.chat.persistence.rejected")
raise DatabaseWorkerOverloadedError(
f"AgentChat 写入容量已用尽(上限 {self._capacity}"
f"AgentChat 写入容量已用尽(全局上限 {self._capacity}"
f"单会话上限 {self._session_capacity}"
)
self._pending_writes += 1
record_metric("agent.chat.persistence.pending", self._pending_writes)
self._pending_by_session[session_id] = session_pending + 1
current = asyncio.current_task()
if current is not None:
self._active_tasks.add(current)
record_metric("agent.chat.persistence.pending", 1)
try:
async with self._session_lock(session_id):
def execute() -> None:
"""执行同步写入,不让 ORM 对象越过 worker 边界。"""
operation(self._repository())
"""在单一同步事务中执行写入,不让 ORM 对象越过 worker 边界。"""
self._sync_transaction(
lambda session: operation(self._repository(session))
)
await self._async_executor.run(execute)
finally:
self._pending_writes -= 1
record_metric("agent.chat.persistence.pending", self._pending_writes)
remaining = self._pending_by_session.get(session_id, 1) - 1
if remaining:
self._pending_by_session[session_id] = remaining
else:
self._pending_by_session.pop(session_id, None)
if current is not None:
self._active_tasks.discard(current)
record_metric("agent.chat.persistence.pending", -1)
async def shutdown(self) -> None:
"""拒绝新写入并等待当前会话锁和 worker 操作取得终态。"""
self._closing = True
current = asyncio.current_task()
tasks = tuple(task for task in self._active_tasks if task is not current)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
async def async_append_display_messages(
self,
+2 -1
View File
@@ -269,7 +269,8 @@ class AgentChatOper(DbOper):
)
if not chat:
return None
display_messages = self._normalize_messages(chat.display_messages)
# JSON 列不是 MutableList;必须复制旧列表,原地 extend 会让 SQLAlchemy 误认为字段未变化。
display_messages = list(self._normalize_messages(chat.display_messages))
display_messages.extend(self._normalize_messages(messages))
title = chat.title if self.has_custom_title(chat.title) else None
return self.save_display_messages(
+2
View File
@@ -6,6 +6,7 @@ from typing import Protocol
from app.application.messaging.chat import (
AsyncAgentChatRepository,
AgentChatPersistenceService,
AsyncUnitOfWork,
)
from app.application.outbox import AsyncOutboxTransaction
@@ -113,6 +114,7 @@ class AgentChatRuntime:
async_session: AsyncSessionProvider
repository: AgentChatRepositoryFactory
transaction: AsyncUnitOfWorkFactory
persistence: AgentChatPersistenceService
@dataclass(frozen=True, slots=True)
+15 -7
View File
@@ -67,7 +67,9 @@ from app.application.messaging.chat import (
AgentChatService,
configure_agent_chat_persistence,
configure_agent_chat_service,
get_configured_agent_chat_persistence,
)
from app.application.messaging.agent import shutdown_web_agent_background_tasks
from app.application.security.user import configure_user_lookups
from app.application.security.auth import AuthService, configure_auth_service
from app.application.security.passkeys import PasskeyService, configure_passkey_service
@@ -576,6 +578,11 @@ async def stop_modules():
await run_step("消息服务", stop_message)
await run_step("Redis缓存连接", lambda: RedisHelper().close())
await run_step("异步Redis缓存连接", lambda: AsyncRedisHelper().close())
await run_step("Web Agent后台任务", shutdown_web_agent_background_tasks)
await run_step(
"Agent会话持久化",
lambda: get_configured_agent_chat_persistence().shutdown(),
)
await run_step("数据库任务", stop_database_worker)
if _database_worker is None:
await run_step("数据库连接", close_database)
@@ -642,11 +649,18 @@ async def init_modules() -> HostRuntime:
chain=lambda: build_chain_runtime_config(settings),
)
runtime_settings = _build_runtime_settings_service()
agent_chat_persistence = AgentChatPersistenceService(
repository=lambda session: AgentChatOper(session),
async_executor=database_worker,
sync_transaction=transaction_runner.sync,
capacity=database_worker.snapshot().capacity,
)
host_runtime = HostRuntime(
agent_chat=AgentChatRuntime(
async_session=get_async_db,
repository=AgentChatOper,
transaction=SqlAlchemyAsyncUnitOfWork,
persistence=agent_chat_persistence,
),
persistence=PersistenceRuntime(
sync_session=get_db,
@@ -714,13 +728,7 @@ async def init_modules() -> HostRuntime:
)
configure_database_governance(build_database_governance())
configure_agent_chat_service(AgentChatService(repository=AgentChatOper()))
configure_agent_chat_persistence(
AgentChatPersistenceService(
repository=AgentChatOper,
async_executor=database_worker,
capacity=database_worker.snapshot().capacity,
)
)
configure_agent_chat_persistence(agent_chat_persistence)
configure_user_lookups(
by_id=lambda user_id: UserOper().get_by_id(user_id),
by_name=lambda username: UserOper().get_by_name(username),