hermes-agent/tools/resource_limits.py

"""
Terminal Sandbox Resource Limits — CPU, memory, time.

Provides resource limits for agent terminal commands to prevent
OOM kills, runaway processes, and excessive resource consumption.
"""

import logging
import os
import signal
import subprocess
import time
from dataclasses import dataclass, field
from typing import Optional, Dict, Any
from enum import Enum

logger = logging.getLogger(__name__)


class ResourceViolation(Enum):
    """Types of resource violations."""
    MEMORY = "memory"
    CPU = "cpu"
    TIME = "time"
    NONE = "none"


@dataclass
class ResourceLimits:
    """Resource limits for a subprocess."""
    memory_mb: int = 2048  # 2GB default
    cpu_percent: int = 80  # 80% of one core
    timeout_seconds: int = 300  # 5 minutes
    kill_timeout: int = 10  # SIGKILL after 10s if SIGTERM fails


@dataclass
class ResourceResult:
    """Result of a resource-limited execution."""
    success: bool
    stdout: str
    stderr: str
    exit_code: int
    violation: ResourceViolation
    violation_message: str
    memory_used_mb: float
    cpu_time_seconds: float
    wall_time_seconds: float
    killed: bool = False


class ResourceLimiter:
    """Apply resource limits to subprocess execution."""
    
    def __init__(self, limits: Optional[ResourceLimits] = None):
        self.limits = limits or ResourceLimits()
    
    def _get_resource_rlimit(self) -> Dict[str, Any]:
        """Get resource limits for subprocess (Unix only)."""
        import resource
        
        rlimit = {}
        
        # Memory limit (RSS)
        if self.limits.memory_mb > 0:
            mem_bytes = self.limits.memory_mb * 1024 * 1024
            rlimit[resource.RLIMIT_AS] = (mem_bytes, mem_bytes)
        
        # CPU time limit
        if self.limits.timeout_seconds > 0:
            rlimit[resource.RLIMIT_CPU] = (self.limits.timeout_seconds, self.limits.timeout_seconds)
        
        return rlimit
    
    def _check_resource_usage(self, process: subprocess.Popen) -> Dict[str, float]:
        """Check resource usage of a process (Unix only)."""
        try:
            import resource
            usage = resource.getrusage(resource.RUSAGE_CHILDREN)
            return {
                "user_time": usage.ru_utime,
                "system_time": usage.ru_stime,
                "max_rss_mb": usage.ru_maxrss / 1024,  # KB to MB
            }
        except:
            return {"user_time": 0, "system_time": 0, "max_rss_mb": 0}
    
    def execute(self, command: str, **kwargs) -> ResourceResult:
        """
        Execute a command with resource limits.
        
        Args:
            command: Shell command to execute
            **kwargs: Additional subprocess arguments
            
        Returns:
            ResourceResult with execution details
        """
        start_time = time.time()
        
        # Try to use resource limits (Unix only)
        preexec_fn = None
        try:
            import resource
            rlimit = self._get_resource_rlimit()
            
            def set_limits():
                for res, limits in rlimit.items():
                    resource.setrlimit(res, limits)
            
            preexec_fn = set_limits
        except ImportError:
            logger.debug("resource module not available, skipping limits")
        
        try:
            # Execute with timeout
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=self.limits.timeout_seconds,
                preexec_fn=preexec_fn,
                **kwargs,
            )
            
            wall_time = time.time() - start_time
            usage = self._check_resource_usage(result)
            
            # Check for violations
            violation = ResourceViolation.NONE
            violation_message = ""
            
            # Check memory (if we can get it)
            if usage["max_rss_mb"] > self.limits.memory_mb:
                violation = ResourceViolation.MEMORY
                violation_message = f"Memory limit exceeded: {usage['max_rss_mb']:.0f}MB > {self.limits.memory_mb}MB"
            
            return ResourceResult(
                success=result.returncode == 0,
                stdout=result.stdout,
                stderr=result.stderr,
                exit_code=result.returncode,
                violation=violation,
                violation_message=violation_message,
                memory_used_mb=usage["max_rss_mb"],
                cpu_time_seconds=usage["user_time"] + usage["system_time"],
                wall_time_seconds=wall_time,
            )
        
        except subprocess.TimeoutExpired as e:
            wall_time = time.time() - start_time
            
            # Try to kill gracefully
            if hasattr(e, 'process') and e.process:
                try:
                    e.process.terminate()
                    time.sleep(self.limits.kill_timeout)
                    if e.process.poll() is None:
                        e.process.kill()
                except:
                    pass
            
            return ResourceResult(
                success=False,
                stdout=e.stdout.decode() if e.stdout else "",
                stderr=e.stderr.decode() if e.stderr else "",
                exit_code=-1,
                violation=ResourceViolation.TIME,
                violation_message=f"Timeout after {self.limits.timeout_seconds}s",
                memory_used_mb=0,
                cpu_time_seconds=0,
                wall_time_seconds=wall_time,
                killed=True,
            )
        
        except MemoryError:
            wall_time = time.time() - start_time
            return ResourceResult(
                success=False,
                stdout="",
                stderr=f"Memory limit exceeded ({self.limits.memory_mb}MB)",
                exit_code=-1,
                violation=ResourceViolation.MEMORY,
                violation_message=f"Memory limit exceeded: {self.limits.memory_mb}MB",
                memory_used_mb=self.limits.memory_mb,
                cpu_time_seconds=0,
                wall_time_seconds=wall_time,
                killed=True,
            )
        
        except Exception as e:
            wall_time = time.time() - start_time
            return ResourceResult(
                success=False,
                stdout="",
                stderr=str(e),
                exit_code=-1,
                violation=ResourceViolation.NONE,
                violation_message=f"Execution error: {e}",
                memory_used_mb=0,
                cpu_time_seconds=0,
                wall_time_seconds=wall_time,
            )


def format_resource_report(result: ResourceResult) -> str:
    """Format resource usage as a report string."""
    lines = [
        f"Exit code: {result.exit_code}",
        f"Wall time: {result.wall_time_seconds:.2f}s",
        f"CPU time: {result.cpu_time_seconds:.2f}s",
        f"Memory: {result.memory_used_mb:.0f}MB",
    ]
    
    if result.violation != ResourceViolation.NONE:
        lines.append(f"⚠️ Violation: {result.violation_message}")
    
    if result.killed:
        lines.append("🔴 Process killed")
    
    return " | ".join(lines)


def execute_with_limits(
    command: str,
    memory_mb: int = 2048,
    cpu_percent: int = 80,
    timeout_seconds: int = 300,
) -> ResourceResult:
    """
    Convenience function to execute with resource limits.
    
    Args:
        command: Shell command
        memory_mb: Memory limit in MB
        cpu_percent: CPU limit as percent of one core
        timeout_seconds: Timeout in seconds
        
    Returns:
        ResourceResult
    """
    limits = ResourceLimits(
        memory_mb=memory_mb,
        cpu_percent=cpu_percent,
        timeout_seconds=timeout_seconds,
    )
    limiter = ResourceLimiter(limits)
    return limiter.execute(command)