From fe177f8883325d50b484a98e484ec107893f563e Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 14 Jul 2026 17:30:36 -0700 Subject: [PATCH] FIX: CLI usability updates for scenario runs Several UX/diagnostic improvements to `pyrit_scan` output and scenario execution: - Fix technique progress count: `get_techniques_used()` now aggregates atomic-attack cells by display group, so the progress bar and summary report real technique counts (e.g. 2/14) instead of atomic-attack-cell counts that could exceed the total. - Resolve adversarial chat lazily for simulated-conversation techniques: add `AttackTechniqueFactory.resolve_adversarial_chat()` so matrix-built atomic attacks get the default adversarial target when a seed group carries a simulated conversation, fixing "adversarial_chat is required..." failures. - Include the atomic-attack name in execution error messages for easier triage. - Surface failed attacks and retry pressure in the run summary even on non-failing runs (error type/message per failed attack, total retries), and drop the redundant "Completed" line. - Stream per-attack retry warnings as each result lands during polling, including component role/name and endpoint from the captured RetryEvents. - `--start-server` now prints the log path on success and tails the log on crash/timeout. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 35659834-769b-4088-93cc-bd2603dcc65c --- .../backend/services/scenario_run_service.py | 37 +++++ pyrit/cli/_output.py | 93 +++++++++-- pyrit/cli/_server_launcher.py | 29 +++- pyrit/cli/pyrit_scan.py | 2 + pyrit/models/catalog/__init__.py | 4 + pyrit/models/catalog/scenario.py | 32 ++++ pyrit/models/results/scenario_result.py | 11 +- pyrit/scenario/core/atomic_attack.py | 4 +- .../scenario/core/attack_technique_factory.py | 20 +++ .../core/matrix_atomic_attack_builder.py | 2 +- pyrit/setup/initializers/techniques/core.py | 6 +- .../unit/backend/test_scenario_run_service.py | 98 ++++++++++++ tests/unit/cli/test_output.py | 146 +++++++++++++++++- tests/unit/cli/test_pyrit_scan.py | 61 ++++++++ tests/unit/cli/test_server_launcher.py | 88 +++++++++++ tests/unit/models/test_scenario_result.py | 17 ++ .../unit/scenario/core/test_atomic_attack.py | 2 +- .../core/test_attack_technique_factory.py | 51 ++++++ .../core/test_matrix_atomic_attack_builder.py | 19 +++ 19 files changed, 697 insertions(+), 25 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 3811caa12b..840858adef 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -18,6 +18,8 @@ from pyrit.memory import CentralMemory from pyrit.models import AttackOutcome, ScenarioResult, ScenarioRunState from pyrit.models.catalog.scenario import ( + AttackErrorSummary, + AttackRetrySummary, RunScenarioRequest, ScenarioRunSummary, ) @@ -597,6 +599,38 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari completed_attacks = total_attacks techniques_used = scenario_result.get_techniques_used() + # Surface per-attack errors and retry pressure regardless of overall run status: + # a COMPLETED scenario can still hide errored objectives or rate-limit retries. + failed_attacks: list[AttackErrorSummary] = [] + attack_retries: list[AttackRetrySummary] = [] + total_retries = 0 + for atomic_attack_name, results in scenario_result.attack_results.items(): + for attack_result in results: + retries = getattr(attack_result, "total_retries", 0) + if isinstance(retries, int): + total_retries += retries + + retry_events = getattr(attack_result, "retry_events", None) + if isinstance(retry_events, list) and retry_events: + attack_retries.append( + AttackRetrySummary( + attack_result_id=str(attack_result.attack_result_id), + atomic_attack_name=atomic_attack_name, + retries=retry_events, + ) + ) + + if attack_result.outcome == AttackOutcome.ERROR: + failed_attacks.append( + AttackErrorSummary( + atomic_attack_name=atomic_attack_name, + objective=attack_result.objective, + error_type=attack_result.error_type, + error_message=attack_result.error_message, + total_retries=retries if isinstance(retries, int) else 0, + ) + ) + return ScenarioRunSummary( scenario_result_id=scenario_result_id, scenario_name=scenario_result.scenario_name, @@ -610,6 +644,9 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari total_attacks=total_attacks, completed_attacks=completed_attacks, objective_achieved_rate=scenario_result.objective_achieved_rate(), + failed_attacks=failed_attacks, + attack_retries=attack_retries, + total_retries=total_retries, labels=scenario_result.labels, completed_at=scenario_result.completion_time, ) diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 2c1b57c654..7eb642aec6 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -50,6 +50,18 @@ def _header(text: str) -> None: _cprint(f"\n {text}", color="cyan", bold=True) +def _truncate_text(text: str, max_length: int) -> str: + """ + Truncate *text* to *max_length* characters, appending an ellipsis when cut. + + Returns: + str: The original text, or a truncated copy ending in ``...``. + """ + if len(text) <= max_length: + return text + return text[: max_length - 3].rstrip() + "..." + + def _wrap(*, text: str, indent: str, width: int = 78) -> str: """ Word-wrap *text* with the given *indent*. @@ -262,6 +274,54 @@ def print_dataset_list(*, items: list[dict[str, Any]]) -> None: # --------------------------------------------------------------------------- +def _format_retry_location(retry: Any) -> str: + """ + Build a human-readable "on , endpoint " clause for a retry. + + Returns: + str: The location clause, or an empty string when no context is available. + """ + bits: list[str] = [] + if retry.component_role: + role = retry.component_role.replace("_", " ") + bits.append(f"{role} {retry.component_name}" if retry.component_name else role) + if retry.endpoint: + bits.append(f"endpoint {retry.endpoint}") + return " on " + ", ".join(bits) if bits else "" + + +def print_scenario_retry_warnings(*, run: ScenarioRunSummary, seen_attack_ids: set[str]) -> None: + """ + Print retry warnings for attack results seen for the first time. + + Called during polling so retries stream to the console as each attack result + lands. ``seen_attack_ids`` is mutated to de-duplicate across polls. + + Args: + run: ``ScenarioRunSummary`` from ``GET /api/scenarios/runs/{id}``. + seen_attack_ids: Attack-result IDs already printed; updated in place. + """ + new_attacks = [a for a in run.attack_retries if a.attack_result_id not in seen_attack_ids] + if not new_attacks: + return + + # Finalize the in-place progress line (written with `\r`, no newline) so these + # warnings become persistent scrollback above the next progress redraw. + print() + for attack in new_attacks: + seen_attack_ids.add(attack.attack_result_id) + for retry in attack.retries: + exc = retry.exception_type or "error" + message = (retry.exception_message or "").strip().splitlines() + if message: + exc = f"{exc}: {_truncate_text(message[0], 160)}" + location = _format_retry_location(retry) + _cprint( + f" ! retry #{retry.attempt_number} [{attack.atomic_attack_name}]{location}: {exc}", + color="yellow", + ) + + def print_scenario_run_progress(*, run: ScenarioRunSummary, total_techniques: int = 0) -> None: """ Print a single-line progress update (overwrites the current line). @@ -277,19 +337,17 @@ def print_scenario_run_progress(*, run: ScenarioRunSummary, total_techniques: in parts: list[str] = [] + # The bar tracks techniques completed / total, which is the only ratio we can + # honestly compute mid-run: the server only knows about attacks already persisted, + # so an attacks-based bar would always read 100%. if effective_total > 0: - parts.append(f"techniques: {techniques_done}/{effective_total}") - elif techniques_done > 0: - parts.append(f"techniques: {techniques_done}") - - if run.total_attacks > 0: - pct = int((run.completed_attacks / run.total_attacks) * 100) + pct = int((techniques_done / effective_total) * 100) bar_width = 30 - filled = int(bar_width * run.completed_attacks / run.total_attacks) + filled = int(bar_width * techniques_done / effective_total) bar = "█" * filled + "░" * (bar_width - filled) - parts.append(f"[{bar}] {run.completed_attacks}/{run.total_attacks} attacks ({pct}%)") + parts.append(f"[{bar}] techniques: {techniques_done}/{effective_total} ({pct}%)") else: - parts.append(f"attacks: {run.completed_attacks}") + parts.append(f"techniques: {techniques_done}") parts.append(f"success rate: {run.objective_achieved_rate}%") parts.append(run.status.value) @@ -310,16 +368,29 @@ def print_scenario_run_summary(*, run: ScenarioRunSummary) -> None: print(f"\nScenario: {run.scenario_name}") print(f" Result ID: {run.scenario_result_id}") print(f" Status: {run.status.value}") - print(f" Total Attacks: {run.total_attacks}") - print(f" Completed: {run.completed_attacks}") + # Count of individual attack-result records persisted (one per technique x objective + # that ran), not a planned total. It stops growing wherever a failed run halted. + print(f" Attack Results: {run.total_attacks}") print(f" Success Rate: {run.objective_achieved_rate}%") + if run.total_retries: + print(f" Retries: {run.total_retries} (endpoint-stress signal)") + if run.error: print(f" Error: {run.error}") if run.techniques_used: print(f" Techniques: {', '.join(run.techniques_used)}") + if run.failed_attacks: + print(f"\n Failed Attacks ({len(run.failed_attacks)}):") + for failed in run.failed_attacks: + error_type = failed.error_type or "Error" + message = (failed.error_message or "").strip().splitlines() + detail = _truncate_text(message[0], 200) if message else "no detail" + retry_note = f" [{failed.total_retries} retries]" if failed.total_retries else "" + print(f" - {failed.atomic_attack_name}{retry_note}: {error_type}: {detail}") + # --------------------------------------------------------------------------- # Scenario run detail (full results via output module) diff --git a/pyrit/cli/_server_launcher.py b/pyrit/cli/_server_launcher.py index d541be4fd3..cdfffa8a6f 100644 --- a/pyrit/cli/_server_launcher.py +++ b/pyrit/cli/_server_launcher.py @@ -237,19 +237,46 @@ async def start_async( exit_code = self._process.poll() if exit_code is not None: + self._print_log_tail() raise RuntimeError( f"Server process exited with code {exit_code} during startup. See logs: {self._log_path}" ) if await self.probe_health_async(base_url=base_url): - print(f"Server ready (PID {self._pid})") + print(f"Server ready (PID {self._pid}). Logs: {self._log_path}") return base_url + self._print_log_tail() raise RuntimeError( f"pyrit_backend did not become healthy within {startup_timeout}s. " f"Check the server logs ({self._log_path}) or start it manually with: pyrit_backend" ) + def _read_log_tail(self, *, max_lines: int = 20) -> str: + """ + Read the last ``max_lines`` lines of the backend log file. + + Returns: + str: The tail of the log, or an empty string when the log is + unavailable or empty. + """ + if not self._log_path: + return "" + try: + with open(self._log_path, encoding="utf-8", errors="replace") as handle: + lines = handle.readlines() + except OSError: + return "" + return "".join(lines[-max_lines:]).rstrip() + + def _print_log_tail(self) -> None: + """Echo the tail of the backend log to stderr, if any is available.""" + tail = self._read_log_tail() + if tail: + print(f"\n--- pyrit_backend log ({self._log_path}) ---", file=sys.stderr) + print(tail, file=sys.stderr) + print("--- end of log ---", file=sys.stderr) + # ------------------------------------------------------------------ # Stop # ------------------------------------------------------------------ diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index 8c10d020cd..9c51481bc9 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -682,8 +682,10 @@ async def _poll_until_terminal_async( terminal_states = {ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, ScenarioRunState.CANCELLED} + seen_retry_attack_ids: set[str] = set() while True: run = await client.get_scenario_run_async(scenario_result_id=scenario_result_id) + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=seen_retry_attack_ids) _output.print_scenario_run_progress(run=run, total_techniques=total_techniques) if run.status in terminal_states: return run diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 79db260789..6d8e2e15d3 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -17,6 +17,8 @@ RegisteredInitializer, ) from pyrit.models.catalog.scenario import ( + AttackErrorSummary, + AttackRetrySummary, RegisteredScenario, RunScenarioRequest, ScenarioRunSummary, @@ -26,6 +28,8 @@ ) __all__ = [ + "AttackErrorSummary", + "AttackRetrySummary", "RegisteredInitializer", "RegisteredScenario", "RunScenarioRequest", diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index b33c7063f6..488ccf8c78 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -20,6 +20,7 @@ from pyrit.models.parameter import Parameter from pyrit.models.results.scenario_result import ScenarioRunState +from pyrit.models.retry_event import RetryEvent # Authoritative set of dataset seed filters exposed over the run request surface. Each entry # is used verbatim as a ``MemoryInterface.get_seeds`` keyword argument, so a filter key IS the @@ -115,6 +116,26 @@ def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[s return value +class AttackErrorSummary(BaseModel): + """A single errored attack result surfaced in a run summary.""" + + atomic_attack_name: str = Field(..., description="Atomic-attack cell that errored") + objective: str = Field("", description="Objective that was being attempted") + error_type: str | None = Field(None, description="Exception class name") + error_message: str | None = Field(None, description="Exception message") + total_retries: int = Field(0, ge=0, description="Retry attempts recorded for this attack") + + +class AttackRetrySummary(BaseModel): + """Retry events recorded for one attack result, for near-real-time CLI display.""" + + attack_result_id: str = Field(..., description="Stable ID of the attack result (used to de-duplicate)") + atomic_attack_name: str = Field(..., description="Atomic-attack cell that retried") + retries: list[RetryEvent] = Field( + default_factory=list, description="Retry attempts, each with component role/name, endpoint, and exception" + ) + + class ScenarioRunSummary(BaseModel): """Response for a scenario run (status + result details).""" @@ -130,5 +151,16 @@ class ScenarioRunSummary(BaseModel): total_attacks: int = Field(0, ge=0, description="Total number of attack results persisted for this run") completed_attacks: int = Field(0, ge=0, description="Number of attacks that reached a terminal outcome") objective_achieved_rate: int = Field(0, ge=0, le=100, description="Success rate as percentage (0-100)") + failed_attacks: list[AttackErrorSummary] = Field( + default_factory=list, + description="Individual attack results that errored, surfaced regardless of overall run status", + ) + attack_retries: list[AttackRetrySummary] = Field( + default_factory=list, + description="Per-attack retry events, surfaced as each attack result lands so the CLI can stream warnings", + ) + total_retries: int = Field( + 0, ge=0, description="Total retry attempts recorded across all attack results (endpoint-stress signal)" + ) labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run") completed_at: datetime | None = Field(None, description="When the scenario finished") diff --git a/pyrit/models/results/scenario_result.py b/pyrit/models/results/scenario_result.py index 3113db9288..793d8ce33f 100644 --- a/pyrit/models/results/scenario_result.py +++ b/pyrit/models/results/scenario_result.py @@ -155,13 +155,18 @@ def objective_scorer_identifier(self) -> ScorerIdentifier | None: def get_techniques_used(self) -> list[str]: """ - Get the list of techniques used in this scenario. + Get the list of techniques present in the results. + + Results are aggregated by display group so each technique is counted once, + even when it fans out into multiple atomic-attack cells (e.g. across + datasets or targets in a matrix scenario). When no ``display_group_map`` is + set, atomic-attack names are returned unchanged. Returns: - list[str]: Atomic attack technique names present in the results. + list[str]: Technique (display-group) names that produced results. """ - return list(self.attack_results.keys()) + return list(self.get_display_groups().keys()) def get_display_groups(self) -> dict[str, list[AttackResult]]: """ diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index 3f999e4aad..427b14384e 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -361,8 +361,8 @@ async def run_async( return results except Exception as e: - logger.error(f"Atomic attack execution failed: {str(e)}") - raise ValueError(f"Failed to execute atomic attack: {str(e)}") from e + logger.error(f"Atomic attack '{self.atomic_attack_name}' execution failed: {str(e)}") + raise ValueError(f"Failed to execute atomic attack '{self.atomic_attack_name}': {str(e)}") from e def _enrich_atomic_attack_identifiers(self, *, results: AttackExecutorResult[AttackResult]) -> None: """ diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index d8bbd60fc3..6a37a7d4b5 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -395,6 +395,26 @@ def adversarial_chat(self) -> PromptTarget | None: """The adversarial chat target baked into this factory, or None.""" return self._adversarial_chat + def resolve_adversarial_chat(self) -> PromptTarget | None: + """ + Resolve the adversarial chat target an ``AtomicAttack`` needs to expand this technique. + + A baked ``adversarial_chat`` always wins. Otherwise, when the technique's seed group + carries a simulated conversation (built via ``with_simulated_conversation``), the default + adversarial target is resolved lazily here — mirroring how ``create()`` resolves the target + for the attack's own ``attack_adversarial_config``. Techniques without a simulated + conversation seed do not need one and return ``None``. + + Returns: + PromptTarget | None: The adversarial chat target for the ``AtomicAttack``, or ``None`` + when the technique does not drive a simulated conversation. + """ + if self._adversarial_chat is not None: + return self._adversarial_chat + if self._seed_technique is not None and self._seed_technique.has_simulated_conversation: + return get_default_adversarial_target() + return None + @property def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" diff --git a/pyrit/scenario/core/matrix_atomic_attack_builder.py b/pyrit/scenario/core/matrix_atomic_attack_builder.py index a0f3e790d5..355b7abb92 100644 --- a/pyrit/scenario/core/matrix_atomic_attack_builder.py +++ b/pyrit/scenario/core/matrix_atomic_attack_builder.py @@ -346,7 +346,7 @@ def build( attack_technique=attack_technique, seed_groups=compatible_groups, adversarial_chat=( - target_instance if target_instance is not None else factory.adversarial_chat + target_instance if target_instance is not None else factory.resolve_adversarial_chat() ), objective_scorer=cast("TrueFalseScorer", self._objective_scorer), memory_labels=self._memory_labels, diff --git a/pyrit/setup/initializers/techniques/core.py b/pyrit/setup/initializers/techniques/core.py index e31b054ce7..824727bfbb 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -45,7 +45,11 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: Factories that need an adversarial chat target do not bake one in; the default adversarial target is resolved lazily inside - ``AttackTechniqueFactory.create`` via ``get_default_adversarial_target()``. + ``AttackTechniqueFactory.create`` (for the attack's own + ``attack_adversarial_config``) and inside + ``AttackTechniqueFactory.resolve_adversarial_chat`` (for the ``AtomicAttack`` + that expands a simulated-conversation seed), both via + ``get_default_adversarial_target()``. Returns: list[AttackTechniqueFactory]: The core scenario techniques. diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index fad28737d1..b22b704aee 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -859,6 +859,104 @@ def test_completed_run_still_shows_full_counts(self, mock_memory) -> None: assert fetched.objective_achieved_rate == 100 +class TestScenarioRunServiceFailedAttackReporting: + """Tests that per-attack errors and retry pressure surface in the summary.""" + + def test_error_attacks_and_retries_are_surfaced(self, mock_memory) -> None: + from pyrit.models import AttackOutcome + + success = MagicMock() + success.outcome = AttackOutcome.SUCCESS + success.total_retries = 2 + + errored = MagicMock() + errored.outcome = AttackOutcome.ERROR + errored.objective = "do the bad thing" + errored.error_type = "RateLimitError" + errored.error_message = "429 Too Many Requests" + errored.total_retries = 4 + + db_result = _make_db_scenario_result( + result_id="sr-mixed", + run_state="COMPLETED", + attack_results={"baseline_airt_hate": [success, errored]}, + ) + db_result.objective_achieved_rate.return_value = 50 + mock_memory.get_scenario_results.return_value = [db_result] + + service = ScenarioRunService() + fetched = service.get_run(scenario_result_id="sr-mixed") + + assert fetched is not None + assert fetched.total_retries == 6 + assert len(fetched.failed_attacks) == 1 + failed = fetched.failed_attacks[0] + assert failed.atomic_attack_name == "baseline_airt_hate" + assert failed.error_type == "RateLimitError" + assert failed.error_message == "429 Too Many Requests" + assert failed.total_retries == 4 + + def test_no_failed_attacks_when_all_succeed(self, mock_memory) -> None: + from pyrit.models import AttackOutcome + + success = MagicMock() + success.outcome = AttackOutcome.SUCCESS + success.total_retries = 0 + success.retry_events = [] + + db_result = _make_db_scenario_result( + result_id="sr-clean", + run_state="COMPLETED", + attack_results={"attack_a": [success]}, + ) + mock_memory.get_scenario_results.return_value = [db_result] + + service = ScenarioRunService() + fetched = service.get_run(scenario_result_id="sr-clean") + + assert fetched is not None + assert fetched.failed_attacks == [] + assert fetched.total_retries == 0 + assert fetched.attack_retries == [] + + def test_retry_events_surface_per_attack(self, mock_memory) -> None: + from pyrit.models import AttackOutcome + from pyrit.models.retry_event import RetryEvent + + attack = MagicMock() + attack.outcome = AttackOutcome.SUCCESS + attack.total_retries = 2 + attack.attack_result_id = "ar-9" + attack.retry_events = [ + RetryEvent( + attempt_number=1, + exception_type="RateLimitError", + exception_message="429", + component_role="objective_scorer", + component_name="TrueFalseScorer", + endpoint="https://ep/", + ) + ] + + db_result = _make_db_scenario_result( + result_id="sr-retry", + run_state="COMPLETED", + attack_results={"baseline_airt_hate": [attack]}, + ) + mock_memory.get_scenario_results.return_value = [db_result] + + service = ScenarioRunService() + fetched = service.get_run(scenario_result_id="sr-retry") + + assert fetched is not None + assert len(fetched.attack_retries) == 1 + summary = fetched.attack_retries[0] + assert summary.attack_result_id == "ar-9" + assert summary.atomic_attack_name == "baseline_airt_hate" + assert summary.retries[0].endpoint == "https://ep/" + assert summary.retries[0].component_role == "objective_scorer" + + class TestResolveTechniquesAndConverters: """Tests for per-technique converter resolution from ``--techniques`` tokens.""" diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index 3198bc4ef7..c3fd62a7c2 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -16,11 +16,14 @@ from pyrit.cli import _output from pyrit.models import Parameter, ScenarioRunState, TargetCapabilities, TargetIdentifier from pyrit.models.catalog import ( + AttackErrorSummary, + AttackRetrySummary, RegisteredInitializer, RegisteredScenario, ScenarioRunSummary, TargetInstance, ) +from pyrit.models.retry_event import RetryEvent from unit.mocks import make_scenario_result # --------------------------------------------------------------------------- @@ -379,13 +382,14 @@ def test_print_scenario_run_progress_with_known_totals(capsys): ) _output.print_scenario_run_progress(run=run, total_techniques=4) captured = capsys.readouterr() - assert "techniques: 2/4" in captured.out - assert "5/10" in captured.out + assert "techniques: 2/4 (50%)" in captured.out assert "IN_PROGRESS" in captured.out assert "30%" in captured.out + # Attacks are no longer surfaced in the progress line. + assert "attacks" not in captured.out -def test_print_scenario_run_progress_no_total_attacks(capsys): +def test_print_scenario_run_progress_no_techniques(capsys): run = _make_run( status=ScenarioRunState.CREATED, total_attacks=0, @@ -395,7 +399,7 @@ def test_print_scenario_run_progress_no_total_attacks(capsys): ) _output.print_scenario_run_progress(run=run, total_techniques=0) captured = capsys.readouterr() - assert "attacks: 0" in captured.out + assert "techniques: 0" in captured.out assert "CREATED" in captured.out @@ -413,10 +417,90 @@ def test_print_scenario_run_progress_techniques_done_only(capsys): # --------------------------------------------------------------------------- -# print_scenario_run_summary +# print_scenario_retry_warnings # --------------------------------------------------------------------------- +def _make_retry_event(**overrides) -> RetryEvent: + defaults = { + "attempt_number": 2, + "function_name": "_score_value_with_llm_async", + "exception_type": "RateLimitError", + "exception_message": "429 Too Many Requests\nsecond line", + "component_role": "objective_scorer", + "component_name": "TrueFalseScorer", + "endpoint": "https://example.openai.azure.com/", + } + defaults.update(overrides) + return RetryEvent(**defaults) + + +def test_print_scenario_retry_warnings_prints_new_attacks(capsys): + run = _make_run( + status=ScenarioRunState.IN_PROGRESS, + attack_retries=[ + AttackRetrySummary( + attack_result_id="ar-1", + atomic_attack_name="baseline_airt_hate", + retries=[_make_retry_event()], + ) + ], + ) + seen: set[str] = set() + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=seen) + out = capsys.readouterr().out + assert "retry #2" in out + assert "baseline_airt_hate" in out + assert "objective scorer TrueFalseScorer" in out + assert "endpoint https://example.openai.azure.com/" in out + assert "RateLimitError" in out + assert "429 Too Many Requests" in out + # Only the first line of the exception message is shown. + assert "second line" not in out + assert "ar-1" in seen + + +def test_print_scenario_retry_warnings_dedupes_across_polls(capsys): + attack = AttackRetrySummary( + attack_result_id="ar-1", + atomic_attack_name="baseline", + retries=[_make_retry_event()], + ) + run = _make_run(status=ScenarioRunState.IN_PROGRESS, attack_retries=[attack]) + seen: set[str] = set() + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=seen) + capsys.readouterr() # discard first print + # Second poll returns the same attack; nothing new should print. + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=seen) + assert capsys.readouterr().out == "" + + +def test_print_scenario_retry_warnings_noop_when_empty(capsys): + run = _make_run(status=ScenarioRunState.IN_PROGRESS, attack_retries=[]) + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=set()) + assert capsys.readouterr().out == "" + + +def test_print_scenario_retry_warnings_without_context(capsys): + run = _make_run( + status=ScenarioRunState.IN_PROGRESS, + attack_retries=[ + AttackRetrySummary( + attack_result_id="ar-2", + atomic_attack_name="crescendo", + retries=[_make_retry_event(component_role="", component_name=None, endpoint=None)], + ) + ], + ) + _output.print_scenario_retry_warnings(run=run, seen_attack_ids=set()) + out = capsys.readouterr().out + assert "retry #2 [crescendo]: RateLimitError" in out + assert " on " not in out + + +# --------------------------------------------------------------------------- +# print_scenario_run_summary +# --------------------------------------------------------------------------- def test_print_scenario_run_summary_completed(capsys): run = _make_run( scenario_name="test_sc", @@ -434,6 +518,9 @@ def test_print_scenario_run_summary_completed(capsys): assert "COMPLETED" in captured.out assert "40%" in captured.out assert "s1, s2" in captured.out + # The count is relabeled and the redundant "Completed" line is gone. + assert "Attack Results: 5" in captured.out + assert "Completed:" not in captured.out def test_print_scenario_run_summary_with_error(capsys): @@ -452,6 +539,55 @@ def test_print_scenario_run_summary_with_error(capsys): assert "boom" in captured.out +def test_print_scenario_run_summary_lists_failed_attacks(capsys): + run = _make_run( + scenario_name="failing", + status=ScenarioRunState.COMPLETED, + total_attacks=4, + completed_attacks=4, + objective_achieved_rate=75, + failed_attacks=[ + AttackErrorSummary( + atomic_attack_name="baseline_airt_hate", + objective="do the bad thing", + error_type="RateLimitError", + error_message="429 Too Many Requests\nsecond line ignored", + total_retries=3, + ) + ], + ) + _output.print_scenario_run_summary(run=run) + out = capsys.readouterr().out + assert "Failed Attacks (1):" in out + assert "baseline_airt_hate" in out + assert "RateLimitError" in out + assert "429 Too Many Requests" in out + assert "[3 retries]" in out + # Only the first line of a multi-line message is shown. + assert "second line ignored" not in out + + +def test_print_scenario_run_summary_shows_retry_pressure(capsys): + run = _make_run( + scenario_name="stressed", + status=ScenarioRunState.COMPLETED, + total_attacks=6, + completed_attacks=6, + objective_achieved_rate=100, + total_retries=9, + ) + _output.print_scenario_run_summary(run=run) + out = capsys.readouterr().out + assert "Retries:" in out + assert "9" in out + + +def test_print_scenario_run_summary_hides_retry_line_when_zero(capsys): + run = _make_run(status=ScenarioRunState.COMPLETED, total_retries=0) + _output.print_scenario_run_summary(run=run) + assert "Retries:" not in capsys.readouterr().out + + # --------------------------------------------------------------------------- # print_scenario_result_async # --------------------------------------------------------------------------- diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index cb196b60ce..d78ec4073c 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -1273,3 +1273,64 @@ def test_parse_args_tolerates_scenario_specific_flags(self): assert parsed.scenario_name == "foo" assert parsed.target == "t" assert parsed._unknown_args == ["--max-turns", "7"] + + +class TestPollStreamsRetryWarnings: + """The poll loop should stream retry warnings as attack results land.""" + + async def test_poll_prints_retry_warnings_once(self, capsys): + from datetime import datetime, timezone + + from pyrit.models import ScenarioRunState + from pyrit.models.catalog import AttackRetrySummary, ScenarioRunSummary + from pyrit.models.retry_event import RetryEvent + + now = datetime(2025, 1, 1, tzinfo=timezone.utc) + retry = RetryEvent( + attempt_number=3, + exception_type="RateLimitError", + exception_message="429 Too Many Requests", + component_role="objective_scorer", + component_name="TrueFalseScorer", + endpoint="https://ep/", + ) + + def _summary(*, status, with_retry): + return ScenarioRunSummary( + scenario_result_id="sr-1", + scenario_name="s", + scenario_version=0, + status=status, + created_at=now, + updated_at=now, + techniques_used=["a"], + total_attacks=1, + completed_attacks=1, + objective_achieved_rate=0, + attack_retries=( + [AttackRetrySummary(attack_result_id="ar-1", atomic_attack_name="baseline", retries=[retry])] + if with_retry + else [] + ), + ) + + client = MagicMock() + # First poll: retry present but still running. Second poll: same retry, terminal. + client.get_scenario_run_async = AsyncMock( + side_effect=[ + _summary(status=ScenarioRunState.IN_PROGRESS, with_retry=True), + _summary(status=ScenarioRunState.COMPLETED, with_retry=True), + ] + ) + + with patch("asyncio.sleep", new=AsyncMock(return_value=None)): + final = await pyrit_scan._poll_until_terminal_async( + client=client, scenario_result_id="sr-1", total_techniques=1 + ) + + assert final.status == ScenarioRunState.COMPLETED + out = capsys.readouterr().out + # The warning is printed exactly once despite appearing in both polls. + assert out.count("retry #3") == 1 + assert "RateLimitError" in out + assert "endpoint https://ep/" in out diff --git a/tests/unit/cli/test_server_launcher.py b/tests/unit/cli/test_server_launcher.py index 4036f0960d..7c5201d0b2 100644 --- a/tests/unit/cli/test_server_launcher.py +++ b/tests/unit/cli/test_server_launcher.py @@ -141,6 +141,94 @@ async def test_start_async_raises_when_timeout_exhausted(): await launcher.start_async(host="localhost", port=8000, startup_timeout=2) +async def test_start_async_prints_log_path_on_success(capsys): + launcher = ServerLauncher() + fake_proc = MagicMock() + fake_proc.pid = 4321 + fake_proc.poll.return_value = None + probe = AsyncMock(side_effect=[False, True]) + + with ( + patch.object(ServerLauncher, "probe_health_async", new=probe), + patch("subprocess.Popen", return_value=fake_proc), + patch("asyncio.sleep", new=AsyncMock(return_value=None)), + ): + await launcher.start_async(host="localhost", port=8001, startup_timeout=5) + + out = capsys.readouterr().out + assert "Server ready (PID 4321)" in out + assert "Logs:" in out + assert launcher._log_path in out + + +async def test_start_async_echoes_log_tail_on_crash(capsys): + launcher = ServerLauncher() + fake_proc = MagicMock() + fake_proc.pid = 42 + fake_proc.poll.return_value = 1 # exited + probe = AsyncMock(return_value=False) + + with ( + patch.object(ServerLauncher, "probe_health_async", new=probe), + patch("subprocess.Popen", return_value=fake_proc), + patch("asyncio.sleep", new=AsyncMock(return_value=None)), + patch.object(ServerLauncher, "_read_log_tail", return_value="ERROR: port already in use"), + ): + with pytest.raises(RuntimeError, match="exited with code 1"): + await launcher.start_async(host="localhost", port=8000, startup_timeout=3) + + err = capsys.readouterr().err + assert "pyrit_backend log" in err + assert "ERROR: port already in use" in err + + +async def test_start_async_echoes_log_tail_on_timeout(capsys): + launcher = ServerLauncher() + fake_proc = MagicMock() + fake_proc.pid = 99 + fake_proc.poll.return_value = None # still running + probe = AsyncMock(return_value=False) + + with ( + patch.object(ServerLauncher, "probe_health_async", new=probe), + patch("subprocess.Popen", return_value=fake_proc), + patch("asyncio.sleep", new=AsyncMock(return_value=None)), + patch.object(ServerLauncher, "_read_log_tail", return_value="Traceback: boom"), + ): + with pytest.raises(RuntimeError, match="did not become healthy"): + await launcher.start_async(host="localhost", port=8000, startup_timeout=2) + + err = capsys.readouterr().err + assert "Traceback: boom" in err + + +# --------------------------------------------------------------------------- +# _read_log_tail +# --------------------------------------------------------------------------- + + +def test_read_log_tail_returns_last_lines(tmp_path): + log_file = tmp_path / "pyrit_backend.log" + log_file.write_text("\n".join(f"line {i}" for i in range(50)), encoding="utf-8") + launcher = ServerLauncher() + launcher._log_path = str(log_file) + + tail = launcher._read_log_tail(max_lines=5) + + assert tail.splitlines() == ["line 45", "line 46", "line 47", "line 48", "line 49"] + + +def test_read_log_tail_returns_empty_when_no_log_path(): + launcher = ServerLauncher() + assert launcher._read_log_tail() == "" + + +def test_read_log_tail_returns_empty_when_file_missing(tmp_path): + launcher = ServerLauncher() + launcher._log_path = str(tmp_path / "does_not_exist.log") + assert launcher._read_log_tail() == "" + + # --------------------------------------------------------------------------- # stop # --------------------------------------------------------------------------- diff --git a/tests/unit/models/test_scenario_result.py b/tests/unit/models/test_scenario_result.py index 490602937a..64b2ecb6ec 100644 --- a/tests/unit/models/test_scenario_result.py +++ b/tests/unit/models/test_scenario_result.py @@ -65,6 +65,23 @@ def test_get_techniques_used(self): techniques = result.get_techniques_used() assert sorted(techniques) == ["crescendo", "flip"] + def test_get_techniques_used_aggregates_by_display_group(self): + # A single technique fanned out across two datasets produces two atomic-attack + # cells; get_techniques_used should collapse them back to one technique. + result = make_scenario_result( + scenario_name="TestScenario", + objective_target_identifier=ComponentIdentifier.model_validate({}), + attack_results={"crescendo_advbench": [], "crescendo_harmbench": [], "flip_advbench": []}, + objective_scorer_identifier=ComponentIdentifier.model_validate({}), + display_group_map={ + "crescendo_advbench": "crescendo", + "crescendo_harmbench": "crescendo", + "flip_advbench": "flip", + }, + ) + techniques = result.get_techniques_used() + assert sorted(techniques) == ["crescendo", "flip"] + def test_get_objectives_all(self): ar1 = _make_attack_result(objective="obj1") ar2 = _make_attack_result(objective="obj2") diff --git a/tests/unit/scenario/core/test_atomic_attack.py b/tests/unit/scenario/core/test_atomic_attack.py index b82fb08b95..28fd9d8004 100644 --- a/tests/unit/scenario/core/test_atomic_attack.py +++ b/tests/unit/scenario/core/test_atomic_attack.py @@ -367,7 +367,7 @@ async def test_run_async_handles_execution_failure(self, mock_attack, sample_see with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: mock_exec.side_effect = Exception("Execution error") - with pytest.raises(ValueError, match="Failed to execute atomic attack"): + with pytest.raises(ValueError, match="Failed to execute atomic attack 'Test Attack Run'"): await atomic_attack.run_async() async def test_run_async_passes_return_partial_on_failure_true_by_default( diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index 7ae54e3c70..65574469cd 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -839,6 +839,57 @@ def test_create_custom_prompt_conflicts_with_baked_raises(self): ) +class TestResolveAdversarialChat: + class _AdversarialAttack: + def __init__(self, *, objective_target=None, attack_scoring_config=None, attack_adversarial_config=None): + self.attack_adversarial_config = attack_adversarial_config + + def get_identifier(self): + return ComponentIdentifier(class_name="_AdversarialAttack", class_module="test") + + def test_returns_baked_adversarial_chat(self): + """A baked adversarial_chat is returned without resolving a default.""" + target = MagicMock(spec=PromptTarget) + factory = AttackTechniqueFactory( + name="durian", + attack_class=self._AdversarialAttack, + adversarial_chat=target, + ) + with patch( + "pyrit.scenario.core.attack_technique_factory.get_default_adversarial_target", + ) as mock_default: + assert factory.resolve_adversarial_chat() is target + mock_default.assert_not_called() + + def test_returns_none_for_non_adversarial_technique(self): + """A technique with no baked chat and no simulated conversation needs no adversarial chat.""" + factory = AttackTechniqueFactory(name="durian", attack_class=_StubAttack) + with patch( + "pyrit.scenario.core.attack_technique_factory.get_default_adversarial_target", + ) as mock_default: + assert factory.resolve_adversarial_chat() is None + mock_default.assert_not_called() + + def test_simulated_conversation_resolves_default_lazily(self): + """A simulated-conversation technique with no baked chat resolves the default target.""" + from pyrit.common.path import EXECUTOR_SEED_PROMPT_PATH + + factory = AttackTechniqueFactory.with_simulated_conversation( + name="role_play_movie_script", + adversarial_chat_system_prompt_path=( + EXECUTOR_SEED_PROMPT_PATH / "red_teaming" / "role_play" / "role_play_movie_script.yaml" + ), + num_turns=2, + ) + default_target = MagicMock(spec=PromptTarget) + with patch( + "pyrit.scenario.core.attack_technique_factory.get_default_adversarial_target", + return_value=default_target, + ) as mock_default: + assert factory.resolve_adversarial_chat() is default_target + mock_default.assert_called_once() + + class TestUnwrapOptional: """Tests for AttackTechniqueFactory._unwrap_optional static method.""" diff --git a/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py index c8667489bc..5e36f7a3b1 100644 --- a/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py +++ b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py @@ -44,6 +44,10 @@ def _mock_factory(*, name: str, seed_technique=None, adversarial_chat=None) -> M factory.name = name factory.seed_technique = seed_technique factory.adversarial_chat = adversarial_chat + # Mirror the real factory: with no simulated-conversation seed, resolution returns + # whatever adversarial_chat was baked (possibly None). Tests that exercise the + # simulated-conversation lazy path override this return value explicitly. + factory.resolve_adversarial_chat.return_value = adversarial_chat factory.create.return_value = MagicMock(name=f"{name}_technique") return factory @@ -170,6 +174,21 @@ def test_no_target_axis_uses_factory_adversarial_chat(self): assert "adversarial_chat" not in factory.create.call_args.kwargs assert result[0]._adversarial_chat is baked + def test_no_target_axis_stamps_factory_resolved_adversarial_chat(self): + # A simulated-conversation technique with no baked adversarial_chat resolves the + # default lazily via factory.resolve_adversarial_chat(); the builder must stamp that + # onto the AtomicAttack so seed-group expansion has an adversarial chat to use. + builder = _builder() + resolved = MagicMock(spec=PromptTarget) + factory = _mock_factory(name="tech") + factory.resolve_adversarial_chat.return_value = resolved + result = builder.build( + technique_factories={"tech": factory}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert "adversarial_chat" not in factory.create.call_args.kwargs + assert result[0]._adversarial_chat is resolved + @pytest.mark.usefixtures("patch_central_database") class TestMatrixCustomCallbacks: