From ff73496fe142d972ec3dec1d15b0b13081a690e6 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 1/4] 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 278bf9c466..be7e2c4c55 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 d698593a5b..3402f26f5b 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: """ @@ -3644,6 +3655,12 @@ def update_scenario_run_state( entry.scenario_run_state = scenario_run_state.value entry.error_message = error_message entry.error_type = error_type + if scenario_run_state in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ): + entry.completion_time = datetime.now(tz=timezone.utc) session.commit() @@ -3677,6 +3694,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 2cc9ef0dc7..c10beefba7 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1555,6 +1555,13 @@ class AttackResultEntry(Base): Index("ix_AttackResultEntries_conversation_id", "conversation_id"), # Serves the History recency ORDER BY timestamp DESC, id DESC and its keyset seek. Index("ix_AttackResultEntries_timestamp_id", "timestamp", "id"), + # Serves scenario progress deltas scoped by parent and ordered oldest-first. + Index( + "ix_AttackResultEntries_attribution_parent_timestamp_id", + "attribution_parent_id", + "timestamp", + "id", + ), {"extend_existing": True}, ) id = mapped_column(CustomUUID, nullable=False, primary_key=True) @@ -1864,12 +1871,9 @@ class ScenarioResultEntry(Base): error_message: Mapped[str | None] = mapped_column(Unicode, nullable=True) error_type: Mapped[str | None] = mapped_column(String, nullable=True) - # Free-form JSON metadata stamped by the scenario. Currently used to record - # ``objective_hashes`` — the objective sha256 set chosen on the - # first run, replayed on resume so a fresh ``random.sample`` can't - # silently change which objectives the scenario operates on. Column is - # named ``scenario_metadata`` because SQLAlchemy's ``DeclarativeBase`` - # reserves ``metadata`` as a class attribute on the model. + # Free-form JSON metadata stamped by the scenario. Stores the normalized run + # plan and sampled objective hashes. Column is named ``scenario_metadata`` + # because SQLAlchemy's ``DeclarativeBase`` reserves ``metadata``. scenario_metadata: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) def __init__(self, *, entry: ScenarioResult) -> None: diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 667a2ca033..23470de50c 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, ) @@ -78,6 +89,17 @@ from pyrit.models.results.scenario_result import ScenarioResult, ScenarioRunState from pyrit.models.results.strategy_result import StrategyResult, StrategyResultT from pyrit.models.retry_event import RetryEvent +from pyrit.models.scenario_progress import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + SCENARIO_RUN_PLAN_VERSION, + ScenarioAttackResultDelta, + ScenarioProgressHeader, + ScenarioProgressResult, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunProgress, +) from pyrit.models.score import ( Condition, ContentScorable, @@ -171,6 +193,7 @@ "IdentifierFilter", "IdentifierType", "JSONValue", + "logical_seed_group_fingerprint", "COMMON_JSON_SCHEMAS", "JsonResponseConfig", "get_common_json_schema", @@ -205,8 +228,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 0ad434542a..0785cca63b 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -22,7 +22,13 @@ from pyrit.executor.attack import AttackExecutor, AttackExecutorResult from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.memory import CentralMemory -from pyrit.models import AtomicAttackEvaluationIdentifier, AtomicAttackIdentifier, AttackResult, AttackSeedGroup +from pyrit.models import ( + AtomicAttackEvaluationIdentifier, + AtomicAttackIdentifier, + AttackResult, + AttackSeedGroup, + config_hash, +) if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget @@ -191,6 +197,16 @@ def technique_eval_hash(self) -> str: ) return AtomicAttackEvaluationIdentifier(composite).eval_hash + @property + def logical_group_id(self) -> str: + """The stable identity of this planned atomic-attack group.""" + return config_hash( + { + "atomic_attack_name": self.atomic_attack_name, + "technique_eval_hash": self.technique_eval_hash, + } + ) + @property def objectives(self) -> list[str]: """ @@ -324,13 +340,17 @@ async def run_async( # a Scenario. The same attribution object is stamped on every # per-task AttackContext; per-task identity is reconstructed from # the row's own objective_sha256 (no positional state required). - attribution: AttackResultAttribution | None = None + attributions: list[AttackResultAttribution] | None = None if self._scenario_result_id is not None: - attribution = AttackResultAttribution( - parent_id=self._scenario_result_id, - parent_collection=self.atomic_attack_name, - parent_eval_hash=self.technique_eval_hash, - ) + attributions = [ + AttackResultAttribution( + parent_id=self._scenario_result_id, + parent_collection=self.atomic_attack_name, + parent_eval_hash=self.technique_eval_hash, + seed_group_id=seed_group.logical_id, + ) + for seed_group in self._seed_groups + ] untyped_results = await executor.execute_attack_from_seed_groups_async( attack=technique.attack, @@ -339,7 +359,7 @@ async def run_async( objective_scorer=self._objective_scorer, memory_labels=self._memory_labels, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, + attributions=attributions, **self._attack_execute_params, ) completed_results: list[AttackResult] = [] diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 990957386c..25083ae5ae 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -84,6 +84,7 @@ def __init__( adversarial_seed_prompt: SeedPrompt | str | None = None, seed_technique: AttackTechniqueSeedGroup | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -121,6 +122,9 @@ def __init__( chat during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Authors can override the derivation explicitly. + supports_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. @@ -145,11 +149,13 @@ class constructor signature and seed-technique shape. adversarial_system_prompt is not None or adversarial_seed_prompt is not None ) self._seed_technique = seed_technique + self._supports_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 @@ -168,6 +174,7 @@ def with_simulated_conversation( attack_kwargs: dict[str, Any] | None = None, adversarial_chat: PromptTarget | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -219,6 +226,9 @@ def with_simulated_conversation( during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Forwarded to the factory constructor. + supports_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 @@ -279,6 +289,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, ) @@ -314,6 +325,23 @@ def _validate_adversarial_flags(self) -> None: f"should not have one wired." ) + def _validate_converter_composition(self) -> None: + """ + Validate that an opt-in factory can receive additive request converters. + + Raises: + ValueError: If composition is enabled but the attack constructor does + not accept ``attack_converter_config``. + """ + if ( + self._supports_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. @@ -484,6 +512,11 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial + @property + def supports_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.""" @@ -828,6 +861,7 @@ def _build_identifier(self) -> ComponentIdentifier: "attack_class": self._attack_class.__name__, "kwargs": kwargs_for_id, "uses_adversarial": self._uses_adversarial, + "supports_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 9385c8c97b..2ef11604ff 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]: """ @@ -526,63 +559,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 @@ -591,10 +787,6 @@ async def initialize_async(self) -> None: self._max_retries = params.get("max_retries", 0) self._memory_labels = params.get("memory_labels") or {} - # Resolve the effective include_baseline. Forbidden is checked first so a forbidden - # scenario type never silently inherits a True default; explicit-True on a forbidden - # type is a hard error rather than a silent ignore. For the Enabled / Disabled states, - # a None runtime value defers to the policy. include_baseline = params.get("include_baseline") if self.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden: if include_baseline is True: @@ -605,16 +797,50 @@ async def initialize_async(self) -> None: include_baseline = False elif include_baseline is None: include_baseline = self.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Enabled - self._include_baseline = include_baseline - # Prepare scenario techniques via the resolution hook (subclasses override to widen - # accepted types or expand composites) and stash any per-technique converter overrides. self._scenario_techniques = self._resolve_scenario_techniques( scenario_techniques=params.get("scenario_techniques") ) self._technique_converters = params.get("technique_converters") or {} + @final + async def initialize_async(self) -> None: + """ + Initialize the scenario by populating self._atomic_attacks and creating the ScenarioResult. + + All run inputs are read from the parameter bag (``self.params``), which is populated by + ``set_params_from_args`` from the merged CLI / config / programmatic arguments. Callers + fill the bag then initialize: + + .. code-block:: python + + scenario.set_params_from_args(args={"objective_target": target, "max_concurrency": 8}) + await scenario.initialize_async() + + This method allows scenarios to be initialized with atomic attacks after construction, + which is useful when atomic attacks require async operations to be built. + + If a scenario_result_id was provided in __init__, this method will check if it exists + in memory and validate that the stored scenario matches the current configuration. + If it matches, the scenario will resume from prior progress. If it doesn't match or + doesn't exist, a new scenario result will be created. + + The common run inputs read from the bag are ``objective_target`` (a ``PromptTarget`` + instance or a registered target name resolved against ``TargetRegistry``), + ``scenario_techniques``, ``technique_converters``, ``dataset_config``, + ``max_concurrency``, ``max_retries``, ``memory_labels``, and ``include_baseline`` + (see ``_common_scenario_parameters``). A subclass that removes a common input via + ``supported_parameters`` falls back to that input's default here. + + Raises: + ValueError: If ``objective_target`` is declared but not resolvable (neither supplied + nor registered as a default), if a supplied target name is not registered in + ``TargetRegistry``, or if ``include_baseline=True`` is set for a scenario whose + ``BASELINE_ATTACK_POLICY`` is ``Forbidden``. + """ + self._resolve_runtime_configuration(require_objective_target=True) + # Build atomic attacks: resolve the seed groups once, snapshot the resolved inputs # into a ScenarioContext, and hand it to the subclass extension point. Baseline emission # is the scenario's own responsibility — matrix scenarios get it for free (the matrix @@ -653,7 +879,19 @@ async def initialize_async(self) -> None: stored_result=existing_results[0], current_identifier=scenario_identifier, ) - self._apply_persisted_objectives(stored_result=existing_results[0]) + stored_result = existing_results[0] + stored_plan = self._get_stored_run_plan(stored_result=stored_result) + if stored_plan is not None: + self._apply_persisted_run_plan(stored_plan=stored_plan) + else: + self._apply_persisted_objectives(stored_result=stored_result) + reconstructed_plan = self._build_run_plan() + metadata = dict(stored_result.metadata) + metadata[SCENARIO_RUN_PLAN_METADATA_KEY] = reconstructed_plan.model_dump(mode="json") + self._memory.update_scenario_metadata( + scenario_result_id=self._scenario_result_id, + metadata=metadata, + ) return # Valid resume - skip creating new scenario result # Build display group mapping from atomic attacks @@ -688,28 +926,131 @@ def _build_initial_scenario_metadata(self) -> dict[str, Any]: chosen objective hashes here so the next ``_setup_scenario_async`` can replay them via ``keep_seed_groups_with_hashes``. - When ``max_dataset_size`` is not set, the sample equals the dataset and - nothing needs pinning; the dict is empty. + The normalized run plan is always stored. When ``max_dataset_size`` is not + set, only the run plan is needed because the full dataset is deterministic. Returns: dict[str, Any]: Metadata payload for the new ScenarioResult. """ metadata: dict[str, Any] = {} - if getattr(self._dataset_config, "max_dataset_size", None) is None: - return metadata - hashes: list[str] = [] - seen: set[str] = set() - for aa in self._atomic_attacks: - for sg in aa.seed_groups: - if sg.objective is None: - continue - sha = to_sha256(sg.objective.value) - if sha not in seen: - seen.add(sha) - hashes.append(sha) - metadata["objective_hashes"] = hashes + if getattr(self._dataset_config, "max_dataset_size", None) is not None: + hashes: list[str] = [] + seen: set[str] = set() + for aa in self._atomic_attacks: + for sg in aa.seed_groups: + sha = to_sha256(sg.objective.value) + if sha not in seen: + seen.add(sha) + hashes.append(sha) + metadata["objective_hashes"] = hashes + metadata[SCENARIO_RUN_PLAN_METADATA_KEY] = self._build_run_plan().model_dump(mode="json") return metadata + def _build_run_plan(self) -> ScenarioRunPlan: + """ + Build the normalized persistent plan for the initialized atomic attacks. + + Returns: + ScenarioRunPlan: The versioned run plan. + """ + seed_groups: dict[str, ScenarioRunPlanSeedGroup] = {} + atomic_groups: list[ScenarioRunPlanAtomicGroup] = [] + for atomic_attack in self._atomic_attacks: + seed_group_ids: list[str] = [] + seen_seed_group_ids: set[str] = set() + for seed_group in atomic_attack.seed_groups: + seed_group_id = seed_group.logical_id + if seed_group_id in seen_seed_group_ids: + continue + seen_seed_group_ids.add(seed_group_id) + seed_group_ids.append(seed_group_id) + seed_groups.setdefault( + seed_group_id, + ScenarioRunPlanSeedGroup( + id=seed_group_id, + objective_sha256=to_sha256(seed_group.objective.value), + objective=seed_group.objective.value, + ), + ) + technique_eval_hash = str(atomic_attack.technique_eval_hash) + atomic_group_id = self._get_atomic_group_id(atomic_attack=atomic_attack) + atomic_groups.append( + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name=atomic_attack.atomic_attack_name, + display_group=atomic_attack.display_group, + technique_eval_hash=technique_eval_hash, + seed_group_ids=seed_group_ids, + ) + ) + return ScenarioRunPlan( + scenario_registry_name=self._scenario_registry_name, + atomic_groups=atomic_groups, + seed_groups=list(seed_groups.values()), + ) + + @staticmethod + def _get_atomic_group_id(*, atomic_attack: AtomicAttack) -> str: + """ + Compute the stable ID of an atomic group from its name and technique. + + Returns: + str: The atomic-group ID. + """ + return config_hash( + { + "atomic_attack_name": atomic_attack.atomic_attack_name, + "technique_eval_hash": str(atomic_attack.technique_eval_hash), + } + ) + + @staticmethod + def _get_stored_run_plan(*, stored_result: ScenarioResult) -> ScenarioRunPlan | None: + """ + Load and validate a stored run plan. + + Returns: + ScenarioRunPlan | None: The plan, or None for a legacy row. + """ + raw_plan = (stored_result.metadata or {}).get(SCENARIO_RUN_PLAN_METADATA_KEY) + if raw_plan is None: + return None + return ScenarioRunPlan.model_validate(raw_plan) + + def _apply_persisted_run_plan(self, *, stored_plan: ScenarioRunPlan) -> None: + """ + Validate and replay the exact logical units captured by a stored plan. + + Raises: + ValueError: If a planned atomic or seed group cannot be reconstructed. + """ + current_by_id = { + self._get_atomic_group_id(atomic_attack=atomic_attack): atomic_attack + for atomic_attack in self._atomic_attacks + } + planned_ids = {group.id for group in stored_plan.atomic_groups} + missing_groups = planned_ids - current_by_id.keys() + if missing_groups: + raise ValueError( + f"Scenario result id '{self._scenario_result_id}' cannot resume: " + f"{len(missing_groups)} planned atomic group(s) are no longer reconstructable." + ) + + retained_attacks: list[AtomicAttack] = [] + for planned_group in stored_plan.atomic_groups: + atomic_attack = current_by_id[planned_group.id] + current_seed_groups = {seed_group.logical_id: seed_group for seed_group in atomic_attack.seed_groups} + missing_seed_groups = set(planned_group.seed_group_ids) - current_seed_groups.keys() + if missing_seed_groups: + raise ValueError( + f"Scenario result id '{self._scenario_result_id}' cannot resume: atomic group " + f"'{planned_group.atomic_attack_name}' is missing {len(missing_seed_groups)} planned seed group(s)." + ) + atomic_attack._seed_groups = [current_seed_groups[group_id] for group_id in planned_group.seed_group_ids] + retained_attacks.append(atomic_attack) + self._atomic_attacks = retained_attacks + self._display_group_map = {group.atomic_attack_name: group.display_group for group in stored_plan.atomic_groups} + def _apply_persisted_objectives(self, *, stored_result: ScenarioResult) -> None: """ On resume, replay the originally-sampled objective subset. @@ -1303,6 +1644,8 @@ async def worker_async() -> None: atomic_attack = queue.get_nowait() except asyncio.QueueEmpty: return + atomic_group_id = atomic_attack.logical_group_id + self._active_atomic_groups[atomic_group_id] = atomic_attack.atomic_attack_name try: result = await atomic_attack.run_async( executor=shared_executor, @@ -1315,6 +1658,7 @@ async def worker_async() -> None: outcomes.append(exc) stop_event.set() finally: + self._active_atomic_groups.pop(atomic_group_id, None) pbar.update(1) # Cap workers at max_concurrency: that's also the objective-budget cap, and it's diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index fec03778d2..200aa2119b 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 5d0cc4f235..176215e5f6 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 baef99ac31..f5be3c2caa 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 @@ -223,29 +232,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"), @@ -267,11 +304,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 45ccc895dd..d0602588f7 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -174,6 +174,21 @@ def test_run_schema_migrations_applies_head_revision(): engine.dispose() +def test_scenario_progress_migration_adds_composite_index(): + """The migration head contains the parent/timestamp/id keyset index.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = os.path.join(temp_dir, "scenario-progress-index.db") + engine = create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, "head") + indexes = {index["name"] for index in inspect(connection).get_indexes("AttackResultEntries")} + assert "ix_AttackResultEntries_attribution_parent_timestamp_id" in indexes + finally: + engine.dispose() + + def test_migration_online_mode(): """ Test that online migration configuration is valid. diff --git a/tests/unit/models/test_attack_seed_group.py b/tests/unit/models/test_attack_seed_group.py index c3c6ff1c9c..567337d0a5 100644 --- a/tests/unit/models/test_attack_seed_group.py +++ b/tests/unit/models/test_attack_seed_group.py @@ -4,6 +4,7 @@ import pytest +from pyrit.models import AtomicAttackIdentifier, ComponentIdentifier from pyrit.models.seeds.attack_seed_group import AttackSeedGroup from pyrit.models.seeds.seed_objective import SeedObjective from pyrit.models.seeds.seed_prompt import SeedPrompt @@ -59,6 +60,40 @@ def test_attack_seed_group_consistent_group_id(): assert None not in group_ids +def test_logical_id_ignores_random_prompt_group_id_and_round_trips() -> None: + first = AttackSeedGroup(seeds=[_make_objective(value="goal"), _make_prompt(value="context")]) + second = AttackSeedGroup(seeds=[_make_objective(value="goal"), _make_prompt(value="context")]) + + assert first.seeds[0].prompt_group_id != second.seeds[0].prompt_group_id + assert first.logical_id == second.logical_id + + identifier = AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="Attack", class_module="tests"), + seed_group=first, + ) + restored = AtomicAttackIdentifier.model_validate(identifier.model_dump(mode="json")) + assert restored.logical_seed_group_id == first.logical_id + + +def test_logical_id_preserves_canonical_seed_order() -> None: + first = AttackSeedGroup( + seeds=[ + _make_objective(value="goal"), + _make_prompt(value="first", sequence=0), + _make_prompt(value="second", sequence=1), + ] + ) + second = AttackSeedGroup( + seeds=[ + _make_objective(value="goal"), + _make_prompt(value="second", sequence=0), + _make_prompt(value="first", sequence=1), + ] + ) + + assert first.logical_id != second.logical_id + + def test_attack_seed_group_with_multiple_prompts(): objective = _make_objective() p1 = _make_prompt(value="p1", sequence=0) diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py new file mode 100644 index 0000000000..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 8a04f20457..30df69492f 100644 --- a/tests/unit/scenario/core/test_atomic_attack.py +++ b/tests/unit/scenario/core/test_atomic_attack.py @@ -1119,7 +1119,7 @@ async def test_no_attribution_when_scenario_result_id_unset( self, mock_attack, sample_seed_groups, sample_attack_results ): """Outside a Scenario, ``_scenario_result_id`` is None and the - executor must receive ``attribution=None``.""" + executor must receive ``attributions=None``.""" atomic = AtomicAttack( attack_technique=AttackTechnique(attack=mock_attack), seed_groups=sample_seed_groups, @@ -1131,13 +1131,13 @@ async def test_no_attribution_when_scenario_result_id_unset( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - assert mock_exec.call_args.kwargs["attribution"] is None + assert mock_exec.call_args.kwargs["attributions"] is None async def test_attribution_built_when_scenario_result_id_set( self, mock_attack, sample_seed_groups, sample_attack_results ): """When the Scenario stamps ``_scenario_result_id`` onto the atomic - attack, ``run_async`` must build and pass a single attribution object.""" + attack, ``run_async`` must build and pass per-seed-group attribution.""" from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution atomic = AtomicAttack( @@ -1151,10 +1151,14 @@ async def test_attribution_built_when_scenario_result_id_set( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - attribution = mock_exec.call_args.kwargs["attribution"] - assert isinstance(attribution, AttackResultAttribution) - assert attribution.parent_id == "00000000-0000-0000-0000-000000000abc" - assert attribution.parent_collection == "MyAtomicAttack" + attributions = mock_exec.call_args.kwargs["attributions"] + assert len(attributions) == len(sample_seed_groups) + assert all(isinstance(attribution, AttackResultAttribution) for attribution in attributions) + assert all(attribution.parent_id == "00000000-0000-0000-0000-000000000abc" for attribution in attributions) + assert all(attribution.parent_collection == "MyAtomicAttack" for attribution in attributions) + assert [attribution.seed_group_id for attribution in attributions] == [ + seed_group.logical_id for seed_group in sample_seed_groups + ] async def test_attribution_includes_technique_eval_hash( self, mock_attack, sample_seed_groups, sample_attack_results @@ -1173,9 +1177,9 @@ async def test_attribution_includes_technique_eval_hash( mock_exec.return_value = wrap_results(sample_attack_results) await atomic.run_async() - attribution = mock_exec.call_args.kwargs["attribution"] - assert attribution.parent_eval_hash is not None - assert attribution.parent_eval_hash == atomic.technique_eval_hash + attributions = mock_exec.call_args.kwargs["attributions"] + assert all(attribution.parent_eval_hash is not None for attribution in attributions) + assert all(attribution.parent_eval_hash == atomic.technique_eval_hash for attribution in attributions) @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index 5ea41cd64b..ff4f61b90c 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -207,6 +207,29 @@ def test_cannot_append_text_converter_to_image_chain(self): assert not factory.can_append_request_converter(converter_type=TranslationConverter) + def test_request_converter_composition_requires_supported_constructor(self): + class _NoConverterAttack: + def __init__(self, *, objective_target, attack_scoring_config=None): + self.objective_target = objective_target + + with pytest.raises(ValueError, match="does not accept 'attack_converter_config'"): + AttackTechniqueFactory( + name="test", + attack_class=_NoConverterAttack, + supports_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 27c8cae746..29c2e81168 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.prompt_target import PromptTarget from pyrit.scenario import ( DatasetAttackConfiguration, @@ -43,6 +51,16 @@ def save_attack_results_to_memory(attack_results): memory.add_attack_results_to_memory(attack_results=attack_results) +def _make_identifiable_mock_attack() -> MagicMock: + """Create a mock attack with a valid canonical identifier for run-plan construction.""" + attack = MagicMock() + attack.get_identifier.return_value = ComponentIdentifier( + class_name="MockAttack", + class_module="tests.unit.scenario.core.test_scenario", + ) + return attack + + def _stamp_scenario_linkage(*, attack_results, atomic_attack): """ Stamp attribution_parent_id + attribution_data on each AttackResult the @@ -260,6 +278,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.""" @@ -422,6 +495,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 @@ -515,6 +621,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.""" @@ -1043,7 +1150,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), ) ) @@ -1107,7 +1214,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), ) ) @@ -1130,7 +1237,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) @@ -1166,6 +1273,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"} @@ -1204,6 +1314,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 d6f58e902b..febe218e83 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -17,7 +17,15 @@ from pyrit.exceptions import ScenarioPartialFailureException from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, ScenarioRunState +from pyrit.models import ( + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + ScenarioRunState, + SeedObjective, + config_hash, +) from pyrit.prompt_target import PromptTarget from pyrit.scenario import DatasetConfiguration, ScenarioResult from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy, Scenario, ScenarioTechnique @@ -76,6 +84,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str]) -> MagicMock: attack = MagicMock(spec=AtomicAttack) attack.atomic_attack_name = name attack.display_group = name + attack.technique_eval_hash = config_hash({"name": name, "objectives": objectives}) attack._attack = mock_attack_strategy attack._scenario_result_id = None @@ -85,13 +94,21 @@ def _set_scenario_result_id(scenario_result_id): attack.set_scenario_result_id = MagicMock(side_effect=_set_scenario_result_id) original_objectives = list(objectives) - current_objectives = {"value": list(objectives)} + current_seed_groups = { + "value": [AttackSeedGroup(seeds=[SeedObjective(value=objective)]) for objective in objectives] + } - type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) - type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) + type(attack).objectives = PropertyMock( + side_effect=lambda: [seed_group.objective.value for seed_group in current_seed_groups["value"]] + ) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_seed_groups["value"]) def drop_hashes(*, hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] + current_seed_groups["value"] = [ + seed_group + for seed_group in current_seed_groups["value"] + if to_sha256(seed_group.objective.value) not in hashes + ] attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) attack._original_objectives = original_objectives diff --git a/tests/unit/scenario/core/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py index af4a222d0b..d29bd0cd22 100644 --- a/tests/unit/scenario/core/test_scenario_retry.py +++ b/tests/unit/scenario/core/test_scenario_retry.py @@ -12,7 +12,15 @@ from pyrit.executor.attack import AttackParameters, AttackStrategy, SingleTurnAttackContext from pyrit.executor.attack.core import AttackExecutorResult from pyrit.memory import CentralMemory -from pyrit.models import AttackOutcome, AttackResult, AttackSeedGroup, ComponentIdentifier, Message, SeedObjective +from pyrit.models import ( + AttackOutcome, + AttackResult, + AttackSeedGroup, + ComponentIdentifier, + Message, + SeedObjective, + config_hash, +) from pyrit.prompt_target import PromptTarget from pyrit.scenario import DatasetConfiguration, ScenarioResult from pyrit.scenario.core import AtomicAttack, AttackTechnique, BaselineAttackPolicy, Scenario, ScenarioTechnique @@ -139,6 +147,7 @@ def create_mock_atomic_attack(name: str, objectives: list[str], run_async_mock: attack = MagicMock(spec=AtomicAttack) attack.atomic_attack_name = name attack.display_group = name + attack.technique_eval_hash = config_hash({"name": name, "objectives": objectives}) attack._attack = mock_attack_strategy attack._scenario_result_id = None @@ -151,12 +160,20 @@ def _set_scenario_result_id(scenario_result_id): # behaves correctly in resume tests. from pyrit.common.utils import to_sha256 - current_objectives = {"value": list(objectives)} - type(attack).objectives = PropertyMock(side_effect=lambda: current_objectives["value"]) - type(attack).seed_groups = PropertyMock(side_effect=lambda: current_objectives["value"]) + current_seed_groups = { + "value": [AttackSeedGroup(seeds=[SeedObjective(value=objective)]) for objective in objectives] + } + type(attack).objectives = PropertyMock( + side_effect=lambda: [seed_group.objective.value for seed_group in current_seed_groups["value"]] + ) + type(attack).seed_groups = PropertyMock(side_effect=lambda: current_seed_groups["value"]) def drop_hashes(*, hashes): - current_objectives["value"] = [o for o in current_objectives["value"] if to_sha256(o) not in hashes] + current_seed_groups["value"] = [ + seed_group + for seed_group in current_seed_groups["value"] + if to_sha256(seed_group.objective.value) not in hashes + ] attack.drop_seed_groups_with_hashes = MagicMock(side_effect=drop_hashes) diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py new file mode 100644 index 0000000000..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] From 1cfb607f2fb057c2ab954d2ab6113c22274d954e Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:30:10 -0700 Subject: [PATCH 2/4] Address scenario progress review feedback Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- doc/code/framework.md | 2 +- pyrit/backend/models/scenarios.py | 4 +- pyrit/backend/routes/scenarios.py | 14 +- pyrit/backend/services/attack_service.py | 12 +- .../scenario_configuration_resolver.py | 222 +++++++ .../backend/services/scenario_run_service.py | 619 ++++++------------ pyrit/backend/services/scenario_service.py | 46 +- pyrit/cli/_output.py | 5 +- pyrit/cli/api_client.py | 9 +- .../attack/component/conversation_manager.py | 11 - pyrit/executor/attack/core/attack_executor.py | 25 +- pyrit/memory/__init__.py | 4 +- pyrit/memory/memory_interface.py | 46 +- pyrit/models/__init__.py | 12 +- pyrit/models/catalog/__init__.py | 8 +- pyrit/models/catalog/scenario.py | 133 +--- pyrit/models/identifiers/__init__.py | 4 +- .../identifiers/atomic_attack_identifier.py | 4 +- pyrit/models/identifiers/seed_identifier.py | 4 +- pyrit/models/seeds/attack_seed_group.py | 4 +- .../registry/components/scenario_registry.py | 6 +- .../scenario/core/attack_technique_factory.py | 22 +- pyrit/scenario/core/scenario.py | 43 +- .../scenarios/adaptive/adaptive_scenario.py | 32 +- pyrit/scenario/scenarios/airt/jailbreak.py | 41 +- pyrit/scenario/scenarios/airt/psychosocial.py | 18 +- .../scenarios/benchmark/adversarial.py | 24 +- .../scenarios/foundry/red_team_agent.py | 18 +- pyrit/scenario/scenarios/garak/encoding.py | 22 +- .../scenario/scenarios/garak/web_injection.py | 20 +- pyrit/setup/initializers/techniques/airt.py | 2 +- pyrit/setup/initializers/techniques/core.py | 2 +- tests/unit/backend/test_attack_service.py | 4 +- .../unit/backend/test_scenario_run_routes.py | 20 +- .../unit/backend/test_scenario_run_service.py | 111 +++- tests/unit/backend/test_scenario_service.py | 83 +-- tests/unit/cli/test_api_client.py | 3 +- tests/unit/cli/test_output.py | 10 +- .../component/test_conversation_manager.py | 20 - .../test_interface_attack_results.py | 16 +- .../test_interface_scenario_progress.py | 27 +- tests/unit/models/test_scenario_catalog.py | 144 +--- tests/unit/scenario/airt/test_jailbreak.py | 21 +- .../core/test_attack_technique_factory.py | 8 +- .../test_default_run_size_estimates.py | 83 +-- 45 files changed, 883 insertions(+), 1105 deletions(-) create mode 100644 pyrit/backend/services/scenario_configuration_resolver.py diff --git a/doc/code/framework.md b/doc/code/framework.md index be7e2c4c55..4dca29100e 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -182,7 +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`. +- Factories opt into additive request-converter composition with `supports_additional_request_converters=True`. This is a semantic capability, not just constructor-signature detection; the factory validates that opted-in attacks accept `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. **Framework Plans**: diff --git a/pyrit/backend/models/scenarios.py b/pyrit/backend/models/scenarios.py index 56858f80b1..820133367d 100644 --- a/pyrit/backend/models/scenarios.py +++ b/pyrit/backend/models/scenarios.py @@ -13,7 +13,7 @@ from pydantic import BaseModel, Field from pyrit.backend.models.common import PaginationInfo -from pyrit.models.catalog.scenario import RegisteredScenario, ScenarioRunSummary +from pyrit.models.catalog.scenario import RegisteredScenario, ScenarioRunListItem __all__ = [ "ListRegisteredScenariosResponse", @@ -31,4 +31,4 @@ class ListRegisteredScenariosResponse(BaseModel): class ScenarioRunListResponse(BaseModel): """Response for listing scenario runs.""" - items: list[ScenarioRunSummary] = Field(..., description="List of scenario runs") + items: list[ScenarioRunListItem] = Field(..., description="List of scenario runs") diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 3245b644a8..a6e9ea5e00 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -26,7 +26,7 @@ from pyrit.models.catalog.scenario import ( RegisteredScenario, RunScenarioRequest, - ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, ScenarioRunSummary, ) @@ -92,7 +92,7 @@ async def get_scenario(scenario_name: str) -> RegisteredScenario: # pyrit-async @router.post( "/catalog/{scenario_name}/estimate", - response_model=ScenarioDefaultRunSizeEstimate, + response_model=ScenarioRunSizeEstimate, responses={ 400: {"model": ProblemDetail, "description": "Invalid estimate configuration"}, 404: {"model": ProblemDetail, "description": "Scenario not found"}, @@ -102,7 +102,7 @@ async def estimate_scenario_run_size( # pyrit-async-suffix-exempt *, scenario_name: str, request: ScenarioRunSizeEstimateRequest, -) -> ScenarioDefaultRunSizeEstimate: +) -> ScenarioRunSizeEstimate: """ Estimate a configured scenario without creating or persisting a run. @@ -111,7 +111,7 @@ async def estimate_scenario_run_size( # pyrit-async-suffix-exempt request: Techniques, datasets, baseline choice, and scenario parameters to preview. Returns: - ScenarioDefaultRunSizeEstimate: Structured request-specific planned-unit estimate. + ScenarioRunSizeEstimate: Structured request-specific planned-unit estimate. """ service = get_scenario_service() try: @@ -165,7 +165,9 @@ async def start_scenario_run(request: RunScenarioRequest) -> ScenarioRunSummary: "/runs", response_model=ScenarioRunListResponse, ) -async def list_scenario_runs(limit: int = Query(100, ge=1)) -> ScenarioRunListResponse: # pyrit-async-suffix-exempt +async def list_scenario_runs( + limit: int = Query(100, ge=1, le=100), +) -> ScenarioRunListResponse: # pyrit-async-suffix-exempt """ List tracked scenario runs (most recent first). @@ -176,7 +178,7 @@ async def list_scenario_runs(limit: int = Query(100, ge=1)) -> ScenarioRunListRe ScenarioRunListResponse: Runs, most recent first. """ service = get_scenario_run_service() - return service.list_runs(limit=limit) + return await run_in_threadpool(service.list_runs, limit=limit) @router.get( diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index 53c63ed0d8..cf34838627 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -58,7 +58,7 @@ from pyrit.backend.models.common import PaginationInfo from pyrit.backend.services.converter_service import get_converter_service from pyrit.backend.services.target_service import get_target_service -from pyrit.memory import AttackResultsKeysetCursor, CentralMemory, data_serializer_factory +from pyrit.memory import AttackResultKeysetCursor, CentralMemory, data_serializer_factory from pyrit.models import ( AtomicAttackIdentifier, AttackIdentifier, @@ -182,7 +182,7 @@ async def list_attacks_async( page_results = list(results[:limit]) next_cursor = ( self._encode_attack_cursor( - cursor=AttackResultsKeysetCursor.from_attack_result(page_results[-1]), + cursor=AttackResultKeysetCursor.from_attack_result(page_results[-1]), fingerprint=filter_fingerprint, ) if has_next_page and page_results @@ -912,7 +912,7 @@ def _norm_labels( return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16] @staticmethod - def _encode_attack_cursor(*, cursor: AttackResultsKeysetCursor, fingerprint: str) -> str: + def _encode_attack_cursor(*, cursor: AttackResultKeysetCursor, fingerprint: str) -> str: """ Encode a keyset anchor and its filter fingerprint into an opaque pagination cursor. @@ -933,7 +933,7 @@ def _encode_attack_cursor(*, cursor: AttackResultsKeysetCursor, fingerprint: str return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") @staticmethod - def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResultsKeysetCursor | None: + def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResultKeysetCursor | None: """ Decode the opaque list-attacks cursor into a keyset (seek) anchor. @@ -945,7 +945,7 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu raising or seeking within the wrong result set. Returns: - The decoded ``AttackResultsKeysetCursor``, or ``None`` to start at the first page. + The decoded ``AttackResultKeysetCursor``, or ``None`` to start at the first page. """ if not cursor: return None @@ -978,7 +978,7 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu # A crafted cursor near datetime's min/max with a large UTC offset overflows the # representable range when shifted to UTC; treat it as malformed and restart at page one. return None - return AttackResultsKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) + return AttackResultKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) # ======================================================================== # Private Helper Methods - Duplicate / Branch diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py new file mode 100644 index 0000000000..a0310bf595 --- /dev/null +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Shared scenario launch and estimate configuration resolution.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from pyrit.registry import ConverterRegistry, ScenarioRegistry, TargetRegistry + +if TYPE_CHECKING: + from pyrit.converter import Converter + from pyrit.prompt_target import PromptTarget + from pyrit.scenario import Scenario + +_CONVERTER_MODIFIER_PREFIX = "converter." + + +class ScenarioConfigurationResolver: + """Resolve registry-backed scenario inputs for launch and estimation.""" + + @staticmethod + def resolve_scenario_class(*, scenario_name: str) -> type[Scenario]: + """ + Resolve a registered scenario class. + + Returns: + type[Scenario]: The registered scenario class. + + Raises: + ValueError: If the scenario name is not registered. + """ + try: + return ScenarioRegistry.get_registry_singleton().get_class(scenario_name) + except KeyError as exc: + raise ValueError(str(exc)) from None + + @staticmethod + def resolve_target(*, target_name: str) -> PromptTarget: + """ + Resolve a registered target instance. + + Returns: + PromptTarget: The registered target. + + Raises: + ValueError: If the target name is not registered. + """ + instances = TargetRegistry.get_registry_singleton().instances + objective_target = instances.get(target_name) + if objective_target is not None: + return objective_target + + available_names = instances.get_names() + if not available_names: + raise ValueError( + f"Target '{target_name}' not found. The target registry is empty. " + "Make sure to include an initializer that registers targets " + "(e.g., initializers: ['target'])." + ) + raise ValueError( + f"Target '{target_name}' not found in registry. Available targets: {', '.join(available_names)}" + ) + + @classmethod + def resolve_configuration( + cls, + *, + scenario_name: str, + scenario_class: type[Scenario], + objective_target: Any | None = None, + techniques: list[str] | None = None, + dataset_names: list[str] | None = None, + max_dataset_size: int | None = None, + dataset_filters: dict[str, list[str]] | None = None, + include_baseline: bool | None = None, + max_concurrency: int | None = None, + max_retries: int | None = None, + memory_labels: dict[str, str] | None = None, + ) -> dict[str, Any]: + """ + Resolve shared launch and estimate fields into scenario parameters. + + Returns: + dict[str, Any]: Values accepted by ``Scenario.set_params_from_args``. + + Raises: + ValueError: If techniques or dataset overrides are invalid. + """ + resolved: dict[str, Any] = {} + if objective_target is not None: + resolved["objective_target"] = objective_target + if max_concurrency is not None: + resolved["max_concurrency"] = max_concurrency + if max_retries is not None: + resolved["max_retries"] = max_retries + if include_baseline is not None: + resolved["include_baseline"] = include_baseline + if memory_labels: + resolved["memory_labels"] = memory_labels + + filters = dataset_filters or {} + needs_introspection = bool(techniques) or bool(dataset_names) or max_dataset_size is not None or bool(filters) + if not needs_introspection: + return resolved + + try: + introspection_instance = scenario_class() # type: ignore[ty:missing-argument] + except Exception as exc: + raise ValueError( + f"Cannot resolve runtime configuration for scenario '{scenario_name}': " + f"scenario class is not instantiable without arguments ({exc})." + ) from exc + + if techniques: + technique_enums, technique_converters = cls.resolve_techniques_and_converters( + tokens=techniques, + technique_class=introspection_instance._technique_class, + scenario_name=scenario_name, + ) + resolved["scenario_techniques"] = technique_enums + if technique_converters: + resolved["technique_converters"] = technique_converters + + if dataset_names or max_dataset_size is not None or filters: + default_config = introspection_instance._default_dataset_config + if dataset_names: + default_config_class = type(default_config) + try: + resolved["dataset_config"] = default_config_class( + dataset_names=dataset_names, + max_dataset_size=max_dataset_size, + filters=filters or None, + ) + except TypeError as exc: + raise ValueError( + f"Scenario '{scenario_name}' does not support overriding dataset names through " + f"its {default_config_class.__name__} configuration: {exc}" + ) from exc + else: + if max_dataset_size is not None: + default_config.max_dataset_size = max_dataset_size + if filters: + default_config.update_filters(filters=filters) + resolved["dataset_config"] = default_config + + return resolved + + @classmethod + def resolve_techniques_and_converters( + cls, + *, + tokens: list[str], + technique_class: type[Any], + scenario_name: str, + ) -> tuple[list[Any], dict[str, list[Converter]]]: + """ + Resolve technique tokens and their additive converter modifiers. + + Returns: + tuple[list[Any], dict[str, list[Converter]]]: Selected enum members and + converters keyed by concrete technique name. + + Raises: + ValueError: If a technique or converter modifier is invalid. + """ + technique_enums: list[Any] = [] + technique_converters: dict[str, list[Converter]] = {} + for token in tokens: + base_name, _, remainder = token.partition(":") + modifiers = [modifier for modifier in remainder.split(":") if modifier] if remainder else [] + try: + technique_enum = technique_class(base_name) + except ValueError: + available_techniques = [technique.value for technique in technique_class] + raise ValueError( + f"Technique '{base_name}' not found for scenario '{scenario_name}'. " + f"Available: {', '.join(available_techniques)}" + ) from None + technique_enums.append(technique_enum) + + converters = cls._resolve_converter_modifiers(modifiers=modifiers, token=token) + for concrete in technique_class.expand({technique_enum}) if converters else (): + technique_converters.setdefault(concrete.value, []).extend(converters) + + return technique_enums, technique_converters + + @staticmethod + def _resolve_converter_modifiers(*, modifiers: list[str], token: str) -> list[Converter]: + """ + Resolve converter modifiers from one technique token. + + Returns: + list[Converter]: Registered converter instances in token order. + + Raises: + ValueError: If a modifier is malformed or references an unknown converter. + """ + if not modifiers: + return [] + + instances = ConverterRegistry.get_registry_singleton().instances + converters: list[Converter] = [] + for modifier in modifiers: + if not modifier.startswith(_CONVERTER_MODIFIER_PREFIX): + raise ValueError( + f"Unknown technique modifier '{modifier}' in '{token}'. " + f"Supported modifiers must use the '{_CONVERTER_MODIFIER_PREFIX}' prefix " + f"(e.g. '{_CONVERTER_MODIFIER_PREFIX}translation_spanish')." + ) + converter_name = modifier[len(_CONVERTER_MODIFIER_PREFIX) :] + converter = instances.get(converter_name) + if converter is None: + available = instances.get_names() + available_text = ", ".join(available) if available else "(none registered)" + raise ValueError( + f"Converter '{converter_name}' in '{token}' is not a registered converter " + f"instance. Available converters: {available_text}" + ) + converters.append(converter) + return converters diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index acc289ea80..a470f42fdf 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -17,16 +17,17 @@ from collections.abc import Sequence from dataclasses import dataclass from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any +from typing import Any from pyrit.backend.models.scenarios import ScenarioRunListResponse +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.common.utils import to_sha256 -from pyrit.memory import CentralMemory -from pyrit.memory.memory_interface import ScenarioProgressKeysetCursor +from pyrit.memory import AttackResultKeysetCursor, CentralMemory from pyrit.models import ( SCENARIO_RUN_PLAN_METADATA_KEY, AtomicAttackIdentifier, AttackOutcome, + AttackResult, ComponentIdentifier, ScenarioAttackResultDelta, ScenarioProgressHeader, @@ -43,26 +44,16 @@ AttackErrorSummary, AttackRetrySummary, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, ) -from pyrit.registry import ( - ConverterRegistry, - InitializerRegistry, - ScenarioRegistry, - TargetRegistry, -) +from pyrit.registry import InitializerRegistry, ScenarioRegistry from pyrit.scenario import Scenario -if TYPE_CHECKING: - from pyrit.converter import Converter - from pyrit.prompt_target import PromptTarget - logger = logging.getLogger(__name__) _DEFAULT_MAX_CONCURRENT_RUNS = 3 -_CONVERTER_MODIFIER_PREFIX = "converter." - @dataclass class _ActiveTask: @@ -82,6 +73,87 @@ class _ActiveRunSnapshot: active_group_ids: tuple[str, ...] = () +@dataclass(frozen=True, slots=True) +class _ResultUnitIdentity: + """Stable identity of one planned scenario execution unit.""" + + atomic_group_id: str + seed_group_id: str + + +@dataclass(frozen=True, slots=True) +class _ScenarioPlanLookup: + """Pre-indexed run-plan data used while mapping many attack results.""" + + groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] + groups_by_name: dict[str, tuple[ScenarioRunPlanAtomicGroup, ...]] + seed_ids_by_group_and_objective: dict[tuple[str, str], tuple[str, ...]] + planned_units: frozenset[_ResultUnitIdentity] + + @classmethod + def from_plan(cls, *, plan: ScenarioRunPlan | None) -> "_ScenarioPlanLookup": + """ + Build constant-time lookup tables for one run plan. + + Returns: + _ScenarioPlanLookup: Indexed plan data. + """ + if plan is None: + return cls( + groups_by_identity={}, + groups_by_name={}, + seed_ids_by_group_and_objective={}, + planned_units=frozenset(), + ) + + groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] = {} + grouped_by_name: dict[str, list[ScenarioRunPlanAtomicGroup]] = {} + seeds_by_id = {seed.id: seed for seed in plan.seed_groups} + seed_ids_by_group_and_objective: dict[tuple[str, str], tuple[str, ...]] = {} + planned_units: set[_ResultUnitIdentity] = set() + for group in plan.atomic_groups: + groups_by_identity[(group.atomic_attack_name, group.technique_eval_hash)] = group + grouped_by_name.setdefault(group.atomic_attack_name, []).append(group) + seed_ids_by_objective: dict[str, list[str]] = {} + for seed_id in group.seed_group_ids: + seed = seeds_by_id[seed_id] + seed_ids_by_objective.setdefault(seed.objective_sha256, []).append(seed_id) + seed_ids_by_group_and_objective.update( + { + (group.id, objective_sha256): tuple(seed_ids) + for objective_sha256, seed_ids in seed_ids_by_objective.items() + } + ) + planned_units.update( + _ResultUnitIdentity(atomic_group_id=group.id, seed_group_id=seed_group_id) + for seed_group_id in group.seed_group_ids + ) + + return cls( + groups_by_identity=groups_by_identity, + groups_by_name={name: tuple(groups) for name, groups in grouped_by_name.items()}, + seed_ids_by_group_and_objective=seed_ids_by_group_and_objective, + planned_units=frozenset(planned_units), + ) + + def resolve_group( + self, + *, + atomic_attack_name: str, + technique_eval_hash: str | None, + ) -> ScenarioRunPlanAtomicGroup | None: + """ + Resolve one planned group from persisted attribution. + + Returns: + ScenarioRunPlanAtomicGroup | None: The uniquely matching group. + """ + if technique_eval_hash is not None: + return self.groups_by_identity.get((atomic_attack_name, technique_eval_hash)) + matching_groups = self.groups_by_name.get(atomic_attack_name, ()) + return matching_groups[0] if len(matching_groups) == 1 else None + + class ScenarioRunService: """ Service for managing scenario run lifecycle. @@ -96,6 +168,7 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) - self._memory = CentralMemory.get_memory_instance() self._active_tasks: dict[str, _ActiveTask] = {} self._run_semaphore = asyncio.Semaphore(max_concurrent_runs) + self._configuration_resolver = ScenarioConfigurationResolver() async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary: """ @@ -126,11 +199,21 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu # Perform all initialization eagerly — errors propagate to caller try: - scenario_class = self._resolve_scenario_class(request=request) + scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name) await self._run_initializers_async(request=request) - objective_target = self._resolve_target(request=request) - init_kwargs = self._build_init_kwargs( - request=request, scenario_class=scenario_class, objective_target=objective_target + objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name) + init_kwargs = self._configuration_resolver.resolve_configuration( + scenario_name=request.scenario_name, + scenario_class=scenario_class, + objective_target=objective_target, + techniques=request.techniques, + dataset_names=request.dataset_names, + max_dataset_size=request.max_dataset_size, + dataset_filters=request.dataset_filters, + include_baseline=request.include_baseline, + max_concurrency=request.max_concurrency, + max_retries=request.max_retries, + memory_labels=request.labels, ) scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs) except Exception: @@ -196,18 +279,49 @@ def list_runs(self, *, limit: int = 100) -> ScenarioRunListResponse: Returns: ScenarioRunListResponse with runs. """ - # This is expensive, and we don't need all the data. At some point - # we may want to add a lightweight "list" query to the DB layer that only - results = self._memory.get_scenario_results(limit=limit) - items = [ - self._build_response_from_db( - scenario_result=sr, - active_error=self.snapshot_active_run(scenario_result_id=str(sr.id)).error, - ) - for sr in results - ] + results = self._memory.get_scenario_result_headers(limit=limit) + items = [self._build_list_response_from_header(scenario_result=result) for result in results] return ScenarioRunListResponse(items=items) + def _build_list_response_from_header(self, *, scenario_result: ScenarioResult) -> ScenarioRunListItem: + """ + Build a bounded run-history item without hydrating attack results. + + Returns: + ScenarioRunListItem: Lightweight run metadata. + """ + status = scenario_result.scenario_run_state + terminal = status in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ) + plan = self._load_run_plan(scenario_result=scenario_result) + total_attacks = sum(len(group.seed_group_ids) for group in plan.atomic_groups) if plan is not None else 0 + techniques_used = ( + list(dict.fromkeys(group.display_group for group in plan.atomic_groups)) if plan is not None else [] + ) + updated_at = ( + scenario_result.completion_time + if terminal and scenario_result.completion_time is not None + else scenario_result.creation_time + ) + return ScenarioRunListItem( + scenario_result_id=str(scenario_result.id), + scenario_name=scenario_result.scenario_name, + scenario_registry_name=plan.scenario_registry_name if plan else None, + scenario_version=scenario_result.scenario_version, + status=status, + created_at=scenario_result.creation_time, + updated_at=updated_at, + error=scenario_result.error_message, + error_type=scenario_result.error_type, + techniques_used=techniques_used, + total_attacks=total_attacks, + labels=scenario_result.labels, + completed_at=scenario_result.completion_time if terminal else None, + ) + async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: """ Cancel a running scenario. @@ -249,25 +363,6 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma return self.get_run(scenario_result_id=scenario_result_id) - def _resolve_scenario_class(self, *, request: RunScenarioRequest) -> type[Scenario]: - """ - Validate and resolve the scenario class from the registry. - - Args: - request: The run request containing the scenario name. - - Returns: - The scenario class. - - Raises: - ValueError: If the scenario name is not found in the registry. - """ - scenario_registry = ScenarioRegistry.get_registry_singleton() - try: - return scenario_registry.get_class(request.scenario_name) - except KeyError as e: - raise ValueError(str(e)) from None - async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: """ Validate and execute initializers specified in the request. @@ -292,294 +387,6 @@ async def _run_initializers_async(self, *, request: RunScenarioRequest) -> None: raise ValueError(f"Initializer not found: {e}") from None await instance.initialize_async() - def _resolve_target(self, *, request: RunScenarioRequest) -> "PromptTarget": - """ - Resolve the objective target from the target registry. - - Args: - request: The run request containing the target name. - - Returns: - The resolved PromptTarget instance. - - Raises: - ValueError: If the target is not found in the registry. - """ - 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(target_name) - if objective_target is None: - available_names = target_registry.instances.get_names() - if not available_names: - raise ValueError( - f"Target '{target_name}' not found. The target registry is empty. " - "Make sure to include an initializer that registers targets " - "(e.g., initializers: ['target'])." - ) - raise ValueError( - f"Target '{target_name}' not found in registry. Available targets: {', '.join(available_names)}" - ) - return objective_target - - def _build_init_kwargs( - self, *, request: RunScenarioRequest, scenario_class: type[Scenario], objective_target: Any - ) -> dict[str, Any]: - """ - Build the kwargs dict for scenario.initialize_async. - - Resolves techniques and dataset configuration from the request. - - Dataset configuration is built so that the scenario's default - ``DatasetAttackConfiguration`` *subclass* (e.g. ``EncodingDatasetConfiguration``) - is preserved when the caller overrides ``dataset_names`` or - ``max_dataset_size``. Subclasses commonly override - ``_build_attack_groups()`` to shape seeds into scenario-appropriate - ``AttackSeedGroup`` objects. - - Args: - request: The run request. - scenario_class: The resolved scenario class. - objective_target: The resolved target instance. - - Returns: - Dict of kwargs to pass to scenario.initialize_async. - - Raises: - ValueError: If a technique name is invalid for the scenario, or the - scenario class cannot be instantiated with no arguments when - introspection is required to resolve techniques or dataset - configuration. - """ - 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 resolved - - try: - introspection_instance = scenario_class() # type: ignore[ty:missing-argument] - except Exception as exc: - raise ValueError( - f"Cannot resolve runtime configuration for scenario '{scenario_name}': " - f"scenario class is not instantiable without arguments ({exc})." - ) from exc - - if techniques: - technique_class = introspection_instance._technique_class - technique_enums, technique_converters = cls._resolve_techniques_and_converters( - tokens=techniques, - technique_class=technique_class, - scenario_name=scenario_name, - ) - resolved["scenario_techniques"] = technique_enums - if technique_converters: - resolved["technique_converters"] = technique_converters - - if dataset_names or max_dataset_size is not None or filters: - default_config = introspection_instance._default_dataset_config - - if dataset_names: - # 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: - resolved["dataset_config"] = default_config_class( - dataset_names=dataset_names, - max_dataset_size=max_dataset_size, - filters=filters or None, - ) - except TypeError as exc: - raise ValueError( - f"Scenario '{scenario_name}' does not support overriding dataset names through " - f"its {default_config_class.__name__} configuration: {exc}" - ) from exc - else: - # 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 max_dataset_size is not None: - default_config.max_dataset_size = max_dataset_size - if filters: - default_config.update_filters(filters=filters) - resolved["dataset_config"] = default_config - - return resolved - - @classmethod - def _resolve_techniques_and_converters( - cls, - *, - tokens: list[str], - technique_class: type[Any], - scenario_name: str, - ) -> tuple[list[Any], dict[str, list["Converter"]]]: - """ - Resolve ``--techniques`` tokens into technique enums and per-technique converters. - - Each token has the form ``[:converter.[:converter....]]``. - The base ```` is resolved to a ``ScenarioTechnique`` enum member (which may - be an aggregate). Each ``converter.`` modifier is resolved to a registered - converter instance and appended (in token order) to every concrete technique that the - base technique expands to. - - Args: - tokens: The raw technique tokens from the request. - technique_class: The scenario's ``ScenarioTechnique`` subclass. - scenario_name: The scenario name, used for error messages. - - Returns: - A tuple of (technique enums to pass as ``scenario_techniques``, mapping from concrete - technique name to the list of converters to append for that technique). - - Raises: - ValueError: If a base technique name is unknown, a modifier is malformed, or a - converter name is not registered. - """ - technique_enums: list[Any] = [] - technique_converters: dict[str, list[Converter]] = {} - - for token in tokens: - base_name, _, remainder = token.partition(":") - modifiers = [m for m in remainder.split(":") if m] if remainder else [] - - try: - technique_enum = technique_class(base_name) - except ValueError: - available_techniques = [s.value for s in technique_class] - raise ValueError( - f"Technique '{base_name}' not found for scenario '{scenario_name}'. " - f"Available: {', '.join(available_techniques)}" - ) from None - technique_enums.append(technique_enum) - - converters = cls._resolve_converter_modifiers(modifiers=modifiers, token=token) - if not converters: - continue - - for concrete in technique_class.expand({technique_enum}): - technique_converters.setdefault(concrete.value, []).extend(converters) - - return technique_enums, technique_converters - - @staticmethod - def _resolve_converter_modifiers(*, modifiers: list[str], token: str) -> list["Converter"]: - """ - Resolve the converter modifiers of a single technique token to converter instances. - - Args: - modifiers: The modifier segments of the token (everything after the base technique). - token: The full original token, used for error messages. - - Returns: - The resolved converter instances in token order. - - Raises: - ValueError: If a modifier does not use the ``converter.`` prefix or names a - converter that is not registered. - """ - if not modifiers: - return [] - - instances = ConverterRegistry.get_registry_singleton().instances - converters: list[Converter] = [] - for modifier in modifiers: - if not modifier.startswith(_CONVERTER_MODIFIER_PREFIX): - raise ValueError( - f"Unknown technique modifier '{modifier}' in '{token}'. " - f"Supported modifiers must use the '{_CONVERTER_MODIFIER_PREFIX}' prefix " - f"(e.g. '{_CONVERTER_MODIFIER_PREFIX}translation_spanish')." - ) - converter_name = modifier[len(_CONVERTER_MODIFIER_PREFIX) :] - converter = instances.get(converter_name) - if converter is None: - available = instances.get_names() - available_text = ", ".join(available) if available else "(none registered)" - raise ValueError( - f"Converter '{converter_name}' in '{token}' is not a registered converter " - f"instance. Available converters: {available_text}" - ) - converters.append(converter) - return converters - async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_kwargs: dict[str, Any]) -> Scenario: """ Build and initialize the scenario via the registry. @@ -587,8 +394,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ Delegates the full create + set-parameters + initialize lifecycle to ``ScenarioRegistry.create_and_initialize_async`` so the registry owns scenario creation and initialization. The run-specific common parameters - (target, techniques, dataset config, concurrency) are resolved by - ``_build_init_kwargs`` and forwarded as ``init_kwargs``. + are resolved before this method and forwarded as ``init_kwargs``. Args: request: The run request (for scenario_name, scenario_params, and @@ -703,11 +509,13 @@ def _build_response_from_db( ScenarioRunState.CANCELLED, ) plan = self._load_run_plan(scenario_result=scenario_result) + plan_lookup = _ScenarioPlanLookup.from_plan(plan=plan) # Build result fields from DB (always computed so in-progress runs show progress) total_attacks, completed_attacks, objective_achieved_rate = self._calculate_progress_counts( scenario_result=scenario_result, plan=plan, + plan_lookup=plan_lookup, ) techniques_used = ( list(dict.fromkeys(group.display_group for group in plan.atomic_groups)) @@ -720,15 +528,15 @@ def _build_response_from_db( failed_attacks: list[AttackErrorSummary] = [] attack_retries: list[AttackRetrySummary] = [] total_retries = 0 - attempts_by_unit: dict[tuple[str, str], int] = {} + attempts_by_unit: dict[_ResultUnitIdentity, int] = {} for atomic_attack_name, results in scenario_result.attack_results.items(): for attack_result in results: - unit_key = self._result_unit_key( + unit_identity = self._resolve_result_unit_identity( atomic_attack_name=atomic_attack_name, attack_result=attack_result, - plan=plan, + plan_lookup=plan_lookup, ) - attempts_by_unit[unit_key] = attempts_by_unit.get(unit_key, 0) + 1 + attempts_by_unit[unit_identity] = attempts_by_unit.get(unit_identity, 0) + 1 retries = getattr(attack_result, "total_retries", 0) if isinstance(retries, int): total_retries += retries @@ -813,59 +621,57 @@ def _load_run_plan(*, scenario_result: ScenarioResult) -> ScenarioRunPlan | None return ScenarioRunPlan.model_validate(raw_plan) if raw_plan is not None else None @staticmethod - def _result_unit_key( + def _resolve_result_unit_identity( *, atomic_attack_name: str, - attack_result: Any, - plan: ScenarioRunPlan | None, - ) -> tuple[str, str]: + attack_result: AttackResult, + plan_lookup: _ScenarioPlanLookup, + ) -> _ResultUnitIdentity: """ - Resolve one attack attempt to its stable planned-unit key. + Resolve one attack attempt to its stable planned-unit identity. Returns: - tuple[str, str]: The atomic-group and seed-group IDs. + _ResultUnitIdentity: The atomic-group and seed-group IDs. """ - atomic_identifier = getattr(attack_result, "atomic_attack_identifier", None) + atomic_identifier = attack_result.atomic_attack_identifier 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) + objective = str(attack_result.objective) + attribution_data = attack_result.attribution_data attributed_seed_group_id = attribution_data.get("seed_group_id") if isinstance(attribution_data, dict) else None seed_group_id = str(attributed_seed_group_id) if attributed_seed_group_id else "" if not seed_group_id and typed_identifier is not None and typed_identifier.seed_identifiers: seed_group_id = typed_identifier.logical_seed_group_id + atomic_group_id = atomic_attack_name - 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] + eval_hash = attribution_data.get("parent_eval_hash") if isinstance(attribution_data, dict) else None + planned_group = plan_lookup.resolve_group( + atomic_attack_name=atomic_attack_name, + technique_eval_hash=str(eval_hash) if eval_hash is not None else None, + ) + if planned_group is not None: + atomic_group_id = planned_group.id + if not seed_group_id: + objective_sha256 = to_sha256(objective) + matching_seed_ids = plan_lookup.seed_ids_by_group_and_objective.get( + (planned_group.id, objective_sha256), + (), + ) + if len(matching_seed_ids) == 1: + seed_group_id = matching_seed_ids[0] if not seed_group_id: seed_group_id = config_hash({"objective": objective}) - return atomic_group_id, seed_group_id + return _ResultUnitIdentity(atomic_group_id=atomic_group_id, seed_group_id=seed_group_id) def _calculate_progress_counts( self, *, scenario_result: ScenarioResult, plan: ScenarioRunPlan | None, + plan_lookup: _ScenarioPlanLookup, ) -> tuple[int, int, int]: """ Calculate planned-unit totals without inflating retries or error attempts. @@ -873,43 +679,33 @@ def _calculate_progress_counts( 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] = {} + latest_result_by_unit: dict[_ResultUnitIdentity, AttackResult] = {} for atomic_attack_name, results in scenario_result.attack_results.items(): for attack_result in results: - unit_key = self._result_unit_key( + unit_identity = self._resolve_result_unit_identity( atomic_attack_name=atomic_attack_name, attack_result=attack_result, - plan=plan, + plan_lookup=plan_lookup, ) - attempted_units.add(unit_key) - if attack_result.outcome == AttackOutcome.ERROR: - continue - previous = latest_non_error_by_unit.get(unit_key) + previous = latest_result_by_unit.get(unit_identity) if previous is None or self._result_order_key(attack_result) > self._result_order_key(previous): - latest_non_error_by_unit[unit_key] = attack_result + latest_result_by_unit[unit_identity] = 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 - ) + planned_units = plan_lookup.planned_units if plan is not None else frozenset(latest_result_by_unit) total = len(planned_units) - completed_results = [ - result for unit_key, result in latest_non_error_by_unit.items() if unit_key in planned_units - ] + completed_results = [result for unit, result in latest_result_by_unit.items() if unit in planned_units] completed = len(completed_results) succeeded = sum(result.outcome == AttackOutcome.SUCCESS for result in completed_results) rate = int((succeeded / completed) * 100) if completed else 0 return total, completed, rate @staticmethod - def _result_order_key(attack_result: Any) -> tuple[datetime, str]: + def _result_order_key(attack_result: AttackResult) -> tuple[datetime, str]: """Return a deterministic chronological key for one hydrated result attempt.""" - timestamp = getattr(attack_result, "timestamp", None) + timestamp = attack_result.timestamp if not isinstance(timestamp, datetime): timestamp = datetime.min.replace(tzinfo=timezone.utc) - return timestamp, str(getattr(attack_result, "attack_result_id", "")) + return timestamp, str(attack_result.attack_result_id) def get_run_progress( self, @@ -952,12 +748,14 @@ def get_run_progress_from_storage( limit=limit, ) plan = self._load_run_plan(scenario_result=header_result) + plan_lookup = _ScenarioPlanLookup.from_plan(plan=plan) plan_complete = plan is not None response_plan = plan if since is None else None if plan is None and since is None: response_plan = self._synthesize_legacy_plan(deltas=deltas) - results = [self._map_progress_delta(delta=delta, plan=plan or response_plan) for delta in deltas] + response_plan_lookup = plan_lookup if plan is not None else _ScenarioPlanLookup.from_plan(plan=response_plan) + results = [self._map_progress_delta(delta=delta, plan_lookup=response_plan_lookup) for delta in deltas] next_cursor = ( self._encode_progress_cursor(scenario_result_id=scenario_result_id, delta=deltas[-1]) if deltas else since ) @@ -989,7 +787,7 @@ def get_run_progress_from_storage( def _map_progress_delta( *, delta: ScenarioAttackResultDelta, - plan: ScenarioRunPlan | None, + plan_lookup: _ScenarioPlanLookup, ) -> ScenarioProgressResult: """ Map a lightweight memory row to its REST progress representation. @@ -1002,13 +800,12 @@ def _map_progress_delta( 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 + planned_group = plan_lookup.resolve_group( + atomic_attack_name=atomic_attack_name, + technique_eval_hash=str(eval_hash) if eval_hash is not None else None, + ) + if planned_group is not None: + atomic_group_id = planned_group.id attributed_seed_group_id = delta.attribution_data.get("seed_group_id") seed_group_id = str(attributed_seed_group_id) if attributed_seed_group_id else "" if ( @@ -1017,13 +814,11 @@ def _map_progress_delta( 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 not seed_group_id and delta.objective_sha256: + matching_seed_ids = plan_lookup.seed_ids_by_group_and_objective.get( + (atomic_group_id, delta.objective_sha256), + (), + ) if len(matching_seed_ids) == 1: seed_group_id = matching_seed_ids[0] if not seed_group_id: @@ -1052,8 +847,12 @@ def _synthesize_legacy_plan(*, deltas: list[ScenarioAttackResultDelta]) -> Scena """ seeds: dict[str, ScenarioRunPlanSeedGroup] = {} groups: dict[str, ScenarioRunPlanAtomicGroup] = {} + empty_plan_lookup = _ScenarioPlanLookup.from_plan(plan=None) for delta in deltas: - mapped = ScenarioRunService._map_progress_delta(delta=delta, plan=None) + mapped = ScenarioRunService._map_progress_delta( + delta=delta, + plan_lookup=empty_plan_lookup, + ) seeds.setdefault( mapped.seed_group_id, ScenarioRunPlanSeedGroup( @@ -1091,7 +890,7 @@ def _decode_progress_cursor( *, since: str | None, scenario_result_id: str, - ) -> ScenarioProgressKeysetCursor | None: + ) -> AttackResultKeysetCursor | None: if since is None: return None try: @@ -1110,7 +909,7 @@ def _decode_progress_cursor( 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) + return AttackResultKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) def get_run_results(self, *, scenario_result_id: str) -> ScenarioResult | None: """ diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index e6fcbfccd6..dca9623d07 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -13,12 +13,11 @@ 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.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.models.catalog.scenario import ( RegisteredScenario, - ScenarioDefaultRunSizeEstimate, + ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, - ScenarioRunSizeEstimateStatus, ) from pyrit.registry import ScenarioMetadata, ScenarioRegistry @@ -28,14 +27,14 @@ _ESTIMATE_INFLIGHT_SIZE = 256 _UNAVAILABLE_CACHE_TTL_SECONDS = 30.0 _EstimateCacheKey = tuple[str, int] -_EstimateCacheValue = tuple[ScenarioDefaultRunSizeEstimate, float | None] -_EstimateTask = asyncio.Task[ScenarioDefaultRunSizeEstimate] +_EstimateCacheValue = tuple[ScenarioRunSizeEstimate, float | None] +_EstimateTask = asyncio.Task[ScenarioRunSizeEstimate] def _metadata_to_registered_scenario( *, metadata: ScenarioMetadata, - default_run_size: ScenarioDefaultRunSizeEstimate | None = None, + default_run_size: ScenarioRunSizeEstimate | None = None, ) -> RegisteredScenario: """ Convert a ScenarioMetadata dataclass to a ScenarioSummary Pydantic model. @@ -47,7 +46,7 @@ def _metadata_to_registered_scenario( Returns: RegisteredScenario: Public catalog projection. """ - estimate = default_run_size or ScenarioDefaultRunSizeEstimate.unavailable() + estimate = default_run_size or ScenarioRunSizeEstimate.unavailable() return RegisteredScenario( scenario_name=metadata.registry_name, scenario_type=metadata.class_name, @@ -139,7 +138,7 @@ async def estimate_scenario_run_size_async( *, scenario_name: str, request: ScenarioRunSizeEstimateRequest, - ) -> ScenarioDefaultRunSizeEstimate | None: + ) -> ScenarioRunSizeEstimate | None: """ Estimate one configured scenario without creating a run. @@ -148,7 +147,7 @@ async def estimate_scenario_run_size_async( request: Request-specific techniques, datasets, baseline, and parameters. Returns: - ScenarioDefaultRunSizeEstimate | None: Estimate, or ``None`` when the scenario is unknown. + ScenarioRunSizeEstimate | None: Estimate, or ``None`` when the scenario is unknown. """ metadata = self._registry.get_registered_class_metadata(scenario_name) if metadata is None: @@ -164,9 +163,7 @@ async def estimate_scenario_run_size_async( request=request, ) - async def _get_default_run_size_estimate_async( - self, *, metadata: ScenarioMetadata - ) -> ScenarioDefaultRunSizeEstimate: + async def _get_default_run_size_estimate_async(self, *, metadata: ScenarioMetadata) -> ScenarioRunSizeEstimate: """Return a cached, cancellation-safe scenario-owned estimate.""" cache_key = (metadata.registry_name, metadata.scenario_version) cache = getattr(self, "_estimate_cache", None) @@ -218,7 +215,7 @@ def clear_estimate_task(completed_task: _EstimateTask) -> None: if wait_for_capacity is not None: await asyncio.shield(wait_for_capacity) - def _read_estimate_cache(self, *, cache_key: _EstimateCacheKey) -> ScenarioDefaultRunSizeEstimate | None: + def _read_estimate_cache(self, *, cache_key: _EstimateCacheKey) -> ScenarioRunSizeEstimate | None: """Return a live cached estimate and discard expired unavailable entries.""" cache = self._estimate_cache cached = cache.get(cache_key) @@ -236,12 +233,12 @@ async def _compute_default_run_size_estimate_async( *, scenario_name: str, cache_key: _EstimateCacheKey, - ) -> ScenarioDefaultRunSizeEstimate: + ) -> ScenarioRunSizeEstimate: """ Construct and estimate one scenario on the owning event loop. Returns: - ScenarioDefaultRunSizeEstimate: Scenario-owned estimate. + ScenarioRunSizeEstimate: Scenario-owned estimate. """ semaphore = getattr(self, "_estimate_semaphore", None) if semaphore is None: @@ -253,15 +250,11 @@ async def _compute_default_run_size_estimate_async( 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( + estimate = ScenarioRunSizeEstimate.unavailable( note=f"The scenario could not resolve its default inputs for estimation ({type(exc).__name__})." ) - expires_at = ( - monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS - if estimate.status is ScenarioRunSizeEstimateStatus.Unavailable - else None - ) + expires_at = monotonic() + _UNAVAILABLE_CACHE_TTL_SECONDS if estimate.estimated_attack_count is None else None cache = self._estimate_cache cache[cache_key] = (estimate, expires_at) cache.move_to_end(cache_key) @@ -280,18 +273,17 @@ async def _estimate_configured_run_size_async( *, scenario_name: str, request: ScenarioRunSizeEstimateRequest, - ) -> ScenarioDefaultRunSizeEstimate: + ) -> ScenarioRunSizeEstimate: """ Resolve and estimate one request on the owning event loop. Returns: - ScenarioDefaultRunSizeEstimate: Request-specific scenario estimate. + ScenarioRunSizeEstimate: 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( + resolver = ScenarioConfigurationResolver() + objective_target = resolver.resolve_target(target_name=request.target_name) if request.target_name else None + estimate_kwargs = resolver.resolve_configuration( scenario_name=scenario_name, scenario_class=scenario_class, objective_target=objective_target, diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 6e2ab40b9f..4e63a313ad 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -21,6 +21,7 @@ from pyrit.models.catalog import ( RegisteredInitializer, RegisteredScenario, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -461,7 +462,7 @@ def print_attacks_table(*, payload: AttacksTablePayload) -> None: # --------------------------------------------------------------------------- -def print_scenario_runs_list(*, runs: list[ScenarioRunSummary]) -> None: +def print_scenario_runs_list(*, runs: list[ScenarioRunListItem]) -> None: """ Print a list of scenario run summaries. @@ -478,7 +479,7 @@ def print_scenario_runs_list(*, runs: list[ScenarioRunSummary]) -> None: created = run.created_at.isoformat() if run.created_at else "?" print( f" {idx}) [{run.status.value}] {run.scenario_name} (id: {run.scenario_result_id}) — " - f"{run.total_attacks} attacks, {run.objective_achieved_rate}% success — {created}" + f"{run.total_attacks} planned attacks — {created}" ) print("=" * 80) print(f"\nTotal runs: {len(runs)}") diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 4f43c57e1c..81de602b42 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -21,6 +21,7 @@ RegisteredInitializer, RegisteredScenario, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -310,17 +311,17 @@ async def cancel_scenario_run_async(self, *, scenario_result_id: str) -> Scenari self._raise_for_status(resp) return ScenarioRunSummary.model_validate(resp.json()) - async def list_scenario_runs_async(self, *, limit: int = 100) -> list[ScenarioRunSummary]: + async def list_scenario_runs_async(self, *, limit: int = 100) -> list[ScenarioRunListItem]: """ List tracked scenario runs. Returns: - list[ScenarioRunSummary]: All tracked scenario runs. + list[ScenarioRunListItem]: All tracked scenario runs. """ - from pyrit.models.catalog import ScenarioRunSummary + from pyrit.models.catalog import ScenarioRunListItem payload = await self._get_json_async(path="/api/scenarios/runs", params={"limit": limit}) - return [ScenarioRunSummary.model_validate(item) for item in payload.get("items", [])] + return [ScenarioRunListItem.model_validate(item) for item in payload.get("items", [])] # ------------------------------------------------------------------ # Lifecycle diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index d859e7a12d..d2a27692f5 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -322,17 +322,6 @@ async def initialize_context_async( 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, diff --git a/pyrit/executor/attack/core/attack_executor.py b/pyrit/executor/attack/core/attack_executor.py index 846aa4a4db..6bfbab03c5 100644 --- a/pyrit/executor/attack/core/attack_executor.py +++ b/pyrit/executor/attack/core/attack_executor.py @@ -232,6 +232,13 @@ async def execute_attack_from_seed_groups_async( ) if attribution is not None and attributions is not None: raise ValueError("Provide attribution or attributions, not both") + effective_attributions = ( + list(attributions) + if attributions is not None + else [attribution] * len(seed_groups) + if attribution is not None + else None + ) params_type = attack.params_type @@ -274,13 +281,14 @@ async def build_params_async(i: int, sg: AttackSeedGroup) -> AttackParameters: raise build_failures[0][2] successful_attributions = ( - [attributions[index] for index in successful_input_indices] if attributions is not None else None + [effective_attributions[index] for index in successful_input_indices] + if effective_attributions is not None + else None ) execution_result = await self._execute_with_params_list_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, attributions=successful_attributions, input_indices=successful_input_indices, ) @@ -355,7 +363,7 @@ async def execute_attack_async( attack=attack, params_list=params_list, return_partial_on_failure=return_partial_on_failure, - attribution=attribution, + attributions=[attribution] * len(params_list) if attribution is not None else None, ) async def _execute_with_params_list_async( @@ -364,7 +372,6 @@ async def _execute_with_params_list_async( attack: AttackStrategy[AttackStrategyContextT, AttackStrategyResultT], params_list: Sequence[AttackParameters], return_partial_on_failure: bool = False, - attribution: AttackResultAttribution | None = None, attributions: Sequence[AttackResultAttribution] | None = None, input_indices: Sequence[int] | None = None, ) -> AttackExecutorResult[AttackStrategyResultT]: @@ -378,9 +385,6 @@ async def _execute_with_params_list_async( attack: The attack strategy to execute. params_list: List of AttackParameters, one per execution. return_partial_on_failure: If True, returns partial results on failure. - attribution: Optional ``AttackResultAttribution`` stamped onto every - per-task ``AttackContext`` so the persistence path can record - orchestrator linkage. attributions: Optional per-task attribution matching ``params_list``. input_indices: Original input positions for ``params_list``. Defaults to sequential positions when parameters were constructed directly. @@ -389,21 +393,18 @@ async def _execute_with_params_list_async( 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. + ValueError: If per-task attribution or input-index lengths do not match. """ semaphore = self._get_semaphore() if attributions is not None and len(attributions) != len(params_list): raise ValueError( f"attributions length ({len(attributions)}) must match params_list length ({len(params_list)})" ) - 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) - task_attribution = attributions[index] if attributions is not None else attribution + task_attribution = attributions[index] if attributions is not None else None if task_attribution is not None: context._attribution = task_attribution return await attack.execute_with_context_async(context=context) diff --git a/pyrit/memory/__init__.py b/pyrit/memory/__init__.py index 2c78a4ba2d..ce16b59116 100644 --- a/pyrit/memory/__init__.py +++ b/pyrit/memory/__init__.py @@ -10,7 +10,7 @@ from pyrit.memory.azure_sql_memory import AzureSQLMemory from pyrit.memory.central_memory import CentralMemory from pyrit.memory.memory_embedding import MemoryEmbedding -from pyrit.memory.memory_interface import AttackResultsKeysetCursor, MemoryInterface +from pyrit.memory.memory_interface import AttackResultKeysetCursor, MemoryInterface from pyrit.memory.memory_models import AttackResultEntry, EmbeddingDataEntry, PromptMemoryEntry, SeedEntry from pyrit.memory.sqlite_memory import SQLiteMemory from pyrit.memory.storage import ( @@ -35,7 +35,7 @@ __all__ = [ "AllowedCategories", "AttackResultEntry", - "AttackResultsKeysetCursor", + "AttackResultKeysetCursor", "AudioPathDataTypeSerializer", "AzureBlobStorageIO", "AzureSQLMemory", diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 3402f26f5b..7a0618f17e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -96,7 +96,7 @@ IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) -class AttackResultsKeysetCursor(NamedTuple): +class AttackResultKeysetCursor(NamedTuple): """ Keyset (seek) anchor identifying the last attack result on a page. @@ -112,12 +112,12 @@ class AttackResultsKeysetCursor(NamedTuple): attack_result_id: str @classmethod - def from_attack_result(cls, result: AttackResult) -> "AttackResultsKeysetCursor": + def from_attack_result(cls, result: AttackResult) -> "AttackResultKeysetCursor": """ Build the keyset anchor for ``result`` (typically the last row of a page). Returns: - AttackResultsKeysetCursor: Anchor capturing the result's recency sort key. + AttackResultKeysetCursor: Anchor capturing the result's recency sort key. """ return cls( timestamp=result.timestamp, @@ -125,13 +125,6 @@ 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: """ @@ -168,7 +161,7 @@ class _AttackResultQuery: min_turns: int | None = None max_turns: int | None = None limit: int | None = None - after: AttackResultsKeysetCursor | None = None + after: AttackResultKeysetCursor | None = None def __post_init__(self) -> None: """Snapshot mutable sequence and mapping inputs.""" @@ -424,7 +417,7 @@ def _attack_results_recency_order_by(self) -> list[Any]: """ return [AttackResultEntry.timestamp.desc(), AttackResultEntry.id.desc()] - def _attack_results_keyset_seek_condition(self, *, after: AttackResultsKeysetCursor) -> Any: + def _attack_results_keyset_seek_condition(self, *, after: AttackResultKeysetCursor) -> Any: """ Build the keyset seek predicate selecting rows strictly after ``after``. @@ -3114,7 +3107,7 @@ def get_attack_results( min_turns: int | None = None, max_turns: int | None = None, limit: int | None = None, - after: AttackResultsKeysetCursor | None = None, + after: AttackResultKeysetCursor | None = None, ) -> Sequence[AttackResult]: """ Retrieve a list of AttackResult objects based on the specified filters. @@ -3180,7 +3173,7 @@ def get_attack_results( return, ordered by recency. When either ``limit`` or ``after`` is provided, deduplication and pagination happen in the database (via ``ROW_NUMBER()``) instead of loading every row into memory. Defaults to None (return all). - after (AttackResultsKeysetCursor | None, optional): Keyset (seek) anchor from a + after (AttackResultKeysetCursor | None, optional): Keyset (seek) anchor from a previous page. When provided, only results ordered strictly after the anchor under the recency sort are returned, giving insert/delete-stable pagination without a drifting numeric offset. Defaults to None (start at the first page). @@ -3465,7 +3458,7 @@ def _query_paginated_attack_results( min_turns: int | None, max_turns: int | None, limit: int | None, - after: AttackResultsKeysetCursor | None, + after: AttackResultKeysetCursor | None, ) -> list[AttackResult]: """ Deduplicate in SQL (filter-aware) and return one recency-ordered page of results. @@ -3486,7 +3479,7 @@ def _query_paginated_attack_results( min_turns (int | None): Inclusive lower bound on ``executed_turns`` for winners. max_turns (int | None): Inclusive upper bound on ``executed_turns`` for winners. limit (int | None): Maximum number of results to return. - after (AttackResultsKeysetCursor | None): Keyset anchor; only rows ordered strictly + after (AttackResultKeysetCursor | None): Keyset anchor; only rows ordered strictly after it are returned. ``None`` starts at the first page. Returns: @@ -3700,11 +3693,30 @@ def get_scenario_result_header(self, *, scenario_result_id: str) -> ScenarioResu entry = session.query(ScenarioResultEntry).filter_by(id=scenario_result_id).first() return entry.get_scenario_result() if entry is not None else None + def get_scenario_result_headers(self, *, limit: int = 100) -> Sequence[ScenarioResult]: + """ + Return recent ScenarioResult headers without hydrating linked attack results. + + Returns: + Sequence[ScenarioResult]: Recent scenario metadata ordered newest first. + + Raises: + ValueError: If limit is outside the bounded run-history range. + """ + if limit < 1 or limit > 100: + raise ValueError("Scenario run history limit must be between 1 and 100.") + entries = self._query_scenario_result_entries( + scenario_result_ids=None, + conditions=[], + limit=limit, + ) + return [entry.get_scenario_result() for entry in entries] + def get_scenario_attack_result_deltas( self, *, scenario_result_id: str, - cursor: ScenarioProgressKeysetCursor | None = None, + cursor: AttackResultKeysetCursor | None = None, limit: int = 100, ) -> tuple[list[ScenarioAttackResultDelta], bool]: """ diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 23470de50c..a2bcdea964 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -20,12 +20,10 @@ from pyrit.models.catalog import ( ScenarioDatasetSizeCap, ScenarioDatasetSummary, - ScenarioDefaultRunSizeEstimate, + ScenarioRunListItem, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, ) from pyrit.models.conversation_stats import ConversationStats from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation @@ -56,8 +54,8 @@ TargetIdentifier, class_name_to_snake_case, compute_eval_hash, + compute_seed_group_hash, config_hash, - logical_seed_group_fingerprint, snake_case_to_class_name, validate_registry_name, ) @@ -193,7 +191,7 @@ "IdentifierFilter", "IdentifierType", "JSONValue", - "logical_seed_group_fingerprint", + "compute_seed_group_hash", "COMMON_JSON_SCHEMAS", "JsonResponseConfig", "get_common_json_schema", @@ -230,13 +228,11 @@ "ScenarioIdentifier", "ScenarioDatasetSizeCap", "ScenarioDatasetSummary", - "ScenarioDefaultRunSizeEstimate", "ScenarioRunSizeEstimate", "ScenarioResult", "ScenarioRunSizeComponent", "ScenarioRunSizeEstimateRequest", - "ScenarioRunSizeEstimateStatus", - "ScenarioRunSizeFactor", + "ScenarioRunListItem", "ScenarioRunState", "SCENARIO_RUN_PLAN_METADATA_KEY", "SCENARIO_RUN_PLAN_VERSION", diff --git a/pyrit/models/catalog/__init__.py b/pyrit/models/catalog/__init__.py index 692e3a54d5..2960916be5 100644 --- a/pyrit/models/catalog/__init__.py +++ b/pyrit/models/catalog/__init__.py @@ -23,12 +23,10 @@ RunScenarioRequest, ScenarioDatasetSizeCap, ScenarioDatasetSummary, - ScenarioDefaultRunSizeEstimate, + ScenarioRunListItem, ScenarioRunSizeComponent, ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, ScenarioRunSummary, ) from pyrit.models.catalog.target import ( @@ -43,12 +41,10 @@ "RunScenarioRequest", "ScenarioDatasetSizeCap", "ScenarioDatasetSummary", - "ScenarioDefaultRunSizeEstimate", "ScenarioRunSizeEstimate", "ScenarioRunSizeComponent", "ScenarioRunSizeEstimateRequest", - "ScenarioRunSizeEstimateStatus", - "ScenarioRunSizeFactor", + "ScenarioRunListItem", "ScenarioRunSummary", "TargetInstance", ] diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 549dd3eb37..1f62081dff 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -14,8 +14,6 @@ """ from datetime import datetime -from enum import Enum -from math import prod from typing import Any, Literal from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator @@ -59,49 +57,14 @@ def _validate_dataset_filter_mapping( 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.""" @@ -127,7 +90,7 @@ class ScenarioDatasetSummary(BaseModel): selection_note: str | None = None -class ScenarioDefaultRunSizeEstimate(BaseModel): +class ScenarioRunSizeEstimate(BaseModel): """ Structured estimate of default planned scenario execution units. @@ -135,87 +98,39 @@ class ScenarioDefaultRunSizeEstimate(BaseModel): 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"), - ) + estimated_attack_count: int | None = Field(default=None, ge=0) components: list[ScenarioRunSizeComponent] = Field(default_factory=list) datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) - note: str | None = 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 + note: str | None = None @model_validator(mode="after") - def validate_total(self) -> "ScenarioDefaultRunSizeEstimate": + def validate_estimated_attack_count(self) -> "ScenarioRunSizeEstimate": """ - Ensure exact estimates expose and explain their complete total. + Ensure available estimates expose a complete additive total. Returns: - ScenarioDefaultRunSizeEstimate: The validated estimate. + ScenarioRunSizeEstimate: The validated estimate. Raises: - ValueError: If an exact estimate omits or misstates its total. + ValueError: If an available estimate 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") + if self.estimated_attack_count is not None: component_total = sum(component.count for component in self.components) - if component_total != self.total_attack_count: + if component_total != self.estimated_attack_count: raise ValueError( - f"Exact default-run estimate components total {component_total}, not {self.total_attack_count}" + f"Default-run estimate components total {component_total}, not {self.estimated_attack_count}" ) return self @classmethod - def unavailable( - cls, *, note: str = "Default-run size estimate is unavailable." - ) -> "ScenarioDefaultRunSizeEstimate": + def unavailable(cls, *, note: str = "Default-run size estimate is unavailable.") -> "ScenarioRunSizeEstimate": """ Build an unavailable estimate without presenting a guessed total. Returns: - ScenarioDefaultRunSizeEstimate: An unavailable estimate. + ScenarioRunSizeEstimate: An unavailable estimate. """ - return cls(status=ScenarioRunSizeEstimateStatus.Unavailable, note=note) - - -# Backward-compatible catalog name from the initial run-size DTO. -ScenarioRunSizeEstimate = ScenarioDefaultRunSizeEstimate + return cls(note=note) class RegisteredScenario(BaseModel): @@ -257,8 +172,8 @@ class RegisteredScenario(BaseModel): supported_parameters: list[Parameter] = Field( default_factory=list, description="Scenario-declared custom parameters" ) - default_run_size: ScenarioDefaultRunSizeEstimate = Field( - default_factory=ScenarioDefaultRunSizeEstimate.unavailable, + default_run_size: ScenarioRunSizeEstimate = Field( + default_factory=ScenarioRunSizeEstimate.unavailable, description="Scenario-owned structured estimate of the default planned execution units", ) @@ -403,3 +318,21 @@ class ScenarioRunSummary(BaseModel): ) labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run") completed_at: datetime | None = Field(None, description="When the scenario finished") + + +class ScenarioRunListItem(BaseModel): + """Lightweight scenario run metadata returned by the history endpoint.""" + + scenario_result_id: str = Field(..., description="UUID of the ScenarioResult in memory") + scenario_name: str = Field(..., description="Registry key of the scenario being run") + scenario_registry_name: str | None = Field(None, description="Requested scenario registry key when available") + scenario_version: int = Field(0, ge=0, description="Version of the scenario") + status: ScenarioRunState = Field(..., description="Current run status") + created_at: datetime = Field(..., description="When the run was created") + updated_at: datetime = Field(..., description="When the run status last changed") + error: str | None = Field(None, description="Persisted run-level error message") + error_type: str | None = Field(None, description="Persisted run-level exception class") + techniques_used: list[str] = Field(default_factory=list, description="Planned technique display groups") + total_attacks: int = Field(0, ge=0, description="Number of planned execution units") + labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run") + completed_at: datetime | None = Field(None, description="When the scenario finished") diff --git a/pyrit/models/identifiers/__init__.py b/pyrit/models/identifiers/__init__.py index 2f671a9f13..d0530b1e34 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, logical_seed_group_fingerprint +from pyrit.models.identifiers.seed_identifier import SeedIdentifier, compute_seed_group_hash from pyrit.models.identifiers.target_identifier import TargetIdentifier __all__ = [ @@ -70,7 +70,7 @@ "ScorerIdentifier", "ScenarioIdentifier", "SeedIdentifier", - "logical_seed_group_fingerprint", + "compute_seed_group_hash", "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 6ccdf93f0a..79ba50b186 100644 --- a/pyrit/models/identifiers/atomic_attack_identifier.py +++ b/pyrit/models/identifiers/atomic_attack_identifier.py @@ -22,7 +22,7 @@ from pyrit.models.identifiers.attack_technique_identifier import AttackTechniqueIdentifier from pyrit.models.identifiers.component_identifier import ComponentIdentifier from pyrit.models.identifiers.evaluation_markers import Evaluate -from pyrit.models.identifiers.seed_identifier import SeedIdentifier, logical_seed_group_fingerprint +from pyrit.models.identifiers.seed_identifier import SeedIdentifier, compute_seed_group_hash if TYPE_CHECKING: from pyrit.models.seeds.seed_group import SeedGroup @@ -113,4 +113,4 @@ def build( @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) + return compute_seed_group_hash(self.seed_identifiers) diff --git a/pyrit/models/identifiers/seed_identifier.py b/pyrit/models/identifiers/seed_identifier.py index aa406559f2..acb22bedd1 100644 --- a/pyrit/models/identifiers/seed_identifier.py +++ b/pyrit/models/identifiers/seed_identifier.py @@ -62,8 +62,8 @@ def from_seed(cls, seed: Seed) -> SeedIdentifier: ) -def logical_seed_group_fingerprint(seed_identifiers: Sequence[SeedIdentifier]) -> str: - """Return the deterministic fingerprint of ordered canonical seed identifiers.""" +def compute_seed_group_hash(seed_identifiers: Sequence[SeedIdentifier]) -> str: + """Return the deterministic hash of ordered canonical seed identifiers.""" return config_hash( { "seed_identifiers": [ diff --git a/pyrit/models/seeds/attack_seed_group.py b/pyrit/models/seeds/attack_seed_group.py index 117b65165a..191b9aa7f9 100644 --- a/pyrit/models/seeds/attack_seed_group.py +++ b/pyrit/models/seeds/attack_seed_group.py @@ -12,7 +12,7 @@ import copy from typing import TYPE_CHECKING -from pyrit.models.identifiers import SeedIdentifier, logical_seed_group_fingerprint +from pyrit.models.identifiers import SeedIdentifier, compute_seed_group_hash from pyrit.models.seeds.seed_group import SeedGroup from pyrit.models.seeds.seed_objective import SeedObjective from pyrit.models.seeds.seed_prompt import SeedPrompt @@ -96,7 +96,7 @@ def logical_id(self) -> str: 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]) + return compute_seed_group_hash([SeedIdentifier.from_seed(seed) for seed in self.seeds]) def is_compatible_with_technique(self, *, technique: AttackTechniqueSeedGroup) -> bool: """ diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index 7e8af09f46..bc2c93c241 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal -from pyrit.models import ScenarioDefaultRunSizeEstimate, class_name_to_snake_case +from pyrit.models import ScenarioRunSizeEstimate, class_name_to_snake_case from pyrit.models.identifiers.scenario_identifier import ScenarioIdentifier from pyrit.registry.registry import ParamBagRegistry from pyrit.registry.registry_metadata import RegistryMetadata @@ -201,7 +201,7 @@ async def create_and_estimate_async( scenario_params: dict[str, Any] | None = None, target_is_configured: bool = False, **estimate_kwargs: Any, - ) -> ScenarioDefaultRunSizeEstimate: + ) -> ScenarioRunSizeEstimate: """ Build, parameterize, and estimate a scenario without initializing a run. @@ -213,7 +213,7 @@ async def create_and_estimate_async( configuration, baseline choice, and an optional objective target. Returns: - ScenarioDefaultRunSizeEstimate: Structured configured-run estimate. + ScenarioRunSizeEstimate: Structured configured-run estimate. """ scenario = await asyncio.to_thread(self.create_instance, name) scenario.set_scenario_registry_name(scenario_registry_name=name) diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 25083ae5ae..53151a7c89 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -84,7 +84,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, + supports_additional_request_converters: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -122,7 +122,7 @@ 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 + supports_additional_request_converters: Whether callers may safely append request converters to this technique. This is an explicit semantic opt-in, not merely constructor-signature detection. scorer_override_policy: What to do when a scenario's scorer is @@ -149,7 +149,7 @@ 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._supports_additional_request_converters = supports_additional_request_converters self._scorer_override_policy = scorer_override_policy self._uses_adversarial = uses_adversarial if uses_adversarial is not None else self._derive_uses_adversarial() @@ -174,7 +174,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, + supports_additional_request_converters: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -226,7 +226,7 @@ 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 + supports_additional_request_converters: Whether callers may safely append request converters to this technique. Forwarded to the factory constructor. scorer_override_policy: Policy applied when a scenario's scorer is @@ -289,7 +289,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, + supports_additional_request_converters=supports_additional_request_converters, scorer_override_policy=scorer_override_policy, ) @@ -334,11 +334,11 @@ def _validate_converter_composition(self) -> None: not accept ``attack_converter_config``. """ if ( - self._supports_request_converter_composition + self._supports_additional_request_converters and "attack_converter_config" not in self._get_accepted_params() ): raise ValueError( - f"Factory '{self._name}' declares supports_request_converter_composition=True, " + f"Factory '{self._name}' declares supports_additional_request_converters=True, " f"but {self._attack_class.__name__} does not accept 'attack_converter_config'." ) @@ -513,9 +513,9 @@ def uses_adversarial(self) -> bool: return self._uses_adversarial @property - def supports_request_converter_composition(self) -> bool: + def supports_additional_request_converters(self) -> bool: """Whether callers may safely append request converters to this technique.""" - return self._supports_request_converter_composition + return self._supports_additional_request_converters @property def scoring_config_type(self) -> type | None: @@ -861,7 +861,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, + "supports_additional_request_converters": self._supports_additional_request_converters, } if self._technique_tags: params["technique_tags"] = list(self._technique_tags) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 2ef11604ff..9eebbc86b8 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -39,7 +39,6 @@ AttackSeedGroup, ScenarioDatasetSizeCap, ScenarioDatasetSummary, - ScenarioDefaultRunSizeEstimate, ScenarioEvaluationIdentifier, ScenarioIdentifier, ScenarioResult, @@ -47,8 +46,7 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanSeedGroup, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, ScenarioRunState, config_hash, ) @@ -559,7 +557,7 @@ def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[Scen return self._technique_class.resolve(scenario_techniques, default=self._default_technique) @final - async def get_default_run_size_estimate_async(self) -> ScenarioDefaultRunSizeEstimate: + async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: """ Estimate the scenario's default planned execution units without starting a run. @@ -567,15 +565,13 @@ async def get_default_run_size_estimate_async(self) -> ScenarioDefaultRunSizeEst configured estimate path used by request-specific previews. Returns: - ScenarioDefaultRunSizeEstimate: Structured default-run estimate. + ScenarioRunSizeEstimate: Structured default-run estimate. """ self.set_params_from_args(args={}) return await self.get_run_size_estimate_async(target_is_configured=False) @final - async def get_run_size_estimate_async( - self, *, target_is_configured: bool = False - ) -> ScenarioDefaultRunSizeEstimate: + async def get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: """ Estimate the currently configured run without creating or persisting it. @@ -584,7 +580,7 @@ async def get_run_size_estimate_async( expansion, dataset selection, and baseline policy as ``initialize_async``. Returns: - ScenarioDefaultRunSizeEstimate: Structured configured-run estimate. + ScenarioRunSizeEstimate: Structured configured-run estimate. Raises: ValueError: If target certainty is asserted without a resolved target. @@ -595,7 +591,7 @@ async def get_run_size_estimate_async( 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: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate a standard technique-by-seed-group scenario. @@ -603,7 +599,7 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: synthesizes technique-specific populations, or selects techniques adaptively. Returns: - ScenarioDefaultRunSizeEstimate: Exact default sweep and baseline count. + ScenarioRunSizeEstimate: Exact default sweep and baseline count. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() seed_group_count = sum(len(groups) for groups in selected_groups.values()) @@ -616,28 +612,21 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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 = ( + estimated_attack_count = ( None - if status is ScenarioRunSizeEstimateStatus.Conditional + if self.RUN_SIZE_USES_FACTORY_COMPATIBILITY and self._estimate_has_binding_size_cap else sum(component.count for component in components) ) note = "Counts planned outer execution units; retries and internal attack turns are excluded." - if status is ScenarioRunSizeEstimateStatus.Conditional: + if estimated_attack_count is None: note += " A binding randomized dataset cap may select a different compatibility mix at launch." - return ScenarioDefaultRunSizeEstimate( - status=status, - total_attack_count=total_attack_count, + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, components=components, datasets=datasets, note=note, @@ -661,10 +650,6 @@ def _build_technique_size_components( 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), - ], ) ] @@ -690,10 +675,6 @@ def _build_technique_size_components( 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 diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index 200aa2119b..09e88734d8 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -23,10 +23,8 @@ from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackScoringConfig from pyrit.models import ( - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, ) from pyrit.models.identifiers import compute_inner_attack_eval_hash from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -203,12 +201,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list return atomic_attacks - async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate compatible persisted envelopes, excluding adaptive inner attempts. Returns: - ScenarioDefaultRunSizeEstimate: The adaptive outer-envelope estimate. + ScenarioRunSizeEstimate: The adaptive outer-envelope estimate. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() selected_count = sum(len(groups) for groups in selected_groups.values()) @@ -218,7 +216,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=selected_count, - factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count)], is_baseline=True, ) ] @@ -231,11 +228,9 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: ScenarioRunSizeComponent( label="Adaptive attack-envelope candidates", count=selected_count, - factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=selected_count)], ), ] - return ScenarioDefaultRunSizeEstimate( - status=ScenarioRunSizeEstimateStatus.Conditional, + return ScenarioRunSizeEstimate( components=components, datasets=datasets, note=( @@ -266,28 +261,19 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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) + estimated_attack_count = ( + None if self._estimate_has_binding_size_cap else sum(component.count for component in components) ) note = ( f"Each planned unit is one persisted adaptive envelope. Up to {max_attempts} selected technique " "attempts may run inside that unit; inner attempts and retries are excluded." ) - if status is ScenarioRunSizeEstimateStatus.Conditional: + if estimated_attack_count is None: note += " A binding randomized dataset cap may select a different compatibility mix at launch." - return ScenarioDefaultRunSizeEstimate( - status=status, - total_attack_count=total_attack_count, + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, components=components, datasets=datasets, note=note, diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 2fbece5e6c..07a14fcadb 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -15,10 +15,8 @@ from pyrit.models import ( AttackTechniqueSeedGroup, Parameter, - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, ) from pyrit.prompt_target import CapabilityName from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -79,7 +77,7 @@ def _prompt_sending_factory() -> AttackTechniqueFactory: name=_PROMPT_SENDING, attack_class=PromptSendingAttack, technique_tags=["single_turn"], - supports_request_converter_composition=True, + supports_additional_request_converters=True, ) @@ -101,7 +99,7 @@ def _jailbreak_system_prompt_factory() -> AttackTechniqueFactory: name=_JAILBREAK_SYSTEM_PROMPT, attack_class=PromptSendingAttack, technique_tags=["single_turn"], - supports_request_converter_composition=True, + supports_additional_request_converters=True, ) @@ -121,7 +119,7 @@ def _is_jailbreak_compatible_factory(factory: AttackTechniqueFactory) -> bool: return ( "multi_turn" not in factory.technique_tags and not has_simulated_conversation - and factory.supports_request_converter_composition + and factory.supports_additional_request_converters ) @@ -327,12 +325,12 @@ 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: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate the template and attempt axes, preserving the target capability caveat. Returns: - ScenarioDefaultRunSizeEstimate: Conditional target-aware estimate. + ScenarioRunSizeEstimate: Conditional target-aware estimate. Raises: ValueError: If native system-prompt delivery is the only selected @@ -364,7 +362,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: ScenarioRunSizeComponent( label="Baseline", count=seed_group_count, - factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], is_baseline=True, ) ) @@ -372,12 +369,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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." @@ -389,11 +380,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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 @@ -416,12 +402,10 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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 + estimated_attack_count = ( + None if system_delivery_selected and system_delivery_supported is None else planned_count ) - if status is ScenarioRunSizeEstimateStatus.Conditional: + if estimated_attack_count is None: capability_note = ( f" {target_agnostic_count} total planned units for target-agnostic delivery; " f"{planned_count} when native system-prompt delivery is supported." @@ -432,9 +416,8 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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, + return ScenarioRunSizeEstimate( + estimated_attack_count=estimated_attack_count, components=components, datasets=datasets, note=f"{formula}{baseline_explanation}{capability_note}", @@ -629,7 +612,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, + supports_additional_request_converters=True, ) @staticmethod diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 43aa58ff73..3b06528296 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -31,10 +31,8 @@ CrescendoAttack, ) from pyrit.models import ( - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, SeedPrompt, ) from pyrit.models.parameter import Parameter @@ -489,12 +487,12 @@ 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: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate the independent sub-harm technique sweeps and per-harm baselines. Returns: - ScenarioDefaultRunSizeEstimate: Exact per-sub-harm estimate. + ScenarioRunSizeEstimate: Exact per-sub-harm estimate. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() technique_count = len(self._scenario_techniques) @@ -505,10 +503,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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: @@ -516,14 +510,12 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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), + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Each default sub-harm is planned independently; retries and internal turns are excluded.", diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index cfe23046d0..ed86c8498f 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -15,11 +15,9 @@ AttackOutcome, AttackResult, ObjectiveTargetEvaluationIdentifier, - ScenarioDefaultRunSizeEstimate, ScenarioResult, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, ) from pyrit.models.parameter import Parameter from pyrit.registry import AttackTechniqueRegistry, TargetRegistry @@ -201,18 +199,17 @@ def __init__( scenario_result_id=scenario_result_id, ) - async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate the target-by-technique matrix using execution compatibility. Returns: - ScenarioDefaultRunSizeEstimate: Structured benchmark estimate. + ScenarioRunSizeEstimate: Structured benchmark estimate. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() target_names = self.params.get("adversarial_targets") or [] if not target_names: - return ScenarioDefaultRunSizeEstimate( - status=ScenarioRunSizeEstimateStatus.Conditional, + return ScenarioRunSizeEstimate( datasets=datasets, note=( "A total is unavailable until adversarial_targets is supplied and resolved. Baseline is forbidden." @@ -236,11 +233,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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), - ], ) ) @@ -250,8 +242,7 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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, + return ScenarioRunSizeEstimate( components=components, datasets=datasets, note=( @@ -259,9 +250,8 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: "so the authoritative total is unavailable before launch." ), ) - return ScenarioDefaultRunSizeEstimate( - status=ScenarioRunSizeEstimateStatus.Exact, - total_attack_count=sum(component.count for component in components), + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Baseline is forbidden; retries and internal attack turns are excluded.", diff --git a/pyrit/scenario/scenarios/foundry/red_team_agent.py b/pyrit/scenario/scenarios/foundry/red_team_agent.py index 176215e5f6..958683adc1 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -52,10 +52,8 @@ from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackConverterConfig, AttackScoringConfig from pyrit.models import ( AttackSeedGroup, - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, ) from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.prompt_target import PromptTarget @@ -420,12 +418,12 @@ def _resolve_foundry_techniques( self._scenario_composites = composites return flat - async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate one selected seed population per resolved Foundry composition. Returns: - ScenarioDefaultRunSizeEstimate: The composition population estimate. + ScenarioRunSizeEstimate: The composition population estimate. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() selected_count = sum(len(groups) for groups in selected_groups.values()) @@ -433,10 +431,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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 ] @@ -445,13 +439,11 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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), + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Counts one population per resolved Foundry composite, not per flattened constituent technique.", diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index f5be3c2caa..aaf00a411e 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -26,10 +26,8 @@ from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.models import ( AttackSeedGroup, - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, Seed, SeedObjective, SeedPrompt, @@ -232,12 +230,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list atomic_attacks.extend(self._get_converter_attacks(context=context)) return atomic_attacks - async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate converter variants crossed with raw and decode-template prompt configurations. Returns: - ScenarioDefaultRunSizeEstimate: Exact converter-variant estimate. + ScenarioRunSizeEstimate: Exact converter-variant estimate. """ selected_groups, datasets = await self._resolve_dataset_groups_for_estimate_async() seed_group_count = sum(len(groups) for groups in selected_groups.values()) @@ -248,14 +246,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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." @@ -267,13 +257,11 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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), + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note="Retries are excluded; each converter and decode-template configuration is a planned outer unit.", diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index d081b98b6b..ddf7216d4b 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -15,10 +15,8 @@ from pyrit.models import ( AttackSeedGroup, ScenarioDatasetSummary, - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, - ScenarioRunSizeEstimateStatus, - ScenarioRunSizeFactor, + ScenarioRunSizeEstimate, SeedObjective, SeedPrompt, ) @@ -527,12 +525,12 @@ def _build_synthesized_seed_groups( ) return seed_groups_by_technique - async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: + async def _estimate_run_size_async(self) -> ScenarioRunSizeEstimate: """ Estimate the technique-specific synthesized populations and their shared baseline. Returns: - ScenarioDefaultRunSizeEstimate: Exact synthesized-population estimate. + ScenarioRunSizeEstimate: Exact synthesized-population estimate. """ dataset_values = await asyncio.to_thread(self._load_dataset_values) seed_groups_by_technique = self._build_synthesized_seed_groups(dataset_values=dataset_values) @@ -560,7 +558,6 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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() ] @@ -570,19 +567,12 @@ async def _estimate_run_size_async(self) -> ScenarioDefaultRunSizeEstimate: 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), + return ScenarioRunSizeEstimate( + estimated_attack_count=sum(component.count for component in components), components=components, datasets=datasets, note=( diff --git a/pyrit/setup/initializers/techniques/airt.py b/pyrit/setup/initializers/techniques/airt.py index 5b3c92cec5..0ca857c6fd 100644 --- a/pyrit/setup/initializers/techniques/airt.py +++ b/pyrit/setup/initializers/techniques/airt.py @@ -42,7 +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, + supports_additional_request_converters=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters(converters=[FirstLetterConverter()]) diff --git a/pyrit/setup/initializers/techniques/core.py b/pyrit/setup/initializers/techniques/core.py index e696429af7..46bc31d6af 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -159,7 +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, + supports_additional_request_converters=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters( diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index ba9d919cd7..10e4a57552 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -29,7 +29,7 @@ AttackService, get_attack_service, ) -from pyrit.memory import AttackResultsKeysetCursor +from pyrit.memory import AttackResultKeysetCursor from pyrit.models import ( AtomicAttackIdentifier, AttackOutcome, @@ -160,7 +160,7 @@ def _cursor_for(result: AttackResult, *, fingerprint: str | None = None) -> str: """ effective_fingerprint = fingerprint if fingerprint is not None else AttackService._attack_filter_fingerprint() return AttackService._encode_attack_cursor( - cursor=AttackResultsKeysetCursor.from_attack_result(result), + cursor=AttackResultKeysetCursor.from_attack_result(result), fingerprint=effective_fingerprint, ) diff --git a/tests/unit/backend/test_scenario_run_routes.py b/tests/unit/backend/test_scenario_run_routes.py index 849e050292..627c0ab705 100644 --- a/tests/unit/backend/test_scenario_run_routes.py +++ b/tests/unit/backend/test_scenario_run_routes.py @@ -23,7 +23,7 @@ ScenarioRunProgress, ScenarioRunState, ) -from pyrit.models.catalog.scenario import ScenarioRunSummary +from pyrit.models.catalog.scenario import ScenarioRunListItem, ScenarioRunSummary from unit.mocks import make_scenario_result @@ -166,21 +166,33 @@ class TestListScenarioRunsRoute: def test_list_runs_returns_200(self, client: TestClient) -> None: """Test that list runs returns 200 with empty list.""" + route_thread: list[int] = [] with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: mock_service = MagicMock() - mock_service.list_runs.return_value = ScenarioRunListResponse(items=[]) + mock_service.list_runs.side_effect = lambda **_: ( + route_thread.append(get_ident()) or ScenarioRunListResponse(items=[]) + ) mock_get.return_value = mock_service + request_thread = get_ident() response = client.get("/api/scenarios/runs") assert response.status_code == status.HTTP_200_OK assert response.json()["items"] == [] + assert route_thread[0] != request_thread + + def test_list_runs_rejects_unbounded_limit(self, client: TestClient) -> None: + response = client.get("/api/scenarios/runs?limit=101") + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT def test_list_runs_returns_multiple_runs(self, client: TestClient) -> None: """Test that list runs returns all tracked runs.""" runs = [ - _mock_run_response(run_id="run-1"), - _mock_run_response(run_id="run-2", run_status=ScenarioRunState.IN_PROGRESS), + ScenarioRunListItem.model_validate(_mock_run_response(run_id="run-1").model_dump()), + ScenarioRunListItem.model_validate( + _mock_run_response(run_id="run-2", run_status=ScenarioRunState.IN_PROGRESS).model_dump() + ), ] with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get: diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index f1fe574069..fe237800c0 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -13,7 +13,9 @@ import pytest +import pyrit.backend.services.scenario_configuration_resolver as _resolver_mod import pyrit.backend.services.scenario_run_service as _svc_mod +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_run_service import ( _DEFAULT_MAX_CONCURRENT_RUNS, ScenarioRunService, @@ -59,7 +61,7 @@ def _patch_converter_registry(instances: dict[str, Any]): reg = MagicMock() reg.instances.get.side_effect = lambda name: instances.get(name) reg.instances.get_names.return_value = list(instances.keys()) - return patch.object(_svc_mod.ConverterRegistry, "get_registry_singleton", return_value=reg) + return patch.object(_resolver_mod.ConverterRegistry, "get_registry_singleton", return_value=reg) _REGISTRY_PATCH_BASE = "pyrit.registry" @@ -134,6 +136,7 @@ def mock_memory(): """Patch CentralMemory.get_memory_instance to return a mock.""" mock = MagicMock() mock.get_scenario_results.return_value = [] + mock.get_scenario_result_headers.return_value = [] # Default: no error AttackResults linked to any scenario. Tests that exercise # the error fallback path explicitly set get_attack_results.return_value. mock.get_attack_results.return_value = [] @@ -678,11 +681,11 @@ class TestScenarioRunServiceListRuns: def test_list_runs_empty(self, mock_memory) -> None: """Test that list_runs returns empty list when DB has no results.""" - mock_memory.get_scenario_results.return_value = [] + mock_memory.get_scenario_result_headers.return_value = [] service = ScenarioRunService() result = service.list_runs() assert result.items == [] - mock_memory.get_scenario_results.assert_called_once_with(limit=100) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=100) def test_list_runs_returns_all_runs(self, mock_memory) -> None: """Test that list_runs returns all runs from the database.""" @@ -690,19 +693,19 @@ def test_list_runs_returns_all_runs(self, mock_memory) -> None: _make_db_scenario_result(result_id="sr-1", run_state=ScenarioRunState.COMPLETED), _make_db_scenario_result(result_id="sr-2", run_state=ScenarioRunState.IN_PROGRESS), ] - mock_memory.get_scenario_results.return_value = db_results + mock_memory.get_scenario_result_headers.return_value = db_results service = ScenarioRunService() result = service.list_runs() assert len(result.items) == 2 - mock_memory.get_scenario_results.assert_called_once_with(limit=100) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=100) def test_list_runs_passes_custom_limit(self, mock_memory) -> None: """Test that list_runs passes a custom limit to the memory query.""" - mock_memory.get_scenario_results.return_value = [] + mock_memory.get_scenario_result_headers.return_value = [] service = ScenarioRunService() service.list_runs(limit=10) - mock_memory.get_scenario_results.assert_called_once_with(limit=10) + mock_memory.get_scenario_result_headers.assert_called_once_with(limit=10) class TestScenarioRunServiceCancelRun: @@ -1092,9 +1095,8 @@ class TestResolveTechniquesAndConverters: """Tests for per-technique converter resolution from ``--techniques`` tokens.""" def test_plain_technique_no_converters(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.ROLE_PLAY] @@ -1102,9 +1104,8 @@ def test_plain_technique_no_converters(self, mock_memory) -> None: def test_single_converter_appended(self, mock_memory) -> None: conv = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"translation_spanish": conv}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.translation_spanish"], technique_class=_StubTechnique, scenario_name="x", @@ -1114,9 +1115,8 @@ def test_single_converter_appended(self, mock_memory) -> None: def test_aggregate_token_applies_converter_to_all_concrete(self, mock_memory) -> None: conv = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": conv}): - enums, converters = service._resolve_techniques_and_converters( + enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.EASY] @@ -1125,9 +1125,8 @@ def test_aggregate_token_applies_converter_to_all_concrete(self, mock_memory) -> def test_multiple_converters_preserve_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = service._resolve_techniques_and_converters( + _, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.c1:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1137,9 +1136,8 @@ def test_multiple_converters_preserve_order(self, mock_memory) -> None: def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) - service = ScenarioRunService() with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = service._resolve_techniques_and_converters( + _, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1", "role_play:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1149,30 +1147,27 @@ def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: assert converters["single_turn"] == [c1] def test_unknown_converter_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({"known": MagicMock(spec=Converter)}): with pytest.raises(ValueError, match="not a registered converter"): - service._resolve_techniques_and_converters( + ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.missing"], technique_class=_StubTechnique, scenario_name="x", ) def test_unknown_modifier_prefix_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): with pytest.raises(ValueError, match="Unknown technique modifier"): - service._resolve_techniques_and_converters( + ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:scorer.something"], technique_class=_StubTechnique, scenario_name="x", ) def test_unknown_base_technique_raises(self, mock_memory) -> None: - service = ScenarioRunService() with _patch_converter_registry({}): with pytest.raises(ValueError, match="not found for scenario"): - service._resolve_techniques_and_converters( + ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["nope:converter.c1"], technique_class=_StubTechnique, scenario_name="x", @@ -1193,7 +1188,7 @@ async def test_start_run_forwards_technique_converters(self, mock_all_registries assert init_call.kwargs["technique_converters"] == {"role_play": [conv]} -def test_planned_progress_deduplicates_attempts_and_keeps_latest_non_error(mock_memory) -> None: +def test_planned_progress_deduplicates_attempts_and_keeps_latest_outcome(mock_memory) -> None: seed_group = AttackSeedGroup(seeds=[SeedObjective(value="objective")]) seed_group_id = seed_group.logical_id atomic_group_id = config_hash({"atomic_attack_name": "attack", "technique_eval_hash": "eval"}) @@ -1243,11 +1238,72 @@ def test_planned_progress_deduplicates_attempts_and_keeps_latest_non_error(mock_ assert summary.total_attacks == 1 assert summary.completed_attacks == 1 - assert summary.objective_achieved_rate == 100 + assert summary.objective_achieved_rate == 0 assert len(summary.failed_attacks) == 2 assert summary.total_retries == 3 +def test_planned_progress_includes_latest_errors_in_success_rate_denominator(mock_memory) -> None: + atomic_group_id = config_hash({"atomic_attack_name": "attack", "technique_eval_hash": "eval"}) + plan = ScenarioRunPlan( + scenario_registry_name="test.scenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id=atomic_group_id, + atomic_attack_name="attack", + display_group="Attack", + technique_eval_hash="eval", + seed_group_ids=["seed-success", "seed-error"], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id="seed-success", + objective_sha256="success-sha", + objective="success objective", + ), + ScenarioRunPlanSeedGroup( + id="seed-error", + objective_sha256="error-sha", + objective="error objective", + ), + ], + ) + results = [ + AttackResult( + conversation_id="success-conversation", + objective="success objective", + outcome=AttackOutcome.SUCCESS, + attribution_data={ + "parent_collection": "attack", + "parent_eval_hash": "eval", + "seed_group_id": "seed-success", + }, + ), + AttackResult( + conversation_id="error-conversation", + objective="error objective", + outcome=AttackOutcome.ERROR, + attribution_data={ + "parent_collection": "attack", + "parent_eval_hash": "eval", + "seed_group_id": "seed-error", + }, + ), + ] + scenario_result = make_scenario_result( + attack_results={"attack": results}, + scenario_run_state=ScenarioRunState.COMPLETED, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: plan.model_dump(mode="json")}, + ) + + summary = ScenarioRunService()._build_response_from_db(scenario_result=scenario_result) + + assert summary.total_attacks == 2 + assert summary.completed_attacks == 2 + assert summary.objective_achieved_rate == 50 + + def test_planned_progress_maps_legacy_objective_hash_to_logical_seed_id(mock_memory) -> None: objective = "legacy resumed objective" seed_group = AttackSeedGroup(seeds=[SeedObjective(value=objective)]) @@ -1373,7 +1429,10 @@ def test_progress_prefers_persisted_logical_seed_group_attribution() -> None: }, ) - mapped = ScenarioRunService._map_progress_delta(delta=delta, plan=None) + mapped = ScenarioRunService._map_progress_delta( + delta=delta, + plan_lookup=_svc_mod._ScenarioPlanLookup.from_plan(plan=None), + ) assert mapped.seed_group_id == "canonical-seed-id" diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index f1545a1f1e..e7297fa1cc 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -18,7 +18,7 @@ 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_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_service import ( ScenarioService, get_scenario_service, @@ -27,10 +27,9 @@ Parameter, ScenarioDatasetSizeCap, ScenarioDatasetSummary, - ScenarioDefaultRunSizeEstimate, ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, ScenarioRunSizeEstimateRequest, - ScenarioRunSizeEstimateStatus, ) from pyrit.models.catalog.scenario import RegisteredScenario from pyrit.registry import ScenarioMetadata @@ -152,9 +151,8 @@ async def test_list_scenarios_returns_scenarios_from_registry(self) -> None: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=4, components=[ScenarioRunSizeComponent(label="Default sweep", count=4)], datasets=[ ScenarioDatasetSummary( @@ -195,15 +193,14 @@ async def test_estimate_is_offloaded_and_cached(self) -> None: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], ) started = asyncio.Event() release = asyncio.Event() - async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + async def estimate_async() -> ScenarioRunSizeEstimate: started.set() await release.wait() return estimate @@ -246,15 +243,14 @@ def test_estimate_task_cleanup_preserves_replacement(self) -> None: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], ) started = asyncio.Event() release = asyncio.Event() - async def estimate_async() -> ScenarioDefaultRunSizeEstimate: + async def estimate_async() -> ScenarioRunSizeEstimate: started.set() await release.wait() return estimate @@ -284,9 +280,8 @@ async def estimate_async() -> ScenarioDefaultRunSizeEstimate: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], ) scenario = MagicMock() @@ -316,9 +311,8 @@ async def test_one_failed_estimate_does_not_break_catalog(self) -> None: _make_scenario_metadata(registry_name="test.good"), _make_scenario_metadata(registry_name="test.bad"), ] - estimate = ScenarioDefaultRunSizeEstimate( - status=ScenarioRunSizeEstimateStatus.Exact, - total_attack_count=2, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=2, components=[ScenarioRunSizeComponent(label="Default sweep", count=2)], ) good_scenario = MagicMock() @@ -336,17 +330,13 @@ async def test_one_failed_estimate_does_not_break_catalog(self) -> None: }[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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], ) scenario = MagicMock() @@ -368,15 +358,13 @@ async def test_unavailable_estimate_cache_expires(self) -> None: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Default sweep", count=1)], ) scenario = MagicMock() @@ -476,9 +464,8 @@ class TestScenarioServiceGetScenario: 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=12, components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)], ) introspection_instance = MagicMock() @@ -489,7 +476,9 @@ async def test_configured_estimate_uses_shared_launch_resolution(self) -> None: with ( patch.object(ScenarioService, "__init__", lambda self: None), - patch.object(ScenarioRunService, "resolve_target_name", return_value=objective_target) as resolve_target, + patch.object( + ScenarioConfigurationResolver, "resolve_target", return_value=objective_target + ) as resolve_target, ): service = ScenarioService() service._registry = MagicMock() @@ -557,8 +546,7 @@ async def test_configured_estimate_rejects_incompatible_v4_jailbreak_technique(s 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, + estimate = ScenarioRunSizeEstimate( note="Target compatibility is unknown.", ) introspection_instance = MagicMock() @@ -568,7 +556,7 @@ async def test_configured_estimate_without_target_does_not_resolve_or_send_to_ta with ( patch.object(ScenarioService, "__init__", lambda self: None), - patch.object(ScenarioRunService, "resolve_target_name") as resolve_target, + patch.object(ScenarioConfigurationResolver, "resolve_target") as resolve_target, ): service = ScenarioService() service._registry = MagicMock() @@ -726,9 +714,8 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: aggregate_techniques=["all"], all_techniques=["role_play"], default_datasets=["airt_hate"], - default_run_size=ScenarioDefaultRunSizeEstimate( - status=ScenarioRunSizeEstimateStatus.Exact, - total_attack_count=8, + default_run_size=ScenarioRunSizeEstimate( + estimated_attack_count=8, components=[ ScenarioRunSizeComponent( label="Default technique sweep", @@ -749,9 +736,7 @@ def test_get_scenario_returns_200(self, client: TestClient) -> None: 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 + assert data["default_run_size"]["estimated_attack_count"] == 8 def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> None: """Test that GET /api/scenarios/catalog/{name} returns 404 when not found.""" @@ -766,9 +751,8 @@ def test_get_scenario_returns_404_when_not_found(self, client: TestClient) -> No 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=12, components=[ScenarioRunSizeComponent(label="Configured Jailbreak", count=12)], ) with patch("pyrit.backend.routes.scenarios.get_scenario_service") as mock_get_service: @@ -789,7 +773,7 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien ) assert response.status_code == status.HTTP_200_OK - assert response.json()["total_attack_count"] == 12 + assert response.json()["estimated_attack_count"] == 12 request = mock_service.estimate_scenario_run_size_async.await_args.kwargs["request"] assert request.techniques == ["prompt_sending"] assert request.include_baseline is True @@ -800,9 +784,8 @@ def test_estimate_scenario_returns_configured_projection(self, client: TestClien 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, + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=1, components=[ScenarioRunSizeComponent(label="Configured estimate", count=1)], ) request = ScenarioRunSizeEstimateRequest() diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index f1b83a6c96..bdfa34f6ce 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -17,6 +17,7 @@ RegisteredInitializer, RegisteredScenario, RunScenarioRequest, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -395,7 +396,7 @@ async def test_list_scenario_runs_async(client, mock_httpx_client): mock_httpx_client.get.return_value = _make_response(json_data={"items": [_run_summary_payload()]}) result = await client.list_scenario_runs_async(limit=20) assert len(result) == 1 - assert isinstance(result[0], ScenarioRunSummary) + assert isinstance(result[0], ScenarioRunListItem) mock_httpx_client.get.assert_awaited_once_with("/api/scenarios/runs", params={"limit": 20}) diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index 2c0dd38c99..0cdfd77a83 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -20,6 +20,7 @@ AttackRetrySummary, RegisteredInitializer, RegisteredScenario, + ScenarioRunListItem, ScenarioRunSummary, TargetInstance, ) @@ -740,21 +741,21 @@ def test_print_scenario_runs_list_empty(capsys): def test_print_scenario_runs_list_populated(capsys): runs = [ - _make_run( + ScenarioRunListItem( status=ScenarioRunState.COMPLETED, scenario_name="scen-a", scenario_result_id="abcdefgh1234", total_attacks=4, - objective_achieved_rate=75, created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2024, 1, 1, tzinfo=timezone.utc), ), - _make_run( + ScenarioRunListItem( status=ScenarioRunState.IN_PROGRESS, scenario_name="scen-b", scenario_result_id="ijklmnop5678", total_attacks=0, - objective_achieved_rate=0, created_at=datetime(2024, 2, 2, tzinfo=timezone.utc), + updated_at=datetime(2024, 2, 2, tzinfo=timezone.utc), ), ] _output.print_scenario_runs_list(runs=runs) @@ -763,6 +764,7 @@ def test_print_scenario_runs_list_populated(capsys): assert "scen-b" in captured.out assert "abcdefgh1234" in captured.out assert "ijklmnop5678" in captured.out + assert "success" not in captured.out assert "…" not in captured.out assert "Total runs: 2" in captured.out diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 2a25984ece..d8378e9f98 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -23,7 +23,6 @@ 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 ( @@ -1094,25 +1093,6 @@ 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, diff --git a/tests/unit/memory/memory_interface/test_interface_attack_results.py b/tests/unit/memory/memory_interface/test_interface_attack_results.py index e532a1fe54..8687929ab4 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -11,7 +11,7 @@ import pytest from pyrit.common.utils import to_sha256 -from pyrit.memory import AttackResultsKeysetCursor, MemoryInterface +from pyrit.memory import AttackResultKeysetCursor, MemoryInterface from pyrit.memory.memory_interface import _AttackResultQuery from pyrit.memory.memory_models import AttackResultEntry from pyrit.models import ( @@ -85,15 +85,15 @@ def _make_attack_result( return AttackResult(**kwargs) -def _after(page: "Sequence[AttackResult]") -> AttackResultsKeysetCursor: +def _after(page: "Sequence[AttackResult]") -> AttackResultKeysetCursor: """Build the keyset anchor for the next page from the last row of ``page``.""" - return AttackResultsKeysetCursor.from_attack_result(page[-1]) + return AttackResultKeysetCursor.from_attack_result(page[-1]) def _drain_keyset(memory: MemoryInterface, *, page_size: int, **filters) -> list[AttackResult]: """Page through get_attack_results with the keyset cursor until exhausted.""" drained: list[AttackResult] = [] - after: AttackResultsKeysetCursor | None = None + after: AttackResultKeysetCursor | None = None while True: page = list(memory.get_attack_results(limit=page_size, after=after, **filters)) drained.extend(page) @@ -122,12 +122,12 @@ def test_attack_result_query_snapshots_mutable_inputs(): def test_attack_result_query_requires_keyword_arguments(): """The internal query does not expose field ordering as a positional API.""" with pytest.raises(TypeError): - _AttackResultQuery(["id"]) # type: ignore[misc] + _AttackResultQuery(["id"]) # ty: ignore[too-many-positional-arguments] def test_get_attack_results_forwards_all_parameters_to_query(sqlite_instance: MemoryInterface): """The compatibility API maps every parameter onto the internal query.""" - cursor = AttackResultsKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) + cursor = AttackResultKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) identifier_filter = IdentifierFilter( identifier_type=IdentifierType.ATTACK, property_path="$.hash", @@ -1898,7 +1898,7 @@ def test_get_attack_results_paginated_empty_metadata_orders_newest_first(sqlite_ def test_get_attack_results_pagination_with_ids_raises(sqlite_instance: MemoryInterface): """limit/keyset pagination cannot be combined with id-batched lookups.""" - anchor = AttackResultsKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) + anchor = AttackResultKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) with pytest.raises(ValueError, match="pagination cannot be combined"): sqlite_instance.get_attack_results(attack_result_ids=[str(uuid.uuid4())], limit=10) with pytest.raises(ValueError, match="pagination cannot be combined"): @@ -2055,7 +2055,7 @@ def test_attack_result_keyset_order_matches_sql_order(sqlite_instance: MemoryInt python_order = sorted( sqlite_instance.get_attack_results(), key=lambda ar: ( - AttackResultsKeysetCursor.from_attack_result(ar).timestamp, + AttackResultKeysetCursor.from_attack_result(ar).timestamp, ar.attack_result_id, ), reverse=True, diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_progress.py b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py index d54ee91560..35e68844be 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_progress.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_progress.py @@ -6,10 +6,10 @@ import uuid from datetime import datetime, timezone +import pytest 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.memory import AttackResultKeysetCursor, MemoryInterface from pyrit.models import ( AtomicAttackIdentifier, AttackOutcome, @@ -88,7 +88,7 @@ def test_scenario_progress_deltas_page_equal_timestamps_by_id( ) second_page, second_has_more = sqlite_instance.get_scenario_attack_result_deltas( scenario_result_id=str(scenario.id), - cursor=ScenarioProgressKeysetCursor( + cursor=AttackResultKeysetCursor( timestamp=first_page[0].timestamp, attack_result_id=first_page[0].attack_result_id, ), @@ -127,3 +127,24 @@ def test_scenario_result_header_does_not_hydrate_attack_results( assert header is not None assert header.attack_results == {} + + +def test_scenario_result_headers_are_bounded_without_attack_results( + sqlite_instance: MemoryInterface, +) -> None: + scenarios = [ + make_scenario_result( + scenario_name=f"scenario-{index}", + attack_results={}, + objective_target_identifier=get_mock_target_identifier(), + ) + for index in range(2) + ] + sqlite_instance.add_scenario_results_to_memory(scenario_results=scenarios) + + headers = sqlite_instance.get_scenario_result_headers(limit=1) + + assert len(headers) == 1 + assert headers[0].attack_results == {} + with pytest.raises(ValueError, match="between 1 and 100"): + sqlite_instance.get_scenario_result_headers(limit=101) diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py index c8451ca502..f787eed043 100644 --- a/tests/unit/models/test_scenario_catalog.py +++ b/tests/unit/models/test_scenario_catalog.py @@ -9,138 +9,65 @@ 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.""" +def test_run_size_estimate_requires_available_total_to_match_components() -> None: + """Available estimates require an additive component total.""" with pytest.raises(ValidationError, match="components total 6, not 7"): - 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), - ], - ) - ], + ScenarioRunSizeEstimate( + estimated_attack_count=7, + components=[ScenarioRunSizeComponent(label="Techniques", count=6)], ) -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_run_size_estimate_allows_unavailable_count_with_components() -> None: + """Unavailable estimates retain useful candidate components and an explanatory note.""" + estimate = ScenarioRunSizeEstimate( + components=[ScenarioRunSizeComponent(label="Candidate techniques", count=6)], + note="The final count depends on target capabilities.", + ) + assert estimate.estimated_attack_count is None + assert estimate.components[0].count == 6 -def test_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), - ], - ) - ], + +def test_run_size_estimate_serializes_canonical_api_shape() -> None: + """The estimate exposes only the available count and additive components.""" + estimate = ScenarioRunSizeEstimate( + estimated_attack_count=6, + components=[ScenarioRunSizeComponent(label="Techniques", count=6)], ) assert estimate.model_dump(mode="json") == { - "version": 1, - "status": "exact", - "total_attack_count": 6, + "estimated_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, +def test_unavailable_run_size_estimate_has_no_count() -> None: + """The unavailable factory communicates that a count cannot be calculated.""" + estimate = ScenarioRunSizeEstimate.unavailable() + + assert estimate.estimated_attack_count is None + assert estimate.note == "Default-run size estimate is unavailable." + + +def test_estimate_exposes_dataset_counts_structurally() -> None: + """Effective dataset selection remains machine-readable.""" + estimate = ScenarioRunSizeEstimate( datasets=[ ScenarioDatasetSummary( name="harmbench", @@ -157,13 +84,11 @@ def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: ], ) ], - note="The final total depends on target capabilities.", + note="The final count depends on target capabilities.", ) - payload = estimate.model_dump(mode="json") - assert payload["status"] == "conditional" - assert payload["total_attack_count"] is None - assert payload["datasets"] == [ + assert estimate.estimated_attack_count is None + assert estimate.model_dump(mode="json")["datasets"] == [ { "name": "harmbench", "kind": "dataset", @@ -180,7 +105,6 @@ def test_conditional_estimate_exposes_dataset_counts_structurally() -> None: ], } ] - assert payload["retries_included"] is False def test_estimate_request_reuses_dataset_filter_validation() -> None: diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index b87c44f8da..cbb43ce36d 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -8,7 +8,7 @@ import pytest -from pyrit.backend.services.scenario_run_service import ScenarioRunService +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.common.path import JAILBREAK_TEMPLATES_PATH from pyrit.converter import TextJailbreakConverter from pyrit.datasets import TextJailBreak @@ -16,7 +16,6 @@ from pyrit.models import ( AttackSeedGroup, ComponentIdentifier, - ScenarioRunSizeEstimateStatus, SeedObjective, SeedPrompt, ) @@ -229,16 +228,8 @@ async def test_run_size_prompt_sending_two_templates_four_groups_is_eight( ) 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 estimate.estimated_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)] @@ -260,9 +251,7 @@ async def test_run_size_is_conditional_when_system_delivery_target_is_not_select ) 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 estimate.estimated_attack_count is None assert [component.label for component in estimate.components] == [ "Inline jailbreak delivery", "Native system-prompt jailbreak delivery", @@ -468,7 +457,7 @@ async def test_incompatible_runtime_factory_is_rejected( name="legacy_multi_turn", attack_class=PromptSendingAttack, technique_tags=["multi_turn"], - supports_request_converter_composition=True, + supports_additional_request_converters=True, ) with _patch_seed_groups(mock_memory_seed_groups): with patch( @@ -825,7 +814,7 @@ def test_registry_metadata_omits_incompatible_techniques(self): def test_configured_estimate_rejects_context_compliance(self): with pytest.raises(ValueError, match="Technique 'context_compliance' not found"): - ScenarioRunService.resolve_scenario_configuration( + ScenarioConfigurationResolver.resolve_configuration( scenario_name="airt.jailbreak", scenario_class=Jailbreak, techniques=["context_compliance"], diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index ff4f61b90c..c5324d8d39 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -216,7 +216,7 @@ def __init__(self, *, objective_target, attack_scoring_config=None): AttackTechniqueFactory( name="test", attack_class=_NoConverterAttack, - supports_request_converter_composition=True, + supports_additional_request_converters=True, ) def test_request_converter_composition_is_explicit_opt_in(self): @@ -224,11 +224,11 @@ def test_request_converter_composition_is_explicit_opt_in(self): composable_factory = AttackTechniqueFactory( name="composable", attack_class=_StubAttack, - supports_request_converter_composition=True, + supports_additional_request_converters=True, ) - assert not default_factory.supports_request_converter_composition - assert composable_factory.supports_request_converter_composition + assert not default_factory.supports_additional_request_converters + assert composable_factory.supports_additional_request_converters class TestFactoryCreate: diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py index 509985104f..2d69edf96c 100644 --- a/tests/unit/scenario/test_default_run_size_estimates.py +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -14,7 +14,6 @@ AttackTechniqueSeedGroup, ComponentIdentifier, ScenarioDatasetSummary, - ScenarioRunSizeEstimateStatus, SeedObjective, SeedPrompt, SeedSimulatedConversation, @@ -131,9 +130,7 @@ def _resolved_groups( 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 estimate.estimated_attack_count == 6 assert [component.count for component in estimate.components] == [4, 2] assert estimate.datasets[0].logical_seed_group_count == 3 assert estimate.datasets[0].selected_seed_group_count == 2 @@ -153,11 +150,8 @@ async def test_configured_estimate_reuses_technique_and_baseline_resolution_with ) estimate = await scenario.get_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 2 + assert estimate.estimated_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() == [] @@ -173,10 +167,7 @@ async def test_configured_estimate_expands_requested_aggregate() -> None: ) 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] + assert estimate.estimated_attack_count == 4 @pytest.mark.usefixtures("patch_central_database") @@ -196,7 +187,7 @@ async def test_configured_estimate_applies_dataset_selection_and_cap() -> None: estimate = await scenario.get_run_size_estimate_async() - assert estimate.total_attack_count == 2 + assert estimate.estimated_attack_count == 2 assert len(estimate.datasets) == 1 assert estimate.datasets[0].logical_seed_group_count == 3 assert estimate.datasets[0].selected_seed_group_count == 2 @@ -270,7 +261,7 @@ async def test_matrix_estimate_filters_each_technique_seed_population_like_execu ): estimate = await scenario.get_run_size_estimate_async() - assert estimate.total_attack_count == 3 + assert estimate.estimated_attack_count == 3 assert [(component.label, component.count) for component in estimate.components] == [("one", 2), ("two", 1)] @@ -293,9 +284,7 @@ async def resolve_groups() -> tuple[dict[str, list[AttackSeedGroup]], list[Scena 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 estimate.estimated_attack_count is None assert "binding randomized dataset cap" in estimate.note @@ -307,9 +296,7 @@ async def test_adaptive_estimate_is_target_conditional_and_does_not_multiply_tec 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 estimate.estimated_attack_count is None assert [component.count for component in estimate.components] == [3, 3] @@ -338,17 +325,13 @@ async def test_adaptive_estimate_counts_exact_compatible_outer_envelopes_with_ta ), ): estimate = await scenario.get_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 2 + assert estimate.estimated_attack_count == 2 assert [component.count for component in estimate.components] == [2] assert "7 selected technique attempts" in estimate.note scenario.set_params_from_args(args={"include_baseline": False}) estimate_without_target = await scenario.get_run_size_estimate_async() - - assert estimate_without_target.status is ScenarioRunSizeEstimateStatus.Conditional - assert estimate_without_target.total_attack_count is None + assert estimate_without_target.estimated_attack_count is None @pytest.mark.usefixtures("patch_central_database") @@ -359,11 +342,8 @@ async def test_jailbreak_estimate_exposes_template_attempt_and_target_capability 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 estimate.estimated_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 @@ -386,11 +366,8 @@ async def test_jailbreak_configured_estimate_counts_prompt_sending_without_basel ) estimate = await scenario.get_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 8 + assert estimate.estimated_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 @@ -411,12 +388,9 @@ async def test_jailbreak_configured_estimate_counts_prompt_sending_with_baseline ) estimate = await scenario.get_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 12 + assert estimate.estimated_attack_count == 12 assert [component.count for component in estimate.components] == [4, 8] assert estimate.components[0].is_baseline is True - assert [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 @@ -440,9 +414,7 @@ async def test_jailbreak_configured_estimate_uses_target_capability() -> None: ) estimate = await scenario.get_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 8 + assert estimate.estimated_attack_count == 8 assert [component.count for component in estimate.components] == [0, 8] objective_target.send_prompt_async.assert_not_called() @@ -481,10 +453,7 @@ async def test_encoding_estimate_counts_concrete_converter_and_decode_variants() 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] + assert estimate.estimated_attack_count == 152 @pytest.mark.usefixtures("patch_central_database") @@ -502,9 +471,8 @@ async def test_web_injection_estimate_uses_synthesized_technique_populations() - 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.estimated_attack_count == synthesized_count * 2 assert estimate.components[-1].label == "Baseline" @@ -520,9 +488,7 @@ async def test_psychosocial_estimate_keeps_sub_harm_baselines_separate() -> None ) estimate = await scenario.get_default_run_size_estimate_async() - - assert estimate.status is ScenarioRunSizeEstimateStatus.Exact - assert estimate.total_attack_count == 12 + assert estimate.estimated_attack_count == 12 assert [component.count for component in estimate.components] == [6, 2, 3, 1] @@ -537,25 +503,22 @@ async def test_adversarial_benchmark_estimate_exposes_per_required_target_formul 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.estimated_attack_count is None assert estimate.components == [] assert "adversarial_targets" in estimate.note @pytest.mark.parametrize( - ("use_cached", "expected_status", "expected_total"), + ("use_cached", "expected_total"), [ - (False, ScenarioRunSizeEstimateStatus.Exact, 6), - (True, ScenarioRunSizeEstimateStatus.Conditional, None), + (False, 6), + (True, None), ], ) @pytest.mark.usefixtures("patch_central_database") async def test_adversarial_benchmark_resolves_targets_and_filters_each_technique( *, use_cached: bool, - expected_status: ScenarioRunSizeEstimateStatus, expected_total: int | None, ) -> None: """Benchmark sizing resolves target names and reports uncached compatible candidates.""" @@ -607,8 +570,7 @@ async def test_adversarial_benchmark_resolves_targets_and_filters_each_technique 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 estimate.estimated_attack_count == expected_total assert [(component.label, component.count) for component in estimate.components] == [("one", 4), ("two", 2)] @@ -635,7 +597,6 @@ async def test_foundry_estimate_counts_composites_instead_of_flattened_technique estimate = await scenario.get_run_size_estimate_async() - assert estimate.total_attack_count == 6 + assert estimate.estimated_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] From c8b2b184e283731df402e67e4d22747641685964 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:36:24 -0700 Subject: [PATCH 3/4] Narrow scenario estimate task result Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- pyrit/backend/services/scenario_service.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index dca9623d07..b5d6d3ac5b 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -211,7 +211,9 @@ def clear_estimate_task(completed_task: _EstimateTask) -> None: task.add_done_callback(clear_estimate_task) if task is not None: - return await asyncio.shield(task) + estimate = await asyncio.shield(task) + assert isinstance(estimate, ScenarioRunSizeEstimate) + return estimate if wait_for_capacity is not None: await asyncio.shield(wait_for_capacity) From c770e714ef538e69c320c7b4f073f546efee4cfd Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:42:26 -0700 Subject: [PATCH 4/4] Remove duplicate scenario resolver import Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 47a96cf7-36b1-482c-ab17-87cbf97afbdf --- tests/unit/backend/test_scenario_run_service.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index fe237800c0..f5f652d9bc 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -15,7 +15,6 @@ import pyrit.backend.services.scenario_configuration_resolver as _resolver_mod import pyrit.backend.services.scenario_run_service as _svc_mod -from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_run_service import ( _DEFAULT_MAX_CONCURRENT_RUNS, ScenarioRunService, @@ -1096,7 +1095,7 @@ class TestResolveTechniquesAndConverters: def test_plain_technique_no_converters(self, mock_memory) -> None: with _patch_converter_registry({}): - enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.ROLE_PLAY] @@ -1105,7 +1104,7 @@ def test_plain_technique_no_converters(self, mock_memory) -> None: def test_single_converter_appended(self, mock_memory) -> None: conv = MagicMock(spec=Converter) with _patch_converter_registry({"translation_spanish": conv}): - enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.translation_spanish"], technique_class=_StubTechnique, scenario_name="x", @@ -1116,7 +1115,7 @@ def test_single_converter_appended(self, mock_memory) -> None: def test_aggregate_token_applies_converter_to_all_concrete(self, mock_memory) -> None: conv = MagicMock(spec=Converter) with _patch_converter_registry({"c1": conv}): - enums, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( + enums, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1"], technique_class=_StubTechnique, scenario_name="x" ) assert enums == [_StubTechnique.EASY] @@ -1126,7 +1125,7 @@ def test_multiple_converters_preserve_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( + _, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.c1:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1137,7 +1136,7 @@ def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: c1 = MagicMock(spec=Converter) c2 = MagicMock(spec=Converter) with _patch_converter_registry({"c1": c1, "c2": c2}): - _, converters = ScenarioConfigurationResolver.resolve_techniques_and_converters( + _, converters = _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["easy:converter.c1", "role_play:converter.c2"], technique_class=_StubTechnique, scenario_name="x", @@ -1149,7 +1148,7 @@ def test_overlapping_tokens_append_in_order(self, mock_memory) -> None: def test_unknown_converter_raises(self, mock_memory) -> None: with _patch_converter_registry({"known": MagicMock(spec=Converter)}): with pytest.raises(ValueError, match="not a registered converter"): - ScenarioConfigurationResolver.resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:converter.missing"], technique_class=_StubTechnique, scenario_name="x", @@ -1158,7 +1157,7 @@ def test_unknown_converter_raises(self, mock_memory) -> None: def test_unknown_modifier_prefix_raises(self, mock_memory) -> None: with _patch_converter_registry({}): with pytest.raises(ValueError, match="Unknown technique modifier"): - ScenarioConfigurationResolver.resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["role_play:scorer.something"], technique_class=_StubTechnique, scenario_name="x", @@ -1167,7 +1166,7 @@ def test_unknown_modifier_prefix_raises(self, mock_memory) -> None: def test_unknown_base_technique_raises(self, mock_memory) -> None: with _patch_converter_registry({}): with pytest.raises(ValueError, match="not found for scenario"): - ScenarioConfigurationResolver.resolve_techniques_and_converters( + _resolver_mod.ScenarioConfigurationResolver.resolve_techniques_and_converters( tokens=["nope:converter.c1"], technique_class=_StubTechnique, scenario_name="x",