From 56400e4aa4349d4fee87aebf1172601a59dafec8 Mon Sep 17 00:00:00 2001 From: Raja Sekhar Rao Dheekonda Date: Tue, 8 Sep 2026 19:59:55 -0700 Subject: [PATCH] fix(ai-red-teaming): recover cleanly from transient tool errors + right endpoint per target (ENG-8427) A learner's TUI evasion run surfaced two alarming-but-non-fatal errors while the attack actually completed (assessment 77ab88c7, 100% ASR, 1 finding). Confirmed from prod Logfire + ClickHouse session_events: a transient TLS handshake timeout on the first list_environments call, and a 404 from the agent probing /attack on an ml-extraction classifier that only serves /predict. The agent recovered on its own but the TUI gave no sign of it. Fixes: - safe_tool: retry transient network faults (TLS handshake, timeouts, conn reset, 5xx) up to 2x with backoff, then surface an explicitly non-fatal Note ('does not affect any attack already running') instead of a raw Error. Works for sync + async tools. - provision_environment: return the endpoint that matches the target type - classifier targets get /predict (+ /pool,/members,/nonmembers) and evasion/ extraction/membership/inversion guidance; only meshes get /attack + ATLAS. Stops the agent probing /attack on a classifier (the 404). - fmt_asr: consistent ASR formatting robust to 0-1 fractions vs 0-100 percents, fixing the final message showing '1.0%' instead of '100%'. Applied in results/assessment/session. - agent prompt: require narrating recovery when a tool errors mid-run; never leave a raw non-fatal error as the last thing the user sees. Bumps capability 1.13.0 -> 1.14.0. Adds 29 tests (fmt_asr, safe_tool retry/ classification sync+async, _target_kind). --- .../agents/ai-red-teaming-agent.md | 3 +- capabilities/ai-red-teaming/capability.yaml | 2 +- .../tests/test_environments_teardown.py | 21 +++ .../tests/test_errors_safe_tool.py | 139 ++++++++++++++++++ .../ai-red-teaming/tools/assessment.py | 8 +- .../ai-red-teaming/tools/environments.py | 92 +++++++++--- capabilities/ai-red-teaming/tools/errors.py | 139 ++++++++++++++---- capabilities/ai-red-teaming/tools/results.py | 10 +- capabilities/ai-red-teaming/tools/session.py | 3 +- 9 files changed, 358 insertions(+), 59 deletions(-) create mode 100644 capabilities/ai-red-teaming/tests/test_errors_safe_tool.py 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 03e1e93..2568444 100644 --- a/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md +++ b/capabilities/ai-red-teaming/agents/ai-red-teaming-agent.md @@ -83,6 +83,7 @@ Keep it to a single line; don't pad it. - **Multimodal LLM (vision/audio/video) with media inputs → `generate_multimodal_attack`**. Detect this when the user attaches or points to media and wants to probe a chat/vision model: "attack this vision model", "run these prompts with the images in `./imgs`", "apply an image transform on the images", "test this voice model with the audio in this folder", "visual prompt injection", "typographic jailbreak". Pass `image_dir`/`audio_dir`/`video_dir` for folders or `image_paths`/`audio_paths`/`video_paths` for explicit files. Do NOT confuse with `generate_image_attack` (classifier evasion, not chat). 2. IMMEDIATELY call `execute_workflow` with the filename returned by the generator, in the SAME turn — a `generate_*` call that returns "workflow generated / NEXT STEP: execute_workflow" is NOT done. Never stop after generating; skipping execution leaves the assessment with 0 trials and looks like a silent failure to the user. - **Narrate every step — no black box.** Before each tool call, say what you're about to do and why (e.g. "Generating the workflow… now executing 20 trials (4 prompts × 5 images)… scoring with gpt-4o-mini…"). After execution, report the assessment ID and how many trials ran. If a step errors or a target returns no response, say so plainly (e.g. "target call failed: 401 auth — provider key missing") instead of moving on silently. + - **Always narrate recovery - never leave an error hanging.** Tool results that start with `Note:` are transient/non-fatal (a network blip that was retried); tell the user it was transient, that no running attack or recorded result was affected, and continue. A `404` from probing a wrong endpoint (e.g. fetching `/attack` on an ML classifier that serves `/predict`) is exploratory, not a failure - say what you learned and switch to the correct path. Whenever a step errors and you then succeed another way, explicitly state that you recovered and whether the final result was impacted. Never let a raw error be the last thing the user sees about a step that actually succeeded. 3. Call `register_assessment`, then `update_assessment_status` once execution finishes. 4. Call `validate_attack_results` FIRST. If it surfaces errors, stop and report them — do not call analytics tools. 5. If validation passes, call `get_assessment_status` for platform metrics and report ONLY those raw values. @@ -163,7 +164,7 @@ 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 `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. +- **provision_environment** — Deploy a hosted target and return the endpoint that matches its type. A multi-agent mesh (e.g. `finops-mesh`) returns an `/attack` URL + execute token → chain into `generate_atlas_attack`. A black-box ML classifier (e.g. `ml-extraction-mnist-image`) returns a `/predict` endpoint (plus `/pool`, `/members`, `/nonmembers`) → use `generate_evasion_attack` / `generate_extraction_attack` / `generate_membership_attack` / `generate_inversion_attack` with `api_url=/predict`. **Do not fetch `/attack` on a classifier target - it does not serve it.** 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 38dc7b7..6a4afba 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.13.0" +version: "1.14.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 index e4f3786..8a3fc6b 100644 --- a/capabilities/ai-red-teaming/tests/test_environments_teardown.py +++ b/capabilities/ai-red-teaming/tests/test_environments_teardown.py @@ -241,3 +241,24 @@ def test_real_teardown_on_complete_returns_empty_when_nothing_registered( # platform call and returns "" (no note appended). monkeypatch.setenv("AIRT_ENV_REGISTRY_PATH", str(Path("/nonexistent/dir/registry.json"))) assert assessment._teardown_on_complete() == "" + + +class TestTargetKind: + """provision_environment must return the right endpoint per target type + (ENG-8427: a classifier was steered to /attack and 404'd).""" + + @pytest.mark.parametrize("ref", [ + "ml-extraction-mnist-image", "ml-extraction-fraud-tabular", + "ml-extraction-imdb-text", "some-classifier", "mnist-demo", + ]) + def test_classifier_targets(self, ref): + assert env._target_kind(ref) == "classifier" + + @pytest.mark.parametrize("ref", [ + "finops-mesh", "devsecops-mesh", "healthcare-mesh", "soc-mesh", + ]) + def test_mesh_targets(self, ref): + assert env._target_kind(ref) == "mesh" + + def test_unknown_target(self): + assert env._target_kind("totally-custom-thing") == "unknown" diff --git a/capabilities/ai-red-teaming/tests/test_errors_safe_tool.py b/capabilities/ai-red-teaming/tests/test_errors_safe_tool.py new file mode 100644 index 0000000..730888d --- /dev/null +++ b/capabilities/ai-red-teaming/tests/test_errors_safe_tool.py @@ -0,0 +1,139 @@ +"""Tests for tools/errors.py - safe_tool retry/classification + fmt_asr. + +Regression for ENG-8427: transient network faults (TLS handshake timeout, etc.) +should be retried and surfaced as an explicitly non-fatal Note, and ASR should +render consistently as a percentage (the '1.0%' bug). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("dreadnode.agents.tools") + +ERRORS_PATH = Path(__file__).resolve().parents[1] / "tools" / "errors.py" + + +def _load(): + spec = importlib.util.spec_from_file_location("airt_errors_under_test", ERRORS_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 + + +E = _load() + + +class TestFmtAsr: + @pytest.mark.parametrize("value,expected", [ + (1.0, "100%"), + (0.78, "78%"), + (0.975, "97.5%"), + (78, "78%"), + (100, "100%"), + (0.0, "0%"), + (None, "N/A"), + ]) + def test_formats(self, value, expected): + assert E.fmt_asr(value) == expected + + def test_fraction_one_is_not_one_percent(self): + # The ENG-8427 bug: 1.0 was rendered as "1.0%" instead of "100%". + assert E.fmt_asr(1.0) == "100%" + assert E.fmt_asr(1.0) != "1.0%" + + +class TestTransientClassification: + @pytest.mark.parametrize("exc", [ + ConnectionError("_ssl.c:983: The handshake operation timed out"), + TimeoutError("read timed out"), + OSError("Connection reset by peer"), + RuntimeError("502 Bad Gateway"), + ]) + def test_transient_true(self, exc): + assert E._is_transient(exc) is True + + @pytest.mark.parametrize("exc", [ + ValueError("num_classes must be > 0"), + KeyError("records"), + ]) + def test_transient_false(self, exc): + assert E._is_transient(exc) is False + + +class TestSafeToolRetry: + @pytest.fixture(autouse=True) + def _no_sleep(self, monkeypatch): + monkeypatch.setattr(E, "_BACKOFF_SECONDS", (0, 0)) + + def test_transient_is_retried_then_noted(self): + calls = {"n": 0} + + @E.safe_tool + def flaky() -> str: + calls["n"] += 1 + raise ConnectionError("handshake operation timed out") + + out = flaky() + assert calls["n"] == E._MAX_RETRIES + 1 # retried to exhaustion + assert out.startswith("Note:") + assert "transient" in out + assert "does not affect any attack" in out + + def test_transient_then_success_returns_value(self): + calls = {"n": 0} + + @E.safe_tool + def recovers() -> str: + calls["n"] += 1 + if calls["n"] < 2: + raise TimeoutError("connection timed out") + return "OK" + + assert recovers() == "OK" + assert calls["n"] == 2 + + def test_non_transient_not_retried_and_is_error(self): + calls = {"n": 0} + + @E.safe_tool + def bad() -> str: + calls["n"] += 1 + raise ValueError("bad parameter") + + out = bad() + assert calls["n"] == 1 # no retry for a non-transient error + assert out.startswith("Error:") + assert "could not complete" in out + + def test_no_em_dash_in_messages(self): + @E.safe_tool + def t1() -> str: + raise ValueError("x") + + @E.safe_tool + def t2() -> str: + raise ConnectionError("timed out") + + assert "—" not in t1() + assert "—" not in t2() + + @pytest.mark.asyncio + async def test_async_transient_then_success(self): + calls = {"n": 0} + + @E.safe_tool + async def arec() -> str: + calls["n"] += 1 + if calls["n"] < 2: + raise ConnectionError("ssl handshake timed out") + return "ASYNC_OK" + + assert await arec() == "ASYNC_OK" + assert calls["n"] == 2 diff --git a/capabilities/ai-red-teaming/tools/assessment.py b/capabilities/ai-red-teaming/tools/assessment.py index 43dcec9..58ce567 100644 --- a/capabilities/ai-red-teaming/tools/assessment.py +++ b/capabilities/ai-red-teaming/tools/assessment.py @@ -21,6 +21,7 @@ _errors_mod = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_errors_mod) safe_tool = _errors_mod.safe_tool +fmt_asr = _errors_mod.fmt_asr ASSESSMENT_PATH = Path(os.environ.get("AIRT_ASSESSMENT_PATH", "/tmp/airt_assessment.json")) @@ -96,8 +97,7 @@ def get_assessment_status() -> str: for c in completed: # ASR is the attack success probability (how often the attack # worked). Shown as a percentage; that *is* the probability metric. - asr = c.get("asr") - asr_str = f"{asr}%" if asr is not None else "N/A" + asr_str = fmt_asr(c.get("asr")) line = f" - {c['attack_name']}: success rate (ASR)={asr_str}" if c.get("notes"): line += f" — {c['notes']}" @@ -113,7 +113,7 @@ def get_assessment_status() -> str: def update_assessment_status( attack_name: t.Annotated[str, "Name of the completed attack"], status: t.Annotated[str, "Attack status (e.g., 'completed', 'failed', 'skipped')"] = "completed", - asr: t.Annotated[float | None, "Attack success rate as percentage (0-100)"] = None, + asr: t.Annotated[float | None, "Attack success rate, either a 0-1 fraction (e.g. 1.0 = 100%) or a 0-100 percentage; displayed consistently as a percentage"] = None, risk_score: t.Annotated[ float | None, "Optional severity-weighted risk (0-10), stored for platform parity but " @@ -166,7 +166,7 @@ def update_assessment_status( total = len(planned) done = len(completed) - asr_str = f" (ASR={asr}%)" if asr is not None else "" + asr_str = f" (ASR={fmt_asr(asr)})" if asr is not None else "" return f"Recorded {attack_name}: {status}{asr_str}. Progress: {done}/{total}.{teardown_note}" diff --git a/capabilities/ai-red-teaming/tools/environments.py b/capabilities/ai-red-teaming/tools/environments.py index 4a2d630..2f63d73 100644 --- a/capabilities/ai-red-teaming/tools/environments.py +++ b/capabilities/ai-red-teaming/tools/environments.py @@ -235,21 +235,41 @@ def list_environments() -> str: @safe_tool +def _target_kind(task_ref: str) -> str: + """Classify a provisionable target so we return the right endpoint + guidance. + + - 'classifier': black-box ML target that serves /predict (+ /pool, /members, + /nonmembers) - evasion / extraction / membership / inversion. + - 'mesh': multi-agent environment that serves /attack - ATLAS. + - 'unknown': fall back to a non-prescriptive message. + """ + ref = (task_ref or "").lower() + if ref.endswith("-mesh") or "mesh" in ref: + return "mesh" + if ("ml-extraction" in ref or "classifier" in ref or "extraction" in ref + or any(k in ref for k in ("mnist", "fraud", "imdb", "tabular", "image", "text"))): + return "classifier" + return "unknown" + + def provision_environment( - task_ref: t.Annotated[str, "Environment/task to deploy, e.g. 'finops-mesh'"], + task_ref: t.Annotated[str, "Environment/task to deploy, e.g. 'finops-mesh' or 'ml-extraction-mnist-image'"], model: t.Annotated[ str, "Model the environment's agents use (e.g. 'dn/claude-haiku-4-5', 'groq/llama-3.3-70b-versatile')" ] = "", model_role: t.Annotated[str, "Role key to override with the model (default 'agent')"] = "agent", timeout_sec: t.Annotated[int, "Provision + run budget in seconds"] = 1800, ) -> str: - """Provision a hosted multi-agent environment and return its attack URL. - - Deploys the environment via the platform sandbox provider, passing ``model`` - to the environment's agents (task-environment model capability). Returns the - ``/attack`` base URL and the bearer execute token — pass the URL to - ``generate_atlas_attack`` (``agent_url=/attack``) with - ``agent_auth_type='bearer'`` and the token via the ``AGENT_API_KEY`` env. + """Provision a hosted target environment and return the correct endpoint for it. + + Deploys the environment via the platform sandbox provider. The next step + depends on the target type: + - a black-box ML classifier (e.g. ``ml-extraction-mnist-image``) serves + ``/predict`` - use ``generate_evasion_attack`` / ``generate_extraction_attack`` + / ``generate_membership_attack`` / ``generate_inversion_attack``. + - a multi-agent mesh (e.g. ``finops-mesh``) serves ``/attack`` - use + ``generate_atlas_attack``. + Do not probe ``/attack`` on a classifier target; it does not serve it. """ from dreadnode.core.environment import TaskEnvironment @@ -267,24 +287,56 @@ def provision_environment( 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. + # 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 ( + url = url.rstrip("/") + kind = _target_kind(task_ref) + header = ( 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" - f" Model: {model or '(env default)'}\n\n" - 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.\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." + f" Base URL: {url}\n" + f" Model: {model or '(env default)'}\n" + ) + teardown_note = ( + "\n>>> WHEN DONE: this sandbox bills for its whole lifetime - it is torn down " + "automatically when the assessment completes, or call teardown_environment() now." + ) + + if kind == "classifier": + return ( + header + + f" Predict endpoint: {url}/predict (dataset helpers: {url}/pool, /members, /nonmembers)\n\n" + + ">>> NEXT STEP: this is a black-box ML classifier. Do NOT fetch /attack - it does not " + "serve it. Run the matching attack with the predict endpoint, e.g.:\n" + + f' generate_evasion_attack(attack_type="hopskipjump", api_url="{url}/predict", ...)\n' + + " (or generate_extraction_attack / generate_membership_attack / generate_inversion_attack)\n" + + " These tools query /predict and pull data from /pool, /members, /nonmembers automatically." + + teardown_note + ) + if kind == "mesh": + return ( + header + + f" Attack URL: {url}/attack\n" + + f" Auth: bearer (execute token below)\n" + + f" Execute token: {token}\n\n" + + ">>> 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") ' + + "and set AGENT_API_KEY to the execute token above." + + teardown_note + ) + # unknown: describe both without prescribing a wrong endpoint + return ( + header + + f" Execute token: {token}\n\n" + + ">>> NEXT STEP: inspect the target before attacking. A black-box ML classifier " + f"serves {url}/predict (use generate_evasion_attack / extraction / membership / " + f"inversion); a multi-agent mesh serves {url}/attack (use generate_atlas_attack). " + "Do not assume /attack exists." + + teardown_note ) diff --git a/capabilities/ai-red-teaming/tools/errors.py b/capabilities/ai-red-teaming/tools/errors.py index a7c56a4..aab71ed 100644 --- a/capabilities/ai-red-teaming/tools/errors.py +++ b/capabilities/ai-red-teaming/tools/errors.py @@ -21,38 +21,109 @@ def my_tool(...) -> str: from __future__ import annotations +import asyncio import functools import sys +import time import typing as t from dreadnode.agents.tools import tool -__all__ = ["safe_tool"] +__all__ = ["safe_tool", "fmt_asr"] F = t.TypeVar("F", bound=t.Callable[..., t.Any]) -def _format_error(tool_name: str, exc: BaseException) -> str: - """Build a concise, user-facing error string (no traceback).""" - # Keep it short and actionable; never leak a stack trace to the user. +def fmt_asr(value: t.Any) -> str: + """Format an attack success rate as a percentage string, robust to inputs + given either as a 0-1 fraction (the SDK convention) or an already-scaled + 0-100 percentage. Fixes the '1.0%' display bug where a fraction was printed + with a bare '%'. + + Examples: 1.0 -> '100%', 0.78 -> '78%', 78 -> '78%', 100 -> '100%'. + """ + if value is None: + return "N/A" + try: + v = float(value) + except (TypeError, ValueError): + return str(value) + pct = v * 100.0 if v <= 1.0 else v + return f"{pct:.0f}%" if abs(pct - round(pct)) < 1e-9 else f"{pct:.1f}%" + +# Transient network faults worth retrying: the connection never completed, so a +# retry commonly succeeds and the failure does not affect anything already +# running (e.g. an attack in progress). Matched on the exception's class name and +# message so we do not need to import every client library's error types. +_TRANSIENT_MARKERS = ( + "handshake", + "timed out", + "timeout", + "temporarily unavailable", + "connection reset", + "connection aborted", + "connection refused", + "connection error", + "econnreset", + "broken pipe", + "ssl", + "eof occurred", + "remotedisconnected", + "remoteprotocolerror", + "readtimeout", + "connecttimeout", + "connecterror", + "poolTimeout".lower(), + "max retries exceeded", + "name or service not known", + "temporary failure in name resolution", + "502 bad gateway", + "503 service unavailable", + "504 gateway timeout", +) + +_MAX_RETRIES = 2 +_BACKOFF_SECONDS = (1.5, 3.0) + + +def _clean_msg(exc: BaseException) -> str: msg = str(exc).strip() or exc.__class__.__name__ - # Collapse multi-line / overly long internal messages. msg = " ".join(msg.split()) - if len(msg) > 500: - msg = msg[:500] + "…" + return msg[:500] + "..." if len(msg) > 500 else msg + + +def _is_transient(exc: BaseException) -> bool: + blob = f"{type(exc).__name__} {exc}".lower() + return any(m in blob for m in _TRANSIENT_MARKERS) + + +def _format_error(tool_name: str, exc: BaseException, *, transient: bool, attempts: int) -> str: + """Build a concise, user-facing string (no traceback). + + Transient network faults are labelled non-fatal so the agent (and the user) + know the run was not compromised and the step can simply be retried. + """ + msg = _clean_msg(exc) + if transient: + return ( + f"Note: '{tool_name}' hit a transient network issue after {attempts} " + f"attempt(s) ({msg}). This is not a problem with your request and does " + "not affect any attack already running or already-recorded results. " + "Retry this step; it usually succeeds on the next try." + ) return ( - f"Error: '{tool_name}' could not complete: {msg}. " - "This is an internal issue, not your input — please retry, or adjust " - "parameters if it persists." + f"Error: '{tool_name}' could not complete: {msg}. This is an internal tool " + "issue, not your input; retry, or adjust parameters if it persists." ) def safe_tool(fn: F) -> t.Any: """Wrap a function as a tool that never raises to the user. - Any exception raised inside ``fn`` is caught and returned as a clean - string. Works for both sync and async tool functions. Applies ``@tool`` - after wrapping, so the decorated callable is a fully-formed tool. + Any exception raised inside ``fn`` is caught and returned as a clean string. + Transient network faults (TLS handshake timeout, connection reset, 5xx, etc.) + are retried up to ``_MAX_RETRIES`` times with a short backoff before being + surfaced as an explicitly non-fatal note. Works for sync and async tools. """ tool_name = getattr(fn, "__name__", "tool") @@ -60,21 +131,39 @@ def safe_tool(fn: F) -> t.Any: @functools.wraps(fn) async def _async_wrapper(*args: t.Any, **kwargs: t.Any) -> t.Any: - try: - return await fn(*args, **kwargs) - except Exception as exc: # noqa: BLE001 — deliberate catch-all safety net - _log(tool_name, exc) - return _format_error(tool_name, exc) + last: BaseException | None = None + for attempt in range(_MAX_RETRIES + 1): + try: + return await fn(*args, **kwargs) + except Exception as exc: # noqa: BLE001 - deliberate catch-all safety net + last = exc + _log(tool_name, exc, attempt) + if attempt < _MAX_RETRIES and _is_transient(exc): + await asyncio.sleep(_BACKOFF_SECONDS[attempt]) + continue + return _format_error( + tool_name, exc, transient=_is_transient(exc), attempts=attempt + 1 + ) + return _format_error(tool_name, last, transient=True, attempts=_MAX_RETRIES + 1) # type: ignore[arg-type] return tool(_async_wrapper) @functools.wraps(fn) def _sync_wrapper(*args: t.Any, **kwargs: t.Any) -> t.Any: - try: - return fn(*args, **kwargs) - except Exception as exc: # noqa: BLE001 — deliberate catch-all safety net - _log(tool_name, exc) - return _format_error(tool_name, exc) + last: BaseException | None = None + for attempt in range(_MAX_RETRIES + 1): + try: + return fn(*args, **kwargs) + except Exception as exc: # noqa: BLE001 - deliberate catch-all safety net + last = exc + _log(tool_name, exc, attempt) + if attempt < _MAX_RETRIES and _is_transient(exc): + time.sleep(_BACKOFF_SECONDS[attempt]) + continue + return _format_error( + tool_name, exc, transient=_is_transient(exc), attempts=attempt + 1 + ) + return _format_error(tool_name, last, transient=True, attempts=_MAX_RETRIES + 1) # type: ignore[arg-type] return tool(_sync_wrapper) @@ -85,9 +174,9 @@ def _is_async(fn: t.Callable[..., t.Any]) -> bool: return inspect.iscoroutinefunction(fn) -def _log(tool_name: str, exc: BaseException) -> None: +def _log(tool_name: str, exc: BaseException, attempt: int = 0) -> None: """Best-effort diagnostic to stderr (never to the user-facing return).""" try: - print(f"[AIRT] tool '{tool_name}' raised: {exc!r}", file=sys.stderr) + print(f"[AIRT] tool '{tool_name}' raised (attempt {attempt + 1}): {exc!r}", file=sys.stderr) except Exception: # noqa: BLE001 pass diff --git a/capabilities/ai-red-teaming/tools/results.py b/capabilities/ai-red-teaming/tools/results.py index e379344..26771f6 100644 --- a/capabilities/ai-red-teaming/tools/results.py +++ b/capabilities/ai-red-teaming/tools/results.py @@ -21,6 +21,7 @@ _errors_mod = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_errors_mod) safe_tool = _errors_mod.safe_tool +fmt_asr = _errors_mod.fmt_asr def _resolve_workspace_dir() -> Path: @@ -205,15 +206,10 @@ def get_analytics_summary( # surfaced to users (kept in the raw data for platform parity only). exec_stats = data.get("execution_stats", {}) if isinstance(data.get("execution_stats"), dict) else {} if "asr" in data: - _asr_pct = data["asr"] - lines.append(f"Success rate (ASR): {_asr_pct}% (probability {round(_asr_pct / 100, 3)})") + lines.append(f"Success rate (ASR): {fmt_asr(data['asr'])}") elif "overall_asr" in exec_stats: # SDK stores ASR as a 0-1 fraction under execution_stats. - _asr_frac = exec_stats["overall_asr"] - lines.append( - f"Success rate (ASR): {round(_asr_frac * 100, 1)}% " - f"(probability {round(_asr_frac, 3)})" - ) + lines.append(f"Success rate (ASR): {fmt_asr(exec_stats['overall_asr'])}") severity = data.get("severity_breakdown", data.get("severity", {})) if severity: diff --git a/capabilities/ai-red-teaming/tools/session.py b/capabilities/ai-red-teaming/tools/session.py index 4a92278..9637a01 100644 --- a/capabilities/ai-red-teaming/tools/session.py +++ b/capabilities/ai-red-teaming/tools/session.py @@ -23,6 +23,7 @@ _errors_mod = _ilu.module_from_spec(_spec) _spec.loader.exec_module(_errors_mod) safe_tool = _errors_mod.safe_tool +fmt_asr = _errors_mod.fmt_asr SESSION_PATH = Path( os.environ.get( @@ -143,7 +144,7 @@ def get_session_context() -> str: lines.append("") lines.append("Attack History ({} runs):".format(len(history))) for h in history[-5:]: # Show last 5 - score_str = "ASR={}%".format(h["best_score"]) if h.get("best_score") is not None else "no score" + score_str = "ASR={}".format(fmt_asr(h["best_score"])) if h.get("best_score") is not None else "no score" tx_str = "+{}".format(",".join(h["transforms"])) if h.get("transforms") else "" lines.append( " - {} {}: {} ({})".format(h.get("attack_type", "?"), tx_str, h.get("goal", "")[:40], score_str)