From 2f6c715571c8087a62f26bca630ecc911408f880 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 02:08:09 -0700 Subject: [PATCH] FEAT: Add durable scenario progress Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5d02c2d5-b499-4f78-a04d-03bffa750817 --- doc/code/framework.md | 1 + pyrit/backend/routes/scenarios.py | 90 ++- .../backend/services/scenario_run_service.py | 645 +++++++++++++++--- pyrit/backend/services/scenario_service.py | 277 +++++++- .../attack/component/conversation_manager.py | 30 +- .../prepended_conversation_config.py | 10 +- pyrit/executor/attack/core/attack_executor.py | 33 +- .../attack/core/attack_result_attribution.py | 3 + pyrit/executor/attack/core/attack_strategy.py | 2 + .../attack/single_turn/prompt_sending.py | 2 + ...0f2a4c1e_index_scenario_progress_deltas.py | 35 + pyrit/memory/memory_interface.py | 111 +++ pyrit/memory/memory_models.py | 16 +- pyrit/models/__init__.py | 40 ++ pyrit/models/catalog/__init__.py | 16 + pyrit/models/catalog/scenario.py | 267 +++++++- pyrit/models/identifiers/__init__.py | 3 +- .../identifiers/atomic_attack_identifier.py | 7 +- pyrit/models/identifiers/seed_identifier.py | 16 +- pyrit/models/results/scenario_result.py | 6 +- pyrit/models/scenario_progress.py | 133 ++++ pyrit/models/seeds/attack_seed_group.py | 12 + .../registry/components/scenario_registry.py | 64 +- pyrit/registry/registry_metadata.py | 14 + pyrit/scenario/core/atomic_attack.py | 36 +- .../scenario/core/attack_technique_factory.py | 34 + pyrit/scenario/core/dataset_configuration.py | 96 ++- .../core/matrix_atomic_attack_builder.py | 47 +- pyrit/scenario/core/scenario.py | 476 +++++++++++-- .../scenarios/adaptive/adaptive_scenario.py | 96 +++ pyrit/scenario/scenarios/airt/cyber.py | 3 +- pyrit/scenario/scenarios/airt/jailbreak.py | 208 +++++- pyrit/scenario/scenarios/airt/leakage.py | 7 +- pyrit/scenario/scenarios/airt/psychosocial.py | 48 +- .../scenario/scenarios/airt/rapid_response.py | 3 +- .../scenarios/benchmark/adversarial.py | 84 ++- .../scenarios/foundry/red_team_agent.py | 45 +- pyrit/scenario/scenarios/garak/doctor.py | 5 + pyrit/scenario/scenarios/garak/encoding.py | 101 ++- .../scenario/scenarios/garak/web_injection.py | 127 +++- pyrit/setup/initializers/techniques/airt.py | 1 + pyrit/setup/initializers/techniques/core.py | 1 + .../unit/backend/test_scenario_run_routes.py | 139 +++- .../unit/backend/test_scenario_run_service.py | 378 +++++++++- tests/unit/backend/test_scenario_service.py | 557 ++++++++++++++- .../component/test_conversation_manager.py | 27 +- .../test_prepended_conversation_config.py | 10 +- .../attack/core/test_attack_strategy.py | 4 + .../attack/single_turn/test_prompt_sending.py | 58 +- .../test_interface_scenario_progress.py | 129 ++++ .../test_interface_scenario_results.py | 25 + tests/unit/memory/test_migration.py | 15 + tests/unit/models/test_attack_seed_group.py | 35 + tests/unit/models/test_scenario_catalog.py | 189 +++++ tests/unit/models/test_scenario_progress.py | 41 ++ tests/unit/registry/test_registry_metadata.py | 32 + tests/unit/registry/test_scenario_registry.py | 110 +++ tests/unit/scenario/airt/test_cyber.py | 11 +- tests/unit/scenario/airt/test_jailbreak.py | 204 ++++-- .../unit/scenario/airt/test_rapid_response.py | 11 +- .../unit/scenario/core/test_atomic_attack.py | 24 +- .../core/test_attack_technique_factory.py | 23 + .../core/test_dataset_configuration.py | 27 + tests/unit/scenario/core/test_scenario.py | 167 ++++- .../core/test_scenario_partial_results.py | 27 +- .../unit/scenario/core/test_scenario_retry.py | 27 +- .../test_default_run_size_estimates.py | 641 +++++++++++++++++ 67 files changed, 5720 insertions(+), 442 deletions(-) create mode 100644 pyrit/memory/alembic/versions/6b8d0f2a4c1e_index_scenario_progress_deltas.py create mode 100644 pyrit/models/scenario_progress.py create mode 100644 tests/unit/memory/memory_interface/test_interface_scenario_progress.py create mode 100644 tests/unit/models/test_scenario_catalog.py create mode 100644 tests/unit/models/test_scenario_progress.py create mode 100644 tests/unit/scenario/test_default_run_size_estimates.py diff --git a/doc/code/framework.md b/doc/code/framework.md index a3938ef2dc..3f98f0b69a 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_request_converter_composition=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/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index fa3a5635bb..3245b644a8 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, + ScenarioDefaultRunSizeEstimate, + 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=ScenarioDefaultRunSizeEstimate, + 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, +) -> ScenarioDefaultRunSizeEstimate: + """ + 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: + ScenarioDefaultRunSizeEstimate: 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 # ============================================================================ @@ -154,7 +197,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 +211,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/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 7ba66d0f43..acc289ea80 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -9,14 +9,36 @@ """ import asyncio +import base64 import contextlib +import json import logging +import uuid +from collections.abc import Sequence from dataclasses import dataclass +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any from pyrit.backend.models.scenarios import ScenarioRunListResponse +from pyrit.common.utils import to_sha256 from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, ScenarioResult, ScenarioRunState +from pyrit.memory.memory_interface import ScenarioProgressKeysetCursor +from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + AtomicAttackIdentifier, + AttackOutcome, + ComponentIdentifier, + ScenarioAttackResultDelta, + ScenarioProgressHeader, + ScenarioProgressResult, + ScenarioResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunProgress, + ScenarioRunState, + config_hash, +) from pyrit.models.catalog.scenario import ( AttackErrorSummary, AttackRetrySummary, @@ -30,7 +52,6 @@ TargetRegistry, ) from pyrit.scenario import Scenario -from pyrit.scenario.core import DatasetAttackConfiguration if TYPE_CHECKING: from pyrit.converter import Converter @@ -53,6 +74,14 @@ 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, ...] = () + + class ScenarioRunService: """ Service for managing scenario run lifecycle. @@ -121,7 +150,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 +165,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: """ @@ -151,7 +199,13 @@ def list_runs(self, *, limit: int = 100) -> ScenarioRunListResponse: # 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] + items = [ + self._build_response_from_db( + scenario_result=sr, + active_error=self.snapshot_active_run(scenario_result_id=str(sr.id)).error, + ) + for sr in results + ] return ScenarioRunListResponse(items=items) async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: @@ -193,7 +247,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) + return self.get_run(scenario_result_id=scenario_result_id) def _resolve_scenario_class(self, *, request: RunScenarioRequest) -> type[Scenario]: """ @@ -251,18 +305,34 @@ def _resolve_target(self, *, request: RunScenarioRequest) -> "PromptTarget": Raises: ValueError: If the target is not found in the registry. """ + return self.resolve_target_name(target_name=request.target_name) + + @staticmethod + def resolve_target_name(*, target_name: str) -> "PromptTarget": + """ + Resolve one registered target name for launch or configured estimation. + + Args: + target_name: Registered target instance name. + + Returns: + PromptTarget: The resolved target. + + Raises: + ValueError: If the target is not registered. + """ target_registry = TargetRegistry.get_registry_singleton() - objective_target = target_registry.instances.get(request.target_name) + objective_target = target_registry.instances.get(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. " + 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 '{request.target_name}' not found in registry. Available targets: {', '.join(available_names)}" + f"Target '{target_name}' not found in registry. Available targets: {', '.join(available_names)}" ) return objective_target @@ -295,102 +365,128 @@ def _build_init_kwargs( 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) + return self.resolve_scenario_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, ) + + @classmethod + def resolve_scenario_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/estimate request fields into scenario parameters. + + Args: + scenario_name: Registered scenario name used in validation errors. + scenario_class: Scenario class used for technique and dataset introspection. + objective_target: Optional resolved objective target. + techniques: Requested technique tokens. + dataset_names: Requested dataset names. + max_dataset_size: Requested logical-group selection cap. + dataset_filters: Validated dataset seed filters. + include_baseline: Optional baseline policy override. + max_concurrency: Optional launch concurrency. + max_retries: Optional launch retry count. + memory_labels: Optional launch memory labels. + + 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 init_kwargs + 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 '{request.scenario_name}': " + f"Cannot resolve runtime configuration for scenario '{scenario_name}': " f"scenario class is not instantiable without arguments ({exc})." ) from exc - if request.techniques: + if techniques: technique_class = introspection_instance._technique_class - technique_enums, technique_converters = self._resolve_techniques_and_converters( - tokens=request.techniques, + technique_enums, technique_converters = cls._resolve_techniques_and_converters( + tokens=techniques, technique_class=technique_class, - scenario_name=request.scenario_name, + scenario_name=scenario_name, ) - init_kwargs["scenario_techniques"] = technique_enums + resolved["scenario_techniques"] = technique_enums if technique_converters: - init_kwargs["technique_converters"] = technique_converters + resolved["technique_converters"] = technique_converters - if request.dataset_names or request.max_dataset_size is not None or dataset_filters: + if dataset_names or max_dataset_size is not None or filters: default_config = introspection_instance._default_dataset_config - if request.dataset_names: + if 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, + resolved["dataset_config"] = default_config_class( + dataset_names=dataset_names, + max_dataset_size=max_dataset_size, + filters=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, - ) + raise ValueError( + f"Scenario '{scenario_name}' does not support overriding dataset names through " + f"its {default_config_class.__name__} configuration: {exc}" + ) from exc 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 + 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 init_kwargs + return resolved + @classmethod def _resolve_techniques_and_converters( - self, + cls, *, tokens: list[str], technique_class: type[Any], @@ -435,7 +531,7 @@ def _resolve_techniques_and_converters( ) from None technique_enums.append(technique_enum) - converters = self._resolve_converter_modifiers(modifiers=modifiers, token=token) + converters = cls._resolve_converter_modifiers(modifiers=modifiers, token=token) if not converters: continue @@ -444,7 +540,8 @@ def _resolve_techniques_and_converters( return technique_enums, technique_converters - def _resolve_converter_modifiers(self, *, modifiers: list[str], token: str) -> list["Converter"]: + @staticmethod + def _resolve_converter_modifiers(*, modifiers: list[str], token: str) -> list["Converter"]: """ Resolve the converter modifiers of a single technique token to converter instances. @@ -541,12 +638,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 +657,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 +693,42 @@ 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) # 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, + ) + 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[tuple[str, str], int] = {} for atomic_attack_name, results in scenario_result.attack_results.items(): for attack_result in results: + unit_key = self._result_unit_key( + atomic_attack_name=atomic_attack_name, + attack_result=attack_result, + plan=plan, + ) + attempts_by_unit[unit_key] = attempts_by_unit.get(unit_key, 0) + 1 retries = getattr(attack_result, "total_retries", 0) if isinstance(retries, int): total_retries += retries @@ -630,27 +753,365 @@ 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 _result_unit_key( + *, + atomic_attack_name: str, + attack_result: Any, + plan: ScenarioRunPlan | None, + ) -> tuple[str, str]: + """ + Resolve one attack attempt to its stable planned-unit key. + + Returns: + tuple[str, str]: The atomic-group and seed-group IDs. + """ + atomic_identifier = getattr(attack_result, "atomic_attack_identifier", None) + typed_identifier = ( + AtomicAttackIdentifier.from_component_identifier(atomic_identifier) + if isinstance(atomic_identifier, ComponentIdentifier) + else None + ) + objective = str(getattr(attack_result, "objective", "")) + attribution_data = getattr(attack_result, "attribution_data", None) + 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 + planned_group: ScenarioRunPlanAtomicGroup | None = None + if plan is not None: + eval_hash = attribution_data.get("parent_eval_hash") if isinstance(attribution_data, dict) else None + for group in plan.atomic_groups: + if group.atomic_attack_name == atomic_attack_name and ( + eval_hash is None or group.technique_eval_hash == eval_hash + ): + atomic_group_id = group.id + planned_group = group + break + if not seed_group_id and plan is not None and planned_group is not None: + objective_sha256 = str(getattr(attack_result, "objective_sha256", "") or to_sha256(objective)) + matching_seed_ids = [ + seed.id + for seed in plan.seed_groups + if seed.id in planned_group.seed_group_ids and seed.objective_sha256 == 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 atomic_group_id, seed_group_id + + def _calculate_progress_counts( + self, + *, + scenario_result: ScenarioResult, + plan: ScenarioRunPlan | None, + ) -> 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. + """ + attempted_units: set[tuple[str, str]] = set() + latest_non_error_by_unit: dict[tuple[str, str], Any] = {} + for atomic_attack_name, results in scenario_result.attack_results.items(): + for attack_result in results: + unit_key = self._result_unit_key( + atomic_attack_name=atomic_attack_name, + attack_result=attack_result, + plan=plan, + ) + attempted_units.add(unit_key) + if attack_result.outcome == AttackOutcome.ERROR: + continue + previous = latest_non_error_by_unit.get(unit_key) + if previous is None or self._result_order_key(attack_result) > self._result_order_key(previous): + latest_non_error_by_unit[unit_key] = attack_result + + planned_units = ( + {(group.id, seed_group_id) for group in plan.atomic_groups for seed_group_id in group.seed_group_ids} + if plan is not None + else attempted_units + ) + total = len(planned_units) + completed_results = [ + result for unit_key, result in latest_non_error_by_unit.items() if unit_key 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: Any) -> tuple[datetime, str]: + """Return a deterministic chronological key for one hydrated result attempt.""" + timestamp = getattr(attack_result, "timestamp", None) + if not isinstance(timestamp, datetime): + timestamp = datetime.min.replace(tzinfo=timezone.utc) + return timestamp, str(getattr(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_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) + + results = [self._map_progress_delta(delta=delta, plan=plan or response_plan) 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: ScenarioRunPlan | None, + ) -> 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 ""} + ) + if plan is not None: + for group in plan.atomic_groups: + if group.atomic_attack_name == atomic_attack_name and ( + eval_hash is None or group.technique_eval_hash == eval_hash + ): + atomic_group_id = group.id + break + 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 plan is not None and delta.objective_sha256: + matching_seed_ids = [ + seed.id + for seed in plan.seed_groups + if seed.objective_sha256 == delta.objective_sha256 + and any(seed.id in group.seed_group_ids for group in plan.atomic_groups if group.id == atomic_group_id) + ] + 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] = {} + for delta in deltas: + mapped = ScenarioRunService._map_progress_delta(delta=delta, plan=None) + 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=[], + ), + ) + if mapped.seed_group_id not in group.seed_group_ids: + 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, + ) -> ScenarioProgressKeysetCursor | 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 ScenarioProgressKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) + def get_run_results(self, *, scenario_result_id: str) -> ScenarioResult | None: """ Get the ScenarioResult for a completed scenario run. diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index 46721d8ed1..e6fcbfccd6 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -1,55 +1,85 @@ # 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_run_service import ScenarioRunService from pyrit.models.catalog.scenario import ( RegisteredScenario, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, ) 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[ScenarioDefaultRunSizeEstimate, float | None] +_EstimateTask = asyncio.Task[ScenarioDefaultRunSizeEstimate] -def _metadata_to_registered_scenario(metadata: ScenarioMetadata) -> RegisteredScenario: + +def _metadata_to_registered_scenario( + *, + metadata: ScenarioMetadata, + default_run_size: ScenarioDefaultRunSizeEstimate | 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 ScenarioDefaultRunSizeEstimate.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 +88,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 +123,190 @@ 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, + ) -> ScenarioDefaultRunSizeEstimate | None: + """ + Estimate one configured scenario without creating a run. + + Args: + scenario_name: Registered scenario name. + request: Request-specific techniques, datasets, baseline, and parameters. + + Returns: + ScenarioDefaultRunSizeEstimate | 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 + ) -> ScenarioDefaultRunSizeEstimate: + """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: + return await asyncio.shield(task) + if wait_for_capacity is not None: + await asyncio.shield(wait_for_capacity) + + def _read_estimate_cache(self, *, cache_key: _EstimateCacheKey) -> ScenarioDefaultRunSizeEstimate | 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, + ) -> ScenarioDefaultRunSizeEstimate: + """ + Construct and estimate one scenario on the owning event loop. + + Returns: + ScenarioDefaultRunSizeEstimate: 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 = ScenarioDefaultRunSizeEstimate.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.status is ScenarioRunSizeEstimateStatus.Unavailable + 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, + ) -> ScenarioDefaultRunSizeEstimate: + """ + Resolve and estimate one request on the owning event loop. + + Returns: + ScenarioDefaultRunSizeEstimate: Request-specific scenario estimate. + """ + scenario_class = self._registry.get_class(scenario_name) + objective_target = ( + ScenarioRunService.resolve_target_name(target_name=request.target_name) if request.target_name else None + ) + estimate_kwargs = ScenarioRunService.resolve_scenario_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 +315,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 +326,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 +334,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/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index d92f99bbde..d859e7a12d 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -321,10 +321,22 @@ async def initialize_context_async( # single-string fallback path via capability-based routing. is_chat_target = target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) if not is_chat_target: + config = prepended_conversation_config or PrependedConversationConfig() + if request_converters: + present_roles = { + piece.api_role for message in prepended_conversation for piece in message.message_pieces + } + excluded_roles = present_roles - set(config.apply_converters_to_roles) + if excluded_roles: + raise ValueError( + "Cannot preserve prepended-conversation converter role scoping for a non-chat target: " + f"the flattened context contains excluded roles {sorted(excluded_roles)}. " + "Use a chat target, remove request converters, or explicitly opt into every prepended role." + ) return await self._handle_non_chat_target_async( context=context, prepended_conversation=prepended_conversation, - config=prepended_conversation_config, + config=config, ) # Process prepended conversation for objective target @@ -453,10 +465,10 @@ async def add_prepended_conversation_to_memory_async( conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) ) - # Get roles that should have converters applied - apply_to_roles = ( - prepended_conversation_config.apply_converters_to_roles if prepended_conversation_config else None - ) + # Assistant history represents simulated target output, so the absent-config + # path must use the same safe role default as an explicit default config. + config = prepended_conversation_config or PrependedConversationConfig() + apply_to_roles = config.apply_converters_to_roles turn_count = 0 @@ -575,7 +587,7 @@ async def _apply_converters_async( *, message: Message, request_converters: list[ConverterConfiguration], - apply_to_roles: list[ChatMessageRole] | None, + apply_to_roles: list[ChatMessageRole], ) -> None: """ Apply converters to message pieces. @@ -583,12 +595,10 @@ async def _apply_converters_async( Args: message: The message containing pieces to convert. request_converters: Converter configurations to apply. - apply_to_roles: If provided, only apply to pieces with these roles. - If None, apply to all roles. + apply_to_roles: Only apply to pieces with these roles. """ for piece in message.message_pieces: - # Filter by role if specified - if apply_to_roles is not None and piece.api_role not in apply_to_roles: + if piece.api_role not in apply_to_roles: continue temp_message = Message(message_pieces=[piece]) diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index a0daedfd6e..80bf7681e7 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -4,13 +4,12 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import get_args from pyrit.message_normalizer import ( ConversationContextNormalizer, MessageStringNormalizer, ) -from pyrit.models import ChatMessageRole +from pyrit.models import ChatMessageRole # noqa: TC001 - public annotation must resolve at runtime @dataclass @@ -27,10 +26,9 @@ class PrependedConversationConfig: first turn (via ``message_normalizer``; default: ConversationContextNormalizer). """ - # Roles for which request converters should be applied to prepended messages. - # By default, converters are applied to all roles. - # Example: ["user"] to apply converters only to user messages. - apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: list(get_args(ChatMessageRole))) + # Request converters default to prepended user messages only. Assistant history is + # simulated target output and must be explicitly opted in with ["assistant"]. + apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) # Optional normalizer to format conversation history into a single text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index 64d75d1066..846aa4a4db 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,12 @@ 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") params_type = attack.params_type @@ -263,11 +273,15 @@ 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 = ( + [attributions[index] for index in successful_input_indices] if 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( @@ -351,6 +365,7 @@ async def _execute_with_params_list_async( 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]: """ @@ -366,19 +381,31 @@ async def _execute_with_params_list_async( 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, + or shared and per-task attribution are both provided. """ 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)})" + ) + if attribution is not None and attributions is not None: + raise ValueError("Provide attribution or attributions, not both") 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 attribution + 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 11c58b60ad..7e199aa005 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -286,6 +286,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/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 508b72d924..ffebec0f2a 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -77,6 +77,8 @@ def __init__( prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter application by role, message normalization, and non-chat target behavior. + Request converters apply to prepended user messages by default; include + ``"assistant"`` explicitly to transform simulated assistant history. Raises: ValueError: If the objective scorer is not a true/false scorer. 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 4bbef89456..e4740ab0a9 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, @@ -121,6 +125,13 @@ def from_attack_result(cls, result: AttackResult) -> "AttackResultsKeysetCursor" ) +class ScenarioProgressKeysetCursor(NamedTuple): + """Ascending keyset anchor for scenario-linked attack-result deltas.""" + + timestamp: datetime + attack_result_id: str + + @dataclass(frozen=True, slots=True, kw_only=True) class _AttackResultQuery: """ @@ -3318,6 +3329,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() @@ -3351,6 +3368,100 @@ 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_attack_result_deltas( + self, + *, + scenario_result_id: str, + cursor: ScenarioProgressKeysetCursor | 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 653abf8b48..4df00fa5fa 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1554,6 +1554,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) @@ -1863,12 +1870,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 289187f44c..67b420de17 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -17,6 +17,16 @@ """ from pyrit.models.additional_initializer import AdditionalInitializer +from pyrit.models.catalog import ( + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) 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 @@ -47,6 +57,7 @@ class_name_to_snake_case, compute_eval_hash, config_hash, + logical_seed_group_fingerprint, snake_case_to_class_name, validate_registry_name, ) @@ -89,6 +100,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 Score, ScoreType, UnvalidatedScore # Seeds - import from new seeds submodule for forward compatibility @@ -170,6 +192,7 @@ "IdentifierFilter", "IdentifierType", "JSONValue", + "logical_seed_group_fingerprint", "COMMON_JSON_SCHEMAS", "JsonResponseConfig", "get_common_json_schema", @@ -200,8 +223,25 @@ "ScorerEvaluationIdentifier", "ScorerIdentifier", "ScenarioIdentifier", + "ScenarioDatasetSizeCap", + "ScenarioDatasetSummary", + "ScenarioDefaultRunSizeEstimate", + "ScenarioRunSizeEstimate", "ScenarioResult", + "ScenarioRunSizeComponent", + "ScenarioRunSizeEstimateRequest", + "ScenarioRunSizeEstimateStatus", + "ScenarioRunSizeFactor", "ScenarioRunState", + "SCENARIO_RUN_PLAN_METADATA_KEY", + "SCENARIO_RUN_PLAN_VERSION", + "ScenarioAttackResultDelta", + "ScenarioProgressHeader", + "ScenarioProgressResult", + "ScenarioRunPlan", + "ScenarioRunPlanAtomicGroup", + "ScenarioRunPlanSeedGroup", + "ScenarioRunProgress", "Seed", "AttackSeedGroup", "AttackTechniqueSeedGroup", diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 6d8e2e15d3..692e3a54d5 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -21,6 +21,14 @@ AttackRetrySummary, RegisteredScenario, RunScenarioRequest, + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ScenarioRunSummary, ) from pyrit.models.catalog.target import ( @@ -33,6 +41,14 @@ "RegisteredInitializer", "RegisteredScenario", "RunScenarioRequest", + "ScenarioDatasetSizeCap", + "ScenarioDatasetSummary", + "ScenarioDefaultRunSizeEstimate", + "ScenarioRunSizeEstimate", + "ScenarioRunSizeComponent", + "ScenarioRunSizeEstimateRequest", + "ScenarioRunSizeEstimateStatus", + "ScenarioRunSizeFactor", "ScenarioRunSummary", "TargetInstance", ] diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 488ccf8c78..549dd3eb37 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -14,9 +14,11 @@ """ from datetime import datetime -from typing import Any +from enum import Enum +from math import prod +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 +41,265 @@ 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 ScenarioRunSizeEstimateStatus(str, Enum): + """Confidence level for a catalog default-run size estimate.""" + + Exact = "exact" + Conditional = "conditional" + Unavailable = "unavailable" + + +class ScenarioRunSizeFactor(BaseModel): + """One labeled multiplicative factor in a run-size component.""" + + label: str = Field(..., min_length=1) + count: int = Field(..., ge=0) + + +class ScenarioRunSizeComponent(BaseModel): + """One additive component of a default-run size estimate.""" + + label: str = Field(..., min_length=1) + count: int = Field(..., ge=0) + factors: list[ScenarioRunSizeFactor] = Field(default_factory=list) + is_baseline: bool = False + note: str | None = None + + @model_validator(mode="after") + def validate_factor_product(self) -> "ScenarioRunSizeComponent": + """ + Require known component totals to equal their ordered factor product. + + Returns: + ScenarioRunSizeComponent: The validated component. + + Raises: + ValueError: If a component with factors has an inconsistent count. + """ + if self.factors: + factor_product = prod(factor.count for factor in self.factors) + if self.count != factor_product: + raise ValueError( + f"Component '{self.label}' count ({self.count}) must equal its factor product ({factor_product})" + ) + return self + + +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 ScenarioDefaultRunSizeEstimate(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. + """ + + version: Literal[1] = 1 + status: ScenarioRunSizeEstimateStatus + total_attack_count: int | None = Field( + default=None, + ge=0, + validation_alias=AliasChoices("total_attack_count", "total"), + ) + components: list[ScenarioRunSizeComponent] = Field(default_factory=list) + datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) + note: str | None = Field(default=None, validation_alias=AliasChoices("note", "caveat")) + retries_included: Literal[False] = False + + @property + def total(self) -> int | None: + """The legacy Python attribute for total_attack_count.""" + return self.total_attack_count + + @property + def caveat(self) -> str | None: + """The legacy Python attribute for note.""" + return self.note + + @model_validator(mode="before") + @classmethod + def normalize_legacy_total(cls, data: Any) -> Any: + """ + Explain component-less legacy exact totals in the canonical shape. + + Returns: + Any: The normalized input when it is a legacy exact estimate; otherwise the original input. + """ + if not isinstance(data, dict) or "total" not in data or "components" in data: + return data + if data.get("status") != ScenarioRunSizeEstimateStatus.Exact and data.get("status") != "exact": + return data + normalized = dict(data) + normalized["components"] = [ + { + "label": "Legacy total", + "count": data["total"], + "note": "Normalized from a legacy component-less estimate.", + } + ] + return normalized + + @model_validator(mode="after") + def validate_total(self) -> "ScenarioDefaultRunSizeEstimate": + """ + Ensure exact estimates expose and explain their complete total. + + Returns: + ScenarioDefaultRunSizeEstimate: The validated estimate. + + Raises: + ValueError: If an exact estimate omits or misstates its total. + """ + if self.status is ScenarioRunSizeEstimateStatus.Exact: + if self.total_attack_count is None: + raise ValueError("Exact default-run estimates require total_attack_count") + component_total = sum(component.count for component in self.components) + if component_total != self.total_attack_count: + raise ValueError( + f"Exact default-run estimate components total {component_total}, not {self.total_attack_count}" + ) + return self + + @classmethod + def unavailable( + cls, *, note: str = "Default-run size estimate is unavailable." + ) -> "ScenarioDefaultRunSizeEstimate": + """ + Build an unavailable estimate without presenting a guessed total. + + Returns: + ScenarioDefaultRunSizeEstimate: An unavailable estimate. + """ + return cls(status=ScenarioRunSizeEstimateStatus.Unavailable, note=note) + + +# Backward-compatible catalog name from the initial run-size DTO. +ScenarioRunSizeEstimate = ScenarioDefaultRunSizeEstimate + + 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: ScenarioDefaultRunSizeEstimate = Field( + default_factory=ScenarioDefaultRunSizeEstimate.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 +321,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 +348,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 +379,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") diff --git a/pyrit/models/identifiers/__init__.py b/pyrit/models/identifiers/__init__.py index aea1532581..2f671a9f13 100644 --- a/pyrit/models/identifiers/__init__.py +++ b/pyrit/models/identifiers/__init__.py @@ -39,7 +39,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, logical_seed_group_fingerprint from pyrit.models.identifiers.target_identifier import TargetIdentifier __all__ = [ @@ -70,6 +70,7 @@ "ScorerIdentifier", "ScenarioIdentifier", "SeedIdentifier", + "logical_seed_group_fingerprint", "snake_case_to_class_name", "TARGET_EVAL_PARAM_FALLBACKS", "TARGET_EVAL_PARAMS", diff --git a/pyrit/models/identifiers/atomic_attack_identifier.py b/pyrit/models/identifiers/atomic_attack_identifier.py index c4a59cf37e..6ccdf93f0a 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, logical_seed_group_fingerprint 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 logical_seed_group_fingerprint(self.seed_identifiers) diff --git a/pyrit/models/identifiers/seed_identifier.py b/pyrit/models/identifiers/seed_identifier.py index 2372164662..aa406559f2 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 logical_seed_group_fingerprint(seed_identifiers: Sequence[SeedIdentifier]) -> str: + """Return the deterministic fingerprint 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..117b65165a 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, logical_seed_group_fingerprint 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 logical_seed_group_fingerprint([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 1551148085..7e8af09f46 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 ScenarioDefaultRunSizeEstimate, 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]): """ @@ -129,6 +145,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()) @@ -145,8 +162,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( @@ -154,13 +181,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, + ) -> ScenarioDefaultRunSizeEstimate: + """ + 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: + ScenarioDefaultRunSizeEstimate: 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, @@ -208,5 +267,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 427b14384e..499f5ca847 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 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.executor.attack.core.attack_executor import AttackExecutorResult @@ -192,6 +198,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]: """ @@ -325,13 +341,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 + ] results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, @@ -340,7 +360,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, ) diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 1d168a0545..c74114c770 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -82,6 +82,7 @@ def __init__( adversarial_seed_prompt: SeedPrompt | str | None = None, seed_technique: AttackTechniqueSeedGroup | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -119,6 +120,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_request_converter_composition: 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. @@ -143,11 +147,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_request_converter_composition = supports_request_converter_composition 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 @@ -166,6 +172,7 @@ def with_simulated_conversation( attack_kwargs: dict[str, Any] | None = None, adversarial_chat: PromptTarget | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -217,6 +224,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_request_converter_composition: 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 @@ -277,6 +287,7 @@ def with_simulated_conversation( adversarial_chat=adversarial_chat, seed_technique=seed_technique, uses_adversarial=uses_adversarial, + supports_request_converter_composition=supports_request_converter_composition, scorer_override_policy=scorer_override_policy, ) @@ -312,6 +323,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_request_converter_composition + and "attack_converter_config" not in self._get_accepted_params() + ): + raise ValueError( + f"Factory '{self._name}' declares supports_request_converter_composition=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. @@ -434,6 +462,11 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial + @property + def supports_request_converter_composition(self) -> bool: + """Whether callers may safely append request converters to this technique.""" + return self._supports_request_converter_composition + @property def scoring_config_type(self) -> type | None: """The required ``attack_scoring_config`` subtype, or ``None`` if any config is accepted.""" @@ -774,6 +807,7 @@ def _build_identifier(self) -> ComponentIdentifier: "attack_class": self._attack_class.__name__, "kwargs": kwargs_for_id, "uses_adversarial": self._uses_adversarial, + "supports_request_converter_composition": self._supports_request_converter_composition, } 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 5ca3c2247c..750bfc6405 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -33,13 +33,24 @@ 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, + ScenarioDefaultRunSizeEstimate, ScenarioEvaluationIdentifier, ScenarioIdentifier, ScenarioResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, ScenarioRunState, + config_hash, ) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference from pyrit.prompt_target import PromptTarget @@ -47,7 +58,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 +76,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 +135,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 +217,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 +235,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 +269,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]: """ @@ -518,63 +551,226 @@ 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) -> ScenarioDefaultRunSizeEstimate: """ - 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: + ScenarioDefaultRunSizeEstimate: 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 + ) -> ScenarioDefaultRunSizeEstimate: + """ + 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: + ScenarioDefaultRunSizeEstimate: 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) -> ScenarioDefaultRunSizeEstimate: + """ + 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: + ScenarioDefaultRunSizeEstimate: 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, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], + is_baseline=True, + note="One unmodified prompt-sending unit per selected seed group.", + ) + ) + + status = ( + ScenarioRunSizeEstimateStatus.Conditional + if self.RUN_SIZE_USES_FACTORY_COMPATIBILITY and self._estimate_has_binding_size_cap + else ScenarioRunSizeEstimateStatus.Exact + ) + total_attack_count = ( + None + if status is ScenarioRunSizeEstimateStatus.Conditional + else sum(component.count for component in components) + ) + note = "Counts planned outer execution units; retries and internal attack turns are excluded." + if status is ScenarioRunSizeEstimateStatus.Conditional: + note += " A binding randomized dataset cap may select a different compatibility mix at launch." + return ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=total_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, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="default concrete techniques", 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, + factors=[ + ScenarioRunSizeFactor(label="selected concrete techniques", count=1), + ScenarioRunSizeFactor(label="compatible logical seed groups", 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 @@ -583,10 +779,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: @@ -597,16 +789,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 @@ -645,7 +871,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 @@ -680,28 +918,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. @@ -1295,6 +1636,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, @@ -1307,6 +1650,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 d1289df6cf..8cb393b483 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -22,6 +22,12 @@ from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackScoringConfig +from pyrit.models import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) 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 +203,96 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list return atomic_attacks + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate compatible persisted envelopes, excluding adaptive inner attempts. + + Returns: + ScenarioDefaultRunSizeEstimate: 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, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", 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, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count)], + ), + ] + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + 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, + factors=[ScenarioRunSizeFactor(label="compatible logical seed groups", count=compatible_group_count)], + ), + ] + status = ( + ScenarioRunSizeEstimateStatus.Conditional + if self._estimate_has_binding_size_cap + else ScenarioRunSizeEstimateStatus.Exact + ) + total_attack_count = ( + None + if status is ScenarioRunSizeEstimateStatus.Conditional + 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 status is ScenarioRunSizeEstimateStatus.Conditional: + note += " A binding randomized dataset cap may select a different compatibility mix at launch." + return ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=total_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 710c7eba7c..2fbece5e6c 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -12,7 +12,14 @@ 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, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.prompt_target import CapabilityName from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -72,6 +79,7 @@ def _prompt_sending_factory() -> AttackTechniqueFactory: name=_PROMPT_SENDING, attack_class=PromptSendingAttack, technique_tags=["single_turn"], + supports_request_converter_composition=True, ) @@ -93,6 +101,7 @@ def _jailbreak_system_prompt_factory() -> AttackTechniqueFactory: name=_JAILBREAK_SYSTEM_PROMPT, attack_class=PromptSendingAttack, technique_tags=["single_turn"], + supports_request_converter_composition=True, ) @@ -104,23 +113,34 @@ def _extra_default_factories() -> dict[str, AttackTechniqueFactory]: } +def _is_jailbreak_compatible_factory(factory: AttackTechniqueFactory) -> bool: + """Return whether a factory supports direct jailbreak-converter delivery.""" + has_simulated_conversation = ( + factory.seed_technique is not None and factory.seed_technique.has_simulated_conversation + ) + return ( + "multi_turn" not in factory.technique_tags + and not has_simulated_conversation + and factory.supports_request_converter_composition + ) + + @cache def _build_jailbreak_technique() -> type[ScenarioTechnique]: """ - Build the Jailbreak technique class dynamically from every registered factory plus the - scenario-local defaults. + Build the Jailbreak technique class from compatible direct factories and local deliveries. - The technique axis is the set of *attack techniques* a jailbreak is delivered through: the two - default deliveries (``prompt_sending`` and ``jailbreak_system_prompt``) plus whatever techniques - are registered (``role_play_*``, ``many_shot``, ``tap``, …). Jailbreak templates are a separate - selector (``num_jailbreaks`` / ``jailbreak_names``), so only the two deliveries are on by default - — crossing every template with every registered technique explodes quickly. + Registered multi-turn techniques, simulated-conversation seed techniques, and factories that + have not explicitly opted into additive request-converter composition are excluded. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. """ registry = AttackTechniqueRegistry.get_registry_singleton() - factories = list(registry.get_factories_or_raise().values()) + list(_extra_default_factories().values()) + registered = [ + factory for factory in registry.get_factories_or_raise().values() if _is_jailbreak_compatible_factory(factory) + ] + factories = registered + list(_extra_default_factories().values()) return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[return-value, ty:invalid-return-type] class_name="JailbreakTechnique", factories=factories, @@ -136,24 +156,25 @@ class Jailbreak(Scenario): selectors: - **dataset** — the harmful objectives (HarmBench). - - **techniques** — the *attack techniques* each jailbreak is delivered through. Two deliveries + - **techniques** — compatible direct deliveries for each jailbreak. Two deliveries are on by default: ``prompt_sending`` (the template rendered inline into the user message) and ``jailbreak_system_prompt`` (the template set as the system prompt with the objective sent as - the user turn). The registry techniques (``role_play_*``, ``many_shot``, ``tap``, …) are - opt-in. + the user turn). Registered direct techniques are opt-in only when their factory explicitly + supports request-converter composition. Multi-turn and simulated-conversation techniques are + not offered. - **jailbreaks** — which jailbreak templates to run (a random ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set). ``prompt_sending`` applies each template as a ``TextJailbreakConverter`` on the outgoing request, so the objective is rendered inline into the template's ``{{prompt}}`` slot; this keeps that - delivery target-agnostic and lets it compose with every technique. ``jailbreak_system_prompt`` + delivery target-agnostic and lets it compose with compatible direct techniques. ``jailbreak_system_prompt`` instead sets the template as a native system prompt and sends the objective as its own user turn, so it is only built for targets that natively support editable history and system prompts (it is skipped for incapable targets, or raises if it is the only selected technique). Responses are scored to determine whether the jailbreak succeeded (non-refusal). """ - VERSION: int = 3 + VERSION: int = 4 #: Baseline (an un-jailbroken prompt-send over the objectives) is included by default: a model #: that complies with the bare objective is itself interesting signal. Callers opt out per run @@ -232,6 +253,30 @@ def __init__( scenario_result_id=scenario_result_id, ) + def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[ScenarioTechnique]: + """ + Resolve techniques while rejecting stale or incompatible enum members. + + Args: + scenario_techniques (Any): Requested Jailbreak technique members. + + Returns: + list[ScenarioTechnique]: Compatible concrete techniques. + + Raises: + ValueError: If a caller supplies members from an older or different + technique enum. + """ + if scenario_techniques: + incompatible = [item for item in scenario_techniques if not isinstance(item, self._technique_class)] + if incompatible: + values = [getattr(item, "value", repr(item)) for item in incompatible] + raise ValueError( + "Jailbreak received stale or incompatible techniques " + f"{values}. Select a compatible direct-delivery technique from JailbreakTechnique." + ) + return super()._resolve_scenario_techniques(scenario_techniques=scenario_techniques) + def _resolve_templates(self) -> list[str]: """ Resolve the jailbreak templates for this run, replaying the persisted set on resume. @@ -282,14 +327,127 @@ 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) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate the template and attempt axes, preserving the target capability caveat. + + Returns: + ScenarioDefaultRunSizeEstimate: 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, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], + is_baseline=True, + ) + ) + components.append( + ScenarioRunSizeComponent( + label="Inline jailbreak delivery", + count=seed_group_count * template_count * attempt_count * converter_count, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="jailbreak templates", count=template_count), + ScenarioRunSizeFactor(label="attempts", count=attempt_count), + ScenarioRunSizeFactor(label="inline delivery techniques", 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, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="jailbreak templates", count=template_count), + ScenarioRunSizeFactor(label="attempts", 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)." + ) + status = ( + ScenarioRunSizeEstimateStatus.Conditional + if system_delivery_selected and system_delivery_supported is None + else ScenarioRunSizeEstimateStatus.Exact + ) + if status is ScenarioRunSizeEstimateStatus.Conditional: + 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 ScenarioDefaultRunSizeEstimate( + status=status, + total_attack_count=planned_count if status is ScenarioRunSizeEstimateStatus.Exact else None, + 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). - ``prompt_sending`` (and any opt-in registry techniques) deliver each jailbreak template as a + ``prompt_sending`` (and compatible opt-in direct techniques) deliver each jailbreak template as a ``TextJailbreakConverter`` appended to that technique's request converters, so the objective - is rendered inline into the template's ``{{prompt}}`` slot on the wire — target-agnostic and - composable with every technique. ``jailbreak_system_prompt`` instead delivers the template as + is rendered inline into the template's ``{{prompt}}`` slot on the wire. ``jailbreak_system_prompt`` + instead delivers the template as a native system prompt (no converter) with the objective sent as its own user turn, so it is only built when the objective target natively supports editable history and system prompts. Results group by jailbreak template so per-template ASR rolls up naturally. @@ -314,6 +472,21 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list num_attempts = self.params.get("num_jailbreak_attempts", 1) technique_factories = resolve_technique_factories(context=context, extra_factories=_extra_default_factories()) + selected_names = {technique.value for technique in context.scenario_techniques} + missing = selected_names - set(technique_factories) + if missing: + raise ValueError( + "Jailbreak selected techniques that are no longer available: " + f"{sorted(missing)}. Refresh the plan and select compatible direct-delivery techniques." + ) + incompatible = [ + name for name, factory in technique_factories.items() if not _is_jailbreak_compatible_factory(factory) + ] + if incompatible: + raise ValueError( + "Jailbreak cannot compose with multi-turn, simulated-conversation, or " + f"non-composable techniques: {sorted(incompatible)}." + ) # ``jailbreak_system_prompt`` is delivered separately (native system prompt, no converter); # every other technique goes through the inline converter path. @@ -456,6 +629,7 @@ def _build_system_prompt_factory(self, *, template_file_name: str) -> AttackTech attack_class=PromptSendingAttack, technique_tags=["single_turn"], seed_technique=seed_technique, + supports_request_converter_composition=True, ) @staticmethod 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..43aa58ff73 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -30,7 +30,13 @@ AttackScoringConfig, CrescendoAttack, ) -from pyrit.models import SeedPrompt +from pyrit.models import ( + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, + 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 +489,46 @@ 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) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate the independent sub-harm technique sweeps and per-harm baselines. + + Returns: + ScenarioDefaultRunSizeEstimate: 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, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="default concrete techniques", count=technique_count), + ], + ) + ) + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label=f"{dataset_name} baseline", + count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], + is_baseline=True, + note="Psychosocial uses a distinct baseline and scorer for each sub-harm.", + ) + ) + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 006e695c75..cfe23046d0 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -11,11 +11,25 @@ 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, + ScenarioDefaultRunSizeEstimate, + ScenarioResult, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) 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: @@ -187,6 +201,72 @@ def __init__( scenario_result_id=scenario_result_id, ) + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate the target-by-technique matrix using execution compatibility. + + Returns: + ScenarioDefaultRunSizeEstimate: 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 ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + 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, + factors=[ + ScenarioRunSizeFactor(label="selected concrete techniques", count=1), + ScenarioRunSizeFactor(label="adversarial targets", count=len(resolved_targets)), + ScenarioRunSizeFactor(label="compatible logical seed groups", count=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 ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + components=components, + datasets=datasets, + note=( + f"Components describe the candidate population. {'; '.join(reasons)}, " + "so the authoritative total is unavailable before launch." + ), + ) + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 874ca6d752..e9d58b65a0 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -50,7 +50,13 @@ TreeOfAttacksWithPruningAttack, ) from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackConverterConfig, AttackScoringConfig -from pyrit.models import AttackSeedGroup +from pyrit.models import ( + AttackSeedGroup, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -414,6 +420,43 @@ def _resolve_foundry_techniques( self._scenario_composites = composites return flat + async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate one selected seed population per resolved Foundry composition. + + Returns: + ScenarioDefaultRunSizeEstimate: 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, + factors=[ + ScenarioRunSizeFactor(label="resolved Foundry composites", count=1), + ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count), + ], + ) + for composition in self._scenario_composites + ] + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=selected_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count)], + is_baseline=True, + ) + ) + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 1bf3e78e4a..010daec513 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -24,7 +24,16 @@ 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, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, + 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 @@ -220,29 +229,57 @@ 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) -> ScenarioDefaultRunSizeEstimate: """ - 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. + ScenarioDefaultRunSizeEstimate: 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, + factors=[ + ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count), + ScenarioRunSizeFactor(label="concrete converter variants", count=variant_count), + ScenarioRunSizeFactor( + label="raw plus decode prompt configurations", + 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, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], + is_baseline=True, + ) + ) + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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"), @@ -264,11 +301,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..d081b98b6b 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,16 @@ 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, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, + 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 +492,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 +525,97 @@ 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) -> ScenarioDefaultRunSizeEstimate: + """ + Estimate the technique-specific synthesized populations and their shared baseline. + + Returns: + ScenarioDefaultRunSizeEstimate: 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), + factors=[ScenarioRunSizeFactor(label="synthesized logical seed groups", 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, + factors=[ + ScenarioRunSizeFactor( + label="all synthesized logical seed groups", + count=synthesized_count, + ) + ], + is_baseline=True, + note="The baseline runs over the union of all default technique populations.", + ) + ) + return ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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..5b3c92cec5 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_request_converter_composition=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 b4be579675..e696429af7 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -159,6 +159,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_request_converter_composition=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters( diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py index dc41e698a3..849e050292 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,7 +16,13 @@ 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.backend.routes.scenarios import get_scenario_run_progress +from pyrit.models import ( + ScenarioProgressHeader, + ScenarioRunPlan, + ScenarioRunProgress, + ScenarioRunState, +) from pyrit.models.catalog.scenario import ScenarioRunSummary from unit.mocks import make_scenario_result @@ -121,6 +128,38 @@ 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.""" @@ -164,7 +203,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 +216,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..f1fe574069 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -5,6 +5,8 @@ Tests for ScenarioRunService. """ +import asyncio +import uuid from datetime import datetime, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -17,7 +19,22 @@ 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 @@ -67,6 +84,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 +97,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, ) @@ -292,6 +313,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 +429,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 +444,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.""" @@ -680,6 +740,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 +942,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 +987,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: @@ -1083,3 +1191,237 @@ 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_non_error(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 == 100 + assert len(summary.failed_attacks) == 2 + assert summary.total_retries == 3 + + +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=None) + + assert mapped.seed_group_id == "canonical-seed-id" + + +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..f1545a1f1e 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,39 @@ 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_run_service import ScenarioRunService from pyrit.backend.services.scenario_service import ( ScenarioService, get_scenario_service, ) -from pyrit.models import Parameter +from pyrit.models import ( + Parameter, + ScenarioDatasetSizeCap, + ScenarioDatasetSummary, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, +) 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 +70,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 +91,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 +139,278 @@ 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + started = asyncio.Event() + release = asyncio.Event() + + async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=1, + components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], + ) + started = asyncio.Event() + release = asyncio.Event() + + async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 result.items[0].default_run_size.status is ScenarioRunSizeEstimateStatus.Exact + assert result.items[1].default_run_size.status is ScenarioRunSizeEstimateStatus.Unavailable + 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 first.default_run_size.status is ScenarioRunSizeEstimateStatus.Unavailable + 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 +428,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 +473,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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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(ScenarioRunService, "resolve_target_name", 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + 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(ScenarioRunService, "resolve_target_name") 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 +646,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 +690,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 +722,20 @@ 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=ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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 +748,10 @@ 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"]["version"] == 1 + assert data["default_run_size"]["status"] == "exact" + assert data["default_run_size"]["total_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 +764,93 @@ 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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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()["total_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 = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_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/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 550e4b631d..2a25984ece 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -23,6 +23,7 @@ import pytest from unit.mocks import get_mock_scorer_identifier +from pyrit.converter import Base64Converter from pyrit.executor.attack import ConversationManager, ConversationState from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ( @@ -1093,6 +1094,25 @@ async def test_non_chat_target_behavior_normalize_is_default( text_value = context.next_message.get_piece().original_value assert len(text_value) > 0 + async def test_non_chat_target_rejects_converter_scoping_that_excludes_history_roles( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) + + with pytest.raises(ValueError, match="non-chat target.*excluded roles.*assistant"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( self, attack_identifier: ComponentIdentifier, @@ -1178,13 +1198,13 @@ async def test_non_chat_target_behavior_normalize_returns_empty_state( # apply_converters_to_roles Tests # ------------------------------------------------------------------------- - async def test_apply_converters_to_roles_default_applies_to_all( + async def test_apply_converters_to_roles_default_applies_to_user_only( self, attack_identifier: ComponentIdentifier, mock_chat_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that converters are applied to all roles by default.""" + """Test that converters are applied only to user history by default.""" mock_normalizer = MagicMock(spec=PromptNormalizer) mock_normalizer.convert_values_async = AsyncMock() manager = ConversationManager(prompt_normalizer=mock_normalizer) @@ -1201,8 +1221,7 @@ async def test_apply_converters_to_roles_default_applies_to_all( request_converters=converter_config, ) - # convert_values_async should be called for each message (both user and assistant) - assert mock_normalizer.convert_values_async.call_count == 2 + mock_normalizer.convert_values_async.assert_awaited_once() async def test_apply_converters_to_roles_user_only( self, 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 b1c34a6770..8a6adda76b 100644 --- a/tests/unit/executor/attack/component/test_prepended_conversation_config.py +++ b/tests/unit/executor/attack/component/test_prepended_conversation_config.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import get_args +from typing import get_type_hints from unittest.mock import MagicMock from pyrit.executor.attack.component.prepended_conversation_config import PrependedConversationConfig @@ -9,9 +9,13 @@ from pyrit.models import ChatMessageRole -def test_default_init_apply_converters_to_all_roles(): +def test_default_init_apply_converters_to_user_role(): config = PrependedConversationConfig() - assert config.apply_converters_to_roles == list(get_args(ChatMessageRole)) + assert config.apply_converters_to_roles == ["user"] + + +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(): diff --git a/tests/unit/executor/attack/core/test_attack_strategy.py b/tests/unit/executor/attack/core/test_attack_strategy.py index 6d588c447c..2da9d7464a 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy.py @@ -721,6 +721,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( @@ -735,6 +736,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( @@ -770,6 +772,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( @@ -788,6 +791,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/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 5d4fda2209..7b938b5d35 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -5,16 +5,18 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import get_mock_scorer_identifier, get_mock_target_identifier +from unit.mocks import MockPromptTarget, get_mock_scorer_identifier, get_mock_target_identifier from pyrit.converter import Base64Converter, StringJoinConverter from pyrit.executor.attack import ( AttackConverterConfig, AttackParameters, AttackScoringConfig, + PrependedConversationConfig, PromptSendingAttack, SingleTurnAttackContext, ) +from pyrit.memory import CentralMemory from pyrit.models import ( AttackOutcome, AttackResult, @@ -280,6 +282,60 @@ async def test_setup_updates_conversation_state_with_converters(self, mock_targe memory_labels={}, ) + async def test_default_converter_scoping_preserves_simulated_assistant_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack(objective_target=target, attack_converter_config=converter_config) + prepended_user = "prepended user request" + simulated_response = "simulated assistant response" + final_request = "live final request" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[ + Message.from_prompt(prompt=prepended_user, role="user"), + Message.from_prompt(prompt=simulated_response, role="assistant"), + ], + next_message=Message.from_prompt(prompt=final_request, role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + final_piece = next(piece for piece in pieces if piece.original_value == final_request) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value == assistant_piece.original_value + assert assistant_piece.converter_identifiers == [] + assert final_piece.converted_value != final_piece.original_value + assert [identifier.class_name for identifier in final_piece.converter_identifiers] == ["Base64Converter"] + + async def test_explicit_assistant_role_opt_in_converts_simulated_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack( + objective_target=target, + attack_converter_config=converter_config, + prepended_conversation_config=PrependedConversationConfig(apply_converters_to_roles=["assistant"]), + ) + simulated_response = "assistant history explicitly converted" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[Message.from_prompt(prompt=simulated_response, role="assistant")], + next_message=Message.from_prompt(prompt="live request", role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value != assistant_piece.original_value + assert [identifier.class_name for identifier in assistant_piece.converter_identifiers] == ["Base64Converter"] + @pytest.mark.usefixtures("patch_central_database") class TestPromptPreparation: 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..d54ee91560 --- /dev/null +++ b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py @@ -0,0 +1,129 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for lightweight scenario progress memory queries.""" + +import uuid +from datetime import datetime, timezone + +from unit.mocks import get_mock_target_identifier, make_scenario_result + +from pyrit.memory import MemoryInterface +from pyrit.memory.memory_interface import ScenarioProgressKeysetCursor +from pyrit.models import ( + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + 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=ScenarioProgressKeysetCursor( + 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 == {} 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 aadbaa8e70..0df3bd14df 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..c8451ca502 --- /dev/null +++ b/tests/unit/models/test_scenario_catalog.py @@ -0,0 +1,189 @@ +# 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, + ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateRequest, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, +) + + +def test_run_size_estimate_compatibility_alias_is_canonical_model() -> None: + """The initial DTO name remains an unambiguous alias of the versioned model.""" + assert ScenarioRunSizeEstimate is ScenarioDefaultRunSizeEstimate + + +def test_run_size_estimate_accepts_legacy_fields_and_serializes_canonically() -> None: + """Legacy constructors parse while the wire shape remains singular and versioned.""" + estimate = ScenarioRunSizeEstimate.model_validate( + { + "status": "exact", + "total": 2, + "components": [{"label": "Sweep", "count": 2}], + "datasets": [ + { + "name": "harmbench", + "seed_group_count": 100, + "selected_seed_group_count": 2, + } + ], + "caveat": "Legacy explanation.", + } + ) + + assert estimate.total == 2 + assert estimate.caveat == "Legacy explanation." + payload = estimate.model_dump(mode="json") + assert payload["version"] == 1 + assert payload["total_attack_count"] == 2 + assert payload["note"] == "Legacy explanation." + assert payload["datasets"][0]["logical_seed_group_count"] == 100 + assert "total" not in payload + assert "caveat" not in payload + assert "seed_group_count" not in payload["datasets"][0] + + +def test_run_size_estimate_normalizes_legacy_componentless_exact_total() -> None: + estimate = ScenarioRunSizeEstimate.model_validate({"status": "exact", "total": 2}) + + assert estimate.total_attack_count == 2 + assert estimate.components == [ + ScenarioRunSizeComponent( + label="Legacy total", + count=2, + note="Normalized from a legacy component-less estimate.", + ) + ] + + +def test_exact_default_run_size_requires_component_total() -> None: + """Exact estimates reject totals that do not match their additive components.""" + with pytest.raises(ValidationError, match="components total 6, not 7"): + ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=7, + components=[ + ScenarioRunSizeComponent( + label="Techniques", + count=6, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + ], + ) + + +def test_run_size_component_requires_factor_product() -> None: + """Components reject counts that disagree with their ordered formula factors.""" + with pytest.raises(ValidationError, match="factor product \\(6\\)"): + ScenarioRunSizeComponent( + label="Techniques", + count=7, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + + +def test_default_run_size_serializes_versioned_api_shape() -> None: + """The estimate exposes stable status, total, component, and factor fields.""" + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Exact, + total_attack_count=6, + components=[ + ScenarioRunSizeComponent( + label="Techniques", + count=6, + factors=[ + ScenarioRunSizeFactor(label="seed groups", count=3), + ScenarioRunSizeFactor(label="techniques", count=2), + ], + ) + ], + ) + + assert estimate.model_dump(mode="json") == { + "version": 1, + "status": "exact", + "total_attack_count": 6, + "components": [ + { + "label": "Techniques", + "count": 6, + "factors": [ + {"label": "seed groups", "count": 3}, + {"label": "techniques", "count": 2}, + ], + "note": None, + "is_baseline": False, + } + ], + "datasets": [], + "note": None, + "retries_included": False, + } + + +def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: + """Conditionality and effective dataset selection are machine-readable.""" + estimate = ScenarioDefaultRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Conditional, + 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 total depends on target capabilities.", + ) + + payload = estimate.model_dump(mode="json") + assert payload["status"] == "conditional" + assert payload["total_attack_count"] is None + assert payload["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", + } + ], + } + ] + assert payload["retries_included"] is False + + +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 5cce529015..b87c44f8da 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -8,14 +8,22 @@ import pytest +from pyrit.backend.services.scenario_run_service import ScenarioRunService from pyrit.common.path import JAILBREAK_TEMPLATES_PATH 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, + ScenarioRunSizeEstimateStatus, + SeedObjective, + SeedPrompt, +) from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.scenarios.airt.jailbreak import ( @@ -202,6 +210,65 @@ 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 8 + assert [component.label for component in estimate.components] == ["Inline jailbreak delivery"] + assert [(factor.label, factor.count) for factor in estimate.components[0].factors] == [ + ("selected logical seed groups", 4), + ("jailbreak templates", 2), + ("attempts", 1), + ("inline delivery techniques", 1), + ] + 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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_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 ): @@ -328,9 +395,7 @@ def _spy_create(self, **kwargs): assert captured, "Expected factory.create to be called" converters = [c for extra in captured if extra for cc in extra for c in cc.converters] - assert any(isinstance(c, TextJailbreakConverter) for c in converters), ( - "Expected a TextJailbreakConverter to be threaded to factory.create" - ) + assert sum(isinstance(c, TextJailbreakConverter) for c in converters) == 1 # The objective seed groups themselves carry no jailbreak framing (converter delivery only). for attack in scenario._atomic_attacks: @@ -376,49 +441,61 @@ def _spy_create(self, **kwargs): "Jailbreak converter must be applied before caller-supplied converters" ) - async def test_simulated_conversation_techniques_produce_attacks_with_jailbreak( + async def test_stale_incompatible_technique_is_rejected( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): - """Regression: simulated-conversation techniques (``role_play_*``, ``crescendo_*``) must still - produce atomic attacks when crossed with a jailbreak template, and each must receive the - jailbreak converter. - - Converter delivery leaves the objective seed group unframed, so it stays compatible with the - simulated-conversation seed technique. (Delivering the jailbreak as a system-role framing seed - instead collided with that technique's seed range and silently produced zero attacks.) - """ - technique_class = _build_jailbreak_technique() - techniques = [ - technique_class("role_play_movie_script"), - technique_class("crescendo_simulated"), - ] - captured: list[Any] = [] - original_create = AttackTechniqueFactory.create - - def _spy_create(self, **kwargs): - captured.append(kwargs.get("extra_request_converters")) - return original_create(self, **kwargs) + registry_factories = list(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise().values()) + legacy_class = AttackTechniqueRegistry.build_technique_class_from_factories( + class_name="LegacyJailbreakTechnique", + factories=registry_factories, + ) + with _patch_seed_groups(mock_memory_seed_groups): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args=_default_args( + mock_objective_target, + scenario_techniques=[legacy_class("tap")], + jailbreak_names=["aim.yaml"], + ) + ) + with pytest.raises(ValueError, match="stale or incompatible"): + await scenario.initialize_async() + async def test_incompatible_runtime_factory_is_rejected( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + ): + incompatible = AttackTechniqueFactory( + name="legacy_multi_turn", + attack_class=PromptSendingAttack, + technique_tags=["multi_turn"], + supports_request_converter_composition=True, + ) with _patch_seed_groups(mock_memory_seed_groups): - with patch.object(AttackTechniqueFactory, "create", _spy_create): + with patch( + "pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", + return_value={_PROMPT_SENDING: incompatible}, + ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) + technique_class = _build_jailbreak_technique() scenario.set_params_from_args( args=_default_args( - mock_objective_target, scenario_techniques=techniques, jailbreak_names=["aim.yaml"] + mock_objective_target, + scenario_techniques=[technique_class(_PROMPT_SENDING)], + jailbreak_names=["aim.yaml"], ) ) - await scenario.initialize_async() - names = {a.atomic_attack_name for a in scenario._atomic_attacks} - assert "role_play_movie_script_aim_harmbench" in names - assert "crescendo_simulated_aim_harmbench" in names - # Every build of a simulated-conversation technique must still carry the jailbreak - # converter. Assert both techniques captured a non-empty converter stack (so the check - # can't pass vacuously on a dropped/None stack) and each contains the jailbreak converter. - populated = [extra for extra in captured if extra] - assert len(populated) == 2, "Expected both simulated-conversation techniques to receive converters" - assert all( - any(isinstance(c, TextJailbreakConverter) for cc in extra for c in cc.converters) for extra in populated - ), "Each simulated-conversation technique must receive the jailbreak converter" + with pytest.raises(ValueError, match="cannot compose"): + await scenario.initialize_async() + + async def test_missing_runtime_factory_is_rejected( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + ): + with _patch_seed_groups(mock_memory_seed_groups): + with patch("pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", return_value={}): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args(args=_default_args(mock_objective_target, jailbreak_names=["aim.yaml"])) + with pytest.raises(ValueError, match="no longer available.*prompt_sending"): + await scenario.initialize_async() async def test_all_templates_produce_attacks( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups @@ -708,15 +785,54 @@ def test_default_techniques_are_the_two_deliveries(self): assert default_values == set(_DEFAULT_TECHNIQUES) assert default_values == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} - def test_registry_techniques_are_available(self): + def test_only_compatible_direct_registry_techniques_are_available(self): technique_class = _technique_class() available = {t.value for t in technique_class.get_all_techniques()} - assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT}.issubset(available) - # The "normal ones available like from rapid response" are exposed as opt-in techniques. - assert {"role_play_movie_script", "many_shot", "tap"}.issubset(available) + incompatible = { + "context_compliance", + "role_play_movie_script", + "role_play_video_game", + "role_play_trivia_game", + "role_play_persuasion", + "role_play_persuasion_written", + "crescendo_simulated", + "crescendo_movie_director", + "crescendo_history_lecture", + "crescendo_journalist_interview", + "red_teaming", + "tap", + "many_shot", + "pair", + } + assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT, "flip"}.issubset(available) + assert incompatible.isdisjoint(available) + + def test_registry_metadata_omits_incompatible_techniques(self): + metadata = ScenarioRegistry()._build_metadata("airt.jailbreak", Jailbreak) + incompatible = { + "context_compliance", + "role_play_movie_script", + "crescendo_simulated", + "red_teaming", + "tap", + "many_shot", + } + assert metadata.scenario_version == 4 + assert metadata.default_techniques == (_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT) + assert incompatible.isdisjoint(metadata.default_techniques) + assert "flip" in metadata.all_techniques + assert incompatible.isdisjoint(metadata.all_techniques) + + def test_configured_estimate_rejects_context_compliance(self): + with pytest.raises(ValueError, match="Technique 'context_compliance' not found"): + ScenarioRunService.resolve_scenario_configuration( + scenario_name="airt.jailbreak", + scenario_class=Jailbreak, + techniques=["context_compliance"], + ) - def test_scenario_version_is_three(self): - assert Jailbreak.VERSION == 3 + def test_scenario_version_is_four(self): + assert Jailbreak.VERSION == 4 def test_default_dataset_is_harmbench(self): assert Jailbreak.required_datasets() == ["harmbench"] 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 28fd9d8004..282e363516 100644 --- a/tests/unit/scenario/core/test_atomic_attack.py +++ b/tests/unit/scenario/core/test_atomic_attack.py @@ -1105,7 +1105,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, @@ -1117,13 +1117,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( @@ -1137,10 +1137,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 @@ -1159,9 +1163,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 39e3046685..a52773c301 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -178,6 +178,29 @@ def test_validate_kwargs_rejects_invalid_param_on_real_attack_class(self): attack_kwargs={"nonexistent_param": 42}, ) + 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_request_converter_composition=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_request_converter_composition=True, + ) + + assert not default_factory.supports_request_converter_composition + assert composable_factory.supports_request_converter_composition + 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 73dbc5a048..d7f7fd7cd8 100644 --- a/tests/unit/scenario/core/test_scenario.py +++ b/tests/unit/scenario/core/test_scenario.py @@ -16,7 +16,15 @@ 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, +) from pyrit.scenario import ( DatasetAttackConfiguration, DatasetConfiguration, @@ -42,6 +50,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 @@ -259,6 +277,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.""" @@ -421,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 @@ -514,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.""" @@ -1042,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), ) ) @@ -1106,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), ) ) @@ -1129,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) @@ -1165,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"} @@ -1203,6 +1313,55 @@ 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_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 74bde8b988..8611a9fb93 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -12,7 +12,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.scenario import DatasetConfiguration, ScenarioResult from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy, Scenario, ScenarioTechnique @@ -70,6 +78,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 @@ -79,13 +88,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 f1ed5a1572..34c2a7f21e 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..509985104f --- /dev/null +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -0,0 +1,641 @@ +# 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, + ScenarioRunSizeEstimateStatus, + 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 2 + assert [component.count for component in estimate.components] == [2] + assert [factor.count for factor in estimate.components[0].factors] == [2, 1] + 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 4 + assert [factor.count for factor in estimate.components[0].factors] == [2, 2] + + +@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.total_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.total_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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate_without_target.total_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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_attack_count is None + assert [component.count for component in estimate.components] == [4, 8, 8] + assert [factor.count for factor in estimate.components[1].factors] == [4, 2, 1, 1] + 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 8 + assert [component.count for component in estimate.components] == [8] + assert [factor.count for factor in estimate.components[0].factors] == [4, 2, 1, 1] + 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 12 + assert [component.count for component in estimate.components] == [4, 8] + assert estimate.components[0].is_baseline is True + assert [factor.count for factor in estimate.components[1].factors] == [4, 2, 1, 1] + 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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_attack_count == 152 + assert [factor.count for factor in estimate.components[0].factors] == [2, 15, 5] + + +@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 estimate.status is ScenarioRunSizeEstimateStatus.Exact + assert len(synthesized) == len(scenario._scenario_techniques) + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Exact + assert estimate.total_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.status is ScenarioRunSizeEstimateStatus.Conditional + assert estimate.total_attack_count is None + assert estimate.components == [] + assert "adversarial_targets" in estimate.note + + +@pytest.mark.parametrize( + ("use_cached", "expected_status", "expected_total"), + [ + (False, ScenarioRunSizeEstimateStatus.Exact, 6), + (True, ScenarioRunSizeEstimateStatus.Conditional, None), + ], +) +@pytest.mark.usefixtures("patch_central_database") +async def test_adversarial_benchmark_resolves_targets_and_filters_each_technique( + *, + use_cached: bool, + expected_status: ScenarioRunSizeEstimateStatus, + 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.status is expected_status + assert estimate.total_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.total_attack_count == 6 + assert len(estimate.components) == 2 + assert [component.count for component in estimate.components] == [3, 3] + assert [factor.count for factor in estimate.components[0].factors] == [1, 3]