From 67c581c4ee84aa7356e161d8f98697fecc0f15da Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 9 Jul 2026 09:49:24 -0700 Subject: [PATCH 1/3] Fix non-deterministic adaptive scenario resume under max_dataset_size Resume re-drew an unseeded random.sample before restricting to persisted objective hashes, so whenever max_dataset_size < len(dataset) the fresh draw diverged from the persisted subset and resume aborted with 'persisted objective hash(es) are no longer present in the dataset'. Thread an apply_sampling keyword through the seed-resolution path so resume resolves the full deterministic dataset and lets the persisted hashes reconstruct the exact original subset. Fresh runs still sample as before. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/dataset_configuration.py | 58 ++++++-- pyrit/scenario/core/scenario.py | 33 +++-- pyrit/scenario/scenarios/airt/psychosocial.py | 12 +- .../scenario/scenarios/garak/web_injection.py | 9 +- tests/unit/scenario/core/test_scenario.py | 125 +++++++++++++++++- .../scenario/core/test_scenario_parameters.py | 2 +- .../core/test_scenario_partial_results.py | 2 +- .../unit/scenario/core/test_scenario_retry.py | 2 +- 8 files changed, 212 insertions(+), 31 deletions(-) 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/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): From b9caf586ff29919ec6359c0252533cf72c7ee0f2 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 9 Jul 2026 10:09:52 -0700 Subject: [PATCH 2/3] Fix flaky 'same conversation' error in simulated next-message generation _generate_next_message_async built the request message with Message.from_prompt, which leaves conversation_id unset (None), then passed that None to set_system_prompt. set_system_prompt calls get_conversation_messages, whose conversation_id filter is skipped when the id is falsy, so it returned every message piece in memory and group_conversation_message_pieces_by_sequence raised 'All message pieces must be from the same conversation' as soon as memory held more than one conversation. This surfaced intermittently in adaptive single_turn runs (context_compliance) once the epsilon-greedy selector picked the simulated path and multiple conversations existed. Scope the system prompt and the generated request message to a fresh unique conversation id so the filter is always applied. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../multi_turn/simulated_conversation.py | 13 +++- .../component/test_simulated_conversation.py | 69 +++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) 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/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, From 36a778935c98e435dc1a118820c1f3f89296fc7d Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 9 Jul 2026 16:13:33 -0700 Subject: [PATCH 3/3] Harden get_conversation_messages against falsy conversation_id A falsy conversation_id caused get_message_pieces to skip its filter and silently return pieces from every conversation, which surfaced as the flaky 'same conversation' error. Raise ValueError instead. Two TAP tests were passing only via this footgun (querying an empty conversation_id); retarget them to resolve the real node conversation from memory. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/memory/memory_interface.py | 7 ++++ .../test_attack_parameter_consistency.py | 40 ++++++++++++++----- .../test_interface_prompts.py | 25 ++++++++++++ 3 files changed, 62 insertions(+), 10 deletions(-) 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/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())