fix(tasks): dispose checkpointer pool per celery task

surfsense_backend/app/automations/actions/builtin/agent_task/dependencies.py

@@ -84,16 +84,9 @@ async def build_dependencies(
     connector_service = await setup_connector_service(
         session, workspace_id=workspace_id
     )
-    # Per-task InMemorySaver: the shared Postgres checkpointer's connection
+    # One-shot runs on a fresh thread don't outlive a single execution, so an
-    # pool binds connections to the loop that opened them, but Celery uses a
+    # in-memory checkpointer suffices. Ongoing chat-thread turns that need
-    # fresh loop per task, so the next task hangs 30s on a dead-loop connection
+    # durable memory use the shared Postgres checkpointer via stream_new_chat.
-    # (`PoolTimeout`). InMemorySaver has no pool and dies with the task — fine
-    # while runs are one-shot (the checkpoint only spans one graph execution).
-    #
-    # TODO(checkpointer): when runs need durability (crash-resume or HITL
-    # interrupt/resume across tasks), dispose the checkpointer pool around each
-    # Celery task in `run_async_celery_task` — as `_dispose_shared_db_engine`
-    # already does for the SQLAlchemy pool — then use the shared checkpointer.
     checkpointer = InMemorySaver()
     return AgentDependencies(
         llm=llm,

surfsense_backend/app/tasks/celery_tasks/__init__.py

@@ -94,6 +94,21 @@ def _dispose_shared_db_engine(loop: asyncio.AbstractEventLoop) -> None:
         logger.warning("Shared DB engine dispose() failed", exc_info=True)
 
 
+def _dispose_shared_checkpointer_pool(loop: asyncio.AbstractEventLoop) -> None:
+    """Drop the shared checkpointer pool so the next task opens a fresh one.
+
+    The durable ``AsyncPostgresSaver`` pool binds asyncpg connections to the
+    loop that opened them; a later task on a fresh loop stalls on a stale
+    connection (``PoolTimeout``). Failure is logged, not raised.
+    """
+    try:
+        from app.agents.chat.runtime.checkpointer import close_checkpointer
+
+        loop.run_until_complete(close_checkpointer())
+    except Exception:
+        logger.warning("Shared checkpointer pool dispose() failed", exc_info=True)
+
+
 T = TypeVar("T")
 
 
@@ -129,11 +144,13 @@ def run_async_celery_task[T](coro_factory: Callable[[], Awaitable[T]]) -> T:
         # Defense-in-depth: prior task may have crashed before
         # disposing. Idempotent — no-op if pool is already empty.
         _dispose_shared_db_engine(loop)
+        _dispose_shared_checkpointer_pool(loop)
         return loop.run_until_complete(coro_factory())
     finally:
         # Drop any connections this task opened so they don't leak
         # into the next task's loop.
         _dispose_shared_db_engine(loop)
+        _dispose_shared_checkpointer_pool(loop)
         with contextlib.suppress(Exception):
             loop.run_until_complete(loop.shutdown_asyncgens())
         with contextlib.suppress(Exception):

surfsense_backend/tests/unit/tasks/test_celery_async_runner.py

@@ -293,6 +293,81 @@ def test_runner_runs_shutdown_asyncgens_before_close() -> None:
     gc.collect()
 
 
+@contextmanager
+def _patch_checkpointer_close(recorder: list[int]) -> Iterator[None]:
+    """Patch ``close_checkpointer`` to record the loop it runs on.
+
+    The runner imports it lazily, so we patch the attribute on the
+    already-loaded checkpointer module (same trick as the engine stub).
+    """
+    import app.agents.chat.runtime.checkpointer as cp
+
+    original = cp.close_checkpointer
+
+    async def _stub() -> None:
+        recorder.append(id(asyncio.get_running_loop()))
+
+    cp.close_checkpointer = _stub  # type: ignore[assignment]
+    try:
+        yield
+    finally:
+        cp.close_checkpointer = original  # type: ignore[assignment]
+
+
+def test_runner_disposes_checkpointer_pool_around_call() -> None:
+    """The durable checkpointer's loop-bound pool must be dropped before
+    and after each task, on the task's own fresh loop — the fix that lets
+    a worker use the shared AsyncPostgresSaver instead of InMemorySaver.
+    """
+    from app.tasks.celery_tasks import run_async_celery_task
+
+    engine_stub = _StaleLoopEngine()
+    cp_loops: list[int] = []
+
+    async def _body() -> str:
+        # Both the engine and the checkpointer are disposed once before we run.
+        assert len(cp_loops) == 1
+        return "ok"
+
+    with _patch_shared_engine(engine_stub), _patch_checkpointer_close(cp_loops):
+        assert run_async_celery_task(_body) == "ok"
+
+    # Once before the body, once after (finally).
+    assert len(cp_loops) == 2
+    # Both ran on the SAME fresh loop, matching the engine's dispose loop.
+    assert cp_loops[0] == cp_loops[1]
+    assert cp_loops[0] == engine_stub.dispose_loop_ids[0]
+
+
+def test_runner_swallows_checkpointer_dispose_errors() -> None:
+    """A flaky checkpointer dispose must never take down a celery task."""
+    from app.tasks.celery_tasks import run_async_celery_task
+
+    engine_stub = _StaleLoopEngine()
+
+    import app.agents.chat.runtime.checkpointer as cp
+
+    original = cp.close_checkpointer
+    calls = {"n": 0}
+
+    async def _angry() -> None:
+        calls["n"] += 1
+        raise RuntimeError("checkpointer dispose blew up")
+
+    cp.close_checkpointer = _angry  # type: ignore[assignment]
+
+    async def _body() -> int:
+        return 42
+
+    try:
+        with _patch_shared_engine(engine_stub):
+            assert run_async_celery_task(_body) == 42
+    finally:
+        cp.close_checkpointer = original  # type: ignore[assignment]
+
+    assert calls["n"] == 2  # before + after both attempted, both swallowed
+
+
 def test_runner_uses_proactor_loop_on_windows() -> None:
     """On Windows the celery worker preselects a Proactor policy so
     subprocess (ffmpeg) calls work. The helper must not silently fall