diff --git a/capabilities/ai-red-teaming/capability.yaml b/capabilities/ai-red-teaming/capability.yaml index 2c91ea6..82b49fe 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.3.7" +version: "1.4.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/scripts/attack_runner.py b/capabilities/ai-red-teaming/scripts/attack_runner.py index 55a3204..42f0f8c 100644 --- a/capabilities/ai-red-teaming/scripts/attack_runner.py +++ b/capabilities/ai-red-teaming/scripts/attack_runner.py @@ -55,47 +55,44 @@ def _get_workspace_path() -> Path: def _resolve_platform_env() -> dict[str, str]: """Build env dict with platform credentials for subprocess execution. - In sandbox mode, env vars are already set. In TUI/CLI mode, reads from - the saved profile at ~/.cache/dreadnode/config.yaml. + In sandbox mode, env vars are already set (DREADNODE_SERVER + API_KEY). + In TUI mode, the runtime sets DREADNODE_LLM_BASE + LLM_API_KEY for the + LLM proxy but does NOT set SERVER/API_KEY/ORG/WORKSPACE/PROJECT — those + come from the saved profile (``dn login``). + + We ALWAYS read the saved profile to fill in scope vars (org, workspace, + project) so the subprocess respects the user's TUI context. ``setdefault`` + ensures explicit env vars (sandbox) take precedence over the profile. """ env = os.environ.copy() - # If the runtime already provides platform credentials in any of the - # forms the SDK understands, pass the env through untouched -- the - # generated script self-configures via dn.configure(), whose precedence - # is: explicit args > env vars > saved profile. - # - DREADNODE_SERVER + DREADNODE_API_KEY (classic platform env) - # - DREADNODE_LLM_BASE + DREADNODE_LLM_API_KEY (runtime LLM proxy env) - if env.get("DREADNODE_SERVER") and env.get("DREADNODE_API_KEY"): - return env - if env.get("DREADNODE_LLM_BASE") and env.get("DREADNODE_LLM_API_KEY"): - return env - - # Fall back to saved profile (TUI/CLI mode) - # Profile lives at ~/.dreadnode/config.yaml (YAML format) + # Read saved profile via the SDK's own config system. This respects + # ``dn login``, ``/workspace`` switches, and profile selection — no + # need to parse YAML manually. try: - from pathlib import Path as _Path - import yaml # type: ignore[import-untyped] - - config_path = _Path.home() / ".dreadnode" / "config.yaml" - if config_path.exists(): - config = yaml.safe_load(config_path.read_text()) - active = config.get("active") - servers = config.get("servers", {}) - if active and active in servers: - profile = servers[active] - env.setdefault("DREADNODE_SERVER", profile.get("url", "")) - env.setdefault("DREADNODE_API_KEY", profile.get("api_key", "")) - env.setdefault("DREADNODE_ORGANIZATION", profile.get("default_organization", "")) - env.setdefault("DREADNODE_WORKSPACE", profile.get("default_workspace", "")) - env.setdefault("DREADNODE_PROJECT", profile.get("default_project", "")) + from dreadnode.app.config import UserConfig + + config = UserConfig.read() + profile_data = config.active_profile + if profile_data: + _, profile = profile_data + if profile.url: + env.setdefault("DREADNODE_SERVER", profile.url) + if profile.api_key: + env.setdefault("DREADNODE_API_KEY", profile.api_key) + if profile.organization: + env.setdefault("DREADNODE_ORGANIZATION", profile.organization) + if profile.workspace: + env.setdefault("DREADNODE_WORKSPACE", profile.workspace) + if profile.project: + env.setdefault("DREADNODE_PROJECT", profile.project) except Exception: - pass # Best-effort — will fail later in the script with a clear error + pass # Best-effort — dn.configure() in the script will try again return env -def _auto_execute_workflow(filename: str, timeout: int = 540) -> str: +def _auto_execute_workflow(filename: str, timeout: int = 3600) -> str: """Execute a workflow script and return output. Used for auto-execution after generate.""" filepath = WORKFLOWS_DIR / filename if not filepath.exists(): @@ -775,6 +772,8 @@ def _auto_execute_workflow(filename: str, timeout: int = 540) -> str: ATTACK_ALIASES["tmap"] = "tmap_trajectory_attack" ATTACK_ALIASES["aprt"] = "aprt_progressive_attack" + + _TRANSFORM_DEFS: dict[str, dict] = { # encoding "base64_encode": {"module": "dreadnode.transforms.encoding", "name": "base64_encode", "code": "base64_encode()"}, @@ -3144,7 +3143,8 @@ def _build_attack_params( params.append("{}={}".format(k, v)) if transforms_expr is not None: params.append("transforms={}".format(transforms_expr)) - # AIRT span linkage — links Study-created spans to assessment in ClickHouse + # AIRT span linkage — links Study-created spans to assessment in ClickHouse. + # All SDK attacks accept these kwargs as of dreadnode-tiger#1693. params.append("airt_assessment_id=assessment.assessment_id") params.append("airt_goal_category={}".format(goal_category_expr)) params.append("airt_target_model=TARGET_MODEL") @@ -3358,12 +3358,12 @@ def _generate_transform_study(config: dict) -> str: for k, v in atk.get("extra_defaults", {}).items(): params.append("{}={}".format(k, v)) params.append("transforms=transforms") + # AIRT span linkage — all attacks accept these as of dreadnode-tiger#1693 params.append("airt_assessment_id=assessment.assessment_id") params.append("airt_goal_category=GOAL_CATEGORY.value") params.append("airt_target_model=TARGET_MODEL") - attack_params = ",\n ".join(params) - canon = atk["canonical_name"] + attack_params = ",\n ".join(params) assessment_name = _safe_str(config.get("assessment_name") or "{} Transform Comparison".format(canon)) assessment_kwargs = _build_assessment_kwargs(config, assessment_name, config.get("filename", "")) @@ -3718,6 +3718,7 @@ def _generate_category_attack(config: dict) -> str: if transforms: transforms_expr = "[{}]".format(", ".join(t["code"] for t in transforms)) params.append("transforms={}".format(transforms_expr)) + # AIRT span linkage — all attacks accept these as of dreadnode-tiger#1693 params.append("airt_assessment_id=assessment.assessment_id") params.append("airt_goal_category=goal_cat.value") params.append("airt_category=goal_row.get('category', '')") diff --git a/capabilities/ai-red-teaming/tools/attacks.py b/capabilities/ai-red-teaming/tools/attacks.py index 4e5d8da..f6c3dc3 100644 --- a/capabilities/ai-red-teaming/tools/attacks.py +++ b/capabilities/ai-red-teaming/tools/attacks.py @@ -42,7 +42,7 @@ def _call_runner(name: str, params: dict) -> str: input=payload, capture_output=True, text=True, - timeout=660, # 11min: runner has 9min internal timeout + overhead + timeout=3660, # 61min: runner has 60min internal timeout + overhead ) output = result.stdout.strip() if result.returncode != 0: diff --git a/capabilities/ai-red-teaming/tools/workflows.py b/capabilities/ai-red-teaming/tools/workflows.py index 15fb9e3..af54747 100644 --- a/capabilities/ai-red-teaming/tools/workflows.py +++ b/capabilities/ai-red-teaming/tools/workflows.py @@ -26,6 +26,36 @@ from dreadnode.app.env import resolve_python_executable +def _resolve_platform_env() -> dict[str, str]: + """Build env dict with platform credentials for subprocess execution. + + Mirrors attack_runner._resolve_platform_env so manually-executed + workflows get the same credential resolution as auto-executed ones. + Always reads the saved profile to fill org/workspace/project scope. + """ + env = os.environ.copy() + try: + from dreadnode.app.config import UserConfig + + config = UserConfig.read() + profile_data = config.active_profile + if profile_data: + _, profile = profile_data + if profile.url: + env.setdefault("DREADNODE_SERVER", profile.url) + if profile.api_key: + env.setdefault("DREADNODE_API_KEY", profile.api_key) + if profile.organization: + env.setdefault("DREADNODE_ORGANIZATION", profile.organization) + if profile.workspace: + env.setdefault("DREADNODE_WORKSPACE", profile.workspace) + if profile.project: + env.setdefault("DREADNODE_PROJECT", profile.project) + except Exception: + pass + return env + + # Get org/workspace from active profile, with fallbacks def _get_workspace_path() -> Path: try: @@ -165,7 +195,7 @@ def list_workflows() -> str: @safe_tool def execute_workflow( filename: t.Annotated[str, "Workflow filename to execute"], - timeout: t.Annotated[int, "Max execution time in seconds (max 600)"] = 540, + timeout: t.Annotated[int, "Max execution time in seconds (max 3600)"] = 540, ) -> str: """Execute a saved attack workflow script. @@ -179,7 +209,7 @@ def execute_workflow( if not filepath.exists(): return f"Workflow not found: {filename}. Use list_workflows to see available." - timeout = min(timeout, 600) + timeout = min(timeout, 3600) try: python_executable = resolve_python_executable() @@ -193,6 +223,7 @@ def execute_workflow( text=True, timeout=timeout, cwd=str(WORKFLOWS_DIR.parent), + env=_resolve_platform_env(), ) output = result.stdout if result.stderr: