From 3278f46585bbd70f1e86abb40469a28f1b2ad28d Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 2 Jul 2026 17:54:31 +0200 Subject: [PATCH] fix(tasks): dispose checkpointer pool per celery task --- .../builtin/agent_task/dependencies.py | 13 +--- .../app/tasks/celery_tasks/__init__.py | 17 +++++ .../unit/tasks/test_celery_async_runner.py | 75 +++++++++++++++++++ 3 files changed, 95 insertions(+), 10 deletions(-) diff --git a/surfsense_backend/app/automations/actions/builtin/agent_task/dependencies.py b/surfsense_backend/app/automations/actions/builtin/agent_task/dependencies.py index 7d8ee6fa1..2dca821d4 100644 --- a/surfsense_backend/app/automations/actions/builtin/agent_task/dependencies.py +++ b/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, diff --git a/surfsense_backend/app/tasks/celery_tasks/__init__.py b/surfsense_backend/app/tasks/celery_tasks/__init__.py index a1113884f..f2401d44e 100644 --- a/surfsense_backend/app/tasks/celery_tasks/__init__.py +++ b/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): diff --git a/surfsense_backend/tests/unit/tasks/test_celery_async_runner.py b/surfsense_backend/tests/unit/tasks/test_celery_async_runner.py index 2342dd8da..761c12cd8 100644 --- a/surfsense_backend/tests/unit/tasks/test_celery_async_runner.py +++ b/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