unraid-mcp/unraid_mcp/core/client.py

"""GraphQL client for Unraid API communication.

This module provides the HTTP client interface for making GraphQL requests
to the Unraid API with proper timeout handling and error management.
"""

import asyncio
import hashlib
import json
import time
from typing import Any

import httpx

from ..config.logging import logger
from ..config.settings import (
    TIMEOUT_CONFIG,
    UNRAID_API_KEY,
    UNRAID_API_URL,
    UNRAID_VERIFY_SSL,
    VERSION,
)
from ..core.exceptions import ToolError


# Sensitive keys to redact from debug logs
_SENSITIVE_KEYS = {
    "password",
    "key",
    "secret",
    "token",
    "apikey",
    "authorization",
    "cookie",
    "session",
    "credential",
    "passphrase",
    "jwt",
}


def _is_sensitive_key(key: str) -> bool:
    """Check if a key name contains any sensitive substring."""
    key_lower = key.lower()
    return any(s in key_lower for s in _SENSITIVE_KEYS)


def _redact_sensitive(obj: Any) -> Any:
    """Recursively redact sensitive values from nested dicts/lists."""
    if isinstance(obj, dict):
        return {
            k: ("***" if _is_sensitive_key(k) else _redact_sensitive(v)) for k, v in obj.items()
        }
    if isinstance(obj, list):
        return [_redact_sensitive(item) for item in obj]
    return obj


# HTTP timeout configuration
DEFAULT_TIMEOUT = httpx.Timeout(10.0, read=30.0, connect=5.0)
DISK_TIMEOUT = httpx.Timeout(10.0, read=TIMEOUT_CONFIG["disk_operations"], connect=5.0)

# Named timeout profiles
_TIMEOUT_PROFILES: dict[str, httpx.Timeout] = {
    "default": DEFAULT_TIMEOUT,
    "disk_operations": DISK_TIMEOUT,
}


def get_timeout_for_operation(profile: str) -> httpx.Timeout:
    """Get a timeout configuration by profile name.

    Args:
        profile: Timeout profile name (e.g., "default", "disk_operations")

    Returns:
        The matching httpx.Timeout, falling back to DEFAULT_TIMEOUT for unknown profiles
    """
    return _TIMEOUT_PROFILES.get(profile, DEFAULT_TIMEOUT)


# Global connection pool (module-level singleton)
_http_client: httpx.AsyncClient | None = None
_client_lock: asyncio.Lock | None = None


def _get_client_lock() -> asyncio.Lock:
    """Get or create the client lock (lazy init to avoid event loop issues)."""
    global _client_lock
    if _client_lock is None:
        _client_lock = asyncio.Lock()
    return _client_lock


class _RateLimiter:
    """Token bucket rate limiter for Unraid API (100 req / 10s hard limit).

    Uses 90 tokens with 9.0 tokens/sec refill for 10% safety headroom.
    """

    def __init__(self, max_tokens: int = 90, refill_rate: float = 9.0) -> None:
        self.max_tokens = max_tokens
        self.tokens = float(max_tokens)
        self.refill_rate = refill_rate  # tokens per second
        self.last_refill = time.monotonic()
        self._lock: asyncio.Lock | None = None

    def _get_lock(self) -> asyncio.Lock:
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    def _refill(self) -> None:
        """Refill tokens based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(self.max_tokens, self.tokens + elapsed * self.refill_rate)
        self.last_refill = now

    async def acquire(self) -> None:
        """Consume one token, waiting if necessary for refill."""
        while True:
            async with self._get_lock():
                self._refill()
                if self.tokens >= 1:
                    self.tokens -= 1
                    return
                wait_time = (1 - self.tokens) / self.refill_rate

            # Sleep outside the lock so other coroutines aren't blocked
            await asyncio.sleep(wait_time)


_rate_limiter = _RateLimiter()


# --- TTL Cache for stable read-only queries ---

# Queries whose results change infrequently and are safe to cache.
# Mutations and volatile queries (metrics, docker, array state) are excluded.
_CACHEABLE_QUERY_PREFIXES = frozenset(
    {
        "GetNetworkConfig",
        "GetRegistrationInfo",
        "GetOwner",
        "GetFlash",
    }
)

_CACHE_TTL_SECONDS = 60.0


class _QueryCache:
    """Simple TTL cache for GraphQL query responses.

    Keyed by a hash of (query, variables). Entries expire after _CACHE_TTL_SECONDS.
    Only caches responses for queries whose operation name is in _CACHEABLE_QUERY_PREFIXES.
    Mutation requests always bypass the cache.
    """

    def __init__(self) -> None:
        self._store: dict[str, tuple[float, dict[str, Any]]] = {}

    @staticmethod
    def _cache_key(query: str, variables: dict[str, Any] | None) -> str:
        raw = query + json.dumps(variables or {}, sort_keys=True)
        return hashlib.sha256(raw.encode()).hexdigest()

    @staticmethod
    def is_cacheable(query: str) -> bool:
        """Check if a query is eligible for caching based on its operation name."""
        if query.lstrip().startswith("mutation"):
            return False
        return any(prefix in query for prefix in _CACHEABLE_QUERY_PREFIXES)

    def get(self, query: str, variables: dict[str, Any] | None) -> dict[str, Any] | None:
        """Return cached result if present and not expired, else None."""
        key = self._cache_key(query, variables)
        entry = self._store.get(key)
        if entry is None:
            return None
        expires_at, data = entry
        if time.monotonic() > expires_at:
            del self._store[key]
            return None
        return data

    def put(self, query: str, variables: dict[str, Any] | None, data: dict[str, Any]) -> None:
        """Store a query result with TTL expiry."""
        key = self._cache_key(query, variables)
        self._store[key] = (time.monotonic() + _CACHE_TTL_SECONDS, data)

    def invalidate_all(self) -> None:
        """Clear the entire cache (called after mutations)."""
        self._store.clear()


_query_cache = _QueryCache()


def is_idempotent_error(error_message: str, operation: str) -> bool:
    """Check if a GraphQL error represents an idempotent operation that should be treated as success.

    Args:
        error_message: The error message from GraphQL API
        operation: The operation being performed (e.g., 'start', 'stop')

    Returns:
        True if this is an idempotent error that should be treated as success
    """
    error_lower = error_message.lower()

    # Docker container operation patterns
    if operation == "start":
        return (
            "already started" in error_lower
            or "container already running" in error_lower
            or "http code 304" in error_lower
        )
    if operation == "stop":
        return (
            "already stopped" in error_lower
            or "container already stopped" in error_lower
            or "container not running" in error_lower
            or "http code 304" in error_lower
        )

    return False


async def _create_http_client() -> httpx.AsyncClient:
    """Create a new HTTP client instance with connection pooling.

    Returns:
        A new AsyncClient configured for Unraid API communication
    """
    return httpx.AsyncClient(
        # Connection pool settings
        limits=httpx.Limits(
            max_keepalive_connections=20, max_connections=20, keepalive_expiry=30.0
        ),
        # Default timeout (can be overridden per-request)
        timeout=DEFAULT_TIMEOUT,
        # SSL verification
        verify=UNRAID_VERIFY_SSL,
        # Connection pooling headers
        headers={"Connection": "keep-alive", "User-Agent": f"UnraidMCPServer/{VERSION}"},
    )


async def get_http_client() -> httpx.AsyncClient:
    """Get or create shared HTTP client with connection pooling.

    Uses double-checked locking: fast-path skips the lock when the client
    is already initialized, only acquiring it for initial creation or
    recovery after close.

    Returns:
        Singleton AsyncClient instance with connection pooling enabled
    """
    global _http_client

    # Fast-path: skip lock if client is already initialized and open
    client = _http_client
    if client is not None and not client.is_closed:
        return client

    # Slow-path: acquire lock for initialization
    async with _get_client_lock():
        if _http_client is None or _http_client.is_closed:
            _http_client = await _create_http_client()
            logger.info(
                "Created shared HTTP client with connection pooling (20 keepalive, 20 max connections)"
            )
        return _http_client


async def close_http_client() -> None:
    """Close the shared HTTP client (call on server shutdown)."""
    global _http_client

    async with _get_client_lock():
        if _http_client is not None:
            await _http_client.aclose()
            _http_client = None
            logger.info("Closed shared HTTP client")


async def make_graphql_request(
    query: str,
    variables: dict[str, Any] | None = None,
    custom_timeout: httpx.Timeout | None = None,
    operation_context: dict[str, str] | None = None,
) -> dict[str, Any]:
    """Make GraphQL requests to the Unraid API.

    Args:
        query: GraphQL query string
        variables: Optional query variables
        custom_timeout: Optional custom timeout configuration
        operation_context: Optional context for operation-specific error handling
                          Should contain 'operation' key (e.g., 'start', 'stop')

    Returns:
        Dict containing the GraphQL response data

    Raises:
        ToolError: For HTTP errors, network errors, or non-idempotent GraphQL errors
    """
    if not UNRAID_API_URL:
        raise ToolError("UNRAID_API_URL not configured")

    if not UNRAID_API_KEY:
        raise ToolError("UNRAID_API_KEY not configured")

    # Check TTL cache for stable read-only queries
    is_mutation = query.lstrip().startswith("mutation")
    if not is_mutation and _query_cache.is_cacheable(query):
        cached = _query_cache.get(query, variables)
        if cached is not None:
            logger.debug("Returning cached response for query")
            return cached

    headers = {
        "Content-Type": "application/json",
        "X-API-Key": UNRAID_API_KEY,
    }

    payload: dict[str, Any] = {"query": query}
    if variables:
        payload["variables"] = variables

    logger.debug(f"Making GraphQL request to {UNRAID_API_URL}:")
    logger.debug(f"Query: {query[:200]}{'...' if len(query) > 200 else ''}")  # Log truncated query
    if variables:
        logger.debug(f"Variables: {_redact_sensitive(variables)}")

    try:
        # Rate limit: consume a token before making the request
        await _rate_limiter.acquire()

        # Get the shared HTTP client with connection pooling
        client = await get_http_client()

        # Retry loop for 429 rate limit responses
        post_kwargs: dict[str, Any] = {"json": payload, "headers": headers}
        if custom_timeout is not None:
            post_kwargs["timeout"] = custom_timeout

        response: httpx.Response | None = None
        for attempt in range(3):
            response = await client.post(UNRAID_API_URL, **post_kwargs)
            if response.status_code == 429:
                backoff = 2**attempt
                logger.warning(
                    f"Rate limited (429) by Unraid API, retrying in {backoff}s (attempt {attempt + 1}/3)"
                )
                await asyncio.sleep(backoff)
                continue
            break

        if response is None:  # pragma: no cover — guaranteed by loop
            raise ToolError("No response received after retry attempts")
        response.raise_for_status()  # Raise an exception for HTTP error codes 4xx/5xx

        response_data = response.json()
        if response_data.get("errors"):
            error_details = "; ".join(
                [err.get("message", str(err)) for err in response_data["errors"]]
            )

            # Check if this is an idempotent error that should be treated as success
            if operation_context and operation_context.get("operation"):
                operation = operation_context["operation"]
                if is_idempotent_error(error_details, operation):
                    logger.warning(
                        f"Idempotent operation '{operation}' - treating as success: {error_details}"
                    )
                    # Return a success response with the current state information
                    return {
                        "idempotent_success": True,
                        "operation": operation,
                        "message": error_details,
                        "original_errors": response_data["errors"],
                    }

            logger.error(f"GraphQL API returned errors: {response_data['errors']}")
            # Use ToolError for GraphQL errors to provide better feedback to LLM
            raise ToolError(f"GraphQL API error: {error_details}")

        logger.debug("GraphQL request successful.")
        data = response_data.get("data", {})
        result = data if isinstance(data, dict) else {}  # Ensure we return dict

        # Invalidate cache on mutations; cache eligible query results
        if is_mutation:
            _query_cache.invalidate_all()
        elif _query_cache.is_cacheable(query):
            _query_cache.put(query, variables, result)

        return result

    except httpx.HTTPStatusError as e:
        # Log full details internally; only expose status code to MCP client
        logger.error(f"HTTP error occurred: {e.response.status_code} - {e.response.text}")
        raise ToolError(
            f"Unraid API returned HTTP {e.response.status_code}. Check server logs for details."
        ) from e
    except httpx.RequestError as e:
        # Log full error internally; give safe summary to MCP client
        logger.error(f"Request error occurred: {e}")
        raise ToolError(f"Network error connecting to Unraid API: {type(e).__name__}") from e
    except json.JSONDecodeError as e:
        # Log full decode error; give safe summary to MCP client
        logger.error(f"Failed to decode JSON response: {e}")
        raise ToolError("Unraid API returned an invalid response (not valid JSON)") from e