From 91ea4ed883bac34cada7c81f4e0cfb25ca5cad61 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:47:29 -0700 Subject: [PATCH 1/2] FIX Persist direct scenario cancellation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 058569a0-6218-4cf4-947e-67fa51001d85 --- pyrit/scenario/core/scenario.py | 10 ++ .../core/test_scenario_partial_results.py | 115 ++++++++++++++++++ 2 files changed, 125 insertions(+) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index b607c1dee0..80101a7bdc 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -1018,6 +1018,8 @@ async def run_async(self) -> ScenarioResult: attack results from all atomic attacks. Raises: + asyncio.CancelledError: If the scenario task is cancelled. Completed results remain persisted + and a later call can resume the unfinished objectives. ValueError: If the scenario has no atomic attacks configured. If your scenario requires initialization, call await scenario.initialize() first. ScenarioPartialFailureException: If an atomic attack only partially completes. @@ -1048,6 +1050,14 @@ async def run_async(self) -> ScenarioResult: for retry_attempt in range(self._max_retries + 1): # +1 for initial attempt try: return await self._execute_scenario_async() + except asyncio.CancelledError: + self._memory.update_scenario_run_state( + scenario_result_id=scenario_result_id, + scenario_run_state=ScenarioRunState.CANCELLED, + error_message="Scenario run was cancelled", + error_type="CancelledError", + ) + raise except Exception as e: last_exception = e diff --git a/tests/unit/scenario/core/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py index e7181a89f6..2ad57f2dc2 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -3,6 +3,7 @@ """Additional tests for Scenario retry with AttackExecutorResult functionality.""" +import asyncio from typing import ClassVar from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch @@ -404,6 +405,120 @@ async def mock_run(*args, **kwargs): # All 5 results should be in final scenario result assert len(result.attack_results["resume_attack"]) == 5 + async def test_run_async_cancellation_persists_progress_cleans_workers_and_resumes(self, mock_objective_target): + completed_attack = create_mock_atomic_attack("completed_attack", ["obj1"]) + in_flight_attack = create_mock_atomic_attack("in_flight_attack", ["obj2"]) + queued_attack = create_mock_atomic_attack("queued_attack", ["obj3"]) + + completed_result = AttackResult( + conversation_id="conv-1", + objective="obj1", + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ) + resumed_results = { + "in_flight_attack": AttackResult( + conversation_id="conv-2", + objective="obj2", + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ), + "queued_attack": AttackResult( + conversation_id="conv-3", + objective="obj3", + outcome=AttackOutcome.SUCCESS, + executed_turns=1, + ), + } + + completed_persisted = asyncio.Event() + in_flight_started = asyncio.Event() + completed_worker_exited = asyncio.Event() + in_flight_worker_exited = asyncio.Event() + block_until_cancelled = asyncio.Event() + persisted_objectives: list[str] = [] + + async def run_completed_attack(*args, **kwargs): + save_attack_results_to_memory([completed_result], atomic_attack=completed_attack) + persisted_objectives.append(completed_result.objective) + completed_persisted.set() + try: + await block_until_cancelled.wait() + finally: + completed_worker_exited.set() + + async def run_in_flight_attack(*args, **kwargs): + if in_flight_attack.run_async.call_count == 1: + in_flight_started.set() + try: + await block_until_cancelled.wait() + finally: + in_flight_worker_exited.set() + + result = resumed_results["in_flight_attack"] + save_attack_results_to_memory([result], atomic_attack=in_flight_attack) + persisted_objectives.append(result.objective) + return AttackExecutorResult(completed_results=[result], incomplete_objectives=[]) + + async def run_queued_attack(*args, **kwargs): + result = resumed_results["queued_attack"] + save_attack_results_to_memory([result], atomic_attack=queued_attack) + persisted_objectives.append(result.objective) + return AttackExecutorResult(completed_results=[result], incomplete_objectives=[]) + + completed_attack.run_async = AsyncMock(side_effect=run_completed_attack) + in_flight_attack.run_async = AsyncMock(side_effect=run_in_flight_attack) + queued_attack.run_async = AsyncMock(side_effect=run_queued_attack) + + scenario = ConcreteScenario( + name="Cancellation Test Scenario", + version=1, + atomic_attacks_to_return=[completed_attack, in_flight_attack, queued_attack], + ) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "max_concurrency": 2, + "max_retries": 3, + } + ) + await scenario.initialize_async() + + scenario_task = asyncio.create_task(scenario.run_async()) + await asyncio.wait_for(completed_persisted.wait(), timeout=5.0) + await asyncio.wait_for(in_flight_started.wait(), timeout=5.0) + scenario_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await scenario_task + + assert completed_worker_exited.is_set() + assert in_flight_worker_exited.is_set() + queued_attack.run_async.assert_not_called() + assert persisted_objectives == ["obj1"] + + [cancelled_result] = CentralMemory.get_memory_instance().get_scenario_results( + scenario_result_ids=[scenario._scenario_result_id] + ) + assert cancelled_result.scenario_run_state == ScenarioRunState.CANCELLED + assert cancelled_result.error_type == "CancelledError" + assert cancelled_result.number_tries == 1 + assert [result.objective for result in cancelled_result.attack_results["completed_attack"]] == ["obj1"] + + await asyncio.sleep(0) + assert persisted_objectives == ["obj1"] + + resumed_result = await scenario.run_async() + + assert resumed_result.scenario_run_state == ScenarioRunState.COMPLETED + assert resumed_result.number_tries == 2 + assert completed_attack.run_async.call_count == 1 + assert in_flight_attack.run_async.call_count == 2 + assert queued_attack.run_async.call_count == 1 + assert persisted_objectives == ["obj1", "obj2", "obj3"] + assert sorted(resumed_result.get_objectives()) == ["obj1", "obj2", "obj3"] + assert all(len(results) == 1 for results in resumed_result.attack_results.values()) + async def test_multiple_atomic_attacks_with_partial_results(self, mock_objective_target): """Test scenario with multiple atomic attacks that return partial results.""" # Create 3 atomic attacks From 65cd2df4c63cff2d793644590c0f43aa550ba2d3 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:50:36 -0700 Subject: [PATCH 2/2] FIX Preserve cancellation on persistence failure Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 058569a0-6218-4cf4-947e-67fa51001d85 --- pyrit/scenario/core/scenario.py | 15 ++++++---- .../core/test_scenario_partial_results.py | 28 +++++++++++++++++++ 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 80101a7bdc..5ca3c2247c 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -1051,12 +1051,15 @@ async def run_async(self) -> ScenarioResult: try: return await self._execute_scenario_async() except asyncio.CancelledError: - self._memory.update_scenario_run_state( - scenario_result_id=scenario_result_id, - scenario_run_state=ScenarioRunState.CANCELLED, - error_message="Scenario run was cancelled", - error_type="CancelledError", - ) + try: + self._memory.update_scenario_run_state( + scenario_result_id=scenario_result_id, + scenario_run_state=ScenarioRunState.CANCELLED, + error_message="Scenario run was cancelled", + error_type="CancelledError", + ) + except Exception: + logger.exception(f"Failed to persist cancellation state for scenario '{self._name}'") raise except Exception as e: last_exception = e diff --git a/tests/unit/scenario/core/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py index 2ad57f2dc2..74bde8b988 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -519,6 +519,34 @@ async def run_queued_attack(*args, **kwargs): assert sorted(resumed_result.get_objectives()) == ["obj1", "obj2", "obj3"] assert all(len(results) == 1 for results in resumed_result.attack_results.values()) + async def test_run_async_cancellation_is_not_masked_by_persistence_failure( + self, mock_objective_target: MagicMock + ) -> None: + atomic_attack = create_mock_atomic_attack("cancelled_attack", ["obj1"]) + scenario = ConcreteScenario( + name="Cancellation Persistence Failure Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.initialize_async() + + with ( + patch.object( + scenario, + "_execute_scenario_async", + new_callable=AsyncMock, + side_effect=asyncio.CancelledError, + ), + patch.object( + scenario._memory, + "update_scenario_run_state", + side_effect=RuntimeError("database unavailable"), + ), + ): + with pytest.raises(asyncio.CancelledError): + await scenario.run_async() + async def test_multiple_atomic_attacks_with_partial_results(self, mock_objective_target): """Test scenario with multiple atomic attacks that return partial results.""" # Create 3 atomic attacks