From 519b64eabc680a100ec0c039ab9f33686ac6bb3d Mon Sep 17 00:00:00 2001 From: GangGreenTemperTatum <104169244+GangGreenTemperTatum@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:05:39 -0400 Subject: [PATCH] fix(ai-red-teaming): register provision_environment and generate_agentic_suite_attack as tools (ENG-8467) --- capabilities/ai-red-teaming/capability.yaml | 2 +- .../skills/error-troubleshooting/SKILL.md | 4 +- .../ai-red-teaming/tests/test_tool_surface.py | 233 ++++++++++++++++++ capabilities/ai-red-teaming/tools/attacks.py | 1 + .../ai-red-teaming/tools/environments.py | 28 ++- 5 files changed, 259 insertions(+), 9 deletions(-) create mode 100644 capabilities/ai-red-teaming/tests/test_tool_surface.py diff --git a/capabilities/ai-red-teaming/capability.yaml b/capabilities/ai-red-teaming/capability.yaml index cfcc080..b123bc1 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.15.0" +version: "1.16.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/skills/error-troubleshooting/SKILL.md b/capabilities/ai-red-teaming/skills/error-troubleshooting/SKILL.md index b03379f..f5229a8 100644 --- a/capabilities/ai-red-teaming/skills/error-troubleshooting/SKILL.md +++ b/capabilities/ai-red-teaming/skills/error-troubleshooting/SKILL.md @@ -121,8 +121,8 @@ For the full lifecycle, load the `provisioning-and-lifecycle` skill. ### "Task not found" / 404 provisioning a bundled target - **Cause**: A bare task name resolves only within the caller's org, but bundled targets (`ml-extraction-*`, `*-mesh`) live in the public `dreadnode/` catalog. -- **Fix**: Use `provision_environment` with the BARE name - it catches the 404 and retries as `/` automatically. Never shell out to `dreadnode env`. -- **Do NOT**: qualify a bundled task with your own org (e.g. `aisf-learner-aug-2026/ml-extraction-imdb-text`) - it is not there and there is no fallback. Use the bare name or `dreadnode/`. +- **Fix**: Use `provision_environment` with the BARE name. If the task is not found it returns an `Error:` message explaining the resolution rule โ€” there is no automatic `/` retry, and qualifying a name does not bypass the visibility check. Never shell out to `dreadnode env`. +- **Do NOT**: qualify a bundled task with your own org (e.g. `aisf-learner-aug-2026/ml-extraction-imdb-text`) โ€” it is not there and there is no fallback. A bare name resolves to your org or any public task; a task owned by another org works only when public or owned by you. ### "404" on GET environments//status - **Cause**: The sandbox is already torn down or expired - you are polling a dead environment. diff --git a/capabilities/ai-red-teaming/tests/test_tool_surface.py b/capabilities/ai-red-teaming/tests/test_tool_surface.py new file mode 100644 index 0000000..fc97566 --- /dev/null +++ b/capabilities/ai-red-teaming/tests/test_tool_surface.py @@ -0,0 +1,233 @@ +"""Public tool-surface invariant for the ai-red-teaming capability. + +Regression tests for ENG-8467. A tool written in ``tools/*.py`` that is +documented as a tool (agent prompt, skill ``allowed-tools``) but missing the +``@safe_tool`` (``@tool``) decorator is silently dropped by the capability +loader: ``_discover_python_tools`` (dreadnode-tiger +``packages/sdk/dreadnode/capabilities/loader.py``) only registers module +attributes that are ``Tool`` instances or ``Toolset`` instances. A bare +function is skipped with no error, so the agent is instructed to call a tool +that does not exist (``Tool '...' not found``). + +These tests assert the invariant directly: every documented tool must be a +registered ``Tool`` with the expected name. Loading the modules is enough โ€” +no tool is invoked, so no network or platform access is required. +""" + +from __future__ import annotations + +import importlib.util +import re +from pathlib import Path + +import pytest + +pytest.importorskip("dreadnode.agents.tools") + +from dreadnode.agents.tools import Tool + +TOOLS_DIR = Path(__file__).resolve().parents[1] / "tools" +CAP_DIR = Path(__file__).resolve().parents[1] + + +def _load_all_tool_modules() -> dict[str, object]: + """Load every tools/*.py exactly like the runtime loader discovers files.""" + modules: dict[str, object] = {} + for path in sorted(TOOLS_DIR.glob("*.py")): + if path.name.startswith("_"): + continue + spec = importlib.util.spec_from_file_location( + f"airt_{path.stem}_under_test", path + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + modules[path.name] = mod + return modules + + +_REGISTERED = _load_all_tool_modules() + + +def _tool_names(module) -> set[str]: + return {value.name for value in vars(module).values() if isinstance(value, Tool)} + + +def _skills_with_allowed_tools() -> dict[str, set[str]]: + """Map skill path -> allowed-tools names from frontmatter.""" + import yaml + + out: dict[str, set[str]] = {} + for skill in sorted((CAP_DIR / "skills").rglob("SKILL.md")): + text = skill.read_text() + if not text.startswith("---\n"): + continue + end = text.find("\n---", 4) + if end == -1: + continue + try: + meta = yaml.safe_load(text[4:end]) + except yaml.YAMLError as exc: + # A skill with broken frontmatter must fail the invariant loudly, + # not silently drop its allowed-tools list from coverage. + raise AssertionError(f"{skill}: malformed frontmatter: {exc}") from exc + if not isinstance(meta, dict): + continue + allowed = (meta or {}).get("allowed-tools", "") + if isinstance(allowed, str): + names = {n for n in re.split(r"[,\s]+", allowed) if n} + elif isinstance(allowed, list): + names = set(allowed) + else: + continue + if names: + out[str(skill.relative_to(CAP_DIR))] = names + return out + + +EXPECTED_SURFACE = { + # ENG-8467: provision_environment was defined without @safe_tool and never + # registered, so agents hit `Tool 'ai_red_teaming__provision_environment' + # not found` on the first step of every hosted-target run. + "environments.py": { + "list_environments", + "provision_environment", + "teardown_environment", + }, + # Same root cause: generate_agentic_suite_attack is documented as a tool in + # agents/ai-red-teaming-agent.md but was never decorated, so the "run all + # possible attacks" path shipped with no executable tool. + "attacks.py": {"generate_agentic_suite_attack"}, +} + + +class TestRegressionENG8467: + @pytest.mark.parametrize("module_name, expected", EXPECTED_SURFACE.items()) + def test_documented_tools_are_registered(self, module_name, expected) -> None: + registered = _tool_names(_REGISTERED[module_name]) + missing = expected - registered + assert not missing, ( + f"{module_name}: documented tools not registered: {sorted(missing)}. " + "A tools/*.py function is only registered if it carries @safe_tool " + "(which applies @tool internally) or is a Toolset instance." + ) + + def test_provision_environment_is_a_tool_instance(self) -> None: + attr = getattr(_REGISTERED["environments.py"], "provision_environment") + assert isinstance(attr, Tool), ( + "provision_environment must be a registered Tool; a bare function is " + "invisible to the capability loader and agents get 'Tool not found'." + ) + assert attr.name == "provision_environment" + + +# Skill tools that are intentionally NOT implemented in this capability: they +# are documented for the platform-hosted ClickHouse analytics service +# (trace-analysis-advisor) that the runtime exposes in analytics contexts. +# The allowlist keeps the invariant below able to assert the FULL rule: every +# skill allowed-tools name must be a registered Tool OR be explicitly listed +# here. A phantom tool (documented but defined nowhere) fails the test. +PLATFORM_PROVIDED_TOOLS = frozenset( + { + "analyze_attack_effectiveness", + "suggest_optimal_transforms", + "predict_attack_success", + "identify_vulnerability_patterns", + "get_historical_metrics", + } +) + + +class TestSkillAllowedToolsInvariant: + """Every skill-listed tool name must resolve somewhere: it is either a + registered Tool in this capability's tools/, or explicitly allowlisted as + platform-provided (PLATFORM_PROVIDED_TOOLS). A defined-but-undecorated tool + (the ENG-8467 failure mode) and a documented-but-nowhere tool both fail + immediately.""" + + def test_all_allowed_tools_resolve_or_are_allowlisted(self) -> None: + # Name -> module of every attribute the tools modules expose. + attrs: dict[str, str] = {} + for rel_path, module in _REGISTERED.items(): + for name in vars(module): + if not name.startswith("_"): + attrs.setdefault( + name, rel_path + ) # first module wins (loader dedups) + registered_names = set() + for module in _REGISTERED.values(): + registered_names |= _tool_names(module) + + problems: list[str] = [] + for skill, allowed in _skills_with_allowed_tools().items(): + for name in sorted(allowed): + if name in attrs and name not in registered_names: + problems.append( + f"{skill}: '{name}' (defined in {attrs[name]}) is not a registered Tool" + ) + elif ( + name not in registered_names and name not in PLATFORM_PROVIDED_TOOLS + ): + problems.append( + f"{skill}: '{name}' is documented as a tool but defined " + "nowhere in tools/ and not in PLATFORM_PROVIDED_TOOLS" + ) + assert not problems, ( + "Skill allowed-tools references unresolved tools:\n" + "\n".join(problems) + ) + + +class TestProvisionErrorContract: + """provision_environment's failure modes, now that it is a live tool. + + The not-found path must return a plain ``Error:`` string (skill ยง5 semantics: + adjust params, don't blind-retry) instead of raising, so the @safe_tool + wrapper does not mislabel a user-input error as an internal tool fault. + Non-not-found failures must still be caught by the wrapper and surfaced as a + formatted ``Error:`` string (never a raw traceback). + """ + + def test_not_found_returns_error_string_with_hint(self, monkeypatch) -> None: + env = _REGISTERED["environments.py"] + pytest.importorskip("dreadnode.app.api.client") + pytest.importorskip("dreadnode.core.environment") + from dreadnode.app.api.client import NotFoundError + + class _FakeEnv: + def setup(self): + raise NotFoundError("task not found (404)") + + def _fake_configured(): + return None, object(), "my-org", "main" + + import dreadnode.core.environment as core_env + + monkeypatch.setattr(env, "_configured", _fake_configured) + monkeypatch.setattr(core_env, "TaskEnvironment", lambda api, **kw: _FakeEnv()) + + out = env.provision_environment("ml-extraction-fraud-tabular") + assert isinstance(out, str) + assert out.startswith("Error: Task 'ml-extraction-fraud-tabular' not found.") + assert "/" in out + + def test_other_failures_surface_via_safe_tool_not_raise(self, monkeypatch) -> None: + env = _REGISTERED["environments.py"] + pytest.importorskip("dreadnode.core.environment") + + class _BoomEnv: + def setup(self): + raise ValueError("provider exploded") + + def _fake_configured(): + return None, object(), "my-org", "main" + + import dreadnode.core.environment as core_env + + monkeypatch.setattr(env, "_configured", _fake_configured) + monkeypatch.setattr(core_env, "TaskEnvironment", lambda api, **kw: _BoomEnv()) + + out = env.provision_environment("ml-extraction-fraud-tabular") + assert isinstance(out, str) + assert out.startswith("Error:") + assert "'provision_environment'" in out + assert "could not complete" in out + assert "provider exploded" in out diff --git a/capabilities/ai-red-teaming/tools/attacks.py b/capabilities/ai-red-teaming/tools/attacks.py index fe718e3..fc9f4c5 100644 --- a/capabilities/ai-red-teaming/tools/attacks.py +++ b/capabilities/ai-red-teaming/tools/attacks.py @@ -337,6 +337,7 @@ def generate_multistep_tool_attack( return _call_runner("generate_multistep_tool_attack", params) +@safe_tool def generate_agentic_suite_attack( goal: t.Annotated[str, "Overall red-team goal for the agent"], agent_url: t.Annotated[str, "HTTP endpoint of the target agent"], diff --git a/capabilities/ai-red-teaming/tools/environments.py b/capabilities/ai-red-teaming/tools/environments.py index b830fbc..8cb1733 100644 --- a/capabilities/ai-red-teaming/tools/environments.py +++ b/capabilities/ai-red-teaming/tools/environments.py @@ -170,7 +170,12 @@ 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.""" + completion hook. Short-circuits (no platform call) when nothing is registered. + + Deliberately NOT a tool (no @safe_tool): this is a shared helper, invoked + directly by assessment.py's completion teardown and by tests. Adding @tool + would expose a user-callable teardown that bypasses assessment lifecycle. + """ if not _registry_load(registry_path): return {"torn_down": [], "skipped": [], "errors": {}} _inst, api, org, workspace = _configured() @@ -251,6 +256,7 @@ def _target_kind(task_ref: str) -> str: return "unknown" +@safe_tool def provision_environment( task_ref: t.Annotated[str, "Environment/task to deploy, e.g. 'finops-mesh' or 'ml-extraction-mnist-image'"], model: t.Annotated[ @@ -271,6 +277,10 @@ def provision_environment( Do not probe ``/attack`` on a classifier target; it does not serve it. """ from dreadnode.core.environment import TaskEnvironment + # NOTE: tests/test_tool_surface.py (TestProvisionErrorContract) monkeypatches + # `dreadnode.core.environment.TaskEnvironment` and relies on this lazy import + # to take effect. If this import is ever hoisted to module level, update those + # tests to patch the module-level binding instead. _inst, api, org, workspace = _configured() if not org or not workspace: @@ -295,11 +305,17 @@ def _mk(ref: str) -> t.Any: except Exception as exc: # noqa: BLE001 - add a resolution hint on not-found if not _is_not_found(exc): raise - raise RuntimeError( - f"Task '{task_ref}' not found. A bare name resolves to a task in your org " - f"or any public task; a task owned by another org must be public or owned " - f"by you. Check the name and version, or qualify it as /." - ) from exc + # A user-input error, not an internal fault: return a plain "Error:" + # string (per the skill contract, "Error:" = adjust params, don't + # blind-retry) instead of raising, so safe_tool does not mislabel it + # as an internal tool issue. + return ( + f"Error: Task '{task_ref}' not found. A bare name resolves to a task " + "in your org or any public task, so a not-found here means the task is " + "private to another org (no / qualification can bypass that " + "visibility rule) or the name/version is wrong - check the name and " + "version." + ) 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