MCP Support - Phase 2: Robustness & Testing Improvements

## 🎯 Objective

Enhance the MCP integration added in Phase 1 with comprehensive error handling, testing, monitoring, and advanced configuration options. This phase focuses on making MCP integration production-ready.

## 📋 Prerequisites

- ✅ Phase 1 completed: Issue #50

## 🔧 Implementation Tasks

### 1. Enhanced Error Handling & Validation

#### 1.1 Add Connection Retry Logic
Update `ai/mcp_manager.py` to add retry logic for transient failures:

```python
class MCPManager:
    def __init__(self, config: Config):
        self.config = config
        self._client: Optional[MultiServerMCPClient] = None
        self._tools: list[BaseTool] = []
        self._connected = False
        self.max_retries = 3
        self.retry_delay = 1.0  # seconds
    
    async def connect(self):
        """Initialize MCP client with retry logic."""
        if self._connected:
            return
        
        for attempt in range(self.max_retries):
            try:
                await self._connect_impl()
                return
            except Exception as e:
                if attempt < self.max_retries - 1:
                    logger.warning(f"MCP connection attempt {attempt + 1} failed: {e}, retrying...")
                    await asyncio.sleep(self.retry_delay * (attempt + 1))
                else:
                    logger.error(f"MCP connection failed after {self.max_retries} attempts: {e}")
                    raise
    
    async def _connect_impl(self):
        """Internal connection implementation."""
        # ... existing connection logic ...
```

#### 1.2 Add Configuration Validation Utilities
Create `ai/mcp_utils.py` for validation utilities:

```python
"""Utilities for MCP configuration validation."""
import json
import logging
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


def validate_mcp_config(config_str: str) -> tuple[bool, str | None]:
    """
    Validate MCP configuration JSON.
    
    :param config_str: JSON configuration string
    :return: Tuple of (is_valid, error_message)
    """
    try:
        config = json.loads(config_str)
        
        if not isinstance(config, dict):
            return False, "Configuration must be a JSON object"
        
        if not config:
            return True, None  # Empty config is valid
        
        for server_name, server_config in config.items():
            valid, error = validate_server_config(server_name, server_config)
            if not valid:
                return False, error
        
        return True, None
        
    except json.JSONDecodeError as e:
        return False, f"Invalid JSON: {e}"


def validate_server_config(name: str, config: dict[str, Any]) -> tuple[bool, str | None]:
    """
    Validate a single server configuration.
    
    :param name: Server name
    :param config: Server configuration
    :return: Tuple of (is_valid, error_message)
    """
    if not isinstance(config, dict):
        return False, f"Server '{name}' configuration must be an object"
    
    transport = config.get("transport", "stdio")
    
    if transport not in ["stdio", "streamable_http"]:
        return False, f"Server '{name}' has invalid transport: {transport}"
    
    if transport == "stdio":
        if "command" not in config:
            return False, f"Server '{name}' missing 'command' for stdio transport"
        
        # Validate command exists
        command = config["command"]
        if not _command_exists(command):
            return False, f"Server '{name}' command not found: {command}"
        
        # Validate args if present
        if "args" in config:
            args = config["args"]
            if not isinstance(args, list):
                return False, f"Server '{name}' args must be a list"
            
            # Check if script files exist
            for arg in args:
                if isinstance(arg, str) and (arg.endswith(".py") or "/" in arg):
                    if not Path(arg).exists():
                        logger.warning(f"Server '{name}' script may not exist: {arg}")
    
    elif transport == "streamable_http":
        if "url" not in config:
            return False, f"Server '{name}' missing 'url' for streamable_http transport"
        
        url = config["url"]
        if not url.startswith(("http://", "https://")):
            return False, f"Server '{name}' URL must start with http:// or https://"
    
    return True, None


def _command_exists(command: str) -> bool:
    """Check if a command exists in PATH."""
    import shutil
    return shutil.which(command) is not None
```

#### 1.3 Add Health Check Mechanism
Add health check methods to `MCPManager`:

```python
class MCPManager:
    # ... existing code ...
    
    async def health_check(self) -> dict[str, bool]:
        """
        Check health of all connected MCP servers.
        
        :return: Dictionary mapping server names to health status
        """
        if not self._client:
            return {}
        
        health = {}
        for server_name in self._client._servers.keys():
            try:
                # Try to get tools from this specific server
                await self._client.session(server_name).__aenter__()
                health[server_name] = True
            except Exception as e:
                logger.warning(f"Health check failed for server '{server_name}': {e}")
                health[server_name] = False
        
        return health
    
    def get_server_info(self) -> dict[str, dict]:
        """Get information about configured servers."""
        if not self._client:
            return {}
        
        return {
            name: {
                "transport": config.get("transport", "stdio"),
                "connected": self._connected,
            }
            for name, config in self._client._servers.items()
        }
```

### 2. Comprehensive Testing

#### 2.1 Create Mock MCP Server for Testing
Create `tests/fixtures/mock_mcp_server.py`:

```python
"""Mock MCP server for testing."""
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("TestServer")


@mcp.tool()
def test_add(a: int, b: int) -> int:
    """Add two numbers (test tool)"""
    return a + b


@mcp.tool()
def test_echo(message: str) -> str:
    """Echo a message (test tool)"""
    return f"Echo: {message}"


@mcp.tool()
async def test_async_tool(value: str) -> str:
    """Test async tool"""
    import asyncio
    await asyncio.sleep(0.1)
    return f"Async result: {value}"


if __name__ == "__main__":
    mcp.run(transport="stdio")
```

#### 2.2 Integration Tests
Create `tests/ai_tests/test_mcp_integration.py`:

```python
"""Integration tests for MCP support."""
import asyncio
import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

from ai.llmagent import LLMAgent
from ai.llmbot import LLMBot
from config import Config


class TestMCPIntegration(unittest.IsolatedAsyncioTestCase):
    """Integration tests for MCP functionality."""
    
    @classmethod
    def setUpClass(cls):
        """Set up test fixtures."""
        cls.mock_server_path = Path(__file__).parent.parent / "fixtures" / "mock_mcp_server.py"
    
    def setUp(self):
        """Set up each test."""
        self.config = Config()
        self.config.enable_mcp = True
        # Use a test LLM configuration
        self.config.ollama_model = "llama3.2:latest"  # or mock
    
    async def test_mcp_integration_with_llmbot(self):
        """Test MCP integration with LLMBot."""
        # Configure mock MCP server
        server_config = {
            "test_server": {
                "command": sys.executable,
                "args": [str(self.mock_server_path)],
                "transport": "stdio"
            }
        }
        self.config.mcp_servers_config = json.dumps(server_config)
        self.config.use_tools = True
        
        # Create bot and initialize
        bot = LLMBot(self.config, [])
        await bot.initialize_async_resources()
        
        try:
            # Verify MCP is connected
            self.assertIsNotNone(bot._mcp_manager)
            self.assertTrue(bot._mcp_manager.is_connected)
            
            # Verify tools are loaded
            tools = await bot._mcp_manager.get_tools()
            self.assertGreater(len(tools), 0)
            
            tool_names = [t.name for t in tools]
            self.assertIn("test_add", tool_names)
            self.assertIn("test_echo", tool_names)
            
        finally:
            await bot.close()
    
    async def test_mcp_integration_with_llmagent(self):
        """Test MCP integration with LLMAgent."""
        # Similar to above but with LLMAgent
        server_config = {
            "test_server": {
                "command": sys.executable,
                "args": [str(self.mock_server_path)],
                "transport": "stdio"
            }
        }
        self.config.mcp_servers_config = json.dumps(server_config)
        self.config.agent_mode = True
        
        agent = LLMAgent(self.config, [])
        await agent.initialize_async_resources()
        
        try:
            # Verify agent is created
            self.assertIsNotNone(agent.agent)
            
            # Verify MCP tools are available
            tools = await agent._mcp_manager.get_tools()
            self.assertGreater(len(tools), 0)
            
        finally:
            await agent.close()
    
    async def test_mcp_graceful_degradation(self):
        """Test that bot works without MCP when it fails."""
        # Configure with invalid server
        server_config = {
            "invalid_server": {
                "command": "nonexistent_command",
                "args": [],
                "transport": "stdio"
            }
        }
        self.config.mcp_servers_config = json.dumps(server_config)
        
        bot = LLMBot(self.config, [])
        
        # Should not raise exception
        await bot.initialize_async_resources()
        
        # Bot should work without MCP
        self.assertIsNone(bot._mcp_manager)
        
        await bot.close()
    
    async def test_mcp_disabled(self):
        """Test bot works correctly when MCP is disabled."""
        self.config.enable_mcp = False
        
        bot = LLMBot(self.config, [])
        await bot.initialize_async_resources()
        
        try:
            self.assertIsNone(bot._mcp_manager)
        finally:
            await bot.close()
```

#### 2.3 Tool Conflict Detection Tests
Add to `tests/ai_tests/test_tools.py`:

```python
async def test_tool_name_conflict_detection(self):
    """Test detection of tool name conflicts between custom and MCP tools."""
    from ai.tools import get_all_tools
    from ai.mcp_manager import MCPManager
    from unittest.mock import AsyncMock, MagicMock
    
    # Create mock MCP manager with conflicting tool name
    mock_manager = AsyncMock(spec=MCPManager)
    mock_manager.is_connected = True
    
    conflicting_tool = MagicMock()
    conflicting_tool.name = "multiply"  # Conflicts with custom tool
    
    mock_manager.get_tools = AsyncMock(return_value=[conflicting_tool])
    
    # Should log warning about conflict
    with self.assertLogs(level='WARNING') as log:
        tools = await get_all_tools(mock_manager)
        
        # Check that warning was logged
        self.assertTrue(any("conflict" in message.lower() for message in log.output))
```

### 3. Monitoring & Observability

#### 3.1 Add MCP Metrics Collection
Create `ai/mcp_metrics.py`:

```python
"""MCP metrics and monitoring."""
import logging
import time
from collections import defaultdict
from typing import Dict

logger = logging.getLogger(__name__)


class MCPMetrics:
    """Track MCP usage metrics."""
    
    def __init__(self):
        self.tool_invocations: Dict[str, int] = defaultdict(int)
        self.tool_errors: Dict[str, int] = defaultdict(int)
        self.tool_latencies: Dict[str, list[float]] = defaultdict(list)
        self.server_health_checks: Dict[str, int] = defaultdict(int)
    
    def record_tool_invocation(self, tool_name: str, duration: float, success: bool = True):
        """Record a tool invocation."""
        self.tool_invocations[tool_name] += 1
        self.tool_latencies[tool_name].append(duration)
        
        if not success:
            self.tool_errors[tool_name] += 1
        
        logger.debug(f"MCP tool '{tool_name}' invoked in {duration:.3f}s (success={success})")
    
    def record_health_check(self, server_name: str, healthy: bool):
        """Record a health check result."""
        if healthy:
            self.server_health_checks[server_name] += 1
        
        logger.debug(f"MCP server '{server_name}' health check: {healthy}")
    
    def get_summary(self) -> dict:
        """Get metrics summary."""
        summary = {
            "total_invocations": sum(self.tool_invocations.values()),
            "total_errors": sum(self.tool_errors.values()),
            "tools": {},
        }
        
        for tool_name in self.tool_invocations.keys():
            latencies = self.tool_latencies[tool_name]
            summary["tools"][tool_name] = {
                "invocations": self.tool_invocations[tool_name],
                "errors": self.tool_errors[tool_name],
                "avg_latency": sum(latencies) / len(latencies) if latencies else 0,
                "min_latency": min(latencies) if latencies else 0,
                "max_latency": max(latencies) if latencies else 0,
            }
        
        return summary
    
    def log_summary(self):
        """Log metrics summary."""
        summary = self.get_summary()
        logger.info(f"MCP Metrics Summary: {summary}")
```

Update `MCPManager` to use metrics:

```python
from ai.mcp_metrics import MCPMetrics

class MCPManager:
    def __init__(self, config: Config):
        # ... existing code ...
        self.metrics = MCPMetrics()
```

#### 3.2 Add Periodic Health Checks
Update `main.py` to add periodic health checks:

```python
async def periodic_mcp_health_check(llm_bot, interval: int = 300):
    """
    Periodically check MCP server health.
    
    :param llm_bot: LLMBot or LLMAgent instance
    :param interval: Check interval in seconds (default: 5 minutes)
    """
    while True:
        try:
            await asyncio.sleep(interval)
            
            if llm_bot._mcp_manager and llm_bot._mcp_manager.is_connected:
                health = await llm_bot._mcp_manager.health_check()
                
                unhealthy = [name for name, status in health.items() if not status]
                if unhealthy:
                    logging.warning(f"Unhealthy MCP servers: {unhealthy}")
                else:
                    logging.debug("All MCP servers healthy")
                
        except asyncio.CancelledError:
            break
        except Exception as e:
            logging.error(f"Error in MCP health check: {e}", exc_info=True)


async def main():
    """Main async function to run the bot."""
    processor_task = asyncio.create_task(process_message_queue())
    
    # Start MCP health check if enabled
    health_check_task = None
    if config.enable_mcp:
        health_check_task = asyncio.create_task(periodic_mcp_health_check(llm_bot))
    
    # ... existing code ...
    
    try:
        await dp.start_polling(bot)
    finally:
        processor_task.cancel()
        if health_check_task:
            health_check_task.cancel()
        
        # Log MCP metrics before shutdown
        if llm_bot._mcp_manager:
            llm_bot._mcp_manager.metrics.log_summary()
        
        await llm_bot.close()
        await bot.session.close()
```

### 4. Configuration Enhancements

#### 4.1 Add YAML Configuration Support
Update `config.py` to support YAML file:

```python
from pathlib import Path

class Config(EnvModel):
    # ... existing fields ...
    
    # MCP Configuration
    enable_mcp = BooleanField("ENABLE_MCP", default=False)
    mcp_servers_config = StringField(
        "MCP_SERVERS_CONFIG",
        default="{}",
        warning="MCP_SERVERS_CONFIG not set. MCP will be disabled."
    )
    mcp_config_file = StringField("MCP_CONFIG_FILE")  # Path to YAML file
    
    def get_mcp_config(self) -> dict:
        """Get MCP configuration from file or env var."""
        # Prefer YAML file if provided
        if self.mcp_config_file and Path(self.mcp_config_file).exists():
            import yaml
            with open(self.mcp_config_file) as f:
                config = yaml.safe_load(f)
                return config.get("mcp", {}).get("servers", {})
        
        # Fall back to env var
        import json
        return json.loads(self.mcp_servers_config)
```

Add `pyyaml` dependency to `pyproject.toml`:

```toml
dependencies = [
    # ... existing ...
    "pyyaml>=6.0",
]
```

Create example `mcp_config.example.yaml`:

```yaml
# Example MCP Configuration File
mcp:
  enabled: true
  servers:
    # stdio transport example
    - name: math
      command: python
      args:
        - /path/to/math_server.py
      transport: stdio
      env:
        # Optional environment variables for the server
        DEBUG: "false"
    
    # HTTP transport example
    - name: weather
      url: http://localhost:8000/mcp/
      transport: streamable_http
      headers:
        # Optional HTTP headers
        Authorization: Bearer ${API_TOKEN}
    
    # Another stdio server
    - name: filesystem
      command: npx
      args:
        - "-y"
        - "@modelcontextprotocol/server-filesystem"
        - "/tmp"
      transport: stdio
```

### 5. Logging Improvements

#### 5.1 Add Structured MCP Logging
Create `ai/mcp_logging.py`:

```python
"""Structured logging for MCP operations."""
import logging
from typing import Any

logger = logging.getLogger(__name__)


def log_mcp_event(event_type: str, **kwargs):
    """Log a structured MCP event."""
    event_data = {
        "event_type": event_type,
        **kwargs
    }
    logger.info(f"MCP_EVENT: {event_data}")


def log_server_connection(server_name: str, success: bool, error: str | None = None):
    """Log MCP server connection attempt."""
    log_mcp_event(
        "server_connection",
        server=server_name,
        success=success,
        error=error
    )


def log_tool_load(server_name: str, tool_count: int):
    """Log MCP tool loading."""
    log_mcp_event(
        "tools_loaded",
        server=server_name,
        tool_count=tool_count
    )


def log_tool_invocation(tool_name: str, duration: float, success: bool):
    """Log MCP tool invocation."""
    log_mcp_event(
        "tool_invocation",
        tool=tool_name,
        duration_ms=duration * 1000,
        success=success
    )
```

## ✅ Acceptance Criteria

### Error Handling & Validation
- [ ] Retry logic implemented for MCP connection failures
- [ ] Configuration validation utilities created and working
- [ ] Health check mechanism implemented
- [ ] Graceful degradation tested and working

### Testing
- [ ] Mock MCP server created for testing
- [ ] Integration tests written and passing
- [ ] Tool conflict detection tested
- [ ] All tests pass with MCP enabled and disabled
- [ ] Test coverage >80% for MCP code

### Monitoring & Observability
- [ ] Metrics collection implemented
- [ ] Periodic health checks running
- [ ] Metrics logged on shutdown
- [ ] Structured logging in place

### Configuration
- [ ] YAML configuration support added
- [ ] Example YAML file created
- [ ] Documentation updated with YAML examples
- [ ] Configuration validation working

### Code Quality
- [ ] All code passes `uv run ruff check`
- [ ] All code passes `uv run ruff format`
- [ ] Comprehensive error messages for common issues
- [ ] Logging is clear and actionable

## 🎯 Success Metrics

- MCP integration is production-ready
- Clear visibility into MCP health and performance
- Easy troubleshooting of MCP issues
- Flexible configuration options
- Comprehensive test coverage

## 📚 References

- Phase 1 Issue: #50
- [Model Context Protocol Documentation](https://modelcontextprotocol.io/)
- [Python logging best practices](https://docs.python.org/3/howto/logging.html)

## 🔗 Related Issues

- Phase 1: #50 - Core Integration (prerequisite)
- Phase 3: #[future] - User Experience Enhancements (CLI tools, advanced features)

Navigation Menu

MCP Support - Phase 2: Robustness & Testing Improvements #51

🎯 Objective

📋 Prerequisites

🔧 Implementation Tasks

1. Enhanced Error Handling & Validation

1.1 Add Connection Retry Logic

1.2 Add Configuration Validation Utilities

1.3 Add Health Check Mechanism

2. Comprehensive Testing

2.1 Create Mock MCP Server for Testing

2.2 Integration Tests

2.3 Tool Conflict Detection Tests

3. Monitoring & Observability

3.1 Add MCP Metrics Collection

3.2 Add Periodic Health Checks

4. Configuration Enhancements

4.1 Add YAML Configuration Support

5. Logging Improvements

5.1 Add Structured MCP Logging

✅ Acceptance Criteria

Error Handling & Validation

Testing

Monitoring & Observability

Configuration

Code Quality

🎯 Success Metrics

📚 References

🔗 Related Issues

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

MCP Support - Phase 2: Robustness & Testing Improvements #51

Description

🎯 Objective

📋 Prerequisites

🔧 Implementation Tasks

1. Enhanced Error Handling & Validation

1.1 Add Connection Retry Logic

1.2 Add Configuration Validation Utilities

1.3 Add Health Check Mechanism

2. Comprehensive Testing

2.1 Create Mock MCP Server for Testing

2.2 Integration Tests

2.3 Tool Conflict Detection Tests

3. Monitoring & Observability

3.1 Add MCP Metrics Collection

3.2 Add Periodic Health Checks

4. Configuration Enhancements

4.1 Add YAML Configuration Support

5. Logging Improvements

5.1 Add Structured MCP Logging

✅ Acceptance Criteria

Error Handling & Validation

Testing

Monitoring & Observability

Configuration

Code Quality

🎯 Success Metrics

📚 References

🔗 Related Issues

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions