awslabs.openapi-mcp-server 0.1.1__py3-none-any.whl

awslabs/openapi_mcp_server/utils/cache_provider.py

@@ -0,0 +1,249 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Cache provider for the OpenAPI MCP Server.
+
+This module provides a pluggable caching system with different backends.
+The default is a simple in-memory implementation, but it can be switched
+to use external caching systems via environment variables.
+"""
+
+import time
+from abc import ABC, abstractmethod
+from awslabs.openapi_mcp_server import logger
+from awslabs.openapi_mcp_server.utils.config import CACHE_MAXSIZE, CACHE_TTL, USE_CACHETOOLS
+from typing import Any, Callable, Dict, Generic, Optional, TypeVar
+
+
+# Type variable for generic cache implementation
+T = TypeVar('T')
+
+# Try to import cachetools if enabled
+CACHETOOLS_AVAILABLE = False
+cachetools = None
+if USE_CACHETOOLS:
+    try:
+        import cachetools
+
+        CACHETOOLS_AVAILABLE = True
+        logger.info('cachetools caching enabled')
+    except ImportError:
+        logger.warning(
+            'cachetools requested but not installed. Install with: pip install cachetools'
+        )
+
+
+class CacheProvider(Generic[T], ABC):
+    """Abstract base class for cache providers."""
+
+    @abstractmethod
+    def get(self, key: str) -> Optional[T]:
+        """Get a value from the cache."""
+        pass
+
+    @abstractmethod
+    def set(self, key: str, value: T) -> None:
+        """Set a value in the cache."""
+        pass
+
+    @abstractmethod
+    def invalidate(self, key: str) -> bool:
+        """Invalidate a cache entry."""
+        pass
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all entries from the cache."""
+        pass
+
+
+class InMemoryCacheProvider(CacheProvider[T]):
+    """Simple in-memory cache provider with TTL support."""
+
+    def __init__(self, ttl_seconds: Optional[int] = None):
+        """Initialize the cache provider.
+
+        Args:
+            ttl_seconds: Time-to-live in seconds for cache entries (defaults to config value)
+
+        """
+        self._cache: Dict[str, tuple[T, float]] = {}
+        self._ttl_seconds = ttl_seconds if ttl_seconds is not None else CACHE_TTL
+        logger.debug(f'Created in-memory cache provider with TTL of {self._ttl_seconds} seconds')
+
+    def get(self, key: str) -> Optional[T]:
+        """Get a value from the cache."""
+        if key not in self._cache:
+            return None
+
+        value, expiry = self._cache[key]
+        if time.time() > expiry:
+            # Entry has expired
+            del self._cache[key]
+            logger.debug(f'Cache entry expired: {key}')
+            return None
+
+        logger.debug(f'Cache hit: {key}')
+        return value
+
+    def set(self, key: str, value: T) -> None:
+        """Set a value in the cache."""
+        expiry = time.time() + self._ttl_seconds
+        self._cache[key] = (value, expiry)
+        logger.debug(f'Cache set: {key} (expires in {self._ttl_seconds} seconds)')
+
+    def invalidate(self, key: str) -> bool:
+        """Invalidate a cache entry."""
+        if key in self._cache:
+            del self._cache[key]
+            logger.debug(f'Cache invalidated: {key}')
+            return True
+        return False
+
+    def clear(self) -> None:
+        """Clear all entries from the cache."""
+        count = len(self._cache)
+        self._cache.clear()
+        logger.debug(f'Cache cleared ({count} entries removed)')
+
+    def cleanup(self) -> int:
+        """Remove all expired entries from the cache.
+
+        Returns:
+            int: Number of entries removed
+
+        """
+        now = time.time()
+        expired_keys = [key for key, (_, expiry) in self._cache.items() if now > expiry]
+
+        for key in expired_keys:
+            del self._cache[key]
+
+        if expired_keys:
+            logger.debug(f'Cache cleanup: removed {len(expired_keys)} expired entries')
+
+        return len(expired_keys)
+
+
+class CachetoolsProvider(CacheProvider[T]):
+    """Cache provider using the cachetools library."""
+
+    def __init__(self, ttl_seconds: Optional[int] = None, maxsize: Optional[int] = None):
+        """Initialize the cache provider.
+
+        Args:
+            ttl_seconds: Time-to-live in seconds for cache entries (defaults to config value)
+            maxsize: Maximum number of entries in the cache (defaults to config value)
+
+        """
+        if not CACHETOOLS_AVAILABLE or cachetools is None:
+            raise ImportError('cachetools not available')
+
+        # Use configuration values if not explicitly provided
+        ttl_seconds = ttl_seconds if ttl_seconds is not None else CACHE_TTL
+        maxsize = maxsize if maxsize is not None else CACHE_MAXSIZE
+
+        self._cache = cachetools.TTLCache(maxsize=maxsize, ttl=ttl_seconds)
+        logger.debug(
+            f'Created cachetools cache provider with TTL of {ttl_seconds} seconds and maxsize of {maxsize}'
+        )
+
+    def get(self, key: str) -> Optional[T]:
+        """Get a value from the cache."""
+        try:
+            value = self._cache[key]
+            logger.debug(f'Cache hit: {key}')
+            return value
+        except KeyError:
+            logger.debug(f'Cache miss: {key}')
+            return None
+
+    def set(self, key: str, value: T) -> None:
+        """Set a value in the cache."""
+        self._cache[key] = value
+        logger.debug(f'Cache set: {key}')
+
+    def invalidate(self, key: str) -> bool:
+        """Invalidate a cache entry."""
+        try:
+            del self._cache[key]
+            logger.debug(f'Cache invalidated: {key}')
+            return True
+        except KeyError:
+            return False
+
+    def clear(self) -> None:
+        """Clear all entries from the cache."""
+        count = len(self._cache)
+        self._cache.clear()
+        logger.debug(f'Cache cleared ({count} entries removed)')
+
+
+# Create the appropriate cache provider based on configuration
+def create_cache_provider(ttl_seconds: Optional[int] = None) -> CacheProvider:
+    """Create a cache provider based on configuration.
+
+    Args:
+        ttl_seconds: Time-to-live in seconds for cache entries (defaults to config value)
+
+    Returns:
+        CacheProvider: The cache provider
+
+    """
+    # Use configuration value if not explicitly provided
+    ttl_seconds = ttl_seconds if ttl_seconds is not None else CACHE_TTL
+
+    if USE_CACHETOOLS and CACHETOOLS_AVAILABLE:
+        try:
+            return CachetoolsProvider(ttl_seconds=ttl_seconds, maxsize=CACHE_MAXSIZE)
+        except Exception as e:
+            logger.error(f'Failed to create cachetools cache provider: {e}')
+            logger.info('Falling back to in-memory cache provider')
+
+    # Default to in-memory provider
+    return InMemoryCacheProvider(ttl_seconds=ttl_seconds)
+
+
+def cached(ttl_seconds: Optional[int] = None) -> Callable:
+    """Cache function results.
+
+    Args:
+        ttl_seconds: Time-to-live in seconds for cache entries (defaults to config value)
+
+    Returns:
+        Callable: Decorated function with caching
+
+    """
+    cache = create_cache_provider(ttl_seconds=ttl_seconds)
+
+    def decorator(func: Callable) -> Callable:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Create a cache key from the function name and arguments
+            key_parts = [func.__name__]
+            key_parts.extend(str(arg) for arg in args)
+            key_parts.extend(f'{k}={v}' for k, v in sorted(kwargs.items()))
+            cache_key = ':'.join(key_parts)
+
+            # Try to get from cache first
+            cached_result = cache.get(cache_key)
+            if cached_result is not None:
+                return cached_result
+
+            # Call the function and cache the result
+            result = func(*args, **kwargs)
+            cache.set(cache_key, result)
+            return result
+
+        return wrapper
+
+    return decorator

awslabs/openapi_mcp_server/utils/config.py

@@ -0,0 +1,35 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Configuration utilities for the OpenAPI MCP Server."""
+
+import os
+
+
+# Metrics configuration
+METRICS_MAX_HISTORY = int(os.environ.get('METRICS_MAX_HISTORY', '100'))
+USE_PROMETHEUS = os.environ.get('ENABLE_PROMETHEUS', 'false').lower() == 'true'
+PROMETHEUS_PORT = int(os.environ.get('PROMETHEUS_PORT', '9090'))
+
+# Operation prompts configuration
+ENABLE_OPERATION_PROMPTS = os.environ.get('ENABLE_OPERATION_PROMPTS', 'true').lower() == 'true'
+
+# HTTP client configuration
+HTTP_MAX_CONNECTIONS = int(os.environ.get('HTTP_MAX_CONNECTIONS', '100'))
+HTTP_MAX_KEEPALIVE = int(os.environ.get('HTTP_MAX_KEEPALIVE', '20'))
+USE_TENACITY = os.environ.get('USE_TENACITY', 'true').lower() == 'true'
+
+# Cache configuration
+CACHE_MAXSIZE = int(os.environ.get('CACHE_MAXSIZE', '1000'))
+CACHE_TTL = int(os.environ.get('CACHE_TTL', '3600'))  # 1 hour default
+USE_CACHETOOLS = os.environ.get('USE_CACHETOOLS', 'true').lower() == 'true'

awslabs/openapi_mcp_server/utils/error_handler.py

@@ -0,0 +1,349 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities for error handling in the OpenAPI MCP Server."""
+
+import httpx
+import json
+from awslabs.openapi_mcp_server import logger
+from typing import Any, Dict, Optional, Type
+
+
+class APIError(Exception):
+    """Base exception class for API errors."""
+
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        details: Any = None,
+        original_error: Optional[Exception] = None,
+    ):
+        """Initialize the API error.
+
+        Args:
+            status_code: HTTP status code
+            message: Error message
+            details: Additional error details
+            original_error: Original exception that caused this error
+
+        """
+        self.status_code = status_code
+        self.message = message
+        self.details = {} if details is None else details
+        self.original_error = original_error
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        """Return a string representation of the error."""
+        return f'{self.status_code}: {self.message}'
+
+    def __repr__(self) -> str:
+        """Return a representation of the error."""
+        return f'{self.__class__.__name__}({self.status_code}, {repr(self.message)}, {repr(self.details)})'
+
+
+class AuthenticationError(APIError):
+    """Exception raised for authentication errors (401)."""
+
+    pass
+
+
+class AuthorizationError(APIError):
+    """Exception raised for authorization errors (403)."""
+
+    pass
+
+
+class ResourceNotFoundError(APIError):
+    """Exception raised for resource not found errors (404)."""
+
+    pass
+
+
+class ValidationError(APIError):
+    """Exception raised for validation errors (422)."""
+
+    pass
+
+
+class RateLimitError(APIError):
+    """Exception raised for rate limit errors (429)."""
+
+    pass
+
+
+class ServerError(APIError):
+    """Exception raised for server errors (5xx)."""
+
+    pass
+
+
+class ConnectionError(APIError):
+    """Exception raised for connection errors."""
+
+    pass
+
+
+class NetworkError(APIError):
+    """Exception raised for network errors."""
+
+    pass
+
+
+# Map status codes to error classes
+ERROR_CLASSES: Dict[Any, Type[APIError]] = {
+    400: ValidationError,
+    401: AuthenticationError,
+    403: AuthorizationError,
+    404: ResourceNotFoundError,
+    422: ValidationError,
+    429: RateLimitError,
+    500: ServerError,
+    502: ServerError,
+    503: ServerError,
+    504: ServerError,
+    # Request error types
+    httpx.ConnectTimeout: ConnectionError,
+    httpx.ReadTimeout: ConnectionError,
+    httpx.ConnectError: NetworkError,
+    httpx.RequestError: NetworkError,
+}
+
+
+def extract_error_details(response: httpx.Response) -> Dict[str, Any]:
+    """Extract error details from an HTTP response.
+
+    Args:
+        response: The HTTP response
+
+    Returns:
+        A dictionary of error details
+
+    """
+    details = {}
+
+    # Try to parse JSON response
+    try:
+        if response.headers.get('content-type', '').startswith('application/json'):
+            details = response.json()
+    except Exception as e:
+        logger.debug(f'Failed to parse JSON response: {e}')
+
+    # If we couldn't parse JSON, use the text
+    if not details and response.text:
+        details = {'message': response.text}
+
+    return details
+
+
+def format_error_message(status_code: int, reason: str, details: Dict[str, Any]) -> str:
+    """Format an error message from status code, reason, and details.
+
+    Args:
+        status_code: HTTP status code
+        reason: HTTP reason phrase
+        details: Additional error details
+
+    Returns:
+        A formatted error message
+
+    """
+    # Start with the status code and reason
+    message = f'{status_code} {reason}'
+
+    # Add details if available
+    if details:
+        # Try to extract a message from the details
+        if 'message' in details:
+            message += f': {details["message"]}'
+        elif 'error' in details:
+            if isinstance(details['error'], str):
+                message += f': {details["error"]}'
+            elif isinstance(details['error'], dict) and 'message' in details['error']:
+                message += f': {details["error"]["message"]}'
+
+    # Add troubleshooting tips based on status code
+    if status_code == 401:
+        message += '\n\nTROUBLESHOOTING: Authentication error. Please check your credentials or ensure your token is valid. You may need to refresh your authentication tokens.'
+    elif status_code == 403:
+        message += "\n\nTROUBLESHOOTING: Authorization error. You don't have permission to access this resource. Please check your IAM permissions or API key scope."
+
+    return message
+
+
+def handle_http_error(error: httpx.HTTPStatusError) -> APIError:
+    """Convert an HTTPX error to an appropriate APIError subclass."""
+    status_code = error.response.status_code
+    details = extract_error_details(error.response)
+    message = format_error_message(status_code, error.response.reason_phrase, details)
+
+    # Enhanced logging for auth errors
+    if status_code == 401:
+        # Extract and log authorization header (masked) for debugging
+        request = error.request
+        if request and hasattr(request, 'headers') and 'Authorization' in request.headers:
+            auth_header = request.headers['Authorization']
+            # Safely mask the token
+            if auth_header.startswith('Bearer '):
+                token = auth_header[7:]
+                # Try to decode JWT token for debugging (without validation)
+                try:
+                    # Split the token into parts
+                    parts = token.split('.')
+                    if len(parts) == 3:
+                        # Decode the payload (middle part)
+                        # Add padding if needed
+                        payload = parts[1]
+                        padding = len(payload) % 4
+                        if padding:
+                            payload += '=' * (4 - padding)
+
+                        # Decode base64
+                        import base64
+
+                        decoded = base64.b64decode(payload)
+                        payload_data = json.loads(decoded)
+
+                        # Check for expiration
+                        if 'exp' in payload_data:
+                            import time
+
+                            exp_time = payload_data['exp']
+                            now = int(time.time())
+                            remaining = exp_time - now
+
+                            if remaining < 0:
+                                logger.warning(f'Token expired {abs(remaining)} seconds ago')
+                            else:
+                                logger.debug(
+                                    f'Token expiration: {exp_time}, Current time: {now}, Remaining: {remaining}s'
+                                )
+                except Exception as e:
+                    logger.warning(f'Could not decode token for debugging: {e}')
+
+    # Use the appropriate error class based on status code
+    error_class = ERROR_CLASSES.get(status_code, APIError)
+    return error_class(status_code, message, details=details, original_error=error)
+
+
+def handle_request_error(error: httpx.RequestError) -> APIError:
+    """Convert an HTTPX request error to an appropriate APIError subclass."""
+    # Map different request error types to different messages
+    error_class = ConnectionError
+    for error_type in [
+        httpx.ConnectTimeout,
+        httpx.ReadTimeout,
+        httpx.ConnectError,
+        httpx.RequestError,
+    ]:
+        if isinstance(error, error_type):
+            error_class = ERROR_CLASSES.get(error_type, ConnectionError)
+            break
+
+    # Get more specific error message based on error type
+    if isinstance(error, httpx.ConnectTimeout):
+        message = 'Connection timed out: The server took too long to respond'
+    elif isinstance(error, httpx.ReadTimeout):
+        message = 'Read timed out: The server took too long to send a response'
+    elif isinstance(error, httpx.ConnectError):
+        message = f'Connection error: Could not connect to the server: {error}'
+    else:
+        message = f'Request error: {error}'
+
+    # Create the error
+    return error_class(500, message, original_error=error)
+
+
+async def safe_request(
+    client: httpx.AsyncClient, method: str, url: str, **kwargs
+) -> httpx.Response:
+    """Execute an HTTP request with comprehensive error handling.
+
+    Args:
+        client: The HTTPX client to use for the request
+        method: The HTTP method to use
+        url: The URL to request
+        **kwargs: Additional arguments to pass to the client's request method
+
+    Returns:
+        The HTTP response
+
+    Raises:
+        APIError: If an error occurs during the request
+
+    """
+    try:
+        # Log request details at DEBUG level
+        request_details = {
+            'method': method,
+            'url': url,
+        }
+
+        # Log headers (safely) if present in kwargs
+        if 'headers' in kwargs and kwargs['headers']:
+            sanitized_headers = {}
+            for header, value in kwargs['headers'].items():
+                if header.lower() == 'authorization' and value:
+                    # Mask authorization header for security
+                    if value.startswith('Bearer ') and len(value) > 15:
+                        sanitized_headers[header] = (
+                            'Bearer ' + value[7:15] + '...' + value[-8:]
+                            if len(value) > 30
+                            else 'Bearer ****'
+                        )
+                    else:
+                        sanitized_headers[header] = '[MASKED]'
+                else:
+                    sanitized_headers[header] = str(value)
+            request_details['headers'] = sanitized_headers
+
+        # Log query params if present
+        if 'params' in kwargs and kwargs['params']:
+            # Convert all values to strings to avoid type issues
+            params_dict = {}
+            for k, v in kwargs['params'].items():
+                params_dict[k] = str(v) if v is not None else None
+            request_details['params'] = params_dict
+
+        logger.debug(f'Making HTTP request: {request_details}')
+
+        # Make the request
+        response = await client.request(method=method, url=url, **kwargs)
+
+        # Log response details at DEBUG level
+        logger.debug(f'Response: {response.status_code} {response.reason_phrase}')
+
+        # Raise an exception for 4xx/5xx responses
+        response.raise_for_status()
+
+        return response
+
+    except httpx.HTTPStatusError as e:
+        # Handle HTTP errors (4xx, 5xx)
+        logger.error(
+            f'HTTP error when accessing {url}: {e.response.status_code} {e.response.reason_phrase}'
+        )
+        raise handle_http_error(e)
+
+    except httpx.RequestError as e:
+        # Handle request errors (connection, timeout, etc.)
+        logger.error(f'Request error when accessing {url}: {e}')
+
+        # Create a more specific error
+        raise handle_request_error(e)
+
+    except Exception as e:
+        # Handle unexpected errors
+        logger.error(f'Unexpected error when accessing {url}: {e}')
+        raise APIError(500, f'Unexpected error: {e}', original_error=e)