diff --git a/doc/code/framework.md b/doc/code/framework.md index b75ecf231b..15a313391a 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -182,6 +182,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - `core` stays deliberately small so a default run doesn't print 200 techniques or take forever; the wider catalog lives in `extra` and is selected on demand. Users pick subsets by passing initializer tags (e.g. `core`, `extra`, `all`) or writing their own initializer, so different runs — including from the CLI — can register different technique sets without changing the catalog. - A technique tied to one scenario is fine; if it's pinned and non-reusable it can stay local to that scenario, but if another scenario could reuse it, promote it to a catalog module and tag it. - Tags describe a technique (behavioral tags like `single_turn`/`multi_turn`, owner tags like `airt`); they don't decide what a scenario runs. There is deliberately **no global `default` tag** — a default is scenario-relative, declared per scenario via `build_technique_class_from_factories` (the `factories` list is the pool, catalog tags become named aggregate presets, and `default_tags` / `default_names` set what runs when nothing is chosen). +- Factories opt into additive request-converter composition with `supports_additional_request_converters=True`. This is a semantic capability, not just constructor-signature detection; the factory validates that opted-in attacks accept `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. **Framework Plans**: diff --git a/pyrit/backend/models/scenarios.py b/pyrit/backend/models/scenarios.py index 56858f80b1..820133367d 100644 --- a/pyrit/backend/models/scenarios.py +++ b/pyrit/backend/models/scenarios.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field from pyrit.backend.models.common import PaginationInfo -from pyrit.models.catalog.scenario import RegisteredScenario, ScenarioRunSummary +from pyrit.models.catalog.scenario import RegisteredScenario, ScenarioRunListItem __all__ = [ "ListRegisteredScenariosResponse", @@ -31,4 +31,4 @@ class ListRegisteredScenariosResponse(BaseModel): class ScenarioRunListResponse(BaseModel): """Response for listing scenario runs.""" - items: list[ScenarioRunSummary] = Field(..., description="List of scenario runs") + items: list[ScenarioRunListItem] = Field(..., description="List of scenario runs") diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index fa3a5635bb..a6e9ea5e00 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -13,6 +13,7 @@ """ from fastapi import APIRouter, HTTPException, Query, status +from starlette.concurrency import run_in_threadpool from pyrit.backend.models.common import ProblemDetail from pyrit.backend.models.scenarios import ( @@ -25,8 +26,11 @@ from pyrit.models.catalog.scenario import ( RegisteredScenario, RunScenarioRequest, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, ScenarioRunSummary, ) +from pyrit.models.scenario_progress import ScenarioRunProgress router = APIRouter(prefix="/scenarios", tags=["scenarios"]) @@ -86,6 +90,45 @@ async def get_scenario(scenario_name: str) -> RegisteredScenario: # pyrit-async return scenario +@router.post( + "/catalog/{scenario_name}/estimate", + response_model=ScenarioRunSizeEstimate, + responses={ + 400: {"model": ProblemDetail, "description": "Invalid estimate configuration"}, + 404: {"model": ProblemDetail, "description": "Scenario not found"}, + }, +) +async def estimate_scenario_run_size( # pyrit-async-suffix-exempt + *, + scenario_name: str, + request: ScenarioRunSizeEstimateRequest, +) -> ScenarioRunSizeEstimate: + """ + Estimate a configured scenario without creating or persisting a run. + + Args: + scenario_name: Registry name of the scenario. + request: Techniques, datasets, baseline choice, and scenario parameters to preview. + + Returns: + ScenarioRunSizeEstimate: Structured request-specific planned-unit estimate. + """ + service = get_scenario_service() + try: + estimate = await service.estimate_scenario_run_size_async( + scenario_name=scenario_name, + request=request, + ) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None + if estimate is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Scenario '{scenario_name}' not found", + ) + return estimate + + # ============================================================================ # Scenario Runs # ============================================================================ @@ -122,7 +165,9 @@ async def start_scenario_run(request: RunScenarioRequest) -> ScenarioRunSummary: "/runs", response_model=ScenarioRunListResponse, ) -async def list_scenario_runs(limit: int = Query(100, ge=1)) -> ScenarioRunListResponse: # pyrit-async-suffix-exempt +async def list_scenario_runs( + limit: int = Query(100, ge=1, le=100), +) -> ScenarioRunListResponse: # pyrit-async-suffix-exempt """ List tracked scenario runs (most recent first). @@ -133,7 +178,7 @@ async def list_scenario_runs(limit: int = Query(100, ge=1)) -> ScenarioRunListRe ScenarioRunListResponse: Runs, most recent first. """ service = get_scenario_run_service() - return service.list_runs(limit=limit) + return await run_in_threadpool(service.list_runs, limit=limit) @router.get( @@ -154,7 +199,12 @@ async def get_scenario_run(scenario_result_id: str) -> ScenarioRunSummary: # py ScenarioRunSummary: Current run status (and result if completed). """ service = get_scenario_run_service() - run = service.get_run(scenario_result_id=scenario_result_id) + active_snapshot = service.snapshot_active_run(scenario_result_id=scenario_result_id) + run = await run_in_threadpool( + service.get_run_from_storage, + scenario_result_id=scenario_result_id, + active_error=active_snapshot.error, + ) if run is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -163,6 +213,46 @@ async def get_scenario_run(scenario_result_id: str) -> ScenarioRunSummary: # py return run +@router.get( + "/runs/{scenario_result_id}/progress", + response_model=ScenarioRunProgress, + responses={ + 400: {"model": ProblemDetail, "description": "Invalid progress cursor"}, + 404: {"model": ProblemDetail, "description": "Run not found"}, + }, +) +async def get_scenario_run_progress( # pyrit-async-suffix-exempt + *, + scenario_result_id: str, + since: str | None = Query(None, description="Opaque ascending progress cursor"), + limit: int = Query(100, ge=1, le=500), +) -> ScenarioRunProgress: + """ + Get a compact, refresh-safe page of scenario progress deltas. + + Returns: + ScenarioRunProgress: The run plan and ascending result deltas. + """ + service = get_scenario_run_service() + active_snapshot = service.snapshot_active_run(scenario_result_id=scenario_result_id) + try: + progress = await run_in_threadpool( + service.get_run_progress_from_storage, + scenario_result_id=scenario_result_id, + since=since, + limit=limit, + active_group_ids=active_snapshot.active_group_ids, + ) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from None + if progress is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Scenario run '{scenario_result_id}' not found", + ) + return progress + + @router.post( "/runs/{scenario_result_id}/cancel", response_model=ScenarioRunSummary, diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index 134d7a9209..09d07d211d 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -58,7 +58,7 @@ from pyrit.backend.models.common import PaginationInfo from pyrit.backend.services.converter_service import get_converter_service from pyrit.backend.services.target_service import get_target_service -from pyrit.memory import AttackResultsKeysetCursor, CentralMemory, data_serializer_factory +from pyrit.memory import AttackResultKeysetCursor, CentralMemory, data_serializer_factory from pyrit.models import ( AtomicAttackIdentifier, AttackIdentifier, @@ -182,7 +182,7 @@ async def list_attacks_async( page_results = list(results[:limit]) next_cursor = ( self._encode_attack_cursor( - cursor=AttackResultsKeysetCursor.from_attack_result(page_results[-1]), + cursor=AttackResultKeysetCursor.from_attack_result(page_results[-1]), fingerprint=filter_fingerprint, ) if has_next_page and page_results @@ -919,7 +919,7 @@ def _norm_labels( return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16] @staticmethod - def _encode_attack_cursor(*, cursor: AttackResultsKeysetCursor, fingerprint: str) -> str: + def _encode_attack_cursor(*, cursor: AttackResultKeysetCursor, fingerprint: str) -> str: """ Encode a keyset anchor and its filter fingerprint into an opaque pagination cursor. @@ -940,7 +940,7 @@ def _encode_attack_cursor(*, cursor: AttackResultsKeysetCursor, fingerprint: str return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") @staticmethod - def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResultsKeysetCursor | None: + def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResultKeysetCursor | None: """ Decode the opaque list-attacks cursor into a keyset (seek) anchor. @@ -952,7 +952,7 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu raising or seeking within the wrong result set. Returns: - The decoded ``AttackResultsKeysetCursor``, or ``None`` to start at the first page. + The decoded ``AttackResultKeysetCursor``, or ``None`` to start at the first page. """ if not cursor: return None @@ -985,7 +985,7 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu # A crafted cursor near datetime's min/max with a large UTC offset overflows the # representable range when shifted to UTC; treat it as malformed and restart at page one. return None - return AttackResultsKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) + return AttackResultKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) # ======================================================================== # Private Helper Methods - Duplicate / Branch diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py new file mode 100644 index 0000000000..a0310bf595 --- /dev/null +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Shared scenario launch and estimate configuration resolution.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from pyrit.registry import ConverterRegistry, ScenarioRegistry, TargetRegistry + +if TYPE_CHECKING: + from pyrit.converter import Converter + from pyrit.prompt_target import PromptTarget + from pyrit.scenario import Scenario + +_CONVERTER_MODIFIER_PREFIX = "converter." + + +class ScenarioConfigurationResolver: + """Resolve registry-backed scenario inputs for launch and estimation.""" + + @staticmethod + def resolve_scenario_class(*, scenario_name: str) -> type[Scenario]: + """ + Resolve a registered scenario class. + + Returns: + type[Scenario]: The registered scenario class. + + Raises: + ValueError: If the scenario name is not registered. + """ + try: + return ScenarioRegistry.get_registry_singleton().get_class(scenario_name) + except KeyError as exc: + raise ValueError(str(exc)) from None + + @staticmethod + def resolve_target(*, target_name: str) -> PromptTarget: + """ + Resolve a registered target instance. + + Returns: + PromptTarget: The registered target. + + Raises: + ValueError: If the target name is not registered. + """ + instances = TargetRegistry.get_registry_singleton().instances + objective_target = instances.get(target_name) + if objective_target is not None: + return objective_target + + available_names = instances.get_names() + if not available_names: + raise ValueError( + f"Target '{target_name}' not found. The target registry is empty. " + "Make sure to include an initializer that registers targets " + "(e.g., initializers: ['target'])." + ) + raise ValueError( + f"Target '{target_name}' not found in registry. Available targets: {', '.join(available_names)}" + ) + + @classmethod + def resolve_configuration( + cls, + *, + scenario_name: str, + scenario_class: type[Scenario], + objective_target: Any | None = None, + techniques: list[str] | None = None, + dataset_names: list[str] | None = None, + max_dataset_size: int | None = None, + dataset_filters: dict[str, list[str]] | None = None, + include_baseline: bool | None = None, + max_concurrency: int | None = None, + max_retries: int | None = None, + memory_labels: dict[str, str] | None = None, + ) -> dict[str, Any]: + """ + Resolve shared launch and estimate fields into scenario parameters. + + Returns: + dict[str, Any]: Values accepted by ``Scenario.set_params_from_args``. + + Raises: + ValueError: If techniques or dataset overrides are invalid. + """ + resolved: dict[str, Any] = {} + if objective_target is not None: + resolved["objective_target"] = objective_target + if max_concurrency is not None: + resolved["max_concurrency"] = max_concurrency + if max_retries is not None: + resolved["max_retries"] = max_retries + if include_baseline is not None: + resolved["include_baseline"] = include_baseline + if memory_labels: + resolved["memory_labels"] = memory_labels + + filters = dataset_filters or {} + needs_introspection = bool(techniques) or bool(dataset_names) or max_dataset_size is not None or bool(filters) + if not needs_introspection: + return resolved + + try: + introspection_instance = scenario_class() # type: ignore[ty:missing-argument] + except Exception as exc: + raise ValueError( + f"Cannot resolve runtime configuration for scenario '{scenario_name}': " + f"scenario class is not instantiable without arguments ({exc})." + ) from exc + + if techniques: + technique_enums, technique_converters = cls.resolve_techniques_and_converters( + tokens=techniques, + technique_class=introspection_instance._technique_class, + scenario_name=scenario_name, + ) + resolved["scenario_techniques"] = technique_enums + if technique_converters: + resolved["technique_converters"] = technique_converters + + if dataset_names or max_dataset_size is not None or filters: + default_config = introspection_instance._default_dataset_config + if dataset_names: + default_config_class = type(default_config) + try: + resolved["dataset_config"] = default_config_class( + dataset_names=dataset_names, + max_dataset_size=max_dataset_size, + filters=filters or None, + ) + except TypeError as exc: + raise ValueError( + f"Scenario '{scenario_name}' does not support overriding dataset names through " + f"its {default_config_class.__name__} configuration: {exc}" + ) from exc + else: + if max_dataset_size is not None: + default_config.max_dataset_size = max_dataset_size + if filters: + default_config.update_filters(filters=filters) + resolved["dataset_config"] = default_config + + return resolved + + @classmethod + def resolve_techniques_and_converters( + cls, + *, + tokens: list[str], + technique_class: type[Any], + scenario_name: str, + ) -> tuple[list[Any], dict[str, list[Converter]]]: + """ + Resolve technique tokens and their additive converter modifiers. + + Returns: + tuple[list[Any], dict[str, list[Converter]]]: Selected enum members and + converters keyed by concrete technique name. + + Raises: + ValueError: If a technique or converter modifier is invalid. + """ + technique_enums: list[Any] = [] + technique_converters: dict[str, list[Converter]] = {} + for token in tokens: + base_name, _, remainder = token.partition(":") + modifiers = [modifier for modifier in remainder.split(":") if modifier] if remainder else [] + try: + technique_enum = technique_class(base_name) + except ValueError: + available_techniques = [technique.value for technique in technique_class] + raise ValueError( + f"Technique '{base_name}' not found for scenario '{scenario_name}'. " + f"Available: {', '.join(available_techniques)}" + ) from None + technique_enums.append(technique_enum) + + converters = cls._resolve_converter_modifiers(modifiers=modifiers, token=token) + for concrete in technique_class.expand({technique_enum}) if converters else (): + technique_converters.setdefault(concrete.value, []).extend(converters) + + return technique_enums, technique_converters + + @staticmethod + def _resolve_converter_modifiers(*, modifiers: list[str], token: str) -> list[Converter]: + """ + Resolve converter modifiers from one technique token. + + Returns: + list[Converter]: Registered converter instances in token order. + + Raises: + ValueError: If a modifier is malformed or references an unknown converter. + """ + if not modifiers: + return [] + + instances = ConverterRegistry.get_registry_singleton().instances + converters: list[Converter] = [] + for modifier in modifiers: + if not modifier.startswith(_CONVERTER_MODIFIER_PREFIX): + raise ValueError( + f"Unknown technique modifier '{modifier}' in '{token}'. " + f"Supported modifiers must use the '{_CONVERTER_MODIFIER_PREFIX}' prefix " + f"(e.g. '{_CONVERTER_MODIFIER_PREFIX}translation_spanish')." + ) + converter_name = modifier[len(_CONVERTER_MODIFIER_PREFIX) :] + converter = instances.get(converter_name) + if converter is None: + available = instances.get_names() + available_text = ", ".join(available) if available else "(none registered)" + raise ValueError( + f"Converter '{converter_name}' in '{token}' is not a registered converter " + f"instance. Available converters: {available_text}" + ) + converters.append(converter) + return converters diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 7ba66d0f43..c635dab06f 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -9,39 +9,51 @@ """ import asyncio +import base64 import contextlib +import json import logging +import uuid +from collections.abc import Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from datetime import datetime, timezone +from typing import Any from pyrit.backend.models.scenarios import ScenarioRunListResponse -from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, ScenarioResult, ScenarioRunState +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver +from pyrit.common.utils import to_sha256 +from pyrit.memory import AttackResultKeysetCursor, CentralMemory +from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + ComponentIdentifier, + ScenarioAttackResultDelta, + ScenarioProgressHeader, + ScenarioProgressResult, + ScenarioResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunProgress, + ScenarioRunState, + config_hash, +) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, ) -from pyrit.registry import ( - ConverterRegistry, - InitializerRegistry, - ScenarioRegistry, - TargetRegistry, -) +from pyrit.registry import InitializerRegistry, ScenarioRegistry from pyrit.scenario import Scenario -from pyrit.scenario.core import DatasetAttackConfiguration - -if TYPE_CHECKING: - from pyrit.converter import Converter - from pyrit.prompt_target import PromptTarget logger = logging.getLogger(__name__) _DEFAULT_MAX_CONCURRENT_RUNS = 3 -_CONVERTER_MODIFIER_PREFIX = "converter." - @dataclass class _ActiveTask: @@ -53,6 +65,95 @@ class _ActiveTask: error: str | None = None +@dataclass(frozen=True, slots=True) +class _ActiveRunSnapshot: + """Event-loop-owned state copied before database work moves to a worker thread.""" + + error: str | None = None + active_group_ids: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class _ResultUnitIdentity: + """Stable identity of one planned scenario execution unit.""" + + atomic_group_id: str + seed_group_id: str + + +@dataclass(frozen=True, slots=True) +class _ScenarioPlanLookup: + """Pre-indexed run-plan data used while mapping many attack results.""" + + groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] + groups_by_name: dict[str, tuple[ScenarioRunPlanAtomicGroup, ...]] + seed_ids_by_group_and_objective: dict[tuple[str, str], tuple[str, ...]] + planned_units: frozenset[_ResultUnitIdentity] + + @classmethod + def from_plan(cls, *, plan: ScenarioRunPlan | None) -> "_ScenarioPlanLookup": + """ + Build constant-time lookup tables for one run plan. + + Returns: + _ScenarioPlanLookup: Indexed plan data. + """ + if plan is None: + return cls( + groups_by_identity={}, + groups_by_name={}, + seed_ids_by_group_and_objective={}, + planned_units=frozenset(), + ) + + groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] = {} + grouped_by_name: dict[str, list[ScenarioRunPlanAtomicGroup]] = {} + seeds_by_id = {seed.id: seed for seed in plan.seed_groups} + seed_ids_by_group_and_objective: dict[tuple[str, str], tuple[str, ...]] = {} + planned_units: set[_ResultUnitIdentity] = set() + for group in plan.atomic_groups: + groups_by_identity[(group.atomic_attack_name, group.technique_eval_hash)] = group + grouped_by_name.setdefault(group.atomic_attack_name, []).append(group) + seed_ids_by_objective: dict[str, list[str]] = {} + for seed_id in group.seed_group_ids: + seed = seeds_by_id[seed_id] + seed_ids_by_objective.setdefault(seed.objective_sha256, []).append(seed_id) + seed_ids_by_group_and_objective.update( + { + (group.id, objective_sha256): tuple(seed_ids) + for objective_sha256, seed_ids in seed_ids_by_objective.items() + } + ) + planned_units.update( + _ResultUnitIdentity(atomic_group_id=group.id, seed_group_id=seed_group_id) + for seed_group_id in group.seed_group_ids + ) + + return cls( + groups_by_identity=groups_by_identity, + groups_by_name={name: tuple(groups) for name, groups in grouped_by_name.items()}, + seed_ids_by_group_and_objective=seed_ids_by_group_and_objective, + planned_units=frozenset(planned_units), + ) + + def resolve_group( + self, + *, + atomic_attack_name: str, + technique_eval_hash: str | None, + ) -> ScenarioRunPlanAtomicGroup | None: + """ + Resolve one planned group from persisted attribution. + + Returns: + ScenarioRunPlanAtomicGroup | None: The uniquely matching group. + """ + if technique_eval_hash is not None: + return self.groups_by_identity.get((atomic_attack_name, technique_eval_hash)) + matching_groups = self.groups_by_name.get(atomic_attack_name, ()) + return matching_groups[0] if len(matching_groups) == 1 else None + + class ScenarioRunService: """ Service for managing scenario run lifecycle. @@ -67,6 +168,7 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) - self._memory = CentralMemory.get_memory_instance() self._active_tasks: dict[str, _ActiveTask] = {} self._run_semaphore = asyncio.Semaphore(max_concurrent_runs) + self._configuration_resolver = ScenarioConfigurationResolver() async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary: """ @@ -97,11 +199,21 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu # Perform all initialization eagerly — errors propagate to caller try: - scenario_class = self._resolve_scenario_class(request=request) + scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name) await self._run_initializers_async(request=request) - objective_target = self._resolve_target(request=request) - init_kwargs = self._build_init_kwargs( - request=request, scenario_class=scenario_class, objective_target=objective_target + objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name) + init_kwargs = self._configuration_resolver.resolve_configuration( + scenario_name=request.scenario_name, + scenario_class=scenario_class, + objective_target=objective_target, + techniques=request.techniques, + dataset_names=request.dataset_names, + max_dataset_size=request.max_dataset_size, + dataset_filters=request.dataset_filters, + include_baseline=request.include_baseline, + max_concurrency=request.max_concurrency, + max_retries=request.max_retries, + memory_labels=request.labels, ) scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs) except Exception: @@ -121,7 +233,7 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id)) active.task = task - response = self._build_response(scenario_result_id=scenario_result_id) + response = self.get_run(scenario_result_id=scenario_result_id) if response is None: raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.") return response @@ -136,7 +248,26 @@ def get_run(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: Returns: ScenarioRunSummary if found, None otherwise. """ - return self._build_response(scenario_result_id=scenario_result_id) + snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id) + return self.get_run_from_storage(scenario_result_id=scenario_result_id, active_error=snapshot.error) + + def get_run_from_storage( + self, + *, + scenario_result_id: str, + active_error: str | None, + ) -> ScenarioRunSummary | None: + """ + Build a run summary using database state plus an event-loop snapshot. + + Args: + scenario_result_id: The scenario result ID. + active_error: Error copied from the active asyncio task, if any. + + Returns: + ScenarioRunSummary | None: The run summary when found. + """ + return self._build_response(scenario_result_id=scenario_result_id, active_error=active_error) def list_runs(self, *, limit: int = 100) -> ScenarioRunListResponse: """ @@ -148,12 +279,49 @@ def list_runs(self, *, limit: int = 100) -> ScenarioRunListResponse: Returns: ScenarioRunListResponse with runs. """ - # This is expensive, and we don't need all the data. At some point - # we may want to add a lightweight "list" query to the DB layer that only - results = self._memory.get_scenario_results(limit=limit) - items = [self._build_response_from_db(scenario_result=sr) for sr in results] + results = self._memory.get_scenario_result_headers(limit=limit) + items = [self._build_list_response_from_header(scenario_result=result) for result in results] return ScenarioRunListResponse(items=items) + def _build_list_response_from_header(self, *, scenario_result: ScenarioResult) -> ScenarioRunListItem: + """ + Build a bounded run-history item without hydrating attack results. + + Returns: + ScenarioRunListItem: Lightweight run metadata. + """ + status = scenario_result.scenario_run_state + terminal = status in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ) + plan = self._load_run_plan(scenario_result=scenario_result) + total_attacks = sum(len(group.seed_group_ids) for group in plan.atomic_groups) if plan is not None else None + techniques_used = ( + list(dict.fromkeys(group.display_group for group in plan.atomic_groups)) if plan is not None else [] + ) + updated_at = ( + scenario_result.completion_time + if terminal and scenario_result.completion_time is not None + else scenario_result.creation_time + ) + return ScenarioRunListItem( + scenario_result_id=str(scenario_result.id), + scenario_name=scenario_result.scenario_name, + scenario_registry_name=plan.scenario_registry_name if plan else None, + scenario_version=scenario_result.scenario_version, + status=status, + created_at=scenario_result.creation_time, + updated_at=updated_at, + error=scenario_result.error_message, + error_type=scenario_result.error_type, + techniques_used=techniques_used, + total_attacks=total_attacks, + labels=scenario_result.labels, + completed_at=scenario_result.completion_time if terminal else None, + ) + async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: """ Cancel a running scenario. @@ -193,26 +361,7 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma error_type="CancelledError", ) - return self._build_response(scenario_result_id=scenario_result_id) - - def _resolve_scenario_class(self, *, request: RunScenarioRequest) -> type[Scenario]: - """ - Validate and resolve the scenario class from the registry. - - Args: - request: The run request containing the scenario name. - - Returns: - The scenario class. - - Raises: - ValueError: If the scenario name is not found in the registry. - """ - scenario_registry = ScenarioRegistry.get_registry_singleton() - try: - return scenario_registry.get_class(request.scenario_name) - except KeyError as e: - raise ValueError(str(e)) from None + return self.get_run(scenario_result_id=scenario_result_id) async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: """ @@ -238,251 +387,6 @@ async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: raise ValueError(f"Initializer not found: {e}") from None await instance.initialize_async() - def _resolve_target(self, *, request: RunScenarioRequest) -> "PromptTarget": - """ - Resolve the objective target from the target registry. - - Args: - request: The run request containing the target name. - - Returns: - The resolved PromptTarget instance. - - Raises: - ValueError: If the target is not found in the registry. - """ - target_registry = TargetRegistry.get_registry_singleton() - objective_target = target_registry.instances.get(request.target_name) - if objective_target is None: - available_names = target_registry.instances.get_names() - if not available_names: - raise ValueError( - f"Target '{request.target_name}' not found. The target registry is empty. " - "Make sure to include an initializer that registers targets " - "(e.g., initializers: ['target'])." - ) - raise ValueError( - f"Target '{request.target_name}' not found in registry. Available targets: {', '.join(available_names)}" - ) - return objective_target - - def _build_init_kwargs( - self, *, request: RunScenarioRequest, scenario_class: type[Scenario], objective_target: Any - ) -> dict[str, Any]: - """ - Build the kwargs dict for scenario.initialize_async. - - Resolves techniques and dataset configuration from the request. - - Dataset configuration is built so that the scenario's default - ``DatasetAttackConfiguration`` *subclass* (e.g. ``EncodingDatasetConfiguration``) - is preserved when the caller overrides ``dataset_names`` or - ``max_dataset_size``. Subclasses commonly override - ``_build_attack_groups()`` to shape seeds into scenario-appropriate - ``AttackSeedGroup`` objects. - - Args: - request: The run request. - scenario_class: The resolved scenario class. - objective_target: The resolved target instance. - - Returns: - Dict of kwargs to pass to scenario.initialize_async. - - Raises: - ValueError: If a technique name is invalid for the scenario, or the - scenario class cannot be instantiated with no arguments when - introspection is required to resolve techniques or dataset - configuration. - """ - init_kwargs: dict[str, Any] = { - "objective_target": objective_target, - "max_concurrency": request.max_concurrency, - "max_retries": request.max_retries, - } - - if request.labels: - init_kwargs["memory_labels"] = request.labels - - # The request model has already validated the filter keys and coerced values into - # lists, so the service can consume them directly. - dataset_filters = request.dataset_filters or {} - - # Resolve techniques and dataset config from a temporary instance of the - # scenario. The downstream _initialize_scenario_async builds its own - # instance (so scenario_result_id can be passed), so this is a cheap - # throwaway used only for introspection. Introspection is required - # whenever the caller wants to override techniques, dataset names, the - # sample cap, or dataset filters, because each of those needs the - # scenario's own technique enum or dataset-config subclass to be resolved - # correctly. - needs_introspection = ( - bool(request.techniques) - or bool(request.dataset_names) - or request.max_dataset_size is not None - or bool(dataset_filters) - ) - if not needs_introspection: - return init_kwargs - - try: - introspection_instance = scenario_class() # type: ignore[ty:missing-argument] - except Exception as exc: - raise ValueError( - f"Cannot resolve runtime configuration for scenario '{request.scenario_name}': " - f"scenario class is not instantiable without arguments ({exc})." - ) from exc - - if request.techniques: - technique_class = introspection_instance._technique_class - technique_enums, technique_converters = self._resolve_techniques_and_converters( - tokens=request.techniques, - technique_class=technique_class, - scenario_name=request.scenario_name, - ) - init_kwargs["scenario_techniques"] = technique_enums - if technique_converters: - init_kwargs["technique_converters"] = technique_converters - - if request.dataset_names or request.max_dataset_size is not None or dataset_filters: - default_config = introspection_instance._default_dataset_config - - if request.dataset_names: - # Construct a fresh instance of the scenario's own dataset-config - # class so subclass-specific behavior is preserved. - default_config_class = type(default_config) - try: - init_kwargs["dataset_config"] = default_config_class( - dataset_names=request.dataset_names, - max_dataset_size=request.max_dataset_size, - filters=dataset_filters or None, - ) - except TypeError as exc: - # The subclass __init__ takes extra required kwargs we cannot - # supply from a backend request. Fall back to the base - # DatasetAttackConfiguration so the run can still proceed; downstream - # scenarios that strictly require the subclass should either - # define a no-extra-required-args constructor or surface the - # incompatibility through their own initialize_async validation. - logger.warning( - "Cannot construct %s(dataset_names=..., max_dataset_size=..., filters=...) (%s). " - "Falling back to a generic DatasetAttackConfiguration; scenario-specific " - "dataset-config behavior may be lost.", - default_config_class.__name__, - exc, - ) - init_kwargs["dataset_config"] = DatasetAttackConfiguration( - dataset_names=request.dataset_names, - max_dataset_size=request.max_dataset_size, - filters=dataset_filters or None, - ) - else: - # Reuse the scenario's default dataset config (preserves subtype + - # the scenario's own default dataset names) and override only the - # sample cap and/or filters. Safe because the introspection instance - # is throwaway. - if request.max_dataset_size is not None: - default_config.max_dataset_size = request.max_dataset_size - if dataset_filters: - default_config.update_filters(filters=dataset_filters) - init_kwargs["dataset_config"] = default_config - - return init_kwargs - - def _resolve_techniques_and_converters( - self, - *, - tokens: list[str], - technique_class: type[Any], - scenario_name: str, - ) -> tuple[list[Any], dict[str, list["Converter"]]]: - """ - Resolve ``--techniques`` tokens into technique enums and per-technique converters. - - Each token has the form ``[:converter.[:converter....]]``. - The base ```` is resolved to a ``ScenarioTechnique`` enum member (which may - be an aggregate). Each ``converter.`` modifier is resolved to a registered - converter instance and appended (in token order) to every concrete technique that the - base technique expands to. - - Args: - tokens: The raw technique tokens from the request. - technique_class: The scenario's ``ScenarioTechnique`` subclass. - scenario_name: The scenario name, used for error messages. - - Returns: - A tuple of (technique enums to pass as ``scenario_techniques``, mapping from concrete - technique name to the list of converters to append for that technique). - - Raises: - ValueError: If a base technique name is unknown, a modifier is malformed, or a - converter name is not registered. - """ - technique_enums: list[Any] = [] - technique_converters: dict[str, list[Converter]] = {} - - for token in tokens: - base_name, _, remainder = token.partition(":") - modifiers = [m for m in remainder.split(":") if m] if remainder else [] - - try: - technique_enum = technique_class(base_name) - except ValueError: - available_techniques = [s.value for s in technique_class] - raise ValueError( - f"Technique '{base_name}' not found for scenario '{scenario_name}'. " - f"Available: {', '.join(available_techniques)}" - ) from None - technique_enums.append(technique_enum) - - converters = self._resolve_converter_modifiers(modifiers=modifiers, token=token) - if not converters: - continue - - for concrete in technique_class.expand({technique_enum}): - technique_converters.setdefault(concrete.value, []).extend(converters) - - return technique_enums, technique_converters - - def _resolve_converter_modifiers(self, *, modifiers: list[str], token: str) -> list["Converter"]: - """ - Resolve the converter modifiers of a single technique token to converter instances. - - Args: - modifiers: The modifier segments of the token (everything after the base technique). - token: The full original token, used for error messages. - - Returns: - The resolved converter instances in token order. - - Raises: - ValueError: If a modifier does not use the ``converter.`` prefix or names a - converter that is not registered. - """ - if not modifiers: - return [] - - instances = ConverterRegistry.get_registry_singleton().instances - converters: list[Converter] = [] - for modifier in modifiers: - if not modifier.startswith(_CONVERTER_MODIFIER_PREFIX): - raise ValueError( - f"Unknown technique modifier '{modifier}' in '{token}'. " - f"Supported modifiers must use the '{_CONVERTER_MODIFIER_PREFIX}' prefix " - f"(e.g. '{_CONVERTER_MODIFIER_PREFIX}translation_spanish')." - ) - converter_name = modifier[len(_CONVERTER_MODIFIER_PREFIX) :] - converter = instances.get(converter_name) - if converter is None: - available = instances.get_names() - available_text = ", ".join(available) if available else "(none registered)" - raise ValueError( - f"Converter '{converter_name}' in '{token}' is not a registered converter " - f"instance. Available converters: {available_text}" - ) - converters.append(converter) - return converters - async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_kwargs: dict[str, Any]) -> Scenario: """ Build and initialize the scenario via the registry. @@ -490,8 +394,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ Delegates the full create + set-parameters + initialize lifecycle to ``ScenarioRegistry.create_and_initialize_async`` so the registry owns scenario creation and initialization. The run-specific common parameters - (target, techniques, dataset config, concurrency) are resolved by - ``_build_init_kwargs`` and forwarded as ``init_kwargs``. + are resolved before this method and forwarded as ``init_kwargs``. Args: request: The run request (for scenario_name, scenario_params, and @@ -541,12 +444,18 @@ async def _execute_run_async(self, *, scenario_result_id: str) -> None: finally: self._run_semaphore.release() - def _build_response(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: + def _build_response( + self, + *, + scenario_result_id: str, + active_error: str | None, + ) -> ScenarioRunSummary | None: """ Build a ScenarioRunResponse by querying the database and merging active task state. Args: scenario_result_id: The scenario result ID. + active_error: Error copied from the active asyncio task, if any. Returns: ScenarioRunResponse if found in the database, None otherwise. @@ -554,24 +463,25 @@ def _build_response(self, *, scenario_result_id: str) -> ScenarioRunSummary | No results = self._memory.get_scenario_results(scenario_result_ids=[scenario_result_id]) if not results: return None - return self._build_response_from_db(scenario_result=results[0]) + return self._build_response_from_db(scenario_result=results[0], active_error=active_error) - def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> ScenarioRunSummary: + def _build_response_from_db( + self, + *, + scenario_result: ScenarioResult, + active_error: str | None = None, + ) -> ScenarioRunSummary: """ Build a ScenarioRunResponse from a database ScenarioResult, merged with active task info. Args: scenario_result: A ScenarioResult retrieved from CentralMemory. + active_error: Error copied from the active asyncio task, if any. Returns: The API response model. """ scenario_result_id = str(scenario_result.id) - active = self._active_tasks.get(scenario_result_id) - - # Clean up finished active tasks - if active is not None and active.task is not None and active.task.done(): - del self._active_tasks[scenario_result_id] # Primary source: DB-persisted error fields error = scenario_result.error_message @@ -589,23 +499,44 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari error_type = error_ars[0].error_type # Fallback: in-memory error for in-flight tasks where DB hasn't been updated yet - if not error and active is not None: - error = active.error + if not error: + error = active_error status = scenario_result.scenario_run_state + terminal = status in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ) + plan = self._load_run_plan(scenario_result=scenario_result) + plan_lookup = _ScenarioPlanLookup.from_plan(plan=plan) # Build result fields from DB (always computed so in-progress runs show progress) - total_attacks = sum(len(results) for results in scenario_result.attack_results.values()) - completed_attacks = total_attacks - techniques_used = scenario_result.get_techniques_used() + total_attacks, completed_attacks, objective_achieved_rate = self._calculate_progress_counts( + scenario_result=scenario_result, + plan=plan, + plan_lookup=plan_lookup, + ) + techniques_used = ( + list(dict.fromkeys(group.display_group for group in plan.atomic_groups)) + if plan is not None + else scenario_result.get_techniques_used() + ) # Surface per-attack errors and retry pressure regardless of overall run status: # a COMPLETED scenario can still hide errored objectives or rate-limit retries. failed_attacks: list[AttackErrorSummary] = [] attack_retries: list[AttackRetrySummary] = [] total_retries = 0 + attempts_by_unit: dict[_ResultUnitIdentity, int] = {} for atomic_attack_name, results in scenario_result.attack_results.items(): for attack_result in results: + unit_identity = self._resolve_result_unit_identity( + atomic_attack_name=atomic_attack_name, + attack_result=attack_result, + plan_lookup=plan_lookup, + ) + attempts_by_unit[unit_identity] = attempts_by_unit.get(unit_identity, 0) + 1 retries = getattr(attack_result, "total_retries", 0) if isinstance(retries, int): total_retries += retries @@ -630,26 +561,358 @@ def _build_response_from_db(self, *, scenario_result: ScenarioResult) -> Scenari total_retries=retries if isinstance(retries, int) else 0, ) ) + total_retries += sum(max(0, attempt_count - 1) for attempt_count in attempts_by_unit.values()) + + updated_at = scenario_result.creation_time + if terminal and scenario_result.completion_time is not None: + updated_at = scenario_result.completion_time return ScenarioRunSummary( scenario_result_id=scenario_result_id, scenario_name=scenario_result.scenario_name, + scenario_registry_name=plan.scenario_registry_name if plan else None, scenario_version=scenario_result.scenario_version, status=status, created_at=scenario_result.creation_time, - updated_at=scenario_result.completion_time or scenario_result.creation_time, + updated_at=updated_at, error=error, error_type=error_type, techniques_used=techniques_used, total_attacks=total_attacks, completed_attacks=completed_attacks, - objective_achieved_rate=scenario_result.objective_achieved_rate(), + objective_achieved_rate=objective_achieved_rate, failed_attacks=failed_attacks, attack_retries=attack_retries, total_retries=total_retries, labels=scenario_result.labels, - completed_at=scenario_result.completion_time, + completed_at=scenario_result.completion_time if terminal else None, + ) + + def _get_active_task(self, *, scenario_result_id: str) -> _ActiveTask | None: + """Return a live task and release completed task state.""" + active = self._active_tasks.get(scenario_result_id) + if active is not None and active.task is not None and active.task.done(): + self._active_tasks.pop(scenario_result_id, None) + return active + + def snapshot_active_run(self, *, scenario_result_id: str) -> _ActiveRunSnapshot: + """ + Copy asyncio-owned run state for use by database-only worker-thread methods. + + Returns: + _ActiveRunSnapshot: An immutable copy of the active state. + """ + active = self._get_active_task(scenario_result_id=scenario_result_id) + if active is None: + return _ActiveRunSnapshot() + active_group_ids = tuple(sorted(active.scenario.active_atomic_group_ids)) if active.scenario is not None else () + return _ActiveRunSnapshot(error=active.error, active_group_ids=active_group_ids) + + @staticmethod + def _load_run_plan(*, scenario_result: ScenarioResult) -> ScenarioRunPlan | None: + """ + Load a validated plan from scenario metadata. + + Returns: + ScenarioRunPlan | None: The stored plan, or None for a legacy row. + """ + metadata = getattr(scenario_result, "metadata", None) + raw_plan = (metadata or {}).get(SCENARIO_RUN_PLAN_METADATA_KEY) + return ScenarioRunPlan.model_validate(raw_plan) if raw_plan is not None else None + + @staticmethod + def _resolve_result_unit_identity( + *, + atomic_attack_name: str, + attack_result: AttackResult, + plan_lookup: _ScenarioPlanLookup, + ) -> _ResultUnitIdentity: + """ + Resolve one attack attempt to its stable planned-unit identity. + + Returns: + _ResultUnitIdentity: The atomic-group and seed-group IDs. + """ + atomic_identifier = attack_result.atomic_attack_identifier + typed_identifier = ( + AtomicAttackIdentifier.from_component_identifier(atomic_identifier) + if isinstance(atomic_identifier, ComponentIdentifier) + else None + ) + objective = str(attack_result.objective) + attribution_data = attack_result.attribution_data + attributed_seed_group_id = attribution_data.get("seed_group_id") if isinstance(attribution_data, dict) else None + seed_group_id = str(attributed_seed_group_id) if attributed_seed_group_id else "" + if not seed_group_id and typed_identifier is not None and typed_identifier.seed_identifiers: + seed_group_id = typed_identifier.logical_seed_group_id + + atomic_group_id = atomic_attack_name + eval_hash = attribution_data.get("parent_eval_hash") if isinstance(attribution_data, dict) else None + planned_group = plan_lookup.resolve_group( + atomic_attack_name=atomic_attack_name, + technique_eval_hash=str(eval_hash) if eval_hash is not None else None, + ) + if planned_group is not None: + atomic_group_id = planned_group.id + if not seed_group_id: + objective_sha256 = to_sha256(objective) + matching_seed_ids = plan_lookup.seed_ids_by_group_and_objective.get( + (planned_group.id, objective_sha256), + (), + ) + if len(matching_seed_ids) == 1: + seed_group_id = matching_seed_ids[0] + if not seed_group_id: + seed_group_id = config_hash({"objective": objective}) + return _ResultUnitIdentity(atomic_group_id=atomic_group_id, seed_group_id=seed_group_id) + + def _calculate_progress_counts( + self, + *, + scenario_result: ScenarioResult, + plan: ScenarioRunPlan | None, + plan_lookup: _ScenarioPlanLookup, + ) -> tuple[int, int, int]: + """ + Calculate planned-unit totals without inflating retries or error attempts. + + Returns: + tuple[int, int, int]: Total, completed, and success-rate percentage. + """ + latest_result_by_unit: dict[_ResultUnitIdentity, AttackResult] = {} + for atomic_attack_name, results in scenario_result.attack_results.items(): + for attack_result in results: + unit_identity = self._resolve_result_unit_identity( + atomic_attack_name=atomic_attack_name, + attack_result=attack_result, + plan_lookup=plan_lookup, + ) + previous = latest_result_by_unit.get(unit_identity) + if previous is None or self._result_order_key(attack_result) > self._result_order_key(previous): + latest_result_by_unit[unit_identity] = attack_result + + planned_units = plan_lookup.planned_units if plan is not None else frozenset(latest_result_by_unit) + total = len(planned_units) + completed_results = [result for unit, result in latest_result_by_unit.items() if unit in planned_units] + completed = len(completed_results) + succeeded = sum(result.outcome == AttackOutcome.SUCCESS for result in completed_results) + rate = int((succeeded / completed) * 100) if completed else 0 + return total, completed, rate + + @staticmethod + def _result_order_key(attack_result: AttackResult) -> tuple[datetime, str]: + """Return a deterministic chronological key for one hydrated result attempt.""" + timestamp = attack_result.timestamp + if not isinstance(timestamp, datetime): + timestamp = datetime.min.replace(tzinfo=timezone.utc) + return timestamp, str(attack_result.attack_result_id) + + def get_run_progress( + self, + *, + scenario_result_id: str, + since: str | None, + limit: int, + ) -> ScenarioRunProgress | None: + """ + Snapshot live state and return compact incremental progress. + + Returns: + ScenarioRunProgress | None: Compact progress when the run exists. + """ + snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id) + return self.get_run_progress_from_storage( + scenario_result_id=scenario_result_id, + since=since, + limit=limit, + active_group_ids=snapshot.active_group_ids, + ) + + def get_run_progress_from_storage( + self, + *, + scenario_result_id: str, + since: str | None, + limit: int, + active_group_ids: Sequence[str], + ) -> ScenarioRunProgress | None: + """Return compact database progress using a previously captured live-state snapshot.""" + header_result = self._memory.get_scenario_result_header(scenario_result_id=scenario_result_id) + if header_result is None: + return None + + cursor = self._decode_progress_cursor(since=since, scenario_result_id=scenario_result_id) + deltas, has_more = self._memory.get_scenario_attack_result_deltas( + scenario_result_id=scenario_result_id, + cursor=cursor, + limit=limit, ) + plan = self._load_run_plan(scenario_result=header_result) + plan_lookup = _ScenarioPlanLookup.from_plan(plan=plan) + plan_complete = plan is not None + response_plan = plan if since is None else None + if plan is None and since is None: + response_plan = self._synthesize_legacy_plan(deltas=deltas) + + response_plan_lookup = plan_lookup if plan is not None else _ScenarioPlanLookup.from_plan(plan=response_plan) + results = [self._map_progress_delta(delta=delta, plan_lookup=response_plan_lookup) for delta in deltas] + next_cursor = ( + self._encode_progress_cursor(scenario_result_id=scenario_result_id, delta=deltas[-1]) if deltas else since + ) + terminal = header_result.scenario_run_state in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ) + return ScenarioRunProgress( + run=ScenarioProgressHeader( + scenario_result_id=scenario_result_id, + scenario_name=header_result.scenario_name, + scenario_registry_name=plan.scenario_registry_name if plan else None, + scenario_version=header_result.scenario_version, + status=header_result.scenario_run_state, + created_at=header_result.creation_time, + completed_at=header_result.completion_time if terminal else None, + ), + plan=response_plan, + reset=False, + active_atomic_group_ids=list(active_group_ids), + results=results, + next_cursor=next_cursor, + has_more=has_more, + plan_complete=plan_complete, + ) + + @staticmethod + def _map_progress_delta( + *, + delta: ScenarioAttackResultDelta, + plan_lookup: _ScenarioPlanLookup, + ) -> ScenarioProgressResult: + """ + Map a lightweight memory row to its REST progress representation. + + Returns: + ScenarioProgressResult: The mapped progress delta. + """ + atomic_attack_name = str(delta.attribution_data.get("parent_collection") or "") + eval_hash = delta.attribution_data.get("parent_eval_hash") + atomic_group_id = config_hash( + {"atomic_attack_name": atomic_attack_name, "technique_eval_hash": eval_hash or ""} + ) + planned_group = plan_lookup.resolve_group( + atomic_attack_name=atomic_attack_name, + technique_eval_hash=str(eval_hash) if eval_hash is not None else None, + ) + if planned_group is not None: + atomic_group_id = planned_group.id + attributed_seed_group_id = delta.attribution_data.get("seed_group_id") + seed_group_id = str(attributed_seed_group_id) if attributed_seed_group_id else "" + if ( + not seed_group_id + and delta.atomic_attack_identifier is not None + and delta.atomic_attack_identifier.seed_identifiers + ): + seed_group_id = delta.atomic_attack_identifier.logical_seed_group_id + if not seed_group_id and delta.objective_sha256: + matching_seed_ids = plan_lookup.seed_ids_by_group_and_objective.get( + (atomic_group_id, delta.objective_sha256), + (), + ) + if len(matching_seed_ids) == 1: + seed_group_id = matching_seed_ids[0] + if not seed_group_id: + seed_group_id = config_hash({"objective": delta.objective}) + return ScenarioProgressResult( + attack_result_id=delta.attack_result_id, + atomic_group_id=atomic_group_id, + atomic_attack_name=atomic_attack_name, + seed_group_id=seed_group_id, + outcome=delta.outcome, + execution_time_ms=delta.execution_time_ms, + timestamp=delta.timestamp, + total_retries=delta.total_retries, + retries=delta.retry_events, + error_type=delta.error_type, + error_message=delta.error_message, + ) + + @staticmethod + def _synthesize_legacy_plan(*, deltas: list[ScenarioAttackResultDelta]) -> ScenarioRunPlan: + """ + Synthesize only known completed legacy units without claiming pending totals. + + Returns: + ScenarioRunPlan: An incomplete plan containing only known units. + """ + seeds: dict[str, ScenarioRunPlanSeedGroup] = {} + groups: dict[str, ScenarioRunPlanAtomicGroup] = {} + seen_seed_ids_by_group: dict[str, set[str]] = {} + empty_plan_lookup = _ScenarioPlanLookup.from_plan(plan=None) + for delta in deltas: + mapped = ScenarioRunService._map_progress_delta( + delta=delta, + plan_lookup=empty_plan_lookup, + ) + seeds.setdefault( + mapped.seed_group_id, + ScenarioRunPlanSeedGroup( + id=mapped.seed_group_id, + objective_sha256=delta.objective_sha256 or to_sha256(delta.objective), + objective=delta.objective, + ), + ) + group = groups.setdefault( + mapped.atomic_group_id, + ScenarioRunPlanAtomicGroup( + id=mapped.atomic_group_id, + atomic_attack_name=mapped.atomic_attack_name, + display_group=mapped.atomic_attack_name, + technique_eval_hash=str(delta.attribution_data.get("parent_eval_hash") or ""), + seed_group_ids=[], + ), + ) + seen_seed_ids = seen_seed_ids_by_group.setdefault(mapped.atomic_group_id, set()) + if mapped.seed_group_id not in seen_seed_ids: + seen_seed_ids.add(mapped.seed_group_id) + group.seed_group_ids.append(mapped.seed_group_id) + return ScenarioRunPlan(atomic_groups=list(groups.values()), seed_groups=list(seeds.values())) + + @staticmethod + def _encode_progress_cursor(*, scenario_result_id: str, delta: ScenarioAttackResultDelta) -> str: + payload = { + "v": 1, + "run": scenario_result_id, + "timestamp": delta.timestamp.isoformat(), + "attack_result_id": delta.attack_result_id, + } + return base64.urlsafe_b64encode(json.dumps(payload, separators=(",", ":")).encode()).decode().rstrip("=") + + @staticmethod + def _decode_progress_cursor( + *, + since: str | None, + scenario_result_id: str, + ) -> AttackResultKeysetCursor | None: + if since is None: + return None + try: + padded = since + "=" * (-len(since) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded).decode()) + except Exception as exc: + raise ValueError("Malformed scenario progress cursor.") from exc + if not isinstance(payload, dict): + raise ValueError("Malformed scenario progress cursor.") + if payload.get("v") != 1 or payload.get("run") != scenario_result_id: + raise ValueError("Cursor does not belong to this scenario run.") + try: + timestamp = datetime.fromisoformat(payload["timestamp"]) + attack_result_id = str(uuid.UUID(payload["attack_result_id"])) + except Exception as exc: + raise ValueError("Malformed scenario progress cursor.") from exc + if timestamp.tzinfo is None: + raise ValueError("Cursor timestamp must include a timezone.") + return AttackResultKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) def get_run_results(self, *, scenario_result_id: str) -> ScenarioResult | None: """ diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index 46721d8ed1..b5d6d3ac5b 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -1,55 +1,84 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -""" -Scenario service for listing available scenarios. +"""Scenario catalog and side-effect-free planning service.""" -Provides read-only access to the ScenarioRegistry, exposing scenario metadata -through the REST API. -""" +from __future__ import annotations +import asyncio +import logging +from collections import OrderedDict from functools import lru_cache +from time import monotonic from pyrit.backend.models.common import PaginationInfo from pyrit.backend.models.scenarios import ListRegisteredScenariosResponse +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.models.catalog.scenario import ( RegisteredScenario, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, ) from pyrit.registry import ScenarioMetadata, ScenarioRegistry +logger = logging.getLogger(__name__) +_ESTIMATE_CACHE_SIZE = 128 +_ESTIMATE_CONCURRENCY = 1 +_ESTIMATE_INFLIGHT_SIZE = 256 +_UNAVAILABLE_CACHE_TTL_SECONDS = 30.0 +_EstimateCacheKey = tuple[str, int] +_EstimateCacheValue = tuple[ScenarioRunSizeEstimate, float | None] +_EstimateTask = asyncio.Task[ScenarioRunSizeEstimate] -def _metadata_to_registered_scenario(metadata: ScenarioMetadata) -> RegisteredScenario: + +def _metadata_to_registered_scenario( + *, + metadata: ScenarioMetadata, + default_run_size: ScenarioRunSizeEstimate | None = None, +) -> RegisteredScenario: """ Convert a ScenarioMetadata dataclass to a ScenarioSummary Pydantic model. Args: metadata: The registry metadata for a scenario. + default_run_size: Scenario-owned default-run estimate. Returns: - ScenarioSummary Pydantic model. + RegisteredScenario: Public catalog projection. """ + estimate = default_run_size or ScenarioRunSizeEstimate.unavailable() return RegisteredScenario( scenario_name=metadata.registry_name, scenario_type=metadata.class_name, + scenario_version=metadata.scenario_version, description=metadata.class_description, + description_markdown=metadata.description_markdown, default_technique=metadata.default_technique, + default_techniques=list(metadata.default_techniques), aggregate_techniques=list(metadata.aggregate_techniques), + aggregate_technique_expansions={ + aggregate: list(expansion) for aggregate, expansion in metadata.aggregate_technique_expansions + }, all_techniques=list(metadata.all_techniques), default_datasets=list(metadata.default_datasets), + default_dataset_summaries=estimate.datasets, supported_parameters=list(metadata.supported_parameters), + baseline_policy=metadata.baseline_policy, + include_baseline_by_default=metadata.include_baseline_by_default, + default_run_size=estimate, ) class ScenarioService: - """ - Service for listing available scenarios. - - Uses ScenarioRegistry as the source of truth for scenario metadata. - """ + """Expose Scenario metadata and scenario-owned run-size planning.""" def __init__(self) -> None: - """Initialize the scenario service.""" + """Initialize registry access and the per-scenario default-estimate cache.""" self._registry = ScenarioRegistry.get_registry_singleton() + self._estimate_cache: OrderedDict[_EstimateCacheKey, _EstimateCacheValue] = OrderedDict() + self._estimate_tasks: OrderedDict[_EstimateCacheKey, _EstimateTask] = OrderedDict() + self._estimate_task_lock = asyncio.Lock() + self._estimate_semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) async def list_scenarios_async( self, @@ -58,21 +87,29 @@ async def list_scenarios_async( cursor: str | None = None, ) -> ListRegisteredScenariosResponse: """ - List all available scenarios with pagination. - - Args: - limit: Maximum items to return per page. - cursor: Pagination cursor (scenario_name to start after). + List scenarios with cached default estimates and cursor pagination. Returns: - ScenarioListResponse with paginated scenario summaries. + ListRegisteredScenariosResponse: The requested catalog page. """ all_metadata = self._registry.get_all_registered_class_metadata() - all_summaries = [_metadata_to_registered_scenario(m) for m in all_metadata] + all_summaries = [_metadata_to_registered_scenario(metadata=m) for m in all_metadata] page, has_more = self._paginate(items=all_summaries, cursor=cursor, limit=limit) + metadata_by_name = {metadata.registry_name: metadata for metadata in all_metadata} + estimates = await asyncio.gather( + *(self._get_default_run_size_estimate_async(metadata=metadata_by_name[item.scenario_name]) for item in page) + ) + page = [ + item.model_copy( + update={ + "default_run_size": estimate, + "default_dataset_summaries": estimate.datasets, + } + ) + for item, estimate in zip(page, estimates, strict=True) + ] next_cursor = page[-1].scenario_name if has_more and page else None - return ListRegisteredScenariosResponse( items=page, pagination=PaginationInfo( @@ -85,19 +122,185 @@ async def list_scenarios_async( async def get_scenario_async(self, *, scenario_name: str) -> RegisteredScenario | None: """ - Get a single scenario by registry name. - - Args: - scenario_name: The registry key of the scenario (e.g., 'foundry.red_team_agent'). + Get one scenario and its cached default estimate. Returns: - ScenarioSummary if found, None otherwise. + RegisteredScenario | None: The catalog entry, or None when it is not registered. """ metadata = self._registry.get_registered_class_metadata(scenario_name) if metadata is not None: - return _metadata_to_registered_scenario(metadata) + estimate = await self._get_default_run_size_estimate_async(metadata=metadata) + return _metadata_to_registered_scenario(metadata=metadata, default_run_size=estimate) return None + async def estimate_scenario_run_size_async( + self, + *, + scenario_name: str, + request: ScenarioRunSizeEstimateRequest, + ) -> ScenarioRunSizeEstimate | None: + """ + Estimate one configured scenario without creating a run. + + Args: + scenario_name: Registered scenario name. + request: Request-specific techniques, datasets, baseline, and parameters. + + Returns: + ScenarioRunSizeEstimate | None: Estimate, or ``None`` when the scenario is unknown. + """ + metadata = self._registry.get_registered_class_metadata(scenario_name) + if metadata is None: + return None + + semaphore = getattr(self, "_estimate_semaphore", None) + if semaphore is None: + semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) + self._estimate_semaphore = semaphore + async with semaphore: + return await self._estimate_configured_run_size_async( + scenario_name=scenario_name, + request=request, + ) + + async def _get_default_run_size_estimate_async(self, *, metadata: ScenarioMetadata) -> ScenarioRunSizeEstimate: + """Return a cached, cancellation-safe scenario-owned estimate.""" + cache_key = (metadata.registry_name, metadata.scenario_version) + cache = getattr(self, "_estimate_cache", None) + if cache is None: + cache = OrderedDict() + self._estimate_cache = cache + while True: + cached = self._read_estimate_cache(cache_key=cache_key) + if cached is not None: + return cached + + task_lock = getattr(self, "_estimate_task_lock", None) + if task_lock is None: + task_lock = asyncio.Lock() + self._estimate_task_lock = task_lock + wait_for_capacity: _EstimateTask | None = None + task: _EstimateTask | None = None + async with task_lock: + cached = self._read_estimate_cache(cache_key=cache_key) + if cached is not None: + return cached + + tasks = getattr(self, "_estimate_tasks", None) + if tasks is None: + tasks = OrderedDict() + self._estimate_tasks = tasks + for completed_key in [key for key, candidate in tasks.items() if candidate.done()]: + del tasks[completed_key] + task = tasks.get(cache_key) + if task is None: + if len(tasks) >= _ESTIMATE_INFLIGHT_SIZE: + wait_for_capacity = next(iter(tasks.values())) + else: + task = asyncio.create_task( + self._compute_default_run_size_estimate_async( + scenario_name=metadata.registry_name, + cache_key=cache_key, + ) + ) + tasks[cache_key] = task + + def clear_estimate_task(completed_task: _EstimateTask) -> None: + self._clear_estimate_task(task=completed_task, cache_key=cache_key) + + task.add_done_callback(clear_estimate_task) + + if task is not None: + estimate = await asyncio.shield(task) + assert isinstance(estimate, ScenarioRunSizeEstimate) + return estimate + if wait_for_capacity is not None: + await asyncio.shield(wait_for_capacity) + + def _read_estimate_cache(self, *, cache_key: _EstimateCacheKey) -> ScenarioRunSizeEstimate | None: + """Return a live cached estimate and discard expired unavailable entries.""" + cache = self._estimate_cache + cached = cache.get(cache_key) + if cached is None: + return None + estimate, expires_at = cached + if expires_at is not None and monotonic() >= expires_at: + del cache[cache_key] + return None + cache.move_to_end(cache_key) + return estimate + + async def _compute_default_run_size_estimate_async( + self, + *, + scenario_name: str, + cache_key: _EstimateCacheKey, + ) -> ScenarioRunSizeEstimate: + """ + Construct and estimate one scenario on the owning event loop. + + Returns: + ScenarioRunSizeEstimate: Scenario-owned estimate. + """ + semaphore = getattr(self, "_estimate_semaphore", None) + if semaphore is None: + semaphore = asyncio.Semaphore(_ESTIMATE_CONCURRENCY) + self._estimate_semaphore = semaphore + async with semaphore: + try: + scenario = await asyncio.to_thread(self._registry.create_instance, scenario_name) + estimate = await scenario.get_default_run_size_estimate_async() + except Exception as exc: + logger.warning("Default-run estimate failed for scenario '%s': %s", scenario_name, exc) + estimate = ScenarioRunSizeEstimate.unavailable( + note=f"The scenario could not resolve its default inputs for estimation ({type(exc).__name__})." + ) + + expires_at = monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS if estimate.estimated_attack_count is None else None + cache = self._estimate_cache + cache[cache_key] = (estimate, expires_at) + cache.move_to_end(cache_key) + while len(cache) > _ESTIMATE_CACHE_SIZE: + cache.popitem(last=False) + return estimate + + def _clear_estimate_task(self, *, task: _EstimateTask, cache_key: _EstimateCacheKey) -> None: + """Remove a completed single-flight task without disturbing a replacement.""" + tasks = self._estimate_tasks + if tasks.get(cache_key) is task: + del tasks[cache_key] + + async def _estimate_configured_run_size_async( + self, + *, + scenario_name: str, + request: ScenarioRunSizeEstimateRequest, + ) -> ScenarioRunSizeEstimate: + """ + Resolve and estimate one request on the owning event loop. + + Returns: + ScenarioRunSizeEstimate: Request-specific scenario estimate. + """ + scenario_class = self._registry.get_class(scenario_name) + resolver = ScenarioConfigurationResolver() + objective_target = resolver.resolve_target(target_name=request.target_name) if request.target_name else None + estimate_kwargs = resolver.resolve_configuration( + scenario_name=scenario_name, + scenario_class=scenario_class, + objective_target=objective_target, + techniques=request.techniques, + dataset_names=request.dataset_names, + max_dataset_size=request.max_dataset_size, + dataset_filters=request.dataset_filters, + include_baseline=request.include_baseline, + ) + return await self._registry.create_and_estimate_async( + name=scenario_name, + scenario_params=request.scenario_params or {}, + **estimate_kwargs, + ) + @staticmethod def _paginate( *, @@ -106,15 +309,10 @@ def _paginate( limit: int, ) -> tuple[list[RegisteredScenario], bool]: """ - Apply cursor-based pagination. - - Args: - items: Full list of items. - cursor: Scenario name to start after. - limit: Maximum items per page. + Apply scenario-name cursor pagination. Returns: - Tuple of (paginated items, has_more flag). + tuple[list[RegisteredScenario], bool]: The page and whether another page exists. """ start_idx = 0 if cursor: @@ -122,7 +320,6 @@ def _paginate( if item.scenario_name == cursor: start_idx = i + 1 break - page = items[start_idx : start_idx + limit] has_more = len(items) > start_idx + limit return page, has_more @@ -131,9 +328,9 @@ def _paginate( @lru_cache(maxsize=1) def get_scenario_service() -> ScenarioService: """ - Get the global scenario service instance. + Get the process-wide Scenario service. Returns: - The singleton ScenarioService instance. + ScenarioService: The cached service instance. """ return ScenarioService() diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 5b04d2611f..0883ee5d46 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -21,6 +21,7 @@ from pyrit.models.catalog import ( RegisteredInitializer, RegisteredScenario, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -519,7 +520,7 @@ def _print_transcript(*, messages: list[TranscriptMessage]) -> None: # --------------------------------------------------------------------------- -def print_scenario_runs_list(*, runs: list[ScenarioRunSummary]) -> None: +def print_scenario_runs_list(*, runs: list[ScenarioRunListItem]) -> None: """ Print a list of scenario run summaries. @@ -534,9 +535,12 @@ def print_scenario_runs_list(*, runs: list[ScenarioRunSummary]) -> None: print("=" * 80) for idx, run in enumerate(runs, start=1): created = run.created_at.isoformat() if run.created_at else "?" + planned_attacks = ( + f"{run.total_attacks} planned attacks" if run.total_attacks is not None else "planned attacks unknown" + ) print( f" {idx}) [{run.status.value}] {run.scenario_name} (id: {run.scenario_result_id}) — " - f"{run.total_attacks} attacks, {run.objective_achieved_rate}% success — {created}" + f"{planned_attacks} — {created}" ) print("=" * 80) print(f"\nTotal runs: {len(runs)}") diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 574d10619b..a930a1f09c 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -23,6 +23,7 @@ RegisteredInitializer, RegisteredScenario, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -342,17 +343,17 @@ async def cancel_scenario_run_async(self, *, scenario_result_id: str) -> Scenari self._raise_for_status(resp) return ScenarioRunSummary.model_validate(resp.json()) - async def list_scenario_runs_async(self, *, limit: int = 100) -> list[ScenarioRunSummary]: + async def list_scenario_runs_async(self, *, limit: int = 100) -> list[ScenarioRunListItem]: """ List tracked scenario runs. Returns: - list[ScenarioRunSummary]: All tracked scenario runs. + list[ScenarioRunListItem]: All tracked scenario runs. """ - from pyrit.models.catalog import ScenarioRunSummary + from pyrit.models.catalog import ScenarioRunListItem payload = await self._get_json_async(path="/api/scenarios/runs", params={"limit": limit}) - return [ScenarioRunSummary.model_validate(item) for item in payload.get("items", [])] + return [ScenarioRunListItem.model_validate(item) for item in payload.get("items", [])] # ------------------------------------------------------------------ # Attacks / conversations diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 30c40d1a9e..005f070ffb 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -12,13 +12,14 @@ MessageListNormalizer, MessageStringNormalizer, ) +from pyrit.models import ChatMessageRole # noqa: TC001 - public annotation must resolve at runtime from pyrit.prompt_target.common.target_capabilities import CapabilityName if TYPE_CHECKING: from pyrit.executor.attack.component.prepended_history_send_context import ( PrependedHistorySendContext, ) - from pyrit.models import ChatMessageRole, Message + from pyrit.models import Message from pyrit.prompt_target.common.prompt_target import PromptTarget diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index 64d75d1066..6bfbab03c5 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -176,6 +176,7 @@ async def execute_attack_from_seed_groups_async( field_overrides: Sequence[dict[str, Any]] | None = None, return_partial_on_failure: bool = False, attribution: AttackResultAttribution | None = None, + attributions: Sequence[AttackResultAttribution] | None = None, **broadcast_fields: Any, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -205,6 +206,8 @@ async def execute_attack_from_seed_groups_async( When ``None`` (default), no attribution is applied. The same attribution is shared across all tasks; per-task identity is reconstructed from the row's own ``objective_sha256``. + attributions: Optional per-seed-group attribution. Must match + ``seed_groups`` and cannot be combined with ``attribution``. **broadcast_fields: Fields applied to all seed groups (e.g., memory_labels). Per-seed-group field_overrides take precedence. @@ -212,7 +215,8 @@ async def execute_attack_from_seed_groups_async( AttackExecutorResult with completed results and any incomplete objectives. Raises: - ValueError: If seed_groups is empty or field_overrides length doesn't match. + ValueError: If seed groups are empty, override/attribution lengths do not + match, or shared and per-task attribution are both provided. BaseException: If return_partial_on_failure=False and any objective fails. """ if not seed_groups: @@ -222,6 +226,19 @@ async def execute_attack_from_seed_groups_async( raise ValueError( f"field_overrides length ({len(field_overrides)}) must match seed_groups length ({len(seed_groups)})" ) + if attributions is not None and len(attributions) != len(seed_groups): + raise ValueError( + f"attributions length ({len(attributions)}) must match seed_groups length ({len(seed_groups)})" + ) + if attribution is not None and attributions is not None: + raise ValueError("Provide attribution or attributions, not both") + effective_attributions = ( + list(attributions) + if attributions is not None + else [attribution] * len(seed_groups) + if attribution is not None + else None + ) params_type = attack.params_type @@ -263,11 +280,16 @@ async def build_params_async(i: int, sg: AttackSeedGroup) -> AttackParameters: if build_failures and not return_partial_on_failure: raise build_failures[0][2] + successful_attributions = ( + [effective_attributions[index] for index in successful_input_indices] + if effective_attributions is not None + else None + ) execution_result = await self._execute_with_params_list_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, + attributions=successful_attributions, input_indices=successful_input_indices, ) return self._merge_parameter_build_failures( @@ -341,7 +363,7 @@ async def execute_attack_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, + attributions=[attribution] * len(params_list) if attribution is not None else None, ) async def _execute_with_params_list_async( @@ -350,7 +372,7 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, - attribution: AttackResultAttribution | None = None, + attributions: Sequence[AttackResultAttribution] | None = None, input_indices: Sequence[int] | None = None, ) -> AttackExecutorResult[AttackStrategyResultT]: """ @@ -363,22 +385,28 @@ async def _execute_with_params_list_async( attack: The attack strategy to execute. params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. - attribution: Optional ``AttackResultAttribution`` stamped onto every - per-task ``AttackContext`` so the persistence path can record - orchestrator linkage. + attributions: Optional per-task attribution matching ``params_list``. input_indices: Original input positions for ``params_list``. Defaults to sequential positions when parameters were constructed directly. Returns: AttackExecutorResult with completed results and any incomplete objectives. + + Raises: + ValueError: If per-task attribution or input-index lengths do not match. """ semaphore = self._get_semaphore() + if attributions is not None and len(attributions) != len(params_list): + raise ValueError( + f"attributions length ({len(attributions)}) must match params_list length ({len(params_list)})" + ) async def run_one_async(index: int, params: AttackParameters) -> AttackStrategyResultT: async with semaphore: context = attack._context_type(params=params) - if attribution is not None: - context._attribution = attribution + task_attribution = attributions[index] if attributions is not None else None + if task_attribution is not None: + context._attribution = task_attribution return await attack.execute_with_context_async(context=context) tasks = [run_one_async(i, p) for i, p in enumerate(params_list)] diff --git a/pyrit/executor/attack/core/attack_result_attribution.py b/pyrit/executor/attack/core/attack_result_attribution.py index 2953f7160a..93bb0efb3a 100644 --- a/pyrit/executor/attack/core/attack_result_attribution.py +++ b/pyrit/executor/attack/core/attack_result_attribution.py @@ -44,8 +44,11 @@ class AttackResultAttribution: to the atomic attack's technique evaluation hash, e.g. ``self.technique_eval_hash`` (computed via ``AtomicAttackEvaluationIdentifier``). + seed_group_id (str | None): Optional logical seed-group fingerprint for + per-task progress attribution. """ parent_id: str parent_collection: str parent_eval_hash: str | None = None + seed_group_id: str | None = None diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index c1e260fc12..1661e6a84c 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -315,6 +315,8 @@ def _apply_attribution( } if attribution.parent_eval_hash is not None: attribution_data["parent_eval_hash"] = attribution.parent_eval_hash + if attribution.seed_group_id is not None: + attribution_data["seed_group_id"] = attribution.seed_group_id result.attribution_data = attribution_data @staticmethod diff --git a/pyrit/memory/__init__.py b/pyrit/memory/__init__.py index fd2ee20450..cab9590799 100644 --- a/pyrit/memory/__init__.py +++ b/pyrit/memory/__init__.py @@ -16,7 +16,7 @@ from pyrit.memory.azure_sql_memory import AzureSQLMemory from pyrit.memory.central_memory import CentralMemory from pyrit.memory.memory_embedding import MemoryEmbedding - from pyrit.memory.memory_interface import AttackResultsKeysetCursor, MemoryInterface + from pyrit.memory.memory_interface import AttackResultKeysetCursor, MemoryInterface from pyrit.memory.memory_models import AttackResultEntry, EmbeddingDataEntry, PromptMemoryEntry, SeedEntry from pyrit.memory.sqlite_memory import SQLiteMemory from pyrit.memory.storage import ( @@ -41,7 +41,7 @@ _LAZY_EXPORTS: dict[str, str | tuple[str, str | None]] = { "AllowedCategories": "pyrit.memory.storage", "AttackResultEntry": "pyrit.memory.memory_models", - "AttackResultsKeysetCursor": "pyrit.memory.memory_interface", + "AttackResultKeysetCursor": "pyrit.memory.memory_interface", "AudioPathDataTypeSerializer": "pyrit.memory.storage", "AzureBlobStorageIO": "pyrit.memory.storage", "AzureSQLMemory": "pyrit.memory.azure_sql_memory", diff --git a/pyrit/memory/alembic/versions/6b8d0f2a4c1e_index_scenario_progress_deltas.py b/pyrit/memory/alembic/versions/6b8d0f2a4c1e_index_scenario_progress_deltas.py new file mode 100644 index 0000000000..09f3578648 --- /dev/null +++ b/pyrit/memory/alembic/versions/6b8d0f2a4c1e_index_scenario_progress_deltas.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Index scenario-linked attack results for ascending progress deltas. + +Revision ID: 6b8d0f2a4c1e +Revises: 4c9a6e1f2b7d +Create Date: 2026-08-06 19:41:22.000000 +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "6b8d0f2a4c1e" +down_revision: str | None = "4c9a6e1f2b7d" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_INDEX_NAME = "ix_AttackResultEntries_attribution_parent_timestamp_id" + + +def upgrade() -> None: + """Create the scenario progress keyset index.""" + op.create_index( + _INDEX_NAME, + "AttackResultEntries", + ["attribution_parent_id", "timestamp", "id"], + ) + + +def downgrade() -> None: + """Drop the scenario progress keyset index.""" + op.drop_index(_INDEX_NAME, table_name="AttackResultEntries") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index d698593a5b..1b0263e865 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -3,6 +3,7 @@ import abc import atexit +import json import logging import re import uuid @@ -55,6 +56,7 @@ AdditionalInitializer, AtomicAttackIdentifier, AttackIdentifier, + AttackOutcome, AttackResult, AttackTechniqueIdentifier, ComponentIdentifier, @@ -67,6 +69,8 @@ IdentifierType, Message, MessagePiece, + RetryEvent, + ScenarioAttackResultDelta, ScenarioIdentifier, ScenarioResult, ScenarioRunState, @@ -92,7 +96,7 @@ IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) -class AttackResultsKeysetCursor(NamedTuple): +class AttackResultKeysetCursor(NamedTuple): """ Keyset (seek) anchor identifying the last attack result on a page. @@ -108,12 +112,12 @@ class AttackResultsKeysetCursor(NamedTuple): attack_result_id: str @classmethod - def from_attack_result(cls, result: AttackResult) -> "AttackResultsKeysetCursor": + def from_attack_result(cls, result: AttackResult) -> "AttackResultKeysetCursor": """ Build the keyset anchor for ``result`` (typically the last row of a page). Returns: - AttackResultsKeysetCursor: Anchor capturing the result's recency sort key. + AttackResultKeysetCursor: Anchor capturing the result's recency sort key. """ return cls( timestamp=result.timestamp, @@ -157,7 +161,7 @@ class _AttackResultQuery: min_turns: int | None = None max_turns: int | None = None limit: int | None = None - after: AttackResultsKeysetCursor | None = None + after: AttackResultKeysetCursor | None = None def __post_init__(self) -> None: """Snapshot mutable sequence and mapping inputs.""" @@ -413,7 +417,7 @@ def _attack_results_recency_order_by(self) -> list[Any]: """ return [AttackResultEntry.timestamp.desc(), AttackResultEntry.id.desc()] - def _attack_results_keyset_seek_condition(self, *, after: AttackResultsKeysetCursor) -> Any: + def _attack_results_keyset_seek_condition(self, *, after: AttackResultKeysetCursor) -> Any: """ Build the keyset seek predicate selecting rows strictly after ``after``. @@ -3103,7 +3107,7 @@ def get_attack_results( min_turns: int | None = None, max_turns: int | None = None, limit: int | None = None, - after: AttackResultsKeysetCursor | None = None, + after: AttackResultKeysetCursor | None = None, ) -> Sequence[AttackResult]: """ Retrieve a list of AttackResult objects based on the specified filters. @@ -3169,7 +3173,7 @@ def get_attack_results( return, ordered by recency. When either ``limit`` or ``after`` is provided, deduplication and pagination happen in the database (via ``ROW_NUMBER()``) instead of loading every row into memory. Defaults to None (return all). - after (AttackResultsKeysetCursor | None, optional): Keyset (seek) anchor from a + after (AttackResultKeysetCursor | None, optional): Keyset (seek) anchor from a previous page. When provided, only results ordered strictly after the anchor under the recency sort are returned, giving insert/delete-stable pagination without a drifting numeric offset. Defaults to None (start at the first page). @@ -3454,7 +3458,7 @@ def _query_paginated_attack_results( min_turns: int | None, max_turns: int | None, limit: int | None, - after: AttackResultsKeysetCursor | None, + after: AttackResultKeysetCursor | None, ) -> list[AttackResult]: """ Deduplicate in SQL (filter-aware) and return one recency-ordered page of results. @@ -3475,7 +3479,7 @@ def _query_paginated_attack_results( min_turns (int | None): Inclusive lower bound on ``executed_turns`` for winners. max_turns (int | None): Inclusive upper bound on ``executed_turns`` for winners. limit (int | None): Maximum number of results to return. - after (AttackResultsKeysetCursor | None): Keyset anchor; only rows ordered strictly + after (AttackResultKeysetCursor | None): Keyset anchor; only rows ordered strictly after it are returned. ``None`` starts at the first page. Returns: @@ -3644,6 +3648,12 @@ def update_scenario_run_state( entry.scenario_run_state = scenario_run_state.value entry.error_message = error_message entry.error_type = error_type + if scenario_run_state in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ): + entry.completion_time = datetime.now(tz=timezone.utc) session.commit() @@ -3677,6 +3687,122 @@ def update_scenario_metadata( entry.scenario_metadata = metadata if metadata else None session.commit() + def get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResult | None: + """Return one ScenarioResult header without hydrating linked attack results.""" + with closing(self.get_session()) as session: + entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first() + return entry.get_scenario_result() if entry is not None else None + + def get_scenario_result_headers(self, *, limit: int = 100) -> Sequence[ScenarioResult]: + """ + Return recent ScenarioResult headers without hydrating linked attack results. + + Returns: + Sequence[ScenarioResult]: Recent scenario metadata ordered newest first. + + Raises: + ValueError: If limit is outside the bounded run-history range. + """ + if limit < 1 or limit > 100: + raise ValueError("Scenario run history limit must be between 1 and 100.") + entries = self._query_entries( + ScenarioResultEntry, + order_by=[ + ScenarioResultEntry.timestamp.desc(), + ScenarioResultEntry.id.desc(), + ], + limit=limit, + ) + return [entry.get_scenario_result() for entry in entries] + + def get_scenario_attack_result_deltas( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 100, + ) -> tuple[list[ScenarioAttackResultDelta], bool]: + """ + Return bounded scenario-linked result deltas in ascending keyset order. + + This projection intentionally selects only progress fields and never + hydrates PromptMemoryEntry, ScoreEntry, or a full ScenarioResult. + + Returns: + tuple[list[ScenarioAttackResultDelta], bool]: The page and whether more rows exist. + + Raises: + ValueError: If the limit or cursor identifiers are invalid. + """ + if limit < 1 or limit > 500: + raise ValueError("Scenario progress limit must be between 1 and 500.") + + scenario_uuid = uuid.UUID(scenario_result_id) + conditions: list[Any] = [AttackResultEntry.attribution_parent_id == scenario_uuid] + if cursor is not None: + cursor_uuid = uuid.UUID(cursor.attack_result_id) + conditions.append( + or_( + AttackResultEntry.timestamp > cursor.timestamp, + and_( + AttackResultEntry.timestamp == cursor.timestamp, + AttackResultEntry.id > cursor_uuid, + ), + ) + ) + + statement = ( + select( + AttackResultEntry.id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + ) + .where(and_(*conditions)) + .order_by(AttackResultEntry.timestamp.asc(), AttackResultEntry.id.asc()) + .limit(limit + 1) + ) + with closing(self.get_session()) as session: + rows = session.execute(statement).all() + + has_more = len(rows) > limit + deltas: list[ScenarioAttackResultDelta] = [] + for row in rows[:limit]: + retry_events = [ + RetryEvent.model_validate(event) + for event in (json.loads(row.retry_events_json) if row.retry_events_json else []) + ] + atomic_identifier = ( + AtomicAttackIdentifier.model_validate(row.atomic_attack_identifier) + if row.atomic_attack_identifier + else None + ) + deltas.append( + ScenarioAttackResultDelta( + attack_result_id=str(row.id), + objective=row.objective, + objective_sha256=row.objective_sha256, + atomic_attack_identifier=atomic_identifier, + outcome=AttackOutcome(row.outcome), + execution_time_ms=row.execution_time_ms, + timestamp=row.timestamp, + retry_events=retry_events, + total_retries=row.total_retries or 0, + error_type=row.error_type, + error_message=row.error_message, + attribution_data=row.attribution_data or {}, + ) + ) + return deltas, has_more + def get_scenario_results( self, *, diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 8fc1247bfd..3803743bb9 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1555,6 +1555,13 @@ class AttackResultEntry(Base): Index("ix_AttackResultEntries_conversation_id", "conversation_id"), # Serves the History recency ORDER BY timestamp DESC, id DESC and its keyset seek. Index("ix_AttackResultEntries_timestamp_id", "timestamp", "id"), + # Serves scenario progress deltas scoped by parent and ordered oldest-first. + Index( + "ix_AttackResultEntries_attribution_parent_timestamp_id", + "attribution_parent_id", + "timestamp", + "id", + ), {"extend_existing": True}, ) id = mapped_column(CustomUUID, nullable=False, primary_key=True) @@ -1864,12 +1871,9 @@ class ScenarioResultEntry(Base): error_message: Mapped[str | None] = mapped_column(Unicode, nullable=True) error_type: Mapped[str | None] = mapped_column(String, nullable=True) - # Free-form JSON metadata stamped by the scenario. Currently used to record - # ``objective_hashes`` — the objective sha256 set chosen on the - # first run, replayed on resume so a fresh ``random.sample`` can't - # silently change which objectives the scenario operates on. Column is - # named ``scenario_metadata`` because SQLAlchemy's ``DeclarativeBase`` - # reserves ``metadata`` as a class attribute on the model. + # Free-form JSON metadata stamped by the scenario. Stores the normalized run + # plan and sampled objective hashes. Column is named ``scenario_metadata`` + # because SQLAlchemy's ``DeclarativeBase`` reserves ``metadata``. scenario_metadata: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) def __init__(self, *, entry: ScenarioResult) -> None: diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 4a6dfc36ca..680a434408 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -23,6 +23,14 @@ if TYPE_CHECKING: from pyrit.models.additional_initializer import AdditionalInitializer + from pyrit.models.catalog import ( + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioRunListItem, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, + ) from pyrit.models.conversation_stats import ConversationStats from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation from pyrit.models.harm_definition import HarmDefinition, ScaleDescription, get_all_harm_definitions @@ -52,6 +60,7 @@ TargetIdentifier, class_name_to_snake_case, compute_eval_hash, + compute_seed_group_hash, config_hash, snake_case_to_class_name, validate_registry_name, @@ -95,6 +104,17 @@ from pyrit.models.results.scenario_result import ScenarioResult, ScenarioRunState from pyrit.models.results.strategy_result import StrategyResult, StrategyResultT from pyrit.models.retry_event import RetryEvent + from pyrit.models.scenario_progress import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + SCENARIO_RUN_PLAN_VERSION, + ScenarioAttackResultDelta, + ScenarioProgressHeader, + ScenarioProgressResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunProgress, + ) from pyrit.models.score import ( Condition, ContentScorable, @@ -185,6 +205,7 @@ "IdentifierFilter": "pyrit.models.identifiers", "IdentifierType": "pyrit.models.identifiers", "JSONValue": "pyrit.models.identifiers", + "compute_seed_group_hash": "pyrit.models.identifiers", "COMMON_JSON_SCHEMAS": "pyrit.models.target", "JsonResponseConfig": "pyrit.models.target", "get_common_json_schema": "pyrit.models.target", @@ -219,8 +240,23 @@ "ScorerEvaluationIdentifier": "pyrit.models.identifiers", "ScorerIdentifier": "pyrit.models.identifiers", "ScenarioIdentifier": "pyrit.models.identifiers", + "ScenarioDatasetSizeCap": "pyrit.models.catalog", + "ScenarioDatasetSummary": "pyrit.models.catalog", + "ScenarioRunListItem": "pyrit.models.catalog", + "ScenarioRunSizeComponent": "pyrit.models.catalog", + "ScenarioRunSizeEstimate": "pyrit.models.catalog", + "ScenarioRunSizeEstimateRequest": "pyrit.models.catalog", "ScenarioResult": "pyrit.models.results.scenario_result", "ScenarioRunState": "pyrit.models.results.scenario_result", + "SCENARIO_RUN_PLAN_METADATA_KEY": "pyrit.models.scenario_progress", + "SCENARIO_RUN_PLAN_VERSION": "pyrit.models.scenario_progress", + "ScenarioAttackResultDelta": "pyrit.models.scenario_progress", + "ScenarioProgressHeader": "pyrit.models.scenario_progress", + "ScenarioProgressResult": "pyrit.models.scenario_progress", + "ScenarioRunPlan": "pyrit.models.scenario_progress", + "ScenarioRunPlanAtomicGroup": "pyrit.models.scenario_progress", + "ScenarioRunPlanSeedGroup": "pyrit.models.scenario_progress", + "ScenarioRunProgress": "pyrit.models.scenario_progress", "Seed": "pyrit.models.seeds", "AttackSeedGroup": "pyrit.models.seeds", "AttackTechniqueSeedGroup": "pyrit.models.seeds", diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 2624c812a2..22b8c05084 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -25,6 +25,12 @@ AttackRetrySummary, RegisteredScenario, RunScenarioRequest, + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioRunListItem, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, ScenarioRunSummary, ) from pyrit.models.catalog.target import TargetInstance @@ -35,6 +41,12 @@ "RegisteredInitializer": "pyrit.models.catalog.initializer", "RegisteredScenario": "pyrit.models.catalog.scenario", "RunScenarioRequest": "pyrit.models.catalog.scenario", + "ScenarioDatasetSizeCap": "pyrit.models.catalog.scenario", + "ScenarioDatasetSummary": "pyrit.models.catalog.scenario", + "ScenarioRunListItem": "pyrit.models.catalog.scenario", + "ScenarioRunSizeComponent": "pyrit.models.catalog.scenario", + "ScenarioRunSizeEstimate": "pyrit.models.catalog.scenario", + "ScenarioRunSizeEstimateRequest": "pyrit.models.catalog.scenario", "ScenarioRunSummary": "pyrit.models.catalog.scenario", "TargetInstance": "pyrit.models.catalog.target", } diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 488ccf8c78..e6f9bd3b75 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -14,9 +14,9 @@ """ from datetime import datetime -from typing import Any +from typing import Any, Literal -from pydantic import BaseModel, Field, field_validator +from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator from pyrit.models.parameter import Parameter from pyrit.models.results.scenario_result import ScenarioRunState @@ -39,21 +39,182 @@ DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) +def _validate_dataset_filter_mapping( + value: dict[str, list[str]] | None, +) -> dict[str, list[str]] | None: + """ + Validate dataset filter keys shared by launch and estimate requests. + + Returns: + dict[str, list[str]] | None: Validated filters. + + Raises: + ValueError: If a filter key is not supported. + """ + for key in value or {}: + if key not in DATASET_FILTERS: + raise ValueError(f"Unknown dataset filter '{key}'. Allowed: {', '.join(sorted(DATASET_FILTERS))}.") + return value + + +class ScenarioRunSizeComponent(BaseModel): + """One additive component of a default-run size estimate.""" + + label: str = Field(..., min_length=1) + count: int = Field(..., ge=0) + is_baseline: bool = False + note: str | None = None + + +class ScenarioDatasetSizeCap(BaseModel): + """One configured cap affecting a dataset or compound population.""" + + label: str = Field(..., min_length=1) + count: int = Field(..., ge=1) + configured_on: Literal["dataset", "configuration", "compound"] = "dataset" + dataset_name: str | None = None + + +class ScenarioDatasetSummary(BaseModel): + """Logical seed-group counts for one default dataset or synthesized population.""" + + name: str = Field(..., min_length=1) + kind: Literal["dataset", "synthesized"] = "dataset" + logical_seed_group_count: int = Field( + ..., + ge=0, + validation_alias=AliasChoices("logical_seed_group_count", "seed_group_count"), + ) + selected_seed_group_count: int = Field(..., ge=0) + configured_caps: list[ScenarioDatasetSizeCap] = Field(default_factory=list) + selection_note: str | None = None + + +class ScenarioRunSizeEstimate(BaseModel): + """ + Structured estimate of default planned scenario execution units. + + Counts use the same outer unit as ``ScenarioRunPlan``: one atomic-attack and + logical-seed-group pair. Retries and internal attack turns are excluded. + """ + + estimated_attack_count: int | None = Field(default=None, ge=0) + components: list[ScenarioRunSizeComponent] = Field(default_factory=list) + datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) + note: str | None = None + + @model_validator(mode="after") + def validate_estimated_attack_count(self) -> "ScenarioRunSizeEstimate": + """ + Ensure available estimates expose a complete additive total. + + Returns: + ScenarioRunSizeEstimate: The validated estimate. + + Raises: + ValueError: If an available estimate misstates its total. + """ + if self.estimated_attack_count is not None: + component_total = sum(component.count for component in self.components) + if component_total != self.estimated_attack_count: + raise ValueError( + f"Default-run estimate components total {component_total}, not {self.estimated_attack_count}" + ) + return self + + @classmethod + def unavailable(cls, *, note: str = "Default-run size estimate is unavailable.") -> "ScenarioRunSizeEstimate": + """ + Build an unavailable estimate without presenting a guessed total. + + Returns: + ScenarioRunSizeEstimate: An unavailable estimate. + """ + return cls(note=note) + + class RegisteredScenario(BaseModel): """Summary of a registered scenario.""" scenario_name: str = Field(..., description="Scenario name (e.g., 'foundry.red_team_agent')") scenario_type: str = Field(..., description="Scenario type identifier (e.g., 'RedTeamAgentScenario')") + scenario_version: int = Field(1, ge=1, description="Scenario definition version used for default metadata") description: str = Field(..., description="Human-readable description of the scenario") + description_markdown: str = Field( + "", + description=( + "Dedented Markdown source preserving the scenario docstring structure. " + "Clients must treat embedded HTML as untrusted text." + ), + ) default_technique: str = Field(..., description="Default technique name used when none specified") + default_techniques: list[str] = Field( + default_factory=list, + description="Ordered concrete techniques selected by the scenario's default technique policy", + ) aggregate_techniques: list[str] = Field( ..., description="Aggregate techniques that combine multiple attack approaches" ) + aggregate_technique_expansions: dict[str, list[str]] = Field( + default_factory=dict, + description="Concrete ordered technique expansion for every aggregate selector", + ) all_techniques: list[str] = Field(..., description="All available concrete technique names") default_datasets: list[str] = Field(..., description="Default dataset names used by the scenario") + default_dataset_summaries: list[ScenarioDatasetSummary] = Field( + default_factory=list, + description="Logical and effectively selected attack-group counts for the default configuration", + ) + baseline_policy: Literal["enabled", "disabled", "forbidden"] = Field( + "enabled", description="Whether baseline execution is enabled, disabled, or forbidden" + ) + include_baseline_by_default: bool = Field(True, description="Whether an omitted baseline flag includes it") supported_parameters: list[Parameter] = Field( default_factory=list, description="Scenario-declared custom parameters" ) + default_run_size: ScenarioRunSizeEstimate = Field( + default_factory=ScenarioRunSizeEstimate.unavailable, + description="Scenario-owned structured estimate of the default planned execution units", + ) + + +class ScenarioRunSizeEstimateRequest(BaseModel): + """Request-specific scenario run-size configuration.""" + + target_name: str | None = Field( + None, + description="Optional registered objective target used to resolve target-capability-dependent estimates", + ) + techniques: list[str] | None = Field( + None, description="Technique names to estimate (uses scenario default if omitted)" + ) + dataset_names: list[str] | None = Field( + None, description="Dataset names to estimate (uses scenario default if omitted)" + ) + max_dataset_size: int | None = Field(None, ge=1, description="Maximum selected logical seed groups") + dataset_filters: dict[str, list[str]] | None = Field( + None, + description="Dataset seed filters keyed by field. Accepted keys: harm_categories, data_types.", + ) + include_baseline: bool | None = Field( + None, + description="Override the scenario baseline default; forbidden scenarios reject true", + ) + scenario_params: dict[str, Any] | None = Field( + None, + description="Scenario-declared parameters such as Jailbreak template and attempt counts", + ) + + @field_validator("dataset_filters") + @classmethod + def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: + """ + Validate estimate dataset filters against the shared allow-list. + + Returns: + dict[str, list[str]] | None: Validated filters. + """ + return _validate_dataset_filter_mapping(value) class RunScenarioRequest(BaseModel): @@ -75,6 +236,9 @@ class RunScenarioRequest(BaseModel): ) max_concurrency: int = Field(10, ge=1, le=100, description="Maximum concurrent operations") max_retries: int = Field(0, ge=0, le=20, description="Maximum retry attempts on failure") + include_baseline: bool | None = Field( + None, description="Override the scenario baseline default; forbidden scenarios reject true" + ) labels: dict[str, str] | None = Field(None, description="Labels to attach to memory entries") scenario_params: dict[str, Any] | None = Field( None, @@ -99,21 +263,10 @@ def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[s """ Reject any dataset-filter key not in the exposed ``DATASET_FILTERS`` allow-list. - Runs for every request source (CLI and GUI), so the allow-list is enforced server-side. - - Args: - value (dict[str, list[str]] | None): The submitted dataset filters. - Returns: dict[str, list[str]] | None: The validated filters, unchanged. - - Raises: - ValueError: If any key is not present in ``DATASET_FILTERS``. """ - for key in value or {}: - if key not in DATASET_FILTERS: - raise ValueError(f"Unknown dataset filter '{key}'. Allowed: {', '.join(sorted(DATASET_FILTERS))}.") - return value + return _validate_dataset_filter_mapping(value) class AttackErrorSummary(BaseModel): @@ -141,6 +294,7 @@ class ScenarioRunSummary(BaseModel): scenario_result_id: str = Field(..., description="UUID of the ScenarioResult in memory") scenario_name: str = Field(..., description="Registry key of the scenario being run") + scenario_registry_name: str | None = Field(None, description="Requested scenario registry key when available") scenario_version: int = Field(0, ge=0, description="Version of the scenario") status: ScenarioRunState = Field(..., description="Current run status") created_at: datetime = Field(..., description="When the run was created") @@ -164,3 +318,21 @@ class ScenarioRunSummary(BaseModel): ) labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run") completed_at: datetime | None = Field(None, description="When the scenario finished") + + +class ScenarioRunListItem(BaseModel): + """Lightweight scenario run metadata returned by the history endpoint.""" + + scenario_result_id: str = Field(..., description="UUID of the ScenarioResult in memory") + scenario_name: str = Field(..., description="Registry key of the scenario being run") + scenario_registry_name: str | None = Field(None, description="Requested scenario registry key when available") + scenario_version: int = Field(0, ge=0, description="Version of the scenario") + status: ScenarioRunState = Field(..., description="Current run status") + created_at: datetime = Field(..., description="When the run was created") + updated_at: datetime = Field(..., description="When the run status last changed") + error: str | None = Field(None, description="Persisted run-level error message") + error_type: str | None = Field(None, description="Persisted run-level exception class") + techniques_used: list[str] = Field(default_factory=list, description="Planned technique display groups") + total_attacks: int | None = Field(None, ge=0, description="Number of planned execution units when known") + labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run") + completed_at: datetime | None = Field(None, description="When the scenario finished") diff --git a/pyrit/models/identifiers/__init__.py b/pyrit/models/identifiers/__init__.py index bfff6d3ca9..64b75b071e 100644 --- a/pyrit/models/identifiers/__init__.py +++ b/pyrit/models/identifiers/__init__.py @@ -43,7 +43,7 @@ from pyrit.models.identifiers.param_markers import Param, ParamMarker from pyrit.models.identifiers.scenario_identifier import ScenarioIdentifier from pyrit.models.identifiers.scorer_identifier import ScorerIdentifier - from pyrit.models.identifiers.seed_identifier import SeedIdentifier + from pyrit.models.identifiers.seed_identifier import SeedIdentifier, compute_seed_group_hash from pyrit.models.identifiers.target_identifier import TargetIdentifier _LAZY_EXPORTS: dict[str, str] = { @@ -74,6 +74,7 @@ "ScorerIdentifier": "pyrit.models.identifiers.scorer_identifier", "ScenarioIdentifier": "pyrit.models.identifiers.scenario_identifier", "SeedIdentifier": "pyrit.models.identifiers.seed_identifier", + "compute_seed_group_hash": "pyrit.models.identifiers.seed_identifier", "snake_case_to_class_name": "pyrit.models.identifiers.class_name_utils", "TARGET_EVAL_PARAM_FALLBACKS": "pyrit.models.identifiers.evaluation_identifier", "TARGET_EVAL_PARAMS": "pyrit.models.identifiers.evaluation_identifier", diff --git a/pyrit/models/identifiers/atomic_attack_identifier.py b/pyrit/models/identifiers/atomic_attack_identifier.py index c4a59cf37e..79ba50b186 100644 --- a/pyrit/models/identifiers/atomic_attack_identifier.py +++ b/pyrit/models/identifiers/atomic_attack_identifier.py @@ -22,7 +22,7 @@ from pyrit.models.identifiers.attack_technique_identifier import AttackTechniqueIdentifier from pyrit.models.identifiers.component_identifier import ComponentIdentifier from pyrit.models.identifiers.evaluation_markers import Evaluate -from pyrit.models.identifiers.seed_identifier import SeedIdentifier +from pyrit.models.identifiers.seed_identifier import SeedIdentifier, compute_seed_group_hash if TYPE_CHECKING: from pyrit.models.seeds.seed_group import SeedGroup @@ -109,3 +109,8 @@ def build( attack_technique=technique, seed_identifiers=seed_identifiers, ) + + @property + def logical_seed_group_id(self) -> str: + """The logical seed-group ID represented by the ordered seed identifiers.""" + return compute_seed_group_hash(self.seed_identifiers) diff --git a/pyrit/models/identifiers/seed_identifier.py b/pyrit/models/identifiers/seed_identifier.py index 2372164662..acb22bedd1 100644 --- a/pyrit/models/identifiers/seed_identifier.py +++ b/pyrit/models/identifiers/seed_identifier.py @@ -7,11 +7,13 @@ from typing import TYPE_CHECKING, Annotated -from pyrit.models.identifiers.component_identifier import ComponentIdentifier +from pyrit.models.identifiers.component_identifier import ComponentIdentifier, config_hash from pyrit.models.identifiers.evaluation_markers import Evaluate from pyrit.models.literals import PromptDataType # noqa: TC001 (runtime-required by Pydantic field annotations) if TYPE_CHECKING: + from collections.abc import Sequence + from pyrit.models.seeds.seed import Seed @@ -58,3 +60,15 @@ def from_seed(cls, seed: Seed) -> SeedIdentifier: dataset_name=seed.dataset_name, is_general_technique=seed.is_general_technique, ) + + +def compute_seed_group_hash(seed_identifiers: Sequence[SeedIdentifier]) -> str: + """Return the deterministic hash of ordered canonical seed identifiers.""" + return config_hash( + { + "seed_identifiers": [ + seed_identifier.model_dump(exclude={"hash", "eval_hash", "pyrit_version"}) + for seed_identifier in seed_identifiers + ] + } + ) diff --git a/pyrit/models/results/scenario_result.py b/pyrit/models/results/scenario_result.py index 793d8ce33f..bddfb26f20 100644 --- a/pyrit/models/results/scenario_result.py +++ b/pyrit/models/results/scenario_result.py @@ -94,10 +94,8 @@ class ScenarioResult(BaseModel): error_type: str | None = None #: IDs of attack results that errored during the scenario run. error_attack_result_ids: list[str] = Field(default_factory=list) - #: Free-form JSON metadata persisted with the scenario result. Currently used to record - #: ``objective_hashes`` — the objective ``sha256`` set chosen on the first run, replayed - #: on resume so a fresh ``random.sample`` can't silently change which objectives the - #: scenario operates on. Keys are not part of any public contract and may evolve. + #: Free-form JSON metadata persisted with the scenario result. Stores the normalized + #: run plan and, for sampled runs, ``objective_hashes`` used to replay the original subset. metadata: dict[str, Any] = Field(default_factory=dict) @model_validator(mode="before") diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py new file mode 100644 index 0000000000..89fc6888c3 --- /dev/null +++ b/pyrit/models/scenario_progress.py @@ -0,0 +1,133 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Canonical models for durable scenario run plans and incremental progress.""" + +from datetime import datetime +from typing import Any, Literal + +from pydantic import AwareDatetime, BaseModel, Field, model_validator + +from pyrit.models.identifiers.atomic_attack_identifier import AtomicAttackIdentifier +from pyrit.models.results.attack_result import AttackOutcome +from pyrit.models.results.scenario_result import ScenarioRunState +from pyrit.models.retry_event import RetryEvent + +SCENARIO_RUN_PLAN_METADATA_KEY = "run_plan" +SCENARIO_RUN_PLAN_VERSION = 1 + + +class ScenarioRunPlanSeedGroup(BaseModel): + """A de-duplicated logical seed group in a scenario run plan.""" + + id: str + objective_sha256: str + objective: str + + +class ScenarioRunPlanAtomicGroup(BaseModel): + """A planned atomic-attack group and its ordered units of work.""" + + id: str + atomic_attack_name: str + display_group: str + technique_eval_hash: str + seed_group_ids: list[str] + + +class ScenarioRunPlan(BaseModel): + """Versioned normalized execution plan persisted in ScenarioResult metadata.""" + + version: Literal[1] = 1 + scenario_registry_name: str | None = None + atomic_groups: list[ScenarioRunPlanAtomicGroup] + seed_groups: list[ScenarioRunPlanSeedGroup] + + @model_validator(mode="after") + def _validate_normalized_plan(self) -> "ScenarioRunPlan": + """ + Reject ambiguous IDs and invalid normalized references. + + Returns: + ScenarioRunPlan: The validated normalized plan. + + Raises: + ValueError: If IDs are duplicated or a group references an unknown seed. + """ + atomic_group_ids = [group.id for group in self.atomic_groups] + if len(atomic_group_ids) != len(set(atomic_group_ids)): + raise ValueError("Scenario run plan contains duplicate atomic group IDs.") + + seed_group_ids = [seed.id for seed in self.seed_groups] + if len(seed_group_ids) != len(set(seed_group_ids)): + raise ValueError("Scenario run plan contains duplicate seed group IDs.") + + known_seed_group_ids = set(seed_group_ids) + for group in self.atomic_groups: + if len(group.seed_group_ids) != len(set(group.seed_group_ids)): + raise ValueError(f"Scenario run plan atomic group '{group.id}' contains duplicate seed group IDs.") + missing_seed_group_ids = set(group.seed_group_ids) - known_seed_group_ids + if missing_seed_group_ids: + raise ValueError( + f"Scenario run plan atomic group '{group.id}' references unknown seed group IDs: " + f"{', '.join(sorted(missing_seed_group_ids))}." + ) + return self + + +class ScenarioProgressHeader(BaseModel): + """Compact persisted run header returned by the progress endpoint.""" + + scenario_result_id: str + scenario_name: str + scenario_registry_name: str | None = None + scenario_version: int + status: ScenarioRunState + created_at: datetime + completed_at: datetime | None = None + + +class ScenarioProgressResult(BaseModel): + """One persisted attack attempt in ascending progress order.""" + + attack_result_id: str + atomic_group_id: str + atomic_attack_name: str + seed_group_id: str + outcome: AttackOutcome + execution_time_ms: int + timestamp: AwareDatetime + total_retries: int = 0 + retries: list[RetryEvent] = Field(default_factory=list) + error_type: str | None = None + error_message: str | None = None + + +class ScenarioRunProgress(BaseModel): + """Incremental scenario progress response.""" + + run: ScenarioProgressHeader + plan: ScenarioRunPlan | None = None + reset: bool = False + active_atomic_group_ids: list[str] = Field(default_factory=list) + results: list[ScenarioProgressResult] = Field(default_factory=list) + next_cursor: str | None = None + has_more: bool = False + plan_complete: bool + + +class ScenarioAttackResultDelta(BaseModel): + """Lightweight memory projection used to map one scenario progress delta.""" + + attack_result_id: str + objective: str + objective_sha256: str | None = None + atomic_attack_identifier: AtomicAttackIdentifier | None = None + outcome: AttackOutcome + execution_time_ms: int + timestamp: AwareDatetime + retry_events: list[RetryEvent] = Field(default_factory=list) + total_retries: int = 0 + error_type: str | None = None + error_message: str | None = None + attribution_data: dict[str, Any] = Field(default_factory=dict) diff --git a/pyrit/models/seeds/attack_seed_group.py b/pyrit/models/seeds/attack_seed_group.py index 99d325dd48..191b9aa7f9 100644 --- a/pyrit/models/seeds/attack_seed_group.py +++ b/pyrit/models/seeds/attack_seed_group.py @@ -12,6 +12,7 @@ import copy from typing import TYPE_CHECKING +from pyrit.models.identifiers import SeedIdentifier, compute_seed_group_hash from pyrit.models.seeds.seed_group import SeedGroup from pyrit.models.seeds.seed_objective import SeedObjective from pyrit.models.seeds.seed_prompt import SeedPrompt @@ -86,6 +87,17 @@ def objective(self) -> SeedObjective: raise ValueError("AttackSeedGroup should always have an objective") return obj + @property + def logical_id(self) -> str: + """ + The deterministic identity of this original logical seed group. + + The ordered seed identifiers contain behavioral seed values but omit + random ``prompt_group_id`` values. Call this before technique seeds are + merged so the same ID is recoverable from an enriched attack result. + """ + return compute_seed_group_hash([SeedIdentifier.from_seed(seed) for seed in self.seeds]) + def is_compatible_with_technique(self, *, technique: AttackTechniqueSeedGroup) -> bool: """ Check whether this seed group can be merged with the given technique. diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index caf6bba721..bfe8d2574d 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -14,10 +14,11 @@ from __future__ import annotations +import asyncio from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal -from pyrit.models import class_name_to_snake_case +from pyrit.models import ScenarioRunSizeEstimate, class_name_to_snake_case from pyrit.models.identifiers.scenario_identifier import ScenarioIdentifier from pyrit.registry.registry import ParamBagRegistry from pyrit.registry.registry_metadata import RegistryMetadata @@ -38,21 +39,36 @@ class ScenarioMetadata(RegistryMetadata): Use get_class() to get the actual class. """ + scenario_version: int = field(kw_only=True, default=1) + # The default technique name (e.g., "single_turn") default_technique: str = field(kw_only=True) + # Ordered concrete techniques selected by the default technique policy. + default_techniques: tuple[str, ...] = field(kw_only=True, default=()) + + # Dedented class docstring with Markdown structure preserved. + description_markdown: str = field(kw_only=True, default="") + # All available technique names for this scenario. all_techniques: tuple[str, ...] = field(kw_only=True) # Aggregate techniques that combine multiple attack approaches. aggregate_techniques: tuple[str, ...] = field(kw_only=True) + # Ordered aggregate selector -> concrete technique expansions. + aggregate_technique_expansions: tuple[tuple[str, tuple[str, ...]], ...] = field(kw_only=True, default=()) + # Default dataset names used by this scenario. default_datasets: tuple[str, ...] = field(kw_only=True) # Scenario-declared custom parameters. supported_parameters: tuple[Parameter, ...] = field(kw_only=True, default=()) + baseline_policy: Literal["enabled", "disabled", "forbidden"] = field(kw_only=True, default="enabled") + + include_baseline_by_default: bool = field(kw_only=True, default=True) + class ScenarioRegistry(ParamBagRegistry["Scenario", ScenarioMetadata]): """ @@ -136,6 +152,7 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: TypeError: If ``cls()`` cannot be called with no arguments. """ description = RegistryMetadata.description_from_docstring(cls, fallback="No description available") + description_markdown = RegistryMetadata.markdown_from_docstring(cls, fallback=description) supported_parameters = tuple(cls.supported_parameters()) @@ -152,8 +169,18 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: technique_class = instance._technique_class default_technique_value = instance._default_technique.value + default_techniques = tuple( + technique.value for technique in instance._resolve_scenario_techniques(scenario_techniques=None) + ) all_techniques = tuple(s.value for s in technique_class.get_all_techniques()) aggregate_techniques = tuple(s.value for s in technique_class.get_aggregate_techniques()) + aggregate_technique_expansions = tuple( + ( + aggregate.value, + tuple(technique.value for technique in technique_class.expand({aggregate})), + ) + for aggregate in technique_class.get_aggregate_techniques() + ) default_datasets = tuple(instance._default_dataset_config.dataset_names) return ScenarioMetadata( @@ -161,13 +188,45 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: class_module=cls.__module__, class_description=description, registry_name=name, + scenario_version=instance._version, default_technique=default_technique_value, + default_techniques=default_techniques, + description_markdown=description_markdown, all_techniques=all_techniques, aggregate_techniques=aggregate_techniques, + aggregate_technique_expansions=aggregate_technique_expansions, default_datasets=default_datasets, supported_parameters=supported_parameters, + baseline_policy=instance.BASELINE_ATTACK_POLICY.value, + include_baseline_by_default=instance.BASELINE_ATTACK_POLICY.value == "enabled", ) + async def create_and_estimate_async( + self, + *, + name: str, + scenario_params: dict[str, Any] | None = None, + target_is_configured: bool = False, + **estimate_kwargs: Any, + ) -> ScenarioRunSizeEstimate: + """ + Build, parameterize, and estimate a scenario without initializing a run. + + Args: + name: Registered scenario name. + scenario_params: Scenario-declared parameter values. + target_is_configured: Whether the estimate has a concrete objective target. + **estimate_kwargs: Common resolved values such as techniques, dataset + configuration, baseline choice, and an optional objective target. + + Returns: + ScenarioRunSizeEstimate: Structured configured-run estimate. + """ + scenario = await asyncio.to_thread(self.create_instance, name) + scenario.set_scenario_registry_name(scenario_registry_name=name) + scenario.set_params_from_args(args={**(scenario_params or {}), **estimate_kwargs}) + return await scenario.get_run_size_estimate_async(target_is_configured=target_is_configured) + async def create_and_initialize_async( self, name: str, @@ -215,5 +274,6 @@ async def create_and_initialize_async( merged_args = {**(scenario_params or {}), **initialize_kwargs} scenario = self._create_and_configure(name, params=merged_args, constructor_kwargs=constructor_kwargs) + scenario.set_scenario_registry_name(scenario_registry_name=name) await scenario.initialize_async() return scenario diff --git a/pyrit/registry/registry_metadata.py b/pyrit/registry/registry_metadata.py index ec472dc96d..24d21fd5fb 100644 --- a/pyrit/registry/registry_metadata.py +++ b/pyrit/registry/registry_metadata.py @@ -66,6 +66,20 @@ def description_from_docstring(cls: type, *, fallback: str = "") -> str: cleaned = " ".join(doc.split()) return cleaned or fallback + @staticmethod + def markdown_from_docstring(cls: type, *, fallback: str = "") -> str: + """ + Extract a dedented description while preserving Markdown structure. + + Returns: + str: The dedented docstring or the fallback value. + """ + doc = cls.__doc__ + if not doc: + return fallback + cleaned = inspect.cleandoc(doc) + return cleaned or fallback + @staticmethod def summary_from_docstring(cls: type) -> str: """ diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index 0ad434542a..0785cca63b 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -22,7 +22,13 @@ from pyrit.executor.attack import AttackExecutor, AttackExecutorResult from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.memory import CentralMemory -from pyrit.models import AtomicAttackEvaluationIdentifier, AtomicAttackIdentifier, AttackResult, AttackSeedGroup +from pyrit.models import ( + AtomicAttackEvaluationIdentifier, + AtomicAttackIdentifier, + AttackResult, + AttackSeedGroup, + config_hash, +) if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget @@ -191,6 +197,16 @@ def technique_eval_hash(self) -> str: ) return AtomicAttackEvaluationIdentifier(composite).eval_hash + @property + def logical_group_id(self) -> str: + """The stable identity of this planned atomic-attack group.""" + return config_hash( + { + "atomic_attack_name": self.atomic_attack_name, + "technique_eval_hash": self.technique_eval_hash, + } + ) + @property def objectives(self) -> list[str]: """ @@ -324,13 +340,17 @@ async def run_async( # a Scenario. The same attribution object is stamped on every # per-task AttackContext; per-task identity is reconstructed from # the row's own objective_sha256 (no positional state required). - attribution: AttackResultAttribution | None = None + attributions: list[AttackResultAttribution] | None = None if self._scenario_result_id is not None: - attribution = AttackResultAttribution( - parent_id=self._scenario_result_id, - parent_collection=self.atomic_attack_name, - parent_eval_hash=self.technique_eval_hash, - ) + attributions = [ + AttackResultAttribution( + parent_id=self._scenario_result_id, + parent_collection=self.atomic_attack_name, + parent_eval_hash=self.technique_eval_hash, + seed_group_id=seed_group.logical_id, + ) + for seed_group in self._seed_groups + ] untyped_results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, @@ -339,7 +359,7 @@ async def run_async( objective_scorer=self._objective_scorer, memory_labels=self._memory_labels, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, + attributions=attributions, **self._attack_execute_params, ) completed_results: list[AttackResult] = [] diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 990957386c..53151a7c89 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -84,6 +84,7 @@ def __init__( adversarial_seed_prompt: SeedPrompt | str | None = None, seed_technique: AttackTechniqueSeedGroup | None = None, uses_adversarial: bool | None = None, + supports_additional_request_converters: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -121,6 +122,9 @@ def __init__( chat during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Authors can override the derivation explicitly. + supports_additional_request_converters: Whether callers may safely + append request converters to this technique. This is an explicit + semantic opt-in, not merely constructor-signature detection. scorer_override_policy: What to do when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to WARN. @@ -145,11 +149,13 @@ class constructor signature and seed-technique shape. adversarial_system_prompt is not None or adversarial_seed_prompt is not None ) self._seed_technique = seed_technique + self._supports_additional_request_converters = supports_additional_request_converters self._scorer_override_policy = scorer_override_policy self._uses_adversarial = uses_adversarial if uses_adversarial is not None else self._derive_uses_adversarial() self._validate_kwargs() + self._validate_converter_composition() self._validate_adversarial_flags() @classmethod @@ -168,6 +174,7 @@ def with_simulated_conversation( attack_kwargs: dict[str, Any] | None = None, adversarial_chat: PromptTarget | None = None, uses_adversarial: bool | None = None, + supports_additional_request_converters: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -219,6 +226,9 @@ def with_simulated_conversation( during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Forwarded to the factory constructor. + supports_additional_request_converters: Whether callers may safely + append request converters to this technique. Forwarded to the + factory constructor. scorer_override_policy: Policy applied when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to ``WARN``. Forwarded to the factory @@ -279,6 +289,7 @@ def with_simulated_conversation( adversarial_chat=adversarial_chat, seed_technique=seed_technique, uses_adversarial=uses_adversarial, + supports_additional_request_converters=supports_additional_request_converters, scorer_override_policy=scorer_override_policy, ) @@ -314,6 +325,23 @@ def _validate_adversarial_flags(self) -> None: f"should not have one wired." ) + def _validate_converter_composition(self) -> None: + """ + Validate that an opt-in factory can receive additive request converters. + + Raises: + ValueError: If composition is enabled but the attack constructor does + not accept ``attack_converter_config``. + """ + if ( + self._supports_additional_request_converters + and "attack_converter_config" not in self._get_accepted_params() + ): + raise ValueError( + f"Factory '{self._name}' declares supports_additional_request_converters=True, " + f"but {self._attack_class.__name__} does not accept 'attack_converter_config'." + ) + def _validate_kwargs(self) -> None: """ Validate that all kwargs are valid parameters for the attack class constructor. @@ -484,6 +512,11 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial + @property + def supports_additional_request_converters(self) -> bool: + """Whether callers may safely append request converters to this technique.""" + return self._supports_additional_request_converters + @property def scoring_config_type(self) -> type | None: """The required ``attack_scoring_config`` subtype, or ``None`` if any config is accepted.""" @@ -828,6 +861,7 @@ def _build_identifier(self) -> ComponentIdentifier: "attack_class": self._attack_class.__name__, "kwargs": kwargs_for_id, "uses_adversarial": self._uses_adversarial, + "supports_additional_request_converters": self._supports_additional_request_converters, } if self._technique_tags: params["technique_tags"] = list(self._technique_tags) diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index ae8aaa3580..4f43d5fdc6 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -27,17 +27,20 @@ from __future__ import annotations +import asyncio import random +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast from pyrit.memory import CentralMemory from pyrit.models import AttackSeedGroup, Seed, SeedGroup, group_seeds_into_attack_groups if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Callable, Iterator, Sequence from pyrit.memory import MemoryInterface @@ -48,6 +51,17 @@ # Internal helper TypeVar for size-capping any homogeneous list. _ItemT = TypeVar("_ItemT") +_AUTO_FETCH_ALLOWED: ContextVar[bool] = ContextVar("dataset_auto_fetch_allowed", default=True) + + +@contextmanager +def read_only_dataset_resolution() -> Iterator[None]: + """Disable dataset auto-fetch persistence within the current async context.""" + token = _AUTO_FETCH_ALLOWED.set(False) + try: + yield + finally: + _AUTO_FETCH_ALLOWED.reset(token) class DatasetSourceKind(Enum): @@ -392,6 +406,28 @@ def filters(self) -> dict[str, list[str]]: """ return dict(self._filters) + @property + def has_size_cap(self) -> bool: + """Whether this configuration applies a logical-group selection cap.""" + return self.max_dataset_size is not None + + def size_caps_by_dataset(self) -> dict[str, list[tuple[str, int, Literal["dataset", "configuration", "compound"]]]]: + """ + Describe configured caps for each named dataset or inline source. + + Returns: + dict[str, list[tuple[str, int, Literal]]]: Source name to ordered + ``(cap label, count, provenance)`` entries. + """ + if self.max_dataset_size is None: + return {} + names = self.dataset_names or [INLINE_DATASET_NAME] + if len(names) == 1: + cap = ("per-dataset cap", self.max_dataset_size, "dataset") + else: + cap = ("combined configuration cap", self.max_dataset_size, "configuration") + return {name: [cap] for name in names} + @property def _get_seeds_filters(self) -> dict[str, Any]: """ @@ -454,25 +490,42 @@ async def _collect_seeds_for_dataset_async(self, *, dataset_name: str) -> list[S DatasetConstraintError: If the dataset yields no seeds even after auto-fetch, or if auto-fetch itself fails (the provider error is chained as the cause). """ - found = list(self._memory.get_seeds(dataset_name=dataset_name, **self._get_seeds_filters)) - if not found and self._auto_fetch: + found = list( + await asyncio.to_thread( + self._memory.get_seeds, + dataset_name=dataset_name, + **self._get_seeds_filters, + ) + ) + auto_fetch_allowed = self._auto_fetch and _AUTO_FETCH_ALLOWED.get() + if not found and auto_fetch_allowed: try: await self._fetch_dataset_async(dataset_name=dataset_name) except Exception as exc: raise DatasetConstraintError( f"Dataset '{dataset_name}' could not be loaded: auto-fetch from the registered provider failed." ) from exc - found = list(self._memory.get_seeds(dataset_name=dataset_name, **self._get_seeds_filters)) + found = list( + await asyncio.to_thread( + self._memory.get_seeds, + dataset_name=dataset_name, + **self._get_seeds_filters, + ) + ) if not found: - if self._filters and self._memory.get_seeds(dataset_name=dataset_name): + unfiltered = ( + await asyncio.to_thread(self._memory.get_seeds, dataset_name=dataset_name) if self._filters else [] + ) + if unfiltered: raise DatasetConstraintError( f"Dataset '{dataset_name}' has seeds, but none match the configured filters {self._filters}." ) - hint = ( - "auto-fetch from the registered provider did not populate it" - if self._auto_fetch - else "auto_fetch is disabled" - ) + if auto_fetch_allowed: + hint = "auto-fetch from the registered provider did not populate it" + elif self._auto_fetch: + hint = "auto_fetch is disabled for read-only resolution" + else: + hint = "auto_fetch is disabled" raise DatasetConstraintError( f"Dataset '{dataset_name}' could not be loaded: no seeds found in memory and {hint}." ) @@ -823,6 +876,27 @@ def source_kind(self) -> DatasetSourceKind: return DatasetSourceKind.INLINE return DatasetSourceKind.MEMORY + @property + def has_size_cap(self) -> bool: + """Whether the compound or any child applies a logical-group cap.""" + return self.max_dataset_size is not None or any(child.has_size_cap for child in self._configurations) + + def size_caps_by_dataset(self) -> dict[str, list[tuple[str, int, Literal["dataset", "configuration", "compound"]]]]: + """ + Describe child and combined caps for every contributed dataset. + + Returns: + dict[str, list[tuple[str, int, Literal]]]: Ordered cap labels, counts, and provenance by source. + """ + caps: dict[str, list[tuple[str, int, Literal["dataset", "configuration", "compound"]]]] = {} + for child in self._configurations: + for name, child_caps in child.size_caps_by_dataset().items(): + caps.setdefault(name, []).extend(child_caps) + if self.max_dataset_size is not None: + for name in self.dataset_names or [INLINE_DATASET_NAME]: + caps.setdefault(name, []).append(("combined compound cap", self.max_dataset_size, "compound")) + return caps + def update_filters(self, *, filters: dict[str, list[str]]) -> None: """ Merge filters into the compound and propagate them to every child configuration. diff --git a/pyrit/scenario/core/matrix_atomic_attack_builder.py b/pyrit/scenario/core/matrix_atomic_attack_builder.py index 40a35ca474..148b098be3 100644 --- a/pyrit/scenario/core/matrix_atomic_attack_builder.py +++ b/pyrit/scenario/core/matrix_atomic_attack_builder.py @@ -36,6 +36,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.scenario_context import ScenarioContext + from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.score import Scorer from pyrit.score.true_false.true_false_scorer import TrueFalseScorer @@ -155,6 +156,23 @@ def resolve_technique_factories( dict[str, AttackTechniqueFactory]: Mapping of technique name to factory, ordered by the selected techniques. """ + return resolve_technique_factories_for_techniques( + scenario_techniques=context.scenario_techniques, + extra_factories=extra_factories, + ) + + +def resolve_technique_factories_for_techniques( + *, + scenario_techniques: Sequence[ScenarioTechnique], + extra_factories: dict[str, AttackTechniqueFactory] | None = None, +) -> dict[str, AttackTechniqueFactory]: + """ + Resolve selected concrete techniques to their canonical factories. + + Returns: + dict[str, AttackTechniqueFactory]: Selected factories in technique order. + """ from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry all_factories = dict(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise()) @@ -162,11 +180,30 @@ def resolve_technique_factories( all_factories.update(extra_factories) return { technique.value: all_factories[technique.value] - for technique in context.scenario_techniques + for technique in scenario_techniques if technique.value in all_factories } +def filter_compatible_seed_groups( + *, + factory: AttackTechniqueFactory, + seed_groups: Sequence[AttackSeedGroup], +) -> list[AttackSeedGroup]: + """ + Apply the matrix builder's seed-technique compatibility rule. + + Returns: + list[AttackSeedGroup]: Compatible groups in source order. + """ + if factory.seed_technique is None: + return list(seed_groups) + return AttackSeedGroup.filter_compatible( + seed_groups=list(seed_groups), + technique=factory.seed_technique, + ) + + def build_matrix_atomic_attacks( *, context: ScenarioContext, @@ -404,13 +441,7 @@ def _filter_compatible_groups( list[AttackSeedGroup] | None: The compatible groups, or ``None`` when the ``(technique, dataset)`` pair has no compatible groups and should be skipped. """ - if factory.seed_technique is None: - return list(seed_groups) - - compatible_groups = AttackSeedGroup.filter_compatible( - seed_groups=seed_groups, - technique=factory.seed_technique, - ) + compatible_groups = filter_compatible_seed_groups(factory=factory, seed_groups=seed_groups) skipped = len(seed_groups) - len(compatible_groups) if skipped: logger.info( diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 9385c8c97b..9eebbc86b8 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -33,13 +33,22 @@ from pyrit.memory import CentralMemory from pyrit.memory.memory_models import ScenarioResultEntry from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, AttackOutcome, AttackResult, AttackSeedGroup, + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, ScenarioEvaluationIdentifier, ScenarioIdentifier, ScenarioResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, ScenarioRunState, + config_hash, ) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference from pyrit.prompt_target import PromptTarget @@ -47,7 +56,7 @@ from pyrit.registry import ScorerRegistry from pyrit.registry.resolution import resolve_declared_params, resolve_reference_value from pyrit.scenario.core.atomic_attack import AtomicAttack -from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration +from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, read_only_dataset_resolution from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_target_defaults import get_default_scorer_target from pyrit.scenario.core.scenario_technique import ScenarioTechnique @@ -65,6 +74,7 @@ if TYPE_CHECKING: from pyrit.converter import Converter from pyrit.models import ComponentIdentifier + from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory logger = logging.getLogger(__name__) @@ -123,6 +133,9 @@ class Scenario(ABC): #: caller-supplied ``include_baseline=True`` raises ``ValueError``. BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Enabled + #: Whether the default estimator must mirror matrix-builder seed compatibility. + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: Any) -> None: """ Enforce the keyword-only constructor contract on subclasses. @@ -202,6 +215,8 @@ def __init__( # These will be set in initialize_async self._objective_target: PromptTarget | None = None self._objective_target_identifier: ComponentIdentifier | None = None + self._estimate_target_is_configured = False + self._estimate_has_binding_size_cap = False self._memory_labels: dict[str, str] = {} self._max_concurrency: int | None = None self._max_retries: int = 0 @@ -218,6 +233,8 @@ def __init__( self._memory = CentralMemory.get_memory_instance() self._atomic_attacks: list[AtomicAttack] = [] self._scenario_result_id: str | None = str(scenario_result_id) if scenario_result_id else None + self._scenario_registry_name: str | None = None + self._active_atomic_groups: dict[str, str] = {} # Store prepared techniques for use in _build_atomic_attacks_async self._scenario_techniques: list[ScenarioTechnique] = [] @@ -250,6 +267,20 @@ def atomic_attack_count(self) -> int: """The number of atomic attacks in this scenario.""" return len(self._atomic_attacks) + @property + def active_atomic_group_ids(self) -> frozenset[str]: + """The stable IDs of atomic groups currently executing.""" + return frozenset(self._active_atomic_groups) + + @property + def active_atomic_group_names(self) -> tuple[str, ...]: + """The names of atomic groups currently executing.""" + return tuple(self._active_atomic_groups.values()) + + def set_scenario_registry_name(self, *, scenario_registry_name: str) -> None: + """Record the requested registry name for durable run-plan attribution.""" + self._scenario_registry_name = scenario_registry_name + @classmethod def _common_scenario_parameters(cls) -> list[Parameter]: """ @@ -526,63 +557,209 @@ def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[Scen return self._technique_class.resolve(scenario_techniques, default=self._default_technique) @final - async def initialize_async(self) -> None: + async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: """ - Initialize the scenario by populating self._atomic_attacks and creating the ScenarioResult. + Estimate the scenario's default planned execution units without starting a run. - All run inputs are read from the parameter bag (``self.params``), which is populated by - ``set_params_from_args`` from the merged CLI / config / programmatic arguments. Callers - fill the bag then initialize: + This resolves declared parameter defaults before delegating to the same + configured estimate path used by request-specific previews. - .. code-block:: python + Returns: + ScenarioRunSizeEstimate: Structured default-run estimate. + """ + self.set_params_from_args(args={}) + return await self.get_run_size_estimate_async(target_is_configured=False) - scenario.set_params_from_args(args={"objective_target": target, "max_concurrency": 8}) - await scenario.initialize_async() + @final + async def get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: + """ + Estimate the currently configured run without creating or persisting it. - This method allows scenarios to be initialized with atomic attacks after construction, - which is useful when atomic attacks require async operations to be built. + ``set_params_from_args`` should be called first for a request-specific + estimate. Omitted values use the same declared defaults, aggregate + expansion, dataset selection, and baseline policy as ``initialize_async``. - If a scenario_result_id was provided in __init__, this method will check if it exists - in memory and validate that the stored scenario matches the current configuration. - If it matches, the scenario will resume from prior progress. If it doesn't match or - doesn't exist, a new scenario result will be created. + Returns: + ScenarioRunSizeEstimate: Structured configured-run estimate. - The common run inputs read from the bag are ``objective_target`` (a ``PromptTarget`` - instance or a registered target name resolved against ``TargetRegistry``), - ``scenario_techniques``, ``technique_converters``, ``dataset_config``, - ``max_concurrency``, ``max_retries``, ``memory_labels``, and ``include_baseline`` - (see ``_common_scenario_parameters``). A subclass that removes a common input via - ``supported_parameters`` falls back to that input's default here. + Raises: + ValueError: If target certainty is asserted without a resolved target. + """ + self._resolve_runtime_configuration(require_objective_target=False) + if target_is_configured and self._objective_target is None: + raise ValueError("target_is_configured requires a resolved objective_target") + self._estimate_target_is_configured = self._objective_target is not None + return await self._estimate_run_size_async() + + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate a standard technique-by-seed-group scenario. + + Subclasses override this hook when their outer execution shape adds axes, + synthesizes technique-specific populations, or selects techniques adaptively. + + Returns: + ScenarioRunSizeEstimate: Exact default sweep and baseline count. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + seed_group_count = sum(len(groups) for groups in selected_groups.values()) + components = self._build_technique_size_components( + selected_groups=selected_groups, + seed_group_count=seed_group_count, + ) + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=seed_group_count, + is_baseline=True, + note="One unmodified prompt-sending unit per selected seed group.", + ) + ) + + estimated_attack_count = ( + None + if self.RUN_SIZE_USES_FACTORY_COMPATIBILITY and self._estimate_has_binding_size_cap + else sum(component.count for component in components) + ) + note = "Counts planned outer execution units; retries and internal attack turns are excluded." + if estimated_attack_count is None: + note += " A binding randomized dataset cap may select a different compatibility mix at launch." + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, + components=components, + datasets=datasets, + note=note, + ) + + def _build_technique_size_components( + self, + *, + selected_groups: dict[str, list[AttackSeedGroup]], + seed_group_count: int, + ) -> list[ScenarioRunSizeComponent]: + """ + Build the standard sweep, applying matrix-builder compatibility when declared. + + Returns: + list[ScenarioRunSizeComponent]: Additive technique components. + """ + if not self.RUN_SIZE_USES_FACTORY_COMPATIBILITY: + technique_count = len(self._scenario_techniques) + return [ + ScenarioRunSizeComponent( + label="Default technique sweep", + count=seed_group_count * technique_count, + ) + ] + + from pyrit.scenario.core.matrix_atomic_attack_builder import ( + filter_compatible_seed_groups, + resolve_technique_factories_for_techniques, + ) + + factories = resolve_technique_factories_for_techniques( + scenario_techniques=self._scenario_techniques, + extra_factories=self._get_run_size_extra_factories(), + ) + components: list[ScenarioRunSizeComponent] = [] + for technique in self._scenario_techniques: + factory = factories.get(technique.value) + if factory is None: + continue + compatible_count = sum( + len(filter_compatible_seed_groups(factory=factory, seed_groups=groups)) + for groups in selected_groups.values() + ) + components.append( + ScenarioRunSizeComponent( + label=technique.value, + count=compatible_count, + ) + ) + return components + + def _get_run_size_extra_factories(self) -> dict[str, "AttackTechniqueFactory"] | None: + """Return scenario-local factories used by compatibility-aware sizing.""" + return None + + async def _resolve_dataset_groups_for_estimate_async( + self, + ) -> tuple[dict[str, list[AttackSeedGroup]], list[ScenarioDatasetSummary]]: + """ + Resolve full and effectively selected logical groups for configured datasets. + + Returns: + tuple: Selected groups keyed by population and their catalog summaries. + """ + configured_dataset = self._dataset_config + with read_only_dataset_resolution(): + self._dataset_config = configured_dataset + full_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=False) + self._dataset_config = configured_dataset + selected_groups = await self._resolve_seed_groups_by_dataset_async(apply_sampling=True) + + configured_caps = self._dataset_config.size_caps_by_dataset() + datasets: list[ScenarioDatasetSummary] = [] + for name in dict.fromkeys([*full_groups, *selected_groups]): + logical_count = len(full_groups.get(name, [])) + selected_count = len(selected_groups.get(name, [])) + selection_note = None + if selected_count != logical_count: + selection_note = f"The default selection uses {selected_count} of {logical_count} logical seed groups." + datasets.append( + ScenarioDatasetSummary( + name=name, + logical_seed_group_count=logical_count, + selected_seed_group_count=selected_count, + configured_caps=[ + ScenarioDatasetSizeCap( + label=label, + count=count, + configured_on=configured_on, + dataset_name=name, + ) + for label, count, configured_on in configured_caps.get(name, []) + ], + selection_note=selection_note, + ) + ) + self._estimate_has_binding_size_cap = bool(configured_caps) and sum( + dataset.selected_seed_group_count for dataset in datasets + ) < sum(dataset.logical_seed_group_count for dataset in datasets) + return selected_groups, datasets + + def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> None: + """ + Resolve the common parameter bag shared by initialization and estimation. + + Args: + require_objective_target: Whether an omitted objective target is an error. Raises: - ValueError: If ``objective_target`` is declared but not resolvable (neither supplied - nor registered as a default), if a supplied target name is not registered in - ``TargetRegistry``, or if ``include_baseline=True`` is set for a scenario whose - ``BASELINE_ATTACK_POLICY`` is ``Forbidden``. + ValueError: If required target or baseline constraints are not satisfied. """ - # Resolve declared parameters through the single registry-owned path, materializing - # defaults for programmatic callers that skipped an explicit set_params_from_args. - # Guarded so the bag is resolved exactly once: the registry/CLI flows already call - # set_params_from_args, so this only runs for a direct construct-then-initialize caller - # and avoids a surprising re-validation / self-mutation of an already-resolved bag. if not self._params_resolved: self.set_params_from_args(args=self.params) params = self.params - declared_names = {p.name for p in self.supported_parameters()} + declared_names = {parameter.name for parameter in self.supported_parameters()} - # objective_target is only required when the scenario declares it; a subclass may drop - # it (then self._objective_target stays None and the scenario supplies its own target). if "objective_target" in declared_names: - objective_target = self._resolve_objective_target(value=params.get("objective_target")) - if objective_target is None: - raise ValueError( - "objective_target is required. Provide it via " - "set_params_from_args(args={'objective_target': ...}) or register a default " - "with set_default_value() in an initialization script." - ) - self._objective_target = objective_target - self._objective_target_identifier = objective_target.get_identifier() - type(self).TARGET_REQUIREMENTS.validate(target=objective_target) + raw_objective_target = params.get("objective_target") + if require_objective_target or raw_objective_target is not None: + objective_target = self._resolve_objective_target(value=raw_objective_target) + if objective_target is None: + raise ValueError( + "objective_target is required. Provide it via " + "set_params_from_args(args={'objective_target': ...}) or register a default " + "with set_default_value() in an initialization script." + ) + self._objective_target = objective_target + self._objective_target_identifier = objective_target.get_identifier() + type(self).TARGET_REQUIREMENTS.validate(target=objective_target) + else: + self._objective_target = None + self._objective_target_identifier = None dataset_config = params.get("dataset_config") self._dataset_config_provided = dataset_config is not None @@ -591,10 +768,6 @@ async def initialize_async(self) -> None: self._max_retries = params.get("max_retries", 0) self._memory_labels = params.get("memory_labels") or {} - # Resolve the effective include_baseline. Forbidden is checked first so a forbidden - # scenario type never silently inherits a True default; explicit-True on a forbidden - # type is a hard error rather than a silent ignore. For the Enabled / Disabled states, - # a None runtime value defers to the policy. include_baseline = params.get("include_baseline") if self.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden: if include_baseline is True: @@ -605,16 +778,50 @@ async def initialize_async(self) -> None: include_baseline = False elif include_baseline is None: include_baseline = self.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Enabled - self._include_baseline = include_baseline - # Prepare scenario techniques via the resolution hook (subclasses override to widen - # accepted types or expand composites) and stash any per-technique converter overrides. self._scenario_techniques = self._resolve_scenario_techniques( scenario_techniques=params.get("scenario_techniques") ) self._technique_converters = params.get("technique_converters") or {} + @final + async def initialize_async(self) -> None: + """ + Initialize the scenario by populating self._atomic_attacks and creating the ScenarioResult. + + All run inputs are read from the parameter bag (``self.params``), which is populated by + ``set_params_from_args`` from the merged CLI / config / programmatic arguments. Callers + fill the bag then initialize: + + .. code-block:: python + + scenario.set_params_from_args(args={"objective_target": target, "max_concurrency": 8}) + await scenario.initialize_async() + + This method allows scenarios to be initialized with atomic attacks after construction, + which is useful when atomic attacks require async operations to be built. + + If a scenario_result_id was provided in __init__, this method will check if it exists + in memory and validate that the stored scenario matches the current configuration. + If it matches, the scenario will resume from prior progress. If it doesn't match or + doesn't exist, a new scenario result will be created. + + The common run inputs read from the bag are ``objective_target`` (a ``PromptTarget`` + instance or a registered target name resolved against ``TargetRegistry``), + ``scenario_techniques``, ``technique_converters``, ``dataset_config``, + ``max_concurrency``, ``max_retries``, ``memory_labels``, and ``include_baseline`` + (see ``_common_scenario_parameters``). A subclass that removes a common input via + ``supported_parameters`` falls back to that input's default here. + + Raises: + ValueError: If ``objective_target`` is declared but not resolvable (neither supplied + nor registered as a default), if a supplied target name is not registered in + ``TargetRegistry``, or if ``include_baseline=True`` is set for a scenario whose + ``BASELINE_ATTACK_POLICY`` is ``Forbidden``. + """ + self._resolve_runtime_configuration(require_objective_target=True) + # Build atomic attacks: resolve the seed groups once, snapshot the resolved inputs # into a ScenarioContext, and hand it to the subclass extension point. Baseline emission # is the scenario's own responsibility — matrix scenarios get it for free (the matrix @@ -653,7 +860,19 @@ async def initialize_async(self) -> None: stored_result=existing_results[0], current_identifier=scenario_identifier, ) - self._apply_persisted_objectives(stored_result=existing_results[0]) + stored_result = existing_results[0] + stored_plan = self._get_stored_run_plan(stored_result=stored_result) + if stored_plan is not None: + self._apply_persisted_run_plan(stored_plan=stored_plan) + else: + self._apply_persisted_objectives(stored_result=stored_result) + reconstructed_plan = self._build_run_plan() + metadata = dict(stored_result.metadata) + metadata[SCENARIO_RUN_PLAN_METADATA_KEY] = reconstructed_plan.model_dump(mode="json") + self._memory.update_scenario_metadata( + scenario_result_id=self._scenario_result_id, + metadata=metadata, + ) return # Valid resume - skip creating new scenario result # Build display group mapping from atomic attacks @@ -688,28 +907,131 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: chosen objective hashes here so the next ``_setup_scenario_async`` can replay them via ``keep_seed_groups_with_hashes``. - When ``max_dataset_size`` is not set, the sample equals the dataset and - nothing needs pinning; the dict is empty. + The normalized run plan is always stored. When ``max_dataset_size`` is not + set, only the run plan is needed because the full dataset is deterministic. Returns: dict[str, Any]: Metadata payload for the new ScenarioResult. """ metadata: dict[str, Any] = {} - if getattr(self._dataset_config, "max_dataset_size", None) is None: - return metadata - hashes: list[str] = [] - seen: set[str] = set() - for aa in self._atomic_attacks: - for sg in aa.seed_groups: - if sg.objective is None: - continue - sha = to_sha256(sg.objective.value) - if sha not in seen: - seen.add(sha) - hashes.append(sha) - metadata["objective_hashes"] = hashes + if getattr(self._dataset_config, "max_dataset_size", None) is not None: + hashes: list[str] = [] + seen: set[str] = set() + for aa in self._atomic_attacks: + for sg in aa.seed_groups: + sha = to_sha256(sg.objective.value) + if sha not in seen: + seen.add(sha) + hashes.append(sha) + metadata["objective_hashes"] = hashes + metadata[SCENARIO_RUN_PLAN_METADATA_KEY] = self._build_run_plan().model_dump(mode="json") return metadata + def _build_run_plan(self) -> ScenarioRunPlan: + """ + Build the normalized persistent plan for the initialized atomic attacks. + + Returns: + ScenarioRunPlan: The versioned run plan. + """ + seed_groups: dict[str, ScenarioRunPlanSeedGroup] = {} + atomic_groups: list[ScenarioRunPlanAtomicGroup] = [] + for atomic_attack in self._atomic_attacks: + seed_group_ids: list[str] = [] + seen_seed_group_ids: set[str] = set() + for seed_group in atomic_attack.seed_groups: + seed_group_id = seed_group.logical_id + if seed_group_id in seen_seed_group_ids: + continue + seen_seed_group_ids.add(seed_group_id) + seed_group_ids.append(seed_group_id) + seed_groups.setdefault( + seed_group_id, + ScenarioRunPlanSeedGroup( + id=seed_group_id, + objective_sha256=to_sha256(seed_group.objective.value), + objective=seed_group.objective.value, + ), + ) + technique_eval_hash = str(atomic_attack.technique_eval_hash) + atomic_group_id = self._get_atomic_group_id(atomic_attack=atomic_attack) + atomic_groups.append( + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name=atomic_attack.atomic_attack_name, + display_group=atomic_attack.display_group, + technique_eval_hash=technique_eval_hash, + seed_group_ids=seed_group_ids, + ) + ) + return ScenarioRunPlan( + scenario_registry_name=self._scenario_registry_name, + atomic_groups=atomic_groups, + seed_groups=list(seed_groups.values()), + ) + + @staticmethod + def _get_atomic_group_id(*, atomic_attack: AtomicAttack) -> str: + """ + Compute the stable ID of an atomic group from its name and technique. + + Returns: + str: The atomic-group ID. + """ + return config_hash( + { + "atomic_attack_name": atomic_attack.atomic_attack_name, + "technique_eval_hash": str(atomic_attack.technique_eval_hash), + } + ) + + @staticmethod + def _get_stored_run_plan(*, stored_result: ScenarioResult) -> ScenarioRunPlan | None: + """ + Load and validate a stored run plan. + + Returns: + ScenarioRunPlan | None: The plan, or None for a legacy row. + """ + raw_plan = (stored_result.metadata or {}).get(SCENARIO_RUN_PLAN_METADATA_KEY) + if raw_plan is None: + return None + return ScenarioRunPlan.model_validate(raw_plan) + + def _apply_persisted_run_plan(self, *, stored_plan: ScenarioRunPlan) -> None: + """ + Validate and replay the exact logical units captured by a stored plan. + + Raises: + ValueError: If a planned atomic or seed group cannot be reconstructed. + """ + current_by_id = { + self._get_atomic_group_id(atomic_attack=atomic_attack): atomic_attack + for atomic_attack in self._atomic_attacks + } + planned_ids = {group.id for group in stored_plan.atomic_groups} + missing_groups = planned_ids - current_by_id.keys() + if missing_groups: + raise ValueError( + f"Scenario result id '{self._scenario_result_id}' cannot resume: " + f"{len(missing_groups)} planned atomic group(s) are no longer reconstructable." + ) + + retained_attacks: list[AtomicAttack] = [] + for planned_group in stored_plan.atomic_groups: + atomic_attack = current_by_id[planned_group.id] + current_seed_groups = {seed_group.logical_id: seed_group for seed_group in atomic_attack.seed_groups} + missing_seed_groups = set(planned_group.seed_group_ids) - current_seed_groups.keys() + if missing_seed_groups: + raise ValueError( + f"Scenario result id '{self._scenario_result_id}' cannot resume: atomic group " + f"'{planned_group.atomic_attack_name}' is missing {len(missing_seed_groups)} planned seed group(s)." + ) + atomic_attack._seed_groups = [current_seed_groups[group_id] for group_id in planned_group.seed_group_ids] + retained_attacks.append(atomic_attack) + self._atomic_attacks = retained_attacks + self._display_group_map = {group.atomic_attack_name: group.display_group for group in stored_plan.atomic_groups} + def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: """ On resume, replay the originally-sampled objective subset. @@ -1303,6 +1625,8 @@ async def worker_async() -> None: atomic_attack = queue.get_nowait() except asyncio.QueueEmpty: return + atomic_group_id = atomic_attack.logical_group_id + self._active_atomic_groups[atomic_group_id] = atomic_attack.atomic_attack_name try: result = await atomic_attack.run_async( executor=shared_executor, @@ -1315,6 +1639,7 @@ async def worker_async() -> None: outcomes.append(exc) stop_event.set() finally: + self._active_atomic_groups.pop(atomic_group_id, None) pbar.update(1) # Cap workers at max_concurrency: that's also the objective-budget cap, and it's diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index fec03778d2..09e88734d8 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -22,6 +22,10 @@ from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackScoringConfig +from pyrit.models import ( + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, +) from pyrit.models.identifiers import compute_inner_attack_eval_hash from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -197,6 +201,84 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list return atomic_attacks + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate compatible persisted envelopes, excluding adaptive inner attempts. + + Returns: + ScenarioRunSizeEstimate: The adaptive outer-envelope estimate. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + selected_count = sum(len(groups) for groups in selected_groups.values()) + max_attempts = int(self.params.get("max_attempts_per_objective", 3)) + baseline_components = ( + [ + ScenarioRunSizeComponent( + label="Baseline", + count=selected_count, + is_baseline=True, + ) + ] + if self._include_baseline + else [] + ) + if not self._estimate_target_is_configured: + components = [ + *baseline_components, + ScenarioRunSizeComponent( + label="Adaptive attack-envelope candidates", + count=selected_count, + ), + ] + return ScenarioRunSizeEstimate( + components=components, + datasets=datasets, + note=( + "The authoritative total depends on which selected techniques are compatible with the " + f"configured objective target and each seed group. Up to {max_attempts} inner attempts per " + "envelope and retries are excluded." + ), + ) + + assert self._objective_target is not None + techniques = self._build_techniques_dict(objective_target=self._objective_target) + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=self._objective_target, + techniques=techniques, + selector=self._selector, + objective_scorer=self._objective_scorer, + max_attempts_per_objective=self.params.get("max_attempts_per_objective", 3), + scenario_result_id=self._scenario_result_id, + ) + compatible_group_count = sum( + bool(dispatcher.compatible_techniques(seed_group=seed_group)) + for seed_groups in selected_groups.values() + for seed_group in seed_groups + ) + + components = [ + *baseline_components, + ScenarioRunSizeComponent( + label="Adaptive attack envelopes", + count=compatible_group_count, + ), + ] + estimated_attack_count = ( + None if self._estimate_has_binding_size_cap else sum(component.count for component in components) + ) + note = ( + f"Each planned unit is one persisted adaptive envelope. Up to {max_attempts} selected technique " + "attempts may run inside that unit; inner attempts and retries are excluded." + ) + if estimated_attack_count is None: + note += " A binding randomized dataset cap may select a different compatibility mix at launch." + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, + components=components, + datasets=datasets, + note=note, + ) + def _build_techniques_dict( self, *, diff --git a/pyrit/scenario/scenarios/airt/cyber.py b/pyrit/scenario/scenarios/airt/cyber.py index 2f622c7b41..a983e97ba7 100644 --- a/pyrit/scenario/scenarios/airt/cyber.py +++ b/pyrit/scenario/scenarios/airt/cyber.py @@ -5,7 +5,7 @@ import logging from functools import cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from pyrit.common import apply_defaults from pyrit.common.path import SCORER_SEED_PROMPT_PATH @@ -70,6 +70,7 @@ class Cyber(Scenario): #: technique pool (and the ``all`` aggregate) reflects whatever the initializer #: registered. ``use_cached`` only matches prior runs at the current ``VERSION``. VERSION: int = 3 + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = True @classmethod def get_override_composite_scorer_questions_path(cls) -> list[Path]: diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 9bfcc63f19..cad7fe6f17 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -12,7 +12,12 @@ from pyrit.converter import TextJailbreakConverter from pyrit.datasets import TextJailBreak from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack -from pyrit.models import AttackTechniqueSeedGroup, Parameter +from pyrit.models import ( + AttackTechniqueSeedGroup, + Parameter, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, +) from pyrit.prompt_target import CapabilityName from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -296,6 +301,104 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: metadata[_JAILBREAK_TEMPLATES_METADATA_KEY] = list(self._resolved_jailbreaks) return metadata + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate the template and attempt axes, preserving the target capability caveat. + + Returns: + ScenarioRunSizeEstimate: Conditional target-aware estimate. + + Raises: + ValueError: If native system-prompt delivery is the only selected + technique but the selected target cannot support it. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + seed_group_count = sum(len(groups) for groups in selected_groups.values()) + template_count = len(self.params.get("jailbreak_names") or []) or ( + self.params.get("num_jailbreaks") or _DEFAULT_NUM_JAILBREAKS + ) + attempt_count = self.params.get("num_jailbreak_attempts") or 1 + technique_names = {technique.value for technique in self._scenario_techniques} + converter_count = len(technique_names - {_JAILBREAK_SYSTEM_PROMPT}) + system_delivery_selected = _JAILBREAK_SYSTEM_PROMPT in technique_names + system_delivery_supported = ( + self._target_supports_system_delivery(self._objective_target) + if system_delivery_selected and self._objective_target is not None + else None + ) + if system_delivery_selected and system_delivery_supported is False and converter_count == 0: + raise ValueError( + "Technique 'jailbreak_system_prompt' requires an objective target with editable history " + "and system-prompt support." + ) + + components: list[ScenarioRunSizeComponent] = [] + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=seed_group_count, + is_baseline=True, + ) + ) + components.append( + ScenarioRunSizeComponent( + label="Inline jailbreak delivery", + count=seed_group_count * template_count * attempt_count * converter_count, + note=( + "Each planned unit is one template, one selected delivery technique, and one logical seed group. " + "num_jailbreaks selects templates; it is not a persisted result or attempt count." + ), + ) + ) + if system_delivery_selected and system_delivery_supported is not False: + components.append( + ScenarioRunSizeComponent( + label="Native system-prompt jailbreak delivery", + count=seed_group_count * template_count * attempt_count, + note=( + "The selected objective target supports native system-prompt delivery." + if system_delivery_supported is True + else "Included only when the objective target supports editable history and system prompts." + ), + ) + ) + + target_agnostic_count = sum( + component.count for component in components if component.label != "Native system-prompt jailbreak delivery" + ) + planned_count = sum(component.count for component in components) + baseline_explanation = ( + f" Baseline adds one unit per selected seed group ({seed_group_count} units)." + if self._include_baseline + else " Baseline is disabled." + ) + formula = ( + f"{template_count} template(s) x {seed_group_count} selected logical seed group(s) x " + f"{converter_count} selected target-agnostic technique(s) x {attempt_count} configured attempt(s) " + f"= {seed_group_count * template_count * attempt_count * converter_count} planned unit(s)." + ) + estimated_attack_count = ( + None if system_delivery_selected and system_delivery_supported is None else planned_count + ) + if estimated_attack_count is None: + capability_note = ( + f" {target_agnostic_count} total planned units for target-agnostic delivery; " + f"{planned_count} when native system-prompt delivery is supported." + ) + elif system_delivery_selected and system_delivery_supported is True: + capability_note = " The selected target supports the native system-prompt component." + elif system_delivery_selected: + capability_note = " The selected target does not support native system-prompt delivery, so it is omitted." + else: + capability_note = "" + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, + components=components, + datasets=datasets, + note=f"{formula}{baseline_explanation}{capability_note}", + ) + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build one atomic attack per (technique x jailbreak template x dataset x attempt). diff --git a/pyrit/scenario/scenarios/airt/leakage.py b/pyrit/scenario/scenarios/airt/leakage.py index 264035f0ae..1c0f3b3e05 100644 --- a/pyrit/scenario/scenarios/airt/leakage.py +++ b/pyrit/scenario/scenarios/airt/leakage.py @@ -5,7 +5,7 @@ import logging from functools import cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from pyrit.common import apply_defaults from pyrit.common.path import SCORER_SEED_PROMPT_PATH @@ -80,6 +80,11 @@ class Leakage(Scenario): """ VERSION: int = 2 + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = True + + def _get_run_size_extra_factories(self) -> dict[str, AttackTechniqueFactory]: + """Return Leakage's source-owned factories for matrix sizing.""" + return {factory.name: factory for factory in _leakage_factories()} @classmethod def _get_additional_scoring_questions(cls) -> list[Path]: diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 83df78d6e2..3b06528296 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -30,7 +30,11 @@ AttackScoringConfig, CrescendoAttack, ) -from pyrit.models import SeedPrompt +from pyrit.models import ( + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + SeedPrompt, +) from pyrit.models.parameter import Parameter from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -483,6 +487,40 @@ async def _resolve_seed_groups_by_dataset_async( self._dataset_config = rebuilt return await super()._resolve_seed_groups_by_dataset_async(apply_sampling=apply_sampling) + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate the independent sub-harm technique sweeps and per-harm baselines. + + Returns: + ScenarioRunSizeEstimate: Exact per-sub-harm estimate. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + technique_count = len(self._scenario_techniques) + components: list[ScenarioRunSizeComponent] = [] + for dataset_name, seed_groups in selected_groups.items(): + seed_group_count = len(seed_groups) + components.append( + ScenarioRunSizeComponent( + label=f"{dataset_name} technique sweep", + count=seed_group_count * technique_count, + ) + ) + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label=f"{dataset_name} baseline", + count=seed_group_count, + is_baseline=True, + note="Psychosocial uses a distinct baseline and scorer for each sub-harm.", + ) + ) + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note="Each default sub-harm is planned independently; retries and internal turns are excluded.", + ) + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build atomic attacks as the ``(selected sub-harm x selected technique)`` cross product. diff --git a/pyrit/scenario/scenarios/airt/rapid_response.py b/pyrit/scenario/scenarios/airt/rapid_response.py index 4fd292bbe8..ca6b8d611b 100644 --- a/pyrit/scenario/scenarios/airt/rapid_response.py +++ b/pyrit/scenario/scenarios/airt/rapid_response.py @@ -14,7 +14,7 @@ import logging from functools import cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from pyrit.common import apply_defaults from pyrit.scenario.core.dataset_configuration import CompoundDatasetAttackConfiguration @@ -66,6 +66,7 @@ class RapidResponse(Scenario): #: technique pool (and the ``all`` aggregate) reflects whatever the initializer #: registered. ``use_cached`` only matches prior runs at the current ``VERSION``. VERSION: int = 3 + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = True @apply_defaults def __init__( diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 9270187920..9a38c2acfa 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -11,11 +11,23 @@ from pyrit.analytics import get_cached_results_for_technique from pyrit.common import apply_defaults -from pyrit.models import AttackOutcome, AttackResult, ObjectiveTargetEvaluationIdentifier, ScenarioResult +from pyrit.models import ( + AttackOutcome, + AttackResult, + ObjectiveTargetEvaluationIdentifier, + ScenarioResult, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, +) from pyrit.models.parameter import Parameter from pyrit.registry import AttackTechniqueRegistry, TargetRegistry from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration -from pyrit.scenario.core.matrix_atomic_attack_builder import MatrixAtomicAttackBuilder, resolve_technique_factories +from pyrit.scenario.core.matrix_atomic_attack_builder import ( + MatrixAtomicAttackBuilder, + filter_compatible_seed_groups, + resolve_technique_factories, + resolve_technique_factories_for_techniques, +) from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario if TYPE_CHECKING: @@ -191,6 +203,64 @@ def __init__( scenario_result_id=scenario_result_id, ) + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate the target-by-technique matrix using execution compatibility. + + Returns: + ScenarioRunSizeEstimate: Structured benchmark estimate. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + target_names = self.params.get("adversarial_targets") or [] + if not target_names: + return ScenarioRunSizeEstimate( + datasets=datasets, + note=( + "A total is unavailable until adversarial_targets is supplied and resolved. Baseline is forbidden." + ), + ) + + resolved_targets = self._resolve_adversarial_targets(target_names=target_names) + factories = resolve_technique_factories_for_techniques( + scenario_techniques=self._scenario_techniques, + ) + components: list[ScenarioRunSizeComponent] = [] + for technique in self._scenario_techniques: + factory = factories.get(technique.value) + if factory is None: + continue + compatible_count = sum( + len(filter_compatible_seed_groups(factory=factory, seed_groups=groups)) + for groups in selected_groups.values() + ) + components.append( + ScenarioRunSizeComponent( + label=technique.value, + count=len(resolved_targets) * compatible_count, + ) + ) + + if self._use_cached or self._estimate_has_binding_size_cap: + reasons = [] + if self._use_cached: + reasons.append("Live behavioral-cache hits can suppress work") + if self._estimate_has_binding_size_cap: + reasons.append("a binding randomized dataset cap may select a different compatibility mix at launch") + return ScenarioRunSizeEstimate( + components=components, + datasets=datasets, + note=( + f"Components describe the candidate population. {'; '.join(reasons)}, " + "so the authoritative total is unavailable before launch." + ), + ) + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note="Baseline is forbidden; retries and internal attack turns are excluded.", + ) + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build atomic attacks from (technique × adversarial_target × dataset), then apply caching. diff --git a/pyrit/scenario/scenarios/foundry/red_team_agent.py b/pyrit/scenario/scenarios/foundry/red_team_agent.py index 5d0cc4f235..958683adc1 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -50,7 +50,11 @@ TreeOfAttacksWithPruningAttack, ) from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackConverterConfig, AttackScoringConfig -from pyrit.models import AttackSeedGroup +from pyrit.models import ( + AttackSeedGroup, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, +) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -414,6 +418,37 @@ def _resolve_foundry_techniques( self._scenario_composites = composites return flat + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate one selected seed population per resolved Foundry composition. + + Returns: + ScenarioRunSizeEstimate: The composition population estimate. + """ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + selected_count = sum(len(groups) for groups in selected_groups.values()) + components = [ + ScenarioRunSizeComponent( + label=composition.name, + count=selected_count, + ) + for composition in self._scenario_composites + ] + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=selected_count, + is_baseline=True, + ) + ) + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note="Counts one population per resolved Foundry composite, not per flattened constituent technique.", + ) + @staticmethod def _technique_to_composite(technique: ScenarioTechnique) -> "FoundryComposite": """ diff --git a/pyrit/scenario/scenarios/garak/doctor.py b/pyrit/scenario/scenarios/garak/doctor.py index 9c608674a5..38273f25d6 100644 --- a/pyrit/scenario/scenarios/garak/doctor.py +++ b/pyrit/scenario/scenarios/garak/doctor.py @@ -104,11 +104,16 @@ class Doctor(Scenario): """ VERSION: int = 1 + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = True # Template-dominated like the Jailbreak scenario: baseline is supported but off # by default since the unmodified objective is a weak comparison point here. BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Disabled + def _get_run_size_extra_factories(self) -> dict[str, AttackTechniqueFactory]: + """Return Doctor's local Policy Puppetry factories for matrix sizing.""" + return {factory.name: factory for factory in DOCTOR_FACTORIES} + @classmethod def required_datasets(cls) -> list[str]: """Return a list of dataset names required by this scenario.""" diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index baef99ac31..aaf00a411e 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -24,7 +24,14 @@ from pyrit.converter.nato_converter import NatoConverter from pyrit.executor.attack.core.attack_config import AttackConverterConfig, AttackScoringConfig from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack -from pyrit.models import AttackSeedGroup, Seed, SeedObjective, SeedPrompt +from pyrit.models import ( + AttackSeedGroup, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + Seed, + SeedObjective, + SeedPrompt, +) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -223,29 +230,47 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list atomic_attacks.extend(self._get_converter_attacks(context=context)) return atomic_attacks - # These are the same as Garak encoding attacks - def _get_converter_attacks(self, *, context: ScenarioContext) -> list[AtomicAttack]: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ - Get all converter-based atomic attacks. - - Creates atomic attacks for each encoding scheme specified in the scenario techniques. - Each encoding scheme is tested both with and without explicit decoding instructions. - - Args: - context (ScenarioContext): The resolved runtime inputs for this run. + Estimate converter variants crossed with raw and decode-template prompt configurations. Returns: - list[AtomicAttack]: List of all atomic attacks to execute. + ScenarioRunSizeEstimate: Exact converter-variant estimate. """ - # Map of all available converters with their encoding name and a unique variant slug. - # ``encoding_name`` drives technique selection and user-facing grouping (display_group); - # ``variant_slug`` is unique per row so atomic-attack names stay unique even when one - # encoding name maps to multiple converter variants (e.g. base64, ascii85). - # NOTE: near-duplicate base64 variants were trimmed alongside the VERSION bump - # (``standard_b64encode`` is byte-identical to the default ``b64encode``; ``b2a_base64`` - # only appends a trailing newline). We keep the default encoding plus the url-safe alphabet, - # which is a genuinely distinct representation. - all_converters_with_encodings: list[tuple[list[Converter], str, str]] = [ + selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() + seed_group_count = sum(len(groups) for groups in selected_groups.values()) + selected_encoding_names = {technique.value for technique in self._scenario_techniques} + variant_count = sum(1 for _, name, _ in self._converter_variants() if name in selected_encoding_names) + prompt_configuration_count = 1 + len(self._encoding_templates) + components = [ + ScenarioRunSizeComponent( + label="Encoding converter variants", + count=seed_group_count * variant_count * prompt_configuration_count, + note=( + "Concrete variants are counted separately when one catalog technique maps to " + "multiple encoders, including base64 and ascii85." + ), + ) + ] + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=seed_group_count, + is_baseline=True, + ) + ) + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note="Retries are excluded; each converter and decode-template configuration is a planned outer unit.", + ) + + @staticmethod + def _converter_variants() -> list[tuple[list[Converter], str, str]]: + """Return the canonical converter implementations and their catalog technique names.""" + return [ ([Base64Converter()], "base64", "base64"), ([Base64Converter(encoding_func="urlsafe_b64encode")], "base64", "base64_urlsafe"), ([Base2048Converter()], "base2048", "base2048"), @@ -267,11 +292,33 @@ def _get_converter_attacks(self, *, context: ScenarioContext) -> list[AtomicAtta ([AsciiSmugglerConverter()], "ascii_smuggler", "ascii_smuggler"), ] + # These are the same as Garak encoding attacks + def _get_converter_attacks(self, *, context: ScenarioContext) -> list[AtomicAttack]: + """ + Get all converter-based atomic attacks. + + Creates atomic attacks for each encoding scheme specified in the scenario techniques. + Each encoding scheme is tested both with and without explicit decoding instructions. + + Args: + context (ScenarioContext): The resolved runtime inputs for this run. + + Returns: + list[AtomicAttack]: List of all atomic attacks to execute. + """ + # Map of all available converters with their encoding name and a unique variant slug. + # ``encoding_name`` drives technique selection and user-facing grouping (display_group); + # ``variant_slug`` is unique per row so atomic-attack names stay unique even when one + # encoding name maps to multiple converter variants (e.g. base64, ascii85). + # NOTE: near-duplicate base64 variants were trimmed alongside the VERSION bump + # (``standard_b64encode`` is byte-identical to the default ``b64encode``; ``b2a_base64`` + # only appends a trailing newline). We keep the default encoding plus the url-safe alphabet, + # which is a genuinely distinct representation. # Filter to only include selected techniques selected_encoding_names = {s.value for s in context.scenario_techniques} converters_with_encodings = [ (conv, name, variant_slug) - for conv, name, variant_slug in all_converters_with_encodings + for conv, name, variant_slug in self._converter_variants() if name in selected_encoding_names ] diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 719f1d87eb..ddf7216d4b 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import logging import random from typing import TYPE_CHECKING, ClassVar, cast @@ -11,7 +12,14 @@ from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory import CentralMemory -from pyrit.models import AttackSeedGroup, SeedObjective, SeedPrompt +from pyrit.models import ( + AttackSeedGroup, + ScenarioDatasetSummary, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + SeedObjective, + SeedPrompt, +) from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration @@ -482,35 +490,20 @@ def _scoring_config_for_technique(self, technique: WebInjectionTechnique) -> Att return self._xss_scoring_config return self._exfil_scoring_config - async def _resolve_seed_groups_by_dataset_async( - self, *, apply_sampling: bool = True + def _build_synthesized_seed_groups( + self, *, dataset_values: dict[str, list[str]] ) -> dict[str, list[AttackSeedGroup]]: """ - Generate the injection prompts and wrap them into seed groups, keyed by technique. - - WebInjection synthesizes its seeds (rather than resolving them from a - ``DatasetAttackConfiguration``): each technique renders its own objective and prompt - set from the raw garak datasets. Resolving them here means the base owns the single - seed sample used for both the atomic attacks and the baseline. - - Args: - apply_sampling (bool): Accepted for base-class compatibility but unused — the - synthesized seeds are already deterministic (``random.Random(self._random_seed)``), - so resume reproduces the same set without a ``max_dataset_size`` sampling path. + Build the deterministic, technique-specific logical populations. Returns: - dict[str, list[AttackSeedGroup]]: Seed groups keyed by technique value. + dict[str, list[AttackSeedGroup]]: Synthesized groups keyed by technique. Raises: - ValueError: If no prompts were generated for any selected technique. + ValueError: If the source datasets produce no prompts. """ - dataset_values = self._load_dataset_values() rng = random.Random(self._random_seed) - seed_groups_by_technique: dict[str, list[AttackSeedGroup]] = {} - # ``_scenario_techniques`` is typed as the base ``ScenarioTechnique`` on the - # ``Scenario`` base class, but this scenario only ever populates it with - # ``WebInjectionTechnique`` members (its ``technique_class``). techniques = cast("list[WebInjectionTechnique]", self._scenario_techniques) for technique in techniques: objective, prompts = self._build_prompts_for_technique( @@ -530,9 +523,89 @@ async def _resolve_seed_groups_by_dataset_async( "(garak_example_domains_xss, garak_markdown_js, garak_web_html_js, " "garak_xss_normal_instructions) are loaded into CentralMemory before running." ) - return seed_groups_by_technique + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: + """ + Estimate the technique-specific synthesized populations and their shared baseline. + + Returns: + ScenarioRunSizeEstimate: Exact synthesized-population estimate. + """ + dataset_values = await asyncio.to_thread(self._load_dataset_values) + seed_groups_by_technique = self._build_synthesized_seed_groups(dataset_values=dataset_values) + datasets = [ + ScenarioDatasetSummary( + name=name, + logical_seed_group_count=len(values), + selected_seed_group_count=len(values), + selection_note="Raw source values used to synthesize technique-specific prompt populations.", + ) + for name, values in dataset_values.items() + ] + datasets.extend( + ScenarioDatasetSummary( + name=technique_name, + kind="synthesized", + logical_seed_group_count=len(seed_groups), + selected_seed_group_count=len(seed_groups), + selection_note="Deterministic prompt population after the per-technique cap.", + ) + for technique_name, seed_groups in seed_groups_by_technique.items() + ) + + components = [ + ScenarioRunSizeComponent( + label=f"{technique_name} synthesized prompts", + count=len(seed_groups), + ) + for technique_name, seed_groups in seed_groups_by_technique.items() + ] + synthesized_count = sum(len(groups) for groups in seed_groups_by_technique.values()) + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=synthesized_count, + is_baseline=True, + note="The baseline runs over the union of all default technique populations.", + ) + ) + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note=( + "Each technique owns a distinct synthesized population; " + "no generic dataset-by-technique formula applies." + ), + ) + + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[AttackSeedGroup]]: + """ + Generate the injection prompts and wrap them into seed groups, keyed by technique. + + WebInjection synthesizes its seeds (rather than resolving them from a + ``DatasetAttackConfiguration``): each technique renders its own objective and prompt + set from the raw garak datasets. Resolving them here means the base owns the single + seed sample used for both the atomic attacks and the baseline. + + Args: + apply_sampling (bool): Accepted for base-class compatibility but unused — the + synthesized seeds are already deterministic (``random.Random(self._random_seed)``), + so resume reproduces the same set without a ``max_dataset_size`` sampling path. + + Returns: + dict[str, list[AttackSeedGroup]]: Seed groups keyed by technique value. + + Raises: + ValueError: If no prompts were generated for any selected technique. + """ + dataset_values = await asyncio.to_thread(self._load_dataset_values) + return self._build_synthesized_seed_groups(dataset_values=dataset_values) + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build one AtomicAttack per selected technique from the resolved seed groups. diff --git a/pyrit/setup/initializers/techniques/airt.py b/pyrit/setup/initializers/techniques/airt.py index 1469ea6e19..0ca857c6fd 100644 --- a/pyrit/setup/initializers/techniques/airt.py +++ b/pyrit/setup/initializers/techniques/airt.py @@ -42,6 +42,7 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Obfuscates the objective by asking for it encoded as the first letter of each word.", technique_tags=["single_turn", "airt", "leakage"], + supports_additional_request_converters=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters(converters=[FirstLetterConverter()]) diff --git a/pyrit/setup/initializers/techniques/core.py b/pyrit/setup/initializers/techniques/core.py index 8770f8dece..d4dbc9667c 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -213,6 +213,7 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Reverses the objective text so it slips past filters, then asks the target to flip it back.", technique_tags=["single_turn", "light"], + supports_additional_request_converters=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters( diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index abe283a559..78697d70e9 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -29,7 +29,7 @@ AttackService, get_attack_service, ) -from pyrit.memory import AttackResultsKeysetCursor +from pyrit.memory import AttackResultKeysetCursor from pyrit.models import ( AtomicAttackIdentifier, AttackOutcome, @@ -160,7 +160,7 @@ def _cursor_for(result: AttackResult, *, fingerprint: str | None = None) -> str: """ effective_fingerprint = fingerprint if fingerprint is not None else AttackService._attack_filter_fingerprint() return AttackService._encode_attack_cursor( - cursor=AttackResultsKeysetCursor.from_attack_result(result), + cursor=AttackResultKeysetCursor.from_attack_result(result), fingerprint=effective_fingerprint, ) diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py index dc41e698a3..627c0ab705 100644 --- a/tests/unit/backend/test_scenario_run_routes.py +++ b/tests/unit/backend/test_scenario_run_routes.py @@ -6,6 +6,7 @@ """ from datetime import datetime, timezone +from threading import get_ident from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,8 +16,14 @@ import pyrit.backend.services.scenario_run_service as _svc_mod from pyrit.backend.main import app from pyrit.backend.models.scenarios import ScenarioRunListResponse -from pyrit.models import ScenarioRunState -from pyrit.models.catalog.scenario import ScenarioRunSummary +from pyrit.backend.routes.scenarios import get_scenario_run_progress +from pyrit.models import ( + ScenarioProgressHeader, + ScenarioRunPlan, + ScenarioRunProgress, + ScenarioRunState, +) +from pyrit.models.catalog.scenario import ScenarioRunListItem, ScenarioRunSummary from unit.mocks import make_scenario_result @@ -121,27 +128,71 @@ def test_start_run_with_all_options(self, client: TestClient) -> None: assert response.status_code == status.HTTP_202_ACCEPTED + def test_start_jailbreak_run_preserves_explicit_selection_and_params(self, client: TestClient) -> None: + """The route parses the exact Jailbreak selection without adding catalog defaults.""" + mock_response = _mock_run_response() + + with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: + mock_service = MagicMock() + mock_service.start_run_async = AsyncMock(return_value=mock_response) + mock_get.return_value = mock_service + + response = client.post( + "/api/scenarios/runs", + json={ + "scenario_name": "airt.jailbreak", + "target_name": "my_target", + "techniques": ["prompt_sending"], + "include_baseline": False, + "scenario_params": { + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + }, + }, + ) + + assert response.status_code == status.HTTP_202_ACCEPTED + request = mock_service.start_run_async.await_args.kwargs["request"] + assert request.techniques == ["prompt_sending"] + assert request.include_baseline is False + assert request.scenario_params == { + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + class TestListScenarioRunsRoute: """Tests for GET /api/scenarios/runs.""" def test_list_runs_returns_200(self, client: TestClient) -> None: """Test that list runs returns 200 with empty list.""" + route_thread: list[int] = [] with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: mock_service = MagicMock() - mock_service.list_runs.return_value = ScenarioRunListResponse(items=[]) + mock_service.list_runs.side_effect = lambda **_: ( + route_thread.append(get_ident()) or ScenarioRunListResponse(items=[]) + ) mock_get.return_value = mock_service + request_thread = get_ident() response = client.get("/api/scenarios/runs") assert response.status_code == status.HTTP_200_OK assert response.json()["items"] == [] + assert route_thread[0] != request_thread + + def test_list_runs_rejects_unbounded_limit(self, client: TestClient) -> None: + response = client.get("/api/scenarios/runs?limit=101") + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT def test_list_runs_returns_multiple_runs(self, client: TestClient) -> None: """Test that list runs returns all tracked runs.""" runs = [ - _mock_run_response(run_id="run-1"), - _mock_run_response(run_id="run-2", run_status=ScenarioRunState.IN_PROGRESS), + ScenarioRunListItem.model_validate(_mock_run_response(run_id="run-1").model_dump()), + ScenarioRunListItem.model_validate( + _mock_run_response(run_id="run-2", run_status=ScenarioRunState.IN_PROGRESS).model_dump() + ), ] with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: @@ -164,7 +215,8 @@ def test_get_run_returns_200(self, client: TestClient) -> None: with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: mock_service = MagicMock() - mock_service.get_run.return_value = mock_response + mock_service.snapshot_active_run.return_value = MagicMock(error=None) + mock_service.get_run_from_storage.return_value = mock_response mock_get.return_value = mock_service response = client.get("/api/scenarios/runs/test-run-id") @@ -176,13 +228,106 @@ def test_get_run_not_found_returns_404(self, client: TestClient) -> None: """Test that getting a non-existent run returns 404.""" with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: mock_service = MagicMock() - mock_service.get_run.return_value = None + mock_service.snapshot_active_run.return_value = MagicMock(error=None) + mock_service.get_run_from_storage.return_value = None mock_get.return_value = mock_service response = client.get("/api/scenarios/runs/nonexistent") assert response.status_code == status.HTTP_404_NOT_FOUND + def test_progress_invalid_cursor_returns_400(self, client: TestClient) -> None: + with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: + mock_service = MagicMock() + mock_service.snapshot_active_run.return_value = MagicMock(active_group_ids=()) + mock_service.get_run_progress_from_storage.side_effect = ValueError("Malformed scenario progress cursor.") + mock_get.return_value = mock_service + + response = client.get("/api/scenarios/runs/test-run-id/progress?since=bad") + + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert response.json()["detail"] == "Malformed scenario progress cursor." + + def test_progress_returns_compact_plan_response(self, client: TestClient) -> None: + progress = ScenarioRunProgress( + run=ScenarioProgressHeader( + scenario_result_id="test-run-id", + scenario_name="TestScenario", + scenario_registry_name="test.scenario", + scenario_version=1, + status=ScenarioRunState.IN_PROGRESS, + created_at=datetime(2025, 1, 1, tzinfo=timezone.utc), + ), + plan=ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[], + seed_groups=[], + ), + active_atomic_group_ids=["active-group"], + plan_complete=True, + ) + snapshot_thread: list[int] = [] + storage_thread: list[int] = [] + with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: + mock_service = MagicMock() + mock_service.snapshot_active_run.side_effect = lambda **_: ( + snapshot_thread.append(get_ident()) or MagicMock(active_group_ids=("active-group",)) + ) + mock_service.get_run_progress_from_storage.side_effect = lambda **_: ( + storage_thread.append(get_ident()) or progress + ) + mock_get.return_value = mock_service + + response = client.get("/api/scenarios/runs/test-run-id/progress?limit=25") + + assert response.status_code == status.HTTP_200_OK + assert response.json()["plan"]["scenario_registry_name"] == "test.scenario" + assert response.json()["active_atomic_group_ids"] == ["active-group"] + mock_service.get_run_progress_from_storage.assert_called_once_with( + scenario_result_id="test-run-id", + since=None, + limit=25, + active_group_ids=("active-group",), + ) + assert snapshot_thread[0] != storage_thread[0] + + async def test_progress_supports_direct_keyword_call(self) -> None: + progress = ScenarioRunProgress( + run=ScenarioProgressHeader( + scenario_result_id="test-run-id", + scenario_name="TestScenario", + scenario_registry_name="test.scenario", + scenario_version=1, + status=ScenarioRunState.IN_PROGRESS, + created_at=datetime(2025, 1, 1, tzinfo=timezone.utc), + ), + plan=ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[], + seed_groups=[], + ), + plan_complete=True, + ) + with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: + mock_service = MagicMock() + mock_service.snapshot_active_run.return_value = MagicMock(active_group_ids=()) + mock_service.get_run_progress_from_storage.return_value = progress + mock_get.return_value = mock_service + + result = await get_scenario_run_progress( + scenario_result_id="test-run-id", + since=None, + limit=25, + ) + + assert result == progress + mock_service.get_run_progress_from_storage.assert_called_once_with( + scenario_result_id="test-run-id", + since=None, + limit=25, + active_group_ids=(), + ) + class TestCancelScenarioRunRoute: """Tests for POST /api/scenarios/runs/{id}/cancel.""" diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 8ce480c62f..83eb1b0995 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -5,19 +5,37 @@ Tests for ScenarioRunService. """ +import asyncio +import uuid from datetime import datetime, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest +import pyrit.backend.services.scenario_configuration_resolver as _resolver_mod import pyrit.backend.services.scenario_run_service as _svc_mod from pyrit.backend.services.scenario_run_service import ( _DEFAULT_MAX_CONCURRENT_RUNS, ScenarioRunService, ) from pyrit.converter import Converter -from pyrit.models import AttackOutcome, ScenarioResult, ScenarioRunState +from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + ScenarioAttackResultDelta, + ScenarioResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunState, + SeedObjective, + config_hash, +) from pyrit.models.catalog.scenario import RunScenarioRequest from pyrit.scenario.core import DatasetAttackConfiguration, DatasetConfiguration from pyrit.scenario.core.scenario_technique import ScenarioTechnique @@ -42,7 +60,7 @@ def _patch_converter_registry(instances: dict[str, Any]): reg = MagicMock() reg.instances.get.side_effect = lambda name: instances.get(name) reg.instances.get_names.return_value = list(instances.keys()) - return patch.object(_svc_mod.ConverterRegistry, "get_registry_singleton", return_value=reg) + return patch.object(_resolver_mod.ConverterRegistry, "get_registry_singleton", return_value=reg) _REGISTRY_PATCH_BASE = "pyrit.registry" @@ -67,6 +85,8 @@ def _make_request( dataset_names: list[str] | None = None, max_dataset_size: int | None = None, dataset_filters: dict[str, list[str]] | None = None, + include_baseline: bool | None = None, + scenario_params: dict[str, Any] | None = None, ) -> RunScenarioRequest: """Create a RunScenarioRequest for testing.""" return RunScenarioRequest( @@ -78,6 +98,8 @@ def _make_request( dataset_names=dataset_names, max_dataset_size=max_dataset_size, dataset_filters=dataset_filters, + include_baseline=include_baseline, + scenario_params=scenario_params, ) @@ -113,6 +135,7 @@ def mock_memory(): """Patch CentralMemory.get_memory_instance to return a mock.""" mock = MagicMock() mock.get_scenario_results.return_value = [] + mock.get_scenario_result_headers.return_value = [] # Default: no error AttackResults linked to any scenario. Tests that exercise # the error fallback path explicitly set get_attack_results.return_value. mock.get_attack_results.return_value = [] @@ -292,6 +315,55 @@ def _lookup(name): init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args assert init_call.kwargs["scenario_techniques"] == [technique_a, technique_b] + async def test_jailbreak_explicit_selection_and_params_reach_registry_unchanged(self, mock_all_registries) -> None: + """An explicit Jailbreak technique never adds the default aggregate or other techniques.""" + + class _JailbreakTechnique(ScenarioTechnique): + ALL = ("all", {"all"}) + DEFAULT = ("default", {"default"}) + PROMPT_SENDING = ("prompt_sending", {"default"}) + CONTEXT_COMPLIANCE = ("context_compliance", {"default"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all", "default"} + + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._technique_class = _JailbreakTechnique + objective_target = mock_all_registries["target_registry"].instances.get.return_value + scenario_params = {"num_jailbreaks": 2, "num_jailbreak_attempts": 1} + + service = ScenarioRunService() + await service.start_run_async( + request=_make_request( + scenario_name="airt.jailbreak", + techniques=["prompt_sending"], + include_baseline=False, + scenario_params=scenario_params, + ) + ) + + mock_all_registries["scenario_registry"].create_and_initialize_async.assert_awaited_once_with( + "airt.jailbreak", + scenario_params=scenario_params, + scenario_result_id=None, + objective_target=objective_target, + max_concurrency=10, + max_retries=0, + include_baseline=False, + scenario_techniques=[_JailbreakTechnique.PROMPT_SENDING], + ) + + async def test_start_run_forwards_include_baseline(self, mock_all_registries) -> None: + service = ScenarioRunService() + request = _make_request() + request.include_baseline = False + + await service.start_run_async(request=request) + + init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args + assert init_call.kwargs["include_baseline"] is False + async def test_start_run_max_dataset_size_uses_default_config(self, mock_all_registries) -> None: """``max_dataset_size`` with no ``dataset_names`` reuses the scenario's default config.""" default_config = MagicMock() @@ -359,10 +431,8 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): assert built_config.dataset_names == ["only_this"] assert built_config.max_dataset_size is None - async def test_start_run_dataset_names_falls_back_when_subclass_constructor_incompatible( - self, mock_all_registries, caplog - ) -> None: - """If the subclass __init__ rejects standard kwargs, fall back to plain ``DatasetConfiguration``.""" + async def test_start_run_dataset_names_rejects_incompatible_subclass_constructor(self, mock_all_registries) -> None: + """Reject overrides that cannot preserve scenario-specific dataset configuration.""" class _RequiresExtraArgConfiguration(DatasetConfiguration): def __init__(self, *, required_extra: str, **kwargs: Any) -> None: @@ -376,21 +446,13 @@ def __init__(self, *, required_extra: str, **kwargs: Any) -> None: ) service = ScenarioRunService() - with caplog.at_level("WARNING", logger=_svc_mod.logger.name): + with pytest.raises( + ValueError, + match="does not support overriding dataset names.*_RequiresExtraArgConfiguration", + ): await service.start_run_async(request=_make_request(dataset_names=["custom"])) - init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] - - # Fallback is the generic base class, not the subclass - assert type(built_config) is DatasetAttackConfiguration - assert built_config.dataset_names == ["custom"] - # Warning was logged so the operator can see the silent degradation - assert any( - "_RequiresExtraArgConfiguration" in record.message - and "Falling back to a generic DatasetAttackConfiguration" in record.message - for record in caplog.records - ) + mock_all_registries["scenario_registry"].create_and_initialize_async.assert_not_awaited() async def test_start_run_dataset_filters_new_config(self, mock_all_registries) -> None: """``dataset_filters`` with ``dataset_names`` builds a config carrying the filters.""" @@ -618,11 +680,11 @@ class TestScenarioRunServiceListRuns: def test_list_runs_empty(self, mock_memory) -> None: """Test that list_runs returns empty list when DB has no results.""" - mock_memory.get_scenario_results.return_value = [] + mock_memory.get_scenario_result_headers.return_value = [] service = ScenarioRunService() result = service.list_runs() assert result.items == [] - mock_memory.get_scenario_results.assert_called_once_with(limit=100) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=100) def test_list_runs_returns_all_runs(self, mock_memory) -> None: """Test that list_runs returns all runs from the database.""" @@ -630,19 +692,27 @@ def test_list_runs_returns_all_runs(self, mock_memory) -> None: _make_db_scenario_result(result_id="sr-1", run_state=ScenarioRunState.COMPLETED), _make_db_scenario_result(result_id="sr-2", run_state=ScenarioRunState.IN_PROGRESS), ] - mock_memory.get_scenario_results.return_value = db_results + mock_memory.get_scenario_result_headers.return_value = db_results service = ScenarioRunService() result = service.list_runs() assert len(result.items) == 2 - mock_memory.get_scenario_results.assert_called_once_with(limit=100) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=100) def test_list_runs_passes_custom_limit(self, mock_memory) -> None: """Test that list_runs passes a custom limit to the memory query.""" - mock_memory.get_scenario_results.return_value = [] + mock_memory.get_scenario_result_headers.return_value = [] service = ScenarioRunService() service.list_runs(limit=10) - mock_memory.get_scenario_results.assert_called_once_with(limit=10) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=10) + + def test_list_runs_reports_unknown_total_without_plan(self, mock_memory) -> None: + """Test that legacy runs do not report a false zero planned total.""" + mock_memory.get_scenario_result_headers.return_value = [_make_db_scenario_result()] + + result = ScenarioRunService().list_runs() + + assert result.items[0].total_attacks is None class TestScenarioRunServiceCancelRun: @@ -680,6 +750,52 @@ async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> No assert result is not None assert result.status == ScenarioRunState.CANCELLED + async def test_cancel_waits_for_final_persisted_progress_delta(self, mock_all_registries) -> None: + """Cancellation completes task cleanup before callers can fetch terminal progress.""" + mock_memory = mock_all_registries["memory"] + scenario_instance = mock_all_registries["scenario_instance"] + delta = ScenarioAttackResultDelta( + attack_result_id=str(uuid.uuid4()), + objective="persisted during cancellation", + outcome=AttackOutcome.ERROR, + execution_time_ms=10, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + error_type="CancelledError", + error_message="cancelled", + attribution_data={"parent_collection": "attack"}, + ) + + async def run_until_cancelled() -> None: + try: + await asyncio.Event().wait() + finally: + mock_memory.get_scenario_attack_result_deltas.return_value = ([delta], False) + + scenario_instance.run_async.side_effect = run_until_cancelled + service = ScenarioRunService() + response = await service.start_run_async(request=_make_request()) + await asyncio.sleep(0) + + running_result = mock_all_registries["db_result"] + cancelled_result = _make_db_scenario_result( + result_id=response.scenario_result_id, + run_state=ScenarioRunState.CANCELLED, + ) + cancelled_result.metadata = {} + mock_memory.get_scenario_results.side_effect = [[running_result], [cancelled_result]] + + await service.cancel_run_async(scenario_result_id=response.scenario_result_id) + mock_memory.get_scenario_result_header.return_value = cancelled_result + progress = service.get_run_progress( + scenario_result_id=response.scenario_result_id, + since=None, + limit=25, + ) + + assert progress is not None + assert progress.run.status is ScenarioRunState.CANCELLED + assert [result.attack_result_id for result in progress.results] == [delta.attack_result_id] + async def test_cancel_completed_run_raises_value_error(self, mock_memory) -> None: """Test that cancelling a completed run raises ValueError.""" db_result = _make_db_scenario_result(result_id="sr-done", run_state=ScenarioRunState.COMPLETED) @@ -836,6 +952,7 @@ def test_in_progress_run_shows_partial_attack_counts(self, mock_memory) -> None: assert fetched.completed_attacks == 3 assert fetched.techniques_used == ["attack_a", "attack_b"] assert fetched.objective_achieved_rate == 33 + assert fetched.completed_at is None def test_created_run_shows_zero_counts(self, mock_memory) -> None: """Test that a CREATED run with no results shows zero counts.""" @@ -880,6 +997,7 @@ def test_completed_run_still_shows_full_counts(self, mock_memory) -> None: assert fetched.completed_attacks == 1 assert fetched.techniques_used == ["attack_a"] assert fetched.objective_achieved_rate == 100 + assert fetched.completed_at == db_result.completion_time class TestScenarioRunServiceFailedAttackReporting: @@ -984,9 +1102,8 @@ class TestResolveTechniquesAndConverters: """Tests for per-technique converter resolution from ``--techniques`` tokens.""" def test_plain_technique_no_converters(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.ROLE_PLAY] @@ -994,9 +1111,8 @@ def test_plain_technique_no_converters(self, mock_memory) -> None: def test_single_converter_appended(self, mock_memory) -> None: conv = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"translation_spanish": conv}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.translation_spanish"], technique_class=_StubTechnique, scenario_name="x", @@ -1006,9 +1122,8 @@ def test_single_converter_appended(self, mock_memory) -> None: def test_aggregate_token_applies_converter_to_all_concrete(self, mock_memory) -> None: conv = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": conv}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.EASY] @@ -1017,9 +1132,8 @@ def test_aggregate_token_applies_converter_to_all_concrete(self, mock_memory) -> def test_multiple_converters_preserve_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = service._resolve_techniques_and_converters( + _, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.c1:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1029,9 +1143,8 @@ def test_multiple_converters_preserve_order(self, mock_memory) -> None: def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = service._resolve_techniques_and_converters( + _, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1", "role_play:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1041,30 +1154,27 @@ def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: assert converters["single_turn"] == [c1] def test_unknown_converter_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({"known": MagicMock(spec=Converter)}): with pytest.raises(ValueError, match="not a registered converter"): - service._resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.missing"], technique_class=_StubTechnique, scenario_name="x", ) def test_unknown_modifier_prefix_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): with pytest.raises(ValueError, match="Unknown technique modifier"): - service._resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:scorer.something"], technique_class=_StubTechnique, scenario_name="x", ) def test_unknown_base_technique_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): with pytest.raises(ValueError, match="not found for scenario"): - service._resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["nope:converter.c1"], technique_class=_StubTechnique, scenario_name="x", @@ -1083,3 +1193,332 @@ async def test_start_run_forwards_technique_converters(self, mock_all_registries init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args assert init_call.kwargs["scenario_techniques"] == [_StubTechnique.ROLE_PLAY] assert init_call.kwargs["technique_converters"] == {"role_play": [conv]} + + +def test_planned_progress_deduplicates_attempts_and_keeps_latest_outcome(mock_memory) -> None: + seed_group = AttackSeedGroup(seeds=[SeedObjective(value="objective")]) + seed_group_id = seed_group.logical_id + atomic_group_id = config_hash({"atomic_attack_name": "attack", "technique_eval_hash": "eval"}) + atomic_identifier = AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="TestAttack", class_module="tests"), + seed_group=seed_group, + ) + plan = ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name="attack", + display_group="Attack", + technique_eval_hash="eval", + seed_group_ids=[seed_group_id], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id=seed_group_id, + objective_sha256="objective-sha", + objective="objective", + ) + ], + ) + attempts = [ + AttackResult( + conversation_id=f"conversation-{index}", + objective="objective", + atomic_attack_identifier=atomic_identifier, + outcome=outcome, + timestamp=datetime(2025, 1, 1, 0, index, tzinfo=timezone.utc), + attribution_data={"parent_collection": "attack", "parent_eval_hash": "eval"}, + ) + for index, outcome in enumerate( + (AttackOutcome.ERROR, AttackOutcome.FAILURE, AttackOutcome.SUCCESS, AttackOutcome.ERROR) + ) + ] + scenario_result = make_scenario_result( + attack_results={"attack": attempts}, + scenario_run_state=ScenarioRunState.COMPLETED, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")}, + ) + + summary = ScenarioRunService()._build_response_from_db(scenario_result=scenario_result) + + assert summary.total_attacks == 1 + assert summary.completed_attacks == 1 + assert summary.objective_achieved_rate == 0 + assert len(summary.failed_attacks) == 2 + assert summary.total_retries == 3 + + +def test_planned_progress_includes_latest_errors_in_success_rate_denominator(mock_memory) -> None: + atomic_group_id = config_hash({"atomic_attack_name": "attack", "technique_eval_hash": "eval"}) + plan = ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name="attack", + display_group="Attack", + technique_eval_hash="eval", + seed_group_ids=["seed-success", "seed-error"], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id="seed-success", + objective_sha256="success-sha", + objective="success objective", + ), + ScenarioRunPlanSeedGroup( + id="seed-error", + objective_sha256="error-sha", + objective="error objective", + ), + ], + ) + results = [ + AttackResult( + conversation_id="success-conversation", + objective="success objective", + outcome=AttackOutcome.SUCCESS, + attribution_data={ + "parent_collection": "attack", + "parent_eval_hash": "eval", + "seed_group_id": "seed-success", + }, + ), + AttackResult( + conversation_id="error-conversation", + objective="error objective", + outcome=AttackOutcome.ERROR, + attribution_data={ + "parent_collection": "attack", + "parent_eval_hash": "eval", + "seed_group_id": "seed-error", + }, + ), + ] + scenario_result = make_scenario_result( + attack_results={"attack": results}, + scenario_run_state=ScenarioRunState.COMPLETED, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")}, + ) + + summary = ScenarioRunService()._build_response_from_db(scenario_result=scenario_result) + + assert summary.total_attacks == 2 + assert summary.completed_attacks == 2 + assert summary.objective_achieved_rate == 50 + + +def test_planned_progress_maps_legacy_objective_hash_to_logical_seed_id(mock_memory) -> None: + objective = "legacy resumed objective" + seed_group = AttackSeedGroup(seeds=[SeedObjective(value=objective)]) + seed_group_id = seed_group.logical_id + atomic_group_id = config_hash({"atomic_attack_name": "attack", "technique_eval_hash": "eval"}) + plan = ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name="attack", + display_group="Attack", + technique_eval_hash="eval", + seed_group_ids=[seed_group_id], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id=seed_group_id, + objective_sha256=_svc_mod.to_sha256(objective), + objective=objective, + ) + ], + ) + legacy_attempt = AttackResult( + conversation_id="legacy-conversation", + objective=objective, + outcome=AttackOutcome.SUCCESS, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + attribution_data={"parent_collection": "attack", "parent_eval_hash": "eval"}, + ) + scenario_result = make_scenario_result( + attack_results={"attack": [legacy_attempt]}, + scenario_run_state=ScenarioRunState.COMPLETED, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")}, + ) + + summary = ScenarioRunService()._build_response_from_db(scenario_result=scenario_result) + + assert summary.total_attacks == 1 + assert summary.completed_attacks == 1 + + +def test_get_progress_uses_lightweight_queries_without_full_hydration(mock_memory) -> None: + plan = ScenarioRunPlan(atomic_groups=[], seed_groups=[], scenario_registry_name="test.scenario") + header = make_scenario_result( + attack_results={}, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")}, + ) + mock_memory.get_scenario_result_header.return_value = header + mock_memory.get_scenario_attack_result_deltas.return_value = ([], False) + mock_memory.get_scenario_results.reset_mock() + + service = ScenarioRunService() + completed_task = MagicMock() + completed_task.done.return_value = True + service._active_tasks[str(header.id)] = _svc_mod._ActiveTask( + scenario_result_id=str(header.id), + task=completed_task, + scenario=MagicMock(), + ) + + progress = service.get_run_progress( + scenario_result_id=str(header.id), + since=None, + limit=25, + ) + + assert progress is not None + assert progress.plan == plan + assert progress.plan_complete is True + mock_memory.get_scenario_results.assert_not_called() + assert str(header.id) not in service._active_tasks + + +def test_get_progress_rejects_duplicate_stored_plan_groups(mock_memory) -> None: + group = ScenarioRunPlanAtomicGroup( + id="duplicate", + atomic_attack_name="attack", + display_group="Attack", + technique_eval_hash="eval", + seed_group_ids=["seed-1"], + ).model_dump(mode="json") + header = make_scenario_result( + attack_results={}, + metadata={ + SCENARIO_RUN_PLAN_METADATA_KEY: { + "version": 1, + "atomic_groups": [group, group], + "seed_groups": [ + ScenarioRunPlanSeedGroup( + id="seed-1", + objective_sha256="objective-sha", + objective="objective", + ).model_dump(mode="json") + ], + } + }, + ) + mock_memory.get_scenario_result_header.return_value = header + mock_memory.get_scenario_attack_result_deltas.return_value = ([], False) + + with pytest.raises(ValueError, match="duplicate atomic group IDs"): + ScenarioRunService().get_run_progress( + scenario_result_id=str(header.id), + since=None, + limit=25, + ) + + +def test_progress_prefers_persisted_logical_seed_group_attribution() -> None: + delta = ScenarioAttackResultDelta( + attack_result_id=str(uuid.uuid4()), + objective="objective", + objective_sha256="objective-sha", + outcome=AttackOutcome.SUCCESS, + execution_time_ms=10, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + attribution_data={ + "parent_collection": "attack", + "parent_eval_hash": "eval", + "seed_group_id": "canonical-seed-id", + }, + ) + + mapped = ScenarioRunService._map_progress_delta( + delta=delta, + plan_lookup=_svc_mod._ScenarioPlanLookup.from_plan(plan=None), + ) + + assert mapped.seed_group_id == "canonical-seed-id" + + +def test_synthesize_legacy_plan_deduplicates_seed_ids_in_first_seen_order() -> None: + delta_units = [ + ("attack", "eval", "seed-b", "objective b"), + ("other attack", "other-eval", "seed-b", "objective b"), + ("attack", "eval", "seed-a", "objective a"), + ("attack", "eval", "seed-b", "objective b"), + ] + deltas = [ + ScenarioAttackResultDelta( + attack_result_id=str(uuid.uuid4()), + objective=objective, + outcome=AttackOutcome.SUCCESS, + execution_time_ms=10, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + attribution_data={ + "parent_collection": attack_name, + "parent_eval_hash": eval_hash, + "seed_group_id": seed_group_id, + }, + ) + for attack_name, eval_hash, seed_group_id, objective in delta_units + ] + + plan = ScenarioRunService._synthesize_legacy_plan(deltas=deltas) + + assert [group.atomic_attack_name for group in plan.atomic_groups] == ["attack", "other attack"] + assert plan.atomic_groups[0].seed_group_ids == ["seed-b", "seed-a"] + assert plan.atomic_groups[1].seed_group_ids == ["seed-b"] + assert [seed.id for seed in plan.seed_groups] == ["seed-b", "seed-a"] + + +def test_get_progress_synthesizes_incomplete_legacy_plan(mock_memory) -> None: + header = make_scenario_result( + attack_results={}, + scenario_run_state=ScenarioRunState.COMPLETED, + metadata={}, + ) + delta = ScenarioAttackResultDelta( + attack_result_id=str(uuid.uuid4()), + objective="legacy objective", + outcome=AttackOutcome.FAILURE, + execution_time_ms=10, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + attribution_data={"parent_collection": "legacy attack"}, + ) + mock_memory.get_scenario_result_header.return_value = header + mock_memory.get_scenario_attack_result_deltas.return_value = ([delta], False) + + progress = ScenarioRunService().get_run_progress( + scenario_result_id=str(header.id), + since=None, + limit=25, + ) + + assert progress is not None + assert progress.plan_complete is False + assert progress.plan is not None + assert len(progress.plan.atomic_groups) == 1 + assert len(progress.results) == 1 + + +def test_decode_progress_cursor_rejects_cross_run_cursor() -> None: + delta = ScenarioAttackResultDelta( + attack_result_id=str(uuid.uuid4()), + objective="objective", + outcome=AttackOutcome.SUCCESS, + execution_time_ms=10, + timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc), + ) + cursor = ScenarioRunService._encode_progress_cursor(scenario_result_id="run-a", delta=delta) + + with pytest.raises(ValueError, match="does not belong"): + ScenarioRunService._decode_progress_cursor(since=cursor, scenario_result_id="run-b") + + +def test_decode_progress_cursor_rejects_malformed_cursor() -> None: + with pytest.raises(ValueError, match="Malformed"): + ScenarioRunService._decode_progress_cursor(since="not-a-cursor", scenario_result_id="run-a") diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index b4b31a7b72..e7297fa1cc 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -5,6 +5,8 @@ Tests for backend scenario service and routes. """ +import asyncio +from collections import OrderedDict from typing import Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -15,13 +17,38 @@ from pyrit.backend.main import app from pyrit.backend.models.common import PaginationInfo from pyrit.backend.models.scenarios import ListRegisteredScenariosResponse +from pyrit.backend.routes.scenarios import estimate_scenario_run_size +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_service import ( ScenarioService, get_scenario_service, ) -from pyrit.models import Parameter +from pyrit.models import ( + Parameter, + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, +) from pyrit.models.catalog.scenario import RegisteredScenario from pyrit.registry import ScenarioMetadata +from pyrit.scenario.core import DatasetAttackConfiguration, ScenarioTechnique + + +class _EstimateTechnique(ScenarioTechnique): + """Technique enum for configured catalog estimate tests.""" + + ALL = ("all", {"all"}) + DEFAULT = ("default", {"default"}) + PROMPT_SENDING = ("prompt_sending", {"default"}) + JAILBREAK_SYSTEM_PROMPT = ("jailbreak_system_prompt", {"default"}) + FLIP = ("flip", {"direct"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + """Return aggregate tags.""" + return {"all", "default"} @pytest.fixture @@ -42,11 +69,20 @@ def _make_scenario_metadata( *, registry_name: str = "test.scenario", class_name: str = "TestScenario", + scenario_version: int = 1, description: str = "A test scenario", + description_markdown: str = "A test scenario", default_technique: str = "default", + default_techniques: tuple[str, ...] = ("role_play", "many_shot"), all_techniques: tuple[str, ...] = ("role_play", "many_shot"), aggregate_techniques: tuple[str, ...] = ("all", "default"), + aggregate_technique_expansions: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("all", ("role_play", "many_shot")), + ("default", ("role_play",)), + ), default_datasets: tuple[str, ...] = ("test_dataset",), + baseline_policy: str = "enabled", + include_baseline_by_default: bool = True, ) -> ScenarioMetadata: """Create a ScenarioMetadata instance for testing.""" return ScenarioMetadata( @@ -54,10 +90,16 @@ def _make_scenario_metadata( class_name=class_name, class_module="pyrit.scenario.scenarios.test", class_description=description, + scenario_version=scenario_version, + description_markdown=description_markdown, default_technique=default_technique, + default_techniques=default_techniques, all_techniques=all_techniques, aggregate_techniques=aggregate_techniques, + aggregate_technique_expansions=aggregate_technique_expansions, default_datasets=default_datasets, + baseline_policy=baseline_policy, + include_baseline_by_default=include_baseline_by_default, ) @@ -96,10 +138,267 @@ async def test_list_scenarios_returns_scenarios_from_registry(self) -> None: assert result.items[0].scenario_name == "test.scenario" assert result.items[0].scenario_type == "TestScenario" assert result.items[0].description == "A test scenario" + assert result.items[0].description_markdown == "A test scenario" assert result.items[0].default_technique == "default" + assert result.items[0].default_techniques == ["role_play", "many_shot"] assert result.items[0].aggregate_techniques == ["all", "default"] + assert result.items[0].aggregate_technique_expansions["default"] == ["role_play"] assert result.items[0].all_techniques == ["role_play", "many_shot"] assert result.items[0].default_datasets == ["test_dataset"] + assert result.items[0].baseline_policy == "enabled" + assert result.items[0].include_baseline_by_default is True + + async def test_estimate_is_offloaded_and_cached(self) -> None: + """Scenario-owned estimates run in a worker once and are reused by subsequent reads.""" + metadata = _make_scenario_metadata() + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=4, + components=[ScenarioRunSizeComponent(label="Default sweep", count=4)], + datasets=[ + ScenarioDatasetSummary( + name="test_dataset", + logical_seed_group_count=4, + selected_seed_group_count=2, + configured_caps=[ + ScenarioDatasetSizeCap( + label="per-dataset cap", + count=2, + configured_on="dataset", + dataset_name="test_dataset", + ) + ], + ) + ], + ) + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + + first = await service.get_scenario_async(scenario_name="test.scenario") + second = await service.get_scenario_async(scenario_name="test.scenario") + + assert first is not None + assert second is not None + assert first.default_run_size == estimate + assert second.default_run_size == estimate + assert first.default_dataset_summaries == estimate.datasets + assert second.default_dataset_summaries == estimate.datasets + service._registry.create_instance.assert_called_once_with("test.scenario") + + async def test_concurrent_estimate_reads_share_one_task(self) -> None: + """Concurrent catalog readers share one atomic single-flight estimate.""" + metadata = _make_scenario_metadata() + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + started = asyncio.Event() + release = asyncio.Event() + + async def estimate_async() -> ScenarioRunSizeEstimate: + started.set() + await release.wait() + return estimate + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.create_instance.return_value = scenario + + first = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + await started.wait() + second = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + await asyncio.sleep(0) + assert service._registry.create_instance.call_count == 1 + + release.set() + assert await asyncio.gather(first, second) == [estimate, estimate] + await asyncio.sleep(0) + + assert service._estimate_tasks == {} + + def test_estimate_task_cleanup_preserves_replacement(self) -> None: + """A stale completion callback cannot remove the replacement task for the same key.""" + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + cache_key = ("test.scenario", 1) + completed_task = MagicMock(spec=asyncio.Task) + replacement_task = MagicMock(spec=asyncio.Task) + service._estimate_tasks = OrderedDict([(cache_key, replacement_task)]) + + service._clear_estimate_task(task=completed_task, cache_key=cache_key) + assert service._estimate_tasks[cache_key] is replacement_task + + service._clear_estimate_task(task=replacement_task, cache_key=cache_key) + assert service._estimate_tasks == {} + + async def test_cancelled_estimate_waiter_does_not_cancel_shared_task(self) -> None: + """Cancelling one waiter leaves the shared estimate available to other readers.""" + metadata = _make_scenario_metadata() + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + started = asyncio.Event() + release = asyncio.Event() + + async def estimate_async() -> ScenarioRunSizeEstimate: + started.set() + await release.wait() + return estimate + + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=estimate_async) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.create_instance.return_value = scenario + + cancelled_waiter = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + await started.wait() + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + + surviving_waiter = asyncio.create_task(service._get_default_run_size_estimate_async(metadata=metadata)) + release.set() + assert await surviving_waiter == estimate + await asyncio.sleep(0) + + assert service._registry.create_instance.call_count == 1 + assert service._estimate_tasks == {} + + async def test_completed_stale_task_cannot_block_inflight_capacity(self) -> None: + """A done task is pruned before the bounded inflight capacity check.""" + metadata = _make_scenario_metadata() + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) + + with ( + patch.object(ScenarioService, "__init__", lambda self: None), + patch("pyrit.backend.services.scenario_service._ESTIMATE_INFLIGHT_SIZE", 1), + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.create_instance.return_value = scenario + stale = asyncio.create_task(asyncio.sleep(0, result=estimate)) + await stale + service._estimate_tasks = OrderedDict([(("stale.scenario", 1), stale)]) + + result = await asyncio.wait_for( + service._get_default_run_size_estimate_async(metadata=metadata), + timeout=1, + ) + + assert result == estimate + + async def test_one_failed_estimate_does_not_break_catalog(self) -> None: + """A scenario estimate failure is explicit and isolated from other catalog entries.""" + metadata = [ + _make_scenario_metadata(registry_name="test.good"), + _make_scenario_metadata(registry_name="test.bad"), + ] + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=2, + components=[ScenarioRunSizeComponent(label="Default sweep", count=2)], + ) + good_scenario = MagicMock() + good_scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) + bad_scenario = MagicMock() + bad_scenario.get_default_run_size_estimate_async = AsyncMock(side_effect=RuntimeError("dataset unavailable")) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_all_registered_class_metadata.return_value = metadata + service._registry.create_instance.side_effect = lambda name: { + "test.good": good_scenario, + "test.bad": bad_scenario, + }[name] + + result = await service.list_scenarios_async() + assert "RuntimeError" in result.items[1].default_run_size.note + + async def test_unavailable_estimate_cache_expires(self) -> None: + """A transient estimate failure is retried after the unavailable-result TTL.""" + metadata = _make_scenario_metadata() + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock( + side_effect=[RuntimeError("temporary failure"), estimate] + ) + + with ( + patch.object(ScenarioService, "__init__", lambda self: None), + patch("pyrit.backend.services.scenario_service._UNAVAILABLE_CACHE_TTL_SECONDS", 0), + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.create_instance.return_value = scenario + + first = await service.get_scenario_async(scenario_name="test.scenario") + second = await service.get_scenario_async(scenario_name="test.scenario") + + assert first is not None + assert second is not None + assert second.default_run_size == estimate + assert service._registry.create_instance.call_count == 2 + + async def test_estimate_cache_is_version_aware_and_bounded(self) -> None: + """Scenario version changes invalidate estimates and the LRU stays bounded.""" + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + scenario = MagicMock() + scenario.get_default_run_size_estimate_async = AsyncMock(return_value=estimate) + + with ( + patch.object(ScenarioService, "__init__", lambda self: None), + patch("pyrit.backend.services.scenario_service._ESTIMATE_CACHE_SIZE", 1), + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.create_instance.return_value = scenario + + await service._get_default_run_size_estimate_async(metadata=_make_scenario_metadata(scenario_version=1)) + await service._get_default_run_size_estimate_async(metadata=_make_scenario_metadata(scenario_version=2)) + + assert service._registry.create_instance.call_count == 2 + assert list(service._estimate_cache) == [("test.scenario", 2)] + + async def test_list_scenarios_preserves_disabled_baseline_policy(self) -> None: + metadata = _make_scenario_metadata( + baseline_policy="disabled", + include_baseline_by_default=False, + ) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_all_registered_class_metadata.return_value = [metadata] + + result = await service.list_scenarios_async() + + assert result.items[0].baseline_policy == "disabled" + assert result.items[0].include_baseline_by_default is False async def test_list_scenarios_paginates_with_limit(self) -> None: """Test that list respects the limit parameter.""" @@ -117,6 +416,11 @@ async def test_list_scenarios_paginates_with_limit(self) -> None: assert len(result.items) == 3 assert result.pagination.has_more is True assert result.pagination.next_cursor == "test.scenario_2" + assert [call.args[0] for call in service._registry.create_instance.call_args_list] == [ + "test.scenario_0", + "test.scenario_1", + "test.scenario_2", + ] async def test_list_scenarios_paginates_with_cursor(self) -> None: """Test that list uses cursor for pagination.""" @@ -157,6 +461,120 @@ async def test_list_scenarios_last_page_has_more_false(self) -> None: class TestScenarioServiceGetScenario: """Tests for ScenarioService.get_scenario_async.""" + async def test_configured_estimate_uses_shared_launch_resolution(self) -> None: + """Configured estimates pass typed selections and parameters into the registry lifecycle.""" + metadata = _make_scenario_metadata(registry_name="airt.jailbreak") + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=12, + components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)], + ) + introspection_instance = MagicMock() + introspection_instance._technique_class = _EstimateTechnique + introspection_instance._default_dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"]) + scenario_class = MagicMock(return_value=introspection_instance) + objective_target = MagicMock() + + with ( + patch.object(ScenarioService, "__init__", lambda self: None), + patch.object( + ScenarioConfigurationResolver, "resolve_target", return_value=objective_target + ) as resolve_target, + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.get_class.return_value = scenario_class + service._registry.create_and_estimate_async = AsyncMock(return_value=estimate) + + result = await service.estimate_scenario_run_size_async( + scenario_name="airt.jailbreak", + request=ScenarioRunSizeEstimateRequest( + target_name="preview_target", + techniques=["prompt_sending"], + dataset_names=["harmbench"], + max_dataset_size=3, + dataset_filters={"harm_categories": ["violence"]}, + include_baseline=True, + scenario_params={ + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + }, + ), + ) + + assert result == estimate + resolve_target.assert_called_once_with(target_name="preview_target") + call = service._registry.create_and_estimate_async.await_args + assert call.args == () + assert call.kwargs["name"] == "airt.jailbreak" + assert call.kwargs["scenario_params"] == { + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + assert call.kwargs["scenario_techniques"] == [_EstimateTechnique.PROMPT_SENDING] + assert call.kwargs["include_baseline"] is True + assert call.kwargs["objective_target"] is objective_target + dataset_config = call.kwargs["dataset_config"] + assert type(dataset_config) is DatasetAttackConfiguration + assert dataset_config.dataset_names == ["harmbench"] + assert dataset_config.max_dataset_size == 3 + assert dataset_config.filters == {"harm_categories": ["violence"]} + + async def test_configured_estimate_rejects_incompatible_v4_jailbreak_technique(self) -> None: + """Request previews reject techniques omitted by Jailbreak's v4 compatibility policy.""" + metadata = _make_scenario_metadata(registry_name="airt.jailbreak") + introspection_instance = MagicMock() + introspection_instance._technique_class = _EstimateTechnique + introspection_instance._default_dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"]) + scenario_class = MagicMock(return_value=introspection_instance) + + with patch.object(ScenarioService, "__init__", lambda self: None): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.get_class.return_value = scenario_class + service._registry.create_and_estimate_async = AsyncMock() + + with pytest.raises(ValueError, match="context_compliance"): + await service.estimate_scenario_run_size_async( + scenario_name="airt.jailbreak", + request=ScenarioRunSizeEstimateRequest(techniques=["context_compliance"]), + ) + + service._registry.create_and_estimate_async.assert_not_awaited() + + async def test_configured_estimate_without_target_does_not_resolve_or_send_to_target(self) -> None: + """Target-conditional previews stay side-effect free when no target is configured.""" + metadata = _make_scenario_metadata(registry_name="adaptive.text") + estimate = ScenarioRunSizeEstimate( + note="Target compatibility is unknown.", + ) + introspection_instance = MagicMock() + introspection_instance._technique_class = _EstimateTechnique + introspection_instance._default_dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"]) + scenario_class = MagicMock(return_value=introspection_instance) + + with ( + patch.object(ScenarioService, "__init__", lambda self: None), + patch.object(ScenarioConfigurationResolver, "resolve_target") as resolve_target, + ): + service = ScenarioService() + service._registry = MagicMock() + service._registry.get_registered_class_metadata.return_value = metadata + service._registry.get_class.return_value = scenario_class + service._registry.create_and_estimate_async = AsyncMock(return_value=estimate) + + result = await service.estimate_scenario_run_size_async( + scenario_name="adaptive.text", + request=ScenarioRunSizeEstimateRequest(), + ) + + assert result == estimate + resolve_target.assert_not_called() + call = service._registry.create_and_estimate_async.await_args + assert "target_is_configured" not in call.kwargs + assert "objective_target" not in call.kwargs + async def test_get_scenario_returns_matching_scenario(self) -> None: """Test that get returns the matching scenario.""" metadata = _make_scenario_metadata(registry_name="foundry.red_team_agent") @@ -216,10 +634,30 @@ def test_list_scenarios_with_items(self, client: TestClient) -> None: scenario_name="foundry.red_team_agent", scenario_type="RedTeamAgentScenario", description="Red team agent testing", + description_markdown='', default_technique="default", aggregate_techniques=["all", "default"], + aggregate_technique_expansions={ + "all": ["role_play", "many_shot"], + "default": ["role_play"], + }, all_techniques=["role_play", "many_shot"], default_datasets=["airt_hate"], + default_dataset_summaries=[ + ScenarioDatasetSummary( + name="airt_hate", + logical_seed_group_count=4, + selected_seed_group_count=4, + configured_caps=[ + ScenarioDatasetSizeCap( + label="per-dataset cap", + count=4, + configured_on="dataset", + dataset_name="airt_hate", + ) + ], + ) + ], ) with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: @@ -240,10 +678,13 @@ def test_list_scenarios_with_items(self, client: TestClient) -> None: item = data["items"][0] assert item["scenario_name"] == "foundry.red_team_agent" assert item["scenario_type"] == "RedTeamAgentScenario" + assert item["description_markdown"] == '' assert item["default_technique"] == "default" assert item["aggregate_techniques"] == ["all", "default"] + assert item["aggregate_technique_expansions"]["default"] == ["role_play"] assert item["all_techniques"] == ["role_play", "many_shot"] assert item["default_datasets"] == ["airt_hate"] + assert item["default_dataset_summaries"][0]["configured_caps"][0]["count"] == 4 def test_list_scenarios_passes_pagination_params(self, client: TestClient) -> None: """Test that pagination params are forwarded to service.""" @@ -269,9 +710,19 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: scenario_type="RedTeamAgentScenario", description="Red team agent testing", default_technique="default", + default_techniques=["role_play"], aggregate_techniques=["all"], all_techniques=["role_play"], default_datasets=["airt_hate"], + default_run_size=ScenarioRunSizeEstimate( + estimated_attack_count=8, + components=[ + ScenarioRunSizeComponent( + label="Default technique sweep", + count=8, + ) + ], + ), ) with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: @@ -284,6 +735,8 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: assert response.status_code == status.HTTP_200_OK data = response.json() assert data["scenario_name"] == "foundry.red_team_agent" + assert data["default_techniques"] == ["role_play"] + assert data["default_run_size"]["estimated_attack_count"] == 8 def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> None: """Test that GET /api/scenarios/catalog/{name} returns 404 when not found.""" @@ -296,6 +749,91 @@ def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> No assert response.status_code == status.HTTP_404_NOT_FOUND + def test_estimate_scenario_returns_configured_projection(self, client: TestClient) -> None: + """POST catalog estimate forwards request fields and returns the structured estimate.""" + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=12, + components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)], + ) + with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: + mock_service = MagicMock() + mock_service.estimate_scenario_run_size_async = AsyncMock(return_value=estimate) + mock_get_service.return_value = mock_service + + response = client.post( + "/api/scenarios/catalog/airt.jailbreak/estimate", + json={ + "techniques": ["prompt_sending"], + "include_baseline": True, + "scenario_params": { + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + }, + }, + ) + + assert response.status_code == status.HTTP_200_OK + assert response.json()["estimated_attack_count"] == 12 + request = mock_service.estimate_scenario_run_size_async.await_args.kwargs["request"] + assert request.techniques == ["prompt_sending"] + assert request.include_baseline is True + assert request.scenario_params == { + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + + async def test_estimate_scenario_supports_direct_keyword_call(self) -> None: + """The FastAPI handler remains directly callable through its keyword-only API.""" + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, + components=[ScenarioRunSizeComponent(label="Configured estimate", count=1)], + ) + request = ScenarioRunSizeEstimateRequest() + with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: + mock_service = MagicMock() + mock_service.estimate_scenario_run_size_async = AsyncMock(return_value=estimate) + mock_get_service.return_value = mock_service + + result = await estimate_scenario_run_size( + scenario_name="test.scenario", + request=request, + ) + + assert result == estimate + mock_service.estimate_scenario_run_size_async.assert_awaited_once_with( + scenario_name="test.scenario", + request=request, + ) + + def test_estimate_scenario_returns_400_for_invalid_configuration(self, client: TestClient) -> None: + """Configured estimate validation errors become clear client errors.""" + with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: + mock_service = MagicMock() + mock_service.estimate_scenario_run_size_async = AsyncMock( + side_effect=ValueError("Technique 'unknown' not found") + ) + mock_get_service.return_value = mock_service + + response = client.post( + "/api/scenarios/catalog/airt.jailbreak/estimate", + json={"techniques": ["unknown"]}, + ) + + assert response.status_code == status.HTTP_400_BAD_REQUEST + assert "Technique 'unknown' not found" in response.json()["detail"] + + def test_estimate_scenario_returns_404_for_unknown_scenario(self, client: TestClient) -> None: + """Unknown configured estimates preserve the catalog not-found contract.""" + with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: + mock_service = MagicMock() + mock_service.estimate_scenario_run_size_async = AsyncMock(return_value=None) + mock_get_service.return_value = mock_service + + response = client.post("/api/scenarios/catalog/missing.scenario/estimate", json={}) + + assert response.status_code == status.HTTP_404_NOT_FOUND + assert "missing.scenario" in response.json()["detail"] + def test_get_scenario_with_dotted_name(self, client: TestClient) -> None: """Test that dotted scenario names (e.g., 'foundry.red_team_agent') work in path.""" summary = RegisteredScenario( diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index 96cc41769e..8267b36293 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -18,6 +18,7 @@ RegisteredInitializer, RegisteredScenario, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -493,7 +494,7 @@ async def test_list_scenario_runs_async(client, mock_httpx_client): mock_httpx_client.get.return_value = _make_response(json_data={"items": [_run_summary_payload()]}) result = await client.list_scenario_runs_async(limit=20) assert len(result) == 1 - assert isinstance(result[0], ScenarioRunSummary) + assert isinstance(result[0], ScenarioRunListItem) mock_httpx_client.get.assert_awaited_once_with("/api/scenarios/runs", params={"limit": 20}) diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index 8880c95466..c361d213dd 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -20,6 +20,7 @@ AttackRetrySummary, RegisteredInitializer, RegisteredScenario, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -836,21 +837,20 @@ def test_print_scenario_runs_list_empty(capsys): def test_print_scenario_runs_list_populated(capsys): runs = [ - _make_run( + ScenarioRunListItem( status=ScenarioRunState.COMPLETED, scenario_name="scen-a", scenario_result_id="abcdefgh1234", total_attacks=4, - objective_achieved_rate=75, created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2024, 1, 1, tzinfo=timezone.utc), ), - _make_run( + ScenarioRunListItem( status=ScenarioRunState.IN_PROGRESS, scenario_name="scen-b", scenario_result_id="ijklmnop5678", - total_attacks=0, - objective_achieved_rate=0, created_at=datetime(2024, 2, 2, tzinfo=timezone.utc), + updated_at=datetime(2024, 2, 2, tzinfo=timezone.utc), ), ] _output.print_scenario_runs_list(runs=runs) @@ -859,6 +859,8 @@ def test_print_scenario_runs_list_populated(capsys): assert "scen-b" in captured.out assert "abcdefgh1234" in captured.out assert "ijklmnop5678" in captured.out + assert "success" not in captured.out + assert "planned attacks unknown" in captured.out assert "…" not in captured.out assert "Total runs: 2" in captured.out diff --git a/tests/unit/executor/attack/component/test_prepended_conversation_config.py b/tests/unit/executor/attack/component/test_prepended_conversation_config.py index 5968102b35..2316877388 100644 --- a/tests/unit/executor/attack/component/test_prepended_conversation_config.py +++ b/tests/unit/executor/attack/component/test_prepended_conversation_config.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from typing import get_type_hints from unittest.mock import MagicMock from pyrit.executor.attack.component.prepended_conversation_config import PrependedConversationConfig from pyrit.message_normalizer import ConversationContextNormalizer +from pyrit.models import ChatMessageRole def test_default_init_apply_converters_to_user_role(): @@ -17,6 +19,10 @@ def test_simulated_assistant_converter_role_normalizes_to_assistant(): assert config.apply_converters_to_roles == ["assistant"] +def test_public_type_hints_resolve_at_runtime(): + assert get_type_hints(PrependedConversationConfig)["apply_converters_to_roles"] == list[ChatMessageRole] + + def test_default_init_message_normalizer_is_none(): config = PrependedConversationConfig() assert config.message_normalizer is None diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 9cf0fbeb9d..f95888bc90 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -726,6 +726,7 @@ async def test_on_post_execute_stamps_scenario_attribution_when_present( sample_attack_context._attribution = AttackResultAttribution( parent_id="scenario-1", parent_collection="atomic_a", + seed_group_id="seed-a", ) event_data = StrategyEventData( @@ -740,6 +741,7 @@ async def test_on_post_execute_stamps_scenario_attribution_when_present( assert sample_attack_result.attribution_parent_id == "scenario-1" assert sample_attack_result.attribution_data == { "parent_collection": "atomic_a", + "seed_group_id": "seed-a", } async def test_on_post_execute_no_attribution_leaves_fields_none( @@ -775,6 +777,7 @@ async def test_on_error_stamps_scenario_attribution_when_present(self, sample_at sample_attack_context._attribution = AttackResultAttribution( parent_id="scenario-err", parent_collection="atomic_err", + seed_group_id="seed-error", ) event_data = StrategyEventData( @@ -793,6 +796,7 @@ async def test_on_error_stamps_scenario_attribution_when_present(self, sample_at assert persisted.attribution_parent_id == "scenario-err" assert persisted.attribution_data == { "parent_collection": "atomic_err", + "seed_group_id": "seed-error", } async def test_on_post_execute_stamps_targeted_harm_categories(self, sample_attack_result, mock_memory): diff --git a/tests/unit/memory/memory_interface/test_interface_attack_results.py b/tests/unit/memory/memory_interface/test_interface_attack_results.py index e532a1fe54..8687929ab4 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -11,7 +11,7 @@ import pytest from pyrit.common.utils import to_sha256 -from pyrit.memory import AttackResultsKeysetCursor, MemoryInterface +from pyrit.memory import AttackResultKeysetCursor, MemoryInterface from pyrit.memory.memory_interface import _AttackResultQuery from pyrit.memory.memory_models import AttackResultEntry from pyrit.models import ( @@ -85,15 +85,15 @@ def _make_attack_result( return AttackResult(**kwargs) -def _after(page: "Sequence[AttackResult]") -> AttackResultsKeysetCursor: +def _after(page: "Sequence[AttackResult]") -> AttackResultKeysetCursor: """Build the keyset anchor for the next page from the last row of ``page``.""" - return AttackResultsKeysetCursor.from_attack_result(page[-1]) + return AttackResultKeysetCursor.from_attack_result(page[-1]) def _drain_keyset(memory: MemoryInterface, *, page_size: int, **filters) -> list[AttackResult]: """Page through get_attack_results with the keyset cursor until exhausted.""" drained: list[AttackResult] = [] - after: AttackResultsKeysetCursor | None = None + after: AttackResultKeysetCursor | None = None while True: page = list(memory.get_attack_results(limit=page_size, after=after, **filters)) drained.extend(page) @@ -122,12 +122,12 @@ def test_attack_result_query_snapshots_mutable_inputs(): def test_attack_result_query_requires_keyword_arguments(): """The internal query does not expose field ordering as a positional API.""" with pytest.raises(TypeError): - _AttackResultQuery(["id"]) # type: ignore[misc] + _AttackResultQuery(["id"]) # ty: ignore[too-many-positional-arguments] def test_get_attack_results_forwards_all_parameters_to_query(sqlite_instance: MemoryInterface): """The compatibility API maps every parameter onto the internal query.""" - cursor = AttackResultsKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) + cursor = AttackResultKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) identifier_filter = IdentifierFilter( identifier_type=IdentifierType.ATTACK, property_path="$.hash", @@ -1898,7 +1898,7 @@ def test_get_attack_results_paginated_empty_metadata_orders_newest_first(sqlite_ def test_get_attack_results_pagination_with_ids_raises(sqlite_instance: MemoryInterface): """limit/keyset pagination cannot be combined with id-batched lookups.""" - anchor = AttackResultsKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) + anchor = AttackResultKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) with pytest.raises(ValueError, match="pagination cannot be combined"): sqlite_instance.get_attack_results(attack_result_ids=[str(uuid.uuid4())], limit=10) with pytest.raises(ValueError, match="pagination cannot be combined"): @@ -2055,7 +2055,7 @@ def test_attack_result_keyset_order_matches_sql_order(sqlite_instance: MemoryInt python_order = sorted( sqlite_instance.get_attack_results(), key=lambda ar: ( - AttackResultsKeysetCursor.from_attack_result(ar).timestamp, + AttackResultKeysetCursor.from_attack_result(ar).timestamp, ar.attack_result_id, ), reverse=True, diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_progress.py b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py new file mode 100644 index 0000000000..8a4e677d55 --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for lightweight scenario progress memory queries.""" + +import uuid +from contextlib import closing +from datetime import datetime, timezone + +import pytest +from unit.mocks import get_mock_target_identifier, make_scenario_result + +from pyrit.memory import AttackResultKeysetCursor, MemoryInterface +from pyrit.memory.memory_models import ScenarioResultEntry +from pyrit.models import ( + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + ScenarioRunState, + SeedObjective, +) + + +def _make_delta_result( + *, + scenario_result_id: str, + attack_result_id: uuid.UUID, + timestamp: datetime, + objective: str, +) -> AttackResult: + seed_group = AttackSeedGroup(seeds=[SeedObjective(value=objective)]) + identifier = AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="TestAttack", class_module="tests"), + seed_group=seed_group, + ) + return AttackResult( + attack_result_id=str(attack_result_id), + conversation_id=f"conversation-{attack_result_id}", + objective=objective, + atomic_attack_identifier=identifier, + outcome=AttackOutcome.SUCCESS, + execution_time_ms=12, + timestamp=timestamp, + attribution_parent_id=scenario_result_id, + attribution_data={"parent_collection": "attack", "parent_eval_hash": "eval"}, + ) + + +def test_scenario_progress_deltas_page_equal_timestamps_by_id( + sqlite_instance: MemoryInterface, +) -> None: + scenario = make_scenario_result( + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + ) + unrelated = make_scenario_result( + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + ) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario, unrelated]) + timestamp = datetime(2026, 8, 6, tzinfo=timezone.utc) + first_id = uuid.UUID(int=1) + second_id = uuid.UUID(int=2) + rows = [ + _make_delta_result( + scenario_result_id=str(scenario.id), + attack_result_id=first_id, + timestamp=timestamp, + objective="first", + ), + _make_delta_result( + scenario_result_id=str(scenario.id), + attack_result_id=second_id, + timestamp=timestamp, + objective="second", + ), + _make_delta_result( + scenario_result_id=str(unrelated.id), + attack_result_id=uuid.UUID(int=3), + timestamp=timestamp, + objective="unrelated", + ), + ] + sqlite_instance.add_attack_results_to_memory(attack_results=rows) + + first_page, has_more = sqlite_instance.get_scenario_attack_result_deltas( + scenario_result_id=str(scenario.id), + limit=1, + ) + second_page, second_has_more = sqlite_instance.get_scenario_attack_result_deltas( + scenario_result_id=str(scenario.id), + cursor=AttackResultKeysetCursor( + timestamp=first_page[0].timestamp, + attack_result_id=first_page[0].attack_result_id, + ), + limit=1, + ) + + assert [row.attack_result_id for row in first_page] == [str(first_id)] + assert has_more is True + assert [row.attack_result_id for row in second_page] == [str(second_id)] + assert second_has_more is False + assert second_page[0].atomic_attack_identifier is not None + source_identifier = AtomicAttackIdentifier.from_component_identifier(rows[1].atomic_attack_identifier) + assert second_page[0].atomic_attack_identifier.logical_seed_group_id == source_identifier.logical_seed_group_id + + +def test_scenario_result_header_does_not_hydrate_attack_results( + sqlite_instance: MemoryInterface, +) -> None: + scenario = make_scenario_result( + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + ) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario]) + sqlite_instance.add_attack_results_to_memory( + attack_results=[ + _make_delta_result( + scenario_result_id=str(scenario.id), + attack_result_id=uuid.UUID(int=4), + timestamp=datetime(2026, 8, 6, tzinfo=timezone.utc), + objective="objective", + ) + ] + ) + + header = sqlite_instance.get_scenario_result_header(scenario_result_id=str(scenario.id)) + + assert header is not None + assert header.attack_results == {} + + +def test_scenario_result_headers_are_bounded_without_attack_results( + sqlite_instance: MemoryInterface, +) -> None: + scenarios = [ + make_scenario_result( + scenario_name=f"scenario-{index}", + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + ) + for index in range(2) + ] + sqlite_instance.add_scenario_results_to_memory(scenario_results=scenarios) + + headers = sqlite_instance.get_scenario_result_headers(limit=1) + + assert len(headers) == 1 + assert headers[0].attack_results == {} + with pytest.raises(ValueError, match="between 1 and 100"): + sqlite_instance.get_scenario_result_headers(limit=101) + + +def test_scenario_result_headers_include_recent_active_runs( + sqlite_instance: MemoryInterface, +) -> None: + completed = make_scenario_result( + scenario_name="completed", + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + scenario_run_state=ScenarioRunState.COMPLETED, + completion_time=datetime(2026, 8, 20, tzinfo=timezone.utc), + ) + active = make_scenario_result( + scenario_name="active", + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + scenario_run_state=ScenarioRunState.IN_PROGRESS, + completion_time=datetime(2026, 8, 10, tzinfo=timezone.utc), + ) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[completed, active]) + with closing(sqlite_instance.get_session()) as session: + completed_entry = session.get(ScenarioResultEntry, completed.id) + active_entry = session.get(ScenarioResultEntry, active.id) + assert completed_entry is not None + assert active_entry is not None + completed_entry.timestamp = datetime(2026, 8, 1, tzinfo=timezone.utc) + active_entry.timestamp = datetime(2026, 8, 10, tzinfo=timezone.utc) + session.commit() + + headers = sqlite_instance.get_scenario_result_headers(limit=1) + + assert headers[0].scenario_name == "active" diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index d03b413f78..ad75bc73e7 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -313,6 +313,31 @@ def test_handles_empty_attack_results(sqlite_instance: MemoryInterface): assert len(results[0].attack_results) == 0 +def test_terminal_state_updates_completion_time_only_on_terminal_transition( + sqlite_instance: MemoryInterface, +) -> None: + old_completion = datetime(2020, 1, 1, tzinfo=timezone.utc) + scenario_result = create_scenario_result(name="Timing Scenario") + scenario_result.completion_time = old_completion + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + + sqlite_instance.update_scenario_run_state( + scenario_result_id=str(scenario_result.id), + scenario_run_state=ScenarioRunState.IN_PROGRESS, + ) + in_progress = sqlite_instance.get_scenario_result_header(scenario_result_id=str(scenario_result.id)) + assert in_progress is not None + assert in_progress.completion_time == old_completion + + sqlite_instance.update_scenario_run_state( + scenario_result_id=str(scenario_result.id), + scenario_run_state=ScenarioRunState.COMPLETED, + ) + completed = sqlite_instance.get_scenario_result_header(scenario_result_id=str(scenario_result.id)) + assert completed is not None + assert completed.completion_time > old_completion + + def test_preserves_metadata(sqlite_instance: MemoryInterface): """Test that scenario metadata is preserved correctly.""" diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 45ccc895dd..d0602588f7 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -174,6 +174,21 @@ def test_run_schema_migrations_applies_head_revision(): engine.dispose() +def test_scenario_progress_migration_adds_composite_index(): + """The migration head contains the parent/timestamp/id keyset index.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "scenario-progress-index.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "head") + indexes = {index["name"] for index in inspect(connection).get_indexes("AttackResultEntries")} + assert "ix_AttackResultEntries_attribution_parent_timestamp_id" in indexes + finally: + engine.dispose() + + def test_migration_online_mode(): """ Test that online migration configuration is valid. diff --git a/tests/unit/models/test_attack_seed_group.py b/tests/unit/models/test_attack_seed_group.py index c3c6ff1c9c..567337d0a5 100644 --- a/tests/unit/models/test_attack_seed_group.py +++ b/tests/unit/models/test_attack_seed_group.py @@ -4,6 +4,7 @@ import pytest +from pyrit.models import AtomicAttackIdentifier, ComponentIdentifier from pyrit.models.seeds.attack_seed_group import AttackSeedGroup from pyrit.models.seeds.seed_objective import SeedObjective from pyrit.models.seeds.seed_prompt import SeedPrompt @@ -59,6 +60,40 @@ def test_attack_seed_group_consistent_group_id(): assert None not in group_ids +def test_logical_id_ignores_random_prompt_group_id_and_round_trips() -> None: + first = AttackSeedGroup(seeds=[_make_objective(value="goal"), _make_prompt(value="context")]) + second = AttackSeedGroup(seeds=[_make_objective(value="goal"), _make_prompt(value="context")]) + + assert first.seeds[0].prompt_group_id != second.seeds[0].prompt_group_id + assert first.logical_id == second.logical_id + + identifier = AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="Attack", class_module="tests"), + seed_group=first, + ) + restored = AtomicAttackIdentifier.model_validate(identifier.model_dump(mode="json")) + assert restored.logical_seed_group_id == first.logical_id + + +def test_logical_id_preserves_canonical_seed_order() -> None: + first = AttackSeedGroup( + seeds=[ + _make_objective(value="goal"), + _make_prompt(value="first", sequence=0), + _make_prompt(value="second", sequence=1), + ] + ) + second = AttackSeedGroup( + seeds=[ + _make_objective(value="goal"), + _make_prompt(value="second", sequence=0), + _make_prompt(value="first", sequence=1), + ] + ) + + assert first.logical_id != second.logical_id + + def test_attack_seed_group_with_multiple_prompts(): objective = _make_objective() p1 = _make_prompt(value="p1", sequence=0) diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py new file mode 100644 index 0000000000..f787eed043 --- /dev/null +++ b/tests/unit/models/test_scenario_catalog.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for canonical scenario catalog models.""" + +import pytest +from pydantic import ValidationError + +from pyrit.models import ( + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, +) + + +def test_run_size_estimate_requires_available_total_to_match_components() -> None: + """Available estimates require an additive component total.""" + with pytest.raises(ValidationError, match="components total 6, not 7"): + ScenarioRunSizeEstimate( + estimated_attack_count=7, + components=[ScenarioRunSizeComponent(label="Techniques", count=6)], + ) + + +def test_run_size_estimate_allows_unavailable_count_with_components() -> None: + """Unavailable estimates retain useful candidate components and an explanatory note.""" + estimate = ScenarioRunSizeEstimate( + components=[ScenarioRunSizeComponent(label="Candidate techniques", count=6)], + note="The final count depends on target capabilities.", + ) + + assert estimate.estimated_attack_count is None + assert estimate.components[0].count == 6 + + +def test_run_size_estimate_serializes_canonical_api_shape() -> None: + """The estimate exposes only the available count and additive components.""" + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=6, + components=[ScenarioRunSizeComponent(label="Techniques", count=6)], + ) + + assert estimate.model_dump(mode="json") == { + "estimated_attack_count": 6, + "components": [ + { + "label": "Techniques", + "count": 6, + "note": None, + "is_baseline": False, + } + ], + "datasets": [], + "note": None, + } + + +def test_unavailable_run_size_estimate_has_no_count() -> None: + """The unavailable factory communicates that a count cannot be calculated.""" + estimate = ScenarioRunSizeEstimate.unavailable() + + assert estimate.estimated_attack_count is None + assert estimate.note == "Default-run size estimate is unavailable." + + +def test_estimate_exposes_dataset_counts_structurally() -> None: + """Effective dataset selection remains machine-readable.""" + estimate = ScenarioRunSizeEstimate( + datasets=[ + ScenarioDatasetSummary( + name="harmbench", + logical_seed_group_count=100, + selected_seed_group_count=4, + selection_note="The default selection uses 4 of 100 logical seed groups.", + configured_caps=[ + ScenarioDatasetSizeCap( + label="per-dataset cap", + count=4, + configured_on="dataset", + dataset_name="harmbench", + ) + ], + ) + ], + note="The final count depends on target capabilities.", + ) + + assert estimate.estimated_attack_count is None + assert estimate.model_dump(mode="json")["datasets"] == [ + { + "name": "harmbench", + "kind": "dataset", + "logical_seed_group_count": 100, + "selected_seed_group_count": 4, + "selection_note": "The default selection uses 4 of 100 logical seed groups.", + "configured_caps": [ + { + "label": "per-dataset cap", + "count": 4, + "configured_on": "dataset", + "dataset_name": "harmbench", + } + ], + } + ] + + +def test_estimate_request_reuses_dataset_filter_validation() -> None: + """Configured estimates reject the same unsupported dataset filters as launches.""" + with pytest.raises(ValidationError, match="Unknown dataset filter 'unknown'"): + ScenarioRunSizeEstimateRequest(dataset_filters={"unknown": ["value"]}) diff --git a/tests/unit/models/test_scenario_progress.py b/tests/unit/models/test_scenario_progress.py new file mode 100644 index 0000000000..65acbfe6af --- /dev/null +++ b/tests/unit/models/test_scenario_progress.py @@ -0,0 +1,41 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for scenario progress plan validation.""" + +import pytest +from pydantic import ValidationError + +from pyrit.models import ScenarioRunPlan, ScenarioRunPlanAtomicGroup, ScenarioRunPlanSeedGroup + + +def _seed(*, seed_id: str = "seed-1") -> ScenarioRunPlanSeedGroup: + return ScenarioRunPlanSeedGroup(id=seed_id, objective_sha256=f"sha-{seed_id}", objective=seed_id) + + +def _group(*, group_id: str = "group-1", seed_group_ids: list[str] | None = None) -> ScenarioRunPlanAtomicGroup: + return ScenarioRunPlanAtomicGroup( + id=group_id, + atomic_attack_name=group_id, + display_group=group_id, + technique_eval_hash=f"eval-{group_id}", + seed_group_ids=seed_group_ids or ["seed-1"], + ) + + +@pytest.mark.parametrize( + ("atomic_groups", "seed_groups", "match"), + [ + ([_group(), _group()], [_seed()], "duplicate atomic group IDs"), + ([_group()], [_seed(), _seed()], "duplicate seed group IDs"), + ([_group(seed_group_ids=["seed-1", "seed-1"])], [_seed()], "duplicate seed group IDs"), + ([_group(seed_group_ids=["missing"])], [_seed()], "unknown seed group IDs"), + ], +) +def test_run_plan_rejects_ambiguous_or_invalid_normalized_ids( + atomic_groups: list[ScenarioRunPlanAtomicGroup], + seed_groups: list[ScenarioRunPlanSeedGroup], + match: str, +) -> None: + with pytest.raises(ValidationError, match=match): + ScenarioRunPlan(atomic_groups=atomic_groups, seed_groups=seed_groups) diff --git a/tests/unit/registry/test_registry_metadata.py b/tests/unit/registry/test_registry_metadata.py index a5599a8b18..40007babfe 100644 --- a/tests/unit/registry/test_registry_metadata.py +++ b/tests/unit/registry/test_registry_metadata.py @@ -46,6 +46,38 @@ class NoDoc: assert result == "" +class TestMarkdownFromDocstring: + """Tests for structurally preserved catalog descriptions.""" + + def test_preserves_markdown_and_untrusted_html_as_source_text(self) -> None: + class MarkdownDoc: + """ + First paragraph with ``literal`` text. + + - First item + - [Split link]( + https://example.com) + + + """ + + result = RegistryMetadata.markdown_from_docstring(MarkdownDoc) + + assert result == ( + "First paragraph with ``literal`` text.\n\n" + "- First item\n" + "- [Split link](\n" + " https://example.com)\n\n" + '' + ) + + def test_returns_fallback_for_missing_docstring(self) -> None: + class NoDoc: + pass + + assert RegistryMetadata.markdown_from_docstring(NoDoc, fallback="fallback") == "fallback" + + class TestMatchesFilters: """Tests for the _matches_filters function.""" diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py index 209fc70381..393d6c5c17 100644 --- a/tests/unit/registry/test_scenario_registry.py +++ b/tests/unit/registry/test_scenario_registry.py @@ -8,6 +8,7 @@ import pytest from pyrit.registry.components.scenario_registry import ScenarioRegistry +from pyrit.scenario.core import BaselineAttackPolicy, ScenarioTechnique class _NotNoArgScenario: @@ -21,6 +22,58 @@ def __init__(self, *, required_arg) -> None: self.required_arg = required_arg +class _MetadataTechnique(ScenarioTechnique): + """Technique catalog for metadata expansion.""" + + ALL = ("all", {"all"}) + DEFAULT = ("default", {"default"}) + ONE = ("one", {"default"}) + TWO = ("two", {"default"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + """Return aggregate tags.""" + return {"all", "default"} + + @classmethod + def default(cls) -> "_MetadataTechnique": + """Return the default aggregate.""" + return cls.DEFAULT + + +class _MetadataScenario: + """Minimal scenario-shaped metadata source.""" + + BASELINE_ATTACK_POLICY = BaselineAttackPolicy.Enabled + + @classmethod + def supported_parameters(cls): + """Return no custom parameters.""" + return [] + + def __init__(self) -> None: + self._version = 1 + self._technique_class = _MetadataTechnique + self._default_technique = _MetadataTechnique.DEFAULT + self._default_dataset_config = MagicMock(dataset_names=["sample"]) + + def _resolve_scenario_techniques(self, *, scenario_techniques): + """Resolve the concrete defaults.""" + return _MetadataTechnique.resolve(scenario_techniques, default=self._default_technique) + + +class _MarkdownMetadataScenario(_MetadataScenario): + """ + First paragraph with ``literal`` text. + + - Item one + - [Split link]( + https://example.com) + + + """ + + def test_build_metadata_raises_when_scenario_requires_constructor_args() -> None: """Scenarios that cannot be instantiated with no args must surface a clear error.""" registry = ScenarioRegistry() @@ -29,6 +82,32 @@ def test_build_metadata_raises_when_scenario_requires_constructor_args() -> None registry._build_metadata("not_no_arg", _NotNoArgScenario) +def test_build_metadata_expands_ordered_default_techniques() -> None: + """Catalog metadata exposes concrete defaults rather than only the aggregate name.""" + metadata = ScenarioRegistry()._build_metadata("sample", _MetadataScenario) + + assert metadata.default_technique == "default" + assert metadata.default_techniques == ("one", "two") + assert dict(metadata.aggregate_technique_expansions) == { + "all": ("one", "two"), + "default": ("one", "two"), + } + + +def test_build_metadata_preserves_structured_markdown_separately() -> None: + """Scenario metadata keeps plain compatibility text and Markdown source.""" + metadata = ScenarioRegistry()._build_metadata("markdown", _MarkdownMetadataScenario) + + assert "\n" not in metadata.class_description + assert metadata.description_markdown == ( + "First paragraph with ``literal`` text.\n\n" + "- Item one\n" + "- [Split link](\n" + " https://example.com)\n\n" + '' + ) + + async def test_create_and_initialize_async_creates_sets_params_and_initializes() -> None: """The registry owns build + set-params + initialize and returns the scenario.""" registry = ScenarioRegistry() @@ -49,12 +128,42 @@ async def test_create_and_initialize_async_creates_sets_params_and_initializes() assert result is scenario registry.create_instance.assert_called_once_with("my.scenario", scenario_result_id="sr-1") + scenario.set_scenario_registry_name.assert_called_once_with(scenario_registry_name="my.scenario") scenario.set_params_from_args.assert_called_once_with( args={"foo": "bar", "objective_target": target, "max_concurrency": 2} ) scenario.initialize_async.assert_awaited_once_with() +async def test_create_and_estimate_async_configures_without_initializing() -> None: + """Configured estimation uses the registry parameter lifecycle without creating a run.""" + registry = ScenarioRegistry() + scenario = MagicMock() + estimate = MagicMock() + scenario.get_run_size_estimate_async = AsyncMock(return_value=estimate) + registry.create_instance = MagicMock(return_value=scenario) # type: ignore[method-assign] + + result = await registry.create_and_estimate_async( + name="my.scenario", + scenario_params={"num_jailbreaks": 2}, + scenario_techniques=["prompt_sending"], + include_baseline=False, + ) + + assert result is estimate + registry.create_instance.assert_called_once_with("my.scenario") + scenario.set_scenario_registry_name.assert_called_once_with(scenario_registry_name="my.scenario") + scenario.set_params_from_args.assert_called_once_with( + args={ + "num_jailbreaks": 2, + "scenario_techniques": ["prompt_sending"], + "include_baseline": False, + } + ) + scenario.get_run_size_estimate_async.assert_awaited_once_with(target_is_configured=False) + scenario.initialize_async.assert_not_called() + + async def test_create_and_initialize_async_omits_result_id_when_none() -> None: """When no scenario_result_id is supplied, it is not forwarded to the constructor.""" registry = ScenarioRegistry() @@ -67,5 +176,6 @@ async def test_create_and_initialize_async_omits_result_id_when_none() -> None: await registry.create_and_initialize_async("my.scenario", objective_target=target) registry.create_instance.assert_called_once_with("my.scenario") + scenario.set_scenario_registry_name.assert_called_once_with(scenario_registry_name="my.scenario") scenario.set_params_from_args.assert_called_once_with(args={"objective_target": target}) scenario.initialize_async.assert_awaited_once_with() diff --git a/tests/unit/scenario/airt/test_cyber.py b/tests/unit/scenario/airt/test_cyber.py index d29d3b5caa..7a94fa653c 100644 --- a/tests/unit/scenario/airt/test_cyber.py +++ b/tests/unit/scenario/airt/test_cyber.py @@ -8,7 +8,7 @@ import pytest from pyrit.executor.attack import RedTeamingAttack -from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, SeedPrompt +from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, SeedPrompt, TargetIdentifier from pyrit.prompt_target import PromptTarget from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration @@ -27,6 +27,10 @@ def _mock_id(name: str) -> ComponentIdentifier: return ComponentIdentifier(class_name=name, class_module="test") +def _mock_target_id(name: str) -> TargetIdentifier: + return TargetIdentifier(class_name=name, class_module="test") + + def _technique_class(): """Get the dynamically-generated CyberTechnique class.""" from pyrit.scenario.scenarios.airt.cyber import _build_cyber_technique @@ -42,14 +46,14 @@ def _technique_class(): @pytest.fixture def mock_objective_target(): mock = MagicMock(spec=PromptTarget) - mock.get_identifier.return_value = _mock_id("MockObjectiveTarget") + mock.get_identifier.return_value = _mock_target_id("MockObjectiveTarget") return mock @pytest.fixture def mock_adversarial_target(): mock = MagicMock(spec=PromptTarget) - mock.get_identifier.return_value = _mock_id("MockAdversarialTarget") + mock.get_identifier.return_value = _mock_target_id("MockAdversarialTarget") return mock @@ -78,6 +82,7 @@ def reset_technique_registry(): adv_target = MagicMock(spec=PromptTarget) adv_target.capabilities.includes.return_value = True + adv_target.get_identifier.return_value = _mock_target_id("MockAdversarialTarget") target_registry = TargetRegistry.get_registry_singleton() target_registry.instances.register(adv_target, name="adversarial_chat") diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index b055e087aa..cf0f25c436 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -12,7 +12,12 @@ from pyrit.converter import TextJailbreakConverter from pyrit.datasets import TextJailBreak from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack -from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, SeedPrompt +from pyrit.models import ( + AttackSeedGroup, + ComponentIdentifier, + SeedObjective, + SeedPrompt, +) from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -197,6 +202,55 @@ async def test_num_jailbreaks_samples_that_many( await scenario.initialize_async() assert len(scenario._resolved_jailbreaks) == 3 + async def test_run_size_prompt_sending_two_templates_four_groups_is_eight( + self, mock_objective_target, mock_objective_scorer + ) -> None: + """The launch-aligned GUI selection has exactly eight persisted outer units.""" + seed_groups = [AttackSeedGroup(seeds=[SeedObjective(value=f"objective {index}")]) for index in range(4)] + technique_class = _build_jailbreak_technique() + with _patch_seed_groups(seed_groups): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_techniques": [technique_class(_PROMPT_SENDING)], + "include_baseline": False, + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + ) + + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + assert estimate.estimated_attack_count == 8 + assert [component.label for component in estimate.components] == ["Inline jailbreak delivery"] + assert estimate.datasets[0].logical_seed_group_count == 4 + assert estimate.datasets[0].selected_seed_group_count == 4 + assert [(cap.label, cap.count) for cap in estimate.datasets[0].configured_caps] == [("per-dataset cap", 4)] + + async def test_run_size_is_conditional_when_system_delivery_target_is_not_selected( + self, mock_objective_scorer + ) -> None: + """The default system-prompt axis does not claim a total before target capability is known.""" + seed_groups = [AttackSeedGroup(seeds=[SeedObjective(value="objective")])] + technique_class = _build_jailbreak_technique() + with _patch_seed_groups(seed_groups): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args={ + "scenario_techniques": [technique_class("default")], + "include_baseline": False, + "num_jailbreaks": 2, + } + ) + + estimate = await scenario.get_run_size_estimate_async(target_is_configured=False) + assert estimate.estimated_attack_count is None + assert [component.label for component in estimate.components] == [ + "Inline jailbreak delivery", + "Native system-prompt jailbreak delivery", + ] + assert "native system-prompt delivery is supported" in (estimate.note or "") + async def test_mutually_exclusive_selectors_raise( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): diff --git a/tests/unit/scenario/airt/test_rapid_response.py b/tests/unit/scenario/airt/test_rapid_response.py index 66259563e0..c32c63f483 100644 --- a/tests/unit/scenario/airt/test_rapid_response.py +++ b/tests/unit/scenario/airt/test_rapid_response.py @@ -14,7 +14,7 @@ PromptSendingAttack, TreeOfAttacksWithPruningAttack, ) -from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective +from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, TargetIdentifier from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -41,6 +41,10 @@ def _mock_id(name: str) -> ComponentIdentifier: return ComponentIdentifier(class_name=name, class_module="test") +def _mock_target_id(name: str) -> TargetIdentifier: + return TargetIdentifier(class_name=name, class_module="test") + + def _technique_class(): """Get the dynamically-generated RapidResponseTechnique class.""" from pyrit.scenario.scenarios.airt.rapid_response import _build_rapid_response_technique @@ -56,14 +60,14 @@ def _technique_class(): @pytest.fixture def mock_objective_target(): mock = MagicMock(spec=PromptTarget) - mock.get_identifier.return_value = _mock_id("MockObjectiveTarget") + mock.get_identifier.return_value = _mock_target_id("MockObjectiveTarget") return mock @pytest.fixture def mock_adversarial_target(): mock = MagicMock(spec=PromptTarget) - mock.get_identifier.return_value = _mock_id("MockAdversarialTarget") + mock.get_identifier.return_value = _mock_target_id("MockAdversarialTarget") return mock @@ -90,6 +94,7 @@ def reset_technique_registry(): adv_target = MagicMock(spec=PromptTarget) adv_target.capabilities.includes.return_value = True + adv_target.get_identifier.return_value = _mock_target_id("MockAdversarialTarget") TargetRegistry.get_registry_singleton().instances.register(adv_target, name="adversarial_chat") technique_registry = AttackTechniqueRegistry.get_registry_singleton() diff --git a/tests/unit/scenario/core/test_atomic_attack.py b/tests/unit/scenario/core/test_atomic_attack.py index 8a04f20457..30df69492f 100644 --- a/tests/unit/scenario/core/test_atomic_attack.py +++ b/tests/unit/scenario/core/test_atomic_attack.py @@ -1119,7 +1119,7 @@ async def test_no_attribution_when_scenario_result_id_unset( self, mock_attack, sample_seed_groups, sample_attack_results ): """Outside a Scenario, ``_scenario_result_id`` is None and the - executor must receive ``attribution=None``.""" + executor must receive ``attributions=None``.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, @@ -1131,13 +1131,13 @@ async def test_no_attribution_when_scenario_result_id_unset( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - assert mock_exec.call_args.kwargs["attribution"] is None + assert mock_exec.call_args.kwargs["attributions"] is None async def test_attribution_built_when_scenario_result_id_set( self, mock_attack, sample_seed_groups, sample_attack_results ): """When the Scenario stamps ``_scenario_result_id`` onto the atomic - attack, ``run_async`` must build and pass a single attribution object.""" + attack, ``run_async`` must build and pass per-seed-group attribution.""" from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution atomic = AtomicAttack( @@ -1151,10 +1151,14 @@ async def test_attribution_built_when_scenario_result_id_set( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - attribution = mock_exec.call_args.kwargs["attribution"] - assert isinstance(attribution, AttackResultAttribution) - assert attribution.parent_id == "00000000-0000-0000-0000-000000000abc" - assert attribution.parent_collection == "MyAtomicAttack" + attributions = mock_exec.call_args.kwargs["attributions"] + assert len(attributions) == len(sample_seed_groups) + assert all(isinstance(attribution, AttackResultAttribution) for attribution in attributions) + assert all(attribution.parent_id == "00000000-0000-0000-0000-000000000abc" for attribution in attributions) + assert all(attribution.parent_collection == "MyAtomicAttack" for attribution in attributions) + assert [attribution.seed_group_id for attribution in attributions] == [ + seed_group.logical_id for seed_group in sample_seed_groups + ] async def test_attribution_includes_technique_eval_hash( self, mock_attack, sample_seed_groups, sample_attack_results @@ -1173,9 +1177,9 @@ async def test_attribution_includes_technique_eval_hash( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - attribution = mock_exec.call_args.kwargs["attribution"] - assert attribution.parent_eval_hash is not None - assert attribution.parent_eval_hash == atomic.technique_eval_hash + attributions = mock_exec.call_args.kwargs["attributions"] + assert all(attribution.parent_eval_hash is not None for attribution in attributions) + assert all(attribution.parent_eval_hash == atomic.technique_eval_hash for attribution in attributions) @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index 5ea41cd64b..c5324d8d39 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -207,6 +207,29 @@ def test_cannot_append_text_converter_to_image_chain(self): assert not factory.can_append_request_converter(converter_type=TranslationConverter) + def test_request_converter_composition_requires_supported_constructor(self): + class _NoConverterAttack: + def __init__(self, *, objective_target, attack_scoring_config=None): + self.objective_target = objective_target + + with pytest.raises(ValueError, match="does not accept 'attack_converter_config'"): + AttackTechniqueFactory( + name="test", + attack_class=_NoConverterAttack, + supports_additional_request_converters=True, + ) + + def test_request_converter_composition_is_explicit_opt_in(self): + default_factory = AttackTechniqueFactory(name="default", attack_class=_StubAttack) + composable_factory = AttackTechniqueFactory( + name="composable", + attack_class=_StubAttack, + supports_additional_request_converters=True, + ) + + assert not default_factory.supports_additional_request_converters + assert composable_factory.supports_additional_request_converters + class TestFactoryCreate: """Tests for AttackTechniqueFactory.create().""" diff --git a/tests/unit/scenario/core/test_dataset_configuration.py b/tests/unit/scenario/core/test_dataset_configuration.py index 14e914b4d5..3c3ad9c25c 100644 --- a/tests/unit/scenario/core/test_dataset_configuration.py +++ b/tests/unit/scenario/core/test_dataset_configuration.py @@ -17,6 +17,7 @@ DatasetSourceKind, ResolvedDataset, forbid_inline_seeds, + read_only_dataset_resolution, require_harm_categories, require_inline_seeds, require_min_size, @@ -327,6 +328,22 @@ async def test_fetch_failure_chains_root_cause(self, mock_memory: MagicMock) -> await config.get_attack_seed_groups_async() assert isinstance(exc_info.value.__cause__, RuntimeError) + async def test_read_only_resolution_does_not_fetch_or_persist(self, mock_memory: MagicMock) -> None: + """Estimate resolution reports missing data without mutating central memory.""" + config = DatasetAttackConfiguration(dataset_names=["d1"]) + with ( + patch(PROVIDER_PATCH_TARGET) as provider, + read_only_dataset_resolution(), + pytest.raises(DatasetConstraintError, match="read-only resolution"), + ): + provider.get_all_dataset_names_async = AsyncMock(return_value=["d1"]) + provider.fetch_datasets_async = AsyncMock() + await config.get_attack_seed_groups_async() + + provider.get_all_dataset_names_async.assert_not_awaited() + provider.fetch_datasets_async.assert_not_awaited() + mock_memory.add_seed_datasets_to_memory_async.assert_not_awaited() + class TestValidators: """The standalone validator builders and base ``validate``.""" @@ -490,6 +507,16 @@ def test_per_dataset_builds_one_child_per_name(self) -> None: assert [child.dataset_names for child in config._configurations] == [["d1"], ["d2"]] assert all(child.max_dataset_size == 4 for child in config._configurations) + def test_size_caps_report_child_and_combined_limits(self) -> None: + """Planning metadata explains independent child caps and the final compound cap.""" + config = CompoundDatasetAttackConfiguration.per_dataset(dataset_names=["d1", "d2"], max_dataset_size=4) + config.max_dataset_size = 6 + + assert config.size_caps_by_dataset() == { + "d1": [("per-dataset cap", 4, "dataset"), ("combined compound cap", 6, "compound")], + "d2": [("per-dataset cap", 4, "dataset"), ("combined compound cap", 6, "compound")], + } + def test_dataset_names_aggregates_and_dedups(self) -> None: config = CompoundDatasetAttackConfiguration( configurations=[ diff --git a/tests/unit/scenario/core/test_scenario.py b/tests/unit/scenario/core/test_scenario.py index d5ce2ae34e..e8d15c4ce1 100644 --- a/tests/unit/scenario/core/test_scenario.py +++ b/tests/unit/scenario/core/test_scenario.py @@ -16,7 +16,16 @@ from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, ScenarioRunState +from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + ScenarioRunState, + SeedObjective, + SeedPrompt, +) from pyrit.prompt_target import PromptTarget from pyrit.scenario import ( DatasetAttackConfiguration, @@ -43,6 +52,16 @@ def save_attack_results_to_memory(attack_results): memory.add_attack_results_to_memory(attack_results=attack_results) +def _make_identifiable_mock_attack() -> MagicMock: + """Create a mock attack with a valid canonical identifier for run-plan construction.""" + attack = MagicMock() + attack.get_identifier.return_value = ComponentIdentifier( + class_name="MockAttack", + class_module="tests.unit.scenario.core.test_scenario", + ) + return attack + + def _stamp_scenario_linkage(*, attack_results, atomic_attack): """ Stamp attribution_parent_id + attribution_data on each AttackResult the @@ -260,6 +279,61 @@ async def test_initialize_async_populates_atomic_attacks(self, mock_atomic_attac assert scenario.atomic_attack_count == len(mock_atomic_attacks) assert scenario._atomic_attacks == mock_atomic_attacks + [stored] = scenario._memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id]) + assert stored.metadata["run_plan"]["version"] == 1 + assert len(stored.metadata["run_plan"]["atomic_groups"]) == len(mock_atomic_attacks) + + async def test_initialize_async_deduplicates_logical_seed_groups_in_run_plan(self, mock_objective_target) -> None: + duplicate_seed_groups = [ + AttackSeedGroup(seeds=[SeedObjective(value="duplicate objective")]), + AttackSeedGroup(seeds=[SeedObjective(value="duplicate objective")]), + ] + atomic_attack = MagicMock(spec=AtomicAttack) + atomic_attack.atomic_attack_name = "duplicate_attack" + atomic_attack.display_group = "duplicate_attack" + atomic_attack.technique_eval_hash = "duplicate-technique" + type(atomic_attack).seed_groups = PropertyMock(return_value=duplicate_seed_groups) + scenario = ConcreteScenario( + name="Duplicate Seed Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.initialize_async() + + [stored] = scenario._memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id]) + persisted_plan = stored.metadata[SCENARIO_RUN_PLAN_METADATA_KEY] + expected_seed_id = duplicate_seed_groups[0].logical_id + assert persisted_plan["atomic_groups"][0]["seed_group_ids"] == [expected_seed_id] + assert [seed_group["id"] for seed_group in persisted_plan["seed_groups"]] == [expected_seed_id] + assert scenario._build_run_plan().model_dump(mode="json") == persisted_plan + assert atomic_attack.seed_groups is duplicate_seed_groups + assert len(atomic_attack.seed_groups) == 2 + + async def test_build_run_plan_preserves_unique_seed_group_order(self, mock_objective_target) -> None: + seed_groups = [ + AttackSeedGroup(seeds=[SeedObjective(value="first objective")]), + AttackSeedGroup(seeds=[SeedObjective(value="second objective")]), + ] + atomic_attack = MagicMock(spec=AtomicAttack) + atomic_attack.atomic_attack_name = "unique_attack" + atomic_attack.display_group = "unique_attack" + atomic_attack.technique_eval_hash = "unique-technique" + type(atomic_attack).seed_groups = PropertyMock(return_value=seed_groups) + scenario = ConcreteScenario( + name="Unique Seed Scenario", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.initialize_async() + + plan = scenario._build_run_plan() + expected_seed_ids = [seed_group.logical_id for seed_group in seed_groups] + assert plan.atomic_groups[0].seed_group_ids == expected_seed_ids + assert [seed_group.id for seed_group in plan.seed_groups] == expected_seed_ids async def test_initialize_async_sets_objective_target(self, mock_objective_target): """Test that initialize_async sets objective_target properly.""" @@ -420,6 +494,39 @@ async def test_run_async_executes_all_runs(self, mock_atomic_attacks, sample_att assert result.attack_results["attack_run_1"][0] == sample_attack_results[0] assert result.attack_results["attack_run_2"][0] == sample_attack_results[1] assert result.attack_results["attack_run_3"][0] == sample_attack_results[2] + assert scenario.active_atomic_group_ids == frozenset() + + async def test_active_atomic_group_is_cleared_when_execution_is_cancelled( + self, + mock_atomic_attacks, + mock_objective_target, + ): + started = asyncio.Event() + blocked = asyncio.Event() + + async def run_until_cancelled(**_kwargs): + started.set() + await blocked.wait() + + atomic_attack = mock_atomic_attacks[0] + atomic_attack.run_async = AsyncMock(side_effect=run_until_cancelled) + scenario = ConcreteScenario( + name="Cancellation cleanup", + version=1, + atomic_attacks_to_return=[atomic_attack], + ) + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.initialize_async() + + task = asyncio.create_task(scenario.run_async()) + await started.wait() + assert scenario.active_atomic_group_ids + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert scenario.active_atomic_group_ids == frozenset() async def test_run_async_with_custom_concurrency( self, mock_atomic_attacks, sample_attack_results, mock_objective_target @@ -513,6 +620,7 @@ async def test_run_async_stops_on_error(self, mock_atomic_attacks, sample_attack mock_atomic_attacks[1].run_async.assert_called_once() # Third run should not have been executed (worker stops pulling after failure) mock_atomic_attacks[2].run_async.assert_not_called() + assert scenario.active_atomic_group_ids == frozenset() async def test_run_async_fails_without_initialization(self, mock_objective_target): """Test that run_async fails if initialize_async was not called.""" @@ -1041,7 +1149,7 @@ async def _build_atomic_attacks_async(self, *, context): attacks.append( AtomicAttack( atomic_attack_name="technique", - attack_technique=AttackTechnique(attack=MagicMock()), + attack_technique=AttackTechnique(attack=_make_identifiable_mock_attack()), seed_groups=list(context.seed_groups), ) ) @@ -1105,7 +1213,7 @@ async def _build_atomic_attacks_async(self, *, context): attacks.append( AtomicAttack( atomic_attack_name="strategy", - attack_technique=AttackTechnique(attack=MagicMock()), + attack_technique=AttackTechnique(attack=_make_identifiable_mock_attack()), seed_groups=list(context.seed_groups), ) ) @@ -1128,7 +1236,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list attacks.extend( AtomicAttack( atomic_attack_name=f"strategy-{index}", - attack_technique=AttackTechnique(attack=MagicMock()), + attack_technique=AttackTechnique(attack=_make_identifiable_mock_attack()), seed_groups=[seed_group], ) for index, seed_group in enumerate(context.seed_groups) @@ -1164,6 +1272,9 @@ def _sample_first_k(population, k): original_id = scenario._scenario_result_id assert original_id is not None + original_header = scenario._memory.get_scenario_result_header(scenario_result_id=original_id) + assert original_header is not None + original_plan = original_header.metadata[SCENARIO_RUN_PLAN_METADATA_KEY] _, first_strategy = scenario._atomic_attacks persisted_objectives = set(first_strategy.objectives) assert persisted_objectives == {"obj0", "obj1", "obj2"} @@ -1202,6 +1313,100 @@ def _sample_last_k(population, k): # Exactly the originally-persisted subset, not the divergent "last 3" draw. assert set(strategy.objectives) == persisted_objectives assert set(baseline.objectives) == persisted_objectives + resumed_header = resumed._memory.get_scenario_result_header(scenario_result_id=original_id) + assert resumed_header is not None + assert resumed_header.metadata[SCENARIO_RUN_PLAN_METADATA_KEY] == original_plan + + async def test_resume_rejects_changed_companion_seed_with_same_objective(self, mock_objective_target): + objective = "unchanged objective" + original_seed_group = AttackSeedGroup( + seeds=[SeedObjective(value=objective), SeedPrompt(value="original context")] + ) + scenario = self._StrategyScenario(name="Changed seed-group resume", version=1) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": DatasetAttackConfiguration(seed_groups=[original_seed_group]), + "include_baseline": False, + } + ) + await scenario.initialize_async() + + scenario_result_id = scenario._scenario_result_id + assert scenario_result_id is not None + header = scenario._memory.get_scenario_result_header(scenario_result_id=scenario_result_id) + assert header is not None + persisted_plan = header.metadata[SCENARIO_RUN_PLAN_METADATA_KEY] + assert persisted_plan["atomic_groups"][0]["seed_group_ids"] == [original_seed_group.logical_id] + + changed_seed_group = AttackSeedGroup( + seeds=[SeedObjective(value=objective), SeedPrompt(value="changed context")] + ) + assert changed_seed_group.logical_id != original_seed_group.logical_id + resumed = self._StrategyScenario( + name="Changed seed-group resume", + version=1, + scenario_result_id=scenario_result_id, + ) + resumed.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": DatasetAttackConfiguration(seed_groups=[changed_seed_group]), + "include_baseline": False, + } + ) + + with pytest.raises( + ValueError, + match=r"cannot resume: atomic group 'strategy' is missing 1 planned seed group", + ): + await resumed.initialize_async() + + async def test_resume_reconstructs_plan_for_legacy_resumable_run(self, mock_objective_target): + config = self._make_config() + with patch( + "pyrit.scenario.core.dataset_configuration.random.sample", + side_effect=lambda population, k: list(population)[:k], + ): + scenario = self._StrategyScenario(name="Legacy resume", version=1) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": config, + } + ) + await scenario.initialize_async() + + scenario_result_id = scenario._scenario_result_id + assert scenario_result_id is not None + header = scenario._memory.get_scenario_result_header(scenario_result_id=scenario_result_id) + assert header is not None + legacy_metadata = dict(header.metadata) + legacy_metadata.pop(SCENARIO_RUN_PLAN_METADATA_KEY) + scenario._memory.update_scenario_metadata( + scenario_result_id=scenario_result_id, + metadata=legacy_metadata, + ) + + resumed = self._StrategyScenario( + name="Legacy resume", + version=1, + scenario_result_id=scenario_result_id, + ) + resumed.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "scenario_strategies": None, + "dataset_config": self._make_config(), + } + ) + await resumed.initialize_async() + + reconstructed = resumed._memory.get_scenario_result_header(scenario_result_id=scenario_result_id) + assert reconstructed is not None + assert SCENARIO_RUN_PLAN_METADATA_KEY in reconstructed.metadata + assert reconstructed.metadata["objective_hashes"] == legacy_metadata["objective_hashes"] async def test_resume_discards_per_objective_attacks_outside_persisted_subset(self, mock_objective_target): def _sample_first_k(population, k): diff --git a/tests/unit/scenario/core/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py index d6f58e902b..febe218e83 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -17,7 +17,15 @@ from pyrit.exceptions import ScenarioPartialFailureException from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, ScenarioRunState +from pyrit.models import ( + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + ScenarioRunState, + SeedObjective, + config_hash, +) from pyrit.prompt_target import PromptTarget from pyrit.scenario import DatasetConfiguration, ScenarioResult from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy, Scenario, ScenarioTechnique @@ -76,6 +84,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: attack = MagicMock(spec=AtomicAttack) attack.atomic_attack_name = name attack.display_group = name + attack.technique_eval_hash = config_hash({"name": name, "objectives": objectives}) attack._attack = mock_attack_strategy attack._scenario_result_id = None @@ -85,13 +94,21 @@ def _set_scenario_result_id(scenario_result_id): attack.set_scenario_result_id = MagicMock(side_effect=_set_scenario_result_id) original_objectives = list(objectives) - current_objectives = {"value": list(objectives)} + current_seed_groups = { + "value": [AttackSeedGroup(seeds=[SeedObjective(value=objective)]) for objective in objectives] + } - type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) - type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) + type(attack).objectives = PropertyMock( + side_effect=lambda: [seed_group.objective.value for seed_group in current_seed_groups["value"]] + ) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_seed_groups["value"]) def drop_hashes(*, hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] + current_seed_groups["value"] = [ + seed_group + for seed_group in current_seed_groups["value"] + if to_sha256(seed_group.objective.value) not in hashes + ] attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) attack._original_objectives = original_objectives diff --git a/tests/unit/scenario/core/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py index af4a222d0b..d29bd0cd22 100644 --- a/tests/unit/scenario/core/test_scenario_retry.py +++ b/tests/unit/scenario/core/test_scenario_retry.py @@ -12,7 +12,15 @@ from pyrit.executor.attack import AttackParameters, AttackStrategy, SingleTurnAttackContext from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, AttackSeedGroup, ComponentIdentifier, Message, SeedObjective +from pyrit.models import ( + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + Message, + SeedObjective, + config_hash, +) from pyrit.prompt_target import PromptTarget from pyrit.scenario import DatasetConfiguration, ScenarioResult from pyrit.scenario.core import AtomicAttack, AttackTechnique, BaselineAttackPolicy, Scenario, ScenarioTechnique @@ -139,6 +147,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack = MagicMock(spec=AtomicAttack) attack.atomic_attack_name = name attack.display_group = name + attack.technique_eval_hash = config_hash({"name": name, "objectives": objectives}) attack._attack = mock_attack_strategy attack._scenario_result_id = None @@ -151,12 +160,20 @@ def _set_scenario_result_id(scenario_result_id): # behaves correctly in resume tests. from pyrit.common.utils import to_sha256 - current_objectives = {"value": list(objectives)} - type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) - type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) + current_seed_groups = { + "value": [AttackSeedGroup(seeds=[SeedObjective(value=objective)]) for objective in objectives] + } + type(attack).objectives = PropertyMock( + side_effect=lambda: [seed_group.objective.value for seed_group in current_seed_groups["value"]] + ) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_seed_groups["value"]) def drop_hashes(*, hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] + current_seed_groups["value"] = [ + seed_group + for seed_group in current_seed_groups["value"] + if to_sha256(seed_group.objective.value) not in hashes + ] attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py new file mode 100644 index 0000000000..2d69edf96c --- /dev/null +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -0,0 +1,602 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for scenario-owned default-run size estimates.""" + +from typing import ClassVar +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.executor.attack.core.attack_config import AttackScoringConfig +from pyrit.models import ( + AttackSeedGroup, + AttackTechniqueSeedGroup, + ComponentIdentifier, + ScenarioDatasetSummary, + SeedObjective, + SeedPrompt, + SeedSimulatedConversation, +) +from pyrit.prompt_target import PromptTarget +from pyrit.scenario.core import BaselineAttackPolicy, DatasetAttackConfiguration, Scenario, ScenarioTechnique +from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive +from pyrit.scenario.scenarios.airt.jailbreak import Jailbreak +from pyrit.scenario.scenarios.airt.psychosocial import Psychosocial +from pyrit.scenario.scenarios.benchmark.adversarial import AdversarialBenchmark +from pyrit.scenario.scenarios.foundry.red_team_agent import FoundryComposite, FoundryTechnique, RedTeamAgent +from pyrit.scenario.scenarios.garak.encoding import Encoding +from pyrit.scenario.scenarios.garak.web_injection import WebInjection +from pyrit.score import TrueFalseScorer + + +class _TwoTechniqueDefault(ScenarioTechnique): + """Two concrete defaults used by estimate-only test scenarios.""" + + ALL = ("all", {"all"}) + DEFAULT = ("default", {"default"}) + ONE = ("one", {"default"}) + TWO = ("two", {"default"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + """Return aggregate tags.""" + return {"all", "default"} + + @classmethod + def default(cls) -> "_TwoTechniqueDefault": + """Return the default aggregate.""" + return cls.DEFAULT + + +class _JailbreakDefault(ScenarioTechnique): + """Jailbreak's two default delivery techniques.""" + + ALL = ("all", {"all"}) + DEFAULT = ("default", {"default"}) + PROMPT_SENDING = ("prompt_sending", {"default"}) + SYSTEM_PROMPT = ("jailbreak_system_prompt", {"default"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + """Return aggregate tags.""" + return {"all", "default"} + + @classmethod + def default(cls) -> "_JailbreakDefault": + """Return the default aggregate.""" + return cls.DEFAULT + + +class _MatrixEstimateScenario(Scenario): + """Minimal ordinary default technique sweep.""" + + BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Enabled + + def __init__(self, *, objective_scorer: TrueFalseScorer) -> None: + super().__init__( + version=1, + technique_class=_TwoTechniqueDefault, + default_dataset_config=DatasetAttackConfiguration(dataset_names=["sample"]), + objective_scorer=objective_scorer, + ) + + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[AttackSeedGroup]]: + """Return three logical groups before selection and two after.""" + if self._dataset_config.dataset_names == ["sample"]: + values = ["one", "two"] if apply_sampling else ["one", "two", "three"] + return {"sample": [_seed_group(value) for value in values]} + return await super()._resolve_seed_groups_by_dataset_async(apply_sampling=apply_sampling) + + async def _build_atomic_attacks_async(self, *, context): + """Return no attacks; only estimation is exercised.""" + return [] + + +class _CompatibilityMatrixEstimateScenario(_MatrixEstimateScenario): + """Matrix scenario whose estimates mirror execution compatibility filtering.""" + + RUN_SIZE_USES_FACTORY_COMPATIBILITY: ClassVar[bool] = True + + +def _scorer() -> MagicMock: + scorer = MagicMock(spec=TrueFalseScorer) + scorer.get_identifier.return_value = ComponentIdentifier(class_name="MockScorer", class_module="test") + return scorer + + +def _seed_group(value: str) -> AttackSeedGroup: + return AttackSeedGroup(seeds=[SeedObjective(value=value)]) + + +def _resolved_groups( + counts: dict[str, int], +) -> tuple[dict[str, list[AttackSeedGroup]], list[ScenarioDatasetSummary]]: + groups = {name: [_seed_group(f"{name}-{index}") for index in range(count)] for name, count in counts.items()} + summaries = [ + ScenarioDatasetSummary( + name=name, + logical_seed_group_count=count, + selected_seed_group_count=count, + ) + for name, count in counts.items() + ] + return groups, summaries + + +@pytest.mark.usefixtures("patch_central_database") +async def test_ordinary_matrix_estimate_uses_planned_seed_units_and_baseline() -> None: + """The base estimate is selected seed groups times concrete defaults plus baseline.""" + estimate = await _MatrixEstimateScenario(objective_scorer=_scorer()).get_default_run_size_estimate_async() + assert estimate.estimated_attack_count == 6 + assert [component.count for component in estimate.components] == [4, 2] + assert estimate.datasets[0].logical_seed_group_count == 3 + assert estimate.datasets[0].selected_seed_group_count == 2 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_configured_estimate_reuses_technique_and_baseline_resolution_without_persistence( + patch_central_database, +) -> None: + """A configured estimate expands only selected inputs and creates no ScenarioResult.""" + scenario = _MatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args( + args={ + "scenario_techniques": [_TwoTechniqueDefault.ONE], + "include_baseline": False, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 2 + assert [component.count for component in estimate.components] == [2] + assert patch_central_database.return_value.get_scenario_results() == [] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_configured_estimate_expands_requested_aggregate() -> None: + """Configured previews expand aggregate technique tokens through the scenario path.""" + scenario = _MatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args( + args={ + "scenario_techniques": [_TwoTechniqueDefault.DEFAULT], + "include_baseline": False, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 4 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_configured_estimate_applies_dataset_selection_and_cap() -> None: + """Configured estimates use the requested dataset population rather than scenario defaults.""" + scenario = _MatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args( + args={ + "dataset_config": DatasetAttackConfiguration( + seed_groups=[_seed_group("one"), _seed_group("two"), _seed_group("three")], + max_dataset_size=2, + ), + "scenario_techniques": [_TwoTechniqueDefault.ONE], + "include_baseline": False, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.estimated_attack_count == 2 + assert len(estimate.datasets) == 1 + assert estimate.datasets[0].logical_seed_group_count == 3 + assert estimate.datasets[0].selected_seed_group_count == 2 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_configured_estimate_exposes_nonbinding_cap_provenance() -> None: + """Configured caps remain visible even when they do not reduce the population.""" + scenario = _MatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args( + args={ + "dataset_config": DatasetAttackConfiguration( + seed_groups=[_seed_group(str(index)) for index in range(4)], + max_dataset_size=4, + ), + "scenario_techniques": [_TwoTechniqueDefault.ONE], + "include_baseline": False, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.datasets[0].logical_seed_group_count == 4 + assert estimate.datasets[0].selected_seed_group_count == 4 + assert [(cap.label, cap.count, cap.configured_on) for cap in estimate.datasets[0].configured_caps] == [ + ("per-dataset cap", 4, "dataset") + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_matrix_estimate_filters_each_technique_seed_population_like_execution() -> None: + """A mixed seed matrix does not use naive technique-by-group multiplication.""" + compatible = _seed_group("compatible") + incompatible = AttackSeedGroup( + seeds=[ + SeedObjective(value="incompatible"), + SeedPrompt(value="user", data_type="text", role="user", sequence=0), + SeedPrompt(value="assistant", data_type="text", role="assistant", sequence=1), + SeedPrompt(value="user again", data_type="text", role="user", sequence=2), + ] + ) + plain_factory = MagicMock() + plain_factory.seed_technique = None + conversation_factory = MagicMock() + conversation_factory.seed_technique = AttackTechniqueSeedGroup( + seeds=[ + SeedSimulatedConversation( + adversarial_chat_system_prompt_path="fake.yaml", + num_turns=3, + ) + ] + ) + scenario = _CompatibilityMatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args(args={"include_baseline": False}) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock( + return_value=( + {"sample": [compatible, incompatible]}, + [ + ScenarioDatasetSummary( + name="sample", + logical_seed_group_count=2, + selected_seed_group_count=2, + ) + ], + ) + ) + + with patch( + "pyrit.scenario.core.matrix_atomic_attack_builder.resolve_technique_factories_for_techniques", + return_value={"one": plain_factory, "two": conversation_factory}, + ): + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.estimated_attack_count == 3 + assert [(component.label, component.count) for component in estimate.components] == [("one", 2), ("two", 1)] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_matrix_estimate_with_binding_cap_and_compatibility_is_conditional() -> None: + """A randomized binding cap cannot promise the same compatibility mix at launch.""" + scenario = _CompatibilityMatrixEstimateScenario(objective_scorer=_scorer()) + scenario.set_params_from_args(args={"include_baseline": False}) + + async def resolve_groups() -> tuple[dict[str, list[AttackSeedGroup]], list[ScenarioDatasetSummary]]: + scenario._estimate_has_binding_size_cap = True + return _resolved_groups({"sample": 1}) + + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(side_effect=resolve_groups) + factory = MagicMock() + factory.seed_technique = None + + with patch( + "pyrit.scenario.core.matrix_atomic_attack_builder.resolve_technique_factories_for_techniques", + return_value={"one": factory, "two": factory}, + ): + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count is None + assert "binding randomized dataset cap" in estimate.note + + +@pytest.mark.usefixtures("patch_central_database") +async def test_adaptive_estimate_is_target_conditional_and_does_not_multiply_techniques() -> None: + """Adaptive techniques are selected internally rather than forming an outer axis.""" + with patch.object(TextAdaptive, "get_technique_class", return_value=_TwoTechniqueDefault): + scenario = TextAdaptive(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"adaptive": 3})) + + estimate = await scenario.get_default_run_size_estimate_async() + assert estimate.estimated_attack_count is None + assert [component.count for component in estimate.components] == [3, 3] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_adaptive_estimate_counts_exact_compatible_outer_envelopes_with_target() -> None: + """A concrete target makes the compatible outer population exact without counting attempts.""" + with patch.object(TextAdaptive, "get_technique_class", return_value=_TwoTechniqueDefault): + scenario = TextAdaptive(objective_scorer=_scorer()) + target = MagicMock(spec=PromptTarget) + scenario.set_params_from_args( + args={ + "objective_target": target, + "include_baseline": False, + "max_attempts_per_objective": 7, + } + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"adaptive": 3})) + dispatcher = MagicMock() + dispatcher.compatible_techniques.side_effect = [["one"], [], ["two"]] + + with ( + patch.object(scenario, "_build_techniques_dict", return_value={"one": MagicMock()}), + patch( + "pyrit.scenario.scenarios.adaptive.adaptive_scenario.AdaptiveTechniqueDispatcher", + return_value=dispatcher, + ), + ): + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 2 + assert [component.count for component in estimate.components] == [2] + assert "7 selected technique attempts" in estimate.note + + scenario.set_params_from_args(args={"include_baseline": False}) + estimate_without_target = await scenario.get_run_size_estimate_async() + assert estimate_without_target.estimated_attack_count is None + + +@pytest.mark.usefixtures("patch_central_database") +async def test_jailbreak_estimate_exposes_template_attempt_and_target_capability_axes() -> None: + """Jailbreak reports guaranteed inline work separately from conditional system delivery.""" + with patch("pyrit.scenario.scenarios.airt.jailbreak._build_jailbreak_technique", return_value=_JailbreakDefault): + scenario = Jailbreak(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) + + estimate = await scenario.get_default_run_size_estimate_async() + assert estimate.estimated_attack_count is None + assert [component.count for component in estimate.components] == [4, 8, 8] + assert "2 template(s) x 4 selected logical seed group(s) x 1 selected" in estimate.note + assert "Baseline adds one unit per selected seed group (4 units)" in estimate.note + assert "num_jailbreaks selects templates" in estimate.components[1].note + assert "20" in estimate.note + + +@pytest.mark.usefixtures("patch_central_database") +async def test_jailbreak_configured_estimate_counts_prompt_sending_without_baseline() -> None: + """Two templates over four groups produce eight units when baseline is disabled.""" + with patch("pyrit.scenario.scenarios.airt.jailbreak._build_jailbreak_technique", return_value=_JailbreakDefault): + scenario = Jailbreak(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) + scenario.set_params_from_args( + args={ + "scenario_techniques": [_JailbreakDefault.PROMPT_SENDING], + "include_baseline": False, + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 8 + assert [component.count for component in estimate.components] == [8] + assert "2 template(s) x 4 selected logical seed group(s) x 1 selected" in estimate.note + assert "Baseline is disabled" in estimate.note + + +@pytest.mark.usefixtures("patch_central_database") +async def test_jailbreak_configured_estimate_counts_prompt_sending_with_baseline() -> None: + """Two templates over four groups plus baseline produce twelve planned units.""" + with patch("pyrit.scenario.scenarios.airt.jailbreak._build_jailbreak_technique", return_value=_JailbreakDefault): + scenario = Jailbreak(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) + scenario.set_params_from_args( + args={ + "scenario_techniques": [_JailbreakDefault.PROMPT_SENDING], + "include_baseline": True, + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 12 + assert [component.count for component in estimate.components] == [4, 8] + assert estimate.components[0].is_baseline is True + assert "Baseline adds one unit per selected seed group (4 units)" in estimate.note + + +@pytest.mark.usefixtures("patch_central_database") +async def test_jailbreak_configured_estimate_uses_target_capability() -> None: + """A capable selected target makes native system-prompt delivery exact.""" + with patch("pyrit.scenario.scenarios.airt.jailbreak._build_jailbreak_technique", return_value=_JailbreakDefault): + scenario = Jailbreak(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) + objective_target = MagicMock(spec=PromptTarget) + objective_target.get_identifier.return_value = ComponentIdentifier(class_name="CapableTarget", class_module="test") + objective_target.configuration.includes.return_value = True + scenario.set_params_from_args( + args={ + "objective_target": objective_target, + "scenario_techniques": [_JailbreakDefault.SYSTEM_PROMPT], + "include_baseline": False, + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + ) + + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 8 + assert [component.count for component in estimate.components] == [0, 8] + objective_target.send_prompt_async.assert_not_called() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_jailbreak_configured_estimate_rejects_incapable_system_delivery() -> None: + """System-only delivery is invalid when the selected target lacks native capabilities.""" + with patch("pyrit.scenario.scenarios.airt.jailbreak._build_jailbreak_technique", return_value=_JailbreakDefault): + scenario = Jailbreak(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 4})) + objective_target = MagicMock(spec=PromptTarget) + objective_target.get_identifier.return_value = ComponentIdentifier( + class_name="IncapableTarget", class_module="test" + ) + objective_target.configuration.includes.return_value = False + scenario.set_params_from_args( + args={ + "objective_target": objective_target, + "scenario_techniques": [_JailbreakDefault.SYSTEM_PROMPT], + "include_baseline": False, + "num_jailbreaks": 2, + "num_jailbreak_attempts": 1, + } + ) + + with pytest.raises(ValueError, match="requires an objective target with editable history"): + await scenario.get_run_size_estimate_async() + + objective_target.send_prompt_async.assert_not_called() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_encoding_estimate_counts_concrete_converter_and_decode_variants() -> None: + """Encoding expands thirteen catalog techniques into fifteen concrete converter variants.""" + scenario = Encoding(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"encoding": 2})) + + estimate = await scenario.get_default_run_size_estimate_async() + assert estimate.estimated_attack_count == 152 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_web_injection_estimate_uses_synthesized_technique_populations() -> None: + """Web injection reports raw sources and capped synthesized populations separately.""" + scenario = WebInjection() + dataset_values = { + scenario.DATASET_EXAMPLE_DOMAINS: ["example.com", "contoso.com"], + scenario.DATASET_MARKDOWN_JS: ["javascript:alert(1)"], + scenario.DATASET_WEB_HTML_JS: [""], + scenario.DATASET_NORMAL_INSTRUCTIONS: ["Write a poem.", "Explain gravity."], + } + with patch.object(scenario, "_load_dataset_values", return_value=dataset_values): + estimate = await scenario.get_default_run_size_estimate_async() + + synthesized = [dataset for dataset in estimate.datasets if dataset.kind == "synthesized"] + synthesized_count = sum(dataset.selected_seed_group_count for dataset in synthesized) + assert len(synthesized) == len(scenario._scenario_techniques) + assert estimate.estimated_attack_count == synthesized_count * 2 + assert estimate.components[-1].label == "Baseline" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_psychosocial_estimate_keeps_sub_harm_baselines_separate() -> None: + """Psychosocial plans each sub-harm's technique cells and baseline independently.""" + scenario = Psychosocial( + imminent_crisis_scorer=_scorer(), + licensed_therapist_scorer=_scorer(), + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock( + return_value=_resolved_groups({"airt_imminent_crisis": 2, "airt_licensed_therapist": 1}) + ) + + estimate = await scenario.get_default_run_size_estimate_async() + assert estimate.estimated_attack_count == 12 + assert [component.count for component in estimate.components] == [6, 2, 3, 1] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_adversarial_benchmark_estimate_exposes_per_required_target_formula() -> None: + """Adversarial benchmark cannot claim a total before its required target count is known.""" + with patch( + "pyrit.scenario.scenarios.benchmark.adversarial._build_benchmark_technique", + return_value=_TwoTechniqueDefault, + ): + scenario = AdversarialBenchmark(objective_scorer=_scorer()) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 3})) + + estimate = await scenario.get_default_run_size_estimate_async() + assert estimate.estimated_attack_count is None + assert estimate.components == [] + assert "adversarial_targets" in estimate.note + + +@pytest.mark.parametrize( + ("use_cached", "expected_total"), + [ + (False, 6), + (True, None), + ], +) +@pytest.mark.usefixtures("patch_central_database") +async def test_adversarial_benchmark_resolves_targets_and_filters_each_technique( + *, + use_cached: bool, + expected_total: int | None, +) -> None: + """Benchmark sizing resolves target names and reports uncached compatible candidates.""" + with patch( + "pyrit.scenario.scenarios.benchmark.adversarial._build_benchmark_technique", + return_value=_TwoTechniqueDefault, + ): + scenario = AdversarialBenchmark(objective_scorer=_scorer(), use_cached=use_cached) + scenario.set_params_from_args(args={"adversarial_targets": ["target-a", "target-b"]}) + compatible = _seed_group("compatible") + incompatible = AttackSeedGroup( + seeds=[ + SeedObjective(value="incompatible"), + SeedPrompt(value="user", data_type="text", role="user", sequence=0), + SeedPrompt(value="assistant", data_type="text", role="assistant", sequence=1), + SeedPrompt(value="user again", data_type="text", role="user", sequence=2), + ] + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock( + return_value=( + {"harmbench": [compatible, incompatible]}, + [ + ScenarioDatasetSummary( + name="harmbench", + logical_seed_group_count=2, + selected_seed_group_count=2, + ) + ], + ) + ) + resolve_targets = MagicMock(return_value=[MagicMock(spec=PromptTarget), MagicMock(spec=PromptTarget)]) + scenario._resolve_adversarial_targets = resolve_targets + plain_factory = MagicMock() + plain_factory.seed_technique = None + conversation_factory = MagicMock() + conversation_factory.seed_technique = AttackTechniqueSeedGroup( + seeds=[ + SeedSimulatedConversation( + adversarial_chat_system_prompt_path="fake.yaml", + num_turns=3, + ) + ] + ) + + with patch( + "pyrit.scenario.scenarios.benchmark.adversarial.resolve_technique_factories_for_techniques", + return_value={"one": plain_factory, "two": conversation_factory}, + ): + estimate = await scenario.get_run_size_estimate_async() + + resolve_targets.assert_called_once_with(target_names=["target-a", "target-b"]) + assert estimate.estimated_attack_count == expected_total + assert [(component.label, component.count) for component in estimate.components] == [("one", 4), ("two", 2)] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_foundry_estimate_counts_composites_instead_of_flattened_techniques() -> None: + """Each Foundry composite contributes one selected seed population.""" + scenario = RedTeamAgent( + adversarial_chat=MagicMock(spec=PromptTarget), + attack_scoring_config=AttackScoringConfig(objective_scorer=_scorer()), + ) + scenario.set_params_from_args( + args={ + "scenario_techniques": [ + FoundryComposite( + attack=FoundryTechnique.Crescendo, + converters=[FoundryTechnique.Base64, FoundryTechnique.ROT13], + ), + FoundryComposite(attack=None, converters=[FoundryTechnique.Tense]), + ], + "include_baseline": False, + } + ) + scenario._resolve_dataset_groups_for_estimate_async = AsyncMock(return_value=_resolved_groups({"harmbench": 3})) + + estimate = await scenario.get_run_size_estimate_async() + + assert estimate.estimated_attack_count == 6 + assert len(estimate.components) == 2 + assert [component.count for component in estimate.components] == [3, 3]