vent_agent_2.0/agents/agent_memory.py

"""图状态（checkpoint）存储。

与 1.5 的关键差异：1.5 用了一个自定义的 ``_PerLoopSaver`` 去桥接
「异步按事件循环分发 / 同步每次开新连接」——那套桥接是为了应付多事件
循环的同进程场景，代价是每个同步调用都要新开一次 sqlite 连接（见 1.5
agents/agent_memory.py 的注释）。2.0 的单进程约束下（uvicorn workers=1）
只有一个事件循环，不需要这层桥接：启动时创建一次 saver，全进程复用。

代价是 **不能多 worker**：多进程各自持有自己的 sqlite 文件句柄与图状态，
并发锁与调度器也仍使用进程级状态。RuntimeEnvironment 启动时提示单进程约束。
"""

from __future__ import annotations

import asyncio
from typing import Any

import config
from obs import get_logger

logger = get_logger(__name__)

# 进程级单例。初始化前为 None，避免在无事件循环时（import 期）创建连接。
_saver: Any | None = None
_init_lock = asyncio.Lock()
_kind: str = config.CHECKPOINTER


class CheckpointerConfigError(RuntimeError):
    """checkpointer 配置不可用。刻意 fail-fast，不静默回落到内存。"""


def saver_kind() -> str:
    """当前档位：sqlite | memory。"""
    return _kind


def _serializer():
    """允许应用声明的业务结果类型，其他自定义类不从 checkpoint 动态反序列化。"""
    from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
    from extensions.registry import load_extensions
    allowed = [(spec.result_schema.__module__, spec.result_schema.__name__)
               for spec in load_extensions().business_agents.values()]
    return JsonPlusSerializer(allowed_msgpack_modules=allowed)


def get_checkpointer() -> Any:
    """取进程级 checkpointer。

    ``memory`` 档在首次调用时**惰性创建**——InMemorySaver 是同步可构造的，
    没有任何初始化依赖，却要求调用方先进 lifespan 是没必要的仪式，也会让
    脚本、测试、调度器这些非 HTTP 入口平白多一层耦合。

    ``sqlite`` 档仍然要求先 ``await init_checkpointer()``：它需要打开异步
    连接，没法在同步函数里做，而且**绝不能静默回落到内存**——那会让
    「重启后中断恢复失效」变成一个没人发现的行为变化。
    """
    global _saver
    if _saver is None and _kind == "memory":
        from langgraph.checkpoint.memory import InMemorySaver

        _saver = InMemorySaver(serde=_serializer())
        logger.info("checkpointer=memory 惰性创建（图状态不落盘）")
    if _saver is None:
        raise CheckpointerConfigError(
            "checkpointer 尚未初始化：请在 async with RuntimeEnvironment() 内执行任务"
        )
    return _saver


def is_initialized() -> bool:
    """是否已初始化（/health 用）。"""
    return _saver is not None


async def init_checkpointer() -> Any:
    """初始化 checkpointer。幂等，重复调用直接返回已有实例。"""
    global _saver, _kind

    async with _init_lock:
        if _saver is not None:
            return _saver

        if _kind == "memory":
            from langgraph.checkpoint.memory import InMemorySaver

            _saver = InMemorySaver(serde=_serializer())
            logger.warning(
                "checkpointer=memory：图状态不落盘，重启后中断恢复失效"
                "（会话历史仍从业务数据库补种）"
            )
            return _saver

        if _kind != "sqlite":
            raise CheckpointerConfigError(
                f"不支持的 CHECKPOINTER={_kind!r}，可选值：sqlite | memory"
            )

        try:
            import aiosqlite
            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
        except ImportError as exc:
            raise CheckpointerConfigError(
                f"checkpointer=sqlite 需要 aiosqlite 与 langgraph-checkpoint-sqlite：{exc}"
            ) from exc

        path = config.resolve_path(config.CHECKPOINT_DB_PATH)
        conn = None
        try:
            conn = await aiosqlite.connect(str(path))
            saver = AsyncSqliteSaver(conn, serde=_serializer())
            await saver.setup()
        except BaseException as exc:
            if conn is not None:
                await conn.close()
            if not isinstance(exc, Exception):
                raise
            raise CheckpointerConfigError(
                f"打开 checkpoint 库失败 path={path}：{type(exc).__name__}: {exc}"
            ) from exc

        _saver = saver
        logger.info("checkpointer=sqlite 就绪 path=%s", path)
        return _saver


async def close_checkpointer() -> None:
    """关闭底层连接（RuntimeEnvironment 退出时调用）。"""
    global _saver
    saver = _saver
    _saver = None
    if saver is None:
        return
    conn = getattr(saver, "conn", None)
    if conn is not None and hasattr(conn, "close"):
        try:
            await conn.close()
        except Exception as exc:  # pragma: no cover - 关闭失败不影响进程退出
            logger.warning("关闭 checkpoint 连接失败：%s: %s", type(exc).__name__, exc)
    logger.info("checkpointer 已关闭")


def thread_config(thread_id: str) -> dict[str, Any]:
    """构造 LangGraph 的 configurable config。

    ``thread_id`` 恒等于 ``session_id``（由 api/chat_routes 强制），
    这是防止越权读写别人会话状态的第一道闸。
    """
    return {"configurable": {"thread_id": thread_id}}


async def has_checkpoint(thread_id: str) -> bool:
    """该 thread 是否已有图状态。

    冷线程需要从业务数据库补种历史消息；热线程只喂当轮新消息，否则历史会被
    重复追加，上下文按 O(N²) 膨胀。
    """
    saver = get_checkpointer()
    try:
        tup = await saver.aget_tuple(thread_config(thread_id))
    except Exception as exc:  # pragma: no cover - 取决于底层库
        logger.warning("探测 thread 状态失败 thread=%s：%s: %s", thread_id, type(exc).__name__, exc)
        raise
    return tup is not None


async def delete_thread(thread_id: str) -> None:
    """删除某 thread 的全部图状态。删会话时必须同步调用，否则图状态残留。"""
    saver = get_checkpointer()
    delete = getattr(saver, "adelete_thread", None)
    if delete is None:
        raise RuntimeError("当前 checkpointer 不支持删除 thread")
    try:
        await delete(thread_id)
        logger.info("已清除图状态 thread=%s", thread_id)
    except Exception as exc:
        logger.warning("清除图状态失败 thread=%s：%s: %s", thread_id, type(exc).__name__, exc)
        raise