From 2c89a56175cbd3c26ba565e372d8094fc905e742 Mon Sep 17 00:00:00 2001 From: biefan <70761325+biefan@users.noreply.github.com> Date: Fri, 25 Sep 2026 21:42:21 +0000 Subject: [PATCH] FIX Read scenario resume progress once and fail safely on storage errors --- pyrit/scenario/core/scenario.py | 66 +++++------ .../unit/scenario/core/test_scenario_retry.py | 110 +++++++++++++----- 2 files changed, 111 insertions(+), 65 deletions(-) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 8fb8ce704c..3984f62c0a 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -1310,12 +1310,11 @@ def _validate_stored_scenario( f"(ID: {self._scenario_result_id}, state: {stored_result.scenario_run_state})" ) - def _get_completed_objective_hashes_for_attack(self, *, atomic_attack: AtomicAttack) -> set[str]: + def _get_completed_objective_hashes_by_attack(self) -> dict[tuple[str, str | None], set[str]]: """ - Return the set of ``objective_sha256`` values already completed (non-error) - for a specific atomic attack inside this scenario. + Index completed objective hashes for every atomic attack in this scenario. - Queries ``AttackResultEntry`` rows directly by ``attribution_parent_id`` — + Read the persisted attack results once by ``attribution_parent_id`` — which is stamped at write-time by the attack persistence path — so results from an interrupted run are visible even though the ``ScenarioResult.attack_results`` aggregate may not yet reflect them. @@ -1329,40 +1328,29 @@ def _get_completed_objective_hashes_for_attack(self, *, atomic_attack: AtomicAtt ``parent_eval_hash`` was introduced (or by callers that don't supply one) match name-only as a backward-compatible fallback. - Args: - atomic_attack (AtomicAttack): The live atomic attack whose - ``atomic_attack_name`` and technique identifier scope the query. - Returns: - set[str]: ``objective_sha256`` hex strings for completed-without-error rows. + dict[tuple[str, str | None], set[str]]: Completed objective hashes keyed by + collection name and technique eval hash. A missing eval hash retains + the legacy name-only matching behavior. + + Raises: + Exception: If persisted progress cannot be read. Treating a failed read as + empty progress would re-execute already-completed objectives. """ if not self._scenario_result_id: - return set() - - atomic_attack_name = atomic_attack.atomic_attack_name - expected_eval_hash = atomic_attack.technique_eval_hash - - completed_hashes: set[str] = set() - try: - rows = self._memory.get_attack_results(scenario_result_id=self._scenario_result_id) - for row in rows: - if row.outcome == AttackOutcome.ERROR: - continue - if row.attribution_data is None: - continue - if row.attribution_data.get("parent_collection") != atomic_attack_name: - continue - row_eval_hash = row.attribution_data.get("parent_eval_hash") - if row_eval_hash is not None and row_eval_hash != expected_eval_hash: - continue - if row.objective: - completed_hashes.add(to_sha256(row.objective)) - except Exception as e: - logger.warning( - f"Failed to retrieve completed objective hashes for atomic attack '{atomic_attack_name}': {str(e)}" - ) + return {} - return completed_hashes + rows = self._memory.get_attack_results(scenario_result_id=self._scenario_result_id) + completed_by_attack: dict[tuple[str, str | None], set[str]] = {} + for row in rows: + if row.outcome == AttackOutcome.ERROR or not row.attribution_data or not row.objective: + continue + name = row.attribution_data.get("parent_collection") + eval_hash = row.attribution_data.get("parent_eval_hash") + if not isinstance(name, str) or (eval_hash is not None and not isinstance(eval_hash, str)): + continue + completed_by_attack.setdefault((name, eval_hash), set()).add(to_sha256(row.objective)) + return completed_by_attack async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: """ @@ -1372,7 +1360,8 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: atomic attack enforces uniqueness of objective hashes at construction time, and the executor stamps ``attribution_parent_id`` + ``attribution_data["parent_collection"]`` on the row so a content-hash - join is sufficient. + join is sufficient. Each call reads a fresh snapshot before filtering any + seed groups; a read failure propagates to the scenario's retry policy. Returns: list[AtomicAttack]: List of atomic attacks with uncompleted objectives. @@ -1382,9 +1371,14 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: return self._atomic_attacks remaining_attacks: list[AtomicAttack] = [] + # Read and index one snapshot before changing any attack's remaining work. + completed_by_attack = self._get_completed_objective_hashes_by_attack() for atomic_attack in self._atomic_attacks: - completed_hashes = self._get_completed_objective_hashes_for_attack(atomic_attack=atomic_attack) + name = atomic_attack.atomic_attack_name + completed_hashes = completed_by_attack.get((name, atomic_attack.technique_eval_hash), set()) | ( + completed_by_attack.get((name, None), set()) + ) if completed_hashes: original_count = len(atomic_attack.seed_groups) diff --git a/tests/unit/scenario/core/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py index d29bd0cd22..1a5307cbb0 100644 --- a/tests/unit/scenario/core/test_scenario_retry.py +++ b/tests/unit/scenario/core/test_scenario_retry.py @@ -18,6 +18,7 @@ AttackSeedGroup, ComponentIdentifier, Message, + ScenarioRunState, SeedObjective, config_hash, ) @@ -484,6 +485,70 @@ async def mock_run_with_logged_failure(*args, **kwargs): class TestScenarioResumption: """Tests for Scenario resumption after partial failure.""" + @pytest.mark.parametrize("persistent_failure", [False, True], ids=["transient-read", "persistent-read"]) + async def test_resume_read_failure_never_reexecutes_completed_objectives( + self, mock_objective_target, persistent_failure + ): + completed = create_mock_atomic_attack("completed", ["objective1"]) + pending = create_mock_atomic_attack("pending", ["objective2"]) + completed.run_async = create_mock_run_async([create_attack_result(1)], atomic_attack=completed) + pending.run_async = create_mock_run_async([create_attack_result(2)], atomic_attack=pending) + scenario = ConcreteScenario( + name="Resume Read Failure", version=1, atomic_attacks_to_return=[completed, pending] + ) + scenario.set_params_from_args(args={"objective_target": mock_objective_target, "max_retries": 1}) + await scenario.initialize_async() + completed.set_scenario_result_id(scenario._scenario_result_id) + save_attack_results_to_memory([create_attack_result(1)], atomic_attack=completed) + + memory = CentralMemory.get_memory_instance() + original_get_results = memory.get_attack_results + reads = 0 + + def read_results(**kwargs): + nonlocal reads + reads += 1 + if persistent_failure or reads == 1: + raise RuntimeError("resume storage unavailable") + return original_get_results(**kwargs) + + with patch.object(memory, "get_attack_results", side_effect=read_results): + if persistent_failure: + with pytest.raises(RuntimeError, match="resume storage unavailable"): + await scenario.run_async() + else: + await scenario.run_async() + + completed.run_async.assert_not_called() + assert pending.run_async.call_count == (0 if persistent_failure else 1) + assert reads == 2 + [stored] = memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id]) + assert stored.number_tries == 2 + assert stored.scenario_run_state == ( + ScenarioRunState.FAILED if persistent_failure else ScenarioRunState.COMPLETED + ) + assert len(stored.attack_results["completed"]) == 1 + if persistent_failure: + assert completed.objectives == ["objective1"] + assert pending.objectives == ["objective2"] + assert stored.error_message == "resume storage unavailable" + + async def test_resume_loads_one_result_snapshot_for_all_atomic_attacks(self, mock_objective_target): + attacks = [create_mock_atomic_attack(f"attack_{i}", [f"objective{i}"]) for i in range(10)] + scenario = ConcreteScenario(name="Resume Snapshot", version=1, atomic_attacks_to_return=attacks) + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.initialize_async() + for i, attack in enumerate(attacks[:5]): + attack.set_scenario_result_id(scenario._scenario_result_id) + save_attack_results_to_memory([create_attack_result(i)], atomic_attack=attack) + + memory = CentralMemory.get_memory_instance() + with patch.object(memory, "get_attack_results", wraps=memory.get_attack_results) as read_results: + remaining = await scenario._get_remaining_atomic_attacks_async() + + assert remaining == attacks[5:] + read_results.assert_called_once_with(scenario_result_id=scenario._scenario_result_id) + async def test_parameter_build_partial_result_persists_linkage_before_retry( self, mock_objective_target: MagicMock, @@ -858,8 +923,8 @@ async def noop_run(*args, **kwargs): @pytest.mark.usefixtures("patch_central_database") -class TestGetCompletedObjectiveHashesForAttack: - """Direct tests for ``Scenario._get_completed_objective_hashes_for_attack`` +class TestGetCompletedObjectiveHashesByAttack: + """Direct tests for ``Scenario._get_completed_objective_hashes_by_attack`` — the filter that excludes already-completed objectives on resume. Covers the row-filtering branches: outcome=ERROR rows, rows without @@ -873,12 +938,6 @@ def _make_scenario(self, scenario_result_id="scn-1"): scenario._memory = MagicMock() return scenario - def _make_atomic(self, name, eval_hash="hash-A"): - atomic = MagicMock(spec=AtomicAttack) - atomic.atomic_attack_name = name - type(atomic).technique_eval_hash = PropertyMock(return_value=eval_hash) - return atomic - def _row(self, *, objective, outcome=AttackOutcome.SUCCESS, attribution_data=None): row = MagicMock() row.outcome = outcome @@ -889,10 +948,8 @@ def _row(self, *, objective, outcome=AttackOutcome.SUCCESS, attribution_data=Non def test_returns_empty_when_scenario_result_id_unset(self): scenario = ConcreteScenario(name="S", version=1, atomic_attacks_to_return=[]) scenario._scenario_result_id = None - result = scenario._get_completed_objective_hashes_for_attack( - atomic_attack=self._make_atomic("a"), - ) - assert result == set() + result = scenario._get_completed_objective_hashes_by_attack() + assert result == {} def test_skips_error_rows(self): from pyrit.common.utils import to_sha256 @@ -910,10 +967,8 @@ def test_skips_error_rows(self): attribution_data={"parent_collection": "a", "parent_eval_hash": "hash-A"}, ), ] - result = scenario._get_completed_objective_hashes_for_attack( - atomic_attack=self._make_atomic("a"), - ) - assert result == {to_sha256("ok")} + result = scenario._get_completed_objective_hashes_by_attack() + assert result == {("a", "hash-A"): {to_sha256("ok")}} def test_skips_rows_without_attribution_data(self): from pyrit.common.utils import to_sha256 @@ -926,12 +981,10 @@ def test_skips_rows_without_attribution_data(self): attribution_data={"parent_collection": "a", "parent_eval_hash": "hash-A"}, ), ] - result = scenario._get_completed_objective_hashes_for_attack( - atomic_attack=self._make_atomic("a"), - ) - assert result == {to_sha256("new")} + result = scenario._get_completed_objective_hashes_by_attack() + assert result == {("a", "hash-A"): {to_sha256("new")}} - def test_skips_rows_with_mismatched_eval_hash(self): + def test_indexes_same_name_techniques_separately(self): """Two atomic attacks with the same name but different techniques must not cross-pollinate completed hashes. This is the core Option-B guarantee.""" @@ -948,10 +1001,11 @@ def test_skips_rows_with_mismatched_eval_hash(self): attribution_data={"parent_collection": "encoding", "parent_eval_hash": "hash-hex"}, ), ] - result = scenario._get_completed_objective_hashes_for_attack( - atomic_attack=self._make_atomic("encoding", eval_hash="hash-base64"), - ) - assert result == {to_sha256("mine")} + result = scenario._get_completed_objective_hashes_by_attack() + assert result == { + ("encoding", "hash-base64"): {to_sha256("mine")}, + ("encoding", "hash-hex"): {to_sha256("theirs")}, + } def test_backward_compat_matches_name_only_when_eval_hash_missing(self): """Rows persisted before ``parent_eval_hash`` shipped match name-only @@ -965,10 +1019,8 @@ def test_backward_compat_matches_name_only_when_eval_hash_missing(self): attribution_data={"parent_collection": "a"}, # no parent_eval_hash ), ] - result = scenario._get_completed_objective_hashes_for_attack( - atomic_attack=self._make_atomic("a", eval_hash="hash-A"), - ) - assert result == {to_sha256("old")} + result = scenario._get_completed_objective_hashes_by_attack() + assert result == {("a", None): {to_sha256("old")}} @pytest.mark.usefixtures("patch_central_database")