MCP Client

class MCPClientMixin:
    """Mixin to add MCP client capabilities to agents.

    This mixin allows any agent to connect to MCP servers and use their tools.
    MCP tools are automatically registered in the agent's _TOOL_REGISTRY.

    The mixin initializes an MCPClientManager instance and provides convenience
    methods for server management and tool registration.

    Attributes:
        _mcp_manager (MCPClientManager): Internal manager for MCP connections

    Usage:
        class MyAgent(Agent, MCPClientMixin):
            def __init__(self, ...):
                super().__init__(...)
                self.connect_mcp_server("github", {
                    "command": "npx",
                    "args": ["-y", "@modelcontextprotocol/server-github"],
                    "env": {"GITHUB_TOKEN": "ghp_xxx"}
                })
    """

def __init__(self, *args, **kwargs):
    """Initialize the mixin.

    Creates an MCPClientManager instance. Inherits debug mode from agent if available.

    Args:
        *args: Passed to parent class
        **kwargs: Passed to parent class (debug: bool checked for manager)
    """

def connect_mcp_server(self, name: str, config: Dict) -> bool:
    """Connect to an MCP server and register its tools.

    This method:
    1. Adds the server via MCPClientManager
    2. Registers all tools from the server in _TOOL_REGISTRY
    3. Rebuilds the agent's system prompt (if rebuild_system_prompt exists)

    Args:
        name (str): Friendly name for the server. Used for namespacing tools
            (e.g., "filesystem" → tools become "mcp_filesystem_*")
        config (Dict): Server configuration dict with:
            - command (required): Base command to run (e.g., "npx", "uvx", "python")
            - args (optional): List of arguments to pass to the command
            - env (optional): Environment variables dict to pass to subprocess

    Returns:
        bool: True if connection and tool registration successful, False otherwise

    Raises:
        ValueError: If config is not a dict or missing 'command' field

    Example:
        >>> success = agent.connect_mcp_server("filesystem", {
        ...     "command": "npx",
        ...     "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        ... })
        >>> if success:
        ...     print("Connected! Tools available.")

        # With environment variables:
        >>> agent.connect_mcp_server("github", {
        ...     "command": "npx",
        ...     "args": ["-y", "@modelcontextprotocol/server-github"],
        ...     "env": {"GITHUB_TOKEN": "ghp_xxx"}
        ... })
    """

def load_mcp_servers_from_config(self) -> int:
    """Load and register MCP servers from configuration file.

    This is a convenience method that:
    1. Loads servers from ~/.gaia/mcp_servers.json via the manager
    2. Registers all tools from loaded servers in _TOOL_REGISTRY
    3. Rebuilds system prompt once (after all servers loaded)

    Returns:
        int: Number of servers successfully loaded and registered

    Example:
        >>> agent = MyAgent()
        >>> count = agent.load_mcp_servers_from_config()
        >>> print(f"Loaded {count} MCP servers")
    """

def disconnect_mcp_server(self, name: str) -> None:
    """Disconnect from an MCP server and unregister its tools.

    This method:
    1. Unregisters all tools from the server from _TOOL_REGISTRY
    2. Disconnects and removes the server via MCPClientManager

    Args:
        name (str): Server name to disconnect

    Note:
        Logs a warning if server not found; does not raise an exception.
    """

def list_mcp_servers(self) -> List[str]:
    """List all connected MCP servers.

    Returns:
        List[str]: Server names currently connected
    """

def get_mcp_client(self, name: str) -> MCPClient:
    """Get an MCP client by name.

    Args:
        name (str): Server name

    Returns:
        MCPClient: Client instance if exists, None otherwise
    """

def _register_mcp_tools(self, client: MCPClient) -> None:
    """Register all tools from an MCP server into _TOOL_REGISTRY.

    MCP tool responses are wrapped in GAIA-style format with
    status/message/data/instruction fields to help the LLM
    understand how to interpret and respond to the data.

    Response wrapping format:
    - Success: {"status": "success", "message": "...", "data": {...}, "instruction": "..."}
    - Error: {"status": "error", "error": "...", "data": {...}}
    - Non-dict responses: Passed through unchanged

    Args:
        client (MCPClient): MCPClient instance to register tools from
    """

def _unregister_mcp_tools(self, client: MCPClient) -> None:
    """Unregister all tools from an MCP server.

    Args:
        client (MCPClient): MCPClient instance to unregister tools from
    """

class MCPClient:
    """Client for interacting with an MCP server.

    Handles connection, tool discovery, and tool execution via the transport layer.

    Attributes:
        name (str): Friendly name for this server
        transport (MCPTransport): Transport implementation
        debug (bool): Debug logging enabled
        server_info (Dict[str, Any]): Server information from initialize response
    """

def __init__(self, name: str, transport: MCPTransport, debug: bool = False):
    """Create an MCP client with specified transport.

    Args:
        name (str): Friendly name for this server (used in tool namespacing)
        transport (MCPTransport): Transport implementation to use
        debug (bool): Enable debug logging (default: False)

    Attributes initialized:
        server_info: Empty dict (populated on connect)
        _tools: None (populated on list_tools)
    """

@classmethod
def from_config(
    cls,
    name: str,
    config: Dict[str, Any],
    timeout: int = 30,
    debug: bool = False
) -> "MCPClient":
    """Create an MCP client from a config dict (Anthropic format).

    Factory method for the common case of subprocess-based MCP servers.
    Uses the same configuration format as Claude Desktop.

    Args:
        name (str): Friendly name for this server
        config (Dict[str, Any]): Server configuration dict with:
            - command (required): Base command to run (e.g., "npx", "uvx")
            - args (optional): List of arguments
            - env (optional): Environment variables dict
        timeout (int): Request timeout in seconds (default: 30)
        debug (bool): Enable debug logging (default: False)

    Returns:
        MCPClient: Configured client instance (not yet connected)

    Raises:
        ValueError: If config is missing required 'command' field

    Example:
        >>> client = MCPClient.from_config(
        ...     "filesystem",
        ...     {
        ...         "command": "npx",
        ...         "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        ...     },
        ...     timeout=60,
        ...     debug=True
        ... )
        >>> client.connect()

        # With environment variables:
        >>> client = MCPClient.from_config(
        ...     "github",
        ...     {
        ...         "command": "npx",
        ...         "args": ["-y", "@modelcontextprotocol/server-github"],
        ...         "env": {"GITHUB_TOKEN": "ghp_xxx"}
        ...     }
        ... )
    """

@classmethod
def from_command(
    cls,
    name: str,
    command: str,
    timeout: int = 30,
    debug: bool = False
) -> "MCPClient":
    """Create an MCP client using stdio transport (legacy method).

    Note: Prefer from_config() for new code. This method uses shell string
    parsing which doesn't support environment variables or complex arguments.

    Args:
        name (str): Friendly name for this server
        command (str): Shell command to start the server
        timeout (int): Request timeout in seconds (default: 30)
        debug (bool): Enable debug logging (default: False)

    Returns:
        MCPClient: Configured client instance (not yet connected)
    """

def connect(self) -> bool:
    """Connect to the MCP server and initialize.

    This method:
    1. Establishes transport connection (starts subprocess)
    2. Sends MCP initialize request with client info
    3. Stores server_info from response

    Initialize request format:
        {
            "protocolVersion": "1.0.0",
            "clientInfo": {"name": "GAIA MCP Client", "version": "0.15.2"},
            "capabilities": {}
        }

    Returns:
        bool: True if connection and initialization successful

    Note:
        Disconnects automatically if initialization fails.
    """

def disconnect(self) -> None:
    """Disconnect from the MCP server.

    Disconnects transport and clears cached tools.
    """

def is_connected(self) -> bool:
    """Check if connected to server.

    Returns:
        bool: True if transport is connected
    """

def list_tools(self, refresh: bool = False) -> List[MCPTool]:
    """List all available tools from the server.

    Sends tools/list request and caches the result.

    Args:
        refresh (bool): Force refresh from server (default: False, uses cache)

    Returns:
        List[MCPTool]: Available tools (empty list on error)

    Note:
        First call fetches from server; subsequent calls return cached result
        unless refresh=True.
    """

def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Call a tool on the MCP server.

    Sends tools/call request and handles response.

    Args:
        tool_name (str): Name of the tool to call (original MCP name, not namespaced)
        arguments (Dict[str, Any]): Tool arguments

    Returns:
        Dict[str, Any]: Tool response. Format:
            - Success: {"content": [...], ...} (MCP result format)
            - Error: {"error": "error message"}

    Raises:
        RuntimeError: If not connected to server

    Note:
        Validation errors (code -32602) are enhanced with schema context.
    """

def create_tool_wrapper(self, tool: MCPTool) -> Callable[..., Dict[str, Any]]:
    """Create a callable wrapper for an MCP tool.

    The wrapper accepts **kwargs to be compatible with GAIA's
    agent._execute_tool() which calls tool(**tool_args).

    Args:
        tool (MCPTool): MCPTool to wrap

    Returns:
        Callable[..., Dict[str, Any]]: Function that calls the MCP tool
            Signature: (**kwargs) -> Dict[str, Any]
    """

class MCPClientManager:
    """Manages multiple MCP client connections.

    Handles configuration loading/saving, connection management, and
    routing tool calls to the appropriate server.

    Attributes:
        config (MCPConfig): Configuration manager instance
        debug (bool): Debug logging enabled
        _clients (Dict[str, MCPClient]): Connected clients by name
    """

def __init__(self, config: Optional[MCPConfig] = None, debug: bool = False):
    """Create manager with optional configuration.

    Args:
        config (MCPConfig, optional): Configuration instance.
            Creates default MCPConfig if not provided.
        debug (bool): Enable debug logging (default: False)
    """

def add_server(self, name: str, config: Dict) -> MCPClient:
    """Add and connect to an MCP server.

    This method:
    1. Creates an MCPClient from the config
    2. Connects to the server
    3. Stores the client
    4. Saves configuration to file

    Args:
        name (str): Friendly name for the server
        config (Dict): Server configuration dict with:
            - command (required): Base command to run
            - args (optional): List of arguments
            - env (optional): Environment variables dict

    Returns:
        MCPClient: Connected client instance

    Raises:
        ValueError: If server with this name already exists
        ValueError: If config is missing 'command' field
        RuntimeError: If connection fails

    Example:
        >>> client = manager.add_server("filesystem", {
        ...     "command": "npx",
        ...     "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        ... })

        # With environment variables:
        >>> client = manager.add_server("github", {
        ...     "command": "npx",
        ...     "args": ["-y", "@modelcontextprotocol/server-github"],
        ...     "env": {"GITHUB_TOKEN": "ghp_xxx"}
        ... })
    """

def remove_server(self, name: str) -> None:
    """Remove and disconnect from an MCP server.

    This method:
    1. Disconnects the client
    2. Removes from internal storage
    3. Removes from configuration file

    Args:
        name (str): Name of the server to remove

    Note:
        Logs warning if server not found; does not raise exception.
    """

def get_client(self, name: str) -> Optional[MCPClient]:
    """Get a client by name.

    Args:
        name (str): Server name

    Returns:
        MCPClient: Client instance if exists, None otherwise
    """

def list_servers(self) -> List[str]:
    """List all registered server names.

    Returns:
        List[str]: Server names (connected clients only)
    """

def disconnect_all(self) -> None:
    """Disconnect from all MCP servers.

    Iterates through all clients and disconnects. Logs errors but continues
    to disconnect remaining servers.
    """

def load_from_config(self) -> None:
    """Load and connect to all servers from configuration.

    Reads servers from configuration file and attempts to connect to each.
    Skips servers that:
    - Are already connected
    - Have no command specified
    - Fail to connect (logs warning/error)

    Note:
        Does not raise exceptions; logs errors for individual failures.
    """

class MCPConfig:
    """Configuration manager for MCP servers.

    Stores server configurations in a JSON file for persistence.
    Auto-saves on modification.

    Attributes:
        config_file (Path): Path to configuration file
        _servers (Dict[str, Dict[str, Any]]): Server configurations

    File format (Anthropic-compatible):
        {
            "mcpServers": {
                "server_name": {
                    "command": "npx",
                    "args": ["-y", "@modelcontextprotocol/server-..."],
                    "env": {"TOKEN": "xxx"}  // optional
                }
            }
        }

    Note: Also supports legacy "servers" key for backwards compatibility.
    """

def __init__(self, config_file: str = None):
    """Create config manager.

    Args:
        config_file (str, optional): Path to configuration file.
            Default: ~/.gaia/mcp_servers.json

    Note:
        Creates ~/.gaia directory if it doesn't exist.
        Loads existing configuration on initialization.
    """

def add_server(self, name: str, config: Dict[str, Any]) -> None:
    """Add or update a server configuration.

    Auto-saves to file after modification.

    Args:
        name (str): Server name
        config (Dict[str, Any]): Server configuration dictionary
            Required keys: "command" (str)
    """

def remove_server(self, name: str) -> None:
    """Remove a server configuration.

    Auto-saves to file after modification (if server existed).

    Args:
        name (str): Server name
    """

def get_server(self, name: str) -> Dict[str, Any]:
    """Get a server configuration.

    Args:
        name (str): Server name

    Returns:
        Dict[str, Any]: Server configuration, or empty dict if not found
    """

def get_servers(self) -> Dict[str, Dict[str, Any]]:
    """Get all server configurations.

    Returns:
        Dict[str, Dict[str, Any]]: Copy of all server configurations
    """

def server_exists(self, name: str) -> bool:
    """Check if a server exists in configuration.

    Args:
        name (str): Server name

    Returns:
        bool: True if server exists
    """

class MCPTransport(ABC):
    """Abstract base class for MCP transport implementations.

    Transport layer handles the communication protocol between GAIA and MCP servers.
    Different implementations support different connection methods (stdio, HTTP, etc.).
    """

@abstractmethod
def connect(self) -> bool:
    """Establish connection to the MCP server.

    Returns:
        bool: True if connection successful, False otherwise
    """

@abstractmethod
def disconnect(self) -> None:
    """Close the connection to the MCP server."""

@abstractmethod
def send_request(
    self, method: str, params: Dict[str, Any] = None
) -> Dict[str, Any]:
    """Send a JSON-RPC request to the server.

    Args:
        method (str): The JSON-RPC method name (e.g., "initialize", "tools/list")
        params (Dict[str, Any], optional): Parameters dictionary

    Returns:
        Dict[str, Any]: JSON-RPC response containing "result" or "error"

    Raises:
        RuntimeError: If connection is not established
        TimeoutError: If request times out
    """

@abstractmethod
def is_connected(self) -> bool:
    """Check if transport is currently connected.

    Returns:
        bool: True if connected, False otherwise
    """

class StdioTransport(MCPTransport):
    """Stdio-based transport using subprocess for MCP servers.

    This transport launches MCP servers as subprocesses and communicates
    via stdin/stdout using JSON-RPC messages.

    Attributes:
        command (str): Program to start the server
        args (list[str] | None): Arguments passed to the program
        timeout (int): Request timeout in seconds
        debug (bool): Debug logging enabled
        _process (subprocess.Popen): Subprocess instance
        _request_id (int): Auto-incrementing request ID
    """

def __init__(
    self,
    command: str,
    args: list[str] | None = None,
    env: dict[str, str] | None = None,
    timeout: int = 30,
    debug: bool = False,
):
    """Create stdio transport.

    Args:
        command (str): Program to start the MCP server, e.g. "npx".
            A full command line ("npx -y @scope/server") is also accepted and
            tokenized into an argv list — it is never run through a shell.
            `~` and environment variables in that string are still expanded.
        args (list[str] | None): Arguments for the program, e.g.
            ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        env (dict[str, str] | None): Extra environment variables, merged
            over the current process environment
        timeout (int): Request timeout in seconds (default: 30)
        debug (bool): Enable debug logging (default: False)
    """

def connect(self) -> bool:
    """Launch the MCP server subprocess.

    Uses subprocess.Popen with:
    - shell=False, always — the command is resolved to an argv list first
    - stdin/stdout/stderr pipes
    - text mode with UTF-8 encoding
    - Line buffering (bufsize=1)

    Returns:
        bool: True if process started successfully

    Raises:
        ValueError: If a single-string command is blank, has unbalanced
            quotes, or carries unquoted shell syntax (";", "|", "&&", ">",
            "$(", ...). Such syntax is rejected rather than passed through
            as literal argv text — split the program from its arguments
            using `args` instead. The same characters *inside a quoted
            argument* are ordinary text and are accepted. A backtick is the
            exception: it is rejected wherever it appears, because it never
            forms a token of its own and quoting state cannot be recovered
            after tokenizing. Pass such a value via `args`.

    Note:
        Returns True without starting new process if already connected.
    """

def disconnect(self) -> None:
    """Terminate the MCP server subprocess.

    Process cleanup sequence:
    1. Send SIGTERM
    2. Wait up to 5 seconds
    3. If still running, send SIGKILL
    4. Wait for process to exit

    Note:
        Safe to call multiple times; no-op if not connected.
    """

def send_request(
    self, method: str, params: Dict[str, Any] = None
) -> Dict[str, Any]:
    """Send a JSON-RPC request via stdin and read response from stdout.

    Request format (JSON-RPC 2.0):
        {"jsonrpc": "2.0", "id": <int>, "method": <str>, "params": <dict>}

    Args:
        method (str): JSON-RPC method name
        params (Dict[str, Any], optional): Parameters dictionary

    Returns:
        Dict[str, Any]: JSON-RPC response

    Raises:
        RuntimeError: If not connected or process died
        ValueError: If response is invalid JSON
    """

def is_connected(self) -> bool:
    """Check if the subprocess is running.

    Returns:
        bool: True if process is alive (poll() returns None)
    """

@dataclass
class MCPTool:
    """Represents an MCP tool with its schema.

    Attributes:
        name (str): Tool name from MCP server (e.g., "read_file")
        description (str): Tool description
        input_schema (Dict[str, Any]): MCP inputSchema (JSON Schema format)
        annotations (Dict[str, Any]): MCP tool annotations (readOnlyHint,
            destructiveHint, …). Server-supplied and therefore advisory —
            consumed by the confirmation classifier below.
    """

    name: str
    description: str
    input_schema: Dict[str, Any]
    annotations: Dict[str, Any] = field(default_factory=dict)

def mcp_tool_requires_confirmation(name: str, annotations: Dict[str, Any]) -> bool:
    """Whether an MCP tool must be user-confirmed before GAIA executes it.

    Fails closed. Exempt only when the server proves the tool is read-only:
    ``annotations.readOnlyHint`` is exactly ``True`` AND the tool name carries
    no state-changing verb token. No annotations, a non-boolean hint, or a
    read-only claim contradicted by a name like ``delete_file`` all require
    confirmation.
    """

def to_gaia_format(self, prefix: str, raw_server_name: Optional[str] = None) -> Dict[str, Any]:
    """Convert MCP tool schema to GAIA _TOOL_REGISTRY format.

    Conversion details:
    - Name: "mcp_{prefix}_{tool_name}"
    - Description: "[MCP:{prefix}] {description}"
    - Parameters: Converted from JSON Schema to GAIA format
    - Metadata: _mcp_server, requires_confirmation added

    Args:
        prefix (str): Sanitised server-name prefix (from MCPClient.prefix),
            used for the registry name and the description tag
        raw_server_name (Optional[str]): Original mcp_servers.json key, stored
            in _mcp_server for routing. Defaults to prefix.

    Returns:
        Dict[str, Any]: GAIA tool registry entry (without function field)

    Example:
        MCP tool:
            {"name": "read_file", "description": "Read file", "inputSchema": {...}}

        GAIA format:
            {
                "name": "mcp_filesystem_read_file",
                "display_name": "read_file (filesystem)",
                "description": "[MCP:filesystem] Read file",
                "parameters": {...},
                "atomic": True,
                "requires_confirmation": True,
                "_mcp_server": "filesystem"
            }
    """

# MCP tool format (from server)
{
    "name": "read_file",
    "description": "Read file contents",
    "inputSchema": {
        "type": "object",
        "properties": {
            "path": {"type": "string", "description": "File path"},
            "encoding": {"type": "string", "description": "Encoding"}
        },
        "required": ["path"]
    }
}

# GAIA tool format (in _TOOL_REGISTRY)
{
    "name": "mcp_filesystem_read_file",
    "display_name": "read_file (filesystem)",
    "description": "[MCP:filesystem] Read file contents",
    "parameters": {
        "path": {
            "type": "string",
            "required": True,
            "description": "File path"
        },
        "encoding": {
            "type": "string",
            "required": False,
            "description": "Encoding"
        }
    },
    "function": <wrapper_callable>,
    "atomic": True,
    "requires_confirmation": True,
    "_mcp_server": "filesystem"
}

# Both work without conflicts:
# - mcp_fs1_read_file → reads from filesystem
# - mcp_github_read_file → reads from GitHub

# Successful response
{
    "status": "success",
    "message": "Tool 'read_file' returned data",
    "data": {"content": [{"type": "text", "text": "file contents..."}]},
    "instruction": "Parse this JSON data and provide a human-readable summary. "
                   "Your answer must be a plain text string, not a JSON object."
}

# Error response
{
    "status": "error",
    "error": "File not found: /nonexistent",
    "data": {"error": "File not found: /nonexistent"}
}

{
    "jsonrpc": "2.0",
    "id": 1,
    "method": "tools/call",
    "params": {
        "name": "read_file",
        "arguments": {"path": "/tmp/file.txt"}
    }
}

{
    "jsonrpc": "2.0",
    "id": 1,
    "result": {
        "content": [{"type": "text", "text": "file contents"}]
    }
}

{
    "jsonrpc": "2.0",
    "id": 1,
    "error": {
        "code": -32602,
        "message": "Invalid params: path is required"
    }
}

MCP tool 'read_file' input validation failed.

Error details:
Invalid params: missing required parameter 'path'

Expected tool schema:
{
  "type": "object",
  "properties": {
    "path": {"type": "string", "description": "File path"},
    "encoding": {"type": "string", "description": "Encoding"}
  },
  "required": ["path"]
}

Your arguments:
{
  "encoding": "utf-8"
}

{
  "mcpServers": {
    "filesystem": {
      "command": "npx",
      "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
    },
    "github": {
      "command": "npx",
      "args": ["-y", "@modelcontextprotocol/server-github"],
      "env": {
        "GITHUB_TOKEN": "ghp_xxx"
      }
    },
    "time": {
      "command": "uvx",
      "args": ["mcp-server-time"]
    }
  }
}

agent.connect_mcp_server("github", {
    "command": "npx",
    "args": ["-y", "@modelcontextprotocol/server-github"],
    "env": {"GITHUB_TOKEN": "ghp_xxx"}
})

{
  "mcpServers": {
    "github": {
      "command": "npx",
      "args": ["-y", "@modelcontextprotocol/server-github"],
      "env": {
        "GITHUB_TOKEN": "ghp_xxx"
      }
    }
  }
}

# Add server (command line - tokenized into the config format internally)
gaia mcp add <name> "<command>"

# Examples
gaia mcp add time "uvx mcp-server-time"
gaia mcp add filesystem "npx -y @modelcontextprotocol/server-filesystem /tmp"

# List servers
gaia mcp list

# List tools from a server
gaia mcp tools <name>

# Test connection
gaia mcp test-client <name>

# Remove server
gaia mcp remove <name>

# Run all MCP client unit tests
python -m pytest tests/unit/mcp/client/ -v

# Run specific test file
python -m pytest tests/unit/mcp/client/test_mcp_client.py -v

# Run with coverage
python -m pytest tests/unit/mcp/client/ --cov=src/gaia/mcp/client

@pytest.fixture
def mock_transport(mocker):
    """Create a mock transport."""
    transport = mocker.Mock(spec=MCPTransport)
    transport.connect.return_value = True
    transport.is_connected.return_value = True
    transport.send_request.return_value = {
        "result": {"tools": [{"name": "test_tool", "description": "Test", "inputSchema": {}}]}
    }
    return transport

def test_client_list_tools(mock_transport):
    client = MCPClient("test", mock_transport)
    client.connect()
    tools = client.list_tools()
    assert len(tools) == 1
    assert tools[0].name == "test_tool"

from gaia.agents.base.agent import Agent
from gaia.mcp import MCPClientMixin

class MyAgent(Agent, MCPClientMixin):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Connect to filesystem MCP server
        self.connect_mcp_server("filesystem", {
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        })

# Use agent - MCP tools available immediately
agent = MyAgent()
result = agent.run("List all files in /tmp")
# LLM can now use: mcp_filesystem_read_file, mcp_filesystem_write_file, etc.

from gaia.agents.base.agent import Agent
from gaia.mcp import MCPClientMixin

class MultiToolAgent(Agent, MCPClientMixin):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Connect to multiple MCP servers
        self.connect_mcp_server("filesystem", {
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        })

        self.connect_mcp_server("github", {
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-github"],
            "env": {"GITHUB_TOKEN": "ghp_xxx"}
        })

# All tools from both servers are available
agent = MultiToolAgent()
response = agent.run("List files in /tmp, then create a GitHub issue about them")

from gaia.agents.base.agent import Agent
from gaia.mcp import MCPClientMixin

class ConfigAgent(Agent, MCPClientMixin):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Load all servers from ~/.gaia/mcp_servers.json
        count = self.load_mcp_servers_from_config()
        print(f"Loaded {count} MCP servers")

agent = ConfigAgent()
servers = agent.list_mcp_servers()
print(f"Connected to: {', '.join(servers)}")

from gaia.mcp.client import MCPClient

# Create client from config
client = MCPClient.from_config(
    name="filesystem",
    config={
        "command": "npx",
        "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
    },
    timeout=30,
    debug=False
)

# Connect to server
if client.connect():
    # List available tools
    tools = client.list_tools()
    for tool in tools:
        print(f"- {tool.name}: {tool.description}")

    # Call tool directly
    result = client.call_tool("read_file", {"path": "/tmp/example.txt"})
    if "content" in result:
        print(result["content"][0]["text"])
    elif "error" in result:
        print(f"Error: {result['error']}")

    # Disconnect when done
    client.disconnect()

from gaia.mcp.client import MCPClientManager
from gaia.mcp.client.config import MCPConfig

# Create manager with config file
config = MCPConfig("~/.gaia/mcp_servers.json")
manager = MCPClientManager(config=config)

# Add servers programmatically
client = manager.add_server("filesystem", {
    "command": "npx",
    "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
})

# List all connected servers
servers = manager.list_servers()
print(f"Servers: {servers}")

# Get specific client
fs_client = manager.get_client("filesystem")
if fs_client and fs_client.is_connected():
    tools = fs_client.list_tools()
    print(f"Tools available: {len(tools)}")

# Load additional servers from config file
manager.load_from_config()

# Cleanup all connections
manager.disconnect_all()

from gaia.agents.base.agent import Agent
from gaia.mcp import MCPClientMixin

class SafeAgent(Agent, MCPClientMixin):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        # Check connection success
        success = self.connect_mcp_server("filesystem", {
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        })

        if not success:
            print("Warning: Failed to connect to MCP server")
        else:
            print("Successfully connected to MCP server")

# Tool execution errors are returned in result
agent = SafeAgent()
client = agent.get_mcp_client("filesystem")
result = client.call_tool("read_file", {"path": "/nonexistent"})

if "error" in result:
    print(f"Tool failed: {result['error']}")
elif "content" in result:
    print(f"Success: {result['content']}")

# Via MCPClient
client = MCPClient.from_config("server", {
    "command": "npx",
    "args": ["-y", "@mcp/server"]
}, debug=True)

# Via MCPClientManager
manager = MCPClientManager(debug=True)

# Via MCPClientMixin (in agent)
class DebugAgent(Agent, MCPClientMixin):
    def __init__(self, **kwargs):
        kwargs["debug"] = True  # Enables MCP debug logging
        super().__init__(**kwargs)
        self.connect_mcp_server("server", {
            "command": "npx",
            "args": ["-y", "@mcp/server"]
        })

from gaia.mcp.client.transports.base import MCPTransport

class HttpTransport(MCPTransport):
    def __init__(self, url: str, timeout: int = 30):
        self.url = url
        self.timeout = timeout
        self._connected = False

    def connect(self) -> bool:
        # Implement connection logic
        self._connected = True
        return True

    def disconnect(self) -> None:
        self._connected = False

    def send_request(self, method: str, params: Dict = None) -> Dict:
        # Implement HTTP request
        pass

    def is_connected(self) -> bool:
        return self._connected

# Override _register_mcp_tools in your agent
class CustomAgent(Agent, MCPClientMixin):
    def _register_mcp_tools(self, client: MCPClient) -> None:
        tools = client.list_tools()
        for tool in tools:
            gaia_tool = tool.to_gaia_format(client.name)

            # Custom wrapper with additional logic
            def custom_wrapper(tool=tool, **kwargs):
                result = client.call_tool(tool.name, kwargs)
                # Add custom processing
                return result

            gaia_tool["function"] = custom_wrapper
            _TOOL_REGISTRY[gaia_tool["name"]] = gaia_tool