nanobot/tests/tools/test_mcp_tool.py

from __future__ import annotations

import asyncio
import json
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from types import ModuleType, SimpleNamespace

import httpx
import pytest

import nanobot.agent.tools.mcp as mcp_mod
from nanobot.agent.tools.mcp import (
    MCPPromptWrapper,
    MCPProvider,
    MCPResourceWrapper,
    MCPToolWrapper,
    _normalize_windows_stdio_command,
    _sanitize_mcp_tool_name,
    _sanitize_name,
    connect_mcp_servers,
)
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.config.schema import MCPServerConfig

_PROXY_ENV_VARS = ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy")


def test_type_checking_only_mcp_annotations_are_deferred() -> None:
    assert mcp_mod._MCPWrapperBase.__annotations__["_session"] == "ClientSession"
    assert MCPToolWrapper.__init__.__annotations__["session"] == "ClientSession"
    assert MCPResourceWrapper.__init__.__annotations__["resource_def"] == "Resource"
    assert MCPPromptWrapper.__init__.__annotations__["prompt_def"] == "Prompt"
    assert connect_mcp_servers.__annotations__["mcp_servers"] == "dict[str, MCPServerConfig]"


class _FakeTextContent:
    def __init__(self, text: str) -> None:
        self.text = text


class _FakeTextResourceContents:
    def __init__(self, text: str) -> None:
        self.text = text


class _FakeBlobResourceContents:
    def __init__(self, blob: bytes) -> None:
        self.blob = blob


class _FakeImageContent:
    def __init__(self, data: str, mime_type: str = "image/png") -> None:
        self.data = data
        self.mimeType = mime_type


@pytest.fixture
def fake_mcp_runtime() -> dict[str, object | None]:
    return {"session": None}


@pytest.fixture(autouse=True)
def _clear_proxy_env(monkeypatch: pytest.MonkeyPatch) -> None:
    for name in (*_PROXY_ENV_VARS, "NO_PROXY", "no_proxy"):
        monkeypatch.delenv(name, raising=False)


@pytest.fixture(autouse=True)
def _fake_mcp_module(
    monkeypatch: pytest.MonkeyPatch, fake_mcp_runtime: dict[str, object | None]
) -> None:
    mod = ModuleType("mcp")
    mod.types = SimpleNamespace(
        TextContent=_FakeTextContent,
        TextResourceContents=_FakeTextResourceContents,
        BlobResourceContents=_FakeBlobResourceContents,
        ImageContent=_FakeImageContent,
    )

    class _FakeStdioServerParameters:
        def __init__(
            self,
            command: str,
            args: list[str],
            env: dict | None = None,
            cwd: str | None = None,
        ) -> None:
            self.command = command
            self.args = args
            self.env = env
            self.cwd = cwd

    class _FakeClientSession:
        def __init__(self, _read: object, _write: object) -> None:
            self._session = fake_mcp_runtime["session"]

        async def __aenter__(self) -> object:
            return self._session

        async def __aexit__(self, exc_type, exc, tb) -> bool:
            return False

    @asynccontextmanager
    async def _fake_stdio_client(_params: object):
        yield object(), object()

    @asynccontextmanager
    async def _fake_sse_client(_url: str, httpx_client_factory=None):
        yield object(), object()

    @asynccontextmanager
    async def _fake_streamable_http_client(_url: str, http_client=None):
        yield object(), object(), object()

    mod.ClientSession = _FakeClientSession
    mod.StdioServerParameters = _FakeStdioServerParameters
    monkeypatch.setitem(sys.modules, "mcp", mod)

    client_mod = ModuleType("mcp.client")
    stdio_mod = ModuleType("mcp.client.stdio")
    stdio_mod.stdio_client = _fake_stdio_client
    sse_mod = ModuleType("mcp.client.sse")
    sse_mod.sse_client = _fake_sse_client
    streamable_http_mod = ModuleType("mcp.client.streamable_http")
    streamable_http_mod.streamable_http_client = _fake_streamable_http_client

    monkeypatch.setitem(sys.modules, "mcp.client", client_mod)
    monkeypatch.setitem(sys.modules, "mcp.client.stdio", stdio_mod)
    monkeypatch.setitem(sys.modules, "mcp.client.sse", sse_mod)
    monkeypatch.setitem(sys.modules, "mcp.client.streamable_http", streamable_http_mod)

    shared_mod = ModuleType("mcp.shared")
    exc_mod = ModuleType("mcp.shared.exceptions")

    class _FakeMcpError(Exception):
        def __init__(self, code: int = -1, message: str = "error"):
            self.error = SimpleNamespace(code=code, message=message)
            super().__init__(message)

    exc_mod.McpError = _FakeMcpError
    monkeypatch.setitem(sys.modules, "mcp.shared", shared_mod)
    monkeypatch.setitem(sys.modules, "mcp.shared.exceptions", exc_mod)


def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={"type": "object", "properties": {}},
    )
    return MCPToolWrapper(session, "test", tool_def, tool_timeout=timeout)


@pytest.mark.asyncio
async def test_mcp_provider_connect_propagates_external_cancellation(monkeypatch) -> None:
    started = asyncio.Event()

    async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict:
        started.set()
        await asyncio.sleep(60)
        return {}

    provider = MCPProvider(
        {"test": MCPServerConfig(command="fake")},
        ToolRegistry(),
    )
    monkeypatch.setattr(mcp_mod, "connect_mcp_servers", connect_mcp_servers)

    task = asyncio.create_task(provider.connect())
    await asyncio.wait_for(started.wait(), timeout=1.0)
    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task

    assert provider.connected_server_names == set()
    assert provider.runtime_status() == {"test": "failed"}


@pytest.mark.asyncio
async def test_saved_oauth_http_403_projects_failed_runtime_without_details(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class OAuthAuth(httpx.Auth):
        def auth_flow(self, request: httpx.Request):
            request.headers["Authorization"] = "Bearer saved-oauth-secret"
            yield request

    class AuthorizationRequiredError(RuntimeError):
        pass

    async def create_auth(_name: str, _url: str, _handlers=None) -> httpx.Auth:
        return OAuthAuth()

    @asynccontextmanager
    async def rejected_streamable_http(url: str, http_client=None):
        request = httpx.Request("POST", f"{url}?access_token=saved-oauth-secret")
        response = httpx.Response(403, request=request)
        raise httpx.HTTPStatusError(
            "403 Forbidden for saved-oauth-secret",
            request=request,
            response=response,
        )
        yield object(), object(), object()

    async def reachable(_url: str) -> bool:
        return True

    oauth_mod = ModuleType("nanobot.agent.tools.mcp_oauth")
    oauth_mod.MCPAuthorizationRequiredError = AuthorizationRequiredError  # type: ignore[attr-defined]
    oauth_mod.create_mcp_oauth_auth = create_auth  # type: ignore[attr-defined]
    oauth_mod.mcp_oauth_has_credentials = lambda _name, _url: True  # type: ignore[attr-defined]
    monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", reachable)
    monkeypatch.setattr(
        sys.modules["mcp.client.streamable_http"],
        "streamable_http_client",
        rejected_streamable_http,
    )

    provider = MCPProvider(
        {
            "xmind": MCPServerConfig(
            type="streamableHttp",
            auth="oauth",
            url="https://app.xmind.com/api/mcp",
            )
        },
        ToolRegistry(),
    )

    await provider.connect()

    snapshot = provider.runtime_status()
    assert snapshot == {"xmind": "failed"}
    assert "saved-oauth-secret" not in str(snapshot)
    assert "app.xmind.com" not in str(snapshot)


def test_wrapper_preserves_non_nullable_unions() -> None:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={
            "type": "object",
            "properties": {
                "value": {
                    "anyOf": [{"type": "string"}, {"type": "integer"}],
                }
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    assert wrapper.parameters["properties"]["value"]["anyOf"] == [
        {"type": "string"},
        {"type": "integer"},
    ]


def test_wrapper_normalizes_nullable_property_type_union() -> None:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={
            "type": "object",
            "properties": {
                "name": {"type": ["string", "null"]},
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    assert wrapper.parameters["properties"]["name"] == {"type": "string", "nullable": True}


def test_wrapper_normalizes_nullable_property_anyof() -> None:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={
            "type": "object",
            "properties": {
                "name": {
                    "anyOf": [{"type": "string"}, {"type": "null"}],
                    "description": "optional name",
                },
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    assert wrapper.parameters["properties"]["name"] == {
        "type": "string",
        "description": "optional name",
        "nullable": True,
    }


def test_wrapper_hoists_recursive_local_refs_into_defs() -> None:
    recursive_items_ref = "#/properties/filter/properties/items"
    tool_def = SimpleNamespace(
        name="search_dataset",
        description="search tool",
        inputSchema={
            "type": "object",
            "properties": {
                "filter": {
                    "type": "object",
                    "properties": {
                        "items": {
                            "type": "array",
                            "items": {"$ref": recursive_items_ref},
                        }
                    },
                    "required": ["items"],
                }
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    generated_ref = wrapper.parameters["properties"]["filter"]["properties"]["items"][
        "items"
    ]["$ref"]
    assert generated_ref.startswith("#/$defs/ref_")
    generated_name = generated_ref.removeprefix("#/$defs/")
    generated_schema = wrapper.parameters["$defs"][generated_name]
    assert generated_schema["type"] == "array"
    assert generated_schema["items"]["$ref"] == generated_ref


def test_wrapper_hoists_root_self_ref_into_defs() -> None:
    tool_def = SimpleNamespace(
        name="tree",
        description="tree tool",
        inputSchema={
            "type": "object",
            "properties": {
                "children": {"type": "array", "items": {"$ref": "#"}},
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    generated_ref = wrapper.parameters["properties"]["children"]["items"]["$ref"]
    assert generated_ref.startswith("#/$defs/ref_")
    generated_name = generated_ref.removeprefix("#/$defs/")
    assert wrapper.parameters["$defs"][generated_name]["properties"]["children"]["items"] == {
        "$ref": generated_ref
    }


def test_wrapper_preserves_existing_defs_refs() -> None:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={
            "type": "object",
            "$defs": {"value": {"type": "string"}},
            "properties": {"value": {"$ref": "#/$defs/value"}},
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    assert wrapper.parameters["properties"]["value"]["$ref"] == "#/$defs/value"
    assert wrapper.parameters["$defs"]["value"]["type"] == "string"


def test_wrapper_resolves_uri_encoded_json_pointer() -> None:
    tool_def = SimpleNamespace(
        name="demo",
        description="demo tool",
        inputSchema={
            "type": "object",
            "properties": {
                "space name/value": {"type": "string"},
                "alias": {"$ref": "#/properties/space%20name~1value"},
            },
        },
    )

    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)

    generated_ref = wrapper.parameters["properties"]["alias"]["$ref"]
    assert generated_ref.startswith("#/$defs/ref_")
    generated_name = generated_ref.removeprefix("#/$defs/")
    assert wrapper.parameters["$defs"][generated_name] == {"type": "string"}


def test_normalize_windows_stdio_command_is_noop_off_windows(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(mcp_mod.os, "name", "posix", raising=False)

    command, args, env = _normalize_windows_stdio_command(
        "npx",
        ["-y", "chrome-devtools-mcp@latest"],
        {"FOO": "bar"},
    )

    assert command == "npx"
    assert args == ["-y", "chrome-devtools-mcp@latest"]
    assert env == {"FOO": "bar"}


def test_normalize_windows_stdio_command_wraps_npx_on_windows(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(mcp_mod.os, "name", "nt", raising=False)
    monkeypatch.setattr(
        mcp_mod.shutil,
        "which",
        lambda command, path=None: r"C:\Program Files\nodejs\npx.cmd",
    )
    monkeypatch.setenv("COMSPEC", r"C:\Windows\System32\cmd.exe")

    command, args, env = _normalize_windows_stdio_command(
        "npx",
        ["-y", "chrome-devtools-mcp@latest"],
        None,
    )

    assert command == r"C:\Windows\System32\cmd.exe"
    assert args == ["/d", "/c", "npx", "-y", "chrome-devtools-mcp@latest"]
    assert env is None


def test_normalize_windows_stdio_command_wraps_resolved_cmd_launcher(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(mcp_mod.os, "name", "nt", raising=False)

    def _fake_which(command: str, path: str | None = None) -> str:
        assert command == "custom-launcher"
        assert path == r"C:\Tools"
        return r"C:\Tools\custom-launcher.cmd"

    monkeypatch.setattr(mcp_mod.shutil, "which", _fake_which)
    monkeypatch.setenv("COMSPEC", r"C:\Windows\System32\cmd.exe")

    command, args, _env = _normalize_windows_stdio_command(
        "custom-launcher",
        ["serve"],
        {"PATH": r"C:\Tools"},
    )

    assert command == r"C:\Windows\System32\cmd.exe"
    assert args == ["/d", "/c", "custom-launcher", "serve"]


def test_normalize_windows_stdio_command_keeps_real_executables_unchanged(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(mcp_mod.os, "name", "nt", raising=False)

    command, args, env = _normalize_windows_stdio_command(
        "python.exe",
        ["-m", "http.server"],
        {"FOO": "bar"},
    )

    assert command == "python.exe"
    assert args == ["-m", "http.server"]
    assert env == {"FOO": "bar"}


def test_normalize_windows_stdio_command_skips_existing_shells(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(mcp_mod.os, "name", "nt", raising=False)

    command, args, env = _normalize_windows_stdio_command(
        "cmd.exe",
        ["/c", "echo", "hello"],
        None,
    )

    assert command == "cmd.exe"
    assert args == ["/c", "echo", "hello"]
    assert env is None


@pytest.mark.asyncio
async def test_execute_returns_text_blocks() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        assert arguments == {"value": 1}
        return SimpleNamespace(content=[_FakeTextContent("hello"), 42])

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute(value=1)

    assert result == "hello\n42"


@pytest.mark.asyncio
async def test_execute_wraps_mcp_is_error_result() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(
            content=[_FakeTextContent("Error: server-side MCP failure")],
            isError=True,
        )

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "Error: server-side MCP failure"
    assert is_tool_error_result(result)


@pytest.mark.asyncio
async def test_execute_contains_malformed_success_result() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(content=None)

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "(MCP tool returned malformed content: TypeError)"
    assert is_tool_error_result(result)


@pytest.mark.asyncio
async def test_registry_adds_retry_hint_to_malformed_mcp_result() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(content=None)

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
    registry = ToolRegistry()
    registry.register(wrapper)

    result = await registry.execute(wrapper.name, {})

    assert is_tool_error_result(result)
    assert "MCP tool returned malformed content" in result
    assert "Analyze the error above and try a different approach" in result


@pytest.mark.asyncio
async def test_execute_preserves_success_text_that_starts_with_error() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(
            content=[_FakeTextContent("Error: generated report successfully")],
            isError=False,
        )

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "Error: generated report successfully"
    assert not is_tool_error_result(result)


# Smallest valid 1x1 PNG, base64 without the data: prefix.
_PNG_B64 = (
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8"
    "/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
)


@pytest.mark.asyncio
async def test_execute_persists_image_block_as_artifact(tmp_path: Path) -> None:
    from nanobot.config.loader import set_config_path

    set_config_path(tmp_path / "config.json")

    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(
            content=[
                _FakeTextContent("here you go"),
                _FakeImageContent(_PNG_B64, "image/png"),
            ]
        )

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute(prompt="a cat", model="sdxl")

    payload = json.loads(result)
    assert payload["text"] == "here you go"
    assert len(payload["artifacts"]) == 1
    artifact = payload["artifacts"][0]
    assert artifact["mime"] == "image/png"
    assert artifact["prompt"] == "a cat"
    assert artifact["provider"] == "mcp:test"
    assert Path(artifact["path"]).is_file()
    # The base64 payload must NOT leak into the model-facing result.
    assert _PNG_B64 not in result
    assert "message tool" in payload["next_step"]


@pytest.mark.asyncio
async def test_execute_notes_unstorable_image_block(tmp_path: Path) -> None:
    from nanobot.config.loader import set_config_path

    set_config_path(tmp_path / "config.json")

    async def call_tool(_name: str, arguments: dict) -> object:
        return SimpleNamespace(content=[_FakeImageContent("not-valid-base64!!", "image/png")])

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "(MCP tool returned an image that could not be stored)"


@pytest.mark.asyncio
async def test_execute_returns_timeout_message() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        await asyncio.sleep(1)
        return SimpleNamespace(content=[])

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool), timeout=0.01)

    result = await wrapper.execute()

    assert result == "(MCP tool call timed out after 0.01s)"
    assert is_tool_error_result(result)


@pytest.mark.asyncio
async def test_execute_handles_server_cancelled_error() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        raise asyncio.CancelledError()

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "(MCP tool call was cancelled)"
    assert is_tool_error_result(result)


@pytest.mark.asyncio
async def test_execute_re_raises_external_cancellation() -> None:
    started = asyncio.Event()

    async def call_tool(_name: str, arguments: dict) -> object:
        started.set()
        await asyncio.sleep(60)
        return SimpleNamespace(content=[])

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool), timeout=10)
    task = asyncio.create_task(wrapper.execute())
    await asyncio.wait_for(started.wait(), timeout=1.0)

    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task


@pytest.mark.asyncio
async def test_execute_handles_generic_exception() -> None:
    async def call_tool(_name: str, arguments: dict) -> object:
        raise RuntimeError("boom")

    wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))

    result = await wrapper.execute()

    assert result == "(MCP tool call failed: RuntimeError)"
    assert is_tool_error_result(result)


def _make_tool_def(name: str) -> SimpleNamespace:
    return SimpleNamespace(
        name=name,
        description=f"{name} tool",
        inputSchema={"type": "object", "properties": {}},
    )


def _make_fake_session(tool_names: list[str]) -> SimpleNamespace:
    async def initialize() -> None:
        return None

    async def list_tools() -> SimpleNamespace:
        return SimpleNamespace(tools=[_make_tool_def(name) for name in tool_names])

    return SimpleNamespace(initialize=initialize, list_tools=list_tools)


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_supports_raw_names(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == ["mcp_test_demo"]


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_defaults_to_all(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake")},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == ["mcp_test_demo", "mcp_test_other"]


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_supports_wrapped_names(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_demo"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == ["mcp_test_demo"]


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_supports_limited_wrapped_names(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    long_tool_name = "tool-" + "very-long-name-" * 8
    wrapped_name = _sanitize_mcp_tool_name(f"mcp_test_{long_tool_name}")
    assert len(wrapped_name) == 64

    fake_mcp_runtime["session"] = _make_fake_session([long_tool_name, "other"])
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=[wrapped_name])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == [wrapped_name]


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_empty_list_registers_none(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo", "other"])
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=[])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == []


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_empty_list_blocks_resources_and_prompts(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    """enabledTools: [] (deny-all) must also block resource and prompt registration."""
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=["demo"],
        resource_names=["secret_data"],
        prompt_names=["admin_prompt"],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=[])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == []
    # Resources and prompts must also be blocked
    assert not any("secret_data" in name for name in registry.tool_names)
    assert not any("admin_prompt" in name for name in registry.tool_names)


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_specific_list_blocks_resources_and_prompts(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    """enabledTools with specific tool names must not leak resources or prompts."""
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=["demo", "other"],
        resource_names=["secret_data"],
        prompt_names=["admin_prompt"],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["demo"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    # Only the allowed tool should be registered
    assert "mcp_test_demo" in registry.tool_names
    assert "mcp_test_other" not in registry.tool_names
    # Resources and prompts must not leak
    assert not any("secret_data" in name for name in registry.tool_names)
    assert not any("admin_prompt" in name for name in registry.tool_names)


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_wildcard_allows_resources_and_prompts(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    """enabledTools: ['*'] should allow all tools, resources, and prompts."""
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=["demo"],
        resource_names=["public_data"],
        prompt_names=["help_prompt"],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["*"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert "mcp_test_demo" in registry.tool_names
    assert any("public_data" in name for name in registry.tool_names)
    assert any("help_prompt" in name for name in registry.tool_names)


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_warns_on_unknown_entries(
    fake_mcp_runtime: dict[str, object | None], monkeypatch: pytest.MonkeyPatch
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    registry = ToolRegistry()
    warnings: list[str] = []

    def _warning(message: str, *args: object) -> None:
        warnings.append(message.format(*args))

    monkeypatch.setattr("nanobot.agent.tools.mcp.logger.warning", _warning)

    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["unknown"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == []
    assert warnings
    assert "enabledTools entries not found: unknown" in warnings[-1]
    assert "Available raw names: demo" in warnings[-1]
    assert "Available wrapped names: mcp_test_demo" in warnings[-1]


@pytest.mark.asyncio
async def test_connect_mcp_servers_logs_stdio_pollution_hint(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    messages: list[str] = []

    @asynccontextmanager
    async def _broken_stdio_client(_params: object):
        raise RuntimeError("Parse error: Unexpected token 'INFO' before JSON-RPC headers")
        yield  # pragma: no cover

    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _broken_stdio_client)
    sink = mcp_mod.logger.add(
        lambda message: messages.append(message.record["message"]), level="ERROR"
    )

    registry = ToolRegistry()
    try:
        stacks = await connect_mcp_servers(
            {"gh": MCPServerConfig(command="github-mcp")}, registry
        )
    finally:
        mcp_mod.logger.remove(sink)

    assert stacks == {}
    assert messages
    assert "stdio protocol pollution" in messages[-1]
    assert "stdout" in messages[-1]
    assert "stderr" in messages[-1]


def test_transient_connection_group_logs_brief_warning_and_debug_trace() -> None:
    records: list[dict] = []
    sink = mcp_mod.logger.add(lambda message: records.append(message.record), level="DEBUG")
    error = ExceptionGroup("transport failed", [httpx.ConnectError("")])
    try:
        mcp_mod._log_mcp_connection_failure("notion", error)
    finally:
        mcp_mod.logger.remove(sink)

    warning = next(record for record in records if record["level"].name == "WARNING")
    debug = next(record for record in records if record["level"].name == "DEBUG")
    assert warning["exception"] is None
    assert "transient connection failure" in warning["message"]
    assert debug["exception"] is not None
    assert not any(record["level"].name == "ERROR" for record in records)


def test_unexpected_connection_failure_keeps_error_trace() -> None:
    records: list[dict] = []
    sink = mcp_mod.logger.add(lambda message: records.append(message.record), level="DEBUG")
    try:
        mcp_mod._log_mcp_connection_failure("notion", RuntimeError("boom"))
    finally:
        mcp_mod.logger.remove(sink)

    error = next(record for record in records if record["level"].name == "ERROR")
    assert error["exception"] is not None
    assert not any(record["level"].name == "WARNING" for record in records)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "config",
    [
        MCPServerConfig(url="http://127.0.0.1:9/sse"),
        MCPServerConfig(type="streamableHttp", url="http://127.0.0.1:9/mcp"),
        MCPServerConfig(url="http://169.254.169.254/sse"),
        MCPServerConfig(type="streamableHttp", url="http://169.254.169.254/mcp"),
    ],
)
async def test_connect_mcp_servers_rejects_unsafe_http_urls_before_probe(
    config: MCPServerConfig,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    attempted_connections: list[tuple[object, ...]] = []
    warnings: list[str] = []

    async def _open_connection(*args: object, **_kwargs: object):
        attempted_connections.append(args)
        raise AssertionError("unsafe MCP URL should be rejected before TCP probe")

    def _warning(message: str, *args: object) -> None:
        warnings.append(message.format(*args))

    monkeypatch.setattr(mcp_mod.asyncio, "open_connection", _open_connection)
    monkeypatch.setattr("nanobot.agent.tools.mcp.logger.warning", _warning)

    registry = ToolRegistry()
    stacks = await connect_mcp_servers({"local": config}, registry)

    assert stacks == {}
    assert registry.tool_names == []
    assert attempted_connections == []
    assert any("blocked unsafe URL" in warning for warning in warnings)


@pytest.mark.asyncio
async def test_validate_mcp_request_url_rejects_loopback_without_whitelist() -> None:
    from nanobot.security.network import configure_ssrf_whitelist

    configure_ssrf_whitelist([])
    request = httpx.Request("GET", "http://127.0.0.1/private")

    with pytest.raises(httpx.RequestError, match="Blocked unsafe MCP URL"):
        await mcp_mod._validate_mcp_request_url(request)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "config",
    [
        MCPServerConfig(type="sse", url="https://mcp.example.com/sse"),
        MCPServerConfig(type="streamableHttp", url="https://mcp.example.com/mcp"),
    ],
)
async def test_connect_mcp_servers_env_proxy_adds_proxy_mounts_and_keeps_pinned_transport(
    config: MCPServerConfig,
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    client_kwargs: list[dict[str, object]] = []

    async def _reachable(_url: str) -> bool:
        return True

    async def _validate(_url: str) -> tuple[bool, str]:
        return True, ""

    class FakeAsyncClient:
        def __init__(self, *args: object, **kwargs: object) -> None:
            client_kwargs.append(kwargs)

        async def __aenter__(self) -> object:
            return self

        async def __aexit__(self, exc_type: object, exc: object, tb: object) -> bool:
            return False

    @asynccontextmanager
    async def _capturing_sse_client(_url: str, httpx_client_factory=None):
        assert httpx_client_factory is not None
        async with httpx_client_factory():
            pass
        yield object(), object()

    @asynccontextmanager
    async def _capturing_streamable_http_client(_url: str, http_client=None):
        assert http_client is not None
        yield object(), object(), object()

    monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
    monkeypatch.setenv("NO_PROXY", "localhost,127.0.0.1,::1")
    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
    monkeypatch.setattr(
        mcp_mod,
        "PinnedDNSAsyncTransport",
        lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
    )
    monkeypatch.setattr(
        "nanobot.security.network.httpx.AsyncHTTPTransport",
        lambda **_kwargs: httpx.MockTransport(
            lambda request: httpx.Response(200, request=request)
        ),
    )
    monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", FakeAsyncClient)
    monkeypatch.setattr(sys.modules["mcp.client.sse"], "sse_client", _capturing_sse_client)
    monkeypatch.setattr(
        sys.modules["mcp.client.streamable_http"],
        "streamable_http_client",
        _capturing_streamable_http_client,
    )

    registry = ToolRegistry()
    stacks = await connect_mcp_servers({"remote": config}, registry)
    for stack in stacks.values():
        await stack.aclose()

    assert client_kwargs
    assert all("transport" in kwargs for kwargs in client_kwargs)
    assert all("mounts" in kwargs for kwargs in client_kwargs)


def test_mcp_http_clients_no_proxy_env_keeps_pinned_direct_route(monkeypatch):
    monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:8080")
    monkeypatch.setenv("NO_PROXY", "mcp.example.com")
    monkeypatch.setattr(
        mcp_mod,
        "PinnedDNSAsyncTransport",
        lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
    )
    monkeypatch.setattr(
        "nanobot.security.network.httpx.AsyncHTTPTransport",
        lambda **_kwargs: httpx.MockTransport(
            lambda request: httpx.Response(200, request=request)
        ),
    )

    kwargs = mcp_mod._pinned_transport_kwargs()

    assert "transport" in kwargs
    assert any(transport is None for transport in kwargs["mounts"].values())


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("config", "expected_transport"),
    [
        (MCPServerConfig(type="sse", url="https://mcp.example.com/sse"), "sse"),
        (
            MCPServerConfig(type="streamableHttp", url="https://mcp.example.com/mcp"),
            "streamableHttp",
        ),
    ],
)
async def test_connect_mcp_servers_http_clients_reject_unsafe_redirect_targets(
    config: MCPServerConfig,
    expected_transport: str,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    checked_urls: list[str] = []
    sent_urls: list[str] = []
    used_transports: list[str] = []

    async def _validate(url: str, **_kwargs: object) -> tuple[bool, str]:
        checked_urls.append(url)
        if url == "http://127.0.0.1/private":
            return False, "loopback blocked"
        return True, ""

    async def _reachable(_url: str) -> bool:
        return True

    def _handler(request: httpx.Request) -> httpx.Response:
        sent_urls.append(str(request.url))
        if str(request.url) == "https://example.com/start":
            return httpx.Response(
                302,
                headers={"Location": "http://127.0.0.1/private"},
                request=request,
            )
        raise AssertionError("unsafe redirect target should be blocked before transport")

    original_async_client = httpx.AsyncClient

    def _async_client_with_mock_transport(*args: object, **kwargs: object) -> httpx.AsyncClient:
        kwargs.setdefault("transport", httpx.MockTransport(_handler))
        return original_async_client(*args, **kwargs)

    @asynccontextmanager
    async def _fake_sse_client(_url: str, httpx_client_factory=None):
        assert httpx_client_factory is not None
        used_transports.append("sse")
        async with httpx_client_factory() as client:
            await client.get("https://example.com/start")
        yield object(), object()

    @asynccontextmanager
    async def _fake_streamable_http_client(_url: str, http_client=None):
        assert http_client is not None
        used_transports.append("streamableHttp")
        await http_client.get("https://example.com/start")
        yield object(), object(), object()

    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
    monkeypatch.setattr(
        mcp_mod,
        "PinnedDNSAsyncTransport",
        lambda **_kwargs: httpx.MockTransport(_handler),
    )
    monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", _async_client_with_mock_transport)
    monkeypatch.setattr(sys.modules["mcp.client.sse"], "sse_client", _fake_sse_client)
    monkeypatch.setattr(
        sys.modules["mcp.client.streamable_http"],
        "streamable_http_client",
        _fake_streamable_http_client,
    )

    registry = ToolRegistry()
    stacks = await connect_mcp_servers({"remote": config}, registry)

    assert stacks == {}
    assert registry.tool_names == []
    assert used_transports == [expected_transport]
    assert checked_urls == [
        config.url,
        "https://example.com/start",
        "http://127.0.0.1/private",
    ]
    assert sent_urls == ["https://example.com/start"]


@pytest.mark.asyncio
@pytest.mark.parametrize("failure_mode", ["exception", "cancellation"])
async def test_connect_mcp_servers_one_failure_does_not_block_others(
    monkeypatch: pytest.MonkeyPatch,
    failure_mode: str,
) -> None:
    bad_session = _make_fake_session([])

    async def _cancel_initialize() -> None:
        raise asyncio.CancelledError("cancelled by SDK")

    if failure_mode == "cancellation":
        bad_session.initialize = _cancel_initialize
    sessions = {
        "bad": bad_session,
        "good": _make_fake_session(["demo"]),
    }

    class _SelectiveClientSession:
        def __init__(self, read: object, _write: object) -> None:
            self._session = sessions[read]

        async def __aenter__(self) -> object:
            return self._session

        async def __aexit__(self, exc_type, exc, tb) -> bool:
            return False

    @asynccontextmanager
    async def _selective_stdio_client(params: object):
        if params.command == "bad" and failure_mode == "exception":
            raise RuntimeError("boom")
        yield params.command, object()

    monkeypatch.setattr(sys.modules["mcp"], "ClientSession", _SelectiveClientSession)
    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _selective_stdio_client)

    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {
            "bad": MCPServerConfig(command="bad"),
            "good": MCPServerConfig(command="good"),
        },
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == ["mcp_good_demo"]
    assert set(stacks) == {"good"}


@pytest.mark.asyncio
async def test_connect_mcp_servers_propagates_external_cancellation(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    started = asyncio.Event()
    closed = asyncio.Event()

    @asynccontextmanager
    async def _blocking_stdio_client(_params: object):
        try:
            started.set()
            await asyncio.Event().wait()
            yield object(), object()
        finally:
            closed.set()

    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _blocking_stdio_client)

    task = asyncio.create_task(
        connect_mcp_servers({"slow": MCPServerConfig(command="slow")}, ToolRegistry())
    )
    await asyncio.wait_for(started.wait(), timeout=1.0)
    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task

    await asyncio.wait_for(closed.wait(), timeout=1.0)


@pytest.mark.asyncio
async def test_connect_mcp_servers_rolls_back_completed_batch_on_cancellation(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    slow_started = asyncio.Event()
    closed: list[str] = []
    sessions = {"fast": _make_fake_session(["demo"])}

    class _SelectiveClientSession:
        def __init__(self, read: object, _write: object) -> None:
            self._session = sessions[str(read)]

        async def __aenter__(self) -> object:
            return self._session

        async def __aexit__(self, exc_type, exc, tb) -> bool:
            return False

    @asynccontextmanager
    async def _selective_stdio_client(params: object):
        command = str(params.command)
        try:
            if command == "slow":
                slow_started.set()
                await asyncio.Event().wait()
            yield command, object()
        finally:
            closed.append(command)

    monkeypatch.setattr(sys.modules["mcp"], "ClientSession", _SelectiveClientSession)
    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _selective_stdio_client)

    registry = ToolRegistry()
    task = asyncio.create_task(
        connect_mcp_servers(
            {
                "fast": MCPServerConfig(command="fast"),
                "slow": MCPServerConfig(command="slow"),
            },
            registry,
        )
    )
    await asyncio.wait_for(slow_started.wait(), timeout=1.0)
    assert registry.tool_names == ["mcp_fast_demo"]

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    assert registry.tool_names == []
    assert sorted(closed) == ["fast", "slow"]


@pytest.mark.asyncio
async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    captured: dict[str, object] = {}

    async def _reachable(_url: str) -> bool:
        return True

    async def _validate(_url: str) -> tuple[bool, str]:
        return True, ""

    @asynccontextmanager
    async def _capturing_streamable_http_client(_url: str, http_client=None):
        captured["timeout"] = http_client.timeout
        yield object(), object(), object()

    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
    monkeypatch.setattr(
        mcp_mod,
        "PinnedDNSAsyncTransport",
        lambda: httpx.MockTransport(lambda request: httpx.Response(200, request=request)),
    )
    monkeypatch.setattr(
        sys.modules["mcp.client.streamable_http"],
        "streamable_http_client",
        _capturing_streamable_http_client,
    )

    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(url="https://mcp.example.com/mcp")},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    timeout = captured["timeout"]
    assert timeout.connect == 10.0
    assert timeout.read == 30.0
    assert timeout.write == 30.0
    assert timeout.pool == 30.0


@pytest.mark.parametrize("transport", ["sse", "streamableHttp"])
@pytest.mark.asyncio
async def test_connect_mcp_servers_attaches_oauth_to_remote_http_client(
    transport: str,
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    oauth_auth = object()
    oauth_handlers = object()
    captured: dict[str, object] = {}

    async def _reachable(_url: str) -> bool:
        return True

    async def _validate(_url: str) -> tuple[bool, str]:
        return True, ""

    async def _create_auth(name: str, url: str, handlers: object) -> object:
        captured.update(name=name, url=url, handlers=handlers)
        return oauth_auth

    oauth_mod = ModuleType("nanobot.agent.tools.mcp_oauth")
    oauth_mod.MCPAuthorizationRequiredError = RuntimeError  # type: ignore[attr-defined]
    oauth_mod.create_mcp_oauth_auth = _create_auth  # type: ignore[attr-defined]
    monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)

    class FakeAsyncClient:
        def __init__(self, *args: object, **kwargs: object) -> None:
            captured["client_kwargs"] = kwargs

        async def __aenter__(self) -> object:
            return self

        async def __aexit__(self, exc_type: object, exc: object, tb: object) -> bool:
            return False

    @asynccontextmanager
    async def _capturing_sse_client(
        _url: str,
        httpx_client_factory=None,
        auth=None,
    ):
        captured["transport_auth"] = auth
        yield object(), object()

    @asynccontextmanager
    async def _capturing_streamable_http_client(_url: str, http_client=None):
        assert http_client is not None
        yield object(), object(), object()

    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _validate)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", _reachable)
    monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", FakeAsyncClient)
    monkeypatch.setattr(sys.modules["mcp.client.sse"], "sse_client", _capturing_sse_client)
    monkeypatch.setattr(
        sys.modules["mcp.client.streamable_http"],
        "streamable_http_client",
        _capturing_streamable_http_client,
    )

    url = "https://mcp.example.com/sse" if transport == "sse" else "https://mcp.example.com/mcp"
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"remote": MCPServerConfig(type=transport, url=url, auth="oauth")},
        registry,
        oauth_handlers={"remote": oauth_handlers},  # type: ignore[arg-type]
    )
    for stack in stacks.values():
        await stack.aclose()

    assert captured["name"] == "remote"
    assert captured["url"] == url
    assert captured["handlers"] is oauth_handlers
    if transport == "sse":
        assert captured["transport_auth"] is oauth_auth
    else:
        client_kwargs = captured["client_kwargs"]
        assert isinstance(client_kwargs, dict)
        assert client_kwargs["auth"] is oauth_auth
        assert client_kwargs["event_hooks"] == {"request": [mcp_mod._validate_mcp_request_url]}


@pytest.mark.asyncio
async def test_connect_mcp_servers_skips_background_oauth_without_credentials(
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    class AuthorizationRequiredError(RuntimeError):
        pass

    async def _create_auth(*_args: object) -> object:
        raise AuthorizationRequiredError

    probe_called = False

    async def _probe(_url: str) -> bool:
        nonlocal probe_called
        probe_called = True
        return True

    oauth_mod = ModuleType("nanobot.agent.tools.mcp_oauth")
    oauth_mod.MCPAuthorizationRequiredError = AuthorizationRequiredError  # type: ignore[attr-defined]
    oauth_mod.create_mcp_oauth_auth = _create_auth  # type: ignore[attr-defined]
    monkeypatch.setitem(sys.modules, "nanobot.agent.tools.mcp_oauth", oauth_mod)
    monkeypatch.setattr(mcp_mod, "async_validate_url_target", _allow_url)
    monkeypatch.setattr(mcp_mod, "_probe_http_url", _probe)

    stacks = await connect_mcp_servers(
        {
            "remote": MCPServerConfig(
                type="streamableHttp",
                url="https://mcp.example.com/mcp",
                auth="oauth",
            )
        },
        ToolRegistry(),
    )

    assert stacks == {}
    assert not probe_called


@pytest.mark.asyncio
async def test_connect_mcp_servers_wraps_windows_stdio_launchers(
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    captured: dict[str, object] = {}

    @asynccontextmanager
    async def _capturing_stdio_client(params: object):
        captured["command"] = params.command
        captured["args"] = params.args
        captured["env"] = params.env
        yield object(), object()

    monkeypatch.setattr(mcp_mod.os, "name", "nt", raising=False)
    monkeypatch.setattr(
        mcp_mod.shutil,
        "which",
        lambda command, path=None: r"C:\Program Files\nodejs\npx.cmd",
    )
    monkeypatch.setenv("COMSPEC", r"C:\Windows\System32\cmd.exe")
    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _capturing_stdio_client)

    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {
            "test": MCPServerConfig(
                command="npx",
                args=["-y", "chrome-devtools-mcp@latest"],
            )
        },
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert captured["command"] == r"C:\Windows\System32\cmd.exe"
    assert captured["args"] == ["/d", "/c", "npx", "-y", "chrome-devtools-mcp@latest"]
    assert captured["env"] is None


@pytest.mark.asyncio
async def test_connect_mcp_servers_passes_stdio_cwd(
    fake_mcp_runtime: dict[str, object | None],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session(["demo"])
    captured: dict[str, object] = {}

    @asynccontextmanager
    async def _capturing_stdio_client(params: object):
        captured["cwd"] = params.cwd
        yield object(), object()

    monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _capturing_stdio_client)

    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", cwd="/tmp/nanobot-mcp-test")},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert captured["cwd"] == "/tmp/nanobot-mcp-test"


# ---------------------------------------------------------------------------
# MCPResourceWrapper tests
# ---------------------------------------------------------------------------


def _make_resource_def(
    name: str = "myres",
    uri: str = "file:///tmp/data.txt",
    description: str = "A test resource",
) -> SimpleNamespace:
    return SimpleNamespace(name=name, uri=uri, description=description)


def _make_resource_wrapper(session: object, *, timeout: float = 0.1) -> MCPResourceWrapper:
    return MCPResourceWrapper(session, "srv", _make_resource_def(), resource_timeout=timeout)


def test_resource_wrapper_properties() -> None:
    wrapper = MCPResourceWrapper(None, "myserver", _make_resource_def())
    assert wrapper.name == "mcp_myserver_resource_myres"
    assert "[MCP Resource]" in wrapper.description
    assert "A test resource" in wrapper.description
    assert "file:///tmp/data.txt" in wrapper.description
    assert wrapper.parameters == {"type": "object", "properties": {}, "required": []}
    assert wrapper.read_only is True


@pytest.mark.asyncio
async def test_resource_wrapper_execute_returns_text() -> None:
    async def read_resource(uri: str) -> object:
        assert uri == "file:///tmp/data.txt"
        return SimpleNamespace(
            contents=[_FakeTextResourceContents("line1"), _FakeTextResourceContents("line2")]
        )

    wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource))
    result = await wrapper.execute()
    assert result == "line1\nline2"


@pytest.mark.asyncio
async def test_resource_wrapper_execute_handles_blob() -> None:
    async def read_resource(uri: str) -> object:
        return SimpleNamespace(contents=[_FakeBlobResourceContents(b"\x00\x01\x02")])

    wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource))
    result = await wrapper.execute()
    assert "[Binary resource: 3 bytes]" in result


@pytest.mark.asyncio
async def test_resource_wrapper_execute_handles_timeout() -> None:
    async def read_resource(uri: str) -> object:
        await asyncio.sleep(1)
        return SimpleNamespace(contents=[])

    wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource), timeout=0.01)
    result = await wrapper.execute()
    assert result == "(MCP resource read timed out after 0.01s)"


@pytest.mark.asyncio
async def test_resource_wrapper_execute_handles_error() -> None:
    async def read_resource(uri: str) -> object:
        raise RuntimeError("boom")

    wrapper = _make_resource_wrapper(SimpleNamespace(read_resource=read_resource))
    result = await wrapper.execute()
    assert result == "(MCP resource read failed: RuntimeError)"


# ---------------------------------------------------------------------------
# MCPPromptWrapper tests
# ---------------------------------------------------------------------------


def _make_prompt_def(
    name: str = "myprompt",
    description: str = "A test prompt",
    arguments: list | None = None,
) -> SimpleNamespace:
    return SimpleNamespace(name=name, description=description, arguments=arguments)


def _make_prompt_wrapper(session: object, *, timeout: float = 0.1) -> MCPPromptWrapper:
    return MCPPromptWrapper(session, "srv", _make_prompt_def(), prompt_timeout=timeout)


def test_prompt_wrapper_properties() -> None:
    arg1 = SimpleNamespace(name="topic", required=True)
    arg2 = SimpleNamespace(name="style", required=False)
    wrapper = MCPPromptWrapper(None, "myserver", _make_prompt_def(arguments=[arg1, arg2]))
    assert wrapper.name == "mcp_myserver_prompt_myprompt"
    assert "[MCP Prompt]" in wrapper.description
    assert "A test prompt" in wrapper.description
    assert "workflow guide" in wrapper.description
    assert wrapper.parameters["properties"]["topic"] == {"type": "string"}
    assert wrapper.parameters["properties"]["style"] == {"type": "string"}
    assert wrapper.parameters["required"] == ["topic"]
    assert wrapper.read_only is True


def test_prompt_wrapper_no_arguments() -> None:
    wrapper = MCPPromptWrapper(None, "myserver", _make_prompt_def())
    assert wrapper.parameters == {"type": "object", "properties": {}, "required": []}


def test_prompt_wrapper_preserves_argument_descriptions() -> None:
    arg = SimpleNamespace(name="topic", required=True, description="The subject to discuss")
    wrapper = MCPPromptWrapper(None, "srv", _make_prompt_def(arguments=[arg]))
    assert wrapper.parameters["properties"]["topic"] == {
        "type": "string",
        "description": "The subject to discuss",
    }


@pytest.mark.asyncio
async def test_prompt_wrapper_execute_returns_text() -> None:
    async def get_prompt(name: str, arguments: dict | None = None) -> object:
        assert name == "myprompt"
        msg1 = SimpleNamespace(
            role="user",
            content=[_FakeTextContent("You are an expert on {{topic}}.")],
        )
        msg2 = SimpleNamespace(
            role="assistant",
            content=[_FakeTextContent("Understood. Ask me anything.")],
        )
        return SimpleNamespace(messages=[msg1, msg2])

    wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt))
    result = await wrapper.execute(topic="AI")
    assert "You are an expert on {{topic}}." in result
    assert "Understood. Ask me anything." in result


@pytest.mark.asyncio
async def test_prompt_wrapper_execute_handles_timeout() -> None:
    async def get_prompt(name: str, arguments: dict | None = None) -> object:
        await asyncio.sleep(1)
        return SimpleNamespace(messages=[])

    wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt), timeout=0.01)
    result = await wrapper.execute()
    assert result == "(MCP prompt call timed out after 0.01s)"


@pytest.mark.asyncio
async def test_prompt_wrapper_execute_handles_mcp_error() -> None:
    from mcp.shared.exceptions import McpError

    async def get_prompt(name: str, arguments: dict | None = None) -> object:
        raise McpError(code=42, message="invalid argument")

    wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt))
    result = await wrapper.execute()
    assert "invalid argument" in result
    assert "code 42" in result


@pytest.mark.asyncio
async def test_prompt_wrapper_execute_handles_error() -> None:
    async def get_prompt(name: str, arguments: dict | None = None) -> object:
        raise RuntimeError("boom")

    wrapper = _make_prompt_wrapper(SimpleNamespace(get_prompt=get_prompt))
    result = await wrapper.execute()
    assert result == "(MCP prompt call failed: RuntimeError)"


# ---------------------------------------------------------------------------
# connect_mcp_servers: resources + prompts integration
# ---------------------------------------------------------------------------


def _make_fake_session_with_capabilities(
    tool_names: list[str],
    resource_names: list[str] | None = None,
    prompt_names: list[str] | None = None,
) -> SimpleNamespace:
    async def initialize() -> None:
        return None

    async def list_tools() -> SimpleNamespace:
        return SimpleNamespace(tools=[_make_tool_def(name) for name in tool_names])

    async def list_resources() -> SimpleNamespace:
        resources = []
        for rname in resource_names or []:
            resources.append(
                SimpleNamespace(
                    name=rname,
                    uri=f"file:///{rname}",
                    description=f"{rname} resource",
                )
            )
        return SimpleNamespace(resources=resources)

    async def list_prompts() -> SimpleNamespace:
        prompts = []
        for pname in prompt_names or []:
            prompts.append(
                SimpleNamespace(
                    name=pname,
                    description=f"{pname} prompt",
                    arguments=None,
                )
            )
        return SimpleNamespace(prompts=prompts)

    return SimpleNamespace(
        initialize=initialize,
        list_tools=list_tools,
        list_resources=list_resources,
        list_prompts=list_prompts,
    )


@pytest.mark.asyncio
async def test_connect_registers_resources_and_prompts(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=["tool_a"],
        resource_names=["res_b"],
        prompt_names=["prompt_c"],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake")},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert "mcp_test_tool_a" in registry.tool_names
    assert "mcp_test_resource_res_b" in registry.tool_names
    assert "mcp_test_prompt_prompt_c" in registry.tool_names


# ---------------------------------------------------------------------------
# _sanitize_name tests
# ---------------------------------------------------------------------------


def test_sanitize_name_replaces_spaces() -> None:
    assert _sanitize_name("PostgreSQL System Information") == "PostgreSQL_System_Information"


def test_sanitize_name_replaces_special_characters() -> None:
    assert _sanitize_name("foo.bar@baz!") == "foo_bar_baz_"


def test_sanitize_name_collapses_consecutive_underscores() -> None:
    assert _sanitize_name("a   b") == "a_b"


def test_sanitize_name_preserves_valid_characters() -> None:
    assert _sanitize_name("my-tool_v2") == "my-tool_v2"


def test_sanitize_name_noop_for_already_clean_names() -> None:
    assert _sanitize_name("mcp_server_tool") == "mcp_server_tool"


# ---------------------------------------------------------------------------
# Wrapper sanitization tests
# ---------------------------------------------------------------------------


def test_tool_wrapper_sanitizes_name() -> None:
    tool_def = SimpleNamespace(
        name="My Tool",
        description="tool with spaces",
        inputSchema={"type": "object", "properties": {}},
    )
    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
    assert wrapper.name == "mcp_srv_My_Tool"


def test_resource_wrapper_sanitizes_name() -> None:
    resource_def = SimpleNamespace(
        name="PostgreSQL System Information",
        uri="file:///pg/info",
        description="PG info",
    )
    wrapper = MCPResourceWrapper(None, "srv", resource_def)
    assert wrapper.name == "mcp_srv_resource_PostgreSQL_System_Information"


def test_prompt_wrapper_sanitizes_name() -> None:
    prompt_def = SimpleNamespace(
        name="design-schema",
        description="Design schema",
        arguments=None,
    )
    # Hyphens are allowed, so this should pass through unchanged
    wrapper = MCPPromptWrapper(None, "my server", prompt_def)
    assert wrapper.name == "mcp_my_server_prompt_design-schema"


def test_tool_wrapper_preserves_original_name_for_mcp_call() -> None:
    tool_def = SimpleNamespace(
        name="My Tool",
        description="tool with spaces",
        inputSchema={"type": "object", "properties": {}},
    )
    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "srv", tool_def)
    # The sanitized API-facing name differs from the original MCP name
    assert wrapper.name == "mcp_srv_My_Tool"
    assert wrapper._original_name == "My Tool"


@pytest.mark.asyncio
async def test_connect_mcp_servers_sanitizes_resource_names(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=[],
        resource_names=["PostgreSQL System Information"],
        prompt_names=[],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake")},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert "mcp_test_resource_PostgreSQL_System_Information" in registry.tool_names


@pytest.mark.asyncio
async def test_connect_mcp_servers_enabled_tools_matches_sanitized_name(
    fake_mcp_runtime: dict[str, object | None],
) -> None:
    fake_mcp_runtime["session"] = _make_fake_session_with_capabilities(
        tool_names=["My Tool", "other"],
    )
    registry = ToolRegistry()
    stacks = await connect_mcp_servers(
        {"test": MCPServerConfig(command="fake", enabled_tools=["mcp_test_My_Tool"])},
        registry,
    )
    for stack in stacks.values():
        await stack.aclose()

    assert registry.tool_names == ["mcp_test_My_Tool"]


@pytest.mark.parametrize(
    "url, expected",
    [
        ("https://user:secret@host.example/sse", "https://host.example/..."),
        ("https://host.example:8443/mcp?token=abc#frag", "https://host.example:8443/..."),
        ("https://user:secret@[::1]:8443/sse?token=abc", "https://[::1]:8443/..."),
        ("https://host.example/sse", "https://host.example/..."),
        ("https://host.example", "https://host.example"),
        ("https://host.example/", "https://host.example/"),
    ],
)
def test_redact_url_strips_credentials_and_query(url: str, expected: str) -> None:
    assert mcp_mod._redact_url(url) == expected

def test_mcp_tool_name_keeps_short_name():
    name = _sanitize_mcp_tool_name("mcp_myserver_resource_myres")
    assert name == "mcp_myserver_resource_myres"


def test_mcp_tool_name_limits_long_name():
    long_name = "mcp_" + "a" * 100
    name = _sanitize_mcp_tool_name(long_name)

    assert len(name) <= 64
    assert name.startswith("mcp_")


def test_long_server_name_tools_are_matched_by_server_name() -> None:
    server_name = "very-long-server-name-" * 4
    tool_def = SimpleNamespace(
        name="search",
        description="search tool",
        inputSchema={"type": "object", "properties": {}},
    )
    other_tool_def = SimpleNamespace(
        name="search",
        description="other search tool",
        inputSchema={"type": "object", "properties": {}},
    )
    wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), server_name, tool_def)
    other_wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "other", other_tool_def)
    registry = ToolRegistry()
    registry.register(wrapper)
    registry.register(other_wrapper)

    assert len(wrapper.name) == 64
    assert not wrapper.name.startswith(mcp_mod._tool_prefix(server_name))

    provider = MCPProvider(
        {server_name: MCPServerConfig(command="fake")},
        registry,
    )
    provider._attach_reconnect_handlers({server_name})
    assert wrapper._reconnect is not None
    assert other_wrapper._reconnect is None

    removed = mcp_mod._unregister_server_tools(registry, server_name)

    assert removed == 1
    assert wrapper.name not in registry.tool_names
    assert other_wrapper.name in registry.tool_names


async def _allow_url(
    _url: str,
    *,
    allow_loopback: bool = False,
) -> tuple[bool, str]:
    del allow_loopback
    return True, ""