From fbabf7e696963953de505938cbef4d595858de24 Mon Sep 17 00:00:00 2001 From: Abhishek Jha Date: Mon, 2 Nov 2026 01:48:02 -0600 Subject: [PATCH 1/2] Add files via upload --- .../backend/services/scenario_run_service.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 0a4cbf1201..6ad80d8966 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1456,13 +1456,35 @@ async def _execute_run_async(self, *, scenario_result_id: str) -> None: except asyncio.CancelledError: try: - await self._memory.try_update_scenario_run_state_async( + # The service owns the reason a task was cancelled: shutdown sets + # cancellation_state to FAILED before cancelling, while a user + # cancel leaves it CANCELLED. Scenario.run_async already persisted + # CANCELLED by the time this handler runs, so the gated write + # below is a no-op for a shutdown-cancelled run: the row is + # already terminal and the compare-and-swap refuses to overwrite + # it. In exactly that case — the stored state is CANCELLED but + # the service's cancellation intent differs — overwrite it with + # the service's state, so a shutdown-interrupted run lands as + # FAILED/ScenarioInterruptedError and stays eligible for + # saved-config resume. Any other terminal state is left alone. + updated = await self._memory.try_update_scenario_run_state_async( scenario_result_id=scenario_result_id, expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, scenario_run_state=active.cancellation_state, error_message=active.cancellation_reason, error_type=active.cancellation_error_type, ) + if not updated and active.cancellation_state != ScenarioRunState.CANCELLED: + header = await self._memory.get_scenario_result_header_async( + scenario_result_id=scenario_result_id + ) + if header is not None and header.scenario_run_state == ScenarioRunState.CANCELLED: + await self._memory.update_scenario_run_state_async( + scenario_result_id=scenario_result_id, + scenario_run_state=active.cancellation_state, + error_message=active.cancellation_reason, + error_type=active.cancellation_error_type, + ) except Exception as exc: handoff_ready = False active.error = str(exc) From e4118a92ad33014bde396a2b6465b40b606dca2c Mon Sep 17 00:00:00 2001 From: Abhishek Jha Date: Mon, 2 Nov 2026 01:49:09 -0600 Subject: [PATCH 2/2] Add files via upload --- .../unit/backend/test_scenario_run_service.py | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 0005eb15f4..056c0c9b50 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -3299,6 +3299,93 @@ def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunSta assert all(call.kwargs["error_type"] == "ScenarioInterruptedError" for call in failure_calls) assert all("shut down" in call.kwargs["error_message"] for call in failure_calls) + async def test_shutdown_overwrites_scenario_cancelled_write_with_interrupted(self, mock_all_registries) -> None: + # Regression test for microsoft/PyRIT#3746: the real Scenario.run_async + # persists CANCELLED when its task is cancelled, before the service's + # cancellation handler runs. A shutdown must still land the active run + # as FAILED/ScenarioInterruptedError (which keeps saved-config resume + # eligible); a gated compare-and-swap write would silently keep + # CANCELLED because the row is already terminal by then. + mock_scenario_registry = mock_all_registries["scenario_registry"] + mock_memory = mock_all_registries["memory"] + records: dict[str, MagicMock] = {} + + async def _create_scenario(*args: object, **kwargs: object) -> MagicMock: + run_id = "shutdown-real-scenario" + records[run_id] = _make_db_scenario_result( + result_id=run_id, + run_state=ScenarioRunState.CREATED, + ) + + async def _run_async() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + # Mirrors Scenario.run_async: persist CANCELLED, then re-raise. + await mock_memory.update_scenario_run_state_async( + scenario_result_id=run_id, + scenario_run_state=ScenarioRunState.CANCELLED, + error_message="Scenario run was cancelled", + error_type="CancelledError", + ) + raise + + scenario = MagicMock() + scenario._scenario_result_id = run_id + scenario.run_async = AsyncMock(side_effect=_run_async) + return scenario + + def _get_results(*, scenario_result_ids: list[str] | None = None) -> list[MagicMock]: + if scenario_result_ids is None: + return list(records.values()) + return [records[run_id] for run_id in scenario_result_ids if run_id in records] + + def _update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunState, **_: object) -> None: + records[scenario_result_id].scenario_run_state = scenario_run_state + + def _try_update_state( + *, + scenario_result_id: str, + expected_states: object, + scenario_run_state: ScenarioRunState, + **_: object, + ) -> bool: + # Faithful to the real compare-and-swap: a row already in a + # terminal state is NOT overwritten. (The pre-existing shutdown + # test mocks this as an unconditional write, which is exactly why + # it could not catch #3746.) + current = records[scenario_result_id].scenario_run_state + if current in expected_states: # type: ignore[operator] + records[scenario_result_id].scenario_run_state = scenario_run_state + return True + return False + + mock_scenario_registry.create_and_initialize_async = AsyncMock(side_effect=_create_scenario) + mock_memory.get_scenario_results_async = AsyncMock(side_effect=_get_results) + mock_memory.get_scenario_result_header_async = AsyncMock( + side_effect=lambda *, scenario_result_id: records.get(scenario_result_id) + ) + mock_memory.update_scenario_run_state_async = AsyncMock(side_effect=_update_state) + mock_memory.try_update_scenario_run_state_async = AsyncMock(side_effect=_try_update_state) + mock_memory.update_scenario_run_state_and_metadata_fields_async = AsyncMock(side_effect=_update_state) + service = ScenarioRunService() + await service.start_run_async(request=_make_request()) + await asyncio.sleep(0) + + await service.shutdown_async() + + assert records["shutdown-real-scenario"].scenario_run_state == ScenarioRunState.FAILED + # The gated try_update is attempted first (and is a no-op once the row + # is CANCELLED); the effective write is the unconditional update below. + failure_writes = [ + call + for call in mock_memory.update_scenario_run_state_async.call_args_list + if call.kwargs.get("scenario_run_state") == ScenarioRunState.FAILED + ] + assert len(failure_writes) == 1 + assert failure_writes[0].kwargs["error_type"] == "ScenarioInterruptedError" + assert "shut down" in failure_writes[0].kwargs["error_message"] + async def test_shutdown_waits_for_preparation_before_terminalizing_run(self, mock_all_registries) -> None: mock_memory = mock_all_registries["memory"] records = {