turnstone/tests/test_mcp_client.py

"""Tests for turnstone.core.mcp_client — MCP client manager and config loading."""

from __future__ import annotations

import asyncio
import concurrent.futures
import json
import time
from contextlib import AsyncExitStack
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from tests.conftest import _seed_static_state
from turnstone.core.mcp_client import (
    MCPClientManager,
    _db_servers_to_config,
    _mcp_to_openai,
    load_mcp_config,
)
from turnstone.core.tools import INTERACTIVE_TOOLS, TOOLS, merge_mcp_tools

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _fake_mcp_tool(name: str = "search", description: str = "Search stuff") -> MagicMock:
    """Create a mock MCP tool object matching the SDK's Tool type."""
    tool = MagicMock()
    tool.name = name
    tool.description = description
    tool.inputSchema = {
        "type": "object",
        "properties": {"query": {"type": "string"}},
        "required": ["query"],
    }
    return tool


def _fake_openai_tool(name: str = "mcp__test__search") -> dict[str, Any]:
    """Create a fake OpenAI-format tool dict."""
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": "[MCP: test] Search stuff",
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string"}},
                "required": ["query"],
            },
        },
    }


def _fake_mcp_resource(
    uri: str = "file:///README.md",
    name: str = "readme",
    description: str = "Project readme",
    mime_type: str = "text/plain",
) -> MagicMock:
    """Create a mock MCP Resource object matching the SDK's Resource type."""
    res = MagicMock()
    res.uri = uri
    res.name = name
    res.description = description
    res.mimeType = mime_type
    return res


def _fake_resource_dict(
    uri: str = "file:///README.md",
    name: str = "readme",
    description: str = "Project readme",
    mime_type: str = "text/plain",
    server: str = "test",
) -> dict[str, Any]:
    """Create a fake resource dict as stored in per-server state."""
    return {
        "uri": uri,
        "name": name,
        "description": description,
        "mimeType": mime_type,
        "server": server,
    }


def _fake_mcp_prompt(
    name: str = "code_review",
    description: str = "Generate a code review",
    arguments: list[dict[str, Any]] | None = None,
) -> MagicMock:
    """Create a mock MCP Prompt object matching the SDK's Prompt type."""
    prompt = MagicMock()
    prompt.name = name
    prompt.description = description
    if arguments is None:
        arg = MagicMock()
        arg.name = "language"
        arg.description = "Programming language"
        arg.required = True
        prompt.arguments = [arg]
    else:
        mock_args = []
        for a in arguments:
            arg = MagicMock()
            arg.name = a["name"]
            arg.description = a.get("description", "")
            arg.required = a.get("required", False)
            mock_args.append(arg)
        prompt.arguments = mock_args
    return prompt


def _fake_prompt_dict(
    name: str = "mcp__test__code_review",
    original_name: str = "code_review",
    server: str = "test",
    description: str = "Generate a code review",
) -> dict[str, Any]:
    """Create a fake prompt dict as stored in per-server state."""
    return {
        "name": name,
        "original_name": original_name,
        "server": server,
        "description": description,
        "arguments": [
            {"name": "language", "description": "Programming language", "required": True}
        ],
    }


@pytest.fixture
def running_loop_mgr():
    """Yield a (mgr, loop, thread) triple with a background loop already running.

    Spawns an MCPClientManager with a default config of {"srv": stdio echo}
    and a fresh asyncio loop driven by a daemon thread.  Tests that need a
    different config can mutate ``mgr._server_configs`` directly.  The
    fixture stops the loop and joins the thread on teardown so each test
    leaves a clean slate.
    """
    import threading as _threading

    cfg = {"srv": {"type": "stdio", "command": "echo"}}
    mgr = MCPClientManager(cfg)
    loop = asyncio.new_event_loop()
    thread = _threading.Thread(target=loop.run_forever, daemon=True)
    thread.start()
    mgr._loop = loop
    try:
        yield mgr, loop, thread
    finally:
        loop.call_soon_threadsafe(loop.stop)
        thread.join(timeout=2)


# ---------------------------------------------------------------------------
# Schema conversion
# ---------------------------------------------------------------------------


class TestMcpToOpenai:
    def test_basic_conversion(self):
        tool = _fake_mcp_tool("search_repos", "Search GitHub repos")
        result = _mcp_to_openai("github", tool)

        assert result["type"] == "function"
        func = result["function"]
        assert func["name"] == "mcp__github__search_repos"
        assert func["description"] == "Search GitHub repos"
        assert func["parameters"]["type"] == "object"
        assert "query" in func["parameters"]["properties"]

    def test_name_prefixing(self):
        tool = _fake_mcp_tool("list_files")
        result = _mcp_to_openai("fs", tool)
        assert result["function"]["name"] == "mcp__fs__list_files"

    def test_missing_input_schema(self):
        tool = MagicMock()
        tool.name = "ping"
        tool.description = "Ping the server"
        tool.inputSchema = None
        result = _mcp_to_openai("test", tool)
        assert result["function"]["parameters"] == {"type": "object", "properties": {}}

    def test_empty_description(self):
        tool = MagicMock()
        tool.name = "noop"
        tool.description = ""
        tool.inputSchema = {"type": "object", "properties": {}}
        result = _mcp_to_openai("test", tool)
        assert result["function"]["description"] == ""


# ---------------------------------------------------------------------------
# Config loading
# ---------------------------------------------------------------------------


class TestLoadMcpConfig:
    def test_load_from_json_file(self, tmp_path):
        config_file = tmp_path / "mcp.json"
        config_file.write_text(
            json.dumps(
                {
                    "mcpServers": {
                        "github": {
                            "command": "npx",
                            "args": ["-y", "@modelcontextprotocol/server-github"],
                            "env": {"GITHUB_TOKEN": "test"},
                        }
                    }
                }
            )
        )
        result = load_mcp_config(str(config_file))
        assert "github" in result
        assert result["github"]["command"] == "npx"
        assert result["github"]["env"]["GITHUB_TOKEN"] == "test"

    def test_load_from_toml(self):
        mock_config = {
            "servers": {
                "postgres": {
                    "type": "http",
                    "url": "https://mcp.example.com/mcp",
                }
            }
        }
        with patch("turnstone.core.mcp_client.load_config", return_value=mock_config):
            result = load_mcp_config(None)
        assert "postgres" in result
        assert result["postgres"]["url"] == "https://mcp.example.com/mcp"

    def test_empty_when_no_config(self):
        with patch("turnstone.core.mcp_client.load_config", return_value={}):
            result = load_mcp_config(None)
        assert result == {}

    def test_json_file_not_found(self, tmp_path):
        with patch("turnstone.core.mcp_client.load_config", return_value={}):
            result = load_mcp_config(str(tmp_path / "nonexistent.json"))
        assert result == {}

    def test_toml_config_path_redirect(self):
        """TOML [mcp] config_path redirects to JSON file."""
        # load_config returns a section with config_path pointing to a nonexistent file
        mock_config = {"config_path": "/tmp/nonexistent_mcp.json"}
        with patch("turnstone.core.mcp_client.load_config", return_value=mock_config):
            result = load_mcp_config(None)
        assert result == {}

    def test_invalid_json(self, tmp_path):
        config_file = tmp_path / "bad.json"
        config_file.write_text("not json")
        with patch("turnstone.core.mcp_client.load_config", return_value={}):
            result = load_mcp_config(str(config_file))
        assert result == {}


class TestDBServersToConfig:
    """``_db_servers_to_config`` shapes DB rows for the MCP client."""

    def test_static_streamable_http_row_passes_through(self) -> None:
        rows = [
            {
                "name": "static-srv",
                "transport": "streamable-http",
                "url": "https://mcp.example.com",
                "headers": '{"Authorization": "Bearer token"}',
                "auth_type": "static",
            }
        ]
        result = _db_servers_to_config(rows)
        assert "static-srv" in result
        assert result["static-srv"]["url"] == "https://mcp.example.com"
        assert result["static-srv"]["headers"] == {"Authorization": "Bearer token"}

    def test_db_servers_to_config_skips_oauth_user_rows(self) -> None:
        """Rows with auth_type=oauth_user must be invisible to the static
        auto-connect path.

        Auto-connecting these with empty headers fails the AS check and
        trips the circuit breaker on startup. Per-user OAuth servers
        come online lazily once the user has consented.
        """
        rows = [
            {
                "name": "static-srv",
                "transport": "streamable-http",
                "url": "https://static.example.com",
                "headers": "{}",
                "auth_type": "static",
            },
            {
                "name": "oauth-srv",
                "transport": "streamable-http",
                "url": "https://oauth.example.com",
                "headers": "{}",
                "auth_type": "oauth_user",
            },
            {
                "name": "stdio-srv",
                "transport": "stdio",
                "command": "echo",
                "args": "[]",
                "env": "{}",
                "auth_type": "none",
            },
        ]
        result = _db_servers_to_config(rows)
        assert set(result) == {"static-srv", "stdio-srv"}
        assert "oauth-srv" not in result


# ---------------------------------------------------------------------------
# merge_mcp_tools
# ---------------------------------------------------------------------------


class TestMergeTools:
    def test_merge_preserves_builtin(self):
        mcp_tools = [_fake_openai_tool()]
        merged = merge_mcp_tools(TOOLS, mcp_tools)
        # First N should be built-in
        for i, t in enumerate(TOOLS):
            assert merged[i] is t

    def test_merge_appends_mcp(self):
        mcp_tools = [_fake_openai_tool("mcp__a__x"), _fake_openai_tool("mcp__b__y")]
        merged = merge_mcp_tools(TOOLS, mcp_tools)
        assert len(merged) == len(TOOLS) + 2
        assert merged[-2]["function"]["name"] == "mcp__a__x"
        assert merged[-1]["function"]["name"] == "mcp__b__y"

    def test_merge_empty_mcp(self):
        merged = merge_mcp_tools(TOOLS, [])
        assert merged == TOOLS

    def test_merge_does_not_mutate_input(self):
        mcp_tools = [_fake_openai_tool()]
        original_len = len(TOOLS)
        merge_mcp_tools(TOOLS, mcp_tools)
        assert len(TOOLS) == original_len


# ---------------------------------------------------------------------------
# MCPClientManager unit tests (no real MCP servers)
# ---------------------------------------------------------------------------


class TestMCPClientManager:
    def test_init_state(self):
        mgr = MCPClientManager({"test": {"command": "echo"}})
        assert mgr.get_tools() == []
        assert mgr.is_mcp_tool("anything") is False
        assert mgr.server_count == 0

    def test_get_tools_returns_copy(self):
        mgr = MCPClientManager({})
        mgr._tools = [_fake_openai_tool()]
        tools = mgr.get_tools()
        assert len(tools) == 1
        tools.clear()  # mutate the copy
        assert len(mgr.get_tools()) == 1  # original unchanged

    def test_is_mcp_tool(self):
        mgr = MCPClientManager({})
        mgr._tool_map["mcp__gh__search"] = ("gh", "search")
        assert mgr.is_mcp_tool("mcp__gh__search") is True
        assert mgr.is_mcp_tool("bash") is False

    def test_server_count(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", session=MagicMock())
        _seed_static_state(mgr, "b", session=MagicMock())
        assert mgr.server_count == 2

    def test_call_tool_sync_unknown_tool(self):
        mgr = MCPClientManager({})
        with pytest.raises(ValueError, match="Unknown MCP tool"):
            mgr.call_tool_sync("mcp__no__such", {})

    def test_call_tool_sync_disconnected_server(self):
        mgr = MCPClientManager({})
        mgr._tool_map["mcp__dead__ping"] = ("dead", "ping")
        # No session registered for "dead", no config/loop → reconnect fails
        with pytest.raises(RuntimeError, match="not connected"):
            mgr.call_tool_sync("mcp__dead__ping", {})

    def test_shutdown_on_unstarted_manager(self):
        """shutdown() should not raise when called on a manager that was never started."""
        mgr = MCPClientManager({})
        mgr.shutdown()  # should be a no-op

    # -- Phase 7: per-user catalog scoping ---------------------------------

    def test_is_mcp_tool_user_id_default_none_unchanged(self):
        """Sanity: default ``user_id=None`` answers static-only.

        The legacy single-arg call still works, and unknown names still
        return False — Phase 7 adds an optional keyword without
        rewriting the static-path semantics.
        """
        mgr = MCPClientManager({})
        mgr._tool_map["mcp__static__list"] = ("static", "list")
        # Legacy single-arg call still works.
        assert mgr.is_mcp_tool("mcp__static__list") is True
        assert mgr.is_mcp_tool("mcp__static__list", user_id=None) is True
        assert mgr.is_mcp_tool("nonexistent") is False
        assert mgr.is_mcp_tool("nonexistent", user_id=None) is False

    def test_is_mcp_tool_user_keyed_pool_tool(self):
        """A name visible only via ``_user_tool_map`` resolves only for
        the matching ``user_id``.

        Verifies the new branch: ``_tool_map`` miss + ``user_id`` hit.
        """
        mgr = MCPClientManager({})
        mgr._user_tool_map["user-1"] = {
            "mcp__pool-srv__do": ("pool-srv", "do"),
        }
        # Visible to user-1.
        assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-1") is True
        # Invisible to None caller (admin / web-search backend resolution).
        assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id=None) is False
        # Invisible to a different user.
        assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-2") is False

    def test_is_mcp_tool_static_wins_for_any_user(self):
        """Static-path tools are visible regardless of ``user_id`` —
        the merged view is ``static ∪ user-pool``."""
        mgr = MCPClientManager({})
        mgr._tool_map["mcp__static__list"] = ("static", "list")
        # Even for an unknown user, a static tool is still reachable —
        # static-path is process-global.
        assert mgr.is_mcp_tool("mcp__static__list", user_id="user-1") is True
        assert mgr.is_mcp_tool("mcp__static__list", user_id="anybody") is True

    def test_get_tools_user_id_none_returns_static_only(self):
        """``get_tools(user_id=None)`` returns the global static catalog.

        Pool tools are NEVER included in the default-arg view — that's
        the legacy contract every pre-Phase-7 caller relies on.
        """
        mgr = MCPClientManager({})
        mgr._tools = [_fake_openai_tool("mcp__static__list")]
        # Seed a pool entry that should NOT appear in the default view.
        from turnstone.core.mcp_client import PoolEntryState

        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry

        tools = mgr.get_tools()
        names = [t["function"]["name"] for t in tools]
        assert names == ["mcp__static__list"]

    def test_get_tools_user_id_merges_pool(self):
        """``get_tools(user_id='user-1')`` merges static + that user's pool tools.

        Other users' pool entries MUST NOT leak into the result —
        privacy / RBAC invariant.
        """
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        mgr._tools = [_fake_openai_tool("mcp__static__list")]
        e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        e1.tools = [_fake_openai_tool("mcp__pool-srv__do")]
        e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
        e2.tools = [_fake_openai_tool("mcp__pool-srv__other")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = e1
        mgr._user_pool_entries[("user-2", "pool-srv")] = e2
        # Production invariant: ``_connect_one_pool`` /
        # ``_refresh_pool_server_tools`` / ``_evict_session`` /
        # ``_close_pool_entry_if_idle`` all call ``_rebuild_user_tool_map``
        # immediately after mutating ``_user_pool_entries``. Tests that
        # seed pool entries directly must mirror that invariant —
        # ``get_tools(user_id=...)`` reads from the ``_user_tools``
        # snapshot (built by ``_rebuild_user_tool_map``), never iterating
        # ``_user_pool_entries`` directly.
        mgr._rebuild_user_tool_map("user-1")
        mgr._rebuild_user_tool_map("user-2")

        u1_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
        assert sorted(u1_names) == ["mcp__pool-srv__do", "mcp__static__list"]

        u2_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-2")]
        assert sorted(u2_names) == ["mcp__pool-srv__other", "mcp__static__list"]

        # Default still global-only — unaffected by either user's entries.
        default_names = [t["function"]["name"] for t in mgr.get_tools()]
        assert default_names == ["mcp__static__list"]

    def test_get_tools_user_id_returns_copies(self):
        """Mirror existing ``test_get_tools_returns_copy``: caller mutation
        of the returned list MUST NOT affect the manager's catalog."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        mgr._tools = [_fake_openai_tool("mcp__static__a")]
        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        entry.tools = [_fake_openai_tool("mcp__pool-srv__b")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry
        # See note in ``test_get_tools_user_id_merges_pool``.
        mgr._rebuild_user_tool_map("user-1")

        tools = mgr.get_tools(user_id="user-1")
        assert len(tools) == 2
        tools.clear()
        # Re-fetch — original catalog unchanged.
        assert len(mgr.get_tools(user_id="user-1")) == 2

    def test_get_tools_user_with_none_tools_skipped(self):
        """A pool entry that hasn't completed discovery (``entry.tools is None``)
        contributes no tools — the merged view skips it cleanly."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        mgr._tools = [_fake_openai_tool("mcp__static__a")]
        # Brand-new pool entry, discovery not yet run.
        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        assert entry.tools is None
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry
        # Rebuild observes ``entry.tools is None`` and skips this entry.
        mgr._rebuild_user_tool_map("user-1")

        names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
        assert names == ["mcp__static__a"]

    def test_rebuild_user_tool_map_populates(self):
        """``_rebuild_user_tool_map`` materializes the per-user index from
        pool entries owned by that user."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry

        mgr._rebuild_user_tool_map("user-1")
        assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__do": ("pool-srv", "do")}
        # Sibling _user_tools cache (bug-1 fix) MUST be populated alongside
        # the map — otherwise get_tools(user_id="user-1") would silently
        # return the static-only view despite is_mcp_tool returning True.
        assert mgr._user_tools["user-1"] == [_fake_openai_tool("mcp__pool-srv__do")]

    def test_rebuild_user_tool_map_drops_empty_user(self):
        """Rebuilding for a user with no pool entries removes the key
        rather than retaining an empty-dict sentinel."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry
        mgr._rebuild_user_tool_map("user-1")
        assert "user-1" in mgr._user_tool_map
        assert "user-1" in mgr._user_tools

        # Drop the entry, rebuild — user_id key should be removed from BOTH
        # the map and the sibling tool list (bug-1 fix). A drop in only one
        # would leave get_tools and is_mcp_tool out of sync.
        mgr._user_pool_entries.clear()
        mgr._rebuild_user_tool_map("user-1")
        assert "user-1" not in mgr._user_tool_map
        assert "user-1" not in mgr._user_tools

    def test_rebuild_user_tool_map_isolates_users(self):
        """Rebuilding for ``user-1`` MUST NOT touch ``user-2``'s entry."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        e1.tools = [_fake_openai_tool("mcp__pool-srv__one")]
        e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
        e2.tools = [_fake_openai_tool("mcp__pool-srv__two")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = e1
        mgr._user_pool_entries[("user-2", "pool-srv")] = e2

        mgr._rebuild_user_tool_map("user-1")
        mgr._rebuild_user_tool_map("user-2")
        assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__one": ("pool-srv", "one")}
        assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}

        # Clear user-1's entry only; rebuild user-1; user-2 must remain.
        mgr._user_pool_entries.pop(("user-1", "pool-srv"))
        mgr._rebuild_user_tool_map("user-1")
        assert "user-1" not in mgr._user_tool_map
        assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}

    def test_rebuild_user_tool_map_does_not_touch_static(self):
        """Invariant 1: per-user rebuild must NOT mutate ``_tool_map``."""
        from turnstone.core.mcp_client import PoolEntryState

        mgr = MCPClientManager({})
        mgr._tool_map["mcp__static__list"] = ("static", "list")
        entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
        entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
        mgr._user_pool_entries[("user-1", "pool-srv")] = entry

        before = dict(mgr._tool_map)
        mgr._rebuild_user_tool_map("user-1")
        assert mgr._tool_map == before


# ---------------------------------------------------------------------------
# Session integration (mock MCP client)
# ---------------------------------------------------------------------------


class TestSessionIntegration:
    @pytest.fixture()
    def tmp_db(self, tmp_path):
        from turnstone.core.storage import init_storage, reset_storage

        reset_storage()
        init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False)
        yield
        reset_storage()

    def _make_session(self, mcp_client=None, **kwargs):
        from turnstone.core.session import ChatSession

        defaults: dict[str, Any] = dict(
            client=MagicMock(),
            model="test-model",
            ui=MagicMock(),
            instructions=None,
            temperature=0.5,
            max_tokens=4096,
            tool_timeout=30,
            mcp_client=mcp_client,
        )
        defaults.update(kwargs)
        return ChatSession(**defaults)

    def test_session_without_mcp(self, tmp_db):
        session = self._make_session(mcp_client=None)
        # Interactive session surface — coordinator tools excluded.
        assert session._tools is INTERACTIVE_TOOLS
        assert session._mcp_client is None

    def test_session_with_mcp(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        session = self._make_session(mcp_client=mock_mcp)
        assert len(session._tools) == len(INTERACTIVE_TOOLS) + 1
        assert session._tools[-1]["function"]["name"] == "mcp__test__search"

    def test_task_tools_include_mcp(self, tmp_db):
        from turnstone.core.tools import TASK_AGENT_TOOLS

        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        session = self._make_session(mcp_client=mock_mcp)
        assert len(session._task_tools) == len(TASK_AGENT_TOOLS) + 1

    def test_agent_tools_include_mcp(self, tmp_db):
        from turnstone.core.tools import AGENT_TOOLS

        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        session = self._make_session(mcp_client=mock_mcp)
        assert len(session._agent_tools) == len(AGENT_TOOLS) + 1

    def test_prepare_mcp_tool(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.is_mcp_tool.return_value = True
        session = self._make_session(mcp_client=mock_mcp)

        tc = {
            "id": "call_123",
            "function": {
                "name": "mcp__test__search",
                "arguments": '{"query": "hello"}',
            },
        }
        prepared = session._prepare_tool(tc)
        assert prepared["func_name"] == "mcp__test__search"
        assert prepared["needs_approval"] is True
        assert "mcp:test/search" in prepared["header"]
        assert callable(prepared["execute"])

    def test_unknown_tool_without_mcp(self, tmp_db):
        session = self._make_session(mcp_client=None)
        tc = {
            "id": "call_456",
            "function": {"name": "nonexistent", "arguments": "{}"},
        }
        prepared = session._prepare_tool(tc)
        assert "error" in prepared
        assert "Unknown tool" in prepared["error"]
        # Error lists available tools so the model can self-correct
        assert "bash" in prepared["error"]
        # Surfaces warning to user
        session.ui.on_error.assert_called_once()
        assert "nonexistent" in session.ui.on_error.call_args[0][0]

    def test_prepare_tool_strips_whitespace_from_name(self, tmp_db):
        """Local models may produce tool names with leading/trailing whitespace."""
        session = self._make_session(mcp_client=None)
        tc = {
            "id": "call_strip",
            "function": {"name": "  bash\n", "arguments": '{"command": "echo hi"}'},
        }
        prepared = session._prepare_tool(tc)
        assert prepared["func_name"] == "bash"
        assert "error" not in prepared

    def test_prepare_tool_malformed_json_surfaces_error(self, tmp_db):
        """Malformed JSON args should surface a warning to the user and
        give the model a hint about expected format."""
        session = self._make_session(mcp_client=None)
        tc = {
            "id": "call_bad",
            "function": {"name": "bash", "arguments": "{command: echo hi}"},
        }
        prepared = session._prepare_tool(tc)
        assert "error" in prepared
        assert "JSON parse error" in prepared["error"]
        assert "command" in prepared["error"]  # hint about expected key
        assert "Please retry" in prepared["error"]
        # User-facing warning
        session.ui.on_error.assert_called_once()
        assert "Malformed tool call" in session.ui.on_error.call_args[0][0]

    def test_ensure_tool_call_ids_dict(self, tmp_db):
        """_ensure_tool_call_ids fills empty IDs on streaming-style dict."""
        from turnstone.core.session import ChatSession

        tool_calls_acc = {
            0: {"id": "", "function": {"name": "bash", "arguments": "{}"}},
            1: {"id": "", "function": {"name": "read_file", "arguments": "{}"}},
        }
        ChatSession._ensure_tool_call_ids(tool_calls_acc)
        ids = [tc["id"] for tc in tool_calls_acc.values()]
        assert all(id_.startswith("call_") for id_ in ids)
        assert len(set(ids)) == 2  # unique

    def test_ensure_tool_call_ids_list(self, tmp_db):
        """_ensure_tool_call_ids fills empty IDs on list (agent path)."""
        from turnstone.core.session import ChatSession

        tool_calls = [
            {"id": None, "function": {"name": "bash", "arguments": "{}"}},
            {"id": "call_existing", "function": {"name": "bash", "arguments": "{}"}},
        ]
        ChatSession._ensure_tool_call_ids(tool_calls)
        assert tool_calls[0]["id"].startswith("call_")
        assert tool_calls[1]["id"] == "call_existing"  # preserved

    def test_mcp_command_no_client(self, tmp_db):
        session = self._make_session(mcp_client=None)
        session.handle_command("/mcp")
        session.ui.on_info.assert_called_once()
        assert "No MCP servers" in session.ui.on_info.call_args[0][0]

    def test_mcp_command_with_tools(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        session = self._make_session(mcp_client=mock_mcp)
        session.handle_command("/mcp")
        session.ui.on_info.assert_called_once()
        output = session.ui.on_info.call_args[0][0]
        assert "MCP tools (1)" in output
        assert "mcp__test__search" in output

    def test_exec_mcp_tool(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.is_mcp_tool.return_value = True
        mock_mcp.call_tool_sync.return_value = "result text"
        session = self._make_session(mcp_client=mock_mcp)

        item = {
            "call_id": "call_789",
            "mcp_func_name": "mcp__test__search",
            "mcp_args": {"query": "hello"},
        }
        call_id, output = session._exec_mcp_tool(item)
        assert call_id == "call_789"
        assert output == "result text"
        mock_mcp.call_tool_sync.assert_called_once_with(
            "mcp__test__search",
            {"query": "hello"},
            user_id=None,
            timeout=30,
            is_interactive_for_consent=True,
        )

    def test_exec_mcp_tool_error(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.is_mcp_tool.return_value = True
        mock_mcp.call_tool_sync.side_effect = RuntimeError("server crashed")
        session = self._make_session(mcp_client=mock_mcp)

        item = {
            "call_id": "call_err",
            "mcp_func_name": "mcp__test__search",
            "mcp_args": {"query": "hello"},
        }
        call_id, output = session._exec_mcp_tool(item)
        assert call_id == "call_err"
        assert "MCP tool error" in output
        assert "server crashed" in output

    # -- Phase 7: per-user catalog scoping ---------------------------------

    def test_session_passes_user_id_to_get_tools(self, tmp_db):
        """ChatSession threads its ``user_id`` into ``get_tools`` so the
        merged static + pool view is scoped to the session's user."""
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        self._make_session(mcp_client=mock_mcp, user_id="user-7")
        mock_mcp.get_tools.assert_called_with(user_id="user-7")

    def test_session_get_tools_empty_user_id_passes_none(self, tmp_db):
        """Sentinel ``user_id=""`` (CLI / service / unknown) collapses to
        ``user_id=None`` so the static-only view is returned."""
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        self._make_session(mcp_client=mock_mcp, user_id="")
        mock_mcp.get_tools.assert_called_with(user_id=None)

    def test_session_passes_user_id_to_add_listener(self, tmp_db):
        """ChatSession registers its tool-change listener under its own
        ``user_id`` so pool-only changes for OTHER users do not fire it."""
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        self._make_session(mcp_client=mock_mcp, user_id="user-7")
        # ``add_listener`` was called with ``user_id="user-7"``.
        listener_calls = mock_mcp.add_listener.call_args_list
        assert listener_calls, "ChatSession did not register a tool listener"
        first_call = listener_calls[0]
        assert first_call.kwargs.get("user_id") == "user-7"

    def test_session_close_removes_listener_with_same_user_id(self, tmp_db):
        """R4 critical: register and remove MUST agree on ``user_id`` —
        the listener identity is ``(user_id, callback)``."""
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        session = self._make_session(mcp_client=mock_mcp, user_id="user-7")

        session.close()
        # ``remove_listener`` must be called with the same ``user_id``.
        remove_calls = mock_mcp.remove_listener.call_args_list
        assert remove_calls, "ChatSession.close did not unregister a tool listener"
        first_remove = remove_calls[0]
        assert first_remove.kwargs.get("user_id") == "user-7"
        # And the callback identity must match what was registered.
        registered_cb = mock_mcp.add_listener.call_args_list[0].args[0]
        removed_cb = first_remove.args[0]
        assert registered_cb is removed_cb

    def test_session_unknown_tool_lists_user_scoped_catalog(self, tmp_db):
        """The "Unknown tool" error message lists tools the session can
        actually invoke — drawn from the merged user-scoped catalog,
        not the manager's private static-only ``_tool_map``."""
        mock_mcp = MagicMock()
        # Pretend the user's merged view contains a static + pool entry.
        mock_mcp.get_tools.return_value = [
            _fake_openai_tool("mcp__static__list"),
            _fake_openai_tool("mcp__pool-srv__do"),
        ]
        mock_mcp.is_mcp_tool.return_value = False
        session = self._make_session(mcp_client=mock_mcp, user_id="user-7")
        # Reset the call counter so we observe only the _prepare_tool call.
        mock_mcp.get_tools.reset_mock()

        tc = {
            "id": "call_unknown",
            "function": {"name": "no_such_tool", "arguments": "{}"},
        }
        prepared = session._prepare_tool(tc)
        assert "error" in prepared
        # The error mentions both static and pool tools — proves we're
        # consulting the merged catalog rather than ``_tool_map``.
        assert "mcp__static__list" in prepared["error"]
        assert "mcp__pool-srv__do" in prepared["error"]
        # And the catalog request was scoped to this session's user.
        assert any(
            call.kwargs.get("user_id") == "user-7" for call in mock_mcp.get_tools.call_args_list
        )


# ---------------------------------------------------------------------------
# Server name validation
# ---------------------------------------------------------------------------


class TestServerNameValidation:
    def test_double_underscore_in_name(self):
        """Server names with __ should be rejected during _connect_one."""
        import asyncio

        async def _run() -> None:
            mgr = MCPClientManager({"my__bad": {"command": "echo"}})
            async with AsyncExitStack() as stack:
                mgr._exit_stack = stack
                await mgr._connect_one("my__bad", {"command": "echo"})
            # Should not have connected
            assert "my__bad" not in mgr._static_servers
            assert mgr.get_tools() == []

        asyncio.run(_run())


# ---------------------------------------------------------------------------
# create_mcp_client guard
# ---------------------------------------------------------------------------


class TestCreateMcpClient:
    def test_returns_none_when_no_config(self):
        with patch("turnstone.core.mcp_client.load_mcp_config", return_value={}):
            from turnstone.core.mcp_client import create_mcp_client

            result = create_mcp_client()
            assert result is None


# ---------------------------------------------------------------------------
# Tool refresh — _rebuild_tools, _refresh_server, listeners
# ---------------------------------------------------------------------------


class TestRebuildTools:
    def test_rebuild_from_per_server(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "github", tools=[_fake_openai_tool("mcp__github__search")])
        _seed_static_state(mgr, "slack", tools=[_fake_openai_tool("mcp__slack__send")])
        mgr._rebuild_tools()
        assert len(mgr._tools) == 2
        names = {t["function"]["name"] for t in mgr._tools}
        assert names == {"mcp__github__search", "mcp__slack__send"}
        assert mgr._tool_map["mcp__github__search"] == ("github", "search")
        assert mgr._tool_map["mcp__slack__send"] == ("slack", "send")

    def test_rebuild_copy_on_write(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
        mgr._rebuild_tools()
        old_tools = mgr._tools
        old_map = mgr._tool_map
        _seed_static_state(mgr, "b", tools=[_fake_openai_tool("mcp__b__y")])
        mgr._rebuild_tools()
        assert mgr._tools is not old_tools
        assert mgr._tool_map is not old_map

    def test_rebuild_empty(self):
        mgr = MCPClientManager({})
        mgr._static_servers = {}
        mgr._rebuild_tools()
        assert mgr._tools == []
        assert mgr._tool_map == {}


class TestRefreshServer:
    @staticmethod
    def _add_empty_resource_prompt_mocks(
        mgr: MCPClientManager, server_name: str, mock_session: MagicMock
    ) -> None:
        """Add empty list_resources/list_prompts mocks so _refresh_server works."""
        _seed_static_state(mgr, server_name, supports_resources=True, supports_prompts=True)
        empty_res = MagicMock()
        empty_res.resources = []
        mock_session.list_resources = AsyncMock(return_value=empty_res)
        empty_tmpl = MagicMock()
        empty_tmpl.resourceTemplates = []
        mock_session.list_resource_templates = AsyncMock(return_value=empty_tmpl)
        empty_prompts = MagicMock()
        empty_prompts.prompts = []
        mock_session.list_prompts = AsyncMock(return_value=empty_prompts)

    def test_refresh_detects_added_tools(self):
        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = MagicMock()
            mock_result = MagicMock()
            mock_result.tools = [
                _fake_mcp_tool("search"),
                _fake_mcp_tool("create"),  # new tool
            ]
            mock_session.list_tools = AsyncMock(return_value=mock_result)
            self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
            _seed_static_state(
                mgr,
                "github",
                session=mock_session,
                tools=[_fake_openai_tool("mcp__github__search")],
            )
            mgr._rebuild_tools()

            added, removed = await mgr._refresh_server("github")
            assert "mcp__github__create" in added
            assert removed == []
            assert len(mgr._tools) == 2

        asyncio.run(_run())

    def test_refresh_detects_removed_tools(self):
        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = MagicMock()
            mock_result = MagicMock()
            mock_result.tools = []  # all tools removed
            mock_session.list_tools = AsyncMock(return_value=mock_result)
            self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
            _seed_static_state(
                mgr,
                "github",
                session=mock_session,
                tools=[_fake_openai_tool("mcp__github__search")],
            )
            mgr._rebuild_tools()

            added, removed = await mgr._refresh_server("github")
            assert added == []
            assert "mcp__github__search" in removed
            assert mgr._tools == []

        asyncio.run(_run())

    def test_refresh_no_changes(self):
        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = MagicMock()
            mock_result = MagicMock()
            mock_result.tools = [_fake_mcp_tool("search")]
            mock_session.list_tools = AsyncMock(return_value=mock_result)
            self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
            _seed_static_state(
                mgr,
                "github",
                session=mock_session,
                tools=[_fake_openai_tool("mcp__github__search")],
            )
            mgr._rebuild_tools()

            added, removed = await mgr._refresh_server("github")
            assert added == []
            assert removed == []

        asyncio.run(_run())

    def test_refresh_disconnected_raises(self):
        async def _run() -> None:
            mgr = MCPClientManager({})
            with pytest.raises(RuntimeError, match="not connected"):
                await mgr._refresh_server_tools("ghost")

        asyncio.run(_run())


class TestLastRefreshTracking:
    """Phase 9 admin status pill — ``_last_refresh`` is written on every
    refresh path so the admin UI reflects manual-refresh AND auto-
    reconnect outcomes uniformly.  This test class pins the contract.
    """

    @staticmethod
    def _seed_minimal(mgr: MCPClientManager, name: str = "srv") -> MagicMock:
        mock_session = MagicMock()
        mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=[]))
        mock_session.list_resources = AsyncMock(return_value=MagicMock(resources=[]))
        mock_session.list_resource_templates = AsyncMock(
            return_value=MagicMock(resourceTemplates=[])
        )
        mock_session.list_prompts = AsyncMock(return_value=MagicMock(prompts=[]))
        _seed_static_state(
            mgr,
            name,
            session=mock_session,
            tools=[],
            supports_resources=True,
            supports_prompts=True,
        )
        return mock_session

    def test_last_refresh_written_on_success(self) -> None:
        async def _run() -> None:
            mgr = MCPClientManager({})
            self._seed_minimal(mgr)
            assert "srv" not in mgr._last_refresh

            await mgr._refresh_server("srv")

            entry = mgr._last_refresh.get("srv")
            assert entry is not None
            ts, outcome = entry
            assert outcome == "ok"
            assert isinstance(ts, float) and ts > 0

        asyncio.run(_run())

    def test_last_refresh_written_on_tool_refresh_failure(self) -> None:
        """When ``_refresh_server_tools`` raises, the outcome reflects
        the exception class and the exception still propagates."""

        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = self._seed_minimal(mgr)
            mock_session.list_tools = AsyncMock(side_effect=RuntimeError("upstream down"))

            with pytest.raises(RuntimeError, match="upstream down"):
                await mgr._refresh_server("srv")

            entry = mgr._last_refresh.get("srv")
            assert entry is not None
            _, outcome = entry
            assert outcome == "error:RuntimeError"

        asyncio.run(_run())

    def test_last_refresh_records_first_exception_when_multiple_fail(
        self,
    ) -> None:
        """``return_exceptions=True`` lets sibling tasks complete; the
        outcome reflects the FIRST exception encountered."""

        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = self._seed_minimal(mgr)
            # Tools succeeds; resources raises first (gather preserves
            # argument order in its results list, so resources is the
            # first failure regardless of which awaitable finished first
            # in wall-clock terms).
            mock_session.list_resources = AsyncMock(side_effect=ValueError("res boom"))
            mock_session.list_prompts = AsyncMock(side_effect=KeyError("prompts boom"))

            with pytest.raises((ValueError, KeyError)):
                await mgr._refresh_server("srv")

            entry = mgr._last_refresh.get("srv")
            assert entry is not None
            _, outcome = entry
            # Either of the two failing tasks could be "first" in
            # gather's results list ordering — the order is positional
            # so resources (arg #2) comes before prompts (arg #3).
            assert outcome == "error:ValueError"

        asyncio.run(_run())

    def test_refresh_all_overwrites_stale_ok_on_reconnect_failure(
        self,
    ) -> None:
        """The chokepoint bug-1 fix: a prior successful refresh's ``'ok'``
        entry MUST be overwritten when a subsequent reconnect fails —
        otherwise the admin pill shows misleading "ok" while the server
        is in fact broken."""

        async def _run() -> None:
            mgr = MCPClientManager({})
            # Server is configured but has no live session — _refresh_all
            # routes to the reconnect branch.
            mgr._server_configs["srv"] = {"type": "stdio", "command": "x"}
            # Pre-seed a stale "ok" from an earlier successful refresh.
            mgr._last_refresh["srv"] = (1000.0, "ok")

            async def _raise(*_a: object, **_kw: object) -> None:
                raise ConnectionError("reconnect failed")

            mgr._connect_one = _raise  # type: ignore[assignment]

            await mgr._refresh_all("srv")

            entry = mgr._last_refresh.get("srv")
            assert entry is not None
            ts, outcome = entry
            # Outcome reflects the new failure, not the stale ok.
            assert outcome == "error:ConnectionError"
            assert ts > 1000.0

        asyncio.run(_run())

    def test_get_server_status_surfaces_last_refresh_fields(self) -> None:
        """``get_server_status`` surfaces ``last_refresh_at`` and
        ``last_refresh_outcome`` for the admin pill — null when no
        refresh has occurred yet, populated after one."""
        mgr = MCPClientManager({})
        mgr._server_configs["srv"] = {"type": "stdio", "command": "x"}

        # No refresh yet — fields must be present and null so the JS
        # renderer can branch on absence cleanly.
        status = mgr.get_server_status("srv")
        assert status["last_refresh_at"] is None
        assert status["last_refresh_outcome"] is None

        # Populate the tuple directly and re-read.
        mgr._last_refresh["srv"] = (12345.5, "ok")
        status = mgr.get_server_status("srv")
        assert status["last_refresh_at"] == 12345.5
        assert status["last_refresh_outcome"] == "ok"


class TestListeners:
    def test_add_and_notify(self):
        mgr = MCPClientManager({})
        calls: list[int] = []
        mgr.add_listener(lambda: calls.append(1))
        _seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
        mgr._rebuild_tools()
        assert len(calls) == 1

    def test_remove_listener(self):
        mgr = MCPClientManager({})
        calls: list[int] = []
        cb = lambda: calls.append(1)  # noqa: E731
        mgr.add_listener(cb)
        mgr.remove_listener(cb)
        mgr._rebuild_tools()
        assert calls == []

    def test_remove_nonexistent_listener(self):
        mgr = MCPClientManager({})
        mgr.remove_listener(lambda: None)  # should not raise

    def test_listener_error_does_not_propagate(self):
        mgr = MCPClientManager({})
        mgr.add_listener(lambda: 1 / 0)  # will raise ZeroDivisionError
        mgr._rebuild_tools()  # should not raise

    # -- Phase 7: user-keyed listener fan-out ------------------------------

    def test_add_listener_records_user_id(self):
        """``add_listener`` stores ``(user_id, callback)`` tuples — the
        listener identity carries the user_id."""
        mgr = MCPClientManager({})
        cb_admin = lambda: None  # noqa: E731
        cb_user = lambda: None  # noqa: E731
        mgr.add_listener(cb_admin)  # default: user_id=None (admin)
        mgr.add_listener(cb_user, user_id="user-1")
        assert (None, cb_admin) in mgr._listeners
        assert ("user-1", cb_user) in mgr._listeners

    def test_remove_listener_requires_matching_user_id(self):
        """Removing with a different ``user_id`` must NOT remove the
        original registration — listener identity is the pair."""
        mgr = MCPClientManager({})
        calls: list[int] = []
        cb = lambda: calls.append(1)  # noqa: E731
        mgr.add_listener(cb, user_id="user-1")

        # Try to remove with the wrong user_id — should be a no-op.
        mgr.remove_listener(cb, user_id="user-2")
        # The user-1 listener should still be live.
        mgr._notify_user_tool_listeners("user-1")
        assert calls == [1]

        # Now remove with the right user_id.
        mgr.remove_listener(cb, user_id="user-1")
        mgr._notify_user_tool_listeners("user-1")
        assert calls == [1]  # not invoked again

    def test_static_change_fires_all_listeners(self):
        """``_rebuild_tools`` (static-path change) fires ALL registered
        listeners — admin + every user. RFC §3.3."""
        mgr = MCPClientManager({})
        admin_calls: list[int] = []
        u1_calls: list[int] = []
        u2_calls: list[int] = []
        mgr.add_listener(lambda: admin_calls.append(1))
        mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
        mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")
        _seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
        mgr._rebuild_tools()
        assert admin_calls == [1]
        assert u1_calls == [1]
        assert u2_calls == [1]

    def test_user_tool_listeners_only_fire_for_matching_user(self):
        """``_notify_user_tool_listeners('user-1')`` fires admin (None)
        and user-1 listeners; user-2's listener is silent."""
        mgr = MCPClientManager({})
        admin_calls: list[int] = []
        u1_calls: list[int] = []
        u2_calls: list[int] = []
        mgr.add_listener(lambda: admin_calls.append(1))
        mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
        mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")

        mgr._notify_user_tool_listeners("user-1")
        assert admin_calls == [1]
        assert u1_calls == [1]
        assert u2_calls == []

        mgr._notify_user_tool_listeners("user-2")
        assert admin_calls == [1, 1]
        assert u1_calls == [1]
        assert u2_calls == [1]


class TestServerNames:
    def test_server_names_property(self):
        mgr = MCPClientManager({"github": {}, "slack": {}})
        assert sorted(mgr.server_names) == ["github", "slack"]

    def test_server_names_empty(self):
        mgr = MCPClientManager({})
        assert mgr.server_names == []


# ---------------------------------------------------------------------------
# Session integration — tool refresh propagation
# ---------------------------------------------------------------------------


class TestSessionRefresh:
    @pytest.fixture()
    def tmp_db(self, tmp_path):
        from turnstone.core.storage import init_storage, reset_storage

        reset_storage()
        init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False)
        yield
        reset_storage()

    def _make_session(self, mcp_client=None, **kwargs):
        from turnstone.core.session import ChatSession

        defaults: dict[str, Any] = dict(
            client=MagicMock(),
            model="test-model",
            ui=MagicMock(),
            instructions=None,
            temperature=0.5,
            max_tokens=4096,
            tool_timeout=30,
            mcp_client=mcp_client,
        )
        defaults.update(kwargs)
        return ChatSession(**defaults)

    def test_listener_registered_on_init(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = []
        session = self._make_session(mcp_client=mock_mcp)
        mock_mcp.add_listener.assert_called_once()
        assert session._mcp_refresh_cb is not None

    def test_no_listener_without_mcp(self, tmp_db):
        session = self._make_session(mcp_client=None)
        assert session._mcp_refresh_cb is None

    def test_close_removes_listener(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = []
        session = self._make_session(mcp_client=mock_mcp)
        session.close()
        mock_mcp.remove_listener.assert_called_once()
        assert session._mcp_refresh_cb is None

    def test_close_idempotent(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = []
        session = self._make_session(mcp_client=mock_mcp)
        session.close()
        session.close()  # should not raise
        assert mock_mcp.remove_listener.call_count == 1

    def test_on_mcp_tools_changed_rebuilds_tools(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool("mcp__test__a")]
        session = self._make_session(mcp_client=mock_mcp)
        initial_count = len(session._tools)

        # Simulate a tool refresh — MCP now has 2 tools
        mock_mcp.get_tools.return_value = [
            _fake_openai_tool("mcp__test__a"),
            _fake_openai_tool("mcp__test__b"),
        ]
        session._on_mcp_tools_changed()
        assert len(session._tools) == initial_count + 1

    def test_tool_search_preserved_across_refresh(self, tmp_db):
        # Create enough MCP tools to trigger tool search
        mcp_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(25)]
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = mcp_tools
        session = self._make_session(
            mcp_client=mock_mcp,
            tool_search="auto",
            tool_search_threshold=20,
        )
        assert session._tool_search is not None

        # Expand a tool
        session._tool_search.expand_visible(["mcp__srv__tool0"])
        assert "mcp__srv__tool0" in session._tool_search.get_expanded_names()

        # Refresh with same tools
        session._on_mcp_tools_changed()
        assert session._tool_search is not None
        assert "mcp__srv__tool0" in session._tool_search.get_expanded_names()

    def test_tool_search_prunes_removed_from_expanded(self, tmp_db):
        mcp_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(25)]
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = mcp_tools
        session = self._make_session(
            mcp_client=mock_mcp,
            tool_search="auto",
            tool_search_threshold=20,
        )
        session._tool_search.expand_visible(["mcp__srv__tool0"])

        # Refresh with tool0 removed
        new_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(1, 25)]
        mock_mcp.get_tools.return_value = new_tools
        session._on_mcp_tools_changed()
        # tool0 was removed, so it should no longer be in expanded
        expanded = session._tool_search.get_expanded_names()
        assert "mcp__srv__tool0" not in expanded

    def test_mcp_refresh_command(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.server_names = ["test"]
        mock_mcp.refresh_sync.return_value = {"test": (["mcp__test__new"], [])}
        session = self._make_session(mcp_client=mock_mcp)

        session.handle_command("/mcp refresh")
        mock_mcp.refresh_sync.assert_called_once_with(None)
        session.ui.on_info.assert_called()

    def test_mcp_refresh_specific_server(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.server_names = ["github", "slack"]
        mock_mcp.refresh_sync.return_value = {"github": ([], [])}
        session = self._make_session(mcp_client=mock_mcp)

        session.handle_command("/mcp refresh github")
        mock_mcp.refresh_sync.assert_called_once_with("github")

    def test_mcp_refresh_unknown_server(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.server_names = ["github"]
        session = self._make_session(mcp_client=mock_mcp)

        session.handle_command("/mcp refresh nonexistent")
        session.ui.on_error.assert_called_once()
        assert "Unknown MCP server" in session.ui.on_error.call_args[0][0]

    def test_mcp_refresh_error_handling(self, tmp_db):
        mock_mcp = MagicMock()
        mock_mcp.get_tools.return_value = [_fake_openai_tool()]
        mock_mcp.server_names = ["test"]
        mock_mcp.refresh_sync.side_effect = TimeoutError("timed out")
        session = self._make_session(mcp_client=mock_mcp)

        session.handle_command("/mcp refresh")
        session.ui.on_error.assert_called_once()
        assert "MCP refresh failed" in session.ui.on_error.call_args[0][0]


# ---------------------------------------------------------------------------
# MCP Resources
# ---------------------------------------------------------------------------


class TestMCPResources:
    def test_resource_discovery(self):
        """Mock list_resources() returning 2 resources, verify get_resources()."""
        mgr = MCPClientManager({})
        _seed_static_state(
            mgr,
            "fs",
            resources=[
                _fake_resource_dict("file:///a.txt", "a", "File A", "text/plain", "fs"),
                _fake_resource_dict("file:///b.txt", "b", "File B", "text/plain", "fs"),
            ],
        )
        mgr._rebuild_resources()
        resources = mgr.get_resources()
        assert len(resources) == 2
        uris = {r["uri"] for r in resources}
        assert uris == {"file:///a.txt", "file:///b.txt"}
        assert all(r["server"] == "fs" for r in resources)

    def test_rebuild_resources_copy_on_write(self):
        """Verify mutation safety — get_resources() returns independent copy."""
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", resources=[_fake_resource_dict("file:///x", "x", "", "", "a")])
        mgr._rebuild_resources()
        old_resources = mgr._resources
        old_map = mgr._resource_map
        _seed_static_state(mgr, "b", resources=[_fake_resource_dict("file:///y", "y", "", "", "b")])
        mgr._rebuild_resources()
        assert mgr._resources is not old_resources
        assert mgr._resource_map is not old_map

    def test_get_resources_returns_copy(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", resources=[_fake_resource_dict("file:///x", "x", "", "", "a")])
        mgr._rebuild_resources()
        resources = mgr.get_resources()
        assert len(resources) == 1
        resources.clear()
        assert len(mgr.get_resources()) == 1

    def test_read_resource_sync(self):
        """Mock session.read_resource(), verify text extraction."""
        mgr = MCPClientManager({})
        mgr._resource_map = {"file:///readme": ("fs", "file:///readme")}
        mock_session = MagicMock()
        _seed_static_state(mgr, "fs", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        # Mock the read_resource result
        text_content = MagicMock(spec=["text"])
        text_content.text = "Hello, world!"
        mock_result = MagicMock()
        mock_result.contents = [text_content]
        mock_session.read_resource = AsyncMock(return_value=mock_result)

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            output = mgr.read_resource_sync("file:///readme", timeout=5)
            assert output == "Hello, world!"
            mock_session.read_resource.assert_awaited_once_with("file:///readme")
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()

    def test_read_resource_sync_blob(self):
        """Verify base64 blob extraction."""
        mgr = MCPClientManager({})
        mgr._resource_map = {"file:///img.png": ("fs", "file:///img.png")}
        mock_session = MagicMock()
        _seed_static_state(mgr, "fs", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        blob_content = MagicMock(spec=["blob"])
        blob_content.blob = "aGVsbG8="
        mock_result = MagicMock()
        mock_result.contents = [blob_content]
        mock_session.read_resource = AsyncMock(return_value=mock_result)

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            output = mgr.read_resource_sync("file:///img.png", timeout=5)
            assert output == "aGVsbG8="
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()

    def test_read_resource_sync_unknown_uri(self):
        mgr = MCPClientManager({})
        with pytest.raises(ValueError, match="Unknown MCP resource"):
            mgr.read_resource_sync("file:///nonexistent")

    def test_read_resource_sync_disconnected(self):
        mgr = MCPClientManager({})
        mgr._resource_map = {"file:///x": ("dead", "file:///x")}
        with pytest.raises(RuntimeError, match="not connected"):
            mgr.read_resource_sync("file:///x")

    def test_read_resource_sync_timeout(self):
        """Verify timeout handling."""
        mgr = MCPClientManager({})
        mgr._resource_map = {"file:///x": ("fs", "file:///x")}
        mock_session = MagicMock()
        _seed_static_state(mgr, "fs", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        async def _slow_read(_uri: str) -> None:
            await asyncio.sleep(10)

        mock_session.read_resource = _slow_read

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            with pytest.raises(TimeoutError):
                mgr.read_resource_sync("file:///x", timeout=1)
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()

    def test_resource_listener_notification(self):
        """Verify callback fires on rebuild."""
        mgr = MCPClientManager({})
        calls: list[int] = []
        mgr.add_resource_listener(lambda: calls.append(1))
        _seed_static_state(mgr, "a", resources=[_fake_resource_dict()])
        mgr._rebuild_resources()
        assert len(calls) == 1

    def test_resource_listener_remove(self):
        mgr = MCPClientManager({})
        calls: list[int] = []
        cb = lambda: calls.append(1)  # noqa: E731
        mgr.add_resource_listener(cb)
        mgr.remove_resource_listener(cb)
        mgr._rebuild_resources()
        assert calls == []

    def test_resource_listener_error_does_not_propagate(self):
        mgr = MCPClientManager({})
        mgr.add_resource_listener(lambda: 1 / 0)
        mgr._rebuild_resources()  # should not raise

    def test_resource_refresh_on_notification(self):
        """Mock notification, verify re-fetch of resources."""

        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = MagicMock()
            _seed_static_state(
                mgr,
                "fs",
                session=mock_session,
                supports_resources=True,
                resources=[_fake_resource_dict("file:///old", server="fs")],
            )
            mgr._rebuild_resources()
            assert len(mgr.get_resources()) == 1

            # Mock the re-fetch returning a new resource
            new_res = _fake_mcp_resource("file:///new", "new")
            mock_res_result = MagicMock()
            mock_res_result.resources = [new_res]
            mock_session.list_resources = AsyncMock(return_value=mock_res_result)
            mock_tmpl_result = MagicMock()
            mock_tmpl_result.resourceTemplates = []
            mock_session.list_resource_templates = AsyncMock(return_value=mock_tmpl_result)

            await mgr._refresh_server_resources("fs")
            resources = mgr.get_resources()
            assert len(resources) == 1
            assert resources[0]["uri"] == "file:///new"

        asyncio.run(_run())

    def test_rebuild_resources_empty(self):
        mgr = MCPClientManager({})
        mgr._static_servers = {}
        mgr._rebuild_resources()
        assert mgr._resources == []
        assert mgr._resource_map == {}

    def test_rebuild_resources_multi_server(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "fs", resources=[_fake_resource_dict("file:///a", server="fs")])
        _seed_static_state(
            mgr,
            "db",
            resources=[_fake_resource_dict("db://table", name="table", server="db")],
        )
        mgr._rebuild_resources()
        assert len(mgr._resources) == 2
        assert mgr._resource_map["file:///a"] == ("fs", "file:///a")
        assert mgr._resource_map["db://table"] == ("db", "db://table")

    def test_template_prefix_matching(self):
        """Expanded URI matches template by prefix."""
        mgr = MCPClientManager({})
        _seed_static_state(
            mgr,
            "db",
            resources=[
                {
                    "uri": "db://tables/{table}/rows/{id}",
                    "name": "row",
                    "description": "A row",
                    "mimeType": "application/json",
                    "server": "db",
                    "template": True,
                },
            ],
        )
        mgr._rebuild_resources()
        # Template should not be in resource_map
        assert "db://tables/{table}/rows/{id}" not in mgr._resource_map
        # But prefix matching should find it
        result = mgr._match_template("db://tables/users/rows/1")
        assert result is not None
        server, template_uri = result
        assert server == "db"
        assert template_uri == "db://tables/{table}/rows/{id}"

    def test_template_longest_prefix_wins(self):
        """When two templates have overlapping prefixes, the longer one wins."""
        mgr = MCPClientManager({})
        # Use templates with genuinely different prefix lengths:
        # "db://data/" (6 chars after scheme) vs "db://data/tables/" (13 chars after scheme)
        _seed_static_state(
            mgr,
            "short",
            resources=[
                {
                    "uri": "db://data/{collection}",
                    "name": "collection",
                    "description": "",
                    "mimeType": "",
                    "server": "short",
                    "template": True,
                },
            ],
        )
        _seed_static_state(
            mgr,
            "long",
            resources=[
                {
                    "uri": "db://data/tables/{table}",
                    "name": "table",
                    "description": "",
                    "mimeType": "",
                    "server": "long",
                    "template": True,
                },
            ],
        )
        mgr._rebuild_resources()
        # "db://data/tables/users" matches both prefixes ("db://data/" and
        # "db://data/tables/") — the longer one should win
        result = mgr._match_template("db://data/tables/users")
        assert result is not None
        server, template_uri = result
        assert server == "long"
        assert template_uri == "db://data/tables/{table}"
        # URI that only matches the short prefix
        result2 = mgr._match_template("db://data/views/active")
        assert result2 is not None
        assert result2[0] == "short"

    def test_template_no_match_raises(self):
        """Completely unrelated URI still raises ValueError."""
        mgr = MCPClientManager({})
        _seed_static_state(
            mgr,
            "db",
            resources=[
                {
                    "uri": "db://tables/{table}",
                    "name": "table",
                    "description": "",
                    "mimeType": "",
                    "server": "db",
                    "template": True,
                },
            ],
        )
        mgr._rebuild_resources()
        assert mgr._match_template("file:///something") is None
        with pytest.raises(ValueError, match="Unknown MCP resource"):
            mgr.read_resource_sync("file:///something")

    def test_read_resource_sync_with_template_uri(self):
        """End-to-end: template discovered, expanded URI dispatched to correct server."""
        mgr = MCPClientManager({})
        _seed_static_state(
            mgr,
            "db",
            resources=[
                {
                    "uri": "db://tables/{table}/rows/{id}",
                    "name": "row",
                    "description": "A row",
                    "mimeType": "application/json",
                    "server": "db",
                    "template": True,
                },
            ],
        )
        mgr._rebuild_resources()

        mock_session = MagicMock()
        _seed_static_state(mgr, "db", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        text_content = MagicMock(spec=["text"])
        text_content.text = '{"name": "Alice"}'
        mock_result = MagicMock()
        mock_result.contents = [text_content]
        mock_session.read_resource = AsyncMock(return_value=mock_result)

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            output = mgr.read_resource_sync("db://tables/users/rows/1", timeout=5)
            assert output == '{"name": "Alice"}'
            mock_session.read_resource.assert_awaited_once_with("db://tables/users/rows/1")
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()


# ---------------------------------------------------------------------------
# MCP Prompts
# ---------------------------------------------------------------------------


class TestMCPPrompts:
    def test_prompt_discovery(self):
        """Mock list_prompts(), verify get_prompts() with correct prefixed names."""
        mgr = MCPClientManager({})
        _seed_static_state(
            mgr,
            "tmpl",
            prompts=[
                _fake_prompt_dict("mcp__tmpl__code_review", "code_review", "tmpl"),
                _fake_prompt_dict("mcp__tmpl__summarize", "summarize", "tmpl"),
            ],
        )
        mgr._rebuild_prompts()
        prompts = mgr.get_prompts()
        assert len(prompts) == 2
        names = {p["name"] for p in prompts}
        assert names == {"mcp__tmpl__code_review", "mcp__tmpl__summarize"}
        # Verify map entries
        assert mgr._prompt_map["mcp__tmpl__code_review"] == ("tmpl", "code_review")
        assert mgr._prompt_map["mcp__tmpl__summarize"] == ("tmpl", "summarize")

    def test_rebuild_prompts_copy_on_write(self):
        """Verify mutation safety."""
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", prompts=[_fake_prompt_dict("mcp__a__p1", "p1", "a")])
        mgr._rebuild_prompts()
        old_prompts = mgr._prompts
        old_map = mgr._prompt_map
        _seed_static_state(mgr, "b", prompts=[_fake_prompt_dict("mcp__b__p2", "p2", "b")])
        mgr._rebuild_prompts()
        assert mgr._prompts is not old_prompts
        assert mgr._prompt_map is not old_map

    def test_get_prompts_returns_copy(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", prompts=[_fake_prompt_dict("mcp__a__p1", "p1", "a")])
        mgr._rebuild_prompts()
        prompts = mgr.get_prompts()
        assert len(prompts) == 1
        prompts.clear()
        assert len(mgr.get_prompts()) == 1

    def test_get_prompt_sync(self):
        """Mock session.get_prompt(), verify message conversion."""
        mgr = MCPClientManager({})
        mgr._prompt_map = {"mcp__tmpl__review": ("tmpl", "review")}
        mock_session = MagicMock()
        _seed_static_state(mgr, "tmpl", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        # Build mock PromptMessage
        msg1 = MagicMock()
        msg1.role = "user"
        msg1.content = MagicMock()
        msg1.content.text = "Review this code"
        msg2 = MagicMock()
        msg2.role = "assistant"
        msg2.content = MagicMock()
        msg2.content.text = "Looks good!"
        mock_result = MagicMock()
        mock_result.messages = [msg1, msg2]
        mock_session.get_prompt = AsyncMock(return_value=mock_result)

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            messages = mgr.get_prompt_sync(
                "mcp__tmpl__review", arguments={"language": "python"}, timeout=5
            )
            assert len(messages) == 2
            assert messages[0] == {"role": "user", "content": "Review this code"}
            assert messages[1] == {"role": "assistant", "content": "Looks good!"}
            mock_session.get_prompt.assert_awaited_once_with(
                "review", arguments={"language": "python"}
            )
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()

    def test_get_prompt_sync_unknown(self):
        mgr = MCPClientManager({})
        with pytest.raises(ValueError, match="Unknown MCP prompt"):
            mgr.get_prompt_sync("mcp__no__such")

    def test_get_prompt_sync_disconnected(self):
        mgr = MCPClientManager({})
        mgr._prompt_map = {"mcp__dead__p": ("dead", "p")}
        with pytest.raises(RuntimeError, match="not connected"):
            mgr.get_prompt_sync("mcp__dead__p")

    def test_get_prompt_sync_timeout(self):
        """Verify timeout handling."""
        mgr = MCPClientManager({})
        mgr._prompt_map = {"mcp__tmpl__slow": ("tmpl", "slow")}
        mock_session = MagicMock()
        _seed_static_state(mgr, "tmpl", session=mock_session)
        mgr._loop = asyncio.new_event_loop()

        async def _slow_prompt(_name: str, *, arguments: dict[str, str] | None = None) -> None:
            await asyncio.sleep(10)

        mock_session.get_prompt = _slow_prompt

        thread = None
        try:
            thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
            thread.start()
            with pytest.raises(TimeoutError):
                mgr.get_prompt_sync("mcp__tmpl__slow", timeout=1)
        finally:
            mgr._loop.call_soon_threadsafe(mgr._loop.stop)
            if thread:
                thread.join(timeout=5)
            mgr._loop.close()

    def test_prompt_listener_notification(self):
        """Verify callback fires on rebuild."""
        mgr = MCPClientManager({})
        calls: list[int] = []
        mgr.add_prompt_listener(lambda: calls.append(1))
        _seed_static_state(mgr, "a", prompts=[_fake_prompt_dict()])
        mgr._rebuild_prompts()
        assert len(calls) == 1

    def test_prompt_listener_remove(self):
        mgr = MCPClientManager({})
        calls: list[int] = []
        cb = lambda: calls.append(1)  # noqa: E731
        mgr.add_prompt_listener(cb)
        mgr.remove_prompt_listener(cb)
        mgr._rebuild_prompts()
        assert calls == []

    def test_prompt_listener_error_does_not_propagate(self):
        mgr = MCPClientManager({})
        mgr.add_prompt_listener(lambda: 1 / 0)
        mgr._rebuild_prompts()  # should not raise

    def test_is_mcp_prompt(self):
        """Verify name lookup."""
        mgr = MCPClientManager({})
        mgr._prompt_map["mcp__tmpl__review"] = ("tmpl", "review")
        assert mgr.is_mcp_prompt("mcp__tmpl__review") is True
        assert mgr.is_mcp_prompt("nonexistent") is False

    def test_prompt_refresh_on_notification(self):
        """Mock notification, verify re-fetch of prompts."""

        async def _run() -> None:
            mgr = MCPClientManager({})
            mock_session = MagicMock()
            _seed_static_state(
                mgr,
                "tmpl",
                session=mock_session,
                supports_prompts=True,
                prompts=[_fake_prompt_dict("mcp__tmpl__old", "old", "tmpl")],
            )
            mgr._rebuild_prompts()
            assert len(mgr.get_prompts()) == 1

            # Mock re-fetch returning a new prompt
            new_prompt = _fake_mcp_prompt("new_prompt", "A new prompt")
            mock_prompt_result = MagicMock()
            mock_prompt_result.prompts = [new_prompt]
            mock_session.list_prompts = AsyncMock(return_value=mock_prompt_result)

            await mgr._refresh_server_prompts("tmpl")
            prompts = mgr.get_prompts()
            assert len(prompts) == 1
            assert prompts[0]["name"] == "mcp__tmpl__new_prompt"
            assert prompts[0]["original_name"] == "new_prompt"

        asyncio.run(_run())

    def test_rebuild_prompts_empty(self):
        mgr = MCPClientManager({})
        mgr._static_servers = {}
        mgr._rebuild_prompts()
        assert mgr._prompts == []
        assert mgr._prompt_map == {}

    def test_rebuild_prompts_multi_server(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", prompts=[_fake_prompt_dict("mcp__a__p1", "p1", "a")])
        _seed_static_state(mgr, "b", prompts=[_fake_prompt_dict("mcp__b__p2", "p2", "b")])
        mgr._rebuild_prompts()
        assert len(mgr._prompts) == 2
        assert mgr._prompt_map["mcp__a__p1"] == ("a", "p1")
        assert mgr._prompt_map["mcp__b__p2"] == ("b", "p2")


# ---------------------------------------------------------------------------
# Shutdown cleans up new state
# ---------------------------------------------------------------------------


class TestShutdownCleanup:
    def test_shutdown_clears_resources_and_prompts(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "a", resources=[_fake_resource_dict()])
        mgr._rebuild_resources()
        _seed_static_state(mgr, "a", prompts=[_fake_prompt_dict()])
        mgr._rebuild_prompts()
        assert mgr.get_resources() != []
        assert mgr.get_prompts() != []

        mgr.shutdown()
        assert mgr.get_resources() == []
        assert mgr.get_prompts() == []
        assert mgr._resource_map == {}
        assert mgr._prompt_map == {}


# ---------------------------------------------------------------------------
# TCP probe and unreachable server handling
# ---------------------------------------------------------------------------


class TestTCPProbe:
    """MCPClientManager._tcp_probe should fail fast on unreachable servers."""

    def test_tcp_probe_unreachable_raises_connection_error(self):
        """Unreachable host raises ConnectionError, not TimeoutError."""
        mgr = MCPClientManager({})

        async def _run():
            with pytest.raises(ConnectionError, match="unreachable"):
                await mgr._tcp_probe("test-server", "http://127.0.0.1:1")

        asyncio.run(_run())

    def test_tcp_probe_parses_url_correctly(self):
        """Port and host are extracted from the URL."""
        mgr = MCPClientManager({})

        async def _run():
            # Non-routable port — should fail with ConnectionError
            with pytest.raises(ConnectionError):
                await mgr._tcp_probe("srv", "https://127.0.0.1:1/mcp")

        asyncio.run(_run())

    def test_tcp_probe_default_port_http(self):
        """Default port 80 used for http:// URLs without explicit port."""
        mgr = MCPClientManager({})

        async def _run():
            # Will fail (nothing on port 80), but should not crash on parsing
            with pytest.raises(ConnectionError):
                await mgr._tcp_probe("srv", "http://127.0.0.1")

        asyncio.run(_run())

    def test_tcp_probe_dns_failure(self):
        """Unresolvable hostname raises ConnectionError."""
        mgr = MCPClientManager({})

        async def _run():
            with pytest.raises(ConnectionError):
                await mgr._tcp_probe("srv", "http://this.host.does.not.exist.invalid:8080/mcp")

        asyncio.run(_run())


class TestConnectOneUnreachable:
    """_connect_one should handle unreachable HTTP servers gracefully."""

    def test_unreachable_http_server_raises_connection_error(self):
        """Unreachable HTTP MCP server raises ConnectionError without spinning."""
        mgr = MCPClientManager({})
        mgr._loop = asyncio.new_event_loop()

        async def _run():
            with pytest.raises(ConnectionError, match="unreachable"):
                await mgr._connect_one(
                    "bad-server",
                    {
                        "type": "http",
                        "url": "http://127.0.0.1:1/mcp",
                    },
                )

        mgr._loop.run_until_complete(_run())
        mgr._loop.close()

        # Server should NOT have a live session (connection failed)
        bad_state = mgr._static_servers.get("bad-server")
        assert bad_state is None or bad_state.session is None

    def test_connect_all_continues_after_unreachable_server(self):
        """_connect_all logs error and continues to next server."""
        mgr = MCPClientManager(
            {
                "bad": {"type": "http", "url": "http://127.0.0.1:1/mcp"},
            }
        )

        loop = asyncio.new_event_loop()
        loop.run_until_complete(mgr._connect_all())
        loop.close()

        bad_state = mgr._static_servers.get("bad")
        assert bad_state is None or bad_state.session is None
        assert "bad" in mgr._last_error


class TestSafeCloseStack:
    """_safe_close_stack should suppress errors from broken anyio scopes."""

    def test_suppresses_runtime_error(self):
        """RuntimeError from broken cancel scope is suppressed."""

        async def _run():
            stack = AsyncExitStack()
            await stack.__aenter__()

            # Simulate a broken close that raises RuntimeError
            async def _broken_close():
                raise RuntimeError("Attempted to exit cancel scope in a different task")

            stack.aclose = _broken_close
            # Should not raise
            await MCPClientManager._safe_close_stack(stack)

        asyncio.run(_run())

    def test_suppresses_cancelled_error(self):
        """CancelledError during close is suppressed."""

        async def _run():
            stack = AsyncExitStack()
            await stack.__aenter__()

            async def _cancel_close():
                raise asyncio.CancelledError()

            stack.aclose = _cancel_close
            await MCPClientManager._safe_close_stack(stack)

        asyncio.run(_run())


# ---------------------------------------------------------------------------
# Fix 1: Cancel orphaned futures on timeout
# ---------------------------------------------------------------------------


class TestFutureCancellation:
    """Verify future.cancel() is called when sync bridge methods time out."""

    def _make_manager_with_session(self) -> MCPClientManager:
        mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
        mock_session = MagicMock()
        # Prevent auto-spec from creating async coroutines that trigger warnings
        mock_session.call_tool = MagicMock(return_value="sentinel")
        mock_session.read_resource = MagicMock(return_value="sentinel")
        mock_session.get_prompt = MagicMock(return_value="sentinel")
        _seed_static_state(mgr, "test", session=mock_session)
        mgr._loop = MagicMock()
        mgr._tool_map["mcp__test__search"] = ("test", "search")
        mgr._resource_map["file:///a.txt"] = ("test", "file:///a.txt")
        mgr._prompt_map["mcp__test__review"] = ("test", "review")
        return mgr

    def test_call_tool_sync_cancels_future_on_timeout(self):
        mgr = self._make_manager_with_session()
        mock_future = MagicMock()
        mock_future.result.side_effect = concurrent.futures.TimeoutError()
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(TimeoutError, match="timed out"),
        ):
            mgr.call_tool_sync("mcp__test__search", {"query": "x"}, timeout=1)
        mock_future.cancel.assert_called_once()

    def test_read_resource_sync_cancels_future_on_timeout(self):
        mgr = self._make_manager_with_session()
        mock_future = MagicMock()
        mock_future.result.side_effect = concurrent.futures.TimeoutError()
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(TimeoutError, match="timed out"),
        ):
            mgr.read_resource_sync("file:///a.txt", timeout=1)
        mock_future.cancel.assert_called_once()

    def test_get_prompt_sync_cancels_future_on_timeout(self):
        mgr = self._make_manager_with_session()
        mock_future = MagicMock()
        mock_future.result.side_effect = concurrent.futures.TimeoutError()
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(TimeoutError, match="timed out"),
        ):
            mgr.get_prompt_sync("mcp__test__review", timeout=1)
        mock_future.cancel.assert_called_once()

    def test_refresh_sync_cancels_future_on_timeout(self):
        mgr = MCPClientManager({})
        mgr._loop = MagicMock()
        mock_future = MagicMock()
        mock_future.result.side_effect = concurrent.futures.TimeoutError()
        with (
            patch.object(mgr, "_refresh_all", return_value=MagicMock()),
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(TimeoutError, match="timed out"),
        ):
            mgr.refresh_sync(timeout=1)
        mock_future.cancel.assert_called_once()


# ---------------------------------------------------------------------------
# Fix 2: Per-server circuit breaker
# ---------------------------------------------------------------------------


class TestCircuitBreaker:
    """Verify per-server circuit breaker behavior."""

    def test_circuit_stays_closed_below_threshold(self):
        mgr = MCPClientManager({})
        mgr._cb_record_failure("srv")
        mgr._cb_record_failure("srv")
        is_open, _ = mgr._cb_check("srv")
        assert not is_open

    def test_circuit_opens_at_threshold(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        is_open, cooldown_expired = mgr._cb_check("srv")
        assert is_open
        assert not cooldown_expired  # just opened, cooldown not expired

    def test_circuit_half_open_after_cooldown(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        # Simulate cooldown expiry
        mgr._circuit_open_until["srv"] = time.monotonic() - 1
        is_open, cooldown_expired = mgr._cb_check("srv")
        assert is_open
        assert cooldown_expired

    def test_circuit_resets_on_success(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        assert "srv" in mgr._circuit_open_until
        mgr._cb_record_success("srv")
        is_open, _ = mgr._cb_check("srv")
        assert not is_open
        assert mgr._consecutive_failures.get("srv") is None

    def test_success_decays_trip_count(self):
        """Success decays trip_count by 1 so flapping servers escalate backoff."""
        mgr = MCPClientManager({})
        mgr._circuit_trip_count["srv"] = 3
        mgr._cb_record_success("srv")
        assert mgr._circuit_trip_count["srv"] == 2
        mgr._cb_record_success("srv")
        assert mgr._circuit_trip_count["srv"] == 1
        mgr._cb_record_success("srv")
        assert "srv" not in mgr._circuit_trip_count

    def test_cooldown_is_exponential(self):
        mgr = MCPClientManager({})
        # First trip (trip_count starts at 0)
        for _ in range(3):
            mgr._cb_record_failure("srv")
        deadline1 = mgr._circuit_open_until["srv"]
        base1 = deadline1 - time.monotonic()
        # Reset circuit but keep trip_count at 1 (set by first trip)
        mgr._cb_record_success("srv")
        # trip_count decayed from 1 to 0 — manually set to 1 for test
        mgr._circuit_trip_count["srv"] = 1
        for _ in range(3):
            mgr._cb_record_failure("srv")
        deadline2 = mgr._circuit_open_until["srv"]
        base2 = deadline2 - time.monotonic()
        # Second trip should have longer cooldown (roughly 2x, within jitter)
        assert base2 > base1 * 1.5

    def test_cooldown_capped_at_max(self):
        mgr = MCPClientManager({})
        mgr._circuit_trip_count["srv"] = 100  # very high trip count
        for _ in range(3):
            mgr._cb_record_failure("srv")
        deadline = mgr._circuit_open_until["srv"]
        cooldown = deadline - time.monotonic()
        # Should not exceed max (300s) + 10% jitter = 330s
        assert cooldown <= mgr._CB_MAX_COOLDOWN * 1.11

    def test_cb_gate_rejects_when_open(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        with pytest.raises(RuntimeError, match="circuit open"):
            mgr._cb_gate("srv")

    def test_cb_gate_allows_after_cooldown(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        mgr._circuit_open_until["srv"] = time.monotonic() - 1
        # Should not raise
        mgr._cb_gate("srv")
        # Deadline should be removed (half-open probe allowed)
        assert "srv" not in mgr._circuit_open_until

    def test_cb_clear_removes_all_state(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("srv")
        mgr._cb_clear("srv")
        assert "srv" not in mgr._consecutive_failures
        assert "srv" not in mgr._circuit_open_until
        assert "srv" not in mgr._circuit_trip_count

    @pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
    @pytest.mark.filterwarnings("ignore:coroutine.*was never awaited:RuntimeWarning")
    def test_call_tool_sync_records_failure_on_timeout(self):
        mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
        mock_session = MagicMock()
        mock_session.call_tool = MagicMock(return_value="sentinel")
        _seed_static_state(mgr, "test", session=mock_session)
        mgr._loop = MagicMock()
        mgr._tool_map["mcp__test__ping"] = ("test", "ping")
        mock_future = MagicMock()
        mock_future.result.side_effect = concurrent.futures.TimeoutError()
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(TimeoutError),
        ):
            mgr.call_tool_sync("mcp__test__ping", {}, timeout=1)
        assert mgr._consecutive_failures.get("test", 0) == 1

    def test_call_tool_sync_records_success(self):
        mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
        mock_session = MagicMock()
        mock_session.call_tool = MagicMock(return_value="sentinel")
        _seed_static_state(mgr, "test", session=mock_session)
        mgr._loop = MagicMock()
        mgr._tool_map["mcp__test__ping"] = ("test", "ping")
        # Pre-set a failure
        mgr._consecutive_failures["test"] = 2
        mock_result = MagicMock()
        mock_result.content = []
        mock_result.isError = False
        mock_future = MagicMock()
        mock_future.result.return_value = mock_result
        with patch("asyncio.run_coroutine_threadsafe", return_value=mock_future):
            mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
        assert mgr._consecutive_failures.get("test") is None

    def test_connection_error_evicts_session(self):
        mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
        mock_session = MagicMock()
        mock_session.call_tool = MagicMock(return_value="sentinel")
        # Seed both session and stack so the test can verify stack survives.
        old_stack = MagicMock()
        old_streams = (MagicMock(), MagicMock())
        _seed_static_state(mgr, "test", session=mock_session, stack=old_stack, streams=old_streams)
        mgr._loop = MagicMock()
        mgr._tool_map["mcp__test__ping"] = ("test", "ping")
        mock_future = MagicMock()
        mock_future.result.side_effect = BrokenPipeError("dead")
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(BrokenPipeError),
        ):
            mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
        # Session evicted, but stack/streams remain for the stale-and-stack
        # guard in _connect_one to clean up on next reconnect attempt.
        state = mgr._static_servers["test"]
        assert state.session is None
        assert state.stack is old_stack
        assert state.streams is old_streams

    def test_independent_circuits_per_server(self):
        mgr = MCPClientManager({})
        for _ in range(3):
            mgr._cb_record_failure("a")
        is_open_a, _ = mgr._cb_check("a")
        is_open_b, _ = mgr._cb_check("b")
        assert is_open_a
        assert not is_open_b

    def test_mcp_error_does_not_trip_circuit(self):
        """Protocol errors (McpError) should not count as transport failures."""
        from mcp import McpError
        from mcp.types import ErrorData

        mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
        mock_session = MagicMock()
        mock_session.call_tool = MagicMock(return_value="sentinel")
        _seed_static_state(mgr, "test", session=mock_session)
        mgr._loop = MagicMock()
        mgr._tool_map["mcp__test__ping"] = ("test", "ping")
        mock_future = MagicMock()
        mock_future.result.side_effect = McpError(ErrorData(code=-32601, message="tool not found"))
        with (
            patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
            pytest.raises(McpError),
        ):
            mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
        # Circuit should NOT have recorded a failure
        assert mgr._consecutive_failures.get("test", 0) == 0


# ---------------------------------------------------------------------------
# Fix 3: Safe transport stream pre-close
# ---------------------------------------------------------------------------


class TestSafeTransportStreams:
    """Verify stream references are stored and pre-closed."""

    def test_pre_close_streams_closes_both(self):
        mgr = MCPClientManager({})
        stream_a = MagicMock()
        stream_b = MagicMock()
        _seed_static_state(mgr, "srv", streams=(stream_a, stream_b))

        async def _run():
            await mgr._pre_close_streams("srv")

        asyncio.run(_run())
        stream_a.aclose.assert_called_once()
        stream_b.aclose.assert_called_once()
        # Streams cleared, but the state entry itself can remain.
        assert mgr._static_servers["srv"].streams is None

    def test_pre_close_streams_ignores_missing(self):
        mgr = MCPClientManager({})

        async def _run():
            await mgr._pre_close_streams("nonexistent")

        asyncio.run(_run())  # should not raise

    def test_pre_close_streams_suppresses_errors(self):
        mgr = MCPClientManager({})
        stream_a = MagicMock()
        stream_a.aclose.side_effect = RuntimeError("boom")
        stream_b = MagicMock()
        _seed_static_state(mgr, "srv", streams=(stream_a, stream_b))

        async def _run():
            await mgr._pre_close_streams("srv")

        asyncio.run(_run())  # should not raise despite stream_a error
        stream_b.aclose.assert_called_once()

    def test_shutdown_clears_stream_refs(self):
        mgr = MCPClientManager({})
        _seed_static_state(mgr, "srv", streams=(MagicMock(), MagicMock()))
        mgr.shutdown()
        assert len(mgr._static_servers) == 0


# ---------------------------------------------------------------------------
# Fix 4: Notification debounce
# ---------------------------------------------------------------------------


class TestNotificationDebounce:
    """Verify notification-triggered refreshes are debounced."""

    def test_debounce_within_window(self):
        mgr = MCPClientManager({})
        mgr._last_notification_refresh["srv"] = time.monotonic()
        # We can't easily call _on_notification (it's a closure), so test
        # the debounce logic directly via the timestamp check
        now = time.monotonic()
        last = mgr._last_notification_refresh.get("srv", 0.0)
        assert now - last < mgr._NOTIFICATION_DEBOUNCE

    def test_debounce_passes_after_window(self):
        mgr = MCPClientManager({})
        # Set timestamp well in the past
        mgr._last_notification_refresh["srv"] = time.monotonic() - 10
        now = time.monotonic()
        last = mgr._last_notification_refresh.get("srv", 0.0)
        assert now - last >= mgr._NOTIFICATION_DEBOUNCE

    def test_debounce_is_per_server(self):
        mgr = MCPClientManager({})
        mgr._last_notification_refresh["srv_a"] = time.monotonic()
        # srv_b has no timestamp — should pass debounce
        now = time.monotonic()
        last_b = mgr._last_notification_refresh.get("srv_b", 0.0)
        assert now - last_b >= mgr._NOTIFICATION_DEBOUNCE


# ---------------------------------------------------------------------------
# reconnect_sync — operator-driven full reconnect
# ---------------------------------------------------------------------------


class TestReconnectSync:
    """Verify reconnect_sync tears down old session, clears CB, calls _connect_one."""

    def test_reconnect_unknown_server_returns_error(self):
        mgr = MCPClientManager({})
        result = mgr.reconnect_sync("missing")
        assert result == {
            "connected": False,
            "tools": 0,
            "resources": 0,
            "prompts": 0,
            "error": "unknown server",
        }

    def test_reconnect_clears_circuit_breaker(self, running_loop_mgr):
        mgr, _loop, _thread = running_loop_mgr

        async def _fake_connect_one(name: str, _cfg: dict[str, Any]) -> None:
            _seed_static_state(mgr, name, session=MagicMock())

        # Pre-trip the breaker
        for _ in range(3):
            mgr._cb_record_failure("srv")
        assert "srv" in mgr._circuit_open_until

        with (
            patch.object(mgr, "_connect_one", side_effect=_fake_connect_one),
            patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
        ):
            result = mgr.reconnect_sync("srv")
        assert result["connected"] is True
        assert result["error"] == ""
        assert "srv" not in mgr._circuit_open_until
        assert "srv" not in mgr._consecutive_failures

    def test_reconnect_closes_old_session_then_calls_connect_one(self, running_loop_mgr):
        mgr, _loop, _thread = running_loop_mgr

        order: list[str] = []
        old_stack = MagicMock(spec=AsyncExitStack)

        async def _pre_close(name: str) -> None:
            order.append("pre_close")

        async def _safe_close(stack: Any) -> None:
            order.append("safe_close")
            assert stack is old_stack

        async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
            order.append("connect_one")
            _seed_static_state(mgr, name, session=MagicMock())

        # Seed the old session/stack/streams that the guard should clear.
        _seed_static_state(
            mgr,
            "srv",
            session=MagicMock(),
            stack=old_stack,
            streams=(MagicMock(), MagicMock()),
        )

        with (
            patch.object(mgr, "_pre_close_streams", side_effect=_pre_close),
            patch.object(mgr, "_safe_close_stack", side_effect=_safe_close),
            patch.object(mgr, "_connect_one", side_effect=_connect_one),
        ):
            result = mgr.reconnect_sync("srv")
        assert result["connected"] is True
        assert order == ["pre_close", "safe_close", "connect_one"]
        # The old stack reference should have been cleared from state.
        assert mgr._static_servers["srv"].stack is not old_stack

    def test_reconnect_failure_returns_error_dict(self, running_loop_mgr):
        mgr, _loop, _thread = running_loop_mgr

        async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
            raise RuntimeError("handshake failed")

        with (
            patch.object(mgr, "_connect_one", side_effect=_connect_one),
            patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
        ):
            result = mgr.reconnect_sync("srv")
        assert result["connected"] is False
        assert "handshake failed" in result["error"]

    def test_reconnect_failure_clears_stale_catalog(self, running_loop_mgr):
        # bug-2: when _connect_one fails mid-reconnect, the per-server
        # catalog must be dropped so the merged tool/resource/prompt maps
        # don't keep advertising entries with no live session.
        mgr, _loop, _thread = running_loop_mgr

        # Seed catalog state from a previous successful connect.
        _seed_static_state(
            mgr,
            "srv",
            tools=[_fake_openai_tool("mcp__srv__t")],
            resources=[_fake_resource_dict(server="srv")],
            prompts=[_fake_prompt_dict(server="srv")],
        )
        mgr._rebuild_tools()
        mgr._rebuild_resources()
        mgr._rebuild_prompts()

        async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
            raise RuntimeError("handshake failed")

        with (
            patch.object(mgr, "_connect_one", side_effect=_connect_one),
            patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
        ):
            result = mgr.reconnect_sync("srv")
        assert result["connected"] is False
        # Per-server catalog should be cleared and merged maps drained.
        srv_state = mgr._static_servers.get("srv")
        assert srv_state is not None
        assert srv_state.tools == []
        assert srv_state.resources == []
        assert srv_state.prompts == []
        assert "mcp__srv__t" not in mgr._tool_map

    def test_reconnect_preserves_static_state_identity(self, running_loop_mgr):
        # q-3: PR #296 invariant 5 — _static_servers[name] must be the SAME
        # object across a connect → transient-failure → reconnect cycle.
        # Guards against future refactors that pop-and-repopulate the entry,
        # which would invalidate any references held by concurrent readers.
        mgr, _loop, _thread = running_loop_mgr

        # First connect: seed an initial entry as if _connect_one succeeded.
        async def _first_connect(name: str, _cfg: dict[str, Any]) -> None:
            _seed_static_state(mgr, name, session=MagicMock())

        with (
            patch.object(mgr, "_connect_one", side_effect=_first_connect),
            patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
        ):
            mgr.reconnect_sync("srv")

        state_before = mgr._static_servers["srv"]
        id_before = id(state_before)

        # Simulate a transient transport failure: evict the session (as
        # call_tool_sync would on BrokenPipeError) but keep the entry.
        state_before.session = None

        # Reconnect.
        async def _reconnect(name: str, _cfg: dict[str, Any]) -> None:
            _seed_static_state(mgr, name, session=MagicMock())

        with (
            patch.object(mgr, "_connect_one", side_effect=_reconnect),
            patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
        ):
            result = mgr.reconnect_sync("srv")
        assert result["connected"] is True

        state_after = mgr._static_servers["srv"]
        assert id(state_after) == id_before
        assert state_after is state_before


# ---------------------------------------------------------------------------
# _cb_auto_reconnect — refresh-on-reconnect
# ---------------------------------------------------------------------------


class TestCBAutoReconnectRefresh:
    """Verify _cb_auto_reconnect schedules catalog refresh after a successful reconnect.

    The refresh runs as a fire-and-forget background task on the loop so it
    doesn't block the caller (perf-1).  Tests wait briefly for the scheduled
    task to run and observe its effect.
    """

    def test_auto_reconnect_schedules_refresh_server_on_success(self, running_loop_mgr):
        import threading as _threading

        mgr, _loop, _thread = running_loop_mgr

        new_session = MagicMock()
        refresh_event = _threading.Event()

        async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
            _seed_static_state(mgr, name, session=new_session)

        async def _refresh(name: str) -> tuple[list[str], list[str]]:
            refresh_event.set()
            return [], []

        with (
            patch.object(mgr, "_connect_one", side_effect=_connect_one),
            patch.object(mgr, "_refresh_server", side_effect=_refresh),
        ):
            session = mgr._cb_auto_reconnect("srv")
            # Wait for the scheduled refresh task to actually run on the loop.
            assert refresh_event.wait(timeout=5), "refresh task was not scheduled"
        assert session is new_session

    def test_auto_reconnect_swallows_refresh_failure(self, running_loop_mgr):
        import threading as _threading

        mgr, _loop, _thread = running_loop_mgr

        new_session = MagicMock()
        refresh_started = _threading.Event()

        async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
            _seed_static_state(mgr, name, session=new_session)

        async def _refresh_failing(name: str) -> tuple[list[str], list[str]]:
            refresh_started.set()
            raise RuntimeError("catalog fetch broke")

        with (
            patch.object(mgr, "_connect_one", side_effect=_connect_one),
            patch.object(mgr, "_refresh_server", side_effect=_refresh_failing),
        ):
            # Must not raise — refresh failures are non-fatal.
            session = mgr._cb_auto_reconnect("srv")
            # Background refresh actually started and exception was swallowed
            # by the task without affecting the synchronous caller.
            assert refresh_started.wait(timeout=5), "refresh task was not scheduled"
        assert session is new_session