FIX Preserve lifecycle on result persistence failures

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,

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