deer-flow/backend/tests/test_mcp_session_pool_loop_isolation.py

"""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
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("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()