core/decorators.py

"""Shared decorators for MCP tools."""

import functools
import json
import logging
from typing import Any, Callable, TypeVar
from tenacity import retry, stop_after_attempt, wait_exponential

logger = logging.getLogger(__name__)

F = TypeVar('F', bound=Callable[..., Any])


def safe_json_return(func: F) -> F:
    """
    Decorator that ensures MCP tools always return JSON strings.
    Handles dicts, lists, None, and exceptions consistently.
    """
    @functools.wraps(func)
    async def async_wrapper(*args, **kwargs):
        try:
            result = await func(*args, **kwargs)

            # Handle None
            if result is None:
                return json.dumps({"result": None})

            # Handle dict or list
            if isinstance(result, (dict, list)):
                return json.dumps(result, ensure_ascii=False)

            # Handle string (might already be JSON)
            if isinstance(result, str):
                return result

            # Handle other types
            return json.dumps({"result": str(result)})

        except Exception as e:
            logger.exception(f"Error in {func.__name__}")
            return json.dumps({
                "error": str(e),
                "error_type": type(e).__name__
            })

    @functools.wraps(func)
    def sync_wrapper(*args, **kwargs):
        try:
            result = func(*args, **kwargs)

            if result is None:
                return json.dumps({"result": None})

            if isinstance(result, (dict, list)):
                return json.dumps(result, ensure_ascii=False)

            if isinstance(result, str):
                return result

            return json.dumps({"result": str(result)})

        except Exception as e:
            logger.exception(f"Error in {func.__name__}")
            return json.dumps({
                "error": str(e),
                "error_type": type(e).__name__
            })

    # Return appropriate wrapper based on function type
    import inspect
    if inspect.iscoroutinefunction(func):
        return async_wrapper
    else:
        return sync_wrapper


def with_retry(func: F) -> F:
    """
    Decorator for retrying failed HTTP requests with exponential backoff.
    """
    return retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=8),
        reraise=True
    )(func)