"""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.runtime.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.runtime.checkpointer.provider`. """ from __future__ import annotations import asyncio import contextlib import logging from collections.abc import AsyncIterator import os from pathlib import Path from langgraph.types import Checkpointer from deerflow.config.app_config import AppConfig, get_app_config from deerflow.runtime.checkpointer.provider import ( POSTGRES_CONN_REQUIRED, POSTGRES_INSTALL, SQLITE_INSTALL, ) from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str logger = logging.getLogger(__name__) def _csv_env(name: str) -> list[str]: raw = os.getenv(name, "") return [part.strip() for part in raw.split(",") if part.strip()] def _dual_read_enabled(config) -> bool: if getattr(config, "dual_read_enabled", False): return True return os.getenv("DEER_FLOW_CHECKPOINT_DUAL_READ", "").strip().lower() in {"1", "true", "yes", "on"} def _use_postgres_enabled(config) -> bool: raw = os.getenv("DEER_FLOW_USE_PG_CHECKPOINTER") if raw: return raw.strip().lower() not in {"0", "false", "no", "off"} return bool(getattr(config, "use_postgres", True)) def _dual_read_paths(config) -> list[str]: paths = list(getattr(config, "dual_read_sqlite_paths", []) or []) paths.extend(_csv_env("DEER_FLOW_CHECKPOINT_SQLITE_PATHS")) resolved: list[str] = [] seen: set[str] = set() for path in paths: conn_str = resolve_sqlite_conn_str(path) if conn_str == ":memory:" or conn_str.startswith("file:"): logger.warning("Skipping unsupported checkpoint migration SQLite path: %s", path) continue if conn_str not in seen: resolved.append(conn_str) seen.add(conn_str) return resolved def _sqlite_conn_str_for_postgres_fallback(config) -> str: configured = getattr(config, "sqlite_connection_string", None) if configured: return configured paths = _dual_read_paths(config) return paths[0] if paths else "store.db" def _pool_size_env(name: str, default: int) -> int: raw = os.getenv(name, "").strip() if not raw: return default try: return max(0, int(raw)) except ValueError: logger.warning("Invalid %s=%r; using default %d", name, raw, default) return default @contextlib.asynccontextmanager async def _postgres_saver_from_pool(conn_string: str) -> AsyncIterator[Checkpointer]: """Yield an ``AsyncPostgresSaver`` backed by a self-healing connection pool. ``AsyncPostgresSaver.from_conn_string`` opens a *single* long-lived connection. If that one connection dies (Postgres restart, idle timeout, network blip, or a prior broken pipeline), every later checkpointer call fails permanently with "the connection is closed" until the process is restarted. A pool with ``check=check_connection`` validates each connection on checkout and transparently replaces dead ones, so the runtime recovers on its own. """ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from psycopg.rows import dict_row from psycopg_pool import AsyncConnectionPool pool = AsyncConnectionPool( conninfo=conn_string, min_size=_pool_size_env("DEER_FLOW_PG_POOL_MIN", 1), max_size=_pool_size_env("DEER_FLOW_PG_POOL_MAX", 10), max_idle=300.0, max_lifetime=3600.0, check=AsyncConnectionPool.check_connection, kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row}, open=False, ) await pool.open(wait=True) try: saver = AsyncPostgresSaver(pool) await saver.setup() yield saver finally: await pool.close() @contextlib.asynccontextmanager async def _sqlite_checkpointer_from_conn_str(raw: str) -> AsyncIterator[Checkpointer]: try: from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver except ImportError as exc: raise ImportError(SQLITE_INSTALL) from exc conn_str = resolve_sqlite_conn_str(raw) await asyncio.to_thread(ensure_sqlite_parent_dir, conn_str) async with AsyncSqliteSaver.from_conn_string(conn_str) as saver: await saver.setup() yield saver # --------------------------------------------------------------------------- # 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": async with _sqlite_checkpointer_from_conn_str(config.connection_string or "store.db") as saver: yield saver return if config.type == "postgres": if not _use_postgres_enabled(config): conn_str = _sqlite_conn_str_for_postgres_fallback(config) logger.info("Checkpointer: use_postgres=false; using SQLite fallback (%s)", conn_str) async with _sqlite_checkpointer_from_conn_str(conn_str) as saver: yield saver return try: from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # noqa: F401 except ImportError as exc: raise ImportError(POSTGRES_INSTALL) from exc if not config.connection_string: raise ValueError(POSTGRES_CONN_REQUIRED) async with _postgres_saver_from_pool(config.connection_string) as saver: if not _dual_read_enabled(config): yield saver return from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from deerflow.runtime.checkpointer.checkpoint_migration import ( DualReadCheckpointer, SQLiteCheckpointSource, ) sources: list[SQLiteCheckpointSource] = [] for sqlite_path in _dual_read_paths(config): if not Path(sqlite_path).exists(): logger.info("Checkpoint dual-read source does not exist, skipping: %s", sqlite_path) continue ctx = AsyncSqliteSaver.from_conn_string(sqlite_path) source_saver = await ctx.__aenter__() await source_saver.setup() sources.append( SQLiteCheckpointSource( path=sqlite_path, saver=source_saver, context=ctx, progress_path=Path(sqlite_path).with_suffix(Path(sqlite_path).suffix + ".pg-migration.json"), ) ) if not sources: logger.info("Checkpoint dual-read enabled but no SQLite sources were opened; using Postgres only") yield saver return proxy = DualReadCheckpointer( saver, sources, ) try: yield proxy finally: await proxy.aclose_sources() return raise ValueError(f"Unknown checkpointer type: {config.type!r}") # --------------------------------------------------------------------------- # Public async context manager # --------------------------------------------------------------------------- @contextlib.asynccontextmanager async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpointer]: """Async context manager that constructs a checkpointer from unified DatabaseConfig.""" if db_config.backend == "memory": from langgraph.checkpoint.memory import InMemorySaver yield InMemorySaver() return if db_config.backend == "sqlite": try: from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver except ImportError as exc: raise ImportError(SQLITE_INSTALL) from exc conn_str = db_config.checkpointer_sqlite_path ensure_sqlite_parent_dir(conn_str) async with AsyncSqliteSaver.from_conn_string(conn_str) as saver: await saver.setup() yield saver return if db_config.backend == "postgres": try: from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # noqa: F401 except ImportError as exc: raise ImportError(POSTGRES_INSTALL) from exc if not db_config.postgres_url: raise ValueError("database.postgres_url is required for the postgres backend") async with _postgres_saver_from_pool(db_config.postgres_url) as saver: yield saver return raise ValueError(f"Unknown database backend: {db_config.backend!r}") @contextlib.asynccontextmanager async def make_checkpointer(app_config: AppConfig | None = None) -> 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(app_config) as checkpointer: app.state.checkpointer = checkpointer Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*. Priority: 1. Legacy ``checkpointer:`` config section (backward compatible) 2. Unified ``database:`` config section 3. Default InMemorySaver """ if app_config is None: app_config = get_app_config() # Legacy: standalone checkpointer config takes precedence if app_config.checkpointer is not None: async with _async_checkpointer(app_config.checkpointer) as saver: yield saver return # Unified database config db_config = getattr(app_config, "database", None) if db_config is not None and db_config.backend != "memory": async with _async_checkpointer_from_database(db_config) as saver: yield saver return # Default: in-memory from langgraph.checkpoint.memory import InMemorySaver yield InMemorySaver()