diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index bfcf6bb296..1515a5fb36 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -46,6 +46,37 @@ from pyrit.models.parameter import Parameter +def _is_read_timeout(exc: BaseException) -> bool: + """ + Report whether an exception is an ``httpx.ReadTimeout``. + + Args: + exc (BaseException): The exception to classify. + + Returns: + bool: True when httpx is importable and the exception is a read timeout. + """ + try: + import httpx + + return isinstance(exc, httpx.ReadTimeout) + except Exception: + return False + + +def _print_debug_traceback(exc: BaseException) -> None: + """ + Print the traceback for an exception when the log level is ``DEBUG``. + + Args: + exc (BaseException): The exception caught by the CLI. + """ + if logging.getLogger().isEnabledFor(logging.DEBUG): + import traceback + + traceback.print_exception(type(exc), exc, exc.__traceback__) + + def _print_cli_exception(*, exc: BaseException) -> None: """ Print a user-facing error line for an exception that bubbled out of the CLI. @@ -56,32 +87,24 @@ def _print_cli_exception(*, exc: BaseException) -> None: the server is taking longer than ``--request-timeout`` to respond and the default bare ``str(exc)`` is empty. + Only ``pyrit_scan`` verbs reach this with a timeout; ``pyrit_shell`` has no + ``--request-timeout`` and handles that case before it gets here. + Args: exc (BaseException): The exception caught by the CLI. """ - import traceback - - try: - import httpx - - is_read_timeout = isinstance(exc, httpx.ReadTimeout) - except Exception: - is_read_timeout = False - cls_name = type(exc).__name__ detail = str(exc) or repr(exc) - if is_read_timeout: + if _is_read_timeout(exc): print( - "\nError (ReadTimeout): server did not respond in time. " - "Pass '--request-timeout ' to wait longer, or check the " - "server logs for a blocked event loop." + "\nError (ReadTimeout): server did not respond in time. Pass '--request-timeout " + "' to wait longer, or check the server logs for a blocked event loop." ) else: print(f"\nError ({cls_name}): {detail}") - if logging.getLogger().isEnabledFor(logging.DEBUG): - traceback.print_exception(type(exc), exc, exc.__traceback__) + _print_debug_traceback(exc) _DESCRIPTION = """PyRIT Scanner - Run AI security scenarios from the command line. @@ -1015,7 +1038,15 @@ async def _run_scenario_async( try: run = await client.start_scenario_run_async(request=request) except Exception as exc: - print(f"Error starting scenario: {exc}") + if _is_read_timeout(exc): + # The server keeps initializing after the client stops waiting, so whether the run + # started is genuinely unknown here and must not be reported as a failure to start. + print("\nERROR: The scenario start request timed out, so it is unknown whether the run started.") + _print_cli_exception(exc=exc) + print("Check 'pyrit_scan scenario-history' before retrying, or the run may be started twice.") + else: + print("\nERROR: The scenario could not be started.") + _print_cli_exception(exc=exc) return 1 scenario_result_id = run.scenario_result_id @@ -1045,6 +1076,8 @@ async def _run_scenario_async( "retrieved or parsed from the server." ) _print_cli_exception(exc=exc) + if _is_read_timeout(exc): + print(f"Retry with 'pyrit_scan scenario-results {scenario_result_id}'.") _output.print_scenario_run_summary(run=run) return 1 return 0 diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 80ea4419b8..b6147f38f7 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -63,6 +63,27 @@ def _strip_surrounding_quotes(token: str) -> str: return token +def _print_shell_exception(*, exc: BaseException) -> None: + """ + Print a user-facing error line for an exception raised by a shell command. + + Mirrors ``pyrit_scan._print_cli_exception`` but never suggests ``--request-timeout``, + which the shell does not accept. A bare ``httpx.ReadTimeout`` stringifies to nothing, + so it needs its own line or the command reports an empty error. + + Args: + exc (BaseException): The exception caught by the shell command. + """ + from pyrit.cli.pyrit_scan import _is_read_timeout, _print_debug_traceback + + if _is_read_timeout(exc): + print("\nError (ReadTimeout): server did not respond in time. Check the server logs for a blocked event loop.") + else: + print(f"\nError ({type(exc).__name__}): {str(exc) or repr(exc)}") + + _print_debug_traceback(exc) + + class PyRITShell(cmd.Cmd): """ Interactive shell for PyRIT (thin REST client). @@ -437,6 +458,7 @@ def do_run(self, line: str) -> None: print_scenario_run_progress, print_scenario_run_summary, ) + from pyrit.cli.pyrit_scan import _is_read_timeout, _print_cli_exception, _print_debug_traceback from pyrit.models import ScenarioRunState from pyrit.models.catalog import RunScenarioRequest @@ -513,7 +535,20 @@ def do_run(self, line: str) -> None: try: run = self._run_async(self._api_client.start_scenario_run_async(request=request)) except Exception as exc: - print(f"Error starting scenario: {exc}") + if _is_read_timeout(exc): + # The server keeps initializing after the client stops waiting, so whether the + # run started is unknown. The shell has no --request-timeout, so it handles the + # timeout itself rather than going through the shared printer. + print("\nERROR: The scenario start request timed out, so it is unknown whether the run started.") + print( + "\nError (ReadTimeout): server did not respond in time. Check " + "'scenario-history' before retrying, or the run may be started twice. Check " + "the server logs for a blocked event loop." + ) + _print_debug_traceback(exc) + else: + print("\nERROR: The scenario could not be started.") + _print_cli_exception(exc=exc) return scenario_result_id = run.scenario_result_id @@ -550,13 +585,21 @@ def do_run(self, line: str) -> None: ) self._run_async(print_scenario_result_async(result=detail)) except Exception as exc: - from pyrit.cli.pyrit_scan import _print_cli_exception - print( "\nERROR: The scenario completed, but its detailed results could not be " "retrieved or parsed from the server." ) - _print_cli_exception(exc=exc) + if _is_read_timeout(exc): + # The shell has no --request-timeout, so it must not reach the shared + # printer, which advises it. + print( + f"\nError (ReadTimeout): server did not respond in time. Retry with " + f"'scenario-results {scenario_result_id}', or check the server logs for a " + "blocked event loop." + ) + _print_debug_traceback(exc) + else: + _print_cli_exception(exc=exc) print_scenario_run_summary(run=run) else: print_scenario_run_summary(run=run) @@ -591,7 +634,7 @@ def do_scenario_history(self, arg: str) -> None: runs = self._run_async(self._api_client.list_scenario_runs_async(limit=limit)) print_scenario_runs_list(runs=runs) except Exception as e: - print(f"Error: {e}") + _print_shell_exception(exc=e) def do_scenario_results(self, arg: str) -> None: """ @@ -655,7 +698,7 @@ def do_scenario_results(self, arg: str) -> None: self._api_client.get_scenario_run_results_async(scenario_result_id=parsed.scenario_result_id) ) except Exception as exc: - print(f"Error: {exc}") + _print_shell_exception(exc=exc) return if view is ScenarioResultView.OVERVIEW: @@ -684,7 +727,7 @@ def do_scenario_results(self, arg: str) -> None: ) ) except Exception as exc: - print(f"Error: {exc}") + _print_shell_exception(exc=exc) return print_conversations(payload=conversations_payload) diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 9fd5aaaa8b..828c47f669 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -1262,6 +1262,31 @@ def test_main_start_scenario_failure(self, mock_client_class, _mock_probe, capsy captured = capsys.readouterr() assert "server full" in captured.out + @patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ) + @patch("pyrit.cli.api_client.PyRITApiClient") + def test_main_start_scenario_read_timeout_reports_type_and_hint(self, mock_client_class, _mock_probe, capsys): + """A ReadTimeout stringifies to '', so the type and a hint have to carry the message.""" + import httpx + + mock_client = _mock_api_client() + mock_client.start_scenario_run_async.side_effect = httpx.ReadTimeout("") + mock_client_class.return_value = mock_client + + result = pyrit_scan.main(["run", "test_scenario", "--target", "t"]) + assert result == 1 + captured = capsys.readouterr() + assert "ReadTimeout" in captured.out + assert "--request-timeout" in captured.out + # The server keeps initializing after the client gives up, so the outcome is unknown + # and must not be reported as a definite failure to start. + assert "unknown whether the run started" in captured.out + assert "could not be started" not in captured.out + assert "scenario-history" in captured.out + @patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, @@ -1283,6 +1308,38 @@ def test_main_run_results_failure_is_hard_error(self, mock_client_class, _mock_p # The summary printer should still be used as a fallback for context. assert "test_scenario" in captured.out + @patch( + "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", + new_callable=AsyncMock, + return_value=True, + ) + @patch("pyrit.cli.api_client.PyRITApiClient") + def test_main_run_results_read_timeout_points_at_scenario_results(self, mock_client_class, _mock_probe, capsys): + """Results are fetched with the request timeout, unlike polling, so they can time out.""" + import httpx + + mock_client = _mock_api_client() + mock_client.get_scenario_run_results_async.side_effect = httpx.ReadTimeout("") + mock_client_class.return_value = mock_client + + result = pyrit_scan.main(["run", "test_scenario", "--target", "t"]) + assert result == 1 + captured = capsys.readouterr() + assert "ReadTimeout" in captured.out + assert "scenario-results" in captured.out + assert "--request-timeout" in captured.out + + def test_print_cli_exception_surfaces_empty_read_timeout(self, capsys): + """A bare ReadTimeout stringifies to '', so the helper has to carry the message.""" + import httpx + + pyrit_scan._print_cli_exception(exc=httpx.ReadTimeout("")) + captured = capsys.readouterr() + assert "ReadTimeout" in captured.out + assert "did not respond in time" in captured.out + # Only pyrit_scan verbs reach the helper with a timeout, and they all take the flag. + assert "--request-timeout" in captured.out + @patch( "pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index 25be9e4d4e..339d36c619 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -617,7 +617,47 @@ def test_run_start_failure(self, shell, capsys): return_value={"scenario_name": "foo", "target": "t"}, ): s.do_run("foo --target t") - assert "Error starting scenario: nope" in capsys.readouterr().out + out = capsys.readouterr().out + assert "The scenario could not be started." in out + assert "Error (RuntimeError): nope" in out + + def test_run_start_failure_read_timeout_reports_type_and_hint(self, shell, capsys): + """A ReadTimeout stringifies to '', so the type and a hint have to carry the message.""" + import httpx + + s, client = shell + client.start_scenario_run_async = AsyncMock(side_effect=httpx.ReadTimeout("")) + with patch( + "pyrit.cli._cli_args.parse_run_arguments", + return_value={"scenario_name": "foo", "target": "t"}, + ): + s.do_run("foo --target t") + out = capsys.readouterr().out + assert "ReadTimeout" in out + assert "blocked event loop" in out + assert "unknown whether the run started" in out + assert "could not be started" not in out + assert "scenario-history" in out + # The shell has no --request-timeout option, so it must not advise passing one. + assert "--request-timeout" not in out + + def test_run_results_failure_read_timeout_omits_unsupported_flag(self, shell, capsys): + """The results fetch uses the request timeout, so it needs its own shell-specific hint.""" + import httpx + + s, client = shell + client.start_scenario_run_async = AsyncMock(return_value=self._run_payload()) + client.get_scenario_run_async = AsyncMock(return_value=self._run_payload("COMPLETED")) + client.get_scenario_run_results_async = AsyncMock(side_effect=httpx.ReadTimeout("")) + with patch( + "pyrit.cli._cli_args.parse_run_arguments", + return_value={"scenario_name": "foo", "target": "t"}, + ): + s.do_run("foo --target t") + out = capsys.readouterr().out + assert "ReadTimeout" in out + assert "scenario-results" in out + assert "--request-timeout" not in out def test_run_completed_path_with_results(self, shell, capsys): s, client = shell @@ -780,7 +820,29 @@ def test_print_scenario_error(self, shell, capsys): s, client = shell client.get_scenario_run_results_async = AsyncMock(side_effect=RuntimeError("oops")) s.do_print_scenario("rid-1") - assert "Error: oops" in capsys.readouterr().out + assert "Error (RuntimeError): oops" in capsys.readouterr().out + + def test_scenario_history_read_timeout_is_not_blank(self, shell, capsys): + """The run hints send users here, so a bare ReadTimeout must not print an empty error.""" + import httpx + + s, client = shell + client.list_scenario_runs_async = AsyncMock(side_effect=httpx.ReadTimeout("")) + s.do_scenario_history("") + out = capsys.readouterr().out + assert "ReadTimeout" in out + assert "--request-timeout" not in out + + def test_scenario_results_read_timeout_is_not_blank(self, shell, capsys): + """Same for the results command the completed-run hint points at.""" + import httpx + + s, client = shell + client.get_scenario_run_results_async = AsyncMock(side_effect=httpx.ReadTimeout("")) + s.do_scenario_results("rid-1") + out = capsys.readouterr().out + assert "ReadTimeout" in out + assert "--request-timeout" not in out def test_do_help_with_arg_normalizes_hyphen(self, shell): s, _ = shell @@ -1019,13 +1081,13 @@ def test_conversations_view_reports_fetch_error(self, shell, capsys): client.get_scenario_run_results_async = AsyncMock(return_value=_attacks_scenario_result()) client.get_conversation_messages_async = AsyncMock(side_effect=RuntimeError("nope")) s.do_scenario_results("rid-1 --view conversations") - assert "Error: nope" in capsys.readouterr().out + assert "Error (RuntimeError): nope" in capsys.readouterr().out def test_fetch_error_is_reported(self, shell, capsys): s, client = shell client.get_scenario_run_results_async = AsyncMock(side_effect=RuntimeError("nope")) s.do_scenario_results("rid-1") - assert "Error: nope" in capsys.readouterr().out + assert "Error (RuntimeError): nope" in capsys.readouterr().out def test_print_scenario_alias_warns_and_delegates(self, shell, capsys): s, client = shell