177 lines
7.3 KiB
Python
177 lines
7.3 KiB
Python
"""Async tool consumers must leave their MCP discovery loop runnable."""
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import threading
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.mcp import cache
|
|
from deerflow.mcp import tools as mcp_tools
|
|
from deerflow.tools import get_available_tools_async
|
|
from deerflow.tools import tools as tool_factory
|
|
from deerflow.tools.builtins.tool_search import (
|
|
DeferredToolRegistry,
|
|
get_deferred_registry,
|
|
set_deferred_registry,
|
|
)
|
|
from deerflow.utils.background_init import run_background_init
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_tools_preserve_options_and_deferred_registry(monkeypatch):
|
|
owner_thread = threading.get_ident()
|
|
user = contextvars.ContextVar("async_tool_test_user")
|
|
user.set("caller")
|
|
registry = DeferredToolRegistry()
|
|
expected = [object()]
|
|
options = {
|
|
"groups": ["files"],
|
|
"include_mcp": False,
|
|
"model_name": "test-model",
|
|
"subagent_enabled": True,
|
|
"is_scheduled_run": True,
|
|
"excluded_tools": ["bash"],
|
|
"allowed_tools": ["read_file"],
|
|
"allow_agent_orchestration": True,
|
|
"app_config": object(),
|
|
"write_file_compatibility": True,
|
|
}
|
|
|
|
def load(**kwargs):
|
|
assert threading.get_ident() != owner_thread
|
|
assert user.get() == "caller"
|
|
assert kwargs == options
|
|
set_deferred_registry(registry)
|
|
return expected
|
|
|
|
monkeypatch.setattr(tool_factory, "get_available_tools", load)
|
|
assert await get_available_tools_async(**options) is expected
|
|
assert get_deferred_registry() is registry
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_loading_during_agent_mcp_discovery_keeps_owner_loop_free(monkeypatch):
|
|
loop = asyncio.get_running_loop()
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
discovery_calls = 0
|
|
expected = [object()]
|
|
identity = contextvars.ContextVar("async_tool_concurrent_identity")
|
|
registries = {name: DeferredToolRegistry() for name in ("first", "second")}
|
|
monkeypatch.setattr(cache, "_mcp_tools_cache", None)
|
|
monkeypatch.setattr(cache, "_cache_initialized", False)
|
|
monkeypatch.setattr(cache, "_config_mtime", None)
|
|
monkeypatch.setattr(cache, "_initialization_lock", asyncio.Lock())
|
|
monkeypatch.setattr(cache, "_get_config_mtime", lambda: 1.0)
|
|
|
|
async def discover():
|
|
nonlocal discovery_calls
|
|
discovery_calls += 1
|
|
assert asyncio.get_running_loop() is loop
|
|
started.set()
|
|
await release.wait()
|
|
return expected
|
|
|
|
def load(**kwargs):
|
|
name = kwargs["model_name"]
|
|
assert identity.get() == name
|
|
result = cache.get_cached_mcp_tools()
|
|
set_deferred_registry(registries[name])
|
|
return result
|
|
|
|
async def request(name):
|
|
identity.set(name)
|
|
result = await get_available_tools_async(model_name=name)
|
|
assert get_deferred_registry() is registries[name]
|
|
return result
|
|
|
|
monkeypatch.setattr(mcp_tools, "get_mcp_tools", discover)
|
|
monkeypatch.setattr(tool_factory, "get_available_tools", load)
|
|
# An agent initializer already owns MCP discovery when two tool consumers arrive.
|
|
agent = asyncio.create_task(run_background_init(cache.get_cached_mcp_tools))
|
|
consumers = []
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=3)
|
|
consumers = [asyncio.create_task(request(name)) for name in registries]
|
|
await asyncio.sleep(0.02)
|
|
assert not agent.done() and all(not task.done() for task in consumers)
|
|
assert discovery_calls == 1
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.wait_for(asyncio.gather(agent, *consumers), timeout=3)
|
|
assert len(results) == 3 and all(result is expected for result in results)
|
|
assert discovery_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("entry_point", ["capabilities", "executor"])
|
|
async def test_deep_research_provider_initialization_leaves_loop_free(monkeypatch, entry_point):
|
|
from app.gateway import deep_research_job_executor as executor_module
|
|
from app.gateway.routers import deep_research as router
|
|
from deerflow.agents.deep_research.adapters import artifacts, context, event_sink, image, llm, material_provider
|
|
from deerflow.agents.deep_research.cancellation import CancellationToken
|
|
from deerflow.agents.deep_research.config import DeepResearchConfig
|
|
from deerflow.agents.deep_research.types import DeepResearchResult
|
|
|
|
loop = asyncio.get_running_loop()
|
|
owner_thread = threading.get_ident()
|
|
started = asyncio.Event()
|
|
release = threading.Event()
|
|
|
|
def load_tools(_model_name):
|
|
assert threading.get_ident() != owner_thread
|
|
loop.call_soon_threadsafe(started.set)
|
|
assert release.wait(3), "the owner loop did not release initialization"
|
|
return {"web_search"}
|
|
|
|
monkeypatch.setattr(material_provider, "_load_question_answer_tool_names", load_tools)
|
|
monkeypatch.setattr(material_provider, "_load_question_answer_research_skills", lambda: [])
|
|
monkeypatch.setattr(image.ImageGenerationSettings, "from_app_config", lambda: SimpleNamespace(is_configured=False))
|
|
monkeypatch.setattr(image, "build_image_provider", lambda: None)
|
|
monkeypatch.setattr(llm, "DeerFlowCompletionBackend", lambda *args, **kwargs: object())
|
|
monkeypatch.setattr(context, "DeerFlowContextCompressor", lambda **kwargs: object())
|
|
monkeypatch.setattr(artifacts, "ThreadOutputsArtifactWriter", lambda *args, **kwargs: object())
|
|
sink = SimpleNamespace(diagnostics=[], publish_live=AsyncMock())
|
|
monkeypatch.setattr(event_sink, "PersistingEventSink", lambda *args, **kwargs: sink)
|
|
monkeypatch.setattr(executor_module, "RecordingEventSink", lambda *args, **kwargs: sink)
|
|
engine_run = AsyncMock(return_value=DeepResearchResult(report_markdown="completed"))
|
|
monkeypatch.setattr(executor_module, "DeepResearchEngine", lambda: SimpleNamespace(run=engine_run))
|
|
executor = executor_module.DeepResearchJobExecutor(None, None, event_store=None)
|
|
monkeypatch.setattr(executor, "_ensure_runtime_thread", AsyncMock(return_value="open-test-research"))
|
|
monkeypatch.setattr(executor, "_load_interrupted_plan", AsyncMock(return_value=None))
|
|
|
|
if entry_point == "capabilities":
|
|
pending = router.get_capabilities()
|
|
else:
|
|
pending = executor._execute(
|
|
job_id="test-job",
|
|
session_id="test-session",
|
|
user_id="test-user",
|
|
query="test query",
|
|
title="test title",
|
|
config=DeepResearchConfig(),
|
|
lease_owner=None,
|
|
cancel_token=CancellationToken(),
|
|
)
|
|
task = asyncio.create_task(pending)
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=3)
|
|
await asyncio.sleep(0)
|
|
assert not task.done()
|
|
finally:
|
|
release.set()
|
|
result = await asyncio.wait_for(task, timeout=3)
|
|
if entry_point == "capabilities":
|
|
assert "materialChannels" in result
|
|
else:
|
|
provider = engine_run.call_args.kwargs["adapters"].materials
|
|
assert isinstance(provider, material_provider.DeerFlowMaterialProvider)
|
|
assert provider._question_answer_tool_names == {"web_search"}
|
|
assert provider._online_search_semaphore is None
|
|
assert provider._user_id == "test-user"
|
|
assert provider._runtime_thread_id == "open-test-research"
|
|
assert result.report_markdown == "completed"
|