mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-13 23:48:53 +00:00
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>
This commit is contained in:
Zhengcy05
GitHub
parent
1b76ab9060
commit
81f2015fe6
@ -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
|
||||
|
||||
174
backend/tests/blocking_io/test_agent_factory_construction.py
Normal file
174
backend/tests/blocking_io/test_agent_factory_construction.py
Normal file
@ -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
|
||||
@ -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",
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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, {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user