diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index e4ede355e7..529c1fe234 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -135,7 +135,7 @@ async def get_snapshot_async( async with self._cache_lock: entry = self._get_cache_entry(scenario_result_id=scenario_result_id, plan=plan, terminal=terminal) first_new_index = len(entry.results) - await self._hydrate_new_deltas_async(scenario_result_id=scenario_result_id, entry=entry) + deltas_updated = await self._hydrate_new_deltas_async(scenario_result_id=scenario_result_id, entry=entry) return self._build_snapshot( entry=entry, first_new_index=first_new_index, @@ -144,6 +144,7 @@ async def get_snapshot_async( active_group_ids=active_group_ids, terminal=terminal, objective_scorer_identifier=objective_scorer_identifier, + deltas_updated=deltas_updated, ) def _get_cache_entry( @@ -178,6 +179,7 @@ def _build_snapshot( active_group_ids: Sequence[str], terminal: bool, objective_scorer_identifier: ComponentIdentifier | None, + deltas_updated: bool = False, ) -> ScenarioProgressSnapshot: summary_plan = plan or self._synthesize_legacy_plan(deltas=entry.deltas) if len(entry.results) < len(entry.deltas): @@ -192,7 +194,12 @@ def _build_snapshot( terminal=terminal, plan_complete=plan_complete, ) - if entry.summary is None or first_new_index < len(entry.deltas) or entry.summary_state != summary_state: + if ( + entry.summary is None + or deltas_updated + or first_new_index < len(entry.deltas) + or entry.summary_state != summary_state + ): technique_details_by_group = self._build_technique_details_by_group( deltas=entry.deltas, results=entry.results, @@ -215,13 +222,40 @@ def _build_snapshot( plan=summary_plan, ) - async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _ProgressCacheEntry) -> None: - """Hydrate every persisted delta after the cached keyset cursor.""" + async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _ProgressCacheEntry) -> bool: + """ + Hydrate persisted deltas while preserving their timestamp/ID keyset order. + + Returns: + bool: Whether any delta was appended or replaced, so the summary rebuilds. + """ + changed = False while True: page, has_more = await self._memory.get_scenario_attack_result_deltas_async( scenario_result_id=scenario_result_id, cursor=entry.cursor, limit=self._STORAGE_PAGE_SIZE ) - entry.deltas.extend(page) + # The keyset is ordered by (timestamp, id), so an edited row whose timestamp + # moved forward is emitted again after the cursor. Such a re-emission is an + # update of an attempt already cached, never a second attempt (#3183). + positions = {delta.attack_result_id: index for index, delta in enumerate(entry.deltas)} + replaced_from: int | None = None + for delta in page: + position = positions.get(delta.attack_result_id) + if position is None: + positions[delta.attack_result_id] = len(entry.deltas) + entry.deltas.append(delta) + changed = True + else: + entry.deltas[position] = delta + changed = True + replaced_from = position if replaced_from is None else min(replaced_from, position) + if replaced_from is not None: + # Timestamp bumps move attempts forward; paging and latest-attempt + # selection need the affected suffix in keyset order. + entry.deltas[replaced_from:] = sorted( + entry.deltas[replaced_from:], key=lambda delta: (delta.timestamp, delta.attack_result_id) + ) + del entry.results[replaced_from:] if page: last = page[-1] entry.cursor = AttackResultKeysetCursor( @@ -229,7 +263,7 @@ async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _Pr attack_result_id=last.attack_result_id, ) if not has_more: - return + return changed if not page: raise RuntimeError("Scenario progress storage returned an empty page with has_more=True.") diff --git a/tests/unit/backend/test_scenario_progress_read_model.py b/tests/unit/backend/test_scenario_progress_read_model.py index bd073aec58..01988e402b 100644 --- a/tests/unit/backend/test_scenario_progress_read_model.py +++ b/tests/unit/backend/test_scenario_progress_read_model.py @@ -12,11 +12,15 @@ import pytest +from pyrit.analytics.scenario_statistics import compute_scenario_statistics +from pyrit.backend.models.attacks import UpdateAttackRequest +from pyrit.backend.services.attack_service import AttackService from pyrit.backend.services.scenario_progress_read_model import ( ScenarioPlanLookup, ScenarioProgressReadModel, ScenarioProgressSnapshot, ) +from pyrit.backend.services.scenario_run_service import ScenarioRunService from pyrit.exceptions import ScenarioPartialFailureException from pyrit.executor.attack import ( AttackScoringConfig, @@ -41,6 +45,7 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanGroupKind, ScenarioRunPlanSeedGroup, + ScenarioRunState, SeedObjective, ) from pyrit.prompt_target import PromptTarget @@ -850,3 +855,225 @@ async def test_preparation_failure_stays_separate_from_result_role_async(sqlite_ assert mapped.outcome is AttackOutcome.UNDETERMINED assert mapped.child_attack_result_ids == [] assert AttackPreparationFailure.from_result(result=stored) == failure + + +async def test_get_snapshot_replaces_and_reorders_reemitted_delta_async() -> None: + memory = MagicMock(spec=MemoryInterface) + before = _make_delta(run_id="run-state", index=-1) + delta = _make_delta(run_id="run-state", index=0) + after = _make_delta(run_id="run-state", index=1) + bumped = delta.model_copy( + update={"outcome": AttackOutcome.FAILURE, "timestamp": delta.timestamp + timedelta(seconds=99)} + ) + memory.get_scenario_attack_result_deltas_async = AsyncMock( + side_effect=[([before, delta, after], False), ([bumped], False), ([], False)] + ) + read_model = ScenarioProgressReadModel(memory=memory) + + first = await _get_snapshot_async(read_model=read_model, run_id="run-state") + second = await _get_snapshot_async(read_model=read_model, run_id="run-state") + unchanged = await _get_snapshot_async(read_model=read_model, run_id="run-state") + + assert first.deltas == (before, delta, after) + assert first.results[1].outcome == AttackOutcome.SUCCESS + assert first.summary.overall.succeeded == 3 + assert second.deltas == (before, after, bumped) + assert [result.attack_result_id for result in second.results] == [ + before.attack_result_id, + after.attack_result_id, + delta.attack_result_id, + ] + assert second.results[0] is first.results[0] + assert second.results[-1].outcome == AttackOutcome.FAILURE + assert second.summary.overall.succeeded == 2 + assert second.summary.overall.retries == 0 + assert unchanged.results == second.results + assert unchanged.summary is second.summary + + +async def test_get_snapshot_orders_replacements_and_new_rows_across_pages_async() -> None: + memory = MagicMock(spec=MemoryInterface) + deltas = [ + _make_delta(run_id="run-pages", index=index).model_copy( + update={"attack_result_id": str(uuid.UUID(int=result_id))} + ) + for index, result_id in enumerate((1, 3, 4)) + ] + timestamp = deltas[-1].timestamp + timedelta(seconds=1) + first_bumped, second_bumped = [delta.model_copy(update={"timestamp": timestamp}) for delta in deltas[:2]] + appended = first_bumped.model_copy(update={"attack_result_id": str(uuid.UUID(int=2))}) + memory.get_scenario_attack_result_deltas_async = AsyncMock( + side_effect=[ + (deltas, False), + ([first_bumped, appended], True), + ([second_bumped], False), + ] + ) + read_model = ScenarioProgressReadModel(memory=memory) + await _get_snapshot_async(read_model=read_model, run_id="run-pages") + snapshot = await _get_snapshot_async(read_model=read_model, run_id="run-pages") + + assert snapshot.deltas == (deltas[2], first_bumped, appended, second_bumped) + assert [result.attack_result_id for result in snapshot.results] == [ + delta.attack_result_id for delta in snapshot.deltas + ] + assert memory.get_scenario_attack_result_deltas_async.call_args.kwargs["cursor"] == AttackResultKeysetCursor( + timestamp=timestamp, attack_result_id=appended.attack_result_id + ) + + +@pytest.mark.parametrize("error_type", [RuntimeError, asyncio.CancelledError]) +async def test_get_snapshot_recovers_reordered_replacement_after_interrupted_refresh_async( + error_type: type[BaseException], +) -> None: + memory = MagicMock(spec=MemoryInterface) + first = _make_delta(run_id="run-interrupted", index=0) + second = _make_delta(run_id="run-interrupted", index=1) + bumped = first.model_copy( + update={"timestamp": second.timestamp + timedelta(seconds=1), "outcome": AttackOutcome.FAILURE} + ) + memory.get_scenario_attack_result_deltas_async = AsyncMock( + side_effect=[([first, second], False), ([bumped], True), error_type(), ([], False)] + ) + read_model = ScenarioProgressReadModel(memory=memory) + initial = await _get_snapshot_async(read_model=read_model, run_id="run-interrupted") + + with pytest.raises(error_type): + await _get_snapshot_async(read_model=read_model, run_id="run-interrupted") + assert not read_model._cache_lock.locked() + refreshed = await _get_snapshot_async(read_model=read_model, run_id="run-interrupted") + + assert initial.deltas == (first, second) + assert initial.summary.overall.succeeded == 2 + assert refreshed.deltas == (second, bumped) + assert [result.outcome for result in refreshed.results] == [AttackOutcome.SUCCESS, AttackOutcome.FAILURE] + assert refreshed.summary.overall.succeeded == 1 + assert memory.get_scenario_attack_result_deltas_async.call_args.kwargs["cursor"] == AttackResultKeysetCursor( + timestamp=bumped.timestamp, attack_result_id=bumped.attack_result_id + ) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("limit", [1, 500]) +async def test_progress_pages_keep_all_results_after_real_attack_edit_async( + *, sqlite_instance: MemoryInterface, limit: int +) -> None: + scenario = make_scenario_result( + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + scenario_run_state=ScenarioRunState.COMPLETED, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario]) + rows = [ + AttackResult( + attack_result_id=str(uuid.UUID(int=index + 1)), + conversation_id=str(uuid.uuid4()), + objective=f"describe everyday object {index}", + outcome=AttackOutcome.FAILURE, + timestamp=datetime(2025, 1, 1, tzinfo=UTC) + timedelta(seconds=index), + attribution_parent_id=str(scenario.id), + attribution_data={"parent_collection": "attack", "seed_group_id": f"seed-{index}"}, + ) + for index in range(limit + 1) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=rows) + service = ScenarioRunService() + try: + initial = await service.get_run_progress_async(scenario_result_id=str(scenario.id), since=None, limit=500) + assert initial is not None + initial_payload = initial.model_dump(mode="json") + edited = rows[limit - 1] + updated = await AttackService().update_attack_async( + attack_result_id=edited.attack_result_id, request=UpdateAttackRequest(outcome="success") + ) + assert updated is not None + + first_page = await service.get_run_progress_async(scenario_result_id=str(scenario.id), since=None, limit=limit) + assert first_page is not None + expected_ids = [row.attack_result_id for row in rows if row is not edited] + [edited.attack_result_id] + assert [result.attack_result_id for result in first_page.results] == expected_ids[:limit] + assert first_page.has_more + last_page = await service.get_run_progress_async( + scenario_result_id=str(scenario.id), since=first_page.next_cursor, limit=limit + ) + assert last_page is not None + assert [result.attack_result_id for result in last_page.results] == expected_ids[limit:] + assert last_page.results[0].outcome == AttackOutcome.SUCCESS + assert not last_page.has_more + assert last_page.summary.overall.completed == limit + 1 + assert last_page.summary.overall.succeeded == 1 + assert last_page.summary.overall.retries == 0 + assert last_page.summary.overall.producer_counts.unknown.attempts == limit + 1 + + replay = await service.get_run_progress_async( + scenario_result_id=str(scenario.id), since=first_page.next_cursor, limit=limit + ) + assert replay is not None + assert replay.results == last_page.results + assert replay.next_cursor == last_page.next_cursor + empty = await service.get_run_progress_async( + scenario_result_id=str(scenario.id), since=last_page.next_cursor, limit=limit + ) + assert empty is not None + assert empty.results == [] + assert not empty.has_more + assert empty.next_cursor == last_page.next_cursor + assert initial.model_dump(mode="json") == initial_payload + finally: + await service.close_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_progress_latest_edited_attempt_matches_sdk_and_cold_cache_async( + sqlite_instance: MemoryInterface, +) -> None: + scenario = make_scenario_result( + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + scenario_run_state=ScenarioRunState.COMPLETED, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scenario]) + rows = [ + AttackResult( + conversation_id=str(uuid.uuid4()), + objective="describe an everyday object", + outcome=AttackOutcome.FAILURE, + timestamp=datetime(2025, 1, 1, tzinfo=UTC) + timedelta(seconds=index), + attribution_parent_id=str(scenario.id), + attribution_data={"parent_collection": "attack", "seed_group_id": "seed"}, + ) + for index in range(2) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=rows) + warm_service = ScenarioRunService() + cold_service = ScenarioRunService() + try: + initial = await warm_service.get_run_progress_async(scenario_result_id=str(scenario.id), since=None, limit=500) + assert initial is not None + assert initial.summary.overall.succeeded == 0 + updated = await AttackService().update_attack_async( + attack_result_id=rows[0].attack_result_id, request=UpdateAttackRequest(outcome="success") + ) + assert updated is not None + warm = await warm_service.get_run_progress_async( + scenario_result_id=str(scenario.id), since=initial.next_cursor, limit=500 + ) + cold = await cold_service.get_run_progress_async(scenario_result_id=str(scenario.id), since=None, limit=500) + [stored] = await sqlite_instance.get_scenario_results_async(scenario_result_ids=[str(scenario.id)]) + sdk = compute_scenario_statistics(stored) + + assert warm is not None and cold is not None + assert [result.attack_result_id for result in warm.results] == [rows[0].attack_result_id] + assert [result.attack_result_id for result in cold.results] == [ + rows[1].attack_result_id, + rows[0].attack_result_id, + ] + assert warm.summary == cold.summary + assert warm.summary.overall == sdk.overall + assert warm.summary.overall.completed == 1 + assert warm.summary.overall.succeeded == 1 + assert warm.summary.overall.retries == 1 + assert warm.summary.overall.producer_counts.unknown.attempts == 2 + finally: + await warm_service.close_async() + await cold_service.close_async()