"""Runtime-owned queued delivery for one messaging platform.""" import asyncio from collections import deque from collections.abc import Awaitable, Callable from typing import Any from loguru import logger from free_claude_code.core.rate_limit import ( StrictSlidingWindowLimiter as SlidingWindowLimiter, ) from .safe_diagnostics import format_exception_for_log class MessagingRateLimiter: """ Rate limiter and compacting work queue for one messaging runtime. Uses a custom queue with task compaction (deduplication) to ensure only the latest version of a message update is processed. """ def __init__( self, *, rate_limit: int, rate_window: float, log_error_details: bool = False, ) -> None: self.limiter = SlidingWindowLimiter(rate_limit, rate_window) self._log_error_details = log_error_details # 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 = None self._background_tasks: set[asyncio.Task[None]] = set() self._active_futures: list[asyncio.Future[Any]] = [] self._closed = False self._paused_until = 0 logger.info( f"MessagingRateLimiter initialized ({rate_limit} req / {rate_window}s with Task Compaction)" ) def start(self) -> None: """Start the owned worker on the current event loop.""" if self._closed: raise RuntimeError("Messaging rate limiter is closed.") if self._worker_task and not self._worker_task.done(): return self._worker_task = asyncio.create_task( self._worker(), name="msg-limiter-worker" ) async def _worker(self) -> None: """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) self._active_futures = futures # 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 asyncio.CancelledError: for f in futures: if not f.done(): f.cancel() worker = asyncio.current_task() if self._shutdown.is_set() or ( worker is not None and worker.cancelling() ): raise logger.debug( "Messaging operation cancelled for key {}; worker remains active", dedup_key, ) 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: logger.error( "Error in limiter worker for key {}: {}", dedup_key, format_exception_for_log( e, log_full_message=self._log_error_details, ), ) finally: self._active_futures = [] except asyncio.CancelledError: for future in self._active_futures: if not future.done(): future.cancel() self._active_futures = [] if self._shutdown.is_set(): break raise except Exception as e: if self._log_error_details: 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 | None = None) -> None: """Cancel queued work and stop every task owned by this limiter.""" self._closed = True self._shutdown.set() async with self._condition: queued_futures = [ future for _func, futures in self._queue_map.values() for future in futures ] self._queue_list.clear() self._queue_map.clear() for future in queued_futures: if not future.done(): future.cancel() for future in self._active_futures: if not future.done(): future.cancel() self._condition.notify_all() cancellation: asyncio.CancelledError | None = None timeout_error: TimeoutError | None = None task = self._worker_task if task and not task.done(): task.cancel() try: drain = asyncio.gather(task, return_exceptions=True) if timeout is None: await drain else: await asyncio.wait_for(drain, timeout=timeout) except TimeoutError as exc: timeout_error = exc except asyncio.CancelledError as exc: cancellation = exc if task is None or task.done(): self._worker_task = None background_tasks = tuple(self._background_tasks) for background_task in background_tasks: background_task.cancel() if background_tasks: try: await asyncio.gather(*background_tasks, return_exceptions=True) except asyncio.CancelledError as exc: cancellation = exc self._background_tasks.difference_update( task for task in background_tasks if task.done() ) if cancellation is not None: raise cancellation if timeout_error is not None: raise TimeoutError( "MessagingRateLimiter worker did not stop before timeout" ) from timeout_error async def _enqueue_internal( self, func: Callable[[], Awaitable[Any]], future: asyncio.Future[Any], dedup_key: str, front: bool = False, ) -> None: await self._enqueue_internal_multi(func, [future], dedup_key, front) async def _enqueue_internal_multi( self, func: Callable[[], Awaitable[Any]], futures: list[asyncio.Future[Any]], dedup_key: str, front: bool = False, ) -> None: async with self._condition: if self._closed: raise RuntimeError("Messaging rate limiter is closed.") if self._worker_task is None or self._worker_task.done(): raise RuntimeError("Messaging rate limiter has not been started.") 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. """ self._require_running() if dedup_key is None: # Unique key to avoid deduplication dedup_key = f"task_{id(func)}_{asyncio.get_running_loop().time()}" future = asyncio.get_running_loop().create_future() try: await self._enqueue_internal(func, future, dedup_key) except BaseException: future.cancel() raise return await future def fire_and_forget( self, func: Callable[[], Awaitable[Any]], dedup_key: str | None = None ) -> None: """Enqueue a task without waiting for the result.""" self._require_running() if dedup_key is None: dedup_key = f"task_{id(func)}_{asyncio.get_running_loop().time()}" async def _wrapped() -> None: max_retries = 2 for attempt in range(max_retries + 1): try: await self.enqueue(func, dedup_key) return 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 if self._log_error_details: 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 logger.error( "Final error in fire_and_forget for key {}: {}", dedup_key, format_exception_for_log( e, log_full_message=self._log_error_details, ), ) break task = asyncio.create_task(_wrapped(), name=f"msg-limiter:{dedup_key}") self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) def _require_running(self) -> None: if self._closed: raise RuntimeError("Messaging rate limiter is closed.") if self._worker_task is None or self._worker_task.done(): raise RuntimeError("Messaging rate limiter has not been started.")