fix(runtime): keep agent construction off event loop (#5217)

* fix(runtime): keep agent construction off event loop

* fix:
- offload checkpoint state accessor graph construction to a worker thread
- update test

* import AsyncKeyedLockTable

* update Agents.md

* fix: update test

* fix: preserve single-flight builds after cancellation

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

backend/app/gateway/services.py

@@ -1067,7 +1067,13 @@ def build_checkpoint_state_accessor(
 
     agent_factory = resolve_agent_factory(assistant_id)
     try:
-        graph = _state_accessor_graph(agent_factory, assistant_id, ctx.checkpoint_channel_mode, getattr(ctx, "checkpoint_snapshot_frequency", None), config)
+        graph = _state_accessor_graph(
+            agent_factory,
+            assistant_id,
+            ctx.checkpoint_channel_mode,
+            getattr(ctx, "checkpoint_snapshot_frequency", None),
+            config,
+        )
     except Exception:
         if ctx.checkpoint_channel_mode != "full":
             # Delta materialization needs the graph's channel table; there is

backend/tests/blocking_io/test_agent_factory_construction.py

@@ -0,0 +1,174 @@
+"""Regression coverage for off-loop Gateway agent construction (#5172)."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from app.gateway import services
+from deerflow.runtime.runs.manager import RunManager
+from deerflow.runtime.runs.worker import RunContext, run_agent
+
+
+class _Agent:
+    async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
+        yield {"messages": []}
+
+
+def _bridge() -> SimpleNamespace:
+    return SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
+
+
+pytestmark = pytest.mark.asyncio
+
+
+async def _assert_factory_runs_off_the_event_loop(invoke) -> None:
+    """The release waits for an event-loop heartbeat while assembly is blocked."""
+    factory_started = threading.Event()
+    release_factory = threading.Event()
+    heartbeat_during_factory = threading.Event()
+    factory_thread_ids: list[int] = []
+    stop_heartbeat = asyncio.Event()
+
+    def agent_factory(*, config):
+        factory_thread_ids.append(threading.get_ident())
+        factory_started.set()
+        release_factory.wait(timeout=1)
+        return _Agent()
+
+    def release_after_factory_starts() -> None:
+        factory_started.wait(timeout=1)
+        heartbeat_during_factory.wait(timeout=1)
+        release_factory.set()
+
+    async def ticker() -> None:
+        while not stop_heartbeat.is_set():
+            if factory_started.is_set() and not release_factory.is_set():
+                heartbeat_during_factory.set()
+            await asyncio.sleep(0.01)
+
+    releaser = threading.Thread(target=release_after_factory_starts, daemon=True)
+    releaser.start()
+    ticker_task = asyncio.create_task(ticker())
+    try:
+        await invoke(agent_factory)
+    finally:
+        stop_heartbeat.set()
+        await ticker_task
+        await asyncio.to_thread(releaser.join, 1)
+
+    assert len(factory_thread_ids) == 1
+    assert factory_thread_ids[0] != threading.get_ident()
+    assert heartbeat_during_factory.is_set()
+
+
+async def test_gateway_agent_factory_runs_off_the_event_loop() -> None:
+    """Run execution keeps synchronous MCP/tool assembly off Gateway's loop."""
+    run_manager = RunManager()
+    record = await run_manager.create("thread-agent-construction")
+
+    async def invoke(agent_factory) -> None:
+        await run_agent(
+            _bridge(),
+            run_manager,
+            record,
+            ctx=RunContext(checkpointer=None),
+            agent_factory=agent_factory,
+            graph_input={},
+            config={},
+        )
+
+    await _assert_factory_runs_off_the_event_loop(invoke)
+
+
+async def test_gateway_checkpoint_state_factory_runs_off_the_event_loop() -> None:
+    """State/history reads must not rebuild MCP tools on Gateway's loop."""
+    request = SimpleNamespace(state=SimpleNamespace(checkpoint_channel_mode="full"))
+    ctx = SimpleNamespace(checkpointer=object(), store=None, checkpoint_channel_mode="full", app_config=None)
+
+    async def invoke(agent_factory) -> None:
+        with (
+            patch.object(services, "get_run_context", return_value=ctx),
+            patch.object(services, "resolve_agent_factory", return_value=agent_factory),
+        ):
+            await services.abuild_checkpoint_state_accessor(request, thread_id="thread-checkpoint-state")
+
+    try:
+        await _assert_factory_runs_off_the_event_loop(invoke)
+    finally:
+        services._state_accessor_graph_cache.clear()
+
+
+async def test_gateway_checkpoint_state_factory_is_single_flight() -> None:
+    """Concurrent cold-cache reads build one graph without occupying extra workers."""
+    request = SimpleNamespace(state=SimpleNamespace(checkpoint_channel_mode="full"))
+    ctx = SimpleNamespace(checkpointer=object(), store=None, checkpoint_channel_mode="full", app_config=None)
+    factory_started = threading.Event()
+    release_factory = threading.Event()
+    factory_calls: list[int] = []
+
+    def agent_factory(*, config):
+        factory_calls.append(threading.get_ident())
+        factory_started.set()
+        release_factory.wait(timeout=1)
+        return _Agent()
+
+    with (
+        patch.object(services, "get_run_context", return_value=ctx),
+        patch.object(services, "resolve_agent_factory", return_value=agent_factory),
+    ):
+        try:
+            first = asyncio.create_task(services.abuild_checkpoint_state_accessor(request, thread_id="thread-single-flight"))
+            assert await asyncio.to_thread(factory_started.wait, 1)
+            second = asyncio.create_task(services.abuild_checkpoint_state_accessor(request, thread_id="thread-single-flight"))
+            await asyncio.sleep(0)
+            release_factory.set()
+            first_accessor, second_accessor = await asyncio.gather(first, second)
+        finally:
+            release_factory.set()
+            services._state_accessor_graph_cache.clear()
+
+    assert len(factory_calls) == 1
+    assert first_accessor[0].graph is second_accessor[0].graph
+
+
+async def test_gateway_checkpoint_state_factory_survives_waiter_cancellation() -> None:
+    """Cancelling one reader cannot make a same-key reader rebuild the graph."""
+    request = SimpleNamespace(state=SimpleNamespace(checkpoint_channel_mode="full"))
+    ctx = SimpleNamespace(checkpointer=object(), store=None, checkpoint_channel_mode="full", app_config=None)
+    factory_started = threading.Event()
+    release_factory = threading.Event()
+    factory_calls: list[int] = []
+
+    def agent_factory(*, config):
+        factory_calls.append(threading.get_ident())
+        factory_started.set()
+        release_factory.wait(timeout=1)
+        return _Agent()
+
+    with (
+        patch.object(services, "get_run_context", return_value=ctx),
+        patch.object(services, "resolve_agent_factory", return_value=agent_factory),
+    ):
+        try:
+            first = asyncio.create_task(services.abuild_checkpoint_state_accessor(request, thread_id="thread-cancelled-single-flight"))
+            assert await asyncio.to_thread(factory_started.wait, 1)
+            first.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await first
+
+            second = asyncio.create_task(services.abuild_checkpoint_state_accessor(request, thread_id="thread-cancelled-single-flight"))
+            await asyncio.sleep(0)
+            assert len(factory_calls) == 1
+            release_factory.set()
+            second_accessor = await second
+        finally:
+            release_factory.set()
+            services._state_accessor_graph_cache.clear()
+
+    assert len(factory_calls) == 1
+    assert second_accessor[0].graph is not None

backend/tests/test_bench_checkpoint_production.py

@@ -206,14 +206,18 @@ def test_run_case_fails_when_warm_cache_is_not_hit(tmp_path, monkeypatch) -> Non
     from deerflow.config.app_config import reset_app_config
 
     gateway_services._state_accessor_graph_cache.clear()
+
     # Bypass the cache entirely: every accessor resolution rebuilds the graph,
     # so warm reads trip the contract assertion. The factory returns a
     # LeadAgentAssembly, so this stub unwraps it exactly as the real accessor
     # does — the point here is the missing cache, not a different return shape.
+    def uncached_state_accessor_graph(agent_factory, assistant_id, mode, snapshot_frequency, config):
+        return agent_factory(config=config).graph
+
     monkeypatch.setattr(
         gateway_services,
         "_state_accessor_graph",
-        lambda agent_factory, assistant_id, mode, snapshot_frequency, config: agent_factory(config=config).graph,
+        uncached_state_accessor_graph,
     )
     case = bench.ProductionCase(
         mode="full",

backend/tests/test_gateway_services.py

@@ -1202,7 +1202,7 @@ def test_apply_checkpoint_to_run_config_writes_checkpoint_fields():
 
 @pytest.mark.anyio
 async def test_seeded_checkpoint_messages_precede_the_first_new_run_messages():
-    from unittest.mock import AsyncMock, patch
+    from unittest.mock import AsyncMock, MagicMock, patch
 
     from langchain_core.messages import AIMessage, HumanMessage
 
@@ -1232,7 +1232,7 @@ async def test_seeded_checkpoint_messages_precede_the_first_new_run_messages():
 
     with patch(
         "app.gateway.services.build_checkpoint_state_accessor",
-        return_value=(accessor, {"configurable": {"thread_id": "thread-1"}}),
+        new=MagicMock(return_value=(accessor, {"configurable": {"thread_id": "thread-1"}})),
     ):
         await ensure_checkpoint_history_seeded(
             request,
@@ -1272,7 +1272,7 @@ async def test_seeded_checkpoint_messages_precede_the_first_new_run_messages():
 
 @pytest.mark.anyio
 async def test_checkpoint_history_seed_skips_new_thread_without_checkpoint():
-    from unittest.mock import AsyncMock, patch
+    from unittest.mock import AsyncMock, MagicMock, patch
 
     from app.gateway.services import ensure_checkpoint_history_seeded
 
@@ -1292,7 +1292,7 @@ async def test_checkpoint_history_seed_skips_new_thread_without_checkpoint():
 
     with patch(
         "app.gateway.services.build_checkpoint_state_accessor",
-        side_effect=AssertionError("new threads should not build an accessor"),
+        new=MagicMock(side_effect=AssertionError("new threads should not build an accessor")),
     ):
         await ensure_checkpoint_history_seeded(
             request,
@@ -1305,7 +1305,7 @@ async def test_checkpoint_history_seed_skips_new_thread_without_checkpoint():
 
 @pytest.mark.anyio
 async def test_checkpoint_history_seed_is_skipped_when_journal_already_has_messages():
-    from unittest.mock import AsyncMock, patch
+    from unittest.mock import AsyncMock, MagicMock, patch
 
     from app.gateway.services import ensure_checkpoint_history_seeded
 
@@ -1317,7 +1317,7 @@ async def test_checkpoint_history_seed_is_skipped_when_journal_already_has_messa
 
     with patch(
         "app.gateway.services.build_checkpoint_state_accessor",
-        side_effect=AssertionError("checkpoint state should not be loaded"),
+        new=MagicMock(side_effect=AssertionError("checkpoint state should not be loaded")),
     ):
         await ensure_checkpoint_history_seeded(
             request,
@@ -1376,7 +1376,7 @@ async def test_checkpoint_history_seed_guard_is_thread_scoped_under_user_context
     even when a user is authenticated. Seed rows stamped by another principal
     (or NULL) are invisible to a user-scoped query, which would re-seed a
     duplicate history per principal."""
-    from unittest.mock import AsyncMock, patch
+    from unittest.mock import AsyncMock, MagicMock, patch
 
     from app.gateway.services import ensure_checkpoint_history_seeded
     from deerflow.runtime.user_context import AUTO
@@ -1392,7 +1392,7 @@ async def test_checkpoint_history_seed_guard_is_thread_scoped_under_user_context
 
     with patch(
         "app.gateway.services.build_checkpoint_state_accessor",
-        side_effect=AssertionError("checkpoint state should not be loaded"),
+        new=MagicMock(side_effect=AssertionError("checkpoint state should not be loaded")),
     ):
         await ensure_checkpoint_history_seeded(
             request,
@@ -1411,7 +1411,7 @@ async def test_checkpoint_history_seed_runs_exactly_once_across_principals(tmp_p
     user_id=NULL; a later authenticated run on the same thread must still
     see them and skip re-seeding (the MemoryRunEventStore-based tests above
     cannot catch this because the memory store ignores user_id)."""
-    from unittest.mock import AsyncMock, patch
+    from unittest.mock import AsyncMock, MagicMock, patch
 
     from langchain_core.messages import AIMessage, HumanMessage
 
@@ -1447,7 +1447,7 @@ async def test_checkpoint_history_seed_runs_exactly_once_across_principals(tmp_p
 
         with patch(
             "app.gateway.services.build_checkpoint_state_accessor",
-            return_value=(accessor, {"configurable": {"thread_id": "thread-1"}}),
+            new=MagicMock(return_value=(accessor, {"configurable": {"thread_id": "thread-1"}})),
         ):
             # First seed: ownerless (no user contextvar) — rows stamped NULL.
             await ensure_checkpoint_history_seeded(

backend/tests/test_thread_regenerate_prepare.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 import copy
 from types import SimpleNamespace
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from _router_auth_helpers import call_unwrapped
@@ -242,7 +242,7 @@ def test_run_wait_readers_return_materialized_final_values() -> None:
                 services,
                 "build_checkpoint_state_accessor",
                 create=True,
-                return_value=(accessor, snapshot.config),
+                new=MagicMock(return_value=(accessor, snapshot.config)),
             ),
             patch.object(runs, "get_stream_bridge", return_value=object()),
             patch.object(runs, "get_run_manager", return_value=object()),
@@ -251,7 +251,7 @@ def test_run_wait_readers_return_materialized_final_values() -> None:
                 services,
                 "build_checkpoint_state_accessor",
                 create=True,
-                return_value=(accessor, snapshot.config),
+                new=MagicMock(return_value=(accessor, snapshot.config)),
             ),
         ):
             thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
@@ -296,7 +296,7 @@ def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None:
             patch.object(
                 services,
                 "build_checkpoint_state_accessor",
-                return_value=(accessor, snapshot.config),
+                new=MagicMock(return_value=(accessor, snapshot.config)),
             ),
             patch.object(runs, "get_stream_bridge", return_value=object()),
             patch.object(runs, "get_run_manager", return_value=object()),
@@ -304,7 +304,7 @@ def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None:
             patch.object(
                 services,
                 "build_checkpoint_state_accessor",
-                return_value=(accessor, snapshot.config),
+                new=MagicMock(return_value=(accessor, snapshot.config)),
             ),
         ):
             thread_result = await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
@@ -342,7 +342,7 @@ def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(ro
                 patch.object(
                     services,
                     "build_checkpoint_state_accessor",
-                    side_effect=RuntimeError("graph construction failed"),
+                    new=MagicMock(side_effect=RuntimeError("graph construction failed")),
                 ),
             ):
                 return await call_unwrapped(thread_runs.wait_run, "thread-1", body, request)
@@ -354,7 +354,7 @@ def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(ro
             patch.object(
                 services,
                 "build_checkpoint_state_accessor",
-                side_effect=RuntimeError("graph construction failed"),
+                new=MagicMock(side_effect=RuntimeError("graph construction failed")),
             ),
         ):
             return await call_unwrapped(runs.stateless_wait, body, request)

backend/tests/test_threads_router.py

@@ -2730,7 +2730,7 @@ def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monke
     source_accessor.aget = source_aget
     branch_accessor = SimpleNamespace(aupdate=branch_aupdate)
 
-    def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
+    async def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
         assert thread_id == source_thread_id
         return source_accessor, {
             "configurable": {
@@ -2849,7 +2849,7 @@ def test_branch_thread_preserves_unlinked_legacy_histories(
     source_accessor = SimpleNamespace(ahistory=source_ahistory, aget=unexpected_lineage_read)
     branch_accessor = SimpleNamespace(aupdate=branch_aupdate)
 
-    def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
+    async def build_accessor(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
         assert thread_id == source_thread_id
         return source_accessor, {
             "configurable": {
@@ -2984,7 +2984,7 @@ def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypat
         aget=AsyncMock(side_effect=lambda config: next(item for item in source_history if item.config["configurable"]["checkpoint_id"] == config["configurable"]["checkpoint_id"])),
     )
 
-    def source_builder(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
+    async def source_builder(_request, *, thread_id, assistant_id=None, checkpoint_id=None):
         if thread_id != source_thread_id:
             raise AssertionError("fresh branches must use the dedicated mutation graph")
         return source_accessor, {