sync update from LAN repo (source dhk-1 @ 3b1ec5b4)

2026-11-05 10:56:34 +08:00
1 parent da6521b6cb
commit 264e2200d6
@@ -20,13 +20,14 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import logging
-from collections.abc import AsyncIterator
 import os
+from collections.abc import AsyncIterator
 from pathlib import Path

 from langgraph.types import Checkpointer

 from deerflow.config.app_config import AppConfig, get_app_config
+from deerflow.runtime.checkpointer.chat_archive import archive_root
 from deerflow.runtime.checkpointer.provider import (
    POSTGRES_CONN_REQUIRED,
    POSTGRES_INSTALL,
@@ -135,7 +136,11 @@ async def _sqlite_checkpointer_from_conn_str(raw: str) -> AsyncIterator[Checkpoi
    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()
+        from deerflow.runtime.store._sqlite_setup import sqlite_setup_lock
+
+        async with sqlite_setup_lock(conn_str):
+            await saver.conn.execute("PRAGMA busy_timeout=30000")
+            await saver.setup()
        yield saver


@@ -174,15 +179,20 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
        if not config.connection_string:
            raise ValueError(POSTGRES_CONN_REQUIRED)

+        from deerflow.runtime.checkpointer.checkpoint_migration import DualReadCheckpointer
+
        async with _postgres_saver_from_pool(config.connection_string) as saver:
            if not _dual_read_enabled(config):
-                yield saver
+                proxy = DualReadCheckpointer(saver, [], archive_dir=archive_root())
+                try:
+                    yield proxy
+                finally:
+                    await proxy.aclose_sources()
                return

            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver

            from deerflow.runtime.checkpointer.checkpoint_migration import (
-                DualReadCheckpointer,
                SQLiteCheckpointSource,
            )

@@ -205,12 +215,17 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:

            if not sources:
                logger.info("Checkpoint dual-read enabled but no SQLite sources were opened; using Postgres only")
-                yield saver
+                proxy = DualReadCheckpointer(saver, [], archive_dir=archive_root())
+                try:
+                    yield proxy
+                finally:
+                    await proxy.aclose_sources()
                return

            proxy = DualReadCheckpointer(
                saver,
                sources,
+                archive_dir=archive_root(),
            )
            try:
                yield proxy
@@ -257,8 +272,14 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
        if not db_config.postgres_url:
            raise ValueError("database.postgres_url is required for the postgres backend")

+        from deerflow.runtime.checkpointer.checkpoint_migration import DualReadCheckpointer
+
        async with _postgres_saver_from_pool(db_config.postgres_url) as saver:
-            yield saver
+            proxy = DualReadCheckpointer(saver, [], archive_dir=archive_root())
+            try:
+                yield proxy
+            finally:
+                await proxy.aclose_sources()
        return

    raise ValueError(f"Unknown database backend: {db_config.backend!r}")
@@ -0,0 +1,222 @@
+"""Persistent local catalog for monthly PostgreSQL chat archives.
+
+The archive directory is deliberately independent from the configured
+checkpointer SQLite file: existing SQLite deployments and legacy dual-read
+paths keep their current behavior.
+"""
+
+from __future__ import annotations
+
+import errno
+import os
+import re
+import sqlite3
+import threading
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import BinaryIO
+
+from deerflow.config.runtime_paths import runtime_home
+
+_MONTH_RE = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$")
+_WINDOWS_LOCKS: dict[str, threading.Lock] = {}
+_WINDOWS_LOCKS_GUARD = threading.Lock()
+_WINDOWS_LOCK_TYPE = type(threading.Lock())
+
+
+def archive_root() -> Path:
+    """Return the persistent archive directory used by the gateway/workers."""
+
+    configured = os.getenv("DEER_FLOW_CHAT_ARCHIVE_DIR", "").strip()
+    if configured:
+        return Path(configured).expanduser().resolve()
+    # The standard deployment keeps ``<DEER_FLOW_HOME>/data`` on a shared,
+    # persistent Docker volume. Workers and the gateway therefore see the
+    # same archive files and cross-process lock files after restarts.
+    return runtime_home() / "data" / "chat-archives"
+
+
+def archive_db_path(root: Path, month: str) -> Path:
+    if not _MONTH_RE.fullmatch(month):
+        raise ValueError(f"Invalid archive month: {month!r}")
+    return root / f"chat-{month}.db"
+
+
+def catalog_path(root: Path) -> Path:
+    return root / "catalog.db"
+
+
+def try_acquire_archive_job_lock(root: Path) -> BinaryIO | None:
+    """Hold one archive operation across gateway processes and worker containers.
+
+    The kernel releases this lock on process exit. Recovery can therefore
+    distinguish a dead runner from an older but still active gateway process.
+    """
+    root.mkdir(parents=True, exist_ok=True)
+    handle = (root / ".archive-operation.lock").open("a+b")
+    try:
+        if os.name == "nt":
+            import msvcrt
+
+            if handle.seek(0, os.SEEK_END) == 0:
+                handle.write(b"\0")
+                handle.flush()
+            handle.seek(0)
+            msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
+        else:
+            import fcntl
+
+            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        return handle
+    except OSError as exc:
+        handle.close()
+        if exc.errno in {errno.EACCES, errno.EAGAIN, errno.EDEADLK}:
+            return None
+        raise
+    except BaseException:
+        handle.close()
+        raise
+
+
+def ensure_catalog(root: Path) -> Path:
+    root.mkdir(parents=True, exist_ok=True)
+    path = catalog_path(root)
+    with sqlite3.connect(path, timeout=30) as conn:
+        conn.execute("PRAGMA journal_mode=WAL")
+        conn.execute("PRAGMA busy_timeout=30000")
+        conn.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS archive_months (
+                month TEXT PRIMARY KEY,
+                status TEXT NOT NULL DEFAULT 'empty',
+                total_threads INTEGER NOT NULL DEFAULT 0,
+                copied_threads INTEGER NOT NULL DEFAULT 0,
+                deleted_threads INTEGER NOT NULL DEFAULT 0,
+                restored_threads INTEGER NOT NULL DEFAULT 0,
+                current_job_id TEXT,
+                updated_at TEXT NOT NULL
+            );
+            CREATE TABLE IF NOT EXISTS archive_threads (
+                thread_id TEXT PRIMARY KEY,
+                month TEXT NOT NULL,
+                status TEXT NOT NULL,
+                checkpoint_count INTEGER NOT NULL DEFAULT 0,
+                archived_at TEXT NOT NULL,
+                deleted_at TEXT,
+                restored_at TEXT
+            );
+            CREATE INDEX IF NOT EXISTS ix_archive_threads_month_status
+                ON archive_threads(month, status);
+            CREATE TABLE IF NOT EXISTS archive_jobs (
+                job_id TEXT PRIMARY KEY,
+                month TEXT NOT NULL,
+                action TEXT NOT NULL,
+                status TEXT NOT NULL,
+                total_threads INTEGER NOT NULL DEFAULT 0,
+                processed_threads INTEGER NOT NULL DEFAULT 0,
+                failed_threads INTEGER NOT NULL DEFAULT 0,
+                current_thread_id TEXT,
+                snapshot_at TEXT NOT NULL DEFAULT '',
+                processed_through TEXT,
+                message TEXT NOT NULL DEFAULT '',
+                error TEXT,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL
+            );
+            CREATE INDEX IF NOT EXISTS ix_archive_jobs_month_created
+                ON archive_jobs(month, created_at DESC);
+            """
+        )
+        # Upgrade catalogs created by earlier releases without replacing data.
+        columns = {str(row[1]) for row in conn.execute("PRAGMA table_info(archive_jobs)")}
+        if "snapshot_at" not in columns:
+            conn.execute("ALTER TABLE archive_jobs ADD COLUMN snapshot_at TEXT NOT NULL DEFAULT ''")
+        if "processed_through" not in columns:
+            conn.execute("ALTER TABLE archive_jobs ADD COLUMN processed_through TEXT")
+    return path
+
+
+def archived_month_for_thread(root: Path, thread_id: str) -> str | None:
+    path = catalog_path(root)
+    if not path.is_file():
+        return None
+    try:
+        with sqlite3.connect(f"file:{path.as_posix()}?mode=ro", uri=True, timeout=5) as conn:
+            row = conn.execute(
+                "SELECT month FROM archive_threads WHERE thread_id = ? AND status IN ('copied', 'deleted', 'restored')",
+                (thread_id,),
+            ).fetchone()
+    except sqlite3.Error:
+        return None
+    return str(row[0]) if row else None
+
+
+def _acquire_thread_file_lock(root: Path, thread_id: str, exclusive: bool) -> int | _WINDOWS_LOCK_TYPE:
+    lock_dir = root / ".locks"
+    lock_dir.mkdir(parents=True, exist_ok=True)
+    import hashlib
+
+    # Stripe locks keep the lock directory bounded instead of creating one
+    # permanent inode for every conversation ever written.
+    name = hashlib.sha256(thread_id.encode("utf-8")).hexdigest()[:3] + ".lock"
+    path = lock_dir / name
+    try:
+        import fcntl
+    except ImportError:  # pragma: no cover - Windows test/development fallback
+        key = str(path.resolve())
+        with _WINDOWS_LOCKS_GUARD:
+            lock = _WINDOWS_LOCKS.setdefault(key, threading.Lock())
+        lock.acquire()
+        return lock
+    descriptor = os.open(path, os.O_CREAT | os.O_RDWR, 0o600)
+    fcntl.flock(descriptor, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
+    return descriptor
+
+
+def _release_thread_file_lock(lock: int | _WINDOWS_LOCK_TYPE) -> None:
+    if isinstance(lock, int):
+        import fcntl
+
+        fcntl.flock(lock, fcntl.LOCK_UN)
+        os.close(lock)
+    else:  # pragma: no cover - Windows test/development fallback
+        lock.release()
+
+
+@asynccontextmanager
+async def thread_archive_lock(root: Path, thread_id: str, *, exclusive: bool = False) -> AsyncIterator[None]:
+    """Cross-process read/write gate shared by the main service and workers.
+
+    Linux ``flock`` readers do not serialize one another. A purge acquires an
+    exclusive lock, waits for current checkpoint writes to finish, and holds
+    off new writes only for the short copy/verify/delete section. Kernel locks
+    are released automatically if a process exits unexpectedly.
+    """
+
+    import asyncio
+
+    lock = await asyncio.to_thread(_acquire_thread_file_lock, root, thread_id, exclusive)
+    try:
+        yield
+    finally:
+        await asyncio.to_thread(_release_thread_file_lock, lock)
+
+
+def remove_archived_thread(root: Path, thread_id: str) -> None:
+    path = catalog_path(root)
+    if not path.is_file():
+        return
+    with sqlite3.connect(path, timeout=30) as conn:
+        conn.execute("PRAGMA busy_timeout=30000")
+        row = conn.execute("SELECT month, status FROM archive_threads WHERE thread_id = ?", (thread_id,)).fetchone()
+        if not row:
+            return
+        month = str(row[0])
+        status = str(row[1])
+        conn.execute("DELETE FROM archive_threads WHERE thread_id = ?", (thread_id,))
+        conn.execute(
+            "UPDATE archive_months SET copied_threads = MAX(0, copied_threads - 1), "
+            "deleted_threads = MAX(0, deleted_threads - ?), updated_at = CURRENT_TIMESTAMP WHERE month = ?",
+            (1 if status == "deleted" else 0, month),
+        )
@@ -16,6 +16,14 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any

+from deerflow.runtime.checkpointer.chat_archive import (
+    archive_db_path,
+    archived_month_for_thread,
+    catalog_path,
+    remove_archived_thread,
+    thread_archive_lock,
+)
+
 logger = logging.getLogger(__name__)


@@ -166,25 +174,97 @@ async def target_has_thread(target, thread_id: str) -> bool:


 class DualReadCheckpointer:
-    """Postgres-first checkpointer with SQLite read fallback."""
+    """Postgres-first checkpointer with legacy and monthly archive fallbacks."""

-    def __init__(self, primary, sources: list[SQLiteCheckpointSource], *, sleep_seconds: float = 0.05):
+    _PRIMARY_ONLY_MAX = 8192
+
+    def __init__(
+        self,
+        primary,
+        sources: list[SQLiteCheckpointSource],
+        *,
+        archive_dir: Path | None = None,
+        sleep_seconds: float = 0.05,
+    ):
        self.primary = primary
        self.sources = sources
+        self.archive_dir = archive_dir
        self.sleep_seconds = max(0.0, sleep_seconds)
        self._closed = False
+        self._archive_sources: dict[str, SQLiteCheckpointSource] = {}
+        self._archive_sources_lock = asyncio.Lock()
+        # Threads created or written after dual-read started cannot exist in
+        # the legacy SQLite files. Skipping that fallback avoids opening
+        # multi-GB sqlite sources on every empty /state or /history read.
+        self._primary_only: set[str] = set()

    def __getattr__(self, name: str):
        return getattr(self.primary, name)

+    def remember_primary_thread(self, thread_id: str | None) -> None:
+        """Mark *thread_id* as primary-only so empty reads skip SQLite."""
+        if not thread_id:
+            return
+        self._primary_only.add(str(thread_id))
+        overflow = len(self._primary_only) - self._PRIMARY_ONLY_MAX
+        if overflow > 0:
+            for stale in list(self._primary_only)[:overflow]:
+                self._primary_only.discard(stale)
+
+    def _skip_legacy_fallback(self, config) -> bool:
+        thread_id = ((config or {}).get("configurable") or {}).get("thread_id")
+        return bool(thread_id) and str(thread_id) in self._primary_only
+
    async def aclose_sources(self) -> None:
        self._closed = True
-        for source in self.sources:
+        for source in [*self.sources, *self._archive_sources.values()]:
            await source.close()
+        self._archive_sources.clear()
+
+    @contextlib.asynccontextmanager
+    async def thread_write_lock(self, thread_id: str, *, exclusive: bool = False):
+        """Coordinate archive deletion with writes across all PG-backed workers."""
+        if not self.archive_dir or not catalog_path(self.archive_dir).is_file():
+            yield
+            return
+        async with thread_archive_lock(self.archive_dir, str(thread_id), exclusive=exclusive):
+            yield
+
+    async def _archive_source_for_thread(self, thread_id: str | None) -> SQLiteCheckpointSource | None:
+        if not self.archive_dir or not thread_id or self._closed:
+            return None
+        month = await asyncio.to_thread(archived_month_for_thread, self.archive_dir, thread_id)
+        if not month:
+            return None
+        cached = self._archive_sources.get(month)
+        if cached is not None and not cached.closed:
+            return cached
+        async with self._archive_sources_lock:
+            cached = self._archive_sources.get(month)
+            if cached is not None and not cached.closed:
+                return cached
+            path = archive_db_path(self.archive_dir, month)
+            if not path.is_file():
+                return None
+            from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+
+            context = AsyncSqliteSaver.from_conn_string(str(path))
+            saver = await context.__aenter__()
+            await saver.conn.execute("PRAGMA busy_timeout=30000")
+            if not await asyncio.to_thread(_sqlite_has_table, str(path), "checkpoints"):
+                await saver.setup()
+            source = SQLiteCheckpointSource(
+                path=str(path),
+                saver=saver,
+                context=context,
+                progress_path=path.with_suffix(path.suffix + ".archive-progress.json"),
+            )
+            self._archive_sources[month] = source
+            return source

    async def aget_tuple(self, config):
        found = await self.primary.aget_tuple(config)
-        if found is not None:
+        if found is not None or self._skip_legacy_fallback(config):
            return found
        for source in self.sources:
            if source.closed:
@@ -192,11 +272,17 @@ class DualReadCheckpointer:
            found = await source.saver.aget_tuple(config)
            if found is not None:
                return found
+        configurable = (config or {}).get("configurable", {}) or {}
+        archive = await self._archive_source_for_thread(configurable.get("thread_id"))
+        if archive is not None:
+            found = await archive.saver.aget_tuple(config)
+            if found is not None:
+                return found
        return None

    async def aget(self, config):
        found = await self.primary.aget(config)
-        if found is not None:
+        if found is not None or self._skip_legacy_fallback(config):
            return found
        for source in self.sources:
            if source.closed:
@@ -204,6 +290,12 @@ class DualReadCheckpointer:
            found = await source.saver.aget(config)
            if found is not None:
                return found
+        configurable = (config or {}).get("configurable", {}) or {}
+        archive = await self._archive_source_for_thread(configurable.get("thread_id"))
+        if archive is not None:
+            found = await archive.saver.aget(config)
+            if found is not None:
+                return found
        return None

    async def alist(self, config, *, filter=None, before=None, limit=None):
@@ -218,6 +310,9 @@ class DualReadCheckpointer:
            if limit is not None and yielded >= limit:
                return

+        if self._skip_legacy_fallback(config):
+            return
+
        remaining = None if limit is None else max(0, limit - yielded)
        for source in self.sources:
            if source.closed:
@@ -233,23 +328,50 @@ class DualReadCheckpointer:
                    remaining = max(0, limit - yielded)
                    if remaining <= 0:
                        return
+        configurable = (config or {}).get("configurable", {}) or {}
+        archive = await self._archive_source_for_thread(configurable.get("thread_id"))
+        if archive is not None and not archive.closed:
+            async for tup in archive.saver.alist(config, filter=filter, before=before, limit=remaining):
+                key = self._tuple_key(tup)
+                if key in seen:
+                    continue
+                yield tup
+                yielded += 1
+                if limit is not None and yielded >= limit:
+                    return

    async def aput(self, *args, **kwargs):
+        config = args[0] if args else kwargs.get("config")
+        if isinstance(config, dict):
+            self.remember_primary_thread(((config.get("configurable") or {}).get("thread_id")))
        return await self.primary.aput(*args, **kwargs)

    async def aput_writes(self, *args, **kwargs):
-        return await self.primary.aput_writes(*args, **kwargs)
+        config = args[0] if args else kwargs.get("config")
+        thread_id = ((config or {}).get("configurable", {}) or {}).get("thread_id")
+        if not thread_id:
+            return await self.primary.aput_writes(*args, **kwargs)
+        async with self.thread_write_lock(str(thread_id)):
+            return await self.primary.aput_writes(*args, **kwargs)

    async def adelete_thread(self, thread_id: str) -> None:
-        delete = getattr(self.primary, "adelete_thread", None)
-        if delete is not None:
-            await delete(thread_id)
-        for source in self.sources:
-            if source.closed:
-                continue
-            delete = getattr(source.saver, "adelete_thread", None)
+        async with self.thread_write_lock(str(thread_id), exclusive=True):
+            delete = getattr(self.primary, "adelete_thread", None)
            if delete is not None:
                await delete(thread_id)
+            for source in self.sources:
+                if source.closed:
+                    continue
+                delete = getattr(source.saver, "adelete_thread", None)
+                if delete is not None:
+                    await delete(thread_id)
+            archive = await self._archive_source_for_thread(str(thread_id))
+            if archive is not None:
+                delete = getattr(archive.saver, "adelete_thread", None)
+                if delete is not None:
+                    await delete(thread_id)
+            if self.archive_dir and catalog_path(self.archive_dir).is_file():
+                await asyncio.to_thread(remove_archived_thread, self.archive_dir, str(thread_id))

    @staticmethod
    def _tuple_key(tup) -> tuple[str, str, str]: