exaai-agent 2.0.5__py3-none-any.whl → 2.0.8__py3-none-any.whl

exaaiagnt/llm/llm_traffic_controller.py

@@ -0,0 +1,351 @@
+"""
+Adaptive LLM Traffic Controller - Intelligent rate limiting and queue management.
+
+Features:
+- Single LLM request at a time (serialized)
+- Non-blocking queue for waiting agents
+- Automatic rate limit detection and delay
+- Tool-first execution mode
+- Automatic recovery with smart retry
+"""
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Optional
+from collections import deque
+
+
+logger = logging.getLogger(__name__)
+
+
+class RequestPriority(Enum):
+    """Priority levels for LLM requests."""
+    CRITICAL = 3  # Must execute ASAP (auth, security critical)
+    NORMAL = 2    # Standard agent requests
+    LOW = 1       # Background tasks, summaries
+
+
+@dataclass
+class QueuedRequest:
+    """A queued LLM request with metadata."""
+    request_id: str
+    agent_id: str
+    priority: RequestPriority
+    request_func: Callable
+    args: tuple
+    kwargs: dict
+    created_at: float = field(default_factory=time.time)
+    future: asyncio.Future = field(default_factory=lambda: asyncio.get_event_loop().create_future())
+
+
+class AdaptiveLLMController:
+    """
+    Adaptive Multi-Agent LLM Traffic Controller.
+    
+    Implements:
+    1. Single concurrent LLM request (serialized calls)
+    2. Non-blocking queue for agents
+    3. Intelligent throttling with adaptive delays
+    4. Tool-first execution mode
+    5. Automatic recovery from rate limits
+    """
+    
+    _instance: Optional["AdaptiveLLMController"] = None
+    
+    def __new__(cls) -> "AdaptiveLLMController":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+    
+    def __init__(self):
+        if self._initialized:
+            return
+        
+        # Core state
+        self._queue: deque[QueuedRequest] = deque()
+        self._is_processing = False
+        self._lock = asyncio.Lock()
+        
+        # Rate limiting state
+        self._last_request_time = 0.0
+        self._consecutive_rate_limits = 0
+        self._base_delay = 2.0  # Base delay between requests
+        self._current_delay = 2.0
+        self._max_delay = 30.0
+        
+        # Statistics
+        self._total_requests = 0
+        self._successful_requests = 0
+        self._rate_limit_hits = 0
+        self._retries = 0
+        
+        # Tool execution mode
+        self._tool_executing = False
+        self._tool_execution_lock = asyncio.Lock()
+        
+        # Configuration
+        self._max_retries = 3
+        self._rate_limit_wait = 6.0  # Wait time after rate limit
+        self._enable_verbose_logging = False
+        
+        self._initialized = True
+        logger.info("AdaptiveLLMController initialized - Traffic Control Enabled")
+    
+    async def queue_request(
+        self,
+        request_func: Callable,
+        *args,
+        agent_id: str = "unknown",
+        priority: RequestPriority = RequestPriority.NORMAL,
+        **kwargs
+    ) -> Any:
+        """
+        Queue an LLM request and wait for result.
+        
+        Non-blocking for the caller - they just await the result.
+        Internally, requests are processed one at a time.
+        """
+        request_id = f"{agent_id}_{int(time.time() * 1000)}"
+        
+        # Create future for this request
+        loop = asyncio.get_event_loop()
+        future = loop.create_future()
+        
+        queued = QueuedRequest(
+            request_id=request_id,
+            agent_id=agent_id,
+            priority=priority,
+            request_func=request_func,
+            args=args,
+            kwargs=kwargs,
+            future=future
+        )
+        
+        # Add to queue
+        async with self._lock:
+            # Insert by priority
+            if priority == RequestPriority.CRITICAL:
+                self._queue.appendleft(queued)
+            else:
+                self._queue.append(queued)
+            
+            queue_size = len(self._queue)
+            if queue_size > 1 and not self._enable_verbose_logging:
+                logger.debug(f"Request queued: {request_id}, queue size: {queue_size}")
+        
+        # Start processing if not already running
+        asyncio.create_task(self._process_queue())
+        
+        # Wait for result
+        return await future
+    
+    async def _process_queue(self):
+        """Process queued requests one at a time."""
+        async with self._lock:
+            if self._is_processing:
+                return
+            self._is_processing = True
+        
+        try:
+            while True:
+                # Get next request
+                async with self._lock:
+                    if not self._queue:
+                        break
+                    request = self._queue.popleft()
+                
+                # Wait if tool is executing
+                if self._tool_executing:
+                    async with self._tool_execution_lock:
+                        pass  # Wait for tool to finish
+                
+                # Execute request with rate limiting
+                await self._execute_request(request)
+                
+        finally:
+            async with self._lock:
+                self._is_processing = False
+    
+    async def _execute_request(self, request: QueuedRequest):
+        """Execute a single request with rate limiting and retry."""
+        self._total_requests += 1
+        
+        # Adaptive delay
+        await self._apply_rate_limit_delay()
+        
+        # Try request with retry
+        last_error = None
+        for attempt in range(self._max_retries + 1):
+            try:
+                # Execute the request
+                result = await self._call_with_timeout(
+                    request.request_func,
+                    *request.args,
+                    **request.kwargs
+                )
+                
+                # Success!
+                self._successful_requests += 1
+                self._consecutive_rate_limits = 0
+                self._current_delay = self._base_delay  # Reset delay
+                
+                request.future.set_result(result)
+                return
+                
+            except Exception as e:
+                last_error = e
+                error_str = str(e).lower()
+                
+                # Check for rate limit
+                if "rate" in error_str and "limit" in error_str:
+                    self._rate_limit_hits += 1
+                    self._consecutive_rate_limits += 1
+                    self._retries += 1
+                    
+                    # Increase delay exponentially
+                    self._current_delay = min(
+                        self._current_delay * 1.5,
+                        self._max_delay
+                    )
+                    
+                    if attempt < self._max_retries:
+                        wait_time = self._rate_limit_wait * (attempt + 1)
+                        logger.warning(
+                            f"Rate limit hit (attempt {attempt + 1}/{self._max_retries + 1}), "
+                            f"waiting {wait_time}s before retry"
+                        )
+                        await asyncio.sleep(wait_time)
+                        continue
+                
+                # Other errors - log and continue to next attempt
+                if attempt < self._max_retries:
+                    logger.warning(f"Request failed (attempt {attempt + 1}): {e}")
+                    self._retries += 1
+                    await asyncio.sleep(2.0 * (attempt + 1))
+                    continue
+        
+        # All retries failed
+        request.future.set_exception(last_error or Exception("Request failed"))
+    
+    async def _call_with_timeout(self, func: Callable, *args, **kwargs) -> Any:
+        """Call function with timeout."""
+        timeout = kwargs.pop('timeout', 300)
+        
+        if asyncio.iscoroutinefunction(func):
+            return await asyncio.wait_for(func(*args, **kwargs), timeout=timeout)
+        else:
+            # Run sync function in thread
+            loop = asyncio.get_event_loop()
+            return await asyncio.wait_for(
+                loop.run_in_executor(None, lambda: func(*args, **kwargs)),
+                timeout=timeout
+            )
+    
+    async def _apply_rate_limit_delay(self):
+        """Apply intelligent rate limiting delay."""
+        now = time.time()
+        time_since_last = now - self._last_request_time
+        
+        # Calculate required delay
+        required_delay = self._current_delay - time_since_last
+        
+        if required_delay > 0:
+            # Add jitter to prevent thundering herd
+            jitter = required_delay * 0.1 * (0.5 - asyncio.get_event_loop().time() % 1)
+            total_delay = max(0, required_delay + jitter)
+            
+            if total_delay > 0.5 and not self._enable_verbose_logging:
+                logger.debug(f"Rate limiting: waiting {total_delay:.2f}s")
+            
+            await asyncio.sleep(total_delay)
+        
+        self._last_request_time = time.time()
+    
+    # Tool execution mode
+    async def enter_tool_mode(self):
+        """Enter tool-first execution mode - pause LLM calls."""
+        await self._tool_execution_lock.acquire()
+        self._tool_executing = True
+    
+    async def exit_tool_mode(self):
+        """Exit tool mode - resume LLM calls."""
+        self._tool_executing = False
+        if self._tool_execution_lock.locked():
+            self._tool_execution_lock.release()
+    
+    def is_tool_executing(self) -> bool:
+        """Check if tool is currently executing."""
+        return self._tool_executing
+    
+    # Statistics
+    def get_stats(self) -> dict[str, Any]:
+        """Get traffic controller statistics."""
+        return {
+            "total_requests": self._total_requests,
+            "successful_requests": self._successful_requests,
+            "rate_limit_hits": self._rate_limit_hits,
+            "retries": self._retries,
+            "queue_size": len(self._queue),
+            "current_delay": self._current_delay,
+            "consecutive_rate_limits": self._consecutive_rate_limits,
+            "is_processing": self._is_processing,
+            "tool_executing": self._tool_executing,
+        }
+    
+    def reset_stats(self):
+        """Reset statistics."""
+        self._total_requests = 0
+        self._successful_requests = 0
+        self._rate_limit_hits = 0
+        self._retries = 0
+    
+    def set_base_delay(self, delay: float):
+        """Set base delay between requests."""
+        self._base_delay = delay
+        self._current_delay = delay
+    
+    def set_verbose(self, enabled: bool):
+        """Enable/disable verbose logging."""
+        self._enable_verbose_logging = enabled
+
+
+# Global instance
+_controller: Optional[AdaptiveLLMController] = None
+
+
+def get_traffic_controller() -> AdaptiveLLMController:
+    """Get or create the global traffic controller."""
+    global _controller
+    if _controller is None:
+        _controller = AdaptiveLLMController()
+    return _controller
+
+
+def reset_traffic_controller():
+    """Reset the global traffic controller."""
+    global _controller
+    _controller = None
+
+
+# Convenience decorator for LLM requests
+def with_traffic_control(priority: RequestPriority = RequestPriority.NORMAL):
+    """Decorator to route LLM requests through traffic controller."""
+    def decorator(func: Callable):
+        async def wrapper(*args, agent_id: str = "unknown", **kwargs):
+            controller = get_traffic_controller()
+            return await controller.queue_request(
+                func, *args,
+                agent_id=agent_id,
+                priority=priority,
+                **kwargs
+            )
+        return wrapper
+    return decorator
+
+
+# Export confirmation
+logger.info("Adaptive LLM Traffic Control Module Loaded")

exaaiagnt/prompts/auto_loader.py

@@ -172,6 +172,44 @@ MODULE_PATTERNS = {
         "keywords": ["recon", "reconnaissance", "enumerate", "discover", "fingerprint", "osint"],
         "domain_only": True,
     },
+    
+    # React2Shell - RSC Deserialization RCE (CVE-2025-55182)
+    "react2shell": {
+        "url_patterns": [
+            r"/_next/",
+            r"/_next/static",
+            r"/_actions",
+            r"/api/",
+        ],
+        "keywords": ["react", "next.js", "nextjs", "vercel", "rsc", "server components", "app router"],
+        "header_patterns": ["x-nextjs", "x-vercel", "x-powered-by.*next"],
+    },
+    
+    # Modern JS Frameworks Security
+    "modern_js_frameworks": {
+        "url_patterns": [
+            r"/_next/",
+            r"/_nuxt/",
+            r"/_svelte",
+            r"/__remix",
+            r"/_astro/",
+        ],
+        "keywords": ["next.js", "nuxt", "sveltekit", "remix", "astro", "react", "vue", "svelte"],
+    },
+    
+    # AWS/Cloud Security
+    "aws_cloud_security": {
+        "url_patterns": [
+            r"\.amazonaws\.com",
+            r"\.s3\.",
+            r"\.azure\.",
+            r"\.blob\.core",
+            r"\.cloudfront\.",
+            r"\.appspot\.com",
+            r"\.storage\.googleapis",
+        ],
+        "keywords": ["aws", "s3", "ec2", "lambda", "azure", "gcp", "cloud", "bucket", "metadata"],
+    },
 }
 
 
@@ -179,6 +217,10 @@ def detect_modules_from_target(target: str, instruction: str = "") -> List[str]:
     """
     Automatically detect which prompt modules should be loaded based on target URL and instruction.
     
+    This function:
+    1. Checks defined patterns in MODULE_PATTERNS
+    2. Auto-discovers ALL .jinja files and matches by filename keywords
+    
     Args:
         target: The target URL or domain
         instruction: The user's instruction/task description
@@ -197,6 +239,7 @@ def detect_modules_from_target(target: str, instruction: str = "") -> List[str]:
     parsed = urlparse(target if "://" in target else f"https://{target}")
     url_path = parsed.path.lower()
     
+    # 1. Check defined patterns (MODULE_PATTERNS)
     for module_name, patterns in MODULE_PATTERNS.items():
         should_load = False
         
@@ -220,6 +263,9 @@ def detect_modules_from_target(target: str, instruction: str = "") -> List[str]:
         if should_load:
             detected_modules.add(module_name)
     
+    # 2. AUTO-DISCOVER: Scan all .jinja files and match by filename
+    detected_modules.update(_auto_discover_modules(combined_text))
+    
     # Always include base modules for comprehensive scans
     if any(kw in instruction_lower for kw in ["full", "comprehensive", "thorough", "complete"]):
         detected_modules.update(["sql_injection", "xss", "authentication_jwt"])
@@ -227,6 +273,64 @@ def detect_modules_from_target(target: str, instruction: str = "") -> List[str]:
     return list(detected_modules)
 
 
+def _auto_discover_modules(search_text: str) -> Set[str]:
+    """
+    Auto-discover modules by scanning all .jinja files and matching by filename.
+    
+    Any new .jinja file will be automatically discovered and used if its name
+    or keywords from name appear in the target/instruction.
+    
+    Example: 'react2shell.jinja' will be loaded if 'react' or 'shell' appears in target.
+    """
+    from pathlib import Path
+    
+    discovered = set()
+    prompts_dir = Path(__file__).parent
+    
+    # Get all available modules
+    for category_dir in prompts_dir.iterdir():
+        if not category_dir.is_dir() or category_dir.name.startswith("__"):
+            continue
+        
+        for jinja_file in category_dir.glob("*.jinja"):
+            module_name = jinja_file.stem  # e.g., "react2shell"
+            
+            # Skip if already in MODULE_PATTERNS (already handled)
+            if module_name in MODULE_PATTERNS:
+                continue
+            
+            # Generate keywords from filename
+            # "react2shell" -> ["react", "shell", "react2shell"]
+            # "sql_injection" -> ["sql", "injection"]
+            # "modern_js_frameworks" -> ["modern", "js", "frameworks"]
+            keywords = _extract_keywords_from_name(module_name)
+            
+            # Check if any keyword matches
+            for keyword in keywords:
+                if len(keyword) >= 3 and keyword in search_text:  # Min 3 chars to avoid false positives
+                    discovered.add(module_name)
+                    break
+    
+    return discovered
+
+
+def _extract_keywords_from_name(name: str) -> List[str]:
+    """Extract searchable keywords from a module name."""
+    keywords = [name]  # Full name
+    
+    # Split by underscore
+    parts = name.split("_")
+    keywords.extend(parts)
+    
+    # Split by numbers (react2shell -> react, shell)
+    import re
+    alpha_parts = re.split(r'\d+', name)
+    keywords.extend([p for p in alpha_parts if p])
+    
+    # Lowercase all
+    return [k.lower() for k in keywords if len(k) >= 3]
+
+
 def get_recommended_modules(target: str, instruction: str = "") -> dict:
     """
     Get recommended modules with confidence scores.