From 94fa42c24228ba52f2bf99e93f17c42dfa1f1d52 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:13:14 -0400 Subject: [PATCH 01/14] feat(agent): enforce per-sandbox cgroup v2 memory limits Ten concurrent agent sandboxes on one executor pod crossed the pod memory limit and the kernel group-killed the container, failing all in-flight sessions at once. rlimit_as bounds neither the sandbox process tree nor tmpfs, so aggregate memory was unbounded. - emit use_cgroupv2/cgroup_mem_max nsjail config (fields verified against the pinned nsjail commit) so each sandbox tree gets its own kernel-enforced memory.max and a runaway session fails alone - prepare the cgroup root at worker startup (move PIDs to a leaf, enable the memory controller, verify with a probe cgroup); any failure logs one warning and falls back to no cgroup limits - clamp max_concurrent_activities to the container memory budget (limit - reserve) / per-sandbox cap instead of overcommitting - write a best-effort readiness sentinel for a future exec probe --- .../test_agent_executor_worker_resources.py | 126 ++++++++++ tests/unit/test_agent_sandbox_cgroup.py | 229 +++++++++++++++++ tests/unit/test_agent_sandbox_config.py | 51 ++++ tracecat/agent/executor_worker.py | 50 +++- tracecat/agent/sandbox/cgroup.py | 234 ++++++++++++++++++ tracecat/agent/sandbox/config.py | 13 + tracecat/agent/sandbox/nsjail.py | 4 + tracecat/config.py | 16 ++ 8 files changed, 722 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_agent_executor_worker_resources.py create mode 100644 tests/unit/test_agent_sandbox_cgroup.py create mode 100644 tracecat/agent/sandbox/cgroup.py diff --git a/tests/unit/test_agent_executor_worker_resources.py b/tests/unit/test_agent_executor_worker_resources.py new file mode 100644 index 0000000000..2bccfad564 --- /dev/null +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import asyncio +import errno +from datetime import UTC, datetime +from pathlib import Path +from unittest.mock import AsyncMock, Mock + +import pytest + +from tracecat.agent.sandbox.cgroup import CgroupAvailability + + +@pytest.mark.anyio +async def test_agent_executor_readiness_sentinel_exists_only_while_running( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + ready_file = tmp_path / "run" / "agent-executor-ready" + shutdown_event = asyncio.Event() + observed_contents: list[str] = [] + + class _FakeWorker: + def __init__(self, *args: object, **kwargs: object) -> None: + del args, kwargs + + async def __aenter__(self) -> _FakeWorker: + return self + + async def __aexit__( + self, + exc_type: object, + exc: object, + tb: object, + ) -> None: + del exc_type, exc, tb + + def keep_concurrency( + max_concurrent: int, + *, + reserve_mb: int, + sandbox_memory_mb: int, + ) -> int: + del reserve_mb, sandbox_memory_mb + return max_concurrent + + async def observe_readiness() -> None: + try: + for _ in range(1000): + if ready_file.exists(): + observed_contents.append(ready_file.read_text().strip()) + return + await asyncio.sleep(0) + pytest.fail("readiness sentinel was not created") + finally: + shutdown_event.set() + + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "2") + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + lambda: CgroupAvailability.DISABLED, + ) + monkeypatch.setattr( + executor_worker, + "clamp_agent_executor_concurrency", + keep_concurrency, + ) + monkeypatch.setattr( + executor_worker, + "_start_runtime_services", + AsyncMock(return_value=object()), + ) + monkeypatch.setattr(executor_worker, "_stop_runtime_services", AsyncMock()) + monkeypatch.setattr(executor_worker, "close_storage_client_cache", AsyncMock()) + monkeypatch.setattr(executor_worker, "Worker", _FakeWorker) + monkeypatch.setattr(executor_worker, "new_sandbox_runner", lambda: object()) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_READY_FILE", + str(ready_file), + ) + + observer = asyncio.create_task(observe_readiness()) + await asyncio.wait_for( + executor_worker.main(shutdown_event=shutdown_event), + timeout=2, + ) + await observer + + assert not ready_file.exists() + assert len(observed_contents) == 1 + started_at = datetime.fromisoformat(observed_contents[0]) + assert started_at.tzinfo is UTC + + +def test_agent_executor_readiness_sentinel_is_best_effort_on_read_only_fs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + mock_logger = Mock() + monkeypatch.setattr(executor_worker, "logger", mock_logger) + + def deny_mkdir( + self: Path, + mode: int = 0o777, + parents: bool = False, + exist_ok: bool = False, + ) -> None: + del self, mode, parents, exist_ok + raise PermissionError(errno.EROFS, "read-only filesystem") + + monkeypatch.setattr(Path, "mkdir", deny_mkdir) + + created = executor_worker._write_readiness_file( + tmp_path / "run" / "agent-executor-ready", + datetime(2026, 7, 27, tzinfo=UTC), + ) + + assert created is False + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args.kwargs["errno"] == errno.EROFS diff --git a/tests/unit/test_agent_sandbox_cgroup.py b/tests/unit/test_agent_sandbox_cgroup.py new file mode 100644 index 0000000000..80e131ad63 --- /dev/null +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import errno +from collections.abc import Iterator +from pathlib import Path +from unittest.mock import Mock + +import pytest + +import tracecat.agent.sandbox.cgroup as cgroup_module +from tracecat.agent.sandbox.cgroup import ( + BYTES_PER_MEBIBYTE, + CgroupAvailability, + CgroupMemoryLimitKind, + clamp_agent_executor_concurrency, + prepare_agent_sandbox_cgroup, + read_cgroup_memory_limit, +) + + +@pytest.fixture(autouse=True) +def reset_cgroup_availability(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(cgroup_module, "_cgroup_availability", None) + yield + + +def _create_fake_cgroup_root( + tmp_path: Path, + *, + pids: str = "101\n202\n", +) -> Path: + cgroup_root = tmp_path / "cgroup" + cgroup_root.mkdir() + (cgroup_root / "cgroup.controllers").write_text("cpu memory pids\n") + (cgroup_root / "cgroup.procs").write_text(pids) + (cgroup_root / "cgroup.subtree_control").write_text("") + return cgroup_root + + +def test_prepare_agent_sandbox_cgroup_returns_cached_disabled_state( + tmp_path: Path, +) -> None: + cgroup_root = tmp_path / "missing" + + assert ( + prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=False) + is CgroupAvailability.DISABLED + ) + assert ( + prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + is CgroupAvailability.DISABLED + ) + + +def test_prepare_agent_sandbox_cgroup_happy_path(tmp_path: Path) -> None: + cgroup_root = _create_fake_cgroup_root(tmp_path) + + result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + + assert result is CgroupAvailability.AVAILABLE + assert (cgroup_root / "main" / "cgroup.procs").read_text() == "202\n" + assert (cgroup_root / "cgroup.subtree_control").read_text() == "+memory\n" + assert not list(cgroup_root.glob("tracecat-agent-probe-*")) + + +def test_prepare_agent_sandbox_cgroup_without_v2_controllers_warns_once( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cgroup_root = tmp_path / "cgroup" + cgroup_root.mkdir() + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + + result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + + assert result is CgroupAvailability.UNAVAILABLE + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args.kwargs["step"] == "read cgroup v2 controllers" + assert mock_logger.warning.call_args.kwargs["errno"] == errno.ENOENT + + +def test_prepare_agent_sandbox_cgroup_subtree_permission_error_falls_back( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cgroup_root = _create_fake_cgroup_root(tmp_path) + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + write_cgroup_file = cgroup_module._write_cgroup_file + + def deny_subtree_control(path: Path, value: str) -> None: + if path == cgroup_root / "cgroup.subtree_control": + raise PermissionError(errno.EACCES, "permission denied") + write_cgroup_file(path, value) + + monkeypatch.setattr( + cgroup_module, + "_write_cgroup_file", + deny_subtree_control, + ) + + result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + + assert result is CgroupAvailability.UNAVAILABLE + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args.kwargs["step"] == "enable memory controller" + assert mock_logger.warning.call_args.kwargs["errno"] == errno.EACCES + + +def test_prepare_agent_sandbox_cgroup_tolerates_pid_vanishing_during_move( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cgroup_root = _create_fake_cgroup_root(tmp_path) + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + write_cgroup_file = cgroup_module._write_cgroup_file + + def vanish_first_pid(path: Path, value: str) -> None: + if path == cgroup_root / "main" / "cgroup.procs" and value == "101\n": + raise ProcessLookupError(errno.ESRCH, "process vanished") + write_cgroup_file(path, value) + + monkeypatch.setattr(cgroup_module, "_write_cgroup_file", vanish_first_pid) + + result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + + assert result is CgroupAvailability.AVAILABLE + assert (cgroup_root / "main" / "cgroup.procs").read_text() == "202\n" + mock_logger.warning.assert_not_called() + + +@pytest.mark.parametrize( + ("contents", "expected_kind", "expected_bytes"), + [ + pytest.param("max\n", CgroupMemoryLimitKind.UNLIMITED, None, id="max"), + pytest.param( + f"{16 * 1024**3}\n", + CgroupMemoryLimitKind.LIMITED, + 16 * 1024**3, + id="numeric", + ), + ], +) +def test_read_cgroup_memory_limit_parses_valid_values( + tmp_path: Path, + contents: str, + expected_kind: CgroupMemoryLimitKind, + expected_bytes: int | None, +) -> None: + memory_max_path = tmp_path / "memory.max" + memory_max_path.write_text(contents) + + result = read_cgroup_memory_limit(memory_max_path) + + assert result.kind is expected_kind + assert result.limit_bytes == expected_bytes + + +def test_read_cgroup_memory_limit_missing_is_unavailable(tmp_path: Path) -> None: + result = read_cgroup_memory_limit(tmp_path / "memory.max") + + assert result.kind is CgroupMemoryLimitKind.UNAVAILABLE + assert result.limit_bytes is None + + +def test_read_cgroup_memory_limit_garbage_warns_and_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + memory_max_path = tmp_path / "memory.max" + memory_max_path.write_text("not-a-limit\n") + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + + result = read_cgroup_memory_limit(memory_max_path) + + assert result.kind is CgroupMemoryLimitKind.UNAVAILABLE + assert result.limit_bytes is None + mock_logger.warning.assert_called_once() + + +@pytest.mark.parametrize( + ( + "memory_max_contents", + "wanted", + "expected", + "expected_error", + ), + [ + pytest.param(f"{16 * 1024**3}\n", 10, 3, True, id="over-budget"), + pytest.param("max\n", 10, 10, False, id="unlimited"), + pytest.param(None, 10, 10, False, id="missing"), + pytest.param(f"{4 * 1024**3}\n", 10, 1, True, id="allowed-zero"), + pytest.param(f"{16 * 1024**3}\n", 2, 2, False, id="under-budget"), + ], +) +def test_clamp_agent_executor_concurrency( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + memory_max_contents: str | None, + wanted: int, + expected: int, + expected_error: bool, +) -> None: + memory_max_path = tmp_path / "memory.max" + if memory_max_contents is not None: + memory_max_path.write_text(memory_max_contents) + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + + result = clamp_agent_executor_concurrency( + wanted, + reserve_mb=4096, + sandbox_memory_mb=4096, + memory_max_path=memory_max_path, + ) + + assert result == expected + assert mock_logger.error.called is expected_error + if memory_max_contents == f"{16 * 1024**3}\n" and wanted == 10: + assert mock_logger.error.call_args.kwargs == { + "container_limit_mb": 16 * 1024**3 // BYTES_PER_MEBIBYTE, + "reserve_mb": 4096, + "sandbox_memory_mb": 4096, + "configured_max_concurrent_activities": 10, + "clamped_max_concurrent_activities": 3, + } diff --git a/tests/unit/test_agent_sandbox_config.py b/tests/unit/test_agent_sandbox_config.py index 18913f8bb6..176ad0cdc7 100644 --- a/tests/unit/test_agent_sandbox_config.py +++ b/tests/unit/test_agent_sandbox_config.py @@ -2,7 +2,10 @@ from pathlib import Path +import pytest + from tracecat.agent.sandbox.config import ( + AgentResourceLimits, AgentSandboxConfig, build_agent_nsjail_config, ) @@ -121,3 +124,51 @@ def test_build_agent_nsjail_config_mounts_fresh_procfs() -> None: assert 'src: "/proc"' not in config_text assert 'mount { dst: "/proc" fstype: "proc" rw: false }' in config_text + + +def test_build_agent_nsjail_config_adds_available_cgroup_v2_memory_limit() -> None: + config_text = build_agent_nsjail_config( + rootfs=Path("/var/lib/tracecat/sandbox-rootfs"), + job_dir=Path("/tmp/agent-job"), + socket_dir=Path("/tmp/agent-job/sockets"), + config=AgentSandboxConfig( + resources=AgentResourceLimits(memory_mb=3072), + ), + site_packages_dir=Path("/app/.venv/lib/python3.12/site-packages"), + llm_socket_path=Path("/tmp/agent-job/sockets/llm.sock"), + cgroup_v2_enabled=True, + cgroup_v2_available=True, + ) + + assert "use_cgroupv2: true" in config_text + assert 'cgroupv2_mount: "/sys/fs/cgroup"' in config_text + assert f"cgroup_mem_max: {3072 * 1024 * 1024}" in config_text + assert "cgroup_mem_swap_max: 0" in config_text + + +@pytest.mark.parametrize( + ("cgroup_v2_enabled", "cgroup_v2_available"), + [ + pytest.param(False, True, id="disabled"), + pytest.param(True, False, id="unavailable"), + ], +) +def test_build_agent_nsjail_config_omits_inactive_cgroup_v2_memory_limit( + cgroup_v2_enabled: bool, + cgroup_v2_available: bool, +) -> None: + config_text = build_agent_nsjail_config( + rootfs=Path("/var/lib/tracecat/sandbox-rootfs"), + job_dir=Path("/tmp/agent-job"), + socket_dir=Path("/tmp/agent-job/sockets"), + config=AgentSandboxConfig(), + site_packages_dir=Path("/app/.venv/lib/python3.12/site-packages"), + llm_socket_path=Path("/tmp/agent-job/sockets/llm.sock"), + cgroup_v2_enabled=cgroup_v2_enabled, + cgroup_v2_available=cgroup_v2_available, + ) + + assert "use_cgroupv2:" not in config_text + assert "cgroupv2_mount:" not in config_text + assert "cgroup_mem_max:" not in config_text + assert "cgroup_mem_swap_max:" not in config_text diff --git a/tracecat/agent/executor_worker.py b/tracecat/agent/executor_worker.py index c4d8b52899..a726459e44 100644 --- a/tracecat/agent/executor_worker.py +++ b/tracecat/agent/executor_worker.py @@ -5,7 +5,8 @@ import asyncio import os from concurrent.futures import ThreadPoolExecutor -from datetime import timedelta +from datetime import UTC, datetime, timedelta +from pathlib import Path from typing import TYPE_CHECKING import uvloop @@ -22,6 +23,10 @@ stop_claude_runtime_broker, stop_mcp_server, ) +from tracecat.agent.sandbox.cgroup import ( + clamp_agent_executor_concurrency, + prepare_agent_sandbox_cgroup, +) from tracecat.agent.worker import new_sandbox_runner from tracecat.dsl.client import get_temporal_client from tracecat.logger import logger @@ -34,6 +39,35 @@ runtime_failure_reason: str | None = None +def _write_readiness_file(path: Path, started_at: datetime) -> bool: + """Write the best-effort agent executor readiness sentinel.""" + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"{started_at.isoformat()}\n") + except OSError as exc: + logger.warning( + "Unable to write AgentExecutorWorker readiness sentinel; continuing", + path=str(path), + errno=exc.errno, + error=str(exc), + ) + return False + return True + + +def _remove_readiness_file(path: Path) -> None: + """Remove the best-effort agent executor readiness sentinel.""" + try: + path.unlink(missing_ok=True) + except OSError as exc: + logger.warning( + "Unable to remove AgentExecutorWorker readiness sentinel; continuing", + path=str(path), + errno=exc.errno, + error=str(exc), + ) + + def get_activities() -> list: """Load runtime activities registered by the agent-executor worker.""" return [ @@ -86,6 +120,14 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: threadpool_max_workers = int( os.environ.get("TEMPORAL__THREADPOOL_MAX_WORKERS") or 100 ) + prepare_agent_sandbox_cgroup() + max_concurrent = clamp_agent_executor_concurrency( + max_concurrent, + reserve_mb=config.TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB, + sandbox_memory_mb=config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, + ) + readiness_file = Path(config.TRACECAT__AGENT_EXECUTOR_READY_FILE) + readiness_file_created = False logger.info( "Starting AgentExecutorWorker", @@ -117,10 +159,16 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: ), ): logger.info("AgentExecutorWorker started, ctrl+c to exit") + readiness_file_created = _write_readiness_file( + readiness_file, + datetime.now(UTC), + ) await shutdown_event.wait() logger.info("AgentExecutorWorker shutdown requested") logger.info("Temporal Worker context exited") finally: + if readiness_file_created: + _remove_readiness_file(readiness_file) await close_storage_client_cache() await _stop_runtime_services() if runtime_failure_reason is not None: diff --git a/tracecat/agent/sandbox/cgroup.py b/tracecat/agent/sandbox/cgroup.py new file mode 100644 index 0000000000..659121d764 --- /dev/null +++ b/tracecat/agent/sandbox/cgroup.py @@ -0,0 +1,234 @@ +"""Cgroup v2 preparation and agent executor memory-budget helpers.""" + +from __future__ import annotations + +import errno +import uuid +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path + +from tracecat import config +from tracecat.logger import logger + +CGROUP_ROOT = Path("/sys/fs/cgroup") +CGROUP_MEMORY_MAX_PATH = CGROUP_ROOT / "memory.max" +BYTES_PER_MEBIBYTE = 1024 * 1024 + + +class CgroupAvailability(StrEnum): + """Process-wide availability of agent sandbox cgroup limits.""" + + DISABLED = "disabled" + UNAVAILABLE = "unavailable" + AVAILABLE = "available" + + +class CgroupMemoryLimitKind(StrEnum): + """Kind of container memory limit reported by cgroup v2.""" + + LIMITED = "limited" + UNLIMITED = "unlimited" + UNAVAILABLE = "unavailable" + + +@dataclass(frozen=True, slots=True) +class CgroupMemoryLimit: + """Parsed cgroup v2 container memory limit.""" + + kind: CgroupMemoryLimitKind + limit_bytes: int | None = None + + +_cgroup_availability: CgroupAvailability | None = None + + +def _write_cgroup_file(path: Path, value: str) -> None: + """Write a value to a cgroup control file.""" + path.write_text(value) + + +def _remove_probe_cgroup(path: Path) -> None: + """Remove a probe cgroup, including its synthetic file in unit tests.""" + try: + path.rmdir() + except OSError as exc: + if exc.errno != errno.ENOTEMPTY: + raise + (path / "memory.max").unlink() + path.rmdir() + + +def _warn_cgroup_unavailable( + *, + step: str, + error: OSError | ValueError, +) -> None: + """Log the single startup warning for a failed cgroup preparation.""" + if isinstance(error, OSError): + logger.warning( + "Agent sandbox cgroup memory limits unavailable; " + "continuing without per-sandbox cgroup limits", + step=step, + errno=error.errno, + error=str(error), + ) + return + logger.warning( + "Agent sandbox cgroup memory limits unavailable; " + "continuing without per-sandbox cgroup limits", + step=step, + error=str(error), + ) + + +def prepare_agent_sandbox_cgroup( + *, + cgroup_root: Path = CGROUP_ROOT, + enabled: bool = config.TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, +) -> CgroupAvailability: + """Prepare the cgroup v2 root for nsjail child memory cgroups. + + The result is cached process-wide. Setup failures are reduced to one warning + and an unavailable result so agent sandbox launches can continue without + cgroup limits. + """ + global _cgroup_availability + + if _cgroup_availability is not None: + return _cgroup_availability + if not enabled: + _cgroup_availability = CgroupAvailability.DISABLED + return _cgroup_availability + + probe_cgroup: Path | None = None + step = "read cgroup v2 controllers" + try: + (cgroup_root / "cgroup.controllers").read_text() + + step = "create main cgroup" + main_cgroup = cgroup_root / "main" + main_cgroup.mkdir(exist_ok=True) + + step = "read root cgroup processes" + root_pids = (cgroup_root / "cgroup.procs").read_text().splitlines() + for raw_pid in root_pids: + if not raw_pid: + continue + pid = int(raw_pid) + step = f"move PID {pid} to main cgroup" + try: + _write_cgroup_file(main_cgroup / "cgroup.procs", f"{pid}\n") + except OSError as exc: + if exc.errno in {errno.ENOENT, errno.ESRCH}: + continue + raise + + step = "enable memory controller" + _write_cgroup_file(cgroup_root / "cgroup.subtree_control", "+memory\n") + + step = "create verification cgroup" + probe_cgroup = cgroup_root / f"tracecat-agent-probe-{uuid.uuid4().hex}" + probe_cgroup.mkdir() + + step = "write verification memory.max" + _write_cgroup_file(probe_cgroup / "memory.max", "max\n") + + step = "remove verification cgroup" + _remove_probe_cgroup(probe_cgroup) + probe_cgroup = None + except (OSError, ValueError) as exc: + if probe_cgroup is not None: + try: + _remove_probe_cgroup(probe_cgroup) + except OSError: + pass + _warn_cgroup_unavailable(step=step, error=exc) + _cgroup_availability = CgroupAvailability.UNAVAILABLE + return _cgroup_availability + + _cgroup_availability = CgroupAvailability.AVAILABLE + return _cgroup_availability + + +def agent_sandbox_cgroup_is_available() -> bool: + """Return whether the startup probe enabled agent sandbox cgroups.""" + return _cgroup_availability is CgroupAvailability.AVAILABLE + + +def read_cgroup_memory_limit( + memory_max_path: Path = CGROUP_MEMORY_MAX_PATH, +) -> CgroupMemoryLimit: + """Read and parse the cgroup v2 container memory limit.""" + try: + raw_limit = memory_max_path.read_text().strip() + except FileNotFoundError: + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + except OSError as exc: + logger.warning( + "Unable to read cgroup memory limit; skipping worker memory-budget " + "validation", + path=str(memory_max_path), + errno=exc.errno, + error=str(exc), + ) + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + + if raw_limit == "max": + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNLIMITED) + + try: + limit_bytes = int(raw_limit) + if limit_bytes < 0: + raise ValueError("memory limit cannot be negative") + except ValueError as exc: + logger.warning( + "Invalid cgroup memory limit; skipping worker memory-budget validation", + path=str(memory_max_path), + value=raw_limit, + error=str(exc), + ) + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + + return CgroupMemoryLimit( + CgroupMemoryLimitKind.LIMITED, + limit_bytes=limit_bytes, + ) + + +def clamp_agent_executor_concurrency( + max_concurrent: int, + *, + reserve_mb: int, + sandbox_memory_mb: int, + memory_max_path: Path = CGROUP_MEMORY_MAX_PATH, +) -> int: + """Clamp agent executor concurrency to the container memory budget.""" + memory_limit = read_cgroup_memory_limit(memory_max_path) + if memory_limit.kind is CgroupMemoryLimitKind.UNLIMITED: + logger.debug( + "Container memory is unlimited; skipping worker memory-budget validation" + ) + return max_concurrent + if memory_limit.kind is CgroupMemoryLimitKind.UNAVAILABLE: + return max_concurrent + + limit_bytes = memory_limit.limit_bytes + if limit_bytes is None: + return max_concurrent + + limit_mb = limit_bytes // BYTES_PER_MEBIBYTE + allowed = (limit_mb - reserve_mb) // sandbox_memory_mb + if max_concurrent <= allowed: + return max_concurrent + + clamped = max(allowed, 1) + logger.error( + "Agent executor concurrency exceeds the container memory budget; clamping", + container_limit_mb=limit_mb, + reserve_mb=reserve_mb, + sandbox_memory_mb=sandbox_memory_mb, + configured_max_concurrent_activities=max_concurrent, + clamped_max_concurrent_activities=clamped, + ) + return clamped diff --git a/tracecat/agent/sandbox/config.py b/tracecat/agent/sandbox/config.py index 4e1caab8f6..7cfc8f5a09 100644 --- a/tracecat/agent/sandbox/config.py +++ b/tracecat/agent/sandbox/config.py @@ -222,6 +222,8 @@ def build_agent_nsjail_config( session_work_dir: Path | None = None, enable_internet_access: bool = False, skills_dir: Path | None = None, + cgroup_v2_enabled: bool = False, + cgroup_v2_available: bool = False, ) -> str: """Build nsjail protobuf config for agent runtime execution. @@ -245,6 +247,8 @@ def build_agent_nsjail_config( outbound internet access. Default is False (network isolated with private loopback only). skills_dir: Optional host path containing staged workspace skills. + cgroup_v2_enabled: Whether agent sandbox cgroup limits are configured. + cgroup_v2_available: Whether the process-wide startup probe succeeded. Returns: nsjail protobuf configuration as a string. @@ -472,6 +476,15 @@ def build_agent_nsjail_config( f"time_limit: {config.resources.timeout_seconds}", ] ) + if cgroup_v2_enabled and cgroup_v2_available: + lines.extend( + [ + "use_cgroupv2: true", + 'cgroupv2_mount: "/sys/fs/cgroup"', + f"cgroup_mem_max: {config.resources.memory_mb * 1024 * 1024}", + "cgroup_mem_swap_max: 0", + ] + ) # Execution settings. lines.extend( diff --git a/tracecat/agent/sandbox/nsjail.py b/tracecat/agent/sandbox/nsjail.py index 901d9e6bab..a04abf8700 100644 --- a/tracecat/agent/sandbox/nsjail.py +++ b/tracecat/agent/sandbox/nsjail.py @@ -46,6 +46,7 @@ JAILED_AGENT_JOB_DIR, JAILED_AGENT_WORK_DIR, ) +from tracecat.agent.sandbox.cgroup import agent_sandbox_cgroup_is_available from tracecat.agent.sandbox.config import ( JAILED_SHIM_ENTRYPOINT_PATH, AgentSandboxConfig, @@ -53,6 +54,7 @@ build_agent_nsjail_config, ) from tracecat.config import ( + TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, TRACECAT__SANDBOX_NSJAIL_PATH, TRACECAT__SANDBOX_ROOTFS_PATH, ) @@ -439,6 +441,8 @@ async def _spawn_nsjail_runtime( session_work_dir=session_work_dir, enable_internet_access=enable_internet_access, skills_dir=skills_dir, + cgroup_v2_enabled=TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, + cgroup_v2_available=agent_sandbox_cgroup_is_available(), ) # Write config to job directory diff --git a/tracecat/config.py b/tracecat/config.py index 83a0a63f43..797d3e76e5 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -734,6 +734,22 @@ def _parse_auth_types() -> set[AuthType]: ) """Default memory limit for agent sandbox execution in megabytes (4 GiB).""" +TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED = env_bool( + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", default=True +) +"""Enable cgroup v2 memory limits for agent sandboxes when available.""" + +TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB = int( + os.environ.get("TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB") or 4096 +) +"""Memory reserved for the agent executor worker and shared services.""" + +TRACECAT__AGENT_EXECUTOR_READY_FILE = os.environ.get( + "TRACECAT__AGENT_EXECUTOR_READY_FILE", + "/var/run/tracecat/agent-executor-ready", +) +"""Best-effort readiness sentinel written after the agent executor starts.""" + TRACECAT__LITELLM_PORT = int(os.environ.get("TRACECAT__LITELLM_PORT") or 4000) """Bind port for the managed LiteLLM service.""" From 842e26f357c5ddd5b852b0700f67c4e0f80a3595 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:40:47 -0400 Subject: [PATCH 02/14] fix(agent): clear stale readiness sentinel on startup A SIGKILLed worker cannot run its cleanup, and the sentinel may live on a filesystem that survives container restarts, so a file-existence readiness probe would report the restarted worker ready during startup. Remove any existing sentinel before startup work begins and remove it unconditionally on exit. --- .../test_agent_executor_worker_resources.py | 50 +++++++++++++++++++ tracecat/agent/executor_worker.py | 13 +++-- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_agent_executor_worker_resources.py b/tests/unit/test_agent_executor_worker_resources.py index 2bccfad564..190ba594de 100644 --- a/tests/unit/test_agent_executor_worker_resources.py +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -124,3 +124,53 @@ def deny_mkdir( assert created is False mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.kwargs["errno"] == errno.EROFS + + +@pytest.mark.anyio +async def test_agent_executor_removes_stale_readiness_sentinel_on_startup( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + ready_file = tmp_path / "run" / "agent-executor-ready" + ready_file.parent.mkdir(parents=True) + ready_file.write_text("stale\n") + + def keep_concurrency( + max_concurrent: int, + *, + reserve_mb: int, + sandbox_memory_mb: int, + ) -> int: + del reserve_mb, sandbox_memory_mb + return max_concurrent + + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + lambda: CgroupAvailability.DISABLED, + ) + monkeypatch.setattr( + executor_worker, + "clamp_agent_executor_concurrency", + keep_concurrency, + ) + monkeypatch.setattr( + executor_worker, + "_start_runtime_services", + AsyncMock(side_effect=RuntimeError("startup failed")), + ) + monkeypatch.setattr(executor_worker, "_stop_runtime_services", AsyncMock()) + monkeypatch.setattr(executor_worker, "close_storage_client_cache", AsyncMock()) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_READY_FILE", + str(ready_file), + ) + + with pytest.raises(RuntimeError, match="startup failed"): + await executor_worker.main(shutdown_event=asyncio.Event()) + + assert not ready_file.exists() diff --git a/tracecat/agent/executor_worker.py b/tracecat/agent/executor_worker.py index a726459e44..3e3a72320b 100644 --- a/tracecat/agent/executor_worker.py +++ b/tracecat/agent/executor_worker.py @@ -127,7 +127,10 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: sandbox_memory_mb=config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, ) readiness_file = Path(config.TRACECAT__AGENT_EXECUTOR_READY_FILE) - readiness_file_created = False + # A SIGKILLed predecessor cannot run its cleanup, and the sentinel may live + # on a filesystem that survives container restarts; clear it before any + # startup work so a readiness probe never sees a stale file. + _remove_readiness_file(readiness_file) logger.info( "Starting AgentExecutorWorker", @@ -159,16 +162,12 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: ), ): logger.info("AgentExecutorWorker started, ctrl+c to exit") - readiness_file_created = _write_readiness_file( - readiness_file, - datetime.now(UTC), - ) + _write_readiness_file(readiness_file, datetime.now(UTC)) await shutdown_event.wait() logger.info("AgentExecutorWorker shutdown requested") logger.info("Temporal Worker context exited") finally: - if readiness_file_created: - _remove_readiness_file(readiness_file) + _remove_readiness_file(readiness_file) await close_storage_client_cache() await _stop_runtime_services() if runtime_failure_reason is not None: From 886454af3b80d27468c9d4301a10baffd514c5ae Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:50:58 -0400 Subject: [PATCH 03/14] fix(config): treat blank agent executor ready-file env as unset --- tracecat/config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tracecat/config.py b/tracecat/config.py index 797d3e76e5..a7184e7062 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -744,9 +744,9 @@ def _parse_auth_types() -> set[AuthType]: ) """Memory reserved for the agent executor worker and shared services.""" -TRACECAT__AGENT_EXECUTOR_READY_FILE = os.environ.get( - "TRACECAT__AGENT_EXECUTOR_READY_FILE", - "/var/run/tracecat/agent-executor-ready", +TRACECAT__AGENT_EXECUTOR_READY_FILE = ( + os.environ.get("TRACECAT__AGENT_EXECUTOR_READY_FILE") + or "/var/run/tracecat/agent-executor-ready" ) """Best-effort readiness sentinel written after the agent executor starts.""" From 72679fb114458a6582ddf73d2e83ee170408804a Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:32:09 -0400 Subject: [PATCH 04/14] fix(agent): derive cgroup root and enforce sandbox memory budget Address review blockers on the per-sandbox cgroup limits: - Derive the cgroup v2 root from /proc/self/cgroup instead of assuming /sys/fs/cgroup is the delegated root, so privileged host-namespace containers (Kubernetes/containerd) never prepare children in the node hierarchy. Preparation returns a typed PreparedCgroup; the root is retained on preparation failure so budget validation still reads the container's memory.max. - Fail fast at worker startup when reserve + one sandbox exceeds the container memory limit instead of forcing concurrency to one. - Apply the prepared cgroup to stdio MCP probe sandboxes and cap probe memory at the per-slot sandbox limit. - Validate sandbox memory (>=1), executor reserve (>=0), and worker concurrency (>=1) at the configuration boundary. - Expose the sandbox memory env vars in compose files and add an agent-executor healthcheck on the readiness sentinel. --- .env.example | 4 + docker-compose.dev.yml | 8 + docker-compose.local.yml | 8 + docker-compose.sandbox.yml | 9 + docker-compose.yml | 8 + .../test_agent_executor_worker_resources.py | 86 +++- tests/unit/test_agent_sandbox_cgroup.py | 432 +++++++++++++++--- tests/unit/test_agent_sandbox_config.py | 21 +- tests/unit/test_config.py | 39 ++ tests/unit/test_nsjail_seccomp.py | 50 +- tests/unit/test_stdio_probe.py | 20 +- tracecat/agent/executor_worker.py | 8 +- tracecat/agent/mcp/stdio_probe.py | 11 +- tracecat/agent/sandbox/cgroup.py | 127 ++++- tracecat/agent/sandbox/config.py | 12 +- tracecat/agent/sandbox/nsjail.py | 11 +- tracecat/config.py | 28 +- tracecat/sandbox/executor.py | 13 + 18 files changed, 750 insertions(+), 145 deletions(-) diff --git a/.env.example b/.env.example index aa2d4c110c..0948d20b9d 100644 --- a/.env.example +++ b/.env.example @@ -153,6 +153,10 @@ TRACECAT__RESULT_EXTERNALIZATION_THRESHOLD_BYTES=128000 # --- Agent filesystem persistence --- # Maximum compressed bytes retained in each worker's local agent snapshot archive cache. TRACECAT__AGENT_FS_ARCHIVE_CACHE_MAX_BYTES=10737418240 +# Per-sandbox cgroup memory limit and executor memory reserved outside activity slots. +TRACECAT__AGENT_SANDBOX_MEMORY_MB=4096 +TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED=true +TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB=4096 # --- Local registry --- # Enable this only for local-registry development. Leave disabled for normal remote/builtin registry use. diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index ccf9b41df4..a37b342961 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -394,6 +394,9 @@ services: TRACECAT__AGENT_EXECUTOR_QUEUE: ${TRACECAT__AGENT_EXECUTOR_QUEUE:-shared-agent-executor-queue} TRACECAT__EXECUTOR_QUEUE: ${TRACECAT__EXECUTOR_QUEUE:-shared-action-queue} TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES: ${TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES:-1} + TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} + TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} + TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -418,6 +421,11 @@ services: - ${TRACECAT__LOCAL_REPOSITORY_PATH}:/app/local_registry - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] + healthcheck: + test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + interval: 30s + retries: 3 + start_period: 120s depends_on: litellm: condition: service_healthy diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 2dbb96b45b..dc3311c3af 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -416,6 +416,9 @@ services: TRACECAT__AGENT_EXECUTOR_QUEUE: ${TRACECAT__AGENT_EXECUTOR_QUEUE:-shared-agent-executor-queue} TRACECAT__EXECUTOR_QUEUE: ${TRACECAT__EXECUTOR_QUEUE:-shared-action-queue} TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES: ${TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES:-1} + TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} + TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} + TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -438,6 +441,11 @@ services: - ${TRACECAT__LOCAL_REPOSITORY_PATH}:/app/local_registry - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] + healthcheck: + test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + interval: 30s + retries: 3 + start_period: 120s depends_on: litellm: condition: service_healthy diff --git a/docker-compose.sandbox.yml b/docker-compose.sandbox.yml index d978d21606..69ddf6104a 100644 --- a/docker-compose.sandbox.yml +++ b/docker-compose.sandbox.yml @@ -21,8 +21,17 @@ services: # Required for nsjail sandbox execution. cap_add: - SYS_ADMIN + environment: + TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} + TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} + TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} security_opt: - seccomp:unconfined - systempaths=unconfined devices: - /dev/net/tun:/dev/net/tun + healthcheck: + test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + interval: 30s + retries: 3 + start_period: 120s diff --git a/docker-compose.yml b/docker-compose.yml index 1cd322159f..dff9dab451 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -407,6 +407,9 @@ services: TRACECAT__AGENT_EXECUTOR_QUEUE: ${TRACECAT__AGENT_EXECUTOR_QUEUE:-shared-agent-executor-queue} TRACECAT__EXECUTOR_QUEUE: ${TRACECAT__EXECUTOR_QUEUE:-shared-action-queue} TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES: ${TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES:-1} + TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} + TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} + TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -433,6 +436,11 @@ services: - ${TRACECAT__LOCAL_REPOSITORY_PATH}:/app/local_registry - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] + healthcheck: + test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + interval: 30s + retries: 3 + start_period: 120s depends_on: litellm: condition: service_healthy diff --git a/tests/unit/test_agent_executor_worker_resources.py b/tests/unit/test_agent_executor_worker_resources.py index 190ba594de..a2bfc8c08e 100644 --- a/tests/unit/test_agent_executor_worker_resources.py +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -8,7 +8,11 @@ import pytest -from tracecat.agent.sandbox.cgroup import CgroupAvailability +from tracecat.agent.sandbox.cgroup import ( + AgentExecutorMemoryBudgetError, + CgroupAvailability, + PreparedCgroup, +) @pytest.mark.anyio @@ -39,11 +43,12 @@ async def __aexit__( def keep_concurrency( max_concurrent: int, + prepared_cgroup: PreparedCgroup, *, reserve_mb: int, sandbox_memory_mb: int, ) -> int: - del reserve_mb, sandbox_memory_mb + del prepared_cgroup, reserve_mb, sandbox_memory_mb return max_concurrent async def observe_readiness() -> None: @@ -61,7 +66,7 @@ async def observe_readiness() -> None: monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: CgroupAvailability.DISABLED, + lambda: PreparedCgroup(CgroupAvailability.DISABLED, None), ) monkeypatch.setattr( executor_worker, @@ -139,18 +144,19 @@ async def test_agent_executor_removes_stale_readiness_sentinel_on_startup( def keep_concurrency( max_concurrent: int, + prepared_cgroup: PreparedCgroup, *, reserve_mb: int, sandbox_memory_mb: int, ) -> int: - del reserve_mb, sandbox_memory_mb + del prepared_cgroup, reserve_mb, sandbox_memory_mb return max_concurrent monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: CgroupAvailability.DISABLED, + lambda: PreparedCgroup(CgroupAvailability.DISABLED, None), ) monkeypatch.setattr( executor_worker, @@ -174,3 +180,73 @@ def keep_concurrency( await executor_worker.main(shutdown_event=asyncio.Event()) assert not ready_file.exists() + + +@pytest.mark.anyio +async def test_agent_executor_memory_budget_error_propagates_from_main( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + cgroup_root = tmp_path / "cgroup" + cgroup_root.mkdir() + (cgroup_root / "memory.max").write_text(f"{4096 * 1024 * 1024}\n") + start_runtime_services = AsyncMock() + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB", + 4096, + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_SANDBOX_MEMORY_MB", + 4096, + ) + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + lambda: PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root), + ) + monkeypatch.setattr( + executor_worker, + "_start_runtime_services", + start_runtime_services, + ) + + with pytest.raises( + AgentExecutorMemoryBudgetError, + match="container_limit_mb=4096", + ): + await executor_worker.main(shutdown_event=asyncio.Event()) + + start_runtime_services.assert_not_awaited() + + +@pytest.mark.anyio +@pytest.mark.parametrize("max_concurrent", ["0", "-1"]) +async def test_agent_executor_rejects_nonpositive_concurrency( + monkeypatch: pytest.MonkeyPatch, + max_concurrent: str, +) -> None: + from tracecat.agent import executor_worker + + prepare_cgroup = Mock() + monkeypatch.setenv( + "TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", + max_concurrent, + ) + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + prepare_cgroup, + ) + + with pytest.raises( + ValueError, + match="TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", + ): + await executor_worker.main(shutdown_event=asyncio.Event()) + + prepare_cgroup.assert_not_called() diff --git a/tests/unit/test_agent_sandbox_cgroup.py b/tests/unit/test_agent_sandbox_cgroup.py index 80e131ad63..d82955270e 100644 --- a/tests/unit/test_agent_sandbox_cgroup.py +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -10,81 +10,236 @@ import tracecat.agent.sandbox.cgroup as cgroup_module from tracecat.agent.sandbox.cgroup import ( BYTES_PER_MEBIBYTE, + AgentExecutorMemoryBudgetError, CgroupAvailability, CgroupMemoryLimitKind, + PreparedCgroup, clamp_agent_executor_concurrency, + detect_cgroup_root, + get_agent_sandbox_cgroup, prepare_agent_sandbox_cgroup, read_cgroup_memory_limit, ) @pytest.fixture(autouse=True) -def reset_cgroup_availability(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - monkeypatch.setattr(cgroup_module, "_cgroup_availability", None) +def reset_prepared_cgroup(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(cgroup_module, "_prepared_cgroup", None) yield +def _write_proc_cgroup(tmp_path: Path, contents: str) -> Path: + proc_cgroup_path = tmp_path / "proc-self-cgroup" + proc_cgroup_path.write_text(contents) + return proc_cgroup_path + + def _create_fake_cgroup_root( tmp_path: Path, *, + relative_path: str = "", + controllers: str = "cpu memory pids\n", pids: str = "101\n202\n", -) -> Path: - cgroup_root = tmp_path / "cgroup" - cgroup_root.mkdir() - (cgroup_root / "cgroup.controllers").write_text("cpu memory pids\n") + memory_max: str = f"{16 * 1024**3}\n", +) -> tuple[Path, Path]: + cgroupfs = tmp_path / "cgroupfs" + cgroup_root = cgroupfs / relative_path + cgroup_root.mkdir(parents=True) + (cgroup_root / "cgroup.controllers").write_text(controllers) (cgroup_root / "cgroup.procs").write_text(pids) (cgroup_root / "cgroup.subtree_control").write_text("") - return cgroup_root + (cgroup_root / "memory.max").write_text(memory_max) + return cgroupfs, cgroup_root + + +@pytest.mark.parametrize( + ("contents", "relative_root"), + [ + pytest.param("0::/\n", "", id="private-cgroup-namespace"), + pytest.param( + "0::/kubepods.slice/kubepods-burstable.slice/" + "kubepods-burstable-podabc.slice/" + "cri-containerd-deadbeef.scope\n", + "kubepods.slice/kubepods-burstable.slice/" + "kubepods-burstable-podabc.slice/" + "cri-containerd-deadbeef.scope", + id="containerd-host-namespace", + ), + pytest.param( + "9:cpu,cpuacct:/kubepods/legacy\n" + "7:memory:/kubepods/legacy\n" + "0::/user.slice/tracecat.scope\n", + "user.slice/tracecat.scope", + id="hybrid-ignores-v1", + ), + ], +) +def test_detect_cgroup_root_from_v2_entry( + tmp_path: Path, + contents: str, + relative_root: str, +) -> None: + proc_cgroup_path = _write_proc_cgroup(tmp_path, contents) + cgroupfs = tmp_path / "cgroupfs" + + result = detect_cgroup_root(proc_cgroup_path, cgroupfs) + + assert result == cgroupfs / relative_root + + +def test_detect_cgroup_root_without_v2_entry_returns_none(tmp_path: Path) -> None: + proc_cgroup_path = _write_proc_cgroup( + tmp_path, + "9:cpu,cpuacct:/kubepods/legacy\n7:memory:/kubepods/legacy\n", + ) + + assert detect_cgroup_root(proc_cgroup_path, tmp_path / "cgroupfs") is None + + +def test_detect_cgroup_root_missing_file_returns_none(tmp_path: Path) -> None: + assert detect_cgroup_root(tmp_path / "missing", tmp_path / "cgroupfs") is None + + +def test_detect_cgroup_root_unreadable_file_returns_none( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + proc_cgroup_path = _write_proc_cgroup(tmp_path, "0::/\n") + + def deny_read( + self: Path, + encoding: str | None = None, + errors: str | None = None, + ) -> str: + del self, encoding, errors + raise PermissionError(errno.EACCES, "permission denied") + + monkeypatch.setattr(Path, "read_text", deny_read) + + assert detect_cgroup_root(proc_cgroup_path, tmp_path / "cgroupfs") is None + + +def test_detect_cgroup_root_uses_parent_when_process_is_in_main( + tmp_path: Path, +) -> None: + proc_cgroup_path = _write_proc_cgroup( + tmp_path, + "0::/kubepods.slice/pod.scope/main\n", + ) + cgroupfs = tmp_path / "cgroupfs" + + result = detect_cgroup_root(proc_cgroup_path, cgroupfs) + + assert result == cgroupfs / "kubepods.slice/pod.scope" + + +def test_get_agent_sandbox_cgroup_returns_unprepared_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + cgroup_module.config, + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", + True, + ) + + result = get_agent_sandbox_cgroup() + + assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + assert result.sandbox_mount is None def test_prepare_agent_sandbox_cgroup_returns_cached_disabled_state( tmp_path: Path, ) -> None: - cgroup_root = tmp_path / "missing" + proc_cgroup_path = tmp_path / "missing" + cgroupfs = tmp_path / "cgroupfs" - assert ( - prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=False) - is CgroupAvailability.DISABLED + first = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=False, ) - assert ( - prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) - is CgroupAvailability.DISABLED + second = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, ) + assert first == PreparedCgroup(CgroupAvailability.DISABLED, None) + assert second is first + assert get_agent_sandbox_cgroup() is first + -def test_prepare_agent_sandbox_cgroup_happy_path(tmp_path: Path) -> None: - cgroup_root = _create_fake_cgroup_root(tmp_path) +def test_prepare_agent_sandbox_cgroup_beneath_detected_root( + tmp_path: Path, +) -> None: + relative_root = "kubepods.slice/pod.scope/container.scope" + cgroupfs, cgroup_root = _create_fake_cgroup_root( + tmp_path, + relative_path=relative_root, + ) + proc_cgroup_path = _write_proc_cgroup(tmp_path, f"0::/{relative_root}\n") - result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + result = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) - assert result is CgroupAvailability.AVAILABLE + assert result == PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + assert result.sandbox_mount == cgroup_root assert (cgroup_root / "main" / "cgroup.procs").read_text() == "202\n" assert (cgroup_root / "cgroup.subtree_control").read_text() == "+memory\n" assert not list(cgroup_root.glob("tracecat-agent-probe-*")) -def test_prepare_agent_sandbox_cgroup_without_v2_controllers_warns_once( +def test_prepare_agent_sandbox_cgroup_detection_failure_warns_once( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + mock_logger = Mock() + monkeypatch.setattr(cgroup_module, "logger", mock_logger) + + result = prepare_agent_sandbox_cgroup( + proc_cgroup_path=tmp_path / "missing", + cgroupfs=tmp_path / "cgroupfs", + enabled=True, + ) + + assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args.kwargs["step"] == "detect cgroup v2 root" + + +def test_prepare_agent_sandbox_cgroup_requires_memory_controller( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - cgroup_root = tmp_path / "cgroup" - cgroup_root.mkdir() + cgroupfs, cgroup_root = _create_fake_cgroup_root( + tmp_path, + controllers="cpu pids\n", + ) + proc_cgroup_path = _write_proc_cgroup(tmp_path, "0::/\n") mock_logger = Mock() monkeypatch.setattr(cgroup_module, "logger", mock_logger) - result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + result = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) - assert result is CgroupAvailability.UNAVAILABLE + assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root) mock_logger.warning.assert_called_once() - assert mock_logger.warning.call_args.kwargs["step"] == "read cgroup v2 controllers" - assert mock_logger.warning.call_args.kwargs["errno"] == errno.ENOENT + assert mock_logger.warning.call_args.kwargs["step"] == "verify memory controller" -def test_prepare_agent_sandbox_cgroup_subtree_permission_error_falls_back( +def test_prepare_agent_sandbox_cgroup_retains_root_on_permission_error( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - cgroup_root = _create_fake_cgroup_root(tmp_path) + cgroupfs, cgroup_root = _create_fake_cgroup_root(tmp_path) + proc_cgroup_path = _write_proc_cgroup(tmp_path, "0::/\n") mock_logger = Mock() monkeypatch.setattr(cgroup_module, "logger", mock_logger) write_cgroup_file = cgroup_module._write_cgroup_file @@ -100,9 +255,15 @@ def deny_subtree_control(path: Path, value: str) -> None: deny_subtree_control, ) - result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + result = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) - assert result is CgroupAvailability.UNAVAILABLE + assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root) + assert result.sandbox_mount is None + assert read_cgroup_memory_limit(result.root).kind is CgroupMemoryLimitKind.LIMITED mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.kwargs["step"] == "enable memory controller" assert mock_logger.warning.call_args.kwargs["errno"] == errno.EACCES @@ -112,7 +273,8 @@ def test_prepare_agent_sandbox_cgroup_tolerates_pid_vanishing_during_move( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - cgroup_root = _create_fake_cgroup_root(tmp_path) + cgroupfs, cgroup_root = _create_fake_cgroup_root(tmp_path) + proc_cgroup_path = _write_proc_cgroup(tmp_path, "0::/\n") mock_logger = Mock() monkeypatch.setattr(cgroup_module, "logger", mock_logger) write_cgroup_file = cgroup_module._write_cgroup_file @@ -124,9 +286,13 @@ def vanish_first_pid(path: Path, value: str) -> None: monkeypatch.setattr(cgroup_module, "_write_cgroup_file", vanish_first_pid) - result = prepare_agent_sandbox_cgroup(cgroup_root=cgroup_root, enabled=True) + result = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) - assert result is CgroupAvailability.AVAILABLE + assert result.availability is CgroupAvailability.AVAILABLE assert (cgroup_root / "main" / "cgroup.procs").read_text() == "202\n" mock_logger.warning.assert_not_called() @@ -143,87 +309,203 @@ def vanish_first_pid(path: Path, value: str) -> None: ), ], ) -def test_read_cgroup_memory_limit_parses_valid_values( +def test_read_cgroup_memory_limit_from_detected_root( tmp_path: Path, contents: str, expected_kind: CgroupMemoryLimitKind, expected_bytes: int | None, ) -> None: - memory_max_path = tmp_path / "memory.max" - memory_max_path.write_text(contents) + cgroupfs, cgroup_root = _create_fake_cgroup_root( + tmp_path, + memory_max=contents, + ) + proc_cgroup_path = _write_proc_cgroup(tmp_path, "0::/\n") + detected_root = detect_cgroup_root(proc_cgroup_path, cgroupfs) - result = read_cgroup_memory_limit(memory_max_path) + result = read_cgroup_memory_limit(detected_root) + assert detected_root == cgroup_root assert result.kind is expected_kind assert result.limit_bytes == expected_bytes -def test_read_cgroup_memory_limit_missing_is_unavailable(tmp_path: Path) -> None: - result = read_cgroup_memory_limit(tmp_path / "memory.max") +def test_read_cgroup_memory_limit_without_root_is_unavailable() -> None: + result = read_cgroup_memory_limit(None) assert result.kind is CgroupMemoryLimitKind.UNAVAILABLE assert result.limit_bytes is None -def test_read_cgroup_memory_limit_garbage_warns_and_is_unavailable( +@pytest.mark.parametrize("contents", ["not-a-limit\n", "-1\n"]) +def test_read_cgroup_memory_limit_invalid_value_warns_and_is_unavailable( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + contents: str, ) -> None: - memory_max_path = tmp_path / "memory.max" - memory_max_path.write_text("not-a-limit\n") + _, cgroup_root = _create_fake_cgroup_root(tmp_path, memory_max=contents) mock_logger = Mock() monkeypatch.setattr(cgroup_module, "logger", mock_logger) - result = read_cgroup_memory_limit(memory_max_path) + result = read_cgroup_memory_limit(cgroup_root) assert result.kind is CgroupMemoryLimitKind.UNAVAILABLE assert result.limit_bytes is None mock_logger.warning.assert_called_once() +def test_prepare_move_does_not_change_cached_root_used_for_budget( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + relative_root = "kubepods.slice/pod.scope/container.scope" + cgroupfs, cgroup_root = _create_fake_cgroup_root( + tmp_path, + relative_path=relative_root, + pids="101\n", + memory_max=f"{12 * 1024**3}\n", + ) + proc_cgroup_path = _write_proc_cgroup(tmp_path, f"0::/{relative_root}\n") + write_cgroup_file = cgroup_module._write_cgroup_file + + def move_worker_into_main(path: Path, value: str) -> None: + write_cgroup_file(path, value) + if path == cgroup_root / "main" / "cgroup.procs": + proc_cgroup_path.write_text(f"0::/{relative_root}/main\n") + + monkeypatch.setattr( + cgroup_module, + "_write_cgroup_file", + move_worker_into_main, + ) + + prepared_cgroup = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + (cgroup_root / "main" / "memory.max").write_text("max\n") + + assert detect_cgroup_root(proc_cgroup_path, cgroupfs) == cgroup_root + assert prepared_cgroup.root == cgroup_root + assert ( + clamp_agent_executor_concurrency( + 10, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + == 2 + ) + + @pytest.mark.parametrize( - ( - "memory_max_contents", - "wanted", - "expected", - "expected_error", - ), + ("container_limit_mb", "reserve_mb"), [ - pytest.param(f"{16 * 1024**3}\n", 10, 3, True, id="over-budget"), - pytest.param("max\n", 10, 10, False, id="unlimited"), - pytest.param(None, 10, 10, False, id="missing"), - pytest.param(f"{4 * 1024**3}\n", 10, 1, True, id="allowed-zero"), - pytest.param(f"{16 * 1024**3}\n", 2, 2, False, id="under-budget"), + pytest.param(4096, 4096, id="exactly-zero-slots"), + pytest.param(4096, 8192, id="reserve-exceeds-limit"), ], ) -def test_clamp_agent_executor_concurrency( +def test_clamp_agent_executor_concurrency_fails_when_no_sandbox_fits( + tmp_path: Path, + container_limit_mb: int, + reserve_mb: int, +) -> None: + _, cgroup_root = _create_fake_cgroup_root( + tmp_path, + memory_max=f"{container_limit_mb * BYTES_PER_MEBIBYTE}\n", + ) + prepared_cgroup = PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + ) + + with pytest.raises(AgentExecutorMemoryBudgetError) as exc_info: + clamp_agent_executor_concurrency( + 10, + prepared_cgroup, + reserve_mb=reserve_mb, + sandbox_memory_mb=4096, + ) + + message = str(exc_info.value) + assert f"container_limit_mb={container_limit_mb}" in message + assert f"reserve_mb={reserve_mb}" in message + assert "sandbox_memory_mb=4096" in message + assert "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB" in message + assert "TRACECAT__AGENT_SANDBOX_MEMORY_MB" in message + assert "container memory limit" in message + + +def test_clamp_agent_executor_concurrency_clamps_and_logs( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, - memory_max_contents: str | None, - wanted: int, - expected: int, - expected_error: bool, ) -> None: - memory_max_path = tmp_path / "memory.max" - if memory_max_contents is not None: - memory_max_path.write_text(memory_max_contents) + _, cgroup_root = _create_fake_cgroup_root(tmp_path) + prepared_cgroup = PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) mock_logger = Mock() monkeypatch.setattr(cgroup_module, "logger", mock_logger) result = clamp_agent_executor_concurrency( - wanted, + 10, + prepared_cgroup, reserve_mb=4096, sandbox_memory_mb=4096, - memory_max_path=memory_max_path, - ) - - assert result == expected - assert mock_logger.error.called is expected_error - if memory_max_contents == f"{16 * 1024**3}\n" and wanted == 10: - assert mock_logger.error.call_args.kwargs == { - "container_limit_mb": 16 * 1024**3 // BYTES_PER_MEBIBYTE, - "reserve_mb": 4096, - "sandbox_memory_mb": 4096, - "configured_max_concurrent_activities": 10, - "clamped_max_concurrent_activities": 3, - } + ) + + assert result == 3 + mock_logger.error.assert_called_once() + assert mock_logger.error.call_args.kwargs == { + "container_limit_mb": 16 * 1024**3 // BYTES_PER_MEBIBYTE, + "reserve_mb": 4096, + "sandbox_memory_mb": 4096, + "configured_max_concurrent_activities": 10, + "clamped_max_concurrent_activities": 3, + } + + +def test_clamp_agent_executor_concurrency_keeps_under_budget_value( + tmp_path: Path, +) -> None: + _, cgroup_root = _create_fake_cgroup_root(tmp_path) + prepared_cgroup = PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + + assert ( + clamp_agent_executor_concurrency( + 2, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + == 2 + ) + + +def test_clamp_agent_executor_concurrency_passes_through_unlimited( + tmp_path: Path, +) -> None: + _, cgroup_root = _create_fake_cgroup_root(tmp_path, memory_max="max\n") + prepared_cgroup = PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + + assert ( + clamp_agent_executor_concurrency( + 10, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + == 10 + ) + + +def test_clamp_agent_executor_concurrency_passes_through_unavailable() -> None: + prepared_cgroup = PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + + assert ( + clamp_agent_executor_concurrency( + 10, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + == 10 + ) diff --git a/tests/unit/test_agent_sandbox_config.py b/tests/unit/test_agent_sandbox_config.py index 176ad0cdc7..ce08ca1506 100644 --- a/tests/unit/test_agent_sandbox_config.py +++ b/tests/unit/test_agent_sandbox_config.py @@ -2,8 +2,6 @@ from pathlib import Path -import pytest - from tracecat.agent.sandbox.config import ( AgentResourceLimits, AgentSandboxConfig, @@ -136,27 +134,16 @@ def test_build_agent_nsjail_config_adds_available_cgroup_v2_memory_limit() -> No ), site_packages_dir=Path("/app/.venv/lib/python3.12/site-packages"), llm_socket_path=Path("/tmp/agent-job/sockets/llm.sock"), - cgroup_v2_enabled=True, - cgroup_v2_available=True, + cgroup_mount=Path("/sys/fs/cgroup/kubepods.slice/pod.scope"), ) assert "use_cgroupv2: true" in config_text - assert 'cgroupv2_mount: "/sys/fs/cgroup"' in config_text + assert 'cgroupv2_mount: "/sys/fs/cgroup/kubepods.slice/pod.scope"' in config_text assert f"cgroup_mem_max: {3072 * 1024 * 1024}" in config_text assert "cgroup_mem_swap_max: 0" in config_text -@pytest.mark.parametrize( - ("cgroup_v2_enabled", "cgroup_v2_available"), - [ - pytest.param(False, True, id="disabled"), - pytest.param(True, False, id="unavailable"), - ], -) -def test_build_agent_nsjail_config_omits_inactive_cgroup_v2_memory_limit( - cgroup_v2_enabled: bool, - cgroup_v2_available: bool, -) -> None: +def test_build_agent_nsjail_config_omits_cgroup_v2_memory_limit_by_default() -> None: config_text = build_agent_nsjail_config( rootfs=Path("/var/lib/tracecat/sandbox-rootfs"), job_dir=Path("/tmp/agent-job"), @@ -164,8 +151,6 @@ def test_build_agent_nsjail_config_omits_inactive_cgroup_v2_memory_limit( config=AgentSandboxConfig(), site_packages_dir=Path("/app/.venv/lib/python3.12/site-packages"), llm_socket_path=Path("/tmp/agent-job/sockets/llm.sock"), - cgroup_v2_enabled=cgroup_v2_enabled, - cgroup_v2_available=cgroup_v2_available, ) assert "use_cgroupv2:" not in config_text diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 3ee60d7d1f..97a397d6c0 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -16,6 +16,7 @@ REPO_ROOT / "docker-compose.yml", REPO_ROOT / "docker-compose.dev.yml", REPO_ROOT / "docker-compose.local.yml", + REPO_ROOT / "docker-compose.sandbox.yml", ) ENV_EXAMPLE_FILES = (REPO_ROOT / ".env.example",) DEPLOYMENT_ENV_FILES = (*COMPOSE_ENV_FILES, *ENV_EXAMPLE_FILES) @@ -220,6 +221,44 @@ def test_action_gateway_socket_uses_default_for_empty_string( importlib.reload(tracecat_config) +@pytest.mark.parametrize( + ("env_var", "offending_value", "minimum"), + [ + pytest.param( + "TRACECAT__AGENT_SANDBOX_MEMORY_MB", + "0", + 1, + id="sandbox-memory", + ), + pytest.param( + "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB", + "-1", + 0, + id="executor-reserve", + ), + ], +) +def test_agent_memory_config_rejects_values_below_minimum( + monkeypatch: pytest.MonkeyPatch, + env_var: str, + offending_value: str, + minimum: int, +) -> None: + try: + with monkeypatch.context() as env: + env.setenv(env_var, offending_value) + + with pytest.raises(ValueError) as exc_info: + importlib.reload(tracecat_config) + + message = str(exc_info.value) + assert env_var in message + assert f"at least {minimum}" in message + assert f"got {offending_value}" in message + finally: + importlib.reload(tracecat_config) + + def test_bound_env_rejects_invalid_numeric_value( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/test_nsjail_seccomp.py b/tests/unit/test_nsjail_seccomp.py index 66f0f91460..1c09b8a1f7 100644 --- a/tests/unit/test_nsjail_seccomp.py +++ b/tests/unit/test_nsjail_seccomp.py @@ -8,7 +8,7 @@ from tracecat.executor.backends.pool import WorkerPool from tracecat.sandbox.executor import ActionSandboxConfig, NsjailExecutor from tracecat.sandbox.seccomp import build_untrusted_seccomp_policy -from tracecat.sandbox.types import SandboxConfig +from tracecat.sandbox.types import ResourceLimits, SandboxConfig _EXPECTED_BLOCKED_SYSCALLS = ( "ptrace", @@ -52,6 +52,32 @@ def test_python_sandbox_config_includes_seccomp_policy(tmp_path: Path): ) _assert_seccomp_config(config_text) + assert "use_cgroupv2:" not in config_text + assert "cgroupv2_mount:" not in config_text + assert "cgroup_mem_max:" not in config_text + assert "cgroup_mem_swap_max:" not in config_text + + +def test_stdio_probe_sandbox_config_includes_cgroup_memory_limit( + tmp_path: Path, +) -> None: + cgroup_mount = tmp_path / "cgroup" / "container.scope" + executor = NsjailExecutor( + rootfs_path=str(tmp_path / "rootfs"), + cgroup_mount=cgroup_mount, + ) + + config_text = executor._build_config( + job_dir=tmp_path / "job", + phase="execute", + config=SandboxConfig(resources=ResourceLimits(memory_mb=512)), + script_name="probe.py", + ) + + assert "use_cgroupv2: true" in config_text + assert f'cgroupv2_mount: "{cgroup_mount}"' in config_text + assert f"cgroup_mem_max: {512 * 1024 * 1024}" in config_text + assert "cgroup_mem_swap_max: 0" in config_text def test_action_sandbox_config_includes_seccomp_policy(tmp_path: Path): @@ -69,6 +95,28 @@ def test_action_sandbox_config_includes_seccomp_policy(tmp_path: Path): _assert_seccomp_config(config_text) +def test_action_sandbox_config_ignores_general_executor_cgroup_mount( + tmp_path: Path, +) -> None: + executor = NsjailExecutor( + rootfs_path=str(tmp_path / "rootfs"), + cgroup_mount=tmp_path / "cgroup", + ) + + config_text = executor._build_action_config( + job_dir=tmp_path / "job", + config=ActionSandboxConfig( + registry_paths=[tmp_path / "registry"], + tracecat_app_dir=tmp_path / "app", + ), + ) + + assert "use_cgroupv2:" not in config_text + assert "cgroupv2_mount:" not in config_text + assert "cgroup_mem_max:" not in config_text + assert "cgroup_mem_swap_max:" not in config_text + + def test_agent_sandbox_config_includes_seccomp_policy(tmp_path: Path): """Agent sandbox configs should emit the shared seccomp policy.""" config_text = build_agent_nsjail_config( diff --git a/tests/unit/test_stdio_probe.py b/tests/unit/test_stdio_probe.py index 57429cd111..55496a41f2 100644 --- a/tests/unit/test_stdio_probe.py +++ b/tests/unit/test_stdio_probe.py @@ -14,6 +14,7 @@ probe_stdio_mcp_tools_in_sandbox, sanitize_stdio_probe_error, ) +from tracecat.agent.sandbox.cgroup import CgroupAvailability, PreparedCgroup from tracecat.sandbox.exceptions import SandboxTimeoutError from tracecat.sandbox.types import SandboxErrorCode, SandboxResult @@ -171,7 +172,9 @@ async def test_probe_falls_back_to_pid_isolation_without_nsjail() -> None: @pytest.mark.anyio -async def test_probe_runs_in_sandbox_when_nsjail_available() -> None: +async def test_probe_runs_in_sandbox_with_cgroup_slot_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: """With nsjail available, the probe executes inside the sandbox.""" sandbox_result = MagicMock( success=True, @@ -184,6 +187,11 @@ async def test_probe_runs_in_sandbox_when_nsjail_available() -> None: stderr="", ) executor = MagicMock(execute=AsyncMock(return_value=sandbox_result)) + cgroup_mount = Path("/sys/fs/cgroup/kubepods.slice/pod.scope") + monkeypatch.setattr( + "tracecat.agent.mcp.stdio_probe.config.TRACECAT__AGENT_SANDBOX_MEMORY_MB", + 512, + ) with ( patch( @@ -193,6 +201,13 @@ async def test_probe_runs_in_sandbox_when_nsjail_available() -> None: patch( "tracecat.agent.mcp.stdio_probe.NsjailExecutor", return_value=executor, + ) as executor_cls, + patch( + "tracecat.agent.mcp.stdio_probe.get_agent_sandbox_cgroup", + return_value=PreparedCgroup( + CgroupAvailability.AVAILABLE, + cgroup_mount, + ), ), ): result = await probe_stdio_mcp_tools_in_sandbox( @@ -204,7 +219,10 @@ async def test_probe_runs_in_sandbox_when_nsjail_available() -> None: assert result.success is True assert [tool.name for tool in result.tools] == ["list_alerts"] + executor_cls.assert_called_once_with(cgroup_mount=cgroup_mount) executor.execute.assert_awaited_once() + sandbox_config = executor.execute.await_args.args[1] + assert sandbox_config.resources.memory_mb == 512 @pytest.mark.anyio diff --git a/tracecat/agent/executor_worker.py b/tracecat/agent/executor_worker.py index 3e3a72320b..1870c6ecd1 100644 --- a/tracecat/agent/executor_worker.py +++ b/tracecat/agent/executor_worker.py @@ -117,12 +117,18 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: max_concurrent = int( os.environ.get("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES") or 1 ) + if max_concurrent < 1: + raise ValueError( + "TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES must be at " + f"least 1 (got {max_concurrent})" + ) threadpool_max_workers = int( os.environ.get("TEMPORAL__THREADPOOL_MAX_WORKERS") or 100 ) - prepare_agent_sandbox_cgroup() + prepared_cgroup = prepare_agent_sandbox_cgroup() max_concurrent = clamp_agent_executor_concurrency( max_concurrent, + prepared_cgroup, reserve_mb=config.TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB, sandbox_memory_mb=config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, ) diff --git a/tracecat/agent/mcp/stdio_probe.py b/tracecat/agent/mcp/stdio_probe.py index d5134c3c70..b26cbd0a33 100644 --- a/tracecat/agent/mcp/stdio_probe.py +++ b/tracecat/agent/mcp/stdio_probe.py @@ -14,6 +14,7 @@ import orjson +from tracecat import config from tracecat.agent.mcp.stdio_probe_types import ( MCP_STDIO_PERSIST_ACTIVITY_NAME, MCP_STDIO_PROBE_ACTIVITY_NAME, @@ -28,6 +29,7 @@ sanitize_stdio_probe_error, ) from tracecat.agent.mcp.utils import STDIO_MCP_TOOL_NAME_RE +from tracecat.agent.sandbox.cgroup import get_agent_sandbox_cgroup from tracecat.integrations.schemas import MCPToolSummary from tracecat.logger import logger from tracecat.sandbox.exceptions import SandboxTimeoutError @@ -364,13 +366,18 @@ async def probe_stdio_mcp_tools_in_sandbox( (job_dir / "input.json").write_bytes(orjson.dumps(payload)) if is_nsjail_available(): - sandbox = NsjailExecutor() + sandbox = NsjailExecutor( + cgroup_mount=get_agent_sandbox_cgroup().sandbox_mount + ) result = await sandbox.execute( job_dir, SandboxConfig( network_enabled=True, resources=ResourceLimits( - memory_mb=1024, + memory_mb=min( + 1024, + config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, + ), cpu_seconds=hard_timeout_seconds, max_open_files=512, max_processes=128, diff --git a/tracecat/agent/sandbox/cgroup.py b/tracecat/agent/sandbox/cgroup.py index 659121d764..8bf628f32e 100644 --- a/tracecat/agent/sandbox/cgroup.py +++ b/tracecat/agent/sandbox/cgroup.py @@ -11,8 +11,8 @@ from tracecat import config from tracecat.logger import logger -CGROUP_ROOT = Path("/sys/fs/cgroup") -CGROUP_MEMORY_MAX_PATH = CGROUP_ROOT / "memory.max" +PROC_SELF_CGROUP = Path("/proc/self/cgroup") +CGROUPFS = Path("/sys/fs/cgroup") BYTES_PER_MEBIBYTE = 1024 * 1024 @@ -32,6 +32,25 @@ class CgroupMemoryLimitKind(StrEnum): UNAVAILABLE = "unavailable" +class AgentExecutorMemoryBudgetError(RuntimeError): + """Raised when the container memory budget cannot fit one sandbox.""" + + +@dataclass(frozen=True, slots=True) +class PreparedCgroup: + """Process-wide agent sandbox cgroup preparation result.""" + + availability: CgroupAvailability + root: Path | None + + @property + def sandbox_mount(self) -> Path | None: + """Return the nsjail cgroup mount when preparation succeeded.""" + if self.availability is CgroupAvailability.AVAILABLE: + return self.root + return None + + @dataclass(frozen=True, slots=True) class CgroupMemoryLimit: """Parsed cgroup v2 container memory limit.""" @@ -40,7 +59,32 @@ class CgroupMemoryLimit: limit_bytes: int | None = None -_cgroup_availability: CgroupAvailability | None = None +_prepared_cgroup: PreparedCgroup | None = None + + +def detect_cgroup_root( + proc_cgroup_path: Path = PROC_SELF_CGROUP, + cgroupfs: Path = CGROUPFS, +) -> Path | None: + """Detect the executor's cgroup v2 directory from /proc/self/cgroup.""" + try: + lines = proc_cgroup_path.read_text().splitlines() + except OSError: + return None + + for line in lines: + fields = line.split(":", maxsplit=2) + if len(fields) != 3: + continue + hierarchy_id, controllers, cgroup_path = fields + if hierarchy_id != "0" or controllers: + continue + + root = cgroupfs if cgroup_path == "/" else cgroupfs / cgroup_path.lstrip("/") + if root.name == "main": + return root.parent + return root + return None def _write_cgroup_file(path: Path, value: str) -> None: @@ -84,27 +128,43 @@ def _warn_cgroup_unavailable( def prepare_agent_sandbox_cgroup( *, - cgroup_root: Path = CGROUP_ROOT, + proc_cgroup_path: Path = PROC_SELF_CGROUP, + cgroupfs: Path = CGROUPFS, enabled: bool = config.TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, -) -> CgroupAvailability: +) -> PreparedCgroup: """Prepare the cgroup v2 root for nsjail child memory cgroups. The result is cached process-wide. Setup failures are reduced to one warning and an unavailable result so agent sandbox launches can continue without cgroup limits. """ - global _cgroup_availability + global _prepared_cgroup - if _cgroup_availability is not None: - return _cgroup_availability + if _prepared_cgroup is not None: + return _prepared_cgroup if not enabled: - _cgroup_availability = CgroupAvailability.DISABLED - return _cgroup_availability + _prepared_cgroup = PreparedCgroup(CgroupAvailability.DISABLED, None) + return _prepared_cgroup + + cgroup_root = detect_cgroup_root( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + ) + if cgroup_root is None: + _warn_cgroup_unavailable( + step="detect cgroup v2 root", + error=ValueError(f"no cgroup v2 entry found in {proc_cgroup_path}"), + ) + _prepared_cgroup = PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + return _prepared_cgroup probe_cgroup: Path | None = None step = "read cgroup v2 controllers" try: - (cgroup_root / "cgroup.controllers").read_text() + controllers = (cgroup_root / "cgroup.controllers").read_text().split() + step = "verify memory controller" + if "memory" not in controllers: + raise ValueError("memory controller is not available") step = "create main cgroup" main_cgroup = cgroup_root / "main" @@ -144,22 +204,33 @@ def prepare_agent_sandbox_cgroup( except OSError: pass _warn_cgroup_unavailable(step=step, error=exc) - _cgroup_availability = CgroupAvailability.UNAVAILABLE - return _cgroup_availability + _prepared_cgroup = PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + ) + return _prepared_cgroup - _cgroup_availability = CgroupAvailability.AVAILABLE - return _cgroup_availability + _prepared_cgroup = PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + return _prepared_cgroup -def agent_sandbox_cgroup_is_available() -> bool: - """Return whether the startup probe enabled agent sandbox cgroups.""" - return _cgroup_availability is CgroupAvailability.AVAILABLE +def get_agent_sandbox_cgroup() -> PreparedCgroup: + """Return the cached process-wide cgroup preparation result.""" + if _prepared_cgroup is not None: + return _prepared_cgroup + if config.TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: + return PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + return PreparedCgroup(CgroupAvailability.DISABLED, None) def read_cgroup_memory_limit( - memory_max_path: Path = CGROUP_MEMORY_MAX_PATH, + cgroup_root: Path | None, ) -> CgroupMemoryLimit: """Read and parse the cgroup v2 container memory limit.""" + if cgroup_root is None: + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + + memory_max_path = cgroup_root / "memory.max" try: raw_limit = memory_max_path.read_text().strip() except FileNotFoundError: @@ -198,13 +269,13 @@ def read_cgroup_memory_limit( def clamp_agent_executor_concurrency( max_concurrent: int, + prepared_cgroup: PreparedCgroup, *, reserve_mb: int, sandbox_memory_mb: int, - memory_max_path: Path = CGROUP_MEMORY_MAX_PATH, ) -> int: """Clamp agent executor concurrency to the container memory budget.""" - memory_limit = read_cgroup_memory_limit(memory_max_path) + memory_limit = read_cgroup_memory_limit(prepared_cgroup.root) if memory_limit.kind is CgroupMemoryLimitKind.UNLIMITED: logger.debug( "Container memory is unlimited; skipping worker memory-budget validation" @@ -219,16 +290,24 @@ def clamp_agent_executor_concurrency( limit_mb = limit_bytes // BYTES_PER_MEBIBYTE allowed = (limit_mb - reserve_mb) // sandbox_memory_mb + if allowed < 1: + raise AgentExecutorMemoryBudgetError( + "Agent executor memory budget cannot fit one sandbox: " + f"container_limit_mb={limit_mb}, reserve_mb={reserve_mb}, " + f"sandbox_memory_mb={sandbox_memory_mb}. Increase the container " + "memory limit or reduce " + "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB or " + "TRACECAT__AGENT_SANDBOX_MEMORY_MB." + ) if max_concurrent <= allowed: return max_concurrent - clamped = max(allowed, 1) logger.error( "Agent executor concurrency exceeds the container memory budget; clamping", container_limit_mb=limit_mb, reserve_mb=reserve_mb, sandbox_memory_mb=sandbox_memory_mb, configured_max_concurrent_activities=max_concurrent, - clamped_max_concurrent_activities=clamped, + clamped_max_concurrent_activities=allowed, ) - return clamped + return allowed diff --git a/tracecat/agent/sandbox/config.py b/tracecat/agent/sandbox/config.py index 7cfc8f5a09..655404aa8c 100644 --- a/tracecat/agent/sandbox/config.py +++ b/tracecat/agent/sandbox/config.py @@ -222,8 +222,7 @@ def build_agent_nsjail_config( session_work_dir: Path | None = None, enable_internet_access: bool = False, skills_dir: Path | None = None, - cgroup_v2_enabled: bool = False, - cgroup_v2_available: bool = False, + cgroup_mount: Path | None = None, ) -> str: """Build nsjail protobuf config for agent runtime execution. @@ -247,8 +246,7 @@ def build_agent_nsjail_config( outbound internet access. Default is False (network isolated with private loopback only). skills_dir: Optional host path containing staged workspace skills. - cgroup_v2_enabled: Whether agent sandbox cgroup limits are configured. - cgroup_v2_available: Whether the process-wide startup probe succeeded. + cgroup_mount: Optional cgroup v2 directory for nsjail child cgroups. Returns: nsjail protobuf configuration as a string. @@ -276,6 +274,8 @@ def build_agent_nsjail_config( _validate_path(mcp_socket_path, "mcp_socket_path") if skills_dir is not None: _validate_path(skills_dir, "skills_dir") + if cgroup_mount is not None: + _validate_path(cgroup_mount, "cgroup_mount") # Derive control socket path from socket_dir using well-known name when enabled. resolved_control_socket_path: Path | None @@ -476,11 +476,11 @@ def build_agent_nsjail_config( f"time_limit: {config.resources.timeout_seconds}", ] ) - if cgroup_v2_enabled and cgroup_v2_available: + if cgroup_mount is not None: lines.extend( [ "use_cgroupv2: true", - 'cgroupv2_mount: "/sys/fs/cgroup"', + f'cgroupv2_mount: "{cgroup_mount}"', f"cgroup_mem_max: {config.resources.memory_mb * 1024 * 1024}", "cgroup_mem_swap_max: 0", ] diff --git a/tracecat/agent/sandbox/nsjail.py b/tracecat/agent/sandbox/nsjail.py index a04abf8700..5a26aab355 100644 --- a/tracecat/agent/sandbox/nsjail.py +++ b/tracecat/agent/sandbox/nsjail.py @@ -46,18 +46,14 @@ JAILED_AGENT_JOB_DIR, JAILED_AGENT_WORK_DIR, ) -from tracecat.agent.sandbox.cgroup import agent_sandbox_cgroup_is_available +from tracecat.agent.sandbox.cgroup import get_agent_sandbox_cgroup from tracecat.agent.sandbox.config import ( JAILED_SHIM_ENTRYPOINT_PATH, AgentSandboxConfig, build_agent_env_map, build_agent_nsjail_config, ) -from tracecat.config import ( - TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, - TRACECAT__SANDBOX_NSJAIL_PATH, - TRACECAT__SANDBOX_ROOTFS_PATH, -) +from tracecat.config import TRACECAT__SANDBOX_NSJAIL_PATH, TRACECAT__SANDBOX_ROOTFS_PATH from tracecat.logger import logger BROKER_SHIM_SCRIPT_NAME = Path(JAILED_SHIM_ENTRYPOINT_PATH).name @@ -441,8 +437,7 @@ async def _spawn_nsjail_runtime( session_work_dir=session_work_dir, enable_internet_access=enable_internet_access, skills_dir=skills_dir, - cgroup_v2_enabled=TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, - cgroup_v2_available=agent_sandbox_cgroup_is_available(), + cgroup_mount=get_agent_sandbox_cgroup().sandbox_mount, ) # Write config to job directory diff --git a/tracecat/config.py b/tracecat/config.py index a7184e7062..3c819b143f 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -73,6 +73,22 @@ def env_bool(var: str, *, default: bool) -> bool: raise ValueError(f"{var} must be a boolean value (got {raw_value!r})") +def _env_int(var: str, default: int, *, min_value: int) -> int: + """Read an integer env var and reject values below the minimum.""" + raw_value = os.environ.get(var) + if raw_value is None or not raw_value.strip(): + value = default + else: + try: + value = int(raw_value) + except ValueError as exc: + raise ValueError(f"{var} must be an integer (got {raw_value!r})") from exc + + if value < min_value: + raise ValueError(f"{var} must be at least {min_value} (got {value})") + return value + + class RLSMode(StrEnum): """Runtime mode for application-level RLS session behavior.""" @@ -729,8 +745,10 @@ def _parse_auth_types() -> set[AuthType]: sandbox. """ -TRACECAT__AGENT_SANDBOX_MEMORY_MB = int( - os.environ.get("TRACECAT__AGENT_SANDBOX_MEMORY_MB") or 4096 +TRACECAT__AGENT_SANDBOX_MEMORY_MB = _env_int( + "TRACECAT__AGENT_SANDBOX_MEMORY_MB", + 4096, + min_value=1, ) """Default memory limit for agent sandbox execution in megabytes (4 GiB).""" @@ -739,8 +757,10 @@ def _parse_auth_types() -> set[AuthType]: ) """Enable cgroup v2 memory limits for agent sandboxes when available.""" -TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB = int( - os.environ.get("TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB") or 4096 +TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB = _env_int( + "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB", + 4096, + min_value=0, ) """Memory reserved for the agent executor worker and shared services.""" diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index 5ad8b8c4d1..6826b5df23 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -179,10 +179,12 @@ def __init__( nsjail_path: str = TRACECAT__SANDBOX_NSJAIL_PATH, rootfs_path: str = TRACECAT__SANDBOX_ROOTFS_PATH, cache_dir: str = TRACECAT__SANDBOX_CACHE_DIR, + cgroup_mount: Path | None = None, ): self.nsjail_path = Path(nsjail_path) self.rootfs = Path(rootfs_path) self.cache_dir = Path(cache_dir) + self.cgroup_mount = cgroup_mount self.package_cache = self.cache_dir / "packages" self.uv_cache = self.cache_dir / "uv-cache" @@ -212,6 +214,8 @@ def _build_config( # Validate inputs to prevent injection into protobuf config _validate_path(job_dir, "job_dir") _validate_path(self.rootfs, "rootfs") + if self.cgroup_mount is not None: + _validate_path(self.cgroup_mount, "cgroup_mount") for i, python_path_dir in enumerate(config.python_path_dirs): _validate_path(python_path_dir, f"python_path_dir_{i}") if cache_key: @@ -356,6 +360,15 @@ def _build_config( f"time_limit: {config.resources.timeout_seconds}", ] ) + if self.cgroup_mount is not None: + lines.extend( + [ + "use_cgroupv2: true", + f'cgroupv2_mount: "{self.cgroup_mount}"', + f"cgroup_mem_max: {config.resources.memory_mb * 1024 * 1024}", + "cgroup_mem_swap_max: 0", + ] + ) # Execution settings - script path must be in exec_bin for config file mode script_path = f"/work/{script_name}" From 7a037f5ce02577d4fe5db698dc58492cf1aea6cf Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:09:27 -0400 Subject: [PATCH 05/14] fix(agent): honor ready-file override in compose healthchecks Interpolate TRACECAT__AGENT_EXECUTOR_READY_FILE into both the agent-executor environment and the healthcheck test so an operator override keeps the worker and healthcheck pointed at the same sentinel path. --- .env.example | 2 ++ docker-compose.dev.yml | 3 ++- docker-compose.local.yml | 3 ++- docker-compose.sandbox.yml | 3 ++- docker-compose.yml | 3 ++- 5 files changed, 10 insertions(+), 4 deletions(-) diff --git a/.env.example b/.env.example index 0948d20b9d..b5b202d934 100644 --- a/.env.example +++ b/.env.example @@ -157,6 +157,8 @@ TRACECAT__AGENT_FS_ARCHIVE_CACHE_MAX_BYTES=10737418240 TRACECAT__AGENT_SANDBOX_MEMORY_MB=4096 TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED=true TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB=4096 +# Readiness sentinel path shared by the worker and the compose healthcheck. +TRACECAT__AGENT_EXECUTOR_READY_FILE=/var/run/tracecat/agent-executor-ready # --- Local registry --- # Enable this only for local-registry development. Leave disabled for normal remote/builtin registry use. diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index a37b342961..d1fb778afa 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -397,6 +397,7 @@ services: TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} + TRACECAT__AGENT_EXECUTOR_READY_FILE: ${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -422,7 +423,7 @@ services: - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] healthcheck: - test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + test: ["CMD", "test", "-f", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready}"] interval: 30s retries: 3 start_period: 120s diff --git a/docker-compose.local.yml b/docker-compose.local.yml index dc3311c3af..c36bf403a2 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -419,6 +419,7 @@ services: TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} + TRACECAT__AGENT_EXECUTOR_READY_FILE: ${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -442,7 +443,7 @@ services: - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] healthcheck: - test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + test: ["CMD", "test", "-f", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready}"] interval: 30s retries: 3 start_period: 120s diff --git a/docker-compose.sandbox.yml b/docker-compose.sandbox.yml index 69ddf6104a..1234527070 100644 --- a/docker-compose.sandbox.yml +++ b/docker-compose.sandbox.yml @@ -25,13 +25,14 @@ services: TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} + TRACECAT__AGENT_EXECUTOR_READY_FILE: ${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready} security_opt: - seccomp:unconfined - systempaths=unconfined devices: - /dev/net/tun:/dev/net/tun healthcheck: - test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + test: ["CMD", "test", "-f", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready}"] interval: 30s retries: 3 start_period: 120s diff --git a/docker-compose.yml b/docker-compose.yml index dff9dab451..d7b8e66dbb 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -410,6 +410,7 @@ services: TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} + TRACECAT__AGENT_EXECUTOR_READY_FILE: ${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready} TRACECAT__LLM_PROXY_READ_TIMEOUT: ${TRACECAT__LLM_PROXY_READ_TIMEOUT:-600} TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS: ${TRACECAT__LLM_GATEWAY_CREDENTIAL_CACHE_TTL_SECONDS:-60} TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS: ${TRACECAT__LLM_GATEWAY_HEALTHCHECK_INTERVAL_SECONDS:-30} @@ -437,7 +438,7 @@ services: - sandbox-cache:/var/lib/tracecat/sandbox-cache command: ["python", "-m", "tracecat.agent.executor_worker"] healthcheck: - test: ["CMD", "test", "-f", "/var/run/tracecat/agent-executor-ready"] + test: ["CMD", "test", "-f", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready}"] interval: 30s retries: 3 start_period: 120s From ed5ce8b2b48c0bd4daaf66f7513f36e82f8ec6a5 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:27:37 -0400 Subject: [PATCH 06/14] fix(agent): clear stale readiness sentinel before startup validation Move the stale-sentinel cleanup ahead of concurrency validation and the memory-budget check so a fail-fast crashloop cannot leave a SIGKILLed predecessor's sentinel visible to healthchecks. --- .../test_agent_executor_worker_resources.py | 42 +++++++++++++++++++ tracecat/agent/executor_worker.py | 11 ++--- 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_agent_executor_worker_resources.py b/tests/unit/test_agent_executor_worker_resources.py index a2bfc8c08e..28b7440c20 100644 --- a/tests/unit/test_agent_executor_worker_resources.py +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -224,6 +224,48 @@ async def test_agent_executor_memory_budget_error_propagates_from_main( start_runtime_services.assert_not_awaited() +@pytest.mark.anyio +async def test_agent_executor_clears_stale_sentinel_before_budget_validation( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + ready_file = tmp_path / "run" / "agent-executor-ready" + ready_file.parent.mkdir(parents=True) + ready_file.write_text("stale\n") + + cgroup_root = tmp_path / "cgroup" + cgroup_root.mkdir() + (cgroup_root / "memory.max").write_text(f"{4096 * 1024 * 1024}\n") + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB", + 4096, + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_SANDBOX_MEMORY_MB", + 4096, + ) + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + lambda: PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root), + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_READY_FILE", + str(ready_file), + ) + + with pytest.raises(AgentExecutorMemoryBudgetError): + await executor_worker.main(shutdown_event=asyncio.Event()) + + assert not ready_file.exists() + + @pytest.mark.anyio @pytest.mark.parametrize("max_concurrent", ["0", "-1"]) async def test_agent_executor_rejects_nonpositive_concurrency( diff --git a/tracecat/agent/executor_worker.py b/tracecat/agent/executor_worker.py index 1870c6ecd1..f4976892ac 100644 --- a/tracecat/agent/executor_worker.py +++ b/tracecat/agent/executor_worker.py @@ -114,6 +114,12 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: if shutdown_event is None: shutdown_event = asyncio.Event() runtime_failure_reason = None + readiness_file = Path(config.TRACECAT__AGENT_EXECUTOR_READY_FILE) + # A SIGKILLed predecessor cannot run its cleanup, and the sentinel may live + # on a filesystem that survives container restarts; clear it before any + # startup work — including validation that can raise — so a readiness + # probe never sees a stale file. + _remove_readiness_file(readiness_file) max_concurrent = int( os.environ.get("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES") or 1 ) @@ -132,11 +138,6 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: reserve_mb=config.TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB, sandbox_memory_mb=config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, ) - readiness_file = Path(config.TRACECAT__AGENT_EXECUTOR_READY_FILE) - # A SIGKILLed predecessor cannot run its cleanup, and the sentinel may live - # on a filesystem that survives container restarts; clear it before any - # startup work so a readiness probe never sees a stale file. - _remove_readiness_file(readiness_file) logger.info( "Starting AgentExecutorWorker", From ffbe671d1bc5efdbd8a75b08321f647a6c063203 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:49:36 -0400 Subject: [PATCH 07/14] feat(agent): delegate cgroup subtree to apiuser in sandbox overlay Start the agent-executor entrypoint as root in the sandbox Compose overlay, remount /sys/fs/cgroup read-write, chown the cgroup v2 delegation files to apiuser, and drop privileges via setpriv before the worker starts. Without delegation the root-owned cgroupfs made preparation fail with EACCES as UID 1001, silently disabling per-sandbox memory limits in the only compose deployment that enables nsjail. --- docker-compose.sandbox.yml | 6 ++++++ docker/scripts/entrypoint.sh | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/docker-compose.sandbox.yml b/docker-compose.sandbox.yml index 1234527070..a422bbe340 100644 --- a/docker-compose.sandbox.yml +++ b/docker-compose.sandbox.yml @@ -21,7 +21,13 @@ services: # Required for nsjail sandbox execution. cap_add: - SYS_ADMIN + # Start the entrypoint as root so it can delegate the container's cgroup + # v2 subtree to apiuser for per-sandbox memory limits; it drops privileges + # before the worker starts. The flag is coupled to user "0:0" — override + # neither independently. + user: "0:0" environment: + TRACECAT__ENTRYPOINT_CGROUP_DELEGATE: "true" TRACECAT__AGENT_SANDBOX_MEMORY_MB: ${TRACECAT__AGENT_SANDBOX_MEMORY_MB:-4096} TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: ${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true} TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB: ${TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB:-4096} diff --git a/docker/scripts/entrypoint.sh b/docker/scripts/entrypoint.sh index 332a968ad0..16ff537d38 100755 --- a/docker/scripts/entrypoint.sh +++ b/docker/scripts/entrypoint.sh @@ -1,6 +1,22 @@ #!/usr/bin/env bash set -euo pipefail +# When started as root with delegation requested (the sandbox Compose overlay +# pairs TRACECAT__ENTRYPOINT_CGROUP_DELEGATE with user: "0:0"), hand the +# container's cgroup v2 subtree to apiuser so agent sandbox memory limits can +# be prepared without root, then drop privileges before doing anything else. +if [[ "$(id -u)" == "0" && "${TRACECAT__ENTRYPOINT_CGROUP_DELEGATE:-false}" == "true" ]]; then + if mount -o remount,rw /sys/fs/cgroup && + chown apiuser:apiuser /sys/fs/cgroup /sys/fs/cgroup/cgroup.procs \ + /sys/fs/cgroup/cgroup.subtree_control /sys/fs/cgroup/cgroup.threads; then + echo "Delegated /sys/fs/cgroup to apiuser." + else + echo "Unable to delegate /sys/fs/cgroup to apiuser; agent sandbox" \ + "cgroup limits will be unavailable." >&2 + fi + exec setpriv --reuid=apiuser --regid=apiuser --init-groups /app/entrypoint.sh "$@" +fi + # Function to run migrations run_migrations() { echo "Running database migrations..." From 17012bb936a181eb9186488e4a0033d190ff4fd9 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:06:47 -0400 Subject: [PATCH 08/14] fix(agent): set apiuser identity env when dropping privileges setpriv changes only IDs, so the delegation path left HOME=/root for the UID 1001 worker, breaking HOME-dependent tooling (npx MCP probes, SSH helpers). Set HOME/USER/LOGNAME explicitly instead of --reset-env, which would clear the service configuration environment. --- docker/scripts/entrypoint.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docker/scripts/entrypoint.sh b/docker/scripts/entrypoint.sh index 16ff537d38..7a913b9752 100755 --- a/docker/scripts/entrypoint.sh +++ b/docker/scripts/entrypoint.sh @@ -14,6 +14,9 @@ if [[ "$(id -u)" == "0" && "${TRACECAT__ENTRYPOINT_CGROUP_DELEGATE:-false}" == " echo "Unable to delegate /sys/fs/cgroup to apiuser; agent sandbox" \ "cgroup limits will be unavailable." >&2 fi + # setpriv changes only IDs; fix the identity env vars ourselves instead of + # --reset-env, which would clear the service configuration environment. + export HOME=/home/apiuser USER=apiuser LOGNAME=apiuser exec setpriv --reuid=apiuser --regid=apiuser --init-groups /app/entrypoint.sh "$@" fi From f544e7e22774b74ae287b836c04ad3371d9f364d Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:38:57 -0400 Subject: [PATCH 09/14] fix(agent): harden cgroup detection edge cases Require a unified cgroup v2 entry and an existing cgroup.controllers file before the entrypoint delegates ownership, so cgroup v1 hosts never get their cgroupfs root chowned, and drop the basename-based main guard in detect_cgroup_root: a runtime-assigned cgroup legitimately named main must be treated as the container boundary, not escaped to its parent. The cached pre-move root remains authoritative for budget validation. --- docker/scripts/agent-executor-entrypoint.sh | 28 +++++++++++++-------- tests/unit/test_agent_sandbox_cgroup.py | 12 ++++++--- tracecat/agent/sandbox/cgroup.py | 5 +--- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/docker/scripts/agent-executor-entrypoint.sh b/docker/scripts/agent-executor-entrypoint.sh index 3afc753673..61dd738d2f 100755 --- a/docker/scripts/agent-executor-entrypoint.sh +++ b/docker/scripts/agent-executor-entrypoint.sh @@ -11,17 +11,25 @@ set -euo pipefail if [[ "$(id -u)" == "0" ]]; then # Resolve this container's cgroup v2 directory: the mount root under a # private cgroup namespace, a subpath of the host cgroupfs otherwise - # (e.g. privileged Kubernetes). Never touch anything above it. - cgroup_rel="$(sed -n 's/^0:://p' /proc/self/cgroup)" - cgroup_rel="${cgroup_rel#/}" - cgroup_dir="/sys/fs/cgroup${cgroup_rel:+/$cgroup_rel}" - mount -o remount,rw /sys/fs/cgroup 2>/dev/null || true - if chown apiuser:apiuser "$cgroup_dir" "$cgroup_dir/cgroup.procs" \ - "$cgroup_dir/cgroup.subtree_control" "$cgroup_dir/cgroup.threads"; then - echo "Delegated $cgroup_dir to apiuser." + # (e.g. privileged Kubernetes). Never touch anything above it, and never + # touch anything at all without a unified v2 entry — on a cgroup v1 host + # an empty match would otherwise point at the cgroupfs root. + cgroup_path="$(sed -n 's/^0:://p' /proc/self/cgroup | head -n 1)" + if [[ -z "$cgroup_path" ]]; then + echo "No cgroup v2 entry in /proc/self/cgroup; agent sandbox cgroup" \ + "limits will be unavailable." >&2 else - echo "Unable to delegate $cgroup_dir to apiuser; agent sandbox" \ - "cgroup limits will be unavailable." >&2 + cgroup_rel="${cgroup_path#/}" + cgroup_dir="/sys/fs/cgroup${cgroup_rel:+/$cgroup_rel}" + mount -o remount,rw /sys/fs/cgroup 2>/dev/null || true + if [[ -f "$cgroup_dir/cgroup.controllers" ]] && + chown apiuser:apiuser "$cgroup_dir" "$cgroup_dir/cgroup.procs" \ + "$cgroup_dir/cgroup.subtree_control" "$cgroup_dir/cgroup.threads"; then + echo "Delegated $cgroup_dir to apiuser." + else + echo "Unable to delegate $cgroup_dir to apiuser; agent sandbox" \ + "cgroup limits will be unavailable." >&2 + fi fi # setpriv changes only IDs; fix the identity env vars ourselves instead of # --reset-env, which would clear the service configuration environment. diff --git a/tests/unit/test_agent_sandbox_cgroup.py b/tests/unit/test_agent_sandbox_cgroup.py index d82955270e..54af2a630e 100644 --- a/tests/unit/test_agent_sandbox_cgroup.py +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -119,9 +119,12 @@ def deny_read( assert detect_cgroup_root(proc_cgroup_path, tmp_path / "cgroupfs") is None -def test_detect_cgroup_root_uses_parent_when_process_is_in_main( +def test_detect_cgroup_root_keeps_cgroup_legitimately_named_main( tmp_path: Path, ) -> None: + # A runtime-assigned cgroup that happens to be named "main" is the + # container's real boundary; escaping to its parent would prepare cgroups + # beside the container's memory limit instead of beneath it. proc_cgroup_path = _write_proc_cgroup( tmp_path, "0::/kubepods.slice/pod.scope/main\n", @@ -130,7 +133,7 @@ def test_detect_cgroup_root_uses_parent_when_process_is_in_main( result = detect_cgroup_root(proc_cgroup_path, cgroupfs) - assert result == cgroupfs / "kubepods.slice/pod.scope" + assert result == cgroupfs / "kubepods.slice/pod.scope/main" def test_get_agent_sandbox_cgroup_returns_unprepared_default( @@ -385,7 +388,10 @@ def move_worker_into_main(path: Path, value: str) -> None: ) (cgroup_root / "main" / "memory.max").write_text("max\n") - assert detect_cgroup_root(proc_cgroup_path, cgroupfs) == cgroup_root + # A fresh detection after the move resolves to the main leaf (whose + # memory.max is unlimited) — which is exactly why budget validation must + # consume the root captured by prepare before any PID moves. + assert detect_cgroup_root(proc_cgroup_path, cgroupfs) == cgroup_root / "main" assert prepared_cgroup.root == cgroup_root assert ( clamp_agent_executor_concurrency( diff --git a/tracecat/agent/sandbox/cgroup.py b/tracecat/agent/sandbox/cgroup.py index 8bf628f32e..7881cbc832 100644 --- a/tracecat/agent/sandbox/cgroup.py +++ b/tracecat/agent/sandbox/cgroup.py @@ -80,10 +80,7 @@ def detect_cgroup_root( if hierarchy_id != "0" or controllers: continue - root = cgroupfs if cgroup_path == "/" else cgroupfs / cgroup_path.lstrip("/") - if root.name == "main": - return root.parent - return root + return cgroupfs if cgroup_path == "/" else cgroupfs / cgroup_path.lstrip("/") return None From 17c69a92feeaaa6cd02c0b7e25e6150c3d6942d9 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:56:34 -0400 Subject: [PATCH 10/14] fix(agent): honor cgroup-enabled flag in delegation entrypoint Skip the root-side remount/chown when TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED is falsy (mirroring env_bool), matching the Python-side preparation skip. The privilege drop stays unconditional so a root-started container never reaches the worker as root regardless of the flag. --- docker/scripts/agent-executor-entrypoint.sh | 47 ++++++++++++--------- 1 file changed, 28 insertions(+), 19 deletions(-) diff --git a/docker/scripts/agent-executor-entrypoint.sh b/docker/scripts/agent-executor-entrypoint.sh index 61dd738d2f..d716746ae2 100755 --- a/docker/scripts/agent-executor-entrypoint.sh +++ b/docker/scripts/agent-executor-entrypoint.sh @@ -9,26 +9,35 @@ set -euo pipefail # needed privileges are absent the worker starts unchanged and falls back to # rlimit-only sandbox limits. Non-root invocations pass straight through. if [[ "$(id -u)" == "0" ]]; then - # Resolve this container's cgroup v2 directory: the mount root under a - # private cgroup namespace, a subpath of the host cgroupfs otherwise - # (e.g. privileged Kubernetes). Never touch anything above it, and never - # touch anything at all without a unified v2 entry — on a cgroup v1 host - # an empty match would otherwise point at the cgroupfs root. - cgroup_path="$(sed -n 's/^0:://p' /proc/self/cgroup | head -n 1)" - if [[ -z "$cgroup_path" ]]; then - echo "No cgroup v2 entry in /proc/self/cgroup; agent sandbox cgroup" \ - "limits will be unavailable." >&2 - else - cgroup_rel="${cgroup_path#/}" - cgroup_dir="/sys/fs/cgroup${cgroup_rel:+/$cgroup_rel}" - mount -o remount,rw /sys/fs/cgroup 2>/dev/null || true - if [[ -f "$cgroup_dir/cgroup.controllers" ]] && - chown apiuser:apiuser "$cgroup_dir" "$cgroup_dir/cgroup.procs" \ - "$cgroup_dir/cgroup.subtree_control" "$cgroup_dir/cgroup.threads"; then - echo "Delegated $cgroup_dir to apiuser." - else - echo "Unable to delegate $cgroup_dir to apiuser; agent sandbox" \ + # Mirror config.env_bool falsy values so disabling the feature also skips + # the root-side delegation, not just the Python-side preparation. The + # privilege drop below is unconditional: root never reaches the worker. + cgroup_enabled="$(printf '%s' "${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true}" | tr '[:upper:]' '[:lower:]')" + case "$cgroup_enabled" in 0 | false | no | off) cgroup_enabled=false ;; *) cgroup_enabled=true ;; esac + + if [[ "$cgroup_enabled" == "true" ]]; then + # Resolve this container's cgroup v2 directory: the mount root under a + # private cgroup namespace, a subpath of the host cgroupfs otherwise + # (e.g. privileged Kubernetes). Never touch anything above it, and + # never touch anything at all without a unified v2 entry — on a + # cgroup v1 host an empty match would otherwise point at the + # cgroupfs root. + cgroup_path="$(sed -n 's/^0:://p' /proc/self/cgroup | head -n 1)" + if [[ -z "$cgroup_path" ]]; then + echo "No cgroup v2 entry in /proc/self/cgroup; agent sandbox" \ "cgroup limits will be unavailable." >&2 + else + cgroup_rel="${cgroup_path#/}" + cgroup_dir="/sys/fs/cgroup${cgroup_rel:+/$cgroup_rel}" + mount -o remount,rw /sys/fs/cgroup 2>/dev/null || true + if [[ -f "$cgroup_dir/cgroup.controllers" ]] && + chown apiuser:apiuser "$cgroup_dir" "$cgroup_dir/cgroup.procs" \ + "$cgroup_dir/cgroup.subtree_control" "$cgroup_dir/cgroup.threads"; then + echo "Delegated $cgroup_dir to apiuser." + else + echo "Unable to delegate $cgroup_dir to apiuser; agent sandbox" \ + "cgroup limits will be unavailable." >&2 + fi fi fi # setpriv changes only IDs; fix the identity env vars ourselves instead of From ab6c868ce9006c30d7607c887d177ad79ddb2143 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:13:55 -0400 Subject: [PATCH 11/14] fix(agent): honor ancestor cgroup limits in memory budget Walk every visible hierarchy level from the executor's cgroup up to the cgroupfs mount and take the minimum finite memory.max, so task-scoped limits (e.g. an ECS task limit with no container-level limit) are not misread as unlimited when the leaf reports max. PreparedCgroup carries the cgroupfs boundary so the walk never leaves the mount. --- tests/unit/test_agent_sandbox_cgroup.py | 56 +++++++++++- tracecat/agent/sandbox/cgroup.py | 111 ++++++++++++++++-------- 2 files changed, 128 insertions(+), 39 deletions(-) diff --git a/tests/unit/test_agent_sandbox_cgroup.py b/tests/unit/test_agent_sandbox_cgroup.py index 54af2a630e..b87a094e2f 100644 --- a/tests/unit/test_agent_sandbox_cgroup.py +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -189,7 +189,11 @@ def test_prepare_agent_sandbox_cgroup_beneath_detected_root( enabled=True, ) - assert result == PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + assert result == PreparedCgroup( + CgroupAvailability.AVAILABLE, + cgroup_root, + cgroupfs, + ) assert result.sandbox_mount == cgroup_root assert (cgroup_root / "main" / "cgroup.procs").read_text() == "202\n" assert (cgroup_root / "cgroup.subtree_control").read_text() == "+memory\n" @@ -209,7 +213,11 @@ def test_prepare_agent_sandbox_cgroup_detection_failure_warns_once( enabled=True, ) - assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + assert result == PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + None, + tmp_path / "cgroupfs", + ) mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.kwargs["step"] == "detect cgroup v2 root" @@ -232,7 +240,11 @@ def test_prepare_agent_sandbox_cgroup_requires_memory_controller( enabled=True, ) - assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root) + assert result == PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + cgroupfs, + ) mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.kwargs["step"] == "verify memory controller" @@ -264,7 +276,11 @@ def deny_subtree_control(path: Path, value: str) -> None: enabled=True, ) - assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root) + assert result == PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + cgroupfs, + ) assert result.sandbox_mount is None assert read_cgroup_memory_limit(result.root).kind is CgroupMemoryLimitKind.LIMITED mock_logger.warning.assert_called_once() @@ -332,6 +348,38 @@ def test_read_cgroup_memory_limit_from_detected_root( assert result.limit_bytes == expected_bytes +def test_read_cgroup_memory_limit_honors_ancestor_limit(tmp_path: Path) -> None: + # Task-scoped hierarchies (e.g. ECS) can enforce the budget on an ancestor + # while the container leaf reads "max"; the effective limit is the minimum + # finite value across the visible hierarchy. + cgroupfs = tmp_path / "cgroupfs" + leaf = cgroupfs / "ecs-task" / "container" + leaf.mkdir(parents=True) + (cgroupfs / "memory.max").write_text("max\n") + (cgroupfs / "ecs-task" / "memory.max").write_text(f"{8 * 1024**3}\n") + (leaf / "memory.max").write_text("max\n") + + result = read_cgroup_memory_limit(leaf, cgroupfs) + + assert result.kind is CgroupMemoryLimitKind.LIMITED + assert result.limit_bytes == 8 * 1024**3 + + +def test_read_cgroup_memory_limit_takes_minimum_across_levels( + tmp_path: Path, +) -> None: + cgroupfs = tmp_path / "cgroupfs" + leaf = cgroupfs / "task" / "container" + leaf.mkdir(parents=True) + (cgroupfs / "task" / "memory.max").write_text(f"{16 * 1024**3}\n") + (leaf / "memory.max").write_text(f"{4 * 1024**3}\n") + + result = read_cgroup_memory_limit(leaf, cgroupfs) + + assert result.kind is CgroupMemoryLimitKind.LIMITED + assert result.limit_bytes == 4 * 1024**3 + + def test_read_cgroup_memory_limit_without_root_is_unavailable() -> None: result = read_cgroup_memory_limit(None) diff --git a/tracecat/agent/sandbox/cgroup.py b/tracecat/agent/sandbox/cgroup.py index 7881cbc832..5a689c5938 100644 --- a/tracecat/agent/sandbox/cgroup.py +++ b/tracecat/agent/sandbox/cgroup.py @@ -42,6 +42,7 @@ class PreparedCgroup: availability: CgroupAvailability root: Path | None + cgroupfs: Path = CGROUPFS @property def sandbox_mount(self) -> Path | None: @@ -152,7 +153,11 @@ def prepare_agent_sandbox_cgroup( step="detect cgroup v2 root", error=ValueError(f"no cgroup v2 entry found in {proc_cgroup_path}"), ) - _prepared_cgroup = PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + _prepared_cgroup = PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + None, + cgroupfs, + ) return _prepared_cgroup probe_cgroup: Path | None = None @@ -204,10 +209,15 @@ def prepare_agent_sandbox_cgroup( _prepared_cgroup = PreparedCgroup( CgroupAvailability.UNAVAILABLE, cgroup_root, + cgroupfs, ) return _prepared_cgroup - _prepared_cgroup = PreparedCgroup(CgroupAvailability.AVAILABLE, cgroup_root) + _prepared_cgroup = PreparedCgroup( + CgroupAvailability.AVAILABLE, + cgroup_root, + cgroupfs, + ) return _prepared_cgroup @@ -222,46 +232,74 @@ def get_agent_sandbox_cgroup() -> PreparedCgroup: def read_cgroup_memory_limit( cgroup_root: Path | None, + cgroupfs: Path = CGROUPFS, ) -> CgroupMemoryLimit: - """Read and parse the cgroup v2 container memory limit.""" + """Read the effective cgroup v2 memory limit for the executor's cgroup.""" if cgroup_root is None: return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) - memory_max_path = cgroup_root / "memory.max" - try: - raw_limit = memory_max_path.read_text().strip() - except FileNotFoundError: - return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) - except OSError as exc: - logger.warning( - "Unable to read cgroup memory limit; skipping worker memory-budget " - "validation", - path=str(memory_max_path), - errno=exc.errno, - error=str(exc), - ) - return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + # The effective limit can live on an ancestor while the leaf reads "max" + # (e.g. an ECS task-level limit with no container-level limit), so take + # the minimum finite memory.max across every hierarchy level visible + # between the leaf and the cgroupfs mount. + directories = [cgroup_root] + if cgroup_root != cgroupfs and cgroupfs in cgroup_root.parents: + for parent in cgroup_root.parents: + directories.append(parent) + if parent == cgroupfs: + break + + finite_limits: list[int] = [] + unreadable = False + unlimited_seen = False + for directory in directories: + memory_max_path = directory / "memory.max" + try: + raw_limit = memory_max_path.read_text().strip() + except FileNotFoundError: + continue + except OSError as exc: + logger.warning( + "Unable to read cgroup memory limit; skipping worker " + "memory-budget validation for this level", + path=str(memory_max_path), + errno=exc.errno, + error=str(exc), + ) + unreadable = True + continue - if raw_limit == "max": - return CgroupMemoryLimit(CgroupMemoryLimitKind.UNLIMITED) + if raw_limit == "max": + unlimited_seen = True + continue - try: - limit_bytes = int(raw_limit) - if limit_bytes < 0: - raise ValueError("memory limit cannot be negative") - except ValueError as exc: - logger.warning( - "Invalid cgroup memory limit; skipping worker memory-budget validation", - path=str(memory_max_path), - value=raw_limit, - error=str(exc), + try: + limit_bytes = int(raw_limit) + if limit_bytes < 0: + raise ValueError("memory limit cannot be negative") + except ValueError as exc: + logger.warning( + "Invalid cgroup memory limit; skipping worker memory-budget " + "validation for this level", + path=str(memory_max_path), + value=raw_limit, + error=str(exc), + ) + unreadable = True + continue + + finite_limits.append(limit_bytes) + + if finite_limits: + return CgroupMemoryLimit( + CgroupMemoryLimitKind.LIMITED, + limit_bytes=min(finite_limits), ) + if unreadable: return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) - - return CgroupMemoryLimit( - CgroupMemoryLimitKind.LIMITED, - limit_bytes=limit_bytes, - ) + if unlimited_seen: + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNLIMITED) + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) def clamp_agent_executor_concurrency( @@ -272,7 +310,10 @@ def clamp_agent_executor_concurrency( sandbox_memory_mb: int, ) -> int: """Clamp agent executor concurrency to the container memory budget.""" - memory_limit = read_cgroup_memory_limit(prepared_cgroup.root) + memory_limit = read_cgroup_memory_limit( + prepared_cgroup.root, + prepared_cgroup.cgroupfs, + ) if memory_limit.kind is CgroupMemoryLimitKind.UNLIMITED: logger.debug( "Container memory is unlimited; skipping worker memory-budget validation" From 7276d2fb67cdf93dc9f62e9e65d33df488237bfa Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:37:00 -0400 Subject: [PATCH 12/14] fix(agent): enforce sandbox cgroup limits --- docker/scripts/agent-executor-entrypoint.sh | 25 +- .../test_agent_executor_worker_resources.py | 102 +++++++- tests/unit/test_agent_sandbox_cgroup.py | 10 +- tests/unit/test_agent_sandbox_config.py | 2 + tests/unit/test_agent_sandbox_litellm.py | 232 +++++++++++++++--- tests/unit/test_config.py | 33 +++ tests/unit/test_nsjail_seccomp.py | 8 + tests/unit/test_stdio_probe.py | 12 + tracecat/agent/common/config.py | 20 +- tracecat/agent/executor_worker.py | 12 +- tracecat/agent/sandbox/cgroup.py | 36 ++- tracecat/agent/sandbox/config.py | 4 +- tracecat/config.py | 2 +- tracecat/sandbox/executor.py | 12 +- 14 files changed, 450 insertions(+), 60 deletions(-) diff --git a/docker/scripts/agent-executor-entrypoint.sh b/docker/scripts/agent-executor-entrypoint.sh index d716746ae2..d521b33835 100755 --- a/docker/scripts/agent-executor-entrypoint.sh +++ b/docker/scripts/agent-executor-entrypoint.sh @@ -5,15 +5,26 @@ set -euo pipefail # memory limits. Start the container as root with this entrypoint (compose: # user "0:0"; Kubernetes: runAsUser 0) and it hands the container's own # cgroup v2 directory to apiuser so the worker can prepare nsjail child -# cgroups without root, then drops privileges before starting it. Where the -# needed privileges are absent the worker starts unchanged and falls back to -# rlimit-only sandbox limits. Non-root invocations pass straight through. +# cgroups without root, then drops privileges before starting it. The worker +# rejects startup if cgroup enforcement is enabled but delegation failed. +# Non-root invocations pass straight through to that same fail-closed check. if [[ "$(id -u)" == "0" ]]; then # Mirror config.env_bool falsy values so disabling the feature also skips # the root-side delegation, not just the Python-side preparation. The # privilege drop below is unconditional: root never reaches the worker. - cgroup_enabled="$(printf '%s' "${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true}" | tr '[:upper:]' '[:lower:]')" - case "$cgroup_enabled" in 0 | false | no | off) cgroup_enabled=false ;; *) cgroup_enabled=true ;; esac + cgroup_enabled="$( + printf '%s' "${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true}" | + tr '[:upper:]' '[:lower:]' | + sed 's/^[[:space:]]*//;s/[[:space:]]*$//' + )" + case "$cgroup_enabled" in + 1 | true | yes | on) cgroup_enabled=true ;; + 0 | false | no | off) cgroup_enabled=false ;; + *) + echo "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED must be a boolean value" >&2 + exit 64 + ;; + esac if [[ "$cgroup_enabled" == "true" ]]; then # Resolve this container's cgroup v2 directory: the mount root under a @@ -25,7 +36,7 @@ if [[ "$(id -u)" == "0" ]]; then cgroup_path="$(sed -n 's/^0:://p' /proc/self/cgroup | head -n 1)" if [[ -z "$cgroup_path" ]]; then echo "No cgroup v2 entry in /proc/self/cgroup; agent sandbox" \ - "cgroup limits will be unavailable." >&2 + "cgroup limits cannot be enforced." >&2 else cgroup_rel="${cgroup_path#/}" cgroup_dir="/sys/fs/cgroup${cgroup_rel:+/$cgroup_rel}" @@ -36,7 +47,7 @@ if [[ "$(id -u)" == "0" ]]; then echo "Delegated $cgroup_dir to apiuser." else echo "Unable to delegate $cgroup_dir to apiuser; agent sandbox" \ - "cgroup limits will be unavailable." >&2 + "cgroup limits cannot be enforced." >&2 fi fi fi diff --git a/tests/unit/test_agent_executor_worker_resources.py b/tests/unit/test_agent_executor_worker_resources.py index 28b7440c20..4d5735c9c6 100644 --- a/tests/unit/test_agent_executor_worker_resources.py +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -10,6 +10,7 @@ from tracecat.agent.sandbox.cgroup import ( AgentExecutorMemoryBudgetError, + AgentSandboxCgroupUnavailableError, CgroupAvailability, PreparedCgroup, ) @@ -66,7 +67,7 @@ async def observe_readiness() -> None: monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: PreparedCgroup(CgroupAvailability.DISABLED, None), + lambda *, enabled: PreparedCgroup(CgroupAvailability.DISABLED, None), ) monkeypatch.setattr( executor_worker, @@ -156,7 +157,7 @@ def keep_concurrency( monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: PreparedCgroup(CgroupAvailability.DISABLED, None), + lambda *, enabled: PreparedCgroup(CgroupAvailability.DISABLED, None), ) monkeypatch.setattr( executor_worker, @@ -207,7 +208,10 @@ async def test_agent_executor_memory_budget_error_propagates_from_main( monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root), + lambda *, enabled: PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + ), ) monkeypatch.setattr( executor_worker, @@ -252,7 +256,10 @@ async def test_agent_executor_clears_stale_sentinel_before_budget_validation( monkeypatch.setattr( executor_worker, "prepare_agent_sandbox_cgroup", - lambda: PreparedCgroup(CgroupAvailability.UNAVAILABLE, cgroup_root), + lambda *, enabled: PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + ), ) monkeypatch.setattr( executor_worker.config, @@ -266,6 +273,93 @@ async def test_agent_executor_clears_stale_sentinel_before_budget_validation( assert not ready_file.exists() +@pytest.mark.anyio +async def test_agent_executor_fails_closed_when_required_cgroup_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from tracecat.agent import executor_worker + + ready_file = tmp_path / "run" / "agent-executor-ready" + prepare_cgroup = Mock( + return_value=PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) + ) + start_runtime_services = AsyncMock() + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", + True, + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__DISABLE_NSJAIL", + False, + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_EXECUTOR_READY_FILE", + str(ready_file), + ) + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + prepare_cgroup, + ) + monkeypatch.setattr( + executor_worker, + "_start_runtime_services", + start_runtime_services, + ) + + with pytest.raises(AgentSandboxCgroupUnavailableError): + await executor_worker.main(shutdown_event=asyncio.Event()) + + prepare_cgroup.assert_called_once_with(enabled=True) + start_runtime_services.assert_not_awaited() + assert not ready_file.exists() + + +@pytest.mark.anyio +async def test_agent_executor_disables_cgroups_when_nsjail_is_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tracecat.agent import executor_worker + + prepare_cgroup = Mock( + return_value=PreparedCgroup(CgroupAvailability.DISABLED, None) + ) + start_runtime_services = AsyncMock(side_effect=RuntimeError("stop after setup")) + monkeypatch.setenv("TRACECAT__AGENT_EXECUTOR_MAX_CONCURRENT_ACTIVITIES", "1") + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", + True, + ) + monkeypatch.setattr( + executor_worker.config, + "TRACECAT__DISABLE_NSJAIL", + True, + ) + monkeypatch.setattr( + executor_worker, + "prepare_agent_sandbox_cgroup", + prepare_cgroup, + ) + monkeypatch.setattr( + executor_worker, + "_start_runtime_services", + start_runtime_services, + ) + monkeypatch.setattr(executor_worker, "_stop_runtime_services", AsyncMock()) + monkeypatch.setattr(executor_worker, "close_storage_client_cache", AsyncMock()) + + with pytest.raises(RuntimeError, match="stop after setup"): + await executor_worker.main(shutdown_event=asyncio.Event()) + + prepare_cgroup.assert_called_once_with(enabled=False) + + @pytest.mark.anyio @pytest.mark.parametrize("max_concurrent", ["0", "-1"]) async def test_agent_executor_rejects_nonpositive_concurrency( diff --git a/tests/unit/test_agent_sandbox_cgroup.py b/tests/unit/test_agent_sandbox_cgroup.py index b87a094e2f..bed9beeaed 100644 --- a/tests/unit/test_agent_sandbox_cgroup.py +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -11,6 +11,7 @@ from tracecat.agent.sandbox.cgroup import ( BYTES_PER_MEBIBYTE, AgentExecutorMemoryBudgetError, + AgentSandboxCgroupUnavailableError, CgroupAvailability, CgroupMemoryLimitKind, PreparedCgroup, @@ -148,7 +149,8 @@ def test_get_agent_sandbox_cgroup_returns_unprepared_default( result = get_agent_sandbox_cgroup() assert result == PreparedCgroup(CgroupAvailability.UNAVAILABLE, None) - assert result.sandbox_mount is None + with pytest.raises(AgentSandboxCgroupUnavailableError): + _ = result.sandbox_mount def test_prepare_agent_sandbox_cgroup_returns_cached_disabled_state( @@ -169,6 +171,9 @@ def test_prepare_agent_sandbox_cgroup_returns_cached_disabled_state( ) assert first == PreparedCgroup(CgroupAvailability.DISABLED, None) + assert first.sandbox_mount is None + with pytest.raises(AgentSandboxCgroupUnavailableError): + first.require_sandbox_mount() assert second is first assert get_agent_sandbox_cgroup() is first @@ -281,7 +286,8 @@ def deny_subtree_control(path: Path, value: str) -> None: cgroup_root, cgroupfs, ) - assert result.sandbox_mount is None + with pytest.raises(AgentSandboxCgroupUnavailableError): + _ = result.sandbox_mount assert read_cgroup_memory_limit(result.root).kind is CgroupMemoryLimitKind.LIMITED mock_logger.warning.assert_called_once() assert mock_logger.warning.call_args.kwargs["step"] == "enable memory controller" diff --git a/tests/unit/test_agent_sandbox_config.py b/tests/unit/test_agent_sandbox_config.py index ce08ca1506..28a4cabce2 100644 --- a/tests/unit/test_agent_sandbox_config.py +++ b/tests/unit/test_agent_sandbox_config.py @@ -139,6 +139,8 @@ def test_build_agent_nsjail_config_adds_available_cgroup_v2_memory_limit() -> No assert "use_cgroupv2: true" in config_text assert 'cgroupv2_mount: "/sys/fs/cgroup/kubepods.slice/pod.scope"' in config_text + assert "rlimit_as: 3072" in config_text + assert "rlimit_fsize: 256" in config_text assert f"cgroup_mem_max: {3072 * 1024 * 1024}" in config_text assert "cgroup_mem_swap_max: 0" in config_text diff --git a/tests/unit/test_agent_sandbox_litellm.py b/tests/unit/test_agent_sandbox_litellm.py index 1afa198f3e..8c4356bef5 100644 --- a/tests/unit/test_agent_sandbox_litellm.py +++ b/tests/unit/test_agent_sandbox_litellm.py @@ -55,6 +55,16 @@ ClaudeTurnRequest, ) from tracecat.agent.runtime.claude_code.transport import SandboxedCLITransport +from tracecat.agent.sandbox.cgroup import ( + CgroupAvailability, + prepare_agent_sandbox_cgroup, +) +from tracecat.agent.sandbox.config import ( + JAILED_SHIM_ENTRYPOINT_PATH, + AgentResourceLimits, + AgentSandboxConfig, + build_agent_nsjail_config, +) from tracecat.agent.sandbox.llm_proxy import ( LLM_SOCKET_NAME, LLMRoute, @@ -95,6 +105,13 @@ class _DuckDBSmokeMessage(TypedDict): duckdb_path: str +_CGROUP_SMOKE_SENTINEL = "TRACE_CAT_AGENT_CGROUP_SMOKE uid=1001 availability=available" +_CGROUP_OOM_SENTINEL = ( + "TRACE_CAT_AGENT_CGROUP_OOM sandbox_killed=true parent_survived=true" +) +_CGROUP_RECOVERY_SENTINEL = "TRACE_CAT_AGENT_CGROUP_RECOVERY sandbox_succeeded=true" + + @dataclass(slots=True) class _FakeClaudeOptions: env: dict[str, str] @@ -1125,6 +1142,7 @@ def _run_nsjail_harness_in_docker_or_skip( cli_flag: str = "--run-nsjail-harness-smoke", failure_label: str = "Dockerized nsjail harness fallback failed.", requires_tun: bool = False, + enable_cgroups: bool = False, ) -> None: if os.environ.get("TRACECAT__AGENT_NSJAIL_DOCKER_FALLBACK_CHILD") == "1": pytest.skip("nsjail unavailable inside Docker fallback child") @@ -1155,6 +1173,7 @@ def _run_nsjail_harness_in_docker_or_skip( compose_env.setdefault("ADDRESS", "0.0.0.0") compose_env.setdefault("LOG_LEVEL", "INFO") compose_env.setdefault("TRACECAT__APP_ENV", "development") + compose_project_name = f"tracecat-agent-nsjail-{uuid.uuid4().hex[:12]}" tests_mount = f"{repo_root / 'tests'}:/app/tests:ro" device_lines = ( [ @@ -1164,6 +1183,20 @@ def _run_nsjail_harness_in_docker_or_skip( if requires_tun else [] ) + cgroup_service_lines = ( + [ + " cgroup: private", + " privileged: true", + ' user: "0:0"', + ] + if enable_cgroups + else [] + ) + cgroup_enabled = "true" if enable_cgroups else "false" + entrypoint = ( + "/usr/local/bin/agent-executor-entrypoint.sh" if enable_cgroups else "sh" + ) + entrypoint_args = ["sh"] if enable_cgroups else [] override_path = Path( tempfile.mkstemp(prefix="tracecat-agent-nsjail-test-", suffix=".yml")[1] ) @@ -1174,6 +1207,7 @@ def _run_nsjail_harness_in_docker_or_skip( " api:", " build:", " target: test", + *cgroup_service_lines, " cap_add:", " - SYS_ADMIN", " security_opt:", @@ -1185,6 +1219,7 @@ def _run_nsjail_harness_in_docker_or_skip( " environment:", ' TRACECAT__AGENT_NSJAIL_DOCKER_FALLBACK_CHILD: "1"', ' TRACECAT__DISABLE_NSJAIL: "false"', + f' TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED: "{cgroup_enabled}"', ' TRACECAT__SANDBOX_NSJAIL_PATH: "/usr/local/bin/nsjail"', ' TRACECAT__SANDBOX_ROOTFS_PATH: "/var/lib/tracecat/sandbox-rootfs"', ' PYTHONDONTWRITEBYTECODE: "1"', @@ -1192,23 +1227,29 @@ def _run_nsjail_harness_in_docker_or_skip( ] ) ) + compose_command = [ + "docker", + "compose", + "--project-name", + compose_project_name, + "-f", + str(repo_root / "docker-compose.dev.yml"), + "-f", + str(override_path), + ] try: result = subprocess.run( [ - "docker", - "compose", - "-f", - str(repo_root / "docker-compose.dev.yml"), - "-f", - str(override_path), + *compose_command, "run", "--rm", "--no-deps", "--build", "-T", "--entrypoint", - "sh", + entrypoint, "api", + *entrypoint_args, "-lc", f"uv run python -m tests.unit.test_agent_sandbox_litellm {cli_flag}", ], @@ -1220,12 +1261,32 @@ def _run_nsjail_harness_in_docker_or_skip( check=False, ) finally: + cleanup_result = subprocess.run( + [*compose_command, "down", "--remove-orphans", "--rmi", "local"], + cwd=repo_root, + env=compose_env, + capture_output=True, + text=True, + timeout=60, + check=False, + ) override_path.unlink(missing_ok=True) + if cleanup_result.returncode != 0: + pytest.fail( + "Dockerized nsjail harness cleanup failed." + f"\n\nstdout:\n{cleanup_result.stdout}" + f"\n\nstderr:\n{cleanup_result.stderr}" + ) if result.returncode != 0: pytest.fail( f"{failure_label}\n\nstdout:\n{result.stdout}\n\nstderr:\n{result.stderr}" ) + if enable_cgroups: + assert "Delegated " in result.stdout + assert _CGROUP_SMOKE_SENTINEL in result.stdout + assert _CGROUP_OOM_SENTINEL in result.stdout + assert _CGROUP_RECOVERY_SENTINEL in result.stdout def _run_nsjail_harness_smoke_from_cli() -> None: @@ -1297,16 +1358,140 @@ async def run() -> None: asyncio.run(run()) -def _run_nsjail_duckdb_smoke_from_cli() -> None: +def _read_cgroup_memory_event(cgroup_root: Path, event: str) -> int: + for line in (cgroup_root / "memory.events").read_text().splitlines(): + key, raw_value = line.split() + if key == event: + return int(raw_value) + raise AssertionError(f"Missing {event!r} in {cgroup_root / 'memory.events'}") + + +async def _run_agent_cgroup_oom_isolation_case( + *, + cgroup_root: Path, + tmp_path: Path, +) -> None: + """Trigger a cgroup OOM kill below the sandbox's per-process RLIMIT_AS.""" + tmp_path.mkdir(parents=True) + workload_path = tmp_path / Path(JAILED_SHIM_ENTRYPOINT_PATH).name + workload_path.write_text( + "\n".join( + [ + "allocations = []", + "", + "while True:", + " allocation = bytearray(16 * 1024 * 1024)", + " for offset in range(0, len(allocation), 4096):", + " allocation[offset] = 1", + " allocations.append(allocation)", + ] + ) + ) + + site_packages_dir = next( + ( + Path(path) + for path in sys.path + if "site-packages" in path and Path(path).is_dir() + ), + None, + ) + assert site_packages_dir is not None + limits = AgentResourceLimits( + memory_mb=512, + cpu_seconds=30, + timeout_seconds=15, + ) + nsjail_config = build_agent_nsjail_config( + rootfs=Path(app_config.TRACECAT__SANDBOX_ROOTFS_PATH), + job_dir=tmp_path, + socket_dir=tmp_path, + config=AgentSandboxConfig(resources=limits), + site_packages_dir=site_packages_dir, + llm_socket_path=None, + mount_control_socket=False, + cgroup_mount=cgroup_root, + ) + configured_cgroup_limit = f"cgroup_mem_max: {limits.memory_mb * 1024 * 1024}" + test_cgroup_limit = f"cgroup_mem_max: {128 * 1024 * 1024}" + assert nsjail_config.count(f"rlimit_as: {limits.memory_mb}") == 1 + assert nsjail_config.count(configured_cgroup_limit) == 1 + # Keep the normal 512 MiB per-process RLIMIT_AS while lowering only the + # child cgroup to 128 MiB. A resulting oom_kill therefore proves that the + # kernel cgroup limit, rather than RLIMIT_AS, terminated the sandbox. + nsjail_config = nsjail_config.replace( + configured_cgroup_limit, + test_cgroup_limit, + 1, + ) + config_path = tmp_path / "nsjail.cfg" + config_path.write_text(nsjail_config) + config_path.chmod(0o600) + + parent_pid = os.getpid() + parent_cgroup = Path("/proc/self/cgroup").read_text() + oom_kills_before = _read_cgroup_memory_event(cgroup_root, "oom_kill") + process = await asyncio.create_subprocess_exec( + app_config.TRACECAT__SANDBOX_NSJAIL_PATH, + "--config", + str(config_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tmp_path, + ) + try: + stdout_bytes, stderr_bytes = await asyncio.wait_for( + process.communicate(), + timeout=30, + ) + except TimeoutError as exc: + process.kill() + await process.wait() + raise AssertionError( + "Cgroup OOM sandbox did not terminate within 30 seconds" + ) from exc + + stdout = stdout_bytes.decode(errors="replace") + stderr = stderr_bytes.decode(errors="replace") + oom_kills_after = _read_cgroup_memory_event(cgroup_root, "oom_kill") + if process.returncode != 137: + raise AssertionError( + f"Cgroup OOM sandbox exited with {process.returncode}, expected 137" + f"\n\nstdout:\n{stdout}\n\nstderr:\n{stderr}" + ) + if oom_kills_after <= oom_kills_before: + raise AssertionError( + "Sandbox failed without a kernel cgroup OOM kill" + f"\n\nstdout:\n{stdout}\n\nstderr:\n{stderr}" + ) + assert os.getpid() == parent_pid + assert Path("/proc/self/cgroup").read_text() == parent_cgroup + print( + f"{_CGROUP_OOM_SENTINEL} oom_kill_delta={oom_kills_after - oom_kills_before}", + flush=True, + ) + + +def _run_nsjail_cgroup_smoke_from_cli() -> None: async def run() -> None: monkeypatch = pytest.MonkeyPatch() - tmp_path = Path(tempfile.mkdtemp(prefix="tracecat-agent-nsjail-duckdb-")) + tmp_path = Path(tempfile.mkdtemp(prefix="tracecat-agent-nsjail-cgroup-")) try: + prepared_cgroup = prepare_agent_sandbox_cgroup() + assert os.getuid() == 1001 + assert prepared_cgroup.availability is CgroupAvailability.AVAILABLE + cgroup_root = prepared_cgroup.require_sandbox_mount() + print(_CGROUP_SMOKE_SENTINEL, flush=True) + await _run_agent_cgroup_oom_isolation_case( + cgroup_root=cgroup_root, + tmp_path=tmp_path / "oom", + ) _set_disable_nsjail_mode(monkeypatch, False) await _run_duckdb_cli_available_case( monkeypatch=monkeypatch, - tmp_path=tmp_path, + tmp_path=tmp_path / "recovery", ) + print(_CGROUP_RECOVERY_SENTINEL, flush=True) finally: monkeypatch.undo() shutil.rmtree(tmp_path, ignore_errors=True) @@ -2137,22 +2322,11 @@ async def test_run_agent_activity_makes_attached_skills_visible_in_each_sandbox_ ) -@pytest.mark.anyio -async def test_agent_nsjail_runtime_has_duckdb_cli( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - if not _agent_nsjail_available(): - _run_nsjail_harness_in_docker_or_skip( - cli_flag="--run-nsjail-duckdb-smoke", - failure_label="Dockerized nsjail DuckDB smoke fallback failed.", - ) - return - - _set_disable_nsjail_mode(monkeypatch, False) - await _run_duckdb_cli_available_case( - monkeypatch=monkeypatch, - tmp_path=tmp_path, +def test_agent_nsjail_cgroup_oom_isolated_and_runtime_recovers() -> None: + _run_nsjail_harness_in_docker_or_skip( + cli_flag="--run-nsjail-cgroup-smoke", + failure_label="Dockerized nsjail cgroup OOM containment smoke failed.", + enable_cgroups=True, ) @@ -2563,12 +2737,12 @@ async def fake_pump_stdin_to_process(_stdin: object) -> None: _run_nsjail_skills_smoke_from_cli() elif sys.argv[1:] == ["--run-nsjail-mcp-compression-smoke"]: _run_nsjail_mcp_compression_smoke_from_cli() - elif sys.argv[1:] == ["--run-nsjail-duckdb-smoke"]: - _run_nsjail_duckdb_smoke_from_cli() + elif sys.argv[1:] == ["--run-nsjail-cgroup-smoke"]: + _run_nsjail_cgroup_smoke_from_cli() else: raise SystemExit( "Usage: python -m tests.unit.test_agent_sandbox_litellm " "[--run-nsjail-harness-smoke|--run-nsjail-pasta-smoke|" "--run-nsjail-skills-smoke|--run-nsjail-mcp-compression-smoke|" - "--run-nsjail-duckdb-smoke]" + "--run-nsjail-cgroup-smoke]" ) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 97a397d6c0..61bb128bde 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -8,6 +8,7 @@ import pytest import tracecat.config as tracecat_config +from tracecat.agent.common.config import _env_bool as agent_env_bool from tracecat.config import bound_env, env_bool REPO_ROOT = Path(__file__).resolve().parents[2] @@ -98,6 +99,38 @@ def test_env_bool_rejects_invalid_value(monkeypatch: pytest.MonkeyPatch) -> None env_bool("TEST_BOOL_ENV", default=True) +@pytest.mark.parametrize( + ("raw_value", "expected"), + [ + ("1", True), + ("true", True), + (" YES ", True), + ("on", True), + ("0", False), + ("false", False), + (" NO ", False), + ("off", False), + ], +) +def test_agent_env_bool_matches_application_tokens( + monkeypatch: pytest.MonkeyPatch, + raw_value: str, + expected: bool, +) -> None: + monkeypatch.setenv("TEST_AGENT_BOOL_ENV", raw_value) + + assert agent_env_bool("TEST_AGENT_BOOL_ENV", default=not expected) is expected + + +def test_agent_env_bool_rejects_invalid_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TEST_AGENT_BOOL_ENV", "not-a-bool") + + with pytest.raises(ValueError, match="TEST_AGENT_BOOL_ENV must be a boolean"): + agent_env_bool("TEST_AGENT_BOOL_ENV", default=True) + + def test_config_boolean_env_values_use_env_bool() -> None: source = CONFIG_PATH.read_text() forbidden_patterns = { diff --git a/tests/unit/test_nsjail_seccomp.py b/tests/unit/test_nsjail_seccomp.py index 1c09b8a1f7..3ac9255d0e 100644 --- a/tests/unit/test_nsjail_seccomp.py +++ b/tests/unit/test_nsjail_seccomp.py @@ -76,6 +76,8 @@ def test_stdio_probe_sandbox_config_includes_cgroup_memory_limit( assert "use_cgroupv2: true" in config_text assert f'cgroupv2_mount: "{cgroup_mount}"' in config_text + assert "rlimit_as: 512" in config_text + assert "rlimit_fsize: 256" in config_text assert f"cgroup_mem_max: {512 * 1024 * 1024}" in config_text assert "cgroup_mem_swap_max: 0" in config_text @@ -89,10 +91,16 @@ def test_action_sandbox_config_includes_seccomp_policy(tmp_path: Path): config=ActionSandboxConfig( registry_paths=[tmp_path / "registry"], tracecat_app_dir=tmp_path / "app", + resources=ResourceLimits( + memory_mb=768, + max_file_size_mb=64, + ), ), ) _assert_seccomp_config(config_text) + assert "rlimit_as: 768" in config_text + assert "rlimit_fsize: 64" in config_text def test_action_sandbox_config_ignores_general_executor_cgroup_mount( diff --git a/tests/unit/test_stdio_probe.py b/tests/unit/test_stdio_probe.py index 55496a41f2..f7ae945ee0 100644 --- a/tests/unit/test_stdio_probe.py +++ b/tests/unit/test_stdio_probe.py @@ -245,6 +245,10 @@ async def test_probe_timeout_leaves_buffer_before_nsjail_limit() -> None: "tracecat.agent.mcp.stdio_probe.NsjailExecutor", return_value=executor, ), + patch( + "tracecat.agent.mcp.stdio_probe.get_agent_sandbox_cgroup", + return_value=PreparedCgroup(CgroupAvailability.DISABLED, None), + ), ): await probe_stdio_mcp_tools_in_sandbox( command="python", @@ -285,6 +289,10 @@ async def test_probe_returns_friendly_structured_timeout() -> None: "tracecat.agent.mcp.stdio_probe.NsjailExecutor", return_value=executor, ), + patch( + "tracecat.agent.mcp.stdio_probe.get_agent_sandbox_cgroup", + return_value=PreparedCgroup(CgroupAvailability.DISABLED, None), + ), ): result = await probe_stdio_mcp_tools_in_sandbox( command="python", @@ -314,6 +322,10 @@ async def test_probe_returns_friendly_sandbox_timeout() -> None: "tracecat.agent.mcp.stdio_probe.NsjailExecutor", return_value=executor, ), + patch( + "tracecat.agent.mcp.stdio_probe.get_agent_sandbox_cgroup", + return_value=PreparedCgroup(CgroupAvailability.DISABLED, None), + ), ): result = await probe_stdio_mcp_tools_in_sandbox( command="python", diff --git a/tracecat/agent/common/config.py b/tracecat/agent/common/config.py index e7c2f0d93d..72cb45b1ab 100644 --- a/tracecat/agent/common/config.py +++ b/tracecat/agent/common/config.py @@ -10,6 +10,22 @@ import os from pathlib import Path + +def _env_bool(var: str, *, default: bool) -> bool: + """Read a boolean env var without importing the full application config.""" + raw_value = os.environ.get(var) + if raw_value is None or not raw_value.strip(): + return default + + value = raw_value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + + raise ValueError(f"{var} must be a boolean value (got {raw_value!r})") + + # === Agent Sandbox Config (read directly from env) === # TRACECAT__AGENT_SANDBOX_TIMEOUT = int( @@ -22,9 +38,7 @@ ) """Default memory limit for agent sandbox execution in megabytes (4 GiB).""" -TRACECAT__DISABLE_NSJAIL = os.environ.get( - "TRACECAT__DISABLE_NSJAIL", "true" -).lower() in ("true", "1") +TRACECAT__DISABLE_NSJAIL = _env_bool("TRACECAT__DISABLE_NSJAIL", default=True) """Disable nsjail sandbox and use the unsafe PID executor instead.""" # === Well-known runtime paths (internal to agent worker) === # diff --git a/tracecat/agent/executor_worker.py b/tracecat/agent/executor_worker.py index f4976892ac..520682e0cf 100644 --- a/tracecat/agent/executor_worker.py +++ b/tracecat/agent/executor_worker.py @@ -131,7 +131,17 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: threadpool_max_workers = int( os.environ.get("TEMPORAL__THREADPOOL_MAX_WORKERS") or 100 ) - prepared_cgroup = prepare_agent_sandbox_cgroup() + cgroup_required = ( + config.TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED + and not config.TRACECAT__DISABLE_NSJAIL + ) + prepared_cgroup = prepare_agent_sandbox_cgroup(enabled=cgroup_required) + if cgroup_required: + cgroup_mount = prepared_cgroup.require_sandbox_mount() + logger.info( + "Agent sandbox cgroup memory limits ready", + cgroup_mount=str(cgroup_mount), + ) max_concurrent = clamp_agent_executor_concurrency( max_concurrent, prepared_cgroup, diff --git a/tracecat/agent/sandbox/cgroup.py b/tracecat/agent/sandbox/cgroup.py index 5a689c5938..cef169548c 100644 --- a/tracecat/agent/sandbox/cgroup.py +++ b/tracecat/agent/sandbox/cgroup.py @@ -36,6 +36,10 @@ class AgentExecutorMemoryBudgetError(RuntimeError): """Raised when the container memory budget cannot fit one sandbox.""" +class AgentSandboxCgroupUnavailableError(RuntimeError): + """Raised when enabled agent sandbox cgroup enforcement is unavailable.""" + + @dataclass(frozen=True, slots=True) class PreparedCgroup: """Process-wide agent sandbox cgroup preparation result.""" @@ -46,11 +50,31 @@ class PreparedCgroup: @property def sandbox_mount(self) -> Path | None: - """Return the nsjail cgroup mount when preparation succeeded.""" + """Return the nsjail cgroup mount unless cgroups were disabled. + + Raises: + AgentSandboxCgroupUnavailableError: If cgroups were enabled but + preparation failed. + """ if self.availability is CgroupAvailability.AVAILABLE: return self.root + if self.availability is CgroupAvailability.UNAVAILABLE: + raise AgentSandboxCgroupUnavailableError( + "Agent sandbox cgroup limits are enabled but unavailable. " + "Ensure the agent-executor entrypoint delegates its cgroup v2 " + "subtree before dropping privileges, or explicitly set " + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED=false." + ) return None + def require_sandbox_mount(self) -> Path: + """Return the prepared cgroup mount required by an enabled sandbox.""" + if (mount := self.sandbox_mount) is not None: + return mount + raise AgentSandboxCgroupUnavailableError( + "Agent sandbox cgroup limits are required but were disabled." + ) + @dataclass(frozen=True, slots=True) class CgroupMemoryLimit: @@ -109,16 +133,14 @@ def _warn_cgroup_unavailable( """Log the single startup warning for a failed cgroup preparation.""" if isinstance(error, OSError): logger.warning( - "Agent sandbox cgroup memory limits unavailable; " - "continuing without per-sandbox cgroup limits", + "Agent sandbox cgroup memory limits unavailable", step=step, errno=error.errno, error=str(error), ) return logger.warning( - "Agent sandbox cgroup memory limits unavailable; " - "continuing without per-sandbox cgroup limits", + "Agent sandbox cgroup memory limits unavailable", step=step, error=str(error), ) @@ -133,8 +155,8 @@ def prepare_agent_sandbox_cgroup( """Prepare the cgroup v2 root for nsjail child memory cgroups. The result is cached process-wide. Setup failures are reduced to one warning - and an unavailable result so agent sandbox launches can continue without - cgroup limits. + and an unavailable result. Callers that enabled cgroup enforcement must + reject that state before accepting sandbox work. """ global _prepared_cgroup diff --git a/tracecat/agent/sandbox/config.py b/tracecat/agent/sandbox/config.py index 655404aa8c..e19cbd7169 100644 --- a/tracecat/agent/sandbox/config.py +++ b/tracecat/agent/sandbox/config.py @@ -468,9 +468,9 @@ def build_agent_nsjail_config( [ "", "# Resource limits", - f"rlimit_as: {config.resources.memory_mb * 1024 * 1024}", + f"rlimit_as: {config.resources.memory_mb}", f"rlimit_cpu: {config.resources.cpu_seconds}", - f"rlimit_fsize: {config.resources.max_file_size_mb * 1024 * 1024}", + f"rlimit_fsize: {config.resources.max_file_size_mb}", f"rlimit_nofile: {config.resources.max_open_files}", f"rlimit_nproc: {config.resources.max_processes}", f"time_limit: {config.resources.timeout_seconds}", diff --git a/tracecat/config.py b/tracecat/config.py index 0a1d8e4194..6253486f70 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -764,7 +764,7 @@ def _parse_auth_types() -> set[AuthType]: TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED = env_bool( "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", default=True ) -"""Enable cgroup v2 memory limits for agent sandboxes when available.""" +"""Require cgroup v2 memory limits for agent sandboxes when nsjail is enabled.""" TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB = _env_int( "TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB", diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index 6826b5df23..d373a7bcb5 100644 --- a/tracecat/sandbox/executor.py +++ b/tracecat/sandbox/executor.py @@ -352,9 +352,9 @@ def _build_config( [ "", "# Resource limits", - f"rlimit_as: {config.resources.memory_mb * 1024 * 1024}", + f"rlimit_as: {config.resources.memory_mb}", f"rlimit_cpu: {config.resources.cpu_seconds}", - f"rlimit_fsize: {config.resources.max_file_size_mb * 1024 * 1024}", + f"rlimit_fsize: {config.resources.max_file_size_mb}", f"rlimit_nofile: {config.resources.max_open_files}", f"rlimit_nproc: {config.resources.max_processes}", f"time_limit: {config.resources.timeout_seconds}", @@ -794,15 +794,19 @@ def _build_action_config( [ "", "# Resource limits", - f"rlimit_as: {config.resources.memory_mb * 1024 * 1024}", + f"rlimit_as: {config.resources.memory_mb}", f"rlimit_cpu: {config.resources.cpu_seconds}", - f"rlimit_fsize: {config.resources.max_file_size_mb * 1024 * 1024}", + f"rlimit_fsize: {config.resources.max_file_size_mb}", f"rlimit_nofile: {config.resources.max_open_files}", f"rlimit_nproc: {config.resources.max_processes}", f"time_limit: {int(config.timeout_seconds)}", ] ) + # Regular action executors do not yet own a delegated cgroup subtree. + # Their aggregate process-tree limit remains a separate deployment + # project; retain the now-correct per-process MiB rlimits here. + # Execution settings - always use minimal_runner.py (untrusted mode) # minimal_runner.py is copied to /work and doesn't need tracecat imports lines.extend( From cb99c561ef03fc802d6772e8dba23a99ece44527 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:21:54 -0400 Subject: [PATCH 13/14] fix(agent): use sandbox overlay for compose nsjail --- docker-compose.local.yml | 11 ++---- docker-compose.yml | 11 ++---- tests/unit/test_config.py | 80 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/docker-compose.local.yml b/docker-compose.local.yml index e012c8768b..ce7b1c64b4 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -371,14 +371,9 @@ services: - core - core-db - temporal - # Required for nsjail sandbox (agent runtime isolation) - # cap_add: - # - SYS_ADMIN - # security_opt: - # - seccomp:unconfined - # # Required for pasta userspace networking (creates TAP device in sandbox netns) - # devices: - # - /dev/net/tun:/dev/net/tun + # To enable nsjail, set TRACECAT__DISABLE_NSJAIL=false and layer + # docker-compose.sandbox.yml; it configures the required privileges and + # cgroup delegation together. environment: # Common LOG_LEVEL: ${LOG_LEVEL} diff --git a/docker-compose.yml b/docker-compose.yml index 44d829fff2..e9c29e1657 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -362,14 +362,9 @@ services: - core - core-db - temporal - # Required for nsjail sandbox (uncomment if using nsjail) - # cap_add: - # - SYS_ADMIN - # security_opt: - # - seccomp:unconfined - # # Required for pasta userspace networking (creates TAP device in sandbox netns) - # devices: - # - /dev/net/tun:/dev/net/tun + # To enable nsjail, set TRACECAT__DISABLE_NSJAIL=false and layer + # docker-compose.sandbox.yml; it configures the required privileges and + # cgroup delegation together. environment: # Common LOG_LEVEL: ${LOG_LEVEL} diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 61bb128bde..ee016694a8 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -2,8 +2,12 @@ import ast import importlib +import json import re +import shutil +import subprocess from pathlib import Path +from typing import TypedDict, cast import pytest @@ -13,16 +17,32 @@ REPO_ROOT = Path(__file__).resolve().parents[2] CONFIG_PATH = REPO_ROOT / "tracecat" / "config.py" -COMPOSE_ENV_FILES = ( +SANDBOX_COMPOSE_PATH = REPO_ROOT / "docker-compose.sandbox.yml" +AGENT_EXECUTOR_BASE_COMPOSE_FILES = ( REPO_ROOT / "docker-compose.yml", REPO_ROOT / "docker-compose.dev.yml", REPO_ROOT / "docker-compose.local.yml", - REPO_ROOT / "docker-compose.sandbox.yml", +) +COMPOSE_ENV_FILES = ( + *AGENT_EXECUTOR_BASE_COMPOSE_FILES, + SANDBOX_COMPOSE_PATH, ) ENV_EXAMPLE_FILES = (REPO_ROOT / ".env.example",) DEPLOYMENT_ENV_FILES = (*COMPOSE_ENV_FILES, *ENV_EXAMPLE_FILES) +class _AgentExecutorComposeService(TypedDict): + user: str + entrypoint: list[str] + environment: list[str] + cap_add: list[str] + security_opt: list[str] + + +class _ComposeConfig(TypedDict): + services: dict[str, _AgentExecutorComposeService] + + def _config_bool_env_vars() -> set[str]: tree = ast.parse(CONFIG_PATH.read_text()) env_vars: set[str] = set() @@ -203,6 +223,62 @@ def test_boolean_env_values_preserve_defaults_and_compose_overrides() -> None: ) +@pytest.mark.parametrize( + "base_compose_path", + AGENT_EXECUTOR_BASE_COMPOSE_FILES, + ids=lambda path: path.name, +) +def test_agent_executor_sandbox_overlay_delegates_cgroups( + base_compose_path: Path, +) -> None: + docker_path = shutil.which("docker") + if docker_path is None: + pytest.skip("Docker CLI unavailable for Compose configuration test") + + compose_version = subprocess.run( + [docker_path, "compose", "version"], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + if compose_version.returncode != 0: + pytest.skip("Docker Compose plugin unavailable") + + result = subprocess.run( + [ + docker_path, + "compose", + "-f", + str(base_compose_path), + "-f", + str(SANDBOX_COMPOSE_PATH), + "config", + "--no-interpolate", + "--format", + "json", + ], + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, result.stderr + config = cast(_ComposeConfig, json.loads(result.stdout)) + service = config["services"]["agent-executor"] + + assert service["user"] == "0:0" + assert service["entrypoint"] == ["/usr/local/bin/agent-executor-entrypoint.sh"] + assert "SYS_ADMIN" in service["cap_add"] + assert "systempaths=unconfined" in service["security_opt"] + assert ( + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED=" + "${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true}" + ) in service["environment"] + + def test_bound_env_clamps_below_lower(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TEST_BOUND_ENV", "4") From 594c9c5a4fecc04cc707fd9c585d4b412cc01ff3 Mon Sep 17 00:00:00 2001 From: Daryl Lim <5508348+daryllimyt@users.noreply.github.com> Date: Tue, 28 Jul 2026 18:42:39 -0400 Subject: [PATCH 14/14] fix(agent): gate root cgroup delegation on nsjail being enabled --- docker/scripts/agent-executor-entrypoint.sh | 36 +++++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/docker/scripts/agent-executor-entrypoint.sh b/docker/scripts/agent-executor-entrypoint.sh index d521b33835..7fdfae63d5 100755 --- a/docker/scripts/agent-executor-entrypoint.sh +++ b/docker/scripts/agent-executor-entrypoint.sh @@ -12,21 +12,29 @@ if [[ "$(id -u)" == "0" ]]; then # Mirror config.env_bool falsy values so disabling the feature also skips # the root-side delegation, not just the Python-side preparation. The # privilege drop below is unconditional: root never reaches the worker. - cgroup_enabled="$( - printf '%s' "${TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED:-true}" | - tr '[:upper:]' '[:lower:]' | - sed 's/^[[:space:]]*//;s/[[:space:]]*$//' - )" - case "$cgroup_enabled" in - 1 | true | yes | on) cgroup_enabled=true ;; - 0 | false | no | off) cgroup_enabled=false ;; - *) - echo "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED must be a boolean value" >&2 - exit 64 - ;; - esac + parse_bool_env() { + local name="$1" default="$2" value + value="$( + printf '%s' "${!name:-$default}" | + tr '[:upper:]' '[:lower:]' | + sed 's/^[[:space:]]*//;s/[[:space:]]*$//' + )" + case "$value" in + 1 | true | yes | on) echo true ;; + 0 | false | no | off) echo false ;; + *) + echo "$name must be a boolean value" >&2 + exit 64 + ;; + esac + } + cgroup_enabled="$(parse_bool_env TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED true)" + nsjail_disabled="$(parse_bool_env TRACECAT__DISABLE_NSJAIL true)" - if [[ "$cgroup_enabled" == "true" ]]; then + # The worker only uses the delegated cgroup when nsjail sandboxing is on + # (executor_worker.py computes cgroup_required the same way), so skip the + # remount + chown entirely when nsjail is disabled. + if [[ "$cgroup_enabled" == "true" && "$nsjail_disabled" == "false" ]]; then # Resolve this container's cgroup v2 directory: the mount root under a # private cgroup namespace, a subpath of the host cgroupfs otherwise # (e.g. privileged Kubernetes). Never touch anything above it, and