From a9bc7de6da575da0972adb1ef78f3c0dd1a51cb1 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Wed, 21 Oct 2026 18:39:40 -0400 Subject: [PATCH 01/20] Sort and count findings without FindingSeverity in memory Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/models/operations.py | 5 +- pyrit/backend/services/operation_service.py | 13 ++++- pyrit/memory/memory_interface.py | 45 +++++++++++---- tests/unit/backend/test_finding_export.py | 2 +- tests/unit/backend/test_operations.py | 9 +++ tests/unit/memory/test_findings.py | 63 +++++++++++++++++---- 6 files changed, 109 insertions(+), 28 deletions(-) diff --git a/pyrit/backend/models/operations.py b/pyrit/backend/models/operations.py index 727b13f081..4805101c5c 100644 --- a/pyrit/backend/models/operations.py +++ b/pyrit/backend/models/operations.py @@ -12,7 +12,6 @@ Finding, FindingCreate, FindingEvidence, - FindingSeverity, Operation, OperationUsage, ) @@ -27,9 +26,9 @@ class FindingOptionsResponse(BaseModel): class OperationListItem(Operation): - """An operation with its finding counts by severity. Severities with no findings are omitted.""" + """An operation with its finding counts by stored severity value. Severities with no findings are omitted.""" - finding_counts: dict[FindingSeverity, int] + finding_counts: dict[str, int] class OperationListResponse(BaseModel): diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py index 8fa55438f0..79e3889348 100644 --- a/pyrit/backend/services/operation_service.py +++ b/pyrit/backend/services/operation_service.py @@ -25,6 +25,7 @@ Finding, FindingCreate, FindingEvidence, + FindingSeverity, Operation, OperationCreate, ) @@ -196,7 +197,11 @@ async def list_findings_async( """ await self.get_async(operation_id) findings = await CentralMemory.get_memory_instance().get_findings_async( - operation_id=operation_id, limit=limit + 1, offset=offset, title_query=title_query + operation_id=operation_id, + limit=limit + 1, + offset=offset, + title_query=title_query, + severity_order=[severity.value for severity in FindingSeverity], ) has_more = len(findings) > limit page = findings[:limit] @@ -269,7 +274,11 @@ async def get_findings_with_evidence_async( """ await self.get_async(operation_id) memory = CentralMemory.get_memory_instance() - findings = await memory.get_findings_by_ids_async(operation_id=operation_id, finding_ids=finding_ids) + findings = await memory.get_findings_by_ids_async( + operation_id=operation_id, + finding_ids=finding_ids, + severity_order=[severity.value for severity in FindingSeverity], + ) found = {finding.id for finding in findings} missing = next((finding_id for finding_id in finding_ids if finding_id not in found), None) if missing is not None: diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index e230165ae0..3596046b3f 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -120,7 +120,6 @@ Finding, FindingCreate, FindingEvidence, - FindingSeverity, IdentifierFilter, IdentifierType, Message, @@ -1044,10 +1043,21 @@ async def get_findings_async( limit: int = 20, offset: int = 0, title_query: str | None = None, + severity_order: Sequence[str] = (), ) -> list[Finding]: """ Read one page of findings, highest severity first, then newest first. + Args: + operation_id (uuid.UUID | None): The operation whose findings to read. Defaults to None, which reads + findings from every operation. + limit (int): Maximum number of findings to return. Defaults to 20. + offset (int): Number of findings to skip. Defaults to 0. + title_query (str | None): Case-insensitive text the title must contain. Wildcard characters match + literally. Defaults to None. + severity_order (Sequence[str]): Severity keys from highest to lowest. Stored values that aren't listed + sort after them. Empty orders newest first. + Returns: list[Finding]: The findings on the page. @@ -1063,18 +1073,24 @@ async def get_findings_async( # SQL Server also treats "[" as a LIKE wildcard. escaped_title = title_query.replace("/", "//").replace("%", "/%").replace("_", "/_").replace("[", "/[") statement = statement.where(FindingEntry.title.ilike(f"%{escaped_title}%", escape="/")) - statement = statement.order_by(*self._finding_order()).limit(limit).offset(offset) + statement = statement.order_by(*self._finding_order(severity_order)).limit(limit).offset(offset) async with await self.get_session_async() as session: return [entry.get_finding() for entry in (await session.scalars(statement)).all()] async def get_findings_by_ids_async( - self, *, operation_id: uuid.UUID, finding_ids: Sequence[uuid.UUID] + self, *, operation_id: uuid.UUID, finding_ids: Sequence[uuid.UUID], severity_order: Sequence[str] = () ) -> list[Finding]: """ Read selected findings of one operation, highest severity first, then newest first. IDs that are not in the operation are omitted. + Args: + operation_id (uuid.UUID): The operation whose findings to read. + finding_ids (Sequence[uuid.UUID]): The findings to read. + severity_order (Sequence[str]): Severity keys from highest to lowest, as for + ``get_findings_async``. Defaults to an empty order, which sorts newest first. + Returns: list[Finding]: The selected findings, in the same order as ``get_findings_async``. """ @@ -1087,7 +1103,7 @@ async def get_findings_by_ids_async( ordered_ids = await session.scalars( select(FindingEntry.id) .where(FindingEntry.operation_id == operation_id) - .order_by(*self._finding_order()) + .order_by(*self._finding_order(severity_order)) ) selected_ids = [finding_id for finding_id in ordered_ids if finding_id in requested] findings: dict[uuid.UUID, Finding] = {} @@ -1251,29 +1267,34 @@ async def get_finding_evidence_counts_async(self, *, finding_ids: Sequence[uuid. async with await self.get_session_async() as session: return dict((await session.execute(statement)).all()) - async def get_finding_severity_counts_async(self) -> dict[uuid.UUID, dict[FindingSeverity, int]]: + async def get_finding_severity_counts_async(self) -> dict[uuid.UUID, dict[str, int]]: """ Count every operation's findings by severity in one query. Returns: - dict[uuid.UUID, dict[FindingSeverity, int]]: Counts for operations with findings. + dict[uuid.UUID, dict[str, int]]: Counts for operations with findings. """ statement = select(FindingEntry.operation_id, FindingEntry.severity, func.count()).group_by( FindingEntry.operation_id, FindingEntry.severity ) - counts: dict[uuid.UUID, dict[FindingSeverity, int]] = {} + counts: dict[uuid.UUID, dict[str, int]] = {} async with await self.get_session_async() as session: for operation_id, severity, count in (await session.execute(statement)).all(): - counts.setdefault(operation_id, {})[FindingSeverity(severity)] = count + counts.setdefault(operation_id, {})[severity] = count return counts @staticmethod - def _finding_order() -> "tuple[ColumnElement[Any], ...]": - severity_order = case( - {severity.value: position for position, severity in enumerate(FindingSeverity)}, + def _finding_order(severity_order: Sequence[str]) -> "tuple[ColumnElement[Any], ...]": + newest_first = (FindingEntry.created_at.desc(), FindingEntry.id.desc()) + if not severity_order: + # case() with no branches compiles to SQL that both dialects reject. + return newest_first + rank = case( + {key: position for position, key in enumerate(severity_order)}, value=FindingEntry.severity, + else_=len(severity_order), ) - return severity_order, FindingEntry.created_at.desc(), FindingEntry.id.desc() + return rank, *newest_first def _uses_legacy_memory_override(self) -> bool: return any( diff --git a/tests/unit/backend/test_finding_export.py b/tests/unit/backend/test_finding_export.py index e78411741c..1b8ea16598 100644 --- a/tests/unit/backend/test_finding_export.py +++ b/tests/unit/backend/test_finding_export.py @@ -543,7 +543,7 @@ async def test_report_matches_the_dashboard_listing_order(sqlite_instance: SQLit ) for index, severity in enumerate(["low", "critical", "moderate", "critical"]) ] - dashboard = await sqlite_instance.get_findings_async(operation_id=operation.id, limit=10) + dashboard = (await OperationService().list_findings_async(operation_id=operation.id, limit=10, offset=0)).items report = await _report_async(operation_id=operation.id, finding_ids=[finding.id for finding in findings]) diff --git a/tests/unit/backend/test_operations.py b/tests/unit/backend/test_operations.py index 2a907c9431..b87fd81a66 100644 --- a/tests/unit/backend/test_operations.py +++ b/tests/unit/backend/test_operations.py @@ -147,6 +147,15 @@ def test_findings_are_created_and_listed_under_an_operation(compatibility_header assert counts == {"Case": {"informational": 1}, "Other": {"low": 1}} +def test_findings_list_highest_severity_first(compatibility_headers: dict[str, str]) -> None: + client = TestClient(app, headers=compatibility_headers) + operation_id = client.post("/api/operations", json={"name": "Sorted"}).json()["id"] + for title, severity in [("First", "critical"), ("Second", "low")]: + client.post(f"/api/operations/{operation_id}/findings", json={"title": title, "severity": severity}) + listed = client.get(f"/api/operations/{operation_id}/findings", params={"limit": 20}).json()["items"] + assert [item["title"] for item in listed] == ["First", "Second"] + + def test_finding_options_exposes_canonical_harm_types(compatibility_headers: dict[str, str]) -> None: client = TestClient(app, headers=compatibility_headers) response = client.get("/api/operations/finding-options") diff --git a/tests/unit/memory/test_findings.py b/tests/unit/memory/test_findings.py index 94d45c0e7a..45958f993a 100644 --- a/tests/unit/memory/test_findings.py +++ b/tests/unit/memory/test_findings.py @@ -24,11 +24,13 @@ PromptMemoryEntry, ScoreEntry, ) -from pyrit.models import Finding, FindingCreate, FindingEvidence, FindingSeverity, Operation +from pyrit.models import Finding, FindingCreate, FindingEvidence, Operation from pyrit.models.harm_category import HarmCategory pytestmark = pytest.mark.usefixtures("patch_central_database") +SEVERITY_ORDER = ["critical", "important", "moderate", "low", "informational"] + async def test_operations_are_unique_by_trimmed_case_folded_name(sqlite_instance: SQLiteMemory) -> None: saved = await sqlite_instance.add_operation_async(Operation(name=" Operation A ")) @@ -164,8 +166,10 @@ async def test_findings_paginate_within_one_operation(sqlite_instance: SQLiteMem ) ) await sqlite_instance.add_finding_async(finding=Finding(operation_id=other.id, title="Other", severity="critical")) - first = await sqlite_instance.get_findings_async(operation_id=operation.id, limit=2) - second = await sqlite_instance.get_findings_async(operation_id=operation.id, limit=2, offset=2) + first = await sqlite_instance.get_findings_async(operation_id=operation.id, limit=2, severity_order=SEVERITY_ORDER) + second = await sqlite_instance.get_findings_async( + operation_id=operation.id, limit=2, offset=2, severity_order=SEVERITY_ORDER + ) assert [finding.title for finding in first + second] == ["3", "2", "1", "4"] async with await sqlite_instance.get_session_async() as session: for table in [AttackResultEntry, PromptMemoryEntry, ScoreEntry]: @@ -218,7 +222,10 @@ async def test_update_finding_preserves_identity_and_reorders_severity(sqlite_in assert updated.created_at == original.created_at assert updated.title == "Edited" assert updated.description == "Notes" - assert [finding.title for finding in await sqlite_instance.get_findings_async()] == ["Edited", "Other"] + assert [finding.title for finding in await sqlite_instance.get_findings_async(severity_order=SEVERITY_ORDER)] == [ + "Edited", + "Other", + ] async def test_finding_classifications_persist_and_clear_without_losing_evidence(sqlite_instance: SQLiteMemory) -> None: @@ -272,7 +279,7 @@ async def test_custom_severity_sorts_after_every_preset(sqlite_instance: SQLiteM } ) ) - assert [finding.title for finding in await sqlite_instance.get_findings_async()] == [ + assert [finding.title for finding in await sqlite_instance.get_findings_async(severity_order=SEVERITY_ORDER)] == [ "critical", "important", "moderate", @@ -282,6 +289,41 @@ async def test_custom_severity_sorts_after_every_preset(sqlite_instance: SQLiteM ] +async def test_get_findings_sorts_unlisted_severities_after_listed_ones(sqlite_instance: SQLiteMemory) -> None: + operation = await sqlite_instance.add_operation_async(Operation(name="Partial order")) + stamp = datetime(2026, 10, 9, tzinfo=UTC) + for number, severity in enumerate(["low", "critical", "moderate"], start=1): + await sqlite_instance.add_finding_async( + finding=Finding.model_validate( + { + "id": UUID(int=number), + "created_at": stamp, + "operation_id": operation.id, + "title": severity, + "severity": severity, + } + ) + ) + found = await sqlite_instance.get_findings_async(severity_order=["low"]) + assert [finding.title for finding in found] == ["low", "moderate", "critical"] + + +async def test_get_findings_without_severity_order_sorts_newest_first(sqlite_instance: SQLiteMemory) -> None: + operation = await sqlite_instance.add_operation_async(Operation(name="Newest first")) + for number, (title, severity) in enumerate([("Old critical", "critical"), ("New low", "low")], start=1): + await sqlite_instance.add_finding_async( + finding=Finding.model_validate( + { + "operation_id": operation.id, + "title": title, + "severity": severity, + "created_at": datetime(2026, 10, number, tzinfo=UTC), + } + ) + ) + assert [finding.title for finding in await sqlite_instance.get_findings_async()] == ["New low", "Old critical"] + + async def test_finding_mutations_are_scoped_and_delete_only_the_finding(sqlite_instance: SQLiteMemory) -> None: operation = await sqlite_instance.add_operation_async(Operation(name="Case")) other = await sqlite_instance.add_operation_async(Operation(name="Other")) @@ -430,7 +472,9 @@ async def test_get_findings_by_ids_follows_dashboard_order_and_scope(sqlite_inst ) requested = [saved["Moderate"].id, foreign.id, saved["Low"].id, uuid4(), saved["Newer critical"].id] requested += [saved["Older critical"].id, saved["Low"].id] - findings = await sqlite_instance.get_findings_by_ids_async(operation_id=operation.id, finding_ids=requested) + findings = await sqlite_instance.get_findings_by_ids_async( + operation_id=operation.id, finding_ids=requested, severity_order=SEVERITY_ORDER + ) assert [finding.title for finding in findings] == ["Newer critical", "Older critical", "Moderate", "Low"] @@ -494,10 +538,9 @@ async def test_finding_severity_counts_group_by_operation(sqlite_instance: SQLit finding=Finding(operation_id=operation.id, title="A", severity=severity) ) await sqlite_instance.add_finding_async(finding=Finding(operation_id=other.id, title="B", severity="informational")) - assert await sqlite_instance.get_finding_severity_counts_async() == { - operation.id: {FindingSeverity.CRITICAL: 2, FindingSeverity.LOW: 1}, - other.id: {FindingSeverity.INFORMATIONAL: 1}, - } + result = await sqlite_instance.get_finding_severity_counts_async() + assert result == {operation.id: {"critical": 2, "low": 1}, other.id: {"informational": 1}} + assert all(type(key) is str for counts in result.values() for key in counts) async def test_delete_finding_removes_evidence_without_deleting_source(sqlite_instance: SQLiteMemory) -> None: From 8850791e09706f3097fd977eb886b33e3bbb4b4d Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Sat, 24 Oct 2026 13:14:50 -0400 Subject: [PATCH 02/20] Add FindingFields models and finding value check Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/models/__init__.py | 10 + pyrit/models/finding_fields.py | 551 +++++++++++++++++++++++ tests/unit/models/test_finding_fields.py | 320 +++++++++++++ 3 files changed, 881 insertions(+) create mode 100644 pyrit/models/finding_fields.py create mode 100644 tests/unit/models/test_finding_fields.py diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index f45b928039..9415a50a99 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -74,6 +74,12 @@ from pyrit.models.conversation_summary import ConversationSummary from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation from pyrit.models.finding import Finding, FindingCreate, FindingEvidence, FindingSeverity + from pyrit.models.finding_fields import ( + DEFAULT_FINDING_FIELDS, + FindingField, + FindingFields, + FindingValuesError, + ) from pyrit.models.harm_definition import HarmDefinition, ScaleDescription, get_all_harm_definitions from pyrit.models.identifiers import ( REGISTRY_NAME_PATTERN, @@ -444,6 +450,10 @@ "FindingCreate": "pyrit.models.finding", "FindingEvidence": "pyrit.models.finding", "FindingSeverity": "pyrit.models.finding", + "DEFAULT_FINDING_FIELDS": "pyrit.models.finding_fields", + "FindingField": "pyrit.models.finding_fields", + "FindingFields": "pyrit.models.finding_fields", + "FindingValuesError": "pyrit.models.finding_fields", "ConversationType": "pyrit.models.messages.conversation_reference", "ContentEntryScorable": "pyrit.models.score", "ContentScorable": "pyrit.models.score", diff --git a/pyrit/models/finding_fields.py b/pyrit/models/finding_fields.py new file mode 100644 index 0000000000..ea789cf10a --- /dev/null +++ b/pyrit/models/finding_fields.py @@ -0,0 +1,551 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import re +from collections.abc import Iterable, Iterator +from enum import StrEnum +from typing import Any, ClassVar, Self + +from pydantic import BaseModel, ConfigDict, field_validator, model_validator + +from pyrit.models.harm_category import HarmCategory + + +class FindingFieldType(StrEnum): + """How the finding dialog collects a field's value.""" + + TEXT = "text" + MULTILINE = "multiline" + CHOICE = "choice" + + +class FindingSeverityColor(StrEnum): + """Badge color of a severity option. CoPyRIT renders each one with a Fluent UI badge color.""" + + RED = "red" + ORANGE = "orange" + YELLOW = "yellow" + GREEN = "green" + BLUE = "blue" + GRAY = "gray" + DARK = "dark" + NEUTRAL = "neutral" + + +class FindingValuesError(ValueError): + """A finding's values do not fit the deployment's finding fields.""" + + +class FindingFieldOption(BaseModel): + """One option of a choice field, with the options of the next dropdown under it.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + key: str + label: str + children: list["FindingFieldOption"] = [] + + +class FindingField(BaseModel): + """ + One field of the finding dialog, as a deployment defines it. + + Attributes: + key (str): The stored name of the value. Keep it fixed; change ``label`` to rename what users see. + label (str): The name the finding dialog and cards show. + type (FindingFieldType): ``text``, ``multiline`` or ``choice``. + required (bool): Whether every finding needs a value. For a cascade, every level with options needs one. + options (list[FindingFieldOption]): A ``choice`` field's options. Write ``(key, label)`` pairs for one + dropdown, or a dict from each pair to its children for a cascade. + levels (list[tuple[str, str]]): The ``(key, label)`` of each dropdown under a cascade's first one. + Each level's value is stored under its own key. + default (str | None): A top-level option the finding dialog preselects for new findings. + allow_other (bool): Whether a single-level ``choice`` field also accepts typed text. + colors (dict[str, FindingSeverityColor]): The severity slot's badge color for each top-level option. + hint (str | None): Help text under the field in the finding dialog. + section (str | None): A heading the finding dialog shows where the section changes. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + _KEY_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r"[a-z0-9_]+") + + key: str + label: str + type: FindingFieldType + required: bool = False + options: list[FindingFieldOption] = [] + levels: list[tuple[str, str]] = [] + default: str | None = None + allow_other: bool = False + colors: dict[str, FindingSeverityColor] = {} + hint: str | None = None + section: str | None = None + + @field_validator("options", mode="before") + @classmethod + def _normalize_options(cls, value: Any) -> Any: + """ + Turn ``(key, label)`` pairs and cascade dicts into option dicts. + + Returns: + Any: The options in a form pydantic validates as ``FindingFieldOption`` objects. + """ + return _parse_options(value) + + @model_validator(mode="after") + def _validate_definition(self) -> Self: + """ + Check keys, labels and the settings that the field's type allows. + + Returns: + Self: The validated field. + """ + self._check_key(self.key) + self._check_not_blank() + for level_key, _ in self.levels: + self._check_key(level_key) + if self.type == FindingFieldType.CHOICE: + self._check_choice() + else: + self._check_no_choice_settings() + return self + + def _check_key(self, key: str) -> None: + """ + Require a key of lowercase letters, digits and underscores. + + Args: + key (str): The field or level key. + + Raises: + ValueError: If the key uses any other character or is empty. + """ + if not self._KEY_PATTERN.fullmatch(key): + raise ValueError(f"Key '{key}' must use only lowercase letters, digits and underscores.") + + def _check_not_blank(self) -> None: + """ + Require non-blank labels, option keys and option labels. + + Raises: + ValueError: If any of them is blank. + """ + texts = [self.label, *(label for _, label in self.levels)] + for option in _walk(self.options): + texts += [option.key, option.label] + if any(not text.strip() for text in texts): + raise ValueError(f"Labels, option keys and option labels of '{self.key}' must not be blank.") + + def _check_no_choice_settings(self) -> None: + """ + Reject the settings that only a choice field can use. + + Raises: + ValueError: If the field sets options, levels, default, allow_other or colors. + """ + settings = { + "options": bool(self.options), + "levels": bool(self.levels), + "default": self.default is not None, + "allow_other": self.allow_other, + "colors": bool(self.colors), + } + for name, is_set in settings.items(): + if is_set: + raise ValueError(f"Field '{self.key}' is a {self.type} field and cannot have {name}.") + + def _check_choice(self) -> None: + """ + Check the option tree of a choice field and the settings that refer to it. + + Raises: + ValueError: If the options are missing, repeat a sibling key, nest to a depth that + ``levels`` does not match, or ``allow_other`` is combined with levels. + """ + if not self.options: + raise ValueError(f"Choice field '{self.key}' needs at least one option.") + self._check_unique_siblings() + depth = _depth(self.options) + if depth != len(self.levels) + 1: + raise ValueError( + f"Field '{self.key}' needs one entry in levels for each dropdown below the field; " + f"its options imply {depth - 1}, but levels has {len(self.levels)}." + ) + self._check_top_level_references() + if self.allow_other and self.levels: + raise ValueError(f"Field '{self.key}' can take Other text only without levels.") + + def _check_unique_siblings(self) -> None: + """ + Require unique option keys among the options of each dropdown. + + Raises: + ValueError: If two sibling options share a key. + """ + for siblings in [self.options, *(option.children for option in _walk(self.options))]: + repeated = _first_repeated(option.key for option in siblings) + if repeated is not None: + raise ValueError(f"Field '{self.key}' repeats the option key '{repeated}' among sibling options.") + + def _check_top_level_references(self) -> None: + """ + Require ``default`` and every ``colors`` key to name a top-level option. + + Raises: + ValueError: If either names anything else. + """ + top_level = {option.key for option in self.options} + defaults = [] if self.default is None else [self.default] + for key in [*defaults, *self.colors]: + if key not in top_level: + raise ValueError(f"'{key}' is not a top-level option of '{self.key}'.") + + +class FindingFields(BaseModel): + """A deployment's finding fields: the severity slot and the fields shown after it.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + severity: FindingField + fields: list[FindingField] = [] + + @property + def severity_order(self) -> list[str]: + """The top-level option keys of the severity slot, from highest to lowest severity.""" + return [option.key for option in self.severity.options] + + @field_validator("severity") + @classmethod + def _validate_severity_slot(cls, value: FindingField) -> FindingField: + """ + Check the severity slot and mark it required. + + Returns: + FindingField: The slot with ``required`` set to True. + + Raises: + ValueError: If the slot's key is not ``severity``, it is not a choice field, or it sets + ``required=False``. + """ + if value.key != "severity": + raise ValueError(f"The severity slot's key must be 'severity', not '{value.key}'.") + if value.type != FindingFieldType.CHOICE: + raise ValueError("The severity slot must be a choice field.") + if "required" in value.model_fields_set and not value.required: + raise ValueError("The severity slot is always required; leave required out or set it to True.") + return value.model_copy(update={"required": True}) + + @model_validator(mode="after") + def _validate_keys(self) -> Self: + """ + Check colors and keys across the severity slot, the fields and their levels. + + Returns: + Self: The validated fields. + + Raises: + ValueError: If a field other than severity sets colors, a key is used more than once, + or a key is ``title``. + """ + colored = next((field.key for field in self.fields if field.colors), None) + if colored is not None: + raise ValueError(f"Only the severity field can set colors, not '{colored}'.") + keys = [key for field in [self.severity, *self.fields] for key in _keys_of(field)] + repeated = _first_repeated(keys) + if repeated is not None: + raise ValueError(f"The key '{repeated}' is used more than once across fields and levels.") + if "title" in keys: + raise ValueError("The key 'title' is reserved for the finding title.") + return self + + def check_values( + self, + *, + severity: str, + field_values: dict[str, str], + previous_severity: str | None = None, + previous_values: dict[str, str] | None = None, + ) -> None: + """ + Check a finding's severity and field values against these fields. + + On an edit, pass the stored finding's severity and values. A value that has not changed + since then may name an option or key that no longer exists. + + Args: + severity (str): The severity slot's value. + field_values (dict[str, str]): The values of the other fields and of every level, by key. + previous_severity (str | None): The stored finding's severity, or None for a new finding. + previous_values (dict[str, str] | None): The stored finding's field values, or None for a + new finding. + + Raises: + FindingValuesError: For the first unknown key, missing required value, or value that is + not an option. + """ + previous_values = previous_values or {} + self._check_known_keys(field_values=field_values, previous_values=previous_values) + values = {**field_values, "severity": severity} + previous = {**previous_values, "severity": previous_severity or ""} + for field in [self.severity, *self.fields]: + _check_field_values(field=field, values=values, previous_values=previous) + + def _check_known_keys(self, *, field_values: dict[str, str], previous_values: dict[str, str]) -> None: + """ + Reject a key that no field or level uses, unless its value has not changed. + + Args: + field_values (dict[str, str]): The values to check, by key. + previous_values (dict[str, str]): The stored finding's field values. + + Raises: + FindingValuesError: If a key is unknown and its value differs from the stored one. + """ + known = {key for field in [self.severity, *self.fields] for key in _keys_of(field)} - {"severity"} + for key, value in field_values.items(): + if key not in known and previous_values.get(key) != value: + raise FindingValuesError(f"Unknown finding field {key!r}.") + + +def _parse_options(value: Any) -> Any: + """ + Turn authoring shorthand for options into a list of option dicts. + + A dict maps each ``(key, label)`` pair to its children, written the same way. In a list, a + ``(key, label)`` pair becomes an option without children, and anything else passes through. + + Args: + value (Any): The options as written. + + Returns: + Any: The options as a list, or ``value`` unchanged if it is neither a dict nor a list. + """ + if isinstance(value, dict): + return [_parse_branch(pair=pair, children=children) for pair, children in value.items()] + if isinstance(value, list): + return [{"key": item[0], "label": item[1]} if _is_pair(item) else item for item in value] + return value + + +def _parse_branch(*, pair: Any, children: Any) -> dict[str, Any]: + """ + Turn one cascade dict entry into an option dict with its children. + + Args: + pair (Any): The entry's key, which must be a ``(key, label)`` pair. + children (Any): The options under it, as a list or a dict. + + Returns: + dict[str, Any]: The option with its parsed children. + + Raises: + ValueError: If ``pair`` is not a ``(key, label)`` pair. + """ + if not _is_pair(pair): + raise ValueError(f"Option {pair!r} must be a (key, label) pair.") + return {"key": pair[0], "label": pair[1], "children": _parse_options(children)} + + +def _is_pair(value: Any) -> bool: + """ + Check whether a value is a 2-item tuple or list of strings. + + Args: + value (Any): The value to check. + + Returns: + bool: True if the value is a ``(key, label)`` pair. + """ + return isinstance(value, tuple | list) and len(value) == 2 and all(isinstance(part, str) for part in value) + + +def _walk(options: list[FindingFieldOption]) -> Iterator[FindingFieldOption]: + """ + Yield every option in a tree, each parent before its children. + + Args: + options (list[FindingFieldOption]): The top-level options. + + Yields: + FindingFieldOption: Each option in the tree. + """ + for option in options: + yield option + yield from _walk(option.children) + + +def _depth(options: list[FindingFieldOption]) -> int: + """ + Count the dropdowns that the deepest branch of an option tree needs. + + Args: + options (list[FindingFieldOption]): The top-level options. + + Returns: + int: The depth of the tree, or 0 for no options. + """ + if not options: + return 0 + return 1 + max(_depth(option.children) for option in options) + + +def _first_repeated(keys: Iterable[str]) -> str | None: + """ + Find the first key that appears a second time. + + Args: + keys (Iterable[str]): The keys in written order. + + Returns: + str | None: The first repeated key, or None if all keys are unique. + """ + seen: set[str] = set() + for key in keys: + if key in seen: + return key + seen.add(key) + return None + + +def _keys_of(field: FindingField) -> list[str]: + """ + List a field's key followed by its level keys. + + Args: + field (FindingField): The field. + + Returns: + list[str]: The field key and level keys in written order. + """ + return [field.key, *(level_key for level_key, _ in field.levels)] + + +def _check_field_values(*, field: FindingField, values: dict[str, str], previous_values: dict[str, str]) -> None: + """ + Check the value of one field and of its levels. + + Args: + field (FindingField): The field. + values (dict[str, str]): Every value of the finding, by key. + previous_values (dict[str, str]): Every value of the stored finding, by key. + + Raises: + FindingValuesError: If a required value is missing or a value is not an option. + """ + if field.type == FindingFieldType.CHOICE: + _check_choice_values(field=field, values=values, previous_values=previous_values) + elif field.required and _non_blank(values.get(field.key)) is None: + raise FindingValuesError(f"{field.label} is required.") + + +def _check_choice_values(*, field: FindingField, values: dict[str, str], previous_values: dict[str, str]) -> None: + """ + Walk a choice field's dropdowns, checking each value against the options under its parent. + + A value that is not an option passes if it and every value above it are unchanged. + + Args: + field (FindingField): The choice field. + values (dict[str, str]): Every value of the finding, by key. + previous_values (dict[str, str]): Every value of the stored finding, by key. + + Raises: + FindingValuesError: If a required dropdown with options has no value, or a value is not an + option and is neither unchanged nor allowed as Other text. + """ + options = field.options + path_unchanged = True + for key, label in [(field.key, field.label), *field.levels]: + value = _non_blank(values.get(key)) + path_unchanged = path_unchanged and value == _non_blank(previous_values.get(key)) + options = _check_level( + label=label, + value=value, + options=options, + required=field.required, + accepts_unmatched=path_unchanged or field.allow_other, + ) + + +def _check_level( + *, label: str, value: str | None, options: list[FindingFieldOption], required: bool, accepts_unmatched: bool +) -> list[FindingFieldOption]: + """ + Check one dropdown's value and find the options of the next dropdown. + + Args: + label (str): The dropdown's label. + value (str | None): The dropdown's value, or None if it is missing. + options (list[FindingFieldOption]): The options the dropdown offers. + required (bool): Whether the field is required. + accepts_unmatched (bool): Whether a value that is not an option passes. + + Returns: + list[FindingFieldOption]: The children of the chosen option, or an empty list if there is + no matching option. + + Raises: + FindingValuesError: If a required dropdown with options has no value, or the value is not + an option and ``accepts_unmatched`` is False. + """ + if value is None: + if required and options: + raise FindingValuesError(f"{label} is required.") + return [] + match = next((option for option in options if option.key == value), None) + if match is None and not accepts_unmatched: + raise FindingValuesError(f"{label} has no option {value!r}.") + return [] if match is None else match.children + + +def _non_blank(value: str | None) -> str | None: + """ + Treat a blank value as missing. + + Args: + value (str | None): The value, or None if it is missing. + + Returns: + str | None: The value, or None if it is missing or blank. + """ + return value if value and value.strip() else None + + +DEFAULT_FINDING_FIELDS = FindingFields( + severity=FindingField( + key="severity", + label="Severity", + type=FindingFieldType.CHOICE, + options=[ + FindingFieldOption(key="critical", label="Critical"), + FindingFieldOption(key="important", label="Important"), + FindingFieldOption(key="moderate", label="Moderate"), + FindingFieldOption(key="low", label="Low"), + FindingFieldOption(key="informational", label="Informational"), + ], + default="moderate", + allow_other=True, + colors={ + "critical": FindingSeverityColor.RED, + "important": FindingSeverityColor.YELLOW, + "moderate": FindingSeverityColor.YELLOW, + "low": FindingSeverityColor.GRAY, + }, + ), + fields=[ + FindingField( + key="harm_type", + label="Harm-type", + type=FindingFieldType.CHOICE, + options=[ + FindingFieldOption(key=category.value, label=category.value) + for category in HarmCategory + if category is not HarmCategory.OTHER + ], + allow_other=True, + hint="Optional. Select a category, or choose Other for custom text.", + ), + FindingField(key="description", label="Description", type=FindingFieldType.MULTILINE, hint="Optional"), + ], +) diff --git a/tests/unit/models/test_finding_fields.py b/tests/unit/models/test_finding_fields.py new file mode 100644 index 0000000000..e3d3fdade2 --- /dev/null +++ b/tests/unit/models/test_finding_fields.py @@ -0,0 +1,320 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest +from pydantic import ValidationError + +from pyrit.models import ( + DEFAULT_FINDING_FIELDS, + FindingField, + FindingFields, + FindingValuesError, +) +from pyrit.models.finding_fields import FindingFieldOption +from pyrit.models.harm_category import HarmCategory + +CONTOSO = FindingFields( + severity=FindingField( + key="severity", + label="Severity", + type="choice", + section="Severity", + colors={"sev1": "red", "sev2": "yellow", "sev3": "blue"}, + levels=[("severity_rationale", "Severity rationale")], + options={ + ("sev1", "Sev 1"): [("remote_no_auth", "Remote, no sign-in needed")], + ("sev2", "Sev 2"): [("signed_in_user", "Needs a signed-in user")], + ("sev3", "Sev 3"): [("admin_only", "Needs admin access")], + }, + ), + fields=[ + FindingField( + key="category", + label="Category", + type="choice", + required=True, + section="Classification", + levels=[("subcategory", "Subcategory")], + options={ + ("model_security", "Model security"): [ + ("system_prompt_leak", "System prompt leak"), + ("tool_misuse", "Tool misuse"), + ], + ("content_safety", "Content safety"): [("hate_speech", "Hate speech"), ("self_harm", "Self-harm")], + }, + ), + FindingField( + key="technique", + label="Technique", + type="choice", + section="Reproduction", + allow_other=True, + options=[("role_play", "Role play"), ("multi_turn", "Multi-turn escalation")], + ), + FindingField( + key="steps", + label="Steps to reproduce", + type="multiline", + required=True, + section="Reproduction", + hint="List the prompts in the order you sent them.", + ), + FindingField(key="owner", label="Owner", type="text", section="Tracking"), + ], +) + + +def test_finding_field_normalizes_pairs_and_nested_options() -> None: + assert CONTOSO.severity.options[0] == FindingFieldOption( + key="sev1", + label="Sev 1", + children=[FindingFieldOption(key="remote_no_auth", label="Remote, no sign-in needed")], + ) + technique = CONTOSO.fields[1] + assert [(option.key, option.label, option.children) for option in technique.options] == [ + ("role_play", "Role play", []), + ("multi_turn", "Multi-turn escalation", []), + ] + + +def test_finding_fields_round_trip_through_json() -> None: + assert FindingFields.model_validate_json(CONTOSO.model_dump_json()) == CONTOSO + + +def test_severity_slot_is_required_and_ordered() -> None: + assert CONTOSO.severity.required is True + assert CONTOSO.severity_order == ["sev1", "sev2", "sev3"] + + +def test_default_finding_fields_match_3066() -> None: + severity, (harm, description) = DEFAULT_FINDING_FIELDS.severity, DEFAULT_FINDING_FIELDS.fields + assert DEFAULT_FINDING_FIELDS.severity_order == ["critical", "important", "moderate", "low", "informational"] + assert [option.label for option in severity.options] == [ + "Critical", + "Important", + "Moderate", + "Low", + "Informational", + ] + assert (severity.default, severity.allow_other) == ("moderate", True) + assert severity.colors == {"critical": "red", "important": "yellow", "moderate": "yellow", "low": "gray"} + assert (harm.key, harm.label, harm.type, harm.required, harm.allow_other) == ( + "harm_type", + "Harm-type", + "choice", + False, + True, + ) + assert [option.key for option in harm.options] == [c.value for c in HarmCategory if c is not HarmCategory.OTHER] + assert harm.hint == "Optional. Select a category, or choose Other for custom text." + assert (description.key, description.label, description.type, description.hint) == ( + "description", + "Description", + "multiline", + "Optional", + ) + + +SEVERITY = {"key": "severity", "label": "Severity", "type": "choice", "options": [("high", "High")]} + + +@pytest.mark.parametrize("color", ["red", "orange", "yellow", "green", "blue", "gray", "dark", "neutral"]) +def test_severity_slot_accepts_every_badge_color(color: str) -> None: + fields = FindingFields.model_validate({"severity": {**SEVERITY, "colors": {"high": color}}, "fields": []}) + assert fields.severity.colors == {"high": color} + + +def test_severity_slot_rejects_an_unknown_color() -> None: + with pytest.raises(ValidationError, match="Input should be"): + FindingFields.model_validate({"severity": {**SEVERITY, "colors": {"high": "pink"}}, "fields": []}) + + +def _choice(**overrides: object) -> dict[str, object]: + return {"key": "area", "label": "Area", "type": "choice", "options": [("chat", "Chat")], **overrides} + + +NESTED = {("chat", "Chat"): [("history", "History")]} + + +@pytest.mark.parametrize( + ("severity", "fields", "match"), + [ + (SEVERITY, [_choice(key="Bad-Key")], "lowercase letters, digits and underscores"), + ( + SEVERITY, + [{"key": "notes", "label": "Notes", "type": "text", "options": [("a", "A")]}], + "cannot have options", + ), + (SEVERITY, [_choice(options=[])], "needs at least one option"), + (SEVERITY, [_choice(options=[("chat", "Chat"), ("chat", "Chat again")])], "repeats the option key 'chat'"), + (SEVERITY, [_choice(options=NESTED)], "one entry in levels"), + (SEVERITY, [_choice(levels=[("feature", "Feature")])], "one entry in levels"), + (SEVERITY, [_choice(default="search")], "'search' is not a top-level option of 'area'"), + ( + SEVERITY, + [_choice(allow_other=True, levels=[("feature", "Feature")], options=NESTED)], + "Other text only without levels", + ), + (SEVERITY, [_choice(colors={"chat": "red"})], "Only the severity field can set colors"), + ({**SEVERITY, "colors": {"low": "red"}}, [], "'low' is not a top-level option of 'severity'"), + ({**SEVERITY, "key": "sev"}, [], "key must be 'severity'"), + ({**SEVERITY, "type": "text", "options": []}, [], "must be a choice field"), + ({**SEVERITY, "required": False}, [], "always required"), + (SEVERITY, [_choice(), _choice(label="Area again")], "'area' is used more than once"), + (SEVERITY, [_choice(levels=[("severity", "Severity")], options=NESTED)], "'severity' is used more than once"), + (SEVERITY, [_choice(key="title")], "'title' is reserved"), + (SEVERITY, [_choice(label=" ")], "must not be blank"), + ], +) +def test_finding_fields_rejects_invalid_definitions(severity: dict, fields: list, match: str) -> None: + with pytest.raises(ValidationError, match=match): + FindingFields.model_validate({"severity": severity, "fields": fields}) + + +@pytest.mark.parametrize( + ("field", "match"), + [ + (_choice(levels=[("Feature", "Feature")], options=NESTED), "lowercase letters, digits and underscores"), + (_choice(options=[("chat", " ")]), "must not be blank"), + (_choice(levels=[("feature", "Feature")], options={("chat", "Chat"): [(" ", "History")]}), "must not be blank"), + ({"key": "notes", "label": "Notes", "type": "text", "default": "a"}, "cannot have default"), + (_choice(options={"chat": []}), r"\(key, label\) pair"), + ], +) +def test_finding_field_rejects_invalid_parts(field: dict, match: str) -> None: + with pytest.raises(ValidationError, match=match): + FindingField.model_validate(field) + + +def test_severity_slot_accepts_explicit_required() -> None: + assert FindingFields.model_validate({"severity": {**SEVERITY, "required": True}}).severity.required is True + + +VALID = { + "severity_rationale": "remote_no_auth", + "category": "model_security", + "subcategory": "system_prompt_leak", + "steps": "1. Open a new chat", +} + + +def _without(*keys: str) -> dict[str, str]: + return {key: value for key, value in VALID.items() if key not in keys} + + +@pytest.mark.parametrize( + ("severity", "values", "previous_severity", "previous_values"), + [ + pytest.param("sev1", VALID, None, None, id="complete"), + pytest.param("sev1", {**VALID, "technique": "Many-shot priming", "owner": "dana"}, None, None, id="other text"), + pytest.param("sev4", _without("severity_rationale"), "sev4", {}, id="unchanged unmatched severity"), + pytest.param("sev1", {**VALID, "repro": "old"}, "sev1", {"repro": "old"}, id="unchanged retired field"), + pytest.param( + "sev1", + {**VALID, "category": "gone", "subcategory": "old"}, + "sev1", + {"category": "gone", "subcategory": "old"}, + id="unchanged unmatched path", + ), + ], +) +def test_check_values_accepts( + severity: str, values: dict, previous_severity: str | None, previous_values: dict | None +) -> None: + CONTOSO.check_values( + severity=severity, field_values=values, previous_severity=previous_severity, previous_values=previous_values + ) + + +@pytest.mark.parametrize( + ("severity", "values", "previous_severity", "previous_values", "message"), + [ + pytest.param("", VALID, None, None, "Severity is required.", id="blank severity"), + pytest.param("sev9", VALID, None, None, "Severity has no option 'sev9'.", id="unknown severity"), + pytest.param("sev4", VALID, "sev1", VALID, "Severity has no option 'sev4'.", id="changed to unmatched"), + pytest.param( + "sev1", _without("severity_rationale"), None, None, "Severity rationale is required.", id="missing level" + ), + pytest.param( + "sev1", _without("category", "subcategory"), None, None, "Category is required.", id="missing field" + ), + pytest.param( + "sev1", + {**VALID, "subcategory": "hate_speech"}, + None, + None, + "Subcategory has no option 'hate_speech'.", + id="child under wrong parent", + ), + pytest.param( + "sev1", + {**VALID, "category": "content_safety"}, + "sev1", + VALID, + "Subcategory has no option 'system_prompt_leak'.", + id="stale child", + ), + pytest.param("sev1", {**VALID, "steps": " "}, None, None, "Steps to reproduce is required.", id="blank text"), + pytest.param( + "sev1", + _without("steps"), + "sev1", + _without("steps"), + "Steps to reproduce is required.", + id="required field added after saving", + ), + pytest.param( + "sev1", {**VALID, "reviewer": "kim"}, None, None, "Unknown finding field 'reviewer'.", id="unknown key" + ), + pytest.param( + "sev1", + {**VALID, "severity": "sev1"}, + None, + None, + "Unknown finding field 'severity'.", + id="severity in values", + ), + pytest.param( + "sev1", + {**VALID, "repro": "new"}, + "sev1", + {**VALID, "repro": "old"}, + "Unknown finding field 'repro'.", + id="changed retired field", + ), + ], +) +def test_check_values_rejects( + severity: str, values: dict, previous_severity: str | None, previous_values: dict | None, message: str +) -> None: + with pytest.raises(FindingValuesError) as error: + CONTOSO.check_values( + severity=severity, field_values=values, previous_severity=previous_severity, previous_values=previous_values + ) + assert str(error.value) == message + + +AREA = _choice( + levels=[("feature", "Feature")], + options={("chat", "Chat"): [("history", "History")], ("search", "Search"): []}, +) +AREAS = FindingFields.model_validate({"severity": SEVERITY, "fields": [AREA]}) +REQUIRED_AREAS = FindingFields.model_validate({"severity": SEVERITY, "fields": [{**AREA, "required": True}]}) + + +@pytest.mark.parametrize("values", [{"area": "chat"}, {}]) +def test_check_values_optional_cascade_allows_a_partial_path(values: dict) -> None: + AREAS.check_values(severity="high", field_values=values) + + +def test_check_values_rejects_a_child_without_its_parent() -> None: + with pytest.raises(FindingValuesError) as error: + AREAS.check_values(severity="high", field_values={"feature": "history"}) + assert str(error.value) == "Feature has no option 'history'." + + +def test_check_values_required_cascade_skips_a_level_with_no_options() -> None: + REQUIRED_AREAS.check_values(severity="high", field_values={"area": "search"}) + with pytest.raises(FindingValuesError) as error: + REQUIRED_AREAS.check_values(severity="high", field_values={"area": "chat"}) + assert str(error.value) == "Feature is required." From ca3d47dcc05589b007358e6a804848735dd6f723 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Mon, 26 Oct 2026 12:00:54 -0400 Subject: [PATCH 03/20] Add FindingFieldsInitializer and active finding fields Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/memory/memory_interface.py | 2 +- pyrit/setup/finding_fields.py | 44 +++++++++ pyrit/setup/initialization.py | 3 + pyrit/setup/initializers/__init__.py | 2 + pyrit/setup/initializers/finding_fields.py | 39 ++++++++ tests/unit/setup/test_finding_fields.py | 45 +++++++++ .../setup/test_finding_fields_initializer.py | 96 +++++++++++++++++++ tests/unit/setup/test_initialization.py | 87 ++++++++++++++++- 8 files changed, 315 insertions(+), 3 deletions(-) create mode 100644 pyrit/setup/finding_fields.py create mode 100644 pyrit/setup/initializers/finding_fields.py create mode 100644 tests/unit/setup/test_finding_fields.py create mode 100644 tests/unit/setup/test_finding_fields_initializer.py diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 3596046b3f..ca6381f859 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1046,7 +1046,7 @@ async def get_findings_async( severity_order: Sequence[str] = (), ) -> list[Finding]: """ - Read one page of findings, highest severity first, then newest first. + Read one page of findings in ``severity_order``, then newest first. Args: operation_id (uuid.UUID | None): The operation whose findings to read. Defaults to None, which reads diff --git a/pyrit/setup/finding_fields.py b/pyrit/setup/finding_fields.py new file mode 100644 index 0000000000..6cc4607c67 --- /dev/null +++ b/pyrit/setup/finding_fields.py @@ -0,0 +1,44 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Hold the finding fields that the current setup pass configured.""" + +from pyrit.models import DEFAULT_FINDING_FIELDS, FindingFields + +_active: tuple[FindingFields, str] | None = None + + +def get_active_finding_fields() -> FindingFields: + """ + Return the finding fields of the current setup pass. + + Returns: + FindingFields: The fields set in this pass, or ``DEFAULT_FINDING_FIELDS`` when none were set. + """ + return _active[0] if _active is not None else DEFAULT_FINDING_FIELDS + + +def set_active_finding_fields(*, finding_fields: FindingFields, source: str) -> None: + """ + Set the finding fields for the current setup pass. + + Args: + finding_fields (FindingFields): The deployment's finding fields. + source (str): The setter's name, which the error names when two setters run in one pass. + + Raises: + ValueError: If finding fields were already set in this setup pass. + """ + global _active + if _active is not None: + raise ValueError( + f"{source} cannot set the finding fields: {_active[1]} already set them in this setup pass. " + "Use one FindingFieldsInitializer per deployment." + ) + _active = (finding_fields, source) + + +def reset_active_finding_fields() -> None: + """Restore the default finding fields at the start of a setup pass.""" + global _active + _active = None diff --git a/pyrit/setup/initialization.py b/pyrit/setup/initialization.py index be385594f2..686dea2365 100644 --- a/pyrit/setup/initialization.py +++ b/pyrit/setup/initialization.py @@ -16,6 +16,7 @@ load_environment_files, validate_env_akv_strict, ) +from pyrit.setup.finding_fields import reset_active_finding_fields if TYPE_CHECKING: from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -206,6 +207,7 @@ async def initialize_pyrit_async( # Reset all default values before executing initialization scripts # This ensures a clean state for each initialization reset_default_values() + reset_active_finding_fields() # Set up memory BEFORE executing initialization scripts # This is critical because initialization scripts may instantiate objects @@ -302,6 +304,7 @@ async def reinitialize_pyrit_async( os.environ.update(environment_values) configure_random_seed(seed=seed) reset_default_values() + reset_active_finding_fields() CentralMemory.set_memory_instance(memory) initializers = await asyncio.to_thread(initializer_factory) await _execute_initializers_async(initializers=initializers, raise_on_initializer_error=True) diff --git a/pyrit/setup/initializers/__init__.py b/pyrit/setup/initializers/__init__.py index d5693e3c80..6bd9964d5f 100644 --- a/pyrit/setup/initializers/__init__.py +++ b/pyrit/setup/initializers/__init__.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from pyrit.models.parameter import Parameter from pyrit.setup.initializers.converters import ConverterInitializer + from pyrit.setup.initializers.finding_fields import FindingFieldsInitializer from pyrit.setup.initializers.load_default_datasets import LoadDefaultDatasets from pyrit.setup.initializers.preload_scenario_metadata import PreloadScenarioMetadata from pyrit.setup.initializers.refresh_datasets import RefreshDatasets @@ -24,6 +25,7 @@ "Parameter": "pyrit.models.parameter", "PyRITInitializer": "pyrit.setup.pyrit_initializer", "ConverterInitializer": "pyrit.setup.initializers.converters", + "FindingFieldsInitializer": "pyrit.setup.initializers.finding_fields", "TechniqueInitializer": "pyrit.setup.initializers.techniques", "ToolsInitializer": "pyrit.setup.initializers.tools", "ScorerInitializer": "pyrit.setup.initializers.scorers", diff --git a/pyrit/setup/initializers/finding_fields.py b/pyrit/setup/initializers/finding_fields.py new file mode 100644 index 0000000000..d4025b0d3d --- /dev/null +++ b/pyrit/setup/initializers/finding_fields.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Base initializer for a deployment's finding fields.""" + +from abc import abstractmethod +from typing import Any + +from pyrit.models import FindingFields +from pyrit.setup.finding_fields import set_active_finding_fields +from pyrit.setup.pyrit_initializer import PyRITInitializer + + +class FindingFieldsInitializer(PyRITInitializer): + """Set a deployment's finding fields from ``get_finding_fields``.""" + + @abstractmethod + def get_finding_fields(self) -> FindingFields: + """ + Return the deployment's finding fields. + + Returns: + FindingFields: The fields that setup makes active. + """ + + async def initialize_async(self) -> None: + """Make the fields from ``get_finding_fields`` the active finding fields.""" + set_active_finding_fields(finding_fields=self.get_finding_fields(), source=type(self).__name__) + + async def get_dynamic_default_values_info_async(self) -> dict[str, Any]: + """ + Report no default values or global variables without running ``initialize_async``. + + The base class runs ``initialize_async`` to find out, which would change the active finding fields. + + Returns: + dict[str, Any]: Empty ``default_values`` and ``global_variables`` lists. + """ + return {"default_values": [], "global_variables": []} diff --git a/tests/unit/setup/test_finding_fields.py b/tests/unit/setup/test_finding_fields.py new file mode 100644 index 0000000000..be8dc16a04 --- /dev/null +++ b/tests/unit/setup/test_finding_fields.py @@ -0,0 +1,45 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from collections.abc import Iterator + +import pytest + +from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields +from pyrit.setup.finding_fields import ( + get_active_finding_fields, + reset_active_finding_fields, + set_active_finding_fields, +) + +CUSTOM = FindingFields( + severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) +) + + +@pytest.fixture(autouse=True) +def reset_finding_fields() -> Iterator[None]: + yield + reset_active_finding_fields() + + +def test_active_finding_fields_default_to_3066() -> None: + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + +def test_set_active_finding_fields_replaces_the_default() -> None: + set_active_finding_fields(finding_fields=CUSTOM, source="First") + assert get_active_finding_fields() is CUSTOM + + +def test_set_active_finding_fields_twice_names_both_sources() -> None: + set_active_finding_fields(finding_fields=CUSTOM, source="First") + with pytest.raises(ValueError, match="Second cannot set the finding fields: First already set them"): + set_active_finding_fields(finding_fields=DEFAULT_FINDING_FIELDS, source="Second") + assert get_active_finding_fields() is CUSTOM + + +def test_reset_restores_the_default() -> None: + set_active_finding_fields(finding_fields=CUSTOM, source="First") + reset_active_finding_fields() + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS diff --git a/tests/unit/setup/test_finding_fields_initializer.py b/tests/unit/setup/test_finding_fields_initializer.py new file mode 100644 index 0000000000..955ffc7e4a --- /dev/null +++ b/tests/unit/setup/test_finding_fields_initializer.py @@ -0,0 +1,96 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import importlib.util +from collections.abc import Iterator +from pathlib import Path + +import pytest + +import pyrit.setup.initializers +from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields +from pyrit.registry import InitializerRegistry +from pyrit.registry.discovery import discover_in_directory +from pyrit.setup.finding_fields import get_active_finding_fields, reset_active_finding_fields +from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.pyrit_initializer import PyRITInitializer + +EXAMPLE = FindingFields( + severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) +) + +CONTOSO_SCRIPT = ''' +from pyrit.models import FindingField, FindingFields +from pyrit.setup.initializers import FindingFieldsInitializer + + +class ContosoFindingFields(FindingFieldsInitializer): + """Contoso finding fields.""" + + def get_finding_fields(self) -> FindingFields: + return FindingFields( + severity=FindingField(key="severity", label="Severity", type="choice", options=[("sev1", "Sev 1")]) + ) +''' + + +class ExampleFindingFields(FindingFieldsInitializer): + """Example finding fields.""" + + def get_finding_fields(self) -> FindingFields: + return EXAMPLE + + +@pytest.fixture(autouse=True) +def reset_finding_fields() -> Iterator[None]: + yield + reset_active_finding_fields() + + +async def test_finding_fields_initializer_sets_the_active_fields() -> None: + await ExampleFindingFields().initialize_async() + assert get_active_finding_fields() is EXAMPLE + + +@pytest.mark.usefixtures("patch_central_database") +async def test_get_info_leaves_the_active_fields_alone() -> None: + await ExampleFindingFields.get_info_async() + info = await ExampleFindingFields.get_info_async() + + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + assert info["default_values"] == [] + assert info["global_variables"] == [] + + +def test_finding_fields_initializer_is_skipped_by_builtin_discovery() -> None: + directory = Path(pyrit.setup.initializers.__file__).parent + names = [ + cls.__name__ + for _, _, cls in discover_in_directory(directory=directory, base_class=PyRITInitializer, recursive=True) + ] + assert "FindingFieldsInitializer" not in names + + +async def test_discovery_copy_shares_the_active_finding_fields() -> None: + path = Path(pyrit.setup.initializers.__file__).parent / "finding_fields.py" + spec = importlib.util.spec_from_file_location("discovered_module.finding_fields", path) + copy = importlib.util.module_from_spec(spec) + spec.loader.exec_module(copy) + + class CopiedFindingFields(copy.FindingFieldsInitializer): + def get_finding_fields(self) -> FindingFields: + return EXAMPLE + + await CopiedFindingFields().initialize_async() + assert get_active_finding_fields() is EXAMPLE + + +async def test_custom_script_subclass_sets_the_active_fields(tmp_path: Path) -> None: + script = tmp_path / "contoso_fields.py" + script.write_text(CONTOSO_SCRIPT) + [initializer] = InitializerRegistry.get_registry_singleton().create_from_script_paths( + script_paths=[script], strict=True + ) + assert type(initializer).__name__ == "ContosoFindingFields" + await initializer.initialize_async() + assert get_active_finding_fields().severity_order == ["sev1"] diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 73c275d49f..2525b79883 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -5,19 +5,32 @@ import os import pathlib import tempfile +from collections.abc import Iterator from unittest import mock import pytest +from pydantic import ValidationError from pyrit.common.apply_defaults import reset_default_values from pyrit.common.random_context import get_configured_random_seed from pyrit.common.singleton import Singleton from pyrit.memory import AzureSQLMemory, CentralMemory, SQLiteMemory -from pyrit.models import MessagePiece +from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields, MessagePiece from pyrit.registry import InitializerRegistry -from pyrit.setup import IN_MEMORY, initialize_pyrit_async +from pyrit.setup import IN_MEMORY, initialization, initialize_pyrit_async +from pyrit.setup.finding_fields import ( + get_active_finding_fields, + reset_active_finding_fields, + set_active_finding_fields, +) +from pyrit.setup.initialization import reinitialize_pyrit_async +from pyrit.setup.initializers import FindingFieldsInitializer from pyrit.setup.pyrit_initializer import PyRITInitializer +CUSTOM = FindingFields( + severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) +) + @pytest.mark.parametrize("existing_memory", [False, True]) async def test_initializer_failure_closes_only_newly_installed_memory( @@ -440,3 +453,73 @@ async def test_initialize_forwards_silent_to_memory(self, mock_load_environment, await initialize_pyrit_async(memory_db_type=IN_MEMORY, silent=silent, load_defaults=False) assert CentralMemory.get_memory_instance()._silent is silent + + +class _SetsCustomFields(PyRITInitializer): + async def initialize_async(self) -> None: + set_active_finding_fields(finding_fields=CUSTOM, source="pass") + + +@pytest.fixture +def reset_finding_fields() -> Iterator[None]: + yield + reset_active_finding_fields() + + +@pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") +async def test_each_setup_pass_resets_finding_fields() -> None: + with mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, initializers=[_SetsCustomFields()]) + assert get_active_finding_fields() is CUSTOM + + await initialize_pyrit_async(memory_db_type=IN_MEMORY, initializers=[_SetsCustomFields()]) + assert get_active_finding_fields() is CUSTOM + + await initialize_pyrit_async(memory_db_type=IN_MEMORY, load_defaults=False) + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + +@pytest.mark.usefixtures("reset_finding_fields") +async def test_reinitialize_resets_finding_fields() -> None: + set_active_finding_fields(finding_fields=CUSTOM, source="pass") + with ( + mock.patch.object(CentralMemory, "set_memory_instance"), + mock.patch.object(initialization, "_execute_initializers_async", new_callable=mock.AsyncMock), + ): + await reinitialize_pyrit_async( + memory=mock.MagicMock(spec=SQLiteMemory), initializer_factory=list, environment_values={}, seed=None + ) + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + +@pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") +async def test_second_finding_fields_initializer_stops_setup() -> None: + class FirstFindingFields(FindingFieldsInitializer): + def get_finding_fields(self) -> FindingFields: + return CUSTOM + + class SecondFindingFields(FindingFieldsInitializer): + def get_finding_fields(self) -> FindingFields: + return CUSTOM + + with mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock): + with pytest.raises( + ValueError, match="SecondFindingFields cannot set the finding fields: FirstFindingFields already set them" + ): + await initialize_pyrit_async( + memory_db_type=IN_MEMORY, initializers=[FirstFindingFields(), SecondFindingFields()] + ) + + +@pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") +async def test_invalid_finding_fields_stop_setup() -> None: + class BrokenFindingFields(FindingFieldsInitializer): + def get_finding_fields(self) -> FindingFields: + return FindingFields( + severity=CUSTOM.severity, + fields=[FindingField(key="broken", label="Broken", type="choice", options=[])], + ) + + with mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock): + with pytest.raises(ValidationError, match="needs at least one option"): + await initialize_pyrit_async(memory_db_type=IN_MEMORY, initializers=[BrokenFindingFields()]) From 627cba2beae489206daa151d64e63d90891cce76 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Tue, 27 Oct 2026 10:52:19 -0400 Subject: [PATCH 04/20] Store finding severity and field values without CHECK constraints Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/services/operation_service.py | 6 +- ...d43b1bcf9c56_store_finding_field_values.py | 95 +++++++++ pyrit/memory/memory_models.py | 52 +---- pyrit/models/__init__.py | 3 +- pyrit/models/finding.py | 61 ++---- tests/unit/backend/test_operations.py | 48 ++--- .../test_finding_field_values_migration.py | 185 ++++++++++++++++++ tests/unit/memory/test_findings.py | 87 +++----- .../test_operations_and_findings_migration.py | 34 +++- tests/unit/models/test_finding.py | 74 ++----- 10 files changed, 404 insertions(+), 241 deletions(-) create mode 100644 pyrit/memory/alembic/versions/d43b1bcf9c56_store_finding_field_values.py create mode 100644 tests/unit/memory/test_finding_field_values_migration.py diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py index 79e3889348..4b1621b310 100644 --- a/pyrit/backend/services/operation_service.py +++ b/pyrit/backend/services/operation_service.py @@ -25,10 +25,10 @@ Finding, FindingCreate, FindingEvidence, - FindingSeverity, Operation, OperationCreate, ) +from pyrit.setup.finding_fields import get_active_finding_fields # Keeps every ID list a batch reads within SQL Server's 2100-parameter limit. _EVIDENCE_BATCH_SIZE = 500 @@ -201,7 +201,7 @@ async def list_findings_async( limit=limit + 1, offset=offset, title_query=title_query, - severity_order=[severity.value for severity in FindingSeverity], + severity_order=get_active_finding_fields().severity_order, ) has_more = len(findings) > limit page = findings[:limit] @@ -277,7 +277,7 @@ async def get_findings_with_evidence_async( findings = await memory.get_findings_by_ids_async( operation_id=operation_id, finding_ids=finding_ids, - severity_order=[severity.value for severity in FindingSeverity], + severity_order=get_active_finding_fields().severity_order, ) found = {finding.id for finding in findings} missing = next((finding_id for finding_id in finding_ids if finding_id not in found), None) diff --git a/pyrit/memory/alembic/versions/d43b1bcf9c56_store_finding_field_values.py b/pyrit/memory/alembic/versions/d43b1bcf9c56_store_finding_field_values.py new file mode 100644 index 0000000000..b1ee65fb17 --- /dev/null +++ b/pyrit/memory/alembic/versions/d43b1bcf9c56_store_finding_field_values.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Store each finding as a severity plus a JSON object of field values. + +Drops the three classification CHECK constraints on ``FindingEntries``, makes ``severity`` unbounded Unicode +(binary-collated on SQL Server), replaces an Other severity with its custom text, and moves the harm type (its +custom text when it is Other) and any description that isn't empty or only spaces into the new ``field_values`` +column. Finding IDs, Operation links, creation times, and evidence links do not change. Downgrading restores the +old columns only while no findings exist. + +Revision ID: d43b1bcf9c56 +Revises: c8d3e5f7a901 +Create Date: 2026-10-27 12:00:00.000000 +""" + +from collections.abc import Sequence # noqa: TC003 + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "d43b1bcf9c56" +down_revision: str | None = "c8d3e5f7a901" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_SEVERITY_TYPE = sa.Unicode().with_variant(sa.Unicode(collation="Latin1_General_100_BIN2"), "mssql") +_HARM_TYPE_SQL = "CASE WHEN harm_type = 'Other' THEN harm_type_other ELSE harm_type END" +_DESCRIPTION_SQL = "CASE WHEN LTRIM(RTRIM(description)) = '' THEN NULL ELSE description END" +_OLD_COLUMNS = ("description", "severity_other", "harm_type", "harm_type_other") +_CHECKS = { + "ck_findings_severity": "severity IN ('critical', 'important', 'moderate', 'low', 'informational', 'other')", + "ck_findings_severity_other": ( + "(severity = 'other' AND severity_other IS NOT NULL AND LTRIM(RTRIM(severity_other)) <> '') " + "OR (severity <> 'other' AND severity_other IS NULL)" + ), + "ck_findings_harm_type_other": ( + "(harm_type IS NOT NULL AND harm_type = 'Other' AND harm_type_other IS NOT NULL " + "AND LTRIM(RTRIM(harm_type_other)) <> '') " + "OR ((harm_type IS NULL OR harm_type <> 'Other') AND harm_type_other IS NULL)" + ), +} + + +def upgrade() -> None: + """Drop the classification checks and move each finding's values into ``severity`` and ``field_values``.""" + # SQL Server rejects ALTER COLUMN while a CHECK references the column, so the checks go first. + with op.batch_alter_table("FindingEntries") as batch_op: + for name in _CHECKS: + batch_op.drop_constraint(name, type_="check") + batch_op.add_column(sa.Column("field_values", sa.JSON(), nullable=True)) + batch_op.alter_column("severity", existing_type=sa.String(32), type_=_SEVERITY_TYPE, existing_nullable=False) + op.execute(sa.text("UPDATE FindingEntries SET severity = severity_other WHERE LOWER(severity) = 'other'")) + op.execute(sa.text(f"UPDATE FindingEntries SET field_values = {_field_values_sql()}")) + with op.batch_alter_table("FindingEntries") as batch_op: + batch_op.alter_column("field_values", existing_type=sa.JSON(), nullable=False) + for column in _OLD_COLUMNS: + batch_op.drop_column(column) + + +def _field_values_sql() -> str: + """ + Build the dialect's SQL expression for a JSON object of harm type and description, leaving out NULL values. + + Returns: + str: The SQL expression. + """ + if op.get_context().dialect.name == "mssql": + return ( + f"JSON_MODIFY(JSON_MODIFY(N'{{}}', '$.harm_type', {_HARM_TYPE_SQL}), '$.description', {_DESCRIPTION_SQL})" + ) + return f"json_patch('{{}}', json_object('harm_type', {_HARM_TYPE_SQL}, 'description', {_DESCRIPTION_SQL}))" + + +def downgrade() -> None: + """ + Restore the classification columns and checks only when no findings exist. + + Raises: + ValueError: If the downgrade is offline or findings would be lost. + """ + if op.get_context().as_sql: + raise ValueError("Finding downgrade requires online data validation.") + if op.get_bind().execute(sa.text("SELECT COUNT(*) FROM FindingEntries")).scalar(): + raise ValueError("Remove retained findings explicitly before downgrading.") + with op.batch_alter_table("FindingEntries") as batch_op: + batch_op.add_column(sa.Column("description", sa.Unicode(), nullable=False)) + for column in _OLD_COLUMNS[1:]: + batch_op.add_column(sa.Column(column, sa.Unicode(), nullable=True)) + batch_op.alter_column("severity", existing_type=_SEVERITY_TYPE, type_=sa.String(32), existing_nullable=False) + batch_op.drop_column("field_values") + for name, condition in _CHECKS.items(): + batch_op.create_check_constraint(name, condition) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a693c7d71b..33cabd6639 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -16,7 +16,6 @@ INTEGER, JSON, Boolean, - CheckConstraint, Computed, DateTime, Float, @@ -84,7 +83,6 @@ TargetIdentifier, scorable_from_dict, ) -from pyrit.models.harm_category import HarmCategory from pyrit.models.operation import operation_name_key from pyrit.models.results.attack_result import normalize_legacy_attack_attribution @@ -295,36 +293,18 @@ class FindingEntry(Base): """A saved finding, kept separate from scores and attack results.""" __tablename__ = "FindingEntries" - # Harm types are deliberately unconstrained: HarmCategory changes over time, so reads map unknown values instead. - __table_args__ = ( - Index("ix_FindingEntries_operation_id", "operation_id"), - CheckConstraint( - "severity IN ('critical', 'important', 'moderate', 'low', 'informational', 'other')", - name="ck_findings_severity", - ), - CheckConstraint( - "(severity = 'other' AND severity_other IS NOT NULL AND LTRIM(RTRIM(severity_other)) <> '') " - "OR (severity <> 'other' AND severity_other IS NULL)", - name="ck_findings_severity_other", - ), - CheckConstraint( - "(harm_type IS NOT NULL AND harm_type = 'Other' AND harm_type_other IS NOT NULL " - "AND LTRIM(RTRIM(harm_type_other)) <> '') " - "OR ((harm_type IS NULL OR harm_type <> 'Other') AND harm_type_other IS NULL)", - name="ck_findings_harm_type_other", - ), - ) + __table_args__ = (Index("ix_FindingEntries_operation_id", "operation_id"),) id: Mapped[uuid.UUID] = mapped_column(CustomUUID, primary_key=True) operation_id: Mapped[uuid.UUID] = mapped_column( CustomUUID, ForeignKey(f"{OperationEntry.__tablename__}.id"), nullable=False ) title: Mapped[str] = mapped_column(Unicode, nullable=False) - description: Mapped[str] = mapped_column(Unicode, nullable=False) - severity: Mapped[str] = mapped_column(String(32), nullable=False) - severity_other: Mapped[str | None] = mapped_column(Unicode, nullable=True) - harm_type: Mapped[str | None] = mapped_column(Unicode, nullable=True) - harm_type_other: Mapped[str | None] = mapped_column(Unicode, nullable=True) + # Binary collation makes SQL Server sort and count severities exactly, as SQLite does. + severity: Mapped[str] = mapped_column( + Unicode().with_variant(Unicode(collation="Latin1_General_100_BIN2"), "mssql"), nullable=False + ) + field_values: Mapped[dict[str, str]] = mapped_column(JSON, nullable=False) created_at: Mapped[datetime] = mapped_column(UTCDateTime, nullable=False) def __init__(self, finding: Finding) -> None: @@ -333,37 +313,23 @@ def __init__(self, finding: Finding) -> None: self.id = validated.id self.operation_id = validated.operation_id self.title = validated.title - self.description = validated.description - self.severity = validated.severity.value - self.severity_other = validated.severity_other - self.harm_type = validated.harm_type.value if validated.harm_type is not None else None - self.harm_type_other = validated.harm_type_other + self.severity = validated.severity + self.field_values = dict(validated.field_values) self.created_at = validated.created_at.astimezone(UTC) def get_finding(self) -> Finding: """ Rebuild the finding model. - A stored harm type that is no longer a ``HarmCategory`` maps through its aliases, or to - Other with the stored text, rather than making the finding unreadable. - Returns: Finding: The stored finding. """ - harm_type, harm_type_other = self.harm_type, self.harm_type_other - if harm_type is not None and harm_type not in set(HarmCategory): - # parse() logs a warning when it falls back to Other. - parsed = HarmCategory.parse(harm_type) - harm_type, harm_type_other = parsed, harm_type if parsed == HarmCategory.OTHER else None return Finding( id=self.id, operation_id=self.operation_id, title=self.title, - description=self.description, severity=self.severity, - severity_other=self.severity_other, - harm_type=harm_type, - harm_type_other=harm_type_other, + field_values=self.field_values, created_at=self.created_at, ) diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 9415a50a99..8aca6abb6b 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -73,7 +73,7 @@ from pyrit.models.conversation_stats import ConversationStats from pyrit.models.conversation_summary import ConversationSummary from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation - from pyrit.models.finding import Finding, FindingCreate, FindingEvidence, FindingSeverity + from pyrit.models.finding import Finding, FindingCreate, FindingEvidence from pyrit.models.finding_fields import ( DEFAULT_FINDING_FIELDS, FindingField, @@ -449,7 +449,6 @@ "Finding": "pyrit.models.finding", "FindingCreate": "pyrit.models.finding", "FindingEvidence": "pyrit.models.finding", - "FindingSeverity": "pyrit.models.finding", "DEFAULT_FINDING_FIELDS": "pyrit.models.finding_fields", "FindingField": "pyrit.models.finding_fields", "FindingFields": "pyrit.models.finding_fields", diff --git a/pyrit/models/finding.py b/pyrit/models/finding.py index 09563386b3..0fa2a9c2d8 100644 --- a/pyrit/models/finding.py +++ b/pyrit/models/finding.py @@ -2,24 +2,9 @@ # Licensed under the MIT license. from datetime import UTC, datetime -from enum import Enum -from typing import Self from uuid import UUID, uuid4 -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator - -from pyrit.models.harm_category import HarmCategory - - -class FindingSeverity(str, Enum): - """Operator-selected severity, ordered from highest to lowest.""" - - CRITICAL = "critical" - IMPORTANT = "important" - MODERATE = "moderate" - LOW = "low" - INFORMATIONAL = "informational" - OTHER = "other" +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator class FindingCreate(BaseModel): @@ -28,39 +13,14 @@ class FindingCreate(BaseModel): model_config = ConfigDict(extra="forbid") title: str - description: str = "" - severity: FindingSeverity - severity_other: str | None = None - harm_type: HarmCategory | None = None - harm_type_other: str | None = None - - @model_validator(mode="after") - def _validate_classifications(self) -> Self: - """ - Require custom text only for an Other selection. - - Returns: - Self: The validated finding. + severity: str + field_values: dict[str, str] = {} - Raises: - ValueError: If custom text is missing, blank, or has no Other selection. - """ - for field, is_other, text in [ - ("severity", self.severity == FindingSeverity.OTHER, self.severity_other), - ("harm_type", self.harm_type == HarmCategory.OTHER, self.harm_type_other), - ]: - if is_other: - if text is None or not text.strip(): - raise ValueError(f"{field}_other must not be blank when {field} is Other.") - elif text is not None: - raise ValueError(f"{field}_other is only allowed when {field} is Other.") - return self - - @field_validator("title") + @field_validator("title", "severity") @classmethod def _reject_blank(cls, value: str) -> str: """ - Reject a blank title without normalizing its contents. + Reject blank text without normalizing its contents. Returns: str: The original value. @@ -72,6 +32,17 @@ def _reject_blank(cls, value: str) -> str: raise ValueError("Must not be blank.") return value + @field_validator("field_values") + @classmethod + def _drop_blank_values(cls, value: dict[str, str]) -> dict[str, str]: + """ + Drop empty and whitespace-only values, keeping other text exactly as entered. + + Returns: + dict[str, str]: The non-blank values. + """ + return {key: text for key, text in value.items() if text.strip()} + class Finding(FindingCreate): """A saved finding in one operation.""" diff --git a/tests/unit/backend/test_operations.py b/tests/unit/backend/test_operations.py index b87fd81a66..c631e11d9e 100644 --- a/tests/unit/backend/test_operations.py +++ b/tests/unit/backend/test_operations.py @@ -139,7 +139,7 @@ def test_findings_are_created_and_listed_under_an_operation(compatibility_header assert created.status_code == 201 finding = created.json() assert finding["operation_id"] == operation_id - assert finding["description"] == "" + assert finding["field_values"] == {} client.post(f"/api/operations/{other_id}/findings", json={"title": "Elsewhere", "severity": "low"}) listed = client.get(f"/api/operations/{operation_id}/findings", params={"limit": 1}).json() assert listed == {"items": [{**finding, "evidence_count": 0}], "has_more": False, "next_offset": None} @@ -165,39 +165,34 @@ def test_finding_options_exposes_canonical_harm_types(compatibility_headers: dic assert TestClient(app).get("/api/operations/finding-options").status_code != 200 -def test_finding_api_round_trips_optional_and_custom_classifications(compatibility_headers: dict[str, str]) -> None: +def test_finding_api_round_trips_field_values_and_drops_blanks(compatibility_headers: dict[str, str]) -> None: client = TestClient(app, headers=compatibility_headers) - operation_id = client.post("/api/operations", json={"name": "Classifications"}).json()["id"] + operation_id = client.post("/api/operations", json={"name": "Field values"}).json()["id"] path = f"/api/operations/{operation_id}/findings" body = { "title": "Custom", - "severity": "other", - "severity_other": " Team severity ", - "harm_type": "Other", - "harm_type_other": " Team harm ", + "severity": " Team severity ", + "field_values": {"harm_type": "Malware", "description": " "}, } response = client.post(path, json=body) assert response.status_code == 201 finding = response.json() - assert {key: finding[key] for key in body} == body + assert finding["severity"] == " Team severity " + assert finding["field_values"] == {"harm_type": "Malware"} assert client.get(path).json()["items"] == [{**finding, "evidence_count": 0}] - updated = client.put(f"{path}/{finding['id']}", json={"title": "Edited", "severity": "low", "harm_type": "Malware"}) + updated = client.put(f"{path}/{finding['id']}", json={"title": "Edited", "severity": "low"}) assert updated.status_code == 200 - assert updated.json()["harm_type"] == "Malware" - assert updated.json()["severity_other"] is None - assert updated.json()["harm_type_other"] is None + assert updated.json()["field_values"] == {} assert updated.json()["id"] == finding["id"] assert updated.json()["created_at"] == finding["created_at"] for invalid in [ - {"severity": "other"}, - {"severity": "low", "severity_other": "Stray"}, - {"severity": "low", "harm_type": "Other"}, - {"severity": "low", "harm_type": "Unknown"}, + {"severity": " "}, + {"severity": "low", "harm_type": "Malware"}, + {"severity": "low", "field_values": {"harm_type": 3}}, ]: assert client.post(path, json={"title": "Invalid", **invalid}).status_code == 422 assert client.put(f"{path}/{finding['id']}", json={"title": "Invalid", **invalid}).status_code == 422 - cleared = client.put(f"{path}/{finding['id']}", json={"title": "No harm type", "severity": "moderate"}) - assert cleared.json()["harm_type"] is None + assert client.get(path).json()["items"] == [{**updated.json(), "evidence_count": 0}] def test_findings_require_an_existing_operation(compatibility_headers: dict[str, str]) -> None: @@ -274,7 +269,8 @@ def test_finding_create_saves_nothing_when_its_evidence_cannot_be_linked( ("/api/operations", {"name": " "}), ("/api/operations", {"name": "😀" * 65}), ("findings", {"title": "", "severity": "low"}), - ("findings", {"title": "Assessment", "severity": "urgent"}), + ("findings", {"title": "Assessment", "severity": " "}), + ("findings", {"title": "Assessment", "severity": "low", "description": "x"}), ("findings", {"title": "Assessment", "severity": "low", "operation": "Case"}), ], ) @@ -306,9 +302,16 @@ def test_finding_update_and_delete_preserve_identity_and_operation( path = f"/api/operations/{operation['id']}/findings" original = client.post(path, json={"title": "Original", "severity": "low"}).json() finding_path = f"{path}/{original['id']}" - edited = client.put(finding_path, json={"title": "Edited", "severity": "critical", "description": "Notes"}) + edited = client.put( + finding_path, json={"title": "Edited", "severity": "critical", "field_values": {"description": "Notes"}} + ) assert edited.status_code == 200 - assert edited.json() == {**original, "title": "Edited", "severity": "critical", "description": "Notes"} + assert edited.json() == { + **original, + "title": "Edited", + "severity": "critical", + "field_values": {"description": "Notes"}, + } assert client.get(path).json()["items"] == [{**edited.json(), "evidence_count": 0}] deleted = client.delete(finding_path) assert deleted.status_code == 204 @@ -335,7 +338,8 @@ def test_finding_mutations_reject_wrong_operations_and_invalid_fields( finding_path = f"{path}/{original['id']}" for invalid in [ {**body, "title": " "}, - {**body, "severity": "urgent"}, + {"title": "Assessment", "severity": " "}, + {"title": "Assessment", "severity": "low", "description": "x"}, {**body, "operation_id": other_id}, {**body, "created_at": original["created_at"]}, ]: diff --git a/tests/unit/memory/test_finding_field_values_migration.py b/tests/unit/memory/test_finding_field_values_migration.py new file mode 100644 index 0000000000..0042d39573 --- /dev/null +++ b/tests/unit/memory/test_finding_field_values_migration.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import importlib +import io +import json +import re +from pathlib import Path + +import pytest +import sqlalchemy as sa +from alembic import command +from alembic.config import Config +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.dialects import mssql, sqlite +from sqlalchemy.exc import IntegrityError +from sqlalchemy.schema import CreateTable + +from pyrit.memory.memory_models import FindingEntry +from pyrit.memory.migration import check_schema_migrations + +PREVIOUS_REVISION = "c8d3e5f7a901" +REVISION = "d43b1bcf9c56" +CHECKS = ("ck_findings_severity", "ck_findings_severity_other", "ck_findings_harm_type_other") +OPERATION_ID = "00000000-0000-0000-0000-000000000001" +ATTACK_ID = "00000000-0000-0000-0000-0000000000aa" +A, B, C, D = (f"00000000-0000-0000-0000-00000000010{number}" for number in range(1, 5)) +STAMP = "2026-10-09 12:00:00" + + +def _config_for(connection: sa.Connection) -> Config: + config = Config() + config.set_main_option("script_location", str(Path(__file__).resolve().parents[3] / "pyrit" / "memory" / "alembic")) + config.attributes["connection"] = connection + config.attributes["version_table"] = "pyrit_memory_alembic_version" + return config + + +def _migration(): + return importlib.import_module("pyrit.memory.alembic.versions.d43b1bcf9c56_store_finding_field_values") + + +def _columns(inspector: sa.Inspector, table: str) -> list[tuple[str, str, bool]]: + return [(c["name"], str(c["type"]), c["nullable"]) for c in inspector.get_columns(table)] + + +def _insert_operation(connection: sa.Connection) -> None: + connection.execute( + sa.text( + "INSERT INTO OperationEntries (id, name, name_key, created_at, description) " + "VALUES (:id, 'Contoso review', 'contoso review', :stamp, '')" + ), + {"id": OPERATION_ID, "stamp": STAMP}, + ) + + +def _insert_finding(*, connection: sa.Connection, **values: str) -> None: + row = {"operation_id": OPERATION_ID, "title": "Contoso finding", "created_at": STAMP, **values} + connection.execute( + sa.text(f"INSERT INTO FindingEntries ({', '.join(row)}) VALUES ({', '.join(f':{name}' for name in row)})"), + row, + ) + + +def test_upgrade_moves_existing_findings_into_field_values() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, PREVIOUS_REVISION) + _insert_operation(connection) + _insert_finding(connection=connection, id=A, severity="important", description="") + _insert_finding( + connection=connection, + id=B, + severity="other", + severity_other="Sev 2.5", + harm_type="Other", + harm_type_other="Model theft", + description="Steps to reproduce", + ) + connection.execute( + sa.text( + "INSERT INTO FindingEvidenceEntries (id, finding_id, attack_result_id, attached_at) " + "VALUES ('00000000-0000-0000-0000-000000000201', :finding_id, :attack_id, :stamp)" + ), + {"finding_id": B, "attack_id": ATTACK_ID, "stamp": STAMP}, + ) + _insert_finding(connection=connection, id=C, severity="low", harm_type="Jailbreak", description=" ") + _insert_finding(connection=connection, id=D, severity="critical", description='Quote " and \\ backslash') + + command.upgrade(config, REVISION) + + rows = { + finding_id: (severity, json.loads(field_values)) + for finding_id, severity, field_values in connection.execute( + sa.text("SELECT id, severity, field_values FROM FindingEntries") + ) + } + assert rows == { + A: ("important", {}), + B: ("Sev 2.5", {"harm_type": "Model theft", "description": "Steps to reproduce"}), + C: ("low", {"harm_type": "Jailbreak"}), + D: ("critical", {"description": 'Quote " and \\ backslash'}), + } + evidence_rows = [ + tuple(row) + for row in connection.execute( + sa.text("SELECT finding_id, attack_result_id FROM FindingEvidenceEntries") + ) + ] + assert evidence_rows == [(B, ATTACK_ID)] + check_schema_migrations(engine=engine) + finally: + engine.dispose() + + +def test_upgrade_drops_the_finding_check_constraints() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + command.upgrade(_config_for(connection), REVISION) + _insert_operation(connection) + _insert_finding(connection=connection, id=A, severity="sev1", field_values='{"category": "x"}') + assert connection.execute(sa.text("SELECT severity, field_values FROM FindingEntries")).one() == ( + "sev1", + '{"category": "x"}', + ) + finally: + engine.dispose() + + +def test_downgrade_restores_the_old_columns_only_without_findings() -> None: + engine = sa.create_engine("sqlite:///:memory:") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, PREVIOUS_REVISION) + columns = sorted(_columns(sa.inspect(connection), "FindingEntries")) + command.upgrade(config, REVISION) + command.downgrade(config, PREVIOUS_REVISION) + assert sorted(_columns(sa.inspect(connection), "FindingEntries")) == columns + with pytest.raises(IntegrityError, match="CHECK"): + _insert_finding(connection=connection, id=A, severity="urgent", description="") + + command.upgrade(config, REVISION) + _insert_operation(connection) + _insert_finding(connection=connection, id=A, severity="low", field_values="{}") + with pytest.raises(ValueError, match="findings"): + command.downgrade(config, PREVIOUS_REVISION) + assert connection.scalar(sa.text("SELECT id FROM FindingEntries")) == A + finally: + engine.dispose() + + +def test_downgrade_refuses_offline() -> None: + context = MigrationContext.configure(dialect_name="mssql", opts={"as_sql": True, "output_buffer": io.StringIO()}) + with Operations.context(context), pytest.raises(ValueError, match="online"): + _migration().downgrade() + + +def test_upgrade_compiles_for_sql_server() -> None: + output = io.StringIO() + context = MigrationContext.configure(dialect_name="mssql", opts={"as_sql": True, "output_buffer": output}) + with Operations.context(context): + _migration().upgrade() + ddl = output.getvalue() + assert _migration().down_revision == "c8d3e5f7a901" + for name in CHECKS: + assert f"DROP CONSTRAINT {name}" in ddl + drops = [match.start() for match in re.finditer("DROP CONSTRAINT", ddl)] + assert len(drops) == 3 and max(drops) < ddl.index("ALTER COLUMN severity") + assert "COLLATE Latin1_General_100_BIN2" in ddl + assert "JSON_MODIFY" in ddl + for column in ("description", "severity_other", "harm_type", "harm_type_other"): + assert f"DROP COLUMN {column}" in ddl + + +def test_finding_severity_is_binary_collated_only_on_sql_server() -> None: + table = FindingEntry.__table__ + assert "severity NVARCHAR(max) COLLATE Latin1_General_100_BIN2 NOT NULL" in str( + CreateTable(table).compile(dialect=mssql.dialect()) + ) + assert "COLLATE" not in str(CreateTable(table).compile(dialect=sqlite.dialect())) diff --git a/tests/unit/memory/test_findings.py b/tests/unit/memory/test_findings.py index 45958f993a..85dd032e9c 100644 --- a/tests/unit/memory/test_findings.py +++ b/tests/unit/memory/test_findings.py @@ -9,7 +9,7 @@ import pytest from pydantic import ValidationError -from sqlalchemy import func, insert, select +from sqlalchemy import func, select from sqlalchemy.dialects import mssql from sqlalchemy.engine import ScalarResult from sqlalchemy.exc import IntegrityError @@ -25,7 +25,6 @@ ScoreEntry, ) from pyrit.models import Finding, FindingCreate, FindingEvidence, Operation -from pyrit.models.harm_category import HarmCategory pytestmark = pytest.mark.usefixtures("patch_central_database") @@ -125,34 +124,14 @@ async def test_operation_identity_follows_casefold_not_database_collation( assert len(await sqlite_instance.get_operations_async()) == (1 if same else 2) -@pytest.mark.parametrize( - ("stored", "harm_type", "harm_type_other"), - [("Malwar", HarmCategory.OTHER, "Malwar"), ("malware", HarmCategory.MALWARE, None)], -) -async def test_unknown_stored_harm_type_stays_readable( - sqlite_instance: SQLiteMemory, - caplog: pytest.LogCaptureFixture, - stored: str, - harm_type: HarmCategory, - harm_type_other: str | None, -) -> None: - operation = await sqlite_instance.add_operation_async(Operation(name="Written elsewhere")) - with sqlite_instance.engine.begin() as connection: - connection.execute( - insert(FindingEntry).values( - id=uuid4(), - operation_id=operation.id, - title="Stored outside PyRIT", - description="", - severity="low", - harm_type=stored, - created_at=datetime.now(UTC), - ) - ) - with caplog.at_level("WARNING"): - [finding] = await sqlite_instance.get_findings_async(operation_id=operation.id) - assert (finding.harm_type, finding.harm_type_other) == (harm_type, harm_type_other) - assert ("Unknown harm category 'Malwar'" in caplog.text) is (harm_type_other is not None) +async def test_stored_field_values_read_back_unchanged(sqlite_instance: SQLiteMemory) -> None: + operation = await sqlite_instance.add_operation_async(Operation(name="Retired values")) + field_values = {"harm_type": "Not a current category", "retired_key": "kept"} + saved = await sqlite_instance.add_finding_async( + finding=Finding(operation_id=operation.id, title="Stored", severity="low", field_values=field_values) + ) + read = await sqlite_instance.get_finding_async(operation_id=operation.id, finding_id=saved.id) + assert read.field_values == field_values async def test_findings_paginate_within_one_operation(sqlite_instance: SQLiteMemory) -> None: @@ -215,58 +194,43 @@ async def test_update_finding_preserves_identity_and_reorders_severity(sqlite_in updated = await sqlite_instance.update_finding_async( operation_id=operation.id, finding_id=original.id, - request=FindingCreate(title="Edited", severity="critical", description="Notes"), + request=FindingCreate(title="Edited", severity="critical", field_values={"description": "Notes"}), ) assert updated.id == original.id assert updated.operation_id == original.operation_id assert updated.created_at == original.created_at assert updated.title == "Edited" - assert updated.description == "Notes" + assert updated.field_values == {"description": "Notes"} assert [finding.title for finding in await sqlite_instance.get_findings_async(severity_order=SEVERITY_ORDER)] == [ "Edited", "Other", ] -async def test_finding_classifications_persist_and_clear_without_losing_evidence(sqlite_instance: SQLiteMemory) -> None: - operation = await sqlite_instance.add_operation_async(Operation(name="Classifications")) - original = Finding.model_validate( - { - "operation_id": operation.id, - "title": "Custom", - "severity": "other", - "severity_other": " Review needed ", - "harm_type": "Other", - "harm_type_other": " Team harm ", - } +async def test_update_finding_replaces_field_values_without_losing_evidence(sqlite_instance: SQLiteMemory) -> None: + operation = await sqlite_instance.add_operation_async(Operation(name="Field values")) + original = Finding(operation_id=operation.id, title="Values", severity="low", field_values={"a": "1", "b": "2"}) + saved = await sqlite_instance.add_finding_async( + finding=original, evidence=FindingEvidence(finding_id=original.id, attack_result_id=uuid4()) ) - saved = await sqlite_instance.add_finding_async(finding=original) assert saved == original - evidence = FindingEvidence(finding_id=saved.id, attack_result_id=uuid4()) - await sqlite_instance.add_finding_evidence_async(evidence=evidence) - read = await sqlite_instance.get_finding_async(operation_id=operation.id, finding_id=saved.id) - assert read == original updated = await sqlite_instance.update_finding_async( operation_id=operation.id, finding_id=saved.id, - request=FindingCreate(title="Canonical", severity="important", harm_type="Malware"), - ) - assert updated.id == saved.id - assert updated.created_at == saved.created_at - assert updated.severity_other is None - assert updated.harm_type.value == "Malware" - assert updated.harm_type_other is None - cleared = await sqlite_instance.update_finding_async( - operation_id=operation.id, finding_id=saved.id, request=FindingCreate(title="Optional", severity="low") + request=FindingCreate(title="Values", severity="low", field_values={"a": "1"}), ) - assert cleared.harm_type is None - assert await sqlite_instance.get_finding_evidence_async(finding_id=saved.id, limit=20, offset=0) == [evidence] + assert (updated.id, updated.created_at) == (saved.id, saved.created_at) + read = await sqlite_instance.get_finding_async(operation_id=operation.id, finding_id=saved.id) + assert read.field_values == {"a": "1"} + assert await sqlite_instance.get_finding_evidence_counts_async(finding_ids=[saved.id]) == {saved.id: 1} async def test_custom_severity_sorts_after_every_preset(sqlite_instance: SQLiteMemory) -> None: operation = await sqlite_instance.add_operation_async(Operation(name="Order")) stamp = datetime(2026, 10, 9, tzinfo=UTC) - for number, severity in enumerate(["other", "informational", "low", "moderate", "important", "critical"], start=1): + for number, severity in enumerate( + ["Sev 2.5", "informational", "low", "moderate", "important", "critical"], start=1 + ): await sqlite_instance.add_finding_async( finding=Finding.model_validate( { @@ -275,7 +239,6 @@ async def test_custom_severity_sorts_after_every_preset(sqlite_instance: SQLiteM "operation_id": operation.id, "title": severity, "severity": severity, - "severity_other": "Critical" if severity == "other" else None, } ) ) @@ -285,7 +248,7 @@ async def test_custom_severity_sorts_after_every_preset(sqlite_instance: SQLiteM "moderate", "low", "informational", - "other", + "Sev 2.5", ] diff --git a/tests/unit/memory/test_operations_and_findings_migration.py b/tests/unit/memory/test_operations_and_findings_migration.py index c9b0299cda..d8153542b8 100644 --- a/tests/unit/memory/test_operations_and_findings_migration.py +++ b/tests/unit/memory/test_operations_and_findings_migration.py @@ -19,7 +19,6 @@ from pyrit.memory.memory_models import ( AttackResultEntry, - FindingEntry, FindingEvidenceEntry, OperationEntry, ScenarioResultEntry, @@ -30,6 +29,24 @@ REVISION = "c8d3e5f7a901" TABLES = {"OperationEntries", "FindingEntries", "FindingEvidenceEntries"} ACTIVITY_TABLES = ("AttackResultEntries", "ScenarioResultEntries") +# FindingEntry models the head revision, which no longer has the #3066 finding columns. +FINDINGS = sa.table( + "FindingEntries", + *( + sa.column(name) + for name in ( + "id", + "operation_id", + "title", + "description", + "severity", + "severity_other", + "harm_type", + "harm_type_other", + "created_at", + ) + ), +) def _config_for(connection: sa.Connection) -> Config: @@ -65,6 +82,7 @@ def test_migration_adds_tables_without_changing_existing_ones() -> None: assert {name: columns for name, columns in after.items() if name not in ACTIVITY_TABLES} == { name: columns for name, columns in existing.items() if name not in ACTIVITY_TABLES } + command.upgrade(config, "head") check_schema_migrations(engine=engine) finally: engine.dispose() @@ -128,28 +146,28 @@ def test_migration_enforces_unique_evidence_source() -> None: ], ) def test_migration_rejects_inconsistent_finding_classifications(invalid: dict[str, str]) -> None: - valid = {"title": "Finding", "description": "", "severity": "low", "created_at": datetime.now(UTC)} + valid = {"title": "Finding", "description": "", "severity": "low", "created_at": "2026-10-09 12:00:00"} engine = sa.create_engine("sqlite:///:memory:") try: with engine.begin() as connection: command.upgrade(_config_for(connection), REVISION) - insert = sa.insert(FindingEntry) - connection.execute(insert.values(id=uuid4(), operation_id=uuid4(), **valid)) + insert = sa.insert(FINDINGS) + connection.execute(insert.values(id=str(uuid4()), operation_id=str(uuid4()), **valid)) # Harm types are unconstrained so HarmCategory can change without a migration. connection.execute( - insert.values(id=uuid4(), operation_id=uuid4(), **valid, harm_type="Not a current category") + insert.values(id=str(uuid4()), operation_id=str(uuid4()), **valid, harm_type="Not a current category") ) connection.execute( insert.values( - id=uuid4(), - operation_id=uuid4(), + id=str(uuid4()), + operation_id=str(uuid4()), **{**valid, "severity": "other", "severity_other": "Custom"}, harm_type="Other", harm_type_other="Custom", ) ) with pytest.raises(IntegrityError, match="CHECK"): - connection.execute(insert.values(id=uuid4(), operation_id=uuid4(), **{**valid, **invalid})) + connection.execute(insert.values(id=str(uuid4()), operation_id=str(uuid4()), **{**valid, **invalid})) finally: engine.dispose() diff --git a/tests/unit/models/test_finding.py b/tests/unit/models/test_finding.py index e334e332fb..7a5dbb39d8 100644 --- a/tests/unit/models/test_finding.py +++ b/tests/unit/models/test_finding.py @@ -7,79 +7,41 @@ from pydantic import ValidationError from pyrit.models.finding import Finding, FindingCreate -from pyrit.models.harm_category import HarmCategory OPERATION_ID = UUID(int=1) -@pytest.mark.parametrize("severity", ["critical", "important", "moderate", "low", "informational"]) -def test_finding_round_trips_named_severity(severity: str) -> None: - finding = Finding(operation_id=OPERATION_ID, title="Assessment", severity=severity) +def test_finding_round_trips_severity_and_field_values() -> None: + finding = Finding( + operation_id=OPERATION_ID, title="Assessment", severity="moderate", field_values={"harm_type": "Malware"} + ) assert Finding.model_validate_json(finding.model_dump_json()) == finding - assert finding.description == "" + assert FindingCreate(title="Assessment", severity="low").field_values == {} assert finding.created_at.utcoffset().total_seconds() == 0 -def test_finding_round_trips_custom_classifications_without_normalizing_text() -> None: - finding = Finding.model_validate( - { - "operation_id": OPERATION_ID, - "title": "Assessment", - "severity": "other", - "severity_other": " Team assessment ", - "harm_type": "Other", - "harm_type_other": " Team harm ", - } +def test_finding_drops_blank_field_values_and_keeps_other_text_exactly() -> None: + finding = FindingCreate( + title="Leak", severity=" Sev 2.5 ", field_values={"steps": "", "owner": " ", "notes": " kept "} ) - assert finding.severity.value == "other" - assert finding.severity_other == " Team assessment " - assert finding.harm_type.value == "Other" - assert finding.harm_type_other == " Team harm " - assert Finding.model_validate_json(finding.model_dump_json()) == finding - - -def test_finding_harm_type_is_optional_and_accepts_every_canonical_category() -> None: - assert FindingCreate(title="Optional", severity="low").harm_type is None - for category in HarmCategory: - finding = FindingCreate.model_validate( - { - "title": "Classified", - "severity": "moderate", - "harm_type": category.value, - "harm_type_other": "Custom harm" if category == HarmCategory.OTHER else None, - } - ) - assert finding.harm_type == category + assert finding.severity == " Sev 2.5 " + assert finding.field_values == {"notes": " kept "} @pytest.mark.parametrize( - "fields", + "data", [ - {"severity": "other"}, - {"severity": "other", "severity_other": " "}, - {"severity": "low", "severity_other": "Stray custom severity"}, - {"harm_type": "Other"}, - {"harm_type": "Other", "harm_type_other": " "}, - {"harm_type": "Malware", "harm_type_other": "Stray custom harm"}, - {"harm_type_other": "Custom without category"}, - {"harm_type": "Unknown category"}, + {"title": " ", "severity": "low"}, + {"title": "Leak", "severity": " "}, + {"title": "Leak", "severity": "low", "description": "a #3066 field"}, + {"title": "Leak", "severity": "low", "field_values": {"steps": 3}}, ], ) -def test_finding_rejects_inconsistent_classifications(fields: dict[str, str]) -> None: +def test_finding_rejects_invalid_input(data: dict[str, object]) -> None: with pytest.raises(ValidationError): - FindingCreate.model_validate({"title": "Assessment", "severity": "low", **fields}) + FindingCreate.model_validate(data) -@pytest.mark.parametrize("title", ["", " "]) -def test_finding_rejects_blank_title(title: str) -> None: - with pytest.raises(ValidationError): - FindingCreate(title=title, severity="low") - - -def test_finding_rejects_unknown_input_and_requires_operation() -> None: - with pytest.raises(ValidationError): - FindingCreate(title="Assessment", severity="urgent") - with pytest.raises(ValidationError): - FindingCreate.model_validate({"title": "Assessment", "severity": "low", "operation": "Case"}) +def test_finding_requires_an_operation() -> None: with pytest.raises(ValidationError): Finding(title="Assessment", severity="low") From 1f206df85ed8675bdbf039315209625934104b0e Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Tue, 27 Oct 2026 11:08:29 -0400 Subject: [PATCH 05/20] Check findings against the active finding fields Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/models/operations.py | 7 - pyrit/backend/routes/operations.py | 32 +++-- pyrit/backend/services/operation_service.py | 12 ++ tests/unit/backend/test_operations.py | 126 +++++++++++++++++- tests/unit/conftest.py | 10 +- tests/unit/setup/test_finding_fields.py | 10 +- .../setup/test_finding_fields_initializer.py | 11 +- tests/unit/setup/test_initialization.py | 13 +- 8 files changed, 166 insertions(+), 55 deletions(-) diff --git a/pyrit/backend/models/operations.py b/pyrit/backend/models/operations.py index 4805101c5c..6062ce1f32 100644 --- a/pyrit/backend/models/operations.py +++ b/pyrit/backend/models/operations.py @@ -15,16 +15,9 @@ Operation, OperationUsage, ) -from pyrit.models.harm_category import HarmCategory from pyrit.models.operation import OperationName -class FindingOptionsResponse(BaseModel): - """Harm categories a finding can use.""" - - harm_types: list[HarmCategory] - - class OperationListItem(Operation): """An operation with its finding counts by stored severity value. Severities with no findings are omitted.""" diff --git a/pyrit/backend/routes/operations.py b/pyrit/backend/routes/operations.py index 69894dcab3..610bd54158 100644 --- a/pyrit/backend/routes/operations.py +++ b/pyrit/backend/routes/operations.py @@ -14,7 +14,6 @@ FindingEvidenceListResponse, FindingExportRequest, FindingListResponse, - FindingOptionsResponse, OperationDetail, OperationListResponse, OperationUpdateRequest, @@ -29,8 +28,8 @@ OperationService, ) from pyrit.memory.memory_interface import OperationInUseError -from pyrit.models import Finding, FindingCreate, Operation, OperationCreate -from pyrit.models.harm_category import HarmCategory +from pyrit.models import Finding, FindingCreate, FindingFields, FindingValuesError, Operation, OperationCreate +from pyrit.setup.finding_fields import get_active_finding_fields T = TypeVar("T") @@ -65,15 +64,15 @@ async def list_operations_async() -> OperationListResponse: return OperationListResponse(items=await OperationService().list_async()) -@router.get("/finding-options", response_model=FindingOptionsResponse) -async def get_finding_options_async() -> FindingOptionsResponse: +@router.get("/finding-fields", response_model=FindingFields) +async def get_finding_fields_async() -> FindingFields: """ - List the harm categories a finding can use. + Read the finding fields that this deployment's setup configured. Returns: - FindingOptionsResponse: Selectable categories, including Other. + FindingFields: The severity slot and the fields shown after it. """ - return FindingOptionsResponse(harm_types=list(HarmCategory)) + return get_active_finding_fields() @router.get("/{operation_id}", response_model=OperationDetail) @@ -137,7 +136,8 @@ async def create_operation_finding_async(*, operation_id: UUID, request: Finding Finding: The saved finding. Raises: - HTTPException: 409 when the ID already belongs to a different finding. + HTTPException: 409 when the ID already belongs to a different finding, and 422 when the values + do not fit the active finding fields. """ try: return await _require_operation_async( @@ -145,6 +145,8 @@ async def create_operation_finding_async(*, operation_id: UUID, request: Finding ) except FindingConflictError as exc: raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc + except FindingValuesError as exc: + raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=str(exc)) from exc @router.get("/{operation_id}/findings", response_model=FindingListResponse) @@ -191,10 +193,16 @@ async def update_operation_finding_async(*, operation_id: UUID, finding_id: UUID Returns: Finding: The updated finding. + + Raises: + HTTPException: 422 when the values do not fit the active finding fields. """ - return await _require_operation_async( - OperationService().update_finding_async(operation_id=operation_id, finding_id=finding_id, request=request) - ) + try: + return await _require_operation_async( + OperationService().update_finding_async(operation_id=operation_id, finding_id=finding_id, request=request) + ) + except FindingValuesError as exc: + raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=str(exc)) from exc @router.delete("/{operation_id}/findings/{finding_id}", status_code=status.HTTP_204_NO_CONTENT) diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py index 4b1621b310..6c002e9843 100644 --- a/pyrit/backend/services/operation_service.py +++ b/pyrit/backend/services/operation_service.py @@ -163,8 +163,10 @@ async def create_finding_async(self, *, operation_id: UUID, request: FindingCrea Raises: FindingConflictError: If the ID already belongs to a different finding. + FindingValuesError: If the values do not fit the active finding fields. """ await self.get_async(operation_id) + get_active_finding_fields().check_values(severity=request.severity, field_values=request.field_values) finding = Finding( operation_id=operation_id, id=request.id or uuid4(), @@ -435,13 +437,23 @@ async def update_finding_async(self, *, operation_id: UUID, finding_id: UUID, re """ Replace a finding's editable fields within an existing operation. + Values that have not changed since the last save may use fields or options that no longer exist. + Returns: Finding: The updated finding. Raises: FindingNotFoundError: If the finding is missing or belongs to another operation. + FindingValuesError: If the values do not fit the active finding fields. """ await self.get_async(operation_id) + stored = await self._require_finding_async(operation_id=operation_id, finding_id=finding_id) + get_active_finding_fields().check_values( + severity=request.severity, + field_values=request.field_values, + previous_severity=stored.severity, + previous_values=stored.field_values, + ) try: return await CentralMemory.get_memory_instance().update_finding_async( operation_id=operation_id, finding_id=finding_id, request=request diff --git a/tests/unit/backend/test_operations.py b/tests/unit/backend/test_operations.py index c631e11d9e..3147ad153a 100644 --- a/tests/unit/backend/test_operations.py +++ b/tests/unit/backend/test_operations.py @@ -18,6 +18,7 @@ from pyrit.memory import CentralMemory, SQLiteMemory from pyrit.memory.memory_models import AttackResultEntry, Base, PromptMemoryEntry from pyrit.models import ( + DEFAULT_FINDING_FIELDS, AttackOutcome, AttackResult, ComponentIdentifier, @@ -25,17 +26,25 @@ ConversationType, Finding, FindingEvidence, + FindingField, + FindingFields, Message, MessagePiece, Operation, Score, ) -from pyrit.models.harm_category import HarmCategory +from pyrit.setup.finding_fields import set_active_finding_fields pytestmark = pytest.mark.usefixtures("patch_central_database") MISSING_ID = UUID(int=404) EMPTY_USAGE = {"findings": 0, "attacks": 0, "scenario_runs": 0} +CONTOSO = FindingFields( + severity=FindingField( + key="severity", label="Severity", type="choice", options=[("sev1", "Sev 1"), ("sev2", "Sev 2")] + ), + fields=[FindingField(key="steps", label="Steps to reproduce", type="multiline", required=True)], +) def test_operations_create_list_and_get(compatibility_headers: dict[str, str]) -> None: @@ -156,13 +165,56 @@ def test_findings_list_highest_severity_first(compatibility_headers: dict[str, s assert [item["title"] for item in listed] == ["First", "Second"] -def test_finding_options_exposes_canonical_harm_types(compatibility_headers: dict[str, str]) -> None: +@pytest.mark.usefixtures("reset_finding_fields") +async def test_findings_list_in_the_active_severity_order_async( + sqlite_instance: SQLiteMemory, compatibility_headers: dict[str, str] +) -> None: + set_active_finding_fields(finding_fields=CONTOSO, source="test") client = TestClient(app, headers=compatibility_headers) - response = client.get("/api/operations/finding-options") - assert response.status_code == 200 - assert response.json() == {"harm_types": [category.value for category in HarmCategory]} - assert response.json()["harm_types"].count("Other") == 1 - assert TestClient(app).get("/api/operations/finding-options").status_code != 200 + operation_id = client.post("/api/operations", json={"name": "Contoso"}).json()["id"] + path = f"/api/operations/{operation_id}/findings" + steps = {"steps": "1. Open"} + assert client.post(path, json={"title": "Second", "severity": "sev2", "field_values": steps}).status_code == 201 + await sqlite_instance.add_finding_async( + finding=Finding(operation_id=UUID(operation_id), title="Unmatched", severity="Sev 2.5", field_values=steps) + ) + assert client.post(path, json={"title": "First", "severity": "sev1", "field_values": steps}).status_code == 201 + listed = client.get(path).json()["items"] + assert [item["severity"] for item in listed] == ["sev1", "sev2", "Sev 2.5"] + + +@pytest.mark.usefixtures("reset_finding_fields") +def test_finding_fields_returns_the_active_fields(compatibility_headers: dict[str, str]) -> None: + client = TestClient(app, headers=compatibility_headers) + assert FindingFields.model_validate(client.get("/api/operations/finding-fields").json()) == DEFAULT_FINDING_FIELDS + set_active_finding_fields(finding_fields=CONTOSO, source="test") + assert FindingFields.model_validate(client.get("/api/operations/finding-fields").json()) == CONTOSO + assert client.get("/api/operations/finding-options").status_code != 200 + assert TestClient(app).get("/api/operations/finding-fields").status_code != 200 + + +@pytest.mark.usefixtures("reset_finding_fields") +@pytest.mark.parametrize( + ("body", "detail"), + [ + ({"severity": "sev1"}, "Steps to reproduce is required."), + ({"severity": "sev9", "field_values": {"steps": "1. Open"}}, "Severity has no option 'sev9'."), + ({"severity": "sev1", "field_values": {"steps": "1. Open", "owner": "dana"}}, "Unknown finding field 'owner'."), + ], +) +def test_finding_create_is_checked_against_the_active_fields( + body: dict[str, object], detail: str, compatibility_headers: dict[str, str] +) -> None: + set_active_finding_fields(finding_fields=CONTOSO, source="test") + client = TestClient(app, headers=compatibility_headers) + operation_id = client.post("/api/operations", json={"name": "Contoso"}).json()["id"] + path = f"/api/operations/{operation_id}/findings" + rejected = client.post(path, json={"title": "Leak", **body}) + assert rejected.status_code == 422 + assert rejected.json()["detail"] == detail + assert client.get(path).json()["items"] == [] + created = client.post(path, json={"title": "Leak", "severity": "sev1", "field_values": {"steps": "1. Open"}}) + assert created.status_code == 201 def test_finding_api_round_trips_field_values_and_drops_blanks(compatibility_headers: dict[str, str]) -> None: @@ -349,6 +401,66 @@ def test_finding_mutations_reject_wrong_operations_and_invalid_fields( assert client.get(path).json()["items"] == [{**original, "evidence_count": 0}] +@pytest.fixture +async def retired_value_finding_async(sqlite_instance: SQLiteMemory, reset_finding_fields: None) -> Finding: + """Seed a finding with a ``description`` value, then make CONTOSO, which has no such field, active.""" + operation = await sqlite_instance.add_operation_async(Operation(name="Contoso")) + finding = await sqlite_instance.add_finding_async( + finding=Finding( + operation_id=operation.id, + title="Leak", + severity="sev1", + field_values={"steps": "1. Open", "description": "Old text"}, + ) + ) + set_active_finding_fields(finding_fields=CONTOSO, source="test") + return finding + + +async def test_update_keeps_an_unchanged_retired_value_async( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, compatibility_headers: dict[str, str] +) -> None: + finding = retired_value_finding_async + response = TestClient(app, headers=compatibility_headers).put( + f"/api/operations/{finding.operation_id}/findings/{finding.id}", + json={"title": "Leak", "severity": "sev1", "field_values": {"steps": "1. Open", "description": "Old text"}}, + ) + assert response.status_code == 200 + stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) + assert stored is not None + assert stored.field_values == {"steps": "1. Open", "description": "Old text"} + + +async def test_update_rejects_a_changed_retired_value_async( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, compatibility_headers: dict[str, str] +) -> None: + finding = retired_value_finding_async + client = TestClient(app, headers=compatibility_headers) + path = f"/api/operations/{finding.operation_id}/findings" + body = {"title": "Leak", "severity": "sev1", "field_values": {"steps": "1. Open", "description": "New text"}} + response = client.put(f"{path}/{finding.id}", json=body) + assert response.status_code == 422 + assert response.json()["detail"] == "Unknown finding field 'description'." + stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) + assert stored is not None + assert stored.field_values == {"steps": "1. Open", "description": "Old text"} + assert client.put(f"{path}/{MISSING_ID}", json=body).status_code == 404 + + +async def test_update_removes_an_omitted_retired_value_async( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, compatibility_headers: dict[str, str] +) -> None: + finding = retired_value_finding_async + response = TestClient(app, headers=compatibility_headers).put( + f"/api/operations/{finding.operation_id}/findings/{finding.id}", + json={"title": "Leak", "severity": "sev1", "field_values": {"steps": "1. Open"}}, + ) + assert response.status_code == 200 + stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) + assert stored is not None + assert stored.field_values == {"steps": "1. Open"} + + @pytest.fixture async def evidence_source_async(sqlite_instance: SQLiteMemory) -> tuple[Operation, Finding, AttackResult]: return await _store_evidence_source_async(sqlite_instance) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 544741a1c2..d69bfdf01d 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -7,7 +7,7 @@ import tempfile import time import weakref -from collections.abc import AsyncGenerator, Generator +from collections.abc import AsyncGenerator, Generator, Iterator from contextlib import closing, contextmanager from typing import Any from unittest.mock import AsyncMock, patch @@ -22,6 +22,7 @@ from pyrit.datasets import SeedDatasetProvider from pyrit.memory import CentralMemory, SQLiteMemory, migration from pyrit.models import SeedDataset +from pyrit.setup.finding_fields import reset_active_finding_fields from test_support.fixtures import use_sqlite_memory # This limits retries and speeds up execution @@ -171,3 +172,10 @@ def patch_central_database(sqlite_instance): """Fixture to mock CentralMemory.get_memory_instance""" with patch.object(CentralMemory, "get_memory_instance", return_value=sqlite_instance) as sqlite_memory: yield sqlite_memory + + +@pytest.fixture +def reset_finding_fields() -> Iterator[None]: + """Restore the default finding fields after the test.""" + yield + reset_active_finding_fields() diff --git a/tests/unit/setup/test_finding_fields.py b/tests/unit/setup/test_finding_fields.py index be8dc16a04..6781f7269b 100644 --- a/tests/unit/setup/test_finding_fields.py +++ b/tests/unit/setup/test_finding_fields.py @@ -1,8 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from collections.abc import Iterator - import pytest from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields @@ -12,17 +10,13 @@ set_active_finding_fields, ) +pytestmark = pytest.mark.usefixtures("reset_finding_fields") + CUSTOM = FindingFields( severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) ) -@pytest.fixture(autouse=True) -def reset_finding_fields() -> Iterator[None]: - yield - reset_active_finding_fields() - - def test_active_finding_fields_default_to_3066() -> None: assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS diff --git a/tests/unit/setup/test_finding_fields_initializer.py b/tests/unit/setup/test_finding_fields_initializer.py index 955ffc7e4a..55da1c7b62 100644 --- a/tests/unit/setup/test_finding_fields_initializer.py +++ b/tests/unit/setup/test_finding_fields_initializer.py @@ -2,7 +2,6 @@ # Licensed under the MIT license. import importlib.util -from collections.abc import Iterator from pathlib import Path import pytest @@ -11,10 +10,12 @@ from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields from pyrit.registry import InitializerRegistry from pyrit.registry.discovery import discover_in_directory -from pyrit.setup.finding_fields import get_active_finding_fields, reset_active_finding_fields +from pyrit.setup.finding_fields import get_active_finding_fields from pyrit.setup.initializers import FindingFieldsInitializer from pyrit.setup.pyrit_initializer import PyRITInitializer +pytestmark = pytest.mark.usefixtures("reset_finding_fields") + EXAMPLE = FindingFields( severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) ) @@ -41,12 +42,6 @@ def get_finding_fields(self) -> FindingFields: return EXAMPLE -@pytest.fixture(autouse=True) -def reset_finding_fields() -> Iterator[None]: - yield - reset_active_finding_fields() - - async def test_finding_fields_initializer_sets_the_active_fields() -> None: await ExampleFindingFields().initialize_async() assert get_active_finding_fields() is EXAMPLE diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index 2525b79883..bbb1f80356 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -5,7 +5,6 @@ import os import pathlib import tempfile -from collections.abc import Iterator from unittest import mock import pytest @@ -18,11 +17,7 @@ from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields, MessagePiece from pyrit.registry import InitializerRegistry from pyrit.setup import IN_MEMORY, initialization, initialize_pyrit_async -from pyrit.setup.finding_fields import ( - get_active_finding_fields, - reset_active_finding_fields, - set_active_finding_fields, -) +from pyrit.setup.finding_fields import get_active_finding_fields, set_active_finding_fields from pyrit.setup.initialization import reinitialize_pyrit_async from pyrit.setup.initializers import FindingFieldsInitializer from pyrit.setup.pyrit_initializer import PyRITInitializer @@ -460,12 +455,6 @@ async def initialize_async(self) -> None: set_active_finding_fields(finding_fields=CUSTOM, source="pass") -@pytest.fixture -def reset_finding_fields() -> Iterator[None]: - yield - reset_active_finding_fields() - - @pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") async def test_each_setup_pass_resets_finding_fields() -> None: with mock.patch("pyrit.setup.initialization.load_environment_async", new_callable=mock.AsyncMock): From 4b99c8b95ae07ee18c5336aa7df0d708dbebd9f2 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Tue, 27 Oct 2026 11:17:33 -0400 Subject: [PATCH 06/20] Reject invalid finding fields during live-apply preflight Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/setup/configuration_loader.py | 27 ++++++++++++- tests/unit/setup/test_reinitialization.py | 48 +++++++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index baf670ca3a..f8ec743ffa 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -91,6 +91,29 @@ class _ConfigurationLayer: reset_extensions: bool +def _check_finding_fields_initializers(initializers: Sequence["PyRITInitializer"]) -> None: + """ + Check the finding fields the initializers would set, without making them active. + + Args: + initializers (Sequence[PyRITInitializer]): The initializers of the replacement setup pass. + + Raises: + ValueError: If two or more initializers set the finding fields, or the one that does returns + invalid fields. + """ + from pyrit.setup.initializers import FindingFieldsInitializer + + setters = [initializer for initializer in initializers if isinstance(initializer, FindingFieldsInitializer)] + if len(setters) > 1: + first, second = (type(setter).__name__ for setter in setters[:2]) + raise ValueError( + f"{first} and {second} both set the finding fields. Use one FindingFieldsInitializer per deployment." + ) + if setters: + setters[0].get_finding_fields() + + @dataclass class ConfigurationLoader(YamlLoadable): """ @@ -706,7 +729,8 @@ async def preflight_reinitialization_async( PreparedReinitialization: Immutable inputs ready for the mutation phase. Raises: - ValueError: If an initializer is invalid or a required environment value is missing. + ValueError: If an initializer is invalid, a required environment value is missing, or the + finding fields are invalid or set by more than one initializer. """ from pyrit.registry import InitializerRegistry from pyrit.setup.initializers.targets import TargetInitializer @@ -744,6 +768,7 @@ async def preflight_reinitialization_async( raise ValueError( f"Initializer '{type(initializer).__name__}' has missing required environment variables." ) + _check_finding_fields_initializers(initializers) return PreparedReinitialization( environment_values=dict(environment_values), diff --git a/tests/unit/setup/test_reinitialization.py b/tests/unit/setup/test_reinitialization.py index d8f6129d57..29623278c2 100644 --- a/tests/unit/setup/test_reinitialization.py +++ b/tests/unit/setup/test_reinitialization.py @@ -15,6 +15,7 @@ from pyrit.common.apply_defaults import get_global_default_values from pyrit.common.random_context import configure_random_seed, get_configured_random_seed from pyrit.memory import AzureSQLMemory, CentralMemory, SQLiteMemory +from pyrit.models import DEFAULT_FINDING_FIELDS from pyrit.models.catalog import ScenarioSuite from pyrit.registry import ( AttackRegistry, @@ -26,6 +27,7 @@ ) from pyrit.setup.configuration_loader import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async +from pyrit.setup.finding_fields import get_active_finding_fields from pyrit.setup.initialization import reset_setup_registries, validate_reinitialization_memory from pyrit.setup.initializers.targets import TARGET_CONFIGS @@ -113,6 +115,52 @@ async def test_preflight_uses_isolated_registry_without_changing_live_state() -> assert prepared.script_initializer_types == () +@pytest.mark.usefixtures("reset_finding_fields") +async def test_preflight_rejects_two_finding_fields_initializers(tmp_path: Path) -> None: + script = tmp_path / "two_finding_fields.py" + script.write_text( + """ +from pyrit.models import FindingField, FindingFields +from pyrit.setup.initializers import FindingFieldsInitializer + +SEVERITY = FindingField(key="severity", label="Severity", type="choice", options=[("sev1", "Sev 1")]) + +class ContosoFindingFields(FindingFieldsInitializer): + def get_finding_fields(self): return FindingFields(severity=SEVERITY) + +class ContosoBackupFields(FindingFieldsInitializer): + def get_finding_fields(self): return FindingFields(severity=SEVERITY) +""", + encoding="utf-8", + ) + config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], initialization_scripts=[str(script)]) + + with pytest.raises(ValueError, match="Use one FindingFieldsInitializer per deployment"): + await config.preflight_reinitialization_async(environment_values={}) + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + +@pytest.mark.usefixtures("reset_finding_fields") +async def test_preflight_rejects_invalid_finding_fields(tmp_path: Path) -> None: + script = tmp_path / "invalid_finding_fields.py" + script.write_text( + """ +from pyrit.models import FindingField, FindingFields +from pyrit.setup.initializers import FindingFieldsInitializer + +class ContosoFindingFields(FindingFieldsInitializer): + def get_finding_fields(self): + return FindingFields(severity=FindingField(key="severity", label="Severity", type="choice", options=[])) +""", + encoding="utf-8", + ) + config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], initialization_scripts=[str(script)]) + + with pytest.raises(ValueError, match="needs at least one option"): + await config.preflight_reinitialization_async(environment_values={}) + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize("source", ["script", "configured"]) async def test_apply_reconstructs_initializers_after_environment_change(tmp_path: Path, source: str) -> None: From 1f8419f83eadc2398e275d9bfa56c4ac61034ef6 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Wed, 28 Oct 2026 14:36:18 -0400 Subject: [PATCH 07/20] Address the Council review of the finding fields backend Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/models/operations.py | 8 ++ pyrit/backend/routes/operations.py | 17 ++- pyrit/backend/services/operation_service.py | 49 ++++++-- pyrit/memory/memory_interface.py | 2 +- pyrit/memory/memory_models.py | 2 +- pyrit/models/finding_fields.py | 74 +++++++++--- pyrit/setup/configuration_loader.py | 25 ++-- pyrit/setup/finding_fields.py | 3 + tests/unit/backend/test_operations.py | 36 +++++- .../test_finding_field_values_migration.py | 7 ++ .../memory/test_operation_service_findings.py | 107 ++++++++++++++++++ .../test_operations_and_findings_migration.py | 3 +- tests/unit/models/test_finding.py | 2 +- tests/unit/models/test_finding_fields.py | 59 +++++++++- tests/unit/setup/test_finding_fields.py | 8 ++ tests/unit/setup/test_reinitialization.py | 43 ++++++- 16 files changed, 399 insertions(+), 46 deletions(-) create mode 100644 tests/unit/memory/test_operation_service_findings.py diff --git a/pyrit/backend/models/operations.py b/pyrit/backend/models/operations.py index 6062ce1f32..65f42ba27f 100644 --- a/pyrit/backend/models/operations.py +++ b/pyrit/backend/models/operations.py @@ -70,6 +70,14 @@ class FindingCreateRequest(FindingCreate): attack_result_id: UUID | None = Field(None, description="A saved attack result to link as evidence") +class FindingUpdateRequest(FindingCreate): + """A finding's new editable values, and any retired values to remove.""" + + remove_retired_keys: list[str] = Field( + default_factory=list, description="Retired values to delete from the finding" + ) + + class FindingEvidenceCreateRequest(BaseModel): """The saved attack result to link.""" diff --git a/pyrit/backend/routes/operations.py b/pyrit/backend/routes/operations.py index 610bd54158..510ea80e3a 100644 --- a/pyrit/backend/routes/operations.py +++ b/pyrit/backend/routes/operations.py @@ -14,6 +14,7 @@ FindingEvidenceListResponse, FindingExportRequest, FindingListResponse, + FindingUpdateRequest, OperationDetail, OperationListResponse, OperationUpdateRequest, @@ -28,7 +29,7 @@ OperationService, ) from pyrit.memory.memory_interface import OperationInUseError -from pyrit.models import Finding, FindingCreate, FindingFields, FindingValuesError, Operation, OperationCreate +from pyrit.models import Finding, FindingFields, FindingValuesError, Operation, OperationCreate from pyrit.setup.finding_fields import get_active_finding_fields T = TypeVar("T") @@ -67,7 +68,7 @@ async def list_operations_async() -> OperationListResponse: @router.get("/finding-fields", response_model=FindingFields) async def get_finding_fields_async() -> FindingFields: """ - Read the finding fields that this deployment's setup configured. + Read the finding fields this deployment uses. Returns: FindingFields: The severity slot and the fields shown after it. @@ -130,7 +131,7 @@ async def create_operation_finding_async(*, operation_id: UUID, request: Finding """ Record a finding within an operation, optionally with one saved attack result as evidence. - Retrying with the same ID and content returns the saved finding. + Retrying with the same ID and content returns the saved finding, as long as it still fits the active finding fields. Returns: Finding: The saved finding. @@ -187,9 +188,15 @@ async def export_operation_findings_async(*, operation_id: UUID, request: Findin @router.put("/{operation_id}/findings/{finding_id}", response_model=Finding) -async def update_operation_finding_async(*, operation_id: UUID, finding_id: UUID, request: FindingCreate) -> Finding: +async def update_operation_finding_async( + *, operation_id: UUID, finding_id: UUID, request: FindingUpdateRequest +) -> Finding: """ - Replace the editable fields of a finding in its operation. + Replace a finding's current field values in its operation. + + A current field left out is cleared. Values of fields that are no longer configured are kept and can't + change; list them in ``remove_retired_keys`` to delete them. Values that haven't changed may use options + that no longer exist. Returns: Finding: The updated finding. diff --git a/pyrit/backend/services/operation_service.py b/pyrit/backend/services/operation_service.py index 6c002e9843..d14876071e 100644 --- a/pyrit/backend/services/operation_service.py +++ b/pyrit/backend/services/operation_service.py @@ -14,6 +14,7 @@ FindingEvidenceListResponse, FindingListItem, FindingListResponse, + FindingUpdateRequest, OperationDetail, OperationListItem, OperationUpdateRequest, @@ -25,6 +26,8 @@ Finding, FindingCreate, FindingEvidence, + FindingFields, + FindingValuesError, Operation, OperationCreate, ) @@ -433,34 +436,66 @@ async def _require_evidence_source_async(*, attack_result_id: UUID) -> None: if not any(stat.message_count > 0 for stat in stats.values()): raise FindingEvidenceNotFoundError("The attack result has no saved messages.") - async def update_finding_async(self, *, operation_id: UUID, finding_id: UUID, request: FindingCreate) -> Finding: + async def update_finding_async( + self, *, operation_id: UUID, finding_id: UUID, request: FindingUpdateRequest + ) -> Finding: """ - Replace a finding's editable fields within an existing operation. + Replace a finding's current field values within an existing operation. - Values that have not changed since the last save may use fields or options that no longer exist. + A current field left out is cleared. Values of fields that are no longer configured are kept and can't + change; list them in ``remove_retired_keys`` to delete them. Values that haven't changed may use options + that no longer exist. Returns: Finding: The updated finding. Raises: FindingNotFoundError: If the finding is missing or belongs to another operation. - FindingValuesError: If the values do not fit the active finding fields. + FindingValuesError: If ``remove_retired_keys`` lists a key that isn't a retired value of the finding + or that is also sent, or the values do not fit the active finding fields. """ await self.get_async(operation_id) stored = await self._require_finding_async(operation_id=operation_id, finding_id=finding_id) - get_active_finding_fields().check_values( + fields = get_active_finding_fields() + values = self._merge_retired_values(fields=fields, stored=stored, request=request) + fields.check_values( severity=request.severity, - field_values=request.field_values, + field_values=values, previous_severity=stored.severity, previous_values=stored.field_values, ) + replacement = FindingCreate(title=request.title, severity=request.severity, field_values=values) try: return await CentralMemory.get_memory_instance().update_finding_async( - operation_id=operation_id, finding_id=finding_id, request=request + operation_id=operation_id, finding_id=finding_id, request=replacement ) except LookupError as exc: raise FindingNotFoundError(str(exc)) from exc + @staticmethod + def _merge_retired_values( + *, fields: FindingFields, stored: Finding, request: FindingUpdateRequest + ) -> dict[str, str]: + """ + Add the stored values of fields that are no longer configured to the sent values. + + Returns: + dict[str, str]: The sent values, plus each retired value not listed in ``remove_retired_keys``. + + Raises: + FindingValuesError: If ``remove_retired_keys`` lists a key that isn't a retired value of the finding + or that is also sent. + """ + known = fields.field_keys + retired = {key: value for key, value in stored.field_values.items() if key not in known} + for key in request.remove_retired_keys: + if key not in retired: + raise FindingValuesError(f"{key!r} is not a retired value of this finding.") + if key in request.field_values: + raise FindingValuesError(f"{key!r} can't be both sent and removed.") + kept = {key: value for key, value in retired.items() if key not in request.remove_retired_keys} + return {**kept, **request.field_values} + async def delete_finding_async(self, *, operation_id: UUID, finding_id: UUID) -> None: """ Permanently remove a finding within an existing operation. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index ca6381f859..b1354e7f71 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1056,7 +1056,7 @@ async def get_findings_async( title_query (str | None): Case-insensitive text the title must contain. Wildcard characters match literally. Defaults to None. severity_order (Sequence[str]): Severity keys from highest to lowest. Stored values that aren't listed - sort after them. Empty orders newest first. + sort after them. Defaults to an empty order, which sorts newest first. Returns: list[Finding]: The findings on the page. diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 33cabd6639..0d540f8642 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -300,7 +300,7 @@ class FindingEntry(Base): CustomUUID, ForeignKey(f"{OperationEntry.__tablename__}.id"), nullable=False ) title: Mapped[str] = mapped_column(Unicode, nullable=False) - # Binary collation makes SQL Server sort and count severities exactly, as SQLite does. + # Binary collation keeps SQL Server from merging severities that differ only in case or accents, as SQLite does. severity: Mapped[str] = mapped_column( Unicode().with_variant(Unicode(collation="Latin1_General_100_BIN2"), "mssql"), nullable=False ) diff --git a/pyrit/models/finding_fields.py b/pyrit/models/finding_fields.py index ea789cf10a..f0d6e1a89e 100644 --- a/pyrit/models/finding_fields.py +++ b/pyrit/models/finding_fields.py @@ -215,6 +215,11 @@ def severity_order(self) -> list[str]: """The top-level option keys of the severity slot, from highest to lowest severity.""" return [option.key for option in self.severity.options] + @property + def field_keys(self) -> frozenset[str]: + """The key of every field and level, except ``severity``.""" + return frozenset(key for field in [self.severity, *self.fields] for key in _keys_of(field)) - {"severity"} + @field_validator("severity") @classmethod def _validate_severity_slot(cls, value: FindingField) -> FindingField: @@ -281,8 +286,9 @@ def check_values( new finding. Raises: - FindingValuesError: For the first unknown key, missing required value, or value that is - not an option. + FindingValuesError: For the first unknown key, changed value of a field that is no longer + configured, missing required value, value that is not an option, or Other text that + matches an option. """ previous_values = previous_values or {} self._check_known_keys(field_values=field_values, previous_values=previous_values) @@ -302,10 +308,15 @@ def _check_known_keys(self, *, field_values: dict[str, str], previous_values: di Raises: FindingValuesError: If a key is unknown and its value differs from the stored one. """ - known = {key for field in [self.severity, *self.fields] for key in _keys_of(field)} - {"severity"} + known = self.field_keys for key, value in field_values.items(): - if key not in known and previous_values.get(key) != value: - raise FindingValuesError(f"Unknown finding field {key!r}.") + if key in known or previous_values.get(key) == value: + continue + if key in previous_values: + raise FindingValuesError( + f"Finding field {key!r} is no longer configured, so its saved value can't change." + ) + raise FindingValuesError(f"Unknown finding field {key!r}.") def _parse_options(value: Any) -> Any: @@ -451,8 +462,8 @@ def _check_choice_values(*, field: FindingField, values: dict[str, str], previou previous_values (dict[str, str]): Every value of the stored finding, by key. Raises: - FindingValuesError: If a required dropdown with options has no value, or a value is not an - option and is neither unchanged nor allowed as Other text. + FindingValuesError: If a required dropdown with options has no value, or a changed value is + not an option and is not allowed as Other text. """ options = field.options path_unchanged = True @@ -464,12 +475,19 @@ def _check_choice_values(*, field: FindingField, values: dict[str, str], previou value=value, options=options, required=field.required, - accepts_unmatched=path_unchanged or field.allow_other, + unchanged=path_unchanged, + allow_other=field.allow_other, ) def _check_level( - *, label: str, value: str | None, options: list[FindingFieldOption], required: bool, accepts_unmatched: bool + *, + label: str, + value: str | None, + options: list[FindingFieldOption], + required: bool, + unchanged: bool, + allow_other: bool, ) -> list[FindingFieldOption]: """ Check one dropdown's value and find the options of the next dropdown. @@ -479,24 +497,52 @@ def _check_level( value (str | None): The dropdown's value, or None if it is missing. options (list[FindingFieldOption]): The options the dropdown offers. required (bool): Whether the field is required. - accepts_unmatched (bool): Whether a value that is not an option passes. + unchanged (bool): Whether this value and every value above it are unchanged, so a value + that is not an option passes. + allow_other (bool): Whether the field accepts Other text. Returns: list[FindingFieldOption]: The children of the chosen option, or an empty list if there is no matching option. Raises: - FindingValuesError: If a required dropdown with options has no value, or the value is not - an option and ``accepts_unmatched`` is False. + FindingValuesError: If a required dropdown with options has no value, or a changed value is + not an option and is not allowed as Other text. """ if value is None: if required and options: raise FindingValuesError(f"{label} is required.") return [] match = next((option for option in options if option.key == value), None) - if match is None and not accepts_unmatched: + if match is not None: + return match.children + if not unchanged: + _check_other_text(label=label, value=value, options=options, allow_other=allow_other) + return [] + + +def _check_other_text(*, label: str, value: str, options: list[FindingFieldOption], allow_other: bool) -> None: + """ + Accept a value that is not an option only as Other text that names no option. + + Args: + label (str): The dropdown's label. + value (str): The dropdown's value. + options (list[FindingFieldOption]): The options the dropdown offers. + allow_other (bool): Whether the field accepts Other text. + + Raises: + FindingValuesError: If the field takes no Other text, or the text matches an option's key or + label apart from case and surrounding spaces. + """ + if not allow_other: raise FindingValuesError(f"{label} has no option {value!r}.") - return [] if match is None else match.children + folded = value.strip().casefold() + for option in options: + if folded in (option.key.strip().casefold(), option.label.strip().casefold()): + raise FindingValuesError( + f"{label} {value!r} matches the option {option.key!r}. Send {option.key!r} instead." + ) def _non_blank(value: str | None) -> str | None: diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index f8ec743ffa..2ceceb1f25 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -24,7 +24,7 @@ from pyrit.common.utils import verify_and_resolve_path from pyrit.common.yaml_helper import safe_load_yaml from pyrit.common.yaml_loadable import YamlLoadable -from pyrit.models import class_name_to_snake_case +from pyrit.models import FindingFields, class_name_to_snake_case from pyrit.setup.environment_loading import validate_env_akv_strict from pyrit.setup.initialization import AZURE_SQL, IN_MEMORY, SQLITE, initialize_pyrit_async, reinitialize_pyrit_async @@ -93,14 +93,16 @@ class _ConfigurationLayer: def _check_finding_fields_initializers(initializers: Sequence["PyRITInitializer"]) -> None: """ - Check the finding fields the initializers would set, without making them active. + Check the finding fields that ``FindingFieldsInitializer`` subclasses would set, without making them active. + + Code that calls ``set_active_finding_fields`` directly isn't checked. Args: initializers (Sequence[PyRITInitializer]): The initializers of the replacement setup pass. Raises: - ValueError: If two or more initializers set the finding fields, or the one that does returns - invalid fields. + ValueError: If two or more ``FindingFieldsInitializer`` subclasses are configured, or the configured one + returns invalid fields or something other than ``FindingFields``. """ from pyrit.setup.initializers import FindingFieldsInitializer @@ -110,8 +112,14 @@ def _check_finding_fields_initializers(initializers: Sequence["PyRITInitializer" raise ValueError( f"{first} and {second} both set the finding fields. Use one FindingFieldsInitializer per deployment." ) - if setters: - setters[0].get_finding_fields() + if not setters: + return + initializer = setters[0] + result = initializer.get_finding_fields() + if not isinstance(result, FindingFields): + raise ValueError( + f"{type(initializer).__name__}.get_finding_fields() must return FindingFields, not {type(result).__name__}." + ) @dataclass @@ -729,8 +737,9 @@ async def preflight_reinitialization_async( PreparedReinitialization: Immutable inputs ready for the mutation phase. Raises: - ValueError: If an initializer is invalid, a required environment value is missing, or the - finding fields are invalid or set by more than one initializer. + ValueError: If an initializer is invalid, a required environment value is missing, more than one + ``FindingFieldsInitializer`` subclass is configured, or the configured one returns invalid finding + fields. Code that calls ``set_active_finding_fields`` directly isn't checked. """ from pyrit.registry import InitializerRegistry from pyrit.setup.initializers.targets import TargetInitializer diff --git a/pyrit/setup/finding_fields.py b/pyrit/setup/finding_fields.py index 6cc4607c67..c5b22e007c 100644 --- a/pyrit/setup/finding_fields.py +++ b/pyrit/setup/finding_fields.py @@ -27,9 +27,12 @@ def set_active_finding_fields(*, finding_fields: FindingFields, source: str) -> source (str): The setter's name, which the error names when two setters run in one pass. Raises: + TypeError: If ``finding_fields`` is not a ``FindingFields``. ValueError: If finding fields were already set in this setup pass. """ global _active + if not isinstance(finding_fields, FindingFields): + raise TypeError(f"{source} must set FindingFields, not {type(finding_fields).__name__}.") if _active is not None: raise ValueError( f"{source} cannot set the finding fields: {_active[1]} already set them in this setup pass. " diff --git a/tests/unit/backend/test_operations.py b/tests/unit/backend/test_operations.py index 3147ad153a..668a40b5e9 100644 --- a/tests/unit/backend/test_operations.py +++ b/tests/unit/backend/test_operations.py @@ -217,6 +217,16 @@ def test_finding_create_is_checked_against_the_active_fields( assert created.status_code == 201 +def test_finding_create_rejects_other_text_that_matches_an_option(compatibility_headers: dict[str, str]) -> None: + client = TestClient(app, headers=compatibility_headers) + operation_id = client.post("/api/operations", json={"name": "Case"}).json()["id"] + path = f"/api/operations/{operation_id}/findings" + rejected = client.post(path, json={"title": "Leak", "severity": "Critical"}) + assert rejected.status_code == 422 + assert rejected.json()["detail"] == "Severity 'Critical' matches the option 'critical'. Send 'critical' instead." + assert client.get(path).json()["items"] == [] + + def test_finding_api_round_trips_field_values_and_drops_blanks(compatibility_headers: dict[str, str]) -> None: client = TestClient(app, headers=compatibility_headers) operation_id = client.post("/api/operations", json={"name": "Field values"}).json()["id"] @@ -440,14 +450,17 @@ async def test_update_rejects_a_changed_retired_value_async( body = {"title": "Leak", "severity": "sev1", "field_values": {"steps": "1. Open", "description": "New text"}} response = client.put(f"{path}/{finding.id}", json=body) assert response.status_code == 422 - assert response.json()["detail"] == "Unknown finding field 'description'." + assert ( + response.json()["detail"] + == "Finding field 'description' is no longer configured, so its saved value can't change." + ) stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) assert stored is not None assert stored.field_values == {"steps": "1. Open", "description": "Old text"} assert client.put(f"{path}/{MISSING_ID}", json=body).status_code == 404 -async def test_update_removes_an_omitted_retired_value_async( +async def test_update_keeps_an_omitted_retired_value_async( sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, compatibility_headers: dict[str, str] ) -> None: finding = retired_value_finding_async @@ -458,6 +471,25 @@ async def test_update_removes_an_omitted_retired_value_async( assert response.status_code == 200 stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) assert stored is not None + assert stored.field_values == {"steps": "1. Open", "description": "Old text"} + + +async def test_update_removes_a_listed_retired_value_async( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, compatibility_headers: dict[str, str] +) -> None: + finding = retired_value_finding_async + response = TestClient(app, headers=compatibility_headers).put( + f"/api/operations/{finding.operation_id}/findings/{finding.id}", + json={ + "title": "Leak", + "severity": "sev1", + "field_values": {"steps": "1. Open"}, + "remove_retired_keys": ["description"], + }, + ) + assert response.status_code == 200 + stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) + assert stored is not None assert stored.field_values == {"steps": "1. Open"} diff --git a/tests/unit/memory/test_finding_field_values_migration.py b/tests/unit/memory/test_finding_field_values_migration.py index 0042d39573..6d5433ab9a 100644 --- a/tests/unit/memory/test_finding_field_values_migration.py +++ b/tests/unit/memory/test_finding_field_values_migration.py @@ -104,6 +104,13 @@ def test_upgrade_moves_existing_findings_into_field_values() -> None: C: ("low", {"harm_type": "Jailbreak"}), D: ("critical", {"description": 'Quote " and \\ backslash'}), } + kept = { + finding_id: (operation_id, created_at) + for finding_id, operation_id, created_at in connection.execute( + sa.text("SELECT id, operation_id, created_at FROM FindingEntries") + ) + } + assert kept == dict.fromkeys((A, B, C, D), (OPERATION_ID, STAMP)) evidence_rows = [ tuple(row) for row in connection.execute( diff --git a/tests/unit/memory/test_operation_service_findings.py b/tests/unit/memory/test_operation_service_findings.py new file mode 100644 index 0000000000..e33339af51 --- /dev/null +++ b/tests/unit/memory/test_operation_service_findings.py @@ -0,0 +1,107 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import re + +import pytest + +from pyrit.backend.models.operations import FindingUpdateRequest +from pyrit.backend.services.operation_service import OperationService +from pyrit.memory import SQLiteMemory +from pyrit.models import Finding, FindingField, FindingFields, FindingValuesError, Operation +from pyrit.setup.finding_fields import set_active_finding_fields + +pytestmark = pytest.mark.usefixtures("patch_central_database") + +CONTOSO = FindingFields( + severity=FindingField( + key="severity", label="Severity", type="choice", options=[("sev1", "Sev 1"), ("sev2", "Sev 2")] + ), + fields=[FindingField(key="steps", label="Steps to reproduce", type="multiline", required=True)], +) +STORED_VALUES = {"steps": "1. Open", "description": "Old text"} + + +@pytest.fixture +async def retired_value_finding_async(sqlite_instance: SQLiteMemory, reset_finding_fields: None) -> Finding: + """Seed a finding with a ``description`` value, then make CONTOSO, which has no such field, active.""" + operation = await sqlite_instance.add_operation_async(Operation(name="Contoso")) + finding = await sqlite_instance.add_finding_async( + finding=Finding(operation_id=operation.id, title="Leak", severity="sev1", field_values=STORED_VALUES) + ) + set_active_finding_fields(finding_fields=CONTOSO, source="test") + return finding + + +async def _update_async( + *, finding: Finding, field_values: dict[str, str], remove_retired_keys: list[str] | None = None +) -> Finding: + request = FindingUpdateRequest( + title="Leak", severity="sev1", field_values=field_values, remove_retired_keys=remove_retired_keys or [] + ) + return await OperationService().update_finding_async( + operation_id=finding.operation_id, finding_id=finding.id, request=request + ) + + +async def _stored_values_async(*, sqlite_instance: SQLiteMemory, finding: Finding) -> dict[str, str]: + stored = await sqlite_instance.get_finding_async(operation_id=finding.operation_id, finding_id=finding.id) + assert stored is not None + return stored.field_values + + +async def test_update_keeps_an_omitted_retired_value( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding +) -> None: + finding = retired_value_finding_async + updated = await _update_async(finding=finding, field_values={"steps": "1. Open"}) + assert updated.field_values == STORED_VALUES + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == STORED_VALUES + + +async def test_update_keeps_a_blank_retired_value( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding +) -> None: + finding = retired_value_finding_async + await _update_async(finding=finding, field_values={"steps": "1. Open", "description": " "}) + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == STORED_VALUES + + +async def test_update_removes_a_listed_retired_value( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding +) -> None: + finding = retired_value_finding_async + updated = await _update_async( + finding=finding, field_values={"steps": "1. Open"}, remove_retired_keys=["description"] + ) + assert updated.field_values == {"steps": "1. Open"} + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == {"steps": "1. Open"} + + +@pytest.mark.parametrize("key", ["steps", "nonexistent"]) +async def test_update_rejects_removing_a_value_that_is_not_retired( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding, key: str +) -> None: + finding = retired_value_finding_async + with pytest.raises(FindingValuesError, match=re.escape(f"{key!r} is not a retired value of this finding.")): + await _update_async(finding=finding, field_values={"steps": "1. Open"}, remove_retired_keys=[key]) + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == STORED_VALUES + + +async def test_update_rejects_sending_and_removing_a_value( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding +) -> None: + finding = retired_value_finding_async + with pytest.raises(FindingValuesError, match=re.escape("'description' can't be both sent and removed.")): + await _update_async(finding=finding, field_values=STORED_VALUES, remove_retired_keys=["description"]) + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == STORED_VALUES + + +async def test_update_rejects_changing_a_retired_value( + sqlite_instance: SQLiteMemory, retired_value_finding_async: Finding +) -> None: + finding = retired_value_finding_async + message = "Finding field 'description' is no longer configured, so its saved value can't change." + with pytest.raises(FindingValuesError, match=re.escape(message)): + await _update_async(finding=finding, field_values={"steps": "1. Open", "description": "New text"}) + assert await _stored_values_async(sqlite_instance=sqlite_instance, finding=finding) == STORED_VALUES diff --git a/tests/unit/memory/test_operations_and_findings_migration.py b/tests/unit/memory/test_operations_and_findings_migration.py index d8153542b8..423157bfd5 100644 --- a/tests/unit/memory/test_operations_and_findings_migration.py +++ b/tests/unit/memory/test_operations_and_findings_migration.py @@ -29,7 +29,7 @@ REVISION = "c8d3e5f7a901" TABLES = {"OperationEntries", "FindingEntries", "FindingEvidenceEntries"} ACTIVITY_TABLES = ("AttackResultEntries", "ScenarioResultEntries") -# FindingEntry models the head revision, which no longer has the #3066 finding columns. +# FindingEntry models the head revision, which no longer has this revision's finding columns. FINDINGS = sa.table( "FindingEntries", *( @@ -82,6 +82,7 @@ def test_migration_adds_tables_without_changing_existing_ones() -> None: assert {name: columns for name, columns in after.items() if name not in ACTIVITY_TABLES} == { name: columns for name, columns in existing.items() if name not in ACTIVITY_TABLES } + # check_schema_migrations compares the database with the head models. command.upgrade(config, "head") check_schema_migrations(engine=engine) finally: diff --git a/tests/unit/models/test_finding.py b/tests/unit/models/test_finding.py index 7a5dbb39d8..0fb17a61f7 100644 --- a/tests/unit/models/test_finding.py +++ b/tests/unit/models/test_finding.py @@ -33,7 +33,7 @@ def test_finding_drops_blank_field_values_and_keeps_other_text_exactly() -> None [ {"title": " ", "severity": "low"}, {"title": "Leak", "severity": " "}, - {"title": "Leak", "severity": "low", "description": "a #3066 field"}, + {"title": "Leak", "severity": "low", "description": "x"}, {"title": "Leak", "severity": "low", "field_values": {"steps": 3}}, ], ) diff --git a/tests/unit/models/test_finding_fields.py b/tests/unit/models/test_finding_fields.py index e3d3fdade2..83a18eb144 100644 --- a/tests/unit/models/test_finding_fields.py +++ b/tests/unit/models/test_finding_fields.py @@ -86,6 +86,12 @@ def test_severity_slot_is_required_and_ordered() -> None: assert CONTOSO.severity_order == ["sev1", "sev2", "sev3"] +def test_field_keys_lists_every_field_and_level_key_except_severity() -> None: + assert CONTOSO.field_keys == frozenset( + {"severity_rationale", "category", "subcategory", "technique", "steps", "owner"} + ) + + def test_default_finding_fields_match_3066() -> None: severity, (harm, description) = DEFAULT_FINDING_FIELDS.severity, DEFAULT_FINDING_FIELDS.fields assert DEFAULT_FINDING_FIELDS.severity_order == ["critical", "important", "moderate", "low", "informational"] @@ -266,6 +272,14 @@ def test_check_values_accepts( pytest.param( "sev1", {**VALID, "reviewer": "kim"}, None, None, "Unknown finding field 'reviewer'.", id="unknown key" ), + pytest.param( + "sev1", + {**VALID, "reviewer": "kim"}, + "sev1", + VALID, + "Unknown finding field 'reviewer'.", + id="unknown key the stored finding never had", + ), pytest.param( "sev1", {**VALID, "severity": "sev1"}, @@ -279,7 +293,7 @@ def test_check_values_accepts( {**VALID, "repro": "new"}, "sev1", {**VALID, "repro": "old"}, - "Unknown finding field 'repro'.", + "Finding field 'repro' is no longer configured, so its saved value can't change.", id="changed retired field", ), ], @@ -294,6 +308,49 @@ def test_check_values_rejects( assert str(error.value) == message +@pytest.mark.parametrize( + ("severity", "values", "message"), + [ + pytest.param( + "Critical", + {}, + "Severity 'Critical' matches the option 'critical'. Send 'critical' instead.", + id="severity label", + ), + pytest.param( + "CRITICAL", + {}, + "Severity 'CRITICAL' matches the option 'critical'. Send 'critical' instead.", + id="severity key in capitals", + ), + pytest.param( + "low ", {}, "Severity 'low ' matches the option 'low'. Send 'low' instead.", id="severity key with a space" + ), + pytest.param( + "low", + {"harm_type": "malware"}, + "Harm-type 'malware' matches the option 'Malware'. Send 'Malware' instead.", + id="harm type key in lowercase", + ), + ], +) +def test_check_values_rejects_other_text_that_matches_an_option( + severity: str, values: dict[str, str], message: str +) -> None: + with pytest.raises(FindingValuesError) as error: + DEFAULT_FINDING_FIELDS.check_values(severity=severity, field_values=values) + assert str(error.value) == message + + +@pytest.mark.parametrize("severity", ["Sev 2.5", "critcal"]) +def test_check_values_accepts_other_text_that_matches_no_option(severity: str) -> None: + DEFAULT_FINDING_FIELDS.check_values(severity=severity, field_values={}) + + +def test_check_values_accepts_unchanged_other_text_that_matches_an_option() -> None: + DEFAULT_FINDING_FIELDS.check_values(severity="Critical", field_values={}, previous_severity="Critical") + + AREA = _choice( levels=[("feature", "Feature")], options={("chat", "Chat"): [("history", "History")], ("search", "Search"): []}, diff --git a/tests/unit/setup/test_finding_fields.py b/tests/unit/setup/test_finding_fields.py index 6781f7269b..eb35a3b7bd 100644 --- a/tests/unit/setup/test_finding_fields.py +++ b/tests/unit/setup/test_finding_fields.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import re + import pytest from pyrit.models import DEFAULT_FINDING_FIELDS, FindingField, FindingFields @@ -33,6 +35,12 @@ def test_set_active_finding_fields_twice_names_both_sources() -> None: assert get_active_finding_fields() is CUSTOM +def test_set_active_finding_fields_rejects_a_value_that_is_not_finding_fields() -> None: + with pytest.raises(TypeError, match=re.escape("Contoso must set FindingFields, not dict.")): + set_active_finding_fields(finding_fields={}, source="Contoso") + assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + + def test_reset_restores_the_default() -> None: set_active_finding_fields(finding_fields=CUSTOM, source="First") reset_active_finding_fields() diff --git a/tests/unit/setup/test_reinitialization.py b/tests/unit/setup/test_reinitialization.py index 29623278c2..77a903e355 100644 --- a/tests/unit/setup/test_reinitialization.py +++ b/tests/unit/setup/test_reinitialization.py @@ -4,6 +4,7 @@ """Repeatable setup uses fresh scripts, replacement values, and the same memory.""" import os +import re import uuid from collections.abc import Generator, Iterator from contextlib import contextmanager, nullcontext @@ -15,7 +16,7 @@ from pyrit.common.apply_defaults import get_global_default_values from pyrit.common.random_context import configure_random_seed, get_configured_random_seed from pyrit.memory import AzureSQLMemory, CentralMemory, SQLiteMemory -from pyrit.models import DEFAULT_FINDING_FIELDS +from pyrit.models import FindingField, FindingFields from pyrit.models.catalog import ScenarioSuite from pyrit.registry import ( AttackRegistry, @@ -27,10 +28,14 @@ ) from pyrit.setup.configuration_loader import ConfigurationLoader from pyrit.setup.environment_loading import resolve_environment_async -from pyrit.setup.finding_fields import get_active_finding_fields +from pyrit.setup.finding_fields import get_active_finding_fields, set_active_finding_fields from pyrit.setup.initialization import reset_setup_registries, validate_reinitialization_memory from pyrit.setup.initializers.targets import TARGET_CONFIGS +CUSTOM = FindingFields( + severity=FindingField(key="severity", label="Severity", type="choice", options=[("high", "High")]) +) + @contextmanager def _isolated_reinitialization_state() -> Generator[None, None, None]: @@ -134,10 +139,15 @@ def get_finding_fields(self): return FindingFields(severity=SEVERITY) encoding="utf-8", ) config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], initialization_scripts=[str(script)]) + set_active_finding_fields(finding_fields=CUSTOM, source="ContosoRunning") - with pytest.raises(ValueError, match="Use one FindingFieldsInitializer per deployment"): + message = ( + "ContosoBackupFields and ContosoFindingFields both set the finding fields. " + "Use one FindingFieldsInitializer per deployment." + ) + with pytest.raises(ValueError, match=re.escape(message)): await config.preflight_reinitialization_async(environment_values={}) - assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + assert get_active_finding_fields() is CUSTOM @pytest.mark.usefixtures("reset_finding_fields") @@ -155,10 +165,33 @@ def get_finding_fields(self): encoding="utf-8", ) config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], initialization_scripts=[str(script)]) + set_active_finding_fields(finding_fields=CUSTOM, source="ContosoRunning") with pytest.raises(ValueError, match="needs at least one option"): await config.preflight_reinitialization_async(environment_values={}) - assert get_active_finding_fields() is DEFAULT_FINDING_FIELDS + assert get_active_finding_fields() is CUSTOM + + +@pytest.mark.usefixtures("reset_finding_fields") +async def test_preflight_rejects_finding_fields_of_the_wrong_type(tmp_path: Path) -> None: + script = tmp_path / "wrong_type_finding_fields.py" + script.write_text( + """ +from pyrit.setup.initializers import FindingFieldsInitializer + +class ContosoFindingFields(FindingFieldsInitializer): + def get_finding_fields(self): + return {} +""", + encoding="utf-8", + ) + config = ConfigurationLoader(memory_db_type="in_memory", env_files=[], initialization_scripts=[str(script)]) + set_active_finding_fields(finding_fields=CUSTOM, source="ContosoRunning") + + message = "ContosoFindingFields.get_finding_fields() must return FindingFields, not dict." + with pytest.raises(ValueError, match=re.escape(message)): + await config.preflight_reinitialization_async(environment_values={}) + assert get_active_finding_fields() is CUSTOM @pytest.mark.usefixtures("patch_central_database") From 5943b6b0311a59cd4d1f45c3ab326912c33132d8 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Thu, 29 Oct 2026 09:35:41 -0400 Subject: [PATCH 08/20] Rename FindingsInitializer, export configured fields and add retired values Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Operations/RetiredFieldValues.styles.ts | 52 +++++ .../Operations/RetiredFieldValues.test.tsx | 164 ++++++++++++++++ .../Operations/RetiredFieldValues.tsx | 103 ++++++++++ frontend/src/hooks/useFindingFields.test.tsx | 62 ++++++ frontend/src/hooks/useFindingFields.ts | 41 ++++ frontend/src/services/api.test.ts | 6 + frontend/src/services/api.ts | 4 + .../src/test-utils/findingFieldsFixtures.ts | 111 +++++++++++ frontend/src/types/index.ts | 42 +++++ frontend/src/utils/findingFields.test.ts | 143 ++++++++++++++ frontend/src/utils/findingFields.ts | 178 ++++++++++++++++++ pyrit/backend/services/finding_export.py | 42 +++-- pyrit/models/finding_fields.py | 98 +++++++++- pyrit/setup/configuration_loader.py | 16 +- pyrit/setup/finding_fields.py | 2 +- pyrit/setup/initializers/__init__.py | 4 +- .../{finding_fields.py => findings.py} | 4 +- tests/unit/backend/test_finding_export.py | 87 ++++++++- .../test_operation_service_findings.py | 0 tests/unit/models/test_finding_fields.py | 73 ++++++- ...alizer.py => test_findings_initializer.py} | 20 +- tests/unit/setup/test_initialization.py | 10 +- tests/unit/setup/test_reinitialization.py | 18 +- 23 files changed, 1210 insertions(+), 70 deletions(-) create mode 100644 frontend/src/components/Operations/RetiredFieldValues.styles.ts create mode 100644 frontend/src/components/Operations/RetiredFieldValues.test.tsx create mode 100644 frontend/src/components/Operations/RetiredFieldValues.tsx create mode 100644 frontend/src/hooks/useFindingFields.test.tsx create mode 100644 frontend/src/hooks/useFindingFields.ts create mode 100644 frontend/src/test-utils/findingFieldsFixtures.ts create mode 100644 frontend/src/utils/findingFields.test.ts create mode 100644 frontend/src/utils/findingFields.ts rename pyrit/setup/initializers/{finding_fields.py => findings.py} (92%) rename tests/unit/{memory => backend}/test_operation_service_findings.py (100%) rename tests/unit/setup/{test_finding_fields_initializer.py => test_findings_initializer.py} (83%) diff --git a/frontend/src/components/Operations/RetiredFieldValues.styles.ts b/frontend/src/components/Operations/RetiredFieldValues.styles.ts new file mode 100644 index 0000000000..397b6ea322 --- /dev/null +++ b/frontend/src/components/Operations/RetiredFieldValues.styles.ts @@ -0,0 +1,52 @@ +import { makeStyles, tokens } from '@fluentui/react-components' + +import { mobileTouchTarget } from '@/styles/touchTargets' + +const DASHED_PANEL = { + padding: `${tokens.spacingVerticalS} ${tokens.spacingHorizontalM}`, + backgroundColor: tokens.colorNeutralBackground2, + border: `${tokens.strokeWidthThin} dashed ${tokens.colorNeutralStroke1}`, + borderRadius: tokens.borderRadiusMedium, +} + +export const useRetiredFieldValuesStyles = makeStyles({ + root: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalS }, + group: { ...DASHED_PANEL, display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalXS }, + heading: { + color: tokens.colorNeutralForeground3, + fontSize: tokens.fontSizeBase200, + fontWeight: tokens.fontWeightSemibold, + }, + list: { + margin: 0, + padding: 0, + listStyleType: 'none', + display: 'flex', + flexDirection: 'column', + gap: tokens.spacingVerticalS, + }, + item: { + display: 'flex', + justifyContent: 'space-between', + alignItems: 'flex-start', + gap: tokens.spacingHorizontalM, + }, + entry: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalXXS, minWidth: 0 }, + key: { + fontFamily: tokens.fontFamilyMonospace, + fontSize: tokens.fontSizeBase200, + color: tokens.colorNeutralForeground3, + overflowWrap: 'anywhere', + }, + titleKey: { fontFamily: tokens.fontFamilyMonospace }, + value: { margin: 0, whiteSpace: 'pre-wrap', overflowWrap: 'anywhere', userSelect: 'text' }, + note: { + ...DASHED_PANEL, + display: 'flex', + flexWrap: 'wrap', + alignItems: 'center', + gap: tokens.spacingHorizontalS, + color: tokens.colorNeutralForeground3, + }, + button: { flexShrink: 0, ...mobileTouchTarget }, +}) diff --git a/frontend/src/components/Operations/RetiredFieldValues.test.tsx b/frontend/src/components/Operations/RetiredFieldValues.test.tsx new file mode 100644 index 0000000000..71f3d45ede --- /dev/null +++ b/frontend/src/components/Operations/RetiredFieldValues.test.tsx @@ -0,0 +1,164 @@ +import React, { useState } from 'react' +import { render, screen, waitFor, within } from '@testing-library/react' +import userEvent, { type UserEvent } from '@testing-library/user-event' +import { FluentProvider, webLightTheme } from '@fluentui/react-components' + +import RetiredFieldValues from './RetiredFieldValues' + +const TestWrapper: React.FC<{ children: React.ReactNode }> = ({ children }) => ( + {children} +) + +const RETIRED_VALUES = [ + { key: 'description', value: 'The assistant agreed to play a character and then\nexplained the restricted steps.' }, + { key: 'contoso_ticket', value: 'CONTOSO-42' }, +] + +function removeButtonName(key: string): string { + return `Remove retired value ${key}` +} + +interface RetiredFieldValuesHarnessProps { + onRemove: (key: string) => void + onUndo: () => void +} + +// Holds the removed keys the way the edit dialog will, so a test can follow a removal through to Undo. +function RetiredFieldValuesHarness({ onRemove, onUndo }: RetiredFieldValuesHarnessProps) { + const [removed, setRemoved] = useState>(new Set()) + return ( + { + onRemove(key) + setRemoved((previous) => new Set([...previous, key])) + }} + onUndo={() => { + onUndo() + setRemoved(new Set()) + }} /> + ) +} + +// jsdom has no default stylesheet, so an accessible name would gain a space around the title's key. +async function findConfirmation(key: string): Promise { + const dialog = await screen.findByRole('dialog') + expect(within(dialog).getByRole('heading')).toHaveTextContent(`Remove ${key}?`) + return dialog +} + +async function confirmRemoval(user: UserEvent, key: string): Promise { + await user.click(await screen.findByRole('button', { name: removeButtonName(key) })) + const dialog = await findConfirmation(key) + await user.click(within(dialog).getByRole('button', { name: 'Remove' })) + await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument()) +} + +describe('RetiredFieldValues', () => { + const defaultProps = { + values: RETIRED_VALUES, + removed: new Set(), + onRemove: jest.fn(), + onUndo: jest.fn(), + } + + beforeEach(() => { + jest.clearAllMocks() + }) + + it('renders nothing without retired values', () => { + render() + expect(screen.queryByText('Retired fields · read-only')).not.toBeInTheDocument() + expect(screen.queryByRole('group')).not.toBeInTheDocument() + expect(screen.queryByRole('button')).not.toBeInTheDocument() + }) + + it('lists each retired value with its raw key and a Remove button', () => { + render() + const group = screen.getByRole('group', { name: 'Retired fields · read-only' }) + const items = within(group).getAllByRole('listitem') + expect(items).toHaveLength(2) + const [description, ticket] = items + expect(within(description).getByText('description')).toBeInTheDocument() + expect(within(description).getByText(/The assistant agreed/)).toHaveTextContent( + 'The assistant agreed to play a character and then\nexplained the restricted steps.', + { normalizeWhitespace: false }, + ) + expect(within(description).getByRole('button', { name: removeButtonName('description') })).toBeEnabled() + expect(within(ticket).getByText('contoso_ticket')).toBeInTheDocument() + expect(within(ticket).getByText('CONTOSO-42')).toBeInTheDocument() + expect(within(ticket).getByRole('button', { name: removeButtonName('contoso_ticket') })).toBeEnabled() + expect(screen.queryByRole('button', { name: 'Undo' })).not.toBeInTheDocument() + }) + + it('asks before removing and keeps the value on Keep it', async () => { + const user = userEvent.setup() + render() + await user.click(screen.getByRole('button', { name: removeButtonName('description') })) + const dialog = await findConfirmation('description') + expect(dialog).toHaveTextContent( + 'This deletes its text from the finding when you save. Copy anything you need first. ' + + 'Cancel on the edit dialog still brings it back.', + ) + expect(within(dialog).getByRole('button', { name: 'Remove' })).toBeInTheDocument() + await user.click(within(dialog).getByRole('button', { name: 'Keep it' })) + await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument()) + expect(await screen.findByRole('button', { name: removeButtonName('description') })).toBeInTheDocument() + expect(defaultProps.onRemove).not.toHaveBeenCalled() + }) + + it('reports a confirmed removal and offers Undo', async () => { + const user = userEvent.setup() + const onRemove = jest.fn() + const onUndo = jest.fn() + render() + await confirmRemoval(user, 'description') + expect(onRemove).toHaveBeenCalledTimes(1) + expect(onRemove).toHaveBeenCalledWith('description') + expect(await screen.findByRole('button', { name: removeButtonName('contoso_ticket') })).toBeInTheDocument() + expect(screen.queryByRole('button', { name: removeButtonName('description') })).not.toBeInTheDocument() + expect(screen.getByText('1 retired value is removed when you save.')).toBeInTheDocument() + await user.click(screen.getByRole('button', { name: 'Undo' })) + expect(onUndo).toHaveBeenCalledTimes(1) + expect(screen.getByRole('button', { name: removeButtonName('description') })).toBeInTheDocument() + expect(screen.queryByText(/removed when you save/)).not.toBeInTheDocument() + }) + + it('counts several removals', async () => { + const user = userEvent.setup() + const onRemove = jest.fn() + render() + await confirmRemoval(user, 'description') + await confirmRemoval(user, 'contoso_ticket') + expect(onRemove).toHaveBeenNthCalledWith(1, 'description') + expect(onRemove).toHaveBeenNthCalledWith(2, 'contoso_ticket') + expect(await screen.findByRole('button', { name: 'Undo' })).toBeInTheDocument() + expect(screen.getByText('2 retired values are removed when you save.')).toBeInTheDocument() + expect(screen.queryByRole('group', { name: 'Retired fields · read-only' })).not.toBeInTheDocument() + }) + + it('moves focus to Undo after a removal is confirmed', async () => { + const user = userEvent.setup() + render() + await confirmRemoval(user, 'description') + expect(await screen.findByRole('button', { name: 'Undo' })).toHaveFocus() + }) + + it('moves focus to the first Remove button after Undo', async () => { + const user = userEvent.setup() + render() + await confirmRemoval(user, 'contoso_ticket') + await user.click(await screen.findByRole('button', { name: 'Undo' })) + expect(screen.getByRole('button', { name: removeButtonName('description') })).toHaveFocus() + }) + + it('disables every Remove button and Undo while disabled', () => { + const values = [...RETIRED_VALUES, { key: 'contoso_owner', value: 'Contoso red team' }] + render( + + ) + const removeButtons = screen.getAllByRole('button', { name: /^Remove retired value / }) + expect(removeButtons).toHaveLength(2) + for (const button of removeButtons) expect(button).toBeDisabled() + expect(screen.getByRole('button', { name: 'Undo' })).toBeDisabled() + }) +}) diff --git a/frontend/src/components/Operations/RetiredFieldValues.tsx b/frontend/src/components/Operations/RetiredFieldValues.tsx new file mode 100644 index 0000000000..66403888b4 --- /dev/null +++ b/frontend/src/components/Operations/RetiredFieldValues.tsx @@ -0,0 +1,103 @@ +import { type RefObject, useEffect, useId, useRef, useState, type JSX } from 'react' +import { + Button, Dialog, DialogActions, DialogBody, DialogContent, DialogSurface, DialogTitle, Link, Text, +} from '@fluentui/react-components' + +import { useRetiredFieldValuesStyles } from './RetiredFieldValues.styles' + +interface RetiredFieldValuesProps { + /** Retired values from the stored finding, in stored order. */ + values: Array<{ key: string; value: string }> + /** Keys marked for removal in this edit. */ + removed: ReadonlySet + disabled?: boolean + onRemove: (key: string) => void + /** Restores every marked key. */ + onUndo: () => void +} + +function removalNote(count: number): string { + return count === 1 + ? '1 retired value is removed when you save.' + : `${count} retired values are removed when you save.` +} + +/** + * Lists stored values whose keys the active finding fields no longer define. They are read-only, but each + * can be marked for removal on save, after a confirmation, and Undo restores every marked value. + */ +export default function RetiredFieldValues({ + values, removed, disabled = false, onRemove, onUndo, +}: RetiredFieldValuesProps): JSX.Element | null { + const styles = useRetiredFieldValuesStyles() + const headingId = useId() + // The key stays set while the dialog closes, so its title doesn't lose the key during the exit animation. + const [confirming, setConfirming] = useState({ key: '', open: false }) + const undoRef = useRef(null) + const firstRemoveRef = useRef(null) + const focusAfterChange = useRef | null>(null) + + // Confirming a removal unmounts the row whose button opened the dialog, and Undo unmounts itself, so focus + // moves to the control that replaces them once the parent has applied the change. + useEffect(() => { + focusAfterChange.current?.current?.focus() + focusAfterChange.current = null + }, [removed]) + + if (values.length === 0) return null + const kept = values.filter(({ key }) => !removed.has(key)) + const closeConfirmation = (): void => { + setConfirming((current) => ({ ...current, open: false })) + } + + return ( +
+ {kept.length > 0 &&
+ Retired fields · read-only +
    + {kept.map(({ key, value }, index) => ( +
  • +
    + {key} + {value} +
    + +
  • + ))} +
+
} + {removed.size > 0 &&
+ {removalNote(removed.size)} + { + focusAfterChange.current = firstRemoveRef + onUndo() + }}>Undo +
} + { + if (!data.open) closeConfirmation() + }}> + + + Remove {confirming.key}? + + + This deletes its text from the finding when you save. Copy anything you need first. Cancel on the + edit dialog still brings it back. + + + + + + + + + +
+ ) +} diff --git a/frontend/src/hooks/useFindingFields.test.tsx b/frontend/src/hooks/useFindingFields.test.tsx new file mode 100644 index 0000000000..1d0bbf6b57 --- /dev/null +++ b/frontend/src/hooks/useFindingFields.test.tsx @@ -0,0 +1,62 @@ +import { StrictMode } from 'react' + +import { act, renderHook, waitFor } from '@testing-library/react' + +import { operationsApi } from '@/services/api' +import { deferred } from '@/test-utils/deferred' +import { CONTOSO_FINDING_FIELDS, DEFAULT_FINDING_FIELDS_FIXTURE } from '@/test-utils/findingFieldsFixtures' +import type { FindingFields } from '@/types' + +import { useFindingFields } from './useFindingFields' + +jest.mock('@/services/api', () => ({ operationsApi: { getFindingFields: jest.fn() } })) +const mockGetFindingFields = jest.mocked(operationsApi.getFindingFields) + +describe('useFindingFields', () => { + beforeEach(() => { + jest.clearAllMocks() + }) + + it('loads the finding fields once', async () => { + mockGetFindingFields.mockResolvedValue(CONTOSO_FINDING_FIELDS) + + const { result, rerender } = renderHook(() => useFindingFields()) + expect(result.current.fields).toBeNull() + + await waitFor(() => expect(result.current.fields).toBe(CONTOSO_FINDING_FIELDS)) + rerender() + expect(result.current.error).toBe('') + expect(mockGetFindingFields).toHaveBeenCalledTimes(1) + }) + + it('reports a load failure and retries', async () => { + mockGetFindingFields + .mockRejectedValueOnce(new Error('The Contoso backend is unavailable.')) + .mockResolvedValueOnce(DEFAULT_FINDING_FIELDS_FIXTURE) + + const { result } = renderHook(() => useFindingFields()) + await waitFor(() => expect(result.current.error).toBe('The Contoso backend is unavailable.')) + expect(result.current.fields).toBeNull() + + act(() => { result.current.retry() }) + expect(result.current.error).toBe('') + + await waitFor(() => expect(result.current.fields).toBe(DEFAULT_FINDING_FIELDS_FIXTURE)) + expect(result.current.error).toBe('') + expect(mockGetFindingFields).toHaveBeenCalledTimes(2) + }) + + it('ignores a response after unmount', async () => { + const unmounted = deferred() + const mounted = deferred() + mockGetFindingFields.mockReturnValueOnce(unmounted.promise).mockReturnValueOnce(mounted.promise) + + // StrictMode unmounts and remounts once, so the first request belongs to an unmounted effect. + const { result } = renderHook(() => useFindingFields(), { wrapper: StrictMode }) + await act(async () => { mounted.resolve(CONTOSO_FINDING_FIELDS) }) + await act(async () => { unmounted.resolve(DEFAULT_FINDING_FIELDS_FIXTURE) }) + + expect(mockGetFindingFields).toHaveBeenCalledTimes(2) + expect(result.current.fields).toBe(CONTOSO_FINDING_FIELDS) + }) +}) diff --git a/frontend/src/hooks/useFindingFields.ts b/frontend/src/hooks/useFindingFields.ts new file mode 100644 index 0000000000..5351bf6990 --- /dev/null +++ b/frontend/src/hooks/useFindingFields.ts @@ -0,0 +1,41 @@ +import { useCallback, useEffect, useState } from 'react' + +import { operationsApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { FindingFields } from '@/types' + +interface FindingFieldsState { + /** Null until the fields load. */ + readonly fields: FindingFields | null + /** The last load failure, or an empty string. */ + readonly error: string + readonly retry: () => void +} + +/** Loads the active finding fields once, for the finding dialog, cards and counts. */ +export function useFindingFields(): FindingFieldsState { + const [fields, setFields] = useState(null) + const [error, setError] = useState('') + const [attempt, setAttempt] = useState(0) + + useEffect(() => { + let ignore = false + operationsApi.getFindingFields() + .then((loaded: FindingFields) => { + if (!ignore) setFields(loaded) + }) + .catch((cause: unknown) => { + if (!ignore) setError(toApiError(cause).detail) + }) + return () => { + ignore = true + } + }, [attempt]) + + const retry = useCallback(() => { + setError('') + setAttempt((count: number) => count + 1) + }, []) + + return { fields, error, retry } +} diff --git a/frontend/src/services/api.test.ts b/frontend/src/services/api.test.ts index c40ce0d65b..d330138e4b 100644 --- a/frontend/src/services/api.test.ts +++ b/frontend/src/services/api.test.ts @@ -42,6 +42,7 @@ import { import { compatibility, COMPATIBILITY_HEADER } from "./compatibility"; import { toApiError } from "./errors"; import { makeAnalyticsFacets, makeAnalyticsReport, makeAnalyticsResults } from "@/test-utils/analyticsFixtures"; +import { DEFAULT_FINDING_FIELDS_FIXTURE } from "@/test-utils/findingFieldsFixtures"; import type { FindingExportRequest } from "@/types"; import { DEFAULT_ANALYTICS_VIEW } from "@/utils/attackAnalytics"; @@ -53,6 +54,11 @@ describe("api service", () => { expect(await operationsApi.getFindingOptions()).toEqual(options); expect(apiClient.get).toHaveBeenCalledWith("/operations/finding-options"); }); + it("gets finding fields", async () => { + (apiClient.get as jest.Mock).mockResolvedValueOnce({ data: DEFAULT_FINDING_FIELDS_FIXTURE }); + expect(await operationsApi.getFindingFields()).toEqual(DEFAULT_FINDING_FIELDS_FIXTURE); + expect(apiClient.get).toHaveBeenCalledWith("/operations/finding-fields"); + }); it("encodes evidence scope, searches literally, and preserves duplicate feedback", async () => { const page = { items: [], has_more: false, next_offset: null }; (apiClient.get as jest.Mock).mockResolvedValue({ data: page }); diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index a8ddd7dc4c..c96b777c73 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -82,6 +82,7 @@ import type { Finding, FindingCreate, FindingCreateRequest, + FindingFields, FindingListResponse, FindingOptionsResponse, FindingEvidenceCreateRequest, @@ -109,6 +110,9 @@ export const operationsApi = { getFindingOptions: async (): Promise => { return (await apiClient.get('/operations/finding-options')).data }, + getFindingFields: async (): Promise => { + return (await apiClient.get('/operations/finding-fields')).data + }, list: async (): Promise => { return (await apiClient.get('/operations')).data }, diff --git a/frontend/src/test-utils/findingFieldsFixtures.ts b/frontend/src/test-utils/findingFieldsFixtures.ts new file mode 100644 index 0000000000..65dd7a7387 --- /dev/null +++ b/frontend/src/test-utils/findingFieldsFixtures.ts @@ -0,0 +1,111 @@ +import type { FindingField, FindingFieldOption, FindingFields } from '@/types' + +type FieldDefinition = Pick & Partial + +/** Fills in the backend's defaults, so each fixture field has the full JSON shape. */ +function field(definition: FieldDefinition): FindingField { + return { + required: false, + options: [], + levels: [], + default: null, + allow_other: false, + colors: {}, + hint: null, + section: null, + internal: false, + ...definition, + } +} + +function option(key: string, label: string, children: FindingFieldOption[] = []): FindingFieldOption { + return { key, label, children } +} + +/** A Contoso deployment's fields: a severity level, a cascade, Other text, a required multiline and an internal field. */ +export const CONTOSO_FINDING_FIELDS: FindingFields = { + severity: field({ + key: 'severity', + label: 'Severity', + type: 'choice', + required: true, + options: [ + option('sev1', 'Sev 1', [option('customer_impact', 'Customer impact')]), + option('sev2', 'Sev 2', [option('internal_impact', 'Internal impact')]), + ], + levels: [['severity_rationale', 'Severity rationale']], + colors: { sev1: 'red', sev2: 'yellow' }, + }), + fields: [ + field({ + key: 'category', + label: 'Category', + type: 'choice', + required: true, + options: [ + option('model_security', 'Model security', [ + option('system_prompt_leak', 'System prompt leak'), + option('tool_misuse', 'Tool misuse'), + ]), + option('content_safety', 'Content safety', [ + option('hate_speech', 'Hate speech'), + option('self_harm', 'Self-harm'), + ]), + ], + levels: [['subcategory', 'Subcategory']], + section: 'Classification', + }), + field({ + key: 'technique', + label: 'Technique', + type: 'choice', + options: [option('prompt_injection', 'Prompt injection'), option('role_play', 'Role play')], + allow_other: true, + section: 'Reproduction', + }), + field({ + key: 'steps', + label: 'Steps to reproduce', + type: 'multiline', + required: true, + hint: 'List each prompt you sent to the Contoso assistant, in order.', + section: 'Reproduction', + }), + field({ key: 'owner', label: 'Owner', type: 'text', section: 'Tracking', internal: true }), + ], +} + +/** A copy of the backend's `DEFAULT_FINDING_FIELDS`, with three harm types instead of the whole taxonomy. */ +export const DEFAULT_FINDING_FIELDS_FIXTURE: FindingFields = { + severity: field({ + key: 'severity', + label: 'Severity', + type: 'choice', + required: true, + options: [ + option('critical', 'Critical'), + option('important', 'Important'), + option('moderate', 'Moderate'), + option('low', 'Low'), + option('informational', 'Informational'), + ], + default: 'moderate', + allow_other: true, + colors: { critical: 'red', important: 'yellow', moderate: 'yellow', low: 'gray' }, + }), + fields: [ + field({ + key: 'harm_type', + label: 'Harm-type', + type: 'choice', + options: [ + option('Hate Speech', 'Hate Speech'), + option('Insecure Code', 'Insecure Code'), + option('Malware', 'Malware'), + ], + allow_other: true, + hint: 'Optional. Select a category, or choose Other for custom text.', + }), + field({ key: 'description', label: 'Description', type: 'multiline', hint: 'Optional' }), + ], +} diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 41ae65c7ac..2adf77c6ad 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -2208,10 +2208,52 @@ export interface FindingCreateRequest extends FindingCreate { attack_result_id?: string } +/** A finding edit. Each key in `remove_retired_keys` drops that retired value from the stored finding. */ +export interface FindingUpdateRequest extends FindingCreate { + remove_retired_keys: string[] +} + export interface FindingOptionsResponse { harm_types: string[] } +export type FindingFieldType = 'text' | 'multiline' | 'choice' + +export type FindingSeverityColor = 'red' | 'orange' | 'yellow' | 'green' | 'blue' | 'gray' | 'dark' | 'neutral' + +export interface FindingFieldOption { + key: string + label: string + /** The options of the next dropdown when this option is chosen. */ + children: FindingFieldOption[] +} + +/** One field of the finding dialog, as the deployment defines it. */ +export interface FindingField { + key: string + label: string + type: FindingFieldType + required: boolean + options: FindingFieldOption[] + /** The `[key, label]` of each dropdown under a cascade's first one. */ + levels: Array<[string, string]> + default: string | null + /** Whether a single-level choice also accepts typed text. */ + allow_other: boolean + /** The severity slot's badge color for each top-level option. */ + colors: Record + hint: string | null + section: string | null + /** Shown on the dashboard but left out of exports. */ + internal: boolean +} + +/** The active finding fields: the severity slot and the fields shown after it. */ +export interface FindingFields { + severity: FindingField + fields: FindingField[] +} + export interface Finding extends FindingCreate { id: string operation_id: string diff --git a/frontend/src/utils/findingFields.test.ts b/frontend/src/utils/findingFields.test.ts new file mode 100644 index 0000000000..7a5cdefcbc --- /dev/null +++ b/frontend/src/utils/findingFields.test.ts @@ -0,0 +1,143 @@ +import { CONTOSO_FINDING_FIELDS, DEFAULT_FINDING_FIELDS_FIXTURE } from '@/test-utils/findingFieldsFixtures' +import type { FindingField, FindingFieldOption } from '@/types' + +import { + choiceLevels, + describeFindingValues, + fieldKeys, + firstMissingRequired, + FLUENT_BADGE_COLOR, + type FindingValueRow, + levelOptions, + severityBadge, + severityCountBadges, +} from './findingFields' + +function contosoField(key: string): FindingField { + const match = CONTOSO_FINDING_FIELDS.fields.find((field: FindingField) => field.key === key) + if (match === undefined) throw new Error(`The Contoso fixture has no field ${key}`) + return match +} + +function optionKeys(options: FindingFieldOption[]): string[] { + return options.map((option: FindingFieldOption) => option.key) +} + +describe('findingFields', () => { + it('maps every author color to one Fluent badge color', () => { + expect(FLUENT_BADGE_COLOR).toEqual({ red: 'danger', orange: 'severe', yellow: 'warning', green: 'success', + blue: 'brand', gray: 'informative', dark: 'important', neutral: 'subtle' }) + }) + + it('keeps the #3066 badge colors for the default fields', () => { + expect(['critical', 'important', 'moderate', 'low', 'informational'] + .map((key: string) => severityBadge(DEFAULT_FINDING_FIELDS_FIXTURE, key).color)) + .toEqual(['danger', 'warning', 'warning', 'informative', 'subtle']) + }) + + it('shows an unmatched severity as typed on a neutral badge', () => { + expect(severityBadge(CONTOSO_FINDING_FIELDS, 'Sev 2.5')).toEqual({ label: 'Sev 2.5', color: 'subtle' }) + expect(severityBadge(CONTOSO_FINDING_FIELDS, 'sev1')).toEqual({ label: 'Sev 1', color: 'danger' }) + }) + + it('lists a field with its levels and every key except severity', () => { + expect(choiceLevels(contosoField('category'))).toEqual([ + { key: 'category', label: 'Category' }, { key: 'subcategory', label: 'Subcategory' }]) + expect(choiceLevels(contosoField('steps'))).toEqual([{ key: 'steps', label: 'Steps to reproduce' }]) + expect([...fieldKeys(CONTOSO_FINDING_FIELDS)]).toEqual( + ['severity_rationale', 'category', 'subcategory', 'technique', 'steps', 'owner']) + }) + + it('describes filled values in configured order and lists retired keys', () => { + const { rows, retired } = describeFindingValues(CONTOSO_FINDING_FIELDS, 'sev1', { + owner: 'dana', category: 'model_security', subcategory: 'system_prompt_leak', + technique: 'Many-shot priming', description: 'Old text', + }) + expect(rows.map((row: FindingValueRow) => `${row.label}: ${row.value}`)).toEqual([ + 'Category: Model security', 'Subcategory: System prompt leak', 'Technique: Many-shot priming', 'Owner: dana']) + expect(retired).toEqual([{ key: 'description', value: 'Old text' }]) + }) + + it('shows a cascade child that is not under its stored parent as typed', () => { + const { rows } = describeFindingValues(CONTOSO_FINDING_FIELDS, 'sev1', { + category: 'model_security', subcategory: 'hate_speech', + }) + expect(rows).toEqual([ + { key: 'category', label: 'Category', value: 'Model security', multiline: false }, + { key: 'subcategory', label: 'Subcategory', value: 'hate_speech', multiline: false }, + ]) + }) + + it('labels a severity level from the given severity and skips blank values', () => { + const { rows, retired } = describeFindingValues(CONTOSO_FINDING_FIELDS, 'sev1', { + severity_rationale: 'customer_impact', owner: ' ', steps: 'Send the Contoso prompt twice.', + }) + expect(rows).toEqual([ + { key: 'severity_rationale', label: 'Severity rationale', value: 'Customer impact', multiline: false }, + { key: 'steps', label: 'Steps to reproduce', value: 'Send the Contoso prompt twice.', multiline: true }, + ]) + expect(retired).toEqual([]) + expect(describeFindingValues(CONTOSO_FINDING_FIELDS, 'sev2', { severity_rationale: 'customer_impact' }).rows) + .toEqual([{ key: 'severity_rationale', label: 'Severity rationale', value: 'customer_impact', multiline: false }]) + }) + + it('treats a stored severity key as retired and resolves levels under the given severity', () => { + const { rows, retired } = describeFindingValues(CONTOSO_FINDING_FIELDS, 'sev1', { + severity: 'sev2', severity_rationale: 'customer_impact', + }) + expect(rows).toEqual([ + { key: 'severity_rationale', label: 'Severity rationale', value: 'Customer impact', multiline: false }, + ]) + expect(retired).toEqual([{ key: 'severity', value: 'sev2' }]) + }) + + it('offers each cascade level the options under the chosen parent', () => { + const category = contosoField('category') + expect(optionKeys(levelOptions(category, {}, 0))).toEqual(['model_security', 'content_safety']) + expect(optionKeys(levelOptions(category, { category: 'model_security' }, 1))) + .toEqual(['system_prompt_leak', 'tool_misuse']) + expect(optionKeys(levelOptions(category, { category: 'content_safety' }, 1))).toEqual(['hate_speech', 'self_harm']) + expect(levelOptions(category, {}, 1)).toEqual([]) + expect(levelOptions(category, { category: 'Prompt injection' }, 1)).toEqual([]) + }) + + it('names the first missing required value', () => { + const missingSteps = { + severity_rationale: 'customer_impact', category: 'model_security', subcategory: 'tool_misuse', steps: ' ', + } + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, 'sev1', missingSteps)).toBe('Steps to reproduce') + const missingSubcategory = { + ...missingSteps, category: 'content_safety', subcategory: '', steps: 'Send the Contoso prompt twice.', + } + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, 'sev1', missingSubcategory)).toBe('Subcategory') + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, 'sev1', { ...missingSubcategory, subcategory: 'self_harm' })) + .toBeNull() + }) + + it('always requires the severity and each of its levels that offers options', () => { + const complete = { + severity_rationale: 'customer_impact', category: 'model_security', subcategory: 'tool_misuse', + steps: 'Send the Contoso prompt twice.', + } + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, ' ', complete)).toBe('Severity') + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, 'sev1', { ...complete, severity_rationale: '' })) + .toBe('Severity rationale') + expect(firstMissingRequired(CONTOSO_FINDING_FIELDS, 'Sev 2.5', { ...complete, severity_rationale: '' })).toBeNull() + expect(firstMissingRequired(DEFAULT_FINDING_FIELDS_FIXTURE, 'moderate', {})).toBeNull() + }) + + it('groups unmatched severity counts into one Other badge with a tooltip', () => { + expect(severityCountBadges(DEFAULT_FINDING_FIELDS_FIXTURE, { low: 1, critical: 2, 'Sev 2.5': 1, sev3: 3 })).toEqual([ + { key: 'critical', label: 'Critical', color: 'danger', count: 2, tooltip: null }, + { key: 'low', label: 'Low', color: 'informative', count: 1, tooltip: null }, + { key: 'other', label: 'Other', color: 'subtle', count: 4, tooltip: 'Sev 2.5: 1, sev3: 3' }, + ]) + }) + + it('returns no Other badge when every count matches an option', () => { + expect(severityCountBadges(CONTOSO_FINDING_FIELDS, { sev2: 1, sev1: 0 })).toEqual([ + { key: 'sev2', label: 'Sev 2', color: 'warning', count: 1, tooltip: null }, + ]) + expect(severityCountBadges(CONTOSO_FINDING_FIELDS, {})).toEqual([]) + }) +}) diff --git a/frontend/src/utils/findingFields.ts b/frontend/src/utils/findingFields.ts new file mode 100644 index 0000000000..5b90ef8b61 --- /dev/null +++ b/frontend/src/utils/findingFields.ts @@ -0,0 +1,178 @@ +import type { FindingField, FindingFieldOption, FindingFields, FindingSeverityColor } from '@/types' + +export type BadgeColor = 'danger' | 'severe' | 'warning' | 'success' | 'brand' | 'informative' | 'important' | 'subtle' + +export const FLUENT_BADGE_COLOR: Record = { + red: 'danger', + orange: 'severe', + yellow: 'warning', + green: 'success', + blue: 'brand', + gray: 'informative', + dark: 'important', + neutral: 'subtle', +} + +const SEVERITY_KEY = 'severity' + +export interface FindingValueRow { + key: string + label: string + /** The option label of a matching choice value, otherwise the stored text. */ + value: string + multiline: boolean +} + +export interface SeverityCountBadge { + key: string + label: string + color: BadgeColor + count: number + /** The `value: count` of each unmatched severity, on the Other badge only. */ + tooltip: string | null +} + +const OTHER_SEVERITY_BADGE: Pick = { + key: 'other', + label: 'Other', + color: FLUENT_BADGE_COLOR.neutral, +} + +function filled(value: string | undefined): string | null { + return value !== undefined && value.trim() !== '' ? value : null +} + +function allFields(fields: FindingFields): FindingField[] { + return [fields.severity, ...fields.fields] +} + +/** The severity's option label and badge color, or the stored text on a neutral badge when no option matches. */ +export function severityBadge(fields: FindingFields, severity: string): { label: string; color: BadgeColor } { + const match = fields.severity.options.find((option: FindingFieldOption) => option.key === severity) + if (match === undefined) return { label: severity, color: FLUENT_BADGE_COLOR.neutral } + return { label: match.label, color: FLUENT_BADGE_COLOR[fields.severity.colors[severity] ?? 'neutral'] } +} + +/** The field itself, then each of its levels: one entry per dropdown of a choice field. */ +export function choiceLevels(field: FindingField): Array<{ key: string; label: string }> { + return [ + { key: field.key, label: field.label }, + ...field.levels.map(([key, label]: [string, string]) => ({ key, label })), + ] +} + +/** + * The options a choice field's dropdown offers under the values chosen above it. + * @param values - Values by key. The severity slot's own value goes under `severity`. + * @param depth - 0 for the field's own dropdown, 1 for its first level, and so on. + * @returns No options when a value above is missing or isn't an option. + */ +export function levelOptions(field: FindingField, values: Record, depth: number): FindingFieldOption[] { + let options = field.options + for (const level of choiceLevels(field).slice(0, depth)) { + const value = filled(values[level.key]) + const match = options.find((option: FindingFieldOption) => option.key === value) + if (match === undefined) return [] + options = match.children + } + return options +} + +/** The key of every field and level, except `severity`, which isn't stored in `field_values`. */ +export function fieldKeys(fields: FindingFields): Set { + const keys = new Set() + for (const field of allFields(fields)) { + for (const level of choiceLevels(field)) keys.add(level.key) + } + keys.delete(SEVERITY_KEY) + return keys +} + +function displayValue(field: FindingField, values: Record, depth: number, value: string): string { + if (field.type !== 'choice') return value + const match = levelOptions(field, values, depth).find((option: FindingFieldOption) => option.key === value) + return match?.label ?? value +} + +/** + * Lists a finding's filled values in configured order, and its retired values in stored order. + * The severity itself isn't a row; it is the parent path of the severity slot's levels. + */ +export function describeFindingValues( + fields: FindingFields, + severity: string, + fieldValues: Record, +): { rows: FindingValueRow[]; retired: Array<{ key: string; value: string }> } { + const values: Record = { ...fieldValues, [SEVERITY_KEY]: severity } + const rows: FindingValueRow[] = [] + for (const field of allFields(fields)) { + for (const [depth, level] of choiceLevels(field).entries()) { + if (field === fields.severity && depth === 0) continue + const value = filled(fieldValues[level.key]) + if (value === null) continue + rows.push({ + key: level.key, + label: level.label, + value: displayValue(field, values, depth, value), + multiline: field.type === 'multiline', + }) + } + } + const known = fieldKeys(fields) + const retired = Object.entries(fieldValues) + .filter(([key]: [string, string]) => !known.has(key)) + .map(([key, value]: [string, string]) => ({ key, value })) + return { rows, retired } +} + +function firstMissingLevel(field: FindingField, values: Record): string | null { + for (const [depth, level] of choiceLevels(field).entries()) { + if (levelOptions(field, values, depth).length === 0) return null + if (filled(values[level.key]) === null) return level.label + } + return null +} + +/** + * The label of the first missing required value, as the backend checks them: the severity always, + * a required text field when blank, and each level of a required choice that offers options under + * the chosen parent. + */ +export function firstMissingRequired( + fields: FindingFields, + severity: string, + fieldValues: Record, +): string | null { + const values: Record = { ...fieldValues, [SEVERITY_KEY]: severity } + for (const field of allFields(fields)) { + if (field !== fields.severity && !field.required) continue + if (field.type === 'choice') { + const missing = firstMissingLevel(field, values) + if (missing !== null) return missing + } else if (filled(values[field.key]) === null) { + return field.label + } + } + return null +} + +/** One badge per counted severity option, in option order, then one Other badge for every unmatched severity. */ +export function severityCountBadges(fields: FindingFields, counts: Record): SeverityCountBadge[] { + const badges: SeverityCountBadge[] = [] + for (const option of fields.severity.options) { + const count = counts[option.key] ?? 0 + if (count > 0) badges.push({ key: option.key, ...severityBadge(fields, option.key), count, tooltip: null }) + } + const optionKeys = new Set(fields.severity.options.map((option: FindingFieldOption) => option.key)) + const unmatched = Object.entries(counts) + .filter(([severity, count]: [string, number]) => count > 0 && !optionKeys.has(severity)) + .sort(([left]: [string, number], [right]: [string, number]) => left.localeCompare(right)) + if (unmatched.length > 0) { + badges.push({ + ...OTHER_SEVERITY_BADGE, + count: unmatched.reduce((total: number, [, count]: [string, number]) => total + count, 0), + tooltip: unmatched.map(([severity, count]: [string, number]) => `${severity}: ${count}`).join(', '), + }) + } + return badges +} diff --git a/pyrit/backend/services/finding_export.py b/pyrit/backend/services/finding_export.py index 3679c1d00a..9c061117e3 100644 --- a/pyrit/backend/services/finding_export.py +++ b/pyrit/backend/services/finding_export.py @@ -29,8 +29,7 @@ FindingExportRequest, ) from pyrit.backend.services.operation_service import OperationService -from pyrit.models import Finding, FindingSeverity, ScoreStatus -from pyrit.models.harm_category import HarmCategory +from pyrit.models import Finding, FindingFields, ScoreStatus from pyrit.output.finding_report.markdown import MarkdownFindingReportPrinter from pyrit.output.finding_report.report import ( DETAIL_SEPARATOR, @@ -44,6 +43,7 @@ FindingReportField, plain_text, ) +from pyrit.setup.finding_fields import get_active_finding_fields _MARKDOWN_MEDIA_TYPE = "text/markdown; charset=utf-8" _DOCX_MEDIA_TYPE = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" @@ -113,14 +113,15 @@ async def _build_report_async(*, operation_id: UUID, request: FindingExportReque selected = await service.get_findings_with_evidence_async( operation_id=operation_id, finding_ids=request.finding_ids ) + fields = get_active_finding_fields() report_findings = [] for finding, evidence in selected: - severity, fields = _project_finding_fields(finding) + severity, projected = _project_finding_fields(finding, fields=fields) report_findings.append( FindingReportEntry( title=finding.title, severity=severity, - fields=fields, + fields=projected, evidence=tuple( _project_evidence(item=item, operation_id=operation_id, dashboard_origin=request.dashboard_origin) for item in evidence @@ -130,28 +131,29 @@ async def _build_report_async(*, operation_id: UUID, request: FindingExportReque return FindingReport(operation_name=operation.name, findings=tuple(report_findings)) -def _project_finding_fields(finding: Finding) -> tuple[FindingReportField, tuple[FindingReportField, ...]]: +def _project_finding_fields( + finding: Finding, *, fields: FindingFields +) -> tuple[FindingReportField, tuple[FindingReportField, ...]]: """ - Label a finding's severity and its populated fields, in report order, and place each in a Word column. + Label a finding's severity and its filled fields, in configured order, and place each in a Word column. - Narrative text goes in the Finding column and classifications in the Assessment column. + Multiline text goes in the Finding column; choices and single-line text go in the Assessment column. + Internal fields and retired values are left out. Returns: - tuple[FindingReportField, tuple[FindingReportField, ...]]: The severity, then the populated fields. + tuple[FindingReportField, tuple[FindingReportField, ...]]: The severity, then the filled fields. """ - severity = ( - finding.severity_other or "Other" - if finding.severity == FindingSeverity.OTHER - else finding.severity.value.capitalize() + description = fields.describe_values(severity=finding.severity, field_values=finding.field_values) + projected = tuple( + FindingReportField( + label=row.label, + value=row.value, + column=FindingReportColumn.FINDING if row.multiline else FindingReportColumn.ASSESSMENT, + ) + for row in description.rows + if not row.internal ) - harm_type = finding.harm_type_other if finding.harm_type == HarmCategory.OTHER else finding.harm_type - fields = [ - FindingReportField(label="Harm-type", value=str(harm_type), column=FindingReportColumn.ASSESSMENT) - if harm_type - else None, - FindingReportField(label="Description", value=finding.description) if finding.description else None, - ] - return FindingReportField(label="Severity", value=severity), tuple(field for field in fields if field) + return FindingReportField(label=fields.severity.label, value=description.severity_label), projected def _project_evidence( diff --git a/pyrit/models/finding_fields.py b/pyrit/models/finding_fields.py index f0d6e1a89e..1f2e2dbbe1 100644 --- a/pyrit/models/finding_fields.py +++ b/pyrit/models/finding_fields.py @@ -64,6 +64,7 @@ class FindingField(BaseModel): colors (dict[str, FindingSeverityColor]): The severity slot's badge color for each top-level option. hint (str | None): Help text under the field in the finding dialog. section (str | None): A heading the finding dialog shows where the section changes. + internal (bool): Whether the field and its levels stay out of exports. The dashboard still shows them. """ model_config = ConfigDict(extra="forbid", frozen=True) @@ -81,6 +82,7 @@ class FindingField(BaseModel): colors: dict[str, FindingSeverityColor] = {} hint: str | None = None section: str | None = None + internal: bool = False @field_validator("options", mode="before") @classmethod @@ -202,6 +204,28 @@ def _check_top_level_references(self) -> None: raise ValueError(f"'{key}' is not a top-level option of '{self.key}'.") +class FindingValueRow(BaseModel): + """One filled value of a finding, labelled for display.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + key: str + label: str + value: str + multiline: bool + internal: bool + + +class FindingValueDescription(BaseModel): + """A finding's values labelled for display, with the stored values that no field uses.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + severity_label: str + rows: list[FindingValueRow] + retired: list[tuple[str, str]] + + class FindingFields(BaseModel): """A deployment's finding fields: the severity slot and the fields shown after it.""" @@ -230,8 +254,8 @@ def _validate_severity_slot(cls, value: FindingField) -> FindingField: FindingField: The slot with ``required`` set to True. Raises: - ValueError: If the slot's key is not ``severity``, it is not a choice field, or it sets - ``required=False``. + ValueError: If the slot's key is not ``severity``, it is not a choice field, it sets + ``required=False``, or it is internal. """ if value.key != "severity": raise ValueError(f"The severity slot's key must be 'severity', not '{value.key}'.") @@ -239,6 +263,8 @@ def _validate_severity_slot(cls, value: FindingField) -> FindingField: raise ValueError("The severity slot must be a choice field.") if "required" in value.model_fields_set and not value.required: raise ValueError("The severity slot is always required; leave required out or set it to True.") + if value.internal: + raise ValueError("The severity slot can't be internal.") return value.model_copy(update={"required": True}) @model_validator(mode="after") @@ -297,6 +323,32 @@ def check_values( for field in [self.severity, *self.fields]: _check_field_values(field=field, values=values, previous_values=previous) + def describe_values(self, *, severity: str, field_values: dict[str, str]) -> FindingValueDescription: + """ + Label a finding's filled values for display. + + Rows follow the configured order: the severity slot's levels, then each field and its levels. + A choice value shows the label of the matching option under its stored parent; any other value + shows as typed. + + Args: + severity (str): The severity slot's value. + field_values (dict[str, str]): The values of the other fields and of every level, by key. + + Returns: + FindingValueDescription: The severity's label, one row per filled value, and the + ``(key, value)`` of each stored value that no field or level uses, in stored order. + """ + values = {**field_values, "severity": severity} + rows = [row for field in [self.severity, *self.fields] for row in _value_rows(field=field, values=values)] + severity_option = _find_option(options=self.severity.options, value=severity) + known = self.field_keys + return FindingValueDescription( + severity_label=severity_option.label if severity_option is not None else severity, + rows=[row for row in rows if row.key != "severity"], + retired=[(key, value) for key, value in field_values.items() if key not in known], + ) + def _check_known_keys(self, *, field_values: dict[str, str], previous_values: dict[str, str]) -> None: """ Reject a key that no field or level uses, unless its value has not changed. @@ -513,7 +565,7 @@ def _check_level( if required and options: raise FindingValuesError(f"{label} is required.") return [] - match = next((option for option in options if option.key == value), None) + match = _find_option(options=options, value=value) if match is not None: return match.children if not unchanged: @@ -558,6 +610,46 @@ def _non_blank(value: str | None) -> str | None: return value if value and value.strip() else None +def _find_option(*, options: list[FindingFieldOption], value: str | None) -> FindingFieldOption | None: + """ + Find the option whose key is a dropdown's value. + + Args: + options (list[FindingFieldOption]): The options the dropdown offers. + value (str | None): The dropdown's value, or None if it is missing. + + Returns: + FindingFieldOption | None: The matching option, or None if no option has that key. + """ + return next((option for option in options if option.key == value), None) + + +def _value_rows(*, field: FindingField, values: dict[str, str]) -> Iterator[FindingValueRow]: + """ + Label the filled values of one field and of its levels. + + Args: + field (FindingField): The field. + values (dict[str, str]): Every value of the finding, by key. + + Yields: + FindingValueRow: A row for each filled value, in level order. + """ + options = field.options + for key, label in [(field.key, field.label), *field.levels]: + value = _non_blank(values.get(key)) + match = _find_option(options=options, value=value) + options = match.children if match is not None else [] + if value is not None: + yield FindingValueRow( + key=key, + label=label, + value=match.label if match is not None else value, + multiline=field.type == FindingFieldType.MULTILINE, + internal=field.internal, + ) + + DEFAULT_FINDING_FIELDS = FindingFields( severity=FindingField( key="severity", diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 2ceceb1f25..bcc601d798 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -91,9 +91,9 @@ class _ConfigurationLayer: reset_extensions: bool -def _check_finding_fields_initializers(initializers: Sequence["PyRITInitializer"]) -> None: +def _check_findings_initializers(initializers: Sequence["PyRITInitializer"]) -> None: """ - Check the finding fields that ``FindingFieldsInitializer`` subclasses would set, without making them active. + Check the finding fields that ``FindingsInitializer`` subclasses would set, without making them active. Code that calls ``set_active_finding_fields`` directly isn't checked. @@ -101,16 +101,16 @@ def _check_finding_fields_initializers(initializers: Sequence["PyRITInitializer" initializers (Sequence[PyRITInitializer]): The initializers of the replacement setup pass. Raises: - ValueError: If two or more ``FindingFieldsInitializer`` subclasses are configured, or the configured one + ValueError: If two or more ``FindingsInitializer`` subclasses are configured, or the configured one returns invalid fields or something other than ``FindingFields``. """ - from pyrit.setup.initializers import FindingFieldsInitializer + from pyrit.setup.initializers import FindingsInitializer - setters = [initializer for initializer in initializers if isinstance(initializer, FindingFieldsInitializer)] + setters = [initializer for initializer in initializers if isinstance(initializer, FindingsInitializer)] if len(setters) > 1: first, second = (type(setter).__name__ for setter in setters[:2]) raise ValueError( - f"{first} and {second} both set the finding fields. Use one FindingFieldsInitializer per deployment." + f"{first} and {second} both set the finding fields. Use one FindingsInitializer per deployment." ) if not setters: return @@ -738,7 +738,7 @@ async def preflight_reinitialization_async( Raises: ValueError: If an initializer is invalid, a required environment value is missing, more than one - ``FindingFieldsInitializer`` subclass is configured, or the configured one returns invalid finding + ``FindingsInitializer`` subclass is configured, or the configured one returns invalid finding fields. Code that calls ``set_active_finding_fields`` directly isn't checked. """ from pyrit.registry import InitializerRegistry @@ -777,7 +777,7 @@ async def preflight_reinitialization_async( raise ValueError( f"Initializer '{type(initializer).__name__}' has missing required environment variables." ) - _check_finding_fields_initializers(initializers) + _check_findings_initializers(initializers) return PreparedReinitialization( environment_values=dict(environment_values), diff --git a/pyrit/setup/finding_fields.py b/pyrit/setup/finding_fields.py index c5b22e007c..e911dcfda7 100644 --- a/pyrit/setup/finding_fields.py +++ b/pyrit/setup/finding_fields.py @@ -36,7 +36,7 @@ def set_active_finding_fields(*, finding_fields: FindingFields, source: str) -> if _active is not None: raise ValueError( f"{source} cannot set the finding fields: {_active[1]} already set them in this setup pass. " - "Use one FindingFieldsInitializer per deployment." + "Use one FindingsInitializer per deployment." ) _active = (finding_fields, source) diff --git a/pyrit/setup/initializers/__init__.py b/pyrit/setup/initializers/__init__.py index 6bd9964d5f..a069b659ba 100644 --- a/pyrit/setup/initializers/__init__.py +++ b/pyrit/setup/initializers/__init__.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from pyrit.models.parameter import Parameter from pyrit.setup.initializers.converters import ConverterInitializer - from pyrit.setup.initializers.finding_fields import FindingFieldsInitializer + from pyrit.setup.initializers.findings import FindingsInitializer from pyrit.setup.initializers.load_default_datasets import LoadDefaultDatasets from pyrit.setup.initializers.preload_scenario_metadata import PreloadScenarioMetadata from pyrit.setup.initializers.refresh_datasets import RefreshDatasets @@ -25,7 +25,7 @@ "Parameter": "pyrit.models.parameter", "PyRITInitializer": "pyrit.setup.pyrit_initializer", "ConverterInitializer": "pyrit.setup.initializers.converters", - "FindingFieldsInitializer": "pyrit.setup.initializers.finding_fields", + "FindingsInitializer": "pyrit.setup.initializers.findings", "TechniqueInitializer": "pyrit.setup.initializers.techniques", "ToolsInitializer": "pyrit.setup.initializers.tools", "ScorerInitializer": "pyrit.setup.initializers.scorers", diff --git a/pyrit/setup/initializers/finding_fields.py b/pyrit/setup/initializers/findings.py similarity index 92% rename from pyrit/setup/initializers/finding_fields.py rename to pyrit/setup/initializers/findings.py index d4025b0d3d..b7c31adf39 100644 --- a/pyrit/setup/initializers/finding_fields.py +++ b/pyrit/setup/initializers/findings.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Base initializer for a deployment's finding fields.""" +"""Base initializer for a deployment's findings.""" from abc import abstractmethod from typing import Any @@ -11,7 +11,7 @@ from pyrit.setup.pyrit_initializer import PyRITInitializer -class FindingFieldsInitializer(PyRITInitializer): +class FindingsInitializer(PyRITInitializer): """Set a deployment's finding fields from ``get_finding_fields``.""" @abstractmethod diff --git a/tests/unit/backend/test_finding_export.py b/tests/unit/backend/test_finding_export.py index 1b8ea16598..e846b3ae62 100644 --- a/tests/unit/backend/test_finding_export.py +++ b/tests/unit/backend/test_finding_export.py @@ -36,13 +36,14 @@ ConversationType, Finding, FindingEvidence, + FindingField, + FindingFields, Message, MessagePiece, Operation, Score, ScoreStatus, ) -from pyrit.models.harm_category import HarmCategory from pyrit.output.finding_report.markdown import MarkdownFindingReportPrinter from pyrit.output.finding_report.report import ( FindingReport, @@ -51,6 +52,7 @@ FindingReportEvidence, FindingReportField, ) +from pyrit.setup.finding_fields import set_active_finding_fields pytestmark = pytest.mark.usefixtures("patch_central_database") @@ -80,6 +82,31 @@ ) NOT_SCORED = FindingReportField(label="Score", value="Not scored") NOT_IN_AN_OPERATION = FindingReportField(label="From operation", value="Not in an operation") +CONTOSO_FIELDS = FindingFields( + severity=FindingField( + key="severity", + label="Impact", + type="choice", + levels=[("severity_rationale", "Impact rationale")], + options={ + ("sev1", "Sev 1"): [("remote_no_auth", "Remote, no sign-in needed")], + ("sev2", "Sev 2"): [("signed_in_user", "Needs a signed-in user")], + }, + ), + fields=[ + FindingField( + key="category", + label="Category", + type="choice", + required=True, + levels=[("subcategory", "Subcategory")], + options={("model_security", "Model security"): [("system_prompt_leak", "System prompt leak")]}, + ), + FindingField(key="steps", label="Steps to reproduce", type="multiline"), + FindingField(key="build", label="Build", type="text"), + FindingField(key="owner", label="Owner", type="text", internal=True), + ], +) def _client() -> TestClient: @@ -212,8 +239,7 @@ async def _save_named_findings_async( operation=operation, title="Critical finding", severity="critical", - harm_type=HarmCategory.PROFANITY, - description="Synthetic description\nsecond line", + field_values={"harm_type": "Profanity", "description": "Synthetic description\nsecond line"}, ), "Included finding": await _add_finding_async(memory=memory, operation=operation, title="Included finding"), "Excluded finding": await _add_finding_async(memory=memory, operation=operation, title="Excluded finding"), @@ -743,17 +769,15 @@ async def test_report_projects_optional_and_custom_fields(sqlite_instance: SQLit memory=sqlite_instance, operation=operation, title="Custom finding", - severity="other", - severity_other="Team severity", - harm_type=HarmCategory.OTHER, - harm_type_other="Team harm", + severity="Team severity", + field_values={"harm_type": "Team harm"}, ) plain = await _add_finding_async( memory=sqlite_instance, operation=operation, title="Plain finding", severity="informational", - description="Only text", + field_values={"description": "Only text"}, ) attack = await _add_attack_async(memory=sqlite_instance, conversation_id="main", objective=" ") evidence = await _link_async(memory=sqlite_instance, finding=custom, attack_result_id=attack.attack_result_id) @@ -784,6 +808,53 @@ async def test_report_projects_optional_and_custom_fields(sqlite_instance: SQLit ) +@pytest.mark.usefixtures("reset_finding_fields") +async def test_report_places_configured_fields_by_type(sqlite_instance: SQLiteMemory) -> None: + set_active_finding_fields(finding_fields=CONTOSO_FIELDS, source="test") + operation = await sqlite_instance.add_operation_async(Operation(name="Contoso assistant")) + finding = await _add_finding_async( + memory=sqlite_instance, + operation=operation, + title="System prompt leak", + severity="Sev 2.5", + field_values={ + "owner": "dana", + "description": "Retired text", + "steps": "1. Ask for the instructions\n2. Repeat them", + "build": "2026.10", + "subcategory": "system_prompt_leak", + "category": "model_security", + "severity_rationale": "remote_no_auth", + }, + ) + + report = await _report_async(operation_id=operation.id, finding_ids=[finding.id]) + + # Internal fields and retired values stay out; the rationale has no matching parent, so it shows as typed. + assert report.findings == ( + FindingReportEntry( + title="System prompt leak", + severity=FindingReportField(label="Impact", value="Sev 2.5"), + fields=( + FindingReportField( + label="Impact rationale", value="remote_no_auth", column=FindingReportColumn.ASSESSMENT + ), + FindingReportField(label="Category", value="Model security", column=FindingReportColumn.ASSESSMENT), + FindingReportField( + label="Subcategory", value="System prompt leak", column=FindingReportColumn.ASSESSMENT + ), + FindingReportField( + label="Steps to reproduce", + value="1. Ask for the instructions\n2. Repeat them", + column=FindingReportColumn.FINDING, + ), + FindingReportField(label="Build", value="2026.10", column=FindingReportColumn.ASSESSMENT), + ), + evidence=(), + ), + ) + + @pytest.mark.parametrize( ("score_value", "scorer", "expected"), [ diff --git a/tests/unit/memory/test_operation_service_findings.py b/tests/unit/backend/test_operation_service_findings.py similarity index 100% rename from tests/unit/memory/test_operation_service_findings.py rename to tests/unit/backend/test_operation_service_findings.py diff --git a/tests/unit/models/test_finding_fields.py b/tests/unit/models/test_finding_fields.py index 83a18eb144..5301474f00 100644 --- a/tests/unit/models/test_finding_fields.py +++ b/tests/unit/models/test_finding_fields.py @@ -10,7 +10,7 @@ FindingFields, FindingValuesError, ) -from pyrit.models.finding_fields import FindingFieldOption +from pyrit.models.finding_fields import FindingFieldOption, FindingValueRow from pyrit.models.harm_category import HarmCategory CONTOSO = FindingFields( @@ -59,7 +59,7 @@ section="Reproduction", hint="List the prompts in the order you sent them.", ), - FindingField(key="owner", label="Owner", type="text", section="Tracking"), + FindingField(key="owner", label="Owner", type="text", section="Tracking", internal=True), ], ) @@ -196,6 +196,75 @@ def test_severity_slot_accepts_explicit_required() -> None: assert FindingFields.model_validate({"severity": {**SEVERITY, "required": True}}).severity.required is True +def test_internal_defaults_to_false_and_is_rejected_on_the_severity_slot() -> None: + assert FindingField.model_validate(_choice()).internal is False + with pytest.raises(ValidationError, match="The severity slot can't be internal."): + FindingFields.model_validate({"severity": {**SEVERITY, "internal": True}}) + + +def test_describe_values_lists_filled_values_in_configured_order_and_retired_keys() -> None: + description = CONTOSO.describe_values( + severity="sev1", + field_values={ + "owner": "dana", + "description": "Old text", + "category": "model_security", + "subcategory": "system_prompt_leak", + "technique": "Many-shot priming", + "steps": " ", + "severity_rationale": "remote_no_auth", + "repro": "Older text", + }, + ) + assert description.severity_label == "Sev 1" + assert [(row.label, row.value) for row in description.rows] == [ + ("Severity rationale", "Remote, no sign-in needed"), + ("Category", "Model security"), + ("Subcategory", "System prompt leak"), + ("Technique", "Many-shot priming"), + ("Owner", "dana"), + ] + assert description.retired == [("description", "Old text"), ("repro", "Older text")] + + +def test_describe_values_shows_a_cascade_child_not_under_its_stored_parent_as_typed() -> None: + description = CONTOSO.describe_values( + severity="sev1", field_values={"category": "model_security", "subcategory": "hate_speech"} + ) + assert [(row.label, row.value) for row in description.rows] == [ + ("Category", "Model security"), + ("Subcategory", "hate_speech"), + ] + + +def test_describe_values_shows_an_unmatched_severity_as_typed() -> None: + description = CONTOSO.describe_values(severity="Sev 2.5", field_values={"severity_rationale": "remote_no_auth"}) + assert description.severity_label == "Sev 2.5" + assert [(row.label, row.value) for row in description.rows] == [("Severity rationale", "remote_no_auth")] + + +def test_describe_values_carries_multiline_and_internal_on_each_row() -> None: + description = CONTOSO.describe_values(severity="sev2", field_values={"owner": "dana", "steps": "1. Open\n2. Ask"}) + assert description.rows == [ + FindingValueRow( + key="steps", label="Steps to reproduce", value="1. Open\n2. Ask", multiline=True, internal=False + ), + FindingValueRow(key="owner", label="Owner", value="dana", multiline=False, internal=True), + ] + assert description.retired == [] + + +def test_describe_values_labels_the_default_fields() -> None: + description = DEFAULT_FINDING_FIELDS.describe_values( + severity="critical", field_values={"description": "Synthetic text", "harm_type": "Hate Speech"} + ) + assert description.severity_label == "Critical" + assert description.rows == [ + FindingValueRow(key="harm_type", label="Harm-type", value="Hate Speech", multiline=False, internal=False), + FindingValueRow(key="description", label="Description", value="Synthetic text", multiline=True, internal=False), + ] + + VALID = { "severity_rationale": "remote_no_auth", "category": "model_security", diff --git a/tests/unit/setup/test_finding_fields_initializer.py b/tests/unit/setup/test_findings_initializer.py similarity index 83% rename from tests/unit/setup/test_finding_fields_initializer.py rename to tests/unit/setup/test_findings_initializer.py index 55da1c7b62..5bf8cd1509 100644 --- a/tests/unit/setup/test_finding_fields_initializer.py +++ b/tests/unit/setup/test_findings_initializer.py @@ -11,7 +11,7 @@ from pyrit.registry import InitializerRegistry from pyrit.registry.discovery import discover_in_directory from pyrit.setup.finding_fields import get_active_finding_fields -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer from pyrit.setup.pyrit_initializer import PyRITInitializer pytestmark = pytest.mark.usefixtures("reset_finding_fields") @@ -22,10 +22,10 @@ CONTOSO_SCRIPT = ''' from pyrit.models import FindingField, FindingFields -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer -class ContosoFindingFields(FindingFieldsInitializer): +class ContosoFindingFields(FindingsInitializer): """Contoso finding fields.""" def get_finding_fields(self) -> FindingFields: @@ -35,14 +35,14 @@ def get_finding_fields(self) -> FindingFields: ''' -class ExampleFindingFields(FindingFieldsInitializer): +class ExampleFindingFields(FindingsInitializer): """Example finding fields.""" def get_finding_fields(self) -> FindingFields: return EXAMPLE -async def test_finding_fields_initializer_sets_the_active_fields() -> None: +async def test_findings_initializer_sets_the_active_fields() -> None: await ExampleFindingFields().initialize_async() assert get_active_finding_fields() is EXAMPLE @@ -57,22 +57,22 @@ async def test_get_info_leaves_the_active_fields_alone() -> None: assert info["global_variables"] == [] -def test_finding_fields_initializer_is_skipped_by_builtin_discovery() -> None: +def test_findings_initializer_is_skipped_by_builtin_discovery() -> None: directory = Path(pyrit.setup.initializers.__file__).parent names = [ cls.__name__ for _, _, cls in discover_in_directory(directory=directory, base_class=PyRITInitializer, recursive=True) ] - assert "FindingFieldsInitializer" not in names + assert "FindingsInitializer" not in names async def test_discovery_copy_shares_the_active_finding_fields() -> None: - path = Path(pyrit.setup.initializers.__file__).parent / "finding_fields.py" - spec = importlib.util.spec_from_file_location("discovered_module.finding_fields", path) + path = Path(pyrit.setup.initializers.__file__).parent / "findings.py" + spec = importlib.util.spec_from_file_location("discovered_module.findings", path) copy = importlib.util.module_from_spec(spec) spec.loader.exec_module(copy) - class CopiedFindingFields(copy.FindingFieldsInitializer): + class CopiedFindingFields(copy.FindingsInitializer): def get_finding_fields(self) -> FindingFields: return EXAMPLE diff --git a/tests/unit/setup/test_initialization.py b/tests/unit/setup/test_initialization.py index bbb1f80356..2f94efb109 100644 --- a/tests/unit/setup/test_initialization.py +++ b/tests/unit/setup/test_initialization.py @@ -19,7 +19,7 @@ from pyrit.setup import IN_MEMORY, initialization, initialize_pyrit_async from pyrit.setup.finding_fields import get_active_finding_fields, set_active_finding_fields from pyrit.setup.initialization import reinitialize_pyrit_async -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer from pyrit.setup.pyrit_initializer import PyRITInitializer CUSTOM = FindingFields( @@ -482,12 +482,12 @@ async def test_reinitialize_resets_finding_fields() -> None: @pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") -async def test_second_finding_fields_initializer_stops_setup() -> None: - class FirstFindingFields(FindingFieldsInitializer): +async def test_second_findings_initializer_stops_setup() -> None: + class FirstFindingFields(FindingsInitializer): def get_finding_fields(self) -> FindingFields: return CUSTOM - class SecondFindingFields(FindingFieldsInitializer): + class SecondFindingFields(FindingsInitializer): def get_finding_fields(self) -> FindingFields: return CUSTOM @@ -502,7 +502,7 @@ def get_finding_fields(self) -> FindingFields: @pytest.mark.usefixtures("reset_memory_singletons", "reset_finding_fields") async def test_invalid_finding_fields_stop_setup() -> None: - class BrokenFindingFields(FindingFieldsInitializer): + class BrokenFindingFields(FindingsInitializer): def get_finding_fields(self) -> FindingFields: return FindingFields( severity=CUSTOM.severity, diff --git a/tests/unit/setup/test_reinitialization.py b/tests/unit/setup/test_reinitialization.py index 77a903e355..8950484163 100644 --- a/tests/unit/setup/test_reinitialization.py +++ b/tests/unit/setup/test_reinitialization.py @@ -121,19 +121,19 @@ async def test_preflight_uses_isolated_registry_without_changing_live_state() -> @pytest.mark.usefixtures("reset_finding_fields") -async def test_preflight_rejects_two_finding_fields_initializers(tmp_path: Path) -> None: +async def test_preflight_rejects_two_findings_initializers(tmp_path: Path) -> None: script = tmp_path / "two_finding_fields.py" script.write_text( """ from pyrit.models import FindingField, FindingFields -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer SEVERITY = FindingField(key="severity", label="Severity", type="choice", options=[("sev1", "Sev 1")]) -class ContosoFindingFields(FindingFieldsInitializer): +class ContosoFindingFields(FindingsInitializer): def get_finding_fields(self): return FindingFields(severity=SEVERITY) -class ContosoBackupFields(FindingFieldsInitializer): +class ContosoBackupFields(FindingsInitializer): def get_finding_fields(self): return FindingFields(severity=SEVERITY) """, encoding="utf-8", @@ -143,7 +143,7 @@ def get_finding_fields(self): return FindingFields(severity=SEVERITY) message = ( "ContosoBackupFields and ContosoFindingFields both set the finding fields. " - "Use one FindingFieldsInitializer per deployment." + "Use one FindingsInitializer per deployment." ) with pytest.raises(ValueError, match=re.escape(message)): await config.preflight_reinitialization_async(environment_values={}) @@ -156,9 +156,9 @@ async def test_preflight_rejects_invalid_finding_fields(tmp_path: Path) -> None: script.write_text( """ from pyrit.models import FindingField, FindingFields -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer -class ContosoFindingFields(FindingFieldsInitializer): +class ContosoFindingFields(FindingsInitializer): def get_finding_fields(self): return FindingFields(severity=FindingField(key="severity", label="Severity", type="choice", options=[])) """, @@ -177,9 +177,9 @@ async def test_preflight_rejects_finding_fields_of_the_wrong_type(tmp_path: Path script = tmp_path / "wrong_type_finding_fields.py" script.write_text( """ -from pyrit.setup.initializers import FindingFieldsInitializer +from pyrit.setup.initializers import FindingsInitializer -class ContosoFindingFields(FindingFieldsInitializer): +class ContosoFindingFields(FindingsInitializer): def get_finding_fields(self): return {} """, From f03c0777d5c7fac6dd03f7a7d8376757bbb687b5 Mon Sep 17 00:00:00 2001 From: Adrian Gavrila Date: Thu, 29 Oct 2026 09:35:52 -0400 Subject: [PATCH 09/20] Add the finding field input with cascades and Other text Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Operations/FindingFieldInput.styles.ts | 9 + .../Operations/FindingFieldInput.test.tsx | 371 ++++++++++++++++++ .../Operations/FindingFieldInput.tsx | 226 +++++++++++ 3 files changed, 606 insertions(+) create mode 100644 frontend/src/components/Operations/FindingFieldInput.styles.ts create mode 100644 frontend/src/components/Operations/FindingFieldInput.test.tsx create mode 100644 frontend/src/components/Operations/FindingFieldInput.tsx diff --git a/frontend/src/components/Operations/FindingFieldInput.styles.ts b/frontend/src/components/Operations/FindingFieldInput.styles.ts new file mode 100644 index 0000000000..4e89a28223 --- /dev/null +++ b/frontend/src/components/Operations/FindingFieldInput.styles.ts @@ -0,0 +1,9 @@ +import { makeStyles, tokens } from '@fluentui/react-components' + +const LISTBOX_MAX_HEIGHT = '240px' + +export const useFindingFieldInputStyles = makeStyles({ + levels: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalM }, + picker: { minWidth: 0, width: '100%' }, + listbox: { maxHeight: LISTBOX_MAX_HEIGHT, overflowY: 'auto' }, +}) diff --git a/frontend/src/components/Operations/FindingFieldInput.test.tsx b/frontend/src/components/Operations/FindingFieldInput.test.tsx new file mode 100644 index 0000000000..80ea60feaa --- /dev/null +++ b/frontend/src/components/Operations/FindingFieldInput.test.tsx @@ -0,0 +1,371 @@ +import React, { useState } from 'react' +import { render, screen } from '@testing-library/react' +import userEvent, { type UserEvent } from '@testing-library/user-event' +import { FluentProvider, webLightTheme } from '@fluentui/react-components' + +import { CONTOSO_FINDING_FIELDS } from '@/test-utils/findingFieldsFixtures' +import type { FindingField, FindingFieldOption } from '@/types' + +import FindingFieldInput from './FindingFieldInput' + +const TestWrapper: React.FC<{ children: React.ReactNode }> = ({ children }) => ( + {children} +) + +const UNMATCHED_HINT = 'Not a current option. It stays unless you change it.' + +function contosoField(key: string): FindingField { + const field = CONTOSO_FINDING_FIELDS.fields.find((candidate: FindingField) => candidate.key === key) + if (field === undefined) throw new Error(`CONTOSO_FINDING_FIELDS has no field ${key}`) + return field +} + +function option(key: string, label: string, children: FindingFieldOption[] = []): FindingFieldOption { + return { key, label, children } +} + +const CATEGORY = contosoField('category') +const TECHNIQUE = contosoField('technique') +const STEPS = contosoField('steps') +const OWNER = contosoField('owner') + +// Both surfaces list a `plugins` component, but only the web chat's plugins include email. +const SURFACE: FindingField = { + ...CATEGORY, + key: 'surface', + label: 'Surface', + levels: [['component', 'Component'], ['capability', 'Capability']], + options: [ + option('web_chat', 'Web chat', [ + option('plugins', 'Plugins', [option('calendar', 'Calendar'), option('email', 'Email')]), + ]), + option('mobile_app', 'Mobile app', [option('plugins', 'Plugins', [option('calendar', 'Calendar')])]), + ], +} + +// Outage is a leaf option, so it leaves Impact detail nothing to offer. +const IMPACT: FindingField = { + ...CATEGORY, + key: 'impact', + label: 'Impact', + levels: [['impact_detail', 'Impact detail']], + options: [ + option('data_exposure', 'Data exposure', [option('customer_records', 'Customer records')]), + option('outage', 'Outage'), + ], +} + +const TACTIC: FindingField = { ...TECHNIQUE, key: 'tactic', label: 'Tactic' } + +interface FieldHarnessProps { + field: FindingField + initialValues: Record + onChange: (changes: Record) => void + onOtherBlankChange?: (blank: boolean) => void +} + +// Holds the draft the way the finding dialog will, so a test can follow several changes. +function FieldHarness({ field, initialValues, onChange, onOtherBlankChange }: FieldHarnessProps) { + const [values, setValues] = useState(initialValues) + return ( + ) => { + onChange(changes) + setValues((previous) => ({ ...previous, ...changes })) + }} /> + ) +} + +async function openPicker(user: UserEvent, label: string): Promise { + await user.click(screen.getByRole('combobox', { name: label })) + await screen.findByRole('listbox') +} + +async function pick(user: UserEvent, label: string, optionName: string): Promise { + await openPicker(user, label) + await user.click(screen.getByRole('option', { name: optionName })) +} + +function optionNames(): string[] { + return screen.getAllByRole('option').map((element: HTMLElement) => element.textContent ?? '') +} + +describe('FindingFieldInput', () => { + const onChange = jest.fn() + + beforeEach(() => { + jest.clearAllMocks() + }) + + it('renders a text or multiline field with its label, hint and required marker', async () => { + const user = userEvent.setup() + render( + + + ) + + const steps = screen.getByRole('textbox', { name: 'Steps to reproduce' }) + expect(steps).toBeInstanceOf(HTMLTextAreaElement) + expect(steps).toHaveValue('1. Open a chat') + expect(steps).toBeRequired() + expect(steps).toHaveAccessibleDescription('List each prompt you sent to the Contoso assistant, in order.') + expect(screen.getByText('List each prompt you sent to the Contoso assistant, in order.')).toBeInTheDocument() + + const owner = screen.getByRole('textbox', { name: 'Owner' }) + expect(owner).toBeInstanceOf(HTMLInputElement) + expect(owner).toHaveValue('') + expect(owner).not.toBeRequired() + expect(owner).not.toHaveAccessibleDescription() + + await user.type(steps, '.') + expect(onChange).toHaveBeenLastCalledWith({ steps: '1. Open a chat.' }) + await user.type(owner, 'C') + expect(onChange).toHaveBeenLastCalledWith({ owner: 'C' }) + }) + + it('picks a cascade path level by level and disables a level until its parent has a value', async () => { + const user = userEvent.setup() + render() + + expect(screen.getByRole('combobox', { name: 'Category' })).toBeEnabled() + expect(screen.getByRole('combobox', { name: 'Subcategory' })).toBeDisabled() + expect(screen.getByRole('combobox', { name: 'Subcategory' })).not.toBeRequired() + + await pick(user, 'Category', 'Model security') + expect(onChange).toHaveBeenLastCalledWith({ category: 'model_security' }) + expect(screen.getByRole('combobox', { name: 'Category' })).toHaveValue('Model security') + expect(screen.getByRole('combobox', { name: 'Subcategory' })).toBeEnabled() + expect(screen.getByRole('combobox', { name: 'Subcategory' })).toBeRequired() + + await openPicker(user, 'Subcategory') + expect(optionNames()).toEqual(['System prompt leak', 'Tool misuse']) + await user.click(screen.getByRole('option', { name: 'Tool misuse' })) + expect(onChange).toHaveBeenLastCalledWith({ subcategory: 'tool_misuse' }) + expect(screen.getByRole('combobox', { name: 'Subcategory' })).toHaveValue('Tool misuse') + }) + + it('clears a child that is not valid under the new parent and keeps one that is', async () => { + const user = userEvent.setup() + render( + + + ) + + await pick(user, 'Category', 'Content safety') + expect(onChange).toHaveBeenLastCalledWith({ category: 'content_safety', subcategory: '' }) + expect(screen.getByRole('combobox', { name: 'Subcategory' })).toHaveValue('') + + await pick(user, 'Surface', 'Mobile app') + expect(onChange).toHaveBeenLastCalledWith({ surface: 'mobile_app', capability: '' }) + expect(screen.getByRole('combobox', { name: 'Component' })).toHaveValue('Plugins') + expect(screen.getByRole('combobox', { name: 'Capability' })).toHaveValue('') + + await pick(user, 'Capability', 'Calendar') + await pick(user, 'Surface', 'Web chat') + expect(onChange).toHaveBeenLastCalledWith({ surface: 'web_chat' }) + expect(screen.getByRole('combobox', { name: 'Component' })).toHaveValue('Plugins') + expect(screen.getByRole('combobox', { name: 'Capability' })).toHaveValue('Calendar') + }) + + it('offers Other on a single-level field and reports the typed text as the value', async () => { + const user = userEvent.setup() + render( + + + ) + + await openPicker(user, 'Category') + expect(screen.queryByRole('option', { name: 'Other' })).not.toBeInTheDocument() + await user.keyboard('{Escape}') + + await openPicker(user, 'Technique') + expect(optionNames()).toEqual(['Not set', 'Prompt injection', 'Role play', 'Other']) + await user.click(screen.getByRole('option', { name: 'Other' })) + expect(onChange).toHaveBeenLastCalledWith({ technique: '' }) + expect(screen.getByRole('combobox', { name: 'Technique' })).toHaveValue('Other') + + const otherText = screen.getByRole('textbox', { name: 'Other Technique' }) + expect(otherText).toBeRequired() + await user.type(otherText, 'Many-shot') + expect(onChange).toHaveBeenLastCalledWith({ technique: 'Many-shot' }) + expect(otherText).toHaveValue('Many-shot') + expect(screen.getByRole('combobox', { name: 'Technique' })).toHaveValue('Other') + + await pick(user, 'Technique', 'Role play') + expect(onChange).toHaveBeenLastCalledWith({ technique: 'role_play' }) + expect(screen.queryByRole('textbox', { name: 'Other Technique' })).not.toBeInTheDocument() + }) + + it('opens a stored Other value with its text', async () => { + const user = userEvent.setup() + render( + + ) + + expect(screen.getByRole('combobox', { name: 'Technique' })).toHaveValue('Other') + expect(screen.getByRole('textbox', { name: 'Other Technique' })).toHaveValue('Model extraction') + expect(screen.queryByText(UNMATCHED_HINT)).not.toBeInTheDocument() + + await openPicker(user, 'Technique') + await user.click(screen.getByRole('option', { name: 'Other', selected: true })) + await user.tab() + expect(onChange).not.toHaveBeenCalled() + expect(screen.getByRole('combobox', { name: 'Technique' })).toHaveValue('Other') + expect(screen.getByRole('textbox', { name: 'Other Technique' })).toHaveValue('Model extraction') + }) + + it('shows a stored value that is no longer an option and keeps it until changed', async () => { + const user = userEvent.setup() + render( + + + ) + + const severity = screen.getByRole('combobox', { name: 'Severity' }) + expect(severity).toHaveValue('Sev 0') + expect(severity).toHaveAccessibleDescription(UNMATCHED_HINT) + + const subcategory = screen.getByRole('combobox', { name: 'Subcategory' }) + expect(subcategory).toHaveValue('Training data leak') + expect(subcategory).toHaveAccessibleDescription(UNMATCHED_HINT) + await openPicker(user, 'Subcategory') + expect(optionNames()).toEqual(['Training data leak', 'System prompt leak', 'Tool misuse']) + await user.click(screen.getByRole('option', { name: 'Training data leak', selected: true })) + await user.tab() + expect(onChange).not.toHaveBeenCalled() + expect(subcategory).toHaveValue('Training data leak') + + await pick(user, 'Subcategory', 'Tool misuse') + expect(onChange).toHaveBeenLastCalledWith({ subcategory: 'tool_misuse' }) + expect(subcategory).toHaveValue('Tool misuse') + expect(subcategory).not.toHaveAccessibleDescription() + await openPicker(user, 'Subcategory') + expect(optionNames()).toEqual(['System prompt leak', 'Tool misuse']) + }) + + it('offers Not set only on optional fields', async () => { + const user = userEvent.setup() + render( + + + + ) + + await openPicker(user, 'Category') + expect(screen.queryByRole('option', { name: 'Not set' })).not.toBeInTheDocument() + await user.keyboard('{Escape}') + + await openPicker(user, 'Surface') + expect(screen.getByRole('option', { name: 'Not set' })).toBeInTheDocument() + await user.keyboard('{Escape}') + await openPicker(user, 'Component') + expect(screen.queryByRole('option', { name: 'Not set' })).not.toBeInTheDocument() + await user.keyboard('{Escape}') + + await pick(user, 'Technique', 'Not set') + expect(onChange).toHaveBeenLastCalledWith({ technique: '' }) + expect(screen.getByRole('combobox', { name: 'Technique' })).toHaveValue('') + }) + + it('filters options by the typed text', async () => { + const user = userEvent.setup() + render( + + + ) + + const category = screen.getByRole('combobox', { name: 'Category' }) + await user.type(category, 'safety') + expect(optionNames()).toEqual(['Content safety']) + await user.clear(category) + await user.type(category, 'zzz') + expect(screen.getByRole('option', { name: 'No matching options' })).toHaveAttribute('aria-disabled', 'true') + await user.keyboard('{Escape}') + + await user.type(screen.getByRole('combobox', { name: 'Technique' }), 'zzz') + expect(optionNames()).toEqual(['Not set', 'No matching options', 'Other']) + expect(onChange).not.toHaveBeenCalled() + + await user.click(screen.getByRole('option', { name: 'Other' })) + expect(onChange).toHaveBeenLastCalledWith({ technique: '' }) + expect(screen.getByRole('textbox', { name: 'Other Technique' })).toHaveValue('') + }) + + it('reports whether the Other text is blank', async () => { + const user = userEvent.setup() + const onOtherBlankChange = jest.fn() + const onStoredOtherBlankChange = jest.fn() + render( + + + ) + expect(onOtherBlankChange).not.toHaveBeenCalled() + expect(onStoredOtherBlankChange).not.toHaveBeenCalled() + + await pick(user, 'Technique', 'Other') + expect(onOtherBlankChange).toHaveBeenLastCalledWith(true) + const otherText = screen.getByRole('textbox', { name: 'Other Technique' }) + await user.type(otherText, 'Many-shot') + expect(onOtherBlankChange).toHaveBeenLastCalledWith(false) + await user.clear(otherText) + expect(onOtherBlankChange).toHaveBeenLastCalledWith(true) + await pick(user, 'Technique', 'Role play') + expect(onOtherBlankChange).toHaveBeenLastCalledWith(false) + + await user.clear(screen.getByRole('textbox', { name: 'Other Tactic' })) + expect(onStoredOtherBlankChange).toHaveBeenCalledTimes(1) + expect(onStoredOtherBlankChange).toHaveBeenLastCalledWith(true) + }) + + it('disables a lower level with nothing to offer and leaves it unmarked', async () => { + const user = userEvent.setup() + render( + + ) + const detail = screen.getByRole('combobox', { name: 'Impact detail' }) + expect(screen.getByRole('combobox', { name: 'Impact' })).toHaveValue('Service degradation') + expect(detail).toBeDisabled() + expect(detail).not.toBeRequired() + + await pick(user, 'Impact', 'Outage') + expect(onChange).toHaveBeenLastCalledWith({ impact: 'outage' }) + expect(detail).toBeDisabled() + expect(detail).not.toBeRequired() + + await pick(user, 'Impact', 'Data exposure') + expect(detail).toBeEnabled() + expect(detail).toBeRequired() + }) + + it('disables every control while disabled', () => { + render( + + + + + ) + const comboboxes = screen.getAllByRole('combobox') + const textboxes = screen.getAllByRole('textbox') + expect(comboboxes).toHaveLength(3) + expect(textboxes).toHaveLength(3) + for (const control of [...comboboxes, ...textboxes]) expect(control).toBeDisabled() + }) + + it('clears every lower level when an optional cascade is set to Not set', async () => { + const user = userEvent.setup() + render( + + ) + + await pick(user, 'Surface', 'Not set') + expect(onChange).toHaveBeenLastCalledWith({ surface: '', component: '', capability: '' }) + expect(screen.getByRole('combobox', { name: 'Surface' })).toHaveValue('') + expect(screen.getByRole('combobox', { name: 'Component' })).toBeDisabled() + expect(screen.getByRole('combobox', { name: 'Capability' })).toBeDisabled() + }) +}) diff --git a/frontend/src/components/Operations/FindingFieldInput.tsx b/frontend/src/components/Operations/FindingFieldInput.tsx new file mode 100644 index 0000000000..b047763dee --- /dev/null +++ b/frontend/src/components/Operations/FindingFieldInput.tsx @@ -0,0 +1,226 @@ +import { useState, type JSX } from 'react' +import { Combobox, Field, Input, Option, Textarea } from '@fluentui/react-components' + +import type { FindingField, FindingFieldOption } from '@/types' +import { choiceLevels, levelOptions } from '@/utils/findingFields' + +import { useFindingFieldInputStyles } from './FindingFieldInput.styles' + +// Option keys are never blank (the backend rejects them), so a whitespace-only value can't be one. '' is Not set. +const OTHER_VALUE = ' ' +const OTHER_TEXT = 'Other' +const NOT_SET_TEXT = 'Not set' +const NO_MATCH_TEXT = 'No matching options' +const UNMATCHED_HINT = 'Not a current option. It stays unless you change it.' + +interface FindingFieldInputProps { + field: FindingField + /** The whole draft by key. The severity slot's value goes under `severity`. */ + values: Record + disabled?: boolean + /** Reports the keys to set. '' clears a value. */ + onChange: (changes: Record) => void + /** Reports whether the Other text box is showing with blank text. */ + onOtherBlankChange?: (blank: boolean) => void +} + +interface FieldControlProps { + field: FindingField + values: Record + disabled: boolean + onChange: (changes: Record) => void +} + +interface ChoiceFieldInputProps extends FieldControlProps { + onOtherBlankChange?: (blank: boolean) => void +} + +interface PickerOption { + value: string + text: string +} + +interface LevelState { + options: PickerOption[] + selected: string + hint?: string + /** Required only where the path offers options, as the backend checks a required cascade. */ + required: boolean + unavailable: boolean +} + +function storedValue(values: Record, key: string): string { + const value = values[key] ?? '' + return value.trim() === '' ? '' : value +} + +function isOption(options: FindingFieldOption[], value: string): boolean { + return options.some((option: FindingFieldOption) => option.key === value) +} + +function levelHint(field: FindingField, depth: number, unmatched: boolean): string | undefined { + if (unmatched) return UNMATCHED_HINT + return depth === 0 ? field.hint ?? undefined : undefined +} + +/** + * What one dropdown of a choice field offers and shows. A stored value that isn't an option under its path is + * listed as typed and stays selected, unless the field is showing Other text. + */ +function levelState(field: FindingField, values: Record, depth: number, otherChosen: boolean): LevelState { + const levels = choiceLevels(field) + const fieldOptions = levelOptions(field, values, depth) + const options = fieldOptions.map((option: FindingFieldOption) => ({ value: option.key, text: option.label })) + const value = storedValue(values, levels[depth].key) + const unmatched = !otherChosen && value !== '' && !isOption(fieldOptions, value) + const parentBlank = depth > 0 && storedValue(values, levels[depth - 1].key) === '' + return { + options: unmatched ? [{ value, text: value }, ...options] : options, + selected: otherChosen ? OTHER_VALUE : value, + hint: levelHint(field, depth, unmatched), + required: field.required && fieldOptions.length > 0, + unavailable: depth > 0 && (parentBlank || (options.length === 0 && !unmatched)), + } +} + +/** The chosen value, plus a clear for each lower level whose value isn't an option under the new path. */ +function cascadeChanges( + field: FindingField, + values: Record, + depth: number, + value: string, +): Record { + const levels = choiceLevels(field) + const changes: Record = { [levels[depth].key]: value } + const next = { ...values, ...changes } + for (const [lower, level] of levels.entries()) { + if (lower <= depth) continue + const current = storedValue(next, level.key) + if (current === '' || isOption(levelOptions(field, next, lower), current)) continue + changes[level.key] = '' + next[level.key] = '' + } + return changes +} + +function selectedText(options: PickerOption[], selected: string): string { + if (selected === OTHER_VALUE) return OTHER_TEXT + return options.find((option: PickerOption) => option.value === selected)?.text ?? '' +} + +interface ChoiceLevelPickerProps { + label: string + hint?: string + required: boolean + disabled: boolean + /** The options the typed text filters. Not set and Other always show. */ + options: PickerOption[] + selected: string + offerNotSet: boolean + offerOther: boolean + onSelect: (optionValue: string) => void +} + +function ChoiceLevelPicker({ + label, hint, required, disabled, options, selected, offerNotSet, offerOther, onSelect, +}: ChoiceLevelPickerProps): JSX.Element { + const styles = useFindingFieldInputStyles() + const [open, setOpen] = useState(false) + const [search, setSearch] = useState('') + const matches = options.filter((option: PickerOption) => option.text.toLowerCase().includes(search.toLowerCase())) + + return ( + + { setOpen(data.open); setSearch('') }} + onChange={(event) => { setSearch(event.target.value); setOpen(true) }} + onOptionSelect={(_, data) => { + // Typed text that no longer matches the selection clears Fluent's selection without a value. Clicking the + // selected option, or blurring on its text, selects it again, so only a different option is a change. + if (data.optionValue === undefined) return + setOpen(false) + setSearch('') + if (data.optionValue !== selected) onSelect(data.optionValue) + }} + listbox={{ className: styles.listbox }}> + {offerNotSet && } + {matches.map((option: PickerOption) => )} + {matches.length === 0 && } + {offerOther && } + + + ) +} + +function ChoiceFieldInput({ + field, values, disabled, onChange, onOtherBlankChange, +}: ChoiceFieldInputProps): JSX.Element { + const styles = useFindingFieldInputStyles() + const levels = choiceLevels(field) + const allowOther = field.allow_other && levels.length === 1 + const [otherChosen, setOtherChosen] = useState(() => { + const value = storedValue(values, field.key) + return allowOther && value !== '' && !isOption(field.options, value) + }) + + const select = (depth: number, optionValue: string): void => { + if (allowOther) { + setOtherChosen(optionValue === OTHER_VALUE) + onOtherBlankChange?.(optionValue === OTHER_VALUE) + } + onChange(cascadeChanges(field, values, depth, optionValue === OTHER_VALUE ? '' : optionValue)) + } + + return ( +
+ {levels.map((level: { key: string; label: string }, depth: number) => { + const state = levelState(field, values, depth, depth === 0 && otherChosen) + return ( + { select(depth, optionValue) }} /> + ) + })} + {otherChosen && + { + onOtherBlankChange?.(data.value.trim() === '') + onChange({ [field.key]: data.value }) + }} /> + } +
+ ) +} + +function TextFieldInput({ field, values, disabled, onChange }: FieldControlProps): JSX.Element { + const value = values[field.key] ?? '' + const report = (text: string): void => { onChange({ [field.key]: text }) } + return ( + + {field.type === 'multiline' + ?