"""Async checkpointer factory. Provides an **async context manager** for long-running async servers that need proper resource cleanup. Supported backends: memory, sqlite, postgres. Usage (e.g. FastAPI lifespan):: from deerflow.agents.checkpointer.async_provider import make_checkpointer async with make_checkpointer() as checkpointer: app.state.checkpointer = checkpointer # InMemorySaver if not configured For sync usage see :mod:`deerflow.agents.checkpointer.provider`. """ from __future__ import annotations import contextlib import logging from collections.abc import AsyncIterator from langgraph.types import Checkpointer from deerflow.agents.checkpointer.provider import ( POSTGRES_CONN_REQUIRED, POSTGRES_INSTALL, SQLITE_INSTALL, _resolve_sqlite_conn_str, ) from deerflow.config.app_config import get_app_config logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Async factory # --------------------------------------------------------------------------- @contextlib.asynccontextmanager async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]: """Async context manager that constructs and tears down a checkpointer.""" if config.type == "memory": from langgraph.checkpoint.memory import InMemorySaver yield InMemorySaver() return if config.type == "sqlite": try: from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver except ImportError as exc: raise ImportError(SQLITE_INSTALL) from exc import pathlib conn_str = _resolve_sqlite_conn_str(config.connection_string or "store.db") # Only create parent directories for real filesystem paths if conn_str != ":memory:" and not conn_str.startswith("file:"): pathlib.Path(conn_str).parent.mkdir(parents=True, exist_ok=True) async with AsyncSqliteSaver.from_conn_string(conn_str) as saver: await saver.setup() yield saver return if config.type == "postgres": try: from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver except ImportError as exc: raise ImportError(POSTGRES_INSTALL) from exc if not config.connection_string: raise ValueError(POSTGRES_CONN_REQUIRED) async with AsyncPostgresSaver.from_conn_string(config.connection_string) as saver: await saver.setup() yield saver return raise ValueError(f"Unknown checkpointer type: {config.type!r}") # --------------------------------------------------------------------------- # Public async context manager # --------------------------------------------------------------------------- @contextlib.asynccontextmanager async def make_checkpointer() -> AsyncIterator[Checkpointer]: """Async context manager that yields a checkpointer for the caller's lifetime. Resources are opened on enter and closed on exit — no global state:: async with make_checkpointer() as checkpointer: app.state.checkpointer = checkpointer Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*. """ config = get_app_config() if config.checkpointer is None: from langgraph.checkpoint.memory import InMemorySaver yield InMemorySaver() return async with _async_checkpointer(config.checkpointer) as saver: yield saver