deer-flow/backend/tests/test_mcp_binding_fencing.py

"""Regression tests for stdio MCP binding ownership/fencing."""

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from langchain_core.tools import StructuredTool
from pydantic import BaseModel

from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.mcp import session_pool as session_pool_module
from deerflow.mcp.session_pool import (
    MCPSessionPool,
    StaleMCPBindingError,
    normalized_connection_fingerprint,
)

_CONNECTION = {"transport": "stdio", "command": "x", "args": []}


class _GatedSessionCm:
    def __init__(self, gate: asyncio.Event | None = None) -> None:
        self.gate = gate
        self.initialize_started = asyncio.Event()
        self.closed = False
        self.enter_task = None
        self.exit_task = None
        self.session = MagicMock()
        self.session.initialize = self.initialize

    async def __aenter__(self):
        self.enter_task = asyncio.current_task()
        return self.session

    async def initialize(self):
        self.initialize_started.set()
        if self.gate is not None:
            await self.gate.wait()

    async def __aexit__(self, *_args):
        self.exit_task = asyncio.current_task()
        self.closed = True


def test_binding_digest_tracks_stdio_identity_without_leaking_env_value():
    a = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/a", "env": {"TOKEN": "secret-a"}})
    b = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/a", "env": {"TOKEN": "secret-b"}})
    c = normalized_connection_fingerprint({**_CONNECTION, "cwd": "/b", "env": {"TOKEN": "secret-a"}})

    assert a != b != c
    assert "secret-a" not in a
    assert len(a) == 64


def test_same_binding_is_idempotent_and_same_name_domains_are_independent():
    pool = MCPSessionPool()
    deployment = pool.ensure_binding("A", "fp-a", domain="deployment")

    assert pool.ensure_binding("A", "fp-a", domain="deployment") is deployment
    personal = pool.ensure_binding("A", "fp-a", domain="personal")
    assert personal is not deployment


@pytest.mark.asyncio
async def test_remove_readd_same_fingerprint_never_reauthorizes_old_binding():
    pool = MCPSessionPool()
    old = pool.ensure_binding("A", "same-fp")

    pool.reconcile_bindings({}, {"A"})
    tombstone = pool._bindings[("deployment", "A")]
    pool.reconcile_bindings({"A": "same-fp"}, ())
    current = pool._bindings[("deployment", "A")]

    assert tombstone.fingerprint is None
    assert current is not old
    assert current.epoch > tombstone.epoch > old.epoch

    with pytest.raises(StaleMCPBindingError):
        await pool.get_session("A", "u:t", _CONNECTION, binding=old)


@pytest.mark.asyncio
async def test_reconcile_detaches_only_changed_server_and_preserves_registry_key_shape():
    pool = MCPSessionPool()
    loop = asyncio.get_running_loop()
    a = pool.ensure_binding("A", "a1")
    b = pool.ensure_binding("B", "b1")
    cm_a = _GatedSessionCm()
    cm_b = _GatedSessionCm()

    with patch("langchain_mcp_adapters.sessions.create_session", side_effect=[cm_a, cm_b]):
        session_a = await pool.get_session("A", "u:t", _CONNECTION, binding=a)
        session_b = await pool.get_session("B", "u:t", _CONNECTION, binding=b)

    prepared = pool.reconcile_bindings({"A": "a2", "B": "b1"}, ())
    assert [entry[0] for entry in prepared.entries] == [session_a]
    assert prepared.inflight == ()
    assert pool._entries[("B", "u:t", loop, "deployment")][0] is session_b
    assert cm_b.closed is False

    await asyncio.wait_for(prepared.entries[0][2], timeout=1)
    assert cm_a.closed is True
    assert cm_a.exit_task is cm_a.enter_task


@pytest.mark.asyncio
async def test_reconcile_force_rebind_reepochs_same_fingerprint_and_preserves_others():
    """A shared generation change re-epochs A even though its fingerprint is unchanged.

    ``force_rebind`` lets a caller that already knows a
    server's shared lifecycle version changed install a fresh binding in the
    same non-awaiting critical section, without a second teardown owner and
    without touching any other domain or server.
    """
    pool = MCPSessionPool()
    loop = asyncio.get_running_loop()
    a = pool.ensure_binding("A", "same-fp")
    b = pool.ensure_binding("B", "b-fp")
    personal_a = pool.ensure_binding("A", "same-fp", domain="personal")
    cm_a = _GatedSessionCm()
    cm_b = _GatedSessionCm()
    cm_personal = _GatedSessionCm()

    with patch("langchain_mcp_adapters.sessions.create_session", side_effect=[cm_a, cm_b, cm_personal]):
        session_a = await pool.get_session("A", "u:t", _CONNECTION, binding=a)
        session_b = await pool.get_session("B", "u:t", _CONNECTION, binding=b)
        session_personal = await pool.get_session("A", "u:p", _CONNECTION, binding=personal_a)

    prepared = pool.reconcile_bindings({"A": "same-fp", "B": "b-fp"}, (), force_rebind={"A"})

    current_a = pool._bindings[("deployment", "A")]
    assert current_a is not a
    assert current_a.fingerprint == "same-fp"
    assert current_a.epoch > a.epoch

    # Only the forced server changed; its sibling and the personal same-name
    # binding keep their exact capability objects.
    assert pool._bindings[("deployment", "B")] is b
    assert pool._bindings[("personal", "A")] is personal_a

    # Only the forced server's established session is detached and signalled.
    assert [entry[0] for entry in prepared.entries] == [session_a]
    assert prepared.inflight == ()
    assert pool._entries[("B", "u:t", loop, "deployment")][0] is session_b
    assert pool._entries[("A", "u:p", loop, "personal")][0] is session_personal
    assert cm_b.closed is False
    assert cm_personal.closed is False

    # The old capability can no longer create or return the session.
    with pytest.raises(StaleMCPBindingError):
        await pool.get_session("A", "u:t", _CONNECTION, binding=a)

    await asyncio.wait_for(prepared.entries[0][2], timeout=1)
    assert cm_a.closed is True

    # Re-consuming an already-applied generation is a no-op: A stays stable.
    again = pool.reconcile_bindings({"A": "same-fp", "B": "b-fp"}, ())
    assert again.entries == ()
    assert pool._bindings[("deployment", "A")] is current_a


@pytest.mark.asyncio
async def test_stale_binding_fails_before_session_creation():
    pool = MCPSessionPool()
    old = pool.ensure_binding("A", "a1")
    pool.reconcile_bindings({"A": "a2"}, ())

    with patch("langchain_mcp_adapters.sessions.create_session") as create_session:
        with pytest.raises(StaleMCPBindingError):
            await pool.get_session("A", "u:t", _CONNECTION, binding=old)

    create_session.assert_not_called()


@pytest.mark.asyncio
async def test_reconcile_fences_creator_blocked_in_initialize():
    pool = MCPSessionPool()
    old = pool.ensure_binding("A", "a1")
    gate = asyncio.Event()
    cm = _GatedSessionCm(gate)

    with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm):
        creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old))
        await asyncio.wait_for(cm.initialize_started.wait(), 1)

        prepared = pool.reconcile_bindings({"A": "a2"}, ())
        assert len(prepared.inflight) == 1
        gate.set()

        with pytest.raises(StaleMCPBindingError):
            await asyncio.wait_for(creator, 1)

    assert not pool._entries
    assert not pool._inflight
    assert cm.closed is True
    assert cm.exit_task is cm.enter_task


@pytest.mark.asyncio
async def test_creator_and_joiner_both_observe_superseded_binding():
    pool = MCPSessionPool()
    old = pool.ensure_binding("A", "a1")
    gate = asyncio.Event()
    cm = _GatedSessionCm(gate)

    with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm):
        creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old))
        await asyncio.wait_for(cm.initialize_started.wait(), 1)
        joiner = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old))
        await asyncio.sleep(0)

        pool.reconcile_bindings({"A": "a2"}, ())
        gate.set()

        with pytest.raises(StaleMCPBindingError):
            await asyncio.wait_for(creator, 1)
        with pytest.raises(StaleMCPBindingError):
            await asyncio.wait_for(joiner, 1)


@pytest.mark.asyncio
async def test_reset_fences_old_pool_and_old_binding():
    pool = MCPSessionPool()
    binding = pool.ensure_binding("A", "a1")
    session_pool_module._pool = pool

    retired = session_pool_module.reset_session_pool()
    assert retired is pool

    with patch("langchain_mcp_adapters.sessions.create_session") as create_session:
        with pytest.raises(StaleMCPBindingError):
            await pool.get_session("A", "u:t", _CONNECTION, binding=binding)
    create_session.assert_not_called()


class _EmptyArgs(BaseModel):
    pass


@pytest.mark.asyncio
async def test_discovery_captures_pool_and_binding_before_first_await():
    from deerflow.mcp.tools import get_mcp_tools

    config = ExtensionsConfig.model_validate(
        {
            "mcpServers": {
                "A": {
                    "type": "stdio",
                    "command": "x",
                    "args": [],
                }
            }
        }
    )
    servers = {"A": {"transport": "stdio", "command": "x", "args": []}}
    pool = MagicMock()
    binding = SimpleNamespace(domain="deployment", server_name="A")
    order: list[str] = []

    def ensure_binding(*_args, **_kwargs):
        order.append("bind")
        return binding

    pool.ensure_binding.side_effect = ensure_binding

    async def raw_tool() -> str:
        return "ok"

    tool = StructuredTool(
        name="A_echo",
        description="echo",
        args_schema=_EmptyArgs,
        coroutine=raw_tool,
    )

    class FakeClient:
        def __init__(self, *_args, **_kwargs):
            self.callbacks = None
            self.tool_interceptors = []

        async def get_tools(self, *, server_name=None):
            assert server_name == "A"
            order.append("discover")
            return [tool]

    with (
        patch("deerflow.mcp.tools.get_session_pool", return_value=pool),
        patch("deerflow.mcp.tools.build_servers_config", return_value=servers),
        patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}),
        patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]),
        patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient),
        patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda wrapped, *_args, **_kwargs: wrapped) as wrap,
    ):
        tools = await get_mcp_tools(config)

    assert order[:2] == ["bind", "discover"]
    assert tools == [tool]
    assert wrap.call_args.kwargs["pool"] is pool
    assert wrap.call_args.kwargs["binding"] is binding


@pytest.mark.asyncio
async def test_durable_stdio_call_binds_base_connection_before_workspace_augmentation():
    from deerflow.mcp.task_tool_caller import McpTaskToolCaller

    config = ExtensionsConfig.model_validate(
        {
            "mcpServers": {
                "reports": {
                    "type": "stdio",
                    "command": "report-mcp",
                    "args": [],
                }
            }
        }
    )
    result = SimpleNamespace(structuredContent={"status": "running"}, isError=False)
    session = SimpleNamespace(call_tool=AsyncMock(return_value=result))
    pool = MCPSessionPool()
    pool.get_session = AsyncMock(return_value=session)

    prepared = {
        "transport": "stdio",
        "command": "report-mcp",
        "args": [],
        "cwd": "/workspace",
        "env": {"TMPDIR": "/workspace/tmp"},
    }

    caller = McpTaskToolCaller(config)
    with (
        patch("deerflow.mcp.task_tool_caller.get_session_pool", return_value=pool),
        patch("deerflow.mcp.task_tool_caller._prepare_stdio_connection", return_value=prepared),
    ):
        actual = await caller.call_tool(
            server_name="reports",
            tool_name="status_report",
            arguments={"task_id": "1"},
            user_id="u1",
            thread_id="t1",
        )

    assert actual is result
    binding = pool._bindings[("deployment", "reports")]
    assert binding.fingerprint == normalized_connection_fingerprint({"transport": "stdio", "command": "report-mcp", "args": []})
    pool.get_session.assert_awaited_once_with(
        "reports",
        "u1:t1",
        prepared,
        binding=binding,
    )


def test_ensure_binding_rejects_changed_fingerprint():
    pool = MCPSessionPool()
    pool.ensure_binding("A", "fp-a")

    with pytest.raises(StaleMCPBindingError):
        pool.ensure_binding("A", "fp-b")


def test_ensure_binding_rejects_tombstoned_server():
    pool = MCPSessionPool()
    pool.ensure_binding("A", "same-fp")
    pool.reconcile_bindings({}, ("A",))

    # remove -> identical re-add must not re-authorize from fingerprint equality
    with pytest.raises(StaleMCPBindingError):
        pool.ensure_binding("A", "same-fp")


def test_ensure_binding_rejects_retired_pool():
    pool = MCPSessionPool()
    pool.retire_all()

    with pytest.raises(StaleMCPBindingError):
        pool.ensure_binding("A", "fp-a")


@pytest.mark.asyncio
async def test_reconcile_retains_detached_owners_when_retirement_is_dropped():
    """Detached owners keep a tracked teardown even without PreparedRetirement."""
    pool = MCPSessionPool()
    binding = pool.ensure_binding("A", "a1")
    gate = asyncio.Event()
    cm = _GatedSessionCm(gate)

    with patch("langchain_mcp_adapters.sessions.create_session", return_value=cm):
        creator = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=binding))
        await asyncio.wait_for(cm.initialize_started.wait(), 1)

        # Return value deliberately dropped: the pool itself must retain the
        # detached owner's teardown instead of leaving a dangling task.
        pool.reconcile_bindings({"A": "a2"}, ())
        reapers = set(pool._teardown_tasks)
        assert reapers, "detached owner teardown was not retained"

        gate.set()
        with pytest.raises(StaleMCPBindingError):
            await asyncio.wait_for(creator, 1)

        await asyncio.wait_for(asyncio.gather(*reapers, return_exceptions=True), 1)
        deadline = asyncio.get_running_loop().time() + 1
        while pool._teardown_tasks and asyncio.get_running_loop().time() < deadline:
            await asyncio.sleep(0)

    assert not pool._teardown_tasks
    assert cm.closed is True
    assert cm.exit_task is cm.enter_task


@pytest.mark.asyncio
async def test_joiner_rechecks_binding_after_ready_resolves():
    """A joiner cannot return a session superseded after ready resolved.

    The joiner path must be pinned to the post-await fence: the test waits until
    the joiner has passed admission and is parked inside ``shield(ready)``, lets
    the shared future resolve, and only then supersedes the binding. A joiner
    that merely raced into admission would fail for the wrong reason.
    """
    pool = MCPSessionPool()
    old = pool.ensure_binding("A", "a1")
    loop = asyncio.get_running_loop()
    key = ("A", "u:t", loop, "deployment")
    ready = loop.create_future()
    close_evt = asyncio.Event()

    async def owner_waiter():
        await close_evt.wait()

    owner = asyncio.create_task(owner_waiter())
    pool._inflight[key] = (loop, ready, owner, close_evt)

    real_shield = asyncio.shield
    joined = asyncio.Event()
    resolved = asyncio.Event()
    release = asyncio.Event()

    async def gated_shield(awaitable):
        if awaitable is not ready:
            return await real_shield(awaitable)
        # Past admission, parked on the shared future: report progress so the
        # test can sequence the supersede strictly after the result is consumed.
        joined.set()
        result = await real_shield(awaitable)
        resolved.set()
        await release.wait()
        return result

    with patch("deerflow.mcp.session_pool.asyncio.shield", new=gated_shield):
        joiner = asyncio.create_task(pool.get_session("A", "u:t", _CONNECTION, binding=old))

        # Proves the joiner cleared admission, saw the in-flight entry and
        # captured the join future instead of failing on the admission fence.
        await asyncio.wait_for(joined.wait(), 1)

        ready.set_result(MagicMock())
        # Proves shield already consumed the result while still parked, so the
        # supersede below lands before the post-await fence is evaluated.
        await asyncio.wait_for(resolved.wait(), 1)

        pool.reconcile_bindings({"A": "a2"}, ())
        release.set()

        with pytest.raises(StaleMCPBindingError):
            await asyncio.wait_for(joiner, 1)

    await asyncio.gather(owner, return_exceptions=True)