tools.py

"""
Pythonnet Tool Registry - MCP Client
====================================

Connects to headless Pythonnet MCP server via HTTP.
Tool definitions are fetched from the server (single source of truth).
"""

import requests
import json
import uuid
import time
import logging
import os
from typing import Optional, Dict, Any, Type, List


logger = logging.getLogger(__name__)

# Suppress annoying FastMCP session cleanup errors
logging.getLogger("mcp.server.streamable_http").setLevel(logging.CRITICAL)

# Configuration for Docker MCP server
# Use Docker service name when running in container, localhost when on host
PYTHONNET_MCP_URL = os.getenv("PYTHONNET_MCP_URL")
if not PYTHONNET_MCP_URL:
    if os.path.exists("/.dockerenv"):
        # Running inside Docker container - use service name
        PYTHONNET_MCP_URL = "http://pythonnet-mcp:8001/mcp"
    else:
        # Running on host - use localhost
        PYTHONNET_MCP_URL = "http://localhost:8001/mcp"


# Cache for tool definitions (avoid repeated HTTP calls)
_tools_cache: Optional[List[Dict[str, Any]]] = None
_tools_cache_time: float = 0
_TOOLS_CACHE_TTL: float = 300.0  # 5 minutes


def _list_pythonnet_tools(base_url: str, timeout: int = 10) -> List[Dict[str, Any]]:
    """Fetch tool definitions from MCP server via tools/list.

    This makes tools.py a proper MCP client - schemas come from the server,
    not hardcoded duplicates.

    Args:
        base_url: MCP server URL (e.g., http://localhost:8001/mcp)
        timeout: Request timeout in seconds

    Returns:
        List of tool definitions with name, description, inputSchema
    """
    try:
        response = requests.post(
            base_url,
            headers={"Content-Type": "application/json", "Accept": "application/json"},
            json={"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}},
            timeout=timeout,
        )

        result = response.json()

        if result.get("error"):
            logger.warning(f"[PYTHONNET] tools/list error: {result['error']}")
            return []

        tools = result.get("result", {}).get("tools", [])
        logger.info(f"[PYTHONNET] Fetched {len(tools)} tools from MCP server")
        return tools

    except requests.exceptions.Timeout:
        logger.error(f"[PYTHONNET] tools/list timed out ({timeout}s)")
        return []
    except requests.exceptions.ConnectionError:
        logger.error("[PYTHONNET] Cannot connect to MCP server for tools/list")
        return []
    except Exception as e:
        logger.error(f"[PYTHONNET] tools/list failed: {e}")
        return []


def _call_pythonnet_mcp(
    base_url: str,
    tool_name: str,
    arguments: dict,
    timeout: int = 60,
    session_id: Optional[str] = None,
) -> str:
    """Call pythonnet MCP tool via HTTP."""

    # Generate unique session ID if not provided
    if session_id is None:
        session_id = f"sara-{uuid.uuid4().hex[:8]}"

    try:
        response = requests.post(
            f"{base_url}?session_id={session_id}",
            headers={
                "Content-Type": "application/json",
                "Accept": "application/json, text/event-stream",
            },
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "tools/call",
                "params": {"name": tool_name, "arguments": arguments},
            },
            timeout=timeout,
        )

        result = response.json()

        if result.get("error"):
            error_msg = result["error"].get("message", "Unknown error")
            return f"Error: {error_msg}"

        # Extract result from MCP response
        content = result.get("result", {})

        # Try structuredContent first (cleaner), fall back to text content
        if "structuredContent" in content:
            return json.dumps(content["structuredContent"], indent=2)
        elif "content" in content and isinstance(content["content"], list):
            # Extract text from content blocks
            texts = []
            for block in content["content"]:
                if isinstance(block, dict) and block.get("type") == "text":
                    texts.append(block["text"])
            return "\n".join(texts)

        return str(content)

    except requests.exceptions.Timeout:
        return f"Error: Pythonnet MCP request timed out ({timeout}s)"
    except requests.exceptions.ConnectionError:
        return "Error: Cannot connect to Pythonnet MCP server. Is Docker running on localhost:8001?"
    except Exception as e:
        return f"Error calling Pythonnet MCP: {str(e)}"


class PythonnetToolRegistry:
    """MCP client for Pythonnet tool set via Docker HTTP.

    This class fetches tool definitions from the MCP server (server.py)
    rather than hardcoding them. The server is the single source of truth.

    NOTE: This class is intentionally *instance-based*.
    Each analysis run should create its own instance with its own session_id
    to avoid cross-run lock conflicts.

    Backward compatibility:
    - is_available() and get_tool_definitions() remain static
    - execute_tool() is an instance method
    """

    def __init__(
        self, session_id: Optional[str] = None, base_url: Optional[str] = None
    ):
        self.session_id = session_id or f"sara-{uuid.uuid4().hex[:8]}"
        self.base_url = base_url or PYTHONNET_MCP_URL

    @staticmethod
    def get_policy() -> Optional[Type]:
        """Return the policy class for Pythonnet tools."""
        try:
            from services.pythonnet.policy import PythonnetToolPolicy

            return PythonnetToolPolicy
        except ImportError:
            return None

    @staticmethod
    def is_available() -> bool:
        """Check if Pythonnet Docker server is available"""
        try:
            url = f"{PYTHONNET_MCP_URL}?session_id=health-check"
            logger.debug(f"[PYTHONNET] Health check: {url}")
            response = requests.get(url, timeout=5)
            logger.debug(f"[PYTHONNET] Response: {response.status_code}")
            result = response.status_code in [
                200,
                400,
                406,
            ]  # Any response means server is up
            if not result:
                logger.warning(f"[PYTHONNET] Unexpected status: {response.status_code}")
            return result
        except Exception as e:
            logger.error(f"[PYTHONNET] Health check error: {e}")
            return False

    @staticmethod
    def get_tool_definitions() -> list[dict]:
        """Fetch tool definitions from MCP server.

        This is a proper MCP client - schemas come from the server via tools/list,
        not hardcoded duplicates. The server (server.py) is the single source of truth.

        Returns:
            List of tool definitions with name, description, inputSchema
        """
        global _tools_cache, _tools_cache_time

        # Return cached tools if still valid
        if (
            _tools_cache is not None
            and (time.time() - _tools_cache_time) < _TOOLS_CACHE_TTL
        ):
            return _tools_cache

        # Fetch fresh tools from server
        tools = _list_pythonnet_tools(PYTHONNET_MCP_URL)

        # Add client-side metadata that the server doesn't know about
        for tool in tools:
            # pythonnet_load_binary should never be extracted from LLM narration
            if tool.get("name") == "pythonnet_load_binary":
                tool["extractable"] = False

        # Update cache
        _tools_cache = tools
        _tools_cache_time = time.time()

        return tools

    def execute_tool(self, tool_name: str, tool_input: dict) -> str:
        """
        Execute a Pythonnet tool via Docker HTTP MCP server

        Args:
            tool_name: Tool name with pythonnet_ prefix
            tool_input: Tool parameters

        Returns:
            Execution result as string
        """
        try:
            # Use longer timeout for binary loading (5 minutes)
            timeout = 300 if tool_name == "pythonnet_load_binary" else 60

            # Path normalisation for pythonnet_load_binary.
            # Both SARA and the pythonnet-mcp container mount the host directory
            # ./data/samples at /data/samples, so the path is already correct in
            # the normal case.  We only handle two edge cases:
            #   - bare filename with no directory component → prepend /data/samples/
            #   - legacy relative path "data/samples/file" → make absolute
            if tool_name == "pythonnet_load_binary" and "binary_path" in tool_input:
                path = tool_input["binary_path"]
                # Handle legacy relative path: data/samples/file.dll -> /data/samples/file.dll
                if path.startswith("data/samples/"):
                    tool_input["binary_path"] = f"/{path}"
                    logger.info(
                        f"[PYTHONNET] Path normalised: {path} -> {tool_input['binary_path']}"
                    )
                # Bare filenames are resolved server-side (searches /data/samples/ then /data/pythonnet_projects/_uploads/)

            # Call the MCP server with session tracking
            result = _call_pythonnet_mcp(
                self.base_url,
                tool_name,
                tool_input,
                timeout=timeout,
                session_id=self.session_id,
            )
            return result

        except Exception as e:
            import traceback

            return f"Error executing {tool_name}: {str(e)}\n{traceback.format_exc()}"


# Export the registry
__all__ = ["PythonnetToolRegistry", "PYTHONNET_MCP_URL"]