From 66f20e831a290298a5d0f6a14f768411d66c6286 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:15:20 +0530 Subject: [PATCH 1/2] FIX: replace re-emitted progress deltas instead of counting them twice The progress keyset is ordered by (timestamp, id), so an attack result whose timestamp is bumped by an edit is emitted again after the stored cursor. The read model appended every page with extend(), so the edited attempt appeared twice in the snapshot and inflated attempt counts. Track cached deltas by attack_result_id: a re-emitted row replaces its cached copy, the mapped results from that position are rebuilt, and the summary rebuilds when anything was appended or replaced. Fixes #3183 Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../services/scenario_progress_read_model.py | 44 ++++++++++++++++--- .../test_scenario_progress_read_model.py | 33 ++++++++++++++ 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index e4ede355e7..3f41a15559 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -135,7 +135,9 @@ 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 +146,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 +181,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 +196,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 +224,36 @@ 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 every persisted delta after the cached keyset cursor. + + 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: + # Drop the mapped results from the first replaced attempt onward so + # they are remapped from the updated deltas. + del entry.results[replaced_from:] if page: last = page[-1] entry.cursor = AttackResultKeysetCursor( @@ -229,7 +261,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..b4c7df99e9 100644 --- a/tests/unit/backend/test_scenario_progress_read_model.py +++ b/tests/unit/backend/test_scenario_progress_read_model.py @@ -850,3 +850,36 @@ 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_reemitted_delta_instead_of_duplicating() -> None: + """A row re-emitted after a timestamp bump updates its cached attempt, not adds one (#3183).""" + memory = MagicMock(spec=MemoryInterface) + delta = _make_delta(run_id="run-state", index=0) + bumped = ScenarioAttackResultDelta( + attack_result_id=delta.attack_result_id, + conversation_id=delta.conversation_id, + objective=delta.objective, + objective_sha256=delta.objective_sha256, + outcome=AttackOutcome.FAILURE, + execution_time_ms=delta.execution_time_ms, + timestamp=delta.timestamp + timedelta(seconds=99), + attribution_data=delta.attribution_data, + ) + memory.get_scenario_attack_result_deltas_async = AsyncMock(side_effect=[([delta], False), ([bumped], False)]) + read_model = ScenarioProgressReadModel(memory=memory) + + kwargs: dict = dict( + plan=None, + plan_complete=False, + active_group_ids=(), + terminal=False, + objective_scorer_identifier=None, + ) + first = await read_model.get_snapshot_async(scenario_result_id="run-state", **kwargs) + second = await read_model.get_snapshot_async(scenario_result_id="run-state", **kwargs) + + assert len(first.results) == 1 + assert len(second.results) == 1 + assert [result.attack_result_id for result in second.results] == [delta.attack_result_id] + assert second.results[0].outcome == AttackOutcome.FAILURE From 820a2f8f86216bc3d04c90ccb99cb060575085c9 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Tue, 13 Oct 2026 23:27:28 -0700 Subject: [PATCH 2/2] FIX: preserve keyset order when progress deltas are replaced Sort and remap the affected cached suffix after timestamp-bumped replacements so cursor pagination and latest-attempt summaries agree with storage order. Add regression coverage for real SQLite attack edits, page limits, replay, tied timestamps, multiple storage pages, and interrupted refreshes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2643b541-19cf-48c2-b27a-e1a68cc4a6e7 --- .../services/scenario_progress_read_model.py | 14 +- .../test_scenario_progress_read_model.py | 242 ++++++++++++++++-- 2 files changed, 226 insertions(+), 30 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 3f41a15559..529c1fe234 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -135,9 +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) - deltas_updated = 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, @@ -225,7 +223,8 @@ def _build_snapshot( ) async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _ProgressCacheEntry) -> bool: - """Hydrate every persisted delta after the cached keyset cursor. + """ + Hydrate persisted deltas while preserving their timestamp/ID keyset order. Returns: bool: Whether any delta was appended or replaced, so the summary rebuilds. @@ -251,8 +250,11 @@ async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _Pr changed = True replaced_from = position if replaced_from is None else min(replaced_from, position) if replaced_from is not None: - # Drop the mapped results from the first replaced attempt onward so - # they are remapped from the updated deltas. + # 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] diff --git a/tests/unit/backend/test_scenario_progress_read_model.py b/tests/unit/backend/test_scenario_progress_read_model.py index b4c7df99e9..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 @@ -852,34 +857,223 @@ async def test_preparation_failure_stays_separate_from_result_role_async(sqlite_ assert AttackPreparationFailure.from_result(result=stored) == failure -async def test_get_snapshot_replaces_reemitted_delta_instead_of_duplicating() -> None: - """A row re-emitted after a timestamp bump updates its cached attempt, not adds one (#3183).""" +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) - bumped = ScenarioAttackResultDelta( - attack_result_id=delta.attack_result_id, - conversation_id=delta.conversation_id, - objective=delta.objective, - objective_sha256=delta.objective_sha256, - outcome=AttackOutcome.FAILURE, - execution_time_ms=delta.execution_time_ms, - timestamp=delta.timestamp + timedelta(seconds=99), - attribution_data=delta.attribution_data, + 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)] ) - memory.get_scenario_attack_result_deltas_async = AsyncMock(side_effect=[([delta], False), ([bumped], False)]) read_model = ScenarioProgressReadModel(memory=memory) - kwargs: dict = dict( - plan=None, - plan_complete=False, - active_group_ids=(), - terminal=False, - objective_scorer_identifier=None, + 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 ) - first = await read_model.get_snapshot_async(scenario_result_id="run-state", **kwargs) - second = await read_model.get_snapshot_async(scenario_result_id="run-state", **kwargs) - assert len(first.results) == 1 - assert len(second.results) == 1 - assert [result.attack_result_id for result in second.results] == [delta.attack_result_id] - assert second.results[0].outcome == AttackOutcome.FAILURE + +@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()