459 lines
17 KiB
Python
459 lines
17 KiB
Python
"""Slow agent initialization must not monopolize the gateway event loop."""
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextvars import ContextVar
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.runtime.runs import worker
|
|
from deerflow.runtime.runs.manager import RunManager
|
|
from deerflow.runtime.runs.schemas import RunStatus
|
|
from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge
|
|
from deerflow.tools.builtins.tool_search import DeferredToolRegistry, get_deferred_registry, reset_deferred_registry, set_deferred_registry
|
|
from deerflow.utils import background_init
|
|
from deerflow.utils.file_io import FileIOPool
|
|
|
|
|
|
def _background(factory):
|
|
factory._deerflow_background_build = True
|
|
return factory
|
|
|
|
|
|
@pytest.fixture
|
|
def build_pool(monkeypatch):
|
|
pool = FileIOPool(1, 1, "test-agent-init")
|
|
monkeypatch.setattr(background_init, "_INITIALIZATION_IO", pool)
|
|
yield pool
|
|
pool._pool.shutdown(wait=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shared_initializer_forwards_arguments_and_commits_context(build_pool):
|
|
from deerflow.mcp.cache import _owner_initialization_loop
|
|
|
|
loop = asyncio.get_running_loop()
|
|
caller_thread = threading.get_ident()
|
|
state = ContextVar("shared-initializer-context", default="before")
|
|
|
|
def initialize(value, *, suffix):
|
|
assert threading.get_ident() != caller_thread
|
|
assert _owner_initialization_loop.get() is loop
|
|
assert state.get() == "before"
|
|
state.set(None)
|
|
return value + suffix
|
|
|
|
assert await background_init.run_background_init(initialize, "tools", suffix=" ready") == "tools ready"
|
|
assert state.get() is None
|
|
assert _owner_initialization_loop.get() is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_shared_initializer_does_not_commit_context(build_pool):
|
|
state = ContextVar("failed-shared-initializer-context", default="before")
|
|
|
|
def fail():
|
|
state.set("partially initialized")
|
|
raise ValueError("initialization failed")
|
|
|
|
with pytest.raises(ValueError, match="initialization failed"):
|
|
await background_init.run_background_init(fail)
|
|
assert state.get() == "before"
|
|
assert await background_init.run_background_init(lambda: state.get()) == "before"
|
|
|
|
|
|
class FinishedAgent:
|
|
async def astream(self, *args, **kwargs):
|
|
yield {"messages": []}
|
|
|
|
|
|
async def _make_run(monkeypatch, factory, *, config=None, app_config=None, thread_id="build-responsiveness"):
|
|
manager = RunManager()
|
|
record = await manager.create(thread_id)
|
|
bridge = MemoryStreamBridge()
|
|
monkeypatch.setattr(bridge, "cleanup", AsyncMock())
|
|
task = asyncio.create_task(
|
|
worker.run_agent(
|
|
bridge,
|
|
manager,
|
|
record,
|
|
ctx=worker.RunContext(checkpointer=None, app_config=app_config),
|
|
agent_factory=factory,
|
|
graph_input={},
|
|
config=config or {},
|
|
)
|
|
)
|
|
record.task = task
|
|
return SimpleNamespace(task=task, manager=manager, record=record, bridge=bridge)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_slow_agent_build_keeps_parallel_requests_responsive(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
entered = asyncio.Event()
|
|
factory_finished = threading.Event()
|
|
factory_started_at = []
|
|
|
|
@_background
|
|
def factory(*, config):
|
|
factory_started_at.append(time.perf_counter())
|
|
loop.call_soon_threadsafe(entered.set)
|
|
deadline = time.perf_counter() + 0.18
|
|
while time.perf_counter() < deadline:
|
|
pass
|
|
factory_finished.set()
|
|
return FinishedAgent()
|
|
|
|
run = await _make_run(monkeypatch, factory)
|
|
await asyncio.wait_for(entered.wait(), 3)
|
|
probe_delay = time.perf_counter() - factory_started_at[0]
|
|
responsive = not factory_finished.is_set()
|
|
await asyncio.wait_for(run.task, 3)
|
|
|
|
assert run.record.status == RunStatus.success
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
print(f"agent build concurrent probe delay = {probe_delay:.3f}s")
|
|
assert responsive, f"Agent initialization blocked parallel requests for {probe_delay:.3f}s"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_factory_keeps_caller_loop_and_context_mutations(monkeypatch):
|
|
"""Third-party synchronous factories have no worker-safety declaration."""
|
|
loop = asyncio.get_running_loop()
|
|
caller_thread = threading.get_ident()
|
|
context = ContextVar("custom-build-context", default="before")
|
|
observed = []
|
|
app_config = object()
|
|
|
|
class Agent:
|
|
async def astream(self, *args, **kwargs):
|
|
observed.append(context.get())
|
|
assert asyncio.get_running_loop() is loop
|
|
yield {"messages": []}
|
|
|
|
def factory(*, config, app_config):
|
|
assert asyncio.get_running_loop() is loop
|
|
assert threading.get_ident() == caller_thread
|
|
assert app_config is expected_app_config
|
|
assert config["context"]["app_config"] is app_config
|
|
context.set("set by custom factory")
|
|
return Agent()
|
|
|
|
expected_app_config = app_config
|
|
run = await _make_run(monkeypatch, factory, app_config=app_config)
|
|
await asyncio.wait_for(run.task, 3)
|
|
|
|
assert run.record.status == RunStatus.success
|
|
assert observed == ["set by custom factory"]
|
|
assert context.get() == "before", "Run-local state must not leak to its caller"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_factory_failure_ends_stream_and_does_not_poison_next_run(monkeypatch, build_pool):
|
|
@_background
|
|
def failing_factory(*, config):
|
|
raise ValueError("synthetic initialization failure")
|
|
|
|
failed = await _make_run(monkeypatch, failing_factory)
|
|
await asyncio.wait_for(failed.task, 3)
|
|
assert failed.record.status == RunStatus.error
|
|
assert failed.record.error == "synthetic initialization failure"
|
|
stream = failed.bridge._streams[failed.record.run_id]
|
|
assert stream.ended
|
|
assert any(event.event == "error" for event in stream.events)
|
|
|
|
recovered = await _make_run(monkeypatch, _background(lambda **kwargs: FinishedAgent()))
|
|
await asyncio.wait_for(recovered.task, 3)
|
|
assert recovered.record.status == RunStatus.success
|
|
assert recovered.bridge._streams[recovered.record.run_id].ended
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("clear_registry", [False, True])
|
|
async def test_background_build_preserves_user_config_and_deferred_registry(monkeypatch, build_pool, clear_registry):
|
|
loop = asyncio.get_running_loop()
|
|
loop_thread = threading.get_ident()
|
|
request_user = ContextVar("build-request-user", default="outside")
|
|
inherited_registry = DeferredToolRegistry()
|
|
previous_registry = get_deferred_registry()
|
|
set_deferred_registry(inherited_registry)
|
|
observations = {}
|
|
app_configs = {name: object() for name in ("alice", "bob")}
|
|
registries = {name: DeferredToolRegistry() for name in app_configs}
|
|
|
|
class Agent:
|
|
def __init__(self, user):
|
|
self.user = user
|
|
|
|
async def astream(self, *args, config, **kwargs):
|
|
assert asyncio.get_running_loop() is loop
|
|
assert threading.get_ident() == loop_thread
|
|
observations[self.user] = (request_user.get(), get_deferred_registry(), config["configurable"]["built_for"])
|
|
yield {"messages": []}
|
|
|
|
@_background
|
|
def factory(*, config, app_config):
|
|
user = request_user.get()
|
|
assert threading.get_ident() != loop_thread
|
|
assert app_config is app_configs[user]
|
|
assert config["context"]["app_config"] is app_config
|
|
assert get_deferred_registry() is inherited_registry
|
|
config["configurable"]["built_for"] = user
|
|
if clear_registry:
|
|
reset_deferred_registry()
|
|
else:
|
|
set_deferred_registry(registries[user])
|
|
return Agent(user)
|
|
|
|
runs = []
|
|
try:
|
|
for user in app_configs:
|
|
token = request_user.set(user)
|
|
try:
|
|
runs.append(await _make_run(monkeypatch, factory, app_config=app_configs[user], thread_id=user))
|
|
finally:
|
|
request_user.reset(token)
|
|
await asyncio.wait_for(asyncio.gather(*(run.task for run in runs)), 3)
|
|
assert all(run.record.status == RunStatus.success for run in runs)
|
|
assert request_user.get() == "outside"
|
|
assert get_deferred_registry() is inherited_registry
|
|
for user in app_configs:
|
|
assert observations[user] == (user, None if clear_registry else registries[user], user)
|
|
finally:
|
|
set_deferred_registry(previous_registry)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_build_queue_serializes_initialization_without_consuming_default_executor(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
first_entered = asyncio.Event()
|
|
release = threading.Event()
|
|
counter_lock = threading.Lock()
|
|
active = peak = entered = 0
|
|
default_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="test-general-worker")
|
|
loop.set_default_executor(default_pool)
|
|
|
|
@_background
|
|
def factory(*, config):
|
|
nonlocal active, peak, entered
|
|
with counter_lock:
|
|
active += 1
|
|
entered += 1
|
|
peak = max(active, peak)
|
|
loop.call_soon_threadsafe(first_entered.set)
|
|
try:
|
|
assert release.wait(5), "Test did not release initialization worker"
|
|
return FinishedAgent()
|
|
finally:
|
|
with counter_lock:
|
|
active -= 1
|
|
|
|
runs = [await _make_run(monkeypatch, factory, thread_id=f"queue-{index}") for index in range(4)]
|
|
try:
|
|
await asyncio.wait_for(first_entered.wait(), 3)
|
|
assert await asyncio.wait_for(asyncio.to_thread(lambda: "authentication worker available"), 1) == "authentication worker available"
|
|
await asyncio.sleep(0.03)
|
|
assert entered == 1, "Cold builders must wait without competing for shared initialization state"
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(asyncio.gather(*(run.task for run in runs)), 3)
|
|
default_pool.shutdown(wait=True)
|
|
|
|
assert peak == 1
|
|
assert entered == 4
|
|
assert all(run.record.status == RunStatus.success for run in runs)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_repeated_cancel_during_build_drains_worker_without_starting_graph(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
entered = asyncio.Event()
|
|
release = threading.Event()
|
|
graph_started = []
|
|
finalization_registries = []
|
|
original_finalize = worker._finalize_run
|
|
previous_registry = get_deferred_registry()
|
|
caller_registry = DeferredToolRegistry()
|
|
set_deferred_registry(caller_registry)
|
|
|
|
async def finalize(**kwargs):
|
|
finalization_registries.append(get_deferred_registry())
|
|
await original_finalize(**kwargs)
|
|
|
|
monkeypatch.setattr(worker, "_finalize_run", finalize)
|
|
|
|
class Agent:
|
|
async def astream(self, *args, **kwargs):
|
|
graph_started.append(True)
|
|
yield {"messages": []}
|
|
|
|
@_background
|
|
def factory(*, config):
|
|
set_deferred_registry(DeferredToolRegistry())
|
|
loop.call_soon_threadsafe(entered.set)
|
|
assert release.wait(5), "Test did not release cancelled builder"
|
|
return Agent()
|
|
|
|
run = await _make_run(monkeypatch, factory)
|
|
try:
|
|
await asyncio.wait_for(entered.wait(), 3)
|
|
assert await run.manager.cancel(run.record.run_id)
|
|
await asyncio.sleep(0)
|
|
assert not run.task.done(), "Cancellation must not abandon a live initialization worker"
|
|
run.task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not run.task.done(), "Repeated cancellation must still drain the builder"
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(run.task, 3)
|
|
set_deferred_registry(previous_registry)
|
|
|
|
assert not graph_started
|
|
assert run.record.status == RunStatus.interrupted
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
assert finalization_registries == [caller_registry], "Cancelled build context must not be committed to the run"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queued_cancel_does_not_build_or_wait_for_active_factory(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
entered = asyncio.Event()
|
|
release = threading.Event()
|
|
queued_built = []
|
|
|
|
@_background
|
|
def first_factory(*, config):
|
|
loop.call_soon_threadsafe(entered.set)
|
|
assert release.wait(5)
|
|
return FinishedAgent()
|
|
|
|
@_background
|
|
def queued_factory(*, config):
|
|
queued_built.append(True)
|
|
return FinishedAgent()
|
|
|
|
first = await _make_run(monkeypatch, first_factory, thread_id="active")
|
|
queued = None
|
|
try:
|
|
await asyncio.wait_for(entered.wait(), 3)
|
|
queued = await _make_run(monkeypatch, queued_factory, thread_id="queued")
|
|
await asyncio.sleep(0.02)
|
|
await queued.manager.cancel(queued.record.run_id)
|
|
await asyncio.wait_for(queued.task, 1)
|
|
assert queued.record.status == RunStatus.interrupted
|
|
assert queued.bridge._streams[queued.record.run_id].ended
|
|
assert not queued_built
|
|
assert not first.task.done()
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(asyncio.gather(first.task, *([queued.task] if queued else [])), 3)
|
|
|
|
assert first.record.status == RunStatus.success
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_abort_flag_set_during_build_prevents_late_stream_start(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
entered = asyncio.Event()
|
|
release = threading.Event()
|
|
graph_started = []
|
|
|
|
class Agent:
|
|
async def astream(self, *args, **kwargs):
|
|
graph_started.append(True)
|
|
yield {"messages": []}
|
|
|
|
@_background
|
|
def factory(*, config):
|
|
loop.call_soon_threadsafe(entered.set)
|
|
assert release.wait(5)
|
|
return Agent()
|
|
|
|
run = await _make_run(monkeypatch, factory)
|
|
try:
|
|
await asyncio.wait_for(entered.wait(), 3)
|
|
run.record.abort_event.set()
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(run.task, 3)
|
|
|
|
assert graph_started == []
|
|
assert run.record.status == RunStatus.interrupted
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_existing_conversation_keeps_streaming_while_another_agent_initializes(monkeypatch, build_pool):
|
|
loop = asyncio.get_running_loop()
|
|
stream_started = asyncio.Event()
|
|
continue_stream = asyncio.Event()
|
|
build_entered = asyncio.Event()
|
|
release_build = threading.Event()
|
|
build_finished = threading.Event()
|
|
|
|
class StreamingAgent:
|
|
async def astream(self, *args, **kwargs):
|
|
stream_started.set()
|
|
yield {"messages": [], "answer": "first delta"}
|
|
await continue_stream.wait()
|
|
yield {"messages": [], "answer": "second delta"}
|
|
|
|
@_background
|
|
def slow_factory(*, config):
|
|
loop.call_soon_threadsafe(build_entered.set)
|
|
assert release_build.wait(5), "Test did not release the second conversation's initializer"
|
|
build_finished.set()
|
|
return FinishedAgent()
|
|
|
|
existing = await _make_run(monkeypatch, lambda **kwargs: StreamingAgent(), thread_id="already-streaming")
|
|
initializing = None
|
|
try:
|
|
await asyncio.wait_for(stream_started.wait(), 3)
|
|
initializing = await _make_run(monkeypatch, slow_factory, thread_id="cold-conversation")
|
|
await asyncio.wait_for(build_entered.wait(), 3)
|
|
continue_stream.set()
|
|
await asyncio.wait_for(existing.task, 1)
|
|
|
|
stream = existing.bridge._streams[existing.record.run_id]
|
|
answers = [event.data["answer"] for event in stream.events if event.event == "values"]
|
|
assert answers == ["first delta", "second delta"]
|
|
assert stream.ended
|
|
assert existing.record.status == RunStatus.success
|
|
assert not initializing.task.done()
|
|
assert not build_finished.is_set(), "The existing answer must finish before the cold build is released"
|
|
finally:
|
|
continue_stream.set()
|
|
release_build.set()
|
|
await asyncio.wait_for(asyncio.gather(existing.task, *([initializing.task] if initializing else [])), 3)
|
|
|
|
assert initializing.record.status == RunStatus.success
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_builtin_lead_factory_uses_background_initializer(monkeypatch, build_pool):
|
|
from deerflow.agents.lead_agent import agent as lead_agent
|
|
|
|
loop_thread = threading.get_ident()
|
|
app_config = object()
|
|
observed_build_threads = []
|
|
|
|
def build_body(config, *, app_config):
|
|
assert app_config is expected_config
|
|
observed_build_threads.append(threading.get_ident())
|
|
return FinishedAgent()
|
|
|
|
expected_config = app_config
|
|
monkeypatch.setattr(lead_agent, "_make_lead_agent", build_body)
|
|
# Exercise the real exported factory: a missing production opt-in must fail
|
|
# even if all tests using synthetic marked factories continue to pass.
|
|
run = await _make_run(monkeypatch, lead_agent.make_lead_agent, app_config=app_config)
|
|
await asyncio.wait_for(run.task, 3)
|
|
|
|
assert run.record.status == RunStatus.success
|
|
assert len(observed_build_threads) == 1
|
|
assert observed_build_threads[0] != loop_thread
|