[argus] checkpointer: add AsyncPostgresSaver.aprune

The langgraph postgres checkpointer ships adelete_thread but no aprune.
langgraph_api warns about unbounded checkpoint growth when no prune
method exists, and operators have no in-process way to keep history
bounded short of dropping the whole thread.

This adds an aprune method on import:

  - keep_latest (default) — for each (thread_id, checkpoint_ns) keep
    only the most recent checkpoint_id (ULID, lexicographically
    monotonic) and cascade to checkpoint_writes. checkpoint_blobs are
    left alone because blobs are content-addressed and reused across
    checkpoints; per-checkpoint blob cleanup risks dropping rows the
    surviving checkpoint still references, mirroring how the upstream
    adelete_thread is the only path that touches blobs.
  - delete_all — bulk delete across all three tables.

Implemented as a small monkey-patch module rather than a subclass so
deer-flow code that builds AsyncPostgresSaver via from_conn_string
gets the new method automatically. Idempotent: no-op if the library
ever ships a native aprune (the patch then becomes safe to delete).
Activated by an explicit import in the postgres branch of
async_provider.py — no .pth tricks, no hidden side effects.

Tests skip when langgraph-checkpoint-postgres is not installed (the
package is not in upstream deer-flow's default deps). They run in
environments that install it (e.g. our argus overlay image).

PR-candidate: yes
Upstream-issue: none
Reason: Fills an obvious gap in the postgres checkpointer story. PR
        could move the patch into langgraph-checkpoint-postgres itself
        rather than deer-flow, but having it here keeps deer-flow
        useful in the meantime.

Author: Nichol4s

backend/packages/harness/deerflow/agents/checkpointer/_postgres_aprune.py

@@ -0,0 +1,90 @@
+"""Add ``aprune`` to ``langgraph.checkpoint.postgres.aio.AsyncPostgresSaver``.
+
+LangGraph's postgres checkpointer (as of langgraph-checkpoint-postgres 3.0.x)
+ships ``adelete_thread`` but not ``aprune``. ``langgraph_api`` warns about
+unbounded checkpoint growth when no prune method exists, and operators have no
+in-process way to keep history bounded short of dropping the whole thread.
+
+This module installs an ``aprune`` method on the class at import time. Idempotent:
+no-op if the library already provides one (so a future upstream release that adds
+``aprune`` natively wins automatically and this file becomes safe to delete).
+
+Two strategies:
+
+- ``keep_latest`` (default) — for each ``(thread_id, checkpoint_ns)`` keep only
+  the most recent ``checkpoint_id`` (ULID, lexicographically monotonic) and
+  cascade the deletion to ``checkpoint_writes``. ``checkpoint_blobs`` is left
+  alone because blobs are content-addressed and reused across checkpoints; per-
+  checkpoint blob cleanup risks dropping rows the surviving checkpoint still
+  references, mirroring how the upstream ``adelete_thread`` is the only path
+  that touches blobs.
+- ``delete_all`` — bulk delete across all three tables for the given threads.
+
+Activated by importing this module from the checkpointer providers (see
+``provider.py`` and ``async_provider.py``).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+try:
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+except ImportError:  # pragma: no cover — package not installed in this env
+    AsyncPostgresSaver = None  # type: ignore[assignment]
+
+
+_PRUNE_SQL = """
+WITH keep AS (
+    SELECT thread_id, checkpoint_ns, MAX(checkpoint_id) AS keep_id
+    FROM {table}
+    WHERE thread_id = ANY(%s)
+    GROUP BY thread_id, checkpoint_ns
+)
+DELETE FROM {table} t
+USING keep k
+WHERE t.thread_id = k.thread_id
+  AND t.checkpoint_ns = k.checkpoint_ns
+  AND t.checkpoint_id <> k.keep_id
+"""
+
+
+async def _aprune(
+    self: Any,
+    thread_ids: Sequence[str],
+    *,
+    strategy: str = "keep_latest",
+) -> None:
+    if not thread_ids:
+        return
+
+    ids: list[str] = [str(t) for t in thread_ids]
+
+    if strategy == "delete_all":
+        async with self._cursor(pipeline=True) as cur:
+            for table in ("checkpoints", "checkpoint_blobs", "checkpoint_writes"):
+                await cur.execute(
+                    f"DELETE FROM {table} WHERE thread_id = ANY(%s)",
+                    (ids,),
+                )
+        return
+
+    if strategy != "keep_latest":
+        raise ValueError(f"unknown aprune strategy: {strategy!r}")
+
+    async with self._cursor(pipeline=True) as cur:
+        for table in ("checkpoints", "checkpoint_writes"):
+            await cur.execute(_PRUNE_SQL.format(table=table), (ids,))
+
+
+def _install() -> None:
+    """Install ``aprune`` on AsyncPostgresSaver if the library doesn't already provide one."""
+    if AsyncPostgresSaver is None:
+        return
+    if "aprune" in AsyncPostgresSaver.__dict__:
+        return
+    AsyncPostgresSaver.aprune = _aprune  # type: ignore[attr-defined]
+
+
+_install()

backend/packages/harness/deerflow/agents/checkpointer/async_provider.py

@@ -64,6 +64,10 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
     if config.type == "postgres":
         try:
             from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+            # Install AsyncPostgresSaver.aprune until upstream ships one natively.
+            # See _postgres_aprune.py for rationale.
+            from deerflow.agents.checkpointer import _postgres_aprune  # noqa: F401
         except ImportError as exc:
             raise ImportError(POSTGRES_INSTALL) from exc
 

backend/tests/test_postgres_aprune.py

@@ -0,0 +1,144 @@
+"""Tests for AsyncPostgresSaver.aprune install and behavior.
+
+The patch lives in ``deerflow.agents.checkpointer._postgres_aprune`` and is
+activated by importing it from ``async_provider.py``. These tests verify:
+
+  - ``_install()`` is idempotent.
+  - ``_install()`` does not override a native ``aprune`` if upstream ever
+    ships one (then the patch becomes a no-op and is safe to delete).
+  - ``aprune`` issues the expected SQL for both strategies.
+  - Unknown strategy raises ValueError.
+  - Empty thread list is a no-op.
+"""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+
+import pytest
+
+# Skip the whole module when langgraph-checkpoint-postgres is not installed
+# (e.g. in upstream's default test environment, which only depends on the
+# sqlite + memory checkpointers). Our Argus fork installs the postgres
+# checkpointer at image build time, so the test runs in our CI.
+pytest.importorskip("langgraph.checkpoint.postgres.aio")
+
+from deerflow.agents.checkpointer import _postgres_aprune  # noqa: E402
+
+
+# ---------------------------------------------------------------------------
+# Install logic
+# ---------------------------------------------------------------------------
+
+
+class TestInstall:
+    def test_install_is_idempotent(self):
+        """Calling _install twice should not double-wrap or error."""
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+        before = AsyncPostgresSaver.aprune
+        _postgres_aprune._install()
+        _postgres_aprune._install()
+        after = AsyncPostgresSaver.aprune
+        assert before is after
+
+    def test_install_skips_when_native_method_exists(self, monkeypatch):
+        """If the library ever ships a native aprune, ours must not overwrite it."""
+        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+        sentinel = object()
+        # Place a fake native method directly in the class __dict__
+        original = AsyncPostgresSaver.__dict__.get("aprune")
+        try:
+            AsyncPostgresSaver.aprune = sentinel  # type: ignore[assignment]
+            _postgres_aprune._install()
+            assert AsyncPostgresSaver.aprune is sentinel
+        finally:
+            if original is None:
+                # Restore our installed method
+                _postgres_aprune._install()
+            else:
+                AsyncPostgresSaver.aprune = original  # type: ignore[assignment]
+
+
+# ---------------------------------------------------------------------------
+# aprune behavior — exercised against a fake saver with a mocked cursor
+# ---------------------------------------------------------------------------
+
+
+class _FakeCursor:
+    """Captures executed SQL + parameters for assertions."""
+
+    def __init__(self):
+        self.calls: list[tuple[str, tuple]] = []
+
+    async def execute(self, sql: str, params: tuple) -> None:
+        self.calls.append((sql, params))
+
+
+class _FakeSaver:
+    """Just enough surface to satisfy ``aprune``: a ``_cursor`` async ctx manager."""
+
+    def __init__(self):
+        self.cursor = _FakeCursor()
+
+    @asynccontextmanager
+    async def _cursor(self, pipeline: bool = False):
+        yield self.cursor
+
+
+@pytest.mark.anyio
+class TestApruneBehavior:
+    async def test_empty_thread_list_is_noop(self):
+        saver = _FakeSaver()
+        await _postgres_aprune._aprune(saver, [])
+        assert saver.cursor.calls == []
+
+    async def test_keep_latest_issues_two_deletes(self):
+        """keep_latest must issue DELETEs for checkpoints and checkpoint_writes,
+        but NOT for checkpoint_blobs (blobs are content-addressed and shared)."""
+        saver = _FakeSaver()
+        await _postgres_aprune._aprune(saver, ["t1", "t2"], strategy="keep_latest")
+
+        assert len(saver.cursor.calls) == 2
+        tables_touched = [
+            "checkpoints" if "FROM checkpoints" in sql else "checkpoint_writes"
+            for sql, _ in saver.cursor.calls
+        ]
+        assert set(tables_touched) == {"checkpoints", "checkpoint_writes"}
+        # Blobs must not be touched
+        for sql, _ in saver.cursor.calls:
+            assert "checkpoint_blobs" not in sql
+
+        # Each DELETE must scope to the provided thread ids
+        for _, params in saver.cursor.calls:
+            assert params == (["t1", "t2"],)
+
+    async def test_delete_all_touches_all_three_tables(self):
+        saver = _FakeSaver()
+        await _postgres_aprune._aprune(saver, ["t1"], strategy="delete_all")
+
+        assert len(saver.cursor.calls) == 3
+        tables = {
+            "checkpoints" if "FROM checkpoints " in sql else
+            "checkpoint_blobs" if "FROM checkpoint_blobs " in sql else
+            "checkpoint_writes"
+            for sql, _ in saver.cursor.calls
+        }
+        assert tables == {"checkpoints", "checkpoint_blobs", "checkpoint_writes"}
+
+    async def test_unknown_strategy_raises(self):
+        saver = _FakeSaver()
+        with pytest.raises(ValueError, match="unknown aprune strategy"):
+            await _postgres_aprune._aprune(saver, ["t1"], strategy="bogus")
+
+    async def test_thread_ids_are_stringified(self):
+        """UUID objects (or anything else) should be coerced to str before the query."""
+        import uuid
+
+        saver = _FakeSaver()
+        tid = uuid.uuid4()
+        await _postgres_aprune._aprune(saver, [tid], strategy="keep_latest")
+
+        for _, params in saver.cursor.calls:
+            assert params == ([str(tid)],)