diff --git a/pyrit/executor/attack/multi_turn/simulated_conversation.py b/pyrit/executor/attack/multi_turn/simulated_conversation.py index def7f590ee..9a5ed222bf 100644 --- a/pyrit/executor/attack/multi_turn/simulated_conversation.py +++ b/pyrit/executor/attack/multi_turn/simulated_conversation.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import uuid from typing import TYPE_CHECKING from pyrit.executor.attack.core.attack_config import ( @@ -211,17 +212,23 @@ async def _generate_next_message_async( conversation_context=conversation_context, ) - # Use the adversarial chat to generate the next message - # Create a simple user message asking for generation + # Use the adversarial chat to generate the next message. + # Scope the system prompt and the generated message to their own fresh + # conversation. Message.from_prompt leaves conversation_id unset (None); + # passing that to set_system_prompt would make get_conversation_messages + # skip its conversation filter and return every piece in memory, which + # raises once memory holds more than one conversation. + conversation_id = str(uuid.uuid4()) request_message = Message.from_prompt( role="user", prompt="Generate the next user message based on the instructions above.", ) + request_message.message_pieces[0].conversation_id = conversation_id # Set the system prompt on the target adversarial_chat.set_system_prompt( system_prompt=system_prompt, - conversation_id=request_message.conversation_id, + conversation_id=conversation_id, ) responses: list[Message] = await adversarial_chat.send_prompt_async(message=request_message) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index a576dda830..9e2bf380af 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -958,7 +958,14 @@ def get_conversation_messages(self, *, conversation_id: str) -> MutableSequence[ Returns: MutableSequence[Message]: A list of chat memory entries with the specified conversation ID. + + Raises: + ValueError: If conversation_id is empty or None. A falsy id would cause the underlying + get_message_pieces filter to be skipped, silently returning pieces from every + conversation in memory. """ + if not conversation_id: + raise ValueError("get_conversation_messages requires a non-empty conversation_id") message_pieces = self.get_message_pieces(conversation_id=conversation_id) return group_conversation_message_pieces_by_sequence(message_pieces=message_pieces) diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 5a822fc3fb..06a8c1c522 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -622,7 +622,7 @@ async def _build_groups_by_dataset_async(self) -> tuple[dict[str, list[SeedAttac ) return groups_by_dataset, resolved - async def get_seed_attack_groups_async(self) -> list[SeedAttackGroup]: + async def get_seed_attack_groups_async(self, *, apply_sampling: bool = True) -> list[SeedAttackGroup]: """ Resolve the configured dataset into a flat ``list[SeedAttackGroup]``. @@ -630,8 +630,15 @@ async def get_seed_attack_groups_async(self) -> list[SeedAttackGroup]: validates the full resolved seed set, then samples ``max_dataset_size`` globally over all built groups. + Args: + apply_sampling (bool): When True (default), apply ``max_dataset_size`` sampling. + Pass False to resolve the full, deterministic dataset with no ``random.sample`` + draw -- used on resume so the persisted objective subset can be reconstructed + exactly rather than intersected against a fresh (divergent) sample. + Returns: - list[SeedAttackGroup]: The validated, sampled attack groups. + list[SeedAttackGroup]: The validated attack groups (sampled when ``apply_sampling`` + is True, otherwise the full resolved set). Raises: DatasetConstraintError: If a configured dataset yields no seeds, the resolved @@ -640,13 +647,16 @@ async def get_seed_attack_groups_async(self) -> list[SeedAttackGroup]: groups_by_dataset, resolved = await self._build_groups_by_dataset_async() self.validate(resolved) groups = [group for groups in groups_by_dataset.values() for group in groups] - groups = self._apply_max_dataset_size(groups) + if apply_sampling: + groups = self._apply_max_dataset_size(groups) if not groups: names = ", ".join(self._dataset_names) if self._dataset_names else "" raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).") return groups - async def get_attack_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + async def get_attack_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[SeedAttackGroup]]: """ Resolve attack groups keyed by dataset name, globally sampled. @@ -656,8 +666,15 @@ async def get_attack_groups_by_dataset_async(self) -> dict[str, list[SeedAttackG keyed by their originating dataset. For an independent budget per dataset, compose ``CompoundDatasetAttackConfiguration.per_dataset(...)`` instead. + Args: + apply_sampling (bool): When True (default), apply ``max_dataset_size`` sampling. + Pass False to resolve the full, deterministic dataset with no ``random.sample`` + draw -- used on resume so the persisted objective subset can be reconstructed + exactly rather than intersected against a fresh (divergent) sample. + Returns: - dict[str, list[SeedAttackGroup]]: Dataset name -> sampled attack groups. + dict[str, list[SeedAttackGroup]]: Dataset name -> attack groups (sampled when + ``apply_sampling`` is True, otherwise the full resolved set). Raises: DatasetConstraintError: If a configured dataset yields no seeds, the resolved @@ -665,7 +682,8 @@ async def get_attack_groups_by_dataset_async(self) -> dict[str, list[SeedAttackG """ groups_by_dataset, resolved = await self._build_groups_by_dataset_async() self.validate(resolved) - result = {name: groups for name, groups in self._sample_groups_by_dataset(groups_by_dataset).items() if groups} + sampled = self._sample_groups_by_dataset(groups_by_dataset) if apply_sampling else groups_by_dataset + result = {name: groups for name, groups in sampled.items() if groups} if not result: names = ", ".join(self._dataset_names) if self._dataset_names else "" raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).") @@ -819,29 +837,42 @@ def update_filters(self, *, filters: dict[str, list[str]]) -> None: for child in self._configurations: child.update_filters(filters=filters) - async def get_seed_attack_groups_async(self) -> list[SeedAttackGroup]: + async def get_seed_attack_groups_async(self, *, apply_sampling: bool = True) -> list[SeedAttackGroup]: """ Concatenate every child's flat result, then validate and apply the global cap. Each child validates and samples itself; the combined result is validated against this compound's validators and capped by an optional compound ``max_dataset_size``. + Args: + apply_sampling (bool): When True (default), sample both each child and the combined + result under ``max_dataset_size``. Pass False to resolve the full, deterministic + dataset with no sampling at any level -- used on resume (propagated to children). + Returns: - list[SeedAttackGroup]: The combined, validated, capped attack groups. + list[SeedAttackGroup]: The combined, validated attack groups (capped when + ``apply_sampling`` is True). Raises: DatasetConstraintError: If a child yields nothing, or the combined result fails validation. """ groups: list[SeedAttackGroup] = [] for child in self._configurations: - groups.extend(await child.get_seed_attack_groups_async()) + groups.extend(await child.get_seed_attack_groups_async(apply_sampling=apply_sampling)) self.validate(self._resolved_from_groups(groups)) - return self._apply_max_dataset_size(groups) + return self._apply_max_dataset_size(groups) if apply_sampling else groups - async def get_attack_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + async def get_attack_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[SeedAttackGroup]]: """ Merge each child's by-dataset result, validate, then apply the global cap across the union. + Args: + apply_sampling (bool): When True (default), sample both each child and the merged + union under ``max_dataset_size``. Pass False to resolve the full, deterministic + dataset with no sampling at any level -- used on resume (propagated to children). + Returns: dict[str, list[SeedAttackGroup]]: Combined groups keyed by dataset name. @@ -850,10 +881,11 @@ async def get_attack_groups_by_dataset_async(self) -> dict[str, list[SeedAttackG """ merged: dict[str, list[SeedAttackGroup]] = {} for child in self._configurations: - for name, groups in (await child.get_attack_groups_by_dataset_async()).items(): + child_groups = await child.get_attack_groups_by_dataset_async(apply_sampling=apply_sampling) + for name, groups in child_groups.items(): merged.setdefault(name, []).extend(groups) self.validate(self._resolved_from_groups([group for groups in merged.values() for group in groups])) - return self._sample_groups_by_dataset(merged) + return self._sample_groups_by_dataset(merged) if apply_sampling else merged def _resolved_from_groups(self, groups: list[SeedAttackGroup]) -> ResolvedDataset: """ diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 775eff4662..a0aa73140c 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -609,7 +609,13 @@ async def initialize_async(self) -> None: # into a ScenarioContext, and hand it to the subclass extension point. Baseline is # emitted centrally (from context.seed_groups) so overrides never re-resolve seeds # or hand-roll baseline emission. - seed_groups_by_dataset = await self._resolve_seed_groups_by_dataset_async() + # + # On resume, resolve the full, deterministic dataset (no max_dataset_size sampling): + # the originally-sampled subset was snapshotted into the ScenarioResult metadata and is + # replayed by _apply_persisted_objectives. Re-drawing a fresh random.sample here would + # diverge from the persisted hashes and abort resume whenever max_dataset_size is set. + is_resume = self._scenario_result_id is not None + seed_groups_by_dataset = await self._resolve_seed_groups_by_dataset_async(apply_sampling=not is_resume) context = self._build_scenario_context(seed_groups_by_dataset=seed_groups_by_dataset) self._atomic_attacks = await self._build_atomic_attacks_async(context=context) @@ -696,12 +702,12 @@ def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: On resume, replay the originally-sampled objective subset. When the first run used ``max_dataset_size``, the chosen subset was - recorded in ``ScenarioResult.metadata["objective_hashes"]``. - Restrict each atomic attack's freshly-resolved seed_groups to that set - so a fresh ``random.sample`` draw on resume can't silently shift which - objectives the scenario operates on. If any persisted hash is no longer - present in the dataset, refuse to resume — running a smaller subset - than the user committed to would silently produce different results. + recorded in ``ScenarioResult.metadata["objective_hashes"]``. Resume resolves + the **full, deterministic** dataset (sampling is bypassed on the resume branch of + ``initialize_async``), so restricting each atomic attack's seed_groups to the + persisted set here reconstructs exactly the objectives the first run committed to. + If any persisted hash is no longer present in the dataset, refuse to resume — that + now signals the dataset itself genuinely changed, not a random resample drift. Args: stored_result (ScenarioResult): The scenario result loaded from memory. @@ -926,7 +932,9 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: return remaining_attacks - async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[SeedAttackGroup]]: """ Resolve the seed groups this scenario attacks, keyed by originating dataset. @@ -938,10 +946,17 @@ async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAtta Override to inject seeds from an alternate source (e.g. deprecated ``objectives``) or to filter the resolved groups before attacks are built. + Args: + apply_sampling (bool): When True (default), apply ``max_dataset_size`` sampling. + On resume the base passes False so the full, deterministic dataset is resolved + and the persisted objective subset is reconstructed exactly (see + ``_apply_persisted_objectives``) rather than intersected against a fresh, + divergent ``random.sample`` draw. + Returns: dict[str, list[SeedAttackGroup]]: Seed groups keyed by dataset name. """ - return await self._dataset_config.get_attack_groups_by_dataset_async() + return await self._dataset_config.get_attack_groups_by_dataset_async(apply_sampling=apply_sampling) def _build_scenario_context(self, *, seed_groups_by_dataset: dict[str, list[SeedAttackGroup]]) -> ScenarioContext: """ diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 829e9ebb05..1e086232ef 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -235,7 +235,9 @@ def __init__( # Store deprecated objectives for later resolution in _resolve_seed_groups_by_dataset_async self._deprecated_objectives = objectives - async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[SeedAttackGroup]]: """ Resolve seed groups from deprecated objectives or dataset configuration. @@ -244,6 +246,12 @@ async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAtta category. The base ``Scenario`` flattens the result into ``context.seed_groups`` and reuses it for the strategy attacks and the baseline. + Args: + apply_sampling (bool): When True (default), apply ``max_dataset_size`` sampling. + On resume the base passes False so the full, deterministic dataset is resolved + and the persisted objective subset is reconstructed exactly. Inline deprecated + objectives are never sampled. + Returns: dict[str, list[SeedAttackGroup]]: Seed groups keyed by dataset (or a synthetic key for deprecated inline objectives). @@ -267,7 +275,7 @@ async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAtta harm_category_filter = self._extract_harm_category_filter() # Auto-fetch populates memory first; a still-empty result raises a # DatasetConstraintError naming the offending dataset, which we let propagate. - seed_groups = list(await self._dataset_config.get_seed_attack_groups_async()) + seed_groups = list(await self._dataset_config.get_seed_attack_groups_async(apply_sampling=apply_sampling)) if harm_category_filter: seed_groups = self._filter_by_harm_category( diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 1e387039c1..8ac7c98d44 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -498,7 +498,9 @@ def _scoring_config_for_strategy(self, strategy: WebInjectionStrategy) -> Attack return self._xss_scoring_config return self._exfil_scoring_config - async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[SeedAttackGroup]]: """ Generate the injection prompts and wrap them into seed groups, keyed by strategy. @@ -507,6 +509,11 @@ async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAtta set from the raw garak datasets. Resolving them here means the base owns the single seed sample used for both the atomic attacks and the central baseline. + Args: + apply_sampling (bool): Accepted for base-class compatibility but unused — the + synthesized seeds are already deterministic (``random.Random(self._random_seed)``), + so resume reproduces the same set without a ``max_dataset_size`` sampling path. + Returns: dict[str, list[SeedAttackGroup]]: Seed groups keyed by strategy value. diff --git a/tests/unit/executor/attack/component/test_simulated_conversation.py b/tests/unit/executor/attack/component/test_simulated_conversation.py index 66f91423d0..7b13ab2722 100644 --- a/tests/unit/executor/attack/component/test_simulated_conversation.py +++ b/tests/unit/executor/attack/component/test_simulated_conversation.py @@ -690,6 +690,75 @@ async def test_next_message_system_prompt_path_sets_system_prompt( # Verify set_system_prompt was called on adversarial_chat mock_adversarial_chat.set_system_prompt.assert_called() + async def test_next_message_scopes_system_prompt_to_generated_message_conversation( + self, + mock_adversarial_chat: MagicMock, + mock_objective_scorer: MagicMock, + adversarial_system_prompt_path: Path, + sample_conversation: list[Message], + ): + """Regression: the next-message system prompt must be scoped to a concrete conversation id. + + ``Message.from_prompt`` leaves ``conversation_id`` unset (None). Passing that to + ``set_system_prompt`` makes ``get_conversation_messages`` skip its conversation filter and + return every piece in memory, which raises once memory holds more than one conversation. + The generated request message and the system prompt must share the same non-empty id. + """ + from pyrit.models.seeds import NextMessageSystemPromptPaths + + next_message_response = Message( + message_pieces=[ + MessagePiece( + role="assistant", + original_value="Generated message", + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), + ) + ] + ) + + with patch("pyrit.executor.attack.multi_turn.simulated_conversation.RedTeamingAttack") as mock_attack_class: + mock_attack = MagicMock() + mock_attack.get_identifier.return_value = ComponentIdentifier( + class_name="RedTeamingAttack", class_module="pyrit.executor.attack" + ) + mock_attack.execute_async = AsyncMock( + return_value=AttackResult( + atomic_attack_identifier=ComponentIdentifier( + class_name="RedTeamingAttack", class_module="pyrit.executor.attack" + ), + conversation_id=str(uuid.uuid4()), + objective="Test objective", + outcome=AttackOutcome.SUCCESS, + executed_turns=3, + ) + ) + mock_attack_class.return_value = mock_attack + + with patch("pyrit.executor.attack.multi_turn.simulated_conversation.CentralMemory") as mock_memory_class: + mock_memory = MagicMock() + mock_memory.get_conversation_messages.return_value = iter(sample_conversation) + mock_memory_class.get_memory_instance.return_value = mock_memory + + mock_adversarial_chat.send_prompt_async = AsyncMock(return_value=[next_message_response]) + + await generate_simulated_conversation_async( + objective="Test objective", + adversarial_chat=mock_adversarial_chat, + objective_scorer=mock_objective_scorer, + adversarial_chat_system_prompt_path=adversarial_system_prompt_path, + num_turns=3, + next_message_system_prompt_path=NextMessageSystemPromptPaths.DIRECT.value, + ) + + system_prompt_conversation_id = mock_adversarial_chat.set_system_prompt.call_args.kwargs[ + "conversation_id" + ] + assert system_prompt_conversation_id + + sent_message = mock_adversarial_chat.send_prompt_async.call_args.kwargs["message"] + assert sent_message.conversation_id == system_prompt_conversation_id + async def test_starting_sequence_sets_first_sequence_number( self, mock_adversarial_chat: MagicMock, diff --git a/tests/unit/executor/attack/test_attack_parameter_consistency.py b/tests/unit/executor/attack/test_attack_parameter_consistency.py index 6c27e764ad..6cb7ff47ff 100644 --- a/tests/unit/executor/attack/test_attack_parameter_consistency.py +++ b/tests/unit/executor/attack/test_attack_parameter_consistency.py @@ -758,8 +758,13 @@ async def test_tap_attack_adds_prepended_to_memory( next_message=multimodal_text_message, # Required when prepended_conversation is provided ) + # TAP prunes all branches with these mocks, so result.conversation_id is empty. The prepended + # messages were duplicated into the single node conversation; resolve that id from memory. + assert not result.conversation_id memory = CentralMemory.get_memory_instance() - conversation = list(memory.get_conversation_messages(conversation_id=result.conversation_id)) + node_conversation_ids = {piece.conversation_id for piece in memory.get_message_pieces()} + assert len(node_conversation_ids) == 1, f"Expected one conversation in memory, got {node_conversation_ids}" + conversation = list(memory.get_conversation_messages(conversation_id=node_conversation_ids.pop())) # Should have exactly the prepended messages in memory (mock normalizer doesn't add responses) assert len(conversation) == 2, f"Expected exactly 2 prepended messages, got {len(conversation)}" @@ -1026,16 +1031,19 @@ async def test_crescendo_injects_prepended_into_adversarial_context( adversarial_chat_mock=mock_adversarial_chat, ) - async def test_tap_injects_prepended_into_adversarial_context( + async def test_tap_persists_prepended_conversation_in_memory( self, tap_attack: TreeOfAttacksWithPruningAttack, - mock_adversarial_chat: MagicMock, prepended_conversation_text: list[Message], multimodal_text_message: Message, sqlite_instance, ) -> None: - """Test that TreeOfAttacksWithPruningAttack injects prepended conversation into adversarial context.""" - # TAP may fail due to JSON parsing, but set_system_prompt should be called before the error + """TAP persists the prepended conversation into the node conversation in memory. + + With these mocks TAP prunes every branch before the adversarial chat's system prompt is + set, so the prepended text is only observable in the node conversation written to memory + (not in the adversarial context). Verify the prepended text is preserved there. + """ with suppress(Exception): await tap_attack.execute_async( objective="Test objective", @@ -1043,9 +1051,21 @@ async def test_tap_injects_prepended_into_adversarial_context( next_message=multimodal_text_message, ) - # Verify prepended text appears in adversarial context (checks mock's set_system_prompt calls) - _assert_prepended_text_in_adversarial_context( - prepended_conversation=prepended_conversation_text, - adversarial_chat_conversation_id="", # Empty - will fall back to mock check - adversarial_chat_mock=mock_adversarial_chat, + memory = CentralMemory.get_memory_instance() + node_conversation_ids = {piece.conversation_id for piece in memory.get_message_pieces()} + assert len(node_conversation_ids) == 1, f"Expected one conversation in memory, got {node_conversation_ids}" + conversation = list(memory.get_conversation_messages(conversation_id=node_conversation_ids.pop())) + + node_text = " ".join( + piece.original_value + for msg in conversation + for piece in msg.message_pieces + if piece.original_value_data_type == "text" ) + for msg in prepended_conversation_text: + for piece in msg.message_pieces: + if piece.original_value_data_type == "text": + assert piece.original_value in node_text, ( + f"Prepended text '{piece.original_value}' not found in node conversation. " + f"Available text: {node_text}" + ) diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index f70b6094ff..b3860519e3 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -1324,6 +1324,31 @@ def test_get_request_from_response_success(sqlite_instance: MemoryInterface): assert request.conversation_id == conversation_id +@pytest.mark.parametrize("bad_conversation_id", ["", None]) +def test_get_conversation_messages_rejects_falsy_conversation_id(sqlite_instance: MemoryInterface, bad_conversation_id): + """A falsy conversation_id must raise instead of skipping the filter and returning every conversation.""" + pieces = [ + MessagePiece( + role="user", + original_value="conversation one", + converted_value="conversation one", + conversation_id=str(uuid4()), + sequence=0, + ), + MessagePiece( + role="user", + original_value="conversation two", + converted_value="conversation two", + conversation_id=str(uuid4()), + sequence=0, + ), + ] + sqlite_instance.add_message_pieces_to_memory(message_pieces=pieces) + + with pytest.raises(ValueError, match="requires a non-empty conversation_id"): + sqlite_instance.get_conversation_messages(conversation_id=bad_conversation_id) + + def test_get_request_from_response_multi_turn_conversation(sqlite_instance: MemoryInterface): """Test get_request_from_response in a multi-turn conversation.""" conversation_id = str(uuid4()) diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py index b8a19568db..fc31b28346 100644 --- a/tests/unit/scenario/core/test_scenario.py +++ b/tests/unit/scenario/core/test_scenario.py @@ -169,7 +169,7 @@ def get_aggregate_tags(cls) -> set[str]: super().__init__(**kwargs) self._atomic_attacks_to_return = atomic_attacks_to_return or [] - async def _resolve_seed_groups_by_dataset_async(self): + async def _resolve_seed_groups_by_dataset_async(self, *, apply_sampling: bool = True): return {} async def _build_atomic_attacks_async(self, *, context): @@ -782,8 +782,8 @@ def get_aggregate_tags(cls) -> set[str]: super().__init__(**kwargs) self._atomic_attacks_to_return = atomic_attacks_to_return or [] - async def _resolve_seed_groups_by_dataset_async(self): - return await self._dataset_config.get_attack_groups_by_dataset_async() + async def _resolve_seed_groups_by_dataset_async(self, *, apply_sampling: bool = True): + return await self._dataset_config.get_attack_groups_by_dataset_async(apply_sampling=apply_sampling) async def _build_atomic_attacks_async(self, *, context): return list(self._atomic_attacks_to_return) @@ -1047,6 +1047,125 @@ def _sample_first_k(population, k): assert len(baseline.objectives) == 3 +@pytest.mark.usefixtures("patch_central_database") +class TestScenarioResumeDeterministicUnderMaxDatasetSize: + """Phase H regression: resume must reconstruct the persisted objective subset. + + ``max_dataset_size`` applies an unseeded ``random.sample`` on every seed + resolution. Before the fix, resume re-sampled and intersected the persisted + objective hashes against a *fresh* (divergent) draw, so resume aborted with + "persisted objective hash(es) are no longer present in the dataset" whenever + ``max_dataset_size`` was smaller than the dataset. The fix bypasses sampling on + the resume branch: the full deterministic dataset is resolved and the persisted + hashes drive selection, reconstructing exactly the first run's objectives. + """ + + class _StrategyScenario(ConcreteScenarioWithTrueFalseScorer): + async def _build_atomic_attacks_async(self, *, context): + from pyrit.scenario.core.attack_technique import AttackTechnique + + return [ + AtomicAttack( + atomic_attack_name="strategy", + attack_technique=AttackTechnique(attack=MagicMock()), + seed_groups=list(context.seed_groups), + ) + ] + + def _make_config(self): + from pyrit.models import SeedGroup, SeedObjective + + seed_groups = [SeedGroup(seeds=[SeedObjective(value=f"obj{i}")]) for i in range(10)] + return DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=3) + + async def test_resume_reconstructs_persisted_subset_without_resampling(self, mock_objective_target): + config = self._make_config() + + def _sample_first_k(population, k): + return list(population)[:k] + + # First run: deterministic "first 3" sample persists obj0/obj1/obj2. + with patch( + "pyrit.scenario.core.dataset_configuration.random.sample", + side_effect=_sample_first_k, + ): + scenario = self._StrategyScenario(name="Phase H resume", version=1) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": config, + } + ) + await scenario.initialize_async() + + original_id = scenario._scenario_result_id + assert original_id is not None + _, first_strategy = scenario._atomic_attacks + persisted_objectives = set(first_strategy.objectives) + assert persisted_objectives == {"obj0", "obj1", "obj2"} + + # Resume: a *divergent* sample (last 3) would have broken the pre-fix intersection. + # With the fix, resume never samples, so this side_effect must go uncalled. + def _sample_last_k(population, k): + return list(population)[-k:] + + with patch( + "pyrit.scenario.core.dataset_configuration.random.sample", + side_effect=_sample_last_k, + ) as resume_sample_mock: + resumed = self._StrategyScenario( + name="Phase H resume", + version=1, + scenario_result_id=original_id, + ) + resumed.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": self._make_config(), + } + ) + # Must not raise "persisted objective hash(es) are no longer present in the dataset". + await resumed.initialize_async() + + # Sampling is bypassed on resume — the full dataset is resolved deterministically. + assert resume_sample_mock.call_count == 0 + assert resumed._scenario_result_id == original_id + + baseline, strategy = resumed._atomic_attacks + assert baseline.atomic_attack_name == "baseline" + assert strategy.atomic_attack_name == "strategy" + # Exactly the originally-persisted subset, not the divergent "last 3" draw. + assert set(strategy.objectives) == persisted_objectives + assert set(baseline.objectives) == persisted_objectives + + async def test_fresh_run_still_samples(self, mock_objective_target): + """The resume bypass must not disable sampling for a normal (non-resume) run.""" + config = self._make_config() + + def _sample_first_k(population, k): + return list(population)[:k] + + with patch( + "pyrit.scenario.core.dataset_configuration.random.sample", + side_effect=_sample_first_k, + ) as sample_mock: + scenario = self._StrategyScenario(name="Phase H fresh", version=1) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": config, + } + ) + await scenario.initialize_async() + + assert sample_mock.call_count == 1 + _, strategy = scenario._atomic_attacks + assert len(strategy.objectives) == 3 + + @pytest.mark.usefixtures("patch_central_database") class TestBuildBaselineAtomicAttack: """Unit tests for Scenario._build_baseline_atomic_attack.""" diff --git a/tests/unit/scenario/core/test_scenario_parameters.py b/tests/unit/scenario/core/test_scenario_parameters.py index 5cd953953c..0cbfd43e1d 100644 --- a/tests/unit/scenario/core/test_scenario_parameters.py +++ b/tests/unit/scenario/core/test_scenario_parameters.py @@ -52,7 +52,7 @@ def supported_parameters(cls) -> list[Parameter]: base = [p for p in base if p.name not in remove_common] return base + list(params_to_declare) - async def _resolve_seed_groups_by_dataset_async(self): + async def _resolve_seed_groups_by_dataset_async(self, *, apply_sampling: bool = True): return {} async def _build_atomic_attacks_async(self, *, context): diff --git a/tests/unit/scenario/core/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py index 5ba5cb6bd1..707c759f25 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -109,7 +109,7 @@ def __init__(self, *, atomic_attacks_to_return=None, objective_scorer=None, **kw super().__init__(strategy_class=strategy_class, objective_scorer=objective_scorer, **kwargs) self._test_atomic_attacks = atomic_attacks_to_return or [] - async def _resolve_seed_groups_by_dataset_async(self): + async def _resolve_seed_groups_by_dataset_async(self, *, apply_sampling: bool = True): return {} async def _build_atomic_attacks_async(self, *, context): diff --git a/tests/unit/scenario/core/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py index 318e991dec..955c9462ab 100644 --- a/tests/unit/scenario/core/test_scenario_retry.py +++ b/tests/unit/scenario/core/test_scenario_retry.py @@ -180,7 +180,7 @@ def __init__(self, *, atomic_attacks_to_return=None, objective_scorer=None, **kw super().__init__(strategy_class=strategy_class, objective_scorer=objective_scorer, **kwargs) self._atomic_attacks_to_return = atomic_attacks_to_return or [] - async def _resolve_seed_groups_by_dataset_async(self): + async def _resolve_seed_groups_by_dataset_async(self, *, apply_sampling: bool = True): return {} async def _build_atomic_attacks_async(self, *, context):