diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index abef3f5431..9b7cc6b82d 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -709,6 +709,7 @@ def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: 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. + Per-objective atomic attacks outside that subset are removed before scheduling. 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. @@ -726,8 +727,11 @@ def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: persisted_hashes: set[str] = set(persisted) retained: set[str] = set() + retained_attacks: list[AtomicAttack] = [] for aa in self._atomic_attacks: retained |= aa.keep_seed_groups_with_hashes(hashes=persisted_hashes) + if aa.seed_groups: + retained_attacks.append(aa) missing = persisted_hashes - retained if missing: @@ -739,6 +743,8 @@ def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: f"Either restore the missing objectives or drop scenario_result_id to start a new scenario." ) + self._atomic_attacks = retained_attacks + def _build_baseline_atomic_attack(self, *, seed_groups: list[AttackSeedGroup]) -> AtomicAttack: """ Build the baseline AtomicAttack from pre-resolved seed groups. @@ -946,8 +952,8 @@ async def _resolve_seed_groups_by_dataset_async( reuses the same population for every atomic attack and the baseline — so sampling under ``max_dataset_size`` stays consistent across all of them. - Override to inject seeds from an alternate source or to filter the resolved groups - before attacks are built. + 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. diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py index 8bf9b368ae..6c3aa53411 100644 --- a/tests/unit/scenario/core/test_scenario.py +++ b/tests/unit/scenario/core/test_scenario.py @@ -24,6 +24,7 @@ ScenarioResult, ) from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy, Scenario, ScenarioTechnique +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.score import Scorer from tests.unit.mocks import make_scenario_identifier, make_scenario_result @@ -1072,6 +1073,19 @@ async def _build_atomic_attacks_async(self, *, context): ) ] + class _PerObjectiveScenario(ConcreteScenarioWithTrueFalseScorer): + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + from pyrit.scenario.core.attack_technique import AttackTechnique + + return [ + AtomicAttack( + atomic_attack_name=f"strategy-{index}", + attack_technique=AttackTechnique(attack=MagicMock()), + seed_groups=[seed_group], + ) + for index, seed_group in enumerate(context.seed_groups) + ] + def _make_config(self): from pyrit.models import SeedGroup, SeedObjective @@ -1140,6 +1154,44 @@ def _sample_last_k(population, k): assert set(strategy.objectives) == persisted_objectives assert set(baseline.objectives) == persisted_objectives + async def test_resume_discards_per_objective_attacks_outside_persisted_subset(self, mock_objective_target): + def _sample_first_k(population, k): + return list(population)[:k] + + with patch( + "pyrit.scenario.core.dataset_configuration.random.sample", + side_effect=_sample_first_k, + ): + scenario = self._PerObjectiveScenario(name="Per-objective resume", version=1) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": self._make_config(), + } + ) + await scenario.initialize_async() + + original_id = scenario._scenario_result_id + assert original_id is not None + + resumed = self._PerObjectiveScenario( + name="Per-objective 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(), + } + ) + await resumed.initialize_async() + + assert len(resumed._atomic_attacks) == 4 + assert all(atomic_attack.seed_groups for atomic_attack in resumed._atomic_attacks) + 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()