diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 0df013d11e..c1e260fc12 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, @@ -95,6 +100,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) # Per-execution prepended-history boundary and send lifecycle. Never persisted. prepended_history_send_context: PrependedHistorySendContext | None = field( @@ -240,7 +246,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 @@ -271,7 +277,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( @@ -399,7 +413,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}") @@ -454,13 +471,13 @@ def __init__( history formatting. 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 # Local import avoids the component package's import cycle through attack config. from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, @@ -668,6 +685,34 @@ 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. + + 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 + @overload async def execute_async( self, diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 6d588c447c..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 @@ -470,8 +475,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 +499,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,6 +910,187 @@ 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_completed_result_persistence_failure_is_propagated_after_teardown( + 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() + + 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 + + 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(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 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_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 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