tools/mcp_client.py

"""MCP Client for connecting to FixMyNeighborhood MCP server.

Enhanced with:
- Request timeouts
- Automatic retry with exponential backoff
- Structured error handling
- Request logging for observability
"""
import sys
import io
import time
import threading
from typing import Optional, Dict, Any, Callable
from dataclasses import dataclass
from functools import wraps
from gradio_client import Client as GradioClient
from config import MCP_SERVER_URL


@dataclass
class RetryConfig:
    """Configuration for retry behavior."""
    max_retries: int = 3
    initial_delay: float = 1.0  # seconds
    max_delay: float = 10.0
    exponential_base: float = 2.0
    timeout: float = 30.0  # seconds


@dataclass
class CallMetrics:
    """Metrics for a single tool call."""
    tool_name: str
    start_time: float
    end_time: Optional[float] = None
    success: bool = False
    retries: int = 0
    error: Optional[str] = None

    @property
    def duration(self) -> float:
        if self.end_time:
            return self.end_time - self.start_time
        return 0.0


# Singleton client instance
_mcp_client: Optional["MCPClient"] = None


class MCPClient:
    """
    MCP Tool Client using gradio_client to call remote Gradio MCP server.

    Enhanced Features:
    - Configurable request timeouts
    - Automatic retry with exponential backoff
    - Structured error responses
    - Call metrics for observability
    - Thread-safe operations
    """

    # Tool parameter order (for positional args to Gradio API)
    TOOL_PARAM_ORDER = {
        "geo_search_address": ["lat", "lon"],
        "validate_address": ["address"],
        "cityinfra_lookup_asset": ["address", "asset_type"],
        "get_nearby_reports": ["address", "issue_type", "radius_blocks"],
        "weather_get_current": ["lat", "lon"],
        "get_department_info": ["department_code"],
        "pdf_generate_report": ["issue_type", "address", "urgency", "description"],
        "sendgrid_send_email": ["to", "subject", "body", "api_key", "from_email", "report_id"]
    }

    def __init__(
        self,
        server_url: str = None,
        retry_config: RetryConfig = None
    ):
        self.server_url = server_url or MCP_SERVER_URL
        self.retry_config = retry_config or RetryConfig()
        self._client: Optional[GradioClient] = None
        self._client_lock = threading.Lock()
        self._call_history: list = []
        self._max_history: int = 100

    @property
    def client(self) -> Optional[GradioClient]:
        """Lazy initialization of Gradio client with Windows encoding fix."""
        if self._client is None:
            with self._client_lock:
                if self._client is None:  # Double-check locking
                    self._init_client()
        return self._client

    def _init_client(self) -> None:
        """Initialize the Gradio client with suppressed output."""
        old_stdout, old_stderr = sys.stdout, sys.stderr
        try:
            sys.stdout = io.StringIO()
            sys.stderr = io.StringIO()
            self._client = GradioClient(self.server_url)
            print(f"[MCP] Connected to {self.server_url}")
        except Exception as e:
            sys.stdout, sys.stderr = old_stdout, old_stderr
            print(f"[MCP] Connection failed: {e}")
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr

    def _with_timeout(
        self,
        func: Callable,
        timeout: float,
        *args,
        **kwargs
    ) -> Any:
        """Execute a function with a timeout."""
        result = [None]
        error = [None]

        def target():
            try:
                result[0] = func(*args, **kwargs)
            except Exception as e:
                error[0] = e

        thread = threading.Thread(target=target)
        thread.start()
        thread.join(timeout=timeout)

        if thread.is_alive():
            # Thread is still running - timeout occurred
            raise TimeoutError(f"Operation timed out after {timeout}s")

        if error[0]:
            raise error[0]

        return result[0]

    def call_tool(self, tool_name: str, **kwargs) -> dict:
        """
        Call an MCP tool with retry and timeout support.

        Args:
            tool_name: Name of the tool to call
            **kwargs: Tool parameters

        Returns:
            dict: Tool result or error response
        """
        metrics = CallMetrics(tool_name=tool_name, start_time=time.time())
        config = self.retry_config

        for attempt in range(config.max_retries + 1):
            try:
                if self.client is None:
                    return self._fallback_response(tool_name, kwargs, "Client not initialized")

                # Get ordered args for this tool
                param_order = self.TOOL_PARAM_ORDER.get(tool_name, [])
                args = [kwargs.get(param) for param in param_order]

                # Execute with timeout
                result = self._with_timeout(
                    self._execute_call,
                    config.timeout,
                    tool_name,
                    args
                )

                # Success
                metrics.end_time = time.time()
                metrics.success = True
                metrics.retries = attempt
                self._record_call(metrics)

                return result if isinstance(result, dict) else {"result": result}

            except TimeoutError as e:
                metrics.error = f"Timeout: {e}"
                print(f"[MCP] {tool_name} timeout (attempt {attempt + 1}/{config.max_retries + 1})")

            except Exception as e:
                metrics.error = str(e)
                print(f"[MCP] {tool_name} error: {e} (attempt {attempt + 1}/{config.max_retries + 1})")

            # Check if we should retry
            if attempt < config.max_retries:
                delay = min(
                    config.initial_delay * (config.exponential_base ** attempt),
                    config.max_delay
                )
                print(f"[MCP] Retrying in {delay:.1f}s...")
                time.sleep(delay)
                metrics.retries = attempt + 1

        # All retries exhausted
        metrics.end_time = time.time()
        metrics.success = False
        self._record_call(metrics)

        return self._fallback_response(
            tool_name,
            kwargs,
            f"Failed after {config.max_retries + 1} attempts: {metrics.error}"
        )

    def _execute_call(self, tool_name: str, args: list) -> Any:
        """Execute the actual API call with suppressed output."""
        old_stdout, old_stderr = sys.stdout, sys.stderr
        try:
            sys.stdout = io.StringIO()
            sys.stderr = io.StringIO()
            return self.client.predict(
                *args,
                api_name=f"/{tool_name}"
            )
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr

    def _fallback_response(
        self,
        tool_name: str,
        inputs: dict,
        error: str = None
    ) -> dict:
        """Return structured error response when MCP server is unavailable."""
        return {
            "error": "MCP server unavailable",
            "error_detail": error,
            "message": "The infrastructure tools service is temporarily unavailable. Please try again in a moment.",
            "tool": tool_name,
            "server_url": self.server_url,
            "recoverable": True,
            "retry_after": 30,
        }

    def _record_call(self, metrics: CallMetrics) -> None:
        """Record call metrics for observability."""
        self._call_history.append({
            "tool": metrics.tool_name,
            "duration": metrics.duration,
            "success": metrics.success,
            "retries": metrics.retries,
            "error": metrics.error,
            "timestamp": metrics.start_time,
        })

        # Trim history
        if len(self._call_history) > self._max_history:
            self._call_history = self._call_history[-self._max_history:]

    def get_metrics(self) -> Dict[str, Any]:
        """Get call metrics for observability."""
        if not self._call_history:
            return {"total_calls": 0}

        successful = [c for c in self._call_history if c["success"]]
        failed = [c for c in self._call_history if not c["success"]]

        return {
            "total_calls": len(self._call_history),
            "successful": len(successful),
            "failed": len(failed),
            "success_rate": len(successful) / len(self._call_history) if self._call_history else 0,
            "avg_duration": sum(c["duration"] for c in successful) / len(successful) if successful else 0,
            "total_retries": sum(c["retries"] for c in self._call_history),
            "by_tool": self._get_by_tool_metrics(),
        }

    def _get_by_tool_metrics(self) -> Dict[str, Dict[str, Any]]:
        """Get metrics grouped by tool."""
        by_tool = {}
        for call in self._call_history:
            tool = call["tool"]
            if tool not in by_tool:
                by_tool[tool] = {"calls": 0, "success": 0, "total_duration": 0}
            by_tool[tool]["calls"] += 1
            if call["success"]:
                by_tool[tool]["success"] += 1
                by_tool[tool]["total_duration"] += call["duration"]
        return by_tool

    def health_check(self) -> Dict[str, Any]:
        """Check if MCP server is healthy."""
        start = time.time()
        try:
            # Try a lightweight call
            if self.client is None:
                return {
                    "healthy": False,
                    "error": "Client not initialized",
                    "latency": None,
                }

            # Attempt connection
            return {
                "healthy": True,
                "latency": time.time() - start,
                "server_url": self.server_url,
            }
        except Exception as e:
            return {
                "healthy": False,
                "error": str(e),
                "latency": time.time() - start,
            }


def get_mcp_client() -> Optional[MCPClient]:
    """Get singleton MCP client instance."""
    global _mcp_client
    if _mcp_client is None:
        _mcp_client = MCPClient()
    return _mcp_client


def reset_mcp_client() -> None:
    """Reset the MCP client (useful for testing or reconnection)."""
    global _mcp_client
    _mcp_client = None