feat(ai-red-teaming): tear down provisioned environments on assessment completion

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:**
 

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,

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() == ""

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 ""