From 792c98bfe86c0ec1293b1e2c12c694850874f361 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sun, 23 Aug 2026 01:29:09 -0700 Subject: [PATCH 1/3] FIX Propagate strategy event handler failures Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3772242b-0ee1-4ecd-a028-a38c91c17926 --- pyrit/executor/core/strategy.py | 5 +- .../attack/core/test_attack_strategy.py | 50 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/pyrit/executor/core/strategy.py b/pyrit/executor/core/strategy.py index cb5e619aab..14d8c1f61b 100644 --- a/pyrit/executor/core/strategy.py +++ b/pyrit/executor/core/strategy.py @@ -284,7 +284,10 @@ async def _handle_event_async( tasks = [ asyncio.create_task(handler.on_event_async(event_data)) for handler in self._event_handlers.values() ] - await asyncio.gather(*tasks, return_exceptions=True) + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + for outcome in outcomes: + if isinstance(outcome, BaseException): + raise outcome @asynccontextmanager async def _execution_context_async(self, context: StrategyContextT) -> AsyncIterator[None]: diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 6d588c447c..cb5fd75ca8 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -905,6 +905,56 @@ async def _teardown_async(self, *, context): # Current behavior: execution_time_ms is not modified by event handler assert result.execution_time_ms == 500 + async def test_post_execute_persistence_failure_is_propagated_and_recorded( + self, mock_objective_target: PromptTarget + ) -> None: + teardown_calls = 0 + + class TestStrategy(AttackStrategy): + def _validate_context(self, *, context: AttackContext) -> None: + pass + + async def _setup_async(self, *, context: AttackContext) -> None: + pass + + async def _perform_async(self, *, context: AttackContext) -> AttackResult: + return AttackResult( + conversation_id="test-conversation-id", + objective=context.objective, + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ) + + async def _teardown_async(self, *, context: AttackContext) -> None: + nonlocal teardown_calls + teardown_calls += 1 + + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) + memory = CentralMemory.get_memory_instance() + persist = memory.add_attack_results_to_memory + persist_calls = 0 + + def fail_first_persist(*, attack_results: list[AttackResult]) -> None: + nonlocal persist_calls + persist_calls += 1 + if persist_calls == 1: + raise RuntimeError("database unavailable") + persist(attack_results=attack_results) + + with ( + patch.object(memory, "add_attack_results_to_memory", side_effect=fail_first_persist), + pytest.raises(RuntimeError, match="database unavailable") as exc_info, + ): + await strategy.execute_async(objective="Test objective") + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert teardown_calls == 1 + assert persist_calls == 2 + [stored_result] = memory.get_attack_results(objective="Test objective") + assert stored_result.outcome is AttackOutcome.ERROR + assert stored_result.error_type == "RuntimeError" + assert stored_result.error_message == "database unavailable" + async def test_cancellation_clears_retry_collector(self, mock_objective_target): teardown_calls = 0 From ddda6ba182264e1e5bf4116fa35edf50429d06d8 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sun, 23 Aug 2026 20:32:54 -0700 Subject: [PATCH 2/3] FIX Preserve lifecycle on persistence failures Persist completed attack results after teardown without making generic event observers fatal or creating duplicate error results. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 75013b84-0c8d-4ea4-9c4c-6c58381ac1a1 --- pyrit/executor/attack/core/attack_strategy.py | 32 ++++++- pyrit/executor/core/strategy.py | 5 +- .../attack/core/test_attack_strategy.py | 93 +++++++++++++++---- tests/unit/executor/core/test_strategy.py | 54 ++++++++++- 4 files changed, 155 insertions(+), 29 deletions(-) diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 11c58b60ad..7ed1364045 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -225,7 +225,7 @@ async def _on_post_execute_async( """ Handle post-execution logic after the attack strategy has run. - Attaches retry events to the result and persists it to memory. + Attaches execution metadata to the result and logs its outcome. Args: event_data (StrategyEventData[AttackStrategyContextT, AttackStrategyResultT]): The event data containing @@ -256,7 +256,15 @@ async def _on_post_execute_async( 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]) + + def _persist_result(self, *, result: AttackStrategyResultT) -> None: + """ + Persist a completed attack result. + + Args: + result (AttackStrategyResultT): The completed result to persist. + """ + self._memory.add_attack_results_to_memory(attack_results=[result]) @staticmethod def _apply_attribution( @@ -435,13 +443,13 @@ def __init__( a params type that rejects certain fields. logger (logging.Logger): Logger instance for logging events. """ + event_handler = _DefaultAttackStrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT](logger=logger) super().__init__( context_type=context_type, - event_handler=_DefaultAttackStrategyEventHandler[AttackStrategyContextT, AttackStrategyResultT]( - logger=logger - ), + event_handler=event_handler, logger=logger, ) + self._default_event_handler = event_handler type(self).TARGET_REQUIREMENTS.validate(target=objective_target) self._objective_target = objective_target self._params_type = params_type @@ -617,6 +625,20 @@ def get_request_converters(self) -> list[Any]: """ return self._request_converters + async def execute_with_context_async(self, *, context: AttackStrategyContextT) -> AttackStrategyResultT: + """ + Execute an attack and persist its completed result after teardown. + + Args: + context (AttackStrategyContextT): The attack execution context. + + Returns: + AttackStrategyResultT: The completed and persisted attack result. + """ + result = await super().execute_with_context_async(context=context) + self._default_event_handler._persist_result(result=result) + return result + @overload async def execute_async( self, diff --git a/pyrit/executor/core/strategy.py b/pyrit/executor/core/strategy.py index 14d8c1f61b..cb5e619aab 100644 --- a/pyrit/executor/core/strategy.py +++ b/pyrit/executor/core/strategy.py @@ -284,10 +284,7 @@ async def _handle_event_async( tasks = [ asyncio.create_task(handler.on_event_async(event_data)) for handler in self._event_handlers.values() ] - outcomes = await asyncio.gather(*tasks, return_exceptions=True) - for outcome in outcomes: - if isinstance(outcome, BaseException): - raise outcome + await asyncio.gather(*tasks, return_exceptions=True) @asynccontextmanager async def _execution_context_async(self, context: StrategyContextT) -> AsyncIterator[None]: diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index cb5fd75ca8..5469e7e466 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -470,8 +470,8 @@ async def test_on_post_execute_logs_error_outcome( expected_message = f"{event_handler.__class__.__name__} failed with an error. Reason: Connection timeout" mock_logger.info.assert_called_with(expected_message) - async def test_on_post_execute_adds_results_to_memory(self, mock_memory): - """Test that post-execute handler adds results to memory""" + async def test_on_post_execute_does_not_persist_result(self, mock_memory): + """Completed results are persisted only after the full lifecycle.""" with patch("pyrit.memory.central_memory.CentralMemory.get_memory_instance", return_value=mock_memory): handler = _DefaultAttackStrategyEventHandler() @@ -494,7 +494,7 @@ async def test_on_post_execute_adds_results_to_memory(self, mock_memory): with patch("time.perf_counter", return_value=100.1): await handler.on_event_async(event_data) - mock_memory.add_attack_results_to_memory.assert_called_once_with(attack_results=[sample_result]) + mock_memory.add_attack_results_to_memory.assert_not_called() async def test_on_post_execute_raises_on_none_result(self, event_handler, sample_attack_context, mock_logger): """Test that post-execute handler raises error for None result""" @@ -905,7 +905,7 @@ async def _teardown_async(self, *, context): # Current behavior: execution_time_ms is not modified by event handler assert result.execution_time_ms == 500 - async def test_post_execute_persistence_failure_is_propagated_and_recorded( + async def test_completed_result_persistence_failure_is_propagated_after_teardown( self, mock_objective_target: PromptTarget ) -> None: teardown_calls = 0 @@ -929,31 +929,92 @@ async def _teardown_async(self, *, context: AttackContext) -> None: nonlocal teardown_calls teardown_calls += 1 + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) + memory = CentralMemory.get_memory_instance() + + with ( + patch.object( + memory, + "add_attack_results_to_memory", + side_effect=RuntimeError("database unavailable"), + ) as persist, + pytest.raises(RuntimeError, match="database unavailable"), + ): + await strategy.execute_async(objective="Test objective") + + assert teardown_calls == 1 + persist.assert_called_once() + assert memory.get_attack_results(objective="Test objective") == [] + + async def test_partial_persistence_failure_does_not_create_error_result( + self, mock_objective_target: PromptTarget + ) -> None: + class TestStrategy(AttackStrategy): + def _validate_context(self, *, context: AttackContext) -> None: + pass + + async def _setup_async(self, *, context: AttackContext) -> None: + pass + + async def _perform_async(self, *, context: AttackContext) -> AttackResult: + return AttackResult( + conversation_id="test-conversation-id", + objective=context.objective, + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ) + + async def _teardown_async(self, *, context: AttackContext) -> None: + pass + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) memory = CentralMemory.get_memory_instance() persist = memory.add_attack_results_to_memory - persist_calls = 0 - def fail_first_persist(*, attack_results: list[AttackResult]) -> None: - nonlocal persist_calls - persist_calls += 1 - if persist_calls == 1: - raise RuntimeError("database unavailable") + def persist_then_fail(*, attack_results: list[AttackResult]) -> None: persist(attack_results=attack_results) + raise RuntimeError("commit acknowledgement lost") with ( - patch.object(memory, "add_attack_results_to_memory", side_effect=fail_first_persist), - pytest.raises(RuntimeError, match="database unavailable") as exc_info, + patch.object(memory, "add_attack_results_to_memory", side_effect=persist_then_fail) as persist_mock, + pytest.raises(RuntimeError, match="commit acknowledgement lost"), ): await strategy.execute_async(objective="Test objective") + persist_mock.assert_called_once() + [stored_result] = memory.get_attack_results(objective="Test objective") + assert stored_result.outcome is AttackOutcome.SUCCESS + + async def test_teardown_failure_persists_only_error_result(self, mock_objective_target: PromptTarget) -> None: + class TestStrategy(AttackStrategy): + def _validate_context(self, *, context: AttackContext) -> None: + pass + + async def _setup_async(self, *, context: AttackContext) -> None: + pass + + async def _perform_async(self, *, context: AttackContext) -> AttackResult: + return AttackResult( + conversation_id="test-conversation-id", + objective=context.objective, + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ) + + async def _teardown_async(self, *, context: AttackContext) -> None: + raise RuntimeError("teardown failed") + + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) + memory = CentralMemory.get_memory_instance() + + with pytest.raises(RuntimeError, match="Strategy execution failed") as exc_info: + await strategy.execute_async(objective="Test objective") + assert isinstance(exc_info.value.__cause__, RuntimeError) - assert teardown_calls == 1 - assert persist_calls == 2 + assert str(exc_info.value.__cause__) == "teardown failed" [stored_result] = memory.get_attack_results(objective="Test objective") assert stored_result.outcome is AttackOutcome.ERROR - assert stored_result.error_type == "RuntimeError" - assert stored_result.error_message == "database unavailable" + assert stored_result.error_message == "teardown failed" async def test_cancellation_clears_retry_collector(self, mock_objective_target): teardown_calls = 0 diff --git a/tests/unit/executor/core/test_strategy.py b/tests/unit/executor/core/test_strategy.py index e4a67d93b7..c380fca1e4 100644 --- a/tests/unit/executor/core/test_strategy.py +++ b/tests/unit/executor/core/test_strategy.py @@ -10,7 +10,13 @@ clear_execution_context, execution_context, ) -from pyrit.executor.core.strategy import Strategy, StrategyContext +from pyrit.executor.core.strategy import ( + Strategy, + StrategyContext, + StrategyEvent, + StrategyEventData, + StrategyEventHandler, +) from pyrit.models import ComponentIdentifier @@ -21,14 +27,32 @@ class MockContext(StrategyContext): value: str = "test" +class RaisingEventHandler(StrategyEventHandler[MockContext, str]): + """An event handler that fails on one configured lifecycle event.""" + + def __init__(self, *, event: StrategyEvent) -> None: + self._event = event + + async def on_event_async(self, event_data: StrategyEventData[MockContext, str]) -> None: + if event_data.event is self._event: + raise RuntimeError(f"{self._event.value} observer failed") + + class MockStrategy(Strategy[MockContext, str]): """A mock strategy for testing.""" - def __init__(self, perform_result: str = "success", perform_exception: Exception = None): + def __init__( + self, + *, + perform_result: str = "success", + perform_exception: Exception | None = None, + event_handler: StrategyEventHandler[MockContext, str] | None = None, + ) -> None: # Initialize base class with the context type - super().__init__(context_type=MockContext) + super().__init__(context_type=MockContext, event_handler=event_handler) self._perform_result = perform_result self._perform_exception = perform_exception + self.teardown_calls = 0 async def _setup_async(self, *, context: MockContext) -> None: pass @@ -39,7 +63,7 @@ async def _perform_async(self, *, context: MockContext) -> str: return self._perform_result async def _teardown_async(self, *, context: MockContext) -> None: - pass + self.teardown_calls += 1 def _validate_context(self, *, context: MockContext) -> None: pass @@ -116,6 +140,28 @@ async def test_execute_with_context_preserves_root_cause(self): # The __cause__ should be the original exception assert exc_info.value.__cause__ is original_error + async def test_pre_teardown_observer_failure_does_not_skip_teardown(self): + """Observer failures must not interrupt lifecycle cleanup.""" + strategy = MockStrategy(event_handler=RaisingEventHandler(event=StrategyEvent.ON_PRE_TEARDOWN)) + + result = await strategy.execute_with_context_async(context=MockContext()) + + assert result == "success" + assert strategy.teardown_calls == 1 + + async def test_error_observer_failure_does_not_mask_original_error(self): + """Observer failures must not replace the strategy's root cause.""" + original_error = ValueError("perform failed") + strategy = MockStrategy( + perform_exception=original_error, + event_handler=RaisingEventHandler(event=StrategyEvent.ON_ERROR), + ) + + with pytest.raises(RuntimeError) as exc_info: + await strategy.execute_with_context_async(context=MockContext()) + + assert exc_info.value.__cause__ is original_error + async def test_execute_with_context_extracts_root_cause(self): """Test that chained exceptions show root cause in error message.""" # Create a chain of exceptions From 6807f6fff8daa8a53c45887a5c88f6cdf730a16b Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Mon, 24 Aug 2026 21:28:33 -0700 Subject: [PATCH 3/3] FIX Surface attack and persistence failures Capture failures from the existing error-result write and report them alongside attack failures without retrying ambiguous database commits. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 75013b84-0c8d-4ea4-9c4c-6c58381ac1a1 --- pyrit/executor/attack/core/attack_strategy.py | 29 ++++++- .../attack/core/test_attack_strategy.py | 75 +++++++++++++++++++ 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 7ed1364045..469c746370 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -13,6 +13,11 @@ from enum import Enum from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, overload +try: + from builtins import ExceptionGroup # type: ignore[attr-defined,ty:unresolved-import] +except ImportError: # pragma: no cover - exercised only on 3.10 + from exceptiongroup import ExceptionGroup # type: ignore[no-redef,ty:unresolved-import] + from pyrit.common.logger import logger from pyrit.exceptions.retry_collector import ( get_retry_collector, @@ -87,6 +92,7 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _next_message_override: Message | None | _NextMessageOverrideState = _NextMessageOverrideState.UNSET _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None + _error_result_persistence_error: Exception | None = field(default=None, init=False, repr=False) # Optional attribution from an upstream orchestrator (e.g. Scenario). When # set, the persistence path stamps attribution_parent_id + attribution_data @@ -392,7 +398,10 @@ async def _on_error_async( self._apply_attribution(context=context, result=error_result) self._apply_targeted_harm_categories(context=context, result=error_result) - self._memory.add_attack_results_to_memory(attack_results=[error_result]) + try: + self._memory.add_attack_results_to_memory(attack_results=[error_result]) + except Exception as persistence_error: + context._error_result_persistence_error = persistence_error self._logger.error(f"Attack failed with {type(error).__name__}: {error}") @@ -634,8 +643,22 @@ async def execute_with_context_async(self, *, context: AttackStrategyContextT) - Returns: AttackStrategyResultT: The completed and persisted attack result. - """ - result = await super().execute_with_context_async(context=context) + + Raises: + ExceptionGroup: If attack execution and recording its error result both fail. + """ + context._error_result_persistence_error = None + try: + result = await super().execute_with_context_async(context=context) + except Exception as attack_error: + persistence_error = context._error_result_persistence_error + if persistence_error is not None: + raise ExceptionGroup( + "Attack execution and error result persistence failed", + [attack_error, persistence_error], + ) from None + raise + self._default_event_handler._persist_result(result=result) return result diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 5469e7e466..9cf0fbeb9d 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -8,6 +8,11 @@ import pytest +try: + from builtins import ExceptionGroup # type: ignore[attr-defined,ty:unresolved-import] +except ImportError: # pragma: no cover - exercised only on 3.10 + from exceptiongroup import ExceptionGroup # type: ignore[no-redef,ty:unresolved-import] + from pyrit.exceptions.retry_collector import RetryCollector, get_retry_collector from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters @@ -1016,6 +1021,76 @@ async def _teardown_async(self, *, context: AttackContext) -> None: assert stored_result.outcome is AttackOutcome.ERROR assert stored_result.error_message == "teardown failed" + async def test_attack_and_error_result_persistence_failures_are_both_propagated( + self, mock_objective_target: PromptTarget + ) -> None: + attack_cause = ValueError("attack failed") + persistence_error = RuntimeError("database unavailable") + + class TestStrategy(AttackStrategy): + def _validate_context(self, *, context: AttackContext) -> None: + pass + + async def _setup_async(self, *, context: AttackContext) -> None: + pass + + async def _perform_async(self, *, context: AttackContext) -> AttackResult: + raise attack_cause + + async def _teardown_async(self, *, context: AttackContext) -> None: + pass + + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) + memory = CentralMemory.get_memory_instance() + + with ( + patch.object(memory, "add_attack_results_to_memory", side_effect=persistence_error) as persist, + pytest.raises(ExceptionGroup, match="Attack execution and error result persistence failed") as exc_info, + ): + await strategy.execute_async(objective="Test objective") + + persist.assert_called_once() + attack_error, recorded_persistence_error = exc_info.value.exceptions + assert attack_error.__cause__ is attack_cause + assert recorded_persistence_error is persistence_error + assert memory.get_attack_results(objective="Test objective") == [] + + async def test_partial_error_result_persistence_failure_does_not_retry( + self, mock_objective_target: PromptTarget + ) -> None: + class TestStrategy(AttackStrategy): + def _validate_context(self, *, context: AttackContext) -> None: + pass + + async def _setup_async(self, *, context: AttackContext) -> None: + pass + + async def _perform_async(self, *, context: AttackContext) -> AttackResult: + raise ValueError("attack failed") + + async def _teardown_async(self, *, context: AttackContext) -> None: + pass + + strategy = TestStrategy(context_type=AttackContext, objective_target=mock_objective_target) + memory = CentralMemory.get_memory_instance() + persist = memory.add_attack_results_to_memory + + def persist_then_fail(*, attack_results: list[AttackResult]) -> None: + persist(attack_results=attack_results) + raise RuntimeError("commit acknowledgement lost") + + with ( + patch.object(memory, "add_attack_results_to_memory", side_effect=persist_then_fail) as persist_mock, + pytest.raises(ExceptionGroup) as exc_info, + ): + await strategy.execute_async(objective="Test objective") + + persist_mock.assert_called_once() + assert str(exc_info.value.exceptions[1]) == "commit acknowledgement lost" + [stored_result] = memory.get_attack_results(objective="Test objective") + assert stored_result.outcome is AttackOutcome.ERROR + assert stored_result.error_message == "attack failed" + async def test_cancellation_clears_retry_collector(self, mock_objective_target): teardown_calls = 0