salim0986/graph-bug-ai

1"""2Rate Limiting Middleware3Protects API from abuse and ensures fair usage across installations4"""5 6import time7from typing import Dict, Optional8from collections import defaultdict9from fastapi import Request, HTTPException10from starlette.middleware.base import BaseHTTPMiddleware11from starlette.responses import JSONResponse12from .logger import setup_logger13 14logger = setup_logger(__name__)15 16 17class RateLimiter:18    """19    Token bucket rate limiter with per-installation tracking20    21    Features:22    - Separate limits for different endpoint types23    - Installation-based tracking24    - Graceful degradation (warning logs instead of hard blocks for now)25    """26    27    def __init__(self):28        # Track: installation_id -> {endpoint_type: [(timestamp, tokens_used)]}29        self.buckets: Dict[str, Dict[str, list]] = defaultdict(lambda: defaultdict(list))30        31        # Rate limits: requests per minute32        self.limits = {33            "ingestion": 10,      # Heavy operations34            "search": 60,         # Medium operations35            "review": 30,         # Heavy AI operations  36            "webhook": 100,       # Lightweight webhooks37            "default": 120        # General API calls38        }39        40        # Cleanup old entries every 5 minutes41        self.last_cleanup = time.time()42        self.cleanup_interval = 300  # 5 minutes43    44    def _cleanup_old_entries(self):45        """Remove entries older than 1 minute"""46        current_time = time.time()47        if current_time - self.last_cleanup < self.cleanup_interval:48            return49        50        cutoff_time = current_time - 60  # 1 minute ago51        52        for installation_id in list(self.buckets.keys()):53            for endpoint_type in list(self.buckets[installation_id].keys()):54                self.buckets[installation_id][endpoint_type] = [55                    (ts, tokens) for ts, tokens in self.buckets[installation_id][endpoint_type]56                    if ts > cutoff_time57                ]58                59                # Remove empty endpoint types60                if not self.buckets[installation_id][endpoint_type]:61                    del self.buckets[installation_id][endpoint_type]62            63            # Remove empty installations64            if not self.buckets[installation_id]:65                del self.buckets[installation_id]66        67        self.last_cleanup = current_time68        logger.debug(f"Cleaned up rate limiter, {len(self.buckets)} installations tracked")69    70    def check_rate_limit(71        self, 72        installation_id: str, 73        endpoint_type: str = "default",74        tokens: int = 175    ) -> tuple[bool, Optional[str]]:76        """77        Check if request is within rate limit78        79        Args:80            installation_id: GitHub installation ID (or "anonymous")81            endpoint_type: Type of endpoint (ingestion, search, review, webhook, default)82            tokens: Number of tokens to consume (default 1)83            84        Returns:85            (allowed: bool, error_message: Optional[str])86        """87        self._cleanup_old_entries()88        89        current_time = time.time()90        cutoff_time = current_time - 60  # 1 minute window91        92        # Get recent requests93        recent_requests = [94            (ts, t) for ts, t in self.buckets[installation_id][endpoint_type]95            if ts > cutoff_time96        ]97        98        # Count tokens used in last minute99        tokens_used = sum(t for _, t in recent_requests)100        limit = self.limits.get(endpoint_type, self.limits["default"])101        102        if tokens_used + tokens > limit:103            remaining_time = 60 - (current_time - recent_requests[0][0]) if recent_requests else 0104            error_msg = (105                f"Rate limit exceeded for {endpoint_type}. "106                f"Limit: {limit} requests/min. "107                f"Current: {tokens_used}. "108                f"Retry in {int(remaining_time)}s"109            )110            logger.warning(f"🚫 {error_msg} | Installation: {installation_id}")111            return False, error_msg112        113        # Record this request114        self.buckets[installation_id][endpoint_type].append((current_time, tokens))115        116        # Log if approaching limit (>80%)117        if tokens_used + tokens > limit * 0.8:118            logger.warning(119                f"⚠️ Installation {installation_id} approaching rate limit: "120                f"{tokens_used + tokens}/{limit} for {endpoint_type}"121            )122        123        return True, None124    125    def get_usage_stats(self, installation_id: str) -> Dict[str, int]:126        """Get current usage statistics for an installation"""127        current_time = time.time()128        cutoff_time = current_time - 60129        130        stats = {}131        for endpoint_type, requests in self.buckets[installation_id].items():132            recent = [t for ts, t in requests if ts > cutoff_time]133            stats[endpoint_type] = sum(recent)134        135        return stats136 137 138class RateLimitMiddleware(BaseHTTPMiddleware):139    """140    FastAPI middleware for rate limiting141    142    Applies rate limits based on endpoint path and installation ID143    """144    145    def __init__(self, app, limiter: RateLimiter):146        super().__init__(app)147        self.limiter = limiter148    149    def _get_endpoint_type(self, path: str) -> str:150        """Determine endpoint type from request path"""151        if "/ingest" in path or "/process" in path:152            return "ingestion"153        elif "/search" in path or "/context" in path:154            return "search"155        elif "/review" in path or "/generate" in path:156            return "review"157        elif "/webhook" in path:158            return "webhook"159        else:160            return "default"161    162    async def _get_rate_limit_key(self, request: Request) -> str:163        """164        Return the rate-limit key for this request.165        Keys by installation_id when available, falls back to IP.166        M12: reads POST body (cached by Starlette) to extract installation_id.167        """168        # Query parameter169        iid = request.query_params.get("installation_id")170        if iid:171            return f"inst:{iid}"172 173        # Header174        iid = request.headers.get("X-Installation-ID")175        if iid:176            return f"inst:{iid}"177 178        # POST/PUT body — Starlette caches body after first read179        if request.method in ("POST", "PUT", "PATCH"):180            try:181                import json as _json182                body_bytes = await request.body()183                if body_bytes:184                    body_data = _json.loads(body_bytes)185                    iid = str(body_data.get("installation_id", ""))186                    if iid:187                        return f"inst:{iid}"188            except Exception:189                pass190 191        # Fallback to IP address192        ip = request.client.host if request.client else "anonymous"193        return f"ip:{ip}"194 195    async def dispatch(self, request: Request, call_next):196        """Process request through rate limiter"""197 198        # Skip rate limiting for health checks and docs199        if request.url.path in ["/health", "/", "/docs", "/openapi.json"]:200            return await call_next(request)201 202        endpoint_type = self._get_endpoint_type(request.url.path)203        installation_id = await self._get_rate_limit_key(request)204        205        # Check rate limit206        allowed, error_msg = self.limiter.check_rate_limit(207            installation_id, 208            endpoint_type209        )210        211        if not allowed:212            return JSONResponse(213                status_code=429,214                content={215                    "error": "rate_limit_exceeded",216                    "message": error_msg,217                    "installation_id": installation_id,218                    "endpoint_type": endpoint_type219                },220                headers={221                    "Retry-After": "60",222                    "X-RateLimit-Limit": str(self.limiter.limits.get(endpoint_type, 120)),223                    "X-RateLimit-Remaining": "0"224                }225            )226        227        # Process request228        response = await call_next(request)229        230        # Add rate limit headers to response231        usage = self.limiter.get_usage_stats(installation_id)232        limit = self.limiter.limits.get(endpoint_type, 120)233        current_usage = usage.get(endpoint_type, 0)234        235        response.headers["X-RateLimit-Limit"] = str(limit)236        response.headers["X-RateLimit-Remaining"] = str(max(0, limit - current_usage))237        response.headers["X-RateLimit-Reset"] = str(int(time.time()) + 60)238        239        return response240