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
-    # pool binds connections to the loop that opened them, but Celery uses a
-    # fresh loop per task, so the next task hangs 30s on a dead-loop connection
-    # (`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.
+    # One-shot runs on a fresh thread don't outlive a single execution, so an
+    # in-memory checkpointer suffices. Ongoing chat-thread turns that need
+    # durable memory use the shared Postgres checkpointer via stream_new_chat.
     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):