ciris-agent 1.7.7__py3-none-any.whl

ciris_engine/logic/adapters/discord/discord_rate_limiter.py

@@ -0,0 +1,249 @@
+"""Discord rate limiting component."""
+
+import asyncio
+import logging
+import time
+from collections import defaultdict
+from typing import Any, Dict
+
+from ciris_engine.schemas.types import JSONDict
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimitBucket:
+    """Represents a rate limit bucket for Discord API."""
+
+    def __init__(self, limit: int, window: float):
+        """Initialize a rate limit bucket.
+
+        Args:
+            limit: Maximum requests allowed
+            window: Time window in seconds
+        """
+        self.limit = limit
+        self.window = window
+        self.remaining = limit
+        self.reset_at = time.time() + window
+        self._lock = asyncio.Lock()
+
+    async def acquire(self) -> float:
+        """Acquire permission to make a request.
+
+        Returns:
+            Wait time in seconds (0 if can proceed immediately)
+        """
+        async with self._lock:
+            now = time.time()
+
+            # Reset if window expired
+            if now >= self.reset_at:
+                self.remaining = self.limit
+                self.reset_at = now + self.window
+
+            # Check if we can proceed
+            if self.remaining > 0:
+                self.remaining -= 1
+                return 0.0
+
+            # Need to wait
+            wait_time = self.reset_at - now
+            return wait_time
+
+    def update_from_headers(self, remaining: int, reset_at: float) -> None:
+        """Update bucket from Discord rate limit headers.
+
+        Args:
+            remaining: Remaining requests from header
+            reset_at: Reset timestamp from header
+        """
+        self.remaining = remaining
+        self.reset_at = reset_at
+
+
+class DiscordRateLimiter:
+    """Manages rate limiting for Discord API calls."""
+
+    # Discord's global rate limit
+    GLOBAL_LIMIT = 50
+    GLOBAL_WINDOW = 1.0  # 1 second
+
+    # Known endpoint limits (per 5 minutes unless specified)
+    ENDPOINT_LIMITS = {
+        # Messages
+        "channels/{channel_id}/messages": {"limit": 5, "window": 5.0},
+        "channels/{channel_id}/messages/{message_id}": {"limit": 5, "window": 5.0},
+        "channels/{channel_id}/messages/{message_id}/reactions": {"limit": 1, "window": 0.25},
+        # Channels
+        "channels/{channel_id}": {"limit": 5, "window": 15.0},
+        "guilds/{guild_id}/channels": {"limit": 2, "window": 10.0},
+        # Users/Members
+        "users/{user_id}": {"limit": 5, "window": 15.0},
+        "guilds/{guild_id}/members/{user_id}": {"limit": 5, "window": 15.0},
+        "guilds/{guild_id}/members/{user_id}/roles/{role_id}": {"limit": 10, "window": 10.0},
+        # Webhooks
+        "webhooks/{webhook_id}": {"limit": 5, "window": 15.0},
+        "webhooks/{webhook_id}/{webhook_token}": {"limit": 30, "window": 60.0},
+    }
+
+    def __init__(self, safety_margin: float = 0.1):
+        """Initialize rate limiter.
+
+        Args:
+            safety_margin: Extra wait time as fraction (0.1 = 10% extra)
+        """
+        self.safety_margin = safety_margin
+        self._global_bucket = RateLimitBucket(self.GLOBAL_LIMIT, self.GLOBAL_WINDOW)
+        self._endpoint_buckets: Dict[str, RateLimitBucket] = {}
+        self._bucket_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+
+        # Statistics
+        self._stats = {"requests": 0, "rate_limited": 0, "total_wait_time": 0.0, "max_wait_time": 0.0}
+
+    async def acquire(self, endpoint: str, method: str = "GET") -> None:
+        """Wait if necessary before making an API call.
+
+        Args:
+            endpoint: API endpoint path
+            method: HTTP method
+        """
+        self._stats["requests"] += 1
+
+        # Normalize endpoint for bucket lookup
+        bucket_key = self._normalize_endpoint(endpoint)
+
+        # Check global rate limit
+        global_wait = await self._global_bucket.acquire()
+        if global_wait > 0:
+            wait_time = global_wait * (1 + self.safety_margin)
+            await self._wait_and_log(wait_time, "global", endpoint)
+
+        # Check endpoint-specific rate limit
+        if bucket_key in self.ENDPOINT_LIMITS:
+            bucket = await self._get_or_create_bucket(bucket_key)
+            endpoint_wait = await bucket.acquire()
+            if endpoint_wait > 0:
+                wait_time = endpoint_wait * (1 + self.safety_margin)
+                await self._wait_and_log(wait_time, bucket_key, endpoint)
+
+    def update_from_response(self, endpoint: str, headers: Dict[str, str]) -> None:
+        """Update rate limits from Discord response headers.
+
+        Args:
+            endpoint: API endpoint that was called
+            headers: Response headers
+        """
+        # Check for rate limit headers
+        remaining_str = headers.get("X-RateLimit-Remaining")
+        reset_at_str = headers.get("X-RateLimit-Reset")
+        headers.get("X-RateLimit-Bucket")
+
+        if remaining_str is not None and reset_at_str is not None:
+            try:
+                remaining = int(remaining_str)
+                reset_at = float(reset_at_str)
+
+                # Update the appropriate bucket
+                bucket_key = self._normalize_endpoint(endpoint)
+                if bucket_key in self._endpoint_buckets:
+                    self._endpoint_buckets[bucket_key].update_from_headers(remaining, reset_at)
+
+            except (ValueError, TypeError) as e:
+                logger.debug(f"Failed to parse rate limit headers: {e}")
+
+    async def handle_rate_limit_response(self, endpoint: str, retry_after: float) -> None:
+        """Handle a 429 rate limit response.
+
+        Args:
+            endpoint: Endpoint that was rate limited
+            retry_after: Seconds to wait before retry
+        """
+        self._stats["rate_limited"] += 1
+        wait_time = retry_after * (1 + self.safety_margin)
+
+        logger.warning(f"Rate limited on {endpoint}, waiting {wait_time:.2f}s")
+        await asyncio.sleep(wait_time)
+
+        self._stats["total_wait_time"] += wait_time
+        self._stats["max_wait_time"] = max(self._stats["max_wait_time"], wait_time)
+
+    def _normalize_endpoint(self, endpoint: str) -> str:
+        """Normalize endpoint path to match bucket patterns.
+
+        Args:
+            endpoint: Raw endpoint path
+
+        Returns:
+            Normalized endpoint for bucket lookup
+        """
+        # Remove leading slash
+        if endpoint.startswith("/"):
+            endpoint = endpoint[1:]
+
+        # Replace IDs with placeholders
+        import re
+
+        # Channel endpoints
+        endpoint = re.sub(r"channels/\d+", "channels/{channel_id}", endpoint)
+        endpoint = re.sub(r"messages/\d+", "messages/{message_id}", endpoint)
+
+        # Guild endpoints
+        endpoint = re.sub(r"guilds/\d+", "guilds/{guild_id}", endpoint)
+        endpoint = re.sub(r"members/\d+", "members/{user_id}", endpoint)
+        endpoint = re.sub(r"roles/\d+", "roles/{role_id}", endpoint)
+
+        # User endpoints
+        endpoint = re.sub(r"users/\d+", "users/{user_id}", endpoint)
+
+        # Webhook endpoints
+        endpoint = re.sub(r"webhooks/\d+", "webhooks/{webhook_id}", endpoint)
+        endpoint = re.sub(r"webhooks/\d+/[\w-]+", "webhooks/{webhook_id}/{webhook_token}", endpoint)
+
+        return endpoint
+
+    async def _get_or_create_bucket(self, bucket_key: str) -> RateLimitBucket:
+        """Get or create a rate limit bucket.
+
+        Args:
+            bucket_key: Normalized endpoint key
+
+        Returns:
+            Rate limit bucket
+        """
+        async with self._bucket_locks[bucket_key]:
+            if bucket_key not in self._endpoint_buckets:
+                limits = self.ENDPOINT_LIMITS.get(bucket_key, {"limit": 5, "window": 60.0})
+                self._endpoint_buckets[bucket_key] = RateLimitBucket(int(limits["limit"]), limits["window"])
+            return self._endpoint_buckets[bucket_key]
+
+    async def _wait_and_log(self, wait_time: float, bucket_type: str, endpoint: str) -> None:
+        """Wait and log rate limit delay.
+
+        Args:
+            wait_time: Time to wait in seconds
+            bucket_type: Type of bucket (global or endpoint)
+            endpoint: Endpoint being accessed
+        """
+        logger.info(f"Rate limit ({bucket_type}) - waiting {wait_time:.2f}s for {endpoint}")
+        await asyncio.sleep(wait_time)
+
+        self._stats["total_wait_time"] += wait_time
+        self._stats["max_wait_time"] = max(self._stats["max_wait_time"], wait_time)
+
+    def get_stats(self) -> JSONDict:
+        """Get rate limiter statistics.
+
+        Returns:
+            Statistics dictionary
+        """
+        return {
+            **self._stats,
+            "average_wait_time": (
+                self._stats["total_wait_time"] / self._stats["requests"] if self._stats["requests"] > 0 else 0.0
+            ),
+            "buckets_tracked": len(self._endpoint_buckets),
+        }
+
+    def reset_stats(self) -> None:
+        """Reset statistics."""
+        self._stats = {"requests": 0, "rate_limited": 0, "total_wait_time": 0.0, "max_wait_time": 0.0}

ciris_engine/logic/adapters/discord/discord_reaction_handler.py

@@ -0,0 +1,278 @@
+"""Discord reaction handling component for approval workflows."""
+
+import asyncio
+import logging
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional
+
+import discord
+
+if TYPE_CHECKING:
+    from ciris_engine.protocols.services.lifecycle.time import TimeServiceProtocol
+
+logger = logging.getLogger(__name__)
+
+
+class ApprovalStatus(Enum):
+    """Status of an approval request."""
+
+    PENDING = "pending"
+    APPROVED = "approved"
+    DENIED = "denied"
+    TIMEOUT = "timeout"
+
+
+class ApprovalRequest:
+    """Represents a pending approval request."""
+
+    def __init__(
+        self, message_id: int, channel_id: int, request_type: str, context: Dict[str, str], timeout_seconds: int = 300
+    ):
+        self.message_id = message_id
+        self.channel_id = channel_id
+        self.request_type = request_type
+        self.context = context
+        self.timeout_seconds = timeout_seconds
+        self.status = ApprovalStatus.PENDING
+        self.created_at = datetime.now()
+        self.resolved_at: Optional[datetime] = None
+        self.resolver_id: Optional[str] = None
+        self.resolver_name: Optional[str] = None
+
+
+class DiscordReactionHandler:
+    """Handles Discord reactions for approval workflows."""
+
+    APPROVE_EMOJI = "✅"
+    DENY_EMOJI = "❌"
+
+    def __init__(
+        self, client: Optional[discord.Client] = None, time_service: Optional["TimeServiceProtocol"] = None
+    ) -> None:
+        """Initialize the reaction handler.
+
+        Args:
+            client: Discord client instance
+            time_service: Time service for consistent time operations
+        """
+        self.client = client
+        self._pending_approvals: Dict[int, ApprovalRequest] = {}
+        self._approval_callbacks: Dict[int, Callable[[ApprovalRequest], Awaitable[None]]] = {}
+        self._timeout_task: Optional[asyncio.Task[None]] = None
+        self._time_service: TimeServiceProtocol
+
+        # Ensure we have a time service
+        if time_service is None:
+            from ciris_engine.logic.services.lifecycle.time import TimeService
+
+            self._time_service = TimeService()
+        else:
+            self._time_service = time_service
+
+    def set_client(self, client: discord.Client) -> None:
+        """Set the Discord client after initialization.
+
+        Args:
+            client: Discord client instance
+        """
+        self.client = client
+
+    async def request_approval(
+        self,
+        channel_id: int,
+        message: str,
+        request_type: str,
+        context: Dict[str, str],
+        timeout_seconds: int = 300,
+        callback: Optional[Callable[[ApprovalRequest], Awaitable[None]]] = None,
+    ) -> Optional[ApprovalRequest]:
+        """Send an approval request message and wait for reactions.
+
+        Args:
+            channel_id: Discord channel ID
+            message: Message content to send
+            request_type: Type of approval request
+            context: Additional context for the request
+            timeout_seconds: Timeout in seconds
+            callback: Optional callback when approval is resolved
+
+        Returns:
+            ApprovalRequest object or None if failed
+        """
+        if not self.client:
+            logger.error("Discord client not initialized")
+            return None
+
+        try:
+            # Get channel
+            channel = self.client.get_channel(channel_id)
+            if not channel:
+                channel = await self.client.fetch_channel(channel_id)
+
+            # Check if channel supports sending messages
+            if not isinstance(
+                channel,
+                (discord.TextChannel, discord.DMChannel, discord.Thread, discord.VoiceChannel, discord.StageChannel),
+            ):
+                logger.error(f"Channel {channel_id} does not support sending messages")
+                return None
+
+            # Send message
+            sent_message = await channel.send(message)
+
+            # Add reactions
+            await sent_message.add_reaction(self.APPROVE_EMOJI)
+            await sent_message.add_reaction(self.DENY_EMOJI)
+
+            # Create approval request
+            approval = ApprovalRequest(
+                message_id=sent_message.id,
+                channel_id=channel_id,
+                request_type=request_type,
+                context=context,
+                timeout_seconds=timeout_seconds,
+            )
+
+            # Store in pending
+            self._pending_approvals[sent_message.id] = approval
+            if callback:
+                self._approval_callbacks[sent_message.id] = callback
+
+            # Schedule timeout
+            self._timeout_task = asyncio.create_task(self._handle_timeout(approval))
+
+            return approval
+
+        except Exception as e:
+            logger.exception(f"Failed to create approval request: {e}")
+            return None
+
+    async def handle_reaction(self, payload: discord.RawReactionActionEvent) -> None:
+        """Handle a reaction event.
+
+        Args:
+            payload: Discord reaction event payload
+        """
+        # Check if this is a reaction we care about
+        if payload.message_id not in self._pending_approvals:
+            return
+
+        # Check if it's an approval/denial emoji
+        if str(payload.emoji) not in [self.APPROVE_EMOJI, self.DENY_EMOJI]:
+            return
+
+        # Get the approval request
+        approval = self._pending_approvals[payload.message_id]
+
+        # Check if already resolved
+        if approval.status != ApprovalStatus.PENDING:
+            return
+
+        # Resolve the approval
+        if str(payload.emoji) == self.APPROVE_EMOJI:
+            approval.status = ApprovalStatus.APPROVED
+        else:
+            approval.status = ApprovalStatus.DENIED
+
+        approval.resolved_at = self._time_service.now()
+        approval.resolver_id = str(payload.user_id)
+
+        # Get resolver name if possible
+        if self.client:
+            try:
+                user = self.client.get_user(payload.user_id)
+                if not user:
+                    user = await self.client.fetch_user(payload.user_id)
+                approval.resolver_name = user.name
+            except Exception as e:
+                logger.warning(
+                    f"Failed to fetch Discord user {payload.user_id} for approval resolution: {e}. Resolver name will be omitted."
+                )
+
+        # Remove from pending
+        del self._pending_approvals[payload.message_id]
+
+        # Call callback if registered
+        if payload.message_id in self._approval_callbacks:
+            callback = self._approval_callbacks.pop(payload.message_id)
+            await callback(approval)
+
+        # Update message to show resolution
+        await self._update_approval_message(approval)
+
+    async def _handle_timeout(self, approval: ApprovalRequest) -> None:
+        """Handle timeout for an approval request.
+
+        Args:
+            approval: The approval request
+        """
+        await asyncio.sleep(approval.timeout_seconds)
+
+        # Check if still pending
+        if approval.message_id in self._pending_approvals:
+            approval.status = ApprovalStatus.TIMEOUT
+            approval.resolved_at = self._time_service.now()
+
+            # Remove from pending
+            del self._pending_approvals[approval.message_id]
+
+            # Call callback if registered
+            if approval.message_id in self._approval_callbacks:
+                callback = self._approval_callbacks.pop(approval.message_id)
+                await callback(approval)
+
+            # Update message
+            await self._update_approval_message(approval)
+
+    async def _update_approval_message(self, approval: ApprovalRequest) -> None:
+        """Update the approval message to show resolution.
+
+        Args:
+            approval: The resolved approval request
+        """
+        if not self.client:
+            return
+
+        try:
+            # Get channel and message
+            channel = self.client.get_channel(approval.channel_id)
+            if not channel:
+                channel = await self.client.fetch_channel(approval.channel_id)
+
+            # Check if channel supports fetching messages
+            if not isinstance(
+                channel,
+                (discord.TextChannel, discord.DMChannel, discord.Thread, discord.VoiceChannel, discord.StageChannel),
+            ):
+                logger.error(f"Channel {approval.channel_id} does not support fetching messages")
+                return
+
+            message = await channel.fetch_message(approval.message_id)
+
+            # Build status line
+            status_line = f"\n\n**RESOLVED**: {approval.status.value.upper()}"
+            if approval.resolver_name:
+                status_line += f" by {approval.resolver_name}"
+            if approval.resolved_at:
+                status_line += f" at {approval.resolved_at.isoformat()}"
+
+            # Update message
+            new_content = message.content + status_line
+            await message.edit(content=new_content)
+
+        except Exception as e:
+            logger.error(f"Failed to update approval message: {e}")
+
+    def get_pending_approvals(self) -> Dict[int, ApprovalRequest]:
+        """Get all pending approval requests.
+
+        Returns:
+            Dictionary of message ID to ApprovalRequest
+        """
+        return self._pending_approvals.copy()
+
+    def clear_pending_approvals(self) -> None:
+        """Clear all pending approvals (for cleanup)."""
+        self._pending_approvals.clear()
+        self._approval_callbacks.clear()