FIX: make value kind labels exhaustive

pyrit/analytics/attack_result_analytics.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import threading
-from typing import TYPE_CHECKING, ClassVar, Self, TypeVar
+from typing import TYPE_CHECKING, ClassVar, Self, TypeVar, assert_never
 from weakref import WeakKeyDictionary
 
 from pydantic import BaseModel
@@ -305,20 +305,33 @@ def _option(*, raw: RawAnalyticsOption, dimension: AttackAnalyticsDimension) ->
             label = "Not recorded"
         elif raw.key.kind is AttackAnalyticsValueKind.NO_CONVERTERS:
             label = "No converters"
+        elif raw.key.kind is AttackAnalyticsValueKind.VALUE:
+            label = AttackResultAnalytics._value_label(raw=raw, dimension=dimension)
+        elif raw.key.kind is AttackAnalyticsValueKind.LEGACY_OPERATION_NAME:
+            label = f"{AttackResultAnalytics._value_label(raw=raw, dimension=dimension)} (legacy name)"
         else:
-            label = raw.label if raw.label is not None else raw.key.value or ""
-            if not label.strip():
-                label = "(Blank)"
-            elif (
-                dimension.name in {AttackAnalyticsDimensionName.OBJECTIVE_TARGET, AttackAnalyticsDimensionName.SCENARIO}
-                and raw.key.value is not None
-                and label != raw.key.value
-            ):
-                label = f"{label} ({raw.key.value[-8:]})"
-            if raw.key.kind is AttackAnalyticsValueKind.LEGACY_OPERATION_NAME:
-                label = f"{label} (legacy name)"
+            assert_never(raw.key.kind)
         return AttackAnalyticsOption(key=raw.key, label=label)
 
+    @staticmethod
+    def _value_label(*, raw: RawAnalyticsOption, dimension: AttackAnalyticsDimension) -> str:
+        """
+        Format the text-bearing value kinds without changing their stored keys.
+
+        Returns:
+            str: The SDK display label for a text-bearing value kind.
+        """
+        label = raw.label if raw.label is not None else raw.key.value or ""
+        if not label.strip():
+            return "(Blank)"
+        if (
+            dimension.name in {AttackAnalyticsDimensionName.OBJECTIVE_TARGET, AttackAnalyticsDimensionName.SCENARIO}
+            and raw.key.value is not None
+            and label != raw.key.value
+        ):
+            return f"{label} ({raw.key.value[-8:]})"
+        return label
+
     @staticmethod
     def _drilldown_unavailable_reason(query: AttackAnalyticsQuery) -> str | None:
         """

tests/unit/analytics/test_attack_result_analytics.py

@@ -4,7 +4,12 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
+import subprocess
+import sys
 from datetime import UTC, datetime, timedelta, timezone
+from pathlib import Path
+from textwrap import dedent
 from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -28,6 +33,7 @@
     AttackAnalyticsQuery,
     AttackAnalyticsResultsQuery,
     AttackAnalyticsValue,
+    AttackAnalyticsValueKind,
     AttackIdentifier,
     AttackOutcome,
     AttackResultMetadata,
@@ -327,6 +333,7 @@ async def test_matrix_empty_cells_and_additional_predicates_are_exact(
         ("value", " \t", None, "label", "(Blank)"),
         ("value", "Unknown", "Unknown", "operation", "Unknown"),
         ("value", "Not recorded", "Not recorded", "operation", "Not recorded"),
+        ("legacy_operation_name", "Legacy", "Legacy", "operation", "Legacy (legacy name)"),
         ("value", "content-independent-eval-hash", "model", "objective_target", "model (val-hash)"),
         ("value", "scenario-12345678", "run", "scenario", "run (12345678)"),
         ("value", "scenario-12345678", None, "scenario", "scenario-12345678"),
@@ -344,6 +351,71 @@ def test_display_labels_do_not_change_typed_keys(
     assert option.key == raw.key
 
 
+@pytest.mark.parametrize("extend_enum", [False, True], ids=["current-enum", "extended-enum"])
+def test_option_dispatch_type_exhaustiveness(tmp_path: Path, extend_enum: bool) -> None:
+    enum_source = inspect.getsource(AttackAnalyticsValueKind)
+    if extend_enum:
+        enum_source += '\n    FUTURE_KIND = "future_kind"\n'
+    probe = tmp_path / "dispatch.py"
+    probe.write_text(
+        "\n".join(
+            [
+                "from dataclasses import dataclass",
+                "from enum import Enum",
+                "from typing import assert_never",
+                "from pyrit.models import AttackAnalyticsDimension, AttackAnalyticsDimensionName",
+                enum_source,
+                dedent(
+                    """
+                    @dataclass
+                    class AttackAnalyticsValue:
+                        kind: AttackAnalyticsValueKind
+                        value: str | None
+
+                    @dataclass
+                    class RawAnalyticsOption:
+                        key: AttackAnalyticsValue
+                        label: str | None
+
+                    @dataclass
+                    class AttackAnalyticsOption:
+                        key: AttackAnalyticsValue
+                        label: str
+                    """
+                ),
+                "class AttackResultAnalytics:",
+                inspect.getsource(AttackResultAnalytics._option),
+                inspect.getsource(AttackResultAnalytics._value_label),
+            ]
+        ),
+        encoding="utf-8",
+    )
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "ty",
+            "check",
+            "--project",
+            str(Path(__file__).resolve().parents[3]),
+            "--python-version",
+            "3.11",
+            str(probe),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=30,
+        check=False,
+    )
+    output = result.stdout + result.stderr
+    assert result.returncode == (1 if extend_enum else 0), output
+    if extend_enum:
+        assert output.count("error[") == 1, output
+        assert "type-assertion-failure" in output, output
+        assert "assert_never(raw.key.kind)" in output, output
+        assert "FUTURE_KIND" in output, output
+
+
 async def test_operation_option_labels(analytics: AttackResultAnalytics, sqlite_instance: SQLiteMemory) -> None:
     saved, removed = [await sqlite_instance.add_operation_async(Operation(name=name)) for name in ("Saved", "Gone")]
     await add_legacy_attack_results_async(