Add configurable per-user command cooldown

.env.example

@@ -9,6 +9,9 @@ DATABASE_URL=sqlite+aiosqlite:///fun2oosh.db  # For development
 # Logging
 LOG_LEVEL=INFO
 
+# Command Cooldown
+COMMAND_COOLDOWN_SECONDS=3.0  # Per-user cooldown between commands in seconds
+
 # TTS Voice Channel Leave timeout in seconds
 TTS_VC_LEAVE_TIMEOUT=90
 

.github/workflows/ci.yml

@@ -9,10 +9,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: "3.14"
 

bot.py

@@ -15,6 +15,7 @@
 from dotenv import load_dotenv
 
 from utils.config import Config
+from utils.cooldown import CooldownManager
 
 # Load environment variables
 load_dotenv()
@@ -52,6 +53,9 @@ def __init__(self, config: Config):
 
         self.start_time = discord.utils.utcnow()
         self.config = config
+        self.cooldown_manager = CooldownManager(config.command_cooldown_seconds)
+        self.tree.on_error = self.on_app_command_error
+        self._setup_cooldowns()
         # Discover available cog modules from the cogs directory
         from pathlib import Path
 
@@ -63,6 +67,57 @@ def __init__(self, config: Config):
                     self.available_cogs.append(p.stem)
         logger.info(f"Available cogs discovered: {self.available_cogs}")
 
+    def _setup_cooldowns(self) -> None:
+        """Register global cooldown checks for prefix and slash commands."""
+
+        @self.check
+        async def prefix_cooldown_check(ctx: commands.Context) -> bool:
+            if self.config.owner_id and ctx.author.id == self.config.owner_id:
+                return True
+
+            # If command defines its own cooldown, let discord.py handle it
+            cmd = ctx.command
+            if cmd is not None and cmd._buckets and cmd._buckets.valid:
+                return True
+
+            retry_after = self.cooldown_manager.trigger(ctx.author.id)
+            if retry_after is not None:
+                raise commands.CommandOnCooldown(
+                    commands.Cooldown(1, self.cooldown_manager.cooldown_seconds),
+                    retry_after,
+                    commands.BucketType.user,
+                )
+            return True
+
+        async def app_command_cooldown_check(
+            interaction: discord.Interaction,
+        ) -> bool:
+            # Only rate-limit application command executions (not autocomplete or components)
+            if interaction.type != discord.InteractionType.application_command:
+                return True
+
+            if self.config.owner_id and interaction.user.id == self.config.owner_id:
+                return True
+
+            cmd = interaction.command
+            if cmd is not None:
+                # Let commands with custom checks or hybrid cooldowns handle their own rate limits
+                if getattr(cmd, "checks", None):
+                    return True
+                wrapped = getattr(cmd, "wrapped", None)
+                if wrapped and getattr(wrapped, "_buckets", None) and wrapped._buckets.valid:
+                    return True
+
+            retry_after = self.cooldown_manager.trigger(interaction.user.id)
+            if retry_after is not None:
+                raise app_commands.CommandOnCooldown(
+                    app_commands.Cooldown(1, self.cooldown_manager.cooldown_seconds),
+                    retry_after,
+                )
+            return True
+
+        self.tree.interaction_check = app_command_cooldown_check
+
     async def setup_hook(self) -> None:
         """Setup hook called before the bot starts."""
         # Initialize CodeBuddy database
@@ -230,23 +285,23 @@ async def _safe_ctx_send(message: str) -> None:
             except Exception:
                 return
 
-        if isinstance(error, commands.CommandOnCooldown):
-            await _safe_ctx_send(
-                f"This command is on cooldown. Try again in {error.retry_after:.2f} seconds."
-            )
-        elif isinstance(error, commands.MissingRequiredArgument):
-            command = getattr(ctx, "command", None)
-            cog = getattr(command, "cog", None)
-            if command and command.qualified_name == "resources" and cog is not None:
-                send_usage = getattr(cog, "_send_resource_usage", None)
-                if callable(send_usage):
-                    await send_usage(ctx)
-                    return
-            await _safe_ctx_send("A required argument is missing.")
-        elif isinstance(error, commands.MissingPermissions):
-            await _safe_ctx_send("You don't have permission to use this command.")
-        elif isinstance(error, commands.BadArgument):
-            await _safe_ctx_send("Invalid argument provided.")
+        if isinstance(error, commands.CommandOnCooldown):
+            await _safe_ctx_send(
+                f"This command is on cooldown. Try again in {error.retry_after:.2f} seconds."
+            )
+        elif isinstance(error, commands.MissingRequiredArgument):
+            command = getattr(ctx, "command", None)
+            cog = getattr(command, "cog", None)
+            if command and command.qualified_name == "resources" and cog is not None:
+                send_usage = getattr(cog, "_send_resource_usage", None)
+                if callable(send_usage):
+                    await send_usage(ctx)
+                    return
+            await _safe_ctx_send("A required argument is missing.")
+        elif isinstance(error, commands.MissingPermissions):
+            await _safe_ctx_send("You don't have permission to use this command.")
+        elif isinstance(error, commands.BadArgument):
+            await _safe_ctx_send("Invalid argument provided.")
         else:
             logger.error(f"Command error: {error}")
             await _safe_ctx_send("An error occurred while processing your command.")

mypy.ini

@@ -1,5 +1,4 @@
 [mypy]
 python_version = 3.14
 ignore_missing_imports = True
-disable_error_code = arg-type, assignment, attr-defined, union-attr, var-annotated, override
-
+disable_error_code = arg-type, assignment, attr-defined, union-attr, var-annotated, override, annotation-unchecked

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+filterwarnings =
+    ignore::DeprecationWarning:discord.*

tests/test_cooldown.py

@@ -0,0 +1,89 @@
+"""Tests for the command cooldown system."""
+
+from __future__ import annotations
+
+import unittest
+from unittest.mock import patch
+
+from utils.config import Config
+from utils.cooldown import CooldownManager
+
+
+class TestCooldownManager(unittest.TestCase):
+    def test_first_invocation_is_allowed(self) -> None:
+        manager = CooldownManager(cooldown_seconds=3.0)
+        self.assertIsNone(manager.trigger(12345))
+
+    def test_subsequent_invocation_within_cooldown_is_blocked(self) -> None:
+        manager = CooldownManager(cooldown_seconds=3.0)
+
+        with patch("time.monotonic", side_effect=[100.0, 101.0, 101.0]):
+            manager.trigger(12345)
+            retry_after = manager.get_retry_after(12345)
+            if retry_after is None:
+                self.fail("Expected retry_after to be not None")
+            self.assertAlmostEqual(retry_after, 2.0, places=2)
+
+            retry = manager.trigger(12345)
+            if retry is None:
+                self.fail("Expected retry to be not None")
+            self.assertAlmostEqual(retry, 2.0, places=2)
+
+    def test_invocation_after_cooldown_window_is_allowed(self) -> None:
+        manager = CooldownManager(cooldown_seconds=3.0)
+
+        with patch("time.monotonic", side_effect=[100.0, 103.5]):
+            manager.trigger(12345)
+            self.assertIsNone(manager.trigger(12345))
+
+    def test_different_users_have_independent_cooldowns(self) -> None:
+        manager = CooldownManager(cooldown_seconds=3.0)
+
+        with patch("time.monotonic", side_effect=[100.0, 100.5, 101.0]):
+            self.assertIsNone(manager.trigger(111))
+            self.assertIsNone(manager.trigger(222))
+            self.assertIsNotNone(manager.trigger(111))
+
+    def test_zero_or_negative_cooldown_disables_rate_limiting(self) -> None:
+        zero_manager = CooldownManager(cooldown_seconds=0.0)
+        neg_manager = CooldownManager(cooldown_seconds=-1.0)
+
+        self.assertIsNone(zero_manager.trigger(123))
+        self.assertIsNone(zero_manager.trigger(123))
+        self.assertIsNone(neg_manager.trigger(123))
+        self.assertIsNone(neg_manager.trigger(123))
+
+    def test_reset_and_clear(self) -> None:
+        manager = CooldownManager(cooldown_seconds=5.0)
+        manager.trigger(100)
+        manager.trigger(200)
+
+        manager.reset(100)
+        self.assertIsNone(manager.get_retry_after(100))
+        self.assertIsNotNone(manager.get_retry_after(200))
+
+        manager.clear()
+        self.assertIsNone(manager.get_retry_after(200))
+
+    def test_pruning_expired_entries_when_cache_large(self) -> None:
+        manager = CooldownManager(cooldown_seconds=2.0)
+        # Pre-populate map with 5001 old entries
+        manager._last_used = {i: 10.0 for i in range(5001)}
+
+        with patch("time.monotonic", return_value=50.0):
+            # Triggering a new user should prune the old entries older than 50.0 - 2.0
+            manager.trigger(999999)
+
+        # Only the new user should remain
+        self.assertEqual(len(manager._last_used), 1)
+        self.assertIn(999999, manager._last_used)
+
+
+class TestCooldownConfig(unittest.TestCase):
+    def test_default_cooldown_config(self) -> None:
+        config = Config()
+        self.assertEqual(config.command_cooldown_seconds, 3.0)
+
+    def test_custom_cooldown_config(self) -> None:
+        config = Config(command_cooldown_seconds=5.5)
+        self.assertEqual(config.command_cooldown_seconds, 5.5)

utils/config.py

@@ -22,6 +22,7 @@ class Config(BaseSettings):
     topgg_token: str | None = Field(default=None)
     topgg_webhook_secret: str | None = Field(default=None)
     redis_url: str | None = Field(default=None)
+    command_cooldown_seconds: float = Field(default=3.0)
 
     # CodeBuddy settings
     question_channel_id: int | None = Field(default=None)

utils/cooldown.py

@@ -0,0 +1,57 @@
+"""Command cooldown manager to prevent command spam."""
+
+from __future__ import annotations
+
+import time
+
+
+class CooldownManager:
+    """Tracks per-user command execution timestamps to enforce a cooldown."""
+
+    def __init__(self, cooldown_seconds: float = 3.0) -> None:
+        self.cooldown_seconds = cooldown_seconds
+        self._last_used: dict[int, float] = {}
+
+    def get_retry_after(self, user_id: int, now: float | None = None) -> float | None:
+        """Return remaining cooldown seconds if user is on cooldown, else None."""
+        if self.cooldown_seconds <= 0:
+            return None
+
+        last_time = self._last_used.get(user_id)
+        if last_time is None:
+            return None
+
+        current_time = now if now is not None else time.monotonic()
+        elapsed = current_time - last_time
+        if elapsed < self.cooldown_seconds:
+            return self.cooldown_seconds - elapsed
+
+        return None
+
+    def trigger(self, user_id: int) -> float | None:
+        """Check if user is on cooldown.
+
+        If on cooldown, returns remaining seconds.
+        If not on cooldown, records current timestamp and returns None.
+        """
+        now = time.monotonic()
+        retry_after = self.get_retry_after(user_id, now=now)
+        if retry_after is not None:
+            return retry_after
+
+        self._last_used[user_id] = now
+
+        # Prevent unbounded memory growth by pruning expired entries when large
+        if len(self._last_used) > 5000:
+            cutoff = now - self.cooldown_seconds
+            self._last_used = {uid: ts for uid, ts in self._last_used.items() if ts > cutoff}
+
+        return None
+
+    def reset(self, user_id: int) -> None:
+        """Clear cooldown for a specific user."""
+        self._last_used.pop(user_id, None)
+
+    def clear(self) -> None:
+        """Clear all active cooldowns."""
+        self._last_used.clear()