implement rate limit notification

This commit is contained in:
rambros 2026-02-25 21:36:22 +05:30
parent 1ba8b3e8ea
commit ed0e79d135

View file

@ -2,11 +2,13 @@ import sys
import asyncio import asyncio
import logging import logging
import re import re
import time
from rich.console import Console from rich.console import Console
from rich.prompt import Prompt, Confirm from rich.prompt import Prompt, Confirm
from rich.table import Table from rich.table import Table
from rich.panel import Panel from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn, ProgressColumn
from rich.text import Text
from src.core.configuration import load_config, save_config from src.core.configuration import load_config, save_config
from src.core.base import MigrationContext from src.core.base import MigrationContext
@ -21,8 +23,19 @@ import src.stoat.migrate_message as stoat_migrate
from src.core.audit import log_audit_event from src.core.audit import log_audit_event
global_rate_limit_msg = ""
global_rate_limit_expires = 0.0
class RateLimitColumn(ProgressColumn):
"""Renders the current dynamic rate limit wait."""
def render(self, task) -> Text:
global global_rate_limit_msg, global_rate_limit_expires
if time.time() < global_rate_limit_expires:
return Text.from_markup(f"[dim]\\[wait: {global_rate_limit_msg}][/dim]")
return Text("")
class RateLimitHandler(logging.Handler): class RateLimitHandler(logging.Handler):
"""Intersects library logs to print clean rate limit messages.""" """Intersects library logs to capture clean dynamic rate limit messages."""
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self._last_print = "" self._last_print = ""
@ -30,27 +43,28 @@ class RateLimitHandler(logging.Handler):
def emit(self, record): def emit(self, record):
try: try:
msg = record.getMessage() msg = record.getMessage()
# Detect rate limit messages from discord.py or fluxer.py # Detect rate limit messages from discord.py, fluxer.py, stoat, etc.
if "retry" in msg.lower() and ("rate limit" in msg.lower() or "429" in msg): if "retry" in msg.lower() and ("rate limit" in msg.lower() or "429" in msg):
# Extract seconds using regex: supports "retry in 5.50s" and "Retrying in 5.50 seconds" # Extract seconds using regex: supports "retry in 5.50s" and "Retrying in 5.50 seconds"
match = re.search(r"in ([\d.]+)\s*(?:seconds?|s)", msg, re.IGNORECASE) match = re.search(r"in ([\d.]+)\s*(?:seconds?|s)", msg, re.IGNORECASE)
if match: if match:
seconds = match.group(1) seconds = match.group(1)
platform = "discord" if "discord" in record.name.lower() else "fluxer" if "discord" in record.name.lower():
platform = "Discord"
elif "fluxer" in record.name.lower():
platform = "Fluxer"
elif "stoat" in record.name.lower():
platform = "Stoat"
else:
platform = "API"
# Format the message # Update the global dynamic rate limit info
new_msg = f"{platform} API rate limit: will retry after {seconds}" global global_rate_limit_msg, global_rate_limit_expires
global_rate_limit_msg = f"{platform} rate limit {seconds}s"
# Avoid spamming the exact same message if nothing changed try:
if new_msg == self._last_print: global_rate_limit_expires = time.time() + float(seconds)
return except ValueError:
pass
self._last_print = new_msg
# Use rich console to print on the same line.
# end="\r" works with rich's internal live-update handling.
# We add some padding to clear old text.
console.print(f"{new_msg} ", end="\r")
except Exception: except Exception:
pass pass
@ -76,6 +90,7 @@ class MigrationCLI:
rl_handler = RateLimitHandler() rl_handler = RateLimitHandler()
logging.getLogger("discord").addHandler(rl_handler) logging.getLogger("discord").addHandler(rl_handler)
logging.getLogger("fluxer").addHandler(rl_handler) logging.getLogger("fluxer").addHandler(rl_handler)
logging.getLogger("stoat").addHandler(rl_handler)
async def validate_config(self): async def validate_config(self):
self.validation_results = { self.validation_results = {
@ -454,7 +469,7 @@ class MigrationCLI:
channels = [] channels = []
try: try:
await self.engine.start_connections() await self.engine.start_connections()
with console.status("[yellow]Syncing Fluxer channel state...[/yellow]"): with console.status(f"[yellow]Syncing {self.target_platform.capitalize()} channel state...[/yellow]"):
await sync_channel_state(self.engine) await sync_channel_state(self.engine)
categories = await self.engine.discord_reader.get_categories() categories = await self.engine.discord_reader.get_categories()
channels = await self.engine.discord_reader.get_channels() channels = await self.engine.discord_reader.get_channels()
@ -544,6 +559,7 @@ class MigrationCLI:
TextColumn("[progress.description]{task.description}"), TextColumn("[progress.description]{task.description}"),
BarColumn(), BarColumn(),
TaskProgressColumn(), TaskProgressColumn(),
RateLimitColumn(),
console=console console=console
) as progress: ) as progress:
@ -649,6 +665,7 @@ class MigrationCLI:
TextColumn("[progress.description]{task.description}"), TextColumn("[progress.description]{task.description}"),
BarColumn(), BarColumn(),
TaskProgressColumn(), TaskProgressColumn(),
RateLimitColumn(),
console=console console=console
) as progress: ) as progress:
@ -700,6 +717,7 @@ class MigrationCLI:
TextColumn("[progress.description]{task.description}"), TextColumn("[progress.description]{task.description}"),
BarColumn(), BarColumn(),
TaskProgressColumn(), TaskProgressColumn(),
RateLimitColumn(),
console=console console=console
) as progress: ) as progress:
@ -825,6 +843,7 @@ class MigrationCLI:
TextColumn("[progress.description]{task.description}"), TextColumn("[progress.description]{task.description}"),
BarColumn(), BarColumn(),
TaskProgressColumn(), TaskProgressColumn(),
RateLimitColumn(),
console=console console=console
) as progress: ) as progress:
@ -1233,6 +1252,7 @@ class MigrationCLI:
TextColumn("[progress.description]{task.description}"), TextColumn("[progress.description]{task.description}"),
BarColumn(), BarColumn(),
TaskProgressColumn(), TaskProgressColumn(),
RateLimitColumn(),
console=console console=console
) as progress: ) as progress:
task = progress.add_task(f"[cyan]Migrating 0/{total_messages} messages...", total=total_messages) task = progress.add_task(f"[cyan]Migrating 0/{total_messages} messages...", total=total_messages)