FIX: give adaptive technique selection the run id on fresh runs

pyrit/scenario/core/scenario.py

@@ -916,7 +916,12 @@ async def initialize_async(self) -> None:
         else:
             await self._dataset_config.prepare_async()
         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)
+        # A fresh run's result id is chosen now, so attacks built below can already refer to the run.
+        new_result_id = uuid.uuid4()
+        context = self._build_scenario_context(
+            seed_groups_by_dataset=seed_groups_by_dataset,
+            scenario_result_id=self._scenario_result_id or str(new_result_id),
+        )
         self._atomic_attacks = await self._build_atomic_attacks_async(context=context)
 
         # Build the canonical scenario identifier once params/techniques/datasets
@@ -958,6 +963,7 @@ async def initialize_async(self) -> None:
         }
 
         result = ScenarioResult(
+            id=new_result_id,
             scenario_identifier=scenario_identifier,
             scenario_description=self._description,
             labels=self._memory_labels,
@@ -1385,7 +1391,12 @@ async def _resolve_seed_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[AttackSeedGroup]]) -> ScenarioContext:
+    def _build_scenario_context(
+        self,
+        *,
+        seed_groups_by_dataset: dict[str, list[AttackSeedGroup]],
+        scenario_result_id: str | None = None,
+    ) -> ScenarioContext:
         """
         Snapshot the resolved runtime inputs into a ``ScenarioContext``.
 
@@ -1398,6 +1409,7 @@ def _build_scenario_context(self, *, seed_groups_by_dataset: dict[str, list[Atta
             seed_groups_by_dataset (dict[str, list[AttackSeedGroup]]): Seed groups already
                 resolved once (see ``_resolve_seed_groups_by_dataset_async``). The flat
                 ``context.seed_groups`` is derived from these so both views share one sample.
+            scenario_result_id (str | None): The id of the ``ScenarioResult`` this run writes to.
 
         Returns:
             ScenarioContext: The immutable inputs for atomic-attack construction.
@@ -1420,6 +1432,7 @@ def _build_scenario_context(self, *, seed_groups_by_dataset: dict[str, list[Atta
             include_baseline=self._include_baseline,
             seed_groups=seed_groups,
             seed_groups_by_dataset=seed_groups_by_dataset,
+            scenario_result_id=scenario_result_id,
         )
 
     @abstractmethod

pyrit/scenario/core/scenario_context.py

@@ -51,6 +51,8 @@ class ScenarioContext:
         seed_groups_by_dataset (Mapping[str, list[AttackSeedGroup]]): The same resolved
             seed groups keyed by originating dataset name, for scenarios that map datasets
             onto separate attacks or display groups.
+        scenario_result_id (str | None): The ``ScenarioResult`` id of this run: the stored id
+            on resume, or the id the new result will be created with on a fresh run.
     """
 
     objective_target: PromptTarget
@@ -60,3 +62,4 @@ class ScenarioContext:
     include_baseline: bool = False
     seed_groups: Sequence[AttackSeedGroup] = field(default_factory=tuple)
     seed_groups_by_dataset: Mapping[str, list[AttackSeedGroup]] = field(default_factory=dict)
+    scenario_result_id: str | None = None

pyrit/scenario/scenarios/adaptive/adaptive_scenario.py

@@ -212,6 +212,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list
                     seed_groups=seed_groups,
                     techniques=techniques,
                     selector=self._selector,
+                    scenario_result_id=context.scenario_result_id,
                 )
             )
 
@@ -397,6 +398,7 @@ async def _build_atomics_for_dataset_async(
         seed_groups: list[AttackSeedGroup],
         techniques: dict[str, TechniqueBundle],
         selector: TechniqueSelector,
+        scenario_result_id: str | None = None,
     ) -> list[AtomicAttack]:
         """
         Build one ``AtomicAttack`` per seed group with at least one
@@ -430,7 +432,7 @@ async def _build_atomics_for_dataset_async(
             selector=selector,
             objective_scorer=self._objective_scorer,
             max_attempts_per_objective=self.params.get("max_attempts_per_objective", 3),
-            scenario_result_id=self._scenario_result_id,
+            scenario_result_id=scenario_result_id,
         )
 
         atomics: list[AtomicAttack] = []

tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py

@@ -7,20 +7,25 @@
 
 import uuid
 import warnings
+from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 from pyrit.models import AttackSeedGroup, ScenarioRunPlanGroupKind, SeedObjective
-from pyrit.models.identifiers import ComponentIdentifier
+from pyrit.models.identifiers import ComponentIdentifier, ScenarioEvaluationIdentifier
 from pyrit.prompt_target import PromptTarget
 from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry
 from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration
 from pyrit.scenario.core.scenario import BaselineAttackPolicy
 from pyrit.scenario.scenarios.adaptive.dispatcher import AdaptiveTechniqueDispatcher
+from pyrit.scenario.scenarios.adaptive.selectors import EpsilonGreedyTechniqueSelector, SelectorScope
 from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive
 from pyrit.score import TrueFalseScorer
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 _MOCK_MANY_SHOT_EXAMPLES = [{"question": f"q{i}", "answer": f"a{i}"} for i in range(100)]
 
 
@@ -292,6 +297,77 @@ def _spy_init(self, *args, **kwargs):
         assert len(selectors_seen) == 2
         assert len({id(s) for s in selectors_seen}) == 1
 
+    async def test_selection_uses_this_run_id_and_resume_repeats_it_async(
+        self, *, mock_objective_target: MagicMock, mock_objective_scorer: MagicMock
+    ) -> None:
+        groups = {"violence": [_make_seed_group(value=f"obj-{i}", harm_categories=["violence"]) for i in range(6)]}
+        seen: list[tuple[str | None, tuple[str, ...]]] = []
+        real_select = EpsilonGreedyTechniqueSelector.select_async
+
+        async def _spy_select_async(
+            self: EpsilonGreedyTechniqueSelector,
+            *,
+            technique_identifiers: Sequence[str],
+            objective: str,
+            num_top_techniques: int = 1,
+            scenario_result_id: str | None = None,
+        ) -> Sequence[str]:
+            picks = await real_select(
+                self,
+                technique_identifiers=technique_identifiers,
+                objective=objective,
+                num_top_techniques=num_top_techniques,
+                scenario_result_id=scenario_result_id,
+            )
+            seen.append((scenario_result_id, tuple(picks)))
+            return picks
+
+        async def _initialize_async(scenario_result_id: str | None = None) -> TextAdaptive:
+            scenario = TextAdaptive(
+                objective_scorer=mock_objective_scorer,
+                selector=EpsilonGreedyTechniqueSelector(epsilon=0.5, random_seed=7, scope=SelectorScope.current_run()),
+                scenario_result_id=scenario_result_id,
+            )
+            scenario.set_params_from_args(args={"objective_target": mock_objective_target, "include_baseline": False})
+            await scenario.initialize_async()
+            return scenario
+
+        with (
+            patch.object(DatasetAttackConfiguration, "prepare_async", new_callable=AsyncMock) as mock_prepare,
+            patch.object(
+                DatasetAttackConfiguration,
+                "get_attack_groups_by_dataset_async",
+                new_callable=AsyncMock,
+                return_value=groups,
+            ) as mock_groups,
+            patch.object(EpsilonGreedyTechniqueSelector, "select_async", _spy_select_async),
+        ):
+            fresh = await _initialize_async()
+            fresh_calls = list(seen)
+            seen.clear()
+            resumed = await _initialize_async(scenario_result_id=fresh._scenario_result_id)
+            resume_calls = list(seen)
+            seen.clear()
+            second = await _initialize_async()
+            stored_results = await fresh._memory.get_scenario_results_async(
+                scenario_result_ids=[fresh._scenario_result_id, second._scenario_result_id]
+            )
+
+        assert len(fresh_calls) == len(resume_calls) == len(seen) == 6
+        assert fresh._scenario_result_id is not None
+        assert resumed._scenario_result_id == fresh._scenario_result_id
+        assert second._scenario_result_id is not None
+        assert second._scenario_result_id != fresh._scenario_result_id
+        assert {run_id for run_id, _ in fresh_calls} == {fresh._scenario_result_id}
+        assert resume_calls == fresh_calls
+        assert {run_id for run_id, _ in seen} == {second._scenario_result_id}
+        assert {str(result.id) for result in stored_results} == {fresh._scenario_result_id, second._scenario_result_id}
+        assert (
+            len({ScenarioEvaluationIdentifier(result.scenario_identifier).eval_hash for result in stored_results}) == 1
+        )
+        assert mock_prepare.await_count == 2
+        assert [call.kwargs["apply_sampling"] for call in mock_groups.await_args_list] == [True, False, True]
+
     async def test_atomic_names_contain_dataset_and_objective_hash(self, mock_objective_target, mock_objective_scorer):
         groups = {
             "violence": [_make_seed_group(value=f"obj-{i}", harm_categories=["violence"]) for i in range(5)],