From f3edc80e8015254b85966cd0305d574a3615d02f Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:37:52 +0530 Subject: [PATCH] FIX Reject an example_count that is not a positive integer in ManyShotJailbreakAttack example_count is used as a slice bound, so a negative value dropped examples from the end of the list instead of selecting from the start, and bools, floats and strings were accepted or failed with an unrelated error. The check lives in validate_constructor_parameters, so the REST technique API rejects the same values without constructing the attack. Co-Authored-By: Claude Opus 5.5 --- .../attack/single_turn/many_shot_jailbreak.py | 12 +++++++---- tests/unit/backend/test_technique_service.py | 1 + .../single_turn/test_many_shot_jailbreak.py | 21 +++++++++++++++++++ 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/pyrit/executor/attack/single_turn/many_shot_jailbreak.py b/pyrit/executor/attack/single_turn/many_shot_jailbreak.py index 91ffcb6c44..d9b9d29f8d 100644 --- a/pyrit/executor/attack/single_turn/many_shot_jailbreak.py +++ b/pyrit/executor/attack/single_turn/many_shot_jailbreak.py @@ -65,12 +65,12 @@ def __init__( prompt_normalizer (PromptNormalizer, Optional): Normalizer for handling prompts. max_attempts_on_failure (int, Optional): Maximum number of attempts to retry on failure. Defaults to 0. example_count (int): The number of examples to include from many_shot_examples or the Many - Shot Jailbreaking dataset. Defaults to the first 100. + Shot Jailbreaking dataset. Must be a positive integer. Defaults to the first 100. many_shot_examples (list[dict[str, str]], Optional): The many shot jailbreaking examples to use. If not provided, takes the first `example_count` examples from Many Shot Jailbreaking dataset. Raises: - ValueError: If many_shot_examples is empty. + ValueError: If example_count is not a positive integer or many_shot_examples is empty. """ ManyShotJailbreakAttack.validate_constructor_parameters( { @@ -106,12 +106,16 @@ def validate_constructor_parameters(cls, params: dict[str, Any]) -> None: Validate the supplied example selection without loading the default dataset. Raises: - ValueError: If the supplied example selection is empty. + ValueError: If ``example_count`` is not a positive integer or the supplied example + selection is empty. """ super().validate_constructor_parameters(params) examples = params.get("many_shot_examples") count = params.get("example_count", 100) - if count == 0 or (examples is not None and not examples[:count]): + # The count is a slice bound, so a negative value would drop examples from the end. + if isinstance(count, bool) or not isinstance(count, int) or count < 1: + raise ValueError(f"example_count must be a positive integer, got {count!r}.") + if examples is not None and not examples: raise ValueError("Many shot examples must be provided.") async def _perform_async(self, *, context: SingleTurnAttackContext[Any]) -> AttackResult: diff --git a/tests/unit/backend/test_technique_service.py b/tests/unit/backend/test_technique_service.py index e9c6582fe4..ce3cf9c570 100644 --- a/tests/unit/backend/test_technique_service.py +++ b/tests/unit/backend/test_technique_service.py @@ -496,6 +496,7 @@ async def register_then_initialize_async(*, request: RunScenarioRequest, init_kw ("PromptSendingAttack", {"max_attempts_on_failure": -1}), ("ManyShotJailbreakAttack", {"max_attempts_on_failure": -1}), ("ManyShotJailbreakAttack", {"example_count": 0}), + ("ManyShotJailbreakAttack", {"example_count": -1}), ("ManyShotJailbreakAttack", {"many_shot_examples": []}), ("CrescendoAttack", {"max_backtracks": -1}), ("CrescendoAttack", {"max_turns": 0}), diff --git a/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py b/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py index f56fa46833..eee0962c22 100644 --- a/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py +++ b/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py @@ -188,6 +188,27 @@ def test_init_raises_error_with_empty_custom_examples(self, mock_from_yaml, mock with pytest.raises(ValueError, match="Many shot examples must be provided"): ManyShotJailbreakAttack(objective_target=mock_objective_target, many_shot_examples=[]) + @pytest.mark.parametrize("example_count", [0, -1, -3, True, 2.5, "5", None]) + @patch("pyrit.executor.attack.single_turn.many_shot_jailbreak.SeedPrompt.from_yaml_file") + def test_init_rejects_example_count_that_is_not_a_positive_integer( + self, mock_from_yaml, mock_objective_target, mock_template, sample_many_shot_examples, example_count + ): + """A negative count would slice examples off the end instead of selecting from the start.""" + mock_from_yaml.return_value = mock_template + + with pytest.raises(ValueError, match="example_count must be a positive integer"): + ManyShotJailbreakAttack( + objective_target=mock_objective_target, + example_count=example_count, + many_shot_examples=sample_many_shot_examples, + ) + + @pytest.mark.parametrize("example_count", [0, -1, True, "5"]) + def test_validate_constructor_parameters_rejects_invalid_example_count(self, example_count): + """Deferred construction (e.g. the REST technique API) must reject the same counts.""" + with pytest.raises(ValueError, match="example_count must be a positive integer"): + ManyShotJailbreakAttack.validate_constructor_parameters({"example_count": example_count}) + @pytest.mark.usefixtures("patch_central_database") class TestManyShotJailbreakAttackParamsType: