From a28993615b956dc1d7f9b3a757dd927ba6294919 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:02:01 -0700 Subject: [PATCH] FIX Handle attack parameter build failures Settle parameter construction failures deterministically, preserve original input alignment, and execute successful builds in partial mode so scenario resume can avoid duplicate materialization. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2b3468bc-0737-461b-bb82-18dad97d400e --- pyrit/executor/attack/core/attack_executor.py | 90 +++++++- .../attack/core/test_attack_executor.py | 203 +++++++++++++++++- .../unit/scenario/core/test_scenario_retry.py | 102 ++++++++- 3 files changed, 385 insertions(+), 10 deletions(-) diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index 81ad441e80..64d75d1066 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -195,8 +195,10 @@ async def execute_attack_from_seed_groups_async( field_overrides: Optional per-seed-group field overrides. If provided, must match the length of seed_groups. Each dict is passed to 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. + return_partial_on_failure: If True, executes successfully constructed + objectives and returns parameter-build or execution failures alongside + completed results. If False (default), a parameter-build failure + suppresses execution and raises the first exception by input order. attribution: Optional ``AttackResultAttribution`` stamped onto every per-task ``AttackContext`` so the persisted ``AttackResultEntry`` row carries ``attribution_parent_id`` + ``attribution_data``. @@ -239,13 +241,39 @@ async def build_params_async(i: int, sg: AttackSeedGroup) -> AttackParameters: **combined_overrides, ) - params_list = list(await asyncio.gather(*[build_params_async(i, sg) for i, sg in enumerate(seed_groups)])) + build_results = list( + await asyncio.gather( + *[build_params_async(i, sg) for i, sg in enumerate(seed_groups)], + return_exceptions=True, + ) + ) + params_list: list[AttackParameters] = [] + successful_input_indices: list[int] = [] + build_failures: list[tuple[int, str, Exception]] = [] + for index, (seed_group, build_result) in enumerate(zip(seed_groups, build_results, strict=True)): + if isinstance(build_result, Exception): + assert seed_group.objective is not None + build_failures.append((index, seed_group.objective.value, build_result)) + elif isinstance(build_result, BaseException): + raise build_result + else: + params_list.append(build_result) + successful_input_indices.append(index) + + if build_failures and not return_partial_on_failure: + raise build_failures[0][2] - return await self._execute_with_params_list_async( + execution_result = await self._execute_with_params_list_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, attribution=attribution, + input_indices=successful_input_indices, + ) + return self._merge_parameter_build_failures( + build_failures=build_failures, + successful_input_indices=successful_input_indices, + execution_result=execution_result, ) async def execute_attack_async( @@ -323,6 +351,7 @@ async def _execute_with_params_list_async( params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, attribution: AttackResultAttribution | None = None, + input_indices: Sequence[int] | None = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Execute attacks in parallel with a list of pre-built parameters. @@ -337,6 +366,8 @@ async def _execute_with_params_list_async( attribution: Optional ``AttackResultAttribution`` stamped onto every per-task ``AttackContext`` so the persistence path can record orchestrator linkage. + input_indices: Original input positions for ``params_list``. Defaults + to sequential positions when parameters were constructed directly. Returns: AttackExecutorResult with completed results and any incomplete objectives. @@ -357,6 +388,7 @@ async def run_one_async(index: int, params: AttackParameters) -> AttackStrategyR objectives=[p.objective for p in params_list], results_or_exceptions=list(results_or_exceptions), return_partial_on_failure=return_partial_on_failure, + input_indices=input_indices, ) def _process_execution_results( @@ -365,6 +397,7 @@ def _process_execution_results( objectives: Sequence[str], results_or_exceptions: list[Any], return_partial_on_failure: bool, + input_indices: Sequence[int] | None = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ Process results from parallel execution into an AttackExecutorResult. @@ -373,23 +406,30 @@ def _process_execution_results( objectives: The objectives that were executed. results_or_exceptions: Results or exceptions from asyncio.gather. return_partial_on_failure: Whether to return partial results on failure. + input_indices: Original input positions corresponding to ``objectives``. Returns: AttackExecutorResult with completed and incomplete results. Raises: BaseException: If return_partial_on_failure=False and any failed. + ValueError: If input_indices length doesn't match objectives length. """ completed: list[AttackStrategyResultT] = [] incomplete: list[tuple[str, BaseException]] = [] completed_indices: list[int] = [] + source_indices = list(input_indices) if input_indices is not None else list(range(len(objectives))) + if len(source_indices) != len(objectives): + raise ValueError("input_indices length must match objectives length") + + self._raise_first_fatal_exception(results_or_exceptions=results_or_exceptions) - for i, (objective, result) in enumerate(zip(objectives, results_or_exceptions, strict=False)): + for i, (objective, result) in enumerate(zip(objectives, results_or_exceptions, strict=True)): if isinstance(result, BaseException): incomplete.append((objective, result)) else: completed.append(result) - completed_indices.append(i) + completed_indices.append(source_indices[i]) executor_result: AttackExecutorResult[AttackStrategyResultT] = AttackExecutorResult( completed_results=completed, @@ -401,3 +441,41 @@ def _process_execution_results( executor_result.raise_if_incomplete() return executor_result + + @staticmethod + def _raise_first_fatal_exception(*, results_or_exceptions: Sequence[Any]) -> None: + """Propagate cancellation and other non-recoverable base exceptions.""" + for result in results_or_exceptions: + if isinstance(result, BaseException) and not isinstance(result, Exception): + raise result + + @staticmethod + def _merge_parameter_build_failures( + *, + build_failures: Sequence[tuple[int, str, Exception]], + successful_input_indices: Sequence[int], + execution_result: AttackExecutorResult[AttackStrategyResultT], + ) -> AttackExecutorResult[AttackStrategyResultT]: + """ + Merge parameter-build and execution failures by original input position. + + Returns: + The combined executor result. + """ + if not build_failures: + return execution_result + + incomplete_by_index: dict[int, tuple[str, BaseException]] = { + index: (objective, error) for index, objective, error in build_failures + } + completed_indices = set(execution_result.input_indices) + execution_failures = iter(execution_result.incomplete_objectives) + for input_index in successful_input_indices: + if input_index not in completed_indices: + incomplete_by_index[input_index] = next(execution_failures) + + return AttackExecutorResult( + completed_results=execution_result.completed_results, + incomplete_objectives=[incomplete_by_index[index] for index in sorted(incomplete_by_index)], + input_indices=execution_result.input_indices, + ) diff --git a/tests/unit/executor/attack/core/test_attack_executor.py b/tests/unit/executor/attack/core/test_attack_executor.py index e5ccfb0a88..6ed0e0e109 100644 --- a/tests/unit/executor/attack/core/test_attack_executor.py +++ b/tests/unit/executor/attack/core/test_attack_executor.py @@ -10,7 +10,8 @@ import asyncio import dataclasses import uuid -from unittest.mock import AsyncMock, MagicMock +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -25,6 +26,7 @@ AttackOutcome, AttackResult, AttackSeedGroup, + Message, SeedObjective, SeedPrompt, ) @@ -60,6 +62,65 @@ def create_seed_group(objective: str) -> AttackSeedGroup: ) +class _ParameterBuildAbort(BaseException): + """Controlled fatal parameter-build failure.""" + + +class _ParameterBuildSchedule: + """Event-controlled parameter-build schedule with out-of-order failures.""" + + def __init__(self) -> None: + self.all_started = asyncio.Event() + self.a_materialized = asyncio.Event() + self.b_failed = asyncio.Event() + self.started_count = 0 + self.failure_completion_order: list[str] = [] + self.generated_conversation_ids: list[str] = [] + self.successful_params: dict[str, AttackParameters] = {} + self.b_error = RuntimeError("build B failed") + self.c_error = ValueError("build C failed") + + async def build_async(self, *, seed_group: AttackSeedGroup, **_: Any) -> AttackParameters: + objective = seed_group.objective.value + self.started_count += 1 + if self.started_count == 3: + self.all_started.set() + await self.all_started.wait() + + if objective == "A": + return await self._build_a_async() + if objective == "B": + await self.a_materialized.wait() + self.failure_completion_order.append("B") + self.b_failed.set() + raise self.b_error + if objective == "C": + await self.b_failed.wait() + self.failure_completion_order.append("C") + raise self.c_error + return await self._build_other_async(objective=objective) + + async def _build_a_async(self) -> AttackParameters: + params = self._record_materialization(objective="A") + self.a_materialized.set() + await self.b_failed.wait() + return params + + async def _build_other_async(self, *, objective: str) -> AttackParameters: + await self.b_failed.wait() + return self._record_materialization(objective=objective) + + def _record_materialization(self, *, objective: str) -> AttackParameters: + conversation_id = f"conv-{objective}-1" + self.generated_conversation_ids.append(conversation_id) + params = AttackParameters( + objective=objective, + prepended_conversation=[Message.from_prompt(role="user", prompt=conversation_id)], + ) + self.successful_params[objective] = params + return params + + @pytest.mark.usefixtures("patch_central_database") class TestAttackExecutorInitialization: """Tests for AttackExecutor initialization.""" @@ -404,6 +465,129 @@ async def test_validates_explicit_empty_field_overrides_for_seed_groups(self): field_overrides=[], ) + async def test_parameter_build_failure_returns_partial_results_in_input_order(self) -> None: + """Successful side-effectful builds execute while build failures retain input order.""" + attack = create_mock_attack() + schedule = _ParameterBuildSchedule() + build_mock = AsyncMock(side_effect=schedule.build_async) + executed_params: list[AttackParameters] = [] + + async def execute_async(*, context: SingleTurnAttackContext) -> AttackResult: + executed_params.append(context.params) + return create_attack_result(context.params.objective) + + attack.execute_with_context_async.side_effect = execute_async + seed_groups = [create_seed_group(objective) for objective in ["C", "A", "B"]] + + with patch.object(AttackParameters, "from_seed_group_async", new=build_mock): + result = await AttackExecutor(max_concurrency=3).execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + return_partial_on_failure=True, + ) + + assert schedule.failure_completion_order == ["B", "C"] + assert schedule.generated_conversation_ids == ["conv-A-1"] + assert executed_params == [schedule.successful_params["A"]] + assert [completed.objective for completed in result.completed_results] == ["A"] + assert result.input_indices == [1] + assert [objective for objective, _ in result.incomplete_objectives] == ["C", "B"] + assert result.incomplete_objectives[0][1] is schedule.c_error + assert result.incomplete_objectives[1][1] is schedule.b_error + + async def test_parameter_build_failure_strict_mode_suppresses_execution(self) -> None: + """Strict mode settles all builds and raises the first input-ordered failure.""" + attack = create_mock_attack() + schedule = _ParameterBuildSchedule() + build_mock = AsyncMock(side_effect=schedule.build_async) + seed_groups = [create_seed_group(objective) for objective in ["C", "A", "B"]] + + with ( + patch.object(AttackParameters, "from_seed_group_async", new=build_mock), + pytest.raises(ValueError, match="build C failed") as exc_info, + ): + await AttackExecutor(max_concurrency=3).execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + ) + + assert exc_info.value is schedule.c_error + assert schedule.started_count == 3 + assert schedule.failure_completion_order == ["B", "C"] + assert schedule.generated_conversation_ids == ["conv-A-1"] + attack.execute_with_context_async.assert_not_awaited() + + async def test_build_and_execution_failures_preserve_original_input_order(self) -> None: + """Build and execution failures merge by original seed-group position.""" + attack = create_mock_attack() + schedule = _ParameterBuildSchedule() + build_mock = AsyncMock(side_effect=schedule.build_async) + a_execution_failed = asyncio.Event() + a_error = LookupError("execute A failed") + + async def execute_async(*, context: SingleTurnAttackContext) -> AttackResult: + if context.params.objective == "A": + a_execution_failed.set() + raise a_error + await a_execution_failed.wait() + return create_attack_result(context.params.objective) + + attack.execute_with_context_async.side_effect = execute_async + seed_groups = [create_seed_group(objective) for objective in ["C", "A", "D", "B"]] + + with patch.object(AttackParameters, "from_seed_group_async", new=build_mock): + result = await AttackExecutor(max_concurrency=4).execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=seed_groups, + return_partial_on_failure=True, + ) + + assert schedule.failure_completion_order == ["B", "C"] + assert schedule.generated_conversation_ids == ["conv-A-1", "conv-D-1"] + assert [completed.objective for completed in result.completed_results] == ["D"] + assert result.input_indices == [2] + assert [objective for objective, _ in result.incomplete_objectives] == ["C", "A", "B"] + assert result.incomplete_objectives[0][1] is schedule.c_error + assert result.incomplete_objectives[1][1] is a_error + assert result.incomplete_objectives[2][1] is schedule.b_error + + @pytest.mark.parametrize("fatal_type", [asyncio.CancelledError, _ParameterBuildAbort]) + async def test_parameter_build_base_exception_propagates( + self, + fatal_type: type[BaseException], + ) -> None: + """Cancellation and other fatal base exceptions are never partial results.""" + attack = create_mock_attack() + all_started = asyncio.Event() + started_count = 0 + fatal_error = fatal_type("fatal build") + + async def build_async(*, seed_group: AttackSeedGroup, **_: Any) -> AttackParameters: + nonlocal started_count + started_count += 1 + if started_count == 2: + all_started.set() + await all_started.wait() + if seed_group.objective.value == "B": + raise fatal_error + return AttackParameters(objective=seed_group.objective.value) + + build_mock = AsyncMock(side_effect=build_async) + with ( + patch.object(AttackParameters, "from_seed_group_async", new=build_mock), + pytest.raises(fatal_type) as exc_info, + ): + await AttackExecutor(max_concurrency=2).execute_attack_from_seed_groups_async( + attack=attack, + seed_groups=[create_seed_group("A"), create_seed_group("B")], + return_partial_on_failure=True, + ) + + if fatal_type is not asyncio.CancelledError: + assert exc_info.value is fatal_error + assert build_mock.await_count == 2 + attack.execute_with_context_async.assert_not_awaited() + @pytest.mark.usefixtures("patch_central_database") class TestAttributionPropagation: @@ -577,6 +761,23 @@ async def mock_execute(*, context): objectives=["Test"], ) + @pytest.mark.parametrize("fatal_type", [asyncio.CancelledError, _ParameterBuildAbort]) + async def test_execution_base_exception_propagates( + self, + fatal_type: type[BaseException], + ) -> None: + """Cancellation and other fatal base exceptions are not incomplete objectives.""" + attack = create_mock_attack() + fatal_error = fatal_type("fatal execution") + attack.execute_with_context_async.side_effect = fatal_error + + with pytest.raises(fatal_type): + await AttackExecutor().execute_attack_async( + attack=attack, + objectives=["Test"], + return_partial_on_failure=True, + ) + @pytest.mark.usefixtures("patch_central_database") class TestAttackExecutorResult: diff --git a/tests/unit/scenario/core/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py index efd3651939..f1ed5a1572 100644 --- a/tests/unit/scenario/core/test_scenario_retry.py +++ b/tests/unit/scenario/core/test_scenario_retry.py @@ -3,16 +3,19 @@ """Tests for Scenario retry functionality.""" +import asyncio from typing import ClassVar -from unittest.mock import AsyncMock, MagicMock, PropertyMock +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest +from pyrit.executor.attack import AttackParameters, AttackStrategy, SingleTurnAttackContext from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier +from pyrit.models import AttackOutcome, AttackResult, AttackSeedGroup, ComponentIdentifier, Message, SeedObjective +from pyrit.prompt_target import PromptTarget from pyrit.scenario import DatasetConfiguration, ScenarioResult -from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy, Scenario, ScenarioTechnique +from pyrit.scenario.core import AtomicAttack, AttackTechnique, BaselineAttackPolicy, Scenario, ScenarioTechnique # Test constants TEST_ATTACK_TYPE = "TestAttack" @@ -162,6 +165,33 @@ def drop_hashes(*, hashes): return attack +class _RetryLinkageAttack(AttackStrategy[SingleTurnAttackContext, AttackResult]): + """Minimal real strategy that exercises production result persistence.""" + + def __init__(self, *, objective_target: PromptTarget) -> None: + super().__init__(objective_target=objective_target, context_type=SingleTurnAttackContext) + self.executed_materializations: list[str] = [] + + def _validate_context(self, *, context: SingleTurnAttackContext) -> None: + pass + + async def _setup_async(self, *, context: SingleTurnAttackContext) -> None: + pass + + async def _perform_async(self, *, context: SingleTurnAttackContext) -> AttackResult: + assert context.params.prepended_conversation is not None + self.executed_materializations.append(context.params.prepended_conversation[0].get_value()) + return AttackResult( + conversation_id=f"outer-{context.params.objective}", + objective=context.params.objective, + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ) + + async def _teardown_async(self, *, context: SingleTurnAttackContext) -> None: + pass + + class ConcreteScenario(Scenario): """Concrete implementation of Scenario for testing.""" @@ -437,6 +467,72 @@ async def mock_run_with_logged_failure(*args, **kwargs): class TestScenarioResumption: """Tests for Scenario resumption after partial failure.""" + async def test_parameter_build_partial_result_persists_linkage_before_retry( + self, + mock_objective_target: MagicMock, + ) -> None: + """A successful build is executed, linked, and not materialized again on retry.""" + target = MagicMock(spec=PromptTarget) + target.get_identifier.return_value = ComponentIdentifier( + class_name="RetryLinkageTarget", + class_module=TEST_MODULE, + ) + attack = _RetryLinkageAttack(objective_target=target) + atomic_attack = AtomicAttack( + atomic_attack_name="retry_linkage", + attack_technique=AttackTechnique(attack=attack), + seed_groups=[ + AttackSeedGroup(seeds=[SeedObjective(value="A")]), + AttackSeedGroup(seeds=[SeedObjective(value="B")]), + ], + ) + scenario = ConcreteScenario( + name="Parameter Build Retry", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "max_concurrency": 2, + "max_retries": 1, + } + ) + await scenario.initialize_async() + + a_built = asyncio.Event() + build_counts = {"A": 0, "B": 0} + generated_conversations: list[str] = [] + + async def build_async(*, seed_group: AttackSeedGroup, **_: object) -> AttackParameters: + objective = seed_group.objective.value + build_counts[objective] += 1 + if objective == "A": + a_built.set() + elif build_counts[objective] == 1: + await a_built.wait() + raise RuntimeError("build B failed") + + conversation_id = f"conv-{objective}-{build_counts[objective]}" + generated_conversations.append(conversation_id) + return AttackParameters( + objective=objective, + prepended_conversation=[Message.from_prompt(role="user", prompt=conversation_id)], + ) + + with patch.object(AttackParameters, "from_seed_group_async", new=AsyncMock(side_effect=build_async)): + result = await scenario.run_async() + + assert build_counts == {"A": 1, "B": 2} + assert generated_conversations == ["conv-A-1", "conv-B-2"] + assert attack.executed_materializations == generated_conversations + assert [item.objective for item in result.attack_results["retry_linkage"]] == ["A", "B"] + + persisted_results = CentralMemory.get_memory_instance().get_attack_results( + scenario_result_id=scenario._scenario_result_id + ) + assert [item.objective for item in persisted_results] == ["A", "B"] + async def test_resumes_from_partial_completion_single_attack(self, mock_objective_target): """Test that scenario resumes from where it left off when an atomic attack partially completes.""" objectives = ["obj1", "obj2", "obj3", "obj4"]