From 4d5d5dfdb11225e7b8a32490590c3fb1c161bba9 Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Wed, 15 Jul 2026 18:58:42 -0700 Subject: [PATCH] fix(ai-red-teaming): import get_generator in generated agentic script Same missing-import bug as the ATLAS template (#100), in the separate `_build_agentic_imports`: generated agentic-attack scripts call get_generator (proxy-routing block) but never imported it, raising NameError at runtime. Add the import and a TestAgenticGeneration regression test (the agentic generator was previously untested). --- .../ai-red-teaming/scripts/attack_runner.py | 2 ++ .../tests/test_attack_runner.py | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/capabilities/ai-red-teaming/scripts/attack_runner.py b/capabilities/ai-red-teaming/scripts/attack_runner.py index 7ede610..9dcb463 100644 --- a/capabilities/ai-red-teaming/scripts/attack_runner.py +++ b/capabilities/ai-red-teaming/scripts/attack_runner.py @@ -4324,6 +4324,8 @@ def _build_agentic_imports(attacks: list[dict], transforms: list[dict], has_scor "", "import dreadnode as dn", "from dreadnode import task", + # Required by the proxy-routing block (get_generator/GenerateParams). + "from dreadnode.generators.generator import get_generator, GenerateParams", ] attack_funcs = set() diff --git a/capabilities/ai-red-teaming/tests/test_attack_runner.py b/capabilities/ai-red-teaming/tests/test_attack_runner.py index c30f923..0786b9b 100644 --- a/capabilities/ai-red-teaming/tests/test_attack_runner.py +++ b/capabilities/ai-red-teaming/tests/test_attack_runner.py @@ -1039,6 +1039,33 @@ def test_bearer_auth_injects_header(self): assert "Authorization" in script +_AGENTIC_BASE = { + "attack_type": "goat", + "goal": "get the agent to misuse a privileged tool", + "agent_url": "http://localhost:8000/attack", + "agent_preset": "custom", + "attacker_model": "groq scout", + "n_iterations": 4, + "generate_only": True, +} + + +class TestAgenticGeneration: + def test_generated_script_compiles_and_imports_get_generator(self): + result = _generate_method("generate_agentic_attack", _AGENTIC_BASE) + assert "error" not in result, result.get("error") + script = Path(result["filepath"]).read_text() + compile(script, result["filepath"], "exec") + # Regression: the proxy-routing block calls get_generator(...), so the + # generated script MUST import it (compile() only checks syntax, so a + # missing import is a runtime NameError). + assert "get_generator(" in script + assert ( + "from dreadnode.generators.generator import get_generator, GenerateParams" + in script + ) + + class TestAtlasValidation: def test_missing_agent_url_errors(self): params = {k: v for k, v in _ATLAS_BASE.items() if k != "agent_url"}