mirror of
https://github.com/Alishahryar1/free-claude-code.git
synced 2026-04-28 03:20:01 +00:00
f3a7528d49
Some checks are pending
CI / checks (push) Waiting to run
Consolidates the incremental refactor work into a single change set: modular web tools (api/web_tools), native Anthropic request building and SSE block policy, OpenAI conversion and error handling, provider transports and rate limiting, messaging handler and tree queue, safe logging, smoke tests, and broad test coverage.
300 lines
12 KiB
Python
300 lines
12 KiB
Python
"""
|
|
Global Rate Limiter for Messaging Platforms.
|
|
|
|
Centralizes outgoing message requests and ensures compliance with rate limits
|
|
using a strict sliding window algorithm and a task queue.
|
|
"""
|
|
|
|
import asyncio
|
|
from collections import deque
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from loguru import logger
|
|
|
|
from config.settings import get_settings
|
|
from core.rate_limit import StrictSlidingWindowLimiter as SlidingWindowLimiter
|
|
|
|
from .safe_diagnostics import format_exception_for_log
|
|
|
|
|
|
class MessagingRateLimiter:
|
|
"""
|
|
A thread-safe global rate limiter for messaging.
|
|
|
|
Uses a custom queue with task compaction (deduplication) to ensure
|
|
only the latest version of a message update is processed.
|
|
"""
|
|
|
|
_instance: MessagingRateLimiter | None = None
|
|
_lock = asyncio.Lock()
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
return super().__new__(cls)
|
|
|
|
@classmethod
|
|
async def get_instance(
|
|
cls,
|
|
*,
|
|
rate_limit: int = 1,
|
|
rate_window: float = 1.0,
|
|
) -> MessagingRateLimiter:
|
|
"""Get the singleton instance of the limiter.
|
|
|
|
``rate_limit`` and ``rate_window`` apply only when the singleton is first
|
|
created. Call :meth:`shutdown_instance` before changing parameters.
|
|
"""
|
|
async with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls(rate_limit=rate_limit, rate_window=rate_window)
|
|
# Start the background worker (tracked for graceful shutdown).
|
|
cls._instance._start_worker()
|
|
return cls._instance
|
|
|
|
def __init__(self, *, rate_limit: int, rate_window: float) -> None:
|
|
# Prevent double initialization in singleton
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
|
|
self.limiter = SlidingWindowLimiter(rate_limit, rate_window)
|
|
# Custom queue state - using deque for O(1) popleft
|
|
self._queue_list: deque[str] = deque() # Deque of dedup_keys in order
|
|
self._queue_map: dict[
|
|
str, tuple[Callable[[], Awaitable[Any]], list[asyncio.Future]]
|
|
] = {}
|
|
self._condition = asyncio.Condition()
|
|
self._shutdown = asyncio.Event()
|
|
self._worker_task: asyncio.Task | None = None
|
|
|
|
self._initialized = True
|
|
self._paused_until = 0
|
|
|
|
logger.info(
|
|
f"MessagingRateLimiter initialized ({rate_limit} req / {rate_window}s with Task Compaction)"
|
|
)
|
|
|
|
def _start_worker(self) -> None:
|
|
"""Ensure the worker task exists."""
|
|
if self._worker_task and not self._worker_task.done():
|
|
return
|
|
# Named task helps debugging shutdown hangs.
|
|
self._worker_task = asyncio.create_task(
|
|
self._worker(), name="msg-limiter-worker"
|
|
)
|
|
|
|
async def _worker(self):
|
|
"""Background worker that processes queued messaging tasks."""
|
|
logger.info("MessagingRateLimiter worker started")
|
|
while not self._shutdown.is_set():
|
|
try:
|
|
# Get a task from the queue
|
|
async with self._condition:
|
|
while not self._queue_list and not self._shutdown.is_set():
|
|
await self._condition.wait()
|
|
|
|
if self._shutdown.is_set():
|
|
break
|
|
|
|
dedup_key = self._queue_list.popleft()
|
|
func, futures = self._queue_map.pop(dedup_key)
|
|
|
|
# Check for manual pause (FloodWait)
|
|
now = asyncio.get_event_loop().time()
|
|
if self._paused_until > now:
|
|
wait_time = self._paused_until - now
|
|
logger.warning(
|
|
f"Limiter worker paused, waiting {wait_time:.1f}s more..."
|
|
)
|
|
await asyncio.sleep(wait_time)
|
|
|
|
# Wait for rate limit capacity
|
|
async with self.limiter:
|
|
try:
|
|
result = await func()
|
|
for f in futures:
|
|
if not f.done():
|
|
f.set_result(result)
|
|
except Exception as e:
|
|
# Report error to all futures and log it
|
|
for f in futures:
|
|
if not f.done():
|
|
f.set_exception(e)
|
|
|
|
error_msg = str(e).lower()
|
|
if "flood" in error_msg or "wait" in error_msg:
|
|
seconds = 30
|
|
try:
|
|
if hasattr(e, "seconds"):
|
|
seconds = e.seconds
|
|
elif "after " in error_msg:
|
|
# Try to parse "retry after X"
|
|
parts = error_msg.split("after ")
|
|
if len(parts) > 1:
|
|
seconds = int(parts[1].split()[0])
|
|
except Exception:
|
|
pass
|
|
|
|
logger.error(
|
|
f"FloodWait detected! Pausing worker for {seconds}s"
|
|
)
|
|
wait_secs = (
|
|
float(seconds)
|
|
if isinstance(seconds, (int, float, str))
|
|
else 30.0
|
|
)
|
|
self._paused_until = (
|
|
asyncio.get_event_loop().time() + wait_secs
|
|
)
|
|
else:
|
|
d = get_settings().log_messaging_error_details
|
|
logger.error(
|
|
"Error in limiter worker for key {}: {}",
|
|
dedup_key,
|
|
format_exception_for_log(e, log_full_message=d),
|
|
)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
d = get_settings().log_messaging_error_details
|
|
if d:
|
|
logger.error(
|
|
"MessagingRateLimiter worker critical error: {}",
|
|
e,
|
|
exc_info=True,
|
|
)
|
|
else:
|
|
logger.error(
|
|
"MessagingRateLimiter worker critical error: exc_type={}",
|
|
type(e).__name__,
|
|
)
|
|
await asyncio.sleep(1)
|
|
|
|
async def shutdown(self, timeout: float = 2.0) -> None:
|
|
"""Stop the background worker so process shutdown doesn't hang."""
|
|
self._shutdown.set()
|
|
try:
|
|
async with self._condition:
|
|
self._condition.notify_all()
|
|
except Exception:
|
|
# Best-effort: condition may be bound to a closing loop.
|
|
pass
|
|
|
|
task = self._worker_task
|
|
if not task or task.done():
|
|
self._worker_task = None
|
|
return
|
|
|
|
task.cancel()
|
|
try:
|
|
await asyncio.wait_for(task, timeout=timeout)
|
|
except TimeoutError:
|
|
logger.warning("MessagingRateLimiter worker did not stop before timeout")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
d = get_settings().log_messaging_error_details
|
|
logger.debug(
|
|
"MessagingRateLimiter worker shutdown error: {}",
|
|
format_exception_for_log(e, log_full_message=d),
|
|
)
|
|
finally:
|
|
self._worker_task = None
|
|
|
|
@classmethod
|
|
async def shutdown_instance(cls, timeout: float = 2.0) -> None:
|
|
"""Shutdown and clear the singleton instance (safe to call multiple times)."""
|
|
inst = cls._instance
|
|
if not inst:
|
|
return
|
|
try:
|
|
await inst.shutdown(timeout=timeout)
|
|
finally:
|
|
cls._instance = None
|
|
|
|
async def _enqueue_internal(self, func, future, dedup_key, front=False):
|
|
await self._enqueue_internal_multi(func, [future], dedup_key, front)
|
|
|
|
async def _enqueue_internal_multi(self, func, futures, dedup_key, front=False):
|
|
async with self._condition:
|
|
if dedup_key in self._queue_map:
|
|
# Compaction: Update existing task with new func, append new futures
|
|
_old_func, old_futures = self._queue_map[dedup_key]
|
|
old_futures.extend(futures)
|
|
self._queue_map[dedup_key] = (func, old_futures)
|
|
logger.debug(
|
|
f"Compacted task for key: {dedup_key} (now {len(old_futures)} futures)"
|
|
)
|
|
else:
|
|
self._queue_map[dedup_key] = (func, futures)
|
|
if front:
|
|
self._queue_list.appendleft(dedup_key)
|
|
else:
|
|
self._queue_list.append(dedup_key)
|
|
self._condition.notify_all()
|
|
|
|
async def enqueue(
|
|
self, func: Callable[[], Awaitable[Any]], dedup_key: str | None = None
|
|
) -> Any:
|
|
"""
|
|
Enqueue a messaging task and return its future result.
|
|
If dedup_key is provided, subsequent tasks with the same key will replace this one.
|
|
"""
|
|
if dedup_key is None:
|
|
# Unique key to avoid deduplication
|
|
dedup_key = f"task_{id(func)}_{asyncio.get_event_loop().time()}"
|
|
|
|
future = asyncio.get_event_loop().create_future()
|
|
await self._enqueue_internal(func, future, dedup_key)
|
|
return await future
|
|
|
|
def fire_and_forget(
|
|
self, func: Callable[[], Awaitable[Any]], dedup_key: str | None = None
|
|
):
|
|
"""Enqueue a task without waiting for the result."""
|
|
if dedup_key is None:
|
|
dedup_key = f"task_{id(func)}_{asyncio.get_event_loop().time()}"
|
|
|
|
future = asyncio.get_event_loop().create_future()
|
|
|
|
async def _wrapped():
|
|
max_retries = 2
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return await self.enqueue(func, dedup_key)
|
|
except Exception as e:
|
|
error_msg = str(e).lower()
|
|
# Only retry transient connectivity issues that might have slipped through
|
|
# or occurred between platform checks.
|
|
if attempt < max_retries and any(
|
|
x in error_msg for x in ["connect", "timeout", "broken"]
|
|
):
|
|
wait = 2**attempt
|
|
d = get_settings().log_messaging_error_details
|
|
if d:
|
|
logger.warning(
|
|
"Limiter fire_and_forget transient error (attempt {}): {}. Retrying in {}s...",
|
|
attempt + 1,
|
|
e,
|
|
wait,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Limiter fire_and_forget transient error (attempt {}): exc_type={}. Retrying in {}s...",
|
|
attempt + 1,
|
|
type(e).__name__,
|
|
wait,
|
|
)
|
|
await asyncio.sleep(wait)
|
|
continue
|
|
|
|
d = get_settings().log_messaging_error_details
|
|
logger.error(
|
|
"Final error in fire_and_forget for key {}: {}",
|
|
dedup_key,
|
|
format_exception_for_log(e, log_full_message=d),
|
|
)
|
|
if not future.done():
|
|
future.set_exception(e)
|
|
break
|
|
|
|
_ = asyncio.create_task(_wrapped())
|