From 81f2015fe637c9b90b27b6880463553f3532df02 Mon Sep 17 00:00:00 2001 From: Zhengcy05 <1825478405@qq.com> Date: Sat, 12 Sep 2026 21:09:05 +0800 Subject: [PATCH] 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 --- backend/app/gateway/services.py | 8 +- .../test_agent_factory_construction.py | 174 ++++++++++++++++++ .../tests/test_bench_checkpoint_production.py | 6 +- backend/tests/test_gateway_services.py | 20 +- .../tests/test_thread_regenerate_prepare.py | 14 +- backend/tests/test_threads_router.py | 6 +- 6 files changed, 206 insertions(+), 22 deletions(-) create mode 100644 backend/tests/blocking_io/test_agent_factory_construction.py diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index e9fa986ec..ef23df0c1 100644 --- a/backend/app/gateway/services.py +++ b/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 diff --git a/backend/tests/blocking_io/test_agent_factory_construction.py b/backend/tests/blocking_io/test_agent_factory_construction.py new file mode 100644 index 000000000..248ad06f1 --- /dev/null +++ b/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 diff --git a/backend/tests/test_bench_checkpoint_production.py b/backend/tests/test_bench_checkpoint_production.py index 15a67ae04..029838cc1 100644 --- a/backend/tests/test_bench_checkpoint_production.py +++ b/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", diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index 95af63958..fe988a5d8 100644 --- a/backend/tests/test_gateway_services.py +++ b/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( diff --git a/backend/tests/test_thread_regenerate_prepare.py b/backend/tests/test_thread_regenerate_prepare.py index 1f8a324d8..88e9c3f30 100644 --- a/backend/tests/test_thread_regenerate_prepare.py +++ b/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) diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index 206f85c5c..2ac54382d 100644 --- a/backend/tests/test_threads_router.py +++ b/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, {