MAINT: Updating CLI to use updated estimates

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)

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

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,

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,

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,37 +400,37 @@ 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,
         completed_attacks=0,
         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
 
 
 # ---------------------------------------------------------------------------

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