From 368b137f5100d98b99b2795efc6d4b17ad8c4cec Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 6 Nov 2026 17:50:57 -0800 Subject: [PATCH 1/3] FEAT: Backend APIs for scan results and running scan attacks - GET /api/attacks/{id}/findings lists findings linked to an attack across Operations, backed by a new FindingEvidenceEntries.attack_result_id index (Alembic b5d7f9a1c3e6). - Attack read APIs (summary, metadata, conversations, messages) resolve scan attacks that are still running through a read-only LiveAttackLookup provided by ScenarioRunService. Saved results always win; summaries/metadata expose in_progress and scenario_result_id. - ScenarioActiveConversation.conversation_ids lists every objective conversation so far. - ScenarioProgressScore carries score_category for progress results and hierarchy nodes. - Frontend types and attacksApi.getAttackFindings for later layers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- frontend/src/services/api.test.ts | 24 +- frontend/src/services/api.ts | 7 + frontend/src/types/index.ts | 26 +- pyrit/backend/README.md | 34 ++ pyrit/backend/mappers/attack_mappers.py | 4 + pyrit/backend/models/attacks.py | 8 + pyrit/backend/models/operations.py | 16 + pyrit/backend/routes/attacks.py | 16 + pyrit/backend/services/attack_service.py | 110 +++++-- pyrit/backend/services/live_attack_reader.py | 94 ++++++ pyrit/backend/services/operation_service.py | 26 ++ .../saved_conversation_read_service.py | 40 ++- .../backend/services/scenario_run_service.py | 30 ++ ...e6_index_finding_evidence_attack_result.py | 31 ++ pyrit/memory/memory_interface.py | 111 +++++-- pyrit/memory/memory_models.py | 1 + pyrit/models/scenario_progress.py | 4 + tests/unit/backend/test_live_attack_reads.py | 292 ++++++++++++++++++ tests/unit/backend/test_scenario_resume.py | 108 +++++++ .../test_interface_scenario_progress.py | 3 + ...finding_evidence_attack_index_migration.py | 77 +++++ tests/unit/memory/test_findings.py | 34 ++ .../test_score_participation_migration.py | 2 +- 23 files changed, 1031 insertions(+), 67 deletions(-) create mode 100644 pyrit/backend/services/live_attack_reader.py create mode 100644 pyrit/memory/alembic/versions/b5d7f9a1c3e6_index_finding_evidence_attack_result.py create mode 100644 tests/unit/backend/test_live_attack_reads.py create mode 100644 tests/unit/memory/test_finding_evidence_attack_index_migration.py diff --git a/frontend/src/services/api.test.ts b/frontend/src/services/api.test.ts index 0da6db1553..8e4b98b374 100644 --- a/frontend/src/services/api.test.ts +++ b/frontend/src/services/api.test.ts @@ -42,7 +42,7 @@ import { import { compatibility, COMPATIBILITY_HEADER } from "./compatibility"; import { toApiError } from "./errors"; import { makeAnalyticsFacets, makeAnalyticsReport, makeAnalyticsResults } from "@/test-utils/analyticsFixtures"; -import type { FindingExportRequest } from "@/types"; +import type { AttackFindingListResponse, FindingExportRequest } from "@/types"; import { DEFAULT_ANALYTICS_VIEW } from "@/utils/attackAnalytics"; describe("api service", () => { @@ -823,6 +823,28 @@ describe("api service", () => { }); }); + it("should get the findings linked to an attack", async () => { + const data: AttackFindingListResponse = { + attack_result_id: "attack/id", + items: [{ + id: "finding", operation_id: "operation", operation_name: "Red team", evidence_id: "evidence", + title: "Protected system instructions disclosed", description: "", severity: "important", + created_at: "2026-11-06T00:00:00Z", + }], + }; + (apiClient.get as jest.Mock).mockResolvedValueOnce({ data }); + + expect(await attacksApi.getAttackFindings("attack/id")).toEqual(data); + expect(apiClient.get).toHaveBeenCalledTimes(1); + expect(apiClient.get).toHaveBeenCalledWith("/attacks/attack%2Fid/findings"); + }); + + it("should propagate attack finding errors", async () => { + const error = new Error("Findings unavailable"); + (apiClient.get as jest.Mock).mockRejectedValueOnce(error); + await expect(attacksApi.getAttackFindings("attack")).rejects.toBe(error); + }); + it("should get attack messages", async () => { const mockResponse = { data: { diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 95f9725efc..795c5b794a 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -7,6 +7,7 @@ import { getGraphScopes } from '../auth/msalConfig' import type { ApplyChatScorerRequest, ApplyChatScorerResponse, + AttackFindingListResponse, TargetInstance, UnregisteredTarget, UnregisteredConverter, @@ -502,6 +503,12 @@ export const attacksApi = { return response.data }, + /** Findings in any Operation that link this attack as evidence; empty for an unsaved running attack. */ + getAttackFindings: async (attackResultId: string): Promise => { + const response = await apiClient.get(`/attacks/${encodeURIComponent(attackResultId)}/findings`) + return response.data + }, + updateAttack: async (attackResultId: string, request: UpdateAttackRequest): Promise => { const response = await apiClient.patch(`/attacks/${encodeURIComponent(attackResultId)}`, request) return response.data diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 2a151f454b..62a9a146fa 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -987,6 +987,10 @@ export interface AttackSummary { chat_scorer?: RegisteredComponentReference | null chat_scorer_inherited?: boolean execution_read_only?: boolean + /** The attack is still running in a scan; its result is not saved yet. */ + in_progress?: boolean + /** The scan run that produced this attack; null outside a scan. */ + scenario_result_id?: string | null manual_continuation?: ManualContinuation | null last_score?: BackendScore | null last_response?: BackendMessagePiece | null @@ -1012,7 +1016,7 @@ export interface AttackSummary { export interface AttackMetadata extends Pick { readonly execution_read_only: boolean readonly last_response_message_piece_id: string | null @@ -2082,13 +2086,14 @@ export interface ScenarioQueueSnapshot { queued: ScenarioQueueEntry[] } -/** One persisted attack attempt in ascending progress order. */ +/** The objective score of one persisted scan attack attempt. */ export interface ScenarioProgressScore { scorer_name: string score_type: 'true_false' | 'float_scale' | 'unknown' status: 'complete' | 'undetermined' score_value?: string | null score_rationale?: string | null + score_category?: string[] } export type ScenarioIdentityValue = @@ -2338,7 +2343,10 @@ export interface ScenarioProgressSummary { export interface ScenarioActiveConversation { attack_result_id: string + /** The conversation the attack is working on, or the recorded objective conversation. */ conversation_id: string + /** Every objective conversation the running attack has started so far, in first-use order. */ + conversation_ids?: string[] atomic_attack_name: string seed_group_id?: string | null objective: string @@ -2446,6 +2454,20 @@ export interface FindingListItem extends Finding { evidence_count: number } +/** A finding that links an attack as evidence. */ +export interface AttackFindingItem extends Finding { + /** Current name of the finding's Operation. */ + operation_name: string + /** The evidence link; detach it through the Operation evidence route. */ + evidence_id: string +} + +/** Every finding that links an attack, highest severity first, then newest first. */ +export interface AttackFindingListResponse { + attack_result_id: string + items: AttackFindingItem[] +} + export interface FindingEvidenceCreateRequest { attack_result_id: string } diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index ced275448d..a7dbf409de 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -112,6 +112,40 @@ Both paged and unpaged reads expose only main and pruned conversations. Adversarial and preparation contexts are excluded from lists, reference totals, and transcript reads, including hierarchy node expansion. +### Running scan attacks + +A scan saves each `AttackResult` only when its attack finishes. While an active scan +is still running an attack, `GET /api/attacks/{attack_result_id}`, `/metadata`, +`/conversations`, and `/messages` (paged or not) resolve it from the live run instead +of returning 404. These responses set `in_progress: true`, report outcome +`undetermined`, and set `execution_read_only: true`; write routes still require a +saved result. The conversation the attack is working on is the main conversation, +and every other objective conversation started so far is listed as a pruned +reference. Messages are read from memory as they are saved. A saved result always +wins, so once the attack finishes the same IDs resolve to its saved result; restart +reference paging if a cursor becomes stale at that transition. An attack that has +not started an objective conversation yet still returns 404. + +Attack summaries and metadata also include `scenario_result_id`, the scan run that +produced the attack (`null` for attacks outside a scan). Scan progress +`active_conversations` entries list every objective conversation so far in +`conversation_ids`, alongside the current `conversation_id`. Each scan progress +result and hierarchy node `score` summarizes the objective score: `scorer_name`, +`score_type`, `status`, `score_value`, `score_rationale`, and `score_category`. + +The attack service reads running attacks through a read-only lookup +(`LiveAttackLookup`) that the scenario run service provides; it does not import +scenario scheduling. + +### Findings linked to an attack + +`GET /api/attacks/{attack_result_id}/findings` lists every finding, across all +Operations, that links the attack as evidence: highest severity first, then +newest. Each item is the finding plus `operation_name` and the `evidence_id` that +unlinks it through the Operation evidence route. An attack with no linked +findings, including one that is not saved yet, returns an empty list; a non-UUID +ID returns 422. + The saved Scenario Chat viewer consumes these bounded reads with independent reference/message continuation controls. It merges stable piece IDs into their ordered turns and reconciles page-local partial-sequence markers against the diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index cad7701e40..4a539066f3 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -232,6 +232,7 @@ async def attack_result_to_summary_async( operation_name: str | None = None, scoring_inputs: AttackScoringInputs | None = None, execution_read_only: bool = False, + in_progress: bool = False, ) -> AttackSummary: """ Build an AttackSummary view from an AttackResult. @@ -248,6 +249,7 @@ async def attack_result_to_summary_async( Legacy attacks without an Operation ID show their historical name instead. scoring_inputs: Effective inputs resolved against the current registry by the service. execution_read_only: Whether execution changes require a manual continuation. + in_progress: Whether ``ar`` projects an attack still running in a scan. Returns: AttackSummary view ready for the API response. @@ -264,6 +266,8 @@ async def attack_result_to_summary_async( chat_scorer=scoring_inputs.scorer if scoring_inputs else None, chat_scorer_inherited=scoring_inputs.inherited if scoring_inputs else False, execution_read_only=execution_read_only, + in_progress=in_progress, + scenario_result_id=ar.attribution_parent_id, last_response=await _summary_last_response_async(ar.last_response), automated_score=( ScoreView.from_domain(ar.automated_score, is_objective_score=True) if ar.automated_score else None diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index 63eb8ab0a9..0692643a94 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -262,6 +262,10 @@ class AttackSummary(AttackResult[MessagePieceView, ScoreView]): ) chat_scorer_inherited: bool = Field(default=False, description="Uses the scan's registered original scorer") execution_read_only: bool = Field(default=False, description="Execution changes require a manual continuation") + in_progress: bool = Field( + default=False, description="The attack is still running in a scan; its result is not saved yet" + ) + scenario_result_id: str | None = Field(default=None, description="The scan run that produced this attack") @computed_field @property @@ -360,6 +364,10 @@ class AttackMetadata(BaseModel): chat_scorer: RegisteredComponentReference | None = None chat_scorer_inherited: bool = False execution_read_only: bool = False + in_progress: bool = Field( + default=False, description="The attack is still running in a scan; its result is not saved yet" + ) + scenario_result_id: str | None = Field(default=None, description="The scan run that produced this attack") manual_continuation: ManualContinuation | None = None last_response_message_piece_id: str | None = None related_conversation_count: int = Field(..., ge=0, description="Number of pruned conversations, excluding main") diff --git a/pyrit/backend/models/operations.py b/pyrit/backend/models/operations.py index 723d5b0468..590689f05a 100644 --- a/pyrit/backend/models/operations.py +++ b/pyrit/backend/models/operations.py @@ -70,6 +70,22 @@ class FindingListItem(Finding): evidence_count: int = Field(ge=0) +class AttackFindingItem(Finding): + """A finding that links an attack as evidence, with its Operation's current name and the evidence link.""" + + operation_name: str + evidence_id: UUID = Field(description="The evidence link, which unlinks the attack from the finding") + + +class AttackFindingListResponse(BaseModel): + """Every finding that links an attack, highest severity first, then newest first.""" + + model_config = ConfigDict(extra="forbid") + + attack_result_id: UUID + items: list[AttackFindingItem] + + class FindingListResponse(BaseModel): """One page of an operation's findings.""" diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index 70e38b65a7..73bfd66f81 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -42,6 +42,7 @@ LabelFilterStr, ProblemDetail, ) +from pyrit.backend.models.operations import AttackFindingListResponse from pyrit.backend.models.scores import ApplyChatScorerRequest, ApplyChatScorerResponse from pyrit.backend.routes.common import parse_label_query_params from pyrit.backend.services.attack_service import AttackObjectiveConflictError, get_attack_service @@ -49,6 +50,7 @@ from pyrit.backend.services.history_replay import HistoryReplayError from pyrit.backend.services.manual_send_scheduler import ManualSendConflictError, ManualSendQueueFullError from pyrit.backend.services.media_persistence import MediaAccessDeniedError +from pyrit.backend.services.operation_service import OperationService from pyrit.common.deprecation import print_deprecation_message from pyrit.memory.memory_interface import AttackStateConflictError @@ -380,6 +382,20 @@ async def get_attack_metadata_async( return metadata +@router.get("/{attack_result_id}/findings", response_model=AttackFindingListResponse) +async def get_attack_findings_async(attack_result_id: uuid.UUID) -> AttackFindingListResponse: + """ + List the findings that link this attack as evidence, across all Operations. + + An attack with no linked findings, including a running scan attack that is not saved yet, + returns an empty list rather than 404. + + Returns: + AttackFindingListResponse: The linked findings with their Operation names and evidence links. + """ + return await OperationService().list_attack_findings_async(attack_result_id) + + @router.patch( "/{attack_result_id}", response_model=AttackSummary, diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index a97e6931fc..246d3b7af7 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -22,7 +22,7 @@ from collections.abc import Mapping, Sequence from datetime import UTC, datetime from functools import lru_cache -from typing import Any, Literal +from typing import Any, Literal, NamedTuple from pyrit.backend.mappers import ( attack_result_to_summary_async, @@ -60,6 +60,7 @@ from pyrit.backend.services.chat_scoring_service import get_chat_scoring_service from pyrit.backend.services.component_lifecycle import release_component_async from pyrit.backend.services.history_replay import validate_history_replay +from pyrit.backend.services.live_attack_reader import LiveAttackLookup from pyrit.backend.services.media_persistence import persist_message_pieces_async from pyrit.backend.services.message_send_service import ( MessageSendService, @@ -83,7 +84,7 @@ data_serializer_factory, set_message_piece_sha256_async, ) -from pyrit.memory.memory_interface import AttackStateConflictError +from pyrit.memory.memory_interface import AttackConversationPage, AttackStateConflictError from pyrit.models import ( AtomicAttackIdentifier, AttackIdentifier, @@ -132,19 +133,35 @@ class AttackObjectiveConflictError(Exception): """The shared objective changed while it was being edited.""" +class _ReadableAttack(NamedTuple): + """A saved result, or the projection of an attack still running in a scan.""" + + result: AttackResult + in_progress: bool + + class AttackService: """ Service for managing attacks. - Uses PyRIT memory (database) as the source of truth via AttackResult. + Uses PyRIT memory (database) as the source of truth via AttackResult. Read APIs also + resolve scan attacks that are still running through an optional read-only lookup; + a saved result always wins over that live projection. """ def __init__( - self, *, message_send_service: MessageSendService | None = None, memory: MemoryInterface | None = None + self, + *, + message_send_service: MessageSendService | None = None, + memory: MemoryInterface | None = None, + live_attack_lookup: LiveAttackLookup | None = None, ) -> None: """Initialize the attack service.""" self._memory = memory if memory is not None else CentralMemory.get_memory_instance() - self._saved_conversation_reader = SavedConversationReadService(memory=self._memory) + self._live_attack_lookup = live_attack_lookup + self._saved_conversation_reader = SavedConversationReadService( + memory=self._memory, read_references=self._read_conversation_page_async + ) self._message_send_service = ( message_send_service if message_send_service is not None else get_message_send_service() ) @@ -362,16 +379,16 @@ async def get_attack_async(self, *, attack_result_id: str) -> AttackSummary | No """ Get attack details (high-level metadata, no messages). - Queries the AttackResult from the database by its primary key. + Reads the saved AttackResult by its primary key, or projects a scan attack that is still running. Returns: AttackSummary if found, None otherwise. """ - results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) - if not results: + attack = await self._read_attack_async(attack_result_id=attack_result_id) + if attack is None: return None - return await self._summarize_attack_async(ar=results[0]) + return await self._summarize_attack_async(ar=attack.result, in_progress=attack.in_progress) async def get_attack_metadata_async( self, *, attack_result_id: str, conversation_id: str | None = None @@ -385,24 +402,25 @@ async def get_attack_metadata_async( Returns: AttackMetadata | None: Metadata without reference arrays, or None for an absent attack. """ - results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) - if not results: + attack = await self._read_attack_async(attack_result_id=attack_result_id) + if attack is None: return None - ar = results[0] - references = await self._memory.get_attack_conversation_page_async(attack_result_id=attack_result_id, limit=1) + ar = attack.result + references = await self._read_conversation_page_async(attack_result_id=attack_result_id, limit=1) if references is None: return None selected_id = conversation_id if conversation_id is not None else ar.conversation_id selected_valid = bool(selected_id) and selected_id == references.main_conversation_id if not selected_valid: - selected = await self._memory.get_attack_conversation_page_async( + selected = await self._read_conversation_page_async( attack_result_id=attack_result_id, conversation_id=selected_id, limit=1 ) if selected is None: return None selected_valid = bool(selected.conversations) summary = await self._summarize_attack_async( - ar=ar.model_copy(update={"related_conversations": set(), "last_response": None}) + ar=ar.model_copy(update={"related_conversations": set(), "last_response": None}), + in_progress=attack.in_progress, ) return AttackMetadata( attack_result_id=summary.attack_result_id, @@ -418,6 +436,8 @@ async def get_attack_metadata_async( chat_scorer=summary.chat_scorer, chat_scorer_inherited=summary.chat_scorer_inherited, execution_read_only=summary.execution_read_only, + in_progress=summary.in_progress, + scenario_result_id=summary.scenario_result_id, manual_continuation=summary.manual_continuation, last_response_message_piece_id=str(ar.last_response.id) if ar.last_response is not None else None, related_conversation_count=references.total_count - int(bool(references.main_conversation_id)), @@ -425,7 +445,42 @@ async def get_attack_metadata_async( selected_conversation_valid=selected_valid, ) - async def _summarize_attack_async(self, *, ar: AttackResult) -> AttackSummary: + async def _read_attack_async(self, *, attack_result_id: str) -> _ReadableAttack | None: + """ + Read the saved result, or project the attack while a scan is still running it. + + The live snapshot is taken before the storage read, so an attack that finishes in between + resolves to its saved result. + + Returns: + _ReadableAttack | None: The readable attack, or None when it is neither saved nor running. + """ + live = self._live_attack_lookup(attack_result_id) if self._live_attack_lookup else None + results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) + if results: + return _ReadableAttack(result=results[0], in_progress=False) + if live is None: + return None + return _ReadableAttack(result=live.to_attack_result(), in_progress=True) + + async def _read_conversation_page_async( + self, *, attack_result_id: str, offset: int = 0, limit: int = 50, conversation_id: str | None = None + ) -> AttackConversationPage | None: + """ + Page saved conversation references, or a running scan attack's live references. + + Returns: + AttackConversationPage | None: The reference page, or None when the attack is neither saved nor running. + """ + live = self._live_attack_lookup(attack_result_id) if self._live_attack_lookup else None + page = await self._memory.get_attack_conversation_page_async( + attack_result_id=attack_result_id, offset=offset, limit=limit, conversation_id=conversation_id + ) + if page is not None or live is None: + return page + return live.conversation_page(offset=offset, limit=limit, conversation_id=conversation_id) + + async def _summarize_attack_async(self, *, ar: AttackResult, in_progress: bool = False) -> AttackSummary: """ Map one loaded result with the shared presentation and scoring policy. @@ -440,7 +495,8 @@ async def _summarize_attack_async(self, *, ar: AttackResult) -> AttackSummary: stats=stats, operation_name=operation_names.get(ar.operation_id) if ar.operation_id else None, scoring_inputs=get_chat_scoring_service().inputs(ar), - execution_read_only=get_chat_scoring_service().execution_read_only(ar), + execution_read_only=in_progress or get_chat_scoring_service().execution_read_only(ar), + in_progress=in_progress, ) async def _get_operation_names_async(self, *, results: Sequence[AttackResult]) -> dict[uuid.UUID, str]: @@ -487,12 +543,12 @@ async def get_conversation_messages_async( raise ValueError("A message-page cursor requires a limit.") # Check attack exists - results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) - if not results: + attack = await self._read_attack_async(attack_result_id=attack_result_id) + if attack is None: return None # Verify the conversation belongs to this attack - ar = results[0] + ar = attack.result if conversation_id not in ar.get_active_conversation_ids(): raise ValueError(f"Conversation '{conversation_id}' is not part of attack '{attack_result_id}'") @@ -1276,12 +1332,11 @@ async def get_conversations_async( ) if cursor is not None: raise ValueError("A conversation-page cursor requires a limit.") - results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) - if not results: + attack = await self._read_attack_async(attack_result_id=attack_result_id) + if attack is None: return None - # attack_result_id is a unique primary key, so at most one result is returned. - ar = results[0] + ar = attack.result # Collect all conversation IDs (main + PRUNED related) and fetch stats in one query. active_conv_ids = list(ar.get_active_conversation_ids()) @@ -1597,7 +1652,12 @@ def get_attack_service() -> AttackService: """ Get the global attack service instance. + Wires the scenario run service's read-only lookup for scan attacks that are still running. + Returns: The singleton AttackService instance. """ - return AttackService() + # Composition root: AttackService depends only on the lookup type, not on scenario scheduling. + from pyrit.backend.services.scenario_run_service import find_live_attack + + return AttackService(live_attack_lookup=find_live_attack) diff --git a/pyrit/backend/services/live_attack_reader.py b/pyrit/backend/services/live_attack_reader.py new file mode 100644 index 0000000000..739054b58b --- /dev/null +++ b/pyrit/backend/services/live_attack_reader.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Read-only port for scan attacks that are still running. + +A scan saves each ``AttackResult`` only when its attack finishes. Until then, the attack read APIs +project the live execution through this port. The scenario run service provides the lookup; the +attack service depends only on these types, so it never imports scenario scheduling logic. +""" + +from collections.abc import Callable +from dataclasses import dataclass + +from pyrit.executor.attack.core.attack_executor import ActiveAttackExecution +from pyrit.memory.memory_interface import AttackConversationPage +from pyrit.models import AttackOutcome, AttackResult, ConversationReference, ConversationType + + +@dataclass(frozen=True, slots=True) +class LiveAttack: + """An attack still running in a scan, with at least one objective conversation.""" + + scenario_result_id: str + execution: ActiveAttackExecution + + @classmethod + def from_execution(cls, *, scenario_result_id: str, execution: ActiveAttackExecution) -> "LiveAttack | None": + """ + Wrap an execution once it has an objective conversation to read. + + Returns: + LiveAttack | None: The live attack, or None before its first objective conversation starts. + """ + if not execution.conversation_ids and execution.current_conversation_id is None: + return None + return cls(scenario_result_id=scenario_result_id, execution=execution) + + @property + def main_conversation_id(self) -> str: + """The conversation the attack is working on, otherwise its first objective conversation.""" + return self.execution.current_conversation_id or self.execution.conversation_ids[0] + + @property + def other_conversation_ids(self) -> list[str]: + """The other objective conversations started so far, in first-use order.""" + main = self.main_conversation_id + return [conversation_id for conversation_id in self.execution.conversation_ids if conversation_id != main] + + def to_attack_result(self) -> AttackResult: + """ + Project the unsaved result the attack read APIs show while the attack runs. + + Other objective conversations are user-visible references, like the pruned conversations of a saved result. + + Returns: + AttackResult: An undetermined, unsaved result owned by the running scan. + """ + execution = self.execution + return AttackResult( + attack_result_id=execution.attack_result_id, + conversation_id=self.main_conversation_id, + objective=execution.objective, + atomic_attack_identifier=execution.atomic_attack_identifier, + outcome=AttackOutcome.UNDETERMINED, + timestamp=execution.started_at, + related_conversations={ + ConversationReference(conversation_id=conversation_id, conversation_type=ConversationType.PRUNED) + for conversation_id in self.other_conversation_ids + }, + attribution_parent_id=self.scenario_result_id, + ) + + def conversation_page( + self, *, offset: int = 0, limit: int = 50, conversation_id: str | None = None + ) -> AttackConversationPage: + """ + Page the live conversation references the same way saved references are paged. + + Returns: + AttackConversationPage: The main conversation first, then the others in first-use order. + """ + return AttackConversationPage.from_references( + main_conversation_id=self.main_conversation_id, + pruned_conversation_ids=self.other_conversation_ids, + timestamp=self.execution.started_at, + offset=offset, + limit=limit, + conversation_id=conversation_id, + ) + + +#: Finds a running scan attack by its ``attack_result_id``; None when no active scan is running it. +LiveAttackLookup = Callable[[str], LiveAttack | None] diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py index 8fa55438f0..5417c16285 100644 --- a/pyrit/backend/services/operation_service.py +++ b/pyrit/backend/services/operation_service.py @@ -6,6 +6,8 @@ from pyrit.backend.models.attacks import ScoreView from pyrit.backend.models.operations import ( + AttackFindingItem, + AttackFindingListResponse, EvidenceAvailability, FindingCreateRequest, FindingEvidenceAttachResponse, @@ -382,6 +384,30 @@ async def detach_finding_evidence_async(self, *, operation_id: UUID, finding_id: except LookupError as exc: raise FindingEvidenceNotFoundError(str(exc)) from exc + async def list_attack_findings_async(self, attack_result_id: UUID) -> AttackFindingListResponse: + """ + Read every finding that links an attack as evidence, across all Operations. + + An attack with no linked findings, including one that is not saved yet, has an empty list. + + Returns: + AttackFindingListResponse: The linked findings with their Operation names and evidence links. + """ + memory = CentralMemory.get_memory_instance() + linked = await memory.get_findings_for_attack_result_async(attack_result_id=attack_result_id) + operation_ids = list({finding.operation_id for finding, _ in linked}) + operations = await memory.get_operations_async(operation_ids=operation_ids) if operation_ids else [] + names = {operation.id: operation.name for operation in operations} + return AttackFindingListResponse( + attack_result_id=attack_result_id, + items=[ + AttackFindingItem( + **finding.model_dump(), operation_name=names.get(finding.operation_id, ""), evidence_id=evidence.id + ) + for finding, evidence in linked + ], + ) + async def _require_finding_async(self, *, operation_id: UUID, finding_id: UUID) -> Finding: finding = await CentralMemory.get_memory_instance().get_finding_async( operation_id=operation_id, finding_id=finding_id diff --git a/pyrit/backend/services/saved_conversation_read_service.py b/pyrit/backend/services/saved_conversation_read_service.py index a78bff901d..1558cd79bd 100644 --- a/pyrit/backend/services/saved_conversation_read_service.py +++ b/pyrit/backend/services/saved_conversation_read_service.py @@ -4,6 +4,7 @@ """Bounded, read-only saved conversation views shared by attack and scenario APIs.""" import uuid +from typing import Protocol from pyrit.backend.mappers._preview import format_last_message_preview from pyrit.backend.mappers.attack_mappers import pyrit_messages_to_dto_async @@ -14,16 +15,34 @@ from pyrit.backend.models.common import PaginationInfo from pyrit.common.pagination import decode_keyset_cursor, encode_keyset_cursor, fingerprint_filters from pyrit.memory import MemoryInterface -from pyrit.memory.memory_interface import MessagePieceKeysetCursor +from pyrit.memory.memory_interface import AttackConversationPage, MessagePieceKeysetCursor from pyrit.models import ConversationStats, ConversationSummary +class ConversationReferenceReader(Protocol): + """Reads one page of an attack's main and pruned conversation references.""" + + async def __call__( + self, *, attack_result_id: str, offset: int = 0, limit: int = 50, conversation_id: str | None = None + ) -> AttackConversationPage | None: + """Return the page, or None for an absent attack.""" + ... + + class SavedConversationReadService: """Read canonical saved references/messages without constructing an executable attack.""" - def __init__(self, *, memory: MemoryInterface) -> None: - """Bind all saved reads to the caller's memory instance.""" + def __init__(self, *, memory: MemoryInterface, read_references: ConversationReferenceReader | None = None) -> None: + """ + Bind all saved reads to the caller's memory instance. + + Args: + memory (MemoryInterface): Memory that holds the saved messages. + read_references (ConversationReferenceReader | None): Where attack conversation references come from. + Defaults to the saved references in memory. + """ self._memory = memory + self._read_references = read_references or self._read_saved_references_async async def list_conversations_async( self, *, attack_result_id: str, limit: int = 50, cursor: str | None = None @@ -37,7 +56,7 @@ async def list_conversations_async( Raises: ValueError: If the cursor belongs to another attack or an obsolete reference ordering. """ - page = await self._memory.get_attack_conversation_page_async(attack_result_id=attack_result_id, limit=limit) + page = await self._read_references(attack_result_id=attack_result_id, limit=limit) if page is None: return None fingerprint = fingerprint_filters( @@ -60,9 +79,7 @@ async def list_conversations_async( offset = anchor.position if anchor is not None and anchor.position is not None else 0 if offset: revision = page.reference_revision - page = await self._memory.get_attack_conversation_page_async( - attack_result_id=attack_result_id, offset=offset, limit=limit - ) + page = await self._read_references(attack_result_id=attack_result_id, offset=offset, limit=limit) if page is None: return None if page.reference_revision != revision: @@ -120,7 +137,7 @@ async def get_messages_async( Raises: ValueError: If the conversation is unrelated or the continuation cursor is invalid. """ - references = await self._memory.get_attack_conversation_page_async( + references = await self._read_references( attack_result_id=attack_result_id, conversation_id=conversation_id, limit=1 ) if references is None: @@ -165,3 +182,10 @@ async def get_messages_async( pagination=PaginationInfo(limit=limit, has_more=next_cursor is not None, next_cursor=next_cursor), partial_message_sequences=list(page.partial_message_sequences), ) + + async def _read_saved_references_async( + self, *, attack_result_id: str, offset: int = 0, limit: int = 50, conversation_id: str | None = None + ) -> AttackConversationPage | None: + return await self._memory.get_attack_conversation_page_async( + attack_result_id=attack_result_id, offset=offset, limit=limit, conversation_id=conversation_id + ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index b21ff17b33..9540622d08 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -36,6 +36,7 @@ from pyrit.backend.models.attacks import ConversationMessagesResponse from pyrit.backend.models.common import PaginationInfo, filter_sensitive_fields from pyrit.backend.models.scenarios import ScenarioRunListResponse +from pyrit.backend.services.live_attack_reader import LiveAttack from pyrit.backend.services.pagination import ( decode_keyset_cursor, encode_keyset_cursor, @@ -2562,6 +2563,21 @@ def snapshot_active_run(self, *, scenario_result_id: str) -> _ActiveRunSnapshot: active_scenario_result_id=active_scenario_result_id, ) + def find_live_attack(self, *, attack_result_id: str) -> LiveAttack | None: + """ + Find an attack that an active scan is still running. + + Returns: + LiveAttack | None: The running attack and its scan, or None when no active scan is running it + or it has not started an objective conversation yet. + """ + for scenario_result_id in list(self._active_tasks): + snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id) + for execution in snapshot.executions: + if execution.attack_result_id == attack_result_id: + return LiveAttack.from_execution(scenario_result_id=scenario_result_id, execution=execution) + return None + def _load_run_plan(self, *, scenario_result: ScenarioResult) -> ScenarioRunPlan | None: """ Load a validated plan from scenario metadata. @@ -2701,6 +2717,7 @@ def _active_conversations( ScenarioActiveConversation( attack_result_id=execution.attack_result_id, conversation_id=conversation_id, + conversation_ids=list(execution.conversation_ids), atomic_attack_name=attribution.parent_collection, seed_group_id=attribution.seed_group_id, objective=execution.objective, @@ -3258,6 +3275,19 @@ def peek_scenario_run_service() -> ScenarioRunService | None: return _service_instance +def find_live_attack(attack_result_id: str) -> LiveAttack | None: + """ + Find a running scan attack in the current service, without constructing one. + + Resolves the service on every call, so the lookup stays valid across service resets. + + Returns: + LiveAttack | None: The running attack, or None when no active scan is running it. + """ + service = peek_scenario_run_service() + return service.find_live_attack(attack_result_id=attack_result_id) if service is not None else None + + async def reset_scenario_run_service_async() -> None: """Close and discard the drained singleton.""" global _service_instance diff --git a/pyrit/memory/alembic/versions/b5d7f9a1c3e6_index_finding_evidence_attack_result.py b/pyrit/memory/alembic/versions/b5d7f9a1c3e6_index_finding_evidence_attack_result.py new file mode 100644 index 0000000000..0d09bafd12 --- /dev/null +++ b/pyrit/memory/alembic/versions/b5d7f9a1c3e6_index_finding_evidence_attack_result.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Index finding evidence by attack result so an attack's findings are a direct lookup. + +``attack_result_id`` is a bounded UUID column on every backend, so SQL Server accepts it as an index key. + +Revision ID: b5d7f9a1c3e6 +Revises: e7c4a2d9f013 +Create Date: 2026-11-06 17:30:00.000000 +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "b5d7f9a1c3e6" +down_revision: str | None = "e7c4a2d9f013" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Add the attack-result lookup index.""" + op.create_index("ix_FindingEvidenceEntries_attack_result_id", "FindingEvidenceEntries", ["attack_result_id"]) + + +def downgrade() -> None: + """Remove the attack-result lookup index.""" + op.drop_index("ix_FindingEvidenceEntries_attack_result_id", table_name="FindingEvidenceEntries") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 47e2899955..9c40bab306 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -503,6 +503,58 @@ class AttackConversationPage: reference_revision: str objective_score_id: str | None = None + @classmethod + def from_references( + cls, + *, + main_conversation_id: str, + pruned_conversation_ids: Sequence[str], + timestamp: datetime, + offset: int = 0, + limit: int = 50, + conversation_id: str | None = None, + objective_score_id: str | None = None, + ) -> "AttackConversationPage": + """ + Page main and pruned references in stored order: main first, then pruned in their recorded order. + + The revision fingerprints the full ordering, so continuation cursors detect reference changes. + + Returns: + AttackConversationPage: The selected or paged references and their ordering revision. + """ + references: dict[str, AttackConversationRecord] = {} + if main_conversation_id: + references[main_conversation_id] = AttackConversationRecord(conversation_id=main_conversation_id) + for identifier in pruned_conversation_ids: + if identifier: + references.setdefault( + identifier, + AttackConversationRecord(conversation_id=identifier, conversation_type=ConversationType.PRUNED), + ) + ordered = tuple(references.values()) + revision = hashlib.sha256( + json.dumps( + [ + ( + reference.conversation_id, + reference.conversation_type.value if reference.conversation_type is not None else None, + ) + for reference in ordered + ] + ).encode("utf-8") + ).hexdigest() + if conversation_id is not None: + ordered = tuple(reference for reference in ordered if reference.conversation_id == conversation_id) + return cls( + main_conversation_id=main_conversation_id, + conversations=ordered[offset : offset + limit], + total_count=len(ordered), + timestamp=timestamp, + reference_revision=revision, + objective_score_id=objective_score_id, + ) + @dataclass(frozen=True, slots=True, kw_only=True) class _AttackResultQuery: @@ -1211,6 +1263,25 @@ async def get_finding_evidence_by_source_async( ) return entry.get_evidence() if entry else None + async def get_findings_for_attack_result_async( + self, *, attack_result_id: uuid.UUID + ) -> list[tuple[Finding, FindingEvidence]]: + """ + Read every finding that links an attack result as evidence, highest severity first, then newest first. + + Returns: + list[tuple[Finding, FindingEvidence]]: Each linked finding with the evidence that links it. + """ + statement = ( + select(FindingEntry, FindingEvidenceEntry) + .join(FindingEvidenceEntry, FindingEvidenceEntry.finding_id == FindingEntry.id) + .where(FindingEvidenceEntry.attack_result_id == attack_result_id) + .order_by(*self._finding_order()) + ) + async with await self.get_session_async() as session: + rows = (await session.execute(statement)).all() + return [(finding.get_finding(), evidence.get_evidence()) for finding, evidence in rows] + async def delete_finding_evidence_async(self, *, finding_id: uuid.UUID, evidence_id: uuid.UUID) -> None: """ Unlink evidence from a finding. The attack result is kept. @@ -7978,6 +8049,7 @@ def _row_to_scenario_attack_result_delta(row: Mapping[str, Any]) -> ScenarioAtta status=ScoreStatus(row["score_status"]), score_value=row["score_value"], score_rationale=row["score_rationale"], + score_category=list(row["score_category"] or []), ) related_cids: list[str] = [] for col in ( @@ -8036,6 +8108,7 @@ def _scenario_attack_result_delta_statement() -> Select: ScoreEntry.score_type, ScoreEntry.status.label("score_status"), ScoreEntry.score_rationale, + ScoreEntry.score_category, ScoreEntry.scorer_class_identifier, ).outerjoin( ScoreEntry, @@ -8483,38 +8556,16 @@ def _execute_get_attack_conversation_page( row = session.execute(statement).mappings().one_or_none() if row is None: return None - references: dict[str, AttackConversationRecord] = {} - main_id = row["conversation_id"] or "" - if main_id: - references[main_id] = AttackConversationRecord(conversation_id=main_id) values = row["pruned_conversation_ids"] - if isinstance(values, list): - for identifier in values: - if isinstance(identifier, str) and identifier: - references.setdefault( - identifier, - AttackConversationRecord(conversation_id=identifier, conversation_type=ConversationType.PRUNED), - ) - ordered = tuple(references.values()) - revision = hashlib.sha256( - json.dumps( - [ - ( - reference.conversation_id, - reference.conversation_type.value if reference.conversation_type is not None else None, - ) - for reference in ordered - ] - ).encode("utf-8") - ).hexdigest() - if conversation_id is not None: - ordered = tuple(reference for reference in ordered if reference.conversation_id == conversation_id) - return AttackConversationPage( - main_conversation_id=main_id, - conversations=ordered[offset : offset + limit], - total_count=len(ordered), + return AttackConversationPage.from_references( + main_conversation_id=row["conversation_id"] or "", + pruned_conversation_ids=[value for value in values if isinstance(value, str)] + if isinstance(values, list) + else [], timestamp=row["timestamp"], - reference_revision=revision, + offset=offset, + limit=limit, + conversation_id=conversation_id, objective_score_id=str(row["objective_score_id"]) if row["objective_score_id"] is not None else None, ) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index f9f1c3b383..c1ed3ca7af 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -374,6 +374,7 @@ class FindingEvidenceEntry(Base): __tablename__ = "FindingEvidenceEntries" __table_args__ = ( Index("ix_FindingEvidenceEntries_finding_id", "finding_id"), + Index("ix_FindingEvidenceEntries_attack_result_id", "attack_result_id"), UniqueConstraint("finding_id", "attack_result_id", name="uq_finding_evidence_source"), ) diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py index ed36364e69..57d3006a6e 100644 --- a/pyrit/models/scenario_progress.py +++ b/pyrit/models/scenario_progress.py @@ -156,6 +156,7 @@ class ScenarioProgressScore(BaseModel): status: ScoreStatus score_value: str | None = None score_rationale: str | None = None + score_category: list[str] = Field(default_factory=list) class ScenarioComponentIdentity(BaseModel): @@ -365,7 +366,10 @@ class ScenarioActiveConversation(BaseModel): """A saved objective conversation and the scan execution that produced it.""" attack_result_id: str + #: The conversation the attack is working on, or the recorded objective conversation. conversation_id: str + #: Every objective conversation the attack has started so far, in first-use order. + conversation_ids: list[str] = Field(default_factory=list) atomic_attack_name: str seed_group_id: str | None = None objective: str diff --git a/tests/unit/backend/test_live_attack_reads.py b/tests/unit/backend/test_live_attack_reads.py new file mode 100644 index 0000000000..9de299f023 --- /dev/null +++ b/tests/unit/backend/test_live_attack_reads.py @@ -0,0 +1,292 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Attack reads for scan attacks that are still running, and findings linked to an attack.""" + +import uuid +from datetime import UTC, datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +from pyrit.backend.main import app +from pyrit.backend.models.attacks import AttackSummary +from pyrit.backend.models.operations import AttackFindingItem, AttackFindingListResponse +from pyrit.backend.services import attack_service as attack_service_module +from pyrit.backend.services import scenario_run_service as scenario_run_service_module +from pyrit.backend.services.attack_service import AttackService +from pyrit.backend.services.live_attack_reader import LiveAttack +from pyrit.backend.services.message_send_service import MessageSendService +from pyrit.backend.services.operation_service import OperationService +from pyrit.backend.services.scenario_run_service import ScenarioRunService, _ActiveTask, find_live_attack +from pyrit.executor.attack.core.attack_executor import ActiveAttackExecution +from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution +from pyrit.memory import SQLiteMemory +from pyrit.models import ( + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + ComponentIdentifier, + ConversationType, + Finding, + FindingEvidence, + Operation, +) +from pyrit.scenario.core import Scenario + +RUN_ID = str(uuid.UUID(int=100)) +ATTACK_ID = str(uuid.UUID(int=200)) +STARTED_AT = datetime(2026, 11, 6, 12, tzinfo=UTC) + + +def _execution( + *, conversation_ids: tuple[str, ...] = ("first", "current"), current: str | None = "current" +) -> ActiveAttackExecution: + return ActiveAttackExecution( + attack_result_id=ATTACK_ID, + objective="Reveal the system prompt", + conversation_ids=conversation_ids, + current_conversation_id=current, + started_at=STARTED_AT, + attribution=AttackResultAttribution(parent_id=RUN_ID, parent_collection="live"), + atomic_attack_identifier=AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="CrescendoAttack", class_module="pyrit.executor.attack") + ), + ) + + +def _live_attack(**execution: object) -> LiveAttack: + live = LiveAttack.from_execution(scenario_result_id=RUN_ID, execution=_execution(**execution)) + assert live is not None + return live + + +def _attack_service(*, memory: SQLiteMemory, live: LiveAttack | None) -> AttackService: + return AttackService( + message_send_service=MagicMock(spec=MessageSendService), + memory=memory, + live_attack_lookup=lambda attack_result_id: live if attack_result_id == ATTACK_ID else None, + ) + + +def test_live_attack_requires_an_objective_conversation() -> None: + assert ( + LiveAttack.from_execution(scenario_result_id=RUN_ID, execution=_execution(conversation_ids=(), current=None)) + is None + ) + + +@pytest.mark.parametrize( + "conversation_ids, current, expected_main, expected_others", + [ + (("first", "current"), "current", "current", ["first"]), + (("first", "second"), None, "first", ["second"]), + ((), "current", "current", []), + ], +) +def test_live_attack_main_is_current_otherwise_first( + conversation_ids: tuple[str, ...], current: str | None, expected_main: str, expected_others: list[str] +) -> None: + live = _live_attack(conversation_ids=conversation_ids, current=current) + assert live.main_conversation_id == expected_main + assert live.other_conversation_ids == expected_others + + +def test_live_attack_projects_an_undetermined_unsaved_result() -> None: + result = _live_attack().to_attack_result() + assert result.attack_result_id == ATTACK_ID + assert result.conversation_id == "current" + assert result.objective == "Reveal the system prompt" + assert result.outcome == AttackOutcome.UNDETERMINED + assert result.timestamp == STARTED_AT + assert result.attribution_parent_id == RUN_ID + assert result.get_active_conversation_ids() == {"first", "current"} + assert {ref.conversation_type for ref in result.related_conversations} == {ConversationType.PRUNED} + assert result.get_attack_strategy_identifier().class_name == "CrescendoAttack" + + +def test_live_attack_pages_references_like_saved_references() -> None: + live = _live_attack(conversation_ids=("first", "current", "third")) + page = live.conversation_page(limit=2) + assert [reference.conversation_id for reference in page.conversations] == ["current", "first"] + assert page.total_count == 3 + assert page.timestamp == STARTED_AT + assert live.conversation_page(offset=2, limit=2).conversations[0].conversation_id == "third" + assert live.conversation_page(conversation_id="first").conversations[0].conversation_type == ConversationType.PRUNED + assert live.conversation_page(conversation_id="unrelated").conversations == () + assert live.conversation_page().reference_revision == page.reference_revision + + +@pytest.mark.usefixtures("patch_central_database") +def test_scenario_run_service_finds_only_attacks_its_active_runs_are_running() -> None: + service = ScenarioRunService() + scenario = MagicMock(spec=Scenario) + scenario.active_executions = (_execution(),) + scenario.active_atomic_group_ids = set() + assert service.find_live_attack(attack_result_id=ATTACK_ID) is None + + service._active_tasks[RUN_ID] = _ActiveTask(scenario_result_id=RUN_ID, scenario=scenario) + live = service.find_live_attack(attack_result_id=ATTACK_ID) + assert live is not None + assert (live.scenario_result_id, live.main_conversation_id) == (RUN_ID, "current") + assert service.find_live_attack(attack_result_id=str(uuid.uuid4())) is None + + scenario.active_executions = (_execution(conversation_ids=(), current=None),) + assert service.find_live_attack(attack_result_id=ATTACK_ID) is None + + scenario.active_executions = () + assert service.find_live_attack(attack_result_id=ATTACK_ID) is None + + +def test_module_lookup_does_not_construct_a_scenario_run_service() -> None: + with patch.object(scenario_run_service_module, "_service_instance", None): + assert find_live_attack(ATTACK_ID) is None + assert scenario_run_service_module._service_instance is None + + +def test_module_lookup_reads_the_current_scenario_run_service() -> None: + service = MagicMock(spec=ScenarioRunService) + service.find_live_attack.return_value = _live_attack() + with patch.object(scenario_run_service_module, "_service_instance", service): + assert find_live_attack(ATTACK_ID) is service.find_live_attack.return_value + assert service.find_live_attack.call_args.kwargs == {"attack_result_id": ATTACK_ID} + + +def test_attack_service_singleton_reads_running_attacks_through_the_scenario_lookup() -> None: + attack_service_module.get_attack_service.cache_clear() + try: + with patch.object(attack_service_module, "AttackService") as constructor: + attack_service_module.get_attack_service() + assert constructor.call_args.kwargs == {"live_attack_lookup": find_live_attack} + finally: + attack_service_module.get_attack_service.cache_clear() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_attack_service_reads_running_attack_when_unsaved_async(sqlite_instance: SQLiteMemory) -> None: + service = _attack_service(memory=sqlite_instance, live=_live_attack()) + + summary = await service.get_attack_async(attack_result_id=ATTACK_ID) + metadata = await service.get_attack_metadata_async(attack_result_id=ATTACK_ID, conversation_id="first") + conversations = await service.get_conversations_async(attack_result_id=ATTACK_ID, limit=50) + + assert summary is not None and metadata is not None and conversations is not None + data = summary.model_dump(mode="json") + assert (data["in_progress"], data["execution_read_only"], data["outcome"]) == (True, True, "undetermined") + assert data["scenario_result_id"] == RUN_ID + assert data["related_conversation_ids"] == ["first"] + assert (metadata.in_progress, metadata.scenario_result_id) == (True, RUN_ID) + assert (metadata.related_conversation_count, metadata.selected_conversation_valid) == (1, True) + assert conversations.main_conversation_id == "current" + assert [item.conversation_id for item in conversations.conversations] == ["current", "first"] + with pytest.raises(ValueError, match="not part of attack"): + await service.get_conversation_messages_async(attack_result_id=ATTACK_ID, conversation_id="other") + + +@pytest.mark.usefixtures("patch_central_database") +async def test_attack_service_saved_result_wins_over_running_attack_async(sqlite_instance: SQLiteMemory) -> None: + saved = AttackResult( + attack_result_id=ATTACK_ID, + conversation_id="saved-main", + objective="Saved objective", + outcome=AttackOutcome.SUCCESS, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[saved]) + service = _attack_service(memory=sqlite_instance, live=_live_attack()) + + summary = await service.get_attack_async(attack_result_id=ATTACK_ID) + metadata = await service.get_attack_metadata_async(attack_result_id=ATTACK_ID) + conversations = await service.get_conversations_async(attack_result_id=ATTACK_ID, limit=50) + + assert summary is not None and metadata is not None and conversations is not None + assert (summary.in_progress, summary.outcome, summary.conversation_id) == ( + False, + AttackOutcome.SUCCESS, + "saved-main", + ) + assert summary.scenario_result_id is None + assert (metadata.in_progress, metadata.conversation_id) == (False, "saved-main") + assert [item.conversation_id for item in conversations.conversations] == ["saved-main"] + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("with_lookup", [True, False]) +async def test_attack_service_unknown_attack_is_not_found_async( + sqlite_instance: SQLiteMemory, with_lookup: bool +) -> None: + service = ( + _attack_service(memory=sqlite_instance, live=None) + if with_lookup + else AttackService(message_send_service=MagicMock(spec=MessageSendService), memory=sqlite_instance) + ) + assert await service.get_attack_async(attack_result_id=ATTACK_ID) is None + assert await service.get_attack_metadata_async(attack_result_id=ATTACK_ID) is None + assert await service.get_conversations_async(attack_result_id=ATTACK_ID) is None + assert await service.get_conversations_async(attack_result_id=ATTACK_ID, limit=50) is None + assert await service.get_conversation_messages_async(attack_result_id=ATTACK_ID, conversation_id="current") is None + assert ( + await service.get_conversation_messages_async(attack_result_id=ATTACK_ID, conversation_id="current", limit=100) + is None + ) + + +def test_attack_route_serializes_in_progress_attack(compatibility_headers: dict[str, str]) -> None: + summary = AttackSummary( + attack_result_id=ATTACK_ID, + conversation_id="current", + objective="Running", + in_progress=True, + execution_read_only=True, + scenario_result_id=RUN_ID, + ) + service = MagicMock(spec=AttackService) + service.get_attack_async = AsyncMock(return_value=summary) + with patch("pyrit.backend.routes.attacks.get_attack_service", return_value=service): + response = TestClient(app, headers=compatibility_headers).get(f"/api/attacks/{ATTACK_ID}") + assert response.status_code == 200 + body = response.json() + assert (body["in_progress"], body["scenario_result_id"], body["outcome"]) == (True, RUN_ID, "undetermined") + + +def test_attack_findings_route_returns_linked_findings(compatibility_headers: dict[str, str]) -> None: + finding = Finding(operation_id=uuid.uuid4(), title="Leak", severity="important") + listed = AttackFindingListResponse( + attack_result_id=uuid.UUID(ATTACK_ID), + items=[AttackFindingItem(**finding.model_dump(), operation_name="Red team", evidence_id=uuid.uuid4())], + ) + with patch.object( + OperationService, "list_attack_findings_async", new_callable=AsyncMock, return_value=listed + ) as list_findings: + response = TestClient(app, headers=compatibility_headers).get(f"/api/attacks/{ATTACK_ID}/findings") + assert response.status_code == 200 + assert response.json() == listed.model_dump(mode="json") + assert list_findings.call_args.args == (uuid.UUID(ATTACK_ID),) + + +def test_attack_findings_route_rejects_non_uuid_ids(compatibility_headers: dict[str, str]) -> None: + with patch.object(OperationService, "list_attack_findings_async", new_callable=AsyncMock) as list_findings: + response = TestClient(app, headers=compatibility_headers).get("/api/attacks/not-a-uuid/findings") + assert response.status_code == 422 + list_findings.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_list_attack_findings_names_operations_and_evidence_async(sqlite_instance: SQLiteMemory) -> None: + operation = await sqlite_instance.add_operation_async(Operation(name="Red team")) + finding = await sqlite_instance.add_finding_async( + finding=Finding(operation_id=operation.id, title="Leak", severity="critical") + ) + attack_result_id = uuid.uuid4() + evidence = await sqlite_instance.add_finding_evidence_async( + evidence=FindingEvidence(finding_id=finding.id, attack_result_id=attack_result_id) + ) + + listed = await OperationService().list_attack_findings_async(attack_result_id) + empty = await OperationService().list_attack_findings_async(uuid.uuid4()) + + assert listed.attack_result_id == attack_result_id + assert [item.model_dump(mode="json") for item in listed.items] == [ + {**finding.model_dump(mode="json"), "operation_name": "Red team", "evidence_id": str(evidence.id)} + ] + assert empty.items == [] diff --git a/tests/unit/backend/test_scenario_resume.py b/tests/unit/backend/test_scenario_resume.py index 0d4cb1ec1b..9db290614a 100644 --- a/tests/unit/backend/test_scenario_resume.py +++ b/tests/unit/backend/test_scenario_resume.py @@ -15,6 +15,7 @@ from pyrit.backend.models.attacks import ConversationMessagesResponse from pyrit.backend.services.attack_service import AttackService +from pyrit.backend.services.message_send_service import MessageSendService from pyrit.backend.services.scenario_run_service import ( _LAUNCH_REFERENCES_METADATA_KEY, _LAUNCH_REQUEST_FIELDS, @@ -342,6 +343,113 @@ async def build_async(self: Scenario, *, context: ScenarioContext) -> list[Atomi release.set() +async def test_running_scan_attack_reads_live_until_its_saved_result_wins_async( + resume_environment: tuple[ScenarioRunService, MockPromptTarget], +) -> None: + service, target = resume_environment + entered, release = asyncio.Event(), asyncio.Event() + conversations: list[str] = [] + original_send = target._send_prompt_to_target_async + + async def send_async(*, normalized_conversation: list[Message]) -> list[Message]: + conversations.append(normalized_conversation[-1].get_piece().conversation_id) + if len(conversations) == 2: + entered.set() + await release.wait() + return await original_send(normalized_conversation=normalized_conversation) + + async def build_async(self: Scenario, *, context: ScenarioContext) -> list[AtomicAttack]: + strategy = PromptSendingAttack( + objective_target=context.objective_target, + attack_scoring_config=AttackScoringConfig( + objective_scorer=SubStringScorer(substring="never matches the response") + ), + max_attempts_on_failure=1, + ) + return [ + AtomicAttack( + atomic_attack_name="live", + attack_technique=AttackTechnique(attack=strategy), + seed_groups=[context.seed_groups[0]], + memory_labels=context.memory_labels, + ) + ] + + attacks = AttackService( + message_send_service=MagicMock(spec=MessageSendService), + memory=service._memory, + live_attack_lookup=lambda attack_result_id: service.find_live_attack(attack_result_id=attack_result_id), + ) + with ( + patch.object(_OfflineResumeScenario, "_build_atomic_attacks_async", new=build_async), + patch.object(target, "_send_prompt_to_target_async", side_effect=send_async), + ): + response = await service.start_run_async( + request=RunScenarioRequest( + scenario_name=_SCENARIO_NAME, + target_name=_TARGET_NAME, + include_baseline=False, + max_concurrency=1, + max_retries=0, + ) + ) + run_id = response.scenario_result_id + try: + await asyncio.wait_for(entered.wait(), TEST_WATCHDOG_SECONDS) + first, current = conversations + progress = await service.get_run_progress_async(scenario_result_id=run_id, since=None, limit=500) + assert progress is not None + active = progress.active_conversations[0] + assert active.conversation_id == current + assert active.conversation_ids == [first, current] + attack_id = active.attack_result_id + + live = await attacks.get_attack_async(attack_result_id=attack_id) + assert live is not None + assert live.in_progress is True + assert live.execution_read_only is True + assert live.outcome == AttackOutcome.UNDETERMINED + assert live.scenario_result_id == run_id + assert live.objective == active.objective + assert live.conversation_id == current + assert live.related_conversation_ids == [first] + + metadata = await attacks.get_attack_metadata_async(attack_result_id=attack_id, conversation_id=first) + assert metadata is not None + assert (metadata.in_progress, metadata.scenario_result_id) == (True, run_id) + assert metadata.related_conversation_count == 1 + assert metadata.selected_conversation_valid is True + + unpaged = await attacks.get_conversations_async(attack_result_id=attack_id) + assert unpaged is not None and unpaged.main_conversation_id == current + assert {conversation.conversation_id for conversation in unpaged.conversations} == {first, current} + paged = await attacks.get_conversations_async(attack_result_id=attack_id, limit=50) + assert paged is not None + assert [conversation.conversation_id for conversation in paged.conversations] == [current, first] + + for limit in (None, 100): + messages = await attacks.get_conversation_messages_async( + attack_result_id=attack_id, conversation_id=first, limit=limit + ) + assert messages is not None + assert [message.role for message in messages.messages] == ["user", "assistant"] + with pytest.raises(ValueError, match="not part of attack"): + await attacks.get_conversation_messages_async( + attack_result_id=attack_id, conversation_id="unrelated", limit=100 + ) + + release.set() + await _wait_for_run_completion_async(service=service, scenario_result_id=run_id) + assert service.find_live_attack(attack_result_id=attack_id) is None + saved = await attacks.get_attack_async(attack_result_id=attack_id) + assert saved is not None + assert saved.in_progress is False + assert saved.scenario_result_id == run_id + assert saved.outcome == AttackOutcome.FAILURE + finally: + release.set() + + async def test_cancelled_provenance_preserves_mixed_outcome_denominators_and_resume_async( resume_environment: tuple[ScenarioRunService, MockPromptTarget], ) -> None: diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_progress.py b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py index 5169e3dcbb..54ffe4c103 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_progress.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py @@ -83,6 +83,7 @@ async def test_scenario_progress_deltas_page_equal_timestamps_by_id( score = Score( score_value="true", score_type="true_false", + score_category=["jailbreak"], score_rationale="The objective was achieved.", scorer_class_identifier=ComponentIdentifier( class_name="TestScorer", @@ -112,6 +113,7 @@ async def test_scenario_progress_deltas_page_equal_timestamps_by_id( assert first_page[0].score is not None assert first_page[0].score.scorer_name == "TestScorer" assert first_page[0].score.score_rationale == "The objective was achieved." + assert first_page[0].score.score_category == ["jailbreak"] assert has_more is True assert [row.attack_result_id for row in second_page] == [str(second_id)] assert second_has_more is False @@ -151,6 +153,7 @@ async def test_scenario_progress_delta_uses_unknown_for_empty_scorer_identifier( assert has_more is False assert deltas[0].score is not None assert deltas[0].score.scorer_name == "Unknown" + assert deltas[0].score.score_category == [] async def test_scenario_result_header_does_not_hydrate_attack_results( diff --git a/tests/unit/memory/test_finding_evidence_attack_index_migration.py b/tests/unit/memory/test_finding_evidence_attack_index_migration.py new file mode 100644 index 0000000000..01717beb37 --- /dev/null +++ b/tests/unit/memory/test_finding_evidence_attack_index_migration.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import importlib +from io import StringIO +from pathlib import Path + +import sqlalchemy as sa +from alembic import command +from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.operations import Operations +from alembic.script import ScriptDirectory + +from pyrit.memory.memory_models import FindingEvidenceEntry +from pyrit.memory.migration import check_schema_migrations + +PREVIOUS_REVISION = "e7c4a2d9f013" +REVISION = "b5d7f9a1c3e6" +INDEX = "ix_FindingEvidenceEntries_attack_result_id" + + +def _config(connection: sa.Connection | None = None) -> Config: + config = Config() + config.set_main_option("script_location", str(Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic")) + if connection is not None: + config.attributes["connection"] = connection + config.attributes["version_table"] = "pyrit_memory_alembic_version" + return config + + +def _evidence_indexes(connection: sa.Connection) -> dict[str, list[str]]: + return { + str(index["name"]): [str(column) for column in index["column_names"]] + for index in sa.inspect(connection).get_indexes("FindingEvidenceEntries") + } + + +def test_migration_is_the_single_head_after_score_participation() -> None: + scripts = ScriptDirectory.from_config(_config()) + assert scripts.get_heads() == [REVISION] + assert scripts.get_revision(REVISION).down_revision == PREVIOUS_REVISION + + +def test_migration_adds_and_drops_attack_result_index() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + config = _config(connection) + command.upgrade(config, PREVIOUS_REVISION) + assert INDEX not in _evidence_indexes(connection) + command.upgrade(config, REVISION) + assert _evidence_indexes(connection)[INDEX] == ["attack_result_id"] + command.downgrade(config, PREVIOUS_REVISION) + assert INDEX not in _evidence_indexes(connection) + command.upgrade(config, "head") + check_schema_migrations(engine=engine) + finally: + engine.dispose() + + +def test_migration_compiles_for_sql_server() -> None: + migration = importlib.import_module( + "pyrit.memory.alembic.versions.b5d7f9a1c3e6_index_finding_evidence_attack_result" + ) + output = StringIO() + context = MigrationContext.configure(dialect_name="mssql", opts={"as_sql": True, "output_buffer": output}) + with Operations.context(context): + migration.upgrade() + assert f"CREATE INDEX [{INDEX}] ON [FindingEvidenceEntries] (attack_result_id)" in output.getvalue() + + +def test_model_declares_attack_result_index() -> None: + indexes = { + index.name: [column.name for column in index.columns] for index in FindingEvidenceEntry.__table__.indexes + } + assert indexes[INDEX] == ["attack_result_id"] diff --git a/tests/unit/memory/test_findings.py b/tests/unit/memory/test_findings.py index 94d45c0e7a..be2eccb636 100644 --- a/tests/unit/memory/test_findings.py +++ b/tests/unit/memory/test_findings.py @@ -411,6 +411,40 @@ async def test_one_attack_result_can_support_multiple_findings_and_counts(sqlite assert await sqlite_instance.get_finding_evidence_counts_async(finding_ids=[]) == {} +async def test_get_findings_for_attack_result_spans_operations_in_dashboard_order( + sqlite_instance: SQLiteMemory, +) -> None: + first = await sqlite_instance.add_operation_async(Operation(name="First")) + second = await sqlite_instance.add_operation_async(Operation(name="Second")) + low = await sqlite_instance.add_finding_async(finding=Finding(operation_id=first.id, title="Low", severity="low")) + critical = await sqlite_instance.add_finding_async( + finding=Finding(operation_id=second.id, title="Critical", severity="critical") + ) + unlinked = await sqlite_instance.add_finding_async( + finding=Finding(operation_id=first.id, title="Unlinked", severity="critical") + ) + attack_result_id, other_attack_id = uuid4(), uuid4() + evidence = { + finding.id: await sqlite_instance.add_finding_evidence_async( + evidence=FindingEvidence(finding_id=finding.id, attack_result_id=attack_result_id) + ) + for finding in (low, critical) + } + await sqlite_instance.add_finding_evidence_async( + evidence=FindingEvidence(finding_id=unlinked.id, attack_result_id=other_attack_id) + ) + + linked = await sqlite_instance.get_findings_for_attack_result_async(attack_result_id=attack_result_id) + + assert [(finding.id, item.id) for finding, item in linked] == [ + (critical.id, evidence[critical.id].id), + (low.id, evidence[low.id].id), + ] + assert {finding.operation_id for finding, _ in linked} == {first.id, second.id} + assert all(item.attack_result_id == attack_result_id for _, item in linked) + assert await sqlite_instance.get_findings_for_attack_result_async(attack_result_id=uuid4()) == [] + + async def test_get_findings_by_ids_follows_dashboard_order_and_scope(sqlite_instance: SQLiteMemory) -> None: operation = await sqlite_instance.add_operation_async(Operation(name="Selected")) other = await sqlite_instance.add_operation_async(Operation(name="Other")) diff --git a/tests/unit/memory/test_score_participation_migration.py b/tests/unit/memory/test_score_participation_migration.py index ba0bf556ce..69d9295d7c 100644 --- a/tests/unit/memory/test_score_participation_migration.py +++ b/tests/unit/memory/test_score_participation_migration.py @@ -60,7 +60,7 @@ def test_private_score_participation_has_unique_head_and_upgrades_current_parent config = Config() config.set_main_option("script_location", str(Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic")) scripts = ScriptDirectory.from_config(config) - assert scripts.get_heads() == ["e7c4a2d9f013"] + assert len(scripts.get_heads()) == 1 assert scripts.get_revision("e7c4a2d9f013").down_revision == "c8d3e5f7a901" engine = sa.create_engine("sqlite:///:memory:") try: From 658e94704a124726d5b53fd477a3df2f5d97f3e1 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 6 Nov 2026 20:36:37 -0800 Subject: [PATCH 2/3] DOC: Clarify running attack response state fields Limit running-state guarantees to attack summary and metadata reads without expanding conversation/message schemas or removing score details. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index a7dbf409de..9c0cce11bf 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -117,9 +117,11 @@ and transcript reads, including hierarchy node expansion. A scan saves each `AttackResult` only when its attack finishes. While an active scan is still running an attack, `GET /api/attacks/{attack_result_id}`, `/metadata`, `/conversations`, and `/messages` (paged or not) resolve it from the live run instead -of returning 404. These responses set `in_progress: true`, report outcome -`undetermined`, and set `execution_read_only: true`; write routes still require a -saved result. The conversation the attack is working on is the main conversation, +of returning 404. The attack summary and metadata responses set `in_progress: true`, +report outcome `undetermined`, and set `execution_read_only: true`. Read these +endpoints for running state; conversation and message responses do not include +the running-state flags. Write routes still require a saved result. The +conversation the attack is working on is the main conversation, and every other objective conversation started so far is listed as a pruned reference. Messages are read from memory as they are saved. A saved result always wins, so once the attack finishes the same IDs resolve to its saved result; restart From 3a6dcfee88f8ac3e64bea51d88be31f3446764e6 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sat, 7 Nov 2026 00:36:03 -0800 Subject: [PATCH 3/3] FIX: Merge finding index and HTTP fingerprint migration heads Preserve both published revision histories and restore singular-head backend startup. Cover fresh, common-parent, either-branch and dual-head upgrades with preserved data, idempotency, downgrade/re-upgrade and SQLite/SQL Server merge semantics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- ...erge_finding_index_and_http_fingerprint.py | 25 ++++ ...finding_evidence_attack_index_migration.py | 137 +++++++++++++++++- 2 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 pyrit/memory/alembic/versions/f8a6c2d4e901_merge_finding_index_and_http_fingerprint.py diff --git a/pyrit/memory/alembic/versions/f8a6c2d4e901_merge_finding_index_and_http_fingerprint.py b/pyrit/memory/alembic/versions/f8a6c2d4e901_merge_finding_index_and_http_fingerprint.py new file mode 100644 index 0000000000..eb876aeee3 --- /dev/null +++ b/pyrit/memory/alembic/versions/f8a6c2d4e901_merge_finding_index_and_http_fingerprint.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Merge the finding-evidence index and HTTP request fingerprint heads. + +Revision ID: f8a6c2d4e901 +Revises: b5d7f9a1c3e6, d6de1045cca4 +Create Date: 2026-11-07 08:40:00.000000 +""" + +from collections.abc import Sequence + +revision: str = "f8a6c2d4e901" +down_revision: str | Sequence[str] | None = ("b5d7f9a1c3e6", "d6de1045cca4") +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Apply this schema upgrade.""" + + +def downgrade() -> None: + """Revert this schema upgrade.""" diff --git a/tests/unit/memory/test_finding_evidence_attack_index_migration.py b/tests/unit/memory/test_finding_evidence_attack_index_migration.py index 01717beb37..3bdfad7cf8 100644 --- a/tests/unit/memory/test_finding_evidence_attack_index_migration.py +++ b/tests/unit/memory/test_finding_evidence_attack_index_migration.py @@ -2,9 +2,12 @@ # Licensed under the MIT license. import importlib +from datetime import UTC, datetime from io import StringIO from pathlib import Path +from uuid import uuid4 +import pytest import sqlalchemy as sa from alembic import command from alembic.config import Config @@ -12,11 +15,13 @@ from alembic.operations import Operations from alembic.script import ScriptDirectory -from pyrit.memory.memory_models import FindingEvidenceEntry -from pyrit.memory.migration import check_schema_migrations +from pyrit.memory.memory_models import FindingEntry, FindingEvidenceEntry, OperationEntry +from pyrit.memory.migration import check_schema_migrations, run_schema_migrations PREVIOUS_REVISION = "e7c4a2d9f013" REVISION = "b5d7f9a1c3e6" +HTTP_REVISION = "d6de1045cca4" +MERGE_REVISION = "f8a6c2d4e901" INDEX = "ix_FindingEvidenceEntries_attack_result_id" @@ -36,10 +41,12 @@ def _evidence_indexes(connection: sa.Connection) -> dict[str, list[str]]: } -def test_migration_is_the_single_head_after_score_participation() -> None: +def test_migration_branches_have_one_merge_head_without_rewriting_history() -> None: scripts = ScriptDirectory.from_config(_config()) - assert scripts.get_heads() == [REVISION] + assert scripts.get_heads() == [MERGE_REVISION] + assert set(scripts.get_revision(MERGE_REVISION).down_revision) == {REVISION, HTTP_REVISION} assert scripts.get_revision(REVISION).down_revision == PREVIOUS_REVISION + assert scripts.get_revision(HTTP_REVISION).down_revision == PREVIOUS_REVISION def test_migration_adds_and_drops_attack_result_index() -> None: @@ -75,3 +82,125 @@ def test_model_declares_attack_result_index() -> None: index.name: [column.name for column in index.columns] for index in FindingEvidenceEntry.__table__.indexes } assert indexes[INDEX] == ["attack_result_id"] + + +def _assert_merged_schema(connection: sa.Connection) -> None: + context = MigrationContext.configure(connection, opts={"version_table": "pyrit_memory_alembic_version"}) + assert context.get_current_heads() == (MERGE_REVISION,) + assert _evidence_indexes(connection)[INDEX] == ["attack_result_id"] + columns = {column["name"]: column for column in sa.inspect(connection).get_columns("TargetIdentifiers")} + assert columns["http_request_fingerprint"]["nullable"] is True + + +def _seed_upgrade_rows(connection: sa.Connection) -> None: + operation_id, finding_id = uuid4(), uuid4() + timestamp = datetime.now(UTC) + connection.execute( + sa.insert(OperationEntry).values( + id=operation_id, name="Preserved", name_key="preserved", created_at=timestamp, description="" + ) + ) + connection.execute( + sa.insert(FindingEntry).values( + id=finding_id, + operation_id=operation_id, + title="Preserved finding", + description="", + severity="important", + created_at=timestamp, + ) + ) + connection.execute( + sa.insert(FindingEvidenceEntry).values( + id=uuid4(), finding_id=finding_id, attack_result_id=uuid4(), attached_at=timestamp + ) + ) + targets = sa.Table("TargetIdentifiers", sa.MetaData(), autoload_with=connection) + connection.execute( + targets.insert().values(hash="a" * 64, class_name="HTTPTarget", identifier_json={"marker": "retained"}) + ) + + +@pytest.mark.parametrize( + "revisions", + [(PREVIOUS_REVISION,), (REVISION,), (HTTP_REVISION,), (REVISION, HTTP_REVISION)], + ids=["common-parent", "finding-index", "http-fingerprint", "both-branches"], +) +def test_startup_upgrades_either_migration_branch_without_losing_data(revisions: tuple[str, ...]) -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + config = _config(connection) + command.upgrade(config, PREVIOUS_REVISION) + _seed_upgrade_rows(connection) + evidence = dict(connection.execute(sa.select(FindingEvidenceEntry)).mappings().one()) + for revision in revisions: + command.upgrade(config, revision) + context = MigrationContext.configure(connection, opts={"version_table": "pyrit_memory_alembic_version"}) + assert set(context.get_current_heads()) == set(revisions) + fingerprint = "recorded-fingerprint" if HTTP_REVISION in revisions else None + if fingerprint is not None: + connection.execute( + sa.text("UPDATE TargetIdentifiers SET http_request_fingerprint = :value"), + {"value": fingerprint}, + ) + + run_schema_migrations(engine=engine, silent=True) + run_schema_migrations(engine=engine, silent=True) + with engine.connect() as connection: + _assert_merged_schema(connection) + assert dict(connection.execute(sa.select(FindingEvidenceEntry)).mappings().one()) == evidence + assert connection.scalar(sa.select(FindingEntry.title)) == "Preserved finding" + assert connection.scalar(sa.select(OperationEntry.name)) == "Preserved" + targets = sa.Table("TargetIdentifiers", sa.MetaData(), autoload_with=connection) + row = connection.execute(sa.select(targets)).mappings().one() + assert row["hash"] == "a" * 64 + assert row["identifier_json"] == {"marker": "retained"} + assert row["http_request_fingerprint"] == fingerprint + check_schema_migrations(engine=engine) + finally: + engine.dispose() + + +def test_startup_creates_merged_schema_from_empty_database() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + run_schema_migrations(engine=engine, silent=True) + with engine.connect() as connection: + _assert_merged_schema(connection) + check_schema_migrations(engine=engine) + finally: + engine.dispose() + + +def test_merge_downgrades_both_branches_and_upgrades_again() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + config = _config(connection) + command.upgrade(config, "head") + _assert_merged_schema(connection) + command.downgrade(config, PREVIOUS_REVISION) + assert INDEX not in _evidence_indexes(connection) + assert "http_request_fingerprint" not in { + column["name"] for column in sa.inspect(connection).get_columns("TargetIdentifiers") + } + run_schema_migrations(engine=engine, silent=True) + with engine.connect() as connection: + _assert_merged_schema(connection) + check_schema_migrations(engine=engine) + finally: + engine.dispose() + + +@pytest.mark.parametrize("dialect_name", ["sqlite", "mssql"]) +def test_merge_revision_adds_no_dialect_specific_ddl(dialect_name: str) -> None: + migration = importlib.import_module( + "pyrit.memory.alembic.versions.f8a6c2d4e901_merge_finding_index_and_http_fingerprint" + ) + output = StringIO() + context = MigrationContext.configure(dialect_name=dialect_name, opts={"as_sql": True, "output_buffer": output}) + with Operations.context(context): + migration.upgrade() + migration.downgrade() + assert output.getvalue() == ""