Jack1808/Claude_Code

1"""Global rate limiter for API requests."""2 3import asyncio4import random5import time6from collections import deque7from collections.abc import AsyncIterator, Callable8from contextlib import asynccontextmanager9from typing import Any, ClassVar, TypeVar10 11import openai12from loguru import logger13 14T = TypeVar("T")15 16 17class GlobalRateLimiter:18    """19    Global singleton rate limiter that blocks all requests20    when a rate limit error is encountered (reactive) and21    throttles requests (proactive) using a strict rolling window.22 23    Optionally enforces a max_concurrency cap: at most N provider streams24    may be open simultaneously, independent of the sliding window.25 26    Proactive limits - throttles requests to stay within API limits.27    Reactive limits - pauses all requests when a 429 is hit.28    Concurrency limit - caps simultaneously open streams.29    """30 31    _instance: ClassVar["GlobalRateLimiter | None"] = None32 33    def __new__(cls, *args: Any, **kwargs: Any) -> "GlobalRateLimiter":34        if cls._instance is not None:35            return cls._instance36        instance = super().__new__(cls)37        return instance38 39    def __init__(40        self,41        rate_limit: int = 40,42        rate_window: float = 60.0,43        max_concurrency: int = 5,44    ):45        # Prevent re-initialization on singleton reuse46        if hasattr(self, "_initialized"):47            return48 49        if rate_limit <= 0:50            raise ValueError("rate_limit must be > 0")51        if rate_window <= 0:52            raise ValueError("rate_window must be > 0")53        if max_concurrency <= 0:54            raise ValueError("max_concurrency must be > 0")55 56        self._rate_limit = rate_limit57        self._rate_window = float(rate_window)58        # Monotonic timestamps of the last granted slots.59        self._request_times: deque[float] = deque()60        self._blocked_until: float = 061        self._lock = asyncio.Lock()62        self._concurrency_sem = asyncio.Semaphore(max_concurrency)63        self._initialized = True64 65        logger.info(66            f"GlobalRateLimiter (Provider) initialized ({rate_limit} req / {rate_window}s, max_concurrency={max_concurrency})"67        )68 69    @classmethod70    def get_instance(71        cls,72        rate_limit: int | None = None,73        rate_window: float | None = None,74        max_concurrency: int = 5,75    ) -> "GlobalRateLimiter":76        """Get or create the singleton instance.77 78        Args:79            rate_limit: Requests per window (only used on first creation)80            rate_window: Window in seconds (only used on first creation)81            max_concurrency: Max simultaneous open streams (only used on first creation)82        """83        if cls._instance is None:84            cls._instance = cls(85                rate_limit=rate_limit or 40,86                rate_window=rate_window or 60.0,87                max_concurrency=max_concurrency,88            )89        return cls._instance90 91    @classmethod92    def reset_instance(cls) -> None:93        """Reset singleton (for testing)."""94        cls._instance = None95 96    async def wait_if_blocked(self) -> bool:97        """98        Wait if currently rate limited or throttle to meet quota.99 100        Returns:101            True if was reactively blocked and waited, False otherwise.102        """103        # 1. Reactive check: Wait if someone hit a 429104        waited_reactively = False105        now = time.monotonic()106        if now < self._blocked_until:107            wait_time = self._blocked_until - now108            logger.warning(109                f"Global provider rate limit active (reactive), waiting {wait_time:.1f}s..."110            )111            await asyncio.sleep(wait_time)112            waited_reactively = True113 114        # 2. Proactive check: strict rolling window (no bursts beyond N in last W seconds)115        await self._acquire_proactive_slot()116        return waited_reactively117 118    async def _acquire_proactive_slot(self) -> None:119        """120        Acquire a proactive slot enforcing a strict rolling window.121 122        Guarantees: at most `self._rate_limit` acquisitions in any interval of length123        `self._rate_window` (seconds).124        """125        while True:126            wait_time = 0.0127            async with self._lock:128                now = time.monotonic()129                cutoff = now - self._rate_window130 131                while self._request_times and self._request_times[0] <= cutoff:132                    self._request_times.popleft()133 134                if len(self._request_times) < self._rate_limit:135                    self._request_times.append(now)136                    return137 138                oldest = self._request_times[0]139                wait_time = max(0.0, (oldest + self._rate_window) - now)140 141            # Sleep outside the lock so other tasks can continue to queue.142            if wait_time > 0:143                await asyncio.sleep(wait_time)144            else:145                await asyncio.sleep(0)146 147    def set_blocked(self, seconds: float = 60) -> None:148        """149        Set global block for specified seconds (reactive).150 151        Args:152            seconds: How long to block (default 60s)153        """154        self._blocked_until = time.monotonic() + seconds155        logger.warning(f"Global provider rate limit set for {seconds:.1f}s (reactive)")156 157    def is_blocked(self) -> bool:158        """Check if currently reactively blocked."""159        return time.monotonic() < self._blocked_until160 161    def remaining_wait(self) -> float:162        """Get remaining reactive wait time in seconds."""163        return max(0.0, self._blocked_until - time.monotonic())164 165    @asynccontextmanager166    async def concurrency_slot(self) -> AsyncIterator[None]:167        """Async context manager that holds one concurrency slot for a stream.168 169        Blocks until a slot is available (controlled by max_concurrency).170        """171        await self._concurrency_sem.acquire()172        try:173            yield174        finally:175            self._concurrency_sem.release()176 177    async def execute_with_retry(178        self,179        fn: Callable[..., Any],180        *args: Any,181        max_retries: int = 3,182        base_delay: float = 2.0,183        max_delay: float = 60.0,184        jitter: float = 1.0,185        **kwargs: Any,186    ) -> Any:187        """Execute an async callable with rate limiting and retry on 429.188 189        Waits for the proactive limiter before each attempt. On 429, applies190        exponential backoff with jitter before retrying.191 192        Args:193            fn: Async callable to execute.194            max_retries: Maximum number of retry attempts after the first failure.195            base_delay: Base delay in seconds for exponential backoff.196            max_delay: Maximum delay cap in seconds.197            jitter: Maximum random jitter in seconds added to each delay.198 199        Returns:200            The result of the callable.201 202        Raises:203            The last exception if all retries are exhausted.204        """205        last_exc: Exception | None = None206 207        for attempt in range(1 + max_retries):208            await self.wait_if_blocked()209 210            try:211                return await fn(*args, **kwargs)212            except openai.RateLimitError as e:213                last_exc = e214                if attempt >= max_retries:215                    logger.warning(216                        f"Rate limit retry exhausted after {max_retries} retries"217                    )218                    break219 220                delay = min(base_delay * (2**attempt), max_delay)221                delay += random.uniform(0, jitter)222                logger.warning(223                    f"Rate limited (429), attempt {attempt + 1}/{max_retries + 1}. "224                    f"Retrying in {delay:.1f}s..."225                )226                self.set_blocked(delay)227                await asyncio.sleep(delay)228 229        assert last_exc is not None230        raise last_exc231