diff --git a/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md b/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md index d7efb2c..41179ef 100644 --- a/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md +++ b/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md @@ -162,7 +162,8 @@ The AI Red Teaming capability provides these tools: **Multi-Agent Environments:** - **list_environments** — List the deployable multi-agent environments (e.g. `finops-mesh`, `devsecops-mesh`, `healthcare-mesh`, `soc-mesh`) that ATLAS can target -- **provision_environment** — Deploy a hosted multi-agent environment (passing the model its agents use) and return its `/attack` URL + execute token. Chain into `generate_atlas_attack` to probe it — closing the loop from Environment to ATLAS probe. +- **provision_environment** — Deploy a hosted multi-agent environment (passing the model its agents use) and return its `id`, `/attack` URL + execute token. Chain into `generate_atlas_attack` to probe it — closing the loop from Environment to ATLAS probe. The sandbox is recorded and torn down automatically when the assessment completes. +- **teardown_environment** — Delete provisioned environment sandboxes to stop billing. Hosted sandboxes bill for their whole lifetime. With no id it reaps every environment provisioned this session; pass an id to reap one. Teardown also runs automatically when `update_assessment_status` marks the assessment complete, so call this only to reap early or after a partial run. **Workflow Management:** diff --git a/capabilities/ai-red-teaming/capability.yaml b/capabilities/ai-red-teaming/capability.yaml index 4447e38..1874b3c 100644 --- a/capabilities/ai-red-teaming/capability.yaml +++ b/capabilities/ai-red-teaming/capability.yaml @@ -1,6 +1,6 @@ schema: 1 name: ai-red-teaming -version: "1.11.0" +version: "1.12.0" description: > Probe the security and safety of AI applications, agents, and foundation models. Orchestrates adversarial attack workflows to discover vulnerabilities in LLMs, diff --git a/capabilities/ai-red-teaming/tests/test_environments_teardown.py b/capabilities/ai-red-teaming/tests/test_environments_teardown.py new file mode 100644 index 0000000..e4f3786 --- /dev/null +++ b/capabilities/ai-red-teaming/tests/test_environments_teardown.py @@ -0,0 +1,243 @@ +"""Tests for tools/environments.py — session registry + environment teardown. + +A hosted sandbox bills for its whole lifetime, so every provisioned environment +is registered and reaped when the assessment completes (or via teardown_environment). +These tests exercise the pure teardown/registry logic with an injected fake API +client and a temp registry file — no network. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("dreadnode.agents.tools") + +TOOL_PATH = Path(__file__).resolve().parents[1] / "tools" / "environments.py" +ASSESSMENT_PATH = Path(__file__).resolve().parents[1] / "tools" / "assessment.py" + + +def _load(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +env = _load(TOOL_PATH, "airt_environments_under_test") + + +class _FakeApi: + """Records delete_environment calls; can be told to 404 specific ids.""" + + def __init__(self, missing: set[str] | None = None, boom: set[str] | None = None): + self.deleted: list[tuple[str, str, str]] = [] + self._missing = missing or set() + self._boom = boom or set() + + def delete_environment(self, org: str, workspace: str, env_id: str) -> None: + self.deleted.append((org, workspace, env_id)) + if env_id in self._boom: + raise RuntimeError("provider exploded") + if env_id in self._missing: + raise _NotFound("environment not found (404)") + + +class _NotFound(Exception): + pass + + +class _FakeEnv: + def __init__(self, env_id): + self.id = env_id + + +@pytest.fixture +def registry(tmp_path): + return tmp_path / "environments.json" + + +class TestRegistry: + def test_register_provisioned_adds_and_returns_id(self, registry) -> None: + env_id = env._register_provisioned( + _FakeEnv("env-1"), "ml-extraction-mnist-image", "org", "main", registry + ) + assert env_id == "env-1" + entries = env._registry_load(registry) + assert len(entries) == 1 + assert entries[0]["id"] == "env-1" + assert entries[0]["task_ref"] == "ml-extraction-mnist-image" + assert "provisioned_at_ts" in entries[0] + + def test_register_no_id_is_noop(self, registry) -> None: + assert env._register_provisioned(_FakeEnv(None), "x", "o", "w", registry) == "" + assert env._registry_load(registry) == [] + + def test_register_dedups_same_id(self, registry) -> None: + env._register_provisioned(_FakeEnv("dup"), "a", "o", "w", registry) + env._register_provisioned(_FakeEnv("dup"), "b", "o", "w", registry) + entries = env._registry_load(registry) + assert len(entries) == 1 + assert entries[0]["task_ref"] == "b" + + def test_load_tolerates_missing_and_corrupt(self, registry) -> None: + assert env._registry_load(registry) == [] + registry.write_text("{ not json") + assert env._registry_load(registry) == [] + + +class TestTeardown: + def test_teardown_all_deletes_and_empties_registry(self, registry) -> None: + env._register_provisioned(_FakeEnv("a"), "t", "org", "main", registry) + env._register_provisioned(_FakeEnv("b"), "t", "org", "main", registry) + api = _FakeApi() + result = env._teardown_environments(api, "org", "main", registry_path=registry) + assert set(result["torn_down"]) == {"a", "b"} + assert result["errors"] == {} + assert env._registry_load(registry) == [] + assert {d[2] for d in api.deleted} == {"a", "b"} + + def test_teardown_single_id_leaves_others(self, registry) -> None: + env._register_provisioned(_FakeEnv("a"), "t", "org", "main", registry) + env._register_provisioned(_FakeEnv("b"), "t", "org", "main", registry) + api = _FakeApi() + result = env._teardown_environments( + api, "org", "main", environment_id="a", registry_path=registry + ) + assert result["torn_down"] == ["a"] + remaining = [e["id"] for e in env._registry_load(registry)] + assert remaining == ["b"] + + def test_teardown_missing_env_counts_as_success(self, registry) -> None: + env._register_provisioned(_FakeEnv("gone"), "t", "org", "main", registry) + api = _FakeApi(missing={"gone"}) + # _is_not_found matches on "404" in the message, so no SDK import needed. + result = env._teardown_environments(api, "org", "main", registry_path=registry) + assert result["torn_down"] == ["gone"] + assert result["errors"] == {} + assert env._registry_load(registry) == [] + + def test_teardown_provider_error_is_reported_and_kept(self, registry) -> None: + env._register_provisioned(_FakeEnv("bad"), "t", "org", "main", registry) + api = _FakeApi(boom={"bad"}) + result = env._teardown_environments(api, "org", "main", registry_path=registry) + assert result["torn_down"] == [] + assert "bad" in result["errors"] + # A failed delete stays in the registry so a later sweep retries it. + assert [e["id"] for e in env._registry_load(registry)] == ["bad"] + + def test_grace_window_skips_recent_envs(self, registry) -> None: + env._register_provisioned(_FakeEnv("fresh"), "t", "org", "main", registry) + entries = env._registry_load(registry) + now = entries[0]["provisioned_at_ts"] + 5 # env is 5s old + api = _FakeApi() + result = env._teardown_environments( + api, "org", "main", older_than_sec=60, registry_path=registry, now_ts=now + ) + assert result["skipped"] == ["fresh"] + assert result["torn_down"] == [] + assert api.deleted == [] # never called the platform + + def test_grace_window_reaps_old_envs(self, registry) -> None: + env._register_provisioned(_FakeEnv("old"), "t", "org", "main", registry) + entries = env._registry_load(registry) + now = entries[0]["provisioned_at_ts"] + 120 # 2 minutes old + api = _FakeApi() + result = env._teardown_environments( + api, "org", "main", older_than_sec=60, registry_path=registry, now_ts=now + ) + assert result["torn_down"] == ["old"] + + def test_teardown_empty_registry_is_noop(self, registry) -> None: + api = _FakeApi() + result = env._teardown_environments(api, "org", "main", registry_path=registry) + assert result == {"torn_down": [], "skipped": [], "errors": {}} + assert api.deleted == [] + + +class TestSessionHelper: + def test_session_teardown_short_circuits_without_config(self, registry, monkeypatch) -> None: + # Empty registry: must not call _configured() at all (no network). + called = {"configured": False} + + def _boom(): + called["configured"] = True + raise AssertionError("_configured must not be called for an empty registry") + + monkeypatch.setattr(env, "_configured", _boom) + result = env.teardown_session_environments(registry_path=registry) + assert result["torn_down"] == [] + assert called["configured"] is False + + def test_session_teardown_reaps_via_configured(self, registry, monkeypatch) -> None: + env._register_provisioned(_FakeEnv("s1"), "t", "org", "main", registry) + api = _FakeApi() + monkeypatch.setattr(env, "_configured", lambda: (None, api, "org", "main")) + result = env.teardown_session_environments(registry_path=registry) + assert result["torn_down"] == ["s1"] + assert api.deleted == [("org", "main", "s1")] + + +class TestAssessmentCompletionHook: + """update_assessment_status must reap environments only when the assessment completes.""" + + @pytest.fixture + def assessment(self, tmp_path, monkeypatch): + monkeypatch.setenv("AIRT_ASSESSMENT_PATH", str(tmp_path / "assessment.json")) + return _load(ASSESSMENT_PATH, "airt_assessment_under_test") + + def test_completion_triggers_teardown_and_appends_note(self, assessment, monkeypatch) -> None: + calls = {"n": 0} + + def _spy() -> str: + calls["n"] += 1 + return " Assessment complete - tore down 2 environment(s) to stop billing." + + monkeypatch.setattr(assessment, "_teardown_on_complete", _spy) + assessment.register_assessment( + name="demo", target="ml_classifier", planned_attacks=["hopskipjump_evasion"] + ) + out = assessment.update_assessment_status( + attack_name="hopskipjump_evasion", status="completed" + ) + assert calls["n"] == 1 + assert "tore down 2 environment(s)" in out + + def test_partial_progress_does_not_tear_down(self, assessment, monkeypatch) -> None: + calls = {"n": 0} + monkeypatch.setattr( + assessment, "_teardown_on_complete", lambda: calls.__setitem__("n", calls["n"] + 1) or "" + ) + assessment.register_assessment( + name="demo", + target="ml_classifier", + planned_attacks=["hopskipjump_evasion", "pwws_evasion"], + ) + assessment.update_assessment_status(attack_name="hopskipjump_evasion", status="completed") + assert calls["n"] == 0 # one of two done — not complete yet + + def test_teardown_fires_once_not_on_repeat_updates(self, assessment, monkeypatch) -> None: + calls = {"n": 0} + monkeypatch.setattr( + assessment, "_teardown_on_complete", lambda: calls.__setitem__("n", calls["n"] + 1) or "" + ) + assessment.register_assessment( + name="demo", target="ml_classifier", planned_attacks=["hopskipjump_evasion"] + ) + assessment.update_assessment_status(attack_name="hopskipjump_evasion", status="completed") + # Re-recording the same attack must not re-trigger teardown. + assessment.update_assessment_status(attack_name="hopskipjump_evasion", status="completed") + assert calls["n"] == 1 + + def test_real_teardown_on_complete_returns_empty_when_nothing_registered( + self, assessment, monkeypatch + ) -> None: + # The real hook, with an empty/absent registry, short-circuits with no + # platform call and returns "" (no note appended). + monkeypatch.setenv("AIRT_ENV_REGISTRY_PATH", str(Path("/nonexistent/dir/registry.json"))) + assert assessment._teardown_on_complete() == "" diff --git a/capabilities/ai-red-teaming/tools/assessment.py b/capabilities/ai-red-teaming/tools/assessment.py index e7b0eae..43dcec9 100644 --- a/capabilities/ai-red-teaming/tools/assessment.py +++ b/capabilities/ai-red-teaming/tools/assessment.py @@ -150,12 +150,45 @@ def update_assessment_status( # Auto-complete assessment if all planned attacks are done planned = data.get("planned_attacks", []) completed_names = {c["attack_name"] for c in completed} - if all(a in completed_names for a in planned): + just_completed = ( + bool(planned) + and all(a in completed_names for a in planned) + and data.get("status") != "completed" + ) + if just_completed: data["status"] = "completed" _save(data) + # On completion, reap provisioned sandboxes so they stop billing. Best-effort: + # never let teardown failure break status recording. + teardown_note = _teardown_on_complete() if just_completed else "" + total = len(planned) done = len(completed) asr_str = f" (ASR={asr}%)" if asr is not None else "" - return f"Recorded {attack_name}: {status}{asr_str}. Progress: {done}/{total}." + return f"Recorded {attack_name}: {status}{asr_str}. Progress: {done}/{total}.{teardown_note}" + + +def _teardown_on_complete() -> str: + """Tear down every provisioned environment now that the assessment is complete. + + Loads the environments tool as a flat module (capability tool files have no + parent package) and reaps the session registry. ``AIRT_ENV_TEARDOWN_GRACE_SEC`` + (default 0) leaves very recently provisioned sandboxes alone so a trailing + attack is not killed mid-run. Returns a note to append, or "" on any failure. + """ + try: + grace = float(os.environ.get("AIRT_ENV_TEARDOWN_GRACE_SEC", "0") or "0") + except ValueError: + grace = 0.0 + try: + env_path = _Path(__file__).resolve().parent / "environments.py" + spec = _ilu.spec_from_file_location("airt_tools_environments", env_path) + mod = _ilu.module_from_spec(spec) + spec.loader.exec_module(mod) + result = mod.teardown_session_environments(older_than_sec=grace) + n = len(result.get("torn_down", [])) + return f" Assessment complete - tore down {n} environment(s) to stop billing." if n else "" + except Exception: # noqa: BLE001 - teardown is best-effort, must not break status + return "" diff --git a/capabilities/ai-red-teaming/tools/environments.py b/capabilities/ai-red-teaming/tools/environments.py index 21c71ff..4a2d630 100644 --- a/capabilities/ai-red-teaming/tools/environments.py +++ b/capabilities/ai-red-teaming/tools/environments.py @@ -17,7 +17,11 @@ import asyncio import importlib.util as _ilu +import json as _json +import os as _os +import time as _time import typing as t +from datetime import datetime, timezone from pathlib import Path as _Path # Load the shared safe_tool wrapper by file path (flat-module loading). @@ -27,6 +31,155 @@ _spec.loader.exec_module(_errors_mod) safe_tool = _errors_mod.safe_tool +# Session registry of provisioned environments. Hosted sandboxes bill for their +# whole lifetime, so every provision is recorded here and torn down when the +# assessment completes (see tools/assessment.py) or via teardown_environment. +# File-based so it survives across separate tool invocations in one session; +# path is overridable for tests. +REGISTRY_PATH = _Path(_os.environ.get("AIRT_ENV_REGISTRY_PATH", "/tmp/airt_environments.json")) + + +def _iso_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _registry_load(registry_path: _Path = REGISTRY_PATH) -> list[dict]: + """Return the provisioned-environment records, tolerating a missing/corrupt file.""" + try: + if registry_path.exists(): + data = _json.loads(registry_path.read_text()) + if isinstance(data, list): + return [e for e in data if isinstance(e, dict)] + except (OSError, ValueError): + pass + return [] + + +def _registry_save(entries: list[dict], registry_path: _Path = REGISTRY_PATH) -> None: + registry_path.parent.mkdir(parents=True, exist_ok=True) + registry_path.write_text(_json.dumps(entries, indent=2)) + + +def _registry_add(entry: dict, registry_path: _Path = REGISTRY_PATH) -> None: + entries = [e for e in _registry_load(registry_path) if e.get("id") != entry.get("id")] + entries.append(entry) + _registry_save(entries, registry_path) + + +def _registry_remove(env_id: str, registry_path: _Path = REGISTRY_PATH) -> None: + _registry_save( + [e for e in _registry_load(registry_path) if e.get("id") != env_id], registry_path + ) + + +def _register_provisioned( + env: t.Any, task_ref: str, org: str, workspace: str, registry_path: _Path = REGISTRY_PATH +) -> str: + """Record a just-provisioned environment in the session registry. Returns its id + (empty string if the environment exposed none, e.g. a provision that never became + ready).""" + env_id = getattr(env, "id", None) or "" + if not env_id: + return "" + _registry_add( + { + "id": env_id, + "task_ref": task_ref, + "org": org, + "workspace": workspace, + "provisioned_at": _iso_now(), + "provisioned_at_ts": _time.time(), + }, + registry_path, + ) + return env_id + + +def _entry_age(entry: dict, now: float) -> float | None: + ts = entry.get("provisioned_at_ts") + if isinstance(ts, (int, float)): + return max(0.0, now - float(ts)) + return None + + +def _is_not_found(exc: BaseException) -> bool: + """A delete of an already-gone environment is success, not failure.""" + try: + from dreadnode.app.api.client import NotFoundError + + if isinstance(exc, NotFoundError): + return True + except Exception: # noqa: BLE001 - NotFoundError may be absent on older SDKs + pass + return "notfound" in exc.__class__.__name__.lower() or "404" in str(exc) + + +def _teardown_environments( + api: t.Any, + org: str, + workspace: str, + *, + environment_id: str = "", + older_than_sec: float = 0.0, + registry_path: _Path = REGISTRY_PATH, + now_ts: float | None = None, +) -> dict: + """Delete provisioned environments and prune the registry. + + Pure and testable: the API client is injected. Idempotent - deleting an + already-gone environment counts as torn down. ``older_than_sec`` is a grace + window: an environment provisioned more recently than that is left alone so + an in-flight attack is not killed out from under it. + + Returns ``{"torn_down": [ids], "skipped": [ids], "errors": {id: msg}}``. + """ + now = _time.time() if now_ts is None else now_ts + entries = _registry_load(registry_path) + if environment_id: + targets = [e for e in entries if e.get("id") == environment_id] or [ + {"id": environment_id} + ] + else: + targets = list(entries) + + torn: list[str] = [] + skipped: list[str] = [] + errors: dict[str, str] = {} + for entry in targets: + env_id = entry.get("id") + if not env_id: + continue + age = _entry_age(entry, now) + if older_than_sec > 0 and age is not None and age < older_than_sec: + skipped.append(env_id) + continue + try: + api.delete_environment(org, workspace, env_id) + torn.append(env_id) + _registry_remove(env_id, registry_path) + except Exception as exc: # noqa: BLE001 - tolerate already-gone / provider errors + if _is_not_found(exc): + torn.append(env_id) + _registry_remove(env_id, registry_path) + else: + errors[env_id] = str(exc) + return {"torn_down": torn, "skipped": skipped, "errors": errors} + + +def teardown_session_environments( + older_than_sec: float = 0.0, registry_path: _Path = REGISTRY_PATH +) -> dict: + """Reap every environment in the session registry. Used by the assessment + completion hook. Short-circuits (no platform call) when nothing is registered.""" + if not _registry_load(registry_path): + return {"torn_down": [], "skipped": [], "errors": {}} + _inst, api, org, workspace = _configured() + if not org or not workspace: + return {"torn_down": [], "skipped": [], "errors": {}} + return _teardown_environments( + api, org, workspace, older_than_sec=older_than_sec, registry_path=registry_path + ) + def _run(coro: t.Any) -> t.Any: """Run an async coroutine from a sync tool, whether or not a loop is running.""" @@ -113,11 +266,15 @@ def provision_environment( svc = (ctx.get("service_urls") or {}).get("challenge") url = (svc.get("url") if isinstance(svc, dict) else svc) or "" token = env._execute_token or "" # noqa: SLF001 - one-shot provision token + # Record the sandbox so it is torn down at assessment completion even if the + # attack path forgets — a hosted sandbox bills for its whole lifetime. + env_id = _register_provisioned(env, task_ref, org, workspace) if not url: return f"Environment '{task_ref}' provisioned but exposed no 'challenge' URL: {ctx.get('service_urls')}" return ( f"Environment '{task_ref}' is ready.\n" + f" Environment id: {env_id}\n" f" Attack URL: {url}/attack\n" f" Auth: bearer (execute token below)\n" f" Execute token: {token}\n" @@ -125,5 +282,48 @@ def provision_environment( f">>> NEXT STEP: run ATLAS against it — call generate_atlas_attack(" f"agent_url=\"{url}/attack\", agent_auth_type=\"bearer\", " f"scenario_name=\"{task_ref.replace('-mesh', '')}\", attacker_model=\"groq scout\") " - f"and set AGENT_API_KEY to the execute token above." + f"and set AGENT_API_KEY to the execute token above.\n" + f">>> WHEN DONE: this sandbox bills for its whole lifetime — it is torn down " + f"automatically when the assessment completes, or call teardown_environment() now." + ) + + +@safe_tool +def teardown_environment( + environment_id: t.Annotated[ + str, + "Environment id to tear down. Leave empty to tear down EVERY environment " + "provisioned in this session.", + ] = "", + older_than_sec: t.Annotated[ + int, + "Only tear down environments provisioned at least this many seconds ago " + "(0 = no age filter). Guards a still-running attack from being killed.", + ] = 0, +) -> str: + """Tear down (delete) provisioned environment sandboxes to stop billing. + + Call this once your attacks are done. Hosted sandboxes bill for their whole + lifetime, so leaving them running costs credits. With no ``environment_id`` + this reaps every environment provisioned in the current session; pass an id + to reap just one. Idempotent - an environment already gone counts as torn + down. + """ + _inst, api, org, workspace = _configured() + if not org or not workspace: + return "Not configured for a platform org/workspace. Run `dreadnode login` first." + result = _teardown_environments( + api, + org, + workspace, + environment_id=environment_id, + older_than_sec=float(older_than_sec), ) + parts = [f"Tore down {len(result['torn_down'])} environment(s)."] + if result["skipped"]: + parts.append(f"Skipped {len(result['skipped'])} within the grace window.") + if result["errors"]: + parts.append( + "Errors: " + "; ".join(f"{k}: {v}" for k, v in result["errors"].items()) + ) + return " ".join(parts)