FIX: Adaptive Scenario bugs (#2152)

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: rlundeen2

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)

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)
 

pyrit/scenario/core/dataset_configuration.py

@@ -622,16 +622,23 @@ 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]``.
 
         Builds attack groups (inline or from memory, auto-fetching missing datasets),
         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 "<inline>"
             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,16 +666,24 @@ 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
                 dataset fails validation, or no attack groups could be built.
         """
         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 "<inline>"
             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:
         """

pyrit/scenario/core/scenario.py

@@ -612,7 +612,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)
 
@@ -699,12 +705,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.
@@ -929,7 +935,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.
 
@@ -941,10 +949,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:
         """

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 technique 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(

pyrit/scenario/scenarios/garak/web_injection.py

@@ -498,7 +498,9 @@ def _scoring_config_for_technique(self, technique: WebInjectionTechnique) -> Att
             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 technique.
 
@@ -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 technique value.