mcphub/core/tool_rate_limiter.py

"""Per-tool, per-user rate limiting (F.5a.6).

Provides a small token-bucket limiter keyed by (user_id, tool_name) for
expensive tools such as AI image generation and chunked-upload finish.
Admin / env-fallback callers (``user_id is None``) are exempt.

Intentionally minimal — separate from :mod:`core.rate_limiter` (which is a
global per-client limiter). This one is scoped to specific tools with
small hourly caps.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass

from core.rate_limiter import TokenBucket

logger = logging.getLogger(__name__)


# Default caps (per hour per user). Documented in ROADMAP F.5a.6.
DEFAULT_LIMITS: dict[str, int] = {
    "wordpress_generate_and_upload_image": 10,
    "wordpress_upload_media_chunked_finish": 30,
}


@dataclass
class ToolRateLimitError(Exception):
    """Raised when a per-tool limit is exceeded."""

    tool_name: str
    limit_per_hour: int
    retry_after_seconds: float

    def __post_init__(self) -> None:
        super().__init__(
            f"Per-tool rate limit exceeded for '{self.tool_name}' "
            f"({self.limit_per_hour}/hour per user). "
            f"Retry in {self.retry_after_seconds:.0f}s."
        )

    def to_dict(self) -> dict:
        return {
            "error_code": "TOOL_RATE_LIMITED",
            "message": str(self),
            "details": {
                "tool": self.tool_name,
                "limit_per_hour": self.limit_per_hour,
                "retry_after_seconds": round(self.retry_after_seconds, 2),
            },
        }


class PerToolRateLimiter:
    """Token-bucket limiter keyed by (user_id, tool_name).

    ``check(tool, user_id)`` consumes one token and raises
    :class:`ToolRateLimitError` when the user is over quota. Admin / env
    callers (``user_id`` is None or empty) are exempt.
    """

    def __init__(self, limits: dict[str, int] | None = None) -> None:
        self._limits = dict(limits if limits is not None else DEFAULT_LIMITS)
        self._buckets: dict[tuple[str, str], TokenBucket] = {}
        self._lock = threading.Lock()

    def configure(self, tool_name: str, per_hour: int) -> None:
        """Override the per-hour cap for a tool."""
        self._limits[tool_name] = per_hour
        # Existing buckets keep their old capacity until reset — tests can
        # reset() to re-read the new limit.

    def reset(self) -> None:
        with self._lock:
            self._buckets.clear()

    def check(self, tool_name: str, user_id: str | None) -> None:
        """Consume one token for (user_id, tool_name). Exempt when user_id is falsy."""
        if not user_id:
            return
        limit = self._limits.get(tool_name)
        if limit is None or limit <= 0:
            return

        key = (user_id, tool_name)
        with self._lock:
            bucket = self._buckets.get(key)
            if bucket is None:
                bucket = TokenBucket(capacity=limit, refill_rate=limit / 3600.0)
                self._buckets[key] = bucket

        if not bucket.consume(1):
            wait = bucket.get_wait_time(1)
            logger.warning(
                "Per-tool rate limit hit: user=%s tool=%s limit=%d/h retry_after=%.1fs",
                user_id,
                tool_name,
                limit,
                wait,
            )
            raise ToolRateLimitError(
                tool_name=tool_name, limit_per_hour=limit, retry_after_seconds=wait
            )


_limiter: PerToolRateLimiter | None = None


def get_tool_rate_limiter() -> PerToolRateLimiter:
    global _limiter
    if _limiter is None:
        _limiter = PerToolRateLimiter()
    return _limiter


def set_tool_rate_limiter(limiter: PerToolRateLimiter | None) -> None:
    """Override the singleton (used by tests)."""
    global _limiter
    _limiter = limiter