/aisuite/mcp/client.py

"""
MCP Client for aisuite.

This module provides the MCPClient class that connects to MCP servers and
exposes their tools as Python callables compatible with aisuite's tool system.
"""

import asyncio
import json
from typing import Any, Callable, Dict, List, Optional
from contextlib import contextmanager

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
    import httpx
except ImportError as e:
    if "mcp" in str(e):
        raise ImportError(
            "MCP support requires the 'mcp' package. "
            "Install it with: pip install 'aisuite[mcp]' or pip install mcp"
        )
    elif "httpx" in str(e):
        raise ImportError(
            "HTTP transport requires the 'httpx' package. "
            "Install it with: pip install httpx"
        )
    raise

from .tool_wrapper import create_mcp_tool_wrapper
from .config import MCPConfig, validate_mcp_config, get_transport_type


class MCPClient:
    """
    Client for connecting to MCP servers and using their tools with aisuite.

    This class manages the connection to an MCP server, discovers available tools,
    and creates Python callable wrappers that work seamlessly with aisuite's
    existing tool calling infrastructure.

    Example:
        >>> # Connect to an MCP server
        >>> mcp = MCPClient(
        ...     command="npx",
        ...     args=["-y", "@modelcontextprotocol/server-filesystem", "/path"]
        ... )
        >>>
        >>> # Get tools and use with aisuite
        >>> import aisuite as ai
        >>> client = ai.Client()
        >>> response = client.chat.completions.create(
        ...     model="openai:gpt-4o",
        ...     messages=[{"role": "user", "content": "List files"}],
        ...     tools=mcp.get_callable_tools(),
        ...     max_turns=2
        ... )

    The MCPClient handles:
    - Starting and managing the MCP server process
    - Performing the MCP handshake
    - Discovering available tools
    - Creating callable wrappers for tools
    - Executing tool calls via the MCP protocol
    """

    def __init__(
        self,
        command: Optional[str] = None,
        args: Optional[List[str]] = None,
        env: Optional[Dict[str, str]] = None,
        server_url: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout: float = 30.0,
        name: Optional[str] = None,
    ):
        """
        Initialize the MCP client and connect to an MCP server.

        Supports both stdio and HTTP transports. Provide either stdio parameters
        (command) OR HTTP parameters (server_url), but not both.

        Args:
            command: Command to start the MCP server (e.g., "npx", "python") - for stdio transport
            args: Arguments to pass to the command (e.g., ["-y", "server-package"]) - for stdio transport
            env: Optional environment variables for the server process - for stdio transport
            server_url: Base URL of the MCP server (e.g., "http://localhost:8000") - for HTTP transport
            headers: Optional HTTP headers (e.g., for authentication) - for HTTP transport
            timeout: Request timeout in seconds - for HTTP transport (default: 30.0)
            name: Optional name for this MCP client (used for logging and prefixing)

        Raises:
            ImportError: If the mcp or httpx package is not installed
            ValueError: If both stdio and HTTP parameters are provided, or neither
            RuntimeError: If connection to the MCP server fails
        """
        # Validate transport parameters
        has_stdio = command is not None
        has_http = server_url is not None

        if not (has_stdio ^ has_http):
            raise ValueError(
                "Must provide exactly one transport: either 'command' (stdio) or 'server_url' (HTTP)."
            )

        # Store parameters based on transport type
        if has_stdio:
            self.server_params = StdioServerParameters(
                command=command,
                args=args or [],
                env=env,
            )
            self.name = name or command
            # Stdio-specific state
            self._session: Optional[ClientSession] = None
            self._read = None
            self._write = None
            self._stdio_context = None
        else:  # HTTP
            self.server_url = server_url
            self.headers = headers or {}
            self.timeout = timeout
            self.name = name or server_url
            # HTTP-specific state (initialized in _async_connect_http)
            self._http_client = None
            self._request_id = 0
            self._session_id: Optional[str] = None  # MCP session ID from server

        # Shared state
        self._tools_cache: Optional[List[Dict[str, Any]]] = None
        self._event_loop: Optional[asyncio.AbstractEventLoop] = None

        # Initialize connection
        self._connect()

    @classmethod
    def from_config(cls, config: Dict[str, Any]) -> "MCPClient":
        """
        Create an MCPClient from a configuration dictionary.

        This method validates the config and creates an MCPClient instance.
        It supports both stdio and HTTP transports.

        Args:
            config: MCP configuration dictionary

        Returns:
            MCPClient instance

        Raises:
            ValueError: If configuration is invalid

        Example (stdio):
            >>> config = {
            ...     "type": "mcp",
            ...     "name": "filesystem",
            ...     "command": "npx",
            ...     "args": ["-y", "@modelcontextprotocol/server-filesystem", "/docs"]
            ... }
            >>> mcp = MCPClient.from_config(config)

        Example (HTTP):
            >>> config = {
            ...     "type": "mcp",
            ...     "name": "api-server",
            ...     "server_url": "http://localhost:8000",
            ...     "headers": {"Authorization": "Bearer token"}
            ... }
            >>> mcp = MCPClient.from_config(config)
        """
        # Validate and normalize config
        validated_config = validate_mcp_config(config)

        # Determine transport type
        transport = get_transport_type(validated_config)

        if transport == "stdio":
            return cls(
                command=validated_config["command"],
                args=validated_config.get("args", []),
                env=validated_config.get("env"),
                name=validated_config["name"],
            )
        else:  # http
            return cls(
                server_url=validated_config["server_url"],
                headers=validated_config.get("headers"),
                timeout=validated_config.get("timeout", 30.0),
                name=validated_config["name"],
            )

    @staticmethod
    def get_tools_from_config(config: Dict[str, Any]) -> List[Callable]:
        """
        Convenience method to create MCPClient and get callable tools from config.

        This is a helper that combines from_config() and get_callable_tools()
        in a single call. It respects the config's allowed_tools and use_tool_prefix
        settings.

        Args:
            config: MCP configuration dictionary

        Returns:
            List of callable tool wrappers

        Example:
            >>> config = {
            ...     "type": "mcp",
            ...     "name": "filesystem",
            ...     "command": "npx",
            ...     "args": ["..."],
            ...     "allowed_tools": ["read_file"],
            ...     "use_tool_prefix": True
            ... }
            >>> tools = MCPClient.get_tools_from_config(config)
            >>> # Returns callable tools filtered and prefixed per config
        """
        # Validate config first
        validated_config = validate_mcp_config(config)

        # Create client
        client = MCPClient.from_config(validated_config)

        # Get tools with config settings
        tools = client.get_callable_tools(
            allowed_tools=validated_config.get("allowed_tools"),
            use_tool_prefix=validated_config.get("use_tool_prefix", False),
        )

        return tools

    def _connect(self):
        """
        Establish connection to the MCP server.

        This method:
        1. Creates an event loop if needed
        2. Detects transport type (stdio or HTTP)
        3. Establishes connection via appropriate transport
        4. Performs the MCP initialization handshake
        5. Caches the available tools

        Note: Automatically handles Jupyter/IPython environments where an event loop
        is already running by using nest_asyncio.
        """
        # Get or create event loop
        try:
            self._event_loop = asyncio.get_running_loop()
        except RuntimeError:
            self._event_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._event_loop)

        # Enable nested event loops for Jupyter/IPython compatibility
        # This allows run_until_complete() to work in environments where
        # an event loop is already running (like Jupyter notebooks)
        try:
            import nest_asyncio

            nest_asyncio.apply()
        except ImportError:
            # nest_asyncio not available - will work fine in regular Python
            # but may fail in Jupyter. User should install: pip install nest-asyncio
            pass

        # Detect transport type and run appropriate async connection
        if hasattr(self, "server_url"):
            # HTTP transport
            self._event_loop.run_until_complete(self._async_connect_http())
        else:
            # Stdio transport
            self._event_loop.run_until_complete(self._async_connect())

    async def _async_connect(self):
        """Async connection initialization for stdio transport."""
        # Start the MCP server and store the context manager
        self._stdio_context = stdio_client(self.server_params)
        self._read, self._write = await self._stdio_context.__aenter__()

        # Create session
        self._session = ClientSession(self._read, self._write)
        await self._session.__aenter__()

        # Initialize connection
        await self._session.initialize()

        # List available tools and cache them
        tools_result = await self._session.list_tools()

        # Convert Tool objects to dicts for easier handling
        if hasattr(tools_result, "tools"):
            self._tools_cache = [
                {
                    "name": tool.name,
                    "description": (
                        tool.description if hasattr(tool, "description") else ""
                    ),
                    "inputSchema": (
                        tool.inputSchema if hasattr(tool, "inputSchema") else {}
                    ),
                }
                for tool in tools_result.tools
            ]
        else:
            self._tools_cache = []

    async def _parse_sse_response(
        self, response: httpx.Response, request_id: int
    ) -> Dict[str, Any]:
        """
        Parse SSE stream and extract JSON-RPC response.

        SSE format per spec:
            data: {"jsonrpc": "2.0", "id": 1, "result": {...}}

            data: {"jsonrpc": "2.0", "method": "notification", ...}

        The server may send multiple events (notifications, requests) before
        sending the final response. We collect events until we find the
        response matching our request_id.

        Args:
            response: HTTP response with text/event-stream content type
            request_id: The JSON-RPC request ID to match

        Returns:
            Response result dictionary

        Raises:
            RuntimeError: If server returns an error or no matching response found
        """
        result = None

        async for line in response.aiter_lines():
            line = line.strip()

            # Skip empty lines and comments
            if not line or line.startswith(":"):
                continue

            # Parse SSE data field
            if line.startswith("data: "):
                data = line[6:]  # Remove 'data: ' prefix

                try:
                    message = json.loads(data)

                    # Check if this is the response to our request
                    if message.get("id") == request_id:
                        if "error" in message:
                            error = message["error"]
                            raise RuntimeError(
                                f"MCP server error: {error.get('message', 'Unknown error')} "
                                f"(code: {error.get('code', 'unknown')})"
                            )
                        result = message.get("result", {})
                        # Found our response, can stop parsing
                        break

                    # Note: Server may send other notifications/requests
                    # which we ignore for now (future enhancement for bidirectional comms)

                except json.JSONDecodeError:
                    # Invalid JSON in SSE data, skip this event
                    continue

        if result is None:
            raise RuntimeError(
                f"No response received in SSE stream for request {request_id}"
            )

        return result

    async def _send_http_request(
        self, method: str, params: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Send JSON-RPC request to MCP server via HTTP.

        Args:
            method: JSON-RPC method name
            params: Optional parameters

        Returns:
            Response result

        Raises:
            RuntimeError: If HTTP request fails or server returns an error
        """
        # Increment request ID
        self._request_id += 1

        # Build JSON-RPC 2.0 request
        request_data = {
            "jsonrpc": "2.0",
            "id": self._request_id,
            "method": method,
        }

        if params:
            request_data["params"] = params

        # Use the exact server URL provided by the user
        url = self.server_url.rstrip("/")

        # Build headers: MCP requires Accept header with both content types
        # Merge with any user-provided headers and session ID
        request_headers = {
            "Accept": "application/json, text/event-stream",
        }
        if self._session_id:
            request_headers["Mcp-Session-Id"] = self._session_id
        if self.headers:
            request_headers.update(self.headers)

        try:
            response = await self._http_client.post(
                url, json=request_data, headers=request_headers
            )
            response.raise_for_status()

            # Check for MCP session ID in response headers
            if "Mcp-Session-Id" in response.headers and not self._session_id:
                self._session_id = response.headers["Mcp-Session-Id"]

            # Check Content-Type to determine response format
            content_type = response.headers.get("content-type", "").lower()

            if "application/json" in content_type:
                # Handle JSON response (simple request-response)
                result = response.json()

                # Check for JSON-RPC error
                if "error" in result:
                    error = result["error"]
                    raise RuntimeError(
                        f"MCP server error: {error.get('message', 'Unknown error')} "
                        f"(code: {error.get('code', 'unknown')})"
                    )

                return result.get("result", {})

            elif "text/event-stream" in content_type:
                # Handle SSE stream response
                return await self._parse_sse_response(response, request_data["id"])

            else:
                raise RuntimeError(
                    f"Unexpected Content-Type from MCP server: {content_type}"
                )

        except httpx.HTTPError as e:
            raise RuntimeError(
                f"HTTP request to MCP server failed: {type(e).__name__}: {str(e)}"
            )

    async def _send_notification(
        self, method: str, params: Optional[Dict[str, Any]] = None
    ):
        """
        Send a JSON-RPC notification (no response expected).

        Notifications are JSON-RPC messages without an ID field.
        Per the spec, the server should not send a response.

        Args:
            method: JSON-RPC method name
            params: Optional parameters
        """
        # Build JSON-RPC notification (no id field)
        notification = {
            "jsonrpc": "2.0",
            "method": method,
        }

        if params:
            notification["params"] = params

        # Build headers
        url = self.server_url.rstrip("/")
        request_headers = {
            "Accept": "application/json, text/event-stream",
        }
        if self._session_id:
            request_headers["Mcp-Session-Id"] = self._session_id
        if self.headers:
            request_headers.update(self.headers)

        try:
            # Send notification - don't wait for/expect a response
            await self._http_client.post(
                url, json=notification, headers=request_headers
            )
            # Note: We don't check response for notifications
        except httpx.HTTPError:
            # Notifications may timeout or fail, which is acceptable
            pass

    async def _async_connect_http(self):
        """Async connection initialization for HTTP transport."""
        # Create HTTP client
        self._http_client = httpx.AsyncClient(timeout=self.timeout)

        # Send initialize request
        init_params = {
            "protocolVersion": "2024-11-05",
            "capabilities": {"roots": {"listChanged": True}, "sampling": {}},
            "clientInfo": {"name": "aisuite-mcp-client", "version": "1.0.0"},
        }

        await self._send_http_request("initialize", init_params)

        # Send initialized notification (required by MCP spec)
        await self._send_notification("notifications/initialized")

        # List available tools
        tools_result = await self._send_http_request("tools/list")

        # Cache tools
        self._tools_cache = [
            {
                "name": tool["name"],
                "description": tool.get("description", ""),
                "inputSchema": tool.get("inputSchema", {}),
            }
            for tool in tools_result.get("tools", [])
        ]

    def list_tools(self) -> List[Dict[str, Any]]:
        """
        List all available tools from the MCP server.

        Returns:
            List of tool schemas in MCP format

        Example:
            >>> tools = mcp.list_tools()
            >>> for tool in tools:
            ...     print(tool['name'], '-', tool['description'])
        """
        if self._tools_cache is None:
            raise RuntimeError("Not connected to MCP server")
        return self._tools_cache

    def get_callable_tools(
        self,
        allowed_tools: Optional[List[str]] = None,
        use_tool_prefix: bool = False,
    ) -> List[Callable]:
        """
        Get all MCP tools as Python callables compatible with aisuite.

        This is the primary method for using MCP tools with aisuite. It returns
        a list of callable wrappers that can be passed directly to the `tools`
        parameter of `client.chat.completions.create()`.

        Args:
            allowed_tools: Optional list of tool names to include. If None, all tools are included.
            use_tool_prefix: If True, prefix tool names with "{client_name}__"

        Returns:
            List of callable tool wrappers

        Example:
            >>> # Get all tools
            >>> mcp_tools = mcp.get_callable_tools()
            >>>
            >>> # Get specific tools only
            >>> mcp_tools = mcp.get_callable_tools(allowed_tools=["read_file"])
            >>>
            >>> # Get tools with name prefixing
            >>> mcp_tools = mcp.get_callable_tools(use_tool_prefix=True)
            >>> # Tools will be named "filesystem__read_file", etc.
        """
        all_tools = self.list_tools()

        # Filter tools if allowed_tools is specified
        if allowed_tools is not None:
            all_tools = [t for t in all_tools if t["name"] in allowed_tools]

        # Create wrappers
        wrappers = []
        for tool in all_tools:
            wrapper = create_mcp_tool_wrapper(self, tool["name"], tool)

            # Apply prefix if requested
            if use_tool_prefix:
                original_name = wrapper.__name__
                wrapper.__name__ = f"{self.name}__{original_name}"

            wrappers.append(wrapper)

        return wrappers

    def get_tool(self, tool_name: str) -> Optional[Callable]:
        """
        Get a specific MCP tool by name as a Python callable.

        Args:
            tool_name: Name of the tool to retrieve

        Returns:
            Callable wrapper for the tool, or None if not found

        Example:
            >>> read_file = mcp.get_tool("read_file")
            >>> write_file = mcp.get_tool("write_file")
            >>> tools = [read_file, write_file]
        """
        tools = self.list_tools()
        for tool in tools:
            if tool["name"] == tool_name:
                return create_mcp_tool_wrapper(self, tool_name, tool)
        return None

    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """
        Execute an MCP tool call.

        This method is called by MCPToolWrapper when the LLM requests a tool.
        It handles the async MCP protocol communication and returns the result.
        Automatically routes to the appropriate transport (stdio or HTTP).

        Args:
            tool_name: Name of the tool to call
            arguments: Tool arguments as a dictionary

        Returns:
            The result from the MCP tool execution

        Raises:
            RuntimeError: If not connected or tool call fails
        """
        # Detect transport type and route to appropriate method
        if hasattr(self, "_http_client") and self._http_client is not None:
            # HTTP transport
            if self._http_client is None:
                raise RuntimeError("Not connected to MCP server (HTTP)")
            result = self._event_loop.run_until_complete(
                self._async_call_tool_http(tool_name, arguments)
            )
        else:
            # Stdio transport
            if self._session is None:
                raise RuntimeError("Not connected to MCP server (stdio)")
            result = self._event_loop.run_until_complete(
                self._async_call_tool(tool_name, arguments)
            )
        return result

    async def _async_call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """
        Async implementation of tool calling for stdio transport.

        Args:
            tool_name: Name of the tool
            arguments: Tool arguments

        Returns:
            Tool execution result
        """
        result = await self._session.call_tool(tool_name, arguments)

        # Extract content from MCP result
        # MCP returns results in various formats, we try to extract the most useful content
        if hasattr(result, "content"):
            if isinstance(result.content, list) and len(result.content) > 0:
                # Get first content item
                content_item = result.content[0]
                if hasattr(content_item, "text"):
                    return content_item.text
                elif hasattr(content_item, "data"):
                    return content_item.data
                return str(content_item)
            return result.content

        # If no content attribute, return the whole result
        return str(result)

    async def _async_call_tool_http(
        self, tool_name: str, arguments: Dict[str, Any]
    ) -> Any:
        """
        Async implementation of tool calling for HTTP transport.

        Args:
            tool_name: Name of the tool
            arguments: Tool arguments

        Returns:
            Tool execution result
        """
        params = {"name": tool_name, "arguments": arguments}

        result = await self._send_http_request("tools/call", params)

        # Extract content from MCP result (HTTP format)
        # Similar to stdio, but result is already a dict
        if "content" in result:
            content = result["content"]
            if isinstance(content, list) and len(content) > 0:
                # Get first content item
                content_item = content[0]
                if isinstance(content_item, dict):
                    if "text" in content_item:
                        return content_item["text"]
                    elif "data" in content_item:
                        return content_item["data"]
                return str(content_item)
            return content

        # If no content field, return the whole result
        return json.dumps(result)

    def close(self):
        """
        Close the connection to the MCP server.

        Works for both stdio and HTTP transports. It's recommended to use
        the MCPClient as a context manager to ensure proper cleanup, but
        this method can be called manually if needed.

        Example:
            >>> mcp = MCPClient(command="npx", args=["server"])
            >>> try:
            ...     # Use mcp
            ...     pass
            ... finally:
            ...     mcp.close()
        """
        # Check if we need to cleanup (either stdio or HTTP)
        needs_cleanup = (hasattr(self, "_session") and self._session is not None) or (
            hasattr(self, "_http_client") and self._http_client is not None
        )

        if needs_cleanup:
            self._event_loop.run_until_complete(self._async_close())

    async def _async_close(self):
        """Async cleanup for both stdio and HTTP transports."""
        # Cleanup stdio transport
        try:
            if hasattr(self, "_session") and self._session:
                await self._session.__aexit__(None, None, None)
        except RuntimeError as e:
            # Suppress anyio cancel scope errors that occur in Jupyter/nest_asyncio environments
            # This is a known incompatibility between nest_asyncio and anyio task groups
            if "cancel scope" not in str(e).lower():
                raise
        except Exception:
            pass  # Ignore other errors during session cleanup

        try:
            if hasattr(self, "_stdio_context") and self._stdio_context:
                await self._stdio_context.__aexit__(None, None, None)
        except RuntimeError as e:
            # Suppress anyio cancel scope errors that occur in Jupyter/nest_asyncio environments
            # This is a known incompatibility between nest_asyncio and anyio task groups
            if "cancel scope" not in str(e).lower():
                raise
        except Exception:
            pass  # Ignore other errors during stdio cleanup

        # Cleanup HTTP transport
        try:
            if hasattr(self, "_http_client") and self._http_client:
                await self._http_client.aclose()
        except Exception:
            pass  # Ignore errors during HTTP client cleanup

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()
        return False

    def __repr__(self) -> str:
        """String representation."""
        num_tools = len(self._tools_cache) if self._tools_cache else 0
        if hasattr(self, "server_url"):
            return f"MCPClient(server_url={self.server_url!r}, tools={num_tools})"
        else:
            return (
                f"MCPClient(command={self.server_params.command!r}, tools={num_tools})"
            )