polyrob 0.5.0__py3-none-any.whl

agents/task/runtime_safety.py

@@ -0,0 +1,308 @@
+"""
+Runtime safety checks and validation for Task.
+
+This module provides critical runtime safety checks to prevent the issues
+identified in AUTOV2_RUNTIME_FIXES.md
+"""
+
+import asyncio
+import logging
+from typing import Optional, Dict, Any, List, Callable
+from functools import wraps
+import time
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeSafety:
+    """Runtime safety checks and guards for Task."""
+
+    @staticmethod
+    def add_timeout_guard(timeout_seconds: float = 60.0):
+        """Add timeout protection to async functions.
+
+        Args:
+            timeout_seconds: Timeout in seconds (default 60)
+
+        Usage:
+            @add_timeout_guard(30)
+            async def my_async_function():
+                # Will timeout after 30 seconds
+        """
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                try:
+                    return await asyncio.wait_for(
+                        func(*args, **kwargs),
+                        timeout=timeout_seconds
+                    )
+                except asyncio.TimeoutError:
+                    logger.warning(f"{func.__name__} timed out after {timeout_seconds}s")
+                    # Return sensible default based on function name
+                    if "approval" in func.__name__.lower():
+                        return {"approved": True, "reason": "timeout_auto_approved"}
+                    return None
+
+            return wrapper
+        return decorator
+
+    @staticmethod
+    def validate_tool_call_structure(tool_call: Dict[str, Any]) -> bool:
+        """Validate tool call has required structure.
+
+        Args:
+            tool_call: Tool call dict to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        required = {"name", "id"}
+        if not all(field in tool_call for field in required):
+            return False
+
+        # Validate types
+        if not isinstance(tool_call.get("name"), str):
+            return False
+        if not isinstance(tool_call.get("id"), str):
+            return False
+
+        # Validate args if present
+        if "args" in tool_call and not isinstance(tool_call["args"], dict):
+            return False
+
+        return True
+
+    @staticmethod
+    def protect_deque_operations(deque_obj, operation: str, *args, **kwargs):
+        """Safely perform operations on deque objects.
+
+        Deques don't support index-based removal, so we convert to list,
+        perform operation, and convert back.
+
+        Args:
+            deque_obj: The deque object
+            operation: Operation to perform ('remove', 'pop', etc.)
+            *args, **kwargs: Arguments for the operation
+
+        Returns:
+            Result of the operation
+        """
+        from collections import deque
+
+        if operation == 'remove_by_index':
+            # Convert to list, remove by index, convert back
+            idx = args[0] if args else -1
+            items_list = list(deque_obj)
+            if 0 <= idx < len(items_list):
+                removed = items_list.pop(idx)
+                # Preserve maxlen if it exists
+                maxlen = getattr(deque_obj, 'maxlen', None)
+                new_deque = deque(items_list, maxlen=maxlen)
+                return new_deque, removed
+            return deque_obj, None
+
+        # For other operations, use standard methods
+        return getattr(deque_obj, operation)(*args, **kwargs)
+
+    @staticmethod
+    def validate_session_id(session_id: str) -> bool:
+        """Validate session ID format.
+
+        Args:
+            session_id: Session ID to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        if not session_id or not isinstance(session_id, str):
+            return False
+
+        # Session ID should not be empty after stripping
+        if not session_id.strip():
+            return False
+
+        # Should not contain dangerous characters
+        dangerous_chars = ['/', '\\', '..', '\x00', '\n', '\r']
+        if any(char in session_id for char in dangerous_chars):
+            return False
+
+        return True
+
+    @staticmethod
+    def safe_json_parse(json_str: str, max_retries: int = 3) -> Optional[Dict]:
+        """Safely parse JSON with retry logic.
+
+        Args:
+            json_str: JSON string to parse
+            max_retries: Maximum parse attempts
+
+        Returns:
+            Parsed dict or None if failed
+        """
+        import json
+        import re
+
+        for attempt in range(max_retries):
+            try:
+                # Clean common issues
+                cleaned = json_str.strip()
+
+                # Remove code fences
+                if cleaned.startswith('```'):
+                    cleaned = re.sub(r'^```[a-z]*\n', '', cleaned)
+                    cleaned = re.sub(r'\n```$', '', cleaned)
+
+                # Remove think tags
+                cleaned = re.sub(r'<think>.*?</think>', '', cleaned, flags=re.DOTALL)
+
+                # Try to parse
+                return json.loads(cleaned)
+
+            except json.JSONDecodeError as e:
+                if attempt == max_retries - 1:
+                    logger.error(f"JSON parse failed after {max_retries} attempts: {e}")
+                    return None
+
+                # Try to fix common issues
+                if "Expecting property name" in str(e):
+                    # Fix trailing commas
+                    cleaned = re.sub(r',\s*}', '}', cleaned)
+                    cleaned = re.sub(r',\s*]', ']', cleaned)
+
+                time.sleep(0.1 * (attempt + 1))  # Exponential backoff
+
+        return None
+
+    @staticmethod
+    def memory_guard(max_size_mb: float = 100):
+        """Decorator to guard against memory explosion.
+
+        Args:
+            max_size_mb: Maximum allowed memory growth in MB
+
+        Usage:
+            @memory_guard(50)
+            def process_screenshots():
+                # Will log warning if memory grows by >50MB
+        """
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                import psutil
+                import os
+
+                process = psutil.Process(os.getpid())
+                mem_before = process.memory_info().rss / 1024 / 1024  # MB
+
+                try:
+                    result = func(*args, **kwargs)
+                finally:
+                    mem_after = process.memory_info().rss / 1024 / 1024  # MB
+                    mem_growth = mem_after - mem_before
+
+                    if mem_growth > max_size_mb:
+                        logger.warning(
+                            f"{func.__name__} memory growth: {mem_growth:.1f}MB "
+                            f"(exceeds {max_size_mb}MB limit)"
+                        )
+
+                return result
+
+            # Handle async functions
+            if asyncio.iscoroutinefunction(func):
+                @wraps(func)
+                async def async_wrapper(*args, **kwargs):
+                    import psutil
+                    import os
+
+                    process = psutil.Process(os.getpid())
+                    mem_before = process.memory_info().rss / 1024 / 1024
+
+                    try:
+                        result = await func(*args, **kwargs)
+                    finally:
+                        mem_after = process.memory_info().rss / 1024 / 1024
+                        mem_growth = mem_after - mem_before
+
+                        if mem_growth > max_size_mb:
+                            logger.warning(
+                                f"{func.__name__} memory growth: {mem_growth:.1f}MB "
+                                f"(exceeds {max_size_mb}MB limit)"
+                            )
+
+                    return result
+
+                return async_wrapper
+
+            return wrapper
+        return decorator
+
+    @staticmethod
+    def validate_token_limits(
+        current_tokens: int,
+        max_tokens: int,
+        safety_buffer: int = 1000
+    ) -> bool:
+        """Validate token counts are within safe limits.
+
+        Args:
+            current_tokens: Current token count
+            max_tokens: Maximum allowed tokens
+            safety_buffer: Safety buffer to maintain
+
+        Returns:
+            True if within safe limits, False otherwise
+        """
+        if current_tokens <= 0:
+            logger.warning("Invalid token count: <= 0")
+            return False
+
+        if current_tokens + safety_buffer >= max_tokens:
+            logger.warning(
+                f"Token limit approaching: {current_tokens} + {safety_buffer} >= {max_tokens}"
+            )
+            return False
+
+        return True
+
+
+# Export convenience decorators
+timeout_guard = RuntimeSafety.add_timeout_guard
+memory_guard = RuntimeSafety.memory_guard
+
+
+def apply_runtime_fixes():
+    """Apply all critical runtime fixes identified in AUTOV2_RUNTIME_FIXES.md"""
+
+    logger.info("Applying Task runtime safety fixes...")
+
+    # 1. Patch approval methods to add timeouts
+    # NOTE: Approval system removed - request_approval now auto-approves
+    try:
+        from agents.task.agent import approval
+        if hasattr(approval, 'request_approval'):
+            # No longer needed - approval workflow simplified
+            logger.debug("Approval system simplified - timeout patching skipped")
+    except Exception as e:
+        logger.debug(f"Could not patch approval: {e}")
+
+    # 2. Add memory monitoring for screenshot operations
+    try:
+        from agents.task.agent import screenshots
+        if hasattr(screenshots, 'create_gif'):
+            original = screenshots.create_gif
+            screenshots.create_gif = memory_guard(100)(original)
+            logger.info("✓ Added memory guard to screenshot operations")
+    except Exception as e:
+        logger.debug(f"Could not patch screenshots: {e}")
+
+    logger.info("Runtime safety fixes applied")
+
+
+# Auto-apply fixes when module is imported
+if __name__ != "__main__":
+    try:
+        apply_runtime_fixes()
+    except Exception as e:
+        logger.warning(f"Could not apply all runtime fixes: {e}")

agents/task/session/__init__.py

@@ -0,0 +1,5 @@
+"""Session orchestrator concern mixins.
+
+Concern-group mixins split out of SessionOrchestrator (pure code-motion):
+workspace, feed, multi-agent relationship tracking, and browser-pool access.
+"""

agents/task/session/browser_pool.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+if TYPE_CHECKING:
+    from tools.browser.browser import Browser
+
+
+class BrowserPoolMixin:
+    """Browser-pool access, tools dict compatibility, and context release for SessionOrchestrator."""
+
+    @property
+    def tools(self) -> Dict[str, Any]:
+        """Backward compatibility: Return tools as a dict."""
+        if not self.controller:
+            return {}
+        result = {}
+        for tool_name in self.controller.list_tools():
+            tool = self.controller.get_tool(tool_name)
+            if tool:
+                result[tool_name] = tool
+        return result
+
+    @tools.setter
+    def tools(self, value: Dict[str, Any]) -> None:
+        """Backward compatibility: Set tools from a dict."""
+        if not self.controller:
+            self.logger.warning("Controller not initialized - cannot set tools")
+            return
+        # Clear existing tools
+        for tool_name in self.controller.list_tools():
+            self.controller.remove_tool(tool_name)
+        # Add all new tools
+        for name, tool in value.items():
+            self.controller.add_tool(name, tool)
+            self.logger.debug(f"Added tool '{name}' to controller")
+
+    async def get_browser_context(self, agent_id: str, reuse: bool = True, timeout: float = 30.0) -> Optional[Any]:
+        """Get browser context for agent from BrowserManager.
+
+        Args:
+            agent_id: Agent ID (format: agentname_sessionid)
+            reuse: Whether to reuse pooled contexts
+            timeout: Unused (kept for compatibility)
+
+        Returns:
+            Browser context or None
+        """
+        # No BrowserManager is an EXPECTED steady state now — the default CLI (and
+        # any non-browser session) provisions none, and the step loop asks for a
+        # context every step regardless of loaded tools. Log at debug, not error
+        # (OR-3): an error here is misleading noise on every step of a normal run.
+        if not self.browser_manager:
+            self.logger.debug(f"No BrowserManager - no browser context for {agent_id}")
+            return None
+
+        # Agent ID already contains session (format set in Agent.__init__)
+        # No need to manipulate - use as-is
+        context = await self.browser_manager.get_context(agent_id, reuse=reuse)
+
+        # Track context for cleanup (prevents leaks if agent removed before cleanup)
+        if context and hasattr(self, '_browser_contexts'):
+            self._browser_contexts.add(agent_id)
+            self.logger.debug(f"Tracking browser context for cleanup: {agent_id}")
+
+        # Track context acquisition with telemetry
+        if context and hasattr(self, 'telemetry_manager') and self.telemetry_manager:
+            try:
+                self.telemetry_manager.capture_event(
+                    event_type="browser_context_acquired",
+                    data={
+                        "agent_id": agent_id,
+                        "session_id": self.session_id,
+                        "context_id": id(context),
+                        "reused": reuse
+                    }
+                )
+            except Exception as e:
+                self.logger.debug(f"Failed to emit browser context telemetry: {e}")
+
+        return context
+
+    async def _ensure_browser_available(self) -> Optional['Browser']:
+        """Get browser from BrowserManager (single source of truth).
+
+        Returns:
+            Browser instance from BrowserManager, or None if unavailable
+        """
+        if not self.browser_manager:
+            self.logger.error("BrowserManager required but not available")
+            return None
+
+        # Initialize BrowserManager if needed
+        if not self.browser_manager.is_initialized:
+            try:
+                await self.browser_manager.initialize()
+            except Exception as e:
+                self.logger.error(f"Failed to initialize BrowserManager: {e}")
+                return None
+
+        # Get browser from BrowserManager (ONLY source)
+        try:
+            browser = await self.browser_manager.get_browser()
+            if browser:
+                self.logger.debug("Got browser from BrowserManager")
+                return browser
+            else:
+                self.logger.error("BrowserManager has no browser available")
+                return None
+        except Exception as e:
+            self.logger.error(f"Failed to get browser from BrowserManager: {e}")
+            return None
+
+    async def release_browser_context(self, context_id: str, close: bool = True):
+        """Release a browser context back to the pool or close it.
+
+        Args:
+            context_id: ID of the context to release
+            close: If True, close the context instead of returning to pool
+
+        This delegates to BrowserManager if available, otherwise is a no-op.
+        """
+        if self.browser_manager:
+            try:
+                await self.browser_manager.release_context(context_id, close=close)
+                self.logger.debug(f"Released browser context {context_id} (close={close})")
+            except Exception as e:
+                self.logger.warning(f"Failed to release browser context {context_id}: {e}")
+        else:
+            self.logger.debug(f"No browser_manager available to release context {context_id}")