175 lines
5.7 KiB
Python
175 lines
5.7 KiB
Python
"""Background graph construction keeps MCP discovery on the request loop."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
|
|
import pytest
|
|
|
|
from deerflow.mcp import cache, tools
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def isolated_mcp_cache(monkeypatch):
|
|
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)
|
|
|
|
|
|
def _background_get(loop):
|
|
with cache.mcp_initialization_loop(loop):
|
|
return cache.get_cached_mcp_tools()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_background_discovery_uses_owner_loop_and_does_not_block_it(monkeypatch):
|
|
loop = asyncio.get_running_loop()
|
|
owner_thread = threading.get_ident()
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
expected = [object()]
|
|
|
|
async def discover():
|
|
assert asyncio.get_running_loop() is loop
|
|
assert threading.get_ident() == owner_thread
|
|
started.set()
|
|
await release.wait()
|
|
return expected
|
|
|
|
monkeypatch.setattr(tools, "get_mcp_tools", discover)
|
|
job = asyncio.create_task(asyncio.to_thread(_background_get, loop))
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=3)
|
|
# The owner loop can execute another request while the factory waits.
|
|
await asyncio.sleep(0)
|
|
assert not job.done()
|
|
finally:
|
|
release.set()
|
|
result = await asyncio.wait_for(job, timeout=3)
|
|
assert result is expected
|
|
assert cache._cache_initialized is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_two_background_builds_share_the_owner_loop_initialization(monkeypatch):
|
|
loop = asyncio.get_running_loop()
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
waiter_ready = threading.Event()
|
|
calls = 0
|
|
expected = [object()]
|
|
|
|
async def discover():
|
|
nonlocal calls
|
|
calls += 1
|
|
started.set()
|
|
await release.wait()
|
|
return expected
|
|
|
|
def second_get():
|
|
waiter_ready.set()
|
|
return _background_get(loop)
|
|
|
|
monkeypatch.setattr(tools, "get_mcp_tools", discover)
|
|
first = asyncio.create_task(asyncio.to_thread(_background_get, loop))
|
|
second = None
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=3)
|
|
second = asyncio.create_task(asyncio.to_thread(second_get))
|
|
while not waiter_ready.is_set():
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0.02)
|
|
assert calls == 1
|
|
assert not first.done() and not second.done()
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.wait_for(asyncio.gather(first, *([second] if second else [])), timeout=3)
|
|
assert len(results) == 2 and all(result is expected for result in results)
|
|
assert calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_background_cache_reloads_after_explicit_reset_and_config_change(monkeypatch):
|
|
loop = asyncio.get_running_loop()
|
|
calls = 0
|
|
modified = [1.0]
|
|
|
|
async def discover():
|
|
nonlocal calls
|
|
assert asyncio.get_running_loop() is loop
|
|
calls += 1
|
|
return [calls]
|
|
|
|
monkeypatch.setattr(tools, "get_mcp_tools", discover)
|
|
monkeypatch.setattr(cache, "_get_config_mtime", lambda: modified[0])
|
|
assert await asyncio.to_thread(_background_get, loop) == [1]
|
|
assert await asyncio.to_thread(_background_get, loop) == [1]
|
|
cache.reset_mcp_tools_cache()
|
|
assert await asyncio.to_thread(_background_get, loop) == [2]
|
|
modified[0] = 2.0
|
|
assert await asyncio.to_thread(_background_get, loop) == [3]
|
|
assert cache._config_mtime == 2.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_background_discovery_failure_returns_empty_and_can_retry(monkeypatch):
|
|
loop = asyncio.get_running_loop()
|
|
calls = 0
|
|
|
|
async def discover():
|
|
nonlocal calls
|
|
assert asyncio.get_running_loop() is loop
|
|
calls += 1
|
|
if calls == 1:
|
|
raise RuntimeError("discovery unavailable")
|
|
return ["recovered"]
|
|
|
|
monkeypatch.setattr(tools, "get_mcp_tools", discover)
|
|
assert await asyncio.to_thread(_background_get, loop) == []
|
|
assert cache._cache_initialized is False
|
|
assert await asyncio.to_thread(_background_get, loop) == ["recovered"]
|
|
|
|
|
|
def test_initialization_loop_context_restores_after_nested_failure():
|
|
first = asyncio.new_event_loop()
|
|
second = asyncio.new_event_loop()
|
|
try:
|
|
assert cache._owner_initialization_loop.get() is None
|
|
with cache.mcp_initialization_loop(first):
|
|
with pytest.raises(ValueError), cache.mcp_initialization_loop(second):
|
|
assert cache._owner_initialization_loop.get() is second
|
|
raise ValueError("factory failed")
|
|
assert cache._owner_initialization_loop.get() is first
|
|
assert cache._owner_initialization_loop.get() is None
|
|
finally:
|
|
first.close()
|
|
second.close()
|
|
|
|
|
|
@pytest.mark.parametrize("closed", [False, True])
|
|
def test_inactive_owner_loop_fails_without_submitting_or_falling_back(monkeypatch, closed):
|
|
loop = asyncio.new_event_loop()
|
|
if closed:
|
|
loop.close()
|
|
|
|
async def unexpected():
|
|
raise AssertionError("inactive owner loop must not initialize elsewhere")
|
|
|
|
monkeypatch.setattr(cache, "initialize_mcp_tools", unexpected)
|
|
try:
|
|
assert _background_get(loop) == []
|
|
assert cache._cache_initialized is False
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
def test_unowned_sync_caller_retains_lazy_initialization(monkeypatch):
|
|
async def discover():
|
|
return ["sync tool"]
|
|
|
|
monkeypatch.setattr(tools, "get_mcp_tools", discover)
|
|
assert cache.get_cached_mcp_tools() == ["sync tool"]
|