Rate Limiting, Sandboxing & Execution Controls

from dataclasses import dataclass, field
from datetime import datetime, timedelta
 
@dataclass
class RateLimitConfig:
    requests_per_minute: int = 20
    tokens_per_minute: int = 40_000
    max_input_tokens: int = 4_096
    max_output_tokens: int = 4_096
    max_session_messages: int = 50
    cost_limit_per_hour_usd: float = 10.0
 
class AIRateLimiter:
    def __init__(self, config: RateLimitConfig):
        self.config = config
        self.windows: dict[str, list] = {}
 
    def check_and_record(
        self, user_id: str, input_tokens: int, session_messages: int
    ) -> tuple[bool, str]:
        now = datetime.utcnow()
        window_start = now - timedelta(minutes=1)
        key = f"rpm:{user_id}"
 
        # Request rate check
        recent = [t for t in self.windows.get(key, []) if t > window_start]
        if len(recent) >= self.config.requests_per_minute:
            return False, "Rate limit exceeded: too many requests"
 
        # Input size check
        if input_tokens > self.config.max_input_tokens:
            return False, f"Input too large: {input_tokens} tokens"
 
        # Session length check
        if session_messages > self.config.max_session_messages:
            return False, "Session message limit reached"
 
        recent.append(now)
        self.windows[key] = recent
        return True, "OK"

from e2b_code_interpreter import Sandbox
 
# Create isolated sandbox with 30-second timeout
sandbox = Sandbox(timeout=30)
 
# Execute AI-generated code in isolation
result = sandbox.run_code("""
import os
# This runs in a completely isolated environment
# No access to host filesystem, network restrictions applied
print(os.listdir('/'))  # Only sees sandbox filesystem
""")
 
print(result.text)  # Output from sandboxed execution
sandbox.kill()  # Clean up

import docker
import tempfile
 
client = docker.from_env()
 
def execute_sandboxed(code: str, timeout: int = 10) -> str:
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(code)
        f.flush()
 
        container = client.containers.run(
            "python:3.11-slim",
            command=f"python /code/{f.name.split('/')[-1]}",
            volumes={f.name: {"bind": f"/code/{f.name.split('/')[-1]}", "mode": "ro"}},
            network_disabled=True,      # No network access
            mem_limit="256m",           # Memory cap
            cpu_period=100000,          # CPU throttling
            cpu_quota=50000,
            read_only=True,             # Read-only filesystem
            remove=True,               # Auto-cleanup
            timeout=timeout,
        )
        return container.decode("utf-8")

# Using wasmtime for sandboxed Python execution
# Pyodide runs Python in a WASM sandbox with no system access
from pyodide_sandbox import PySandbox
 
sandbox = PySandbox(
    allowed_imports=["math", "json", "re"],  # Whitelist imports
    max_execution_time_ms=5000,
    max_memory_mb=64,
)
 
result = sandbox.execute(ai_generated_code)

from enum import Enum
 
class ApprovalPolicy(Enum):
    AUTO_APPROVE = "auto"          # Low-risk tools
    NOTIFY_AND_PROCEED = "notify"  # Medium-risk: log but allow
    REQUIRE_APPROVAL = "approve"   # High-risk: block until human approves
    ALWAYS_DENY = "deny"           # Forbidden tools
 
TOOL_POLICIES = {
    "search_web": ApprovalPolicy.AUTO_APPROVE,
    "read_file": ApprovalPolicy.NOTIFY_AND_PROCEED,
    "send_email": ApprovalPolicy.REQUIRE_APPROVAL,
    "execute_sql": ApprovalPolicy.REQUIRE_APPROVAL,
    "delete_record": ApprovalPolicy.ALWAYS_DENY,
    "modify_permissions": ApprovalPolicy.ALWAYS_DENY,
}
 
class ToolGatekeeper:
    def __init__(self, policies: dict[str, ApprovalPolicy]):
        self.policies = policies
 
    async def check_tool_call(
        self, tool_name: str, arguments: dict, user_id: str
    ) -> tuple[bool, str]:
        policy = self.policies.get(tool_name, ApprovalPolicy.ALWAYS_DENY)
 
        if policy == ApprovalPolicy.AUTO_APPROVE:
            return True, "Auto-approved"
        elif policy == ApprovalPolicy.NOTIFY_AND_PROCEED:
            await self.log_notification(tool_name, arguments, user_id)
            return True, "Approved with notification"
        elif policy == ApprovalPolicy.REQUIRE_APPROVAL:
            approved = await self.request_human_approval(
                tool_name, arguments, user_id
            )
            return approved, "Human review"
        else:
            return False, f"Tool '{tool_name}' is forbidden"