diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 0883ee5d46..8fa83950ff 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -324,36 +324,26 @@ def print_scenario_retry_warnings(*, run: ScenarioRunSummary, seen_attack_ids: s ) -def print_scenario_run_progress(*, run: ScenarioRunSummary, total_techniques: int = 0) -> None: +def print_scenario_run_progress(*, run: ScenarioRunSummary) -> None: """ Print a single-line progress update (overwrites the current line). Args: run: ``ScenarioRunSummary`` from ``GET /api/scenarios/runs/{id}``. - total_techniques: Total number of techniques expected (0 if unknown). """ - techniques_done = len(run.techniques_used) - # Techniques the user passed may be aggregates that expand on the server - # (e.g. `single_turn` -> N concrete techniques). Trust whichever count is larger. - effective_total = max(total_techniques, techniques_done) - 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: - pct = int((techniques_done / effective_total) * 100) + if run.total_attacks > 0: + pct = int((run.completed_attacks / run.total_attacks) * 100) bar_width = 30 - filled = int(bar_width * techniques_done / effective_total) + filled = int(bar_width * run.completed_attacks / run.total_attacks) bar = "█" * filled + "░" * (bar_width - filled) try: bar.encode(sys.stdout.encoding or "utf-8") except (LookupError, UnicodeEncodeError): bar = "#" * filled + "-" * (bar_width - filled) - parts.append(f"[{bar}] techniques: {techniques_done}/{effective_total} ({pct}%)") + parts.append(f"[{bar}] units: {run.completed_attacks}/{run.total_attacks} ({pct}%)") else: - parts.append(f"techniques: {techniques_done}") + parts.append(f"units: {run.completed_attacks}") parts.append(f"success rate: {run.objective_achieved_rate}%") parts.append(run.status.value) diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index 1515a5fb36..2930a78f66 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -39,7 +39,6 @@ from collections.abc import Callable from pyrit.models.catalog import ( - RegisteredScenario, RunScenarioRequest, ScenarioRunSummary, ) @@ -990,7 +989,6 @@ async def _poll_until_terminal_async( *, client: Any, scenario_result_id: str, - total_techniques: int, ) -> ScenarioRunSummary: """ Poll the server until the run reaches a terminal status. @@ -1007,7 +1005,7 @@ async def _poll_until_terminal_async( while True: run: ScenarioRunSummary = 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) + _output.print_scenario_run_progress(run=run) if run.status in terminal_states: return run await asyncio.sleep(0.5) @@ -1017,7 +1015,6 @@ async def _run_scenario_async( *, client: Any, parsed_args: Namespace, - scenario_meta: RegisteredScenario, ) -> int: """ Start a scenario run, poll for completion, and print results. @@ -1031,7 +1028,6 @@ async def _run_scenario_async( scenario_name = parsed_args.scenario_name request = _build_run_request(parsed_args=parsed_args, scenario_name=scenario_name) - total_techniques = len(request.techniques or scenario_meta.all_techniques or []) print(f"\nRunning scenario: {scenario_name}") sys.stdout.flush() @@ -1055,7 +1051,6 @@ async def _run_scenario_async( run = await _poll_until_terminal_async( client=client, scenario_result_id=scenario_result_id, - total_techniques=total_techniques, ) except KeyboardInterrupt: print("\n\nCancelling scenario run...") @@ -1110,7 +1105,7 @@ async def _handle_run_async(*, client: Any, parsed_args: Namespace) -> int: if reparsed is None: return 1 - return await _run_scenario_async(client=client, parsed_args=reparsed, scenario_meta=scenario_meta) + return await _run_scenario_async(client=client, parsed_args=reparsed) #: Post-client verbs, each a uniform ``(*, client, parsed_args) -> int`` handler. Reached diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index b6147f38f7..806b7aa092 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -528,7 +528,6 @@ def do_run(self, line: str) -> None: request = RunScenarioRequest(**request_kwargs) # Start run - total_techniques = len(request.techniques or []) print(f"\nRunning scenario: {scenario_name}") sys.stdout.flush() @@ -559,7 +558,7 @@ def do_run(self, line: str) -> None: try: while True: run = self._run_async(self._api_client.get_scenario_run_async(scenario_result_id=scenario_result_id)) - print_scenario_run_progress(run=run, total_techniques=total_techniques) + print_scenario_run_progress(run=run) if run.status in { ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 11c7c2e269..4e72b0e722 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -323,8 +323,10 @@ class ScenarioRunSummary(BaseModel): error: str | None = Field(None, description="Error message if status is FAILED") error_type: str | None = Field(None, description="Exception class name if status is FAILED") techniques_used: list[str] = Field(default_factory=list, description="Technique names that were executed") - 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") + total_attacks: int = Field( + 0, ge=0, description="Planned execution units, or the observed units when no plan is persisted" + ) + completed_attacks: int = Field(0, ge=0, description="Planned execution units 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, diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index c361d213dd..546936bd11 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -381,13 +381,11 @@ def test_print_scenario_run_progress_with_known_totals(capsys): objective_achieved_rate=30, techniques_used=["s1", "s2"], ) - _output.print_scenario_run_progress(run=run, total_techniques=4) + _output.print_scenario_run_progress(run=run) captured = capsys.readouterr() - assert "techniques: 2/4 (50%)" in captured.out + assert "units: 5/10 (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_uses_ascii_for_limited_console(): @@ -402,13 +400,13 @@ def test_print_scenario_run_progress_uses_ascii_for_limited_console(): stdout.encoding = "cp1252" with patch.object(_output.sys, "stdout", stdout): - _output.print_scenario_run_progress(run=run, total_techniques=2) + _output.print_scenario_run_progress(run=run) line = stdout.write.call_args.args[0] assert "[###############---------------]" in line -def test_print_scenario_run_progress_no_techniques(capsys): +def test_print_scenario_run_progress_no_units(capsys): run = _make_run( status=ScenarioRunState.CREATED, total_attacks=0, @@ -416,23 +414,23 @@ def test_print_scenario_run_progress_no_techniques(capsys): objective_achieved_rate=0, techniques_used=[], ) - _output.print_scenario_run_progress(run=run, total_techniques=0) + _output.print_scenario_run_progress(run=run) captured = capsys.readouterr() - assert "techniques: 0" in captured.out + assert "units: 0" in captured.out assert "CREATED" in captured.out -def test_print_scenario_run_progress_techniques_done_only(capsys): +def test_print_scenario_run_progress_completed_units_without_plan(capsys): run = _make_run( status=ScenarioRunState.IN_PROGRESS, total_attacks=0, - completed_attacks=0, + completed_attacks=1, objective_achieved_rate=0, techniques_used=["s1"], ) - _output.print_scenario_run_progress(run=run, total_techniques=0) + _output.print_scenario_run_progress(run=run) captured = capsys.readouterr() - assert "techniques: 1" in captured.out + assert "units: 1" in captured.out # --------------------------------------------------------------------------- diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 828c47f669..eba9f6e5eb 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -1808,9 +1808,7 @@ def _summary(*, status, with_retry): ) 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 - ) + final = await pyrit_scan._poll_until_terminal_async(client=client, scenario_result_id="sr-1") assert final.status == ScenarioRunState.COMPLETED out = capsys.readouterr().out