"""
Manages the lifecycle of multiple MCP server connections.
"""

import asyncio
import traceback
from collections import deque
from contextlib import AbstractAsyncContextManager, suppress
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Callable, Union, cast

import httpx
from anyio import Event, Lock, create_task_group
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from httpx import HTTPStatusError
from mcp import ClientSession
from mcp.client.stdio import (
    StdioServerParameters,
    get_default_environment,
)
from mcp.client.streamable_http import GetSessionIdCallback
from mcp.shared._httpx_utils import (
    MCP_DEFAULT_SSE_READ_TIMEOUT,
    MCP_DEFAULT_TIMEOUT,
    create_mcp_http_client,
)
from mcp.types import Implementation, JSONRPCMessage, ServerCapabilities

from fast_agent.config import MCPServerSettings
from fast_agent.context_dependent import ContextDependent
from fast_agent.core.exceptions import ServerInitializationError
from fast_agent.core.logging.logger import get_logger
from fast_agent.event_progress import ProgressAction
from fast_agent.mcp.logger_textio import get_stderr_handler
from fast_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
from fast_agent.mcp.oauth_client import build_oauth_provider
from fast_agent.mcp.sse_tracking import tracking_sse_client
from fast_agent.mcp.stdio_tracking_simple import tracking_stdio_client
from fast_agent.mcp.streamable_http_tracking import tracking_streamablehttp_client
from fast_agent.mcp.transport_tracking import TransportChannelMetrics

if TYPE_CHECKING:
    from mcp.client.auth import OAuthClientProvider

    from fast_agent.context import Context
    from fast_agent.mcp_server_registry import ServerRegistry

logger = get_logger(__name__)


class StreamingContextAdapter:
    """Adapter to provide a 3-value context from a 2-value context manager"""

    def __init__(self, context_manager):
        self.context_manager = context_manager
        self.cm_instance = None

    async def __aenter__(self):
        self.cm_instance = await self.context_manager.__aenter__()
        read_stream, write_stream = self.cm_instance
        return read_stream, write_stream, None

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return await self.context_manager.__aexit__(exc_type, exc_val, exc_tb)


def _add_none_to_context(context_manager):
    """Helper to add a None value to context managers that return 2 values instead of 3"""
    return StreamingContextAdapter(context_manager)


def _prepare_headers_and_auth(
    server_config: MCPServerSettings,
) -> tuple[dict[str, str], Union["OAuthClientProvider", None], set[str]]:
    """
    Prepare request headers and determine if OAuth authentication should be used.

    Returns a copy of the headers, an OAuth auth provider when applicable, and the set
    of user-supplied authorization header keys.
    """
    headers: dict[str, str] = dict(server_config.headers or {})
    auth_header_keys = {"authorization", "x-hf-authorization"}
    user_provided_auth_keys = {key for key in headers if key.lower() in auth_header_keys}

    # OAuth is only relevant for SSE/HTTP transports and should be skipped when the
    # user has already supplied explicit Authorization headers.
    if server_config.transport not in ("sse", "http") or user_provided_auth_keys:
        return headers, None, user_provided_auth_keys

    oauth_auth = build_oauth_provider(server_config)
    if oauth_auth is not None:
        # Scrub Authorization headers so OAuth-managed credentials are the only ones sent.
        for header_name in (
            "Authorization",
            "authorization",
            "X-HF-Authorization",
            "x-hf-authorization",
        ):
            headers.pop(header_name, None)

    return headers, oauth_auth, user_provided_auth_keys


class ServerConnection:
    """
    Represents a long-lived MCP server connection, including:
    - The ClientSession to the server
    - The transport streams (via stdio/sse, etc.)
    """

    def __init__(
        self,
        server_name: str,
        server_config: MCPServerSettings,
        transport_context_factory: Callable[
            [],
            AbstractAsyncContextManager[
                tuple[
                    MemoryObjectReceiveStream[JSONRPCMessage | Exception],
                    MemoryObjectSendStream[JSONRPCMessage],
                    GetSessionIdCallback | None,
                ]
            ],
        ],
        client_session_factory: Callable[..., ClientSession],
    ) -> None:
        self.server_name = server_name
        self.server_config = server_config
        self.session: ClientSession | None = None
        self._client_session_factory = client_session_factory
        self._transport_context_factory = transport_context_factory
        # Signal that session is fully up and initialized
        self._initialized_event = Event()

        # Signal we want to shut down
        self._shutdown_event = Event()

        # Track error state
        self._error_occurred = False
        self._error_message = None

        # Server instructions from initialization
        self.server_instructions: str | None = None
        self.server_capabilities: ServerCapabilities | None = None
        self.server_implementation: Implementation | None = None
        self.client_capabilities: dict | None = None
        self.server_instructions_available: bool = False
        self.server_instructions_enabled: bool = (
            server_config.include_instructions if server_config else True
        )
        self.session_id: str | None = None
        self._get_session_id_cb: GetSessionIdCallback | None = None
        self.transport_metrics: TransportChannelMetrics | None = None
        self._ping_ok_count = 0
        self._ping_fail_count = 0
        self._ping_consecutive_failures = 0
        self._ping_last_ok_at: datetime | None = None
        self._ping_last_fail_at: datetime | None = None
        self._ping_last_error: str | None = None
        self._ping_history: deque[tuple[datetime, str]] = deque(maxlen=200)

    def is_healthy(self) -> bool:
        """Check if the server connection is healthy and ready to use."""
        return self.session is not None and not self._error_occurred

    def reset_error_state(self) -> None:
        """Reset the error state, allowing reconnection attempts."""
        self._error_occurred = False
        self._error_message = None

    def request_shutdown(self) -> None:
        """
        Request the server to shut down. Signals the server lifecycle task to exit.
        """
        self._shutdown_event.set()

    async def wait_for_shutdown_request(self) -> None:
        """
        Wait until the shutdown event is set.
        """
        await self._shutdown_event.wait()

    async def initialize_session(self) -> None:
        """
        Initializes the server connection and session.
        Must be called within an async context.
        """
        assert self.session, "Session must be created before initialization"
        result = await self.session.initialize()

        self.server_capabilities = result.capabilities
        # InitializeResult exposes server info via `serverInfo`; keep fallback for older fields
        self.server_implementation = result.serverInfo

        raw_instructions = result.instructions
        self.server_instructions_available = bool(raw_instructions)

        # Store instructions if provided by the server and enabled in config
        if self.server_config.include_instructions:
            self.server_instructions = raw_instructions
            if self.server_instructions:
                logger.debug(
                    f"{self.server_name}: Received server instructions",
                    data={"instructions": self.server_instructions},
                )
        else:
            self.server_instructions = None
            if self.server_instructions_available:
                logger.debug(
                    f"{self.server_name}: Server instructions disabled by configuration",
                    data={"instructions": raw_instructions},
                )
            else:
                logger.debug(f"{self.server_name}: No server instructions provided")

        # If there's an init hook, run it

        # Now the session is ready for use
        self._initialized_event.set()

    async def wait_for_initialized(self) -> None:
        """
        Wait until the session is fully initialized.
        """
        await self._initialized_event.wait()

    def record_ping_event(self, state: str) -> None:
        self._ping_history.append((datetime.now(timezone.utc), state))

    def build_ping_activity_buckets(self, bucket_seconds: int, bucket_count: int) -> list[str]:
        try:
            seconds = int(bucket_seconds)
        except (TypeError, ValueError):
            seconds = 30
        if seconds <= 0:
            seconds = 30

        try:
            count = int(bucket_count)
        except (TypeError, ValueError):
            count = 20
        if count <= 0:
            count = 20

        if not self._ping_history:
            return ["none"] * count

        priority = {"error": 2, "ping": 1, "none": 0}
        history_map: dict[int, str] = {}
        for timestamp, state in self._ping_history:
            bucket = int(timestamp.timestamp() // seconds)
            existing = history_map.get(bucket)
            if existing is None or priority.get(state, 0) >= priority.get(existing, 0):
                history_map[bucket] = state

        current_bucket = int(datetime.now(timezone.utc).timestamp() // seconds)
        buckets: list[str] = []
        for offset in range(count - 1, -1, -1):
            bucket_index = current_bucket - offset
            buckets.append(history_map.get(bucket_index, "none"))

        return buckets

    def create_session(
        self,
        read_stream: MemoryObjectReceiveStream,
        send_stream: MemoryObjectSendStream,
    ) -> ClientSession:
        """
        Create a new session instance for this server connection.
        """

        read_timeout = (
            timedelta(seconds=self.server_config.read_timeout_seconds)
            if self.server_config.read_timeout_seconds
            else None
        )

        session = self._client_session_factory(
            read_stream,
            send_stream,
            read_timeout,
            server_config=self.server_config,
            transport_metrics=self.transport_metrics,
        )

        self.session = session
        self.client_capabilities = getattr(session, "client_capabilities", None)

        return session


async def _run_ping_loop(server_conn: ServerConnection) -> None:
    interval = server_conn.server_config.ping_interval_seconds
    if not interval or interval <= 0:
        return

    max_missed = server_conn.server_config.max_missed_pings
    missed = 0
    read_timeout = (
        timedelta(seconds=server_conn.server_config.read_timeout_seconds)
        if server_conn.server_config.read_timeout_seconds
        else None
    )
    if read_timeout is None:
        read_timeout = timedelta(seconds=interval)

    while not server_conn._shutdown_event.is_set():
        await asyncio.sleep(interval)
        if server_conn._shutdown_event.is_set():
            break
        session = server_conn.session
        if session is None:
            break
        if not hasattr(session, "ping"):
            return
        try:
            from fast_agent.human_input.elicitation_state import elicitation_state

            if elicitation_state.is_active(server_conn.server_name):
                continue
        except Exception:
            pass
        try:
            await cast("MCPAgentClientSession", session).ping(read_timeout_seconds=read_timeout)
            missed = 0
            server_conn._ping_ok_count += 1
            server_conn._ping_consecutive_failures = 0
            server_conn._ping_last_ok_at = datetime.now(timezone.utc)
            server_conn._ping_last_error = None
            server_conn.record_ping_event("ping")
        except Exception as exc:
            missed += 1
            server_conn._ping_fail_count += 1
            server_conn._ping_consecutive_failures = missed
            server_conn._ping_last_fail_at = datetime.now(timezone.utc)
            server_conn._ping_last_error = str(exc)
            server_conn.record_ping_event("error")
            logger.warning(
                f"{server_conn.server_name}: Ping failed ({missed}/{max_missed}): {exc}"
            )
            if missed >= max_missed:
                server_conn._error_occurred = True
                server_conn._error_message = (
                    f"Ping failed {missed} time(s); last error: {exc}"
                )
                server_conn.request_shutdown()
                break


async def _server_lifecycle_task(server_conn: ServerConnection) -> None:
    """
    Manage the lifecycle of a single server connection.
    Runs inside the MCPConnectionManager's shared TaskGroup.

    IMPORTANT: This function must NEVER raise an exception, as it runs in a shared
    task group. Any exceptions must be caught and handled gracefully, with errors
    recorded in server_conn._error_occurred and _error_message.
    """
    server_name = server_conn.server_name
    try:
        transport_context = server_conn._transport_context_factory()

        try:
            async with transport_context as (read_stream, write_stream, get_session_id_cb):
                server_conn._get_session_id_cb = get_session_id_cb

                if get_session_id_cb is not None:
                    try:
                        server_conn.session_id = get_session_id_cb()
                    except Exception:
                        logger.debug(f"{server_name}: Unable to retrieve session id from transport")
                elif server_conn.server_config.transport == "stdio":
                    server_conn.session_id = "local"

                server_conn.create_session(read_stream, write_stream)
                assert server_conn.session is not None

                try:
                    async with server_conn.session:
                        await server_conn.initialize_session()

                        if get_session_id_cb is not None:
                            try:
                                server_conn.session_id = get_session_id_cb() or server_conn.session_id
                            except Exception:
                                logger.debug(f"{server_name}: Unable to refresh session id after init")
                        elif server_conn.server_config.transport == "stdio":
                            server_conn.session_id = "local"

                        if (
                            server_conn.server_config.ping_interval_seconds
                            and server_conn.server_config.ping_interval_seconds > 0
                        ):
                            ping_task = asyncio.create_task(_run_ping_loop(server_conn))
                            try:
                                await server_conn.wait_for_shutdown_request()
                            finally:
                                if not ping_task.done():
                                    ping_task.cancel()
                                with suppress(asyncio.CancelledError):
                                    await ping_task
                        else:
                            await server_conn.wait_for_shutdown_request()
                except Exception as session_exit_exc:
                    if server_conn._shutdown_event.is_set():
                        # Cleanup errors can happen when disconnecting a session that was already
                        # terminated; treat as expected during shutdown.
                        logger.debug(
                            f"{server_name}: Exception during session cleanup (expected during shutdown): {session_exit_exc}"
                        )
                        if not server_conn._initialized_event.is_set():
                            server_conn._error_occurred = True
                            server_conn._error_message = "Shutdown requested before initialization"
                            server_conn._initialized_event.set()
                    else:
                        raise
        except Exception as transport_exit_exc:
            if server_conn._shutdown_event.is_set():
                # Cleanup errors can happen when disconnecting a transport that was already
                # terminated; treat as expected during shutdown.
                logger.debug(
                    f"{server_name}: Exception during transport cleanup (expected during shutdown): {transport_exit_exc}"
                )
                if not server_conn._initialized_event.is_set():
                    server_conn._error_occurred = True
                    server_conn._error_message = "Shutdown requested before initialization"
                    server_conn._initialized_event.set()
            else:
                raise

    except HTTPStatusError as http_exc:
        logger.error(
            f"{server_name}: Lifecycle task encountered HTTP error: {http_exc}",
            exc_info=True,
            data={
                "progress_action": ProgressAction.FATAL_ERROR,
                "server_name": server_name,
            },
        )
        server_conn._error_occurred = True
        server_conn._error_message = f"HTTP Error: {http_exc.response.status_code} {http_exc.response.reason_phrase} for URL: {http_exc.request.url}"
        server_conn._initialized_event.set()
        # No raise - let get_server handle it with a friendly message

    except Exception as exc:
        logger.error(
            f"{server_name}: Lifecycle task encountered an error: {exc}",
            exc_info=True,
            data={
                "progress_action": ProgressAction.FATAL_ERROR,
                "server_name": server_name,
            },
        )
        server_conn._error_occurred = True

        if "ExceptionGroup" in type(exc).__name__ and hasattr(exc, "exceptions"):
            # Handle ExceptionGroup better by extracting the actual errors
            def extract_errors(exception_group):
                """Recursively extract meaningful errors from ExceptionGroups"""
                messages = []
                for subexc in exception_group.exceptions:
                    if "ExceptionGroup" in type(subexc).__name__ and hasattr(subexc, "exceptions"):
                        # Recursively handle nested ExceptionGroups
                        messages.extend(extract_errors(subexc))
                    elif isinstance(subexc, HTTPStatusError):
                        # Special handling for HTTP errors to make them more user-friendly
                        messages.append(
                            f"HTTP Error: {subexc.response.status_code} {subexc.response.reason_phrase} for URL: {subexc.request.url}"
                        )
                    else:
                        # Show the exception type and message, plus the root cause if available
                        error_msg = f"{type(subexc).__name__}: {subexc}"
                        messages.append(error_msg)

                        # If there's a root cause, show that too as it's often the most informative
                        if hasattr(subexc, "__cause__") and subexc.__cause__:
                            messages.append(
                                f"Caused by: {type(subexc.__cause__).__name__}: {subexc.__cause__}"
                            )
                return messages

            error_messages = extract_errors(exc)
            # If we didn't extract any meaningful errors, fall back to the original exception
            if not error_messages:
                error_messages = [f"{type(exc).__name__}: {exc}"]
            server_conn._error_message = error_messages
        else:
            # For regular exceptions, keep the traceback but format it more cleanly
            server_conn._error_message = traceback.format_exception(exc)

        # If there's an error, we should also set the event so that
        # 'get_server' won't hang
        server_conn._initialized_event.set()
        # No raise - allow graceful exit


class MCPConnectionManager(ContextDependent):
    """
    Manages the lifecycle of multiple MCP server connections.
    Integrates with the application context system for proper resource management.
    """

    def __init__(
        self, server_registry: "ServerRegistry", context: Union["Context", None] = None
    ) -> None:
        super().__init__(context=context)
        self.server_registry = server_registry
        self.running_servers: dict[str, ServerConnection] = {}
        self._lock = Lock()
        # Manage our own task group - independent of task context
        self._task_group = None
        self._task_group_active = False
        self._mcp_sse_filter_added = False

    async def __aenter__(self):
        # Create a task group that isn't tied to a specific task
        self._task_group = create_task_group()
        # Enter the task group context
        await self._task_group.__aenter__()
        self._task_group_active = True
        self._tg = self._task_group
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Ensure clean shutdown of all connections before exiting."""
        try:
            # First request all servers to shutdown
            await self.disconnect_all()

            # Add a small delay to allow for clean shutdown
            await asyncio.sleep(0.5)

            # Then close the task group if it's active
            if self._task_group_active:
                assert self._task_group is not None
                await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
                self._task_group_active = False
                self._task_group = None
                self._tg = None
        except Exception:
            logger.exception("Error during connection manager shutdown")

    def _suppress_mcp_sse_errors(self) -> None:
        """Suppress MCP library's 'Error in sse_reader' messages."""
        if self._mcp_sse_filter_added:
            return

        import logging

        class MCPSSEErrorFilter(logging.Filter):
            def filter(self, record):
                return not (
                    record.name == "mcp.client.sse" and "Error in sse_reader" in record.getMessage()
                )

        mcp_sse_logger = logging.getLogger("mcp.client.sse")
        mcp_sse_logger.addFilter(MCPSSEErrorFilter())
        self._mcp_sse_filter_added = True

    async def launch_server(
        self,
        server_name: str,
        client_session_factory: Callable[
            [MemoryObjectReceiveStream, MemoryObjectSendStream, timedelta | None],
            ClientSession,
        ],
    ) -> ServerConnection:
        """
        Connect to a server and return a RunningServer instance that will persist
        until explicitly disconnected.
        """
        # Create task group if it doesn't exist yet - make this method more resilient
        if not self._task_group_active:
            self._task_group = create_task_group()
            await self._task_group.__aenter__()
            self._task_group_active = True
            self._tg = self._task_group
            logger.info(f"Auto-created task group for server: {server_name}")

        config = self.server_registry.get_server_config(server_name)
        if not config:
            raise ValueError(f"Server '{server_name}' not found in registry.")

        logger.debug(f"{server_name}: Found server configuration=", data=config.model_dump())

        timeline_steps = 20
        timeline_seconds = 30
        try:
            ctx = self.context
        except RuntimeError:
            ctx = None

        config_obj = getattr(ctx, "config", None)
        timeline_config = getattr(config_obj, "mcp_timeline", None)
        if timeline_config:
            timeline_steps = getattr(timeline_config, "steps", timeline_steps)
            timeline_seconds = getattr(timeline_config, "step_seconds", timeline_seconds)

        transport_metrics = (
            TransportChannelMetrics(
                bucket_seconds=timeline_seconds,
                bucket_count=timeline_steps,
            )
            if config.transport in ("http", "sse", "stdio")
            else None
        )

        def transport_context_factory():
            if config.transport == "stdio":
                if not config.command:
                    raise ValueError(
                        f"Server '{server_name}' uses stdio transport but no command is specified"
                    )
                server_params = StdioServerParameters(
                    command=config.command,
                    args=config.args if config.args is not None else [],
                    env={**get_default_environment(), **(config.env or {})},
                    cwd=config.cwd,
                )
                # Create custom error handler to ensure all output is captured
                error_handler = get_stderr_handler(server_name)
                # Explicitly ensure we're using our custom logger for stderr
                logger.debug(f"{server_name}: Creating stdio client with custom error handler")

                channel_hook = transport_metrics.record_event if transport_metrics else None
                return _add_none_to_context(
                    tracking_stdio_client(
                        server_params, channel_hook=channel_hook, errlog=error_handler
                    )
                )
            elif config.transport == "sse":
                if not config.url:
                    raise ValueError(
                        f"Server '{server_name}' uses sse transport but no url is specified"
                    )
                # Suppress MCP library error spam
                self._suppress_mcp_sse_errors()
                headers, oauth_auth, user_auth_keys = _prepare_headers_and_auth(config)
                if user_auth_keys:
                    logger.debug(
                        f"{server_name}: Using user-specified auth header(s); skipping OAuth provider.",
                        user_auth_headers=sorted(user_auth_keys),
                    )
                channel_hook = None
                if transport_metrics is not None:

                    def channel_hook(event):
                        try:
                            transport_metrics.record_event(event)
                        except Exception:  # pragma: no cover - defensive guard
                            logger.debug(
                                "%s: transport metrics hook failed",
                                server_name,
                                exc_info=True,
                            )

                return tracking_sse_client(
                    config.url,
                    headers,
                    sse_read_timeout=config.read_transport_sse_timeout_seconds,
                    auth=oauth_auth,
                    channel_hook=channel_hook,
                )
            elif config.transport == "http":
                if not config.url:
                    raise ValueError(
                        f"Server '{server_name}' uses http transport but no url is specified"
                    )
                headers, oauth_auth, user_auth_keys = _prepare_headers_and_auth(config)
                if user_auth_keys:
                    logger.debug(
                        f"{server_name}: Using user-specified auth header(s); skipping OAuth provider.",
                        user_auth_headers=sorted(user_auth_keys),
                    )
                channel_hook = None
                if transport_metrics is not None:

                    def channel_hook(event):
                        try:
                            transport_metrics.record_event(event)
                        except Exception:  # pragma: no cover - defensive guard
                            logger.debug(
                                "%s: transport metrics hook failed",
                                server_name,
                                exc_info=True,
                            )

                timeout = None
                if (
                    config.http_timeout_seconds is not None
                    or config.http_read_timeout_seconds is not None
                ):
                    timeout = httpx.Timeout(
                        config.http_timeout_seconds or MCP_DEFAULT_TIMEOUT,
                        read=config.http_read_timeout_seconds or MCP_DEFAULT_SSE_READ_TIMEOUT,
                    )

                http_client = create_mcp_http_client(
                    headers=headers,
                    auth=oauth_auth,
                    timeout=timeout,
                )
                return tracking_streamablehttp_client(
                    config.url,
                    http_client=http_client,
                    channel_hook=channel_hook,
                )
            else:
                raise ValueError(f"Unsupported transport: {config.transport}")

        server_conn = ServerConnection(
            server_name=server_name,
            server_config=config,
            transport_context_factory=transport_context_factory,
            client_session_factory=client_session_factory,
        )

        if transport_metrics is not None:
            server_conn.transport_metrics = transport_metrics

        async with self._lock:
            # Check if already running
            if server_name in self.running_servers:
                return self.running_servers[server_name]

            self.running_servers[server_name] = server_conn
            assert self._tg is not None
            self._tg.start_soon(_server_lifecycle_task, server_conn)

        logger.info(f"{server_name}: Up and running with a persistent connection!")
        return server_conn

    async def get_server(
        self,
        server_name: str,
        client_session_factory: Callable,
    ) -> ServerConnection:
        """
        Get a running server instance, launching it if needed.
        """
        # Get the server connection if it's already running and healthy
        async with self._lock:
            server_conn = self.running_servers.get(server_name)
            if server_conn and server_conn.is_healthy():
                return server_conn

            # If server exists but isn't healthy, remove it so we can create a new one
            if server_conn:
                logger.info(f"{server_name}: Server exists but is unhealthy, recreating...")
                self.running_servers.pop(server_name)
                server_conn.request_shutdown()

        # Launch the connection
        server_conn = await self.launch_server(
            server_name=server_name,
            client_session_factory=client_session_factory,
        )

        # Wait until it's fully initialized, or an error occurs
        await server_conn.wait_for_initialized()

        # Check if the server is healthy after initialization
        if not server_conn.is_healthy():
            error_msg = server_conn._error_message or "Unknown error"

            # Format the error message for better display
            if isinstance(error_msg, list):
                # Join the list with newlines for better readability
                formatted_error = "\n".join(error_msg)
            else:
                formatted_error = str(error_msg)

            raise ServerInitializationError(
                f"MCP Server: '{server_name}': Failed to initialize - see details. Check fastagent.config.yaml?",
                formatted_error,
            )

        return server_conn

    async def get_server_capabilities(self, server_name: str) -> ServerCapabilities | None:
        """Get the capabilities of a specific server."""
        server_conn = await self.get_server(
            server_name, client_session_factory=MCPAgentClientSession
        )
        return server_conn.server_capabilities if server_conn else None

    async def disconnect_server(self, server_name: str) -> None:
        """
        Disconnect a specific server if it's running under this connection manager.
        """
        logger.info(f"{server_name}: Disconnecting persistent connection to server...")

        async with self._lock:
            server_conn = self.running_servers.pop(server_name, None)
        if server_conn:
            server_conn.request_shutdown()
            logger.info(f"{server_name}: Shutdown signal sent (lifecycle task will exit).")
        else:
            logger.info(f"{server_name}: No persistent connection found. Skipping server shutdown")

    async def reconnect_server(
        self,
        server_name: str,
        client_session_factory: Callable,
    ) -> "ServerConnection":
        """
        Force reconnection to a server by disconnecting and re-establishing the connection.

        This is used when a session has been terminated (e.g., 404 from server restart)
        and we need to create a fresh connection with a new session.

        Args:
            server_name: Name of the server to reconnect
            client_session_factory: Factory function to create client sessions

        Returns:
            The new ServerConnection instance
        """
        logger.info(f"{server_name}: Initiating reconnection...")

        # First, disconnect the existing connection
        await self.disconnect_server(server_name)

        # Brief pause to allow cleanup
        await asyncio.sleep(0.1)

        # Launch a fresh connection
        server_conn = await self.launch_server(
            server_name=server_name,
            client_session_factory=client_session_factory,
        )

        # Wait for initialization
        await server_conn.wait_for_initialized()

        # Check if the reconnection was successful
        if not server_conn.is_healthy():
            error_msg = server_conn._error_message or "Unknown error during reconnection"
            if isinstance(error_msg, list):
                formatted_error = "\n".join(error_msg)
            else:
                formatted_error = str(error_msg)

            raise ServerInitializationError(
                f"MCP Server: '{server_name}': Failed to reconnect - see details.",
                formatted_error,
            )

        logger.info(f"{server_name}: Reconnection successful")
        return server_conn

    async def disconnect_all(self) -> None:
        """Disconnect all servers that are running under this connection manager."""
        # Get a copy of servers to shutdown
        servers_to_shutdown = []

        async with self._lock:
            if not self.running_servers:
                return

            # Make a copy of the servers to shut down
            servers_to_shutdown = list(self.running_servers.items())
            # Clear the dict immediately to prevent any new access
            self.running_servers.clear()

        # Release the lock before waiting for servers to shut down
        for name, conn in servers_to_shutdown:
            logger.info(f"{name}: Requesting shutdown...")
            conn.request_shutdown()
