From cb35773ad1d6a90c4428d1b962c98727266f8252 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:11:18 +0530 Subject: [PATCH 1/2] FIX Reject num_templates and num_jailbreaks values below 1 in the jailbreak template sampler TextJailBreak.get_jailbreak_templates only sampled when num_templates was truthy, so num_templates=0 returned every template (650 on main), -1 failed inside random.sample with an unrelated message, True returned one template and 2.5 raised a TypeError. The Jailbreak scenario passed num_jailbreaks through as `num_jailbreaks or default`, so num_jailbreaks=0 silently ran the default sample of 2 and -1 only failed later in initialize_async. get_jailbreak_templates now raises ValueError unless num_templates is None or an integer of at least 1, and Jailbreak.set_params_from_args rejects num_jailbreaks below 1 next to its num_jailbreak_attempts check. Co-Authored-By: Claude Opus 5.5 --- pyrit/datasets/jailbreak/text_jailbreak.py | 8 +++++++- pyrit/scenario/scenarios/airt/jailbreak.py | 6 +++++- tests/unit/datasets/test_jailbreak_text.py | 13 +++++++++++++ tests/unit/scenario/airt/test_jailbreak.py | 7 +++++++ 4 files changed, 32 insertions(+), 2 deletions(-) diff --git a/pyrit/datasets/jailbreak/text_jailbreak.py b/pyrit/datasets/jailbreak/text_jailbreak.py index 5937bc6431..4797114ea2 100644 --- a/pyrit/datasets/jailbreak/text_jailbreak.py +++ b/pyrit/datasets/jailbreak/text_jailbreak.py @@ -217,14 +217,20 @@ def get_jailbreak_templates(cls, num_templates: int | None = None) -> list[str]: list[str]: List of jailbreak template file names. Raises: + ValueError: If ``num_templates`` is not None and not an integer of at least 1. ValueError: If no jailbreak templates are found in the jailbreak directory. ValueError: If n is larger than the number of templates that exist. """ + if num_templates is not None and ( + isinstance(num_templates, bool) or not isinstance(num_templates, int) or num_templates < 1 + ): + raise ValueError(f"num_templates must be an integer of at least 1 or None, got {num_templates!r}") + jailbreak_template_names = sorted(cls._get_template_cache().keys()) if not jailbreak_template_names: raise ValueError("No jailbreak templates found in the jailbreak directory") - if num_templates: + if num_templates is not None: if num_templates > len(jailbreak_template_names): raise ValueError( f"Attempted to pull {num_templates} jailbreaks from a dataset" diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 7c0957982b..7cdc9c1656 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -220,12 +220,16 @@ def set_params_from_args(self, *, args: dict[str, Any]) -> None: args (dict[str, Any]): Raw scenario run parameters. Raises: - ValueError: If ``num_jailbreak_attempts`` is less than one. + ValueError: If ``num_jailbreak_attempts`` is less than one, or if ``num_jailbreaks`` + is provided and is less than one. """ super().set_params_from_args(args=args) num_attempts = self.params["num_jailbreak_attempts"] if num_attempts < 1: raise ValueError("num_jailbreak_attempts must be at least 1") + num_jailbreaks = self.params.get("num_jailbreaks") + if num_jailbreaks is not None and num_jailbreaks < 1: + raise ValueError("num_jailbreaks must be at least 1") @apply_defaults def __init__( diff --git a/tests/unit/datasets/test_jailbreak_text.py b/tests/unit/datasets/test_jailbreak_text.py index cf7ffeebac..7655b20b15 100644 --- a/tests/unit/datasets/test_jailbreak_text.py +++ b/tests/unit/datasets/test_jailbreak_text.py @@ -89,6 +89,19 @@ def test_get_jailbreak_templates_includes_subdirectory_templates(): assert len(templates) > top_level_count, "Subdirectory templates should be included in the listing" +def test_get_jailbreak_templates_returns_requested_count(): + all_templates = TextJailBreak.get_jailbreak_templates() + templates = TextJailBreak.get_jailbreak_templates(num_templates=3) + assert len(templates) == 3 + assert set(templates).issubset(all_templates) + + +@pytest.mark.parametrize("num_templates", [0, -1, True, 2.5, "3"]) +def test_get_jailbreak_templates_rejects_invalid_num_templates(num_templates): + with pytest.raises(ValueError, match="num_templates must be an integer of at least 1 or None"): + TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + + def test_all_templates_render_without_syntax_errors(jailbreak_dir): """Test that all jailbreak templates can be successfully rendered with a test prompt.""" yaml_files = [f for f in jailbreak_dir.rglob("*.yaml") if "multi_parameter" not in f.parts] diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index c6a68bf2c1..a501c0d050 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -208,6 +208,13 @@ def test_rejects_non_positive_num_jailbreak_attempts(self, mock_objective_scorer with pytest.raises(ValueError, match="num_jailbreak_attempts must be at least 1"): scenario.set_params_from_args(args={"num_jailbreak_attempts": num_attempts}) + @pytest.mark.parametrize("num_jailbreaks", [0, -1]) + def test_rejects_non_positive_num_jailbreaks(self, mock_objective_scorer, num_jailbreaks: int) -> None: + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + + with pytest.raises(ValueError, match="num_jailbreaks must be at least 1"): + scenario.set_params_from_args(args={"num_jailbreaks": num_jailbreaks}) + async def test_default_draws_random_template_sample( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): From 9ac2536acbb53cf8e112ce9a2a0cad1b77b2cb4a Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sun, 25 Oct 2026 01:44:29 -0700 Subject: [PATCH 2/2] Simplify jailbreak template count validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/datasets/jailbreak/text_jailbreak.py | 10 ++++------ tests/unit/datasets/test_jailbreak_text.py | 23 ++++++++++++++++------ 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/pyrit/datasets/jailbreak/text_jailbreak.py b/pyrit/datasets/jailbreak/text_jailbreak.py index 4797114ea2..2e988fdfe2 100644 --- a/pyrit/datasets/jailbreak/text_jailbreak.py +++ b/pyrit/datasets/jailbreak/text_jailbreak.py @@ -211,20 +211,18 @@ def get_jailbreak_templates(cls, num_templates: int | None = None) -> list[str]: Retrieve all jailbreaks from the JAILBREAK_TEMPLATES_PATH. Args: - num_templates (int, optional): Number of jailbreak templates to return. None to get all. + num_templates (int | None): Number of jailbreak templates to return. None to get all. Returns: list[str]: List of jailbreak template file names. Raises: - ValueError: If ``num_templates`` is not None and not an integer of at least 1. + ValueError: If ``num_templates`` is provided and is less than 1. ValueError: If no jailbreak templates are found in the jailbreak directory. ValueError: If n is larger than the number of templates that exist. """ - if num_templates is not None and ( - isinstance(num_templates, bool) or not isinstance(num_templates, int) or num_templates < 1 - ): - raise ValueError(f"num_templates must be an integer of at least 1 or None, got {num_templates!r}") + if num_templates is not None and num_templates < 1: + raise ValueError("num_templates must be at least 1") jailbreak_template_names = sorted(cls._get_template_cache().keys()) if not jailbreak_template_names: diff --git a/tests/unit/datasets/test_jailbreak_text.py b/tests/unit/datasets/test_jailbreak_text.py index 7655b20b15..0b78f1bf57 100644 --- a/tests/unit/datasets/test_jailbreak_text.py +++ b/tests/unit/datasets/test_jailbreak_text.py @@ -89,16 +89,27 @@ def test_get_jailbreak_templates_includes_subdirectory_templates(): assert len(templates) > top_level_count, "Subdirectory templates should be included in the listing" -def test_get_jailbreak_templates_returns_requested_count(): +@pytest.mark.parametrize("num_templates", [1, 3]) +def test_get_jailbreak_templates_returns_requested_count(num_templates: int) -> None: all_templates = TextJailBreak.get_jailbreak_templates() - templates = TextJailBreak.get_jailbreak_templates(num_templates=3) - assert len(templates) == 3 + templates = TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + assert len(templates) == num_templates assert set(templates).issubset(all_templates) -@pytest.mark.parametrize("num_templates", [0, -1, True, 2.5, "3"]) -def test_get_jailbreak_templates_rejects_invalid_num_templates(num_templates): - with pytest.raises(ValueError, match="num_templates must be an integer of at least 1 or None"): +@pytest.mark.parametrize("num_templates", [0, -1]) +def test_get_jailbreak_templates_rejects_invalid_num_templates(num_templates: int) -> None: + with pytest.raises(ValueError, match="num_templates must be at least 1"): + TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + + +def test_get_jailbreak_templates_none_returns_all() -> None: + assert TextJailBreak.get_jailbreak_templates(num_templates=None) == sorted(TextJailBreak._get_template_cache()) + + +def test_get_jailbreak_templates_rejects_count_above_population() -> None: + num_templates = len(TextJailBreak.get_jailbreak_templates()) + 1 + with pytest.raises(ValueError, match=f"Attempted to pull {num_templates} jailbreaks"): TextJailBreak.get_jailbreak_templates(num_templates=num_templates)