mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +00:00
243 lines
8.6 KiB
Python
243 lines
8.6 KiB
Python
"""Persistent-loop regressions for concurrent sync callers (#5256)."""
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import sys
|
|
import threading
|
|
from contextlib import asynccontextmanager
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool
|
|
from deerflow.tools.sync import make_sync_tool_wrapper
|
|
|
|
|
|
@pytest.fixture
|
|
def loop_pool(monkeypatch):
|
|
pool = MCPSessionPool()
|
|
loops = [asyncio.new_event_loop(), asyncio.new_event_loop()]
|
|
threads = [threading.Thread(target=loop.run_forever) for loop in loops]
|
|
closed = []
|
|
|
|
@asynccontextmanager
|
|
async def create_session(connection):
|
|
owner = asyncio.current_task()
|
|
session = AsyncMock()
|
|
try:
|
|
yield session
|
|
finally:
|
|
assert asyncio.current_task() is owner
|
|
closed.append(session)
|
|
|
|
monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session)
|
|
for thread in threads:
|
|
thread.start()
|
|
|
|
def run(index, coroutine):
|
|
return asyncio.run_coroutine_threadsafe(coroutine, loops[index]).result(timeout=5)
|
|
|
|
try:
|
|
yield pool, run, closed
|
|
finally:
|
|
run(0, pool.close_all())
|
|
for loop in loops:
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
for thread in threads:
|
|
thread.join(timeout=5)
|
|
assert not thread.is_alive()
|
|
for loop in loops:
|
|
loop.close()
|
|
|
|
|
|
def test_live_loops_reuse_only_their_own_sessions_and_disconnect_is_isolated(loop_pool):
|
|
pool, run, closed = loop_pool
|
|
first = run(0, pool.get_session("s", "u:t", {}))
|
|
sibling = run(1, pool.get_session("s", "u:t", {}))
|
|
assert first is not sibling
|
|
assert not closed
|
|
assert run(0, pool.get_session("s", "u:t", {})) is first
|
|
assert run(1, pool.get_session("s", "u:t", {})) is sibling
|
|
assert run(0, pool.close_session_if_current("s", "u:t", first))
|
|
assert closed == [first]
|
|
replacement = run(0, pool.get_session("s", "u:t", {}))
|
|
assert replacement is not first
|
|
assert not run(0, pool.close_session_if_current("s", "u:t", first))
|
|
assert run(1, pool.get_session("s", "u:t", {})) is sibling
|
|
assert run(0, pool.get_session("s", "u:t", {})) is replacement
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["scope", "server", "session", "all"])
|
|
def test_explicit_cleanup_closes_both_loops(loop_pool, operation):
|
|
pool, run, closed = loop_pool
|
|
sessions = [run(i, pool.get_session("s", "u:t", {})) for i in range(2)]
|
|
other = run(1, pool.get_session("other", "u:other", {}))
|
|
cleanup = {
|
|
"scope": lambda: pool.close_scope("u:t"),
|
|
"server": lambda: pool.close_server("s"),
|
|
"session": lambda: pool.close_session("s", "u:t"),
|
|
"all": pool.close_all,
|
|
}[operation]
|
|
run(0, cleanup())
|
|
assert all(session in closed for session in sessions)
|
|
if operation == "all":
|
|
assert other in closed
|
|
else:
|
|
assert other not in closed
|
|
assert run(1, pool.get_session("other", "u:other", {})) is other
|
|
|
|
|
|
def test_owner_completion_does_not_remove_a_replacement(loop_pool):
|
|
pool, run, _closed = loop_pool
|
|
|
|
async def replace():
|
|
await pool.get_session("s", "u:t", {})
|
|
key = ("s", "u:t", asyncio.get_running_loop(), "deployment")
|
|
old_owner = pool._entries[key][2]
|
|
await pool.close_session("s", "u:t")
|
|
replacement = await pool.get_session("s", "u:t", {})
|
|
pool._discard_owner(key, old_owner)
|
|
assert await pool.get_session("s", "u:t", {}) is replacement
|
|
|
|
run(0, replace())
|
|
|
|
|
|
@pytest.mark.parametrize("capacity", [1, 2])
|
|
def test_capacity_pressure_on_foreign_loop_preserves_active_call(loop_pool, capacity):
|
|
pool, run, closed = loop_pool
|
|
pool.MAX_SESSIONS = capacity
|
|
active = run(0, pool.get_session("s", "u:t", {}))
|
|
idle = run(1, pool.get_session("s", "u:t", {})) if capacity == 2 else None
|
|
started = threading.Event()
|
|
|
|
async def start():
|
|
release = asyncio.Event()
|
|
|
|
async def slow(*args, **kwargs):
|
|
started.set()
|
|
await release.wait()
|
|
return "finished"
|
|
|
|
active.call_tool.side_effect = slow
|
|
task = asyncio.create_task(call_pooled_session_tool(active, pool, server_name="s", scope_key="u:t", tool_name="slow", arguments={}, call_kwargs={}))
|
|
return release, task
|
|
|
|
release, task = run(0, start())
|
|
|
|
async def finish():
|
|
release.set()
|
|
return await task
|
|
|
|
try:
|
|
assert started.wait(5)
|
|
if idle is not None:
|
|
run(1, pool.get_session("s", "new", {}))
|
|
assert idle in closed
|
|
else:
|
|
with pytest.raises(RuntimeError, match="capacity"):
|
|
run(1, pool.get_session("s", "new", {}))
|
|
assert not pool._inflight
|
|
assert active not in closed
|
|
assert run(0, finish()) == "finished"
|
|
assert not pool._active_calls
|
|
finally:
|
|
run(0, finish())
|
|
|
|
|
|
@pytest.mark.parametrize("retirement", ["lru", "explicit"])
|
|
def test_abandoned_closed_loop_entry_can_be_retired(loop_pool, retirement):
|
|
"""Model the registry left by unsupported loop.close() with pending owners.
|
|
|
|
Use a synthetic pending owner so the test itself does not leak a real task
|
|
or transport on a closed loop. Normal owner shutdown is tested separately.
|
|
"""
|
|
pool, run, _closed = loop_pool
|
|
closed_loop = asyncio.new_event_loop()
|
|
closed_loop.close()
|
|
owner = MagicMock(spec=asyncio.Task)
|
|
owner.done.return_value = False
|
|
key = ("s", "u:t", closed_loop, "deployment")
|
|
pool._entries[key] = (MagicMock(), closed_loop, owner, asyncio.Event())
|
|
pool.MAX_SESSIONS = 1
|
|
if retirement == "explicit":
|
|
run(0, pool.close_scope("u:t"))
|
|
assert not pool._entries
|
|
else:
|
|
replacement = run(0, pool.get_session("s", "u:t", {}))
|
|
assert key not in pool._entries
|
|
assert len(pool._entries) == 1
|
|
assert run(0, pool.get_session("s", "u:t", {})) is replacement
|
|
owner.cancel.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["scope", "server", "session", "all"])
|
|
def test_cleanup_cancels_inflight_owners_on_both_loops(loop_pool, monkeypatch, operation):
|
|
pool, run, _closed = loop_pool
|
|
started = [threading.Event(), threading.Event()]
|
|
exited = [threading.Event(), threading.Event()]
|
|
|
|
@asynccontextmanager
|
|
async def create_session(connection):
|
|
index = connection["index"]
|
|
owner = asyncio.current_task()
|
|
|
|
async def initialize():
|
|
started[index].set()
|
|
await asyncio.Future()
|
|
|
|
try:
|
|
yield AsyncMock(initialize=initialize)
|
|
finally:
|
|
assert asyncio.current_task() is owner
|
|
exited[index].set()
|
|
|
|
monkeypatch.setattr("langchain_mcp_adapters.sessions.create_session", create_session)
|
|
|
|
async def start(index):
|
|
return asyncio.create_task(pool.get_session("s", "u:t", {"index": index}))
|
|
|
|
calls = [run(i, start(i)) for i in range(2)]
|
|
assert all(event.wait(5) for event in started)
|
|
cleanup = {
|
|
"scope": lambda: pool.close_scope("u:t"),
|
|
"server": lambda: pool.close_server("s"),
|
|
"session": lambda: pool.close_session("s", "u:t"),
|
|
"all": pool.close_all,
|
|
}[operation]
|
|
run(0, cleanup())
|
|
assert all(event.wait(5) for event in exited)
|
|
|
|
async def cancelled(call):
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await call
|
|
|
|
for i, call in enumerate(calls):
|
|
run(i, cancelled(call))
|
|
assert not pool._inflight and not pool._entries
|
|
|
|
|
|
def test_parallel_sync_wrappers_complete_real_stdio_calls(tmp_path):
|
|
server = tmp_path / "echo_server.py"
|
|
server.write_text(
|
|
'from mcp.server.fastmcp import FastMCP\nmcp = FastMCP("echo")\n@mcp.tool()\ndef echo(text: str) -> str:\n return text\nmcp.run(transport="stdio")\n',
|
|
encoding="utf-8",
|
|
)
|
|
pool = MCPSessionPool()
|
|
barrier = threading.Barrier(2)
|
|
connection = {"transport": "stdio", "command": sys.executable, "args": [str(server)]}
|
|
|
|
async def echo(text):
|
|
await asyncio.to_thread(barrier.wait, 5)
|
|
session = await pool.get_session("echo", "user:thread", connection)
|
|
result = await session.call_tool("echo", {"text": text})
|
|
return result.content[0].text
|
|
|
|
wrapper = make_sync_tool_wrapper(echo, "echo")
|
|
try:
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
|
|
futures = [executor.submit(wrapper, text) for text in ("A", "B")]
|
|
assert [future.result(timeout=30) for future in futures] == ["A", "B"]
|
|
assert not pool._entries and not pool._inflight
|
|
finally:
|
|
pool.close_all_sync()
|