From 5e858ad7e6a831e512967fccec64dac5d5ec49a2 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Mon, 18 May 2026 17:01:41 -0700 Subject: [PATCH 01/10] Stamp scenario_result_id FK on AttackResultEntry for granular resume Previously, if a Scenario was interrupted mid-AtomicAttack (Ctrl-C, OOM, crash), completed AttackResults persisted to the DB became orphaned because the scenario-to-attack-result link only lived in a JSON manifest (attack_results_json) written after the whole AtomicAttack returned. On resume, those objectives re-executed wastefully. This change makes scenario linkage a first-class column on AttackResultEntry: - New columns: scenario_result_id (indexed FK, ON DELETE SET NULL) and scenario_data (JSON with fixed schema {atomic_attack_name, objective_index}). - New ExecutionAttribution dataclass in pyrit/executor/attack/core/ (so the executor never imports from the scenario layer) is set on AttackContext by AttackExecutor per-task before scheduling, and read by the default attack event handler when persisting. - Hydration in get_scenario_results uses the FK with a merge-mode fallback to the legacy manifest for partially-migrated DBs. - Resume uses objective_index (deterministic, parallel-safe; derived from seed_groups input_indices) rather than objective text, so duplicate objective text doesn't collapse two seed groups. - Drops the unreleased error_attack_result_ids_json column outright; error AttackResults are now linkable via get_attack_results(scenario_result_id=..., outcome=ERROR). - attack_results_json stays write-through this release for downgrade safety; future releases will stop populating and then drop. - update_scenario_run_state becomes a targeted UPDATE rather than a full row rebuild (so it doesn't clobber the manifest during the deprecation window). Includes Alembic migration with idempotent backfill, scenario_data round-trip on AttackResultEntry, and tests for: event-handler attribution stamping, executor attribution propagation at max_concurrency>1, FK + manifest + mixed hydration paths, migration backfill correctness/idempotency/downgrade, interruption-recovery regression, duplicate-objective-text resume safety, and duplicate atomic_attack_name validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../backend/services/scenario_run_service.py | 10 +- pyrit/executor/attack/core/attack_executor.py | 33 ++- pyrit/executor/attack/core/attack_strategy.py | 45 ++++ .../attack/core/execution_attribution.py | 44 +++ ..._add_scenario_linkage_to_attack_results.py | 167 ++++++++++++ pyrit/memory/memory_interface.py | 221 ++++++++------- pyrit/memory/memory_models.py | 33 ++- pyrit/models/attack_result.py | 8 + pyrit/models/scenario_result.py | 5 - pyrit/scenario/core/atomic_attack.py | 76 +++++- pyrit/scenario/core/scenario.py | 157 +++++------ .../unit/backend/test_scenario_run_service.py | 18 +- .../attack/core/test_attack_executor.py | 106 ++++++++ .../attack/core/test_attack_strategy.py | 86 ++++++ .../test_interface_scenario_results.py | 233 +++++++++++++--- tests/unit/memory/test_memory_models.py | 14 - tests/unit/memory/test_migration.py | 255 ++++++++++++++++++ tests/unit/models/test_scenario_result.py | 21 -- tests/unit/scenario/test_scenario.py | 53 +++- .../scenario/test_scenario_partial_results.py | 66 +++-- tests/unit/scenario/test_scenario_retry.py | 254 +++++++++++++++-- 21 files changed, 1551 insertions(+), 354 deletions(-) create mode 100644 pyrit/executor/attack/core/execution_attribution.py create mode 100644 pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 37f0ff1b71..85907a430f 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -409,9 +409,13 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari error = scenario_result.error_message error_type = scenario_result.error_type - # Fallback: look up error from persisted error AttackResults - if not error and scenario_result.error_attack_result_ids: - error_ars = self._memory.get_attack_results(attack_result_ids=scenario_result.error_attack_result_ids) + # Fallback: look up error from any persisted error AttackResults linked + # to this scenario via the new FK. + if not error: + error_ars = self._memory.get_attack_results( + scenario_result_id=scenario_result_id, + outcome=AttackOutcome.ERROR, + ) if error_ars: error = error_ars[0].error_message error_type = error_ars[0].error_type diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index fea9315e0d..edc5a212f8 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -8,7 +8,7 @@ """ import asyncio -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, @@ -24,6 +24,7 @@ AttackStrategyContextT, AttackStrategyResultT, ) +from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution from pyrit.models import SeedAttackGroup if TYPE_CHECKING: @@ -142,6 +143,7 @@ async def execute_attack_from_seed_groups_async( objective_scorer: Optional["TrueFalseScorer"] = None, field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, + attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -163,6 +165,12 @@ async def execute_attack_from_seed_groups_async( from_seed_group() as overrides. return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. + attribution_factory: Optional callable that maps an input index (the + seed group's original index, parallel-safe and deterministic) to + an ``ExecutionAttribution``. When provided, each per-task + ``AttackContext`` is stamped with the attribution so the + resulting ``AttackResultEntry`` row carries the scenario FK + + scenario_data. When ``None``, no attribution is applied. **broadcast_fields: Fields applied to all seed groups (e.g., memory_labels). Per-seed-group field_overrides take precedence. @@ -205,6 +213,7 @@ async def build_params(i: int, sg: SeedAttackGroup) -> AttackParameters: attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, + attribution_factory=attribution_factory, ) async def execute_attack_async( @@ -214,6 +223,7 @@ async def execute_attack_async( objectives: Sequence[str], field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, + attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -228,6 +238,10 @@ async def execute_attack_async( must match the length of objectives. return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. + attribution_factory: Optional callable mapping each input index to + an ExecutionAttribution. When provided, the per-task context is + stamped with the attribution so the persistence path can record + scenario linkage. **broadcast_fields: Fields applied to all objectives (e.g., memory_labels). Per-objective field_overrides take precedence. @@ -268,6 +282,7 @@ async def execute_attack_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, + attribution_factory=attribution_factory, ) async def _execute_with_params_list_async( @@ -276,6 +291,7 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, + attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Execute attacks in parallel with a list of pre-built parameters. @@ -287,19 +303,28 @@ async def _execute_with_params_list_async( attack: The attack strategy to execute. params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. + attribution_factory: Optional callable mapping each input index to + an ExecutionAttribution. When provided, the per-task context is + stamped with the attribution so the persistence path can record + scenario linkage. Returns: AttackExecutorResult with completed results and any incomplete objectives. """ semaphore = asyncio.Semaphore(self._max_concurrency) - async def run_one(params: AttackParameters) -> AttackStrategyResultT: + async def run_one(index: int, params: AttackParameters) -> AttackStrategyResultT: async with semaphore: - # Create context with params + # Create context with params and stamp attribution (if any). The + # input index is the seed group's original position and is + # deterministic and parallel-safe — assigned BEFORE the task + # runs, not from completion order. context = attack._context_type(params=params) + if attribution_factory is not None: + context._attribution = attribution_factory(index) return await attack.execute_with_context_async(context=context) - tasks = [run_one(p) for p in params_list] + tasks = [run_one(i, p) for i, p in enumerate(params_list)] results_or_exceptions = await asyncio.gather(*tasks, return_exceptions=True) return self._process_execution_results( diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index a054ace813..278c96ec7f 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from pyrit.executor.attack.core.attack_config import AttackScoringConfig + from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution from pyrit.prompt_target import PromptTarget AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") @@ -73,6 +74,13 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): # Set by the ON_ERROR handler to link error AttackResults to ScenarioResults _error_attack_result_id: str | None = None + # Optional attribution from an upstream orchestrator (e.g. Scenario). When + # set, the persistence path stamps scenario_result_id + scenario_data onto + # the resulting AttackResult so it can be located later for hydration and + # resume. Set by AttackExecutor per-task before scheduling. Stays None for + # ad-hoc/direct attack execution outside any scenario. + _attribution: Optional[ExecutionAttribution] = None + # Convenience properties that delegate to params or overrides @property def objective(self) -> str: @@ -223,11 +231,43 @@ async def _on_post_execute( event_data.result.retry_events = collector.events event_data.result.total_retries = len(collector.events) + # Stamp scenario attribution onto the result before persistence so the + # AttackResultEntry row carries the FK + scenario_data. Outside scenarios + # _attribution is None and both fields stay None. + self._stamp_attribution(context=event_data.context, result=event_data.result) + self._logger.debug(f"Attack execution completed in {execution_time_ms}ms") self._log_attack_outcome(event_data.result) self._memory.add_attack_results_to_memory(attack_results=[event_data.result]) + @staticmethod + def _stamp_attribution( + *, + context: AttackStrategyContextT, + result: AttackStrategyResultT, + ) -> None: + """ + Copy scenario attribution from the AttackContext onto the AttackResult. + + Reads ``context._attribution`` (an ``ExecutionAttribution`` set by the + AttackExecutor when running inside a Scenario). When present, writes + ``scenario_result_id`` and a fixed-schema ``scenario_data`` dict onto + the result so they round-trip into ``AttackResultEntry``. + + Args: + context: The per-task AttackContext. + result: The AttackResult that is about to be persisted. + """ + attribution = getattr(context, "_attribution", None) + if attribution is None: + return + result.scenario_result_id = attribution.scenario_result_id + result.scenario_data = { + "atomic_attack_name": attribution.atomic_attack_name, + "objective_index": attribution.objective_index, + } + def _log_attack_outcome(self, result: AttackResult) -> None: """ Log the outcome of the attack. @@ -295,6 +335,11 @@ async def _on_error_async( if context.start_time: error_result.execution_time_ms = int((end_time - context.start_time) * 1000) + # Stamp scenario attribution onto the error result so it is locatable + # via the scenario FK on resume (rather than via the previous + # error_attack_result_ids_json manifest). + self._stamp_attribution(context=context, result=error_result) + # Persist first, then set the ID on the context so scenario-level code # only sees the reference if the write succeeded. self._memory.add_attack_results_to_memory(attack_results=[error_result]) diff --git a/pyrit/executor/attack/core/execution_attribution.py b/pyrit/executor/attack/core/execution_attribution.py new file mode 100644 index 0000000000..09b6ecc131 --- /dev/null +++ b/pyrit/executor/attack/core/execution_attribution.py @@ -0,0 +1,44 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Typed attribution metadata used to link a persisted ``AttackResult`` to an +upstream orchestrator (e.g. a ``Scenario``). + +The attribution lives in the ``executor`` layer so the executor never imports +from ``scenario``. ``Scenario`` is one producer; future orchestrators may +produce attribution too. The attack persistence path (the default attack event +handler) is the consumer. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ExecutionAttribution: + """ + Attribution metadata produced by an upstream orchestrator and consumed by + the attack persistence path to record linkage on the resulting + ``AttackResultEntry``. + + Attributes: + scenario_result_id (str): The ID of the scenario result that produced + this attack execution. Persisted to + ``AttackResultEntry.scenario_result_id`` so per-scenario hydration + and resume can locate the row directly without relying on a JSON + manifest written at the end of an atomic attack. + atomic_attack_name (str): The unique key of the atomic attack within + the scenario (matches ``AtomicAttack.atomic_attack_name``). + Persisted into ``AttackResultEntry.scenario_data``. + objective_index (int): The 0-based original seed-group index (the + ``input_indices`` value from ``AttackExecutorResult``). Assigned + **before** task execution so it is deterministic and parallel-safe. + Persisted into ``AttackResultEntry.scenario_data`` and used as the + stable resume key (instead of the easily-duplicated objective text). + """ + + scenario_result_id: str + atomic_attack_name: str + objective_index: int diff --git a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py new file mode 100644 index 0000000000..7d39f46d70 --- /dev/null +++ b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py @@ -0,0 +1,167 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Add scenario_result_id FK + scenario_data JSON to AttackResultEntries; drop +ScenarioResultEntries.error_attack_result_ids_json; backfill linkage from the +existing attack_results_json manifest. + +Revision ID: 9c8b7a6d5e4f +Revises: 7a1b2c3d4e5f +Create Date: 2026-05-18 15:00:00.000000 +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Sequence # noqa: TC003 +from typing import Any + +import sqlalchemy as sa +from alembic import op + +from pyrit.memory.memory_models import CustomUUID + +# revision identifiers, used by Alembic. +revision: str = "9c8b7a6d5e4f" +down_revision: str | None = "7a1b2c3d4e5f" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +def upgrade() -> None: + """Apply this schema upgrade.""" + # AttackResultEntries: scenario linkage columns. + op.add_column( + "AttackResultEntries", + sa.Column("scenario_result_id", CustomUUID(), nullable=True), + ) + op.add_column( + "AttackResultEntries", + sa.Column("scenario_data", sa.JSON(), nullable=True), + ) + op.create_index( + "ix_AttackResultEntries_scenario_result_id", + "AttackResultEntries", + ["scenario_result_id"], + ) + + # FK with ON DELETE SET NULL: deleting a scenario nulls the FK on its + # AttackResults; scenario_data is retained as historical provenance. + # Use a batch operation for SQLite portability (no plain ALTER TABLE ADD FK). + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.create_foreign_key( + "fk_attack_results_scenario", + "ScenarioResultEntries", + ["scenario_result_id"], + ["id"], + ondelete="SET NULL", + ) + + # ScenarioResultEntries: drop the not-yet-released error_attack_result_ids_json. + # Error AttackResults are now linkable via the new FK; the per-scenario + # manifest column is no longer used. Wrapped in a batch op for SQLite. + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.drop_column("error_attack_result_ids_json") + + # Backfill scenario linkage from the existing attack_results_json manifest. + _backfill_scenario_linkage() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + # Re-add error_attack_result_ids_json on ScenarioResultEntries. + op.add_column( + "ScenarioResultEntries", + sa.Column("error_attack_result_ids_json", sa.Unicode(), nullable=True), + ) + + # Drop FK + columns from AttackResultEntries. + with op.batch_alter_table("AttackResultEntries") as batch_op: + batch_op.drop_constraint("fk_attack_results_scenario", type_="foreignkey") + + op.drop_index("ix_AttackResultEntries_scenario_result_id", table_name="AttackResultEntries") + op.drop_column("AttackResultEntries", "scenario_data") + op.drop_column("AttackResultEntries", "scenario_result_id") + + +def _backfill_scenario_linkage() -> None: + """ + Walk every ScenarioResultEntry and copy its attack_results_json manifest + into the new FK + scenario_data columns on AttackResultEntries. + + Idempotent: the ``WHERE scenario_result_id IS NULL`` guard prevents + clobbering rows that were already linked (e.g. by a re-run of the + migration, or by code that ran after the schema change but before this + backfill). ``conversation_id`` is logically unique per AttackResult but is + not DB-enforced, so the guard is purely defensive and a WARNING is logged + if any duplicate match is observed in the wild. + """ + bind = op.get_bind() + + scenarios = bind.execute(sa.text('SELECT id, attack_results_json FROM "ScenarioResultEntries"')).fetchall() + + update_stmt = sa.text( + 'UPDATE "AttackResultEntries" ' + "SET scenario_result_id = :sid, scenario_data = :sdata " + "WHERE conversation_id = :cid AND scenario_result_id IS NULL" + ) + + total_updates = 0 + duplicate_warnings = 0 + + for row in scenarios: + scenario_id = row[0] + manifest_json = row[1] + if not manifest_json: + continue + try: + manifest: dict[str, Any] = json.loads(manifest_json) + except (TypeError, ValueError): + logger.warning(f"Skipping scenario {scenario_id}: attack_results_json is not valid JSON") + continue + + for atomic_attack_name, conversation_ids in manifest.items(): + if not isinstance(conversation_ids, list): + continue + for objective_index, conversation_id in enumerate(conversation_ids): + if not isinstance(conversation_id, str): + continue + # Check for duplicate conversation_id matches (data anomaly). + match_count = bind.execute( + sa.text( + 'SELECT COUNT(*) FROM "AttackResultEntries" ' + "WHERE conversation_id = :cid AND scenario_result_id IS NULL" + ), + {"cid": conversation_id}, + ).scalar() + if isinstance(match_count, int) and match_count > 1: + duplicate_warnings += 1 + logger.warning( + f"Backfill: conversation_id {conversation_id!r} matches {match_count} " + f"unlinked AttackResultEntries rows; conversation_id should be unique. " + f"All matching rows will be linked to scenario {scenario_id}." + ) + + scenario_data = json.dumps( + {"atomic_attack_name": atomic_attack_name, "objective_index": objective_index} + ) + result = bind.execute( + update_stmt, + { + "sid": str(scenario_id), + "sdata": scenario_data, + "cid": conversation_id, + }, + ) + total_updates += result.rowcount or 0 + + if total_updates or duplicate_warnings: + logger.info( + f"Scenario linkage backfill: linked {total_updates} AttackResultEntries row(s); " + f"{duplicate_warnings} duplicate-conversation_id warning(s)." + ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 65d4480aaf..2b80b30e4d 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1701,6 +1701,7 @@ def get_attack_results( targeted_harm_categories: Optional[Sequence[str]] = None, labels: Optional[dict[str, str | Sequence[str]]] = None, identifier_filters: Optional[Sequence[IdentifierFilter]] = None, + scenario_result_id: Optional[str] = None, ) -> Sequence[AttackResult]: """ Retrieve a list of AttackResult objects based on the specified filters. @@ -1750,6 +1751,10 @@ def get_attack_results( identifier_filters (Optional[Sequence[IdentifierFilter]], optional): A sequence of IdentifierFilter objects that allows filtering by various attack identifier JSON properties. Defaults to None. + scenario_result_id (Optional[str], optional): Filter to attack results linked to a + specific scenario via the ``AttackResultEntry.scenario_result_id`` FK. Combined + with ``outcome=AttackOutcome.ERROR`` this is the replacement for the removed + per-scenario error_attack_result_ids manifest. Defaults to None. Returns: Sequence[AttackResult]: A list of AttackResult objects that match the specified filters. @@ -1783,6 +1788,8 @@ def get_attack_results( conditions.append(AttackResultEntry.objective.contains(objective)) if outcome: conditions.append(AttackResultEntry.outcome == outcome) + if scenario_result_id: + conditions.append(AttackResultEntry.scenario_result_id == uuid.UUID(scenario_result_id)) if attack_classes: # Case-insensitive to mirror converter_classes; forgives casing drift in @@ -1966,55 +1973,31 @@ def add_attack_results_to_scenario( attack_results: Sequence[AttackResult], ) -> bool: """ - Add attack results to an existing scenario result in memory. + No-op shim retained for backward compatibility. - This method efficiently updates a scenario result by appending new attack results - to a specific atomic attack name without requiring a full retrieve-modify-save cycle. + Scenario→AttackResult linkage is now stamped per-result by the attack + persistence path (via ``ExecutionAttribution`` on the + ``AttackContext``). This method is kept for one release for external + callers; it will be removed. Args: - scenario_result_id (str): The ID of the scenario result to update. - atomic_attack_name (str): The name of the atomic attack to add results for. - attack_results (Sequence[AttackResult]): The attack results to add. + scenario_result_id (str): The ID of the scenario result (ignored). + atomic_attack_name (str): The atomic attack name (ignored). + attack_results (Sequence[AttackResult]): The attack results (ignored). Returns: - bool: True if the update was successful, False otherwise. - - Example: - >>> memory.add_attack_results_to_scenario( - ... scenario_result_id="123e4567-e89b-12d3-a456-426614174000", - ... atomic_attack_name="base64_attack", - ... attack_results=[result1, result2] - ... ) - """ - try: - # Retrieve current scenario result - scenario_results = self.get_scenario_results(scenario_result_ids=[scenario_result_id]) - - if not scenario_results: - logger.error(f"Scenario result with ID {scenario_result_id} not found in memory") - return False - - scenario_result = scenario_results[0] - - # Update attack results for this atomic attack name - if atomic_attack_name not in scenario_result.attack_results: - scenario_result.attack_results[atomic_attack_name] = [] - - scenario_result.attack_results[atomic_attack_name].extend(list(attack_results)) - - # Save updated result back to memory using update - entry = ScenarioResultEntry(entry=scenario_result) - self._update_entry(entry) - - logger.info( - f"Added {len(attack_results)} attack results to scenario {scenario_result_id} " - f"for atomic attack '{atomic_attack_name}'" - ) - return True - - except Exception as e: - logger.exception(f"Failed to add attack results to scenario {scenario_result_id}: {str(e)}") - raise + bool: Always True (no-op success). + """ + print_deprecation_message( + old_item="memory.add_attack_results_to_scenario(...)", + new_item=( + "Scenario→AttackResult linkage is stamped per-result by the attack persistence path. " + "Setting Scenario._scenario_result_id on each AtomicAttack is sufficient; no explicit " + "post-run bulk-write is needed." + ), + removed_in="0.16.0", + ) + return True def update_scenario_run_state( self, @@ -2027,6 +2010,13 @@ def update_scenario_run_state( """ Update the run state of an existing scenario result. + Performs a targeted UPDATE of only the state/error columns instead of + rebuilding the entire ``ScenarioResultEntry`` row. The full-row rebuild + used to read the stored row, mutate the ScenarioResult, and re-serialize + every column — including ``attack_results_json`` which is being phased + out and could be stale during the deprecation window. A targeted UPDATE + avoids clobbering manifest data and is also cheaper. + Args: scenario_result_id (str): The ID of the scenario result to update. scenario_run_state (str): The new state for the scenario @@ -2037,61 +2027,21 @@ def update_scenario_run_state( Raises: ValueError: If the scenario result is not found. """ - scenario_results = self.get_scenario_results(scenario_result_ids=[scenario_result_id]) - - if not scenario_results: - raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory") - - scenario_result = scenario_results[0] - - # Update the scenario run state - scenario_result.scenario_run_state = scenario_run_state # type: ignore[ty:invalid-assignment] - - if error_message is not None: - scenario_result.error_message = error_message - if error_type is not None: - scenario_result.error_type = error_type - - # Save updated result back to memory using update - entry = ScenarioResultEntry(entry=scenario_result) - self._update_entry(entry) - - logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state}'") - - def update_scenario_error_attacks(self, *, scenario_result_id: str, error_attack_result_ids: list[str]) -> None: - """ - Update the error attack result IDs on an existing scenario result. - - This links failed AttackResults to the ScenarioResult so the REST API - can quickly find error details without scanning all attacks. - - Performs the read-modify-write within a single DB session to avoid - inter-session consistency issues. - - Args: - scenario_result_id: The ID of the scenario result to update. - error_attack_result_ids: IDs of AttackResults that contain error information. - - Raises: - ValueError: If the scenario result is not found. - """ - import json - with closing(self.get_session()) as session: entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first() if not entry: raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory") - existing: list[str] = ( - json.loads(entry.error_attack_result_ids_json) if entry.error_attack_result_ids_json else [] - ) - merged = list(dict.fromkeys(existing + error_attack_result_ids)) - entry.error_attack_result_ids_json = json.dumps(merged) + entry.scenario_run_state = scenario_run_state + if error_message is not None: + entry.error_message = error_message + if error_type is not None: + entry.error_type = error_type session.commit() - logger.info(f"Updated scenario {scenario_result_id} with {len(error_attack_result_ids)} error attack result(s)") + logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state}'") def get_scenario_results( self, @@ -2218,36 +2168,83 @@ def get_scenario_results( limit=limit, ) - # Convert entries to ScenarioResults and populate attack_results efficiently + # Convert entries to ScenarioResults and populate attack_results. + # + # Hydration strategy: + # 1. Pull AttackResultEntry rows for this scenario via the new + # scenario_result_id FK (the post-redesign source of truth). + # 2. Fall back to the legacy manifest path for any conversation_id + # present in attack_results_json but NOT in the FK result set. + # This handles partial-migration / pre-FK data. + # 3. Group by atomic_attack_name (from scenario_data for FK rows, + # from manifest keys for legacy rows) and sort each group by + # scenario_data["objective_index"] when available. scenario_results = [] for entry in entries: scenario_result = entry.get_scenario_result() + scenario_id_str = str(entry.id) - # Get conversation IDs grouped by attack name - conversation_ids_by_attack = entry.get_conversation_ids_by_attack_name() - - # Collect all conversation IDs to query in a single batch - all_conversation_ids = [] - for conv_ids in conversation_ids_by_attack.values(): - all_conversation_ids.extend(conv_ids) - - # Query all AttackResults using batched queries if needed - if all_conversation_ids: - attack_entries = self._execute_batched_query( + # FK-linked rows (preferred source). + fk_rows = self._query_entries( + AttackResultEntry, + conditions=AttackResultEntry.scenario_result_id == entry.id, + ) + fk_results: list[AttackResult] = [row.get_attack_result() for row in fk_rows] + fk_attack_result_ids = {r.attack_result_id for r in fk_results} + + # Legacy manifest path — only fill in conversation_ids not + # already covered by the FK path. + manifest_map = entry.get_conversation_ids_by_attack_name() + manifest_missing_ids: list[str] = [] + for conv_ids in manifest_map.values(): + manifest_missing_ids.extend(conv_ids) + + legacy_results_by_conv: dict[str, AttackResult] = {} + if manifest_missing_ids: + legacy_rows = self._execute_batched_query( AttackResultEntry, batch_column=AttackResultEntry.conversation_id, - batch_values=all_conversation_ids, + batch_values=manifest_missing_ids, + ) + for row in legacy_rows: + result = row.get_attack_result() + if result.attack_result_id in fk_attack_result_ids: + continue # already represented by the FK path + legacy_results_by_conv[row.conversation_id] = result + + # Group: FK rows by scenario_data["atomic_attack_name"], legacy + # rows by manifest key. Sort each group by objective_index when + # available (FK rows) and append legacy entries after, preserving + # manifest order. + grouped: dict[str, list[tuple[int, AttackResult]]] = {} + for r in fk_results: + name = (r.scenario_data or {}).get("atomic_attack_name") + if not name: + continue + idx = (r.scenario_data or {}).get("objective_index") + sort_key = idx if isinstance(idx, int) else len(grouped.setdefault(name, [])) + grouped.setdefault(name, []).append((sort_key, r)) + + legacy_count = 0 + for attack_name, conv_ids in manifest_map.items(): + legacy_bucket = grouped.setdefault(attack_name, []) + base_offset = len(legacy_bucket) + for offset, cid in enumerate(conv_ids): + legacy_result = legacy_results_by_conv.get(cid) + if legacy_result is None: + continue + legacy_bucket.append((base_offset + offset, legacy_result)) + legacy_count += 1 + + if legacy_count or fk_results: + logger.debug( + f"Hydrated scenario {scenario_id_str}: " + f"{len(fk_results)} via FK, {legacy_count} via legacy manifest fallback" ) - # Build a dict for quick lookup - attack_results_dict = {entry.conversation_id: entry.get_attack_result() for entry in attack_entries} - - # Populate attack_results by attack name, preserving order - scenario_result.attack_results = {} - for attack_name, conv_ids in conversation_ids_by_attack.items(): - scenario_result.attack_results[attack_name] = [ - attack_results_dict[conv_id] for conv_id in conv_ids if conv_id in attack_results_dict - ] + scenario_result.attack_results = { + name: [r for _, r in sorted(bucket, key=lambda kv: kv[0])] for name, bucket in grouped.items() + } scenario_results.append(scenario_result) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index b4a901b79b..672dc96e79 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -745,6 +745,18 @@ class AttackResultEntry(Base): retry_events_json: Mapped[Optional[str]] = mapped_column(Unicode, nullable=True) total_retries = mapped_column(INTEGER, nullable=True, default=0) + # Scenario linkage (set when the AttackResult is produced inside a Scenario). + # scenario_result_id is an indexed FK so per-scenario hydration and resume + # queries are direct lookups (no JSON manifest required, no orphaning if a + # scenario is interrupted mid-AtomicAttack). + # scenario_data is a documented-fixed-schema JSON blob with two keys: + # atomic_attack_name (str) and objective_index (int). When the AttackResult + # is created outside a Scenario both fields remain NULL. + scenario_result_id: Mapped[Optional[uuid.UUID]] = mapped_column( + CustomUUID, ForeignKey("ScenarioResultEntries.id", ondelete="SET NULL"), nullable=True, index=True + ) + scenario_data: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) + last_response: Mapped[Optional["PromptMemoryEntry"]] = relationship( "PromptMemoryEntry", foreign_keys=[last_response_id], @@ -819,6 +831,11 @@ def __init__(self, *, entry: AttackResult) -> None: ) self.total_retries = entry.total_retries + # Scenario linkage (set by the attack persistence path when an + # ExecutionAttribution is present on the AttackContext; otherwise None) + self.scenario_result_id = uuid.UUID(entry.scenario_result_id) if entry.scenario_result_id else None + self.scenario_data = entry.scenario_data + @staticmethod def _get_id_as_uuid(obj: Any) -> Optional[uuid.UUID]: """ @@ -931,6 +948,8 @@ def get_attack_result(self) -> AttackResult: error_traceback=self.error_traceback, retry_events=retry_events, total_retries=self.total_retries or 0, + scenario_result_id=str(self.scenario_result_id) if self.scenario_result_id else None, + scenario_data=self.scenario_data, ) @@ -992,9 +1011,6 @@ class ScenarioResultEntry(Base): completion_time = mapped_column(DateTime, nullable=False) timestamp = mapped_column(DateTime, nullable=False) - # Pointer to failed attack result(s) — avoids scanning all attacks for error info - error_attack_result_ids_json: Mapped[Optional[str]] = mapped_column(Unicode, nullable=True) - # Scenario-level error info (persisted so it survives process restarts) error_message: Mapped[Optional[str]] = mapped_column(Unicode, nullable=True) error_type: Mapped[Optional[str]] = mapped_column(String, nullable=True) @@ -1048,11 +1064,6 @@ def __init__(self, *, entry: ScenarioResult) -> None: # Serialize display_group_map if present self.display_group_map_json = json.dumps(entry._display_group_map) if entry._display_group_map else None - # Serialize error_attack_result_ids if present - self.error_attack_result_ids_json = ( - json.dumps(entry.error_attack_result_ids) if entry.error_attack_result_ids else None - ) - self.error_message = entry.error_message self.error_type = entry.error_type @@ -1097,11 +1108,6 @@ def get_scenario_result(self) -> ScenarioResult: if self.display_group_map_json: display_group_map = json.loads(self.display_group_map_json) - # Deserialize error_attack_result_ids if stored - error_attack_result_ids: list[str] | None = None - if self.error_attack_result_ids_json: - error_attack_result_ids = json.loads(self.error_attack_result_ids_json) - return ScenarioResult( id=self.id, scenario_identifier=scenario_identifier, @@ -1114,7 +1120,6 @@ def get_scenario_result(self) -> ScenarioResult: number_tries=self.number_tries, completion_time=self.completion_time, display_group_map=display_group_map, - error_attack_result_ids=error_attack_result_ids, error_message=self.error_message, error_type=self.error_type, ) diff --git a/pyrit/models/attack_result.py b/pyrit/models/attack_result.py index 123c83a918..46f8656250 100644 --- a/pyrit/models/attack_result.py +++ b/pyrit/models/attack_result.py @@ -107,6 +107,14 @@ class AttackResult(StrategyResult): retry_events: list[RetryEvent] = field(default_factory=list) total_retries: int = 0 + # Scenario linkage (infrastructure-managed). Set by the attack persistence + # path when an ExecutionAttribution is present on the AttackContext. User + # code should not set these directly; ad-hoc AttackResults created outside + # a Scenario leave both fields as None and the corresponding DB columns + # remain NULL. + scenario_result_id: str | None = None + scenario_data: dict[str, Any] | None = None + @property def attack_identifier(self) -> Optional[ComponentIdentifier]: """ diff --git a/pyrit/models/scenario_result.py b/pyrit/models/scenario_result.py index 6111998ff0..2177fafaa7 100644 --- a/pyrit/models/scenario_result.py +++ b/pyrit/models/scenario_result.py @@ -69,7 +69,6 @@ def __init__( number_tries: int = 0, id: uuid.UUID | None = None, # noqa: A002 display_group_map: dict[str, str] | None = None, - error_attack_result_ids: list[str] | None = None, error_message: str | None = None, error_type: str | None = None, ) -> None: @@ -91,9 +90,6 @@ def __init__( display_group_map (Optional[dict[str, str]]): Optional mapping of atomic_attack_name → display group label. Used by the console printer to aggregate results for user-facing output. - error_attack_result_ids (Optional[list[str]]): IDs of AttackResults that - contain error information. Used for quick error lookup without scanning - all attack results. error_message (Optional[str]): Scenario-level error message when the run fails. error_type (Optional[str]): Exception class name when the run fails. @@ -112,7 +108,6 @@ def __init__( self.completion_time = completion_time if completion_time is not None else datetime.now(timezone.utc) self.number_tries = number_tries self._display_group_map = display_group_map or {} - self.error_attack_result_ids = error_attack_result_ids or [] self.error_message = error_message self.error_type = error_type diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index d63c5e65f5..cdf6681f29 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -19,6 +19,7 @@ from pyrit.executor.attack import AttackExecutor, AttackStrategy from pyrit.executor.attack.core.attack_executor import AttackExecutorResult +from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution from pyrit.identifiers import build_atomic_attack_identifier from pyrit.identifiers.evaluation_identifier import AtomicAttackEvaluationIdentifier from pyrit.memory import CentralMemory @@ -117,10 +118,20 @@ def __init__( sg.validate() self._seed_groups = seed_groups + # Original positions for each currently-active seed group. Used to map a + # per-task input index (from the executor) back to the stable + # objective_index stored in AttackResultEntry.scenario_data. Resume + # filtering preserves the original indices so newly-persisted results + # don't collide with already-persisted ones. + self._original_indices: list[int] = list(range(len(seed_groups))) self._adversarial_chat = adversarial_chat self._objective_scorer = objective_scorer self._memory_labels = memory_labels or {} self._attack_execute_params = attack_execute_params + # Set by Scenario._execute_scenario_async before run_async. When set, + # each persisted AttackResult is linked to this scenario via the FK on + # AttackResultEntry. + self._scenario_result_id: str | None = None logger.info( f"Initialized atomic attack with {len(self._seed_groups)} seed groups, " @@ -156,15 +167,55 @@ def filter_seed_groups_by_objectives(self, *, remaining_objectives: list[str]) - """ Filter seed groups to only those with objectives in the remaining list. - This is used for scenario resumption to skip already completed objectives. + Deprecated — prefer ``filter_seed_groups_by_indices``. Filtering by + objective text collapses two distinct seed groups that happen to share + the same objective text (common when seed groups are programmatically + generated or when a baseline is paired with its strategy variants). + Index-based filtering is the stable identity for resume. Args: remaining_objectives (List[str]): List of objectives that still need to be executed. """ + warnings.warn( + "filter_seed_groups_by_objectives is deprecated; use filter_seed_groups_by_indices " + "for index-stable resume. Duplicate objective text can collapse otherwise-distinct " + "seed groups.", + DeprecationWarning, + stacklevel=2, + ) remaining_set = set(remaining_objectives) - self._seed_groups = [ - sg for sg in self._seed_groups if sg.objective is not None and sg.objective.value in remaining_set + kept: list[tuple[int, SeedAttackGroup]] = [ + (orig_idx, sg) + for orig_idx, sg in zip(self._original_indices, self._seed_groups, strict=True) + if sg.objective is not None and sg.objective.value in remaining_set + ] + self._original_indices = [idx for idx, _ in kept] + self._seed_groups = [sg for _, sg in kept] + + def filter_seed_groups_by_indices(self, *, completed_indices: set[int]) -> None: + """ + Filter out seed groups whose original index is in ``completed_indices``. + + This is the stable-identity replacement for + ``filter_seed_groups_by_objectives``. Each seed group's original index + (its position when the ``AtomicAttack`` was first constructed) is + tracked across filtering, so the executor can later map each per-task + input index back to the original index when stamping + ``scenario_data["objective_index"]``. This prevents newly-persisted + results from colliding with already-persisted ones on resume. + + Args: + completed_indices (set[int]): Original indices that have already + been executed (typically retrieved from + ``Scenario._get_completed_objectives_for_attack``). + """ + kept: list[tuple[int, SeedAttackGroup]] = [ + (orig_idx, sg) + for orig_idx, sg in zip(self._original_indices, self._seed_groups, strict=True) + if orig_idx not in completed_indices ] + self._original_indices = [idx for idx, _ in kept] + self._seed_groups = [sg for _, sg in kept] async def run_async( self, @@ -220,6 +271,24 @@ async def run_async( else: execution_seed_groups = self._seed_groups + # Build an attribution factory when this atomic attack is being + # executed inside a Scenario. The factory maps each per-task input + # index (position in the currently-active seed_groups list) back to + # the original objective_index so resume can locate already-done + # work and newly-persisted results don't collide with old ones. + attribution_factory = None + if self._scenario_result_id is not None: + scenario_id = self._scenario_result_id + name = self.atomic_attack_name + original_indices = list(self._original_indices) + + def attribution_factory(input_index: int) -> ExecutionAttribution: + return ExecutionAttribution( + scenario_result_id=scenario_id, + atomic_attack_name=name, + objective_index=original_indices[input_index], + ) + results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, seed_groups=execution_seed_groups, @@ -227,6 +296,7 @@ async def run_async( objective_scorer=self._objective_scorer, memory_labels=self._memory_labels, return_partial_on_failure=return_partial_on_failure, + attribution_factory=attribution_factory, **self._attack_execute_params, ) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index b82f16af90..000f75d850 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -28,7 +28,7 @@ from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory import CentralMemory from pyrit.memory.memory_models import ScenarioResultEntry -from pyrit.models import AttackResult, SeedAttackGroup +from pyrit.models import AttackOutcome, AttackResult, SeedAttackGroup from pyrit.models.scenario_result import ScenarioIdentifier, ScenarioResult from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements @@ -699,6 +699,21 @@ async def initialize_async( seed_groups = self._dataset_config.get_all_seed_attack_groups() self._atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=seed_groups)) + # Enforce atomic_attack_name uniqueness. Duplicate names cause result + # aggregation maps keyed by name to silently collapse rows, and + # FK-based resume cannot distinguish them. Some existing scenarios + # (e.g. Encoding) construct multiple AtomicAttacks with the same name + # for different converter configs, so this is a warning for now; in a + # future release it will become an error. + names = [aa.atomic_attack_name for aa in self._atomic_attacks] + duplicates = sorted({name for name in names if names.count(name) > 1}) + if duplicates: + logger.warning( + f"Scenario '{self._name}' has duplicate atomic_attack_name values: {duplicates}. " + f"Duplicate names will be collapsed in result aggregation and resume tracking. " + f"Each AtomicAttack within a scenario should have a unique name." + ) + # Store original objectives for each atomic attack (before any mutations during execution) self._original_objectives_map = { atomic_attack.atomic_attack_name: tuple(atomic_attack.objectives) for atomic_attack in self._atomic_attacks @@ -847,45 +862,52 @@ def _validate_stored_scenario(self, *, stored_result: ScenarioResult) -> None: f"(ID: {self._scenario_result_id}, state: {stored_result.scenario_run_state})" ) - def _get_completed_objectives_for_attack(self, *, atomic_attack_name: str) -> set[str]: + def _get_completed_objective_indices_for_attack(self, *, atomic_attack_name: str) -> set[int]: """ - Get the set of objectives that have already been completed for a specific atomic attack. + Get the set of objective_index values that have already been completed + (non-error) for a specific atomic attack inside this scenario. + + Queries AttackResultEntry rows directly by ``scenario_result_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. Args: - atomic_attack_name (str): The name of the atomic attack to check. + atomic_attack_name (str): The atomic attack name to scope to. Returns: - Set[str]: Set of objective strings that have been completed. + Set[int]: Original objective_index values that completed without error. """ if not self._scenario_result_id: return set() - completed_objectives: set[str] = set() - + completed_indices: set[int] = set() try: - # Retrieve the scenario result from memory - scenario_results = self._memory.get_scenario_results(scenario_result_ids=[self._scenario_result_id]) - - if scenario_results: - scenario_result = scenario_results[0] - # Get completed objectives for this atomic attack name - if atomic_attack_name in scenario_result.attack_results: - completed_objectives = { - result.objective for result in scenario_result.attack_results[atomic_attack_name] - } + 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.scenario_data is None: + continue + if row.scenario_data.get("atomic_attack_name") != atomic_attack_name: + continue + objective_index = row.scenario_data.get("objective_index") + if isinstance(objective_index, int): + completed_indices.add(objective_index) except Exception as e: logger.warning( - f"Failed to retrieve completed objectives for atomic attack '{atomic_attack_name}': {str(e)}" + f"Failed to retrieve completed objective indices for atomic attack '{atomic_attack_name}': {str(e)}" ) - return completed_objectives + return completed_indices async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: """ Get the list of atomic attacks that still have objectives to complete. - This method filters out atomic attacks where all objectives have been completed, - and updates the objectives list for atomic attacks that are partially complete. + Uses ``objective_index`` (the original position of each seed group in the + atomic attack) for resume rather than objective text, so duplicate + objective text in two distinct seed groups doesn't collapse on resume. Returns: List[AtomicAttack]: List of atomic attacks with uncompleted objectives. @@ -897,64 +919,29 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: remaining_attacks: list[AtomicAttack] = [] for atomic_attack in self._atomic_attacks: - # Get completed objectives for this atomic attack name - completed_objectives = self._get_completed_objectives_for_attack( + completed_indices = self._get_completed_objective_indices_for_attack( atomic_attack_name=atomic_attack.atomic_attack_name ) - # Get ORIGINAL objectives (before any mutations) from stored map - original_objectives = self._original_objectives_map.get(atomic_attack.atomic_attack_name, ()) - - # Calculate remaining objectives - remaining_objectives = [obj for obj in original_objectives if obj not in completed_objectives] - - if remaining_objectives: - # If there are remaining objectives, update the atomic attack - if len(remaining_objectives) < len(original_objectives): + if completed_indices: + original_count = len(atomic_attack.seed_groups) + atomic_attack.filter_seed_groups_by_indices(completed_indices=completed_indices) + remaining_count = len(atomic_attack.seed_groups) + if remaining_count == 0: + logger.info( + f"Atomic attack '{atomic_attack.atomic_attack_name}' has all objectives completed, skipping" + ) + continue + if remaining_count < original_count: logger.info( f"Atomic attack '{atomic_attack.atomic_attack_name}' has " - f"{len(remaining_objectives)}/{len(original_objectives)} objectives remaining" + f"{remaining_count}/{original_count} objectives remaining" ) - # Update the objectives for this atomic attack to only include remaining ones - atomic_attack.filter_seed_groups_by_objectives(remaining_objectives=remaining_objectives) - remaining_attacks.append(atomic_attack) - else: - logger.info( - f"Atomic attack '{atomic_attack.atomic_attack_name}' has all objectives completed, skipping" - ) + remaining_attacks.append(atomic_attack) return remaining_attacks - async def _update_scenario_result_async( - self, *, atomic_attack_name: str, attack_results: list[AttackResult] - ) -> None: - """ - Update the scenario result in memory with new attack results (thread-safe). - - This method is thread-safe and can be called from parallel executions. - - Args: - atomic_attack_name (str): The name of the atomic attack. - attack_results (List[AttackResult]): The list of new attack results to add. - """ - if not self._scenario_result_id: - logger.warning("Cannot update scenario result: no scenario result ID available") - return - - async with self._result_lock: - success = self._memory.add_attack_results_to_scenario( - scenario_result_id=self._scenario_result_id, - atomic_attack_name=atomic_attack_name, - attack_results=attack_results, - ) - - if not success: - logger.error( - f"Failed to update scenario result with {len(attack_results)} results " - f"for atomic attack '{atomic_attack_name}'" - ) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: """ Build atomic attacks from the cross-product of selected techniques and datasets. @@ -1191,6 +1178,12 @@ async def _execute_scenario_async(self) -> ScenarioResult: ), start=completed_count + 1, ): + # Stamp the scenario id onto the atomic attack so each persisted + # AttackResult carries the FK linkage. This is what enables + # mid-run interruption recovery (results are visible without the + # post-atomic-attack bulk manifest write). + atomic_attack._scenario_result_id = scenario_result_id + logger.info( f"Executing atomic attack {i}/{len(self._atomic_attacks)} " f"('{atomic_attack.atomic_attack_name}') in scenario '{self._name}'" @@ -1202,12 +1195,8 @@ async def _execute_scenario_async(self) -> ScenarioResult: return_partial_on_failure=True, ) - # Always save completed results, even if some objectives didn't complete - if atomic_results.completed_results: - await self._update_scenario_result_async( - atomic_attack_name=atomic_attack.atomic_attack_name, - attack_results=atomic_results.completed_results, - ) + # Per-result scenario linkage is now stamped by the attack + # event handler at write time; no post-atomic bulk update. # Check if there were any incomplete objectives if atomic_results.has_incomplete: @@ -1224,19 +1213,11 @@ async def _execute_scenario_async(self) -> ScenarioResult: for obj, exc in atomic_results.incomplete_objectives: logger.error(f" Incomplete objective '{obj[:50]}...': {str(exc)}") - # Collect error attack result IDs from the exceptions - error_ids = [] - for _, exc in atomic_results.incomplete_objectives: - error_id = getattr(exc, "error_attack_result_id", None) - if error_id: - error_ids.append(error_id) - - # Link error attack results to the scenario result - if error_ids: - self._memory.update_scenario_error_attacks( - scenario_result_id=scenario_result_id, - error_attack_result_ids=error_ids, - ) + # Error AttackResults are linked to this scenario via the + # scenario_result_id FK on AttackResultEntry (stamped by + # the attack event handler when an ExecutionAttribution + # is on the context). The previous per-scenario error_id + # manifest is no longer needed. # Mark scenario as failed error_msg = ( diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 29d2855cdb..645c6acf21 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -19,6 +19,7 @@ _DEFAULT_MAX_CONCURRENT_RUNS, ScenarioRunService, ) +from pyrit.models import AttackOutcome _REGISTRY_PATCH_BASE = "pyrit.registry" _MEMORY_PATCH = "pyrit.memory.CentralMemory.get_memory_instance" @@ -74,7 +75,6 @@ def _make_db_scenario_result( sr.display_group_map = {} sr.error_message = None sr.error_type = None - sr.error_attack_result_ids = [] return sr @@ -83,6 +83,9 @@ def mock_memory(): """Patch CentralMemory.get_memory_instance to return a mock.""" mock = MagicMock() mock.get_scenario_results.return_value = [] + # Default: no error AttackResults linked to any scenario. Tests that exercise + # the error fallback path explicitly set get_attack_results.return_value. + mock.get_attack_results.return_value = [] with patch(_MEMORY_PATCH, return_value=mock): yield mock @@ -298,9 +301,13 @@ def test_get_run_returns_existing_run(self, mock_memory) -> None: assert fetched.status == ScenarioRunStatus.IN_PROGRESS def test_get_run_falls_back_to_persisted_error(self, mock_memory) -> None: - """Test that get_run extracts error from persisted error AttackResult when no active task.""" + """Test that get_run extracts error from persisted error AttackResult when no active task. + + After the FK-based scenario linkage refactor, error AttackResults are + located via ``get_attack_results(scenario_result_id=..., outcome=ERROR)`` + rather than via a per-scenario error_attack_result_ids manifest. + """ db_result = _make_db_scenario_result(result_id="sr-fail", run_state="FAILED") - db_result.error_attack_result_ids = ["err-ar-1"] # Mock the error AttackResult lookup error_ar = MagicMock() @@ -315,7 +322,10 @@ def test_get_run_falls_back_to_persisted_error(self, mock_memory) -> None: assert fetched is not None assert fetched.error == "Connection refused" assert fetched.error_type == "ConnectionError" - mock_memory.get_attack_results.assert_called_once_with(attack_result_ids=["err-ar-1"]) + mock_memory.get_attack_results.assert_called_once_with( + scenario_result_id="sr-fail", + outcome=AttackOutcome.ERROR, + ) class TestScenarioRunServiceListRuns: diff --git a/tests/unit/executor/attack/core/test_attack_executor.py b/tests/unit/executor/attack/core/test_attack_executor.py index f949f69bfc..0b5d0ed55e 100644 --- a/tests/unit/executor/attack/core/test_attack_executor.py +++ b/tests/unit/executor/attack/core/test_attack_executor.py @@ -336,6 +336,112 @@ async def test_validates_explicit_empty_field_overrides_for_seed_groups(self): ) +@pytest.mark.usefixtures("patch_central_database") +class TestAttributionFactoryPropagation: + """Tests for ExecutionAttribution propagation through the AttackExecutor. + + The executor stamps ``context._attribution = attribution_factory(input_index)`` + on each per-task context. The input_index must be the seed group's original + position so attribution is deterministic and parallel-safe — never derived + from completion order. + """ + + async def test_attribution_factory_sets_per_task_attribution_in_order(self): + from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + + attack = create_mock_attack() + seen_indices: list[int] = [] + + async def capture(context): + # Read the attribution stamped by the executor BEFORE the task runs. + attr = getattr(context, "_attribution", None) + assert attr is not None + seen_indices.append(attr.objective_index) + return create_attack_result(context.params.objective) + + attack.execute_with_context_async = AsyncMock(side_effect=capture) + + seed_groups = [create_seed_group(f"obj-{i}") for i in range(4)] + + def factory(idx: int) -> ExecutionAttribution: + return ExecutionAttribution( + scenario_result_id="sid", + atomic_attack_name="atomic", + objective_index=idx, + ) + + executor = AttackExecutor(max_concurrency=1) + result = await executor.execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + attribution_factory=factory, + ) + + # Each per-task context saw a distinct, contiguous objective_index — the + # seed group's original input position. + assert sorted(seen_indices) == [0, 1, 2, 3] + assert len(result.completed_results) == 4 + + async def test_attribution_factory_parallel_safe_with_high_concurrency(self): + """At max_concurrency > 1, each task still receives its own objective_index + regardless of completion order. Out-of-order completion must not produce + stamped indices that point at the wrong seed group. + """ + from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + + attack = create_mock_attack() + seen: dict[str, int] = {} + + async def out_of_order(context): + attr = getattr(context, "_attribution", None) + assert attr is not None + # Reverse-delay the tasks so completion order is the inverse of input + # order. The stamped index must still match the seed group. + await asyncio.sleep(0.005 * (10 - attr.objective_index)) + seen[context.params.objective] = attr.objective_index + return create_attack_result(context.params.objective) + + attack.execute_with_context_async = AsyncMock(side_effect=out_of_order) + + seed_groups = [create_seed_group(f"obj-{i}") for i in range(6)] + + def factory(idx: int) -> ExecutionAttribution: + return ExecutionAttribution( + scenario_result_id="sid", + atomic_attack_name="atomic", + objective_index=idx, + ) + + executor = AttackExecutor(max_concurrency=6) + await executor.execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + attribution_factory=factory, + ) + + # Each objective is stamped with its OWN input index, regardless of + # completion order. + for i in range(6): + assert seen[f"obj-{i}"] == i + + async def test_no_attribution_factory_leaves_context_attribution_none(self): + attack = create_mock_attack() + + async def capture(context): + attr = getattr(context, "_attribution", None) + assert attr is None + return create_attack_result(context.params.objective) + + attack.execute_with_context_async = AsyncMock(side_effect=capture) + + seed_groups = [create_seed_group("obj-0"), create_seed_group("obj-1")] + executor = AttackExecutor(max_concurrency=2) + await executor.execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + ) + + @pytest.mark.usefixtures("patch_central_database") class TestPartialFailureHandling: """Tests for partial failure handling.""" diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 33370f0c9d..7d441b5550 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -610,6 +610,92 @@ async def test_on_event_handles_other_events(self, event_handler, sample_attack_ f"Attack is in '{StrategyEvent.ON_PRE_VALIDATE.value}' stage for {event_handler.__class__.__name__}" ) + async def test_on_post_execute_stamps_scenario_attribution_when_present( + self, sample_attack_context, sample_attack_result, mock_memory + ): + """When the context carries an ExecutionAttribution, the persisted + AttackResult must have scenario_result_id + scenario_data populated.""" + from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + + with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): + handler = _DefaultAttackStrategyEventHandler() + sample_attack_context.start_time = 100.0 + sample_attack_context._attribution = ExecutionAttribution( + scenario_result_id="scenario-1", + atomic_attack_name="atomic_a", + objective_index=3, + ) + + event_data = StrategyEventData( + event=StrategyEvent.ON_POST_EXECUTE, + strategy_name="TestStrategy", + strategy_id="test-id", + context=sample_attack_context, + result=sample_attack_result, + ) + await handler.on_event(event_data) + + assert sample_attack_result.scenario_result_id == "scenario-1" + assert sample_attack_result.scenario_data == { + "atomic_attack_name": "atomic_a", + "objective_index": 3, + } + + async def test_on_post_execute_no_attribution_leaves_fields_none( + self, sample_attack_context, sample_attack_result, mock_memory + ): + """Outside a Scenario, _attribution is None and the FK + scenario_data + fields on the persisted AttackResult must stay None.""" + with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): + handler = _DefaultAttackStrategyEventHandler() + sample_attack_context.start_time = 100.0 + # _attribution defaults to None — no scenario stamping should happen. + + event_data = StrategyEventData( + event=StrategyEvent.ON_POST_EXECUTE, + strategy_name="TestStrategy", + strategy_id="test-id", + context=sample_attack_context, + result=sample_attack_result, + ) + await handler.on_event(event_data) + + assert sample_attack_result.scenario_result_id is None + assert sample_attack_result.scenario_data is None + + async def test_on_error_stamps_scenario_attribution_when_present(self, sample_attack_context, mock_memory): + """Error AttackResults must also carry the scenario FK so error lookups + via get_attack_results(scenario_result_id=..., outcome=ERROR) work.""" + from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + + with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): + handler = _DefaultAttackStrategyEventHandler() + sample_attack_context.start_time = 100.0 + sample_attack_context._attribution = ExecutionAttribution( + scenario_result_id="scenario-err", + atomic_attack_name="atomic_err", + objective_index=7, + ) + + event_data = StrategyEventData( + event=StrategyEvent.ON_ERROR, + strategy_name="TestStrategy", + strategy_id="test-id", + context=sample_attack_context, + error=RuntimeError("boom"), + ) + await handler.on_event(event_data) + + # The error AttackResult was persisted; inspect what was sent to memory. + call = mock_memory.add_attack_results_to_memory.call_args + persisted = call.kwargs["attack_results"][0] + assert persisted.outcome == AttackOutcome.ERROR + assert persisted.scenario_result_id == "scenario-err" + assert persisted.scenario_data == { + "atomic_attack_name": "atomic_err", + "objective_index": 7, + } + @pytest.mark.usefixtures("patch_central_database") class TestAttackStrategyIntegration: diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index 97a240e37c..4ea6258a50 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -648,58 +648,225 @@ def test_combined_filters(sqlite_instance: MemoryInterface): assert "gpt-4" in results[0].objective_target_identifier.params["model_name"] -def test_update_scenario_error_attacks_success(sqlite_instance: MemoryInterface, sample_attack_results): - """Test successfully linking error attack result IDs to a scenario result.""" +# ============================================================================= +# Scenario linkage (FK + scenario_data on AttackResultEntry) hydration tests +# ============================================================================= + + +def _make_attack_result_for_scenario( + *, + scenario_result_id, + atomic_attack_name, + objective_index, + conversation_id=None, + outcome=AttackOutcome.SUCCESS, +): + """Build an AttackResult pre-stamped with scenario linkage (mirrors what + the event handler does when an ExecutionAttribution is on the context).""" + return AttackResult( + conversation_id=conversation_id or f"conv-{atomic_attack_name}-{objective_index}", + objective=f"objective-{atomic_attack_name}-{objective_index}", + outcome=outcome, + executed_turns=1, + scenario_result_id=str(scenario_result_id), + scenario_data={"atomic_attack_name": atomic_attack_name, "objective_index": objective_index}, + ) + + +def test_get_scenario_results_hydrates_via_fk(sqlite_instance: MemoryInterface): + """When AttackResultEntry rows carry the scenario_result_id FK, hydration + picks them up directly — without needing the legacy attack_results_json + manifest. This is the path that makes mid-AtomicAttack interruption-recovery + work.""" scenario_result = create_scenario_result( - name="Error Scenario", - attack_results={"Attack1": [sample_attack_results[0]]}, + name="FK-only Scenario", + attack_results={}, # manifest intentionally empty ) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - error_ids = ["error-ar-1", "error-ar-2"] - sqlite_instance.update_scenario_error_attacks( - scenario_result_id=str(scenario_result.id), - error_attack_result_ids=error_ids, + sid = scenario_result.id + ar1 = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="a", objective_index=0) + ar2 = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="a", objective_index=1) + ar3 = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="b", objective_index=0) + sqlite_instance.add_attack_results_to_memory(attack_results=[ar1, ar2, ar3]) + + [result] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(sid)]) + assert set(result.attack_results.keys()) == {"a", "b"} + assert [r.conversation_id for r in result.attack_results["a"]] == [ + "conv-a-0", + "conv-a-1", + ] + assert [r.conversation_id for r in result.attack_results["b"]] == ["conv-b-0"] + + +def test_get_scenario_results_hydrates_via_legacy_manifest(sqlite_instance: MemoryInterface, sample_attack_results): + """When a scenario predates the FK column (its AttackResults have no FK set + but the manifest is populated), hydration falls back to the manifest path.""" + scenario_result = create_scenario_result( + name="Manifest-only Scenario", + attack_results={"legacy_attack": sample_attack_results[:2]}, ) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - # Verify the error IDs were persisted - results = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) - assert len(results) == 1 - assert results[0].error_attack_result_ids == error_ids + [result] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) + assert list(result.attack_results.keys()) == ["legacy_attack"] + assert len(result.attack_results["legacy_attack"]) == 2 -def test_update_scenario_error_attacks_appends_to_existing(sqlite_instance: MemoryInterface, sample_attack_results): - """Test that updating error attacks appends to existing IDs without duplicates.""" +def test_get_scenario_results_mixed_hydration_merges_and_dedupes( + sqlite_instance: MemoryInterface, sample_attack_results +): + """A scenario with BOTH FK-linked rows and legacy manifest rows merges + both sources. FK rows are authoritative; manifest entries already present + via FK are not duplicated.""" + import uuid as _uuid + from contextlib import closing + + from pyrit.memory.memory_models import AttackResultEntry + + # Build manifest from existing sample_attack_results. scenario_result = create_scenario_result( - name="Error Scenario", - attack_results={"Attack1": [sample_attack_results[0]]}, + name="Mixed Scenario", + attack_results={"legacy_attack": sample_attack_results[:2]}, ) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - # First update - sqlite_instance.update_scenario_error_attacks( - scenario_result_id=str(scenario_result.id), - error_attack_result_ids=["error-ar-1"], + sid = scenario_result.id + # Add a new FK-linked row not in the manifest. + fk_only = _make_attack_result_for_scenario( + scenario_result_id=sid, atomic_attack_name="fk_attack", objective_index=0 + ) + sqlite_instance.add_attack_results_to_memory(attack_results=[fk_only]) + + # Retroactively stamp ONE of the existing manifest rows with the FK to + # verify dedup-by-attack_result_id: the row appears in both sources but + # the hydration must yield exactly one copy. UPDATE in place rather than + # re-INSERT (which would trip the PK unique constraint). + target_id = _uuid.UUID(sample_attack_results[0].attack_result_id) + with closing(sqlite_instance.get_session()) as session: + row = session.query(AttackResultEntry).filter_by(id=target_id).one() + row.scenario_result_id = sid + row.scenario_data = {"atomic_attack_name": "legacy_attack", "objective_index": 0} + session.commit() + + [result] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(sid)]) + assert "fk_attack" in result.attack_results + assert "legacy_attack" in result.attack_results + # legacy_attack should have 2 rows (one upgraded to FK + one still manifest-only), + # with no duplicates from the merge. + legacy_ids = [r.attack_result_id for r in result.attack_results["legacy_attack"]] + assert len(legacy_ids) == 2 + assert len(set(legacy_ids)) == 2 + + +def test_get_attack_results_filters_by_scenario_result_id(sqlite_instance: MemoryInterface): + """get_attack_results gains a scenario_result_id filter — replaces the + removed error_attack_result_ids_json lookup path.""" + scenario_result = create_scenario_result(name="Filter Scenario") + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + sid = scenario_result.id + + ok = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="a", objective_index=0) + err = _make_attack_result_for_scenario( + scenario_result_id=sid, + atomic_attack_name="a", + objective_index=1, + outcome=AttackOutcome.ERROR, + ) + # An unrelated AttackResult NOT linked to this scenario should be excluded. + unrelated = create_attack_result("unrelated-conv", "unrelated-obj") + sqlite_instance.add_attack_results_to_memory(attack_results=[ok, err, unrelated]) + + all_for_scenario = sqlite_instance.get_attack_results(scenario_result_id=str(sid)) + assert {r.conversation_id for r in all_for_scenario} == {ok.conversation_id, err.conversation_id} + + only_errors = sqlite_instance.get_attack_results( + scenario_result_id=str(sid), + outcome=AttackOutcome.ERROR.value, ) + assert [r.conversation_id for r in only_errors] == [err.conversation_id] + + +def test_delete_scenario_sets_attack_result_fk_to_null(sqlite_instance: MemoryInterface): + """ON DELETE SET NULL: deleting the parent ScenarioResultEntry nulls the FK + on its linked AttackResultEntries but the AttackResultEntries survive + (scenario_data is retained as historical provenance). + + Note: SQLite does not enforce foreign keys by default; this test enables + them on the session for the duration of the delete to verify the + ON DELETE SET NULL clause works. Production deployments using SQL Server + enforce FKs by default. + """ + from contextlib import closing + + from sqlalchemy import text as _sql_text + + from pyrit.memory.memory_models import AttackResultEntry, ScenarioResultEntry - # Second update with overlap and new ID - sqlite_instance.update_scenario_error_attacks( + scenario_result = create_scenario_result(name="To Be Deleted") + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + sid = scenario_result.id + + ar = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="a", objective_index=0) + sqlite_instance.add_attack_results_to_memory(attack_results=[ar]) + + # Enable FKs for the delete and verify the SET NULL clause fires. + with closing(sqlite_instance.get_session()) as session: + session.execute(_sql_text("PRAGMA foreign_keys = ON")) + session.query(ScenarioResultEntry).filter_by(id=sid).delete() + session.commit() + + # The AttackResult survives, but its FK is now NULL. scenario_data is + # retained as historical provenance. + with closing(sqlite_instance.get_session()) as session: + entry = session.query(AttackResultEntry).filter_by(conversation_id=ar.conversation_id).one() + assert entry.scenario_result_id is None + assert entry.scenario_data == {"atomic_attack_name": "a", "objective_index": 0} + + +def test_add_attack_results_to_scenario_is_deprecated_noop(sqlite_instance: MemoryInterface): + """add_attack_results_to_scenario is now a deprecated no-op shim — it + returns True without writing anything. Linkage is stamped per-result by + the attack persistence path instead.""" + scenario_result = create_scenario_result(name="Shim") + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + + # Should NOT raise and should return True (no-op success). + result = sqlite_instance.add_attack_results_to_scenario( scenario_result_id=str(scenario_result.id), - error_attack_result_ids=["error-ar-1", "error-ar-2"], + atomic_attack_name="ignored", + attack_results=[], ) + assert result is True + + # And no rows were added. + [hydrated] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) + assert hydrated.attack_results == {} - results = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) - # Should be deduplicated: ["error-ar-1", "error-ar-2"] - assert results[0].error_attack_result_ids == ["error-ar-1", "error-ar-2"] +def test_update_scenario_run_state_targeted_update_preserves_manifest(sqlite_instance: MemoryInterface): + """update_scenario_run_state must be a targeted UPDATE — it must not + re-serialize the whole row and clobber the manifest column during the + deprecation window.""" + scenario_result = create_scenario_result( + name="Targeted Update", + attack_results={"a": []}, # baseline manifest + ) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + sid = str(scenario_result.id) + + sqlite_instance.update_scenario_run_state( + scenario_result_id=sid, + scenario_run_state="FAILED", + error_message="boom", + error_type="RuntimeError", + ) -def test_update_scenario_error_attacks_not_found(sqlite_instance: MemoryInterface): - """Test that updating a nonexistent scenario result raises ValueError.""" - with pytest.raises(ValueError, match="not found in memory"): - sqlite_instance.update_scenario_error_attacks( - scenario_result_id="nonexistent-id", - error_attack_result_ids=["error-ar-1"], - ) + # State and error fields updated. + [hydrated] = sqlite_instance.get_scenario_results(scenario_result_ids=[sid]) + assert hydrated.scenario_run_state == "FAILED" + assert hydrated.error_message == "boom" + assert hydrated.error_type == "RuntimeError" def test_get_scenario_results_by_target_identifier_filter_hash(sqlite_instance: MemoryInterface): diff --git a/tests/unit/memory/test_memory_models.py b/tests/unit/memory/test_memory_models.py index b382bf8646..f143a69e98 100644 --- a/tests/unit/memory/test_memory_models.py +++ b/tests/unit/memory/test_memory_models.py @@ -474,17 +474,3 @@ def test_init_with_empty_attack_results(self): entry = ScenarioResultEntry(entry=sr) conv_ids = entry.get_conversation_ids_by_attack_name() assert conv_ids == {} - - def test_roundtrip_error_attack_result_ids(self): - sr = self._make_scenario_result(error_attack_result_ids=["err-1", "err-2"]) - entry = ScenarioResultEntry(entry=sr) - assert entry.error_attack_result_ids_json is not None - recovered = entry.get_scenario_result() - assert recovered.error_attack_result_ids == ["err-1", "err-2"] - - def test_roundtrip_error_attack_result_ids_none(self): - sr = self._make_scenario_result() - entry = ScenarioResultEntry(entry=sr) - assert entry.error_attack_result_ids_json is None - recovered = entry.get_scenario_result() - assert recovered.error_attack_result_ids == [] diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 0140ce5b1a..7a7d94e655 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -205,6 +205,261 @@ def test_migration_downgrade_creates_proper_structure(): engine.dispose() +# ============================================================================= +# Backfill tests for the scenario_result_id FK migration +# ============================================================================= + + +_SCENARIO_LINKAGE_REV = "9c8b7a6d5e4f" +_PREV_REV = "7a1b2c3d4e5f" + + +def _seed_pre_migration_scenario(connection, *, scenario_id, manifest_json): + """Insert a ScenarioResultEntry row at the pre-migration revision.""" + connection.execute( + text( + 'INSERT INTO "ScenarioResultEntries" ' + "(id, scenario_name, scenario_description, scenario_version, pyrit_version, " + "objective_target_identifier, scenario_run_state, attack_results_json, " + "number_tries, completion_time, timestamp) " + "VALUES (:id, :name, '', 1, '0.14.0.dev0', '{}', 'COMPLETED', :manifest, 0, '2026-05-18', '2026-05-18')" + ), + {"id": scenario_id, "name": "Backfill Test", "manifest": manifest_json}, + ) + + +def _seed_pre_migration_attack_result(connection, *, attack_id, conversation_id): + """Insert an AttackResultEntry row at the pre-migration revision.""" + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, attack_identifier, objective_sha256, executed_turns, " + "execution_time_ms, outcome, timestamp) " + "VALUES (:id, :conv, 'obj', '{}', 'sha', 1, 0, 'success', '2026-05-18')" + ), + {"id": attack_id, "conv": conversation_id}, + ) + + +def _config_for(connection): + pyrit_root = Path(__file__).resolve().parent.parent.parent.parent / "pyrit" + script_location = pyrit_root / "memory" / "alembic" + config = Config() + config.set_main_option("script_location", str(script_location)) + config.attributes["connection"] = connection + config.attributes["version_table"] = "pyrit_memory_alembic_version" + return config + + +def test_backfill_links_attack_results_via_conversation_id(): + """Upgrading from the pre-FK revision backfills scenario_result_id + + scenario_data on AttackResultEntries by matching conversation_id.""" + import json + + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "backfill-test.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + sid = str(uuid.uuid4()) + ar1_id = str(uuid.uuid4()) + ar2_id = str(uuid.uuid4()) + ar3_id = str(uuid.uuid4()) + + with engine.begin() as connection: + config = _config_for(connection) + # Step the schema up to JUST before the linkage migration. + command.upgrade(config, _PREV_REV) + + _seed_pre_migration_attack_result(connection, attack_id=ar1_id, conversation_id="conv-a-0") + _seed_pre_migration_attack_result(connection, attack_id=ar2_id, conversation_id="conv-a-1") + _seed_pre_migration_attack_result(connection, attack_id=ar3_id, conversation_id="conv-b-0") + _seed_pre_migration_scenario( + connection, + scenario_id=sid, + manifest_json=json.dumps({"a": ["conv-a-0", "conv-a-1"], "b": ["conv-b-0"]}), + ) + + command.upgrade(config, _SCENARIO_LINKAGE_REV) + + rows = connection.execute( + text( + "SELECT conversation_id, scenario_result_id, scenario_data " + 'FROM "AttackResultEntries" ORDER BY conversation_id' + ) + ).fetchall() + + results_by_conv = {r[0]: (r[1], r[2]) for r in rows} + + # All three rows now point at the scenario via the new FK. + for conv in ("conv-a-0", "conv-a-1", "conv-b-0"): + assert results_by_conv[conv][0] == sid, f"{conv} should be backfilled" + + # scenario_data carries atomic_attack_name and the 0-based manifest index. + sd_a0 = json.loads(results_by_conv["conv-a-0"][1]) + sd_a1 = json.loads(results_by_conv["conv-a-1"][1]) + sd_b0 = json.loads(results_by_conv["conv-b-0"][1]) + + assert sd_a0 == {"atomic_attack_name": "a", "objective_index": 0} + assert sd_a1 == {"atomic_attack_name": "a", "objective_index": 1} + assert sd_b0 == {"atomic_attack_name": "b", "objective_index": 0} + finally: + engine.dispose() + + +def test_backfill_is_idempotent_and_does_not_clobber_existing_linkage(): + """The backfill is safe to re-run: rows that already carry a + ``scenario_result_id`` are not overwritten (the WHERE IS NULL guard). We + verify by upgrading, manually retargeting a row, then downgrading + + re-upgrading and asserting the manual retarget survives.""" + import json + + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "idempotent.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + sid_old = str(uuid.uuid4()) + sid_manual = str(uuid.uuid4()) + ar_id = str(uuid.uuid4()) + + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _PREV_REV) + _seed_pre_migration_attack_result(connection, attack_id=ar_id, conversation_id="conv-shared") + _seed_pre_migration_scenario( + connection, scenario_id=sid_old, manifest_json=json.dumps({"a": ["conv-shared"]}) + ) + command.upgrade(config, _SCENARIO_LINKAGE_REV) + + # Manually retarget the row to a DIFFERENT scenario_result_id — + # simulate code that already linked it post-upgrade. + connection.execute( + text('UPDATE "AttackResultEntries" SET scenario_result_id = :sid WHERE conversation_id = :conv'), + {"sid": sid_manual, "conv": "conv-shared"}, + ) + + # Downgrade then re-upgrade to re-run the backfill. + command.downgrade(config, _PREV_REV) + + # After downgrade the FK column is gone, but the manifest still + # references conv-shared. On re-upgrade, the backfill should + # NOT clobber sid_manual because the column was just re-added + # as NULL — actually downgrade DROPS the column data, so on + # re-upgrade the row will start at NULL and get linked again. + # The test we want is: re-running the backfill while a row + # already has a non-NULL FK does not overwrite it. We exercise + # that with a fresh second upgrade-then-no-op-re-upgrade. + command.upgrade(config, _SCENARIO_LINKAGE_REV) + + # First upgrade after downgrade re-links it to sid_old (the + # manifest source). Now manually retarget again. + connection.execute( + text('UPDATE "AttackResultEntries" SET scenario_result_id = :sid WHERE conversation_id = :conv'), + {"sid": sid_manual, "conv": "conv-shared"}, + ) + + # Stamping should NOT happen on re-invocation since the column + # is already non-NULL. We verify by re-running the backfill + # logic via downgrade+upgrade is NOT what we want here — we + # want the IS NULL guard. Simulate by adding another scenario + # referencing the same conversation_id and re-running the + # backfill function only. + connection.execute( + text( + 'INSERT INTO "ScenarioResultEntries" ' + "(id, scenario_name, scenario_description, scenario_version, pyrit_version, " + "objective_target_identifier, scenario_run_state, attack_results_json, " + "number_tries, completion_time, timestamp) " + "VALUES (:id, 'Other', '', 1, '0.14.0.dev0', '{}', 'COMPLETED', :manifest, 0, " + "'2026-05-18', '2026-05-18')" + ), + {"id": str(uuid.uuid4()), "manifest": json.dumps({"x": ["conv-shared"]})}, + ) + + # Manually call the backfill function (loaded via the alembic + # script directory — modules with leading-digit filenames are + # not importable through normal Python import). + from importlib.util import module_from_spec, spec_from_file_location + + migration_path = ( + Path(__file__).resolve().parent.parent.parent.parent + / "pyrit" + / "memory" + / "alembic" + / "versions" + / "9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py" + ) + spec = spec_from_file_location("scenario_linkage_migration", migration_path) + assert spec is not None and spec.loader is not None + mig = module_from_spec(spec) + spec.loader.exec_module(mig) + + from alembic import op as _op_mod + + _original_get_bind = _op_mod.get_bind + _op_mod.get_bind = lambda: connection + try: + mig._backfill_scenario_linkage() + finally: + _op_mod.get_bind = _original_get_bind + + # The row's manual retarget MUST survive — the IS NULL guard + # prevents the backfill from overwriting it. + row = connection.execute( + text('SELECT scenario_result_id FROM "AttackResultEntries" WHERE conversation_id = :conv'), + {"conv": "conv-shared"}, + ).scalar_one() + assert row == sid_manual + finally: + engine.dispose() + + +def test_migration_drops_error_attack_result_ids_json_column(): + """The not-yet-released error_attack_result_ids_json column is removed + in this migration (no deprecation window needed).""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "drop-col.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _PREV_REV) + cols_before = {c["name"] for c in inspect(connection).get_columns("ScenarioResultEntries")} + assert "error_attack_result_ids_json" in cols_before + + command.upgrade(config, _SCENARIO_LINKAGE_REV) + cols_after = {c["name"] for c in inspect(connection).get_columns("ScenarioResultEntries")} + assert "error_attack_result_ids_json" not in cols_after + finally: + engine.dispose() + + +def test_migration_downgrade_restores_dropped_column(): + """Downgrading from the linkage revision re-adds error_attack_result_ids_json + and removes the new AttackResultEntries columns.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "downgrade-linkage.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _SCENARIO_LINKAGE_REV) + + attack_cols_up = {c["name"] for c in inspect(connection).get_columns("AttackResultEntries")} + assert "scenario_result_id" in attack_cols_up + assert "scenario_data" in attack_cols_up + + command.downgrade(config, _PREV_REV) + + attack_cols_down = {c["name"] for c in inspect(connection).get_columns("AttackResultEntries")} + assert "scenario_result_id" not in attack_cols_down + assert "scenario_data" not in attack_cols_down + + scenario_cols = {c["name"] for c in inspect(connection).get_columns("ScenarioResultEntries")} + assert "error_attack_result_ids_json" in scenario_cols + finally: + engine.dispose() + + def test_check_schema_migrations_calls_alembic_check(): with tempfile.TemporaryDirectory() as temp_dir: db_path = os.path.join(temp_dir, "check-test.db") diff --git a/tests/unit/models/test_scenario_result.py b/tests/unit/models/test_scenario_result.py index 02af031429..68a1cef1a9 100644 --- a/tests/unit/models/test_scenario_result.py +++ b/tests/unit/models/test_scenario_result.py @@ -165,24 +165,3 @@ def test_normalize_scenario_name_already_pascal(self): def test_normalize_scenario_name_mixed_case_with_underscore(self): assert ScenarioResult.normalize_scenario_name("Content_harms") == "Content_harms" - - def test_error_attack_result_ids_defaults_to_empty(self): - """error_attack_result_ids defaults to empty list.""" - sr = ScenarioResult( - scenario_identifier=_make_scenario_identifier(), - objective_target_identifier=ComponentIdentifier.from_dict({}), - attack_results={}, - objective_scorer_identifier=ComponentIdentifier.from_dict({}), - ) - assert sr.error_attack_result_ids == [] - - def test_error_attack_result_ids_stored(self): - """error_attack_result_ids are stored correctly.""" - sr = ScenarioResult( - scenario_identifier=_make_scenario_identifier(), - objective_target_identifier=ComponentIdentifier.from_dict({}), - attack_results={}, - objective_scorer_identifier=ComponentIdentifier.from_dict({}), - error_attack_result_ids=["id-1", "id-2"], - ) - assert sr.error_attack_result_ids == ["id-1", "id-2"] diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/test_scenario.py index e7042183d6..df234dcbce 100644 --- a/tests/unit/scenario/test_scenario.py +++ b/tests/unit/scenario/test_scenario.py @@ -29,11 +29,36 @@ def save_attack_results_to_memory(attack_results): memory.add_attack_results_to_memory(attack_results=attack_results) -def create_mock_run_async(attack_results): - """Create a mock run_async that saves results to memory before returning.""" +def _stamp_scenario_linkage(*, attack_results, atomic_attack): + """ + Stamp scenario_result_id + scenario_data on each AttackResult the same way + the real attack persistence path does. Mirrors what + ``_DefaultAttackStrategyEventHandler._stamp_attribution`` does at runtime + so test fixtures that mock out the executor still produce DB rows the new + FK-based hydration can find. + """ + sid = getattr(atomic_attack, "_scenario_result_id", None) + name = getattr(atomic_attack, "atomic_attack_name", None) + if not sid or not name: + return + for i, r in enumerate(attack_results): + r.scenario_result_id = sid + r.scenario_data = {"atomic_attack_name": name, "objective_index": i} + + +def create_mock_run_async(attack_results, *, atomic_attack=None): + """ + Create a mock ``run_async`` that stamps + saves results to memory. + + Pass ``atomic_attack`` (the AtomicAttack MagicMock) so the helper can copy + its ``_scenario_result_id`` (set by ``Scenario._execute_scenario_async``) + and ``atomic_attack_name`` onto each result. Without those the FK-based + hydration in ``get_scenario_results`` won't see the rows. + """ async def mock_run_async(*args, **kwargs): - # Save results to memory (mimics what real attacks do) + if atomic_attack is not None: + _stamp_scenario_linkage(attack_results=attack_results, atomic_attack=atomic_attack) save_attack_results_to_memory(attack_results) return AttackExecutorResult(completed_results=attack_results, incomplete_objectives=[]) @@ -327,7 +352,7 @@ async def test_run_async_executes_all_runs(self, mock_atomic_attacks, sample_att """Test that run_async executes all atomic attacks sequentially.""" # Configure each run to return different results for i, run in enumerate(mock_atomic_attacks): - run.run_async = create_mock_run_async([sample_attack_results[i]]) + run.run_async = create_mock_run_async([sample_attack_results[i]], atomic_attack=run) scenario = ConcreteScenario( name="Test Scenario", @@ -359,7 +384,7 @@ async def test_run_async_with_custom_concurrency( ): """Test that max_concurrency from init is passed to each atomic attack.""" for i, run in enumerate(mock_atomic_attacks): - run.run_async = create_mock_run_async([sample_attack_results[i]]) + run.run_async = create_mock_run_async([sample_attack_results[i]], atomic_attack=run) scenario = ConcreteScenario( name="Test Scenario", @@ -383,9 +408,15 @@ async def test_run_async_aggregates_multiple_results( ): """Test that results from multiple atomic attacks are properly aggregated.""" # Configure runs to return different numbers of results - mock_atomic_attacks[0].run_async = create_mock_run_async(sample_attack_results[0:2]) - mock_atomic_attacks[1].run_async = create_mock_run_async(sample_attack_results[2:4]) - mock_atomic_attacks[2].run_async = create_mock_run_async(sample_attack_results[4:5]) + mock_atomic_attacks[0].run_async = create_mock_run_async( + sample_attack_results[0:2], atomic_attack=mock_atomic_attacks[0] + ) + mock_atomic_attacks[1].run_async = create_mock_run_async( + sample_attack_results[2:4], atomic_attack=mock_atomic_attacks[1] + ) + mock_atomic_attacks[2].run_async = create_mock_run_async( + sample_attack_results[4:5], atomic_attack=mock_atomic_attacks[2] + ) scenario = ConcreteScenario( name="Test Scenario", @@ -441,7 +472,7 @@ async def test_run_async_returns_scenario_result_with_identifier( ): """Test that run_async returns ScenarioResult with proper identifier.""" for i, run in enumerate(mock_atomic_attacks): - run.run_async = create_mock_run_async([sample_attack_results[i]]) + run.run_async = create_mock_run_async([sample_attack_results[i]], atomic_attack=run) scenario = ConcreteScenario( name="Test Scenario", @@ -779,7 +810,9 @@ async def test_baseline_only_execution_runs_successfully(self, mock_objective_ta ) # Mock the baseline attack's run_async - scenario._atomic_attacks[0].run_async = create_mock_run_async([sample_attack_results[0]]) + scenario._atomic_attacks[0].run_async = create_mock_run_async( + [sample_attack_results[0]], atomic_attack=scenario._atomic_attacks[0] + ) # Run the scenario result = await scenario.run_async() diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/test_scenario_partial_results.py index a18625dc8b..223b435c8f 100644 --- a/tests/unit/scenario/test_scenario_partial_results.py +++ b/tests/unit/scenario/test_scenario_partial_results.py @@ -35,8 +35,20 @@ def mock_objective_target(): return target -def save_attack_results_to_memory(attack_results): - """Helper function to save attack results to memory.""" +def save_attack_results_to_memory(attack_results, *, atomic_attack=None): + """ + Helper function to save attack results to memory. When ``atomic_attack`` is + provided, also stamps ``scenario_result_id`` and ``scenario_data`` on each + result the same way the real attack persistence path does — so FK-based + hydration in ``get_scenario_results`` finds them. + """ + if atomic_attack is not None: + sid = getattr(atomic_attack, "_scenario_result_id", None) + name = getattr(atomic_attack, "atomic_attack_name", None) + if sid and name: + for i, r in enumerate(attack_results): + r.scenario_result_id = sid + r.scenario_data = {"atomic_attack_name": name, "objective_index": i} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) @@ -44,7 +56,8 @@ def save_attack_results_to_memory(attack_results): def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: """Create a mock AtomicAttack with required attributes for baseline creation. - The mock tracks its objectives and properly updates when filter_seed_groups_by_objectives is called. + The mock tracks its objectives and properly updates when + filter_seed_groups_by_objectives OR filter_seed_groups_by_indices is called. """ mock_attack_strategy = MagicMock() mock_attack_strategy.get_objective_target.return_value = MagicMock() @@ -54,19 +67,38 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: attack.atomic_attack_name = name attack.display_group = name attack._attack = mock_attack_strategy + attack._scenario_result_id = None - # Track current objectives in a mutable container so it can be updated + # Track current objectives and their original indices in mutable containers. + original_objectives = list(objectives) current_objectives = {"value": list(objectives)} + current_indices = {"value": list(range(len(objectives)))} - # Configure objectives property to return current objectives type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) - # Configure filter_seed_groups_by_objectives to update the tracked objectives def filter_objectives(*, remaining_objectives): remaining_set = set(remaining_objectives) - current_objectives["value"] = [obj for obj in current_objectives["value"] if obj in remaining_set] + kept = [ + (idx, obj) + for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) + if obj in remaining_set + ] + current_indices["value"] = [i for i, _ in kept] + current_objectives["value"] = [o for _, o in kept] + + def filter_indices(*, completed_indices): + kept = [ + (idx, obj) + for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) + if idx not in completed_indices + ] + current_indices["value"] = [i for i, _ in kept] + current_objectives["value"] = [o for _, o in kept] attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) + attack.filter_seed_groups_by_indices = MagicMock(side_effect=filter_indices) + attack._original_objectives = original_objectives return attack @@ -142,7 +174,7 @@ async def mock_run(*args, **kwargs): incomplete = [("obj3", ValueError("Failed to complete obj3"))] # Save completed results to memory - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=incomplete) # Retry: complete the remaining objective @@ -154,7 +186,7 @@ async def mock_run(*args, **kwargs): executed_turns=1, ) ] - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=[]) atomic_attack.run_async = mock_run @@ -200,7 +232,7 @@ async def mock_run(*args, **kwargs): incomplete = [("obj3", RuntimeError("Failed obj3")), ("obj4", RuntimeError("Failed obj4"))] # Save completed results to memory - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=incomplete) @@ -257,7 +289,7 @@ async def mock_run(*args, **kwargs): ] incomplete = [("obj4", Exception("Failed obj4")), ("obj5", Exception("Failed obj5"))] - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=incomplete) # Retry: complete remaining objectives @@ -271,7 +303,7 @@ async def mock_run(*args, **kwargs): for i in [4, 5] ] - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=[]) @@ -313,10 +345,12 @@ async def test_multiple_atomic_attacks_with_partial_results(self, mock_objective attack3 = create_mock_atomic_attack("attack_3", ["a3_obj1"]) call_counts = {"attack_1": 0, "attack_2": 0, "attack_3": 0} + attacks_by_name = {"attack_1": attack1, "attack_2": attack2, "attack_3": attack3} async def make_mock_run(attack_name, objectives): async def mock_run(*args, **kwargs): call_counts[attack_name] += 1 + this_attack = attacks_by_name[attack_name] if attack_name == "attack_2" and call_counts[attack_name] == 1: # Attack 2 fails partially on first attempt @@ -330,7 +364,7 @@ async def mock_run(*args, **kwargs): ] incomplete = [("a2_obj2", Exception("Failed a2_obj2")), ("a2_obj3", Exception("Failed a2_obj3"))] - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=this_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=incomplete) # All other attempts succeed fully @@ -341,12 +375,10 @@ async def mock_run(*args, **kwargs): outcome=AttackOutcome.SUCCESS, executed_turns=1, ) - for obj in ( - attack1 if attack_name == "attack_1" else (attack2 if attack_name == "attack_2" else attack3) - ).objectives + for obj in this_attack.objectives ] - save_attack_results_to_memory(completed) + save_attack_results_to_memory(completed, atomic_attack=this_attack) return AttackExecutorResult(completed_results=completed, incomplete_objectives=[]) diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index 836503ff5f..5c13a639a9 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -40,8 +40,20 @@ def mock_objective_scorer(): # Helper functions -def save_attack_results_to_memory(attack_results): - """Helper function to save attack results to memory (mimics what real attacks do).""" +def save_attack_results_to_memory(attack_results, *, atomic_attack=None): + """Helper function to save attack results to memory. + + When ``atomic_attack`` is provided, stamps ``scenario_result_id`` and + ``scenario_data`` onto each result (mirrors the real attack persistence + path so FK-based hydration sees the rows). + """ + if atomic_attack is not None: + sid = getattr(atomic_attack, "_scenario_result_id", None) + name = getattr(atomic_attack, "atomic_attack_name", None) + if sid and name: + for i, r in enumerate(attack_results): + r.scenario_result_id = sid + r.scenario_data = {"atomic_attack_name": name, "objective_index": i} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) @@ -86,19 +98,21 @@ def create_attack_results_list(count: int, start_index: int = 1) -> list[AttackR return [create_attack_result(i) for i in range(start_index, start_index + count)] -def create_mock_run_async(attack_results): - """Create a mock run_async that saves results to memory before returning. +def create_mock_run_async(attack_results, *, atomic_attack=None): + """Create a mock run_async that stamps + saves results to memory before returning. Args: attack_results: List of AttackResult objects to return + atomic_attack: Optional AtomicAttack mock. When provided, results are + stamped with scenario_result_id and scenario_data so FK-based + hydration finds them. Returns: AsyncMock configured to return the results """ async def mock_run_async(*args, **kwargs): - # Save results to memory (mimics what real attacks do) - save_attack_results_to_memory(attack_results) + save_attack_results_to_memory(attack_results, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=attack_results, incomplete_objectives=[]) return AsyncMock(side_effect=mock_run_async) @@ -124,10 +138,37 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack.atomic_attack_name = name attack.display_group = name attack._attack = mock_attack_strategy - type(attack).objectives = PropertyMock(return_value=objectives) - - # Configure filter_seed_groups_by_objectives - needed for scenario retry filtering - attack.filter_seed_groups_by_objectives = MagicMock() + attack._scenario_result_id = None + + # Track objectives + their original indices so the new index-based filtering + # mirrors the new resume path. The legacy text-based filter is wired too + # to support callers that still exercise it. + current_objectives = {"value": list(objectives)} + current_indices = {"value": list(range(len(objectives)))} + type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) + + def filter_objectives(*, remaining_objectives): + remaining_set = set(remaining_objectives) + kept = [ + (idx, obj) + for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) + if obj in remaining_set + ] + current_indices["value"] = [i for i, _ in kept] + current_objectives["value"] = [o for _, o in kept] + + def filter_indices(*, completed_indices): + kept = [ + (idx, obj) + for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) + if idx not in completed_indices + ] + current_indices["value"] = [i for i, _ in kept] + current_objectives["value"] = [o for _, o in kept] + + attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) + attack.filter_seed_groups_by_indices = MagicMock(side_effect=filter_indices) if run_async_mock: attack.run_async = run_async_mock @@ -214,7 +255,7 @@ async def test_no_retry_on_success(self, mock_atomic_attacks, sample_attack_resu """Test that scenario doesn't retry when execution succeeds.""" # Configure successful execution for i, run in enumerate(mock_atomic_attacks): - run.run_async = create_mock_run_async([sample_attack_results[i]]) + run.run_async = create_mock_run_async([sample_attack_results[i]], atomic_attack=run) scenario = ConcreteScenario( name="Test Scenario", @@ -247,11 +288,13 @@ async def mock_run_with_retry(*args, **kwargs): raise Exception("Test failure") # Retry succeeds results = [sample_attack_results[0]] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=mock_atomic_attacks[0]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) mock_atomic_attacks[0].run_async = mock_run_with_retry - mock_atomic_attacks[1].run_async = create_mock_run_async([sample_attack_results[1]]) + mock_atomic_attacks[1].run_async = create_mock_run_async( + [sample_attack_results[1]], atomic_attack=mock_atomic_attacks[1] + ) scenario = ConcreteScenario( name="Test Scenario", @@ -326,11 +369,13 @@ async def mock_run_with_multiple_retries(*args, **kwargs): raise Exception("Test failure") # Third attempt succeeds results = [sample_attack_results[0]] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=mock_atomic_attacks[0]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) mock_atomic_attacks[0].run_async = mock_run_with_multiple_retries - mock_atomic_attacks[1].run_async = create_mock_run_async([sample_attack_results[1]]) + mock_atomic_attacks[1].run_async = create_mock_run_async( + [sample_attack_results[1]], atomic_attack=mock_atomic_attacks[1] + ) scenario = ConcreteScenario( name="Test Scenario", @@ -360,11 +405,13 @@ async def mock_run_with_logged_failure(*args, **kwargs): raise ValueError("First failure") # Retry succeeds results = [sample_attack_results[0]] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=mock_atomic_attacks[0]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) mock_atomic_attacks[0].run_async = mock_run_with_logged_failure - mock_atomic_attacks[1].run_async = create_mock_run_async([sample_attack_results[1]]) + mock_atomic_attacks[1].run_async = create_mock_run_async( + [sample_attack_results[1]], atomic_attack=mock_atomic_attacks[1] + ) scenario = ConcreteScenario( name="Test Scenario", @@ -407,12 +454,12 @@ async def mock_run_with_partial_completion(*args, **kwargs): # First attempt: complete 2 objectives, then fail executed_objectives.extend(["obj1", "obj2"]) results = [create_attack_result(i, objective=f"obj{i}") for i in [1, 2]] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=atomic_attack) raise Exception("Failed after 2 objectives") # Retry: should only execute remaining objectives (obj3, obj4) executed_objectives.extend(["obj3", "obj4"]) results = [create_attack_result(i, objective=f"obj{i}") for i in [3, 4]] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=atomic_attack) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) atomic_attack.run_async = mock_run_with_partial_completion @@ -448,7 +495,7 @@ async def test_resumes_skipping_completed_atomic_attacks(self, mock_objective_ta async def mock_run_attack1(*args, **kwargs): call_count["attack_1"] += 1 results = [create_attack_result(1, objective="objective1")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attack1) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) # Attack 2: Succeeds on first attempt, should not be retried @@ -456,7 +503,7 @@ async def mock_run_attack2(*args, **kwargs): call_count["attack_2"] += 1 if call_count["attack_2"] == 1: results = [create_attack_result(2, objective="objective2")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attack2) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) raise AssertionError("Attack 2 should not be retried after completion") @@ -466,7 +513,7 @@ async def mock_run_attack3(*args, **kwargs): if call_count["attack_3"] == 1: raise Exception("Attack 3 failed on first attempt") results = [create_attack_result(3, objective="objective3")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attack3) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) attack1.run_async = mock_run_attack1 @@ -509,7 +556,7 @@ async def test_resumes_with_multiple_failures_across_attacks(self, mock_objectiv async def mock_run_attack1(*args, **kwargs): call_count["attack_1"] += 1 results = [create_attack_result(1, objective="objective1")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attacks[0]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) # Attack 2: Fails on first attempt, succeeds on retry @@ -518,21 +565,21 @@ async def mock_run_attack2(*args, **kwargs): if call_count["attack_2"] == 1: raise Exception("Attack 2 failed") results = [create_attack_result(2, objective="objective2")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attacks[1]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) # Attack 3: Only called on retry (after attack 2 succeeds) async def mock_run_attack3(*args, **kwargs): call_count["attack_3"] += 1 results = [create_attack_result(3, objective="objective3")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attacks[2]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) # Attack 4: Only called on retry async def mock_run_attack4(*args, **kwargs): call_count["attack_4"] += 1 results = [create_attack_result(4, objective="objective4")] - save_attack_results_to_memory(results) + save_attack_results_to_memory(results, atomic_attack=attacks[3]) return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) attacks[0].run_async = mock_run_attack1 @@ -564,3 +611,158 @@ async def mock_run_attack4(*args, **kwargs): assert call_count["attack_4"] == 1 # All four attacks should be in results assert len(result.attack_results) == 4 + + +@pytest.mark.usefixtures("patch_central_database") +class TestScenarioFKResumeRegression: + """Regression tests for the FK-based scenario linkage resume path. + + The bug being regression-tested: when a Scenario is interrupted mid- + AtomicAttack (Ctrl-C, OOM, crash), AttackResults already persisted to the + DB used to be invisible to the scenario because the scenario→attack-result + link only lived in a JSON manifest written after the whole AtomicAttack + returned. On resume, those objectives were re-executed (wasted compute). + + After the refactor, ``scenario_result_id`` is stamped on each + ``AttackResultEntry`` at write time, so resume reads them directly and + skips the already-done work even when the manifest was never updated. + """ + + async def test_resume_skips_objectives_persisted_before_interruption(self, mock_objective_target): + """Simulate Ctrl-C after some objectives in an atomic attack persisted + results but before the manifest was bulk-written. On resume, only the + missing objectives are re-executed.""" + atomic_attack = create_mock_atomic_attack("partial", ["o1", "o2", "o3", "o4"]) + + async def first_run(*args, **kwargs): + partials = [ + create_attack_result(0, conversation_id="c1", objective="o1"), + create_attack_result(1, conversation_id="c2", objective="o2"), + ] + save_attack_results_to_memory(partials, atomic_attack=atomic_attack) + raise Exception("simulated crash after partial persistence") + + atomic_attack.run_async = first_run + + scenario = ConcreteScenario( + name="Interrupted Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + await scenario.initialize_async(objective_target=mock_objective_target, max_retries=0) + + with pytest.raises(Exception, match="simulated crash"): + await scenario.run_async() + + scenario_result_id = scenario._scenario_result_id + assert scenario_result_id is not None + + # === Resume by scenario_result_id === + atomic_attack_resume = create_mock_atomic_attack("partial", ["o1", "o2", "o3", "o4"]) + executed: list[str] = [] + + async def second_run(*args, **kwargs): + executed.extend(atomic_attack_resume.objectives) + results = [ + create_attack_result(i, conversation_id=f"c{i + 1}", objective=obj) + for i, obj in enumerate(atomic_attack_resume.objectives, start=2) + ] + save_attack_results_to_memory(results, atomic_attack=atomic_attack_resume) + return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) + + atomic_attack_resume.run_async = second_run + + scenario_resumed = ConcreteScenario( + name="Interrupted Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack_resume], + scenario_result_id=scenario_result_id, + ) + await scenario_resumed.initialize_async(objective_target=mock_objective_target, max_retries=0) + await scenario_resumed.run_async() + + # Resume executed only the missing objectives — the core fix. + assert executed == ["o3", "o4"] + + async def test_duplicate_objective_text_does_not_collapse_on_resume(self, mock_objective_target): + """Two seed groups with the SAME objective text are distinct slots — + resume must treat them as separate (index-based) rather than collapse + them by text.""" + atomic_attack = create_mock_atomic_attack("dup_attack", ["dup-obj", "dup-obj"]) + + async def first(*args, **kwargs): + partials = [create_attack_result(0, conversation_id="dup-c0", objective="dup-obj")] + save_attack_results_to_memory(partials, atomic_attack=atomic_attack) + raise Exception("crash mid-attack") + + atomic_attack.run_async = first + + scenario = ConcreteScenario( + name="Dup Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + await scenario.initialize_async(objective_target=mock_objective_target, max_retries=0) + + with pytest.raises(Exception, match="crash mid-attack"): + await scenario.run_async() + + scenario_result_id = scenario._scenario_result_id + assert scenario_result_id is not None + + # === Resume === + atomic_attack_resume = create_mock_atomic_attack("dup_attack", ["dup-obj", "dup-obj"]) + executed: list[str] = [] + + async def second(*args, **kwargs): + executed.extend(atomic_attack_resume.objectives) + results = [ + create_attack_result(i, conversation_id=f"dup-c-resume-{i}", objective=obj) + for i, obj in enumerate(atomic_attack_resume.objectives, start=1) + ] + save_attack_results_to_memory(results, atomic_attack=atomic_attack_resume) + return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) + + atomic_attack_resume.run_async = second + + scenario_resumed = ConcreteScenario( + name="Dup Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack_resume], + scenario_result_id=scenario_result_id, + ) + await scenario_resumed.initialize_async(objective_target=mock_objective_target, max_retries=0) + await scenario_resumed.run_async() + + # Resume executed exactly ONE "dup-obj" — the missing index — even + # though the text matched the already-done seed group. Index-based + # resume is what makes this work; the deprecated text-based filter + # would have collapsed both seed groups. + assert executed == ["dup-obj"] + + async def test_duplicate_atomic_attack_name_warns_but_does_not_raise(self, mock_objective_target, caplog): + """Duplicate atomic_attack_name within a scenario degrades resume to + collapse-by-name. The validator logs a WARNING but does not raise — + some existing scenarios (Encoding, Jailbreak) intentionally use + duplicate names for different converter configs.""" + dup1 = create_mock_atomic_attack("dup_name", ["objA"]) + dup2 = create_mock_atomic_attack("dup_name", ["objB"]) + + async def noop_run(*args, **kwargs): + return AttackExecutorResult(completed_results=[], incomplete_objectives=[]) + + dup1.run_async = noop_run + dup2.run_async = noop_run + + scenario = ConcreteScenario( + name="Dup Name Scenario", + version=1, + atomic_attacks_to_return=[dup1, dup2], + ) + + with caplog.at_level("WARNING"): + await scenario.initialize_async(objective_target=mock_objective_target) + + assert any("duplicate atomic_attack_name" in record.message for record in caplog.records), ( + "Expected a duplicate-name warning" + ) From 3e844c5d51c615a1775b7776ab5b14431f649b53 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 19 May 2026 11:49:45 -0700 Subject: [PATCH 02/10] Address review: cleanup, rename to ScenarioExecutionAttribution, simplify hydration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Delete dishonest no-op add_attack_results_to_scenario shim. - Standardize on print_deprecation_message (drop ad-hoc warnings.warn). Style guide gains a concise Deprecations section with the `removed_in = current minor + 2` rule. - Remove stale per-scenario state left over from the manifest era: AttackContext._error_attack_result_id, _StrategyRuntimeError.error_attack_result_id, Scenario._result_lock, Scenario._original_objectives_map, and the stray `import asyncio`. Replace defensive `getattr(context, '_attribution', None)` with direct attribute access — the contract is mandatory. - Rename ExecutionAttribution -> ScenarioExecutionAttribution (and the module file) to match its scenario-specific schema. - Refactor MemoryInterface.get_scenario_results: split into _build_scenario_result_query_conditions, _query_scenario_result_entries, _hydrate_scenario_attack_results. The hydrator now issues a single batched IN-query on AttackResultEntry.scenario_result_id (fixes the previous N+1) and drops the legacy attack_results_json manifest fallback entirely — the FK is the sole source of truth. - Narrow _stamp_attribution(result=) to AttackResult to satisfy ty. - Update affected tests; rewrite the four hydration tests that incidentally relied on the manifest fallback to use the production FK write path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../instructions/style-guide.instructions.md | 20 ++ pyrit/executor/attack/core/attack_executor.py | 14 +- pyrit/executor/attack/core/attack_strategy.py | 22 +- ...n.py => scenario_execution_attribution.py} | 19 +- pyrit/executor/core/strategy.py | 10 +- pyrit/memory/memory_interface.py | 270 +++++++++--------- pyrit/memory/memory_models.py | 4 +- pyrit/models/attack_result.py | 2 +- pyrit/scenario/core/atomic_attack.py | 17 +- pyrit/scenario/core/scenario.py | 13 +- .../attack/core/test_attack_executor.py | 20 +- .../attack/core/test_attack_strategy.py | 10 +- .../test_interface_scenario_results.py | 194 ++++--------- 13 files changed, 268 insertions(+), 347 deletions(-) rename pyrit/executor/attack/core/{execution_attribution.py => scenario_execution_attribution.py} (71%) diff --git a/.github/instructions/style-guide.instructions.md b/.github/instructions/style-guide.instructions.md index f5f245f9f5..42fdab3315 100644 --- a/.github/instructions/style-guide.instructions.md +++ b/.github/instructions/style-guide.instructions.md @@ -293,6 +293,26 @@ def process_items(self, *, items: list[str]) -> list[str]: return [] ``` +## Deprecations + +When deprecating a public class, function, method, parameter, or module path, use `pyrit.common.deprecation.print_deprecation_message` — never `warnings.warn` directly. It wraps `warnings.warn(..., DeprecationWarning, stacklevel=3)` with a consistent format so filtering still works. + +Set `removed_in` to **current version + 2 minor versions** (e.g. `0.14.x` → `removed_in="0.16.0"`). This gives one full release cycle of warning before removal. + +```python +from pyrit.common.deprecation import print_deprecation_message + +def old_method(self, *, foo: str) -> None: + print_deprecation_message( + old_item="MyClass.old_method", + new_item="MyClass.new_method", + removed_in="0.16.0", + ) + ... +``` + +`old_item` / `new_item` accept a class/callable (qualified name is generated) or a string. + ## Pythonic Patterns ### List Comprehensions diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index edc5a212f8..7dee98f4ae 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -24,7 +24,7 @@ AttackStrategyContextT, AttackStrategyResultT, ) -from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution +from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution from pyrit.models import SeedAttackGroup if TYPE_CHECKING: @@ -143,7 +143,7 @@ async def execute_attack_from_seed_groups_async( objective_scorer: Optional["TrueFalseScorer"] = None, field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -167,7 +167,7 @@ async def execute_attack_from_seed_groups_async( objectives fail. If False (default), raises the first exception. attribution_factory: Optional callable that maps an input index (the seed group's original index, parallel-safe and deterministic) to - an ``ExecutionAttribution``. When provided, each per-task + a ``ScenarioExecutionAttribution``. When provided, each per-task ``AttackContext`` is stamped with the attribution so the resulting ``AttackResultEntry`` row carries the scenario FK + scenario_data. When ``None``, no attribution is applied. @@ -223,7 +223,7 @@ async def execute_attack_async( objectives: Sequence[str], field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -239,7 +239,7 @@ async def execute_attack_async( return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. attribution_factory: Optional callable mapping each input index to - an ExecutionAttribution. When provided, the per-task context is + a ScenarioExecutionAttribution. When provided, the per-task context is stamped with the attribution so the persistence path can record scenario linkage. **broadcast_fields: Fields applied to all objectives (e.g., memory_labels). @@ -291,7 +291,7 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Execute attacks in parallel with a list of pre-built parameters. @@ -304,7 +304,7 @@ async def _execute_with_params_list_async( params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. attribution_factory: Optional callable mapping each input index to - an ExecutionAttribution. When provided, the per-task context is + a ScenarioExecutionAttribution. When provided, the per-task context is stamped with the attribution so the persistence path can record scenario linkage. diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 278c96ec7f..b92a3be203 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -36,7 +36,7 @@ if TYPE_CHECKING: from pyrit.executor.attack.core.attack_config import AttackScoringConfig - from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution from pyrit.prompt_target import PromptTarget AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") @@ -71,15 +71,12 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: Optional[list[Message]] = None _memory_labels_override: Optional[dict[str, str]] = None - # Set by the ON_ERROR handler to link error AttackResults to ScenarioResults - _error_attack_result_id: str | None = None - # Optional attribution from an upstream orchestrator (e.g. Scenario). When # set, the persistence path stamps scenario_result_id + scenario_data onto # the resulting AttackResult so it can be located later for hydration and # resume. Set by AttackExecutor per-task before scheduling. Stays None for # ad-hoc/direct attack execution outside any scenario. - _attribution: Optional[ExecutionAttribution] = None + _attribution: Optional[ScenarioExecutionAttribution] = None # Convenience properties that delegate to params or overrides @property @@ -245,12 +242,12 @@ async def _on_post_execute( def _stamp_attribution( *, context: AttackStrategyContextT, - result: AttackStrategyResultT, + result: AttackResult, ) -> None: """ Copy scenario attribution from the AttackContext onto the AttackResult. - Reads ``context._attribution`` (an ``ExecutionAttribution`` set by the + Reads ``context._attribution`` (a ``ScenarioExecutionAttribution`` set by the AttackExecutor when running inside a Scenario). When present, writes ``scenario_result_id`` and a fixed-schema ``scenario_data`` dict onto the result so they round-trip into ``AttackResultEntry``. @@ -259,7 +256,7 @@ def _stamp_attribution( context: The per-task AttackContext. result: The AttackResult that is about to be persisted. """ - attribution = getattr(context, "_attribution", None) + attribution = context._attribution if attribution is None: return result.scenario_result_id = attribution.scenario_result_id @@ -307,9 +304,6 @@ async def _on_error_async( if not error or not context: return - # Clear any stale ID from a previous execution - context._error_attack_result_id = None - # Collect retry events (visible via inherited ContextVar copy) collector = get_retry_collector() retry_events = collector.events if collector else [] @@ -336,14 +330,10 @@ async def _on_error_async( error_result.execution_time_ms = int((end_time - context.start_time) * 1000) # Stamp scenario attribution onto the error result so it is locatable - # via the scenario FK on resume (rather than via the previous - # error_attack_result_ids_json manifest). + # via the scenario FK on resume. self._stamp_attribution(context=context, result=error_result) - # Persist first, then set the ID on the context so scenario-level code - # only sees the reference if the write succeeded. self._memory.add_attack_results_to_memory(attack_results=[error_result]) - context._error_attack_result_id = error_result.attack_result_id self._logger.error(f"Attack failed with {type(error).__name__}: {error}") diff --git a/pyrit/executor/attack/core/execution_attribution.py b/pyrit/executor/attack/core/scenario_execution_attribution.py similarity index 71% rename from pyrit/executor/attack/core/execution_attribution.py rename to pyrit/executor/attack/core/scenario_execution_attribution.py index 09b6ecc131..07b97d9d49 100644 --- a/pyrit/executor/attack/core/execution_attribution.py +++ b/pyrit/executor/attack/core/scenario_execution_attribution.py @@ -2,13 +2,12 @@ # Licensed under the MIT license. """ -Typed attribution metadata used to link a persisted ``AttackResult`` to an -upstream orchestrator (e.g. a ``Scenario``). +Typed attribution metadata used to link a persisted ``AttackResult`` to the +``Scenario`` that produced it. -The attribution lives in the ``executor`` layer so the executor never imports -from ``scenario``. ``Scenario`` is one producer; future orchestrators may -produce attribution too. The attack persistence path (the default attack event -handler) is the consumer. +Lives in the ``executor`` layer (rather than ``scenario``) so the attack +persistence path — the consumer — does not introduce a dependency on +``pyrit.scenario``. """ from __future__ import annotations @@ -17,11 +16,11 @@ @dataclass(frozen=True) -class ExecutionAttribution: +class ScenarioExecutionAttribution: """ - Attribution metadata produced by an upstream orchestrator and consumed by - the attack persistence path to record linkage on the resulting - ``AttackResultEntry``. + Scenario-linkage metadata stamped onto an ``AttackContext`` by the + ``AttackExecutor`` and copied onto the resulting ``AttackResult`` by the + attack persistence path so the row carries the scenario FK + scenario_data. Attributes: scenario_result_id (str): The ID of the scenario result that produced diff --git a/pyrit/executor/core/strategy.py b/pyrit/executor/core/strategy.py index 3aa9ede03b..fe299e6405 100644 --- a/pyrit/executor/core/strategy.py +++ b/pyrit/executor/core/strategy.py @@ -31,11 +31,7 @@ class _StrategyRuntimeError(RuntimeError): - """RuntimeError subclass that carries an optional error_attack_result_id.""" - - def __init__(self, message: str, *, error_attack_result_id: str | None = None) -> None: - super().__init__(message) - self.error_attack_result_id = error_attack_result_id + """RuntimeError subclass for strategy execution failures.""" @dataclass @@ -386,9 +382,7 @@ async def execute_with_context_async(self, *, context: StrategyContextT) -> Stra else: error_message = f"Strategy execution failed for {self.__class__.__name__}: {str(e)}" - # Attach the error attack result ID if the ON_ERROR handler created one - error_attack_result_id = getattr(context, "_error_attack_result_id", None) - runtime_error = _StrategyRuntimeError(error_message, error_attack_result_id=error_attack_result_id) + runtime_error = _StrategyRuntimeError(error_message) raise runtime_error from e async def execute_async(self, **kwargs: Any) -> StrategyResultT: diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 2b80b30e4d..5c28856d33 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1965,40 +1965,6 @@ def add_scenario_results_to_memory(self, *, scenario_results: Sequence[ScenarioR entries=[ScenarioResultEntry(entry=scenario_result) for scenario_result in scenario_results] ) - def add_attack_results_to_scenario( - self, - *, - scenario_result_id: str, - atomic_attack_name: str, - attack_results: Sequence[AttackResult], - ) -> bool: - """ - No-op shim retained for backward compatibility. - - Scenario→AttackResult linkage is now stamped per-result by the attack - persistence path (via ``ExecutionAttribution`` on the - ``AttackContext``). This method is kept for one release for external - callers; it will be removed. - - Args: - scenario_result_id (str): The ID of the scenario result (ignored). - atomic_attack_name (str): The atomic attack name (ignored). - attack_results (Sequence[AttackResult]): The attack results (ignored). - - Returns: - bool: Always True (no-op success). - """ - print_deprecation_message( - old_item="memory.add_attack_results_to_scenario(...)", - new_item=( - "Scenario→AttackResult linkage is stamped per-result by the attack persistence path. " - "Setting Scenario._scenario_result_id on each AtomicAttack is sufficient; no explicit " - "post-run bulk-write is needed." - ), - removed_in="0.16.0", - ) - return True - def update_scenario_run_state( self, *, @@ -2094,6 +2060,57 @@ def get_scenario_results( if scenario_result_ids is not None and len(scenario_result_ids) == 0: return [] + conditions = self._build_scenario_result_query_conditions( + scenario_name=scenario_name, + scenario_version=scenario_version, + pyrit_version=pyrit_version, + added_after=added_after, + added_before=added_before, + labels=labels, + objective_target_endpoint=objective_target_endpoint, + objective_target_model_name=objective_target_model_name, + identifier_filters=identifier_filters, + ) + + try: + entries = self._query_scenario_result_entries( + scenario_result_ids=scenario_result_ids, + conditions=conditions, + limit=limit, + ) + + attack_results_by_scenario = self._hydrate_scenario_attack_results(entries=entries) + + scenario_results: list[ScenarioResult] = [] + for entry in entries: + scenario_result = entry.get_scenario_result() + scenario_result.attack_results = attack_results_by_scenario.get(entry.id, {}) + scenario_results.append(scenario_result) + + return scenario_results + except Exception as e: + logger.exception(f"Failed to retrieve scenario results with error {e}") + raise + + def _build_scenario_result_query_conditions( + self, + *, + scenario_name: str | None, + scenario_version: int | None, + pyrit_version: str | None, + added_after: datetime | None, + added_before: datetime | None, + labels: dict[str, str] | None, + objective_target_endpoint: str | None, + objective_target_model_name: str | None, + identifier_filters: Sequence[IdentifierFilter] | None, + ) -> "list[ColumnElement[bool]]": + """ + Build the WHERE conditions for ``get_scenario_results``. + + Returns: + list[ColumnElement[bool]]: SQLAlchemy WHERE clauses derived from the supplied filters. + """ conditions: list[ColumnElement[bool]] = [] if scenario_name: @@ -2147,111 +2164,96 @@ def get_scenario_results( ) ) - try: - order_by_clause = ScenarioResultEntry.completion_time.desc() - - # Handle scenario_result_ids with batched queries if needed - if scenario_result_ids: - entries = self._execute_batched_query( - ScenarioResultEntry, - batch_column=ScenarioResultEntry.id, - batch_values=list(scenario_result_ids), - other_conditions=conditions, - order_by=order_by_clause, - limit=limit, - ) - else: - entries = self._query_entries( - ScenarioResultEntry, - conditions=and_(*conditions) if conditions else None, - order_by=order_by_clause, - limit=limit, - ) + return conditions - # Convert entries to ScenarioResults and populate attack_results. - # - # Hydration strategy: - # 1. Pull AttackResultEntry rows for this scenario via the new - # scenario_result_id FK (the post-redesign source of truth). - # 2. Fall back to the legacy manifest path for any conversation_id - # present in attack_results_json but NOT in the FK result set. - # This handles partial-migration / pre-FK data. - # 3. Group by atomic_attack_name (from scenario_data for FK rows, - # from manifest keys for legacy rows) and sort each group by - # scenario_data["objective_index"] when available. - scenario_results = [] - for entry in entries: - scenario_result = entry.get_scenario_result() - scenario_id_str = str(entry.id) + def _query_scenario_result_entries( + self, + *, + scenario_result_ids: Sequence[str] | None, + conditions: "list[ColumnElement[bool]]", + limit: int | None, + ) -> Sequence[ScenarioResultEntry]: + """ + Run the (possibly batched) ScenarioResultEntry query. - # FK-linked rows (preferred source). - fk_rows = self._query_entries( - AttackResultEntry, - conditions=AttackResultEntry.scenario_result_id == entry.id, - ) - fk_results: list[AttackResult] = [row.get_attack_result() for row in fk_rows] - fk_attack_result_ids = {r.attack_result_id for r in fk_results} - - # Legacy manifest path — only fill in conversation_ids not - # already covered by the FK path. - manifest_map = entry.get_conversation_ids_by_attack_name() - manifest_missing_ids: list[str] = [] - for conv_ids in manifest_map.values(): - manifest_missing_ids.extend(conv_ids) - - legacy_results_by_conv: dict[str, AttackResult] = {} - if manifest_missing_ids: - legacy_rows = self._execute_batched_query( - AttackResultEntry, - batch_column=AttackResultEntry.conversation_id, - batch_values=manifest_missing_ids, - ) - for row in legacy_rows: - result = row.get_attack_result() - if result.attack_result_id in fk_attack_result_ids: - continue # already represented by the FK path - legacy_results_by_conv[row.conversation_id] = result - - # Group: FK rows by scenario_data["atomic_attack_name"], legacy - # rows by manifest key. Sort each group by objective_index when - # available (FK rows) and append legacy entries after, preserving - # manifest order. - grouped: dict[str, list[tuple[int, AttackResult]]] = {} - for r in fk_results: - name = (r.scenario_data or {}).get("atomic_attack_name") - if not name: - continue - idx = (r.scenario_data or {}).get("objective_index") - sort_key = idx if isinstance(idx, int) else len(grouped.setdefault(name, [])) - grouped.setdefault(name, []).append((sort_key, r)) - - legacy_count = 0 - for attack_name, conv_ids in manifest_map.items(): - legacy_bucket = grouped.setdefault(attack_name, []) - base_offset = len(legacy_bucket) - for offset, cid in enumerate(conv_ids): - legacy_result = legacy_results_by_conv.get(cid) - if legacy_result is None: - continue - legacy_bucket.append((base_offset + offset, legacy_result)) - legacy_count += 1 - - if legacy_count or fk_results: - logger.debug( - f"Hydrated scenario {scenario_id_str}: " - f"{len(fk_results)} via FK, {legacy_count} via legacy manifest fallback" - ) + Returns: + Sequence[ScenarioResultEntry]: The matching rows ordered by completion_time descending. + """ + order_by_clause = ScenarioResultEntry.completion_time.desc() - scenario_result.attack_results = { - name: [r for _, r in sorted(bucket, key=lambda kv: kv[0])] for name, bucket in grouped.items() - } + if scenario_result_ids: + return self._execute_batched_query( + ScenarioResultEntry, + batch_column=ScenarioResultEntry.id, + batch_values=list(scenario_result_ids), + other_conditions=conditions, + order_by=order_by_clause, + limit=limit, + ) - scenario_results.append(scenario_result) + return self._query_entries( + ScenarioResultEntry, + conditions=and_(*conditions) if conditions else None, + order_by=order_by_clause, + limit=limit, + ) - return scenario_results - except Exception as e: - logger.exception(f"Failed to retrieve scenario results with error {e}") - raise + def _hydrate_scenario_attack_results( + self, + *, + entries: Sequence[ScenarioResultEntry], + ) -> dict[uuid.UUID, dict[str, list[AttackResult]]]: + """ + Fetch every ``AttackResult`` linked to the given scenarios via the + ``AttackResultEntry.scenario_result_id`` FK in a single batched query, + then group by scenario + ``atomic_attack_name`` and sort each group by + ``scenario_data["objective_index"]``. + + FK linkage is the sole source of truth — set at write-time by the + attack persistence path when a ``ScenarioExecutionAttribution`` is on + the context. Rows without a valid ``scenario_data`` payload are skipped + (and logged) rather than guessed at. + + Returns: + dict[uuid.UUID, dict[str, list[AttackResult]]]: Mapping of + ``scenario_result_id`` → ``atomic_attack_name`` → ordered list of + ``AttackResult`` objects. Scenarios with no linked rows map to ``{}``. + """ + if not entries: + return {} + + scenario_ids = [entry.id for entry in entries] + attack_rows = self._execute_batched_query( + AttackResultEntry, + batch_column=AttackResultEntry.scenario_result_id, + batch_values=scenario_ids, + ) + + grouped: dict[uuid.UUID, dict[str, list[tuple[int, AttackResult]]]] = {entry.id: {} for entry in entries} + + for row in attack_rows: + scenario_id = row.scenario_result_id + if scenario_id is None or scenario_id not in grouped: + continue + + data = row.scenario_data or {} + name = data.get("atomic_attack_name") + if not name: + logger.debug( + f"Skipping AttackResultEntry {row.id} during hydration: scenario_data missing atomic_attack_name" + ) + continue + + objective_index = data.get("objective_index") + sort_key = objective_index if isinstance(objective_index, int) else 0 + grouped[scenario_id].setdefault(name, []).append((sort_key, row.get_attack_result())) + + return { + scenario_id: { + name: [ar for _, ar in sorted(bucket, key=lambda kv: kv[0])] for name, bucket in name_buckets.items() + } + for scenario_id, name_buckets in grouped.items() + } def print_schema(self) -> None: """ diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 672dc96e79..562733fadd 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -831,8 +831,8 @@ def __init__(self, *, entry: AttackResult) -> None: ) self.total_retries = entry.total_retries - # Scenario linkage (set by the attack persistence path when an - # ExecutionAttribution is present on the AttackContext; otherwise None) + # Scenario linkage (set by the attack persistence path when a + # ScenarioExecutionAttribution is present on the AttackContext; otherwise None) self.scenario_result_id = uuid.UUID(entry.scenario_result_id) if entry.scenario_result_id else None self.scenario_data = entry.scenario_data diff --git a/pyrit/models/attack_result.py b/pyrit/models/attack_result.py index 46f8656250..d8c6cf9a98 100644 --- a/pyrit/models/attack_result.py +++ b/pyrit/models/attack_result.py @@ -108,7 +108,7 @@ class AttackResult(StrategyResult): total_retries: int = 0 # Scenario linkage (infrastructure-managed). Set by the attack persistence - # path when an ExecutionAttribution is present on the AttackContext. User + # path when a ScenarioExecutionAttribution is present on the AttackContext. User # code should not set these directly; ad-hoc AttackResults created outside # a Scenario leave both fields as None and the corresponding DB columns # remain NULL. diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index cdf6681f29..6ef19d11f3 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -17,9 +17,10 @@ import warnings from typing import TYPE_CHECKING, Any, Optional +from pyrit.common.deprecation import print_deprecation_message from pyrit.executor.attack import AttackExecutor, AttackStrategy from pyrit.executor.attack.core.attack_executor import AttackExecutorResult -from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution +from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution from pyrit.identifiers import build_atomic_attack_identifier from pyrit.identifiers.evaluation_identifier import AtomicAttackEvaluationIdentifier from pyrit.memory import CentralMemory @@ -176,12 +177,10 @@ def filter_seed_groups_by_objectives(self, *, remaining_objectives: list[str]) - Args: remaining_objectives (List[str]): List of objectives that still need to be executed. """ - warnings.warn( - "filter_seed_groups_by_objectives is deprecated; use filter_seed_groups_by_indices " - "for index-stable resume. Duplicate objective text can collapse otherwise-distinct " - "seed groups.", - DeprecationWarning, - stacklevel=2, + print_deprecation_message( + old_item="AtomicAttack.filter_seed_groups_by_objectives", + new_item="AtomicAttack.filter_seed_groups_by_indices", + removed_in="0.16.0", ) remaining_set = set(remaining_objectives) kept: list[tuple[int, SeedAttackGroup]] = [ @@ -282,8 +281,8 @@ async def run_async( name = self.atomic_attack_name original_indices = list(self._original_indices) - def attribution_factory(input_index: int) -> ExecutionAttribution: - return ExecutionAttribution( + def attribution_factory(input_index: int) -> ScenarioExecutionAttribution: + return ScenarioExecutionAttribution( scenario_result_id=scenario_id, atomic_attack_name=name, objective_index=original_indices[input_index], diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 000f75d850..4120f05d41 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -8,7 +8,6 @@ AtomicAttack instances sequentially, enabling comprehensive security testing campaigns. """ -import asyncio import copy import json import logging @@ -218,15 +217,10 @@ def __init__( self._memory = CentralMemory.get_memory_instance() self._atomic_attacks: list[AtomicAttack] = [] self._scenario_result_id: Optional[str] = str(scenario_result_id) if scenario_result_id else None - self._result_lock = asyncio.Lock() # Store prepared strategies for use in _get_atomic_attacks_async self._scenario_strategies: list[ScenarioStrategy] = [] - # Store original objectives for each atomic attack (before any mutations) - # Key: atomic_attack_name, Value: tuple of original objectives - self._original_objectives_map: dict[str, tuple[str, ...]] = {} - # Maps atomic_attack_name → display_group for user-facing aggregation self._display_group_map: dict[str, str] = {} @@ -714,11 +708,6 @@ async def initialize_async( f"Each AtomicAttack within a scenario should have a unique name." ) - # Store original objectives for each atomic attack (before any mutations during execution) - self._original_objectives_map = { - atomic_attack.atomic_attack_name: tuple(atomic_attack.objectives) for atomic_attack in self._atomic_attacks - } - # Snapshot params onto the identifier before the resume branch so the identifier # is fully populated regardless of which branch we take. Deep-copy avoids sharing # mutable state with self.params. @@ -1215,7 +1204,7 @@ async def _execute_scenario_async(self) -> ScenarioResult: # Error AttackResults are linked to this scenario via the # scenario_result_id FK on AttackResultEntry (stamped by - # the attack event handler when an ExecutionAttribution + # the attack event handler when a ScenarioExecutionAttribution # is on the context). The previous per-scenario error_id # manifest is no longer needed. diff --git a/tests/unit/executor/attack/core/test_attack_executor.py b/tests/unit/executor/attack/core/test_attack_executor.py index 0b5d0ed55e..ff81e703c9 100644 --- a/tests/unit/executor/attack/core/test_attack_executor.py +++ b/tests/unit/executor/attack/core/test_attack_executor.py @@ -338,7 +338,7 @@ async def test_validates_explicit_empty_field_overrides_for_seed_groups(self): @pytest.mark.usefixtures("patch_central_database") class TestAttributionFactoryPropagation: - """Tests for ExecutionAttribution propagation through the AttackExecutor. + """Tests for ScenarioExecutionAttribution propagation through the AttackExecutor. The executor stamps ``context._attribution = attribution_factory(input_index)`` on each per-task context. The input_index must be the seed group's original @@ -347,14 +347,14 @@ class TestAttributionFactoryPropagation: """ async def test_attribution_factory_sets_per_task_attribution_in_order(self): - from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution attack = create_mock_attack() seen_indices: list[int] = [] async def capture(context): # Read the attribution stamped by the executor BEFORE the task runs. - attr = getattr(context, "_attribution", None) + attr = context._attribution assert attr is not None seen_indices.append(attr.objective_index) return create_attack_result(context.params.objective) @@ -363,8 +363,8 @@ async def capture(context): seed_groups = [create_seed_group(f"obj-{i}") for i in range(4)] - def factory(idx: int) -> ExecutionAttribution: - return ExecutionAttribution( + def factory(idx: int) -> ScenarioExecutionAttribution: + return ScenarioExecutionAttribution( scenario_result_id="sid", atomic_attack_name="atomic", objective_index=idx, @@ -387,13 +387,13 @@ async def test_attribution_factory_parallel_safe_with_high_concurrency(self): regardless of completion order. Out-of-order completion must not produce stamped indices that point at the wrong seed group. """ - from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution attack = create_mock_attack() seen: dict[str, int] = {} async def out_of_order(context): - attr = getattr(context, "_attribution", None) + attr = context._attribution assert attr is not None # Reverse-delay the tasks so completion order is the inverse of input # order. The stamped index must still match the seed group. @@ -405,8 +405,8 @@ async def out_of_order(context): seed_groups = [create_seed_group(f"obj-{i}") for i in range(6)] - def factory(idx: int) -> ExecutionAttribution: - return ExecutionAttribution( + def factory(idx: int) -> ScenarioExecutionAttribution: + return ScenarioExecutionAttribution( scenario_result_id="sid", atomic_attack_name="atomic", objective_index=idx, @@ -428,7 +428,7 @@ async def test_no_attribution_factory_leaves_context_attribution_none(self): attack = create_mock_attack() async def capture(context): - attr = getattr(context, "_attribution", None) + attr = context._attribution assert attr is None return create_attack_result(context.params.objective) diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 7d441b5550..22ed8bee0c 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -613,14 +613,14 @@ async def test_on_event_handles_other_events(self, event_handler, sample_attack_ async def test_on_post_execute_stamps_scenario_attribution_when_present( self, sample_attack_context, sample_attack_result, mock_memory ): - """When the context carries an ExecutionAttribution, the persisted + """When the context carries a ScenarioExecutionAttribution, the persisted AttackResult must have scenario_result_id + scenario_data populated.""" - from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() sample_attack_context.start_time = 100.0 - sample_attack_context._attribution = ExecutionAttribution( + sample_attack_context._attribution = ScenarioExecutionAttribution( scenario_result_id="scenario-1", atomic_attack_name="atomic_a", objective_index=3, @@ -666,12 +666,12 @@ async def test_on_post_execute_no_attribution_leaves_fields_none( async def test_on_error_stamps_scenario_attribution_when_present(self, sample_attack_context, mock_memory): """Error AttackResults must also carry the scenario FK so error lookups via get_attack_results(scenario_result_id=..., outcome=ERROR) work.""" - from pyrit.executor.attack.core.execution_attribution import ExecutionAttribution + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() sample_attack_context.start_time = 100.0 - sample_attack_context._attribution = ExecutionAttribution( + sample_attack_context._attribution = ScenarioExecutionAttribution( scenario_result_id="scenario-err", atomic_attack_name="atomic_err", objective_index=7, diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index 4ea6258a50..f56516b2ba 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -160,21 +160,24 @@ def test_empty_ids_returns_empty(sqlite_instance: MemoryInterface): def test_attack_results_populated_correctly(sqlite_instance: MemoryInterface): """Test that retrieving scenario results populates attack_results correctly.""" - # Create and add attack results - attack_result1 = create_attack_result("conv_1", "Objective 1", AttackOutcome.SUCCESS) - attack_result2 = create_attack_result("conv_2", "Objective 2", AttackOutcome.FAILURE) - attack_result3 = create_attack_result("conv_3", "Objective 3", AttackOutcome.SUCCESS) - sqlite_instance.add_attack_results_to_memory(attack_results=[attack_result1, attack_result2, attack_result3]) + scenario_result = create_scenario_result(name="Multi-Attack Scenario", attack_results={}) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - # Create scenario result with multiple attacks - scenario_result = create_scenario_result( - name="Multi-Attack Scenario", - attack_results={ - "PromptInjection": [attack_result1, attack_result2], - "Crescendo": [attack_result3], - }, + sid = scenario_result.id + attack_result1 = _make_attack_result_for_scenario( + scenario_result_id=sid, atomic_attack_name="PromptInjection", objective_index=0, conversation_id="conv_1" ) - sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + attack_result2 = _make_attack_result_for_scenario( + scenario_result_id=sid, + atomic_attack_name="PromptInjection", + objective_index=1, + conversation_id="conv_2", + outcome=AttackOutcome.FAILURE, + ) + attack_result3 = _make_attack_result_for_scenario( + scenario_result_id=sid, atomic_attack_name="Crescendo", objective_index=0, conversation_id="conv_3" + ) + sqlite_instance.add_attack_results_to_memory(attack_results=[attack_result1, attack_result2, attack_result3]) # Retrieve and verify attack_results are populated results = sqlite_instance.get_scenario_results() @@ -198,40 +201,40 @@ def test_attack_results_populated_correctly(sqlite_instance: MemoryInterface): def test_attack_order_preserved(sqlite_instance: MemoryInterface): - """Test that attack results maintain their order within each attack name.""" - # Create and add attack results - attack_results = [create_attack_result(f"conv_{i}", f"Objective {i}") for i in range(5)] - sqlite_instance.add_attack_results_to_memory(attack_results=attack_results) - - # Create scenario result with ordered attacks - scenario_result = create_scenario_result( - name="Ordered Scenario", - attack_results={ - "Attack1": attack_results, - }, - ) + """Test that attack results are sorted by objective_index within each attack name.""" + scenario_result = create_scenario_result(name="Ordered Scenario", attack_results={}) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - # Retrieve and verify order is preserved + sid = scenario_result.id + # Insert in reverse order to prove the hydration sorts by objective_index, not insert order. + attack_results = [ + _make_attack_result_for_scenario( + scenario_result_id=sid, atomic_attack_name="Attack1", objective_index=i, conversation_id=f"conv_{i}" + ) + for i in reversed(range(5)) + ] + sqlite_instance.add_attack_results_to_memory(attack_results=attack_results) + results = sqlite_instance.get_scenario_results() retrieved_attacks = results[0].attack_results["Attack1"] - # Verify the conversation IDs are in the same order retrieved_conv_ids = [ar.conversation_id for ar in retrieved_attacks] - original_conv_ids = [ar.conversation_id for ar in attack_results] - assert retrieved_conv_ids == original_conv_ids + assert retrieved_conv_ids == [f"conv_{i}" for i in range(5)] -def test_stores_conversation_ids_only(sqlite_instance: MemoryInterface, sample_attack_results): - """Test that scenario results store only conversation IDs, not full AttackResult objects.""" - # Create and add scenario result - scenario_result = create_scenario_result( - name="Test Scenario", - attack_results={"Attack1": [sample_attack_results[0]]}, - ) +def test_stores_conversation_ids_only(sqlite_instance: MemoryInterface): + """Test that scenario results expose AttackResult objects with conversation IDs after hydration.""" + scenario_result = create_scenario_result(name="Test Scenario", attack_results={}) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - # Retrieve the scenario result to verify structure + ar = _make_attack_result_for_scenario( + scenario_result_id=scenario_result.id, + atomic_attack_name="Attack1", + objective_index=0, + conversation_id="conv_1", + ) + sqlite_instance.add_attack_results_to_memory(attack_results=[ar]) + results = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) assert len(results) == 1 @@ -299,24 +302,29 @@ def test_preserves_metadata(sqlite_instance: MemoryInterface): def test_multiple_scenarios_with_attacks(sqlite_instance: MemoryInterface): """Test retrieving multiple scenarios with their attack results populated.""" - # Create attack results for multiple scenarios - attack_results_scenario1 = [create_attack_result(f"conv_s1_{i}", f"S1 Objective {i}") for i in range(5)] - attack_results_scenario2 = [create_attack_result(f"conv_s2_{i}", f"S2 Objective {i}") for i in range(3)] + scenario1 = create_scenario_result(name="Scenario 1", attack_results={}) + scenario2 = create_scenario_result(name="Scenario 2", attack_results={}) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario1, scenario2]) - all_attack_results = attack_results_scenario1 + attack_results_scenario2 + all_attack_results = [ + _make_attack_result_for_scenario( + scenario_result_id=scenario1.id, + atomic_attack_name="Attack1", + objective_index=i, + conversation_id=f"conv_s1_{i}", + ) + for i in range(5) + ] + [ + _make_attack_result_for_scenario( + scenario_result_id=scenario2.id, + atomic_attack_name="Attack2", + objective_index=i, + conversation_id=f"conv_s2_{i}", + ) + for i in range(3) + ] sqlite_instance.add_attack_results_to_memory(attack_results=all_attack_results) - # Create multiple scenario results - scenario1 = create_scenario_result( - name="Scenario 1", - attack_results={"Attack1": attack_results_scenario1}, - ) - scenario2 = create_scenario_result( - name="Scenario 2", - attack_results={"Attack2": attack_results_scenario2}, - ) - sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario1, scenario2]) - # Retrieve all scenarios results = sqlite_instance.get_scenario_results() assert len(results) == 2 @@ -662,7 +670,7 @@ def _make_attack_result_for_scenario( outcome=AttackOutcome.SUCCESS, ): """Build an AttackResult pre-stamped with scenario linkage (mirrors what - the event handler does when an ExecutionAttribution is on the context).""" + the event handler does when a ScenarioExecutionAttribution is on the context).""" return AttackResult( conversation_id=conversation_id or f"conv-{atomic_attack_name}-{objective_index}", objective=f"objective-{atomic_attack_name}-{objective_index}", @@ -699,66 +707,6 @@ def test_get_scenario_results_hydrates_via_fk(sqlite_instance: MemoryInterface): assert [r.conversation_id for r in result.attack_results["b"]] == ["conv-b-0"] -def test_get_scenario_results_hydrates_via_legacy_manifest(sqlite_instance: MemoryInterface, sample_attack_results): - """When a scenario predates the FK column (its AttackResults have no FK set - but the manifest is populated), hydration falls back to the manifest path.""" - scenario_result = create_scenario_result( - name="Manifest-only Scenario", - attack_results={"legacy_attack": sample_attack_results[:2]}, - ) - sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - - [result] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) - assert list(result.attack_results.keys()) == ["legacy_attack"] - assert len(result.attack_results["legacy_attack"]) == 2 - - -def test_get_scenario_results_mixed_hydration_merges_and_dedupes( - sqlite_instance: MemoryInterface, sample_attack_results -): - """A scenario with BOTH FK-linked rows and legacy manifest rows merges - both sources. FK rows are authoritative; manifest entries already present - via FK are not duplicated.""" - import uuid as _uuid - from contextlib import closing - - from pyrit.memory.memory_models import AttackResultEntry - - # Build manifest from existing sample_attack_results. - scenario_result = create_scenario_result( - name="Mixed Scenario", - attack_results={"legacy_attack": sample_attack_results[:2]}, - ) - sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - - sid = scenario_result.id - # Add a new FK-linked row not in the manifest. - fk_only = _make_attack_result_for_scenario( - scenario_result_id=sid, atomic_attack_name="fk_attack", objective_index=0 - ) - sqlite_instance.add_attack_results_to_memory(attack_results=[fk_only]) - - # Retroactively stamp ONE of the existing manifest rows with the FK to - # verify dedup-by-attack_result_id: the row appears in both sources but - # the hydration must yield exactly one copy. UPDATE in place rather than - # re-INSERT (which would trip the PK unique constraint). - target_id = _uuid.UUID(sample_attack_results[0].attack_result_id) - with closing(sqlite_instance.get_session()) as session: - row = session.query(AttackResultEntry).filter_by(id=target_id).one() - row.scenario_result_id = sid - row.scenario_data = {"atomic_attack_name": "legacy_attack", "objective_index": 0} - session.commit() - - [result] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(sid)]) - assert "fk_attack" in result.attack_results - assert "legacy_attack" in result.attack_results - # legacy_attack should have 2 rows (one upgraded to FK + one still manifest-only), - # with no duplicates from the merge. - legacy_ids = [r.attack_result_id for r in result.attack_results["legacy_attack"]] - assert len(legacy_ids) == 2 - assert len(set(legacy_ids)) == 2 - - def test_get_attack_results_filters_by_scenario_result_id(sqlite_instance: MemoryInterface): """get_attack_results gains a scenario_result_id filter — replaces the removed error_attack_result_ids_json lookup path.""" @@ -824,26 +772,6 @@ def test_delete_scenario_sets_attack_result_fk_to_null(sqlite_instance: MemoryIn assert entry.scenario_data == {"atomic_attack_name": "a", "objective_index": 0} -def test_add_attack_results_to_scenario_is_deprecated_noop(sqlite_instance: MemoryInterface): - """add_attack_results_to_scenario is now a deprecated no-op shim — it - returns True without writing anything. Linkage is stamped per-result by - the attack persistence path instead.""" - scenario_result = create_scenario_result(name="Shim") - sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) - - # Should NOT raise and should return True (no-op success). - result = sqlite_instance.add_attack_results_to_scenario( - scenario_result_id=str(scenario_result.id), - atomic_attack_name="ignored", - attack_results=[], - ) - assert result is True - - # And no rows were added. - [hydrated] = sqlite_instance.get_scenario_results(scenario_result_ids=[str(scenario_result.id)]) - assert hydrated.attack_results == {} - - def test_update_scenario_run_state_targeted_update_preserves_manifest(sqlite_instance: MemoryInterface): """update_scenario_run_state must be a targeted UPDATE — it must not re-serialize the whole row and clobber the manifest column during the From 752d082d7313630c6435d91bdb850afc11b71072 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 19 May 2026 13:22:18 -0700 Subject: [PATCH 03/10] Add AtomicAttack tests for resume-path filter and attribution factory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps pyrit/scenario/core/atomic_attack.py coverage from 37% to 94% by exercising the three resume-critical surfaces the PR introduced or changed but that had no dedicated tests: - TestAtomicAttackFilterSeedGroupsByIndices: stable-identity filter that drops completed seeds while preserving each survivor's original index across successive filter calls. - TestAtomicAttackFilterSeedGroupsByObjectives: keeps the deprecated legacy path under test and asserts the DeprecationWarning fires until removed_in=0.16.0. - TestAtomicAttackAttributionFactory: the closure built in run_async when _scenario_result_id is set — no factory outside a Scenario, factory maps input_index -> original objective_index after filtering, and the snapshot is taken by value so post-call mutations cannot poison in-flight attributions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/unit/scenario/test_atomic_attack.py | 197 ++++++++++++++++++++++ 1 file changed, 197 insertions(+) diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index 5f86391759..0fd836d390 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -980,3 +980,200 @@ async def test_enrichment_skips_db_update_when_no_attack_result_id(self, mock_at await atomic.run_async() mock_memory.update_attack_result_by_id.assert_not_called() + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackFilterSeedGroupsByIndices: + """Tests for ``filter_seed_groups_by_indices`` — the stable-identity filter + that powers resume.""" + + def test_filters_out_completed_indices(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + atomic.filter_seed_groups_by_indices(completed_indices={0, 2}) + + assert atomic.seed_groups == [sample_seed_groups[1]] + assert atomic._original_indices == [1] + + def test_empty_completed_indices_is_noop(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + atomic.filter_seed_groups_by_indices(completed_indices=set()) + + assert atomic.seed_groups == sample_seed_groups + assert atomic._original_indices == [0, 1, 2] + + def test_all_indices_completed_clears_seed_groups(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + atomic.filter_seed_groups_by_indices(completed_indices={0, 1, 2}) + + assert atomic.seed_groups == [] + assert atomic._original_indices == [] + + def test_filter_preserves_original_indices_across_successive_calls(self, mock_attack, sample_seed_groups): + """After two successive filters, ``_original_indices`` must still reference + the *initial* construction positions, not the post-filter positions.""" + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + atomic.filter_seed_groups_by_indices(completed_indices={0}) + assert atomic._original_indices == [1, 2] + + atomic.filter_seed_groups_by_indices(completed_indices={1}) + # Index 1 (original position) must be dropped; index 2 (original) remains. + assert atomic._original_indices == [2] + assert atomic.seed_groups == [sample_seed_groups[2]] + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackFilterSeedGroupsByObjectives: + """Tests for the deprecated ``filter_seed_groups_by_objectives`` path — + still covered so we know the deprecation warning fires and the legacy + semantics keep working until ``removed_in=0.16.0``.""" + + def test_emits_deprecation_warning(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + with pytest.warns(DeprecationWarning, match="filter_seed_groups_by_objectives"): + atomic.filter_seed_groups_by_objectives(remaining_objectives=["objective2"]) + + assert [sg.objective.value for sg in atomic.seed_groups] == ["objective2"] + # The kept seed group was at original position 1. + assert atomic._original_indices == [1] + + def test_drops_all_when_no_objectives_match(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + + with pytest.warns(DeprecationWarning): + atomic.filter_seed_groups_by_objectives(remaining_objectives=["nope"]) + + assert atomic.seed_groups == [] + assert atomic._original_indices == [] + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackAttributionFactory: + """Tests for the ``attribution_factory`` closure built in ``run_async`` — + the bridge that lets the AttackExecutor stamp each per-task input index + back onto its *original* objective_index when persisting attack results.""" + + async def test_no_factory_passed_when_scenario_result_id_unset( + self, mock_attack, sample_seed_groups, sample_attack_results + ): + """Outside a Scenario, ``_scenario_result_id`` is None and the + executor must receive ``attribution_factory=None``.""" + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + assert atomic._scenario_result_id is None + + with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: + mock_exec.return_value = wrap_results(sample_attack_results) + await atomic.run_async() + + assert mock_exec.call_args.kwargs["attribution_factory"] is None + + async def test_factory_built_when_scenario_result_id_set( + self, mock_attack, sample_seed_groups, sample_attack_results + ): + """When the Scenario stamps ``_scenario_result_id`` onto the atomic + attack, ``run_async`` must build and pass a factory.""" + from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="MyAtomicAttack", + ) + atomic._scenario_result_id = "00000000-0000-0000-0000-000000000abc" + + with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: + mock_exec.return_value = wrap_results(sample_attack_results) + await atomic.run_async() + + factory = mock_exec.call_args.kwargs["attribution_factory"] + assert factory is not None + + # input_index 0 → original_indices[0] == 0 + attribution = factory(0) + assert isinstance(attribution, ScenarioExecutionAttribution) + assert attribution.scenario_result_id == "00000000-0000-0000-0000-000000000abc" + assert attribution.atomic_attack_name == "MyAtomicAttack" + assert attribution.objective_index == 0 + + # input_index 2 → original_indices[2] == 2 + assert factory(2).objective_index == 2 + + async def test_factory_maps_input_index_to_original_index_after_filtering( + self, mock_attack, sample_seed_groups, sample_attack_results + ): + """The whole point of the factory: after filtering out a completed + seed group, the per-task input index 0 must still map back to the + *original* objective_index of the surviving seed group, not 0.""" + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="ResumeAttack", + ) + atomic._scenario_result_id = "scenario-1" + # Simulate resume: original positions 0 and 1 are done; only position 2 remains. + atomic.filter_seed_groups_by_indices(completed_indices={0, 1}) + + with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: + mock_exec.return_value = wrap_results(sample_attack_results[:1]) + await atomic.run_async() + + factory = mock_exec.call_args.kwargs["attribution_factory"] + assert factory is not None + # input_index 0 (only remaining task) → original index 2. + assert factory(0).objective_index == 2 + + async def test_factory_captures_indices_by_value_not_reference( + self, mock_attack, sample_seed_groups, sample_attack_results + ): + """The closure must snapshot ``_original_indices`` at build time so + mutating the AtomicAttack after ``run_async`` started doesn't poison + in-flight attributions.""" + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + atomic._scenario_result_id = "scenario-1" + + with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: + mock_exec.return_value = wrap_results(sample_attack_results) + await atomic.run_async() + + factory = mock_exec.call_args.kwargs["attribution_factory"] + # Mutate after run_async has captured the snapshot. + atomic._original_indices = [999, 999, 999] + + # Factory still returns the snapshotted original indices. + assert [factory(i).objective_index for i in range(3)] == [0, 1, 2] From ddaa42c9414760ef915db4c305a2e5d04b8b0531 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 19 May 2026 17:49:02 -0700 Subject: [PATCH 04/10] Rename ScenarioExecutionAttribution to generic AttackResultAttribution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decouple the attack persistence path from scenario vocabulary. The attack layer now ships an opaque attribution dataclass (parent_id, parent_collection, position) — the scenario layer interprets those fields to mean (scenario_result_id, atomic_attack_name, objective_index). - ScenarioExecutionAttribution -> AttackResultAttribution (renamed module and class) - AttackResult.scenario_result_id / scenario_data -> attribution_parent_id / attribution_data - AttackResultEntry columns, index, and foreign key constraint renamed; migration 9c8b7a6d5e4f rewritten in place (still unreleased on this branch) - Replaced FK abbreviation with foreign key / ForeignKey in comments and docstrings The DB foreign key still targets ScenarioResultEntries.id; that is a relational fact, not a layering violation. The attack layer has no scenario-specific identifiers in its type signatures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../backend/services/scenario_run_service.py | 2 +- pyrit/executor/attack/core/attack_executor.py | 19 +++--- .../attack/core/attack_result_attribution.py | 51 ++++++++++++++ pyrit/executor/attack/core/attack_strategy.py | 41 ++++++------ .../core/scenario_execution_attribution.py | 43 ------------ ..._add_scenario_linkage_to_attack_results.py | 67 ++++++++++--------- pyrit/memory/memory_interface.py | 37 +++++----- pyrit/memory/memory_models.py | 31 ++++----- pyrit/models/attack_result.py | 14 ++-- pyrit/scenario/core/atomic_attack.py | 33 ++++----- pyrit/scenario/core/scenario.py | 35 +++++----- .../unit/backend/test_scenario_run_service.py | 7 +- .../attack/core/test_attack_executor.py | 40 +++++------ .../attack/core/test_attack_strategy.py | 52 +++++++------- .../test_interface_scenario_results.py | 42 ++++++------ tests/unit/memory/test_migration.py | 59 ++++++++-------- tests/unit/scenario/test_atomic_attack.py | 20 +++--- tests/unit/scenario/test_scenario.py | 14 ++-- .../scenario/test_scenario_partial_results.py | 9 +-- tests/unit/scenario/test_scenario_retry.py | 20 +++--- 20 files changed, 330 insertions(+), 306 deletions(-) create mode 100644 pyrit/executor/attack/core/attack_result_attribution.py delete mode 100644 pyrit/executor/attack/core/scenario_execution_attribution.py diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 85907a430f..6b04f2a71c 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -410,7 +410,7 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari error_type = scenario_result.error_type # Fallback: look up error from any persisted error AttackResults linked - # to this scenario via the new FK. + # to this scenario via the new attribution_parent_id foreign key. if not error: error_ars = self._memory.get_attack_results( scenario_result_id=scenario_result_id, diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index 7dee98f4ae..dd4a7140e6 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -19,12 +19,12 @@ ) from pyrit.executor.attack.core.attack_parameters import AttackParameters +from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.executor.attack.core.attack_strategy import ( AttackStrategy, AttackStrategyContextT, AttackStrategyResultT, ) -from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution from pyrit.models import SeedAttackGroup if TYPE_CHECKING: @@ -143,7 +143,7 @@ async def execute_attack_from_seed_groups_async( objective_scorer: Optional["TrueFalseScorer"] = None, field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -167,10 +167,11 @@ async def execute_attack_from_seed_groups_async( objectives fail. If False (default), raises the first exception. attribution_factory: Optional callable that maps an input index (the seed group's original index, parallel-safe and deterministic) to - a ``ScenarioExecutionAttribution``. When provided, each per-task + an ``AttackResultAttribution``. When provided, each per-task ``AttackContext`` is stamped with the attribution so the - resulting ``AttackResultEntry`` row carries the scenario FK + - scenario_data. When ``None``, no attribution is applied. + resulting ``AttackResultEntry`` row carries + ``attribution_parent_id`` + ``attribution_data``. When ``None``, + no attribution is applied. **broadcast_fields: Fields applied to all seed groups (e.g., memory_labels). Per-seed-group field_overrides take precedence. @@ -223,7 +224,7 @@ async def execute_attack_async( objectives: Sequence[str], field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -239,7 +240,7 @@ async def execute_attack_async( return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. attribution_factory: Optional callable mapping each input index to - a ScenarioExecutionAttribution. When provided, the per-task context is + a AttackResultAttribution. When provided, the per-task context is stamped with the attribution so the persistence path can record scenario linkage. **broadcast_fields: Fields applied to all objectives (e.g., memory_labels). @@ -291,7 +292,7 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], ScenarioExecutionAttribution]] = None, + attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Execute attacks in parallel with a list of pre-built parameters. @@ -304,7 +305,7 @@ async def _execute_with_params_list_async( params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. attribution_factory: Optional callable mapping each input index to - a ScenarioExecutionAttribution. When provided, the per-task context is + a AttackResultAttribution. When provided, the per-task context is stamped with the attribution so the persistence path can record scenario linkage. diff --git a/pyrit/executor/attack/core/attack_result_attribution.py b/pyrit/executor/attack/core/attack_result_attribution.py new file mode 100644 index 0000000000..bcabaefc9c --- /dev/null +++ b/pyrit/executor/attack/core/attack_result_attribution.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Generic attribution metadata that an upstream orchestrator can stamp onto an +``AttackContext`` so the persisted ``AttackResult`` carries the linkage back +to whatever produced it. + +The attack layer treats this as opaque infrastructure — three string-typed +fields, no scenario semantics. The orchestrator (e.g. ``Scenario``) interprets +them however it likes. Keeping the type in ``executor`` rather than +``scenario`` means the persistence path has no dependency on the +``pyrit.scenario`` package. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AttackResultAttribution: + """ + Attribution stamped onto an ``AttackContext`` by an upstream caller (the + ``AttackExecutor`` populates it via an ``attribution_factory``) and copied + onto the resulting ``AttackResult`` by the attack persistence path so the + DB row records its lineage. + + All three fields are opaque to the attack layer. The orchestrator chooses + what they mean and how to query them back later. For example, + ``Scenario`` uses ``parent_id`` for the scenario result UUID, + ``parent_collection`` for the atomic attack name, and ``position`` for + the original 0-based seed-group index. + + Attributes: + parent_id (str): The ID of the parent entity that owns this attack + execution. Persisted to ``AttackResultEntry.attribution_parent_id`` + and indexed (with a foreign key to ``ScenarioResultEntries.id``) + so per-parent hydration and resume lookups are direct. + parent_collection (str): A free-form label naming the per-parent + collection this result belongs to. Persisted into + ``AttackResultEntry.attribution_data``. + position (int): The 0-based position of this result within its + ``parent_collection``. Assigned **before** task execution so it is + deterministic and parallel-safe, and used as the stable resume key. + Persisted into ``AttackResultEntry.attribution_data``. + """ + + parent_id: str + parent_collection: str + position: int diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index b92a3be203..e4a4e58b59 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -36,7 +36,7 @@ if TYPE_CHECKING: from pyrit.executor.attack.core.attack_config import AttackScoringConfig - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.prompt_target import PromptTarget AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") @@ -72,11 +72,11 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _memory_labels_override: Optional[dict[str, str]] = None # Optional attribution from an upstream orchestrator (e.g. Scenario). When - # set, the persistence path stamps scenario_result_id + scenario_data onto - # the resulting AttackResult so it can be located later for hydration and - # resume. Set by AttackExecutor per-task before scheduling. Stays None for - # ad-hoc/direct attack execution outside any scenario. - _attribution: Optional[ScenarioExecutionAttribution] = None + # set, the persistence path stamps attribution_parent_id + attribution_data + # onto the resulting AttackResult so it can be located later for hydration + # and resume. Set by AttackExecutor per-task before scheduling. Stays None + # for ad-hoc/direct attack execution outside any orchestrator. + _attribution: Optional[AttackResultAttribution] = None # Convenience properties that delegate to params or overrides @property @@ -228,9 +228,9 @@ async def _on_post_execute( event_data.result.retry_events = collector.events event_data.result.total_retries = len(collector.events) - # Stamp scenario attribution onto the result before persistence so the - # AttackResultEntry row carries the FK + scenario_data. Outside scenarios - # _attribution is None and both fields stay None. + # Stamp attribution onto the result before persistence so the + # AttackResultEntry row records its lineage. Outside an orchestrator + # _attribution is None and both attribution fields stay None. self._stamp_attribution(context=event_data.context, result=event_data.result) self._logger.debug(f"Attack execution completed in {execution_time_ms}ms") @@ -245,12 +245,13 @@ def _stamp_attribution( result: AttackResult, ) -> None: """ - Copy scenario attribution from the AttackContext onto the AttackResult. + Copy attribution from the AttackContext onto the AttackResult. - Reads ``context._attribution`` (a ``ScenarioExecutionAttribution`` set by the - AttackExecutor when running inside a Scenario). When present, writes - ``scenario_result_id`` and a fixed-schema ``scenario_data`` dict onto - the result so they round-trip into ``AttackResultEntry``. + Reads ``context._attribution`` (an ``AttackResultAttribution`` set by + the AttackExecutor when an upstream orchestrator supplied a factory). + When present, writes ``attribution_parent_id`` and a fixed-schema + ``attribution_data`` dict onto the result so they round-trip into + ``AttackResultEntry``. Args: context: The per-task AttackContext. @@ -259,10 +260,10 @@ def _stamp_attribution( attribution = context._attribution if attribution is None: return - result.scenario_result_id = attribution.scenario_result_id - result.scenario_data = { - "atomic_attack_name": attribution.atomic_attack_name, - "objective_index": attribution.objective_index, + result.attribution_parent_id = attribution.parent_id + result.attribution_data = { + "parent_collection": attribution.parent_collection, + "position": attribution.position, } def _log_attack_outcome(self, result: AttackResult) -> None: @@ -329,8 +330,8 @@ async def _on_error_async( if context.start_time: error_result.execution_time_ms = int((end_time - context.start_time) * 1000) - # Stamp scenario attribution onto the error result so it is locatable - # via the scenario FK on resume. + # Stamp attribution onto the error result so it is locatable via the + # attribution_parent_id foreign key on resume. self._stamp_attribution(context=context, result=error_result) self._memory.add_attack_results_to_memory(attack_results=[error_result]) diff --git a/pyrit/executor/attack/core/scenario_execution_attribution.py b/pyrit/executor/attack/core/scenario_execution_attribution.py deleted file mode 100644 index 07b97d9d49..0000000000 --- a/pyrit/executor/attack/core/scenario_execution_attribution.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Typed attribution metadata used to link a persisted ``AttackResult`` to the -``Scenario`` that produced it. - -Lives in the ``executor`` layer (rather than ``scenario``) so the attack -persistence path — the consumer — does not introduce a dependency on -``pyrit.scenario``. -""" - -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class ScenarioExecutionAttribution: - """ - Scenario-linkage metadata stamped onto an ``AttackContext`` by the - ``AttackExecutor`` and copied onto the resulting ``AttackResult`` by the - attack persistence path so the row carries the scenario FK + scenario_data. - - Attributes: - scenario_result_id (str): The ID of the scenario result that produced - this attack execution. Persisted to - ``AttackResultEntry.scenario_result_id`` so per-scenario hydration - and resume can locate the row directly without relying on a JSON - manifest written at the end of an atomic attack. - atomic_attack_name (str): The unique key of the atomic attack within - the scenario (matches ``AtomicAttack.atomic_attack_name``). - Persisted into ``AttackResultEntry.scenario_data``. - objective_index (int): The 0-based original seed-group index (the - ``input_indices`` value from ``AttackExecutorResult``). Assigned - **before** task execution so it is deterministic and parallel-safe. - Persisted into ``AttackResultEntry.scenario_data`` and used as the - stable resume key (instead of the easily-duplicated objective text). - """ - - scenario_result_id: str - atomic_attack_name: str - objective_index: int diff --git a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py index 7d39f46d70..d5eeb2b910 100644 --- a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py +++ b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py @@ -2,9 +2,9 @@ # Licensed under the MIT license. """ -Add scenario_result_id FK + scenario_data JSON to AttackResultEntries; drop -ScenarioResultEntries.error_attack_result_ids_json; backfill linkage from the -existing attack_results_json manifest. +Add attribution_parent_id (foreign key) + attribution_data (JSON) to +AttackResultEntries; drop ScenarioResultEntries.error_attack_result_ids_json; +backfill the linkage from the existing attack_results_json manifest. Revision ID: 9c8b7a6d5e4f Revises: 7a1b2c3d4e5f @@ -35,41 +35,43 @@ def upgrade() -> None: """Apply this schema upgrade.""" - # AttackResultEntries: scenario linkage columns. + # AttackResultEntries: attribution / parent linkage columns. op.add_column( "AttackResultEntries", - sa.Column("scenario_result_id", CustomUUID(), nullable=True), + sa.Column("attribution_parent_id", CustomUUID(), nullable=True), ) op.add_column( "AttackResultEntries", - sa.Column("scenario_data", sa.JSON(), nullable=True), + sa.Column("attribution_data", sa.JSON(), nullable=True), ) op.create_index( - "ix_AttackResultEntries_scenario_result_id", + "ix_AttackResultEntries_attribution_parent_id", "AttackResultEntries", - ["scenario_result_id"], + ["attribution_parent_id"], ) - # FK with ON DELETE SET NULL: deleting a scenario nulls the FK on its - # AttackResults; scenario_data is retained as historical provenance. - # Use a batch operation for SQLite portability (no plain ALTER TABLE ADD FK). + # Foreign key with ON DELETE SET NULL: deleting a scenario nulls the + # attribution_parent_id on its AttackResults; attribution_data is retained + # as historical provenance. Use a batch operation for SQLite portability + # (no plain ALTER TABLE ADD CONSTRAINT for foreign keys on SQLite). with op.batch_alter_table("AttackResultEntries") as batch_op: batch_op.create_foreign_key( - "fk_attack_results_scenario", + "fk_attack_results_attribution_parent", "ScenarioResultEntries", - ["scenario_result_id"], + ["attribution_parent_id"], ["id"], ondelete="SET NULL", ) # ScenarioResultEntries: drop the not-yet-released error_attack_result_ids_json. - # Error AttackResults are now linkable via the new FK; the per-scenario - # manifest column is no longer used. Wrapped in a batch op for SQLite. + # Error AttackResults are now linkable via the new attribution_parent_id + # foreign key; the per-scenario manifest column is no longer used. + # Wrapped in a batch op for SQLite. with op.batch_alter_table("ScenarioResultEntries") as batch_op: batch_op.drop_column("error_attack_result_ids_json") - # Backfill scenario linkage from the existing attack_results_json manifest. - _backfill_scenario_linkage() + # Backfill attribution linkage from the existing attack_results_json manifest. + _backfill_attribution_linkage() def downgrade() -> None: @@ -80,21 +82,22 @@ def downgrade() -> None: sa.Column("error_attack_result_ids_json", sa.Unicode(), nullable=True), ) - # Drop FK + columns from AttackResultEntries. + # Drop foreign key + columns from AttackResultEntries. with op.batch_alter_table("AttackResultEntries") as batch_op: - batch_op.drop_constraint("fk_attack_results_scenario", type_="foreignkey") + batch_op.drop_constraint("fk_attack_results_attribution_parent", type_="foreignkey") - op.drop_index("ix_AttackResultEntries_scenario_result_id", table_name="AttackResultEntries") - op.drop_column("AttackResultEntries", "scenario_data") - op.drop_column("AttackResultEntries", "scenario_result_id") + op.drop_index("ix_AttackResultEntries_attribution_parent_id", table_name="AttackResultEntries") + op.drop_column("AttackResultEntries", "attribution_data") + op.drop_column("AttackResultEntries", "attribution_parent_id") -def _backfill_scenario_linkage() -> None: +def _backfill_attribution_linkage() -> None: """ Walk every ScenarioResultEntry and copy its attack_results_json manifest - into the new FK + scenario_data columns on AttackResultEntries. + into the new attribution_parent_id + attribution_data columns on + AttackResultEntries. - Idempotent: the ``WHERE scenario_result_id IS NULL`` guard prevents + Idempotent: the ``WHERE attribution_parent_id IS NULL`` guard prevents clobbering rows that were already linked (e.g. by a re-run of the migration, or by code that ran after the schema change but before this backfill). ``conversation_id`` is logically unique per AttackResult but is @@ -107,8 +110,8 @@ def _backfill_scenario_linkage() -> None: update_stmt = sa.text( 'UPDATE "AttackResultEntries" ' - "SET scenario_result_id = :sid, scenario_data = :sdata " - "WHERE conversation_id = :cid AND scenario_result_id IS NULL" + "SET attribution_parent_id = :sid, attribution_data = :sdata " + "WHERE conversation_id = :cid AND attribution_parent_id IS NULL" ) total_updates = 0 @@ -135,7 +138,7 @@ def _backfill_scenario_linkage() -> None: match_count = bind.execute( sa.text( 'SELECT COUNT(*) FROM "AttackResultEntries" ' - "WHERE conversation_id = :cid AND scenario_result_id IS NULL" + "WHERE conversation_id = :cid AND attribution_parent_id IS NULL" ), {"cid": conversation_id}, ).scalar() @@ -147,14 +150,12 @@ def _backfill_scenario_linkage() -> None: f"All matching rows will be linked to scenario {scenario_id}." ) - scenario_data = json.dumps( - {"atomic_attack_name": atomic_attack_name, "objective_index": objective_index} - ) + attribution_data = json.dumps({"parent_collection": atomic_attack_name, "position": objective_index}) result = bind.execute( update_stmt, { "sid": str(scenario_id), - "sdata": scenario_data, + "sdata": attribution_data, "cid": conversation_id, }, ) @@ -162,6 +163,6 @@ def _backfill_scenario_linkage() -> None: if total_updates or duplicate_warnings: logger.info( - f"Scenario linkage backfill: linked {total_updates} AttackResultEntries row(s); " + f"Attribution linkage backfill: linked {total_updates} AttackResultEntries row(s); " f"{duplicate_warnings} duplicate-conversation_id warning(s)." ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 5c28856d33..70f2ab5b7b 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1752,9 +1752,9 @@ def get_attack_results( A sequence of IdentifierFilter objects that allows filtering by various attack identifier JSON properties. Defaults to None. scenario_result_id (Optional[str], optional): Filter to attack results linked to a - specific scenario via the ``AttackResultEntry.scenario_result_id`` FK. Combined - with ``outcome=AttackOutcome.ERROR`` this is the replacement for the removed - per-scenario error_attack_result_ids manifest. Defaults to None. + specific scenario via the ``AttackResultEntry.attribution_parent_id`` foreign key. + Combined with ``outcome=AttackOutcome.ERROR`` this is the replacement for the + removed per-scenario error_attack_result_ids manifest. Defaults to None. Returns: Sequence[AttackResult]: A list of AttackResult objects that match the specified filters. @@ -1789,7 +1789,7 @@ def get_attack_results( if outcome: conditions.append(AttackResultEntry.outcome == outcome) if scenario_result_id: - conditions.append(AttackResultEntry.scenario_result_id == uuid.UUID(scenario_result_id)) + conditions.append(AttackResultEntry.attribution_parent_id == uuid.UUID(scenario_result_id)) if attack_classes: # Case-insensitive to mirror converter_classes; forgives casing drift in @@ -2205,14 +2205,15 @@ def _hydrate_scenario_attack_results( ) -> dict[uuid.UUID, dict[str, list[AttackResult]]]: """ Fetch every ``AttackResult`` linked to the given scenarios via the - ``AttackResultEntry.scenario_result_id`` FK in a single batched query, - then group by scenario + ``atomic_attack_name`` and sort each group by - ``scenario_data["objective_index"]``. + ``AttackResultEntry.attribution_parent_id`` foreign key in a single + batched query, then group by scenario + ``parent_collection`` (which + the scenario layer uses for the atomic attack name) and sort each + group by ``attribution_data["position"]``. - FK linkage is the sole source of truth — set at write-time by the - attack persistence path when a ``ScenarioExecutionAttribution`` is on - the context. Rows without a valid ``scenario_data`` payload are skipped - (and logged) rather than guessed at. + Foreign-key linkage is the sole source of truth — set at write-time by + the attack persistence path when an ``AttackResultAttribution`` is on + the context. Rows without a valid ``attribution_data`` payload are + skipped (and logged) rather than guessed at. Returns: dict[uuid.UUID, dict[str, list[AttackResult]]]: Mapping of @@ -2225,27 +2226,27 @@ def _hydrate_scenario_attack_results( scenario_ids = [entry.id for entry in entries] attack_rows = self._execute_batched_query( AttackResultEntry, - batch_column=AttackResultEntry.scenario_result_id, + batch_column=AttackResultEntry.attribution_parent_id, batch_values=scenario_ids, ) grouped: dict[uuid.UUID, dict[str, list[tuple[int, AttackResult]]]] = {entry.id: {} for entry in entries} for row in attack_rows: - scenario_id = row.scenario_result_id + scenario_id = row.attribution_parent_id if scenario_id is None or scenario_id not in grouped: continue - data = row.scenario_data or {} - name = data.get("atomic_attack_name") + data = row.attribution_data or {} + name = data.get("parent_collection") if not name: logger.debug( - f"Skipping AttackResultEntry {row.id} during hydration: scenario_data missing atomic_attack_name" + f"Skipping AttackResultEntry {row.id} during hydration: attribution_data missing parent_collection" ) continue - objective_index = data.get("objective_index") - sort_key = objective_index if isinstance(objective_index, int) else 0 + position = data.get("position") + sort_key = position if isinstance(position, int) else 0 grouped[scenario_id].setdefault(name, []).append((sort_key, row.get_attack_result())) return { diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 562733fadd..9400a78cc5 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -745,17 +745,18 @@ class AttackResultEntry(Base): retry_events_json: Mapped[Optional[str]] = mapped_column(Unicode, nullable=True) total_retries = mapped_column(INTEGER, nullable=True, default=0) - # Scenario linkage (set when the AttackResult is produced inside a Scenario). - # scenario_result_id is an indexed FK so per-scenario hydration and resume - # queries are direct lookups (no JSON manifest required, no orphaning if a - # scenario is interrupted mid-AtomicAttack). - # scenario_data is a documented-fixed-schema JSON blob with two keys: - # atomic_attack_name (str) and objective_index (int). When the AttackResult - # is created outside a Scenario both fields remain NULL. - scenario_result_id: Mapped[Optional[uuid.UUID]] = mapped_column( + # Attribution / parent linkage (set when the AttackResult is produced + # inside an orchestrator that supplies an AttackResultAttribution, e.g. a + # Scenario). attribution_parent_id is an indexed foreign key so per-parent + # hydration and resume queries are direct lookups (no JSON manifest + # required, no orphaning if the orchestrator is interrupted mid-run). + # attribution_data is a documented-fixed-schema JSON blob with two keys: + # parent_collection (str) and position (int). When the AttackResult is + # created outside an orchestrator both fields remain NULL. + attribution_parent_id: Mapped[Optional[uuid.UUID]] = mapped_column( CustomUUID, ForeignKey("ScenarioResultEntries.id", ondelete="SET NULL"), nullable=True, index=True ) - scenario_data: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) + attribution_data: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) last_response: Mapped[Optional["PromptMemoryEntry"]] = relationship( "PromptMemoryEntry", @@ -831,10 +832,10 @@ def __init__(self, *, entry: AttackResult) -> None: ) self.total_retries = entry.total_retries - # Scenario linkage (set by the attack persistence path when a - # ScenarioExecutionAttribution is present on the AttackContext; otherwise None) - self.scenario_result_id = uuid.UUID(entry.scenario_result_id) if entry.scenario_result_id else None - self.scenario_data = entry.scenario_data + # Attribution / parent linkage (set by the attack persistence path when + # an AttackResultAttribution is present on the AttackContext; otherwise None) + self.attribution_parent_id = uuid.UUID(entry.attribution_parent_id) if entry.attribution_parent_id else None + self.attribution_data = entry.attribution_data @staticmethod def _get_id_as_uuid(obj: Any) -> Optional[uuid.UUID]: @@ -948,8 +949,8 @@ def get_attack_result(self) -> AttackResult: error_traceback=self.error_traceback, retry_events=retry_events, total_retries=self.total_retries or 0, - scenario_result_id=str(self.scenario_result_id) if self.scenario_result_id else None, - scenario_data=self.scenario_data, + attribution_parent_id=str(self.attribution_parent_id) if self.attribution_parent_id else None, + attribution_data=self.attribution_data, ) diff --git a/pyrit/models/attack_result.py b/pyrit/models/attack_result.py index d8c6cf9a98..b8cfdffcc4 100644 --- a/pyrit/models/attack_result.py +++ b/pyrit/models/attack_result.py @@ -107,13 +107,13 @@ class AttackResult(StrategyResult): retry_events: list[RetryEvent] = field(default_factory=list) total_retries: int = 0 - # Scenario linkage (infrastructure-managed). Set by the attack persistence - # path when a ScenarioExecutionAttribution is present on the AttackContext. User - # code should not set these directly; ad-hoc AttackResults created outside - # a Scenario leave both fields as None and the corresponding DB columns - # remain NULL. - scenario_result_id: str | None = None - scenario_data: dict[str, Any] | None = None + # Attribution / parent linkage (infrastructure-managed). Set by the attack + # persistence path when an AttackResultAttribution is present on the + # AttackContext. User code should not set these directly; ad-hoc + # AttackResults created outside an orchestrator leave both fields as None + # and the corresponding DB columns remain NULL. + attribution_parent_id: str | None = None + attribution_data: dict[str, Any] | None = None @property def attack_identifier(self) -> Optional[ComponentIdentifier]: diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index 6ef19d11f3..4699041844 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -20,7 +20,7 @@ from pyrit.common.deprecation import print_deprecation_message from pyrit.executor.attack import AttackExecutor, AttackStrategy from pyrit.executor.attack.core.attack_executor import AttackExecutorResult -from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution +from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.identifiers import build_atomic_attack_identifier from pyrit.identifiers.evaluation_identifier import AtomicAttackEvaluationIdentifier from pyrit.memory import CentralMemory @@ -120,18 +120,18 @@ def __init__( self._seed_groups = seed_groups # Original positions for each currently-active seed group. Used to map a - # per-task input index (from the executor) back to the stable - # objective_index stored in AttackResultEntry.scenario_data. Resume - # filtering preserves the original indices so newly-persisted results - # don't collide with already-persisted ones. + # per-task input index (from the executor) back to the stable position + # stored in AttackResultEntry.attribution_data. Resume filtering + # preserves the original indices so newly-persisted results don't + # collide with already-persisted ones. self._original_indices: list[int] = list(range(len(seed_groups))) self._adversarial_chat = adversarial_chat self._objective_scorer = objective_scorer self._memory_labels = memory_labels or {} self._attack_execute_params = attack_execute_params # Set by Scenario._execute_scenario_async before run_async. When set, - # each persisted AttackResult is linked to this scenario via the FK on - # AttackResultEntry. + # each persisted AttackResult is linked to this scenario via the + # attribution_parent_id foreign key on AttackResultEntry. self._scenario_result_id: str | None = None logger.info( @@ -200,7 +200,7 @@ def filter_seed_groups_by_indices(self, *, completed_indices: set[int]) -> None: (its position when the ``AtomicAttack`` was first constructed) is tracked across filtering, so the executor can later map each per-task input index back to the original index when stamping - ``scenario_data["objective_index"]``. This prevents newly-persisted + ``attribution_data["position"]``. This prevents newly-persisted results from colliding with already-persisted ones on resume. Args: @@ -273,19 +273,22 @@ async def run_async( # Build an attribution factory when this atomic attack is being # executed inside a Scenario. The factory maps each per-task input # index (position in the currently-active seed_groups list) back to - # the original objective_index so resume can locate already-done - # work and newly-persisted results don't collide with old ones. + # the original position so resume can locate already-done work and + # newly-persisted results don't collide with old ones. The Scenario + # uses parent_collection for the atomic attack name and position + # for the original seed-group index; the attack layer treats both + # as opaque strings/ints. attribution_factory = None if self._scenario_result_id is not None: scenario_id = self._scenario_result_id name = self.atomic_attack_name original_indices = list(self._original_indices) - def attribution_factory(input_index: int) -> ScenarioExecutionAttribution: - return ScenarioExecutionAttribution( - scenario_result_id=scenario_id, - atomic_attack_name=name, - objective_index=original_indices[input_index], + def attribution_factory(input_index: int) -> AttackResultAttribution: + return AttackResultAttribution( + parent_id=scenario_id, + parent_collection=name, + position=original_indices[input_index], ) results = await executor.execute_attack_from_seed_groups_async( diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 4120f05d41..2e505ad90a 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -695,10 +695,10 @@ async def initialize_async( # Enforce atomic_attack_name uniqueness. Duplicate names cause result # aggregation maps keyed by name to silently collapse rows, and - # FK-based resume cannot distinguish them. Some existing scenarios - # (e.g. Encoding) construct multiple AtomicAttacks with the same name - # for different converter configs, so this is a warning for now; in a - # future release it will become an error. + # foreign-key-based resume cannot distinguish them. Some existing + # scenarios (e.g. Encoding) construct multiple AtomicAttacks with the + # same name for different converter configs, so this is a warning for + # now; in a future release it will become an error. names = [aa.atomic_attack_name for aa in self._atomic_attacks] duplicates = sorted({name for name in names if names.count(name) > 1}) if duplicates: @@ -856,7 +856,7 @@ def _get_completed_objective_indices_for_attack(self, *, atomic_attack_name: str Get the set of objective_index values that have already been completed (non-error) for a specific atomic attack inside this scenario. - Queries AttackResultEntry rows directly by ``scenario_result_id`` — + Queries AttackResultEntry rows directly 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. @@ -876,13 +876,13 @@ def _get_completed_objective_indices_for_attack(self, *, atomic_attack_name: str for row in rows: if row.outcome == AttackOutcome.ERROR: continue - if row.scenario_data is None: + if row.attribution_data is None: continue - if row.scenario_data.get("atomic_attack_name") != atomic_attack_name: + if row.attribution_data.get("parent_collection") != atomic_attack_name: continue - objective_index = row.scenario_data.get("objective_index") - if isinstance(objective_index, int): - completed_indices.add(objective_index) + position = row.attribution_data.get("position") + if isinstance(position, int): + completed_indices.add(position) except Exception as e: logger.warning( f"Failed to retrieve completed objective indices for atomic attack '{atomic_attack_name}': {str(e)}" @@ -1168,9 +1168,9 @@ async def _execute_scenario_async(self) -> ScenarioResult: start=completed_count + 1, ): # Stamp the scenario id onto the atomic attack so each persisted - # AttackResult carries the FK linkage. This is what enables - # mid-run interruption recovery (results are visible without the - # post-atomic-attack bulk manifest write). + # AttackResult carries the attribution_parent_id linkage. This + # is what enables mid-run interruption recovery (results are + # visible without the post-atomic-attack bulk manifest write). atomic_attack._scenario_result_id = scenario_result_id logger.info( @@ -1203,10 +1203,11 @@ async def _execute_scenario_async(self) -> ScenarioResult: logger.error(f" Incomplete objective '{obj[:50]}...': {str(exc)}") # Error AttackResults are linked to this scenario via the - # scenario_result_id FK on AttackResultEntry (stamped by - # the attack event handler when a ScenarioExecutionAttribution - # is on the context). The previous per-scenario error_id - # manifest is no longer needed. + # attribution_parent_id foreign key on AttackResultEntry + # (stamped by the attack event handler when an + # AttackResultAttribution is on the context). The + # previous per-scenario error_id manifest is no longer + # needed. # Mark scenario as failed error_msg = ( diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 645c6acf21..65169c90f1 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -303,9 +303,10 @@ def test_get_run_returns_existing_run(self, mock_memory) -> None: def test_get_run_falls_back_to_persisted_error(self, mock_memory) -> None: """Test that get_run extracts error from persisted error AttackResult when no active task. - After the FK-based scenario linkage refactor, error AttackResults are - located via ``get_attack_results(scenario_result_id=..., outcome=ERROR)`` - rather than via a per-scenario error_attack_result_ids manifest. + After the foreign-key-based scenario linkage refactor, error + AttackResults are located via + ``get_attack_results(scenario_result_id=..., outcome=ERROR)`` rather + than via a per-scenario error_attack_result_ids manifest. """ db_result = _make_db_scenario_result(result_id="sr-fail", run_state="FAILED") diff --git a/tests/unit/executor/attack/core/test_attack_executor.py b/tests/unit/executor/attack/core/test_attack_executor.py index ff81e703c9..ddf2af6cf1 100644 --- a/tests/unit/executor/attack/core/test_attack_executor.py +++ b/tests/unit/executor/attack/core/test_attack_executor.py @@ -338,7 +338,7 @@ async def test_validates_explicit_empty_field_overrides_for_seed_groups(self): @pytest.mark.usefixtures("patch_central_database") class TestAttributionFactoryPropagation: - """Tests for ScenarioExecutionAttribution propagation through the AttackExecutor. + """Tests for AttackResultAttribution propagation through the AttackExecutor. The executor stamps ``context._attribution = attribution_factory(input_index)`` on each per-task context. The input_index must be the seed group's original @@ -347,7 +347,7 @@ class TestAttributionFactoryPropagation: """ async def test_attribution_factory_sets_per_task_attribution_in_order(self): - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution attack = create_mock_attack() seen_indices: list[int] = [] @@ -356,18 +356,18 @@ async def capture(context): # Read the attribution stamped by the executor BEFORE the task runs. attr = context._attribution assert attr is not None - seen_indices.append(attr.objective_index) + seen_indices.append(attr.position) return create_attack_result(context.params.objective) attack.execute_with_context_async = AsyncMock(side_effect=capture) seed_groups = [create_seed_group(f"obj-{i}") for i in range(4)] - def factory(idx: int) -> ScenarioExecutionAttribution: - return ScenarioExecutionAttribution( - scenario_result_id="sid", - atomic_attack_name="atomic", - objective_index=idx, + def factory(idx: int) -> AttackResultAttribution: + return AttackResultAttribution( + parent_id="sid", + parent_collection="atomic", + position=idx, ) executor = AttackExecutor(max_concurrency=1) @@ -377,17 +377,17 @@ def factory(idx: int) -> ScenarioExecutionAttribution: attribution_factory=factory, ) - # Each per-task context saw a distinct, contiguous objective_index — the + # Each per-task context saw a distinct, contiguous position — the # seed group's original input position. assert sorted(seen_indices) == [0, 1, 2, 3] assert len(result.completed_results) == 4 async def test_attribution_factory_parallel_safe_with_high_concurrency(self): - """At max_concurrency > 1, each task still receives its own objective_index + """At max_concurrency > 1, each task still receives its own position regardless of completion order. Out-of-order completion must not produce - stamped indices that point at the wrong seed group. + stamped positions that point at the wrong seed group. """ - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution attack = create_mock_attack() seen: dict[str, int] = {} @@ -396,20 +396,20 @@ async def out_of_order(context): attr = context._attribution assert attr is not None # Reverse-delay the tasks so completion order is the inverse of input - # order. The stamped index must still match the seed group. - await asyncio.sleep(0.005 * (10 - attr.objective_index)) - seen[context.params.objective] = attr.objective_index + # order. The stamped position must still match the seed group. + await asyncio.sleep(0.005 * (10 - attr.position)) + seen[context.params.objective] = attr.position return create_attack_result(context.params.objective) attack.execute_with_context_async = AsyncMock(side_effect=out_of_order) seed_groups = [create_seed_group(f"obj-{i}") for i in range(6)] - def factory(idx: int) -> ScenarioExecutionAttribution: - return ScenarioExecutionAttribution( - scenario_result_id="sid", - atomic_attack_name="atomic", - objective_index=idx, + def factory(idx: int) -> AttackResultAttribution: + return AttackResultAttribution( + parent_id="sid", + parent_collection="atomic", + position=idx, ) executor = AttackExecutor(max_concurrency=6) diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 22ed8bee0c..b44656f602 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -613,17 +613,17 @@ async def test_on_event_handles_other_events(self, event_handler, sample_attack_ async def test_on_post_execute_stamps_scenario_attribution_when_present( self, sample_attack_context, sample_attack_result, mock_memory ): - """When the context carries a ScenarioExecutionAttribution, the persisted - AttackResult must have scenario_result_id + scenario_data populated.""" - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + """When the context carries an AttackResultAttribution, the persisted + AttackResult must have attribution_parent_id + attribution_data populated.""" + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() sample_attack_context.start_time = 100.0 - sample_attack_context._attribution = ScenarioExecutionAttribution( - scenario_result_id="scenario-1", - atomic_attack_name="atomic_a", - objective_index=3, + sample_attack_context._attribution = AttackResultAttribution( + parent_id="scenario-1", + parent_collection="atomic_a", + position=3, ) event_data = StrategyEventData( @@ -635,17 +635,17 @@ async def test_on_post_execute_stamps_scenario_attribution_when_present( ) await handler.on_event(event_data) - assert sample_attack_result.scenario_result_id == "scenario-1" - assert sample_attack_result.scenario_data == { - "atomic_attack_name": "atomic_a", - "objective_index": 3, + assert sample_attack_result.attribution_parent_id == "scenario-1" + assert sample_attack_result.attribution_data == { + "parent_collection": "atomic_a", + "position": 3, } async def test_on_post_execute_no_attribution_leaves_fields_none( self, sample_attack_context, sample_attack_result, mock_memory ): - """Outside a Scenario, _attribution is None and the FK + scenario_data - fields on the persisted AttackResult must stay None.""" + """Outside a Scenario, _attribution is None and the attribution fields + on the persisted AttackResult must stay None.""" with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() sample_attack_context.start_time = 100.0 @@ -660,21 +660,21 @@ async def test_on_post_execute_no_attribution_leaves_fields_none( ) await handler.on_event(event_data) - assert sample_attack_result.scenario_result_id is None - assert sample_attack_result.scenario_data is None + assert sample_attack_result.attribution_parent_id is None + assert sample_attack_result.attribution_data is None async def test_on_error_stamps_scenario_attribution_when_present(self, sample_attack_context, mock_memory): - """Error AttackResults must also carry the scenario FK so error lookups - via get_attack_results(scenario_result_id=..., outcome=ERROR) work.""" - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + """Error AttackResults must also carry the attribution foreign key so + error lookups via get_attack_results(scenario_result_id=..., outcome=ERROR) work.""" + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() sample_attack_context.start_time = 100.0 - sample_attack_context._attribution = ScenarioExecutionAttribution( - scenario_result_id="scenario-err", - atomic_attack_name="atomic_err", - objective_index=7, + sample_attack_context._attribution = AttackResultAttribution( + parent_id="scenario-err", + parent_collection="atomic_err", + position=7, ) event_data = StrategyEventData( @@ -690,10 +690,10 @@ async def test_on_error_stamps_scenario_attribution_when_present(self, sample_at call = mock_memory.add_attack_results_to_memory.call_args persisted = call.kwargs["attack_results"][0] assert persisted.outcome == AttackOutcome.ERROR - assert persisted.scenario_result_id == "scenario-err" - assert persisted.scenario_data == { - "atomic_attack_name": "atomic_err", - "objective_index": 7, + assert persisted.attribution_parent_id == "scenario-err" + assert persisted.attribution_data == { + "parent_collection": "atomic_err", + "position": 7, } diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index f56516b2ba..313f5434f0 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -657,7 +657,8 @@ def test_combined_filters(sqlite_instance: MemoryInterface): # ============================================================================= -# Scenario linkage (FK + scenario_data on AttackResultEntry) hydration tests +# Scenario linkage (attribution_parent_id foreign key + attribution_data on +# AttackResultEntry) hydration tests # ============================================================================= @@ -670,24 +671,24 @@ def _make_attack_result_for_scenario( outcome=AttackOutcome.SUCCESS, ): """Build an AttackResult pre-stamped with scenario linkage (mirrors what - the event handler does when a ScenarioExecutionAttribution is on the context).""" + the event handler does when an AttackResultAttribution is on the context).""" return AttackResult( conversation_id=conversation_id or f"conv-{atomic_attack_name}-{objective_index}", objective=f"objective-{atomic_attack_name}-{objective_index}", outcome=outcome, executed_turns=1, - scenario_result_id=str(scenario_result_id), - scenario_data={"atomic_attack_name": atomic_attack_name, "objective_index": objective_index}, + attribution_parent_id=str(scenario_result_id), + attribution_data={"parent_collection": atomic_attack_name, "position": objective_index}, ) -def test_get_scenario_results_hydrates_via_fk(sqlite_instance: MemoryInterface): - """When AttackResultEntry rows carry the scenario_result_id FK, hydration - picks them up directly — without needing the legacy attack_results_json - manifest. This is the path that makes mid-AtomicAttack interruption-recovery - work.""" +def test_get_scenario_results_hydrates_via_foreign_key(sqlite_instance: MemoryInterface): + """When AttackResultEntry rows carry the attribution_parent_id foreign key, + hydration picks them up directly — without needing the legacy + attack_results_json manifest. This is the path that makes mid-AtomicAttack + interruption-recovery work.""" scenario_result = create_scenario_result( - name="FK-only Scenario", + name="ForeignKey-only Scenario", attack_results={}, # manifest intentionally empty ) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) @@ -735,15 +736,16 @@ def test_get_attack_results_filters_by_scenario_result_id(sqlite_instance: Memor assert [r.conversation_id for r in only_errors] == [err.conversation_id] -def test_delete_scenario_sets_attack_result_fk_to_null(sqlite_instance: MemoryInterface): - """ON DELETE SET NULL: deleting the parent ScenarioResultEntry nulls the FK - on its linked AttackResultEntries but the AttackResultEntries survive - (scenario_data is retained as historical provenance). +def test_delete_scenario_sets_attack_result_foreign_key_to_null(sqlite_instance: MemoryInterface): + """ON DELETE SET NULL: deleting the parent ScenarioResultEntry nulls the + attribution_parent_id foreign key on its linked AttackResultEntries but + the AttackResultEntries survive (attribution_data is retained as + historical provenance). Note: SQLite does not enforce foreign keys by default; this test enables them on the session for the duration of the delete to verify the ON DELETE SET NULL clause works. Production deployments using SQL Server - enforce FKs by default. + enforce foreign keys by default. """ from contextlib import closing @@ -758,18 +760,18 @@ def test_delete_scenario_sets_attack_result_fk_to_null(sqlite_instance: MemoryIn ar = _make_attack_result_for_scenario(scenario_result_id=sid, atomic_attack_name="a", objective_index=0) sqlite_instance.add_attack_results_to_memory(attack_results=[ar]) - # Enable FKs for the delete and verify the SET NULL clause fires. + # Enable foreign keys for the delete and verify the SET NULL clause fires. with closing(sqlite_instance.get_session()) as session: session.execute(_sql_text("PRAGMA foreign_keys = ON")) session.query(ScenarioResultEntry).filter_by(id=sid).delete() session.commit() - # The AttackResult survives, but its FK is now NULL. scenario_data is - # retained as historical provenance. + # The AttackResult survives, but its foreign key is now NULL. + # attribution_data is retained as historical provenance. with closing(sqlite_instance.get_session()) as session: entry = session.query(AttackResultEntry).filter_by(conversation_id=ar.conversation_id).one() - assert entry.scenario_result_id is None - assert entry.scenario_data == {"atomic_attack_name": "a", "objective_index": 0} + assert entry.attribution_parent_id is None + assert entry.attribution_data == {"parent_collection": "a", "position": 0} def test_update_scenario_run_state_targeted_update_preserves_manifest(sqlite_instance: MemoryInterface): diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 7a7d94e655..3b283a2347 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -206,7 +206,7 @@ def test_migration_downgrade_creates_proper_structure(): # ============================================================================= -# Backfill tests for the scenario_result_id FK migration +# Backfill tests for the attribution_parent_id foreign key migration # ============================================================================= @@ -252,8 +252,9 @@ def _config_for(connection): def test_backfill_links_attack_results_via_conversation_id(): - """Upgrading from the pre-FK revision backfills scenario_result_id + - scenario_data on AttackResultEntries by matching conversation_id.""" + """Upgrading from the pre-foreign-key revision backfills + attribution_parent_id + attribution_data on AttackResultEntries by + matching conversation_id.""" import json with tempfile.TemporaryDirectory() as temp_dir: @@ -283,32 +284,33 @@ def test_backfill_links_attack_results_via_conversation_id(): rows = connection.execute( text( - "SELECT conversation_id, scenario_result_id, scenario_data " + "SELECT conversation_id, attribution_parent_id, attribution_data " 'FROM "AttackResultEntries" ORDER BY conversation_id' ) ).fetchall() results_by_conv = {r[0]: (r[1], r[2]) for r in rows} - # All three rows now point at the scenario via the new FK. + # All three rows now point at the scenario via the new foreign key. for conv in ("conv-a-0", "conv-a-1", "conv-b-0"): assert results_by_conv[conv][0] == sid, f"{conv} should be backfilled" - # scenario_data carries atomic_attack_name and the 0-based manifest index. + # attribution_data carries parent_collection (the atomic attack name) + # and the 0-based manifest position. sd_a0 = json.loads(results_by_conv["conv-a-0"][1]) sd_a1 = json.loads(results_by_conv["conv-a-1"][1]) sd_b0 = json.loads(results_by_conv["conv-b-0"][1]) - assert sd_a0 == {"atomic_attack_name": "a", "objective_index": 0} - assert sd_a1 == {"atomic_attack_name": "a", "objective_index": 1} - assert sd_b0 == {"atomic_attack_name": "b", "objective_index": 0} + assert sd_a0 == {"parent_collection": "a", "position": 0} + assert sd_a1 == {"parent_collection": "a", "position": 1} + assert sd_b0 == {"parent_collection": "b", "position": 0} finally: engine.dispose() def test_backfill_is_idempotent_and_does_not_clobber_existing_linkage(): - """The backfill is safe to re-run: rows that already carry a - ``scenario_result_id`` are not overwritten (the WHERE IS NULL guard). We + """The backfill is safe to re-run: rows that already carry an + ``attribution_parent_id`` are not overwritten (the WHERE IS NULL guard). We verify by upgrading, manually retargeting a row, then downgrading + re-upgrading and asserting the manual retarget survives.""" import json @@ -330,30 +332,31 @@ def test_backfill_is_idempotent_and_does_not_clobber_existing_linkage(): ) command.upgrade(config, _SCENARIO_LINKAGE_REV) - # Manually retarget the row to a DIFFERENT scenario_result_id — + # Manually retarget the row to a DIFFERENT attribution_parent_id — # simulate code that already linked it post-upgrade. connection.execute( - text('UPDATE "AttackResultEntries" SET scenario_result_id = :sid WHERE conversation_id = :conv'), + text('UPDATE "AttackResultEntries" SET attribution_parent_id = :sid WHERE conversation_id = :conv'), {"sid": sid_manual, "conv": "conv-shared"}, ) # Downgrade then re-upgrade to re-run the backfill. command.downgrade(config, _PREV_REV) - # After downgrade the FK column is gone, but the manifest still - # references conv-shared. On re-upgrade, the backfill should - # NOT clobber sid_manual because the column was just re-added - # as NULL — actually downgrade DROPS the column data, so on - # re-upgrade the row will start at NULL and get linked again. - # The test we want is: re-running the backfill while a row - # already has a non-NULL FK does not overwrite it. We exercise - # that with a fresh second upgrade-then-no-op-re-upgrade. + # After downgrade the foreign key column is gone, but the + # manifest still references conv-shared. On re-upgrade, the + # backfill should NOT clobber sid_manual because the column was + # just re-added as NULL — actually downgrade DROPS the column + # data, so on re-upgrade the row will start at NULL and get + # linked again. The test we want is: re-running the backfill + # while a row already has a non-NULL foreign key does not + # overwrite it. We exercise that with a fresh second + # upgrade-then-no-op-re-upgrade. command.upgrade(config, _SCENARIO_LINKAGE_REV) # First upgrade after downgrade re-links it to sid_old (the # manifest source). Now manually retarget again. connection.execute( - text('UPDATE "AttackResultEntries" SET scenario_result_id = :sid WHERE conversation_id = :conv'), + text('UPDATE "AttackResultEntries" SET attribution_parent_id = :sid WHERE conversation_id = :conv'), {"sid": sid_manual, "conv": "conv-shared"}, ) @@ -398,14 +401,14 @@ def test_backfill_is_idempotent_and_does_not_clobber_existing_linkage(): _original_get_bind = _op_mod.get_bind _op_mod.get_bind = lambda: connection try: - mig._backfill_scenario_linkage() + mig._backfill_attribution_linkage() finally: _op_mod.get_bind = _original_get_bind # The row's manual retarget MUST survive — the IS NULL guard # prevents the backfill from overwriting it. row = connection.execute( - text('SELECT scenario_result_id FROM "AttackResultEntries" WHERE conversation_id = :conv'), + text('SELECT attribution_parent_id FROM "AttackResultEntries" WHERE conversation_id = :conv'), {"conv": "conv-shared"}, ).scalar_one() assert row == sid_manual @@ -445,14 +448,14 @@ def test_migration_downgrade_restores_dropped_column(): command.upgrade(config, _SCENARIO_LINKAGE_REV) attack_cols_up = {c["name"] for c in inspect(connection).get_columns("AttackResultEntries")} - assert "scenario_result_id" in attack_cols_up - assert "scenario_data" in attack_cols_up + assert "attribution_parent_id" in attack_cols_up + assert "attribution_data" in attack_cols_up command.downgrade(config, _PREV_REV) attack_cols_down = {c["name"] for c in inspect(connection).get_columns("AttackResultEntries")} - assert "scenario_result_id" not in attack_cols_down - assert "scenario_data" not in attack_cols_down + assert "attribution_parent_id" not in attack_cols_down + assert "attribution_data" not in attack_cols_down scenario_cols = {c["name"] for c in inspect(connection).get_columns("ScenarioResultEntries")} assert "error_attack_result_ids_json" in scenario_cols diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index 0fd836d390..8ed70647ba 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -1079,7 +1079,7 @@ def test_drops_all_when_no_objectives_match(self, mock_attack, sample_seed_group class TestAtomicAttackAttributionFactory: """Tests for the ``attribution_factory`` closure built in ``run_async`` — the bridge that lets the AttackExecutor stamp each per-task input index - back onto its *original* objective_index when persisting attack results.""" + back onto its *original* position when persisting attack results.""" async def test_no_factory_passed_when_scenario_result_id_unset( self, mock_attack, sample_seed_groups, sample_attack_results @@ -1104,7 +1104,7 @@ async def test_factory_built_when_scenario_result_id_set( ): """When the Scenario stamps ``_scenario_result_id`` onto the atomic attack, ``run_async`` must build and pass a factory.""" - from pyrit.executor.attack.core.scenario_execution_attribution import ScenarioExecutionAttribution + from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), @@ -1122,20 +1122,20 @@ async def test_factory_built_when_scenario_result_id_set( # input_index 0 → original_indices[0] == 0 attribution = factory(0) - assert isinstance(attribution, ScenarioExecutionAttribution) - assert attribution.scenario_result_id == "00000000-0000-0000-0000-000000000abc" - assert attribution.atomic_attack_name == "MyAtomicAttack" - assert attribution.objective_index == 0 + assert isinstance(attribution, AttackResultAttribution) + assert attribution.parent_id == "00000000-0000-0000-0000-000000000abc" + assert attribution.parent_collection == "MyAtomicAttack" + assert attribution.position == 0 # input_index 2 → original_indices[2] == 2 - assert factory(2).objective_index == 2 + assert factory(2).position == 2 async def test_factory_maps_input_index_to_original_index_after_filtering( self, mock_attack, sample_seed_groups, sample_attack_results ): """The whole point of the factory: after filtering out a completed seed group, the per-task input index 0 must still map back to the - *original* objective_index of the surviving seed group, not 0.""" + *original* position of the surviving seed group, not 0.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, @@ -1152,7 +1152,7 @@ async def test_factory_maps_input_index_to_original_index_after_filtering( factory = mock_exec.call_args.kwargs["attribution_factory"] assert factory is not None # input_index 0 (only remaining task) → original index 2. - assert factory(0).objective_index == 2 + assert factory(0).position == 2 async def test_factory_captures_indices_by_value_not_reference( self, mock_attack, sample_seed_groups, sample_attack_results @@ -1176,4 +1176,4 @@ async def test_factory_captures_indices_by_value_not_reference( atomic._original_indices = [999, 999, 999] # Factory still returns the snapshotted original indices. - assert [factory(i).objective_index for i in range(3)] == [0, 1, 2] + assert [factory(i).position for i in range(3)] == [0, 1, 2] diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/test_scenario.py index df234dcbce..01825a4ba2 100644 --- a/tests/unit/scenario/test_scenario.py +++ b/tests/unit/scenario/test_scenario.py @@ -31,19 +31,19 @@ def save_attack_results_to_memory(attack_results): def _stamp_scenario_linkage(*, attack_results, atomic_attack): """ - Stamp scenario_result_id + scenario_data on each AttackResult the same way - the real attack persistence path does. Mirrors what + Stamp attribution_parent_id + attribution_data on each AttackResult the + same way the real attack persistence path does. Mirrors what ``_DefaultAttackStrategyEventHandler._stamp_attribution`` does at runtime so test fixtures that mock out the executor still produce DB rows the new - FK-based hydration can find. + foreign-key-based hydration can find. """ sid = getattr(atomic_attack, "_scenario_result_id", None) name = getattr(atomic_attack, "atomic_attack_name", None) if not sid or not name: return for i, r in enumerate(attack_results): - r.scenario_result_id = sid - r.scenario_data = {"atomic_attack_name": name, "objective_index": i} + r.attribution_parent_id = sid + r.attribution_data = {"parent_collection": name, "position": i} def create_mock_run_async(attack_results, *, atomic_attack=None): @@ -52,8 +52,8 @@ def create_mock_run_async(attack_results, *, atomic_attack=None): Pass ``atomic_attack`` (the AtomicAttack MagicMock) so the helper can copy its ``_scenario_result_id`` (set by ``Scenario._execute_scenario_async``) - and ``atomic_attack_name`` onto each result. Without those the FK-based - hydration in ``get_scenario_results`` won't see the rows. + and ``atomic_attack_name`` onto each result. Without those the foreign-key- + based hydration in ``get_scenario_results`` won't see the rows. """ async def mock_run_async(*args, **kwargs): diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/test_scenario_partial_results.py index 223b435c8f..0b03bf5a65 100644 --- a/tests/unit/scenario/test_scenario_partial_results.py +++ b/tests/unit/scenario/test_scenario_partial_results.py @@ -38,8 +38,9 @@ def mock_objective_target(): def save_attack_results_to_memory(attack_results, *, atomic_attack=None): """ Helper function to save attack results to memory. When ``atomic_attack`` is - provided, also stamps ``scenario_result_id`` and ``scenario_data`` on each - result the same way the real attack persistence path does — so FK-based + provided, also stamps ``attribution_parent_id`` and ``attribution_data`` on + each result the same way the real attack persistence path does — so + foreign-key-based hydration in ``get_scenario_results`` finds them. """ if atomic_attack is not None: @@ -47,8 +48,8 @@ def save_attack_results_to_memory(attack_results, *, atomic_attack=None): name = getattr(atomic_attack, "atomic_attack_name", None) if sid and name: for i, r in enumerate(attack_results): - r.scenario_result_id = sid - r.scenario_data = {"atomic_attack_name": name, "objective_index": i} + r.attribution_parent_id = sid + r.attribution_data = {"parent_collection": name, "position": i} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index 5c13a639a9..26a98161f5 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -43,17 +43,17 @@ def mock_objective_scorer(): def save_attack_results_to_memory(attack_results, *, atomic_attack=None): """Helper function to save attack results to memory. - When ``atomic_attack`` is provided, stamps ``scenario_result_id`` and - ``scenario_data`` onto each result (mirrors the real attack persistence - path so FK-based hydration sees the rows). + When ``atomic_attack`` is provided, stamps ``attribution_parent_id`` and + ``attribution_data`` onto each result (mirrors the real attack persistence + path so foreign-key-based hydration sees the rows). """ if atomic_attack is not None: sid = getattr(atomic_attack, "_scenario_result_id", None) name = getattr(atomic_attack, "atomic_attack_name", None) if sid and name: for i, r in enumerate(attack_results): - r.scenario_result_id = sid - r.scenario_data = {"atomic_attack_name": name, "objective_index": i} + r.attribution_parent_id = sid + r.attribution_data = {"parent_collection": name, "position": i} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) @@ -104,8 +104,8 @@ def create_mock_run_async(attack_results, *, atomic_attack=None): Args: attack_results: List of AttackResult objects to return atomic_attack: Optional AtomicAttack mock. When provided, results are - stamped with scenario_result_id and scenario_data so FK-based - hydration finds them. + stamped with attribution_parent_id and attribution_data so + foreign-key-based hydration finds them. Returns: AsyncMock configured to return the results @@ -614,8 +614,8 @@ async def mock_run_attack4(*args, **kwargs): @pytest.mark.usefixtures("patch_central_database") -class TestScenarioFKResumeRegression: - """Regression tests for the FK-based scenario linkage resume path. +class TestScenarioForeignKeyResumeRegression: + """Regression tests for the foreign-key-based scenario linkage resume path. The bug being regression-tested: when a Scenario is interrupted mid- AtomicAttack (Ctrl-C, OOM, crash), AttackResults already persisted to the @@ -623,7 +623,7 @@ class TestScenarioFKResumeRegression: link only lived in a JSON manifest written after the whole AtomicAttack returned. On resume, those objectives were re-executed (wasted compute). - After the refactor, ``scenario_result_id`` is stamped on each + After the refactor, ``attribution_parent_id`` is stamped on each ``AttackResultEntry`` at write time, so resume reads them directly and skips the already-done work even when the manifest was never updated. """ From 547dff67704eaf79faa004c695643d9505c31c9c Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 10:20:31 -0700 Subject: [PATCH 05/10] FEAT: hash-based scenario resume keys Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/executor/attack/core/attack_executor.py | 47 ++--- .../attack/core/attack_result_attribution.py | 15 +- pyrit/executor/attack/core/attack_strategy.py | 1 - ..._add_scenario_linkage_to_attack_results.py | 18 +- pyrit/memory/memory_interface.py | 35 +++- pyrit/memory/memory_models.py | 10 + pyrit/models/scenario_result.py | 9 + pyrit/scenario/core/atomic_attack.py | 146 ++++++++------ pyrit/scenario/core/scenario.py | 111 ++++++++-- .../attack/core/test_attack_executor.py | 69 +++---- .../attack/core/test_attack_strategy.py | 4 - .../test_interface_scenario_results.py | 14 +- tests/unit/memory/test_migration.py | 9 +- tests/unit/scenario/test_atomic_attack.py | 190 +++++++++--------- tests/unit/scenario/test_scenario.py | 4 +- .../scenario/test_scenario_partial_results.py | 32 +-- tests/unit/scenario/test_scenario_retry.py | 111 +++------- 17 files changed, 451 insertions(+), 374 deletions(-) diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index dd4a7140e6..6405b6816f 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -8,7 +8,7 @@ """ import asyncio -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Iterator, Sequence from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, @@ -143,7 +143,7 @@ async def execute_attack_from_seed_groups_async( objective_scorer: Optional["TrueFalseScorer"] = None, field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, + attribution: Optional[AttackResultAttribution] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -165,13 +165,12 @@ async def execute_attack_from_seed_groups_async( from_seed_group() as overrides. return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. - attribution_factory: Optional callable that maps an input index (the - seed group's original index, parallel-safe and deterministic) to - an ``AttackResultAttribution``. When provided, each per-task - ``AttackContext`` is stamped with the attribution so the - resulting ``AttackResultEntry`` row carries - ``attribution_parent_id`` + ``attribution_data``. When ``None``, - no attribution is applied. + attribution: Optional ``AttackResultAttribution`` stamped onto every + per-task ``AttackContext`` so the persisted ``AttackResultEntry`` + row carries ``attribution_parent_id`` + ``attribution_data``. + When ``None`` (default), no attribution is applied. The same + attribution is shared across all tasks; per-task identity is + reconstructed from the row's own ``objective_sha256``. **broadcast_fields: Fields applied to all seed groups (e.g., memory_labels). Per-seed-group field_overrides take precedence. @@ -214,7 +213,7 @@ async def build_params(i: int, sg: SeedAttackGroup) -> AttackParameters: attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution_factory=attribution_factory, + attribution=attribution, ) async def execute_attack_async( @@ -224,7 +223,7 @@ async def execute_attack_async( objectives: Sequence[str], field_overrides: Optional[Sequence[dict[str, Any]]] = None, return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, + attribution: Optional[AttackResultAttribution] = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -239,10 +238,9 @@ async def execute_attack_async( must match the length of objectives. return_partial_on_failure: If True, returns partial results when some objectives fail. If False (default), raises the first exception. - attribution_factory: Optional callable mapping each input index to - a AttackResultAttribution. When provided, the per-task context is - stamped with the attribution so the persistence path can record - scenario linkage. + attribution: Optional ``AttackResultAttribution`` stamped onto every + per-task ``AttackContext`` so the persistence path can record + orchestrator linkage. When ``None``, no attribution is applied. **broadcast_fields: Fields applied to all objectives (e.g., memory_labels). Per-objective field_overrides take precedence. @@ -283,7 +281,7 @@ async def execute_attack_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution_factory=attribution_factory, + attribution=attribution, ) async def _execute_with_params_list_async( @@ -292,7 +290,7 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, - attribution_factory: Optional[Callable[[int], AttackResultAttribution]] = None, + attribution: Optional[AttackResultAttribution] = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Execute attacks in parallel with a list of pre-built parameters. @@ -304,10 +302,9 @@ async def _execute_with_params_list_async( attack: The attack strategy to execute. params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. - attribution_factory: Optional callable mapping each input index to - a AttackResultAttribution. When provided, the per-task context is - stamped with the attribution so the persistence path can record - scenario linkage. + attribution: Optional ``AttackResultAttribution`` stamped onto every + per-task ``AttackContext`` so the persistence path can record + orchestrator linkage. Returns: AttackExecutorResult with completed results and any incomplete objectives. @@ -316,13 +313,9 @@ async def _execute_with_params_list_async( async def run_one(index: int, params: AttackParameters) -> AttackStrategyResultT: async with semaphore: - # Create context with params and stamp attribution (if any). The - # input index is the seed group's original position and is - # deterministic and parallel-safe — assigned BEFORE the task - # runs, not from completion order. context = attack._context_type(params=params) - if attribution_factory is not None: - context._attribution = attribution_factory(index) + if attribution is not None: + context._attribution = attribution return await attack.execute_with_context_async(context=context) tasks = [run_one(i, p) for i, p in enumerate(params_list)] diff --git a/pyrit/executor/attack/core/attack_result_attribution.py b/pyrit/executor/attack/core/attack_result_attribution.py index bcabaefc9c..208dcf473d 100644 --- a/pyrit/executor/attack/core/attack_result_attribution.py +++ b/pyrit/executor/attack/core/attack_result_attribution.py @@ -26,11 +26,13 @@ class AttackResultAttribution: onto the resulting ``AttackResult`` by the attack persistence path so the DB row records its lineage. - All three fields are opaque to the attack layer. The orchestrator chooses + Both fields are opaque to the attack layer. The orchestrator chooses what they mean and how to query them back later. For example, - ``Scenario`` uses ``parent_id`` for the scenario result UUID, - ``parent_collection`` for the atomic attack name, and ``position`` for - the original 0-based seed-group index. + ``Scenario`` uses ``parent_id`` for the scenario result UUID and + ``parent_collection`` for the atomic attack name. Result identity within + a collection is reconstructed from the row's own ``objective_sha256`` + rather than a positional index, so resume is stable under reordering or + re-sampling. Attributes: parent_id (str): The ID of the parent entity that owns this attack @@ -40,12 +42,7 @@ class AttackResultAttribution: parent_collection (str): A free-form label naming the per-parent collection this result belongs to. Persisted into ``AttackResultEntry.attribution_data``. - position (int): The 0-based position of this result within its - ``parent_collection``. Assigned **before** task execution so it is - deterministic and parallel-safe, and used as the stable resume key. - Persisted into ``AttackResultEntry.attribution_data``. """ parent_id: str parent_collection: str - position: int diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index e4a4e58b59..77fc69dcc4 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -263,7 +263,6 @@ def _stamp_attribution( result.attribution_parent_id = attribution.parent_id result.attribution_data = { "parent_collection": attribution.parent_collection, - "position": attribution.position, } def _log_attack_outcome(self, result: AttackResult) -> None: diff --git a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py index d5eeb2b910..6f5e6e5b3a 100644 --- a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py +++ b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py @@ -63,12 +63,15 @@ def upgrade() -> None: ondelete="SET NULL", ) - # ScenarioResultEntries: drop the not-yet-released error_attack_result_ids_json. + # ScenarioResultEntries: drop the not-yet-released error_attack_result_ids_json, + # and add scenario_metadata for free-form scenario-level JSON state (e.g. + # the persisted sampled_objective_hashes used for resume). # Error AttackResults are now linkable via the new attribution_parent_id # foreign key; the per-scenario manifest column is no longer used. # Wrapped in a batch op for SQLite. with op.batch_alter_table("ScenarioResultEntries") as batch_op: batch_op.drop_column("error_attack_result_ids_json") + batch_op.add_column(sa.Column("scenario_metadata", sa.JSON(), nullable=True)) # Backfill attribution linkage from the existing attack_results_json manifest. _backfill_attribution_linkage() @@ -76,11 +79,10 @@ def upgrade() -> None: def downgrade() -> None: """Revert this schema upgrade.""" - # Re-add error_attack_result_ids_json on ScenarioResultEntries. - op.add_column( - "ScenarioResultEntries", - sa.Column("error_attack_result_ids_json", sa.Unicode(), nullable=True), - ) + # Re-add error_attack_result_ids_json on ScenarioResultEntries and drop scenario_metadata. + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.add_column(sa.Column("error_attack_result_ids_json", sa.Unicode(), nullable=True)) + batch_op.drop_column("scenario_metadata") # Drop foreign key + columns from AttackResultEntries. with op.batch_alter_table("AttackResultEntries") as batch_op: @@ -131,7 +133,7 @@ def _backfill_attribution_linkage() -> None: for atomic_attack_name, conversation_ids in manifest.items(): if not isinstance(conversation_ids, list): continue - for objective_index, conversation_id in enumerate(conversation_ids): + for conversation_id in conversation_ids: if not isinstance(conversation_id, str): continue # Check for duplicate conversation_id matches (data anomaly). @@ -150,7 +152,7 @@ def _backfill_attribution_linkage() -> None: f"All matching rows will be linked to scenario {scenario_id}." ) - attribution_data = json.dumps({"parent_collection": atomic_attack_name, "position": objective_index}) + attribution_data = json.dumps({"parent_collection": atomic_attack_name}) result = bind.execute( update_stmt, { diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 70f2ab5b7b..afcc9386e0 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2009,6 +2009,34 @@ def update_scenario_run_state( logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state}'") + def update_scenario_metadata( + self, + *, + scenario_result_id: str, + metadata: dict[str, Any], + ) -> None: + """ + Replace the ``scenario_metadata`` JSON blob on an existing scenario result. + + Used by the scenario layer to persist first-run state (e.g. + ``sampled_objective_hashes``) that resume needs to replay. Performs a + targeted UPDATE so it doesn't clobber other columns. + + Args: + scenario_result_id (str): The ID of the scenario result to update. + metadata (dict[str, Any]): The full metadata dict to store. Pass the + merged dict, not just the new keys — this writes the whole value. + + Raises: + ValueError: If the scenario result is not found. + """ + with closing(self.get_session()) as session: + entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first() + if not entry: + raise ValueError(f"Scenario result with ID {scenario_result_id} not found in memory") + entry.scenario_metadata = metadata if metadata else None + session.commit() + def get_scenario_results( self, *, @@ -2208,7 +2236,7 @@ def _hydrate_scenario_attack_results( ``AttackResultEntry.attribution_parent_id`` foreign key in a single batched query, then group by scenario + ``parent_collection`` (which the scenario layer uses for the atomic attack name) and sort each - group by ``attribution_data["position"]``. + group by ``AttackResultEntry.timestamp``. Foreign-key linkage is the sole source of truth — set at write-time by the attack persistence path when an ``AttackResultAttribution`` is on @@ -2230,7 +2258,7 @@ def _hydrate_scenario_attack_results( batch_values=scenario_ids, ) - grouped: dict[uuid.UUID, dict[str, list[tuple[int, AttackResult]]]] = {entry.id: {} for entry in entries} + grouped: dict[uuid.UUID, dict[str, list[tuple[datetime, AttackResult]]]] = {entry.id: {} for entry in entries} for row in attack_rows: scenario_id = row.attribution_parent_id @@ -2245,8 +2273,7 @@ def _hydrate_scenario_attack_results( ) continue - position = data.get("position") - sort_key = position if isinstance(position, int) else 0 + sort_key = row.timestamp or datetime.min.replace(tzinfo=timezone.utc) grouped[scenario_id].setdefault(name, []).append((sort_key, row.get_attack_result())) return { diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 9400a78cc5..8eb636d92c 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1016,6 +1016,14 @@ class ScenarioResultEntry(Base): error_message: Mapped[Optional[str]] = mapped_column(Unicode, nullable=True) error_type: Mapped[Optional[str]] = mapped_column(String, nullable=True) + # Free-form JSON metadata stamped by the scenario. Currently used to record + # ``sampled_objective_hashes`` — the objective sha256 set chosen on the + # first run, replayed on resume so a fresh ``random.sample`` can't + # silently change which objectives the scenario operates on. Column is + # named ``scenario_metadata`` because SQLAlchemy's ``DeclarativeBase`` + # reserves ``metadata`` as a class attribute on the model. + scenario_metadata: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) + def __init__(self, *, entry: ScenarioResult) -> None: """ Initialize a ScenarioResultEntry from a ScenarioResult object. @@ -1067,6 +1075,7 @@ def __init__(self, *, entry: ScenarioResult) -> None: self.error_message = entry.error_message self.error_type = entry.error_type + self.scenario_metadata = entry.metadata if entry.metadata else None self.timestamp = datetime.now(tz=timezone.utc) @@ -1123,6 +1132,7 @@ def get_scenario_result(self) -> ScenarioResult: display_group_map=display_group_map, error_message=self.error_message, error_type=self.error_type, + metadata=dict(self.scenario_metadata) if self.scenario_metadata else None, ) def get_conversation_ids_by_attack_name(self) -> dict[str, list[str]]: diff --git a/pyrit/models/scenario_result.py b/pyrit/models/scenario_result.py index 2177fafaa7..aa26133af3 100644 --- a/pyrit/models/scenario_result.py +++ b/pyrit/models/scenario_result.py @@ -71,6 +71,7 @@ def __init__( display_group_map: dict[str, str] | None = None, error_message: str | None = None, error_type: str | None = None, + metadata: dict[str, Any] | None = None, ) -> None: """ Initialize a scenario result. @@ -92,6 +93,13 @@ def __init__( printer to aggregate results for user-facing output. error_message (Optional[str]): Scenario-level error message when the run fails. error_type (Optional[str]): Exception class name when the run fails. + metadata (Optional[dict[str, Any]]): Free-form JSON metadata persisted + with the scenario result. Currently used to record + ``sampled_objective_hashes`` — the objective ``sha256`` set chosen + on the first run, replayed on resume so a fresh + ``random.sample`` can't silently change which objectives the + scenario operates on. Keys are not part of any public contract + and may evolve. """ self.id = id if id is not None else uuid.uuid4() @@ -110,6 +118,7 @@ def __init__( self._display_group_map = display_group_map or {} self.error_message = error_message self.error_type = error_type + self.metadata: dict[str, Any] = metadata if metadata is not None else {} @property def display_group_map(self) -> dict[str, str]: diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index 4699041844..6b422e0f27 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Optional from pyrit.common.deprecation import print_deprecation_message +from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackExecutor, AttackStrategy from pyrit.executor.attack.core.attack_executor import AttackExecutorResult from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution @@ -119,12 +120,7 @@ def __init__( sg.validate() self._seed_groups = seed_groups - # Original positions for each currently-active seed group. Used to map a - # per-task input index (from the executor) back to the stable position - # stored in AttackResultEntry.attribution_data. Resume filtering - # preserves the original indices so newly-persisted results don't - # collide with already-persisted ones. - self._original_indices: list[int] = list(range(len(seed_groups))) + self._validate_unique_objective_hashes() self._adversarial_chat = adversarial_chat self._objective_scorer = objective_scorer self._memory_labels = memory_labels or {} @@ -139,6 +135,33 @@ def __init__( f"attack type: {type(self._attack_technique.attack).__name__}" ) + def _validate_unique_objective_hashes(self) -> None: + """ + Ensure each seed group in this atomic attack has a unique objective hash. + + Within a single ``AtomicAttack`` (one ``atomic_attack_name``, one + technique), the objective text identifies a unit of work. Duplicates + would mean two indistinguishable rows on the write side, which makes + resume reconciliation ambiguous — the new hash-based resume key + treats a set of hashes as already-done, with no way to distinguish + which of two duplicate rows is "the one" that is still outstanding. + Loud failure here beats silent resume corruption later. + + Raises: + ValueError: If two seed groups share the same ``objective_sha256``. + """ + seen: dict[str, int] = {} + for sg in self._seed_groups: + if sg.objective is None: + continue + sha = to_sha256(sg.objective.value) + if sha in seen: + raise ValueError( + f"AtomicAttack '{self.atomic_attack_name}' has duplicate objective hash " + f"{sha[:12]}... across seed_groups; each (objective, technique) pair must be unique." + ) + seen[sha] = 1 + @property def attack_technique(self) -> AttackTechnique: """Get the attack technique for this atomic attack.""" @@ -168,53 +191,74 @@ def filter_seed_groups_by_objectives(self, *, remaining_objectives: list[str]) - """ Filter seed groups to only those with objectives in the remaining list. - Deprecated — prefer ``filter_seed_groups_by_indices``. Filtering by - objective text collapses two distinct seed groups that happen to share - the same objective text (common when seed groups are programmatically - generated or when a baseline is paired with its strategy variants). - Index-based filtering is the stable identity for resume. + Deprecated — prefer ``filter_seed_groups_by_completed_hashes``. Filtering + by objective text collapses two distinct seed groups that happen to share + the same objective text. Hash-based filtering aligns with the + ``objective_sha256`` resume key written to ``AttackResultEntry``. Args: remaining_objectives (List[str]): List of objectives that still need to be executed. """ print_deprecation_message( old_item="AtomicAttack.filter_seed_groups_by_objectives", - new_item="AtomicAttack.filter_seed_groups_by_indices", + new_item="AtomicAttack.filter_seed_groups_by_completed_hashes", removed_in="0.16.0", ) remaining_set = set(remaining_objectives) - kept: list[tuple[int, SeedAttackGroup]] = [ - (orig_idx, sg) - for orig_idx, sg in zip(self._original_indices, self._seed_groups, strict=True) - if sg.objective is not None and sg.objective.value in remaining_set + self._seed_groups = [ + sg for sg in self._seed_groups if sg.objective is not None and sg.objective.value in remaining_set ] - self._original_indices = [idx for idx, _ in kept] - self._seed_groups = [sg for _, sg in kept] - def filter_seed_groups_by_indices(self, *, completed_indices: set[int]) -> None: + def filter_seed_groups_by_completed_hashes(self, *, completed_hashes: set[str]) -> None: """ - Filter out seed groups whose original index is in ``completed_indices``. + Drop seed groups whose ``objective_sha256`` is already in the completed set. - This is the stable-identity replacement for - ``filter_seed_groups_by_objectives``. Each seed group's original index - (its position when the ``AtomicAttack`` was first constructed) is - tracked across filtering, so the executor can later map each per-task - input index back to the original index when stamping - ``attribution_data["position"]``. This prevents newly-persisted - results from colliding with already-persisted ones on resume. + This is the resume filter: within an atomic attack, ``objective_sha256`` + is the stable identity (enforced unique by ``__init__``). Content-derived + keys are robust to reordering and resampling, so resume produces the + right remaining-work set even when ``get_seed_groups()`` is rebuilt + from scratch on each ``run_async()``. Args: - completed_indices (set[int]): Original indices that have already - been executed (typically retrieved from - ``Scenario._get_completed_objectives_for_attack``). + completed_hashes (set[str]): SHA256 hashes of objective text for + seed groups that have already produced a (non-error) ``AttackResult``. """ - kept: list[tuple[int, SeedAttackGroup]] = [ - (orig_idx, sg) - for orig_idx, sg in zip(self._original_indices, self._seed_groups, strict=True) - if orig_idx not in completed_indices + self._seed_groups = [ + sg + for sg in self._seed_groups + if sg.objective is None or to_sha256(sg.objective.value) not in completed_hashes ] - self._original_indices = [idx for idx, _ in kept] - self._seed_groups = [sg for _, sg in kept] + + def restrict_seed_groups_to_hashes(self, *, keep_hashes: set[str]) -> set[str]: + """ + Keep only seed groups whose ``objective_sha256`` is in ``keep_hashes``. + + Inverse of ``filter_seed_groups_by_completed_hashes``: used on resume to + replay the originally-sampled subset and ignore any seed groups that + were added since (or that landed in this run's fresh ``random.sample`` + draw and are no longer in the persisted set). + + Args: + keep_hashes (set[str]): SHA256 hashes of objective text for seed + groups to keep. + + Returns: + set[str]: The hashes that were actually retained (intersection of + ``keep_hashes`` and the current seed_groups' hashes). The caller + can union these across atomic attacks to detect persisted hashes + that no longer exist in the dataset. + """ + retained: set[str] = set() + new_groups: list[SeedAttackGroup] = [] + for sg in self._seed_groups: + if sg.objective is None: + continue + sha = to_sha256(sg.objective.value) + if sha in keep_hashes: + retained.add(sha) + new_groups.append(sg) + self._seed_groups = new_groups + return retained async def run_async( self, @@ -270,26 +314,16 @@ async def run_async( else: execution_seed_groups = self._seed_groups - # Build an attribution factory when this atomic attack is being - # executed inside a Scenario. The factory maps each per-task input - # index (position in the currently-active seed_groups list) back to - # the original position so resume can locate already-done work and - # newly-persisted results don't collide with old ones. The Scenario - # uses parent_collection for the atomic attack name and position - # for the original seed-group index; the attack layer treats both - # as opaque strings/ints. - attribution_factory = None + # Build attribution when this atomic attack is being executed inside + # a Scenario. The same attribution object is stamped on every + # per-task AttackContext; per-task identity is reconstructed from + # the row's own objective_sha256 (no positional state required). + attribution: AttackResultAttribution | None = None if self._scenario_result_id is not None: - scenario_id = self._scenario_result_id - name = self.atomic_attack_name - original_indices = list(self._original_indices) - - def attribution_factory(input_index: int) -> AttackResultAttribution: - return AttackResultAttribution( - parent_id=scenario_id, - parent_collection=name, - position=original_indices[input_index], - ) + attribution = AttackResultAttribution( + parent_id=self._scenario_result_id, + parent_collection=self.atomic_attack_name, + ) results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, @@ -298,7 +332,7 @@ def attribution_factory(input_index: int) -> AttackResultAttribution: objective_scorer=self._objective_scorer, memory_labels=self._memory_labels, return_partial_on_failure=return_partial_on_failure, - attribution_factory=attribution_factory, + attribution=attribution, **self._attack_execute_params, ) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 2e505ad90a..cd35425a31 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -24,6 +24,7 @@ from pyrit.common import REQUIRED_VALUE, Parameter, apply_defaults from pyrit.common.deprecation import print_deprecation_message from pyrit.common.parameter import coerce_value, validate_param_type +from pyrit.common.utils import to_sha256 from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory import CentralMemory from pyrit.memory.memory_models import ScenarioResultEntry @@ -728,6 +729,7 @@ async def initialize_async( ) self._validate_stored_scenario(stored_result=existing_results[0]) + self._apply_persisted_sample(stored_result=existing_results[0]) return # Valid resume - skip creating new scenario result # Build display group mapping from atomic attacks @@ -746,12 +748,84 @@ async def initialize_async( attack_results=attack_results, scenario_run_state="CREATED", display_group_map=self._display_group_map, + metadata=self._build_initial_scenario_metadata(), ) self._memory.add_scenario_results_to_memory(scenario_results=[result]) self._scenario_result_id = str(result.id) logger.info(f"Created new scenario result with ID: {self._scenario_result_id}") + def _build_initial_scenario_metadata(self) -> dict[str, Any]: + """ + Build the metadata dict persisted with a freshly-created ``ScenarioResult``. + + When ``max_dataset_size`` is in effect, the dataset config draws an + unseeded ``random.sample`` and the chosen subset would silently change + on the next run (e.g. a resume). To make resume reliable, snapshot the + chosen objective hashes here so the next ``_setup_scenario_async`` can + replay them via ``restrict_seed_groups_to_hashes``. + + When ``max_dataset_size`` is not set, the sample equals the dataset and + nothing needs pinning; the dict is empty. + + Returns: + dict[str, Any]: Metadata payload for the new ScenarioResult. + """ + metadata: dict[str, Any] = {} + if getattr(self._dataset_config, "max_dataset_size", None) is None: + return metadata + hashes: list[str] = [] + seen: set[str] = set() + for aa in self._atomic_attacks: + for sg in aa.seed_groups: + if sg.objective is None: + continue + sha = to_sha256(sg.objective.value) + if sha not in seen: + seen.add(sha) + hashes.append(sha) + metadata["sampled_objective_hashes"] = hashes + return metadata + + def _apply_persisted_sample(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["sampled_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. + + Args: + stored_result (ScenarioResult): The scenario result loaded from memory. + + Raises: + ValueError: If any persisted objective hash is missing from the + currently-resolved dataset. + """ + metadata = stored_result.metadata or {} + persisted = metadata.get("sampled_objective_hashes") + if not persisted: + return + + keep_hashes: set[str] = set(persisted) + retained: set[str] = set() + for aa in self._atomic_attacks: + retained |= aa.restrict_seed_groups_to_hashes(keep_hashes=keep_hashes) + + missing = keep_hashes - retained + if missing: + sample = sorted(missing)[:3] + raise ValueError( + f"Scenario result id '{self._scenario_result_id}' cannot resume: " + f"{len(missing)} persisted objective hash(es) are no longer present in the dataset " + f"(missing examples: {', '.join(h[:12] + '...' for h in sample)}). " + f"Either restore the missing objectives or drop scenario_result_id to start a new scenario." + ) + def _build_baseline_atomic_attack(self, *, seed_groups: list[SeedAttackGroup]) -> AtomicAttack: """ Build the baseline AtomicAttack from pre-resolved seed groups. @@ -851,26 +925,28 @@ def _validate_stored_scenario(self, *, stored_result: ScenarioResult) -> None: f"(ID: {self._scenario_result_id}, state: {stored_result.scenario_run_state})" ) - def _get_completed_objective_indices_for_attack(self, *, atomic_attack_name: str) -> set[int]: + def _get_completed_objective_hashes_for_attack(self, *, atomic_attack_name: str) -> set[str]: """ - Get the set of objective_index values that have already been completed - (non-error) for a specific atomic attack inside this scenario. + Return the set of ``objective_sha256`` values already completed (non-error) + for a specific atomic attack inside this scenario. - Queries AttackResultEntry rows directly by ``attribution_parent_id`` — + Queries ``AttackResultEntry`` rows directly 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. + Identity is content-derived (``to_sha256(objective)``), so it stays + stable even if ``get_seed_groups()`` reorders or resamples between runs. Args: atomic_attack_name (str): The atomic attack name to scope to. Returns: - Set[int]: Original objective_index values that completed without error. + set[str]: ``objective_sha256`` hex strings for completed-without-error rows. """ if not self._scenario_result_id: return set() - completed_indices: set[int] = set() + completed_hashes: set[str] = set() try: rows = self._memory.get_attack_results(scenario_result_id=self._scenario_result_id) for row in rows: @@ -880,23 +956,24 @@ def _get_completed_objective_indices_for_attack(self, *, atomic_attack_name: str continue if row.attribution_data.get("parent_collection") != atomic_attack_name: continue - position = row.attribution_data.get("position") - if isinstance(position, int): - completed_indices.add(position) + if row.objective: + completed_hashes.add(to_sha256(row.objective)) except Exception as e: logger.warning( - f"Failed to retrieve completed objective indices for atomic attack '{atomic_attack_name}': {str(e)}" + f"Failed to retrieve completed objective hashes for atomic attack '{atomic_attack_name}': {str(e)}" ) - return completed_indices + return completed_hashes async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: """ Get the list of atomic attacks that still have objectives to complete. - Uses ``objective_index`` (the original position of each seed group in the - atomic attack) for resume rather than objective text, so duplicate - objective text in two distinct seed groups doesn't collapse on resume. + Uses ``objective_sha256`` as the stable identity for resume: each + 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. Returns: List[AtomicAttack]: List of atomic attacks with uncompleted objectives. @@ -908,13 +985,13 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: remaining_attacks: list[AtomicAttack] = [] for atomic_attack in self._atomic_attacks: - completed_indices = self._get_completed_objective_indices_for_attack( + completed_hashes = self._get_completed_objective_hashes_for_attack( atomic_attack_name=atomic_attack.atomic_attack_name ) - if completed_indices: + if completed_hashes: original_count = len(atomic_attack.seed_groups) - atomic_attack.filter_seed_groups_by_indices(completed_indices=completed_indices) + atomic_attack.filter_seed_groups_by_completed_hashes(completed_hashes=completed_hashes) remaining_count = len(atomic_attack.seed_groups) if remaining_count == 0: logger.info( diff --git a/tests/unit/executor/attack/core/test_attack_executor.py b/tests/unit/executor/attack/core/test_attack_executor.py index ddf2af6cf1..7034739872 100644 --- a/tests/unit/executor/attack/core/test_attack_executor.py +++ b/tests/unit/executor/attack/core/test_attack_executor.py @@ -337,94 +337,81 @@ async def test_validates_explicit_empty_field_overrides_for_seed_groups(self): @pytest.mark.usefixtures("patch_central_database") -class TestAttributionFactoryPropagation: +class TestAttributionPropagation: """Tests for AttackResultAttribution propagation through the AttackExecutor. - The executor stamps ``context._attribution = attribution_factory(input_index)`` - on each per-task context. The input_index must be the seed group's original - position so attribution is deterministic and parallel-safe — never derived - from completion order. + The executor stamps the same ``AttackResultAttribution`` on every per-task + context. Per-task identity is reconstructed from each row's own + ``objective_sha256`` at hydration/resume time, so no positional state is + threaded through the executor. """ - async def test_attribution_factory_sets_per_task_attribution_in_order(self): + async def test_attribution_stamps_every_per_task_context(self): from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution attack = create_mock_attack() - seen_indices: list[int] = [] + seen_parent_ids: list[str] = [] + seen_collections: list[str] = [] async def capture(context): - # Read the attribution stamped by the executor BEFORE the task runs. attr = context._attribution assert attr is not None - seen_indices.append(attr.position) + seen_parent_ids.append(attr.parent_id) + seen_collections.append(attr.parent_collection) return create_attack_result(context.params.objective) attack.execute_with_context_async = AsyncMock(side_effect=capture) seed_groups = [create_seed_group(f"obj-{i}") for i in range(4)] - - def factory(idx: int) -> AttackResultAttribution: - return AttackResultAttribution( - parent_id="sid", - parent_collection="atomic", - position=idx, - ) + attribution = AttackResultAttribution(parent_id="sid", parent_collection="atomic") executor = AttackExecutor(max_concurrency=1) result = await executor.execute_attack_from_seed_groups_async( attack=attack, seed_groups=seed_groups, - attribution_factory=factory, + attribution=attribution, ) - # Each per-task context saw a distinct, contiguous position — the - # seed group's original input position. - assert sorted(seen_indices) == [0, 1, 2, 3] + assert seen_parent_ids == ["sid"] * 4 + assert seen_collections == ["atomic"] * 4 assert len(result.completed_results) == 4 - async def test_attribution_factory_parallel_safe_with_high_concurrency(self): - """At max_concurrency > 1, each task still receives its own position - regardless of completion order. Out-of-order completion must not produce - stamped positions that point at the wrong seed group. + async def test_attribution_parallel_safe_with_high_concurrency(self): + """At max_concurrency > 1, every task still sees the same attribution + regardless of completion order — there is no per-task positional state. """ from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution attack = create_mock_attack() - seen: dict[str, int] = {} + seen: dict[str, AttackResultAttribution] = {} async def out_of_order(context): attr = context._attribution assert attr is not None - # Reverse-delay the tasks so completion order is the inverse of input - # order. The stamped position must still match the seed group. - await asyncio.sleep(0.005 * (10 - attr.position)) - seen[context.params.objective] = attr.position + # Reverse-delay tasks so completion order is inverse of input order. + i = int(context.params.objective.split("-")[1]) + await asyncio.sleep(0.005 * (10 - i)) + seen[context.params.objective] = attr return create_attack_result(context.params.objective) attack.execute_with_context_async = AsyncMock(side_effect=out_of_order) seed_groups = [create_seed_group(f"obj-{i}") for i in range(6)] - - def factory(idx: int) -> AttackResultAttribution: - return AttackResultAttribution( - parent_id="sid", - parent_collection="atomic", - position=idx, - ) + attribution = AttackResultAttribution(parent_id="sid", parent_collection="atomic") executor = AttackExecutor(max_concurrency=6) await executor.execute_attack_from_seed_groups_async( attack=attack, seed_groups=seed_groups, - attribution_factory=factory, + attribution=attribution, ) - # Each objective is stamped with its OWN input index, regardless of - # completion order. for i in range(6): - assert seen[f"obj-{i}"] == i + attr = seen[f"obj-{i}"] + assert attr.parent_id == "sid" + assert attr.parent_collection == "atomic" - async def test_no_attribution_factory_leaves_context_attribution_none(self): + async def test_no_attribution_leaves_context_attribution_none(self): attack = create_mock_attack() async def capture(context): diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index b44656f602..e932da3da9 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -623,7 +623,6 @@ async def test_on_post_execute_stamps_scenario_attribution_when_present( sample_attack_context._attribution = AttackResultAttribution( parent_id="scenario-1", parent_collection="atomic_a", - position=3, ) event_data = StrategyEventData( @@ -638,7 +637,6 @@ async def test_on_post_execute_stamps_scenario_attribution_when_present( assert sample_attack_result.attribution_parent_id == "scenario-1" assert sample_attack_result.attribution_data == { "parent_collection": "atomic_a", - "position": 3, } async def test_on_post_execute_no_attribution_leaves_fields_none( @@ -674,7 +672,6 @@ async def test_on_error_stamps_scenario_attribution_when_present(self, sample_at sample_attack_context._attribution = AttackResultAttribution( parent_id="scenario-err", parent_collection="atomic_err", - position=7, ) event_data = StrategyEventData( @@ -693,7 +690,6 @@ async def test_on_error_stamps_scenario_attribution_when_present(self, sample_at assert persisted.attribution_parent_id == "scenario-err" assert persisted.attribution_data == { "parent_collection": "atomic_err", - "position": 7, } diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index 313f5434f0..16ec2578d1 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -201,19 +201,21 @@ def test_attack_results_populated_correctly(sqlite_instance: MemoryInterface): def test_attack_order_preserved(sqlite_instance: MemoryInterface): - """Test that attack results are sorted by objective_index within each attack name.""" + """Hydration sorts each atomic attack's results by ``timestamp`` (which + monotonically tracks insertion order under normal sequential execution).""" scenario_result = create_scenario_result(name="Ordered Scenario", attack_results={}) sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) sid = scenario_result.id - # Insert in reverse order to prove the hydration sorts by objective_index, not insert order. + # Insert in a specific order; hydration must surface them in the same order. attack_results = [ _make_attack_result_for_scenario( scenario_result_id=sid, atomic_attack_name="Attack1", objective_index=i, conversation_id=f"conv_{i}" ) - for i in reversed(range(5)) + for i in range(5) ] - sqlite_instance.add_attack_results_to_memory(attack_results=attack_results) + for ar in attack_results: + sqlite_instance.add_attack_results_to_memory(attack_results=[ar]) results = sqlite_instance.get_scenario_results() retrieved_attacks = results[0].attack_results["Attack1"] @@ -678,7 +680,7 @@ def _make_attack_result_for_scenario( outcome=outcome, executed_turns=1, attribution_parent_id=str(scenario_result_id), - attribution_data={"parent_collection": atomic_attack_name, "position": objective_index}, + attribution_data={"parent_collection": atomic_attack_name}, ) @@ -771,7 +773,7 @@ def test_delete_scenario_sets_attack_result_foreign_key_to_null(sqlite_instance: with closing(sqlite_instance.get_session()) as session: entry = session.query(AttackResultEntry).filter_by(conversation_id=ar.conversation_id).one() assert entry.attribution_parent_id is None - assert entry.attribution_data == {"parent_collection": "a", "position": 0} + assert entry.attribution_data == {"parent_collection": "a"} def test_update_scenario_run_state_targeted_update_preserves_manifest(sqlite_instance: MemoryInterface): diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 3b283a2347..2694b08ade 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -295,15 +295,14 @@ def test_backfill_links_attack_results_via_conversation_id(): for conv in ("conv-a-0", "conv-a-1", "conv-b-0"): assert results_by_conv[conv][0] == sid, f"{conv} should be backfilled" - # attribution_data carries parent_collection (the atomic attack name) - # and the 0-based manifest position. + # attribution_data carries parent_collection (the atomic attack name). sd_a0 = json.loads(results_by_conv["conv-a-0"][1]) sd_a1 = json.loads(results_by_conv["conv-a-1"][1]) sd_b0 = json.loads(results_by_conv["conv-b-0"][1]) - assert sd_a0 == {"parent_collection": "a", "position": 0} - assert sd_a1 == {"parent_collection": "a", "position": 1} - assert sd_b0 == {"parent_collection": "b", "position": 0} + assert sd_a0 == {"parent_collection": "a"} + assert sd_a1 == {"parent_collection": "a"} + assert sd_b0 == {"parent_collection": "b"} finally: engine.dispose() diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index 8ed70647ba..1050f1b0d6 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -983,69 +983,105 @@ async def test_enrichment_skips_db_update_when_no_attack_result_id(self, mock_at @pytest.mark.usefixtures("patch_central_database") -class TestAtomicAttackFilterSeedGroupsByIndices: - """Tests for ``filter_seed_groups_by_indices`` — the stable-identity filter - that powers resume.""" +class TestAtomicAttackFilterSeedGroupsByCompletedHashes: + """Tests for ``filter_seed_groups_by_completed_hashes`` — the hash-based + resume filter.""" + + def test_filters_out_completed_hashes(self, mock_attack, sample_seed_groups): + from pyrit.common.utils import to_sha256 - def test_filters_out_completed_indices(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) - - atomic.filter_seed_groups_by_indices(completed_indices={0, 2}) + completed = {to_sha256("objective1"), to_sha256("objective3")} + atomic.filter_seed_groups_by_completed_hashes(completed_hashes=completed) assert atomic.seed_groups == [sample_seed_groups[1]] - assert atomic._original_indices == [1] - def test_empty_completed_indices_is_noop(self, mock_attack, sample_seed_groups): + def test_empty_completed_hashes_is_noop(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) - atomic.filter_seed_groups_by_indices(completed_indices=set()) + atomic.filter_seed_groups_by_completed_hashes(completed_hashes=set()) assert atomic.seed_groups == sample_seed_groups - assert atomic._original_indices == [0, 1, 2] - def test_all_indices_completed_clears_seed_groups(self, mock_attack, sample_seed_groups): + def test_all_hashes_completed_clears_seed_groups(self, mock_attack, sample_seed_groups): + from pyrit.common.utils import to_sha256 + atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) - atomic.filter_seed_groups_by_indices(completed_indices={0, 1, 2}) + atomic.filter_seed_groups_by_completed_hashes( + completed_hashes={to_sha256(f"objective{i}") for i in range(1, 4)} + ) assert atomic.seed_groups == [] - assert atomic._original_indices == [] - def test_filter_preserves_original_indices_across_successive_calls(self, mock_attack, sample_seed_groups): - """After two successive filters, ``_original_indices`` must still reference - the *initial* construction positions, not the post-filter positions.""" + def test_filter_is_stable_across_resampling(self, mock_attack, sample_seed_groups): + """Identity is content-derived, so reordering ``_seed_groups`` between + two calls (e.g. a fresh ``random.sample``) doesn't break the filter.""" + from pyrit.common.utils import to_sha256 + atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, atomic_attack_name="test", ) + # Simulate a re-sample by reversing the internal list. + atomic._seed_groups = list(reversed(atomic._seed_groups)) + + atomic.filter_seed_groups_by_completed_hashes(completed_hashes={to_sha256("objective1")}) + kept_objectives = [sg.objective.value for sg in atomic.seed_groups] + assert "objective1" not in kept_objectives + assert set(kept_objectives) == {"objective2", "objective3"} + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackRestrictSeedGroupsToHashes: + """Tests for ``restrict_seed_groups_to_hashes`` — the keep-set inverse used + on resume to replay the originally-sampled subset.""" + + def test_keeps_only_listed_hashes(self, mock_attack, sample_seed_groups): + from pyrit.common.utils import to_sha256 - atomic.filter_seed_groups_by_indices(completed_indices={0}) - assert atomic._original_indices == [1, 2] + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + keep = {to_sha256("objective1"), to_sha256("objective3")} + retained = atomic.restrict_seed_groups_to_hashes(keep_hashes=keep) - atomic.filter_seed_groups_by_indices(completed_indices={1}) - # Index 1 (original position) must be dropped; index 2 (original) remains. - assert atomic._original_indices == [2] - assert atomic.seed_groups == [sample_seed_groups[2]] + assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1", "objective3"} + assert retained == keep + + def test_retained_set_excludes_missing_hashes(self, mock_attack, sample_seed_groups): + from pyrit.common.utils import to_sha256 + + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + keep = {to_sha256("objective1"), to_sha256("not-in-dataset")} + retained = atomic.restrict_seed_groups_to_hashes(keep_hashes=keep) + + assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1"} + assert retained == {to_sha256("objective1")} @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackFilterSeedGroupsByObjectives: - """Tests for the deprecated ``filter_seed_groups_by_objectives`` path — - still covered so we know the deprecation warning fires and the legacy - semantics keep working until ``removed_in=0.16.0``.""" + """Deprecated path — verifies the warning fires and legacy semantics work.""" def test_emits_deprecation_warning(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( @@ -1058,8 +1094,6 @@ def test_emits_deprecation_warning(self, mock_attack, sample_seed_groups): atomic.filter_seed_groups_by_objectives(remaining_objectives=["objective2"]) assert [sg.objective.value for sg in atomic.seed_groups] == ["objective2"] - # The kept seed group was at original position 1. - assert atomic._original_indices == [1] def test_drops_all_when_no_objectives_match(self, mock_attack, sample_seed_groups): atomic = AtomicAttack( @@ -1072,20 +1106,44 @@ def test_drops_all_when_no_objectives_match(self, mock_attack, sample_seed_group atomic.filter_seed_groups_by_objectives(remaining_objectives=["nope"]) assert atomic.seed_groups == [] - assert atomic._original_indices == [] @pytest.mark.usefixtures("patch_central_database") -class TestAtomicAttackAttributionFactory: - """Tests for the ``attribution_factory`` closure built in ``run_async`` — - the bridge that lets the AttackExecutor stamp each per-task input index - back onto its *original* position when persisting attack results.""" +class TestAtomicAttackDuplicateObjectiveValidation: + """``AtomicAttack.__init__`` enforces objective-hash uniqueness within a + single atomic attack so resume can use the hash as a stable identity.""" + + def test_constructing_with_duplicate_objective_raises(self, mock_attack): + duplicate_groups = [ + SeedAttackGroup(seeds=[SeedObjective(value="same-objective")]), + SeedAttackGroup(seeds=[SeedObjective(value="same-objective")]), + ] + with pytest.raises(ValueError, match="duplicate objective hash"): + AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=duplicate_groups, + atomic_attack_name="dup", + ) + + def test_constructing_with_unique_objectives_succeeds(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="ok", + ) + assert len(atomic.seed_groups) == 3 + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackAttributionStamping: + """Tests for how ``run_async`` builds the ``AttackResultAttribution`` it + passes to the executor.""" - async def test_no_factory_passed_when_scenario_result_id_unset( + async def test_no_attribution_when_scenario_result_id_unset( self, mock_attack, sample_seed_groups, sample_attack_results ): """Outside a Scenario, ``_scenario_result_id`` is None and the - executor must receive ``attribution_factory=None``.""" + executor must receive ``attribution=None``.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, @@ -1097,13 +1155,13 @@ async def test_no_factory_passed_when_scenario_result_id_unset( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - assert mock_exec.call_args.kwargs["attribution_factory"] is None + assert mock_exec.call_args.kwargs["attribution"] is None - async def test_factory_built_when_scenario_result_id_set( + async def test_attribution_built_when_scenario_result_id_set( self, mock_attack, sample_seed_groups, sample_attack_results ): """When the Scenario stamps ``_scenario_result_id`` onto the atomic - attack, ``run_async`` must build and pass a factory.""" + attack, ``run_async`` must build and pass a single attribution object.""" from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution atomic = AtomicAttack( @@ -1117,63 +1175,7 @@ async def test_factory_built_when_scenario_result_id_set( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - factory = mock_exec.call_args.kwargs["attribution_factory"] - assert factory is not None - - # input_index 0 → original_indices[0] == 0 - attribution = factory(0) + attribution = mock_exec.call_args.kwargs["attribution"] assert isinstance(attribution, AttackResultAttribution) assert attribution.parent_id == "00000000-0000-0000-0000-000000000abc" assert attribution.parent_collection == "MyAtomicAttack" - assert attribution.position == 0 - - # input_index 2 → original_indices[2] == 2 - assert factory(2).position == 2 - - async def test_factory_maps_input_index_to_original_index_after_filtering( - self, mock_attack, sample_seed_groups, sample_attack_results - ): - """The whole point of the factory: after filtering out a completed - seed group, the per-task input index 0 must still map back to the - *original* position of the surviving seed group, not 0.""" - atomic = AtomicAttack( - attack_technique=AttackTechnique(attack=mock_attack), - seed_groups=sample_seed_groups, - atomic_attack_name="ResumeAttack", - ) - atomic._scenario_result_id = "scenario-1" - # Simulate resume: original positions 0 and 1 are done; only position 2 remains. - atomic.filter_seed_groups_by_indices(completed_indices={0, 1}) - - with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: - mock_exec.return_value = wrap_results(sample_attack_results[:1]) - await atomic.run_async() - - factory = mock_exec.call_args.kwargs["attribution_factory"] - assert factory is not None - # input_index 0 (only remaining task) → original index 2. - assert factory(0).position == 2 - - async def test_factory_captures_indices_by_value_not_reference( - self, mock_attack, sample_seed_groups, sample_attack_results - ): - """The closure must snapshot ``_original_indices`` at build time so - mutating the AtomicAttack after ``run_async`` started doesn't poison - in-flight attributions.""" - atomic = AtomicAttack( - attack_technique=AttackTechnique(attack=mock_attack), - seed_groups=sample_seed_groups, - atomic_attack_name="test", - ) - atomic._scenario_result_id = "scenario-1" - - with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: - mock_exec.return_value = wrap_results(sample_attack_results) - await atomic.run_async() - - factory = mock_exec.call_args.kwargs["attribution_factory"] - # Mutate after run_async has captured the snapshot. - atomic._original_indices = [999, 999, 999] - - # Factory still returns the snapshotted original indices. - assert [factory(i).position for i in range(3)] == [0, 1, 2] diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/test_scenario.py index 01825a4ba2..4c247df691 100644 --- a/tests/unit/scenario/test_scenario.py +++ b/tests/unit/scenario/test_scenario.py @@ -41,9 +41,9 @@ def _stamp_scenario_linkage(*, attack_results, atomic_attack): name = getattr(atomic_attack, "atomic_attack_name", None) if not sid or not name: return - for i, r in enumerate(attack_results): + for r in attack_results: r.attribution_parent_id = sid - r.attribution_data = {"parent_collection": name, "position": i} + r.attribution_data = {"parent_collection": name} def create_mock_run_async(attack_results, *, atomic_attack=None): diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/test_scenario_partial_results.py index 0b03bf5a65..f15a1d26bd 100644 --- a/tests/unit/scenario/test_scenario_partial_results.py +++ b/tests/unit/scenario/test_scenario_partial_results.py @@ -47,9 +47,9 @@ def save_attack_results_to_memory(attack_results, *, atomic_attack=None): sid = getattr(atomic_attack, "_scenario_result_id", None) name = getattr(atomic_attack, "atomic_attack_name", None) if sid and name: - for i, r in enumerate(attack_results): + for r in attack_results: r.attribution_parent_id = sid - r.attribution_data = {"parent_collection": name, "position": i} + r.attribution_data = {"parent_collection": name} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) @@ -58,8 +58,10 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: """Create a mock AtomicAttack with required attributes for baseline creation. The mock tracks its objectives and properly updates when - filter_seed_groups_by_objectives OR filter_seed_groups_by_indices is called. + filter_seed_groups_by_objectives OR filter_seed_groups_by_completed_hashes is called. """ + from pyrit.common.utils import to_sha256 + mock_attack_strategy = MagicMock() mock_attack_strategy.get_objective_target.return_value = MagicMock() mock_attack_strategy.get_attack_scoring_config.return_value = MagicMock() @@ -70,35 +72,21 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: attack._attack = mock_attack_strategy attack._scenario_result_id = None - # Track current objectives and their original indices in mutable containers. original_objectives = list(objectives) current_objectives = {"value": list(objectives)} - current_indices = {"value": list(range(len(objectives)))} type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) def filter_objectives(*, remaining_objectives): remaining_set = set(remaining_objectives) - kept = [ - (idx, obj) - for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) - if obj in remaining_set - ] - current_indices["value"] = [i for i, _ in kept] - current_objectives["value"] = [o for _, o in kept] - - def filter_indices(*, completed_indices): - kept = [ - (idx, obj) - for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) - if idx not in completed_indices - ] - current_indices["value"] = [i for i, _ in kept] - current_objectives["value"] = [o for _, o in kept] + current_objectives["value"] = [o for o in current_objectives["value"] if o in remaining_set] + + def filter_completed_hashes(*, completed_hashes): + current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) - attack.filter_seed_groups_by_indices = MagicMock(side_effect=filter_indices) + attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) attack._original_objectives = original_objectives return attack diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index 26a98161f5..da8854a5b2 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -51,9 +51,9 @@ def save_attack_results_to_memory(attack_results, *, atomic_attack=None): sid = getattr(atomic_attack, "_scenario_result_id", None) name = getattr(atomic_attack, "atomic_attack_name", None) if sid and name: - for i, r in enumerate(attack_results): + for r in attack_results: r.attribution_parent_id = sid - r.attribution_data = {"parent_collection": name, "position": i} + r.attribution_data = {"parent_collection": name} memory = CentralMemory.get_memory_instance() memory.add_attack_results_to_memory(attack_results=attack_results) @@ -140,35 +140,24 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack._attack = mock_attack_strategy attack._scenario_result_id = None - # Track objectives + their original indices so the new index-based filtering - # mirrors the new resume path. The legacy text-based filter is wired too - # to support callers that still exercise it. + # Track objectives + objective-hash mapping so the new hash-based filter + # behaves correctly in resume tests. The legacy text-based filter is wired + # too to support callers that still exercise it. + from pyrit.common.utils import to_sha256 + current_objectives = {"value": list(objectives)} - current_indices = {"value": list(range(len(objectives)))} type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) def filter_objectives(*, remaining_objectives): remaining_set = set(remaining_objectives) - kept = [ - (idx, obj) - for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) - if obj in remaining_set - ] - current_indices["value"] = [i for i, _ in kept] - current_objectives["value"] = [o for _, o in kept] - - def filter_indices(*, completed_indices): - kept = [ - (idx, obj) - for idx, obj in zip(current_indices["value"], current_objectives["value"], strict=True) - if idx not in completed_indices - ] - current_indices["value"] = [i for i, _ in kept] - current_objectives["value"] = [o for _, o in kept] + current_objectives["value"] = [o for o in current_objectives["value"] if o in remaining_set] + + def filter_completed_hashes(*, completed_hashes): + current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) - attack.filter_seed_groups_by_indices = MagicMock(side_effect=filter_indices) + attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) if run_async_mock: attack.run_async = run_async_mock @@ -684,61 +673,27 @@ async def second_run(*args, **kwargs): # Resume executed only the missing objectives — the core fix. assert executed == ["o3", "o4"] - async def test_duplicate_objective_text_does_not_collapse_on_resume(self, mock_objective_target): - """Two seed groups with the SAME objective text are distinct slots — - resume must treat them as separate (index-based) rather than collapse - them by text.""" - atomic_attack = create_mock_atomic_attack("dup_attack", ["dup-obj", "dup-obj"]) - - async def first(*args, **kwargs): - partials = [create_attack_result(0, conversation_id="dup-c0", objective="dup-obj")] - save_attack_results_to_memory(partials, atomic_attack=atomic_attack) - raise Exception("crash mid-attack") - - atomic_attack.run_async = first - - scenario = ConcreteScenario( - name="Dup Scenario", - version=1, - atomic_attacks_to_return=[atomic_attack], - ) - await scenario.initialize_async(objective_target=mock_objective_target, max_retries=0) - - with pytest.raises(Exception, match="crash mid-attack"): - await scenario.run_async() - - scenario_result_id = scenario._scenario_result_id - assert scenario_result_id is not None - - # === Resume === - atomic_attack_resume = create_mock_atomic_attack("dup_attack", ["dup-obj", "dup-obj"]) - executed: list[str] = [] - - async def second(*args, **kwargs): - executed.extend(atomic_attack_resume.objectives) - results = [ - create_attack_result(i, conversation_id=f"dup-c-resume-{i}", objective=obj) - for i, obj in enumerate(atomic_attack_resume.objectives, start=1) - ] - save_attack_results_to_memory(results, atomic_attack=atomic_attack_resume) - return AttackExecutorResult(completed_results=results, incomplete_objectives=[]) - - atomic_attack_resume.run_async = second - - scenario_resumed = ConcreteScenario( - name="Dup Scenario", - version=1, - atomic_attacks_to_return=[atomic_attack_resume], - scenario_result_id=scenario_result_id, - ) - await scenario_resumed.initialize_async(objective_target=mock_objective_target, max_retries=0) - await scenario_resumed.run_async() - - # Resume executed exactly ONE "dup-obj" — the missing index — even - # though the text matched the already-done seed group. Index-based - # resume is what makes this work; the deprecated text-based filter - # would have collapsed both seed groups. - assert executed == ["dup-obj"] + async def test_duplicate_objective_text_in_atomic_attack_is_rejected(self, mock_objective_target): + """Resume identity is the objective sha256 within an AtomicAttack, so + the real ``AtomicAttack.__init__`` refuses to construct with duplicate + objective text. We exercise the production constructor here to lock + that contract in (the resume mocks bypass it intentionally).""" + from pyrit.executor.attack import AttackStrategy + from pyrit.models import SeedAttackGroup, SeedObjective + from pyrit.scenario import AtomicAttack + from pyrit.scenario.core.attack_technique import AttackTechnique + + mock_attack = MagicMock(spec=AttackStrategy) + duplicate_groups = [ + SeedAttackGroup(seeds=[SeedObjective(value="dup-obj")]), + SeedAttackGroup(seeds=[SeedObjective(value="dup-obj")]), + ] + with pytest.raises(ValueError, match="duplicate objective hash"): + AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=duplicate_groups, + atomic_attack_name="dup_attack", + ) async def test_duplicate_atomic_attack_name_warns_but_does_not_raise(self, mock_objective_target, caplog): """Duplicate atomic_attack_name within a scenario degrades resume to From ead1d1c30a6013f598b338e236c8119ac7799fa9 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 10:59:32 -0700 Subject: [PATCH 06/10] MAINT: address PR review nits Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attack/core/attack_result_attribution.py | 35 ++++++------------- pyrit/executor/attack/core/attack_strategy.py | 6 ++-- ...5e4f_add_attribution_to_attack_results.py} | 0 tests/unit/memory/test_migration.py | 2 +- tests/unit/scenario/test_scenario.py | 2 +- 5 files changed, 16 insertions(+), 29 deletions(-) rename pyrit/memory/alembic/versions/{9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py => 9c8b7a6d5e4f_add_attribution_to_attack_results.py} (100%) diff --git a/pyrit/executor/attack/core/attack_result_attribution.py b/pyrit/executor/attack/core/attack_result_attribution.py index 208dcf473d..b3c107e591 100644 --- a/pyrit/executor/attack/core/attack_result_attribution.py +++ b/pyrit/executor/attack/core/attack_result_attribution.py @@ -2,15 +2,9 @@ # Licensed under the MIT license. """ -Generic attribution metadata that an upstream orchestrator can stamp onto an -``AttackContext`` so the persisted ``AttackResult`` carries the linkage back -to whatever produced it. - -The attack layer treats this as opaque infrastructure — three string-typed -fields, no scenario semantics. The orchestrator (e.g. ``Scenario``) interprets -them however it likes. Keeping the type in ``executor`` rather than -``scenario`` means the persistence path has no dependency on the -``pyrit.scenario`` package. +Generic attribution metadata an orchestrator stamps onto an +``AttackContext`` so the persisted ``AttackResult`` carries linkage back to +whatever produced it. """ from __future__ import annotations @@ -21,25 +15,18 @@ @dataclass(frozen=True) class AttackResultAttribution: """ - Attribution stamped onto an ``AttackContext`` by an upstream caller (the - ``AttackExecutor`` populates it via an ``attribution_factory``) and copied - onto the resulting ``AttackResult`` by the attack persistence path so the + Attribution copied onto an ``AttackResult`` by the persistence path so the DB row records its lineage. - Both fields are opaque to the attack layer. The orchestrator chooses - what they mean and how to query them back later. For example, - ``Scenario`` uses ``parent_id`` for the scenario result UUID and - ``parent_collection`` for the atomic attack name. Result identity within - a collection is reconstructed from the row's own ``objective_sha256`` - rather than a positional index, so resume is stable under reordering or - re-sampling. + Both fields are opaque to the attack layer; the orchestrator chooses what + they mean. ``Scenario`` uses ``parent_id`` for the scenario result UUID and + ``parent_collection`` for the atomic attack name. Attributes: - parent_id (str): The ID of the parent entity that owns this attack - execution. Persisted to ``AttackResultEntry.attribution_parent_id`` - and indexed (with a foreign key to ``ScenarioResultEntries.id``) - so per-parent hydration and resume lookups are direct. - parent_collection (str): A free-form label naming the per-parent + parent_id (str): The ID of the parent entity. Persisted to + ``AttackResultEntry.attribution_parent_id`` (foreign key to + ``ScenarioResultEntries.id``) so per-parent hydration is indexed. + parent_collection (str): Free-form label naming the per-parent collection this result belongs to. Persisted into ``AttackResultEntry.attribution_data``. """ diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 77fc69dcc4..b1dfb7dbb0 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -231,7 +231,7 @@ async def _on_post_execute( # Stamp attribution onto the result before persistence so the # AttackResultEntry row records its lineage. Outside an orchestrator # _attribution is None and both attribution fields stay None. - self._stamp_attribution(context=event_data.context, result=event_data.result) + self._apply_attribution(context=event_data.context, result=event_data.result) self._logger.debug(f"Attack execution completed in {execution_time_ms}ms") @@ -239,7 +239,7 @@ async def _on_post_execute( self._memory.add_attack_results_to_memory(attack_results=[event_data.result]) @staticmethod - def _stamp_attribution( + def _apply_attribution( *, context: AttackStrategyContextT, result: AttackResult, @@ -331,7 +331,7 @@ async def _on_error_async( # Stamp attribution onto the error result so it is locatable via the # attribution_parent_id foreign key on resume. - self._stamp_attribution(context=context, result=error_result) + self._apply_attribution(context=context, result=error_result) self._memory.add_attack_results_to_memory(attack_results=[error_result]) diff --git a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py similarity index 100% rename from pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py rename to pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 2694b08ade..386775d9fd 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -388,7 +388,7 @@ def test_backfill_is_idempotent_and_does_not_clobber_existing_linkage(): / "memory" / "alembic" / "versions" - / "9c8b7a6d5e4f_add_scenario_linkage_to_attack_results.py" + / "9c8b7a6d5e4f_add_attribution_to_attack_results.py" ) spec = spec_from_file_location("scenario_linkage_migration", migration_path) assert spec is not None and spec.loader is not None diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/test_scenario.py index 4c247df691..edb5cf005c 100644 --- a/tests/unit/scenario/test_scenario.py +++ b/tests/unit/scenario/test_scenario.py @@ -33,7 +33,7 @@ def _stamp_scenario_linkage(*, attack_results, atomic_attack): """ Stamp attribution_parent_id + attribution_data on each AttackResult the same way the real attack persistence path does. Mirrors what - ``_DefaultAttackStrategyEventHandler._stamp_attribution`` does at runtime + ``_DefaultAttackStrategyEventHandler._apply_attribution`` does at runtime so test fixtures that mock out the executor still produce DB rows the new foreign-key-based hydration can find. """ From 452a383ae3d148bd99c7696c76e4983428656c24 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 12:34:49 -0700 Subject: [PATCH 07/10] Rename _hydrate_scenario_attack_results, drop deprecated objective filter --- pyrit/memory/memory_interface.py | 7 ++-- pyrit/scenario/core/atomic_attack.py | 35 +++++-------------- .../test_interface_scenario_results.py | 2 +- tests/unit/scenario/test_atomic_attack.py | 29 --------------- .../scenario/test_scenario_partial_results.py | 7 +--- tests/unit/scenario/test_scenario_retry.py | 10 ++---- 6 files changed, 17 insertions(+), 73 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index afcc9386e0..a7d5828cd5 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2107,7 +2107,7 @@ def get_scenario_results( limit=limit, ) - attack_results_by_scenario = self._hydrate_scenario_attack_results(entries=entries) + attack_results_by_scenario = self._get_attack_results_by_scenario(entries=entries) scenario_results: list[ScenarioResult] = [] for entry in entries: @@ -2226,7 +2226,7 @@ def _query_scenario_result_entries( limit=limit, ) - def _hydrate_scenario_attack_results( + def _get_attack_results_by_scenario( self, *, entries: Sequence[ScenarioResultEntry], @@ -2269,7 +2269,8 @@ def _hydrate_scenario_attack_results( name = data.get("parent_collection") if not name: logger.debug( - f"Skipping AttackResultEntry {row.id} during hydration: attribution_data missing parent_collection" + f"Skipping AttackResultEntry {row.id} during scenario load: " + "attribution_data missing parent_collection" ) continue diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index 7197794efb..bed1ad1443 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -141,10 +141,15 @@ def _validate_unique_objective_hashes(self) -> None: Within a single ``AtomicAttack`` (one ``atomic_attack_name``, one technique), the objective text identifies a unit of work. Duplicates would mean two indistinguishable rows on the write side, which makes - resume reconciliation ambiguous — the new hash-based resume key - treats a set of hashes as already-done, with no way to distinguish - which of two duplicate rows is "the one" that is still outstanding. - Loud failure here beats silent resume corruption later. + resume reconciliation ambiguous — the hash-based resume key treats a + set of hashes as already-done, with no way to distinguish which of two + duplicate rows is "the one" that is still outstanding. Loud failure + here beats silent resume corruption later. + + The hash is currently derived from objective text only. A future + iteration may hash the full ``SeedGroup`` (minus technique-specific + fields) so two seed groups that share an objective string but differ + in other inputs can coexist in one atomic attack. Raises: ValueError: If two seed groups share the same ``objective_sha256``. @@ -186,28 +191,6 @@ def seed_groups(self) -> list[SeedAttackGroup]: """ return list(self._seed_groups) - def filter_seed_groups_by_objectives(self, *, remaining_objectives: list[str]) -> None: - """ - Filter seed groups to only those with objectives in the remaining list. - - Deprecated — prefer ``filter_seed_groups_by_completed_hashes``. Filtering - by objective text collapses two distinct seed groups that happen to share - the same objective text. Hash-based filtering aligns with the - ``objective_sha256`` resume key written to ``AttackResultEntry``. - - Args: - remaining_objectives (List[str]): List of objectives that still need to be executed. - """ - print_deprecation_message( - old_item="AtomicAttack.filter_seed_groups_by_objectives", - new_item="AtomicAttack.filter_seed_groups_by_completed_hashes", - removed_in="0.16.0", - ) - remaining_set = set(remaining_objectives) - self._seed_groups = [ - sg for sg in self._seed_groups if sg.objective is not None and sg.objective.value in remaining_set - ] - def filter_seed_groups_by_completed_hashes(self, *, completed_hashes: set[str]) -> None: """ Drop seed groups whose ``objective_sha256`` is already in the completed set. diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index 16ec2578d1..6ce891943a 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -684,7 +684,7 @@ def _make_attack_result_for_scenario( ) -def test_get_scenario_results_hydrates_via_foreign_key(sqlite_instance: MemoryInterface): +def test_get_scenario_results_loads_attack_results_via_foreign_key(sqlite_instance: MemoryInterface): """When AttackResultEntry rows carry the attribution_parent_id foreign key, hydration picks them up directly — without needing the legacy attack_results_json manifest. This is the path that makes mid-AtomicAttack diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index 1050f1b0d6..7ccec0697c 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -1079,35 +1079,6 @@ def test_retained_set_excludes_missing_hashes(self, mock_attack, sample_seed_gro assert retained == {to_sha256("objective1")} -@pytest.mark.usefixtures("patch_central_database") -class TestAtomicAttackFilterSeedGroupsByObjectives: - """Deprecated path — verifies the warning fires and legacy semantics work.""" - - def test_emits_deprecation_warning(self, mock_attack, sample_seed_groups): - atomic = AtomicAttack( - attack_technique=AttackTechnique(attack=mock_attack), - seed_groups=sample_seed_groups, - atomic_attack_name="test", - ) - - with pytest.warns(DeprecationWarning, match="filter_seed_groups_by_objectives"): - atomic.filter_seed_groups_by_objectives(remaining_objectives=["objective2"]) - - assert [sg.objective.value for sg in atomic.seed_groups] == ["objective2"] - - def test_drops_all_when_no_objectives_match(self, mock_attack, sample_seed_groups): - atomic = AtomicAttack( - attack_technique=AttackTechnique(attack=mock_attack), - seed_groups=sample_seed_groups, - atomic_attack_name="test", - ) - - with pytest.warns(DeprecationWarning): - atomic.filter_seed_groups_by_objectives(remaining_objectives=["nope"]) - - assert atomic.seed_groups == [] - - @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackDuplicateObjectiveValidation: """``AtomicAttack.__init__`` enforces objective-hash uniqueness within a diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/test_scenario_partial_results.py index f15a1d26bd..4959dc2c6d 100644 --- a/tests/unit/scenario/test_scenario_partial_results.py +++ b/tests/unit/scenario/test_scenario_partial_results.py @@ -58,7 +58,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: """Create a mock AtomicAttack with required attributes for baseline creation. The mock tracks its objectives and properly updates when - filter_seed_groups_by_objectives OR filter_seed_groups_by_completed_hashes is called. + filter_seed_groups_by_completed_hashes is called. """ from pyrit.common.utils import to_sha256 @@ -78,14 +78,9 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) - def filter_objectives(*, remaining_objectives): - remaining_set = set(remaining_objectives) - current_objectives["value"] = [o for o in current_objectives["value"] if o in remaining_set] - def filter_completed_hashes(*, completed_hashes): current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] - attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) attack._original_objectives = original_objectives diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index da8854a5b2..7cae6af2c3 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -140,23 +140,17 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack._attack = mock_attack_strategy attack._scenario_result_id = None - # Track objectives + objective-hash mapping so the new hash-based filter - # behaves correctly in resume tests. The legacy text-based filter is wired - # too to support callers that still exercise it. + # Track objectives + objective-hash mapping so the hash-based filter + # behaves correctly in resume tests. from pyrit.common.utils import to_sha256 current_objectives = {"value": list(objectives)} type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) - def filter_objectives(*, remaining_objectives): - remaining_set = set(remaining_objectives) - current_objectives["value"] = [o for o in current_objectives["value"] if o in remaining_set] - def filter_completed_hashes(*, completed_hashes): current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] - attack.filter_seed_groups_by_objectives = MagicMock(side_effect=filter_objectives) attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) if run_async_mock: From 5d907b6b1fabb10add0a4bd0e65ff78cca3496ce Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 14:01:27 -0700 Subject: [PATCH 08/10] scope scenario resume by technique eval hash Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attack/core/attack_result_attribution.py | 26 +++++-- pyrit/executor/attack/core/attack_strategy.py | 5 +- ...d5e4f_add_attribution_to_attack_results.py | 2 +- pyrit/memory/memory_interface.py | 2 +- pyrit/memory/memory_models.py | 2 +- pyrit/models/scenario_result.py | 2 +- pyrit/scenario/core/atomic_attack.py | 72 ++++++++++++++----- pyrit/scenario/core/scenario.py | 59 ++++++++------- tests/unit/scenario/test_atomic_attack.py | 18 +++-- tests/unit/scenario/test_scenario.py | 15 ++++ .../scenario/test_scenario_partial_results.py | 13 ++-- tests/unit/scenario/test_scenario_retry.py | 25 ++++--- 12 files changed, 158 insertions(+), 83 deletions(-) diff --git a/pyrit/executor/attack/core/attack_result_attribution.py b/pyrit/executor/attack/core/attack_result_attribution.py index b3c107e591..2953f7160a 100644 --- a/pyrit/executor/attack/core/attack_result_attribution.py +++ b/pyrit/executor/attack/core/attack_result_attribution.py @@ -18,18 +18,34 @@ class AttackResultAttribution: Attribution copied onto an ``AttackResult`` by the persistence path so the DB row records its lineage. - Both fields are opaque to the attack layer; the orchestrator chooses what - they mean. ``Scenario`` uses ``parent_id`` for the scenario result UUID and - ``parent_collection`` for the atomic attack name. + All fields are opaque to the attack layer; the orchestrator chooses what + they mean. Together, ``(parent_collection, parent_eval_hash)`` form the + "what was running" key the orchestrator can use to scope per-parent + work on resume. Two ``AtomicAttack`` instances sharing a name but using + different techniques (e.g. base64 vs hex encoders) get distinct + ``parent_eval_hash`` values, so resume bookkeeping never cross-pollinates. Attributes: parent_id (str): The ID of the parent entity. Persisted to ``AttackResultEntry.attribution_parent_id`` (foreign key to - ``ScenarioResultEntries.id``) so per-parent hydration is indexed. + ``ScenarioResultEntries.id``) so per-parent loading is indexed. + ``Scenario`` sets this to the scenario result UUID, e.g. + ``self._scenario_result_id`` + (``"8a8d17b9-b671-4a3d-8170-e65ea9b44053"``). parent_collection (str): Free-form label naming the per-parent collection this result belongs to. Persisted into - ``AttackResultEntry.attribution_data``. + ``AttackResultEntry.attribution_data``. ``Scenario`` sets this + to the atomic attack name, e.g. ``self.atomic_attack_name`` + (``"encoded_jailbreaks"``). + parent_eval_hash (str | None): Optional content-addressed hash that + disambiguates configurations sharing the same + ``parent_collection``. Persisted into + ``AttackResultEntry.attribution_data``. ``Scenario`` sets this + to the atomic attack's technique evaluation hash, e.g. + ``self.technique_eval_hash`` (computed via + ``AtomicAttackEvaluationIdentifier``). """ parent_id: str parent_collection: str + parent_eval_hash: str | None = None diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index b1dfb7dbb0..ee8ae379ee 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -261,9 +261,12 @@ def _apply_attribution( if attribution is None: return result.attribution_parent_id = attribution.parent_id - result.attribution_data = { + attribution_data: dict[str, Any] = { "parent_collection": attribution.parent_collection, } + if attribution.parent_eval_hash is not None: + attribution_data["parent_eval_hash"] = attribution.parent_eval_hash + result.attribution_data = attribution_data def _log_attack_outcome(self, result: AttackResult) -> None: """ diff --git a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py index 6f5e6e5b3a..06559825cd 100644 --- a/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py +++ b/pyrit/memory/alembic/versions/9c8b7a6d5e4f_add_attribution_to_attack_results.py @@ -65,7 +65,7 @@ def upgrade() -> None: # ScenarioResultEntries: drop the not-yet-released error_attack_result_ids_json, # and add scenario_metadata for free-form scenario-level JSON state (e.g. - # the persisted sampled_objective_hashes used for resume). + # the persisted objective_hashes used for resume). # Error AttackResults are now linkable via the new attribution_parent_id # foreign key; the per-scenario manifest column is no longer used. # Wrapped in a batch op for SQLite. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index a7d5828cd5..cc68f82506 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -2019,7 +2019,7 @@ def update_scenario_metadata( Replace the ``scenario_metadata`` JSON blob on an existing scenario result. Used by the scenario layer to persist first-run state (e.g. - ``sampled_objective_hashes``) that resume needs to replay. Performs a + ``objective_hashes``) that resume needs to replay. Performs a targeted UPDATE so it doesn't clobber other columns. Args: diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index afb360c5df..2afeb4533f 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1013,7 +1013,7 @@ class ScenarioResultEntry(Base): error_type: Mapped[str | None] = mapped_column(String, nullable=True) # Free-form JSON metadata stamped by the scenario. Currently used to record - # ``sampled_objective_hashes`` — the objective sha256 set chosen on the + # ``objective_hashes`` — the objective sha256 set chosen on the # first run, replayed on resume so a fresh ``random.sample`` can't # silently change which objectives the scenario operates on. Column is # named ``scenario_metadata`` because SQLAlchemy's ``DeclarativeBase`` diff --git a/pyrit/models/scenario_result.py b/pyrit/models/scenario_result.py index b627dd601a..7310034986 100644 --- a/pyrit/models/scenario_result.py +++ b/pyrit/models/scenario_result.py @@ -134,7 +134,7 @@ def __init__( errored during the scenario run. Defaults to an empty list. metadata (Optional[dict[str, Any]]): Free-form JSON metadata persisted with the scenario result. Currently used to record - ``sampled_objective_hashes`` — the objective ``sha256`` set chosen + ``objective_hashes`` — the objective ``sha256`` set chosen on the first run, replayed on resume so a fresh ``random.sample`` can't silently change which objectives the scenario operates on. Keys are not part of any public contract diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index bed1ad1443..fd63ee7c78 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -124,9 +124,10 @@ def __init__( self._objective_scorer = objective_scorer self._memory_labels = memory_labels or {} self._attack_execute_params = attack_execute_params - # Set by Scenario._execute_scenario_async before run_async. When set, - # each persisted AttackResult is linked to this scenario via the - # attribution_parent_id foreign key on AttackResultEntry. + # Set via set_scenario_result_id() by Scenario._execute_scenario_async + # before run_async. When set, each persisted AttackResult is linked to + # the scenario via the attribution_parent_id foreign key on + # AttackResultEntry. self._scenario_result_id: str | None = None logger.info( @@ -134,6 +135,22 @@ def __init__( f"attack type: {type(self._attack_technique.attack).__name__}" ) + def set_scenario_result_id(self, scenario_result_id: str | None) -> None: + """ + Bind this atomic attack to a scenario result for attribution. + + Called by ``Scenario._execute_scenario_async`` before each + ``run_async`` so persisted ``AttackResult`` rows carry the + ``attribution_parent_id`` foreign key back to the scenario. Pass + ``None`` to clear the binding (e.g. when running an atomic attack + outside of a scenario). + + Args: + scenario_result_id (str | None): The scenario result UUID this + atomic attack belongs to, or ``None`` to detach. + """ + self._scenario_result_id = scenario_result_id + def _validate_unique_objective_hashes(self) -> None: """ Ensure each seed group in this atomic attack has a unique objective hash. @@ -171,6 +188,25 @@ def attack_technique(self) -> AttackTechnique: """Get the attack technique for this atomic attack.""" return self._attack_technique + @property + def technique_eval_hash(self) -> str: + """ + Behavioral evaluation hash for this atomic attack's technique configuration. + + Builds an ``AtomicAttack`` identifier from this attack's technique + (without any seed group) and runs it through + ``AtomicAttackEvaluationIdentifier`` so target/scorer/seed-identifier + noise is stripped per the standard atomic-attack eval rules. The + result is stable across resume runs and across different seed groups, + which is what makes it usable as the resume disambiguator alongside + ``atomic_attack_name``. + """ + composite = build_atomic_attack_identifier( + technique_identifier=self._attack_technique.get_identifier(), + seed_group=None, + ) + return AtomicAttackEvaluationIdentifier(composite).eval_hash + @property def objectives(self) -> list[str]: """ @@ -191,9 +227,9 @@ def seed_groups(self) -> list[SeedAttackGroup]: """ return list(self._seed_groups) - def filter_seed_groups_by_completed_hashes(self, *, completed_hashes: set[str]) -> None: + def drop_seed_groups_with_hashes(self, *, hashes: set[str]) -> None: """ - Drop seed groups whose ``objective_sha256`` is already in the completed set. + Drop seed groups whose ``objective_sha256`` is in ``hashes``. This is the resume filter: within an atomic attack, ``objective_sha256`` is the stable identity (enforced unique by ``__init__``). Content-derived @@ -202,33 +238,32 @@ def filter_seed_groups_by_completed_hashes(self, *, completed_hashes: set[str]) from scratch on each ``run_async()``. Args: - completed_hashes (set[str]): SHA256 hashes of objective text for - seed groups that have already produced a (non-error) ``AttackResult``. + hashes (set[str]): SHA256 hashes of objective text for seed groups + to drop (typically those that have already produced a + non-error ``AttackResult``). """ self._seed_groups = [ - sg - for sg in self._seed_groups - if sg.objective is None or to_sha256(sg.objective.value) not in completed_hashes + sg for sg in self._seed_groups if sg.objective is None or to_sha256(sg.objective.value) not in hashes ] - def restrict_seed_groups_to_hashes(self, *, keep_hashes: set[str]) -> set[str]: + def keep_seed_groups_with_hashes(self, *, hashes: set[str]) -> set[str]: """ - Keep only seed groups whose ``objective_sha256`` is in ``keep_hashes``. + Keep only seed groups whose ``objective_sha256`` is in ``hashes``. - Inverse of ``filter_seed_groups_by_completed_hashes``: used on resume to + Inverse of ``drop_seed_groups_with_hashes``: used on resume to replay the originally-sampled subset and ignore any seed groups that were added since (or that landed in this run's fresh ``random.sample`` draw and are no longer in the persisted set). Args: - keep_hashes (set[str]): SHA256 hashes of objective text for seed + hashes (set[str]): SHA256 hashes of objective text for seed groups to keep. Returns: set[str]: The hashes that were actually retained (intersection of - ``keep_hashes`` and the current seed_groups' hashes). The caller - can union these across atomic attacks to detect persisted hashes - that no longer exist in the dataset. + ``hashes`` and the current seed_groups' hashes). The caller can + union these across atomic attacks to detect persisted hashes that + no longer exist in the dataset. """ retained: set[str] = set() new_groups: list[SeedAttackGroup] = [] @@ -236,7 +271,7 @@ def restrict_seed_groups_to_hashes(self, *, keep_hashes: set[str]) -> set[str]: if sg.objective is None: continue sha = to_sha256(sg.objective.value) - if sha in keep_hashes: + if sha in hashes: retained.add(sha) new_groups.append(sg) self._seed_groups = new_groups @@ -305,6 +340,7 @@ async def run_async( attribution = AttackResultAttribution( parent_id=self._scenario_result_id, parent_collection=self.atomic_attack_name, + parent_eval_hash=self.technique_eval_hash, ) results = await executor.execute_attack_from_seed_groups_async( diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index cd35425a31..ef5a374cc5 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -694,21 +694,6 @@ async def initialize_async( seed_groups = self._dataset_config.get_all_seed_attack_groups() self._atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=seed_groups)) - # Enforce atomic_attack_name uniqueness. Duplicate names cause result - # aggregation maps keyed by name to silently collapse rows, and - # foreign-key-based resume cannot distinguish them. Some existing - # scenarios (e.g. Encoding) construct multiple AtomicAttacks with the - # same name for different converter configs, so this is a warning for - # now; in a future release it will become an error. - names = [aa.atomic_attack_name for aa in self._atomic_attacks] - duplicates = sorted({name for name in names if names.count(name) > 1}) - if duplicates: - logger.warning( - f"Scenario '{self._name}' has duplicate atomic_attack_name values: {duplicates}. " - f"Duplicate names will be collapsed in result aggregation and resume tracking. " - f"Each AtomicAttack within a scenario should have a unique name." - ) - # Snapshot params onto the identifier before the resume branch so the identifier # is fully populated regardless of which branch we take. Deep-copy avoids sharing # mutable state with self.params. @@ -729,7 +714,7 @@ async def initialize_async( ) self._validate_stored_scenario(stored_result=existing_results[0]) - self._apply_persisted_sample(stored_result=existing_results[0]) + self._apply_persisted_objectives(stored_result=existing_results[0]) return # Valid resume - skip creating new scenario result # Build display group mapping from atomic attacks @@ -763,7 +748,7 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: unseeded ``random.sample`` and the chosen subset would silently change on the next run (e.g. a resume). To make resume reliable, snapshot the chosen objective hashes here so the next ``_setup_scenario_async`` can - replay them via ``restrict_seed_groups_to_hashes``. + replay them via ``keep_seed_groups_with_hashes``. When ``max_dataset_size`` is not set, the sample equals the dataset and nothing needs pinning; the dict is empty. @@ -784,15 +769,15 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: if sha not in seen: seen.add(sha) hashes.append(sha) - metadata["sampled_objective_hashes"] = hashes + metadata["objective_hashes"] = hashes return metadata - def _apply_persisted_sample(self, *, stored_result: ScenarioResult) -> None: + 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["sampled_objective_hashes"]``. + 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 @@ -807,16 +792,16 @@ def _apply_persisted_sample(self, *, stored_result: ScenarioResult) -> None: currently-resolved dataset. """ metadata = stored_result.metadata or {} - persisted = metadata.get("sampled_objective_hashes") + persisted = metadata.get("objective_hashes") if not persisted: return - keep_hashes: set[str] = set(persisted) + persisted_hashes: set[str] = set(persisted) retained: set[str] = set() for aa in self._atomic_attacks: - retained |= aa.restrict_seed_groups_to_hashes(keep_hashes=keep_hashes) + retained |= aa.keep_seed_groups_with_hashes(hashes=persisted_hashes) - missing = keep_hashes - retained + missing = persisted_hashes - retained if missing: sample = sorted(missing)[:3] raise ValueError( @@ -925,7 +910,7 @@ def _validate_stored_scenario(self, *, stored_result: ScenarioResult) -> None: f"(ID: {self._scenario_result_id}, state: {stored_result.scenario_run_state})" ) - def _get_completed_objective_hashes_for_attack(self, *, atomic_attack_name: str) -> set[str]: + def _get_completed_objective_hashes_for_attack(self, *, atomic_attack: AtomicAttack) -> set[str]: """ Return the set of ``objective_sha256`` values already completed (non-error) for a specific atomic attack inside this scenario. @@ -937,8 +922,16 @@ def _get_completed_objective_hashes_for_attack(self, *, atomic_attack_name: str) Identity is content-derived (``to_sha256(objective)``), so it stays stable even if ``get_seed_groups()`` reorders or resamples between runs. + Rows are matched on ``(parent_collection, parent_eval_hash)`` so that + two ``AtomicAttack`` instances sharing a name but using different + techniques (e.g. base64 vs hex encoders) never cross-pollinate their + completed-hash sets on resume. Rows persisted before + ``parent_eval_hash`` was introduced (or by callers that don't supply + one) match name-only as a backward-compatible fallback. + Args: - atomic_attack_name (str): The atomic attack name to scope to. + 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. @@ -946,6 +939,9 @@ def _get_completed_objective_hashes_for_attack(self, *, atomic_attack_name: str) 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) @@ -956,6 +952,9 @@ def _get_completed_objective_hashes_for_attack(self, *, atomic_attack_name: str) 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: @@ -985,13 +984,11 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: remaining_attacks: list[AtomicAttack] = [] for atomic_attack in self._atomic_attacks: - completed_hashes = self._get_completed_objective_hashes_for_attack( - atomic_attack_name=atomic_attack.atomic_attack_name - ) + completed_hashes = self._get_completed_objective_hashes_for_attack(atomic_attack=atomic_attack) if completed_hashes: original_count = len(atomic_attack.seed_groups) - atomic_attack.filter_seed_groups_by_completed_hashes(completed_hashes=completed_hashes) + atomic_attack.drop_seed_groups_with_hashes(hashes=completed_hashes) remaining_count = len(atomic_attack.seed_groups) if remaining_count == 0: logger.info( @@ -1248,7 +1245,7 @@ async def _execute_scenario_async(self) -> ScenarioResult: # AttackResult carries the attribution_parent_id linkage. This # is what enables mid-run interruption recovery (results are # visible without the post-atomic-attack bulk manifest write). - atomic_attack._scenario_result_id = scenario_result_id + atomic_attack.set_scenario_result_id(scenario_result_id) logger.info( f"Executing atomic attack {i}/{len(self._atomic_attacks)} " diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index 7ccec0697c..e962e24897 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -984,7 +984,7 @@ async def test_enrichment_skips_db_update_when_no_attack_result_id(self, mock_at @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackFilterSeedGroupsByCompletedHashes: - """Tests for ``filter_seed_groups_by_completed_hashes`` — the hash-based + """Tests for ``drop_seed_groups_with_hashes`` — the hash-based resume filter.""" def test_filters_out_completed_hashes(self, mock_attack, sample_seed_groups): @@ -996,7 +996,7 @@ def test_filters_out_completed_hashes(self, mock_attack, sample_seed_groups): atomic_attack_name="test", ) completed = {to_sha256("objective1"), to_sha256("objective3")} - atomic.filter_seed_groups_by_completed_hashes(completed_hashes=completed) + atomic.drop_seed_groups_with_hashes(hashes=completed) assert atomic.seed_groups == [sample_seed_groups[1]] @@ -1007,7 +1007,7 @@ def test_empty_completed_hashes_is_noop(self, mock_attack, sample_seed_groups): atomic_attack_name="test", ) - atomic.filter_seed_groups_by_completed_hashes(completed_hashes=set()) + atomic.drop_seed_groups_with_hashes(hashes=set()) assert atomic.seed_groups == sample_seed_groups @@ -1020,9 +1020,7 @@ def test_all_hashes_completed_clears_seed_groups(self, mock_attack, sample_seed_ atomic_attack_name="test", ) - atomic.filter_seed_groups_by_completed_hashes( - completed_hashes={to_sha256(f"objective{i}") for i in range(1, 4)} - ) + atomic.drop_seed_groups_with_hashes(hashes={to_sha256(f"objective{i}") for i in range(1, 4)}) assert atomic.seed_groups == [] @@ -1039,7 +1037,7 @@ def test_filter_is_stable_across_resampling(self, mock_attack, sample_seed_group # Simulate a re-sample by reversing the internal list. atomic._seed_groups = list(reversed(atomic._seed_groups)) - atomic.filter_seed_groups_by_completed_hashes(completed_hashes={to_sha256("objective1")}) + atomic.drop_seed_groups_with_hashes(hashes={to_sha256("objective1")}) kept_objectives = [sg.objective.value for sg in atomic.seed_groups] assert "objective1" not in kept_objectives assert set(kept_objectives) == {"objective2", "objective3"} @@ -1047,7 +1045,7 @@ def test_filter_is_stable_across_resampling(self, mock_attack, sample_seed_group @pytest.mark.usefixtures("patch_central_database") class TestAtomicAttackRestrictSeedGroupsToHashes: - """Tests for ``restrict_seed_groups_to_hashes`` — the keep-set inverse used + """Tests for ``keep_seed_groups_with_hashes`` — the keep-set inverse used on resume to replay the originally-sampled subset.""" def test_keeps_only_listed_hashes(self, mock_attack, sample_seed_groups): @@ -1059,7 +1057,7 @@ def test_keeps_only_listed_hashes(self, mock_attack, sample_seed_groups): atomic_attack_name="test", ) keep = {to_sha256("objective1"), to_sha256("objective3")} - retained = atomic.restrict_seed_groups_to_hashes(keep_hashes=keep) + retained = atomic.keep_seed_groups_with_hashes(hashes=keep) assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1", "objective3"} assert retained == keep @@ -1073,7 +1071,7 @@ def test_retained_set_excludes_missing_hashes(self, mock_attack, sample_seed_gro atomic_attack_name="test", ) keep = {to_sha256("objective1"), to_sha256("not-in-dataset")} - retained = atomic.restrict_seed_groups_to_hashes(keep_hashes=keep) + retained = atomic.keep_seed_groups_with_hashes(hashes=keep) assert {sg.objective.value for sg in atomic.seed_groups} == {"objective1"} assert retained == {to_sha256("objective1")} diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/test_scenario.py index edb5cf005c..9f3eaf9beb 100644 --- a/tests/unit/scenario/test_scenario.py +++ b/tests/unit/scenario/test_scenario.py @@ -77,18 +77,24 @@ def mock_atomic_attacks(): run1.atomic_attack_name = "attack_run_1" run1.display_group = "attack_run_1" run1._attack = mock_attack + run1._scenario_result_id = None + run1.set_scenario_result_id = MagicMock(side_effect=lambda sid: setattr(run1, "_scenario_result_id", sid)) type(run1).objectives = PropertyMock(return_value=["objective1"]) run2 = MagicMock(spec=AtomicAttack) run2.atomic_attack_name = "attack_run_2" run2.display_group = "attack_run_2" run2._attack = mock_attack + run2._scenario_result_id = None + run2.set_scenario_result_id = MagicMock(side_effect=lambda sid: setattr(run2, "_scenario_result_id", sid)) type(run2).objectives = PropertyMock(return_value=["objective2"]) run3 = MagicMock(spec=AtomicAttack) run3.atomic_attack_name = "attack_run_3" run3.display_group = "attack_run_3" run3._attack = mock_attack + run3._scenario_result_id = None + run3.set_scenario_result_id = MagicMock(side_effect=lambda sid: setattr(run3, "_scenario_result_id", sid)) type(run3).objectives = PropertyMock(return_value=["objective3"]) return [run1, run2, run3] @@ -533,6 +539,10 @@ async def test_atomic_attack_count_with_different_sizes(self, mock_objective_tar single_run_mock.atomic_attack_name = "attack_1" single_run_mock.display_group = "attack_1" single_run_mock._attack = mock_attack + single_run_mock._scenario_result_id = None + single_run_mock.set_scenario_result_id = MagicMock( + side_effect=lambda sid: setattr(single_run_mock, "_scenario_result_id", sid) + ) type(single_run_mock).objectives = PropertyMock(return_value=["obj1"]) single_run = [single_run_mock] @@ -550,6 +560,11 @@ async def test_atomic_attack_count_with_different_sizes(self, mock_objective_tar run.atomic_attack_name = f"attack_{i}" run.display_group = f"attack_{i}" run._attack = mock_attack + run._scenario_result_id = None + # Capture run by default arg to avoid late-binding in the closure. + run.set_scenario_result_id = MagicMock( + side_effect=lambda sid, _run=run: setattr(_run, "_scenario_result_id", sid) + ) type(run).objectives = PropertyMock(return_value=[f"obj{i}"]) many_runs.append(run) diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/test_scenario_partial_results.py index 4959dc2c6d..f146830938 100644 --- a/tests/unit/scenario/test_scenario_partial_results.py +++ b/tests/unit/scenario/test_scenario_partial_results.py @@ -58,7 +58,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: """Create a mock AtomicAttack with required attributes for baseline creation. The mock tracks its objectives and properly updates when - filter_seed_groups_by_completed_hashes is called. + drop_seed_groups_with_hashes is called. """ from pyrit.common.utils import to_sha256 @@ -72,16 +72,21 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: attack._attack = mock_attack_strategy attack._scenario_result_id = None + def _set_scenario_result_id(scenario_result_id): + attack._scenario_result_id = scenario_result_id + + attack.set_scenario_result_id = MagicMock(side_effect=_set_scenario_result_id) + original_objectives = list(objectives) current_objectives = {"value": list(objectives)} type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) - def filter_completed_hashes(*, completed_hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] + def drop_hashes(*, hashes): + current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] - attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) + attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) attack._original_objectives = original_objectives return attack diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index 7cae6af2c3..931be6769d 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -140,6 +140,11 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack._attack = mock_attack_strategy attack._scenario_result_id = None + def _set_scenario_result_id(scenario_result_id): + attack._scenario_result_id = scenario_result_id + + attack.set_scenario_result_id = MagicMock(side_effect=_set_scenario_result_id) + # Track objectives + objective-hash mapping so the hash-based filter # behaves correctly in resume tests. from pyrit.common.utils import to_sha256 @@ -148,10 +153,10 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) - def filter_completed_hashes(*, completed_hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in completed_hashes] + def drop_hashes(*, hashes): + current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] - attack.filter_seed_groups_by_completed_hashes = MagicMock(side_effect=filter_completed_hashes) + attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) if run_async_mock: attack.run_async = run_async_mock @@ -689,11 +694,11 @@ async def test_duplicate_objective_text_in_atomic_attack_is_rejected(self, mock_ atomic_attack_name="dup_attack", ) - async def test_duplicate_atomic_attack_name_warns_but_does_not_raise(self, mock_objective_target, caplog): - """Duplicate atomic_attack_name within a scenario degrades resume to - collapse-by-name. The validator logs a WARNING but does not raise — - some existing scenarios (Encoding, Jailbreak) intentionally use - duplicate names for different converter configs.""" + async def test_duplicate_atomic_attack_name_does_not_warn(self, mock_objective_target, caplog): + """Duplicate ``atomic_attack_name`` is supported: resume disambiguates + rows by ``(parent_collection, parent_eval_hash)``, so two atomic + attacks sharing a name with different techniques don't cross-pollinate + their completed-hash sets. No warning is emitted.""" dup1 = create_mock_atomic_attack("dup_name", ["objA"]) dup2 = create_mock_atomic_attack("dup_name", ["objB"]) @@ -712,6 +717,6 @@ async def noop_run(*args, **kwargs): with caplog.at_level("WARNING"): await scenario.initialize_async(objective_target=mock_objective_target) - assert any("duplicate atomic_attack_name" in record.message for record in caplog.records), ( - "Expected a duplicate-name warning" + assert not any("duplicate atomic_attack_name" in record.message for record in caplog.records), ( + "Duplicate atomic_attack_name should be supported without warning" ) From a0a5b715c2b67ccdc82dabb661e54d516bd3a05b Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 14:09:41 -0700 Subject: [PATCH 09/10] restore filter_seed_groups_by_objectives as deprecated shim Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/atomic_attack.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index fd63ee7c78..1138671f35 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -160,8 +160,7 @@ def _validate_unique_objective_hashes(self) -> None: would mean two indistinguishable rows on the write side, which makes resume reconciliation ambiguous — the hash-based resume key treats a set of hashes as already-done, with no way to distinguish which of two - duplicate rows is "the one" that is still outstanding. Loud failure - here beats silent resume corruption later. + duplicate rows is "the one" that is still outstanding. The hash is currently derived from objective text only. A future iteration may hash the full ``SeedGroup`` (minus technique-specific @@ -246,6 +245,28 @@ def drop_seed_groups_with_hashes(self, *, hashes: set[str]) -> None: sg for sg in self._seed_groups if sg.objective is None or to_sha256(sg.objective.value) not in hashes ] + def filter_seed_groups_by_objectives(self, *, remaining_objectives: list[str]) -> None: + """ + Filter seed groups to only those with objectives in the remaining list. + + .. deprecated:: + Use ``drop_seed_groups_with_hashes`` (or ``keep_seed_groups_with_hashes``) + which keys on content-addressed ``objective_sha256`` instead of + objective text. Scheduled for removal in 0.16.0. + + Args: + remaining_objectives (List[str]): List of objectives that still need to be executed. + """ + print_deprecation_message( + old_item="AtomicAttack.filter_seed_groups_by_objectives(remaining_objectives=...)", + new_item="AtomicAttack.keep_seed_groups_with_hashes(hashes=...)", + removed_in="0.16.0", + ) + remaining_set = set(remaining_objectives) + self._seed_groups = [ + sg for sg in self._seed_groups if sg.objective is not None and sg.objective.value in remaining_set + ] + def keep_seed_groups_with_hashes(self, *, hashes: set[str]) -> set[str]: """ Keep only seed groups whose ``objective_sha256`` is in ``hashes``. From e75bc0c63e0b6e9cb566b81928ade409b274d16a Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 20 May 2026 14:51:59 -0700 Subject: [PATCH 10/10] add tests for resume hash filter, eval-hash disambiguation, and deprecation shim Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/unit/scenario/test_atomic_attack.py | 88 ++++++++++++ tests/unit/scenario/test_scenario_retry.py | 159 +++++++++++++++++++++ 2 files changed, 247 insertions(+) diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/test_atomic_attack.py index e962e24897..b17db749ff 100644 --- a/tests/unit/scenario/test_atomic_attack.py +++ b/tests/unit/scenario/test_atomic_attack.py @@ -1148,3 +1148,91 @@ async def test_attribution_built_when_scenario_result_id_set( assert isinstance(attribution, AttackResultAttribution) assert attribution.parent_id == "00000000-0000-0000-0000-000000000abc" assert attribution.parent_collection == "MyAtomicAttack" + + async def test_attribution_includes_technique_eval_hash( + self, mock_attack, sample_seed_groups, sample_attack_results + ): + """The stamped attribution must carry ``parent_eval_hash`` equal to + ``technique_eval_hash`` so resume disambiguates between two atomic + attacks that share a name but use different techniques.""" + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="MyAtomicAttack", + ) + atomic._scenario_result_id = "00000000-0000-0000-0000-000000000abc" + + with patch.object(AttackExecutor, "execute_attack_from_seed_groups_async", new_callable=AsyncMock) as mock_exec: + mock_exec.return_value = wrap_results(sample_attack_results) + await atomic.run_async() + + attribution = mock_exec.call_args.kwargs["attribution"] + assert attribution.parent_eval_hash is not None + assert attribution.parent_eval_hash == atomic.technique_eval_hash + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackTechniqueEvalHash: + """``technique_eval_hash`` must be stable across seed groups and differ + between distinct technique configurations — it's the resume bucket key.""" + + def test_hash_is_independent_of_seed_groups(self, mock_attack, sample_seed_groups): + a1 = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="same", + ) + a2 = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=[SeedAttackGroup(seeds=[SeedObjective(value="different-objective")])], + atomic_attack_name="same", + ) + assert a1.technique_eval_hash == a2.technique_eval_hash + + def test_hash_differs_for_different_attacks(self, sample_seed_groups): + attack_a = MagicMock(spec=AttackStrategy) + attack_a.get_identifier.return_value = ComponentIdentifier(class_name="AttackA", class_module="pyrit.test") + attack_b = MagicMock(spec=AttackStrategy) + attack_b.get_identifier.return_value = ComponentIdentifier(class_name="AttackB", class_module="pyrit.test") + + a1 = AtomicAttack( + attack_technique=AttackTechnique(attack=attack_a), + seed_groups=sample_seed_groups, + atomic_attack_name="same", + ) + a2 = AtomicAttack( + attack_technique=AttackTechnique(attack=attack_b), + seed_groups=sample_seed_groups, + atomic_attack_name="same", + ) + assert a1.technique_eval_hash != a2.technique_eval_hash + + +@pytest.mark.usefixtures("patch_central_database") +class TestAtomicAttackFilterSeedGroupsByObjectivesDeprecation: + """Tests for the deprecated ``filter_seed_groups_by_objectives`` shim + that ships with v0.13.0 → 0.16.0 deprecation.""" + + def test_emits_deprecation_warning(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + with patch("pyrit.scenario.core.atomic_attack.print_deprecation_message") as mock_dep: + atomic.filter_seed_groups_by_objectives(remaining_objectives=["objective1"]) + assert mock_dep.call_count == 1 + kwargs = mock_dep.call_args.kwargs + assert "filter_seed_groups_by_objectives" in kwargs["old_item"] + assert "keep_seed_groups_with_hashes" in kwargs["new_item"] + assert kwargs["removed_in"] == "0.16.0" + + def test_filters_by_text_match(self, mock_attack, sample_seed_groups): + atomic = AtomicAttack( + attack_technique=AttackTechnique(attack=mock_attack), + seed_groups=sample_seed_groups, + atomic_attack_name="test", + ) + with patch("pyrit.scenario.core.atomic_attack.print_deprecation_message"): + atomic.filter_seed_groups_by_objectives(remaining_objectives=["objective2"]) + assert [sg.objective.value for sg in atomic.seed_groups] == ["objective2"] diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/test_scenario_retry.py index 931be6769d..0add8d4e64 100644 --- a/tests/unit/scenario/test_scenario_retry.py +++ b/tests/unit/scenario/test_scenario_retry.py @@ -720,3 +720,162 @@ async def noop_run(*args, **kwargs): assert not any("duplicate atomic_attack_name" in record.message for record in caplog.records), ( "Duplicate atomic_attack_name should be supported without warning" ) + + +@pytest.mark.usefixtures("patch_central_database") +class TestGetCompletedObjectiveHashesForAttack: + """Direct tests for ``Scenario._get_completed_objective_hashes_for_attack`` + — the filter that excludes already-completed objectives on resume. + + Covers the row-filtering branches: outcome=ERROR rows, rows without + attribution_data, and the technique-disambiguation branch where two + atomic attacks share a name but differ in technique eval hash. + """ + + def _make_scenario(self, scenario_result_id="scn-1"): + scenario = ConcreteScenario(name="S", version=1, atomic_attacks_to_return=[]) + scenario._scenario_result_id = scenario_result_id + 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 + row.attribution_data = attribution_data + row.objective = objective + return row + + 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() + + def test_skips_error_rows(self): + from pyrit.common.utils import to_sha256 + + scenario = self._make_scenario() + scenario._memory.get_attack_results.return_value = [ + self._row( + objective="ok", + outcome=AttackOutcome.SUCCESS, + attribution_data={"parent_collection": "a", "parent_eval_hash": "hash-A"}, + ), + self._row( + objective="failed", + outcome=AttackOutcome.ERROR, + 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")} + + def test_skips_rows_without_attribution_data(self): + from pyrit.common.utils import to_sha256 + + scenario = self._make_scenario() + scenario._memory.get_attack_results.return_value = [ + self._row(objective="legacy", attribution_data=None), + self._row( + objective="new", + 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")} + + def test_skips_rows_with_mismatched_eval_hash(self): + """Two atomic attacks with the same name but different techniques + must not cross-pollinate completed hashes. This is the core Option-B + guarantee.""" + from pyrit.common.utils import to_sha256 + + scenario = self._make_scenario() + scenario._memory.get_attack_results.return_value = [ + self._row( + objective="mine", + attribution_data={"parent_collection": "encoding", "parent_eval_hash": "hash-base64"}, + ), + self._row( + objective="theirs", + 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")} + + def test_backward_compat_matches_name_only_when_eval_hash_missing(self): + """Rows persisted before ``parent_eval_hash`` shipped match name-only + so pre-existing resume runs aren't stranded.""" + from pyrit.common.utils import to_sha256 + + scenario = self._make_scenario() + scenario._memory.get_attack_results.return_value = [ + self._row( + objective="old", + 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")} + + +@pytest.mark.usefixtures("patch_central_database") +class TestApplyPersistedObjectives: + """Direct tests for ``Scenario._apply_persisted_objectives`` — the + resume-time replay that locks subsequent runs to the originally-sampled + objective subset.""" + + def _make_scenario_with_atomics(self, atomics): + scenario = ConcreteScenario(name="S", version=1, atomic_attacks_to_return=[]) + scenario._scenario_result_id = "scn-1" + scenario._atomic_attacks = atomics + return scenario + + def test_noop_when_metadata_has_no_persisted_hashes(self): + atomic = MagicMock(spec=AtomicAttack) + scenario = self._make_scenario_with_atomics([atomic]) + stored = MagicMock() + stored.metadata = {} + scenario._apply_persisted_objectives(stored_result=stored) + atomic.keep_seed_groups_with_hashes.assert_not_called() + + def test_replays_persisted_subset_across_atomics(self): + atomic_a = MagicMock(spec=AtomicAttack) + atomic_a.keep_seed_groups_with_hashes.return_value = {"h1", "h2"} + atomic_b = MagicMock(spec=AtomicAttack) + atomic_b.keep_seed_groups_with_hashes.return_value = {"h3"} + scenario = self._make_scenario_with_atomics([atomic_a, atomic_b]) + + stored = MagicMock() + stored.metadata = {"objective_hashes": ["h1", "h2", "h3"]} + scenario._apply_persisted_objectives(stored_result=stored) + + atomic_a.keep_seed_groups_with_hashes.assert_called_once_with(hashes={"h1", "h2", "h3"}) + atomic_b.keep_seed_groups_with_hashes.assert_called_once_with(hashes={"h1", "h2", "h3"}) + + def test_raises_when_persisted_hash_is_missing(self): + atomic = MagicMock(spec=AtomicAttack) + atomic.keep_seed_groups_with_hashes.return_value = {"h1"} # h2 missing + scenario = self._make_scenario_with_atomics([atomic]) + + stored = MagicMock() + stored.metadata = {"objective_hashes": ["h1", "h2"]} + with pytest.raises(ValueError, match="cannot resume"): + scenario._apply_persisted_objectives(stored_result=stored)