TEST Add Crescendo mixed-failure resilience coverage

pyrit/executor/core/strategy.py

@@ -344,6 +344,7 @@ async def execute_with_context_async(self, *, context: StrategyContextT) -> Stra
 
         # Execution with lifecycle management
         # This uses an async context manager to ensure setup and teardown are handled correctly
+        retry_collector_started = False
         try:
             async with self._execution_context_async(context):
                 await self._handle_event_async(event=StrategyEvent.ON_PRE_EXECUTE, context=context)
@@ -356,15 +357,14 @@ async def execute_with_context_async(self, *, context: StrategyContextT) -> Stra
                 # handlers can see it.
                 collector = RetryCollector()
                 set_retry_collector(collector)
+                retry_collector_started = True
 
                 result = await self._perform_async(context=context)
                 await self._handle_event_async(event=StrategyEvent.ON_POST_EXECUTE, context=context, result=result)
-                clear_retry_collector()
                 return result
         except Exception as e:
             # Notify error event
             await self._handle_event_async(event=StrategyEvent.ON_ERROR, context=context, error=e)
-            clear_retry_collector()
 
             # Build enhanced error message with execution context if available
             # Note: The context is preserved on exception by ExecutionContextManager
@@ -394,6 +394,9 @@ async def execute_with_context_async(self, *, context: StrategyContextT) -> Stra
 
             runtime_error = _StrategyRuntimeError(error_message)
             raise runtime_error from e
+        finally:
+            if retry_collector_started:
+                clear_retry_collector()
 
     async def execute_async(self, **kwargs: Any) -> StrategyResultT:
         """

tests/unit/executor/attack/core/test_attack_strategy.py

@@ -1,13 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import asyncio
 import logging
 from dataclasses import replace
 from unittest.mock import MagicMock, patch
 
 import pytest
 
-from pyrit.exceptions.retry_collector import RetryCollector
+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
 from pyrit.executor.attack.core.attack_strategy import (
@@ -859,6 +860,31 @@ 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_cancellation_clears_retry_collector(self, mock_objective_target):
+        teardown_calls = 0
+
+        class CancelledStrategy(AttackStrategy):
+            def _validate_context(self, *, context):
+                pass
+
+            async def _setup_async(self, *, context):
+                pass
+
+            async def _perform_async(self, *, context):
+                raise asyncio.CancelledError
+
+            async def _teardown_async(self, *, context):
+                nonlocal teardown_calls
+                teardown_calls += 1
+
+        strategy = CancelledStrategy(context_type=AttackContext, objective_target=mock_objective_target)
+
+        with pytest.raises(asyncio.CancelledError):
+            await strategy.execute_async(objective="Test objective")
+
+        assert teardown_calls == 1
+        assert get_retry_collector() is None
+
     async def test_attack_strategy_with_custom_event_handler(self, mock_objective_target):
         """Test that AttackStrategy can work with custom event handlers"""
         custom_handler_called = False