FIX: Coerce typed scenario parameters in scan CLI (#2092)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: 3 people

pyrit/cli/pyrit_scan.py

@@ -17,16 +17,22 @@
 import sys
 from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, get_args, get_origin
 
 from pyrit.cli._cli_args import (
     ARG_HELP,
     _parse_initializer_arg,
+    build_parameters_from_api,
     non_negative_int,
     positive_int,
     validate_log_level_argparse,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from pyrit.models.parameter import Parameter
+
 _TERMINAL_STATUSES = {"COMPLETED", "FAILED", "CANCELLED"}
 
 
@@ -248,70 +254,94 @@ def _build_base_parser(*, add_help: bool = True) -> ArgumentParser:
 _SCENARIO_DEST_PREFIX = "scenario__"
 
 
-_SCALAR_TYPE_COERCERS: dict[str, Any] = {
-    "int": int,
-    "float": float,
-    "bool": lambda v: str(v).strip().lower() in ("1", "true", "yes", "y", "on"),
-    "str": str,
-}
+def _scenario_value_coercer(*, name: str, annotation: Any) -> Callable[[Any], Any] | None:
+    """
+    Build an argparse ``type=`` callable that coerces a single CLI token through
+    ``Parameter.coerce_value`` — the same coercion the shell and backend use.
 
+    Returns ``None`` when no coercion is needed (a plain ``str`` or an untyped
+    passthrough). Coercion/validation failures (including ``Literal`` choice
+    membership) are re-raised as ``argparse.ArgumentTypeError`` so argparse renders
+    them as a clean CLI error.
+
+    Args:
+        name: Scenario parameter name (used for the flag in error messages).
+        annotation: Scalar element type to coerce to (e.g. ``int``, ``bool``, or
+            ``Literal[...]`` for choices), or ``None`` / ``str`` for passthrough.
 
-def _scenario_param_kwargs(*, param: dict[str, Any]) -> dict[str, Any]:
+    Returns:
+        Callable[[Any], Any] | None: The coercer, or ``None`` for passthrough.
     """
-    Build argparse ``add_argument`` kwargs for a scenario-declared parameter dict.
+    if annotation is None or annotation is str:
+        return None
 
-    Uses ``param_type``, ``is_list`` and ``choices`` from the catalog payload
-    so list params accept ``nargs='+'`` and scalar params get client-side
-    type coercion and choice validation.
+    from pyrit.models.parameter import Parameter
+
+    element_param = Parameter(name=name, description="", param_type=annotation)
+
+    def _coerce(raw: Any) -> Any:
+        try:
+            return element_param.coerce_value(raw)
+        except (ValueError, TypeError) as exc:
+            raise argparse.ArgumentTypeError(f"--{name.replace('_', '-')}: invalid value {raw!r} ({exc})") from exc
+
+    return _coerce
+
+
+def _scenario_param_kwargs(*, parameter: Parameter) -> dict[str, Any]:
+    """
+    Build argparse ``add_argument`` kwargs for a scenario-declared ``Parameter``.
+
+    List params get ``nargs='+'`` and coerce per element; scalar params coerce the
+    single token. All coercion — including ``Literal`` choice membership — routes
+    through ``Parameter.coerce_value`` so scan, the shell, and the backend agree on
+    accepted values.
 
     Args:
-        param: Single entry from ``RegisteredScenario.supported_parameters``.
+        parameter: Scenario parameter built from the catalog payload via
+            ``build_parameters_from_api``.
 
     Returns:
         dict[str, Any]: kwargs ready to pass to ``ArgumentParser.add_argument``.
     """
     kwargs: dict[str, Any] = {
-        "dest": f"{_SCENARIO_DEST_PREFIX}{param.get('name', '')}",
+        "dest": f"{_SCENARIO_DEST_PREFIX}{parameter.name}",
         "default": argparse.SUPPRESS,
-        "help": param.get("description", ""),
+        "help": parameter.description,
     }
-    if param.get("is_list"):
+    param_type = parameter.param_type
+    element_type: Any
+    if get_origin(param_type) is list:
+        type_args = get_args(param_type)
+        element_type = type_args[0] if type_args else str
         kwargs["nargs"] = "+"
     else:
-        coercer = _SCALAR_TYPE_COERCERS.get(param.get("param_type", ""))
-        if coercer is not None and coercer is not str:
-            param_name = param.get("name", "")
-
-            def _typed(raw: str) -> Any:
-                try:
-                    return coercer(raw)
-                except (ValueError, TypeError) as exc:
-                    raise argparse.ArgumentTypeError(
-                        f"--{param_name.replace('_', '-')}: invalid value {raw!r} ({exc})"
-                    ) from exc
-
-            kwargs["type"] = _typed
-    choices = param.get("choices")
-    if choices:
-        kwargs["choices"] = list(choices)
+        element_type = param_type
+
+    coercer = _scenario_value_coercer(name=parameter.name, annotation=element_type)
+    if coercer is not None:
+        kwargs["type"] = coercer
     return kwargs
 
 
 def _add_scenario_params_from_api(*, parser: ArgumentParser, params: list[dict[str, Any]]) -> None:
     """
     Add scenario-declared parameters (from the API response) as CLI flags.
 
+    Catalog payloads are converted to ``Parameter`` objects via
+    ``build_parameters_from_api`` (shared with the shell) so type coercion and
+    choice handling stay consistent across entry points.
+
     Args:
         parser: Parser to extend.
         params: List of parameter dicts from ``GET /api/scenarios/catalog/{name}``.
     """
     seen_flags: set[str] = set(parser._option_string_actions.keys())
-    for p in params:
-        name = p.get("name", "")
-        flag = f"--{name.replace('_', '-')}"
+    for parameter in build_parameters_from_api(api_params=params) or []:
+        flag = f"--{parameter.name.replace('_', '-')}"
         if flag in seen_flags:
             continue
-        parser.add_argument(flag, **_scenario_param_kwargs(param=p))
+        parser.add_argument(flag, **_scenario_param_kwargs(parameter=parameter))
         seen_flags.add(flag)
 
 

tests/unit/cli/test_pyrit_scan.py

@@ -634,6 +634,56 @@ def test_int_param_invalid_value_rejected_client_side(self, capsys):
             parser.parse_args(["--max-turns", "not-an-int"])
         assert "invalid value" in capsys.readouterr().err
 
+    def test_bool_param_rejects_invalid_value_client_side(self, capsys):
+        from argparse import ArgumentParser
+
+        parser = ArgumentParser()
+        pyrit_scan._add_scenario_params_from_api(
+            parser=parser,
+            params=[{"name": "dry_run", "description": "...", "param_type": "bool"}],
+        )
+
+        parsed = parser.parse_args(["--dry-run", "false"])
+        assert parsed.scenario__dry_run is False
+
+        parsed = parser.parse_args(["--dry-run", "yes"])
+        assert parsed.scenario__dry_run is True
+
+        with pytest.raises(SystemExit):
+            parser.parse_args(["--dry-run", "maybe"])
+        assert "invalid value" in capsys.readouterr().err
+
+        # "on"/"y" are NOT part of the canonical boolean vocabulary the shell and
+        # backend accept (true/false, 1/0, yes/no); scan must reject them too so
+        # the same flag never behaves differently depending on the entry point.
+        with pytest.raises(SystemExit):
+            parser.parse_args(["--dry-run", "on"])
+        assert "invalid value" in capsys.readouterr().err
+
+    def test_list_int_param_coerces_each_value(self):
+        from argparse import ArgumentParser
+
+        parser = ArgumentParser()
+        pyrit_scan._add_scenario_params_from_api(
+            parser=parser,
+            params=[{"name": "sample_ids", "description": "...", "param_type": "list[int]", "is_list": True}],
+        )
+
+        parsed = parser.parse_args(["--sample-ids", "1", "2", "3"])
+        assert parsed.scenario__sample_ids == [1, 2, 3]
+
+    def test_typed_choices_are_compared_after_coercion(self):
+        from argparse import ArgumentParser
+
+        parser = ArgumentParser()
+        pyrit_scan._add_scenario_params_from_api(
+            parser=parser,
+            params=[{"name": "max_turns", "description": "...", "param_type": "int", "choices": ["1", "2"]}],
+        )
+
+        parsed = parser.parse_args(["--max-turns", "1"])
+        assert parsed.scenario__max_turns == 1
+
     def test_choices_validated_client_side(self, capsys):
         from argparse import ArgumentParser
 
@@ -647,7 +697,7 @@ def test_choices_validated_client_side(self, capsys):
 
         with pytest.raises(SystemExit):
             parser.parse_args(["--mode", "warp"])
-        assert "invalid choice" in capsys.readouterr().err
+        assert "invalid value" in capsys.readouterr().err
 
 
 class TestMainExtraPaths:
@@ -810,6 +860,27 @@ def test_scenario_declared_flag_is_forwarded(self, _mock_prog, _mock_print, mock
         sent_request = client.start_scenario_run_async.call_args.kwargs["request"]
         assert sent_request["scenario_params"] == {"max_turns": "7"}
 
+    @patch("pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, return_value=True)
+    @patch("pyrit.cli.api_client.PyRITApiClient")
+    @patch("pyrit.cli._output.print_scenario_result_async", new_callable=AsyncMock)
+    @patch("pyrit.cli._output.print_scenario_run_progress")
+    def test_typed_scenario_flags_are_forwarded_as_typed_values(
+        self, _mock_prog, _mock_print, mock_client_class, _mock_probe
+    ):
+        client = self._build_mock_client(
+            supported_params=[
+                {"name": "dry_run", "description": "...", "param_type": "bool"},
+                {"name": "sample_ids", "description": "...", "param_type": "list[int]", "is_list": True},
+            ]
+        )
+        mock_client_class.return_value = client
+
+        result = pyrit_scan.main(["foo", "--target", "t", "--dry-run", "yes", "--sample-ids", "1", "2"])
+
+        assert result == 0
+        sent_request = client.start_scenario_run_async.call_args.kwargs["request"]
+        assert sent_request["scenario_params"] == {"dry_run": True, "sample_ids": [1, 2]}
+
     @patch("pyrit.cli._server_launcher.ServerLauncher.probe_health_async", new_callable=AsyncMock, return_value=True)
     @patch("pyrit.cli.api_client.PyRITApiClient")
     @patch("pyrit.cli._output.print_scenario_result_async", new_callable=AsyncMock)