diff --git a/.env.example b/.env.example index aa2d4c110..b5b202d93 100644 --- a/.env.example +++ b/.env.example @@ -153,6 +153,12 @@ 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 +# 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/Dockerfile b/Dockerfile index 4a5d1e114..a0dcb6b33 100644 --- a/Dockerfile +++ b/Dockerfile @@ -227,6 +227,9 @@ ENV PYTHONPATH="/home/apiuser/.local" RUN mkdir -p /home/apiuser/.local/bin && ln -s $(which uv) /home/apiuser/.local/bin/uv +COPY docker/scripts/agent-executor-entrypoint.sh /usr/local/bin/agent-executor-entrypoint.sh +RUN chmod +x /usr/local/bin/agent-executor-entrypoint.sh + # Switch to non-root user (matches production, required for pasta userspace networking) USER apiuser @@ -278,6 +281,9 @@ ENV TMPDIR="/home/apiuser/.cache/tmp" TEMP="/home/apiuser/.cache/tmp" TMP="/home RUN mkdir -p /app/.scripts && chown -R apiuser:apiuser /app +COPY docker/scripts/agent-executor-entrypoint.sh /usr/local/bin/agent-executor-entrypoint.sh +RUN chmod +x /usr/local/bin/agent-executor-entrypoint.sh + # Switch to non-root user USER apiuser diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 56a3f8716..a8a7bd717 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -406,6 +406,10 @@ 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__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} @@ -430,6 +434,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", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/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 c178f3619..ce7b1c64b 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} @@ -428,6 +423,10 @@ 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__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} @@ -450,6 +449,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", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/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 d978d2160..9b4268de2 100644 --- a/docker-compose.sandbox.yml +++ b/docker-compose.sandbox.yml @@ -21,8 +21,23 @@ services: # Required for nsjail sandbox execution. cap_add: - SYS_ADMIN + # Start as root so the entrypoint can delegate the container's cgroup v2 + # subtree to apiuser for per-sandbox memory limits; it drops privileges + # before the worker starts. Override neither independently. + user: "0:0" + entrypoint: ["/usr/local/bin/agent-executor-entrypoint.sh"] + 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} + 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", "${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 00d604c1f..e9c29e165 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} @@ -419,6 +414,10 @@ 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__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} @@ -445,6 +444,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", "${TRACECAT__AGENT_EXECUTOR_READY_FILE:-/var/run/tracecat/agent-executor-ready}"] + interval: 30s + retries: 3 + start_period: 120s depends_on: litellm: condition: service_healthy diff --git a/docker/scripts/agent-executor-entrypoint.sh b/docker/scripts/agent-executor-entrypoint.sh new file mode 100755 index 000000000..7fdfae63d --- /dev/null +++ b/docker/scripts/agent-executor-entrypoint.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Agent-executor entrypoint for deployments that enable per-sandbox cgroup +# 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. 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. + 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)" + + # 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 + # 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 cannot be enforced." >&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 cannot be enforced." >&2 + fi + fi + 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 "$@" +fi + +exec "$@" 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 000000000..4d5735c9c --- /dev/null +++ b/tests/unit/test_agent_executor_worker_resources.py @@ -0,0 +1,388 @@ +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 ( + AgentExecutorMemoryBudgetError, + AgentSandboxCgroupUnavailableError, + CgroupAvailability, + PreparedCgroup, +) + + +@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, + prepared_cgroup: PreparedCgroup, + *, + reserve_mb: int, + sandbox_memory_mb: int, + ) -> int: + del prepared_cgroup, 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 *, enabled: PreparedCgroup(CgroupAvailability.DISABLED, None), + ) + 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 + + +@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, + prepared_cgroup: PreparedCgroup, + *, + reserve_mb: int, + sandbox_memory_mb: int, + ) -> int: + 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 *, enabled: PreparedCgroup(CgroupAvailability.DISABLED, None), + ) + 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() + + +@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 *, enabled: 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 +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 *, enabled: 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 +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( + 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 new file mode 100644 index 000000000..bed9beeae --- /dev/null +++ b/tests/unit/test_agent_sandbox_cgroup.py @@ -0,0 +1,571 @@ +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, + AgentExecutorMemoryBudgetError, + AgentSandboxCgroupUnavailableError, + 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_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", + 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("") + (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_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", + ) + cgroupfs = tmp_path / "cgroupfs" + + result = detect_cgroup_root(proc_cgroup_path, cgroupfs) + + assert result == cgroupfs / "kubepods.slice/pod.scope/main" + + +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) + with pytest.raises(AgentSandboxCgroupUnavailableError): + _ = result.sandbox_mount + + +def test_prepare_agent_sandbox_cgroup_returns_cached_disabled_state( + tmp_path: Path, +) -> None: + proc_cgroup_path = tmp_path / "missing" + cgroupfs = tmp_path / "cgroupfs" + + first = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=False, + ) + second = prepare_agent_sandbox_cgroup( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + + 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 + + +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( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + + 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" + assert not list(cgroup_root.glob("tracecat-agent-probe-*")) + + +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, + tmp_path / "cgroupfs", + ) + 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: + 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( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + + 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" + + +def test_prepare_agent_sandbox_cgroup_retains_root_on_permission_error( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + 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 + + 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( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + + assert result == PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + cgroupfs, + ) + 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" + 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: + 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 + + 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( + proc_cgroup_path=proc_cgroup_path, + cgroupfs=cgroupfs, + enabled=True, + ) + + assert result.availability 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_from_detected_root( + tmp_path: Path, + contents: str, + expected_kind: CgroupMemoryLimitKind, + expected_bytes: int | None, +) -> None: + 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(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_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) + + assert result.kind is CgroupMemoryLimitKind.UNAVAILABLE + assert result.limit_bytes is None + + +@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: + _, 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(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") + + # 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( + 10, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + == 2 + ) + + +@pytest.mark.parametrize( + ("container_limit_mb", "reserve_mb"), + [ + pytest.param(4096, 4096, id="exactly-zero-slots"), + pytest.param(4096, 8192, id="reserve-exceeds-limit"), + ], +) +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, +) -> None: + _, 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( + 10, + prepared_cgroup, + reserve_mb=4096, + sandbox_memory_mb=4096, + ) + + 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 18913f8bb..28a4cabce 100644 --- a/tests/unit/test_agent_sandbox_config.py +++ b/tests/unit/test_agent_sandbox_config.py @@ -3,6 +3,7 @@ from pathlib import Path from tracecat.agent.sandbox.config import ( + AgentResourceLimits, AgentSandboxConfig, build_agent_nsjail_config, ) @@ -121,3 +122,40 @@ 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_mount=Path("/sys/fs/cgroup/kubepods.slice/pod.scope"), + ) + + 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 + + +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"), + 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"), + ) + + 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/tests/unit/test_agent_sandbox_litellm.py b/tests/unit/test_agent_sandbox_litellm.py index 1afa198f3..8c4356bef 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 3ee60d7d1..ee016694a 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -2,25 +2,47 @@ import ast import importlib +import json import re +import shutil +import subprocess from pathlib import Path +from typing import TypedDict, cast 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] 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", ) +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() @@ -97,6 +119,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 = { @@ -169,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") @@ -220,6 +330,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 66f0f9146..3ac9255d0 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,34 @@ 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 "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 def test_action_sandbox_config_includes_seccomp_policy(tmp_path: Path): @@ -63,10 +91,38 @@ 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( + 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): diff --git a/tests/unit/test_stdio_probe.py b/tests/unit/test_stdio_probe.py index 57429cd11..f7ae945ee 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 @@ -227,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", @@ -267,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", @@ -296,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 e7c2f0d93..72cb45b1a 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 c4d8b5289..520682e0c 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 [ @@ -80,12 +114,40 @@ 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 ) + 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 ) + 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, + reserve_mb=config.TRACECAT__AGENT_EXECUTOR_MEMORY_RESERVE_MB, + sandbox_memory_mb=config.TRACECAT__AGENT_SANDBOX_MEMORY_MB, + ) logger.info( "Starting AgentExecutorWorker", @@ -117,10 +179,12 @@ async def main(shutdown_event: asyncio.Event | None = None) -> None: ), ): logger.info("AgentExecutorWorker started, ctrl+c to exit") + _write_readiness_file(readiness_file, datetime.now(UTC)) await shutdown_event.wait() logger.info("AgentExecutorWorker shutdown requested") logger.info("Temporal Worker context exited") finally: + _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/mcp/stdio_probe.py b/tracecat/agent/mcp/stdio_probe.py index d5134c3c7..b26cbd0a3 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 new file mode 100644 index 000000000..cef169548 --- /dev/null +++ b/tracecat/agent/sandbox/cgroup.py @@ -0,0 +1,373 @@ +"""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 + +PROC_SELF_CGROUP = Path("/proc/self/cgroup") +CGROUPFS = Path("/sys/fs/cgroup") +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" + + +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.""" + + availability: CgroupAvailability + root: Path | None + cgroupfs: Path = CGROUPFS + + @property + def sandbox_mount(self) -> Path | None: + """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: + """Parsed cgroup v2 container memory limit.""" + + kind: CgroupMemoryLimitKind + limit_bytes: int | 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 + + return cgroupfs if cgroup_path == "/" else cgroupfs / cgroup_path.lstrip("/") + return 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", + step=step, + errno=error.errno, + error=str(error), + ) + return + logger.warning( + "Agent sandbox cgroup memory limits unavailable", + step=step, + error=str(error), + ) + + +def prepare_agent_sandbox_cgroup( + *, + proc_cgroup_path: Path = PROC_SELF_CGROUP, + cgroupfs: Path = CGROUPFS, + enabled: bool = config.TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED, +) -> 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. Callers that enabled cgroup enforcement must + reject that state before accepting sandbox work. + """ + global _prepared_cgroup + + if _prepared_cgroup is not None: + return _prepared_cgroup + if not enabled: + _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, + cgroupfs, + ) + return _prepared_cgroup + + probe_cgroup: Path | None = None + step = "read cgroup v2 controllers" + try: + 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" + 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) + _prepared_cgroup = PreparedCgroup( + CgroupAvailability.UNAVAILABLE, + cgroup_root, + cgroupfs, + ) + return _prepared_cgroup + + _prepared_cgroup = PreparedCgroup( + CgroupAvailability.AVAILABLE, + cgroup_root, + cgroupfs, + ) + return _prepared_cgroup + + +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( + cgroup_root: Path | None, + cgroupfs: Path = CGROUPFS, +) -> CgroupMemoryLimit: + """Read the effective cgroup v2 memory limit for the executor's cgroup.""" + if cgroup_root is None: + 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": + 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 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) + if unlimited_seen: + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNLIMITED) + return CgroupMemoryLimit(CgroupMemoryLimitKind.UNAVAILABLE) + + +def clamp_agent_executor_concurrency( + max_concurrent: int, + prepared_cgroup: PreparedCgroup, + *, + reserve_mb: int, + sandbox_memory_mb: int, +) -> int: + """Clamp agent executor concurrency to the container memory budget.""" + 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" + ) + 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 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 + + 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=allowed, + ) + return allowed diff --git a/tracecat/agent/sandbox/config.py b/tracecat/agent/sandbox/config.py index 4e1caab8f..e19cbd716 100644 --- a/tracecat/agent/sandbox/config.py +++ b/tracecat/agent/sandbox/config.py @@ -222,6 +222,7 @@ def build_agent_nsjail_config( session_work_dir: Path | None = None, enable_internet_access: bool = False, skills_dir: Path | None = None, + cgroup_mount: Path | None = None, ) -> str: """Build nsjail protobuf config for agent runtime execution. @@ -245,6 +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_mount: Optional cgroup v2 directory for nsjail child cgroups. Returns: nsjail protobuf configuration as a string. @@ -272,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 @@ -464,14 +468,23 @@ 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}", ] ) + if cgroup_mount is not None: + lines.extend( + [ + "use_cgroupv2: true", + f'cgroupv2_mount: "{cgroup_mount}"', + 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 901d9e6ba..5a26aab35 100644 --- a/tracecat/agent/sandbox/nsjail.py +++ b/tracecat/agent/sandbox/nsjail.py @@ -46,16 +46,14 @@ JAILED_AGENT_JOB_DIR, JAILED_AGENT_WORK_DIR, ) +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__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 @@ -439,6 +437,7 @@ async def _spawn_nsjail_runtime( session_work_dir=session_work_dir, enable_internet_access=enable_internet_access, skills_dir=skills_dir, + cgroup_mount=get_agent_sandbox_cgroup().sandbox_mount, ) # Write config to job directory diff --git a/tracecat/config.py b/tracecat/config.py index 7f36391f4..6253486f7 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.""" @@ -738,11 +754,31 @@ 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).""" +TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED = env_bool( + "TRACECAT__AGENT_SANDBOX_CGROUP_ENABLED", default=True +) +"""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", + 4096, + min_value=0, +) +"""Memory reserved for the agent executor worker and shared services.""" + +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.""" + TRACECAT__LITELLM_PORT = int(os.environ.get("TRACECAT__LITELLM_PORT") or 4000) """Bind port for the managed LiteLLM service.""" diff --git a/tracecat/sandbox/executor.py b/tracecat/sandbox/executor.py index 5ad8b8c4d..d373a7bcb 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: @@ -348,14 +352,23 @@ 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}", ] ) + 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}" @@ -781,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(