diff --git a/src/agent/__init__.py b/src/agent/__init__.py index dbfe92991..50d6b76d4 100644 --- a/src/agent/__init__.py +++ b/src/agent/__init__.py @@ -1,11 +1,17 @@ """Agent lifecycle management module.""" -from .registry import AgentRegistry +from .registry import AgentRegistry, AgentStatus from .executor import AgentExecutor from .runtime import AgentRuntime from .sandbox import AgentSandbox -__all__ = ["AgentRegistry", "AgentExecutor", "AgentRuntime", "AgentSandbox"] +__all__ = [ + "AgentRegistry", + "AgentStatus", + "AgentExecutor", + "AgentRuntime", + "AgentSandbox", +] # 2019-02-05T12:34:30 update diff --git a/src/agent/sandbox.py b/src/agent/sandbox.py index c6ad6969c..d3481e68f 100644 --- a/src/agent/sandbox.py +++ b/src/agent/sandbox.py @@ -1,14 +1,19 @@ """Agent Sandbox — Isolated execution environment for agents.""" -import os -import tempfile import resource -from typing import Dict, Optional +import shutil +import tempfile from pathlib import Path +from typing import Dict, List, Optional class ResourceLimits: - def __init__(self, cpu_time: int = 60, memory_mb: int = 512, disk_mb: int = 100): + def __init__( + self, + cpu_time: int = 60, + memory_mb: int = 512, + disk_mb: int = 100, + ): self.cpu_time = cpu_time self.memory_mb = memory_mb self.disk_mb = disk_mb @@ -16,37 +21,56 @@ def __init__(self, cpu_time: int = 60, memory_mb: int = 512, disk_mb: int = 100) class AgentSandbox: def __init__(self, base_path: Optional[str] = None): - self.base_path = Path(base_path or tempfile.mkdtemp(prefix="ao_sandbox_")) + self.base_path = Path( + base_path or tempfile.mkdtemp(prefix="ao_sandbox_") + ) self._sandboxes: Dict[str, Path] = {} - def create(self, agent_id: str, limits: Optional[ResourceLimits] = None) -> Path: + def create( + self, + agent_id: str, + limits: Optional[ResourceLimits] = None, + ) -> Path: sandbox_path = self.base_path / agent_id sandbox_path.mkdir(parents=True, exist_ok=True) self._sandboxes[agent_id] = sandbox_path return sandbox_path def destroy(self, agent_id: str) -> bool: - sandbox = self._sandboxes.pop(agent_id, None) - if sandbox and sandbox.exists(): - import shutil - shutil.rmtree(sandbox, ignore_errors=True) - return True - return False + sandbox = self._sandboxes.get(agent_id) + if not sandbox: + return False + if not sandbox.exists(): + self._sandboxes.pop(agent_id, None) + return False + try: + shutil.rmtree(sandbox) + except OSError: + return False + + self._sandboxes.pop(agent_id, None) + return True def get_path(self, agent_id: str) -> Optional[Path]: return self._sandboxes.get(agent_id) def apply_limits(self, agent_id: str, limits: ResourceLimits) -> None: try: - resource.setrlimit(resource.RLIMIT_CPU, (limits.cpu_time, limits.cpu_time)) + resource.setrlimit( + resource.RLIMIT_CPU, + (limits.cpu_time, limits.cpu_time), + ) mem_bytes = limits.memory_mb * 1024 * 1024 resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes)) - except (ValueError, resource.error) as e: + except (ValueError, resource.error): pass - def cleanup_all(self) -> None: + def cleanup_all(self) -> List[str]: + failed = [] for agent_id in list(self._sandboxes.keys()): - self.destroy(agent_id) + if not self.destroy(agent_id): + failed.append(agent_id) + return failed # 2019-01-10T19:56:24 update diff --git a/src/common/metrics.py b/src/common/metrics.py index 99d82f499..c8d063360 100644 --- a/src/common/metrics.py +++ b/src/common/metrics.py @@ -3,12 +3,12 @@ import time from collections import defaultdict from typing import Dict, List -from threading import Lock +from threading import RLock class MetricsCollector: def __init__(self): - self._lock = Lock() + self._lock = RLock() self._counters: Dict[str, int] = defaultdict(int) self._gauges: Dict[str, float] = {} self._histograms: Dict[str, List[float]] = defaultdict(list) @@ -43,8 +43,14 @@ def snapshot(self) -> Dict: return { "counters": dict(self._counters), "gauges": dict(self._gauges), - "histograms": {k: {"count": len(v), "sum": sum(v), "avg": sum(v) / len(v) if v else 0} - for k, v in self._histograms.items()}, + "histograms": { + k: { + "count": len(v), + "sum": sum(v), + "avg": sum(v) / len(v) if v else 0, + } + for k, v in self._histograms.items() + }, } diff --git a/tests/test_agent_sandbox.py b/tests/test_agent_sandbox.py new file mode 100644 index 000000000..64bc937a1 --- /dev/null +++ b/tests/test_agent_sandbox.py @@ -0,0 +1,38 @@ +from pathlib import Path +import shutil + +from src.agent.sandbox import AgentSandbox + + +def test_cleanup_all_returns_failures_and_continues(tmp_path, monkeypatch): + sandbox = AgentSandbox(str(tmp_path)) + ok_path = sandbox.create("ok-agent") + stuck_path = sandbox.create("stuck-agent") + original_rmtree = shutil.rmtree + + def fail_for_stuck(path): + if Path(path) == stuck_path: + raise OSError("locked") + return original_rmtree(path) + + monkeypatch.setattr(shutil, "rmtree", fail_for_stuck) + + failures = sandbox.cleanup_all() + + assert failures == ["stuck-agent"] + assert not ok_path.exists() + assert sandbox.get_path("ok-agent") is None + assert stuck_path.exists() + assert sandbox.get_path("stuck-agent") == stuck_path + + +def test_cleanup_all_returns_empty_list_when_all_destroyed(tmp_path): + sandbox = AgentSandbox(str(tmp_path)) + first = sandbox.create("first") + second = sandbox.create("second") + + assert sandbox.cleanup_all() == [] + assert not first.exists() + assert not second.exists() + assert sandbox.get_path("first") is None + assert sandbox.get_path("second") is None