FIX Prevent duplicate attacks when scenario resume history cannot be read (#2857)

Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: biefan

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)

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")