From 78f57417653fe8468f88d8c892cde266c556de7d Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sun, 11 Oct 2026 15:16:19 -0700 Subject: [PATCH 1/4] [BREAKING] FEAT: Add scenario dataset controls across application surfaces Expose scenario-owned dataset options, shared requests, and source-level controls in Python, REST, CLI, shell, and GUI. Preserve read-only estimates and resume behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16cbdc1b-2bf6-4286-be66-c9f30c7ed72b --- .../instructions/scenarios.instructions.md | 14 ++ doc/code/scenarios/0_scenarios.ipynb | 3 +- doc/code/scenarios/0_scenarios.py | 3 +- .../1_common_scenario_parameters.ipynb | 2 +- .../scenarios/1_common_scenario_parameters.py | 2 +- doc/code/scenarios/3_adaptive_scenarios.ipynb | 2 +- doc/code/scenarios/3_adaptive_scenarios.py | 2 +- doc/gui/0_gui.md | 24 ++- doc/scanner/1_pyrit_scan.ipynb | 10 +- doc/scanner/1_pyrit_scan.py | 10 +- doc/scanner/adaptive.ipynb | 4 +- doc/scanner/adaptive.py | 4 +- doc/scanner/airt.ipynb | 22 +- doc/scanner/airt.py | 22 +- doc/scanner/benchmark.ipynb | 2 +- doc/scanner/benchmark.py | 2 +- doc/scanner/foundry.ipynb | 2 +- doc/scanner/foundry.py | 2 +- doc/scanner/garak.ipynb | 28 +-- doc/scanner/garak.py | 28 +-- .../Scenarios/ScenarioDatasetTable.styles.ts | 13 ++ .../Scenarios/ScenarioDatasetTable.test.tsx | 109 ++++++++++ .../Scenarios/ScenarioDatasetTable.tsx | 194 +++++++++++++++++ .../Scenarios/ScenarioDetail.test.tsx | 113 +++++++--- .../components/Scenarios/ScenarioDetail.tsx | 177 +++++++--------- .../src/hooks/useScenarioDatasets.test.tsx | 58 +++++ frontend/src/hooks/useScenarioDatasets.ts | 50 +++++ frontend/src/services/api.ts | 10 + frontend/src/types/index.ts | 59 ++++++ pyrit/backend/models/datasets.py | 1 + pyrit/backend/services/dataset_service.py | 3 + .../scenario_configuration_resolver.py | 31 +-- .../backend/services/scenario_run_service.py | 22 +- pyrit/backend/services/scenario_service.py | 6 +- pyrit/cli/_cli_args.py | 88 +++++++- pyrit/cli/_output.py | 11 + pyrit/cli/pyrit_scan.py | 29 ++- pyrit/cli/pyrit_shell.py | 27 ++- pyrit/memory/memory_interface.py | 12 +- pyrit/models/__init__.py | 16 ++ pyrit/models/catalog/scenario.py | 142 ++++++------- pyrit/models/scenario_dataset.py | 130 ++++++++++++ pyrit/models/seeds/seed_dataset_summary.py | 1 + .../registry/components/scenario_registry.py | 15 +- pyrit/scenario/core/dataset_configuration.py | 13 +- pyrit/scenario/core/dataset_overrides.py | 122 +++++++++++ pyrit/scenario/core/scenario.py | 55 ++++- pyrit/scenario/scenarios/airt/psychosocial.py | 42 +++- pyrit/scenario/scenarios/garak/api_key.py | 2 + pyrit/scenario/scenarios/garak/encoding.py | 4 + .../scenario/scenarios/garak/exploitation.py | 6 +- pyrit/scenario/scenarios/garak/figstep.py | 4 + .../scenarios/garak/latent_injection.py | 2 + .../scenarios/garak/package_hallucination.py | 15 +- .../scenario/scenarios/garak/prompt_inject.py | 2 + .../garak/system_prompt_extraction.py | 4 + .../scenario/scenarios/garak/web_injection.py | 16 ++ tests/unit/backend/test_dataset_service.py | 3 + .../test_scenario_configuration_resolver.py | 46 ++-- tests/unit/backend/test_scenario_resume.py | 30 ++- .../unit/backend/test_scenario_run_service.py | 44 ++-- tests/unit/backend/test_scenario_service.py | 19 +- tests/unit/cli/test_dataset_request.py | 54 +++++ ...st_interface_seed_dataset_summary_query.py | 16 +- tests/unit/models/test_scenario_catalog.py | 3 + tests/unit/models/test_scenario_request.py | 99 +++++++++ tests/unit/registry/test_scenario_registry.py | 11 + tests/unit/scenario/airt/test_leakage.py | 4 + tests/unit/scenario/airt/test_psychosocial.py | 11 +- tests/unit/scenario/airt/test_scam.py | 4 + .../core/test_scenario_dataset_overrides.py | 198 ++++++++++++++++++ .../scenario/foundry/test_red_team_agent.py | 4 + tests/unit/scenario/garak/test_api_key.py | 4 +- tests/unit/scenario/garak/test_divergence.py | 3 +- tests/unit/scenario/garak/test_encoding.py | 4 + 75 files changed, 1935 insertions(+), 414 deletions(-) create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.tsx create mode 100644 frontend/src/hooks/useScenarioDatasets.test.tsx create mode 100644 frontend/src/hooks/useScenarioDatasets.ts create mode 100644 pyrit/models/scenario_dataset.py create mode 100644 pyrit/scenario/core/dataset_overrides.py create mode 100644 tests/unit/cli/test_dataset_request.py create mode 100644 tests/unit/scenario/core/test_scenario_dataset_overrides.py diff --git a/.github/instructions/scenarios.instructions.md b/.github/instructions/scenarios.instructions.md index bef7214bd7..20e5f4c80b 100644 --- a/.github/instructions/scenarios.instructions.md +++ b/.github/instructions/scenarios.instructions.md @@ -163,6 +163,20 @@ Options: - Ingredients must remain complete. Use `max_per_dataset="all"` and no finite source `max_size`; cap the assembled attack groups with `max_total`, not ingredient rows. - Use `with_overrides()` instead of reconstructing a subclass or mutating its defaults. +- Clients and Python use `ScenarioDatasetRequest` through the common `dataset_overrides` + parameter. `dataset_config` still replaces the default; overrides apply on top. +- Resolve parameter/technique-dependent names in `_resolve_default_dataset_config()`, + before overrides, validation, and preparation. Do not rewrite user selections in validators. +- Declare `DATASET_OPTIONS` (or parameter-aware `get_dataset_options()`): `ANY` only for + generic named populations, `SUBSET` for approved alternatives, `FIXED` for ingredients. + Disable limits that would truncate ingredients. Compound source edits are ambiguous; + preserve child shaping/caps and allow only supported global edits. +- Estimates expose `dataset_defaults`, `effective_datasets`, and `dataset_options`. + These are response projections, not incoming edit state. Never read seeds for projections. + The catalog reads options and defaults from registry metadata, never from an estimate. +- Resume replays the saved launch request; the stored run plan fixes the selected seed groups, + so caps and fetch policies cannot change a resumed run. Dataset names are checked by the + stored scenario identity. Resume never prepares datasets. - Keep custom shaping in `_build_attack_groups()` or `_build_groups_by_dataset_async()`. Reads use `_collect_seeds_for_dataset_async()` and must not fetch. - `dataset_names`, `max_dataset_size`, `auto_fetch`, and `per_dataset()` are deprecated. diff --git a/doc/code/scenarios/0_scenarios.ipynb b/doc/code/scenarios/0_scenarios.ipynb index b088db7741..7b0b110a57 100644 --- a/doc/code/scenarios/0_scenarios.ipynb +++ b/doc/code/scenarios/0_scenarios.ipynb @@ -66,10 +66,11 @@ "\n", "3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box.\n", " - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name=\"my_dataset\")]`)\n", - " - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically\n", + " - Users select sources with `--datasets name[:N|:all]`, or pass `dataset_overrides=ScenarioDatasetRequest(...)` in `set_params_from_args`. A custom `dataset_config` replaces the default; overrides apply on top.\n", " - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`\"default\"` uses the default; `\"all\"` removes that limit.\n", " - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset=\"all\", max_total=10`; to remove both caps, set both to `\"all\"`. `None` now uses the default, not unlimited selection.\n", " - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch\n", + " - Scenarios declare allowed source selections and limit controls. Estimates return resolved `dataset_defaults` and `effective_datasets` separately from requested edits. Name-only selection preserves existing source caps and fetch policies.\n", "\n", "4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:\n", " - `name`: Descriptive name for your scenario\n", diff --git a/doc/code/scenarios/0_scenarios.py b/doc/code/scenarios/0_scenarios.py index 3a951368a4..a9407131da 100644 --- a/doc/code/scenarios/0_scenarios.py +++ b/doc/code/scenarios/0_scenarios.py @@ -68,10 +68,11 @@ # # 3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box. # - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name="my_dataset")]`) -# - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically +# - Users select sources with `--datasets name[:N|:all]`, or pass `dataset_overrides=ScenarioDatasetRequest(...)` in `set_params_from_args`. A custom `dataset_config` replaces the default; overrides apply on top. # - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`"default"` uses the default; `"all"` removes that limit. # - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset="all", max_total=10`; to remove both caps, set both to `"all"`. `None` now uses the default, not unlimited selection. # - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch +# - Scenarios declare allowed source selections and limit controls. Estimates return resolved `dataset_defaults` and `effective_datasets` separately from requested edits. Name-only selection preserves existing source caps and fetch policies. # # 4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata: # - `name`: Descriptive name for your scenario diff --git a/doc/code/scenarios/1_common_scenario_parameters.ipynb b/doc/code/scenarios/1_common_scenario_parameters.ipynb index e6081ccad4..5950f87f7c 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.ipynb +++ b/doc/code/scenarios/1_common_scenario_parameters.ipynb @@ -15,7 +15,7 @@ "\n", "> **Two selection axes**: *Techniques* select attack techniques (*how* attacks run — e.g., prompt\n", "> sending, role play, TAP). *Datasets* select objectives (*what* is tested — e.g., harm categories,\n", - "> compliance topics). Use `--dataset-names` on the CLI to filter by content category.\n", + "> compliance topics). Use `--datasets` on the CLI to filter by content category.\n", "\n", "> **Running scenarios from the command line?** See the [Scanner documentation](../../scanner/0_scanner.md).\n", "\n", diff --git a/doc/code/scenarios/1_common_scenario_parameters.py b/doc/code/scenarios/1_common_scenario_parameters.py index 354df1410e..0ced660ed1 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.py +++ b/doc/code/scenarios/1_common_scenario_parameters.py @@ -19,7 +19,7 @@ # # > **Two selection axes**: *Techniques* select attack techniques (*how* attacks run — e.g., prompt # > sending, role play, TAP). *Datasets* select objectives (*what* is tested — e.g., harm categories, -# > compliance topics). Use `--dataset-names` on the CLI to filter by content category. +# > compliance topics). Use `--datasets` on the CLI to filter by content category. # # > **Running scenarios from the command line?** See the [Scanner documentation](../../scanner/0_scanner.md). # diff --git a/doc/code/scenarios/3_adaptive_scenarios.ipynb b/doc/code/scenarios/3_adaptive_scenarios.ipynb index 2c9b2bc2f5..436b55b09a 100644 --- a/doc/code/scenarios/3_adaptive_scenarios.ipynb +++ b/doc/code/scenarios/3_adaptive_scenarios.ipynb @@ -804,7 +804,7 @@ "# Use specific datasets and limit size\n", "pyrit_scan --scenario TextAdaptive --target openai_chat \\\n", " --datasets airt_hate airt_violence \\\n", - " --max-dataset-size 10\n", + " --max-total 10\n", "```" ] } diff --git a/doc/code/scenarios/3_adaptive_scenarios.py b/doc/code/scenarios/3_adaptive_scenarios.py index 247b24b786..2614eb1d17 100644 --- a/doc/code/scenarios/3_adaptive_scenarios.py +++ b/doc/code/scenarios/3_adaptive_scenarios.py @@ -301,5 +301,5 @@ def _technique_label(result) -> str: # # Use specific datasets and limit size # pyrit_scan --scenario TextAdaptive --target openai_chat \ # --datasets airt_hate airt_violence \ -# --max-dataset-size 10 +# --max-total 10 # ``` diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index d23fe60a8f..7707351f71 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -490,7 +490,29 @@ Defaults are saved in this browser, separately for each signed-in account. They A missing or changed default is shown as unavailable. Select a new default or clear it; CoPyRIT does not silently substitute a different target. If browser storage is unavailable, a warning states that the choice applies only in the current session. -Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters** above **Dataset override**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. +Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters** above **Datasets**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. + +### Scenario datasets + +The **Datasets** table shows each selected name, stored seed origins (including mixed +origins), stored example count, source limit, and fetch policy. Unknown counts for +provider-only datasets are not zero. Provider availability is separate from stored origin. +**Reload dataset metadata** refreshes the page-scoped catalog without fetching seed data. +Runtime changes and launches also invalidate this catalog. + +Use **Add dataset** for allowed alternatives or named memory datasets, **Remove** to +change the selection, and **Use default sources** to restore parameter-aware defaults. +FigStep allows one of `figstep` or `figstep_pro`; ingredient and compound sources are fixed. +Unavailable limit/source controls are not submitted. Parameter and technique changes +refresh resolved defaults without discarding edits; invalid selections display an error. + +**Max per dataset**, source limits, and **Max total** accept positive integers, `all`, +or `default`/blank. Source caps apply before the total; `all` removes only that cap. +Name-only selection retains existing explicit source settings. Fetch is `never` or +`if_missing`: only a new run may prepare missing registered datasets. Preview and launch +send the same edit-only dataset object, with harm/category filters preserved. Resume uses +the saved request and the original seed selection; a different dataset selection is +rejected. Resume never prepares datasets. #### Core Adversarial Default Override diff --git a/doc/scanner/1_pyrit_scan.ipynb b/doc/scanner/1_pyrit_scan.ipynb index ee28f8b9e7..670f985d40 100644 --- a/doc/scanner/1_pyrit_scan.ipynb +++ b/doc/scanner/1_pyrit_scan.ipynb @@ -1194,9 +1194,17 @@ "on demand, so a full dataset preload is not required:\n", "\n", "```shell\n", - "pyrit_scan run garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1\n", + "pyrit_scan run garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 --dataset-fetch if_missing\n", "```\n", "\n", + "Scanner and shell share `--datasets name[:N|:all]`, `--max-per-dataset`, `--max-total`,\n", + "and `--dataset-fetch never|if_missing`. Source limits apply before the combined total.\n", + "Omitted, empty, or `default` keeps the scenario setting; `all` removes only its own cap.\n", + "`list-scenarios` displays allowed sources and defaults. Ingredient/compound scenarios\n", + "restrict controls to preserve their shaping. `--dataset-filters` works with either form.\n", + "Deprecated `--dataset-names` and `--max-dataset-size` cannot be mixed with the new flags.\n", + "Discovery and estimates never fetch; only a new run prepares absent registered datasets.\n", + "\n", "Custom initialization scripts are loaded by the backend at startup: list them in the\n", "`initialization_scripts` section of the config the server loads (paths are relative to your\n", "working directory, but full paths avoid confusion). Once the server is running with that\n", diff --git a/doc/scanner/1_pyrit_scan.py b/doc/scanner/1_pyrit_scan.py index 0aa18722aa..1e5dce2700 100644 --- a/doc/scanner/1_pyrit_scan.py +++ b/doc/scanner/1_pyrit_scan.py @@ -125,9 +125,17 @@ # on demand, so a full dataset preload is not required: # # ```shell -# pyrit_scan run garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1 +# pyrit_scan run garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 --dataset-fetch if_missing # ``` # +# Scanner and shell share `--datasets name[:N|:all]`, `--max-per-dataset`, `--max-total`, +# and `--dataset-fetch never|if_missing`. Source limits apply before the combined total. +# Omitted, empty, or `default` keeps the scenario setting; `all` removes only its own cap. +# `list-scenarios` displays allowed sources and defaults. Ingredient/compound scenarios +# restrict controls to preserve their shaping. `--dataset-filters` works with either form. +# Deprecated `--dataset-names` and `--max-dataset-size` cannot be mixed with the new flags. +# Discovery and estimates never fetch; only a new run prepares absent registered datasets. +# # Custom initialization scripts are loaded by the backend at startup: list them in the # `initialization_scripts` section of the config the server loads (paths are relative to your # working directory, but full paths avoid confusion). Once the server is running with that diff --git a/doc/scanner/adaptive.ipynb b/doc/scanner/adaptive.ipynb index a4bc5849a4..1f7c72fd64 100644 --- a/doc/scanner/adaptive.ipynb +++ b/doc/scanner/adaptive.ipynb @@ -26,8 +26,8 @@ "pyrit_scan run adaptive.text_adaptive \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names airt_hate \\\n", - " --max-dataset-size 2 \\\n", + " --datasets airt_hate \\\n", + " --max-total 2 \\\n", " --max-attempts-per-objective 2\n", "```" ] diff --git a/doc/scanner/adaptive.py b/doc/scanner/adaptive.py index b229e0e55d..0eda326f0b 100644 --- a/doc/scanner/adaptive.py +++ b/doc/scanner/adaptive.py @@ -25,8 +25,8 @@ # pyrit_scan run adaptive.text_adaptive \ # --initializers target \ # --target openai_chat \ -# --dataset-names airt_hate \ -# --max-dataset-size 2 \ +# --datasets airt_hate \ +# --max-total 2 \ # --max-attempts-per-objective 2 # ``` diff --git a/doc/scanner/airt.ipynb b/doc/scanner/airt.ipynb index 09cf12e75e..04866a9005 100644 --- a/doc/scanner/airt.ipynb +++ b/doc/scanner/airt.ipynb @@ -102,8 +102,8 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques role_play_movie_script \\\n", - " --dataset-names airt_hate \\\n", - " --max-dataset-size 1\n", + " --datasets airt_hate \\\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, role_play_movie_script, many_shot, tap" @@ -255,8 +255,8 @@ "\n", "Each sub-harm escalates a simulated multi-turn conversation toward the objective, then layers the\n", "selected converter techniques on top (natural-language rewrites that preserve emotional framing;\n", - "obfuscation converters like base64/morse are excluded). Datasets are bound to the sub-harms, so\n", - "`--dataset-names` is ignored (`--max-dataset-size` still applies).\n", + "obfuscation converters like base64/morse are excluded). `--datasets` can select a subset of the\n", + "selected sub-harms' datasets; `--max-total` must cover each selected sub-harm.\n", "\n", "**Available techniques:** ALL, DEFAULT, tone, language, persuasion, deterministic, crescendo" ] @@ -413,7 +413,7 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques role_play_movie_script \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** Use `pyrit_scan run airt.cyber --list-scenario-parameters` to inspect\n", @@ -570,8 +570,8 @@ "pyrit_scan run airt.jailbreak \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names harmbench \\\n", - " --max-dataset-size 1\n", + " --datasets harmbench \\\n", + " --max-total 1\n", "```\n", "\n", "**Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both\n", @@ -730,8 +730,8 @@ "pyrit_scan airt.multilingual \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names harmbench \\\n", - " --max-dataset-size 1\n", + " --datasets harmbench \\\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** `prompt_sending` is the default. Every registry technique (`role_play_*`,\n", @@ -870,7 +870,7 @@ "plagiarism detection.\n", "\n", "```bash\n", - "pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-dataset-size 1\n", + "pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, SINGLE_TURN, MULTI_TURN, IP, SENSITIVE_DATA, FirstLetter, Image, RolePlay, Crescendo\n", @@ -1049,7 +1049,7 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques context_compliance \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, ContextCompliance, RolePlay,\n", diff --git a/doc/scanner/airt.py b/doc/scanner/airt.py index cb804f14e7..a7fba16ebc 100644 --- a/doc/scanner/airt.py +++ b/doc/scanner/airt.py @@ -43,8 +43,8 @@ # --initializers target \ # --target openai_chat \ # --techniques role_play_movie_script \ -# --dataset-names airt_hate \ -# --max-dataset-size 1 +# --datasets airt_hate \ +# --max-total 1 # ``` # # **Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, role_play_movie_script, many_shot, tap @@ -84,8 +84,8 @@ # # Each sub-harm escalates a simulated multi-turn conversation toward the objective, then layers the # selected converter techniques on top (natural-language rewrites that preserve emotional framing; -# obfuscation converters like base64/morse are excluded). Datasets are bound to the sub-harms, so -# `--dataset-names` is ignored (`--max-dataset-size` still applies). +# obfuscation converters like base64/morse are excluded). `--datasets` can select a subset of the +# selected sub-harms' datasets; `--max-total` must cover each selected sub-harm. # # **Available techniques:** ALL, DEFAULT, tone, language, persuasion, deterministic, crescendo @@ -124,7 +124,7 @@ # --initializers target \ # --target openai_chat \ # --techniques role_play_movie_script \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # **Available techniques:** Use `pyrit_scan run airt.cyber --list-scenario-parameters` to inspect @@ -169,8 +169,8 @@ # pyrit_scan run airt.jailbreak \ # --initializers target \ # --target openai_chat \ -# --dataset-names harmbench \ -# --max-dataset-size 1 +# --datasets harmbench \ +# --max-total 1 # ``` # # **Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both @@ -213,8 +213,8 @@ # pyrit_scan airt.multilingual \ # --initializers target \ # --target openai_chat \ -# --dataset-names harmbench \ -# --max-dataset-size 1 +# --datasets harmbench \ +# --max-total 1 # ``` # # **Available techniques:** `prompt_sending` is the default. Every registry technique (`role_play_*`, @@ -254,7 +254,7 @@ # plagiarism detection. # # ```bash -# pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-dataset-size 1 +# pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-total 1 # ``` # # **Available techniques:** ALL, SINGLE_TURN, MULTI_TURN, IP, SENSITIVE_DATA, FirstLetter, Image, RolePlay, Crescendo @@ -308,7 +308,7 @@ # --initializers target \ # --target openai_chat \ # --techniques context_compliance \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # **Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, ContextCompliance, RolePlay, diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 6abd678fe9..2ed6d84aeb 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -35,7 +35,7 @@ " --target openai_chat \\\n", " --adversarial-targets adversarial_chat \\\n", " --techniques role_play_video_game \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "Pass multiple `--adversarial-targets` values to compare across models in a single run.\n", diff --git a/doc/scanner/benchmark.py b/doc/scanner/benchmark.py index 8adddffbc1..92588c50db 100644 --- a/doc/scanner/benchmark.py +++ b/doc/scanner/benchmark.py @@ -34,7 +34,7 @@ # --target openai_chat \ # --adversarial-targets adversarial_chat \ # --techniques role_play_video_game \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # Pass multiple `--adversarial-targets` values to compare across models in a single run. diff --git a/doc/scanner/foundry.ipynb b/doc/scanner/foundry.ipynb index 78a5dbeaab..dca0cb3342 100644 --- a/doc/scanner/foundry.ipynb +++ b/doc/scanner/foundry.ipynb @@ -73,7 +73,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-dataset-size 1\n", + "pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-total 1\n", "```\n", "\n", "**Available techniques by difficulty:**\n", diff --git a/doc/scanner/foundry.py b/doc/scanner/foundry.py index bbcc9b8bb6..e73747b8ca 100644 --- a/doc/scanner/foundry.py +++ b/doc/scanner/foundry.py @@ -40,7 +40,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-dataset-size 1 +# pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-total 1 # ``` # # **Available techniques by difficulty:** diff --git a/doc/scanner/garak.ipynb b/doc/scanner/garak.ipynb index 01f9dc85fb..e4be9c7077 100644 --- a/doc/scanner/garak.ipynb +++ b/doc/scanner/garak.ipynb @@ -36,7 +36,7 @@ "PromptInject tests whether a target follows injected instructions hidden in benign tasks.\n", "`--techniques` selects the injection wording, and `--goal-texts` selects the requested output.\n", "The default run uses all five techniques, the three Garak goal texts, and up to 60 attempts.\n", - "`--max-dataset-size` must be at least the number of selected goal texts.\n", + "`--max-total` must be at least the number of selected goal texts.\n", "Sampling selects at least one context for every goal, then fills the remaining budget\n", "without replacement. All techniques use the same sample. Injection wording is applied\n", "through request converters, before any caller-supplied converters.\n", @@ -48,7 +48,7 @@ "pyrit_scan run garak.prompt_inject --target openai_chat \\\n", " --techniques ignore_print \\\n", " --goal-texts \"Custom trigger text\" \\\n", - " --max-dataset-size 16\n", + " --max-total 16\n", "```\n", "\n", "**Techniques:** `ignore_print`, `ignore_say`, `ignore_print_upper`, `nevermind`, and\n", @@ -295,7 +295,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-dataset-size 1\n", + "pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-total 1\n", "```\n", "\n", "**Available techniques** (17 encodings): Base64, Base2048, Base16, Base32, ASCII85, Hex,\n", @@ -453,8 +453,8 @@ "**CLI examples:**\n", "\n", "```bash\n", - "pyrit_scan garak.figstep --target openai_chat --dataset-names figstep --max-dataset-size 1\n", - "pyrit_scan garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1\n", + "pyrit_scan garak.figstep --target openai_chat --datasets figstep --max-total 1\n", + "pyrit_scan garak.figstep --target openai_chat --datasets figstep_pro --max-total 1\n", "```\n", "\n", "> **Note:** The objective target must natively support multi-piece user messages and accept text\n", @@ -616,7 +616,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.web_injection --target openai_chat --techniques xss --max-dataset-size 1\n", + "pyrit_scan run garak.web_injection --target openai_chat --techniques xss\n", "```\n", "\n", "**Available techniques** (8 probes): MarkdownImageExfil, ColabAIDataLeakage,\n", @@ -626,7 +626,7 @@ "**Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended\n", "probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS).\n", "\n", - "**Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include\n", + "**Dataset requirements:** `--datasets` replaces the default datasets, so it must include\n", "every dataset the selected techniques read. Otherwise the scenario stops before running and names\n", "each missing dataset. Datasets already loaded into memory never fill the gap.\n", "\n", @@ -641,7 +641,7 @@ "| StringAssemblyDataExfil | None (built-in seeds) |\n", "\n", "```bash\n", - "pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js\n", + "pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --datasets garak_xss_normal_instructions garak_web_html_js\n", "```" ] }, @@ -982,7 +982,7 @@ "pyrit_scan run garak.api_key --target openai_chat\n", "\n", "# Run only GetKey with a smaller sample.\n", - "pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-dataset-size 2\n", + "pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-total 2\n", "```\n", "\n", "**Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two\n", @@ -1157,7 +1157,7 @@ "\n", "```bash\n", "pyrit_scan run garak.latent_injection --target openai_chat --techniques bare \\\n", - " --families whois --max-dataset-size 1\n", + " --families whois --max-total 1\n", "```\n", "\n", "**Available techniques** (14 separator styles): Bare, Newline, Blockquote, BlockquoteInline,\n", @@ -1349,7 +1349,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-dataset-size 1\n", + "pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-total 1\n", "```\n", "\n", "**Available techniques** (2 probes): `PolicyPuppetry` (wraps the objective in the Dr House\n", @@ -1841,13 +1841,13 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-dataset-size 2\n", + "pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-total 2\n", "```\n", "\n", "> **Note:** The objective target must accept `audio_path` input (i.e. be multimodal). The example\n", "> below uses the registered Azure OpenAI Realtime target; non-audio targets such as the default\n", "> `openai_chat` will error when the audio request is sent. The full dataset holds ~350\n", - "> clips, so a default run samples a small subset to finish quickly — raise `--max-dataset-size`\n", + "> clips, so a default run samples a small subset to finish quickly — raise `--max-total`\n", "> for broader coverage." ] }, @@ -2005,7 +2005,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-dataset-size 2\n", + "pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-total 2\n", "```\n", "\n", "**Available techniques:** `Repeat`, `DEFAULT`, and `ALL` all select the same probe.\n", diff --git a/doc/scanner/garak.py b/doc/scanner/garak.py index 766c8a9f83..1d3426f5dd 100644 --- a/doc/scanner/garak.py +++ b/doc/scanner/garak.py @@ -35,7 +35,7 @@ # PromptInject tests whether a target follows injected instructions hidden in benign tasks. # `--techniques` selects the injection wording, and `--goal-texts` selects the requested output. # The default run uses all five techniques, the three Garak goal texts, and up to 60 attempts. -# `--max-dataset-size` must be at least the number of selected goal texts. +# `--max-total` must be at least the number of selected goal texts. # Sampling selects at least one context for every goal, then fills the remaining budget # without replacement. All techniques use the same sample. Injection wording is applied # through request converters, before any caller-supplied converters. @@ -47,7 +47,7 @@ # pyrit_scan run garak.prompt_inject --target openai_chat \ # --techniques ignore_print \ # --goal-texts "Custom trigger text" \ -# --max-dataset-size 16 +# --max-total 16 # ``` # # **Techniques:** `ignore_print`, `ignore_say`, `ignore_print_upper`, `nevermind`, and @@ -128,7 +128,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-dataset-size 1 +# pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-total 1 # ``` # # **Available techniques** (17 encodings): Base64, Base2048, Base16, Base32, ASCII85, Hex, @@ -176,8 +176,8 @@ # **CLI examples:** # # ```bash -# pyrit_scan garak.figstep --target openai_chat --dataset-names figstep --max-dataset-size 1 -# pyrit_scan garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1 +# pyrit_scan garak.figstep --target openai_chat --datasets figstep --max-total 1 +# pyrit_scan garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 # ``` # # > **Note:** The objective target must natively support multi-piece user messages and accept text @@ -217,7 +217,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.web_injection --target openai_chat --techniques xss --max-dataset-size 1 +# pyrit_scan run garak.web_injection --target openai_chat --techniques xss # ``` # # **Available techniques** (8 probes): MarkdownImageExfil, ColabAIDataLeakage, @@ -227,7 +227,7 @@ # **Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended # probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS). # -# **Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include +# **Dataset requirements:** `--datasets` replaces the default datasets, so it must include # every dataset the selected techniques read. Otherwise the scenario stops before running and names # each missing dataset. Datasets already loaded into memory never fill the gap. # @@ -242,7 +242,7 @@ # | StringAssemblyDataExfil | None (built-in seeds) | # # ```bash -# pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js +# pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --datasets garak_xss_normal_instructions garak_web_html_js # ``` # %% @@ -329,7 +329,7 @@ # pyrit_scan run garak.api_key --target openai_chat # # # Run only GetKey with a smaller sample. -# pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-dataset-size 2 +# pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-total 2 # ``` # # **Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two @@ -377,7 +377,7 @@ # # ```bash # pyrit_scan run garak.latent_injection --target openai_chat --techniques bare \ -# --families whois --max-dataset-size 1 +# --families whois --max-total 1 # ``` # # **Available techniques** (14 separator styles): Bare, Newline, Blockquote, BlockquoteInline, @@ -443,7 +443,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-dataset-size 1 +# pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-total 1 # ``` # # **Available techniques** (2 probes): `PolicyPuppetry` (wraps the objective in the Dr House @@ -571,13 +571,13 @@ # **CLI example:** # # ```bash -# pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-dataset-size 2 +# pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-total 2 # ``` # # > **Note:** The objective target must accept `audio_path` input (i.e. be multimodal). The example # > below uses the registered Azure OpenAI Realtime target; non-audio targets such as the default # > `openai_chat` will error when the audio request is sent. The full dataset holds ~350 -# > clips, so a default run samples a small subset to finish quickly — raise `--max-dataset-size` +# > clips, so a default run samples a small subset to finish quickly — raise `--max-total` # > for broader coverage. # %% @@ -616,7 +616,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-dataset-size 2 +# pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-total 2 # ``` # # **Available techniques:** `Repeat`, `DEFAULT`, and `ALL` all select the same probe. diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts new file mode 100644 index 0000000000..e40aea247c --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts @@ -0,0 +1,13 @@ +import { makeStyles, tokens } from '@fluentui/react-components' + +import { mobileTouchTarget } from '@/styles/touchTargets' + +export const useScenarioDatasetTableStyles = makeStyles({ + root: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalM }, + controls: { display: 'flex', flexWrap: 'wrap', gap: tokens.spacingHorizontalM }, + scroll: { overflowX: 'auto' }, + table: { width: '100%', borderCollapse: 'collapse' }, + cell: { padding: tokens.spacingHorizontalS, textAlign: 'left' }, + input: { maxWidth: '12em' }, + button: { ...mobileTouchTarget }, +}) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx new file mode 100644 index 0000000000..5eea65b2dc --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx @@ -0,0 +1,109 @@ +import { useState } from 'react' + +import { render, screen, within } from '@testing-library/react' +import userEvent from '@testing-library/user-event' +import { FluentProvider, webLightTheme } from '@fluentui/react-components' + +import type { ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings } from '@/types' + +import ScenarioDatasetTable from './ScenarioDatasetTable' + +const DEFAULTS: ScenarioDatasetSettings = { + sources: [{ name: 'stored', max_size: 2, fetch: 'never', configured_provider: false }], + max_per_dataset: 5, max_total: 'all', fetch: 'if_missing', filters: {}, +} +const OPTIONS: ScenarioDatasetOptions = { + mode: 'any', allowed_names: [], per_dataset_limit: true, total_limit: true, supports_overrides: true, filters: true, +} + +interface HarnessProps { + options?: ScenarioDatasetOptions + defaults?: ScenarioDatasetSettings + onChange: jest.Mock +} + +function Harness({ options = OPTIONS, defaults = DEFAULTS, onChange }: HarnessProps) { + const [request, setRequest] = useState({}) + return { setRequest(next); onChange(next) }} + /> +} + +describe('ScenarioDatasetTable', () => { + beforeEach(() => jest.clearAllMocks()) + + it('shows mixed origins, stored counts, and unknown provider-only counts without sending defaults', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + render() + expect(screen.getByText('local, unknown')).toBeInTheDocument() + expect(screen.getByText('8')).toBeInTheDocument() + expect(onChange).not.toHaveBeenCalled() + await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'remote') + await user.click(screen.getByRole('button', { name: 'Add source' })) + const row = screen.getByRole('rowheader', { name: 'remote' }).closest('tr') + expect(row).not.toBeNull() + if (row) { + expect(within(row).getByText('Provider available')).toBeInTheDocument() + expect(within(row).getAllByText('Unknown')).toHaveLength(2) + } + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'stored' }, { name: 'remote' }] }) + }) + + it('edits source caps and policies independently of global limits with keyboard-accessible controls', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + render() + await user.type(screen.getByRole('textbox', { name: 'stored limit' }), 'all') + await user.selectOptions(screen.getByRole('combobox', { name: 'stored fetch' }), 'if_missing') + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '12') + await user.selectOptions(screen.getByRole('combobox', { name: 'Default dataset fetch' }), 'never') + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ + sources: [{ name: 'stored', max_size: 'all', fetch: 'if_missing' }], max_total: 12, fetch: 'never', + }) + await user.clear(screen.getByRole('textbox', { name: 'Max total' })) + expect(onChange.mock.calls.at(-1)?.[0]).not.toHaveProperty('max_total') + screen.getByRole('textbox', { name: 'Max per dataset' }).focus() + await user.keyboard('3') + expect(onChange.mock.calls.at(-1)?.[0]).toHaveProperty('max_per_dataset', 3) + }) + + it('does not send fixed or unsupported inputs and supports one-source replacement', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + const view = render() + expect(screen.queryByRole('button', { name: 'Add source' })).not.toBeInTheDocument() + expect(screen.queryByRole('textbox', { name: 'stored limit' })).not.toBeInTheDocument() + expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() + expect(onChange).not.toHaveBeenCalled() + view.rerender() + await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) + await user.click(screen.getByRole('option', { name: 'remote' })) + await user.click(screen.getByRole('button', { name: 'Add source' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'remote' }] }) + expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() + }) + + it('retains edits and displays invalid selections when parameter-dependent options change', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + const view = render() + await user.type(screen.getByRole('textbox', { name: 'stored limit' }), '1') + view.rerender() + expect(screen.getByRole('alert')).toHaveTextContent('Invalid selection for current parameters: stored') + expect(screen.getByRole('textbox', { name: 'stored limit' })).toHaveValue('1') + await user.click(screen.getByRole('button', { name: 'Use default sources' })) + expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) + }) +}) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx new file mode 100644 index 0000000000..51078acc07 --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx @@ -0,0 +1,194 @@ +import { useState } from 'react' + +import { Button, Combobox, Field, Input, MessageBar, MessageBarBody, Option, Select, Text } from '@fluentui/react-components' + +import type { + DatasetInfo, ScenarioDatasetLimit, ScenarioDatasetOptions, ScenarioDatasetRequest, + ScenarioDatasetSettings, ScenarioDatasetSourceInput, +} from '@/types' + +import { useScenarioDatasetTableStyles } from './ScenarioDatasetTable.styles' + +interface ScenarioDatasetTableProps { + request: ScenarioDatasetRequest + defaults: ScenarioDatasetSettings + effective: ScenarioDatasetSettings + options: ScenarioDatasetOptions + datasets: DatasetInfo[] + disabled: boolean + loading: boolean + metadataError: string | null + onReload: () => void + onChange: (request: ScenarioDatasetRequest) => void + onValidationChange: (error: string | null) => void +} + +interface DatasetLimitFieldProps { + label: string + value: ScenarioDatasetLimit | null | undefined + effective: number | 'all' + disabled: boolean + onChange: (value: ScenarioDatasetLimit) => void + onError: (error: string | null) => void +} + +function DatasetLimitField({ label, value, effective, disabled, onChange, onError }: DatasetLimitFieldProps) { + const [invalidText, setInvalidText] = useState(null) + const styles = useScenarioDatasetTableStyles() + return ( + + { + const text = data.value.trim() + const parsed = Number(text) + if (text !== '' && text !== 'default' && text !== 'all' && (!Number.isSafeInteger(parsed) || parsed < 1)) { + setInvalidText(data.value) + onError(`${label} must be a positive integer, all, or default.`) + return + } + setInvalidText(null) + onError(null) + onChange(text === '' || text === 'default' ? 'default' : text === 'all' ? 'all' : parsed) + }} + /> + + ) +} + +/** Displays effective settings while emitting only user edits, never catalog defaults. */ +export default function ScenarioDatasetTable({ + request, defaults, effective, options, datasets, disabled, loading, metadataError, + onReload, onChange, onValidationChange, +}: ScenarioDatasetTableProps) { + const styles = useScenarioDatasetTableStyles() + const [newName, setNewName] = useState('') + const [errors, setErrors] = useState>({}) + const selected: ScenarioDatasetSourceInput[] = request.sources ?? effective.sources.map((source) => ({ name: source.name })) + const canSelect = options.supports_overrides && options.mode !== 'fixed' + const canEdit = options.supports_overrides && options.source_settings !== false + const candidates = datasets.filter((dataset) => options.mode === 'any' || options.allowed_names.includes(dataset.name)) + const setError = (key: string, error: string | null): void => { + const next = { ...errors } + if (error) next[key] = error + else delete next[key] + setErrors(next) + onValidationChange(Object.values(next)[0] ?? null) + } + const updateSource = (name: string, edits: Partial): void => { + onChange({ ...request, sources: selected.map((source) => source.name === name ? { ...source, ...edits } : source) }) + } + const updateGlobal = (key: 'max_per_dataset' | 'max_total', value: ScenarioDatasetLimit): void => { + const edits = { ...request } + if (value === 'default') delete edits[key] + else edits[key] = value + onChange(edits) + } + const invalidNames = selected.filter((source) => + options.mode === 'subset' && !options.allowed_names.includes(source.name), + ) + return ( +
+ Datasets + + {metadataError && {metadataError}} + {invalidNames.length > 0 && + Invalid selection for current parameters: {invalidNames.map((source) => source.name).join(', ')} + } + {Object.values(errors).map((error) => {error})} + {options.supports_overrides &&
+ {options.per_dataset_limit && setError('perDataset', error)} + onChange={(value) => updateGlobal('max_per_dataset', value)} />} + {options.total_limit && setError('total', error)} + onChange={(value) => updateGlobal('max_total', value)} />} + + + +
} +
+ + {['Name', 'Stored origins', 'Stored examples', 'Limit', 'Fetch policy', 'Selection'].map((label) => + )} + {selected.map((source) => { + const metadata = datasets.find((dataset) => dataset.name === source.name) + const settings = effective.sources.find((item) => item.name === source.name) + ?? defaults.sources.find((item) => item.name === source.name) + const hasProvider = metadata?.provider_available || settings?.configured_provider + return + + + + + + + + })} +
{label}
{source.name}{metadata?.origins?.join(', ') || 'Unknown'} +
{hasProvider ? 'Provider available' : metadata ? 'Memory only' : 'Availability unknown'}
{metadata?.logical_examples ?? 'Unknown'}{canEdit && options.per_dataset_limit + ? setError(source.name, error)} + onChange={(value) => updateSource(source.name, { max_size: value })} /> + : settings?.max_size ?? 'all'}{canEdit + ? : settings?.fetch ?? effective.fetch}{canSelect && }
+
+ {canSelect &&
+ + setNewName(event.target.value)} + onOptionSelect={(_, data) => setNewName(data.optionValue ?? '')}> + {candidates.map((dataset) => )} + + + + +
} +
+ ) +} diff --git a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx index 477c78c8fb..d6f47508c1 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx @@ -6,6 +6,7 @@ import { MemoryRouter, Route, Routes } from 'react-router' import { scenariosApi, targetsApi } from '@/services/api' import type { RegisteredScenario, + ScenarioRunSizeEstimateRequest, ScenarioRunSizeEstimateResponse, TargetInstance, } from '@/types' @@ -13,6 +14,7 @@ import type { import ScenarioDetail from './ScenarioDetail' jest.mock('@/services/api', () => ({ + datasetsApi: { list: jest.fn().mockResolvedValue([]) }, scenariosApi: { estimateRun: jest.fn(), getScenario: jest.fn(), @@ -335,7 +337,7 @@ describe('ScenarioDetail', () => { expect(within(screen.getByRole('region', { name: 'Objective Target' })).queryByRole( 'combobox', { name: 'Adversarial Target' }, )).not.toBeInTheDocument() - expect(fallback.compareDocumentPosition(within(parameters).getByRole('textbox', { name: 'Dataset override' }))) + expect(fallback.compareDocumentPosition(within(parameters).getByRole('table', { name: 'Selected datasets' }))) .toBe(Node.DOCUMENT_POSITION_FOLLOWING) expect(fallback).toHaveAccessibleDescription('The registered target this scenario uses to generate attacks.') expect(fallback).toHaveValue('adversary') @@ -942,10 +944,8 @@ describe('ScenarioDetail', () => { expect(await screen.findByRole('heading', { name: 'Parameters' })).toBeInTheDocument() expect(screen.queryByText('Advanced options')).not.toBeInTheDocument() - expect(screen.getByTestId('max-dataset-size-input')).toHaveValue(8) - expect(screen.getByText( - 'The scenario default is 8. Edit it to override the default.', - )).toBeInTheDocument() + expect(screen.getByRole('textbox', { name: 'Max total' })).toHaveValue('') + expect(screen.getByText('Default/effective: 8. Blank or default retains the scenario setting.')).toBeInTheDocument() const estimate = screen.getByTestId('run-estimate') expect(within(estimate).getByText('Dataset size')).toBeInTheDocument() expect(within(estimate).getByText('8')).toBeInTheDocument() @@ -975,7 +975,8 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByTestId('max-dataset-size-input')).toHaveValue(5) + expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') + expect(screen.getByText('Default/effective: 5. Blank or default retains the scenario setting.')).toBeInTheDocument() expect(await screen.findByText('Up to 10')).toBeInTheDocument() expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -989,7 +990,7 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByRole('spinbutton', { name: 'Max dataset size' })).toHaveValue(null) + expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -1003,9 +1004,8 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('spinbutton', { name: 'Max dataset size' }) - expect(input).toBeDisabled() - expect(input).toHaveValue(null) + await screen.findByRole('table', { name: 'Selected datasets' }) + expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() }) @@ -1019,8 +1019,8 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('spinbutton', { name: 'Max dataset size' }) - expect(input).toHaveValue(5) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, '2') await user.clear(input) expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) @@ -1033,13 +1033,12 @@ describe('ScenarioDetail', () => { it('sends an explicit all total for both estimation and launch', async () => { const user = userEvent.setup() renderDetail('/scanner/foundry.red_team_agent') - const unlimited = await screen.findByRole('checkbox', { name: 'No total dataset limit' }) - await user.click(unlimited) - expect(screen.getByRole('spinbutton', { name: 'Max dataset size' })).toBeDisabled() - await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('max_dataset_size', 'all')) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, 'all') + await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) - expect(mockStartRun.mock.calls[0][0]).toHaveProperty('max_dataset_size', 'all') + expect(mockStartRun.mock.calls[0][0]).toHaveProperty('datasets.max_total', 'all') }) it('restores the scenario default when no-total-limit is unchecked', async () => { @@ -1051,11 +1050,11 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const unlimited = await screen.findByRole('checkbox', { name: 'No total dataset limit' }) - await user.click(unlimited) - await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('max_dataset_size', 'all')) - await user.click(unlimited) - expect(screen.getByRole('spinbutton', { name: 'Max dataset size' })).toHaveValue(5) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, 'all') + await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) + await user.clear(input) + expect(input).toHaveValue('') await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) @@ -1067,17 +1066,22 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') await screen.findByTestId('scenario-target-select') - await user.type(screen.getByTestId('dataset-override-input'), 'ds_a, ds_b') - await user.type(screen.getByTestId('max-dataset-size-input'), '25') + const addDataset = screen.getByRole('combobox', { name: 'Add dataset' }) + await user.type(addDataset, 'ds_a') + await user.click(screen.getByRole('button', { name: 'Add source' })) + await user.type(addDataset, 'ds_b') + await user.click(screen.getByRole('button', { name: 'Add source' })) + await user.click(screen.getByRole('button', { name: 'Remove harmbench' })) + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '25') await user.type(screen.getByTestId('harm-categories-filter-input'), 'cyber, violence') await user.type(screen.getByTestId('data-types-filter-input'), 'text, image_path') await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) const request = mockStartRun.mock.calls[0][0] - expect(request.dataset_names).toEqual(['ds_a', 'ds_b']) - expect(request.max_dataset_size).toBe(25) - expect(request.dataset_filters).toEqual({ + expect(request.datasets.sources).toEqual([{ name: 'ds_a' }, { name: 'ds_b' }]) + expect(request.datasets.max_total).toBe(25) + expect(request.datasets.filters).toEqual({ harm_categories: ['cyber', 'violence'], data_types: ['text', 'image_path'], }) @@ -1086,11 +1090,13 @@ describe('ScenarioDetail', () => { expect.objectContaining({ target_name: 'target-a', techniques: ['default_technique'], - dataset_names: ['ds_a', 'ds_b'], - max_dataset_size: 25, - dataset_filters: { + datasets: { + sources: [{ name: 'ds_a' }, { name: 'ds_b' }], + max_total: 25, + filters: { harm_categories: ['cyber', 'violence'], data_types: ['text', 'image_path'], + }, }, include_baseline: true, }), @@ -1099,17 +1105,56 @@ describe('ScenarioDetail', () => { expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('labels') }) + it('refreshes parameter-aware defaults after invalid source edits without silently dropping them', async () => { + const user = userEvent.setup() + mockGetScenario.mockResolvedValueOnce(makeScenario({ + supported_parameters: [ + { name: 'scope', type_name: 'str', required: false, default: 'all', choices: null, is_list: false }, + ], + })) + mockEstimateRun.mockImplementation(async (_name: string, request: ScenarioRunSizeEstimateRequest) => { + const narrowed = request.scenario_params?.scope === 'one' + if (narrowed && request.datasets?.sources) { + throw { isAxiosError: true, response: { status: 400, data: { detail: 'Dataset is outside the selected scope.' } } } + } + const name = narrowed ? 'remote' : 'harmbench' + const settings = { + sources: [{ name, max_size: 5, fetch: 'never', configured_provider: false }], + max_per_dataset: 5, max_total: 5, fetch: 'never', filters: {}, + } + return { + ...makeEstimate(10), dataset_defaults: settings, effective_datasets: settings, + dataset_options: { + mode: 'subset', allowed_names: [name], per_dataset_limit: true, total_limit: true, + supports_overrides: true, filters: true, + }, + } + }) + renderDetail('/scanner/foundry.red_team_agent') + await screen.findByRole('textbox', { name: 'harmbench limit' }) + await user.type(screen.getByRole('textbox', { name: 'harmbench limit' }), '1') + await user.clear(screen.getByTestId('scenario-param-scope')) + await user.type(screen.getByTestId('scenario-param-scope'), 'one') + await screen.findByText('Invalid selection for current parameters: harmbench') + expect(screen.getByRole('textbox', { name: 'harmbench limit' })).toHaveValue('1') + expect(mockEstimateRun.mock.calls.some((call) => + call[1].scenario_params?.scope === 'one' && call[1].datasets === undefined, + )).toBe(true) + await user.click(screen.getByRole('button', { name: 'Use default sources' })) + await screen.findByRole('rowheader', { name: 'remote' }) + }) + it('rejects a non-positive-integer max dataset size', async () => { const user = userEvent.setup() renderDetail('/scanner/foundry.red_team_agent') await screen.findByTestId('scenario-target-select') - await user.type(screen.getByTestId('max-dataset-size-input'), '0') + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '0') await user.click(screen.getByTestId('launch-scenario-btn')) - expect(await screen.findByRole('alert')).toHaveTextContent( - 'Max dataset size must be a positive integer.', - ) + expect(await screen.findAllByRole('alert')).toEqual(expect.arrayContaining([ + expect.objectContaining({ textContent: expect.stringContaining('Max total must be a positive integer, all, or default.') }), + ])) expect(mockStartRun).not.toHaveBeenCalled() }) diff --git a/frontend/src/components/Scenarios/ScenarioDetail.tsx b/frontend/src/components/Scenarios/ScenarioDetail.tsx index f4378c4f3c..096636a64e 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.tsx @@ -32,6 +32,7 @@ import { Link, useNavigate, useParams } from 'react-router' import MarkdownContent from '@/components/Markdown/MarkdownContent' import TargetSelect from '@/components/Config/TargetSelect' import { useRuntime } from '@/hooks/useRuntime' +import { useScenarioDatasets } from '@/hooks/useScenarioDatasets' import ParameterField from '@/components/Parameters/ParameterField' import SingleStepSpinButton from '@/components/Parameters/SingleStepSpinButton' import { @@ -51,6 +52,10 @@ import type { ScenarioRunSizeEstimateRequest, ScenarioRunEstimateState, ScenarioTechniqueSummary, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetOptions, + ScenarioRunSizeEstimateResponse, TargetInstance, } from '@/types' import { routerPathParamValue, scenarioRunRoutePath } from '@/utils/routeParams' @@ -61,6 +66,7 @@ import { ScenarioRunEstimateDetails } from './ScenarioRunEstimate' import { normalizeScenarioMarkdown } from './scenarioMarkdown' import { mapScenarioRunEstimate } from './scenarioRunEstimateAdapter' import { techniqueSetName } from './scenarioTechniqueSets' +import ScenarioDatasetTable from './ScenarioDatasetTable' /** * Common/opaque parameters every scenario declares via @@ -74,6 +80,7 @@ const COMMON_SCENARIO_PARAMETER_NAMES = new Set([ 'scenario_techniques', 'technique_converters', 'dataset_config', + 'dataset_overrides', 'memory_labels', 'max_concurrency', 'max_retries', @@ -95,11 +102,6 @@ function targetOptionLabel(target: TargetInstance): string { : target.target_registry_name } -function defaultMaxDatasetSize(scenario: RegisteredScenario): string { - const limit = scenario.default_run_size.dataset_limit - return limit.state === 'value' ? String(limit.value) : '' -} - /** Resolves a Fluent `SpinButton` change event to a numeric value, preferring the parsed `value` over the raw `displayValue`. */ function resolveSpinButtonValue(data: { value?: number | null; displayValue?: string }, previous: number): number { if (typeof data.value === 'number') { @@ -234,8 +236,8 @@ interface BuildEstimateRequestInput { techniques: string[] dynamicParameters: Parameter[] scenarioParamValues: Record - datasetOverride: string - maxDatasetSize: string + datasetEdits: ScenarioDatasetRequest + datasetError: string | null harmCategoriesFilter: string dataTypesFilter: string includeBaseline: boolean @@ -291,8 +293,8 @@ function buildEstimateRequest({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, includeBaseline, @@ -310,18 +312,7 @@ function buildEstimateRequest({ scenarioParams = result.parameters } - let maxDatasetSizeValue: number | 'all' | undefined - const trimmedMaxDatasetSize = maxDatasetSize.trim() - if (trimmedMaxDatasetSize === 'all') { - maxDatasetSizeValue = 'all' - } else if (trimmedMaxDatasetSize.length > 0) { - const parsed = Number(trimmedMaxDatasetSize) - if (!Number.isInteger(parsed) || parsed < 1) { - return { ok: false, error: 'Max dataset size must be a positive integer.' } - } - maxDatasetSizeValue = parsed - } - const datasetNames = parseDatasetNames(datasetOverride) + if (datasetError) return { ok: false, error: datasetError } const request: ScenarioRunSizeEstimateRequest = { techniques, include_baseline: includeBaseline, @@ -332,19 +323,15 @@ function buildEstimateRequest({ if (scenario.uses_default_adversarial_target && adversarialTargetName) { request.adversarial_target_name = adversarialTargetName } - if (datasetNames.length > 0) { - request.dataset_names = datasetNames - } - if (maxDatasetSizeValue !== undefined) { - request.max_dataset_size = maxDatasetSizeValue - } const harmCategories = parseDatasetNames(harmCategoriesFilter) const dataTypes = parseDatasetNames(dataTypesFilter) if (harmCategories.length > 0 || dataTypes.length > 0) { - request.dataset_filters = { + request.datasets = { ...datasetEdits, filters: { ...(harmCategories.length > 0 ? { harm_categories: harmCategories } : {}), ...(dataTypes.length > 0 ? { data_types: dataTypes } : {}), - } + } } + } else if (Object.keys(datasetEdits).length > 0) { + request.datasets = datasetEdits } if (scenarioParams) { request.scenario_params = scenarioParams @@ -394,15 +381,7 @@ function buildRunRequest(input: BuildRunRequestInput): BuildRunRequestResult { if (estimateRequest.adversarial_target_name !== undefined) { request.adversarial_target_name = estimateRequest.adversarial_target_name } - if (estimateRequest.dataset_names !== undefined) { - request.dataset_names = estimateRequest.dataset_names - } - if (estimateRequest.max_dataset_size !== undefined) { - request.max_dataset_size = estimateRequest.max_dataset_size - } - if (estimateRequest.dataset_filters !== undefined) { - request.dataset_filters = estimateRequest.dataset_filters - } + if (estimateRequest.datasets !== undefined) request.datasets = estimateRequest.datasets if (estimateRequest.scenario_params !== undefined) { request.scenario_params = estimateRequest.scenario_params } @@ -612,12 +591,21 @@ function ScenarioLaunchForm({ const [baselineChecked, setBaselineChecked] = useState( () => !isBaselineForbidden && scenario.include_baseline_by_default, ) - const [datasetOverride, setDatasetOverride] = useState('') - const configuredDefaultMaxDatasetSize = useMemo( - () => defaultMaxDatasetSize(scenario), - [scenario], - ) - const [maxDatasetSize, setMaxDatasetSize] = useState(configuredDefaultMaxDatasetSize) + const [datasetEdits, setDatasetEdits] = useState({}) + const [datasetError, setDatasetError] = useState(null) + const datasetMetadata = useScenarioDatasets() + const [datasetProjection, setDatasetProjection] = useState(null) + const datasetDefaults: ScenarioDatasetSettings = datasetProjection?.dataset_defaults ?? scenario.dataset_defaults ?? { + sources: scenario.default_datasets.map((name) => ({ name, max_size: 'all', fetch: 'if_missing', configured_provider: false })), + max_per_dataset: 'all', max_total: scenario.default_run_size.dataset_limit.value ?? 'all', + fetch: 'if_missing', filters: {}, + } + const effectiveDatasets = datasetProjection?.effective_datasets ?? datasetDefaults + const datasetOptions: ScenarioDatasetOptions = datasetProjection?.dataset_options ?? scenario.dataset_options ?? { + mode: 'any', allowed_names: [], per_dataset_limit: true, + total_limit: scenario.default_run_size.dataset_limit.state !== 'not_applicable', + supports_overrides: true, filters: true, + } const [harmCategoriesFilter, setHarmCategoriesFilter] = useState('') const [dataTypesFilter, setDataTypesFilter] = useState('') const [maxConcurrency, setMaxConcurrency] = useState(DEFAULT_MAX_CONCURRENCY) @@ -670,15 +658,8 @@ function ScenarioLaunchForm({ [isBaselineForbidden, techniqueOptions], ) const techniques = selectedTechniques - const maxDatasetSizeOverride = maxDatasetSize.trim() - && maxDatasetSize !== configuredDefaultMaxDatasetSize - ? maxDatasetSize - : '' - const datasetSizeLabel = scenario.default_run_size.dataset_limit.state === 'not_applicable' - ? 'Not applicable' - : maxDatasetSize === 'all' - ? 'No total limit' - : maxDatasetSize.trim() || configuredDefaultMaxDatasetSize || 'Scenario default' + const datasetSizeLabel = !datasetOptions.total_limit ? 'Not applicable' + : effectiveDatasets.max_total === 'all' ? 'No total limit' : String(effectiveDatasets.max_total) const estimateResult = useMemo( () => buildEstimateRequest({ scenario, @@ -687,8 +668,8 @@ function ScenarioLaunchForm({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, includeBaseline: isBaselineForbidden ? false : baselineChecked, @@ -696,12 +677,12 @@ function ScenarioLaunchForm({ [ adversarialTargetName, baselineChecked, - datasetOverride, + datasetEdits, + datasetError, dataTypesFilter, dynamicParameters, harmCategoriesFilter, isBaselineForbidden, - maxDatasetSizeOverride, scenario, scenarioParamValues, targetName, @@ -716,8 +697,8 @@ function ScenarioLaunchForm({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, maxConcurrency, @@ -728,14 +709,14 @@ function ScenarioLaunchForm({ [ adversarialTargetName, baselineChecked, - datasetOverride, + datasetEdits, + datasetError, dataTypesFilter, dynamicParameters, harmCategoriesFilter, isBaselineForbidden, labels, maxConcurrency, - maxDatasetSizeOverride, maxRetries, scenario, scenarioParamValues, @@ -774,6 +755,7 @@ function ScenarioLaunchForm({ return } const result = mapScenarioRunEstimate(response, 'request') + setDatasetProjection(response) setEstimateRequestState({ status: 'resolved', requestKey: estimateRequestKey, @@ -790,11 +772,26 @@ function ScenarioLaunchForm({ ) { return } + const error = toApiError(err) setEstimateRequestState({ status: 'error', requestKey: estimateRequestKey, - error: toApiError(err).detail, + error: error.detail, }) + if (estimateRequest.datasets && (error.status === 400 || error.status === 422)) { + const { datasets, ...defaultRequest } = estimateRequest + void datasets + // Keep invalid edits visible while refreshing parameter-aware allowed sources. + void scenariosApi.estimateRun(scenario.scenario_name, defaultRequest, controller.signal) + .then((response) => { + if (controller.signal.aborted || requestSequence !== estimateSequenceRef.current) return + setDatasetProjection((previous) => ({ + ...response, + effective_datasets: previous?.effective_datasets ?? response.effective_datasets, + })) + }) + .catch(() => { /* The original validation error remains authoritative. */ }) + } }) }, ESTIMATE_DEBOUNCE_MS) @@ -908,6 +905,7 @@ function ScenarioLaunchForm({ try { const summary = await scenariosApi.startRun(requestResult.request) + datasetMetadata.reload() setPreviewOpen(false) navigate(scenarioRunRoutePath(summary.scenario_result_id), { state: { scenarioName: scenario.scenario_name }, @@ -935,8 +933,8 @@ function ScenarioLaunchForm({ const displayedEstimateNotes = estimateNotes(estimateState) const displayedEstimate = estimateFromState(estimateState) const selectedTechniqueCount = selectedTechniques.length + (baselineChecked ? 1 : 0) - const previewDatasets = parseDatasetNames(datasetOverride) - const effectiveDatasets = previewDatasets.length > 0 ? previewDatasets : scenario.default_datasets + const previewDatasets = (datasetEdits.sources ?? effectiveDatasets.sources).map((source) => source.name) + const previewDatasetNames = previewDatasets.length > 0 ? previewDatasets : scenario.default_datasets const previewHarmCategories = parseDatasetNames(harmCategoriesFilter) const previewDataTypes = parseDatasetNames(dataTypesFilter) @@ -1130,45 +1128,12 @@ function ScenarioLaunchForm({ disabled={submitting} /> )} - - setDatasetOverride(data.value)} - placeholder={scenario.default_datasets.join(', ') || undefined} - data-testid="dataset-override-input" - /> - - - setMaxDatasetSize(data.value)} - data-testid="max-dataset-size-input" - /> - - setMaxDatasetSize(data.checked === true ? 'all' : configuredDefaultMaxDatasetSize)} + - Per-dataset limits still apply. An empty input uses the scenario default total. setHarmCategoriesFilter(data.value)} data-testid="harm-categories-filter-input" @@ -1189,7 +1154,7 @@ function ScenarioLaunchForm({ setDataTypesFilter(data.value)} data-testid="data-types-filter-input" @@ -1341,8 +1306,8 @@ function ScenarioLaunchForm({
- {effectiveDatasets.length > 0 - ? effectiveDatasets.join(', ') + {previewDatasetNames.length > 0 + ? previewDatasetNames.join(', ') : 'No datasets declared'} diff --git a/frontend/src/hooks/useScenarioDatasets.test.tsx b/frontend/src/hooks/useScenarioDatasets.test.tsx new file mode 100644 index 0000000000..c522c620f8 --- /dev/null +++ b/frontend/src/hooks/useScenarioDatasets.test.tsx @@ -0,0 +1,58 @@ +import { act, renderHook, waitFor } from '@testing-library/react' + +import { datasetsApi } from '@/services/api' +import type { DatasetInfo } from '@/types' + +import { useScenarioDatasets } from './useScenarioDatasets' + +jest.mock('@/services/api', () => ({ datasetsApi: { list: jest.fn() } })) +let generation = 'one' +jest.mock('./useRuntime', () => ({ useRuntime: () => ({ ready: true, generation }) })) +const list = datasetsApi.list as jest.Mock +const NAMED: DatasetInfo = { + name: 'named', selection_key: 'dataset:named:named', loaded: true, provider_available: false, logical_examples: 1, origins: ['local'], +} + +describe('useScenarioDatasets', () => { + beforeEach(() => { jest.clearAllMocks(); generation = 'one' }) + + it('caches metadata for the page, excludes unnamed populations, and explicitly reloads', async () => { + list.mockResolvedValue([NAMED, { ...NAMED, name: '(unnamed)', selection_key: 'dataset:unnamed' }]) + const view = renderHook(() => useScenarioDatasets()) + await waitFor(() => expect(view.result.current.items).toEqual([NAMED])) + view.rerender() + expect(list).toHaveBeenCalledTimes(1) + act(() => { view.result.current.reload() }) + await waitFor(() => expect(list).toHaveBeenCalledTimes(2)) + generation = 'two' + view.rerender() + await waitFor(() => expect(list).toHaveBeenCalledTimes(3)) + }) + + it('ignores old-generation responses and surfaces reload failures without generating datasets', async () => { + let completeOld: (value: DatasetInfo[]) => void = () => {} + list.mockReturnValueOnce(new Promise((resolve) => { completeOld = resolve })) + list.mockResolvedValueOnce([{ ...NAMED, name: 'new' }]) + const view = renderHook(() => useScenarioDatasets()) + generation = 'two' + view.rerender() + await waitFor(() => expect(view.result.current.items[0]?.name).toBe('new')) + await act(async () => { completeOld([NAMED]) }) + expect(view.result.current.items[0]?.name).toBe('new') + list.mockRejectedValueOnce(new Error('offline')) + act(() => { view.result.current.reload() }) + await waitFor(() => expect(view.result.current.error).not.toBeNull()) + }) + + it('hides loaded metadata immediately when the runtime changes', async () => { + list.mockResolvedValueOnce([NAMED]) + const view = renderHook(() => useScenarioDatasets()) + await waitFor(() => expect(view.result.current.items).toEqual([NAMED])) + list.mockReturnValueOnce(new Promise(() => {})) + generation = 'two' + view.rerender() + expect(view.result.current.items).toEqual([]) + expect(view.result.current.loading).toBe(true) + view.unmount() + }) +}) diff --git a/frontend/src/hooks/useScenarioDatasets.ts b/frontend/src/hooks/useScenarioDatasets.ts new file mode 100644 index 0000000000..08e1b7503c --- /dev/null +++ b/frontend/src/hooks/useScenarioDatasets.ts @@ -0,0 +1,50 @@ +import { useCallback, useEffect, useState } from 'react' + +import { datasetsApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { DatasetInfo } from '@/types' + +import { useRuntime } from './useRuntime' + +interface DatasetSnapshot { + key: string + items: DatasetInfo[] + error: string | null +} + +/** Page-scoped metadata cache; no dataset loading or generation occurs here. */ +export function useScenarioDatasets() { + const runtime = useRuntime() + const [snapshot, setSnapshot] = useState(null) + const [reloadCount, setReloadCount] = useState(0) + const reload = useCallback(() => { setReloadCount((count) => count + 1) }, []) + const key = JSON.stringify([runtime.generation, runtime.ready, reloadCount]) + + useEffect(() => { + if (!runtime.ready) return + const controller = new AbortController() + const load = async (): Promise => { + try { + const data = await datasetsApi.list(controller.signal) + if (controller.signal.aborted) return + setSnapshot({ + key, + items: data.filter((item) => item.selection_key !== 'dataset:unnamed'), + error: null, + }) + } catch (err: unknown) { + if (!controller.signal.aborted) setSnapshot({ key, items: [], error: toApiError(err).detail }) + } + } + void load() + return () => { controller.abort() } + }, [key, runtime.ready]) + + const current = runtime.ready && snapshot?.key === key ? snapshot : null + return { + items: current?.items ?? [], + error: current?.error ?? null, + loading: runtime.ready && current === null, + reload, + } +} diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 4e4cde69df..f910c41d28 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -5,6 +5,7 @@ import { toApiError } from './errors' import { compatibility, COMPATIBILITY_HEADER } from './compatibility' import { getGraphScopes } from '../auth/msalConfig' import type { + DatasetInfo, TargetInstance, TargetListResponse, TargetTypeListResponse, @@ -462,6 +463,15 @@ export const labelsApi = { }, } +export const datasetsApi = { + list: async (signal?: AbortSignal): Promise => { + const response = await apiClient.get<{ items: DatasetInfo[] }>('/datasets', { + params: { loaded_only: false }, signal, + }) + return response.data.items + }, +} + export const scenariosApi = { /** * Lists one page of the scenario catalog. Callers that need the full diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 938631dc8f..186af95382 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -770,6 +770,58 @@ export interface ChangeMainConversationResponse { // --- Scenarios --- +export type ScenarioDatasetLimit = number | 'all' | 'default' +export type DatasetFetchPolicy = 'never' | 'if_missing' + +export interface ScenarioDatasetSourceInput { + name: string + max_size?: ScenarioDatasetLimit | null + fetch?: DatasetFetchPolicy | 'default' | null +} + +export interface ScenarioDatasetRequest { + sources?: ScenarioDatasetSourceInput[] | null + max_per_dataset?: ScenarioDatasetLimit | null + max_total?: ScenarioDatasetLimit | null + fetch?: DatasetFetchPolicy | 'default' | null + filters?: Record | null +} + +export interface ScenarioDatasetOptions { + mode: 'any' | 'subset' | 'fixed' + allowed_names: string[] + max_selected?: number | null + per_dataset_limit: boolean + total_limit: boolean + supports_overrides: boolean + filters: boolean + source_settings?: boolean +} + +export interface ScenarioDatasetSourceSettings { + name: string + max_size: number | 'all' + fetch: DatasetFetchPolicy + configured_provider: boolean +} + +export interface ScenarioDatasetSettings { + sources: ScenarioDatasetSourceSettings[] + max_per_dataset: number | 'all' + max_total: number | 'all' + fetch: DatasetFetchPolicy + filters: Record +} + +export interface DatasetInfo { + name: string + selection_key: string + loaded: boolean + provider_available: boolean + logical_examples: number | null + origins: string[] +} + export interface RegisteredScenario { scenario_name: string scenario_type: string @@ -783,6 +835,8 @@ export interface RegisteredScenario { all_techniques: string[] technique_summaries: ScenarioTechniqueSummary[] default_datasets: string[] + dataset_options?: ScenarioDatasetOptions + dataset_defaults?: ScenarioDatasetSettings baseline_policy: 'enabled' | 'disabled' | 'forbidden' include_baseline_by_default: boolean uses_default_adversarial_target: boolean @@ -810,6 +864,7 @@ export interface RunScenarioRequest { dataset_names?: string[] | null max_dataset_size?: number | 'all' | 'default' | '' | null dataset_filters?: Record | null + datasets?: ScenarioDatasetRequest max_concurrency?: number max_retries?: number include_baseline?: boolean | null @@ -855,6 +910,9 @@ export interface ScenarioRunSizeEstimateResponse { status?: 'exact' | 'approximate' | 'conditional' | 'unavailable' dataset_size: ScenarioDatasetSizeEstimate dataset_limit: DatasetLimitInput + dataset_options?: ScenarioDatasetOptions + dataset_defaults?: ScenarioDatasetSettings + effective_datasets?: ScenarioDatasetSettings estimated_attack_count: number | null minimum_attack_count?: number | null maximum_attack_count?: number | null @@ -871,6 +929,7 @@ export interface ScenarioRunSizeEstimateRequest { dataset_names?: string[] | null max_dataset_size?: number | 'all' | 'default' | '' | null dataset_filters?: Record | null + datasets?: ScenarioDatasetRequest include_baseline?: boolean | null scenario_params?: Record | null } diff --git a/pyrit/backend/models/datasets.py b/pyrit/backend/models/datasets.py index f3566c340a..97bf5ef246 100644 --- a/pyrit/backend/models/datasets.py +++ b/pyrit/backend/models/datasets.py @@ -30,6 +30,7 @@ class DatasetInfo(BaseModel): ) loaded: bool = Field(..., description="Whether this dataset currently has seeds in memory") provider_available: bool = Field(..., description="Whether a registered provider can load this dataset") + origins: list[str] = Field(default_factory=list, description="Distinct origins of stored seeds, not providers") logical_examples: int | None = Field( None, description="Number of logical examples represented by stored seed groups" ) diff --git a/pyrit/backend/services/dataset_service.py b/pyrit/backend/services/dataset_service.py index 84920ffb85..c41c6fe6cc 100644 --- a/pyrit/backend/services/dataset_service.py +++ b/pyrit/backend/services/dataset_service.py @@ -83,6 +83,7 @@ async def list_datasets_async(self, *, loaded_only: bool = False) -> DatasetList selection_key=f"dataset:named:{dataset_name}", loaded=summary is not None, provider_available=dataset_name in provider_names, + origins=list(summary.origins) if summary else [], logical_examples=summary.logical_examples if summary else None, seed_pieces=summary.seed_pieces if summary else None, objectives=summary.objectives if summary else None, @@ -100,6 +101,7 @@ async def list_datasets_async(self, *, loaded_only: bool = False) -> DatasetList selection_key="dataset:unnamed", loaded=True, provider_available=False, + origins=list(unnamed_summary.origins), logical_examples=unnamed_summary.logical_examples, seed_pieces=unnamed_summary.seed_pieces, objectives=unnamed_summary.objectives, @@ -290,6 +292,7 @@ def _merge_unnamed_summaries(summaries: Sequence[SeedDatasetSummary]) -> SeedDat modalities=tuple(sorted({modality for summary in unnamed for modality in summary.modalities})), harm_categories=tuple(sorted({category for summary in unnamed for category in summary.harm_categories})), has_unlabeled_harm_categories=any(summary.has_unlabeled_harm_categories for summary in unnamed), + origins=tuple(sorted({origin for summary in unnamed for origin in summary.origins})), ) diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py index b0c6e19256..5933763248 100644 --- a/pyrit/backend/services/scenario_configuration_resolver.py +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -7,13 +7,12 @@ from typing import TYPE_CHECKING, Any -from pyrit.models.dataset_limit import DatasetLimit, normalize_dataset_limit from pyrit.registry import ConverterRegistry, ScenarioRegistry, TargetRegistry -from pyrit.scenario import DatasetSource from pyrit.scenario.core.scenario_target_defaults import validate_default_adversarial_target if TYPE_CHECKING: from pyrit.converter import Converter + from pyrit.models.scenario_dataset import ScenarioDatasetRequest from pyrit.prompt_target import PromptTarget from pyrit.scenario import Scenario @@ -91,9 +90,7 @@ def resolve_configuration( scenario_class: type[Scenario], objective_target: Any | None = None, techniques: list[str] | None = None, - dataset_names: list[str] | None = None, - max_dataset_size: DatasetLimit = "default", - dataset_filters: dict[str, list[str]] | None = None, + datasets: ScenarioDatasetRequest | None = None, include_baseline: bool | None = None, max_concurrency: int | None = None, max_retries: int | None = None, @@ -102,7 +99,7 @@ def resolve_configuration( """ Resolve shared launch and estimate fields into scenario parameters. - Omitted, null, empty, and "default" limits retain the default. "all" removes only the total limit. + Dataset edits pass through unchanged; the scenario applies and validates them. Returns: dict[str, Any]: Values accepted by ``Scenario.set_params_from_args``. @@ -122,11 +119,9 @@ def resolve_configuration( if memory_labels: resolved["memory_labels"] = memory_labels - filters = dataset_filters or {} - total_limit = normalize_dataset_limit(max_dataset_size) - has_total_override = total_limit != "default" - needs_introspection = bool(techniques) or bool(dataset_names) or has_total_override or bool(filters) - if not needs_introspection: + if datasets is not None: + resolved["dataset_overrides"] = datasets + if not techniques: return resolved try: @@ -147,20 +142,6 @@ def resolve_configuration( if technique_converters: resolved["technique_converters"] = technique_converters - if dataset_names or has_total_override or filters: - default_config = introspection_instance._default_dataset_config - config = default_config.with_overrides(filters=filters) - if dataset_names: - existing_sources = {source.name: source for source in config.sources} - config = config.with_overrides( - sources=[ - existing_sources[name] if name in existing_sources else DatasetSource(name=name) - for name in dataset_names - ] - ) - config = config.with_overrides(max_total=total_limit) - resolved["dataset_config"] = config - return resolved @classmethod diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index a187b94e8e..f800cf25ba 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -103,6 +103,7 @@ "target_name", "adversarial_target_name", "techniques", + "datasets", "dataset_names", "max_dataset_size", "dataset_filters", @@ -110,6 +111,8 @@ "max_retries", "include_baseline", ) +# Fields added after the first saved launch records; older records omit them. +_OPTIONAL_LAUNCH_FIELDS = frozenset({"adversarial_target_name", "datasets", "max_dataset_size"}) _SAFE_SCENARIO_PARAMETER_NAMES = frozenset( { @@ -324,13 +327,15 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque raw_request = {"adversarial_target_name": None, **raw_request} if ( not isinstance(raw_request, dict) - or any(name not in raw_request for name in _LAUNCH_REQUEST_FIELDS if name != "max_dataset_size") + or any(name not in raw_request for name in _LAUNCH_REQUEST_FIELDS if name not in _OPTIONAL_LAUNCH_FIELDS) or raw_request["include_baseline"] is None ): raise ScenarioRunConflictError("The saved launch configuration is incomplete; resume was not started.") try: - request = RunScenarioRequest.model_validate( - {name: raw_request[name] for name in _LAUNCH_REQUEST_FIELDS if name in raw_request}, strict=True + # JSON validation accepts saved enum values as strings while keeping strict type checks. + request = RunScenarioRequest.model_validate_json( + json.dumps({name: raw_request[name] for name in _LAUNCH_REQUEST_FIELDS if name in raw_request}), + strict=True, ) except ValidationError as exc: raise ScenarioRunConflictError( @@ -340,7 +345,7 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque raise ScenarioRunConflictError("The saved scenario or target registration name is empty.") identifier = stored.scenario_identifier if (request.techniques is None and identifier.techniques is None) or ( - request.dataset_names is None and identifier.datasets is None + request.dataset_request is None and identifier.datasets is None ): raise ScenarioRunConflictError("The saved scenario identity is missing techniques or datasets.") custom_params = { @@ -348,6 +353,8 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque for name, value in identifier.params.items() if name not in {"version", "techniques", "datasets"} } + # Dataset names are not copied from the identity: the saved request reproduces them, and the + # stored-identity check on initialization rejects any difference. return request.model_copy( update={ "scenario_result_id": str(stored.id), @@ -355,7 +362,6 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque "initializer_args": None, "scenario_params": custom_params, "techniques": request.techniques if request.techniques is not None else identifier.techniques, - "dataset_names": request.dataset_names if request.dataset_names is not None else identifier.datasets, "labels": dict(stored.labels), }, deep=True, @@ -643,9 +649,7 @@ async def _prepare_run_async(self, *, request: RunScenarioRequest) -> _PreparedR 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, + datasets=request.dataset_request, include_baseline=request.include_baseline, max_concurrency=request.max_concurrency, max_retries=request.max_retries, @@ -1231,7 +1235,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ The fully initialized Scenario instance ready for run_async. """ scenario_registry = ScenarioRegistry.get_registry_singleton() - launch_request = {name: getattr(request, name) for name in _LAUNCH_REQUEST_FIELDS} + launch_request = request.model_dump(mode="json", include=set(_LAUNCH_REQUEST_FIELDS)) if launch_request["include_baseline"] is None: scenario_class = scenario_registry.get_class(request.scenario_name) baseline_parameter = next( diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index a1b1a620b7..d6e95d3e2e 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -68,6 +68,8 @@ def _metadata_to_registered_scenario( all_techniques=list(metadata.all_techniques), technique_summaries=list(metadata.technique_summaries), default_datasets=list(metadata.default_datasets), + dataset_options=metadata.dataset_options, + dataset_defaults=metadata.dataset_defaults, supported_parameters=[ parameter.for_external_catalog() for parameter in metadata.supported_parameters @@ -498,9 +500,7 @@ async def _estimate_configured_run_size_async( 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, + datasets=request.dataset_request, include_baseline=request.include_baseline, ) return await self._registry.create_and_estimate_async( diff --git a/pyrit/cli/_cli_args.py b/pyrit/cli/_cli_args.py index 3d92084889..3cd2c7b0ea 100644 --- a/pyrit/cli/_cli_args.py +++ b/pyrit/cli/_cli_args.py @@ -173,6 +173,84 @@ def parse_dataset_limit(value: str) -> int | Literal["all", "default"]: return normalize_dataset_limit(value) +def parse_dataset_source(value: str) -> Any: + """ + Parse a named source with an optional source cap. + + Returns: + DatasetSourceInput: Validated named source. + + Raises: + ValueError: If the token is malformed. + """ + from pyrit.models.scenario_dataset import DatasetSourceInput + + parts = value.split(":") + if len(parts) > 2 or not parts[0].strip(): + raise ValueError("Expected dataset name[:positive-limit|all|default].") + return DatasetSourceInput(name=parts[0], max_size=parse_dataset_limit(parts[1]) if len(parts) == 2 else "default") + + +def parse_dataset_fetch(value: str) -> str: + """ + Validate the global preparation policy without importing dataset providers. + + Returns: + str: Validated preparation policy. + + Raises: + ValueError: If the policy is unsupported. + """ + if value not in {"never", "if_missing"}: + raise ValueError("Dataset fetch must be 'never' or 'if_missing'.") + return value + + +def build_dataset_request_fields(args: dict[str, Any]) -> dict[str, Any]: + """ + Build the same edit-only dataset request for scanner and shell. + + Returns: + dict[str, Any]: Either the new request or legacy dataset fields. + + Raises: + ValueError: If new and deprecated flags are mixed. + """ + from pyrit.models.scenario_dataset import ScenarioDatasetRequest + + new_keys = ("datasets", "max_per_dataset", "max_total", "dataset_fetch") + new_form = any(args.get(key) is not None for key in new_keys) + if new_form and any(args.get(key) is not None for key in ("dataset_names", "max_dataset_size")): + raise ValueError( + "Do not mix --datasets/--max-per-dataset/--max-total/--dataset-fetch with deprecated dataset flags." + ) + filters = collapse_dataset_filters(args["dataset_filters"]) if args.get("dataset_filters") else None + if new_form: + return { + "datasets": ScenarioDatasetRequest( + sources=args.get("datasets"), + max_per_dataset=args.get("max_per_dataset") or "default", + max_total=args.get("max_total") or "default", + fetch=args.get("dataset_fetch"), + filters=filters, + ) + } + fields: dict[str, Any] = {} + for key in ("dataset_names", "max_dataset_size"): + if args.get(key) is not None: + from pyrit.common.deprecation import print_deprecation_message + + print_deprecation_message( + old_item="--" + key.replace("_", "-"), + new_item="--datasets" if key == "dataset_names" else "--max-total", + removed_in="1.4.0", + ) + fields[key] = args[key] + if filters: + fields["dataset_filters"] = filters + return fields + + # --------------------------------------------------------------------------- # Argparse adapter # --------------------------------------------------------------------------- @@ -426,9 +504,11 @@ def _coerce_filter_values(value: str) -> list[str]: "memory_labels": 'Additional labels as JSON string (e.g., \'{"experiment": "test1"}\')', "database": "Database type to use for memory storage", "log_level": "Logging level", - "dataset_names": "List of dataset names to use instead of scenario defaults (e.g., harmbench advbench). " + "dataset_names": "Deprecated; use --datasets. List of dataset names instead of scenario defaults " + "(e.g., harmbench advbench). " "Retains source settings for names already selected by the scenario", - "max_dataset_size": "Total attack-group limit (positive integer, 'default', or 'all' for no total limit). " + "max_dataset_size": "Deprecated; use --max-total. Total attack-group limit " + "(positive integer, 'default', or 'all'). " "Omit to retain the scenario default. Per-dataset limits still apply", "dataset_filters": "Dataset seed filters as KEY=VALUE tokens " "(e.g., harm_categories=cyber data_types=text). Keys filter seeds before sizing. " @@ -662,6 +742,10 @@ class _ArgSpec: _DATASET_NAMES_ARG, _MAX_DATASET_SIZE_ARG, _DATASET_FILTERS_ARG, + _ArgSpec(flags=["--datasets"], result_key="datasets", multi_value=True, parser=parse_dataset_source), + _ArgSpec(flags=["--max-per-dataset"], result_key="max_per_dataset", parser=parse_dataset_limit), + _ArgSpec(flags=["--max-total"], result_key="max_total", parser=parse_dataset_limit), + _ArgSpec(flags=["--dataset-fetch"], result_key="dataset_fetch", parser=parse_dataset_fetch), _TARGET_ARG, ] diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 3cdc8c3291..8b7f015680 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -123,6 +123,17 @@ def print_scenario_list(*, items: list[RegisteredScenario]) -> None: if sc.default_datasets: print(f" Default Datasets ({len(sc.default_datasets)}):") print(_wrap(text=", ".join(sc.default_datasets), indent=" ")) + options = sc.dataset_options + print(f" Dataset selection: {options.mode.value}; overrides: {options.supports_overrides}") + if options.allowed_names: + print(_wrap(text="Options: " + ", ".join(options.allowed_names), indent=" ")) + defaults = sc.dataset_defaults + print( + f" Dataset defaults: per-source={defaults.max_per_dataset}, " + f"total={defaults.max_total}, fetch={defaults.fetch.value}" + ) + for source in defaults.sources: + print(f" {source.name}: limit={source.max_size}, fetch={source.fetch.value}") if sc.supported_parameters: print(" Supported Parameters:") for p in sc.supported_parameters: diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index f912bb5bb5..3f2cb8a7cf 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -27,11 +27,12 @@ ARG_HELP, _parse_initializer_arg, add_results_arguments, + build_dataset_request_fields, build_parameters_from_api, - collapse_dataset_filters, non_negative_int, parse_dataset_filter, parse_dataset_limit, + parse_dataset_source, positive_int, validate_log_level_argparse, ) @@ -313,6 +314,19 @@ def _add_run_arguments(*, parser: ArgumentParser, scenario_params: list[Paramete parser.add_argument("--memory-labels", type=str, help=ARG_HELP["memory_labels"]) parser.add_argument("--dataset-names", type=str, nargs="+", help=ARG_HELP["dataset_names"]) parser.add_argument("--max-dataset-size", type=parse_dataset_limit, help=ARG_HELP["max_dataset_size"]) + parser.add_argument( + "--datasets", + type=parse_dataset_source, + nargs="+", + help="Ordered dataset names with optional :N or :all source caps", + ) + parser.add_argument( + "--max-per-dataset", type=parse_dataset_limit, help="Default source cap; use all for no source cap" + ) + parser.add_argument("--max-total", type=parse_dataset_limit, help="Combined attack-group cap") + parser.add_argument( + "--dataset-fetch", choices=["never", "if_missing"], help="Preparation policy for missing datasets" + ) parser.add_argument( "--dataset-filters", type=parse_dataset_filter, @@ -1004,12 +1018,7 @@ def _build_run_request(*, parsed_args: Namespace, scenario_name: str) -> RunScen kwargs["max_concurrency"] = parsed_args.max_concurrency if parsed_args.max_retries is not None: kwargs["max_retries"] = parsed_args.max_retries - if parsed_args.dataset_names: - kwargs["dataset_names"] = parsed_args.dataset_names - if parsed_args.max_dataset_size is not None: - kwargs["max_dataset_size"] = parsed_args.max_dataset_size - if parsed_args.dataset_filters: - kwargs["dataset_filters"] = collapse_dataset_filters(parsed_args.dataset_filters) + kwargs.update(build_dataset_request_fields(vars(parsed_args))) if parsed_args.memory_labels: kwargs["labels"] = parse_memory_labels(json_string=parsed_args.memory_labels) @@ -1062,7 +1071,11 @@ async def _run_scenario_async( from pyrit.models import ScenarioRunState scenario_name = parsed_args.scenario_name - request = _build_run_request(parsed_args=parsed_args, scenario_name=scenario_name) + try: + request = _build_run_request(parsed_args=parsed_args, scenario_name=scenario_name) + except ValueError as exc: + print(f"Error: {exc}") + return 1 print(f"\nRunning scenario: {scenario_name}") sys.stdout.flush() diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 8110dada98..746f36fb1f 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -434,8 +434,12 @@ def do_run(self, line: str) -> None: --max-concurrency Maximum concurrent operations --max-retries Maximum retry attempts --memory-labels JSON string of labels - --dataset-names ... Override default dataset names - --max-dataset-size Maximum items per dataset + --datasets ... Select sources and optional caps + --max-per-dataset Default per-source cap + --max-total Combined attack-group cap + --dataset-fetch never or if_missing + --dataset-names ... Deprecated source selection + --max-dataset-size Deprecated total cap -- Scenario-declared parameters (see list-scenarios) Notes: @@ -452,8 +456,8 @@ def do_run(self, line: str) -> None: return from pyrit.cli._cli_args import ( + build_dataset_request_fields, build_parameters_from_api, - collapse_dataset_filters, extract_scenario_args, parse_run_arguments, ) @@ -517,12 +521,11 @@ def do_run(self, line: str) -> None: request_kwargs["max_concurrency"] = args["max_concurrency"] if args.get("max_retries") is not None: request_kwargs["max_retries"] = args["max_retries"] - if args.get("dataset_names"): - request_kwargs["dataset_names"] = args["dataset_names"] - if args.get("max_dataset_size") is not None: - request_kwargs["max_dataset_size"] = args["max_dataset_size"] - if args.get("dataset_filters"): - request_kwargs["dataset_filters"] = collapse_dataset_filters(args["dataset_filters"]) + try: + request_kwargs.update(build_dataset_request_fields(args)) + except ValueError as exc: + print(f"Error: {exc}") + return if args.get("memory_labels"): request_kwargs["labels"] = args["memory_labels"] @@ -530,7 +533,11 @@ def do_run(self, line: str) -> None: if scenario_params: request_kwargs["scenario_params"] = scenario_params - request = RunScenarioRequest(**request_kwargs) + try: + request = RunScenarioRequest(**request_kwargs) + except ValueError as exc: + print(f"Error: {exc}") + return # Start run print(f"\nRunning scenario: {scenario_name}") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index ca0997e73d..037e2f84c6 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -4295,6 +4295,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: SeedEntry.dataset_name.label("dataset_name"), SeedEntry.data_type, SeedEntry.harm_categories, + SeedEntry.origin, ) .where(named_condition) .distinct() @@ -4307,6 +4308,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: named_aggregate.c.objectives, named_metadata.c.data_type, named_metadata.c.harm_categories, + named_metadata.c.origin, ).join( named_metadata, named_aggregate.c.dataset_name == named_metadata.c.dataset_name, @@ -4323,7 +4325,10 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: .subquery() ) unnamed_metadata = ( - select(SeedEntry.data_type, SeedEntry.harm_categories).where(unnamed_condition).distinct().subquery() + select(SeedEntry.data_type, SeedEntry.harm_categories, SeedEntry.origin) + .where(unnamed_condition) + .distinct() + .subquery() ) unnamed_statement = select( unnamed_aggregate.c.dataset_name, @@ -4332,6 +4337,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: unnamed_aggregate.c.objectives, unnamed_metadata.c.data_type, unnamed_metadata.c.harm_categories, + unnamed_metadata.c.origin, ).select_from(unnamed_aggregate.join(unnamed_metadata, literal(True))) combined_statement = named_statement.union_all(unnamed_statement) @@ -4351,9 +4357,12 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: "modalities": set(), "harm_categories": set(), "has_unlabeled_harm_categories": False, + "origins": set(), } dataset_order.append(dataset_name) summary = summaries_by_dataset[dataset_name] + if row.origin: + summary["origins"].add(row.origin) if row.data_type: summary["modalities"].add(row.data_type) categories = row.harm_categories or [] @@ -4371,6 +4380,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: modalities=tuple(sorted(summaries_by_dataset[dataset_name]["modalities"])), harm_categories=tuple(sorted(summaries_by_dataset[dataset_name]["harm_categories"])), has_unlabeled_harm_categories=summaries_by_dataset[dataset_name]["has_unlabeled_harm_categories"], + origins=tuple(sorted(summaries_by_dataset[dataset_name]["origins"])), ) for dataset_name in dataset_order ] diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 3cfc87e1aa..3bec545932 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -134,6 +134,15 @@ 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_dataset import ( + DatasetFetchPolicy, + DatasetSelectionMode, + DatasetSourceInput, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetSourceSettings, + ) from pyrit.models.scenario_dataset_size_estimate import ( AllAvailableDatasetSize, BoundedDatasetSize, @@ -264,6 +273,13 @@ "ConversationObservationPayload": "pyrit.models.score", "AllAvailableDatasetSize": "pyrit.models.scenario_dataset_size_estimate", "BoundedDatasetSize": "pyrit.models.scenario_dataset_size_estimate", + "DatasetFetchPolicy": "pyrit.models.scenario_dataset", + "DatasetSelectionMode": "pyrit.models.scenario_dataset", + "DatasetSourceInput": "pyrit.models.scenario_dataset", + "ScenarioDatasetOptions": "pyrit.models.scenario_dataset", + "ScenarioDatasetRequest": "pyrit.models.scenario_dataset", + "ScenarioDatasetSettings": "pyrit.models.scenario_dataset", + "ScenarioDatasetSourceSettings": "pyrit.models.scenario_dataset", "DatasetLimitInput": "pyrit.models.scenario_dataset_size_estimate", "DatasetLimitState": "pyrit.models.scenario_dataset_size_estimate", "IndeterminateDatasetSize": "pyrit.models.scenario_dataset_size_estimate", diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 9a9f269a25..8350f7e1a5 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -16,15 +16,29 @@ from datetime import datetime from enum import Enum from math import prod -from typing import Annotated, Any, Literal - -from pydantic import AliasChoices, BaseModel, BeforeValidator, Field, computed_field, field_validator, model_validator +from typing import Annotated, Any, Literal, Self + +from pydantic import ( + AliasChoices, + BaseModel, + Field, + computed_field, + field_validator, + model_validator, +) -from pyrit.models.dataset_limit import normalize_dataset_limit from pyrit.models.parameter import Parameter from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_LABEL_KEY_LENGTH, MAX_LABEL_VALUE_LENGTH from pyrit.models.results.scenario_result import ScenarioRunState from pyrit.models.retry_event import RetryEvent +from pyrit.models.scenario_dataset import ( + DATASET_FILTERS, + DatasetLimitInputValue, + DatasetSourceInput, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, +) from pyrit.models.scenario_dataset_size_estimate import ( DatasetLimitInput, IndeterminateDatasetSize, @@ -48,28 +62,6 @@ Field(max_length=MAX_ITEMS), ] -# Authoritative set of dataset seed filters exposed over the run request surface. Each entry -# is used verbatim as a ``MemoryInterface.get_seeds`` keyword argument, so a filter key IS the -# get_seeds kwarg. Every exposed filter must be a list-valued (Sequence) get_seeds parameter. -# Adding a filterable field is a one-line change here; the CLI ``--dataset-filters`` help text -# describes these keys, and this request model validates them server-side (covering the GUI too). -# -# Comma-list semantics differ per key because ``get_seeds`` treats each field differently, and -# that behavior lives in ``pyrit.memory`` (this layer cannot import it). As of today -# (see ``MemoryInterface.get_seeds`` / ``_add_list_conditions``): -# - harm_categories -> AND + substring: a seed must be tagged with EVERY value, and each value -# is a substring match (``cyber`` matches ``cyber_harm``). So ``harm_categories=cyber,violence`` -# is an intersection, not a union. -# - data_types -> OR + exact: a seed matches ANY value, compared for exact equality. So -# ``data_types=text,image_path`` is a union. -DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) - - -_RequestDatasetLimit = Annotated[ - Annotated[int, Field(ge=1, strict=True)] | Literal["all", "default"] | None, - BeforeValidator(normalize_dataset_limit), -] - def _validate_dataset_filter_mapping( value: dict[str, list[str]] | None, @@ -98,6 +90,45 @@ class ScenarioRunSizeEstimateStatus(str, Enum): Unavailable = "unavailable" +class _DatasetRequestModel(BaseModel): + """Dataset selection shared by launch and estimate requests.""" + + datasets: ScenarioDatasetRequest | None = Field( + None, description="Dataset sources, limits, fetch policy, and filters. Omit to use scenario defaults." + ) + dataset_names: _RequestNames | None = Field(None, description="Deprecated; use datasets.sources.") + max_dataset_size: DatasetLimitInputValue = Field("default", description="Deprecated; use datasets.max_total.") + dataset_filters: _RequestFilters | None = Field( + None, description="Deprecated; use datasets.filters. Accepted keys: harm_categories, data_types." + ) + + @property + def dataset_request(self) -> ScenarioDatasetRequest | None: + """The dataset edits, converted from the deprecated fields when they are used.""" + if self.datasets is not None or not self._uses_deprecated_dataset_fields(): + return self.datasets + return ScenarioDatasetRequest( + sources=[DatasetSourceInput(name=name) for name in self.dataset_names] if self.dataset_names else None, + max_total=self.max_dataset_size, + filters=self.dataset_filters or None, + ) + + @field_validator("dataset_filters") + @classmethod + def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: + return _validate_dataset_filter_mapping(value) + + @model_validator(mode="after") + def _validate_dataset_form(self) -> Self: + if self.datasets is not None and self._uses_deprecated_dataset_fields(): + raise ValueError("Use either 'datasets' or the deprecated dataset fields, not both.") + self.dataset_request # noqa: B018 - converts deprecated fields so invalid values fail here + return self + + def _uses_deprecated_dataset_fields(self) -> bool: + return bool(self.dataset_names) or self.max_dataset_size != "default" or bool(self.dataset_filters) + + class ScenarioRunSizeEstimateCondition(str, Enum): """Reason an estimate remains conditional until launch.""" @@ -194,6 +225,9 @@ class ScenarioRunSizeEstimate(BaseModel): components: list[ScenarioRunSizeComponent] = Field(default_factory=list) datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) dataset_size: ScenarioDatasetSizeEstimate = Field(default_factory=IndeterminateDatasetSize) + dataset_options: ScenarioDatasetOptions = Field(default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) + effective_datasets: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) dataset_limit: DatasetLimitInput = Field(default_factory=DatasetLimitInput) effective_parameters: dict[str, bool | int | float | str | list[str]] = Field( default_factory=dict, @@ -365,6 +399,8 @@ class RegisteredScenario(BaseModel): description="Descriptions and tags for the available concrete techniques", ) default_datasets: list[str] = Field(..., description="Default dataset names used by the scenario") + dataset_options: ScenarioDatasetOptions = Field(default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) baseline_policy: Literal["enabled", "disabled", "forbidden"] = Field( "enabled", description="Whether baseline execution is enabled, disabled, or forbidden" ) @@ -381,7 +417,7 @@ class RegisteredScenario(BaseModel): ) -class ScenarioRunSizeEstimateRequest(BaseModel): +class ScenarioRunSizeEstimateRequest(_DatasetRequestModel): """Request-specific scenario run-size configuration.""" adversarial_target_name: _RequestName | None = Field( @@ -396,19 +432,6 @@ class ScenarioRunSizeEstimateRequest(BaseModel): techniques: _RequestTechniques | None = Field( None, description="Technique names to estimate (uses scenario default if omitted)" ) - dataset_names: _RequestNames | None = Field( - None, description="Dataset names to estimate (uses scenario default if omitted)" - ) - max_dataset_size: _RequestDatasetLimit = Field( - "default", - description="Total selected logical seed-group limit. Omitted, null, empty, or 'default': scenario default. " - "'all': no total limit. " - "Per-dataset limits remain in effect.", - ) - dataset_filters: _RequestFilters | 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", @@ -418,19 +441,8 @@ class ScenarioRunSizeEstimateRequest(BaseModel): 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): +class RunScenarioRequest(_DatasetRequestModel): """Request body for starting a scenario run.""" scenario_name: _RequestName = Field(..., description="Scenario name (e.g., 'foundry.red_team_agent')") @@ -446,21 +458,6 @@ class RunScenarioRequest(BaseModel): techniques: _RequestTechniques | None = Field( None, description="Technique names to use (uses scenario default if omitted)" ) - dataset_names: _RequestNames | None = Field( - None, description="Dataset names to use (uses scenario default if omitted)" - ) - max_dataset_size: _RequestDatasetLimit = Field( - "default", - description="Total selected logical seed-group limit. Omitted, null, empty, or 'default': scenario default. " - "'all': no total limit. " - "Per-dataset limits remain in effect.", - ) - dataset_filters: _RequestFilters | None = Field( - None, - description=( - "Dataset seed filters keyed by field, applied before sampling. Accepted keys: harm_categories, data_types." - ), - ) 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( @@ -485,17 +482,6 @@ class RunScenarioRequest(BaseModel): "If provided, the scenario will resume from prior progress instead of starting fresh.", ) - @field_validator("dataset_filters") - @classmethod - def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: - """ - Reject any dataset-filter key not in the exposed ``DATASET_FILTERS`` allow-list. - - Returns: - dict[str, list[str]] | None: The validated filters, unchanged. - """ - return _validate_dataset_filter_mapping(value) - class AttackErrorSummary(BaseModel): """A single errored attack result surfaced in a run summary.""" diff --git a/pyrit/models/scenario_dataset.py b/pyrit/models/scenario_dataset.py new file mode 100644 index 0000000000..45d93ac303 --- /dev/null +++ b/pyrit/models/scenario_dataset.py @@ -0,0 +1,130 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Shared scenario dataset inputs and selection capabilities.""" + +from enum import Enum +from typing import Annotated, Literal + +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, field_validator + +from pyrit.models.dataset_limit import normalize_dataset_limit +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS + +# Dataset seed filters exposed to clients. Each key is a list-valued ``MemoryInterface.get_seeds`` +# keyword argument. Comma-list semantics differ per key because ``get_seeds`` treats each field +# differently (see ``MemoryInterface.get_seeds`` / ``_add_list_conditions``): +# - harm_categories -> AND + substring: a seed must be tagged with EVERY value, and each value +# is a substring match (``cyber`` matches ``cyber_harm``). +# - data_types -> OR + exact: a seed matches ANY value, compared for exact equality. +DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) +DatasetLimitInputValue = Annotated[ + Annotated[int, Field(ge=1, strict=True)] | Literal["all", "default"], + BeforeValidator(normalize_dataset_limit), +] + + +class DatasetFetchPolicy(str, Enum): + """When scenario preparation may load a dataset.""" + + NEVER = "never" + IF_MISSING = "if_missing" + + +def _normalize_fetch(value: object) -> object: + if isinstance(value, str): + value = value.strip().lower() + return None if value in (None, "", "default") else value + + +class DatasetSelectionMode(str, Enum): + """The permitted named-dataset selection.""" + + ANY = "any" + SUBSET = "subset" + FIXED = "fixed" + + +class ScenarioDatasetOptions(BaseModel): + """Scenario-owned capabilities for dataset overrides.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + mode: DatasetSelectionMode = DatasetSelectionMode.ANY + allowed_names: list[str] = Field(default_factory=list) + max_selected: int | None = Field(None, ge=1) + per_dataset_limit: bool = True + total_limit: bool = True + supports_overrides: bool = True + filters: bool = True + source_settings: bool = True + + +class DatasetSourceInput(BaseModel): + """A named source and optional limit and preparation overrides.""" + + model_config = ConfigDict(extra="forbid") + + name: str = Field(min_length=1, max_length=MAX_IDENTIFIER_LENGTH, strict=True) + max_size: DatasetLimitInputValue = "default" + fetch: Annotated[DatasetFetchPolicy | None, BeforeValidator(_normalize_fetch)] = None + + @field_validator("name") + @classmethod + def _validate_name(cls, value: str) -> str: + if not value.strip(): + raise ValueError("Dataset names must be non-empty.") + return value + + +class ScenarioDatasetRequest(BaseModel): + """Serializable dataset overrides shared by all application surfaces.""" + + model_config = ConfigDict(extra="forbid") + + sources: list[DatasetSourceInput] | None = Field(None, min_length=1, max_length=MAX_ITEMS) + max_per_dataset: DatasetLimitInputValue = "default" + max_total: DatasetLimitInputValue = "default" + fetch: Annotated[DatasetFetchPolicy | None, BeforeValidator(_normalize_fetch)] = None + filters: dict[str, list[str]] | None = Field(None, max_length=MAX_ITEMS) + + @field_validator("sources") + @classmethod + def _validate_sources(cls, value: list[DatasetSourceInput] | None) -> list[DatasetSourceInput] | None: + names = [source.name for source in value or []] + if len(names) != len(set(names)): + raise ValueError("Duplicate dataset source names are not allowed.") + return value + + @field_validator("filters") + @classmethod + def _validate_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: + for key, values in (value or {}).items(): + if key not in DATASET_FILTERS: + raise ValueError(f"Unknown dataset filter '{key}'. Allowed: {', '.join(sorted(DATASET_FILTERS))}.") + if len(values) > MAX_ITEMS or any(len(item) > MAX_IDENTIFIER_LENGTH for item in values): + raise ValueError("Dataset filter values exceed the request limits.") + return value + + +class ScenarioDatasetSourceSettings(BaseModel): + """Effective source settings, not editable request fields.""" + + model_config = ConfigDict(extra="forbid") + + name: str + max_size: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] + fetch: DatasetFetchPolicy + configured_provider: bool = False + + +class ScenarioDatasetSettings(BaseModel): + """Resolved defaults or effective settings for display and persistence.""" + + model_config = ConfigDict(extra="forbid") + + sources: list[ScenarioDatasetSourceSettings] = Field(default_factory=list) + max_per_dataset: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] = "all" + max_total: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] = "all" + fetch: DatasetFetchPolicy = DatasetFetchPolicy.IF_MISSING + filters: dict[str, list[str]] = Field(default_factory=dict) diff --git a/pyrit/models/seeds/seed_dataset_summary.py b/pyrit/models/seeds/seed_dataset_summary.py index bcfb8e85a1..0b5b882303 100644 --- a/pyrit/models/seeds/seed_dataset_summary.py +++ b/pyrit/models/seeds/seed_dataset_summary.py @@ -15,3 +15,4 @@ class SeedDatasetSummary: modalities: tuple[str, ...] harm_categories: tuple[str, ...] has_unlabeled_harm_categories: bool + origins: tuple[str, ...] = () diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index a11afd7828..8a15dcc360 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -18,7 +18,13 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal -from pyrit.models import ScenarioRunSizeEstimate, ScenarioTechniqueSummary, class_name_to_snake_case +from pyrit.models import ( + ScenarioDatasetOptions, + ScenarioDatasetSettings, + ScenarioRunSizeEstimate, + ScenarioTechniqueSummary, + 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 @@ -66,6 +72,11 @@ class ScenarioMetadata(RegistryMetadata): # Default dataset names used by this scenario. default_datasets: tuple[str, ...] = field(kw_only=True) + # Dataset changes the scenario accepts, and its default dataset settings. Both are static: + # they never depend on memory contents or on a run-size estimate succeeding. + dataset_options: ScenarioDatasetOptions = field(kw_only=True, default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = field(kw_only=True, default_factory=ScenarioDatasetSettings) + # Scenario-declared custom parameters. supported_parameters: tuple[Parameter, ...] = field(kw_only=True, default=()) @@ -210,6 +221,8 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: aggregate_techniques=aggregate_techniques, aggregate_technique_expansions=aggregate_technique_expansions, default_datasets=default_datasets, + dataset_options=instance.get_dataset_options(), + dataset_defaults=instance.get_default_dataset_settings(), supported_parameters=supported_parameters, baseline_policy=instance.BASELINE_ATTACK_POLICY.value, include_baseline_by_default=instance.BASELINE_ATTACK_POLICY.value == "enabled", diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 16e6c246b1..f31fe6449a 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -47,6 +47,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import DatasetLimit, ResolvedDatasetLimit, normalize_dataset_limit +from pyrit.models.scenario_dataset import DatasetFetchPolicy if TYPE_CHECKING: from collections.abc import Callable, Sequence @@ -66,13 +67,6 @@ class _Unset(Enum): VALUE = "unset" -class DatasetFetchPolicy(str, Enum): - """When preparation may load a registered dataset into memory.""" - - NEVER = "never" - IF_MISSING = "if_missing" - - def _resolve_limit(*, name: str, value: DatasetLimit | _Unset, default: ResolvedDatasetLimit) -> ResolvedDatasetLimit: try: limit = normalize_dataset_limit(None if isinstance(value, _Unset) else value) @@ -997,6 +991,11 @@ def __init__( ) self._configurations = list(configurations) + @property + def configurations(self) -> tuple[DatasetAttackConfiguration, ...]: + """The child configurations, in order.""" + return tuple(self._configurations) + @property def has_sampling_limits(self) -> bool: """Whether this compound or any child can select a subset.""" diff --git a/pyrit/scenario/core/dataset_overrides.py b/pyrit/scenario/core/dataset_overrides.py new file mode 100644 index 0000000000..69bd297375 --- /dev/null +++ b/pyrit/scenario/core/dataset_overrides.py @@ -0,0 +1,122 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Scenario-owned override validation and lightweight dataset projections.""" + +from dataclasses import replace +from typing import Any + +from pyrit.models.scenario_dataset import ( + DatasetSelectionMode, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetSourceSettings, +) +from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + DatasetConstraintError, + DatasetSource, + DatasetSourceKind, +) + + +def apply_dataset_overrides( + *, + config: DatasetAttackConfiguration, + request: ScenarioDatasetRequest, + options: ScenarioDatasetOptions, +) -> DatasetAttackConfiguration: + """ + Apply named-source edits without rebuilding custom configurations. + + Returns: + DatasetAttackConfiguration: Shallow copy preserving live object identity. + + Raises: + DatasetConstraintError: If a requested edit is unsupported. + """ + if not options.supports_overrides and request != ScenarioDatasetRequest(): + raise DatasetConstraintError("This scenario does not support dataset overrides.") + if not options.per_dataset_limit and ( + request.max_per_dataset != "default" or any(source.max_size != "default" for source in request.sources or []) + ): + raise DatasetConstraintError("This scenario does not support per-dataset limits.") + if not options.total_limit and request.max_total != "default": + raise DatasetConstraintError("This scenario does not support a total dataset limit.") + if not options.filters and request.filters: + raise DatasetConstraintError("This scenario does not support dataset filters.") + kwargs: dict[str, Any] = {} + if request.sources is not None: + if isinstance(config, CompoundDatasetAttackConfiguration): + raise DatasetConstraintError("Compound sources are fixed; source edits are ambiguous.") + existing = {source.name: source for source in config.sources} + selected = [] + for source in request.sources: + base = existing.get(source.name, DatasetSource(name=source.name)) + edits = { + **({"max_size": source.max_size} if source.max_size != "default" else {}), + **({"fetch": source.fetch} if source.fetch is not None else {}), + } + selected.append(replace(base, **edits) if edits else base) + kwargs["sources"] = selected + if request.fetch is not None: + kwargs["fetch"] = request.fetch + effective = config.with_overrides( + **kwargs, + max_per_dataset=request.max_per_dataset, + max_total=request.max_total, + filters=request.filters, + ) + validate_dataset_selection(config=effective, defaults=config, options=options) + return effective + + +def validate_dataset_selection( + *, + config: DatasetAttackConfiguration, + defaults: DatasetAttackConfiguration, + options: ScenarioDatasetOptions, +) -> None: + """ + Validate names before preparation while preserving inline Python inputs. + + Raises: + DatasetConstraintError: If the source selection violates scenario capabilities. + """ + if config.source_kind is DatasetSourceKind.INLINE: + return + names = config.dataset_names + if options.mode is DatasetSelectionMode.FIXED and names != defaults.dataset_names: + raise DatasetConstraintError(f"Dataset sources are fixed: {', '.join(defaults.dataset_names)}.") + if options.mode is DatasetSelectionMode.SUBSET and (not names or set(names) - set(options.allowed_names)): + raise DatasetConstraintError(f"Select datasets from: {', '.join(options.allowed_names)}.") + if options.max_selected is not None and len(names) > options.max_selected: + raise DatasetConstraintError(f"Select at most {options.max_selected} dataset(s).") + + +def project_dataset_settings(config: DatasetAttackConfiguration) -> ScenarioDatasetSettings: + """ + Project configuration settings without memory reads or live wire objects. + + Returns: + ScenarioDatasetSettings: Effective settings for response data. + """ + children = config.configurations if isinstance(config, CompoundDatasetAttackConfiguration) else (config,) + return ScenarioDatasetSettings( + sources=[ + ScenarioDatasetSourceSettings( + name=source.name, + max_size=child.source_limit(source.name), + fetch=source.fetch or child.fetch, + configured_provider=getattr(source, "provider", None) is not None, + ) + for child in children + for source in child.sources + ], + max_per_dataset=config.max_per_dataset, + max_total=config.max_total, + fetch=children[0].fetch, + filters=config.filters, + ) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 2b2ed3fb73..b49f249a42 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -53,6 +53,7 @@ config_hash, ) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.registry import ScorerRegistry @@ -61,6 +62,7 @@ from pyrit.scenario.core.dataset_configuration import ( DatasetAttackConfiguration, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides, project_dataset_settings 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 @@ -348,6 +350,11 @@ def _common_scenario_parameters(cls) -> list[Parameter]: description="Techniques to execute; defaults to the scenario's default aggregate when omitted.", opaque=True, ), + Parameter( + name="dataset_overrides", + description="Serializable dataset edits applied on top of the dataset configuration.", + opaque=True, + ), Parameter( name="technique_converters", description="Mapping of concrete technique name to extra request converters to append.", @@ -636,6 +643,22 @@ async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: @final async def get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: + """ + Estimate without preparation and expose resolved defaults separately from edits. + + Returns: + ScenarioRunSizeEstimate: Estimate with dataset defaults and effective settings. + """ + estimate = await self._get_run_size_estimate_async(target_is_configured=target_is_configured) + return estimate.model_copy( + update={ + "dataset_options": self.get_dataset_options(), + "dataset_defaults": project_dataset_settings(self._resolved_default_dataset_config), + "effective_datasets": project_dataset_settings(self._dataset_config), + } + ) + + async def _get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: """ Estimate the currently configured run without creating or persisting it. @@ -818,7 +841,6 @@ def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> N dataset_config = params.get("dataset_config") self._dataset_config_provided = dataset_config is not None - self._dataset_config = (dataset_config if dataset_config else self._default_dataset_config).with_overrides() self._max_concurrency = params.get("max_concurrency", 4) self._max_retries = params.get("max_retries", 0) self._memory_labels = params.get("memory_labels") or {} @@ -839,8 +861,39 @@ def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> N scenario_techniques=params.get("scenario_techniques") ) self._technique_converters = params.get("technique_converters") or {} + self._resolved_default_dataset_config = self._resolve_default_dataset_config().with_overrides() + self._dataset_config = ( + dataset_config if dataset_config is not None else self._resolved_default_dataset_config + ).with_overrides() + options = self.get_dataset_options() + raw_overrides = params.get("dataset_overrides") + if raw_overrides is not None: + self._dataset_config = apply_dataset_overrides( + config=self._dataset_config, + request=ScenarioDatasetRequest.model_validate(raw_overrides), + options=options, + ) self._validate_runtime_configuration() + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + """ + Resolve parameter-aware defaults before any override or preparation. + + Returns: + DatasetAttackConfiguration: Scenario-owned resolved default configuration. + """ + return self._default_dataset_config + + DATASET_OPTIONS: ClassVar[ScenarioDatasetOptions] = ScenarioDatasetOptions() + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """Return the scenario's named-dataset capabilities for the current parameters.""" + return self.DATASET_OPTIONS + + def get_default_dataset_settings(self) -> ScenarioDatasetSettings: + """Return the default dataset settings for the current parameters, without reading memory.""" + return project_dataset_settings(self._resolve_default_dataset_config()) + def _validate_runtime_configuration(self) -> None: """ Check resolved parameters for preview and launch without reading datasets. diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 1f54d65918..0a006f7f95 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -40,11 +40,13 @@ SeedPrompt, ) from pyrit.models.parameter import Parameter +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, DatasetConstraintError, DatasetSource, @@ -478,25 +480,43 @@ def _selected_sub_harms(self) -> list[_SubHarm]: ) return [_SUB_HARMS_BY_NAME[name]] - def _validate_runtime_configuration(self) -> None: - config = self._dataset_config - allowed = {harm.dataset_name for harm in _SUB_HARMS} - if set(config.dataset_names) - allowed: - raise DatasetConstraintError("Psychosocial sources must match its sub-harm datasets.") + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + config = self._default_dataset_config by_name = {source.name: source for source in config.sources} - self._dataset_config = config.with_overrides( + return config.with_overrides( sources=[ by_name.get(harm.dataset_name, DatasetSource(name=harm.dataset_name)) for harm in self._selected_sub_harms() ] ) + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """ + Describe named datasets allowed for the selected sub-harm parameter. + + Returns: + ScenarioDatasetOptions: Parameter-aware source selection capabilities. + """ + return ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, + allowed_names=[harm.dataset_name for harm in self._selected_sub_harms()], + ) + + def _effective_sub_harms(self) -> list[_SubHarm]: + return [harm for harm in self._selected_sub_harms() if harm.dataset_name in self._dataset_config.dataset_names] + + def _validate_runtime_configuration(self) -> None: + if isinstance(self._dataset_config, CompoundDatasetAttackConfiguration): + raise DatasetConstraintError("Psychosocial requires flat sources, not compound children.") + if set(self._dataset_config.dataset_names) - set(self.get_dataset_options().allowed_names): + raise DatasetConstraintError("Psychosocial sources must match the selected sub-harm datasets.") super()._validate_runtime_configuration() cap = self._dataset_config.max_total - if cap != "all" and cap < len(self._selected_sub_harms()): + if cap != "all" and cap < len(self._effective_sub_harms()): raise DatasetConstraintError( f"Psychosocial max_total ({cap}) must cover every selected sub-harm " - f"({len(self._selected_sub_harms())}). Increase the total limit to at least " - f"{len(self._selected_sub_harms())}, or select one sub-harm." + f"({len(self._effective_sub_harms())}). Increase the total limit to at least " + f"{len(self._effective_sub_harms())}, or select one sub-harm." ) async def _resolve_seed_groups_by_dataset_async( @@ -515,7 +535,7 @@ async def _resolve_seed_groups_by_dataset_async( harm.dataset_name: self._dataset_config._sample_source_groups( name=harm.dataset_name, groups=populations.get(harm.dataset_name, []) ) - for harm in self._selected_sub_harms() + for harm in self._effective_sub_harms() } return sample_with_coverage( groups_by_dataset=selected, @@ -591,7 +611,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list # Resolved lazily so a no-arg ``Psychosocial()`` works for registry metadata introspection. adversarial_chat = self._adversarial_chat or get_default_adversarial_target() - sub_harms = self._selected_sub_harms() + sub_harms = self._effective_sub_harms() techniques = [t for t in context.scenario_techniques if isinstance(t, PsychosocialTechnique)] max_turns = int(self.params.get("max_turns", 5)) seed_groups_by_dataset = context.seed_groups_by_dataset diff --git a/pyrit/scenario/scenarios/garak/api_key.py b/pyrit/scenario/scenarios/garak/api_key.py index 648de88d0a..95ff038900 100644 --- a/pyrit/scenario/scenarios/garak/api_key.py +++ b/pyrit/scenario/scenarios/garak/api_key.py @@ -25,6 +25,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -210,6 +211,7 @@ class ApiKey(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden @classmethod diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index d777cc415d..e70188bd7f 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -34,6 +34,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -168,6 +169,9 @@ class Encoding(Scenario): """ VERSION: int = 2 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, source_settings=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False @apply_defaults diff --git a/pyrit/scenario/scenarios/garak/exploitation.py b/pyrit/scenario/scenarios/garak/exploitation.py index c6d81968b6..4444089eec 100644 --- a/pyrit/scenario/scenarios/garak/exploitation.py +++ b/pyrit/scenario/scenarios/garak/exploitation.py @@ -28,6 +28,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -189,6 +190,9 @@ class Exploitation(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False, supports_overrides=False + ) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden @classmethod @@ -222,7 +226,7 @@ def additional_parameters(cls) -> list[Parameter]: @classmethod def supported_parameters(cls) -> list[Parameter]: """Return parameters, excluding overrides of the fixed reference corpus.""" - unsupported = {"dataset_config"} + unsupported = {"dataset_config", "dataset_overrides"} return [parameter for parameter in super().supported_parameters() if parameter.name not in unsupported] @apply_defaults diff --git a/pyrit/scenario/scenarios/garak/figstep.py b/pyrit/scenario/scenarios/garak/figstep.py index 3dea0f7a3b..a9d8109c2e 100644 --- a/pyrit/scenario/scenarios/garak/figstep.py +++ b/pyrit/scenario/scenarios/garak/figstep.py @@ -11,6 +11,7 @@ from pyrit.common.path import SCORER_SEED_PROMPT_PATH from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack from pyrit.models import AttackSeedGroup, PromptDataType +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.prompt_target import CapabilityName, TargetRequirements from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -68,6 +69,9 @@ class FigStep(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, allowed_names=["figstep", "figstep_pro"], max_selected=1 + ) TARGET_REQUIREMENTS: ClassVar[TargetRequirements] = TargetRequirements( native_required=frozenset({CapabilityName.MULTI_MESSAGE_PIECES}), required_input_modalities=_FIGSTEP_INPUT_MODALITIES, diff --git a/pyrit/scenario/scenarios/garak/latent_injection.py b/pyrit/scenario/scenarios/garak/latent_injection.py index 9615946c51..e62ca764a5 100644 --- a/pyrit/scenario/scenarios/garak/latent_injection.py +++ b/pyrit/scenario/scenarios/garak/latent_injection.py @@ -18,6 +18,7 @@ from pyrit.converter import SearchReplaceConverter from pyrit.executor.attack import AttackConverterConfig, AttackScoringConfig, PromptSendingAttack from pyrit.models import AttackSeedGroup, Parameter, Seed, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -314,6 +315,7 @@ class LatentInjection(Scenario): """Apply each selected separator to the same sample of document injections.""" VERSION: int = 2 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden SEPARATORS: ClassVar[dict[str, tuple[str, str]]] = { "bare": ("", ""), diff --git a/pyrit/scenario/scenarios/garak/package_hallucination.py b/pyrit/scenario/scenarios/garak/package_hallucination.py index e3a502e286..aa821b54ff 100644 --- a/pyrit/scenario/scenarios/garak/package_hallucination.py +++ b/pyrit/scenario/scenarios/garak/package_hallucination.py @@ -23,6 +23,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetSource @@ -139,6 +140,9 @@ class PackageHallucination(Scenario): """ VERSION: int = 3 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # The plain code request is not an adversarial baseline to compare against, so no baseline. @@ -203,15 +207,18 @@ def __init__( USES_DATASET_SIZE_LIMIT: ClassVar[bool] = False - def _validate_runtime_configuration(self) -> None: + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + techniques = self._resolve_scenario_techniques(scenario_techniques=self.params.get("scenario_techniques")) names = [ *_CORPUS_DATASETS, - *(_LANGUAGE_SPECS[technique.value].dataset_name for technique in self._scenario_techniques), + *(_LANGUAGE_SPECS[technique.value].dataset_name for technique in techniques), ] - by_name = {source.name: source for source in self._dataset_config.sources} - self._dataset_config = self._dataset_config.with_overrides( + by_name = {source.name: source for source in self._default_dataset_config.sources} + return self._default_dataset_config.with_overrides( sources=[by_name.get(name, DatasetSource(name=name)) for name in dict.fromkeys(names)] ) + + def _validate_runtime_configuration(self) -> None: super()._validate_runtime_configuration() if self._max_prompts_per_language < 1: raise ValueError("max_prompts_per_language must be greater than zero") diff --git a/pyrit/scenario/scenarios/garak/prompt_inject.py b/pyrit/scenario/scenarios/garak/prompt_inject.py index 7608f5cb7d..533a3e5a57 100644 --- a/pyrit/scenario/scenarios/garak/prompt_inject.py +++ b/pyrit/scenario/scenarios/garak/prompt_inject.py @@ -17,6 +17,7 @@ from pyrit.executor.attack import AttackConverterConfig, AttackScoringConfig, PromptSendingAttack from pyrit.memory import CentralMemory from pyrit.models import AttackSeedGroup, Parameter, Seed, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -217,6 +218,7 @@ class PromptInject(Scenario): """ VERSION: int = 3 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden DEFAULT_GOAL_TEXTS: ClassVar[tuple[str, str, str]] = _DEFAULT_GOAL_TEXTS TECHNIQUE_DATASET_NAME: ClassVar[str] = "prompt_inject_techniques" diff --git a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py index ffc53cd1ce..2eb4e5b65a 100644 --- a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py +++ b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py @@ -22,6 +22,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import DatasetLimit, normalize_dataset_limit +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetSource @@ -98,6 +99,9 @@ class SystemPromptExtraction(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # Template-dominated like the Doctor/Jailbreak scenarios: the bare system prompt with no diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 6ad61d6e65..26f64a8b78 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -23,6 +23,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import ( @@ -104,6 +105,21 @@ class WebInjection(Scenario): """ VERSION: int = 2 + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """ + Describe approved ingredient sources subject to technique coverage validation. + + Returns: + ScenarioDatasetOptions: Approved sources with ingredient limits disabled. + """ + return ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, + allowed_names=self._default_dataset_config.dataset_names, + per_dataset_limit=False, + total_limit=False, + ) + SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False USES_DATASET_SIZE_LIMIT: ClassVar[bool] = False BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Enabled diff --git a/tests/unit/backend/test_dataset_service.py b/tests/unit/backend/test_dataset_service.py index 86cc51f3a8..19da7f43e5 100644 --- a/tests/unit/backend/test_dataset_service.py +++ b/tests/unit/backend/test_dataset_service.py @@ -58,6 +58,7 @@ async def test_list_datasets_includes_memory_summaries(self, dataset_service, mo modalities=("image_path", "text"), harm_categories=("hate", "violence"), has_unlabeled_harm_categories=True, + origins=("local", "unknown"), ), SeedDatasetSummary( dataset_name=None, @@ -87,11 +88,13 @@ async def test_list_datasets_includes_memory_summaries(self, dataset_service, mo assert loaded.modalities == ["image_path", "text"] assert loaded.harm_categories == ["hate", "violence"] assert loaded.has_unlabeled_harm_categories is True + assert loaded.origins == ["local", "unknown"] provider_only = next(item for item in result.items if item.name == "provider_only") assert provider_only.loaded is False assert provider_only.provider_available is True assert provider_only.logical_examples is None + assert provider_only.origins == [] unnamed = next(item for item in result.items if item.name == "(unnamed)") assert unnamed.selection_key == "dataset:unnamed" diff --git a/tests/unit/backend/test_scenario_configuration_resolver.py b/tests/unit/backend/test_scenario_configuration_resolver.py index 22ff1e2c0d..7710f991fb 100644 --- a/tests/unit/backend/test_scenario_configuration_resolver.py +++ b/tests/unit/backend/test_scenario_configuration_resolver.py @@ -11,7 +11,9 @@ from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import SeedObjective +from pyrit.models.catalog.scenario import RunScenarioRequest from pyrit.models.dataset_limit import DatasetLimit, ResolvedDatasetLimit +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.registry import TargetRegistry from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource, Scenario @@ -20,6 +22,7 @@ override_default_adversarial_target, scenario_target_defaults, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse from pyrit.scenario.scenarios.garak.api_key import ApiKey @@ -28,16 +31,22 @@ from unit.mocks import MockPromptTarget +def _datasets(**fields: object) -> ScenarioDatasetRequest | None: + return RunScenarioRequest.model_validate({"scenario_name": "s", "target_name": "t", **fields}).dataset_request + + @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize(("limit", "expected"), [(None, 20), ("", 20), ("default", 20), ("all", "all"), (7, 7)]) def test_dataset_name_override_applies_explicit_total(*, limit: DatasetLimit, expected: ResolvedDatasetLimit) -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="garak.api_key", scenario_class=ApiKey, - dataset_names=ApiKey.required_datasets(), - max_dataset_size=limit, + datasets=_datasets(dataset_names=ApiKey.required_datasets(), max_dataset_size=limit), ) - assert resolved["dataset_config"].max_dataset_size == expected + scenario = ApiKey() + scenario.set_params_from_args(args=resolved) + scenario._resolve_runtime_configuration(require_objective_target=False) + assert scenario._dataset_config.max_total == expected @pytest.mark.usefixtures("patch_central_database") @@ -51,10 +60,13 @@ def test_dataset_overrides_preserve_subclass_state_and_default() -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="garak.prompt_inject", scenario_class=MagicMock(return_value=scenario), - dataset_names=PromptInject.required_datasets(), - max_dataset_size=7, - dataset_filters={"data_types": ["text"]}, - )["dataset_config"] + datasets=_datasets( + dataset_names=PromptInject.required_datasets(), + max_dataset_size=7, + dataset_filters={"data_types": ["text"]}, + ), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) assert type(resolved) is PromptInjectDatasetConfiguration assert resolved._goal_texts == ["first goal", "second goal"] assert resolved.max_total == 7 @@ -83,8 +95,9 @@ async def test_total_override_preserves_source_limits_async( resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="test", scenario_class=MagicMock(return_value=scenario), - max_dataset_size=total, - )["dataset_config"] + datasets=_datasets(max_dataset_size=total), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) populations = { name: [SeedObjective(value=f"{name}-{index}", dataset_name=name) for index in range(20)] for name in original.dataset_names @@ -115,9 +128,9 @@ def test_name_override_retains_source_options_and_adds_defaults_for_new_names(to resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="test", scenario_class=MagicMock(return_value=scenario), - dataset_names=["new", "retained"], - max_dataset_size=total, - )["dataset_config"] + datasets=_datasets(dataset_names=["new", "retained"], max_dataset_size=total), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) assert resolved.sources == (DatasetSource(name="new"), retained) assert resolved.sources[1] is retained assert resolved.source_limit("new") == 4 @@ -130,9 +143,14 @@ def test_name_override_retains_source_options_and_adds_defaults_for_new_names(to @pytest.mark.usefixtures("patch_central_database") def test_omitted_total_retains_default() -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( - scenario_name="garak.api_key", scenario_class=ApiKey, dataset_names=ApiKey.required_datasets() + scenario_name="garak.api_key", + scenario_class=ApiKey, + datasets=_datasets(dataset_names=ApiKey.required_datasets()), ) - assert resolved["dataset_config"].max_total == 20 + scenario = ApiKey() + scenario.set_params_from_args(args=resolved) + scenario._resolve_runtime_configuration(require_objective_target=False) + assert scenario._dataset_config.max_total == 20 @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/backend/test_scenario_resume.py b/tests/unit/backend/test_scenario_resume.py index a85af70e57..aa821b0a45 100644 --- a/tests/unit/backend/test_scenario_resume.py +++ b/tests/unit/backend/test_scenario_resume.py @@ -13,6 +13,7 @@ from pyrit.backend.services.scenario_run_service import ( _LAUNCH_REQUEST_FIELDS, _LAUNCH_REQUEST_METADATA_KEY, + _OPTIONAL_LAUNCH_FIELDS, ScenarioRunConflictError, ScenarioRunNotFoundError, ScenarioRunService, @@ -24,8 +25,11 @@ from pyrit.models import ( SCENARIO_RUN_PLAN_METADATA_KEY, AttackOutcome, + DatasetFetchPolicy, + DatasetSourceInput, Message, Parameter, + ScenarioDatasetRequest, ScenarioResult, ScenarioRunState, SeedObjective, @@ -448,9 +452,7 @@ async def test_resume_missing_scenario_registration_is_explicit_async( await service.resume_run_async(scenario_result_id=str(stored.id)) -@pytest.mark.parametrize( - "missing", [name for name in _LAUNCH_REQUEST_FIELDS if name not in {"adversarial_target_name", "max_dataset_size"}] -) +@pytest.mark.parametrize("missing", [name for name in _LAUNCH_REQUEST_FIELDS if name not in _OPTIONAL_LAUNCH_FIELDS]) async def test_resume_incomplete_saved_configuration_never_uses_defaults_async( *, resume_environment: tuple[ScenarioRunService, MockPromptTarget], missing: str ) -> None: @@ -495,6 +497,28 @@ async def test_saved_total_limit_round_trips_async( assert restored.max_dataset_size == (limit_args.get("max_dataset_size") or "default") +async def test_saved_dataset_request_round_trips_async( + resume_environment: tuple[ScenarioRunService, MockPromptTarget], +) -> None: + service, _ = resume_environment + datasets = ScenarioDatasetRequest( + sources=[DatasetSourceInput(name=_DATASET_NAME, max_size=1, fetch=DatasetFetchPolicy.NEVER)], + max_total=1, + filters={"data_types": ["text"]}, + ) + prepared = await service._prepare_run_async( + request=RunScenarioRequest(scenario_name=_SCENARIO_NAME, target_name=_TARGET_NAME, datasets=datasets) + ) + stored = ( + await CentralMemory.get_memory_instance().get_scenario_results_async( + scenario_result_ids=[prepared.scenario._scenario_result_id] + ) + )[0] + restored = service._restore_launch_request(stored=stored) + assert restored.dataset_request == datasets + assert restored.dataset_names is None + + async def test_resume_older_launch_record_without_adversarial_selection_async( resume_environment: tuple[ScenarioRunService, MockPromptTarget], ) -> None: diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index a3d04ae1e0..86d21bf21d 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -56,6 +56,7 @@ config_hash, ) from pyrit.models.catalog.scenario import RunScenarioRequest, ScenarioTechniqueSummary +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.registry import ScenarioMetadata, ScenarioRegistry from pyrit.scenario import Scenario @@ -64,6 +65,7 @@ DatasetConfiguration, get_default_adversarial_target, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.scenario.scenarios.airt.jailbreak import Jailbreak from pyrit.scenario.scenarios.airt.scam import Scam @@ -74,6 +76,16 @@ from pyrit.prompt_target import PromptTarget +def _apply_mapped_datasets(registries: dict[str, Any]) -> DatasetAttackConfiguration: + call = registries["scenario_registry"].create_and_initialize_async.await_args + assert "dataset_config" not in call.kwargs + return apply_dataset_overrides( + config=registries["scenario_instance"]._default_dataset_config, + request=call.kwargs.get("dataset_overrides", ScenarioDatasetRequest()), + options=ScenarioDatasetOptions(), + ) + + class _StubTechnique(ScenarioTechnique): """Minimal concrete ScenarioTechnique used to exercise converter-token parsing.""" @@ -423,7 +435,8 @@ def update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunStat release_first.set() await asyncio.wait_for(finished.wait(), timeout=5) await asyncio.wait_for(first_task, timeout=5) - assert introspected == initialized == [first, second, fallback, second] + assert introspected == [] + assert initialized == [first, second, fallback, second] assert executed == [first, second, fallback] assert get_default_adversarial_target() is fallback finally: @@ -826,8 +839,8 @@ async def test_start_run_max_dataset_size_uses_default_config(self, mock_all_reg assert default_config.max_total == 100 init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - assert init_call.kwargs["dataset_config"] is not default_config - assert init_call.kwargs["dataset_config"].max_total == 5 + assert init_call.kwargs["dataset_overrides"].max_total == 5 + assert "dataset_config" not in init_call.kwargs @pytest.mark.parametrize( ("limit_args", "expected"), @@ -850,7 +863,7 @@ async def test_start_run_resolves_total_limit_async( request=RunScenarioRequest(scenario_name="test", target_name="my_target", **limit_args) ) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - config = init_call.kwargs.get("dataset_config", default_config) + config = _apply_mapped_datasets(mock_all_registries) assert config.max_total == expected saved = init_call.kwargs["initial_metadata"][_svc_mod._LAUNCH_REQUEST_METADATA_KEY] assert saved["max_dataset_size"] == (limit_args.get("max_dataset_size") or "default") @@ -876,7 +889,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): await service.start_run_async(request=_make_request(dataset_names=["custom_a", "custom_b"], max_dataset_size=3)) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) # Type is preserved (this is the regression assertion) assert type(built_config) is _MarkerDatasetConfiguration @@ -902,7 +915,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): await service.start_run_async(request=_make_request(dataset_names=["only_this"])) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert type(built_config) is _MarkerDatasetConfiguration assert built_config.dataset_names == ["only_this"] assert built_config.max_dataset_size == "all" @@ -926,7 +939,7 @@ def __init__(self, *, required_extra: str, **kwargs: Any) -> None: service = ScenarioRunService() await service.start_run_async(request=_make_request(dataset_names=["custom"])) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - config = init_call.kwargs["dataset_config"] + config = _apply_mapped_datasets(mock_all_registries) assert isinstance(config, _RequiresExtraArgConfiguration) assert config._required_extra == "seeded" assert config.dataset_names == ["custom"] @@ -950,7 +963,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): ) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert built_config.dataset_names == ["custom"] assert built_config.max_dataset_size == 7 assert built_config.filters == {"harm_categories": ["cyber"]} @@ -965,7 +978,7 @@ async def test_start_run_dataset_filters_updates_default_config(self, mock_all_r await service.start_run_async(request=_make_request(dataset_filters={"harm_categories": ["cyber"]})) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert built_config is not default_config assert default_config.filters == {} assert built_config.filters == {"harm_categories": ["cyber"]} @@ -994,8 +1007,13 @@ async def test_start_run_dataset_names_introspection_failure_raises(self, mock_m patch(f"{_REGISTRY_PATCH_BASE}.TargetRegistry.get_registry_singleton", return_value=mock_tr), patch(f"{_REGISTRY_PATCH_BASE}.InitializerRegistry.get_registry_singleton", return_value=mock_ir), ): - with pytest.raises(ValueError, match="not instantiable without arguments"): - await service.start_run_async(request=_make_request(dataset_names=["custom"])) + mapped = _svc_mod.ScenarioConfigurationResolver.resolve_configuration( + scenario_name="test", + scenario_class=mock_scenario_class, + datasets=_make_request(dataset_names=["custom"]).dataset_request, + ) + assert mapped["dataset_overrides"].sources[0].name == "custom" + mock_scenario_class.assert_not_called() async def test_start_run_max_dataset_size_with_dataset_names_uses_subclass_with_both( self, mock_all_registries @@ -1013,9 +1031,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): service = ScenarioRunService() await service.start_run_async(request=_make_request(dataset_names=["a", "b"], max_dataset_size=7)) - built_config = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args.kwargs[ - "dataset_config" - ] + built_config = _apply_mapped_datasets(mock_all_registries) assert type(built_config) is _MarkerDatasetConfiguration assert built_config.dataset_names == ["a", "b"] assert built_config.max_dataset_size == 7 diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index 17c0505402..e99fce1b1f 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -397,7 +397,8 @@ async def estimate_async(**kwargs: object) -> ScenarioRunSizeEstimate: for name in ("first", "second") ) ) - assert introspected == estimated == [first, second] + assert introspected == [] + assert estimated == [first, second] assert get_default_adversarial_target() is fallback assert dict(service._estimate_cache) == cache_before assert await service.get_scenario_async(scenario_name=metadata.registry_name) == before @@ -1541,7 +1542,11 @@ async def test_estimate_resolves_total_limit_async( scenario_name="test.scenario", request=ScenarioRunSizeEstimateRequest(dataset_names=["harmbench"], **limit_args), ) - config = service._registry.create_and_estimate_async.await_args.kwargs["dataset_config"] + from pyrit.models.scenario_dataset import ScenarioDatasetOptions + from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides + + mapped = service._registry.create_and_estimate_async.await_args.kwargs["dataset_overrides"] + config = apply_dataset_overrides(config=original, request=mapped, options=ScenarioDatasetOptions()) assert config.max_total == expected assert config.max_per_dataset == 5 assert original.max_total == 20 @@ -1619,11 +1624,11 @@ async def test_configured_estimate_uses_shared_launch_resolution(self) -> None: 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"]} + overrides = call.kwargs["dataset_overrides"] + assert "dataset_config" not in call.kwargs + assert [source.name for source in overrides.sources] == ["harmbench"] + assert overrides.max_total == 3 + assert overrides.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.""" diff --git a/tests/unit/cli/test_dataset_request.py b/tests/unit/cli/test_dataset_request.py new file mode 100644 index 0000000000..58cdbd63b4 --- /dev/null +++ b/tests/unit/cli/test_dataset_request.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Scanner and shell share strict edit-only dataset parsing.""" + +from contextlib import nullcontext + +import pytest + +from pyrit.cli._cli_args import build_dataset_request_fields, parse_run_arguments +from pyrit.cli.pyrit_scan import _build_run_request, parse_args + + +@pytest.mark.parametrize( + "flags", + [ + "--datasets harmbench --dataset-names other", + "--max-per-dataset 2 --max-dataset-size 4", + "--max-total all --dataset-names harmbench", + "--dataset-fetch never --max-dataset-size default", + ], +) +def test_scanner_and_shell_reject_mixed_dataset_flags(flags: str) -> None: + parsed = parse_args(["run", "test", "--target", "target", *flags.split()]) + with pytest.raises(ValueError, match="Do not mix"): + _build_run_request(parsed_args=parsed, scenario_name="test") + with pytest.raises(ValueError, match="Do not mix"): + build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + + +@pytest.mark.parametrize( + "flags", + [ + "--datasets name:0", + "--datasets name:4:all", + "--datasets name name", + "--dataset-fetch always", + "--max-total -1", + ], +) +def test_shell_rejects_malformed_datasets(flags: str) -> None: + with pytest.raises(ValueError): + build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + + +@pytest.mark.parametrize("source_flags", ["--datasets harmbench:all", "--dataset-names harmbench"]) +def test_common_filters_work_with_either_form(source_flags: str) -> None: + flags = source_flags + " --dataset-filters data_types=text,image_path" + with pytest.warns(DeprecationWarning) if "--dataset-names" in flags else nullcontext(): + fields = build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + if "datasets" in fields: + assert fields["datasets"].filters == {"data_types": ["text", "image_path"]} + else: + assert fields["dataset_filters"] == {"data_types": ["text", "image_path"]} diff --git a/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py b/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py index 820494cd22..24e9bdeede 100644 --- a/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py +++ b/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py @@ -5,7 +5,21 @@ from pyrit.memory import MemoryInterface from pyrit.memory.memory_session import MemorySession -from pyrit.models import SeedPrompt +from pyrit.models import SeedOrigin, SeedPrompt + + +async def test_dataset_origins_aggregate_without_seed_hydration(sqlite_instance: MemoryInterface) -> None: + await sqlite_instance.add_seeds_to_memory_async( + seeds=[ + SeedPrompt(value="local", dataset_name="mixed", origin=SeedOrigin.LOCAL), + SeedPrompt(value="unknown", dataset_name="mixed", origin=SeedOrigin.UNKNOWN), + ], + added_by="test", + ) + with patch.object(sqlite_instance, "get_seeds_async", side_effect=AssertionError("No hydration")): + summaries = await sqlite_instance.get_seed_dataset_summaries_async() + assert summaries[0].origins == ("local", "unknown") + assert summaries[0].logical_examples == summaries[0].seed_pieces == 2 async def test_get_seed_dataset_summaries_avoids_metadata_row_multiplication( diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py index 6910656e1a..01debc3841 100644 --- a/tests/unit/models/test_scenario_catalog.py +++ b/tests/unit/models/test_scenario_catalog.py @@ -66,6 +66,9 @@ def test_run_size_estimate_preserves_legacy_total_and_serializes_additively() -> assert estimate.minimum_attack_count == 6 assert estimate.maximum_attack_count == 6 payload = estimate.model_dump(mode="json") + assert payload.pop("dataset_options")["supports_overrides"] is True + assert payload.pop("dataset_defaults")["sources"] == [] + assert payload.pop("effective_datasets")["sources"] == [] assert list(payload["effective_parameters"]) == ["include_baseline", "techniques"] assert payload == { "status": "exact", diff --git a/tests/unit/models/test_scenario_request.py b/tests/unit/models/test_scenario_request.py index e51c384c7d..4105756438 100644 --- a/tests/unit/models/test_scenario_request.py +++ b/tests/unit/models/test_scenario_request.py @@ -6,9 +6,108 @@ import pytest from pydantic import ValidationError +from pyrit.models import DatasetSourceInput, ScenarioDatasetRequest from pyrit.models.catalog.scenario import DATASET_FILTERS, RunScenarioRequest, ScenarioRunSizeEstimateRequest +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_new_dataset_request_transport_round_trip(model: type) -> None: + data = { + "scenario_name": "s", + "target_name": "t", + "datasets": { + "sources": [{"name": "harmbench"}, {"name": "custom", "max_size": "all", "fetch": "never"}], + "max_per_dataset": 4, + "max_total": 12, + "fetch": "if_missing", + "filters": {"data_types": ["text"]}, + }, + } + request = model.model_validate(data) + assert request.dataset_request == request.datasets + assert model.model_validate(request.model_dump(mode="json", exclude_none=True)) == request + assert model.model_validate_json(request.model_dump_json()) == request + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +@pytest.mark.parametrize( + "legacy", + [{"dataset_names": ["a"]}, {"max_dataset_size": 3}, {"dataset_filters": {"data_types": ["text"]}}], +) +def test_new_dataset_request_rejects_deprecated_values(*, model: type, legacy: dict[str, object]) -> None: + with pytest.raises(ValidationError, match="not both"): + model.model_validate({"scenario_name": "s", "target_name": "t", "datasets": {"max_total": 2}, **legacy}) + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_new_dataset_request_accepts_deprecated_defaults(model: type) -> None: + request = model.model_validate( + { + "scenario_name": "s", + "target_name": "t", + "datasets": {"max_total": 2}, + "dataset_names": None, + "max_dataset_size": "default", + "dataset_filters": None, + } + ) + assert request.dataset_request == ScenarioDatasetRequest(max_total=2) + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_deprecated_dataset_fields_convert_to_one_request(model: type) -> None: + request = model.model_validate( + { + "scenario_name": "s", + "target_name": "t", + "dataset_names": ["a", "b"], + "max_dataset_size": 3, + "dataset_filters": {"data_types": ["text"]}, + } + ) + assert request.datasets is None + assert request.dataset_request == ScenarioDatasetRequest( + sources=[DatasetSourceInput(name="a"), DatasetSourceInput(name="b")], + max_total=3, + filters={"data_types": ["text"]}, + ) + assert model.model_validate({"scenario_name": "s", "target_name": "t"}).dataset_request is None + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_invalid_deprecated_dataset_names_fail_validation(model: type) -> None: + with pytest.raises(ValidationError, match="Duplicate"): + model.model_validate({"scenario_name": "s", "target_name": "t", "dataset_names": ["a", "a"]}) + + +@pytest.mark.parametrize( + "data", + [ + {"sources": []}, + {"sources": [{"name": "a"}, {"name": "a"}]}, + {"sources": [{"name": " "}]}, + {"sources": [{"name": "a", "provider": {}}]}, + {"max_total": True}, + {"max_per_dataset": 0}, + {"max_total": 1.5}, + {"fetch": "always"}, + {"filters": {"authors": ["me"]}}, + {"filters": {"data_types": "text"}}, + {"unexpected": True}, + ], +) +def test_shared_dataset_request_rejects_invalid_inputs(data: dict[str, object]) -> None: + with pytest.raises(ValidationError): + ScenarioDatasetRequest.model_validate(data) + + +@pytest.mark.parametrize("value", [None, "", "default"]) +def test_shared_dataset_request_default_semantics(value: object) -> None: + request = ScenarioDatasetRequest.model_validate({"max_total": value, "max_per_dataset": value, "fetch": value}) + assert request.max_total == request.max_per_dataset == "default" + assert request.fetch is None + + @pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) @pytest.mark.parametrize( ("limit", "expected"), diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py index 6722782574..10d2af2782 100644 --- a/tests/unit/registry/test_scenario_registry.py +++ b/tests/unit/registry/test_scenario_registry.py @@ -7,6 +7,7 @@ import pytest +from pyrit.models import DatasetSelectionMode, ScenarioDatasetOptions, ScenarioDatasetSettings from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy, ScenarioTechnique @@ -62,6 +63,14 @@ def _resolve_scenario_techniques(self, *, scenario_techniques): """Resolve the concrete defaults.""" return _MetadataTechnique.resolve(scenario_techniques, default=self._default_technique) + def get_dataset_options(self) -> ScenarioDatasetOptions: + """Return fixed dataset options.""" + return ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED) + + def get_default_dataset_settings(self) -> ScenarioDatasetSettings: + """Return default dataset settings.""" + return ScenarioDatasetSettings(max_total=4) + class _MarkdownMetadataScenario(_MetadataScenario): """ @@ -104,6 +113,8 @@ def test_build_metadata_expands_ordered_default_techniques() -> None: "all": ("one", "two"), "default": ("one", "two"), } + assert metadata.dataset_options.mode is DatasetSelectionMode.FIXED + assert metadata.dataset_defaults.max_total == 4 def test_build_metadata_reports_scenario_owned_adversarial_usage() -> None: diff --git a/tests/unit/scenario/airt/test_leakage.py b/tests/unit/scenario/airt/test_leakage.py index 1888d3be4f..903538b642 100644 --- a/tests/unit/scenario/airt/test_leakage.py +++ b/tests/unit/scenario/airt/test_leakage.py @@ -10,6 +10,7 @@ from pyrit.common.path import DATASETS_PATH from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedDataset, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -52,6 +53,9 @@ def mock_dataset_config(mock_memory_seeds): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_groups_by_dataset_async = AsyncMock(return_value={"airt_leakage": seed_groups}) mock_config.dataset_names = ["airt_leakage"] return mock_config diff --git a/tests/unit/scenario/airt/test_psychosocial.py b/tests/unit/scenario/airt/test_psychosocial.py index 816fb42e0d..da5ada8929 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -24,6 +24,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import ResolvedDatasetLimit +from pyrit.models.scenario_dataset import ScenarioDatasetRequest from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.scenario.core.dataset_configuration import ( @@ -188,7 +189,9 @@ def get_seeds(*, dataset_name: str, **_: object) -> list[SeedObjective]: return list(seeds_by_dataset[dataset_name]) config = DatasetAttackConfiguration( - sources=[DatasetSource(name=name) for name in seeds_by_dataset], + sources=[ + DatasetSource(name=harm.dataset_name) for harm in _SUB_HARMS if sub_harm == "all" or harm.name == sub_harm + ], max_per_dataset="all", max_total=outer_limit, ) @@ -253,9 +256,9 @@ async def test_backend_total_one_rejects_both_sub_harms_before_reads_async( scenario_name="airt.psychosocial", scenario_class=MagicMock(return_value=scenario), objective_target=mock_objective_target, - max_dataset_size=1, + datasets=ScenarioDatasetRequest(max_total=1), ) - assert args["dataset_config"].max_total == 1 + assert args["dataset_overrides"].max_total == 1 scenario.set_params_from_args(args=args) with ( patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), @@ -310,7 +313,7 @@ async def test_backend_total_one_allows_one_selected_sub_harm_async(mock_objecti scenario_name="airt.psychosocial", scenario_class=MagicMock(return_value=scenario), objective_target=mock_objective_target, - max_dataset_size=1, + datasets=ScenarioDatasetRequest(max_total=1), ) scenario.set_params_from_args( args={**args, "sub_harm": "imminent_crisis", "scenario_techniques": [PsychosocialTechnique.NoConverter]} diff --git a/tests/unit/scenario/airt/test_scam.py b/tests/unit/scenario/airt/test_scam.py index 5c4a99043e..f28390de81 100644 --- a/tests/unit/scenario/airt/test_scam.py +++ b/tests/unit/scenario/airt/test_scam.py @@ -12,6 +12,7 @@ from pyrit.executor.attack import PromptSendingAttack, RedTeamingAttack from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedDataset, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import OpenAIChatTarget, PromptTarget from pyrit.scenario import DatasetAttackConfiguration, DatasetConfiguration from pyrit.scenario.scenarios.airt.scam import Scam, ScamTechnique @@ -57,6 +58,9 @@ def mock_dataset_config(mock_memory_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=attack_seed_groups) mock_config.get_attack_groups_by_dataset_async = AsyncMock(return_value={"airt_scam": attack_seed_groups}) mock_config.dataset_names = ["airt_scam"] diff --git a/tests/unit/scenario/core/test_scenario_dataset_overrides.py b/tests/unit/scenario/core/test_scenario_dataset_overrides.py new file mode 100644 index 0000000000..6b54762149 --- /dev/null +++ b/tests/unit/scenario/core/test_scenario_dataset_overrides.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Offline dataset contract coverage across Python, REST, scanner, and shell.""" + +from unittest.mock import AsyncMock, patch + +import pytest +from unit.mocks import MockPromptTarget + +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver +from pyrit.backend.services.scenario_service import ScenarioService +from pyrit.cli._cli_args import ( + build_dataset_request_fields, + parse_run_arguments, +) +from pyrit.cli.pyrit_scan import _build_run_request, parse_args +from pyrit.memory import CentralMemory +from pyrit.models import ( + DatasetFetchPolicy, + DatasetSelectionMode, + DatasetSourceInput, + ScenarioDatasetRequest, + SeedObjective, +) +from pyrit.models.catalog import RunScenarioRequest +from pyrit.registry import ScenarioRegistry +from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario +from pyrit.scenario.core.atomic_attack import AtomicAttack +from pyrit.scenario.core.matrix_atomic_attack_builder import build_baseline_atomic_attack +from pyrit.scenario.core.scenario_context import ScenarioContext +from pyrit.scenario.core.scenario_technique import ScenarioTechnique +from pyrit.scenario.scenarios.garak.api_key import ApiKey +from pyrit.scenario.scenarios.garak.encoding import Encoding +from pyrit.scenario.scenarios.garak.figstep import FigStep +from pyrit.scenario.scenarios.garak.package_hallucination import PackageHallucination +from pyrit.score import SubStringScorer + + +class _Technique(ScenarioTechnique): + ALL = ("all", {"all"}) + DIRECT = ("direct", {"direct"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + +class _Scenario(Scenario): + def __init__(self, *, scenario_result_id: str | None = None) -> None: + super().__init__( + version=1, + technique_class=_Technique, + default_dataset_config=DatasetAttackConfiguration( + sources=[DatasetSource(name="a", max_size=2, fetch=DatasetFetchPolicy.NEVER), DatasetSource(name="b")], + max_per_dataset=5, + max_total=7, + ), + objective_scorer=SubStringScorer(substring="yes"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + return [ + build_baseline_atomic_attack( + objective_target=context.objective_target, + objective_scorer=self._objective_scorer, + seed_groups=list(context.seed_groups), + memory_labels=context.memory_labels, + ) + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_python_rest_scanner_shell_resolve_identical_settings_without_fetching() -> None: + tokens = "--datasets a b:all --max-per-dataset 3 --max-total 4 --dataset-fetch never" + scanner_request = _build_run_request( + parsed_args=parse_args(["run", "test", "--target", "test", *tokens.split()]), scenario_name="test" + ) + scanner = scanner_request.datasets + shell = build_dataset_request_fields(parse_run_arguments(args_string="test " + tokens))["datasets"] + python = ScenarioDatasetRequest( + sources=[DatasetSourceInput(name="a"), DatasetSourceInput(name="b", max_size="all")], + max_per_dataset=3, + max_total=4, + fetch=DatasetFetchPolicy.NEVER, + ) + assert scanner == shell == python + request = RunScenarioRequest(scenario_name="test", target_name="test", datasets=python) + transported = RunScenarioRequest.model_validate_json(request.model_dump_json(exclude_none=True)) + mapped = ScenarioConfigurationResolver.resolve_configuration( + scenario_name="test", + scenario_class=_Scenario, + datasets=transported.dataset_request, + ) + assert "dataset_config" not in mapped + estimates = [] + memory = CentralMemory.get_memory_instance() + with ( + patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), + patch.object(memory, "get_seeds_async", side_effect=AssertionError("No seed reads")), + ): + for overrides in [python, scanner, shell, mapped["dataset_overrides"]]: + scenario = _Scenario() + scenario.set_params_from_args(args={"dataset_overrides": overrides}) + estimates.append(await scenario.get_run_size_estimate_async()) + assert all(estimate.effective_datasets == estimates[0].effective_datasets for estimate in estimates) + estimate = estimates[0] + assert estimate.dataset_defaults.max_total == 7 + assert estimate.effective_datasets.max_total == 4 + assert [source.max_size for source in estimate.effective_datasets.sources] == [2, "all"] + assert estimate.dataset_defaults.sources[0].fetch is DatasetFetchPolicy.NEVER + + +@pytest.mark.usefixtures("patch_central_database") +async def test_resume_replays_saved_selection_without_preparation() -> None: + memory = CentralMemory.get_memory_instance() + await memory.add_seeds_to_memory_async( + seeds=[SeedObjective(value=f"{name}-{index}", dataset_name=name) for name in ["a", "b"] for index in range(6)], + added_by="test", + ) + target = MockPromptTarget() + overrides = ScenarioDatasetRequest(max_total=4, fetch=DatasetFetchPolicy.NEVER) + scenario = _Scenario() + scenario.set_params_from_args(args={"objective_target": target, "dataset_overrides": overrides}) + with patch.object(DatasetAttackConfiguration, "prepare_async", new_callable=AsyncMock) as prepare: + await scenario.initialize_async() + prepare.assert_awaited_once() + selected = [group.logical_id for group in scenario._atomic_attacks[0].seed_groups] + assert len(selected) == 4 + + # Changed caps do not matter on resume: the stored run plan fixes the selected seed groups. + resumed = _Scenario(scenario_result_id=scenario._scenario_result_id) + resumed._default_dataset_config = resumed._default_dataset_config.with_overrides(max_total=1, max_per_dataset=1) + resumed.set_params_from_args(args={"objective_target": target, "dataset_overrides": overrides}) + with patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No resume preparation")): + await resumed.initialize_async() + assert [group.logical_id for group in resumed._atomic_attacks[0].seed_groups] == selected + + # A different dataset selection is an error, never a silent replacement. + diverged = _Scenario(scenario_result_id=scenario._scenario_result_id) + diverged.set_params_from_args( + args={ + "objective_target": target, + "dataset_overrides": ScenarioDatasetRequest(sources=[DatasetSourceInput(name="a")]), + } + ) + with pytest.raises(ValueError, match="does not match"): + await diverged.initialize_async() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("scenario_class", "overrides", "error"), + [ + (FigStep, {"sources": [{"name": "other"}]}, "Select datasets"), + (FigStep, {"sources": [{"name": "figstep"}, {"name": "figstep_pro"}]}, "at most"), + (ApiKey, {"max_per_dataset": 1}, "per-dataset"), + (ApiKey, {"sources": [{"name": "harmbench"}]}, "fixed"), + (Encoding, {"sources": [{"name": "garak_slur_terms_en"}]}, "ambiguous"), + (PackageHallucination, {"max_total": 1}, "total dataset"), + ], +) +async def test_constrained_scenarios_reject_unsupported_overrides_before_preparation( + *, + scenario_class: type[Scenario], + overrides: dict[str, object], + error: str, +) -> None: + scenario = scenario_class(objective_scorer=SubStringScorer(substring="yes")) + scenario.set_params_from_args(args={"dataset_overrides": overrides}) + with ( + patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), + pytest.raises(ValueError, match=error), + ): + await scenario.get_run_size_estimate_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_encoding_global_overrides_preserve_child_caps() -> None: + scenario = Encoding() + scenario.set_params_from_args(args={"dataset_overrides": {"max_total": 3, "fetch": "never"}}) + estimate = await scenario.get_run_size_estimate_async() + assert [child.max_total for child in scenario._dataset_config.configurations] == [10, 10] + assert estimate.effective_datasets.max_total == 3 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_catalog_dataset_controls_do_not_depend_on_estimates() -> None: + with patch.object(ScenarioRegistry, "_discover"): + registry = ScenarioRegistry() + registry.register_class(ApiKey, name="garak.api_key") + with patch.object(ScenarioRegistry, "get_registry_singleton", return_value=registry): + [item] = (await ScenarioService().list_scenarios_async(include_estimates=False)).items + assert item.dataset_options.mode is DatasetSelectionMode.FIXED + assert item.dataset_options.per_dataset_limit is False + assert [source.name for source in item.dataset_defaults.sources] == ApiKey.required_datasets() + assert item.dataset_defaults.max_total == 20 diff --git a/tests/unit/scenario/foundry/test_red_team_agent.py b/tests/unit/scenario/foundry/test_red_team_agent.py index ad15602a47..cc00251ff3 100644 --- a/tests/unit/scenario/foundry/test_red_team_agent.py +++ b/tests/unit/scenario/foundry/test_red_team_agent.py @@ -41,6 +41,7 @@ from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.scenario import AtomicAttack, DatasetAttackConfiguration from pyrit.scenario.foundry import ( # type: ignore[ty:unresolved-import] @@ -133,6 +134,9 @@ def mock_dataset_config(mock_memory_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=mock_memory_seed_groups) mock_config.dataset_names = ["foundry_red_team"] return mock_config diff --git a/tests/unit/scenario/garak/test_api_key.py b/tests/unit/scenario/garak/test_api_key.py index 0a93b2389c..21175a6600 100644 --- a/tests/unit/scenario/garak/test_api_key.py +++ b/tests/unit/scenario/garak/test_api_key.py @@ -12,7 +12,7 @@ from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.converter import Base64Converter, Converter from pyrit.executor.attack import PromptSendingAttack -from pyrit.models import ComponentIdentifier, ScenarioRunSizeEstimateStatus, Seed, SeedDataset +from pyrit.models import ComponentIdentifier, ScenarioDatasetRequest, ScenarioRunSizeEstimateStatus, Seed, SeedDataset from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetConstraintError from pyrit.scenario.core.scenario import BaselineAttackPolicy @@ -247,7 +247,7 @@ async def test_launch_and_estimate_use_standard_dataset_size( scenario_name="garak.api_key", scenario_class=ApiKey, objective_target=mock_objective_target, - max_dataset_size=size, + datasets=ScenarioDatasetRequest(max_total=size), ) scenario.set_params_from_args(args=args) with patch.object( diff --git a/tests/unit/scenario/garak/test_divergence.py b/tests/unit/scenario/garak/test_divergence.py index bd07e847fa..49730b0fc0 100644 --- a/tests/unit/scenario/garak/test_divergence.py +++ b/tests/unit/scenario/garak/test_divergence.py @@ -18,6 +18,7 @@ DivergesFromRepetition, Message, MessagePiece, + ScenarioDatasetRequest, Seed, SeedDataset, SeedObjective, @@ -154,7 +155,7 @@ async def test_runtime_budget_matches_estimate(self, corpus: list[Seed], size: i scenario_name="garak.divergence", scenario_class=Divergence, objective_target=MockPromptTarget(), - max_dataset_size=size, + datasets=ScenarioDatasetRequest(max_total=size), ) scenario.set_params_from_args(args=args) with patch.object( diff --git a/tests/unit/scenario/garak/test_encoding.py b/tests/unit/scenario/garak/test_encoding.py index f6ef5ea8f1..3668873cc6 100644 --- a/tests/unit/scenario/garak/test_encoding.py +++ b/tests/unit/scenario/garak/test_encoding.py @@ -10,6 +10,7 @@ from pyrit.converter import Base64Converter from pyrit.executor.attack import PromptSendingAttack from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.scenario import CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, DatasetConfiguration from pyrit.scenario.garak import Encoding, EncodingTechnique # type: ignore[ty:unresolved-import] @@ -65,6 +66,9 @@ def mock_dataset_config(mock_attack_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=mock_attack_seed_groups) mock_config.dataset_names = ["garak_slur_terms_en", "garak_web_html_js"] return mock_config From 9373de3ada31b9fb4620987136683a52533a55cc Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 13 Oct 2026 12:25:22 -0700 Subject: [PATCH 2/4] Refine scenario dataset controls and explicit loading Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16cbdc1b-2bf6-4286-be66-c9f30c7ed72b --- doc/getting_started/install_local_dev.md | 9 + doc/gui/0_gui.md | 72 +++- frontend/e2e/numeric-controls.spec.ts | 106 +++-- frontend/e2e/touch-targets.spec.ts | 4 +- frontend/src/App.labels.test.tsx | 4 +- .../Scenarios/DatasetLimitControl.tsx | 63 +++ .../Scenarios/ScenarioDatasetTable.styles.ts | 23 +- .../Scenarios/ScenarioDatasetTable.test.tsx | 331 ++++++++++------ .../Scenarios/ScenarioDatasetTable.tsx | 363 ++++++++++-------- .../Scenarios/ScenarioDetail.test.tsx | 104 +++-- .../components/Scenarios/ScenarioDetail.tsx | 104 ++--- .../Scenarios/datasetPresentation.ts | 34 ++ .../components/Sidebar/Navigation.test.tsx | 13 +- .../src/components/Sidebar/Navigation.tsx | 4 +- frontend/src/components/Tour/tourSteps.ts | 2 +- .../src/hooks/useScenarioDatasets.test.tsx | 49 ++- frontend/src/hooks/useScenarioDatasets.ts | 53 ++- frontend/src/hooks/useTour.test.ts | 1 + frontend/src/services/api.ts | 9 + frontend/src/types/index.ts | 5 + pyrit/backend/models/datasets.py | 8 + pyrit/backend/routes/scenarios.py | 27 ++ pyrit/backend/services/scenario_service.py | 42 ++ .../remote/remote_dataset_loader.py | 3 +- pyrit/scenario/core/dataset_configuration.py | 8 +- pyrit/scenario/core/scenario.py | 41 +- .../scenarios/garak/package_hallucination.py | 20 +- .../garak/system_prompt_extraction.py | 21 +- .../scenario/scenarios/garak/web_injection.py | 11 +- .../backend/test_scenario_dataset_loading.py | 136 +++++++ .../datasets/test_remote_dataset_loader.py | 30 ++ .../core/test_scenario_dataset_overrides.py | 86 ++++- .../test_default_run_size_estimates.py | 21 +- 33 files changed, 1319 insertions(+), 488 deletions(-) create mode 100644 frontend/src/components/Scenarios/DatasetLimitControl.tsx create mode 100644 frontend/src/components/Scenarios/datasetPresentation.ts create mode 100644 tests/unit/backend/test_scenario_dataset_loading.py diff --git a/doc/getting_started/install_local_dev.md b/doc/getting_started/install_local_dev.md index c877bbcee5..3bb4ef97aa 100644 --- a/doc/getting_started/install_local_dev.md +++ b/doc/getting_started/install_local_dev.md @@ -114,6 +114,15 @@ Use `uv run` to execute Python with the virtual environment: uv run python your_script.py ``` +Remote dataset loaders use `HF_DATASETS_CACHE` when it is set. Otherwise they cache +Hugging Face data under `dbdata/huggingface`. On Windows, a long worktree path can +cause `WinError 206` when Hugging Face creates a cache lock file. Set a shorter +cache path before you start the server or script: + +```powershell +$env:HF_DATASETS_CACHE = Join-Path $env:TEMP "pyrit-hf" +``` + #### Running Tests ```bash diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index 61dfc0e020..b85fc8d30b 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -34,7 +34,7 @@ To deploy an isolated instance for an external team, see [Deploy a New Instance] ## Views -CoPyRIT has three main views, accessible from the left sidebar: **Chat**, **Attack History**, and **Target Configuration**. The **Theme** menu is available at the bottom of the sidebar. +CoPyRIT has three main views, accessible from the left sidebar: **Technique.execute** (Chat), **Attack History**, and **Target Configuration**. The **Theme** menu is available at the bottom of the sidebar. ### Themes @@ -505,29 +505,59 @@ Defaults are saved in this browser, separately for each signed-in account. They A missing or changed default is shown as unavailable. Select a new default or clear it; CoPyRIT does not silently substitute a different target. If browser storage is unavailable, a warning states that the choice applies only in the current session. -Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters** above **Datasets**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. +Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. ### Scenario datasets -The **Datasets** table shows each selected name, stored seed origins (including mixed -origins), stored example count, source limit, and fetch policy. Unknown counts for -provider-only datasets are not zero. Provider availability is separate from stored origin. -**Reload dataset metadata** refreshes the page-scoped catalog without fetching seed data. -Runtime changes and launches also invalidate this catalog. - -Use **Add dataset** for allowed alternatives or named memory datasets, **Remove** to -change the selection, and **Use default sources** to restore parameter-aware defaults. -FigStep allows one of `figstep` or `figstep_pro`; ingredient and compound sources are fixed. -Unavailable limit/source controls are not submitted. Parameter and technique changes -refresh resolved defaults without discarding edits; invalid selections display an error. - -**Max per dataset**, source limits, and **Max total** accept positive integers, `all`, -or `default`/blank. Source caps apply before the total; `all` removes only that cap. -Name-only selection retains existing explicit source settings. Fetch is `never` or -`if_missing`: only a new run may prepare missing registered datasets. Preview and launch -send the same edit-only dataset object, with harm/category filters preserved. Resume uses -the saved request and the original seed selection; a different dataset selection is -rejected. Resume never prepares datasets. +**Datasets** is a separate section below **Techniques**. Its table has **Dataset**, +**Size**, and **Include Criteria** columns. **Size** shows an upper bound, such as **Up to 4 +of 200**, when the stored total is known, or **Up to 4** before loading. This is not +an exact selected count: filters and the combined objective limit can reduce it. +Click a supported size to edit its limit, use all available examples, or reset to +the displayed default. Required reference data shows its stored size without a +limit editor. + +Select a name from **Add dataset** to add it immediately. Provider names remain +available before their data is stored. Selecting a name never loads it. The remove +icon is available only when the scenario permits removal and another source will +remain. A scenario that permits only one source uses **Choose dataset** to replace +it. **Browse datasets**, beside **Reset datasets**, opens the dataset browser. +**Reset datasets** restores source defaults and row limits without clearing +the combined objective limit or filters. + +**Load datasets** appears only while selected data is missing. It loads missing +sources and keeps existing data, including manual edits. It may download data or +use a configured generator, but does not start a scan. A missing source without a +provider must be imported first. When all selected data is loaded, the load button +and loading controls are omitted. + +Loaded rows show harm categories and data types as differently colored chips. +All labels start checked when there is no configured filter. Uncheck labels to +narrow the existing filters for **all selected datasets**, not just the row clicked. +Checking all labels restores the default request when no filter was configured. +Partial category selections require every checked label and retain substring +matching; data types match any checked type. An empty selection blocks launch +instead of silently restoring all data. Selection is shown on the metadata chips; +there is no separate filter list below the table. No preview of data values is provided. + +**Maximum objectives** shows its resolved value. Click it to enter a positive +integer, use all available objectives, or reset to the displayed default. This +limits objectives, not prompts sent. Four objectives with three techniques produce +12 attacks, plus four more if the baseline is enabled. Each attack can send several +prompts. The run estimate shows the resulting attacks. + +Scenarios that build objectives from templates or reference data apply the total +limit to the completed objectives, not to the required reference data. Their +generation limits still apply. A scenario that requires one objective per selected +technique rejects a total limit smaller than the number of techniques. + +**Load missing data when the scan starts** is one switch for the selected sources. +Its initial state reflects their resolved loading rules, which normally load +missing data. Turning it off prevents automatic loading on scan start without +disabling the manual load action. Existing source-specific rules are not silently +replaced until the switch is changed. Closing a limit editor does not clear an +invalid draft. Estimate and launch use the same dataset edits. +Resume uses the saved request and original selection and never loads datasets. #### Core Adversarial Default Override diff --git a/frontend/e2e/numeric-controls.spec.ts b/frontend/e2e/numeric-controls.spec.ts index 1b2d277d1a..fb3d8e3262 100644 --- a/frontend/e2e/numeric-controls.spec.ts +++ b/frontend/e2e/numeric-controls.spec.ts @@ -17,7 +17,8 @@ const DATASET_OPTIONS = { const DATASETS = { items: [ { name: 'harmbench', selection_key: 'dataset:named:harmbench', loaded: false, provider_available: true, origins: [], logical_examples: null }, - { name: 'custom_objectives', selection_key: 'dataset:named:custom_objectives', loaded: true, provider_available: false, origins: ['user'], logical_examples: 10 }, + { name: 'custom_objectives', selection_key: 'dataset:named:custom_objectives', loaded: true, provider_available: false, + origins: ['user'], logical_examples: 10, modalities: ['text', 'image_path'], harm_categories: [], has_unlabeled_harm_categories: false }, ], } @@ -58,6 +59,7 @@ const SCENARIO: RegisteredScenario = { } async function mockNumericControlApis(page: Page): Promise { + let datasetItems = DATASETS.items.map((item) => ({ ...item })) // No backend, model calls, or connections to another worktree's HMR server. await page.routeWebSocket('**/*', (socket) => socket.close()) await page.route('**/api/**', async (route) => { @@ -92,10 +94,19 @@ async function mockNumericControlApis(page: Page): Promise { ], }, '/api/labels': { source: 'attacks', labels: {} }, - '/api/datasets': DATASETS, + '/api/datasets': { items: datasetItems }, [`/api/scenarios/catalog/${SCENARIO.scenario_name}`]: SCENARIO, [`/api/scenarios/catalog/${SCENARIO.scenario_name}/estimate`]: SCENARIO.default_run_size, } + if (path.endsWith('/datasets/load')) { + const edits: ScenarioDatasetRequest = route.request().postDataJSON()?.datasets ?? {} + const names = (edits.sources ?? DATASET_DEFAULTS.sources).map((source) => source.name) + datasetItems = datasetItems.map((item) => names.includes(item.name) + ? { ...item, loaded: true, logical_examples: 200, modalities: ['text'], harm_categories: ['malware'] } + : item) + await route.fulfill({ json: { settings: DATASET_DEFAULTS, items: datasetItems.filter((item) => names.includes(item.name)) } }) + return + } if (path.endsWith('/estimate')) { const edits: ScenarioDatasetRequest = route.request().postDataJSON()?.datasets ?? {} const limit = (value: ScenarioDatasetRequest['max_total'], fallback: number | 'all'): number | 'all' => @@ -213,33 +224,41 @@ test('native shared dynamic parameter controls step once', async ({ page }) => { await expectNativeSteps(page.getByRole('spinbutton', { name: 'temperature', exact: true }), 0.5) }) -test('dataset limits support default and all and keep invalid text blocking after source reset', async ({ page }, testInfo) => { +test('dataset sources and explicit loading are separate from Parameters', async ({ page }, testInfo) => { await page.goto(`/scanner/${SCENARIO.scenario_name}`) - await expect(page.getByRole('table', { name: 'Selected datasets' })).toBeVisible() - await page.getByRole('combobox', { name: 'Objective Target' }).selectOption('mock-target') - const total = page.getByRole('textbox', { name: 'Max total', exact: true }) - await expect(total).toHaveValue('') - await page.getByRole('region', { name: 'Dataset configuration' }).screenshot({ path: testInfo.outputPath('datasets-default.png') }) - await total.fill('all') - await page.getByRole('combobox', { name: 'Add dataset' }).fill('custom_objectives') - await page.getByRole('button', { name: 'Add source', exact: true }).click() - await page.getByRole('textbox', { name: 'custom_objectives limit' }).fill('8') - await page.getByRole('combobox', { name: 'custom_objectives fetch' }).selectOption('never') + const section = page.getByRole('region', { name: 'Datasets', exact: true }) + await expect(section.getByRole('table', { name: 'Selected datasets' })).toBeVisible() + await expect(section.getByRole('columnheader', { name: 'Size' })).toBeVisible() + await expect(section.getByRole('button', { name: 'Edit harmbench objective limit' })).toHaveText('Up to 5') + await expect(section.getByText('Load dataset to see details', { exact: true })).toBeVisible() + await expect(page.getByRole('region', { name: 'Parameters' }).getByRole('table')).toHaveCount(0) + await expect(section.getByRole('button', { name: 'Edit Maximum objectives' })).toHaveText('Up to 5') + await expect(section.getByText('Details and advanced settings')).toHaveCount(0) + await section.screenshot({ path: testInfo.outputPath('datasets-default.png') }) + await section.getByRole('button', { name: 'Load datasets' }).click() + await expect(section.getByRole('button', { name: 'Edit harmbench objective limit' })).toHaveText('Up to 5 of 200') + await expect(section.getByRole('button', { name: 'Load datasets' })).toHaveCount(0) + await expect(section.getByRole('switch')).toHaveCount(0) + await expect(section.getByRole('button', { name: 'Filter harm category malware' })).toBeVisible() + await expect(section.getByRole('button', { name: 'Filter data type Text' })).toBeVisible() + await section.screenshot({ path: testInfo.outputPath('datasets-loaded.png') }) + await section.getByRole('button', { name: 'Edit Maximum objectives' }).click() + const total = page.getByRole('textbox', { name: 'Maximum objectives', exact: true }) await total.fill('12') - await expect(page.getByText('Default/effective: 12. Blank or default retains the scenario setting.')).toBeVisible() - await expect(page.getByText('Default/effective: 8. Blank or default retains the scenario setting.')).toBeVisible() - await page.getByRole('region', { name: 'Dataset configuration' }).screenshot({ path: testInfo.outputPath('datasets-edited.png') }) + await page.getByRole('button', { name: 'Done', exact: true }).click() + await section.screenshot({ path: testInfo.outputPath('datasets-custom.png') }) + await section.getByRole('button', { name: 'Edit Maximum objectives' }).click() await total.fill('0') - await expect(page.getByRole('alert')).toContainText('Max total must be a positive integer') - await page.getByRole('button', { name: 'Use default sources' }).click() - await expect(total).toHaveValue('0') - await expect(page.getByRole('alert')).toContainText('Max total must be a positive integer') + await total.press('Escape') + await section.getByRole('button', { name: 'Reset datasets' }).click() + await expect(section.getByRole('alert')).toContainText('Maximum objectives must be a positive integer') await expect(page.getByRole('button', { name: 'Launch scan', exact: true })).toBeDisabled() - await page.getByRole('region', { name: 'Dataset configuration' }).screenshot({ path: testInfo.outputPath('datasets-invalid.png') }) - await total.fill('default') - await expect(page.getByRole('alert')).not.toBeVisible() - await total.fill('') - await expect(total).toHaveValue('') + await section.screenshot({ path: testInfo.outputPath('datasets-invalid.png') }) + await section.getByRole('button', { name: 'Edit Maximum objectives' }).click() + await expect(total).toHaveValue('0') + await page.getByRole('button', { name: 'Reset to 5', exact: true }).click() + await page.getByRole('button', { name: 'Done', exact: true }).click() + await expect(section.getByRole('alert')).not.toBeVisible() }) test('native Azure ML advanced controls retain steps and reset to default guidance', async ({ page }) => { @@ -267,6 +286,43 @@ test('native Azure ML advanced controls retain steps and reset to default guidan await expect(resetTemperature).toHaveAttribute('placeholder', 'Defaults to 1.0') }) +test('dataset rows add immediately, edit in place, share filters, and reset without loading', async ({ page }) => { + await page.goto(`/scanner/${SCENARIO.scenario_name}`) + const section = page.getByRole('region', { name: 'Datasets', exact: true }) + await section.getByRole('combobox', { name: 'Add dataset' }).click() + await page.getByRole('option', { name: 'custom_objectives', exact: true }).click() + await expect(section.getByText('custom_objectives', { exact: true })).toBeVisible() + await section.getByRole('button', { name: 'Remove harmbench' }).click() + await expect(section.getByRole('button', { name: 'Load datasets' })).toHaveCount(0) + await section.getByRole('button', { name: 'Edit custom_objectives objective limit' }).click() + await page.getByRole('textbox', { name: 'custom_objectives objective limit' }).fill('3') + await page.getByRole('button', { name: 'Done', exact: true }).click() + await expect(section.getByRole('button', { name: 'Edit custom_objectives objective limit' })).toHaveText('Up to 3 of 10') + await section.getByRole('button', { name: 'Filter data type Image' }).click() + await expect(section.getByRole('button', { name: 'Filter data type Text' })).toHaveAttribute('aria-pressed', 'true') + await expect(section.getByRole('link', { name: 'Browse datasets' })).toHaveAttribute('href', '/datasets') + await expect(section.getByText('Filters:', { exact: true })).toHaveCount(0) + await section.getByRole('button', { name: 'Reset datasets' }).click() + await expect(section.getByText('harmbench', { exact: true })).toBeVisible() + await expect(section.getByRole('button', { name: 'Load datasets' })).toBeVisible() + await expect(section.getByRole('button', { name: 'Remove type filter text' })).toHaveCount(0) +}) + +test('dataset manual loading keeps an explicit load-on-start switch override', async ({ page }) => { + await page.goto(`/scanner/${SCENARIO.scenario_name}`) + const section = page.getByRole('region', { name: 'Datasets', exact: true }) + const policy = section.getByRole('switch', { name: 'Load missing data when the scan starts' }) + await expect(policy).toBeChecked() + await policy.click() + const load = page.waitForRequest((request) => new URL(request.url()).pathname.endsWith('/datasets/load')) + await section.getByRole('button', { name: 'Load datasets' }).click() + const request = await load + expect(request.postDataJSON().datasets.fetch).toBe('never') + expect(request.postDataJSON().datasets.sources[0].fetch).toBe('never') + await expect(section.getByRole('button', { name: 'Load datasets' })).toHaveCount(0) + await expect(policy).toHaveCount(0) +}) + test('native round-robin weights step once and preserve bounds and integer validation', async ({ page }) => { await page.goto('/registry/targets') await page.getByRole('button', { name: /new target/i }).click() diff --git a/frontend/e2e/touch-targets.spec.ts b/frontend/e2e/touch-targets.spec.ts index d75603ed57..7618194d20 100644 --- a/frontend/e2e/touch-targets.spec.ts +++ b/frontend/e2e/touch-targets.spec.ts @@ -476,7 +476,7 @@ async function startChatWithMessages(page: Page): Promise { await page.getByRole("button", { name: "Registry", exact: true }).click(); await expect(page.getByText("gpt-4o-mobile", { exact: true })).toBeVisible(); await page.getByRole("combobox", { name: "Default objective target", exact: true }).selectOption({ index: 1 }); - await page.getByRole("button", { name: "Chat", exact: true }).click(); + await page.getByRole("button", { name: "Technique.execute", exact: true }).click(); await page.getByTestId("chat-input").fill( "Assess this deterministic mobile prompt" ); @@ -501,7 +501,7 @@ test.describe("Mobile touch targets", () => { await page.getByRole("button", { name: "Registry", exact: true }).click(); await expect(page.getByText("gpt-4o-mobile", { exact: true })).toBeVisible(); await page.getByRole("combobox", { name: "Default objective target", exact: true }).selectOption({ index: 1 }); - await page.getByRole("button", { name: "Chat", exact: true }).click(); + await page.getByRole("button", { name: "Technique.execute", exact: true }).click(); await page.getByRole("button", { name: "Add objective" }).click(); const objectiveInput = page.getByRole("textbox", { diff --git a/frontend/src/App.labels.test.tsx b/frontend/src/App.labels.test.tsx index 9f17d1610f..935503bfc7 100644 --- a/frontend/src/App.labels.test.tsx +++ b/frontend/src/App.labels.test.tsx @@ -208,7 +208,7 @@ describe('Shared new run labels', () => { await chooseOperation(user, 'remembered_op') const bar = currentLabels() - for (const destination of ['Home', 'Chat', 'Scanner']) { + for (const destination of ['Home', 'Technique.execute', 'Scanner']) { await user.click(screen.getByRole('button', { name: destination, exact: true })) expect(screen.getAllByTestId('labels-bar')).toHaveLength(1) expect(currentLabels()).toBe(bar) @@ -273,7 +273,7 @@ describe('Shared new run labels', () => { expect(screen.getAllByTestId('labels-bar')).toHaveLength(1) await user.click(screen.getByRole('button', { name: 'Scanner', exact: true })) expect(screen.queryByRole('group', { name: 'Chat controls' })).not.toBeInTheDocument() - await user.click(screen.getByRole('button', { name: 'Chat', exact: true })) + await user.click(screen.getByRole('button', { name: 'Technique.execute', exact: true })) expect(screen.getAllByRole('group', { name: 'Chat controls' })).toHaveLength(1) expect(within(currentLabels()).getByRole('switch')).toBeChecked() }) diff --git a/frontend/src/components/Scenarios/DatasetLimitControl.tsx b/frontend/src/components/Scenarios/DatasetLimitControl.tsx new file mode 100644 index 0000000000..a3d1fce71d --- /dev/null +++ b/frontend/src/components/Scenarios/DatasetLimitControl.tsx @@ -0,0 +1,63 @@ +import { useState } from 'react' + +import { Button, Checkbox, Field, Input, Popover, PopoverSurface, PopoverTrigger, Text } from '@fluentui/react-components' + +import type { ScenarioDatasetLimit } from '@/types' +import { useScenarioDatasetTableStyles } from './ScenarioDatasetTable.styles' +import { parseCaseLimit } from './datasetPresentation' + +interface DatasetLimitControlProps { + label: string + value: ScenarioDatasetLimit | null | undefined + defaultValue: number | 'all' + text?: string + showLabel?: boolean + hint?: string + draft?: string + disabled: boolean + onChange: (value: ScenarioDatasetLimit) => void + onDraftChange: (text: string | null) => void +} + +export default function DatasetLimitControl({ + label, value, defaultValue, text, showLabel = true, hint, draft, disabled, onChange, onDraftChange, +}: DatasetLimitControlProps) { + const styles = useScenarioDatasetTableStyles() + const [open, setOpen] = useState(false) + const limit = typeof value === 'number' || value === 'all' ? value : defaultValue + const invalid = draft !== undefined && parseCaseLimit(draft) === null + const control = setOpen(data.open)}> + + + + +
+ + { + onDraftChange(data.value) + const parsed = parseCaseLimit(data.value) + if (parsed !== null) onChange(parsed) + }} /> + + { + onDraftChange(null) + onChange(data.checked ? 'all' : typeof defaultValue === 'number' ? defaultValue : 1) + }} /> + {hint && {hint}} +
+ + +
+
+ + + return showLabel ?
{label}:{control}
: control +} diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts index e40aea247c..fdc7ebb74a 100644 --- a/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts @@ -4,10 +4,25 @@ import { mobileTouchTarget } from '@/styles/touchTargets' export const useScenarioDatasetTableStyles = makeStyles({ root: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalM }, - controls: { display: 'flex', flexWrap: 'wrap', gap: tokens.spacingHorizontalM }, + controls: { display: 'flex', flexWrap: 'wrap', alignItems: 'center', gap: tokens.spacingHorizontalM }, + browseLink: { color: tokens.colorBrandForegroundLink, ...mobileTouchTarget, display: 'inline-flex', alignItems: 'center' }, scroll: { overflowX: 'auto' }, - table: { width: '100%', borderCollapse: 'collapse' }, - cell: { padding: tokens.spacingHorizontalS, textAlign: 'left' }, - input: { maxWidth: '12em' }, + table: { width: '100%', minWidth: '36em', borderCollapse: 'collapse' }, + toolbar: { display: 'flex', flexWrap: 'wrap', alignItems: 'center', gap: tokens.spacingHorizontalM }, + metadata: { display: 'flex', flexWrap: 'wrap', alignItems: 'center', gap: tokens.spacingHorizontalXS }, + chip: { ...mobileTouchTarget, padding: tokens.spacingHorizontalXXS, minWidth: 'unset' }, + actionsCell: { width: '3em' }, + datasetName: { overflowWrap: 'anywhere' }, + limitButton: { + ...mobileTouchTarget, + justifyContent: 'flex-start', + paddingLeft: tokens.spacingHorizontalXS, + paddingRight: tokens.spacingHorizontalXS, + color: tokens.colorBrandForeground1, + ':hover': { textDecorationLine: 'underline' }, + }, + limitEditor: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalM, maxWidth: '24em' }, + totalLimit: { display: 'flex', flexWrap: 'wrap', alignItems: 'center', gap: tokens.spacingHorizontalXS }, + loadingRule: { display: 'inline-flex' }, button: { ...mobileTouchTarget }, }) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx index 3de4f9480d..964247c1e0 100644 --- a/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx @@ -3,172 +3,277 @@ import { useState } from 'react' import { render, screen, within } from '@testing-library/react' import userEvent from '@testing-library/user-event' import { FluentProvider, webLightTheme } from '@fluentui/react-components' +import { MemoryRouter } from 'react-router' -import type { ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings } from '@/types' +import type { DatasetInfo, ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings } from '@/types' import ScenarioDatasetTable from './ScenarioDatasetTable' const DEFAULTS: ScenarioDatasetSettings = { sources: [{ name: 'stored', max_size: 2, fetch: 'never', configured_provider: false }], - max_per_dataset: 5, max_total: 'all', fetch: 'if_missing', filters: {}, + max_per_dataset: 5, max_total: 4, fetch: 'if_missing', filters: {}, } const OPTIONS: ScenarioDatasetOptions = { mode: 'any', allowed_names: [], per_dataset_limit: true, total_limit: true, supports_overrides: true, filters: true, } +const STORED: DatasetInfo = { + name: 'stored', selection_key: 'dataset:named:stored', loaded: true, provider_available: false, + origins: ['local'], logical_examples: 8, seed_pieces: 8, objectives: 8, + modalities: ['text'], harm_categories: ['hate', 'violence'], has_unlabeled_harm_categories: false, +} +const REMOTE: DatasetInfo = { + ...STORED, name: 'remote', selection_key: 'dataset:named:remote', loaded: false, + provider_available: true, logical_examples: null, origins: [], modalities: [], harm_categories: [], +} +const DATASETS = [STORED, REMOTE, { ...STORED, name: 'other', selection_key: 'dataset:named:other' }] interface HarnessProps { options?: ScenarioDatasetOptions defaults?: ScenarioDatasetSettings - onChange: jest.Mock + datasets?: DatasetInfo[] + initialRequest?: ScenarioDatasetRequest + onChange?: jest.Mock + onLoad?: jest.Mock } - -function Harness({ options = OPTIONS, defaults = DEFAULTS, onChange }: HarnessProps) { - const [request, setRequest] = useState({}) +function Harness({ + options = OPTIONS, defaults = DEFAULTS, datasets = DATASETS, initialRequest = {}, onChange = jest.fn(), onLoad = jest.fn(), +}: HarnessProps) { + const [request, setRequest] = useState(initialRequest) const [validation, setValidation] = useState(null) - return { setRequest(next); onChange(next) }} - /> + const effective: ScenarioDatasetSettings = { + ...defaults, + sources: (request.sources ?? defaults.sources).map((source) => ({ + name: source.name, + max_size: typeof source.max_size === 'number' || source.max_size === 'all' + ? source.max_size : defaults.sources.find((item) => item.name === source.name)?.max_size ?? defaults.max_per_dataset, + fetch: source.fetch === 'never' || source.fetch === 'if_missing' ? source.fetch + : defaults.sources.find((item) => item.name === source.name)?.fetch ?? defaults.fetch, + configured_provider: false, + })), + } + return + { setRequest(next); onChange(next) }} /> + + } -describe('ScenarioDatasetTable', () => { - beforeEach(() => jest.clearAllMocks()) +async function editLimit(user: ReturnType, label: string, value: string): Promise { + await user.click(screen.getByRole('button', { name: `Edit ${label}` })) + const all = screen.getByRole('checkbox', { name: `All available for ${label}` }) + if (all instanceof HTMLInputElement && all.checked) await user.click(all) + const input = screen.getByRole('textbox', { name: label }) + await user.clear(input) + await user.type(input, value) +} - it('shows mixed origins, stored counts, and unknown provider-only counts without sending defaults', async () => { - const user = userEvent.setup() +async function addDataset(user: ReturnType, name: string): Promise { + await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) + await user.click(screen.getByRole('option', { name, exact: true })) +} + +describe('ScenarioDatasetTable', () => { + it('shows a compact table with resolved sizes and metadata, without a loading banner for loaded data', () => { const onChange = jest.fn() render() - expect(screen.getByText('local, unknown')).toBeInTheDocument() - expect(screen.getByText('8')).toBeInTheDocument() + expect(screen.getByRole('columnheader', { name: 'Dataset', exact: true })).toBeInTheDocument() + expect(screen.getByRole('columnheader', { name: 'Size' })).toBeInTheDocument() + expect(screen.getByRole('columnheader', { name: 'Include Criteria' })).toBeInTheDocument() + expect(screen.getByRole('separator')).toBeInTheDocument() + expect(screen.queryByText(/All selected uses the defaults/)).not.toBeInTheDocument() + expect(screen.getByText('Maximum objectives:', { exact: true })).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Edit stored objective limit' })).toHaveTextContent('Up to 2 of 8') + expect(screen.getByRole('button', { name: 'Filter harm category hate' })).toHaveAttribute('aria-pressed', 'true') + expect(screen.getByRole('button', { name: 'Filter data type Text' })).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Load datasets' })).not.toBeInTheDocument() + expect(screen.queryByRole('switch')).not.toBeInTheDocument() + expect(screen.queryByText('All selected datasets are loaded.')).not.toBeInTheDocument() + expect(screen.queryByText('Details and advanced settings')).not.toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Edit Maximum objectives' })).toHaveTextContent('Up to 4') expect(onChange).not.toHaveBeenCalled() - await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'remote') - await user.click(screen.getByRole('button', { name: 'Add source' })) - const row = screen.getByRole('rowheader', { name: 'remote' }).closest('tr') - expect(row).not.toBeNull() - if (row) { - expect(within(row).getByText('Provider available')).toBeInTheDocument() - expect(within(row).getAllByText('Unknown')).toHaveLength(2) - } + }) + + it('adds a provider immediately from a name-only dropdown without fetching it', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + const onLoad = jest.fn() + render() + await addDataset(user, 'remote') expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'stored' }, { name: 'remote' }] }) + expect(screen.getByRole('button', { name: 'Edit remote objective limit' })).toHaveTextContent('Up to 4') + expect(screen.getByText('Load dataset to see details')).toBeInTheDocument() + expect(onLoad).not.toHaveBeenCalled() + await user.click(screen.getByRole('button', { name: 'Load datasets' })) + expect(onLoad).toHaveBeenCalledTimes(1) + }) + + it('keeps required reference sources complete and does not expose source edits', () => { + render() + expect(screen.getByText('8 examples')).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Edit stored objective limit' })).not.toBeInTheDocument() + expect(screen.queryByRole('combobox', { name: 'Add dataset' })).not.toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Remove stored' })).not.toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Edit Maximum objectives' })).toBeInTheDocument() }) - it('edits source caps and policies independently of global limits with keyboard-accessible controls', async () => { + it('supports inline numeric, all, and inherited objective limits', async () => { const user = userEvent.setup() const onChange = jest.fn() render() - await user.type(screen.getByRole('textbox', { name: 'stored limit' }), 'all') - await user.selectOptions(screen.getByRole('combobox', { name: 'stored fetch' }), 'if_missing') - await user.type(screen.getByRole('textbox', { name: 'Max total' }), '12') - await user.selectOptions(screen.getByRole('combobox', { name: 'Default dataset fetch' }), 'never') - expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ - sources: [{ name: 'stored', max_size: 'all', fetch: 'if_missing' }], max_total: 12, fetch: 'never', - }) - await user.clear(screen.getByRole('textbox', { name: 'Max total' })) - expect(onChange.mock.calls.at(-1)?.[0]).not.toHaveProperty('max_total') - screen.getByRole('textbox', { name: 'Max per dataset' }).focus() - await user.keyboard('3') - expect(onChange.mock.calls.at(-1)?.[0]).toHaveProperty('max_per_dataset', 3) + await editLimit(user, 'Maximum objectives', '12') + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ max_total: 12 }) + await user.click(screen.getByRole('checkbox', { name: 'All available for Maximum objectives' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ max_total: 'all' }) + await user.click(screen.getByRole('button', { name: 'Reset to 4' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) + }) + + it('keeps invalid global drafts blocking when datasets reset', async () => { + const user = userEvent.setup() + render() + await editLimit(user, 'Maximum objectives', '0') + await user.keyboard('{Escape}') + await user.click(screen.getByRole('button', { name: 'Reset datasets' })) + expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() + expect(screen.getByRole('alert')).toHaveTextContent('Maximum objectives must be a positive integer') + await user.click(screen.getByRole('button', { name: 'Edit Maximum objectives' })) + expect(screen.getByRole('textbox', { name: 'Maximum objectives' })).toHaveValue('0') + await user.click(screen.getByRole('button', { name: 'Reset to 4' })) + expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() }) - it('does not send fixed or unsupported inputs and supports one-source replacement', async () => { + it('removes only the invalid row draft and restores other rows and limits on reset', async () => { const user = userEvent.setup() const onChange = jest.fn() - const view = render() - expect(screen.queryByRole('button', { name: 'Add source' })).not.toBeInTheDocument() - expect(screen.queryByRole('textbox', { name: 'stored limit' })).not.toBeInTheDocument() - expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() - expect(onChange).not.toHaveBeenCalled() - view.rerender() - await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) - await user.click(screen.getByRole('option', { name: 'remote' })) - await user.click(screen.getByRole('button', { name: 'Add source' })) + render() + await addDataset(user, 'other') + await editLimit(user, 'other objective limit', '0') + await user.keyboard('{Escape}') + expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() + await user.click(screen.getByRole('button', { name: 'Remove other' })) + expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() + expect(onChange.mock.calls.at(-1)?.[0].sources).toEqual([{ name: 'stored' }]) + expect(screen.queryByRole('button', { name: 'Remove stored' })).not.toBeInTheDocument() + await editLimit(user, 'stored objective limit', '3') + await user.keyboard('{Escape}') + await user.click(screen.getByRole('button', { name: 'Reset datasets' })) + expect(screen.getByRole('button', { name: 'Edit stored objective limit' })).toHaveTextContent('Up to 2 of 8') + }) + + it('replaces a one-source subset without permitting invalid or additional sources', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + render() + await user.click(screen.getByRole('combobox', { name: 'Choose dataset' })) + expect(screen.queryByRole('option', { name: 'other', exact: true })).not.toBeInTheDocument() + await user.click(screen.getByRole('option', { name: 'remote', exact: true })) expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'remote' }] }) - expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() }) - it('retains edits and displays invalid selections when parameter-dependent options change', async () => { + it('shares category and type filters across rows and preserves inherited filters when deselected', async () => { const user = userEvent.setup() const onChange = jest.fn() - const view = render() - await user.type(screen.getByRole('textbox', { name: 'stored limit' }), '1') - view.rerender() - expect(screen.getByRole('alert')).toHaveTextContent('Invalid selection for current parameters: stored') - expect(screen.getByRole('textbox', { name: 'stored limit' })).toHaveValue('1') - await user.click(screen.getByRole('button', { name: 'Use default sources' })) - expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() - expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) + render() + await addDataset(user, 'other') + const hate = screen.getAllByRole('button', { name: 'Filter harm category hate' }) + expect(hate.every((button) => button.getAttribute('aria-pressed') === 'true')).toBe(true) + await user.click(screen.getAllByRole('button', { name: 'Filter harm category violence' })[0]) + await user.click(hate[0]) + expect(onChange.mock.calls.at(-1)?.[0].filters).toEqual({ harm_categories: ['violence'] }) + expect(hate.every((button) => button.getAttribute('aria-pressed') === 'false')).toBe(true) + await user.click(hate[0]) + expect(onChange.mock.calls.at(-1)?.[0].filters.harm_categories).toEqual([]) + expect(screen.queryByText('Filters:')).not.toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Remove type filter text' })).not.toBeInTheDocument() }) - it.each([ - ['Max total', 'max_total', '3'], ['Max total', 'max_total', 'all'], - ['Max per dataset', 'max_per_dataset', '3'], ['Max per dataset', 'max_per_dataset', 'all'], - ])('keeps invalid %s blocking launch after source reset (%s=%s)', async (label, key, initial) => { + + it('preserves custom inherited filters without rendering a separate filter list', async () => { const user = userEvent.setup() const onChange = jest.fn() - render() - const input = screen.getByRole('textbox', { name: label }) - await user.type(input, initial) - await user.click(input) - await user.keyboard('{Control>}a{/Control}0') - expect(input).toHaveValue('0') - expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() - await user.click(screen.getByRole('button', { name: 'Use default sources' })) - expect(input).toHaveValue('0') - expect(screen.getByRole('alert')).toHaveTextContent(`${label} must be a positive integer`) - expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() - expect(onChange.mock.calls.at(-1)?.[0]).toHaveProperty(key, initial === 'all' ? 'all' : 3) - await user.click(input) - await user.keyboard('{Control>}a{/Control}4') - expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() + render() + await user.click(screen.getByRole('button', { name: 'Filter harm category hate' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ filters: { harm_categories: ['custom-substring', 'hate'] } }) + expect(screen.queryByText('Filters:')).not.toBeInTheDocument() }) - it('clears only removed source errors and preserves other invalid fields', async () => { + it('starts with every label checked and restores the default request after reselecting all labels', async () => { const user = userEvent.setup() const onChange = jest.fn() render() - await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'remote') - await user.click(screen.getByRole('button', { name: 'Add source' })) - await user.type(screen.getByRole('textbox', { name: 'remote limit' }), '0') - await user.type(screen.getByRole('textbox', { name: 'Max total' }), '0') - await user.click(screen.getByRole('button', { name: 'Remove remote' })) - expect(screen.queryByRole('textbox', { name: 'remote limit' })).not.toBeInTheDocument() - expect(screen.getByRole('alert')).toHaveTextContent('Max total must be a positive integer') - expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() - await user.clear(screen.getByRole('textbox', { name: 'Max total' })) - expect(screen.queryByRole('alert')).not.toBeInTheDocument() - expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() - await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'remote') - await user.click(screen.getByRole('button', { name: 'Add source' })) - expect(screen.getByRole('textbox', { name: 'remote limit' })).toHaveValue('') + expect(screen.getByRole('button', { name: 'Filter data type Text' })).toHaveAttribute('aria-pressed', 'true') + await user.click(screen.getByRole('button', { name: 'Filter harm category hate' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ filters: { harm_categories: ['violence'] } }) + await user.click(screen.getByRole('button', { name: 'Filter harm category hate' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) }) - it('resets source drafts even when default sources have the same names', async () => { + it('links to the dataset browser next to reset and for required sources', () => { + const view = render() + expect(screen.getByRole('link', { name: 'Browse datasets' })).toHaveAttribute('href', '/datasets') + expect(screen.getByRole('button', { name: 'Reset datasets' })).toBeInTheDocument() + view.rerender() + expect(screen.getByRole('link', { name: 'Browse datasets' })).toHaveAttribute('href', '/datasets') + }) + + it('blocks an empty selection instead of sending an empty list that means include all', async () => { const user = userEvent.setup() - render() - await user.type(screen.getByRole('textbox', { name: 'stored limit' }), '0') + const onChange = jest.fn() + render() + await user.click(screen.getByRole('button', { name: 'Filter data type Text' })) + expect(screen.getByRole('button', { name: 'Filter data type Text' })).toHaveAttribute('aria-pressed', 'false') expect(screen.getByRole('button', { name: 'Launch' })).toBeDisabled() - await user.click(screen.getByRole('button', { name: 'Use default sources' })) - expect(screen.getByRole('textbox', { name: 'stored limit' })).toHaveValue('') + expect(screen.getByRole('alert')).toHaveTextContent('Select at least one data type.') + expect(onChange).not.toHaveBeenCalled() + await user.click(screen.getByRole('button', { name: 'Filter data type Text' })) expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) }) - it('replacing the sole selected source removes its invalid draft', async () => { + it('uses a single switch that updates all selected source loading policies', async () => { const user = userEvent.setup() - render() - await user.type(screen.getByRole('textbox', { name: 'stored limit' }), '0') - await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) - await user.click(screen.getByRole('option', { name: 'remote' })) - await user.click(screen.getByRole('button', { name: 'Add source' })) - expect(screen.queryByRole('alert')).not.toBeInTheDocument() - expect(screen.getByRole('button', { name: 'Launch' })).toBeEnabled() + const onChange = jest.fn() + render() + await addDataset(user, 'remote') + await user.click(screen.getByRole('switch', { name: 'Load missing data when the scan starts' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ + fetch: 'if_missing', sources: [{ name: 'stored', fetch: 'if_missing' }, { name: 'remote', fetch: 'if_missing' }], + }) + await user.click(screen.getByRole('switch', { name: 'Load missing data when the scan starts' })) + expect(onChange.mock.calls.at(-1)?.[0].sources.every((source: { fetch: string }) => source.fetch === 'never')).toBe(true) + expect(screen.getByRole('button', { name: 'Load datasets' })).toBeEnabled() }) + it('keeps missing-data loading on by default and hides all loading controls once data is loaded', () => { + const defaults = { ...DEFAULTS, sources: [{ name: 'remote', max_size: 4 as const, fetch: 'if_missing' as const, configured_provider: false }] } + const view = render() + expect(screen.getByRole('switch', { name: 'Load missing data when the scan starts' })).toBeChecked() + view.rerender() + expect(screen.queryByRole('button', { name: 'Load datasets' })).not.toBeInTheDocument() + expect(screen.queryByRole('switch')).not.toBeInTheDocument() + expect(screen.queryByText(/All selected datasets are loaded/)).not.toBeInTheDocument() + }) + + it('does not assume stored counts are exact selected sizes after filters and a total cap', () => { + render() + expect(screen.getByRole('button', { name: 'Edit stored objective limit' })).toHaveTextContent('Up to 1 of 8') + const row = screen.getByRole('row', { name: /stored/ }) + expect(within(row).queryByText('1 of 8', { exact: true })).not.toBeInTheDocument() + }) + + it('does not expose loading or source selection for an inline-only scenario', () => { + render() + expect(screen.getByText('This scenario has no named datasets.')).toBeInTheDocument() + expect(screen.queryByRole('button', { name: 'Load datasets' })).not.toBeInTheDocument() + expect(screen.queryByRole('combobox', { name: 'Add dataset' })).not.toBeInTheDocument() + }) + + it('does not pretend a global policy can replace protected source-specific loading rules', () => { + render() + expect(screen.getByRole('switch', { name: 'Load missing data when the scan starts' })).not.toBeChecked() + expect(screen.getByRole('switch', { name: 'Load missing data when the scan starts' })).toBeDisabled() + expect(screen.getByRole('button', { name: 'Load datasets' })).toBeEnabled() + }) }) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx index ac23988494..ffe344634c 100644 --- a/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx @@ -1,12 +1,20 @@ import { useEffect, useState } from 'react' -import { Button, Combobox, Field, Input, MessageBar, MessageBarBody, Option, Select, Text } from '@fluentui/react-components' +import { + Badge, Button, Combobox, Divider, MessageBar, MessageBarBody, Option, Switch, + Table, TableBody, TableCell, TableHeader, TableHeaderCell, TableRow, Text, ToggleButton, Tooltip, +} from '@fluentui/react-components' +import { DismissRegular } from '@fluentui/react-icons' +import { Link } from 'react-router' +import { MODALITY_LABELS } from '@/components/Config/targetFilters' import type { - DatasetInfo, ScenarioDatasetLimit, ScenarioDatasetOptions, ScenarioDatasetRequest, + DatasetInfo, DatasetFetchPolicy, ScenarioDatasetLimit, ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings, ScenarioDatasetSourceInput, } from '@/types' +import DatasetLimitControl from './DatasetLimitControl' +import { datasetChoices, datasetFacts, datasetSize, parseCaseLimit } from './datasetPresentation' import { useScenarioDatasetTableStyles } from './ScenarioDatasetTable.styles' interface ScenarioDatasetTableProps { @@ -18,189 +26,228 @@ interface ScenarioDatasetTableProps { disabled: boolean loading: boolean metadataError: string | null + preparing?: boolean + loadError?: string | null + loadEnabled?: boolean + onLoad?: () => void onReload: () => void onChange: (request: ScenarioDatasetRequest) => void onValidationChange: (error: string | null) => void } -interface DatasetLimitFieldProps { - label: string - text: string - effective: number | 'all' - disabled: boolean - onChange: (text: string) => void -} - -function parseLimit(text: string): ScenarioDatasetLimit | null { - const value = text.trim().toLowerCase() - if (value === '' || value === 'default') return 'default' - if (value === 'all') return 'all' - const number = Number(value) - return /^\d+$/.test(value) && Number.isSafeInteger(number) && number > 0 ? number : null -} - -function DatasetLimitField({ label, text, effective, disabled, onChange }: DatasetLimitFieldProps) { - const styles = useScenarioDatasetTableStyles() - return ( - - onChange(data.value)} /> - - ) -} - -/** Displays effective settings while emitting only user edits, never catalog defaults. */ export default function ScenarioDatasetTable({ request, defaults, effective, options, datasets, disabled, loading, metadataError, - onReload, onChange, onValidationChange, + preparing = false, loadError, loadEnabled = false, onLoad, onReload, onChange, onValidationChange, }: ScenarioDatasetTableProps) { const styles = useScenarioDatasetTableStyles() - const [newName, setNewName] = useState('') + const [search, setSearch] = useState('') const [drafts, setDrafts] = useState>({}) - const [selectionError, setSelectionError] = useState(null) - const selected: ScenarioDatasetSourceInput[] = request.sources ?? effective.sources.map((source) => ({ name: source.name })) - const canSelect = options.supports_overrides && options.mode !== 'fixed' - const canEdit = options.supports_overrides && options.source_settings !== false - const candidates = datasets.filter((dataset) => options.mode === 'any' || options.allowed_names.includes(dataset.name)) - const limitText = (key: string, value: ScenarioDatasetLimit | null | undefined): string => - drafts[key] ?? (value == null || value === 'default' ? '' : String(value)) - const updateLimit = (key: string, text: string, apply: (value: ScenarioDatasetLimit) => void): void => { - setDrafts((previous) => ({ ...previous, [key]: text })) - const value = parseLimit(text) - if (value !== null) apply(value) + const [emptySelections, setEmptySelections] = useState([]) + const selected: ScenarioDatasetSourceInput[] = request.sources ?? defaults.sources.map((source) => ({ name: source.name })) + const canSelect = options.supports_overrides && options.mode !== 'fixed' && selected.length > 0 + const canEditSources = options.supports_overrides && options.source_settings !== false + const canLimit = canEditSources && options.per_dataset_limit + const canFilter = options.supports_overrides && options.filters + const canAdd = options.max_selected == null || selected.length < options.max_selected + const replaceSingleSource = options.max_selected === 1 + const busy = disabled || preparing + const total = typeof request.max_total === 'number' || request.max_total === 'all' + ? request.max_total : defaults.max_total + const filters = { ...defaults.filters, ...request.filters } + const candidates = datasetChoices(datasets, defaults.sources, options) + .filter((name) => !selected.some((source) => source.name === name) && name.toLowerCase().includes(search.toLowerCase())) + const { loaded, allLoaded } = datasetFacts(datasets, selected.map((source) => source.name)) + const availableMetadata: Record = { + harm_categories: [...new Set(loaded.flatMap((dataset) => dataset.harm_categories))], + data_types: [...new Set(loaded.flatMap((dataset) => dataset.modalities))], + } + const selectedMetadata = (key: string): string[] => emptySelections.includes(key) + ? [] : filters[key]?.length ? filters[key] : availableMetadata[key] ?? [] + const metadataSelected = (key: string, value: string): boolean => selectedMetadata(key).includes(value) + const settingsFor = (name: string) => effective.sources.find((source) => source.name === name) + ?? defaults.sources.find((source) => source.name === name) + const requiresImport = !loading && !metadataError && selected.some((source) => { + const metadata = datasets.find((item) => item.name === source.name) + return !metadata?.loaded && !metadata?.provider_available && !settingsFor(source.name)?.configured_provider + }) + const updateDraft = (key: string, text: string | null): void => { + setDrafts((previous) => { + const next = { ...previous } + if (text === null) delete next[key] + else next[key] = text + return next + }) } const clearSourceDrafts = (name?: string): void => { setDrafts((previous) => Object.fromEntries(Object.entries(previous).filter(([key]) => name === undefined ? !key.startsWith('source:') : key !== `source:${name}`, ))) - setSelectionError(null) - } - const updateSource = (name: string, edits: Partial): void => { - onChange({ ...request, sources: selected.map((source) => source.name === name ? { ...source, ...edits } : source) }) } - const updateGlobal = (key: 'max_per_dataset' | 'max_total', value: ScenarioDatasetLimit): void => { + const updateTotal = (value: ScenarioDatasetLimit): void => { const edits = { ...request } - if (value === 'default') delete edits[key] - else edits[key] = value + if (value === 'default') delete edits.max_total + else edits.max_total = value onChange(edits) } - const invalidNames = selected.filter((source) => - options.mode === 'subset' && !options.allowed_names.includes(source.name), - ) - const limits: Array<[string, string]> = [] - if (options.supports_overrides) { - if (options.per_dataset_limit) limits.push(['Max per dataset', limitText('perDataset', request.max_per_dataset)]) - if (options.total_limit) limits.push(['Max total', limitText('total', request.max_total)]) - if (canEdit && options.per_dataset_limit) { - for (const source of selected) { - limits.push([`${source.name} limit`, limitText(`source:${source.name}`, source.max_size)]) - } + const toggleFilter = (key: string, value: string): void => { + const current = selectedMetadata(key) + const next = current.includes(value) ? current.filter((item) => item !== value) : [...current, value] + if (next.length === 0) { + setEmptySelections((previous) => [...new Set([...previous, key])]) + return } + setEmptySelections((previous) => previous.filter((item) => item !== key)) + const edits = { ...request.filters } + const allSelected = (availableMetadata[key] ?? []).every((item) => next.includes(item)) + && next.every((item) => (availableMetadata[key] ?? []).includes(item)) + if (allSelected && !(defaults.filters[key]?.length)) delete edits[key] + else edits[key] = allSelected ? [] : next + const result = { ...request } + if (Object.keys(edits).length > 0) result.filters = edits + else delete result.filters + onChange(result) } - const errors = limits.filter(([, text]) => parseLimit(text) === null) - .map(([label]) => `${label} must be a positive integer, all, or default.`) - if (invalidNames.length > 0) { - errors.push(`Invalid selection for current parameters: ${invalidNames.map((source) => source.name).join(', ')}`) + const labels: Record = {} + if (options.supports_overrides && options.total_limit) labels.total = 'Maximum objectives' + if (canLimit) for (const source of selected) labels[`source:${source.name}`] = `${source.name} objective limit` + const errors = Object.entries(labels) + .filter(([key]) => drafts[key] !== undefined && parseCaseLimit(drafts[key]) === null) + .map(([, label]) => `${label} must be a positive integer.`) + if (canFilter) for (const key of emptySelections) { + if (availableMetadata[key]?.length) errors.push(`Select at least one ${key === 'harm_categories' ? 'harm category' : 'data type'}.`) + } + if (options.mode === 'subset') { + const invalid = selected.filter((source) => !options.allowed_names.includes(source.name)) + if (invalid.length > 0) errors.push(`Select a permitted source: ${invalid.map((source) => source.name).join(', ')}`) } - if (selectionError) errors.push(selectionError) const validationError = errors[0] ?? null useEffect(() => { onValidationChange(validationError) }, [validationError, onValidationChange]) + + const policies = selected.map((source) => !canEditSources + ? settingsFor(source.name)?.fetch ?? defaults.fetch + : source.fetch ?? ( + request.fetch === 'never' || request.fetch === 'if_missing' ? request.fetch : settingsFor(source.name)?.fetch ?? defaults.fetch + )) + const mixedPolicies = policies.some((policy) => policy !== policies[0]) + const canSetLoading = options.supports_overrides && (canEditSources || ( + !mixedPolicies && policies.every((policy) => policy === effective.fetch) + )) + const changeLoading = (policy: DatasetFetchPolicy): void => onChange({ + ...request, fetch: policy, + ...(canEditSources ? { sources: selected.map((source) => ({ ...source, fetch: policy })) } : {}), + }) + const loadingDisabledReason = !options.supports_overrides + ? 'Loading rules are fixed by this scenario.' + : !canSetLoading ? 'This scenario uses source-specific loading rules.' : '' + return ( -
- Datasets - - {metadataError && {metadataError}} - {errors.map((error) => {error})} - {options.supports_overrides &&
- {options.per_dataset_limit && updateLimit('perDataset', text, (value) => updateGlobal('max_per_dataset', value))} />} - {options.total_limit && updateLimit('total', text, (value) => updateGlobal('max_total', value))} />} - - - +
+ {!allLoaded && selected.length > 0 &&
+ {onLoad && } + {requiresImport && Import the missing datasets into PyRIT first.} + {metadataError && }
} -
- - {['Name', 'Stored origins', 'Stored examples', 'Limit', 'Fetch policy', 'Selection'].map((label) => - )} - {selected.map((source) => { - const metadata = datasets.find((dataset) => dataset.name === source.name) - const settings = effective.sources.find((item) => item.name === source.name) - ?? defaults.sources.find((item) => item.name === source.name) - const hasProvider = metadata?.provider_available || settings?.configured_provider - return - - - - - - - - })} -
{label}
{source.name}{metadata?.origins?.join(', ') || 'Unknown'} -
{hasProvider ? 'Provider available' : metadata ? 'Memory only' : 'Availability unknown'}
{metadata?.logical_examples ?? 'Unknown'}{canEdit && options.per_dataset_limit - ? updateLimit(`source:${source.name}`, text, (value) => updateSource(source.name, { max_size: value }))} /> - : settings?.max_size ?? 'all'}{canEdit - ? : settings?.fetch ?? effective.fetch}{canSelect && }
-
- {canSelect &&
- - setNewName(event.target.value)} - onOptionSelect={(_, data) => setNewName(data.optionValue ?? '')}> - {candidates.map((dataset) => )} - - - - + setSearch('') + onChange(edits) + }}>Reset datasets + } + Browse datasets
} -
+ {selected.length > 0 && } + {canFilter && emptySelections.length > 0 && } + {!allLoaded && selected.length > 0 && + + policy === 'if_missing')} + disabled={busy || !canSetLoading} onChange={(_, data) => changeLoading(data.checked ? 'if_missing' : 'never')} /> + + } + {!allLoaded && mixedPolicies && The selected datasets currently use different loading rules.} + {options.supports_overrides && options.total_limit && updateDraft('total', text)} onChange={updateTotal} />} + {options.total_limit && Each objective can run through several techniques} +
) } diff --git a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx index 87b100d47e..66ba959d6f 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx @@ -17,7 +17,11 @@ import ScenarioDetail from './ScenarioDetail' jest.mock('@/hooks/useRuntime', () => ({ useRuntime: jest.fn() })) jest.mock('@/services/api', () => ({ - datasetsApi: { listDatasets: jest.fn().mockResolvedValue({ items: [] }) }, + datasetsApi: { listDatasets: jest.fn().mockResolvedValue({ items: ['harmbench', 'ds_a', 'ds_b', 'remote', 'stale-dataset'].map((name) => ({ + name, selection_key: `dataset:named:${name}`, loaded: true, provider_available: true, + logical_examples: 8, seed_pieces: null, objectives: null, modalities: ['text', 'image_path', 'audio_path'], + harm_categories: ['cyber', 'violence', 'other'], has_unlabeled_harm_categories: false, + })) }) }, scenariosApi: { estimateRun: jest.fn(), getScenario: jest.fn(), @@ -164,6 +168,20 @@ async function confirmRunPreview(user: ReturnType): Prom await user.click(screen.getByTestId('confirm-launch-scenario-btn')) } +async function openObjectiveLimit(user: ReturnType): Promise { + await user.click(await screen.findByRole('button', { name: 'Edit Maximum objectives' })) +} + +async function setObjectiveLimit(user: ReturnType, value: string): Promise { + await openObjectiveLimit(user) + const all = screen.getByRole('checkbox', { name: 'All available for Maximum objectives' }) + if (all instanceof HTMLInputElement && all.checked) await user.click(all) + const input = screen.getByRole('textbox', { name: 'Maximum objectives' }) + await user.clear(input) + await user.type(input, value) + await user.keyboard('{Escape}') +} + function renderDetail( path: string, props: Partial> = {}, @@ -308,11 +326,11 @@ describe('ScenarioDetail', () => { await flushRenderedPromises() await user.type(screen.getByLabelText(/^iterations/), '1.5') await user.type(screen.getByLabelText(/^note/), 'Keep this value') - await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'stale-dataset') - await user.click(screen.getByRole('button', { name: 'Add source' })) - await user.type(screen.getByRole('textbox', { name: 'Max total' }), '0') - await user.type(screen.getByLabelText('Harm categories'), 'stale-harm') - await user.type(screen.getByLabelText('Data types'), 'stale-type') + await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) + await user.click(screen.getByRole('option', { name: 'stale-dataset', exact: true })) + await setObjectiveLimit(user, '0') + await user.click(screen.getAllByRole('button', { name: 'Filter harm category cyber' })[0]) + await user.click(screen.getAllByRole('button', { name: 'Filter data type Text' })[0]) for (const label of ['Max concurrency', 'Max retries']) { const input = screen.getByRole('spinbutton', { name: label }) await user.clear(input) @@ -320,7 +338,7 @@ describe('ScenarioDetail', () => { await user.tab() } expect(screen.getByRole('button', { name: 'Launch scan' })).toBeDisabled() - expect(screen.getByRole('alert')).toHaveTextContent('Max total must be a positive integer') + expect(screen.getByRole('alert')).toHaveTextContent('Maximum objectives must be a positive integer') mockGetScenario.mockResolvedValue(makeScenario({ supported_parameter_names: ['note'], @@ -479,8 +497,9 @@ describe('ScenarioDetail', () => { expect(within(screen.getByRole('region', { name: 'Objective Target' })).queryByRole( 'combobox', { name: 'Adversarial Target' }, )).not.toBeInTheDocument() - expect(fallback.compareDocumentPosition(within(parameters).getByRole('table', { name: 'Selected datasets' }))) - .toBe(Node.DOCUMENT_POSITION_FOLLOWING) + const datasetsSection = screen.getByRole('region', { name: 'Datasets' }) + expect(within(datasetsSection).getByRole('table', { name: 'Selected datasets' })).toBeInTheDocument() + expect(datasetsSection.compareDocumentPosition(parameters)).toBe(Node.DOCUMENT_POSITION_FOLLOWING) expect(fallback).toHaveAccessibleDescription('The registered target this scenario uses to generate attacks.') expect(fallback).toHaveValue('adversary') expect(within(fallback).queryByRole('option', { name: 'target-a' })).not.toBeInTheDocument() @@ -1086,8 +1105,7 @@ describe('ScenarioDetail', () => { expect(await screen.findByRole('heading', { name: 'Parameters' })).toBeInTheDocument() expect(screen.queryByText('Advanced options')).not.toBeInTheDocument() - expect(screen.getByRole('textbox', { name: 'Max total' })).toHaveValue('') - expect(screen.getByText('Default/effective: 8. Blank or default retains the scenario setting.')).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Edit Maximum objectives' })).toHaveTextContent('Up to 8') const estimate = screen.getByTestId('run-estimate') expect(within(estimate).getByText('Dataset size')).toBeInTheDocument() expect(within(estimate).getByText('8')).toBeInTheDocument() @@ -1117,8 +1135,7 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') - expect(screen.getByText('Default/effective: 5. Blank or default retains the scenario setting.')).toBeInTheDocument() + expect(await screen.findByRole('button', { name: 'Edit Maximum objectives' })).toHaveTextContent('Up to 5') expect(await screen.findByText('Up to 10')).toBeInTheDocument() expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -1132,7 +1149,7 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') + expect(await screen.findByRole('button', { name: 'Edit Maximum objectives' })).toHaveTextContent('All available') await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -1147,7 +1164,7 @@ describe('ScenarioDetail', () => { })) renderDetail('/scanner/foundry.red_team_agent') await screen.findByRole('table', { name: 'Selected datasets' }) - expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() + expect(screen.queryByRole('textbox', { name: 'Maximum objectives' })).not.toBeInTheDocument() expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() }) @@ -1161,9 +1178,10 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('textbox', { name: 'Max total' }) - await user.type(input, '2') - await user.clear(input) + await setObjectiveLimit(user, '2') + await openObjectiveLimit(user) + await user.click(screen.getByRole('button', { name: 'Reset to 5' })) + await user.keyboard('{Escape}') expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') @@ -1175,8 +1193,10 @@ describe('ScenarioDetail', () => { it('sends an explicit all total for both estimation and launch', async () => { const user = userEvent.setup() renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('textbox', { name: 'Max total' }) - await user.type(input, 'all') + await setObjectiveLimit(user, '2') + await openObjectiveLimit(user) + await user.click(screen.getByRole('checkbox', { name: 'All available for Maximum objectives' })) + await user.keyboard('{Escape}') await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) @@ -1192,11 +1212,11 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('textbox', { name: 'Max total' }) - await user.type(input, 'all') + await openObjectiveLimit(user) + await user.click(screen.getByRole('checkbox', { name: 'All available for Maximum objectives' })) await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) - await user.clear(input) - expect(input).toHaveValue('') + await user.click(screen.getByRole('button', { name: 'Reset to 5' })) + await user.keyboard('{Escape}') await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) @@ -1209,14 +1229,14 @@ describe('ScenarioDetail', () => { await screen.findByTestId('scenario-target-select') const addDataset = screen.getByRole('combobox', { name: 'Add dataset' }) - await user.type(addDataset, 'ds_a') - await user.click(screen.getByRole('button', { name: 'Add source' })) - await user.type(addDataset, 'ds_b') - await user.click(screen.getByRole('button', { name: 'Add source' })) + await user.click(addDataset) + await user.click(screen.getByRole('option', { name: 'ds_a', exact: true })) + await user.click(addDataset) + await user.click(screen.getByRole('option', { name: 'ds_b', exact: true })) await user.click(screen.getByRole('button', { name: 'Remove harmbench' })) - await user.type(screen.getByRole('textbox', { name: 'Max total' }), '25') - await user.type(screen.getByTestId('harm-categories-filter-input'), 'cyber, violence') - await user.type(screen.getByTestId('data-types-filter-input'), 'text, image_path') + await setObjectiveLimit(user, '25') + await user.click(screen.getAllByRole('button', { name: 'Filter harm category other' })[0]) + await user.click(screen.getAllByRole('button', { name: 'Filter data type Audio' })[0]) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) @@ -1273,17 +1293,23 @@ describe('ScenarioDetail', () => { } }) renderDetail('/scanner/foundry.red_team_agent') - await screen.findByRole('textbox', { name: 'harmbench limit' }) - await user.type(screen.getByRole('textbox', { name: 'harmbench limit' }), '1') + await screen.findByRole('table', { name: 'Selected datasets' }) + await waitFor(() => expect(screen.getByRole('button', { name: 'Edit harmbench objective limit' })).toHaveTextContent('Up to 5')) + await user.click(screen.getByRole('button', { name: 'Edit harmbench objective limit' })) + const all = screen.getByRole('checkbox', { name: 'All available for harmbench objective limit' }) + if (all instanceof HTMLInputElement && all.checked) await user.click(all) + await user.clear(screen.getByRole('textbox', { name: 'harmbench objective limit' })) + await user.type(screen.getByRole('textbox', { name: 'harmbench objective limit' }), '1') + await user.keyboard('{Escape}') await user.clear(screen.getByTestId('scenario-param-scope')) await user.type(screen.getByTestId('scenario-param-scope'), 'one') - await screen.findByText('Invalid selection for current parameters: harmbench') - expect(screen.getByRole('textbox', { name: 'harmbench limit' })).toHaveValue('1') + await screen.findByText('Select a permitted source: harmbench') + expect(screen.getByRole('button', { name: 'Edit harmbench objective limit' })).toHaveTextContent('Up to 1') expect(mockEstimateRun.mock.calls.some((call) => call[1].scenario_params?.scope === 'one' && call[1].datasets === undefined, )).toBe(true) - await user.click(screen.getByRole('button', { name: 'Use default sources' })) - await screen.findByRole('rowheader', { name: 'remote' }) + await user.click(screen.getByRole('button', { name: 'Reset datasets' })) + await screen.findByText('remote', { exact: true }) }) it('rejects a non-positive-integer max dataset size', async () => { @@ -1291,11 +1317,11 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') await screen.findByTestId('scenario-target-select') - await user.type(screen.getByRole('textbox', { name: 'Max total' }), '0') + await setObjectiveLimit(user, '0') await user.click(screen.getByTestId('launch-scenario-btn')) expect(await screen.findAllByRole('alert')).toEqual(expect.arrayContaining([ - expect.objectContaining({ textContent: expect.stringContaining('Max total must be a positive integer, all, or default.') }), + expect.objectContaining({ textContent: expect.stringContaining('Maximum objectives must be a positive integer.') }), ])) expect(mockStartRun).not.toHaveBeenCalled() }) diff --git a/frontend/src/components/Scenarios/ScenarioDetail.tsx b/frontend/src/components/Scenarios/ScenarioDetail.tsx index 0b05aa4550..ade7e2a49d 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.tsx @@ -11,7 +11,6 @@ import { DialogSurface, DialogTitle, Field, - Input, MessageBar, MessageBarBody, mergeClasses, @@ -154,13 +153,6 @@ const BASELINE_TECHNIQUE: ScenarioTechniqueSummary = { tags: ['baseline', 'single_turn'], } -function parseDatasetNames(datasetOverride: string): string[] { - return datasetOverride - .split(',') - .map((entry) => entry.trim()) - .filter((entry) => entry.length > 0) -} - function formatParameterPreview(value: ParameterFormValue | undefined): string { if (Array.isArray(value)) { return value.length > 0 ? value.join(', ') : 'Not set' @@ -242,8 +234,6 @@ interface BuildEstimateRequestInput { scenarioParamValues: Record datasetEdits: ScenarioDatasetRequest datasetError: string | null - harmCategoriesFilter: string - dataTypesFilter: string includeBaseline: boolean } @@ -299,8 +289,6 @@ function buildEstimateRequest({ scenarioParamValues, datasetEdits, datasetError, - harmCategoriesFilter, - dataTypesFilter, includeBaseline, }: BuildEstimateRequestInput): BuildEstimateRequestResult { if (supportsParameter(scenario, 'scenario_techniques') && techniques.length === 0) { @@ -331,14 +319,7 @@ function buildEstimateRequest({ if (scenario.uses_default_adversarial_target && adversarialTargetName) { request.adversarial_target_name = adversarialTargetName } - const harmCategories = supportsDatasets ? parseDatasetNames(harmCategoriesFilter) : [] - const dataTypes = supportsDatasets ? parseDatasetNames(dataTypesFilter) : [] - if (harmCategories.length > 0 || dataTypes.length > 0) { - request.datasets = { ...datasetEdits, filters: { - ...(harmCategories.length > 0 ? { harm_categories: harmCategories } : {}), - ...(dataTypes.length > 0 ? { data_types: dataTypes } : {}), - } } - } else if (supportsDatasets && Object.keys(datasetEdits).length > 0) { + if (supportsDatasets && Object.keys(datasetEdits).length > 0) { request.datasets = datasetEdits } if (scenarioParams) { @@ -630,8 +611,16 @@ function ScenarioLaunchForm({ total_limit: scenario.default_run_size.dataset_limit.state !== 'not_applicable', supports_overrides: true, filters: true, } - const [harmCategoriesFilter, setHarmCategoriesFilter] = useState('') - const [dataTypesFilter, setDataTypesFilter] = useState('') + const selectedDatasetNames = (datasetEdits.sources ?? datasetDefaults.sources).map((source) => source.name) + const datasetSelectionKey = JSON.stringify(selectedDatasetNames) + const reloadDatasetMetadata = datasetMetadata.reload + const previousDatasetSelection = useRef(datasetSelectionKey) + useEffect(() => { + if (previousDatasetSelection.current !== datasetSelectionKey) { + previousDatasetSelection.current = datasetSelectionKey + reloadDatasetMetadata() + } + }, [datasetSelectionKey, reloadDatasetMetadata]) const [maxConcurrency, setMaxConcurrency] = useState(DEFAULT_MAX_CONCURRENCY) const [maxRetries, setMaxRetries] = useState(DEFAULT_MAX_RETRIES) const [scenarioParamValues, setScenarioParamValues] = useState>(() => @@ -694,8 +683,6 @@ function ScenarioLaunchForm({ scenarioParamValues, datasetEdits, datasetError, - harmCategoriesFilter, - dataTypesFilter, includeBaseline: isBaselineForbidden ? false : baselineChecked, }), [ @@ -703,9 +690,7 @@ function ScenarioLaunchForm({ baselineChecked, datasetEdits, datasetError, - dataTypesFilter, dynamicParameters, - harmCategoriesFilter, isBaselineForbidden, scenario, scenarioParamValues, @@ -723,8 +708,6 @@ function ScenarioLaunchForm({ scenarioParamValues, datasetEdits, datasetError, - harmCategoriesFilter, - dataTypesFilter, maxConcurrency, maxRetries, includeBaseline: isBaselineForbidden ? false : baselineChecked, @@ -735,9 +718,7 @@ function ScenarioLaunchForm({ baselineChecked, datasetEdits, datasetError, - dataTypesFilter, dynamicParameters, - harmCategoriesFilter, isBaselineForbidden, labels, maxConcurrency, @@ -958,10 +939,11 @@ function ScenarioLaunchForm({ const displayedEstimate = estimateFromState(estimateState) const selectedTechniqueCount = (supportsTechniques ? selectedTechniques.length : 0) + (supportsBaseline && baselineChecked && !isBaselineForbidden ? 1 : 0) - const previewDatasets = (datasetEdits.sources ?? effectiveDatasets.sources).map((source) => source.name) + const previewDatasets = (datasetEdits.sources ?? datasetDefaults.sources).map((source) => source.name) const previewDatasetNames = previewDatasets.length > 0 ? previewDatasets : scenario.default_datasets - const previewHarmCategories = parseDatasetNames(harmCategoriesFilter) - const previewDataTypes = parseDatasetNames(dataTypesFilter) + const previewFilters = { ...datasetDefaults.filters, ...datasetEdits.filters } + const previewHarmCategories = previewFilters.harm_categories ?? [] + const previewDataTypes = previewFilters.data_types ?? [] return (
)} + {(supportsDatasets || datasetDefaults.sources.length > 0) && ( +
+ Datasets + { if (estimateRequest) void datasetMetadata.loadSelected(scenario.scenario_name, estimateRequest) }} + onReload={datasetMetadata.reload} onChange={setDatasetEdits} onValidationChange={setDatasetError} + /> +
+ )} + {(dynamicParameters.length > 0 || scenario.uses_default_adversarial_target - || supportsDatasets || supportsConcurrency || supportsRetries) && ( + || supportsConcurrency || supportsRetries) && (
Parameters @@ -1159,42 +1157,6 @@ function ScenarioLaunchForm({ disabled={submitting} /> )} - {supportsDatasets && ( - <> - - - setHarmCategoriesFilter(data.value)} - data-testid="harm-categories-filter-input" - /> - - - setDataTypesFilter(data.value)} - data-testid="data-types-filter-input" - /> - - - )} {supportsConcurrency && Launch scan diff --git a/frontend/src/components/Scenarios/datasetPresentation.ts b/frontend/src/components/Scenarios/datasetPresentation.ts new file mode 100644 index 0000000000..e0edfc4d27 --- /dev/null +++ b/frontend/src/components/Scenarios/datasetPresentation.ts @@ -0,0 +1,34 @@ +import type { DatasetInfo, ScenarioDatasetOptions, ScenarioDatasetSettings, ScenarioDatasetSourceSettings } from '@/types' + +export function parseCaseLimit(text: string): number | null { + const value = text.trim() + const number = Number(value) + return /^\d+$/.test(value) && Number.isSafeInteger(number) && number > 0 ? number : null +} + +export function datasetSize( + dataset: DatasetInfo | undefined, limit: ScenarioDatasetSettings['max_total'], +): string { + const count = dataset?.loaded ? dataset.logical_examples : null + if (limit === 'all') return count != null + ? `${count.toLocaleString()} examples` + : dataset?.loaded ? 'Loaded' : 'Not loaded' + return count != null + ? `Up to ${Math.min(limit, count).toLocaleString()} of ${count.toLocaleString()}` + : `Up to ${limit.toLocaleString()}` +} + +export function datasetChoices( + datasets: DatasetInfo[], defaults: ScenarioDatasetSourceSettings[], options: ScenarioDatasetOptions, +): string[] { + const names = new Set([...datasets.map((item) => item.name), ...defaults.map((item) => item.name)]) + return [...names].filter((name) => options.mode === 'any' || options.allowed_names.includes(name)).sort() +} + +export function datasetFacts(datasets: DatasetInfo[], names: string[]) { + const loaded = datasets.filter((item) => item.loaded && names.includes(item.name)) + return { + loaded, + allLoaded: names.length > 0 && loaded.length === names.length, + } +} diff --git a/frontend/src/components/Sidebar/Navigation.test.tsx b/frontend/src/components/Sidebar/Navigation.test.tsx index c4b4a0e7b1..6ebaa0943a 100644 --- a/frontend/src/components/Sidebar/Navigation.test.tsx +++ b/frontend/src/components/Sidebar/Navigation.test.tsx @@ -41,7 +41,7 @@ describe("Navigation", () => { renderWithProvider(); expect(screen.getByRole("navigation", { name: "Primary" })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Chat" })).toHaveAttribute( + expect(screen.getByRole("button", { name: "Technique.execute" })).toHaveAttribute( "aria-current", "page" ); @@ -61,9 +61,12 @@ describe("Navigation", () => { expect(onNavigate).toHaveBeenCalledWith("home"); }); - it("renders the chat button", () => { + it("renders the Technique.execute button for chat", () => { renderWithProvider(); - expect(screen.getByRole("button", { name: "Chat" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Technique.execute" })).toHaveAttribute( + "title", + "Technique.execute" + ); }); it("renders the registry button", () => { @@ -80,7 +83,7 @@ describe("Navigation", () => { ); - await user.click(screen.getByRole("button", { name: "Chat" })); + await user.click(screen.getByRole("button", { name: "Technique.execute" })); expect(onNavigate).toHaveBeenCalledWith("chat"); }); @@ -139,7 +142,7 @@ describe("Navigation", () => { expect(labels).toEqual([ "Home", - "Chat", + "Technique.execute", "History", "Scanner", "Datasets", diff --git a/frontend/src/components/Sidebar/Navigation.tsx b/frontend/src/components/Sidebar/Navigation.tsx index 3112aebf6e..ca70d163c3 100644 --- a/frontend/src/components/Sidebar/Navigation.tsx +++ b/frontend/src/components/Sidebar/Navigation.tsx @@ -86,8 +86,8 @@ export default function Navigation({ data-active={currentView === 'chat'} appearance="subtle" icon={} - title="Chat" - aria-label="Chat" + title="Technique.execute" + aria-label="Technique.execute" aria-current={currentView === 'chat' ? 'page' : undefined} onClick={() => onNavigate('chat')} /> diff --git a/frontend/src/components/Tour/tourSteps.ts b/frontend/src/components/Tour/tourSteps.ts index a31652cfe6..315953d477 100644 --- a/frontend/src/components/Tour/tourSteps.ts +++ b/frontend/src/components/Tour/tourSteps.ts @@ -28,7 +28,7 @@ export function createTourSteps({ { target: '[data-tour="sidebar-nav"]', content: - 'Ahoy! Welcome to Co-PyRIT! This is your main navigation panel. Home is your dashboard, Chat is where you send prompts, ' + + 'Ahoy! Welcome to Co-PyRIT! This is your main navigation panel. Home is your dashboard, Technique.execute opens Chat, ' + 'History tracks past attacks and scanner runs, Scanner launches full test campaigns, and Registry is where you manage ' + 'targets and converters.' + (canManageConfiguration diff --git a/frontend/src/hooks/useScenarioDatasets.test.tsx b/frontend/src/hooks/useScenarioDatasets.test.tsx index 3164033fac..8dc7a121f0 100644 --- a/frontend/src/hooks/useScenarioDatasets.test.tsx +++ b/frontend/src/hooks/useScenarioDatasets.test.tsx @@ -1,14 +1,18 @@ import { act, renderHook, waitFor } from '@testing-library/react' -import { datasetsApi } from '@/services/api' +import { datasetsApi, scenariosApi } from '@/services/api' import type { DatasetInfo } from '@/types' import { useScenarioDatasets } from './useScenarioDatasets' -jest.mock('@/services/api', () => ({ datasetsApi: { listDatasets: jest.fn() } })) +jest.mock('@/services/api', () => ({ + datasetsApi: { listDatasets: jest.fn() }, + scenariosApi: { loadDatasets: jest.fn() }, +})) let generation = 'one' jest.mock('./useRuntime', () => ({ useRuntime: () => ({ ready: true, generation }) })) const list = datasetsApi.listDatasets as jest.Mock +const loadDatasets = scenariosApi.loadDatasets as jest.Mock const NAMED: DatasetInfo = { name: 'named', selection_key: 'dataset:named:named', loaded: true, provider_available: false, logical_examples: 1, origins: ['local'], seed_pieces: 1, objectives: 1, modalities: ['text'], harm_categories: [], has_unlabeled_harm_categories: false, } @@ -55,4 +59,45 @@ describe('useScenarioDatasets', () => { expect(view.result.current.loading).toBe(true) view.unmount() }) + + it('loads only on explicit action and refreshes metadata after success', async () => { + list.mockResolvedValue({ items: [NAMED] }) + loadDatasets.mockResolvedValue({}) + const view = renderHook(() => useScenarioDatasets()) + await waitFor(() => expect(view.result.current.items).toEqual([NAMED])) + expect(loadDatasets).not.toHaveBeenCalled() + const request = { datasets: { max_total: 4 } } + await act(async () => { await view.result.current.loadSelected('test', request) }) + expect(loadDatasets).toHaveBeenCalledWith('test', request) + await waitFor(() => expect(list).toHaveBeenCalledTimes(2)) + expect(view.result.current.loadError).toBeNull() + }) + + it('retains load failures and still refreshes current metadata', async () => { + list.mockResolvedValue({ items: [NAMED] }) + loadDatasets.mockRejectedValue(new Error('Load failed')) + const view = renderHook(() => useScenarioDatasets()) + await act(async () => { await view.result.current.loadSelected('test', {}) }) + await waitFor(() => expect(view.result.current.loadError).not.toBeNull()) + expect(view.result.current.preparing).toBe(false) + expect(list).toHaveBeenCalledTimes(2) + }) + + it('deduplicates load clicks and ignores completion from a previous generation', async () => { + list.mockResolvedValue({ items: [NAMED] }) + let finish: () => void = () => {} + loadDatasets.mockReturnValue(new Promise((resolve) => { finish = resolve })) + const view = renderHook(() => useScenarioDatasets()) + let pending: Promise + act(() => { + pending = view.result.current.loadSelected('test', {}) + void view.result.current.loadSelected('test', {}) + }) + expect(loadDatasets).toHaveBeenCalledTimes(1) + generation = 'two' + view.rerender() + await act(async () => { finish(); await pending }) + expect(view.result.current.preparing).toBe(false) + expect(view.result.current.loadError).toBeNull() + }) }) diff --git a/frontend/src/hooks/useScenarioDatasets.ts b/frontend/src/hooks/useScenarioDatasets.ts index cf6071c3e5..3732736f2f 100644 --- a/frontend/src/hooks/useScenarioDatasets.ts +++ b/frontend/src/hooks/useScenarioDatasets.ts @@ -1,25 +1,58 @@ -import { useCallback, useEffect, useState } from 'react' +import { useCallback, useEffect, useRef, useState } from 'react' -import { datasetsApi } from '@/services/api' +import { datasetsApi, scenariosApi } from '@/services/api' import { toApiError } from '@/services/errors' -import type { DatasetInfo } from '@/types' +import type { DatasetInfo, ScenarioRunSizeEstimateRequest } from '@/types' import { useRuntime } from './useRuntime' interface DatasetSnapshot { key: string + generation: string items: DatasetInfo[] error: string | null } -/** Page-scoped metadata cache; no dataset loading or generation occurs here. */ +/** Page-scoped metadata cache with an explicit, missing-only load action. */ export function useScenarioDatasets() { const runtime = useRuntime() const [snapshot, setSnapshot] = useState(null) const [reloadCount, setReloadCount] = useState(0) + const [operation, setOperation] = useState<{ generation: string; busy: boolean; error: string | null } | null>(null) + const generationRef = useRef(runtime.generation) + const mountedRef = useRef(true) + const pendingRef = useRef | null>(null) const reload = useCallback(() => { setReloadCount((count) => count + 1) }, []) const key = JSON.stringify([runtime.generation, runtime.ready, reloadCount]) + useEffect(() => { generationRef.current = runtime.generation }, [runtime.generation]) + useEffect(() => { + mountedRef.current = true + return () => { mountedRef.current = false } + }, []) + + const loadSelected = (scenarioName: string, request: ScenarioRunSizeEstimateRequest): Promise => { + if (pendingRef.current) return pendingRef.current + const generation = runtime.generation + setOperation({ generation, busy: true, error: null }) + const task = async (): Promise => { + let error: string | null = null + try { + await scenariosApi.loadDatasets(scenarioName, request) + } catch (err: unknown) { + error = toApiError(err).detail + } finally { + pendingRef.current = null + if (mountedRef.current && generationRef.current === generation) { + setOperation({ generation, busy: false, error }) + reload() + } + } + } + pendingRef.current = task() + return pendingRef.current + } + useEffect(() => { if (!runtime.ready) return const controller = new AbortController() @@ -29,22 +62,28 @@ export function useScenarioDatasets() { if (controller.signal.aborted) return setSnapshot({ key, + generation: runtime.generation, items: data.items.filter((item) => item.selection_key !== 'dataset:unnamed'), error: null, }) } catch (err: unknown) { - if (!controller.signal.aborted) setSnapshot({ key, items: [], error: toApiError(err).detail }) + if (!controller.signal.aborted) setSnapshot({ + key, generation: runtime.generation, items: [], error: toApiError(err).detail, + }) } } void load() return () => { controller.abort() } - }, [key, runtime.ready]) + }, [key, runtime.ready, runtime.generation]) const current = runtime.ready && snapshot?.key === key ? snapshot : null return { - items: current?.items ?? [], + items: runtime.ready && snapshot?.generation === runtime.generation ? snapshot.items : [], error: current?.error ?? null, loading: runtime.ready && current === null, reload, + loadSelected, + preparing: operation?.generation === runtime.generation && operation.busy, + loadError: operation?.generation === runtime.generation ? operation.error : null, } } diff --git a/frontend/src/hooks/useTour.test.ts b/frontend/src/hooks/useTour.test.ts index 8c23cd0995..81c3baca45 100644 --- a/frontend/src/hooks/useTour.test.ts +++ b/frontend/src/hooks/useTour.test.ts @@ -121,6 +121,7 @@ describe('useTour', () => { '[data-tour="history-tabs"]', '[data-tour="registry-tabs"]', ]) + expect(steps[0].content).toContain('Technique.execute opens Chat') expect(steps[0].content).toContain('Scanner launches full test campaigns') expect(steps[4].content).toContain('Scanner runs a whole campaign') expect(steps[5].content).toContain('The Attacks tab lists individual conversations') diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 6758ef71a8..b0575a8828 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -43,6 +43,7 @@ import type { RunScenarioRequest, ScenarioRunSizeEstimateResponse, ScenarioRunSizeEstimateRequest, + ScenarioDatasetLoadResponse, ScenarioRunSummary, ScenarioRunListResponse, ScenarioRunProgress, @@ -480,6 +481,14 @@ export const labelsApi = { } export const scenariosApi = { + loadDatasets: async ( + scenarioName: string, request: ScenarioRunSizeEstimateRequest, + ): Promise => { + const response = await apiClient.post( + `/scenarios/catalog/${encodeURIComponent(scenarioName)}/datasets/load`, request, { timeout: 300_000 }, + ) + return response.data + }, /** * Lists one page of the scenario catalog. Callers that need the full * catalog should follow `pagination.next_cursor` until `has_more` is false. diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 87d0b534b2..8eb2ea7613 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -883,6 +883,11 @@ export interface ScenarioDatasetSettings { filters: Record } +export interface ScenarioDatasetLoadResponse { + settings: ScenarioDatasetSettings + items: DatasetInfo[] +} + export interface RegisteredScenario { scenario_name: string scenario_type: string diff --git a/pyrit/backend/models/datasets.py b/pyrit/backend/models/datasets.py index 97bf5ef246..ebd1773375 100644 --- a/pyrit/backend/models/datasets.py +++ b/pyrit/backend/models/datasets.py @@ -15,6 +15,7 @@ from pyrit.backend.models.common import PaginationInfo from pyrit.models import PromptDataType, SeedRecord, SeedType +from pyrit.models.scenario_dataset import ScenarioDatasetSettings class DatasetInfo(BaseModel): @@ -49,6 +50,13 @@ class DatasetListResponse(BaseModel): items: list[DatasetInfo] = Field(..., description="List of available datasets") +class ScenarioDatasetLoadResponse(BaseModel): + """Selected dataset settings and metadata after an explicit load.""" + + settings: ScenarioDatasetSettings + items: list[DatasetInfo] + + class SeedExampleSummary(BaseModel): """One logical seed example: the seeds that share a group ID, or one seed without a group.""" diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 70966e8a40..06f8482189 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -15,6 +15,7 @@ from fastapi import APIRouter, HTTPException, Query, Request, status from pyrit.backend.models.common import MAX_ITEMS, CursorStr, IdentifierStr, LabelFilterStr, ProblemDetail +from pyrit.backend.models.datasets import ScenarioDatasetLoadResponse from pyrit.backend.models.scenarios import ( ListRegisteredScenariosResponse, ScenarioRunListResponse, @@ -137,6 +138,32 @@ async def estimate_scenario_run_size( # pyrit-async-suffix-exempt return estimate +@router.post( + "/catalog/{scenario_name}/datasets/load", + response_model=ScenarioDatasetLoadResponse, + responses={ + 400: {"model": ProblemDetail, "description": "Invalid dataset selection"}, + 404: {"model": ProblemDetail, "description": "Scenario not found"}, + }, +) +async def load_scenario_datasets_async( + *, scenario_name: IdentifierStr, request: ScenarioRunSizeEstimateRequest +) -> ScenarioDatasetLoadResponse: + """ + Load missing selected datasets on explicit request, without starting a scan. + + Returns: + ScenarioDatasetLoadResponse: Selected settings and stored metadata. + """ + try: + result = await get_scenario_service().load_datasets_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 result is None: + raise HTTPException(status_code=404, detail=f"Scenario '{scenario_name}' not found") + return result + + # ============================================================================ # Scenario Runs # ============================================================================ diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index 7b2e3a1838..3176dea277 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -13,7 +13,9 @@ from time import monotonic from pyrit.backend.models.common import PaginationInfo +from pyrit.backend.models.datasets import ScenarioDatasetLoadResponse from pyrit.backend.models.scenarios import ListRegisteredScenariosResponse +from pyrit.backend.services.dataset_service import get_dataset_service from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.models.catalog.scenario import ( RegisteredScenario, @@ -239,6 +241,46 @@ def clear_estimate_task(completed_task: _EstimateTask) -> None: assert isinstance(estimate, ScenarioRunSizeEstimate) return estimate + async def load_datasets_async( + self, *, scenario_name: str, request: ScenarioRunSizeEstimateRequest + ) -> ScenarioDatasetLoadResponse | None: + """ + Resolve the same source selection as launch and load only its missing datasets. + + Returns: + ScenarioDatasetLoadResponse | None: Selected metadata, or None for an unknown scenario. + """ + try: + scenario_class = self._registry.get_class(scenario_name) + except KeyError: + return None + reject_non_external_params( + params=request.scenario_params or {}, + declared=scenario_class.supported_parameters(), + owner=scenario_name, + ) + resolver = ScenarioConfigurationResolver() + adversarial_target = resolver.resolve_adversarial_target(target_name=request.adversarial_target_name) + objective_target = resolver.resolve_target(target_name=request.target_name) if request.target_name else None + with override_default_adversarial_target(adversarial_target): + kwargs = await asyncio.to_thread( + resolver.resolve_configuration, + scenario_name=scenario_name, + scenario_class=scenario_class, + techniques=request.techniques, + datasets=request.dataset_request, + objective_target=objective_target, + include_baseline=request.include_baseline, + ) + scenario = await asyncio.to_thread(self._registry.create_instance, scenario_name) + scenario.set_params_from_args(args={**(request.scenario_params or {}), **kwargs}) + settings = await scenario.load_datasets_async() + metadata = await get_dataset_service().list_datasets_async() + names = {source.name for source in settings.sources} + return ScenarioDatasetLoadResponse( + settings=settings, items=[item for item in metadata.items if item.name in names] + ) + async def _get_default_run_size_estimate_async(self, *, metadata: ScenarioMetadata) -> ScenarioRunSizeEstimate: """Return a cached, cancellation-safe scenario-owned estimate.""" cache_key = ( diff --git a/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py b/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py index e15a4cb82c..695eb03850 100644 --- a/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py +++ b/pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py @@ -6,6 +6,7 @@ import hashlib import io import logging +import os import tempfile import zipfile from abc import ABC, abstractmethod @@ -420,7 +421,7 @@ def _load_dataset_sync() -> Any: Returns: Dataset: The loaded dataset from Hugging Face. """ - cache_dir = str(DB_DATA_PATH / "huggingface") if cache else None + cache_dir = (os.getenv("HF_DATASETS_CACHE") or str(DB_DATA_PATH / "huggingface")) if cache else None # Reuse cached data when caching is enabled; force a re-download otherwise so # cache=False genuinely picks up upstream edits instead of silently reusing the cache. diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 50278a1561..72eeac59eb 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -515,19 +515,25 @@ def has_sampling_limits(self) -> bool: def _preparation_sources(self) -> list[tuple[DatasetSource, DatasetFetchPolicy]]: return [(source, source.fetch or self.fetch) for source in self.sources] - async def prepare_async(self) -> None: + async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> None: """ Check all sources, then populate missing datasets from configured or registered providers. + Args: + fetch: Optional policy for this operation only. Source settings remain unchanged. + Raises: DatasetConstraintError: If a missing source cannot be fetched or provider data is invalid. """ self.validate_configuration() + if fetch is not None and not isinstance(fetch, DatasetFetchPolicy): + raise DatasetConstraintError("'fetch' must be a DatasetFetchPolicy.") if self.source_kind is DatasetSourceKind.INLINE: return policies: dict[str, DatasetFetchPolicy] = {} configured_providers: dict[str, SeedDatasetProvider | None] = {} for source, policy in self._preparation_sources(): + policy = fetch if fetch is not None else policy name = source.name if name in policies and policies[name] is not policy: raise DatasetConstraintError(f"Dataset '{name}' has conflicting fetch policies.") diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 19c4ef4597..cd160e3e0a 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -53,7 +53,12 @@ config_hash, ) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference -from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings +from pyrit.models.scenario_dataset import ( + DatasetFetchPolicy, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, +) from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.registry import ScorerRegistry @@ -63,6 +68,7 @@ DatasetAttackConfiguration, ) from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides, project_dataset_settings +from pyrit.scenario.core.dataset_sampling import sample_with_coverage 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 @@ -681,6 +687,9 @@ async def _get_run_size_estimate_async(self, *, target_is_configured: bool = Fal if target_is_configured and self._objective_target is None: raise ValueError("target_is_configured requires a resolved objective_target") budget = self._get_run_size_budget() + cap = self._dataset_config.max_total + if cap != "all" and isinstance(budget, BoundedDatasetSize): + budget = BoundedDatasetSize(value=min(cap, budget.value)) dataset_limit = self._get_dataset_limit_input() if isinstance(budget, IndeterminateDatasetSize): return ScenarioRunSizeEstimate.unavailable( @@ -762,7 +771,7 @@ async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> Scena def _get_dataset_limit_input(self) -> DatasetLimitInput: """Return the editable limit, never an aggregate population budget.""" - if not self.USES_DATASET_SIZE_LIMIT: + if not self.get_dataset_options().total_limit: return DatasetLimitInput(state=DatasetLimitState.NotApplicable) limit = self._dataset_config.max_total return DatasetLimitInput(state=DatasetLimitState.Value, value=limit) if limit != "all" else DatasetLimitInput() @@ -894,6 +903,17 @@ def get_default_dataset_settings(self) -> ScenarioDatasetSettings: """Return the default dataset settings for the current parameters, without reading memory.""" return project_dataset_settings(self._resolve_default_dataset_config()) + async def load_datasets_async(self) -> ScenarioDatasetSettings: + """ + Load missing selected datasets explicitly, without creating or executing a run. + + Returns: + ScenarioDatasetSettings: Resolved settings, with the scan's policies unchanged. + """ + self._resolve_runtime_configuration(require_objective_target=False) + await self._dataset_config.prepare_async(fetch=DatasetFetchPolicy.IF_MISSING) + return project_dataset_settings(self._dataset_config) + def _validate_runtime_configuration(self) -> None: """ Check resolved parameters for preview and launch without reading datasets. @@ -902,11 +922,19 @@ def _validate_runtime_configuration(self) -> None: ValueError: If an ingredient-only scenario receives per-dataset limits. """ self._dataset_config.validate_configuration() + cap = self._dataset_config.max_total + minimum = max(1, len(self._get_required_objective_keys())) + if cap != "all" and cap < minimum: + raise ValueError(f"Select at least {minimum} objectives, or select fewer techniques.") if not self.USES_DATASET_SIZE_LIMIT and any( self._dataset_config.source_limit(source.name) != "all" for source in self._dataset_config.sources ): raise ValueError("This scenario uses ingredient datasets and does not support per-dataset limits.") + def _get_required_objective_keys(self) -> list[str]: + """Return population names that must retain at least one objective.""" + return [] + @final async def initialize_async(self) -> None: """ @@ -969,6 +997,15 @@ async def initialize_async(self) -> None: else: await self._dataset_config.prepare_async() seed_groups_by_dataset = await self._resolve_seed_groups_by_dataset_async(apply_sampling=not is_resume) + if not is_resume and self._dataset_config.max_total != "all": + count = sum(len(groups) for groups in seed_groups_by_dataset.values()) + if count > self._dataset_config.max_total: + seed_groups_by_dataset = sample_with_coverage( + groups_by_dataset=seed_groups_by_dataset, + cap=self._dataset_config.max_total, + required_keys=self._get_required_objective_keys(), + key=lambda name, _: name, + ) # A fresh run's result id is chosen now, so attacks built below can already refer to the run. new_result_id = uuid.uuid4() context = self._build_scenario_context( diff --git a/pyrit/scenario/scenarios/garak/package_hallucination.py b/pyrit/scenario/scenarios/garak/package_hallucination.py index aa821b54ff..d82148cb4c 100644 --- a/pyrit/scenario/scenarios/garak/package_hallucination.py +++ b/pyrit/scenario/scenarios/garak/package_hallucination.py @@ -140,9 +140,7 @@ class PackageHallucination(Scenario): """ VERSION: int = 3 - DATASET_OPTIONS = ScenarioDatasetOptions( - mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False - ) + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # The plain code request is not an adversarial baseline to compare against, so no baseline. @@ -227,6 +225,9 @@ def _get_run_size_budget(self) -> BoundedDatasetSize: """Return the combined generated-prompt cap for the selected languages.""" return BoundedDatasetSize(value=self._max_prompts_per_language * len(self._scenario_techniques)) + def _get_required_objective_keys(self) -> list[str]: + return [technique.value for technique in self._scenario_techniques] + async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> ScenarioRunSizeEstimate: """ Estimate generated prompts from the per-language cap without loading the corpus. @@ -235,18 +236,21 @@ async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> Scena ScenarioRunSizeEstimate: Combined budget for the selected languages. """ assert isinstance(budget, BoundedDatasetSize) + components = [ + ScenarioRunSizeComponent(label=technique.value, count=self._max_prompts_per_language) + for technique in self._scenario_techniques + ] + if budget.value < sum(component.count for component in components): + components = [ScenarioRunSizeComponent(label="Selected generated objectives", count=budget.value)] return ScenarioRunSizeEstimate( total_attack_count=budget.value, - components=[ - ScenarioRunSizeComponent(label=technique.value, count=self._max_prompts_per_language) - for technique in self._scenario_techniques - ], + components=components, datasets=[ ScenarioDatasetSummary(name=technique.value, kind="synthesized") for technique in self._scenario_techniques ], effective_parameters={"max_prompts_per_language": self._max_prompts_per_language}, - note="The generated-prompt cap applies per language. Dataset size limits do not apply.", + note="The objective limit caps all selected languages; the per-language generation cap also applies.", ) @staticmethod diff --git a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py index 9fe3de603a..f9a975a567 100644 --- a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py +++ b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py @@ -99,9 +99,7 @@ class SystemPromptExtraction(Scenario): """ VERSION: int = 1 - DATASET_OPTIONS = ScenarioDatasetOptions( - mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False - ) + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # Template-dominated like the Doctor/Jailbreak scenarios: the bare system prompt with no @@ -159,6 +157,7 @@ def __init__( technique_class=SystemPromptExtractionTechnique, default_dataset_config=DatasetAttackConfiguration( max_per_dataset="all", + max_total=self._prompt_cap, sources=[ DatasetSource(name=name) for name in [ @@ -181,7 +180,7 @@ def _validate_runtime_configuration(self) -> None: def _get_run_size_budget(self) -> ScenarioDatasetSizeEstimate: """Return the shared cap on system-prompt and template combinations.""" - return scenario_dataset_size_from_limit(self._prompt_cap) + return scenario_dataset_size_from_limit(self._dataset_config.max_total) async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> ScenarioRunSizeEstimate: """ @@ -196,7 +195,7 @@ async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> Scena components=[ScenarioRunSizeComponent(label="System-prompt/template combinations", count=count)], datasets=[ScenarioDatasetSummary(name="System-prompt/template combinations", kind="synthesized")], effective_parameters={"prompt_cap": budget.value}, - note="The prompt cap is shared across all selected categories. Dataset size limits do not apply.", + note="The objective limit is shared across all selected categories.", ) def _sample_system_prompts(self, seeds_by_dataset: dict[str, list[Seed]]) -> list[str]: @@ -238,15 +237,14 @@ async def _resolve_seed_groups_by_dataset_async( Build the (system prompt x template) seed groups, keyed by technique category. Mirrors garak: every selected (system prompt x template) combination is enumerated across - all selected categories, then randomly sampled down to ``prompt_cap`` when it is set. + all selected categories, then sampled down to the configured objective limit. Surviving combinations are grouped by category so each becomes one atomic attack. Resolving them here (rather than via the dataset config) means the base owns the single seed sample shared across the atomic attacks. Args: - apply_sampling (bool): Accepted for base-class compatibility but unused — the prompt-cap - sampling is already deterministic (``random.Random(self._random_seed)``), so resume - reproduces the same set without a ``max_dataset_size`` sampling path. + apply_sampling (bool): Apply the objective limit on fresh runs. Resume resolves + the full population so the stored run plan can restore its selected objectives. Returns: dict[str, list[AttackSeedGroup]]: Seed groups keyed by technique category value. @@ -275,8 +273,9 @@ async def _resolve_seed_groups_by_dataset_async( f"{DATASET_EXTRACTION_TEMPLATES}) are loaded into CentralMemory before running." ) - if self._prompt_cap != "all" and len(combinations) > self._prompt_cap: - combinations = random.Random(self._random_seed).sample(combinations, self._prompt_cap) + cap = self._dataset_config.max_total + if apply_sampling and cap != "all" and len(combinations) > cap: + combinations = random.Random(self._random_seed).sample(combinations, cap) seed_groups_by_category: dict[str, list[AttackSeedGroup]] = {} for category, system_prompt, template in combinations: diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 3c4576aabd..4d5f2bfd7c 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -116,7 +116,7 @@ def get_dataset_options(self) -> ScenarioDatasetOptions: mode=DatasetSelectionMode.SUBSET, allowed_names=self._default_dataset_config.dataset_names, per_dataset_limit=False, - total_limit=False, + total_limit=True, ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False @@ -608,6 +608,9 @@ def _get_technique_size_budgets(self) -> dict[WebInjectionTechnique, BoundedData budgets[technique] = BoundedDatasetSize(value=cap) return budgets + def _get_required_objective_keys(self) -> list[str]: + return [technique.value for technique in self._scenario_techniques] + def _get_run_size_budget(self) -> BoundedDatasetSize: """Return the combined size contract without reading source datasets.""" budgets = self._get_technique_size_budgets() @@ -629,7 +632,9 @@ async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> Scena ) for name, count in counts.items() ] - synthesized_count = sum(component.count for component in components) + synthesized_count = budget.value + if budget.value < sum(counts.values()): + components = [ScenarioRunSizeComponent(label="Selected generated objectives", count=budget.value)] if self._include_baseline: components.append( ScenarioRunSizeComponent( @@ -644,7 +649,7 @@ async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> Scena components=components, datasets=datasets, effective_parameters={"max_prompts_per_technique": self._max_prompts_per_technique}, - note=("Each technique owns a distinct synthesized population; dataset size limits do not apply."), + note="Each technique owns a distinct generated population; the objective limit caps their union.", ) async def _resolve_seed_groups_by_dataset_async( diff --git a/tests/unit/backend/test_scenario_dataset_loading.py b/tests/unit/backend/test_scenario_dataset_loading.py new file mode 100644 index 0000000000..ef8cbd9259 --- /dev/null +++ b/tests/unit/backend/test_scenario_dataset_loading.py @@ -0,0 +1,136 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Explicit dataset loading uses scenario selection without creating a scan.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +from pyrit.backend.main import app +from pyrit.backend.services.scenario_service import ScenarioService +from pyrit.datasets import SeedDatasetProvider +from pyrit.memory import CentralMemory +from pyrit.models import SeedDataset, SeedObjective +from pyrit.models.catalog import ScenarioRunSizeEstimateRequest +from pyrit.models.scenario_dataset import DatasetFetchPolicy, ScenarioDatasetRequest +from pyrit.registry import ScenarioRegistry +from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario +from pyrit.scenario.core.atomic_attack import AtomicAttack +from pyrit.scenario.core.scenario_context import ScenarioContext +from pyrit.scenario.core.scenario_technique import ScenarioTechnique +from pyrit.score import SubStringScorer + + +class _Technique(ScenarioTechnique): + ALL = ("all", {"all"}) + DIRECT = ("direct", set()) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + +class _LoadScenario(Scenario): + def __init__(self, *, provider: SeedDatasetProvider | None = None) -> None: + super().__init__( + version=1, + technique_class=_Technique, + objective_scorer=SubStringScorer(substring="yes"), + default_dataset_config=DatasetAttackConfiguration( + sources=[DatasetSource(name="manual_load", fetch=DatasetFetchPolicy.NEVER, provider=provider)], + ), + ) + + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + raise AssertionError("Loading must not build attacks") + + +def _provider() -> MagicMock: + provider = MagicMock(spec=SeedDatasetProvider) + provider.dataset_name = "manual_load" + provider.fetch_dataset_async = AsyncMock( + return_value=SeedDataset( + dataset_name="manual_load", + seeds=[SeedObjective(value=f"Objective {i}", dataset_name="manual_load") for i in range(2)], + ) + ) + return provider + + +@pytest.mark.usefixtures("patch_central_database") +async def test_manual_load_fetches_once_and_preserves_provider_and_scan_policy() -> None: + provider = _provider() + scenario = _LoadScenario(provider=provider) + scenario.set_params_from_args(args={}) + settings = await scenario.load_datasets_async() + await scenario.load_datasets_async() + provider.fetch_dataset_async.assert_awaited_once() + assert settings.sources[0].fetch is DatasetFetchPolicy.NEVER + assert scenario._default_dataset_config.sources[0].provider is provider + assert await CentralMemory.get_memory_instance().get_scenario_results_async() == [] + assert scenario._scenario_result_id is None + + +@pytest.mark.usefixtures("patch_central_database") +async def test_load_action_uses_registry_and_returns_full_stored_metadata() -> None: + provider = _provider() + with patch.object(ScenarioRegistry, "_discover"): + registry = ScenarioRegistry() + registry.register_class(_LoadScenario, name="test.load") + with ( + patch.object(ScenarioRegistry, "get_registry_singleton", return_value=registry), + patch.object(SeedDatasetProvider, "get_providers_by_name_async", return_value={"manual_load": provider}), + patch.object(SeedDatasetProvider, "get_all_dataset_names_async", return_value=["manual_load"]), + ): + result = await ScenarioService().load_datasets_async( + scenario_name="test.load", + request=ScenarioRunSizeEstimateRequest(datasets=ScenarioDatasetRequest(max_total=1)), + ) + assert result is not None + assert result.settings.max_total == 1 + assert result.items[0].name == "manual_load" + assert result.items[0].loaded is True + assert result.items[0].logical_examples == 2 + assert result.settings.sources[0].fetch is DatasetFetchPolicy.NEVER + assert await CentralMemory.get_memory_instance().get_scenario_results_async() == [] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_manual_load_keeps_edited_data_and_does_not_call_provider() -> None: + memory = CentralMemory.get_memory_instance() + await memory.add_seeds_to_memory_async( + seeds=[SeedObjective(value="Edited objective", dataset_name="manual_load")], + added_by="test", + ) + provider = _provider() + scenario = _LoadScenario(provider=provider) + await scenario.load_datasets_async() + provider.fetch_dataset_async.assert_not_awaited() + assert [seed.value for seed in await memory.get_seeds_async(dataset_name="manual_load")] == ["Edited objective"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_load_failure_surfaces_and_stores_no_partial_dataset() -> None: + provider = _provider() + provider.fetch_dataset_async.side_effect = RuntimeError("Provider unavailable") + with pytest.raises(RuntimeError, match="Provider unavailable"): + await _LoadScenario(provider=provider).load_datasets_async() + assert await CentralMemory.get_memory_instance().get_seeds_async(dataset_name="manual_load") == [] + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("error,code", [(None, 404), (ValueError("Import this dataset first"), 400)]) +def test_load_route_errors_are_explicit( + compatibility_headers: dict[str, str], + error: Exception | None, + code: int, +) -> None: + with patch("pyrit.backend.routes.scenarios.get_scenario_service") as get_service: + get_service.return_value.load_datasets_async = AsyncMock(return_value=None, side_effect=error) + response = TestClient(app, headers=compatibility_headers).post( + "/api/scenarios/catalog/test.load/datasets/load", + json={}, + ) + assert response.status_code == code diff --git a/tests/unit/datasets/test_remote_dataset_loader.py b/tests/unit/datasets/test_remote_dataset_loader.py index 8b58a57a30..0c667a7b68 100644 --- a/tests/unit/datasets/test_remote_dataset_loader.py +++ b/tests/unit/datasets/test_remote_dataset_loader.py @@ -435,3 +435,33 @@ async def test_cache_false_forces_redownload(self): assert mock_load.call_args.kwargs["download_mode"] == DownloadMode.FORCE_REDOWNLOAD assert mock_load.call_args.kwargs["cache_dir"] is None + + @pytest.mark.parametrize("cache", [False, True]) + async def test_custom_cache_directory_respects_cache_flag(self, cache: bool) -> None: + from datasets import DownloadMode + + loader = ConcreteRemoteLoader() + with ( + patch.dict("os.environ", {"HF_DATASETS_CACHE": "short-cache"}), + patch("pyrit.datasets.seed_datasets.remote.remote_dataset_loader.load_dataset") as mock_load, + ): + await loader._fetch_from_huggingface_async(dataset_name="owner/ds", split="train", cache=cache) + + assert mock_load.call_args.kwargs["cache_dir"] == ("short-cache" if cache else None) + assert mock_load.call_args.kwargs["download_mode"] == ( + DownloadMode.REUSE_DATASET_IF_EXISTS if cache else DownloadMode.FORCE_REDOWNLOAD + ) + + @pytest.mark.parametrize("value", ["", None]) + async def test_default_cache_directory_without_override(self, value: str | None) -> None: + from pyrit.common.path import DB_DATA_PATH + + environment = {} if value is None else {"HF_DATASETS_CACHE": value} + loader = ConcreteRemoteLoader() + with ( + patch.dict("os.environ", environment, clear=True), + patch("pyrit.datasets.seed_datasets.remote.remote_dataset_loader.load_dataset") as mock_load, + ): + await loader._fetch_from_huggingface_async(dataset_name="owner/ds", split="train", cache=True) + + assert mock_load.call_args.kwargs["cache_dir"] == str(DB_DATA_PATH / "huggingface") diff --git a/tests/unit/scenario/core/test_scenario_dataset_overrides.py b/tests/unit/scenario/core/test_scenario_dataset_overrides.py index 1bf0bb0646..7a69b69a89 100644 --- a/tests/unit/scenario/core/test_scenario_dataset_overrides.py +++ b/tests/unit/scenario/core/test_scenario_dataset_overrides.py @@ -36,7 +36,7 @@ from pyrit.scenario.scenarios.garak.api_key import ApiKey from pyrit.scenario.scenarios.garak.encoding import Encoding from pyrit.scenario.scenarios.garak.figstep import FigStep -from pyrit.scenario.scenarios.garak.package_hallucination import PackageHallucination +from pyrit.scenario.scenarios.garak.package_hallucination import PackageHallucination, PackageHallucinationTechnique from pyrit.scenario.scenarios.garak.prompt_inject import PromptInject, PromptInjectTechnique from pyrit.scenario.scenarios.garak.system_prompt_extraction import ( DATASET_EXTRACTION_TEMPLATES, @@ -168,7 +168,7 @@ async def test_resume_replays_saved_selection_without_preparation() -> None: (ApiKey, {"max_per_dataset": 1}, "per-dataset"), (ApiKey, {"sources": [{"name": "harmbench"}]}, "fixed"), (Encoding, {"sources": [{"name": "garak_slur_terms_en"}]}, "ambiguous"), - (PackageHallucination, {"max_total": 1}, "total dataset"), + (PackageHallucination, {"max_per_dataset": 1}, "per-dataset"), ], ) async def test_constrained_scenarios_reject_unsupported_overrides_before_preparation( @@ -264,6 +264,88 @@ async def test_synthesized_scenarios_build_only_from_matching_seeds(kind: str) - assert scenario._technique_templates["ignore_print"].value.startswith("selected:") +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("kind", ["system_prompt", "web", "packages"]) +async def test_generated_objective_limits_preserve_ingredients_and_resume(kind: str) -> None: + scorer = SubStringScorer(substring="yes") + if kind == "system_prompt": + scenario = SystemPromptExtraction(objective_scorer=scorer) + techniques = [SystemPromptExtractionTechnique.DirectRequests] + elif kind == "web": + scenario = WebInjection(objective_scorer=scorer) + techniques = [WebInjectionTechnique.TaskXSS, WebInjectionTechnique.MarkdownURIImageExfilExtended] + else: + scenario = PackageHallucination(objective_scorer=scorer) + techniques = [PackageHallucinationTechnique.Rust, PackageHallucinationTechnique.Python] + args = { + "objective_target": MockPromptTarget(), + "scenario_techniques": techniques, + "include_baseline": False, + "dataset_overrides": ScenarioDatasetRequest(max_total=3, fetch=DatasetFetchPolicy.NEVER), + } + scenario.set_params_from_args(args=args) + memory = CentralMemory.get_memory_instance() + with patch.object(memory, "get_seeds_async", side_effect=AssertionError("Estimate read seeds")): + estimate = await scenario.get_run_size_estimate_async() + assert estimate.estimated_attack_count == 3 + assert estimate.dataset_limit.value == 3 + names = [source.name for source in estimate.effective_datasets.sources] + seeds = [ + SeedPrompt( + value=f"{name}-{index} " if name.endswith("_stubs") else f"{name}-{index}", + dataset_name=name, + metadata={"technique": "direct_requests"} if name == DATASET_EXTRACTION_TEMPLATES else {}, + ) + for name in names + for index in range(4) + ] + await memory.add_seeds_to_memory_async(seeds=seeds, added_by="test") + await scenario.initialize_async() + assert sum(len(attack.seed_groups) for attack in scenario._atomic_attacks) == 3 + assert all(attack.seed_groups for attack in scenario._atomic_attacks) + for name in names: + assert len(await memory.get_seeds_async(dataset_name=name)) == 4 + plan = scenario._build_run_plan() + resumed = type(scenario)(objective_scorer=scorer, scenario_result_id=scenario._scenario_result_id) + resumed.set_params_from_args(args=args) + with patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("Resume prepared data")): + await resumed.initialize_async() + assert resumed._build_run_plan() == plan + + uncapped = type(scenario)(objective_scorer=scorer) + uncapped.set_params_from_args( + args={**args, "dataset_overrides": ScenarioDatasetRequest(max_total="all", fetch=DatasetFetchPolicy.NEVER)} + ) + await uncapped.initialize_async() + assert sum(len(attack.seed_groups) for attack in uncapped._atomic_attacks) > 3 + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("scenario_class", [WebInjection, PackageHallucination]) +@pytest.mark.parametrize("operation", ["estimate", "load", "initialize"]) +async def test_generated_limits_reject_missing_technique_coverage_before_loading( + *, scenario_class: type[Scenario], operation: str +) -> None: + scenario = scenario_class() + scenario.set_params_from_args( + args={ + "objective_target": MockPromptTarget(), + "scenario_techniques": [scenario._technique_class.ALL], + "dataset_overrides": ScenarioDatasetRequest(max_total=1), + } + ) + with ( + patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("Prepared datasets")), + pytest.raises(ValueError, match="Select at least .* objectives"), + ): + if operation == "estimate": + await scenario.get_run_size_estimate_async() + elif operation == "load": + await scenario.load_datasets_async() + else: + await scenario.initialize_async() + + def test_scenarios_do_not_bypass_configured_seed_reads() -> None: root = Path(__file__).parents[4] / "pyrit" / "scenario" / "scenarios" direct_reads = [] diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py index 653dd36893..bd547e867c 100644 --- a/tests/unit/scenario/test_default_run_size_estimates.py +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -535,7 +535,7 @@ async def test_web_injection_estimate_does_not_load_or_synthesize_populations_as @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize("baseline", [False, True]) -@pytest.mark.parametrize("dataset_limit", [1, None]) +@pytest.mark.parametrize("dataset_limit", [4, None]) async def test_web_injection_capped_techniques_use_generation_limits_async( *, baseline: bool, dataset_limit: int | None ) -> None: @@ -563,6 +563,8 @@ async def test_web_injection_capped_techniques_use_generation_limits_async( ) estimate = await scenario.get_run_size_estimate_async() budget = len(scenario.STRING_ASSEMBLY_SEEDS) + 3 * 7 + if dataset_limit is not None: + budget = min(budget, dataset_limit) assert estimate.status is ScenarioRunSizeEstimateStatus.Approximate assert estimate.estimated_attack_count == budget * (2 if baseline else 1) assert estimate.dataset_size == scenario_dataset_size_from_limit(budget) @@ -572,7 +574,7 @@ async def test_web_injection_capped_techniques_use_generation_limits_async( @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize("prompt_cap", [None, 3]) @pytest.mark.parametrize("technique", [PackageHallucinationTechnique.DEFAULT, PackageHallucinationTechnique.ALL]) -@pytest.mark.parametrize("dataset_limit", [1, None]) +@pytest.mark.parametrize("dataset_limit", [7, None]) async def test_package_hallucination_uses_per_language_generation_cap_async( *, prompt_cap: int | None, technique: PackageHallucinationTechnique, dataset_limit: int | None ) -> None: @@ -588,9 +590,12 @@ async def test_package_hallucination_uses_per_language_generation_cap_async( estimate = await scenario.get_run_size_estimate_async() cap = 12 if prompt_cap is None else prompt_cap language_count = len(PackageHallucinationTechnique.expand({technique})) + expected = cap * language_count + if dataset_limit is not None: + expected = min(expected, dataset_limit) assert estimate.status is ScenarioRunSizeEstimateStatus.Approximate - assert estimate.estimated_attack_count == cap * language_count - assert estimate.dataset_size == scenario_dataset_size_from_limit(cap * language_count) + assert estimate.estimated_attack_count == expected + assert estimate.dataset_size == scenario_dataset_size_from_limit(expected) assert estimate.effective_parameters == {"max_prompts_per_language": cap} @@ -607,13 +612,13 @@ async def test_system_prompt_extraction_uses_one_shared_generation_cap_async( scenario.set_params_from_args( args={ "scenario_techniques": [technique], - "dataset_config": DatasetAttackConfiguration( - dataset_names=["missing"], max_per_dataset="all", max_total=dataset_limit - ), + "dataset_overrides": {"max_total": dataset_limit}, } ) estimate = await scenario.get_run_size_estimate_async() - expected = 256 if prompt_cap in (None, "", "default") else prompt_cap + expected = ( + dataset_limit if dataset_limit is not None else (256 if prompt_cap in (None, "", "default") else prompt_cap) + ) assert estimate.estimated_attack_count == (None if expected == "all" else expected) assert estimate.dataset_size == scenario_dataset_size_from_limit(expected) if expected == "all": From 23854bc4a2902d524cfb786bc4d9ad579ed1c1f9 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 13 Oct 2026 12:49:17 -0700 Subject: [PATCH 3/4] Isolate cached dataset services between backend tests Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16cbdc1b-2bf6-4286-be66-c9f30c7ed72b --- tests/unit/backend/conftest.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/backend/conftest.py b/tests/unit/backend/conftest.py index c1c8dd27a3..7c44cdc546 100644 --- a/tests/unit/backend/conftest.py +++ b/tests/unit/backend/conftest.py @@ -13,6 +13,7 @@ from pyrit.analytics._execution import AnalyticsExecution from pyrit.backend.services.analytics_service import close_analytics_service_async, get_analytics_service from pyrit.backend.services.attack_service import get_attack_service +from pyrit.backend.services.dataset_service import get_dataset_service from pyrit.backend.services.manual_send_scheduler import get_manual_send_scheduler from pyrit.backend.services.message_send_service import get_message_send_service from pyrit.memory import SQLiteMemory @@ -34,8 +35,8 @@ async def analytics_service(sqlite_instance: SQLiteMemory) -> AsyncGenerator[Att @pytest.fixture(autouse=True) def isolated_manual_message_services() -> Iterator[None]: """Keep cached service owners bound to this test's memory and event loop.""" - for factory in (get_attack_service, get_message_send_service, get_manual_send_scheduler): + for factory in (get_attack_service, get_message_send_service, get_manual_send_scheduler, get_dataset_service): factory.cache_clear() yield - for factory in (get_attack_service, get_message_send_service, get_manual_send_scheduler): + for factory in (get_attack_service, get_message_send_service, get_manual_send_scheduler, get_dataset_service): factory.cache_clear() From e80c1ee0af918fe5fc7d2b62cd3263e4e5554c67 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 13 Oct 2026 13:21:05 -0700 Subject: [PATCH 4/4] Update browser tests for Technique.execute navigation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16cbdc1b-2bf6-4286-be66-c9f30c7ed72b --- frontend/e2e/accessibility.spec.ts | 12 ++++++------ frontend/e2e/api.spec.ts | 2 +- frontend/e2e/chat-recovery.spec.ts | 2 +- frontend/e2e/chat.spec.ts | 16 ++++++++-------- frontend/e2e/compatibility.spec.ts | 10 +++++----- frontend/e2e/config.spec.ts | 6 +++--- frontend/e2e/converters.spec.ts | 4 ++-- frontend/e2e/errors.spec.ts | 6 +++--- frontend/e2e/flows.spec.ts | 2 +- frontend/e2e/routing.spec.ts | 2 +- frontend/e2e/scenario-history.spec.ts | 2 +- 11 files changed, 32 insertions(+), 32 deletions(-) diff --git a/frontend/e2e/accessibility.spec.ts b/frontend/e2e/accessibility.spec.ts index ee07c64883..a279a2c6e3 100644 --- a/frontend/e2e/accessibility.spec.ts +++ b/frontend/e2e/accessibility.spec.ts @@ -129,7 +129,7 @@ test.describe("Accessibility", () => { const objectiveDefault = page.getByRole("combobox", { name: "Default objective target", exact: true }); await expect(objectiveDefault).toBeVisible({ timeout: 5000 }); await objectiveDefault.selectOption({ index: 1 }); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); // Input should be accessible const input = page.getByRole("textbox"); @@ -146,7 +146,7 @@ test.describe("Accessibility", () => { test("should have accessible sidebar navigation", async ({ page }) => { // Chat button - const chatBtn = page.getByTitle("Chat"); + const chatBtn = page.getByTitle("Technique.execute"); await expect(chatBtn).toBeVisible(); // Registry button @@ -272,7 +272,7 @@ test.describe("Accessibility", () => { const objectiveDefault = page.getByRole("combobox", { name: "Default objective target", exact: true }); await expect(objectiveDefault).toBeVisible({ timeout: 5000 }); await objectiveDefault.selectOption({ index: 1 }); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); const input = page.getByRole("textbox"); await expect(input).toBeEnabled({ timeout: 5000 }); @@ -337,7 +337,7 @@ test.describe("Accessibility", () => { const views = [ { button: "History", heading: "History" }, { button: "Registry", heading: "Target Registry" }, - { button: "Chat", heading: "Chat" }, + { button: "Technique.execute", heading: "Chat" }, ]; for (const view of views) { @@ -407,8 +407,8 @@ test.describe("Visual Consistency", () => { await page.goto("/"); // Wait for initial render then navigate to chat to measure the chat ribbon - await expect(page.getByTitle("Chat")).toBeVisible(); - await page.getByTitle("Chat").click(); + await expect(page.getByTitle("Technique.execute")).toBeVisible(); + await page.getByTitle("Technique.execute").click(); const anchor = page.getByTestId("new-attack-btn"); await expect(anchor).toBeVisible(); diff --git a/frontend/e2e/api.spec.ts b/frontend/e2e/api.spec.ts index 3739b3fb3e..c15a7dccfe 100644 --- a/frontend/e2e/api.spec.ts +++ b/frontend/e2e/api.spec.ts @@ -141,6 +141,6 @@ test.describe("Error Handling", () => { await page.goto("/"); // UI should be responsive even while APIs are delayed - await expect(page.getByTitle("Chat")).toBeVisible({ timeout: 10000 }); + await expect(page.getByTitle("Technique.execute")).toBeVisible({ timeout: 10000 }); }); }); diff --git a/frontend/e2e/chat-recovery.spec.ts b/frontend/e2e/chat-recovery.spec.ts index b31f3e96c7..ddc5985f4c 100644 --- a/frontend/e2e/chat-recovery.spec.ts +++ b/frontend/e2e/chat-recovery.spec.ts @@ -240,7 +240,7 @@ test.describe("Chat processing recovery @seeded", () => { await page.getByTitle("Registry", { exact: true }).click(); await page.getByRole("combobox", { name: "Default objective target", exact: true }) .selectOption(localTarget.registryName); - await page.getByTitle("Chat", { exact: true }).click(); + await page.getByTitle("Technique.execute", { exact: true }).click(); }); test("keeps the drawer closed when a delayed send refresh discovers another conversation", async ({ diff --git a/frontend/e2e/chat.spec.ts b/frontend/e2e/chat.spec.ts index 4ccfb4241d..d75df13f55 100644 --- a/frontend/e2e/chat.spec.ts +++ b/frontend/e2e/chat.spec.ts @@ -152,7 +152,7 @@ async function activateMockTarget(page: Page) { await objectiveDefault.selectOption({ index: 1 }); // Return to Chat view - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible({ timeout: 5000 }); } @@ -170,18 +170,18 @@ test.describe("Application Smoke Tests", () => { }); test("should display chat ribbon", async ({ page }) => { - await expect(page.getByTitle("Chat")).toBeVisible({ timeout: 10000 }); - await page.getByTitle("Chat").click(); + await expect(page.getByTitle("Technique.execute")).toBeVisible({ timeout: 10000 }); + await page.getByTitle("Technique.execute").click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible({ timeout: 10000 }); }); test("should have New Attack button", async ({ page }) => { - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByRole("button", { name: /new attack/i })).toBeVisible(); }); test("should offer the target picker without a bottom warning", async ({ page }) => { - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByRole("combobox", { name: "Chat target" })).toBeVisible(); await expect(page.getByTestId("no-target-banner")).toHaveCount(0); }); @@ -190,7 +190,7 @@ test.describe("Application Smoke Tests", () => { test.describe("Theme Toggle", () => { test("should toggle dark/light theme", async ({ page }) => { await page.goto("/"); - await expect(page.getByTitle("Chat")).toBeVisible({ timeout: 10000 }); + await expect(page.getByTitle("Technique.execute")).toBeVisible({ timeout: 10000 }); // The app defaults to system mode, so the toggle button title should say "Theme: System" const themeBtn = page.getByTitle("Theme: System"); @@ -374,7 +374,7 @@ test.describe("Multiple Messages", () => { test.describe("Chat without target", () => { test("should disable input when no target is selected", async ({ page }) => { await page.goto("/"); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByRole("textbox")).toBeDisabled(); await expect(page.getByRole("button", { name: "Send message" })).toBeDisabled(); @@ -912,7 +912,7 @@ test.describe("Target type scenarios", () => { .selectOption("dall-e-image-gen"); // Navigate to chat - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); // Scope queries to the badge so we don't also match the (hidden) // copy of the target text that Fluent's Tooltip renders into the DOM. const badge = page.getByTestId("target-badge"); diff --git a/frontend/e2e/compatibility.spec.ts b/frontend/e2e/compatibility.spec.ts index f4983b1640..b49f8c9095 100644 --- a/frontend/e2e/compatibility.spec.ts +++ b/frontend/e2e/compatibility.spec.ts @@ -43,10 +43,10 @@ test("waits for a matching packaged stamp before sending business requests", asy }); await page.goto("/"); await expect(page.getByRole("status")).toContainText("Checking frontend and backend compatibility"); - await expect(page.getByTitle("Chat", { exact: true })).not.toBeVisible(); + await expect(page.getByTitle("Technique.execute", { exact: true })).not.toBeVisible(); expect(requests).toHaveLength(0); releaseVersion(); - await expect(page.getByTitle("Chat", { exact: true })).toBeVisible(); + await expect(page.getByTitle("Technique.execute", { exact: true })).toBeVisible(); await expect.poll(() => requests.length).toBeGreaterThan(0); }); @@ -62,7 +62,7 @@ for (const [name, payload] of [ }); await page.goto("/"); await expect(page.getByRole("heading", { name: "PyRIT compatibility blocked" })).toBeVisible(); - await expect(page.getByTitle("Chat", { exact: true })).not.toBeVisible(); + await expect(page.getByTitle("Technique.execute", { exact: true })).not.toBeVisible(); expect(requests).toHaveLength(0); }); } @@ -81,7 +81,7 @@ for (const [status, type] of [ await route.fulfill({ status, json: { type, expected: otherCompatibilityId(), actual: getCompatibilityId() } }); }); await page.goto("/"); - await expect(page.getByTitle("Chat", { exact: true })).toBeVisible(); + await expect(page.getByTitle("Technique.execute", { exact: true })).toBeVisible(); await page.evaluate(async () => { const modulePath = "/src/services/api.ts"; const { apiClient } = await import(modulePath); @@ -92,7 +92,7 @@ for (const [status, type] of [ await expect(page.getByRole("heading", { name: "PyRIT compatibility blocked" })).toBeVisible(); await expect(page.getByRole("alertdialog")).toContainText(getCompatibilityId()); await expect(page.getByRole("alertdialog")).toContainText(otherCompatibilityId()); - await expect(page.getByTitle("Chat", { exact: true })).not.toBeVisible(); + await expect(page.getByTitle("Technique.execute", { exact: true })).not.toBeVisible(); expect(mutations).toBe(1); }); } diff --git a/frontend/e2e/config.spec.ts b/frontend/e2e/config.spec.ts index b0d048e8cd..a77f466179 100644 --- a/frontend/e2e/config.spec.ts +++ b/frontend/e2e/config.spec.ts @@ -699,7 +699,7 @@ test.describe("Target Config ↔ Chat Navigation", () => { await page.getByRole("combobox", { name: "Default objective target", exact: true }).selectOption({ index: 1 }); // Navigate back to chat - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible(); // Chat should show the selected target type. Scope to the badge to @@ -717,7 +717,7 @@ test.describe("Target Config ↔ Chat Navigation", () => { // Start in chat with the composer disabled until a target is selected. await page.goto("/"); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByRole("textbox")).toBeDisabled(); // Go to targets, set a target @@ -726,7 +726,7 @@ test.describe("Target Config ↔ Chat Navigation", () => { await page.getByRole("combobox", { name: "Default objective target", exact: true }).selectOption({ index: 1 }); // Return to chat — send should be enabled when there's text - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); const input = page.getByRole("textbox"); await input.fill("Hello"); await expect(page.getByRole("button", { name: /send/i })).toBeEnabled(); diff --git a/frontend/e2e/converters.spec.ts b/frontend/e2e/converters.spec.ts index 185e66e8b9..22f2b32933 100644 --- a/frontend/e2e/converters.spec.ts +++ b/frontend/e2e/converters.spec.ts @@ -469,7 +469,7 @@ async function activateMockTarget(page: Page) { await expect(objectiveDefault).toBeVisible({ timeout: 5000 }); await objectiveDefault.selectOption({ index: 1 }); - await page.getByTitle("Chat", { exact: true }).click(); + await page.getByTitle("Technique.execute", { exact: true }).click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible({ timeout: 5000 }); } @@ -568,7 +568,7 @@ test.describe("Shared per-piece converter pipelines @seeded", () => { await page.getByTitle("Registry", { exact: true }).click(); await page.getByRole("combobox", { name: "Default objective target", exact: true }) .selectOption(targetRegistryName); - await page.getByTitle("Chat", { exact: true }).click(); + await page.getByTitle("Technique.execute", { exact: true }).click(); await expect(page.getByTestId("chat-input")).toBeEnabled(); }); diff --git a/frontend/e2e/errors.spec.ts b/frontend/e2e/errors.spec.ts index af764de659..444c6cef12 100644 --- a/frontend/e2e/errors.spec.ts +++ b/frontend/e2e/errors.spec.ts @@ -241,7 +241,7 @@ async function activateMockTarget(page: Page) { const objectiveDefault = page.getByRole("combobox", { name: "Default objective target", exact: true }); await expect(objectiveDefault).toBeVisible({ timeout: 5000 }); await objectiveDefault.selectOption({ index: 1 }); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible({ timeout: 5000 }); } @@ -504,7 +504,7 @@ test.describe("Error: connection banner on health failure", () => { }) => { // Let the page load normally first await page.goto("/"); - await expect(page.getByTitle("Chat")).toBeVisible({ + await expect(page.getByTitle("Technique.execute")).toBeVisible({ timeout: 10000, }); @@ -552,7 +552,7 @@ test.describe("Error: connection banner recovery", () => { }); await page.goto("/"); - await expect(page.getByTitle("Chat")).toBeVisible({ + await expect(page.getByTitle("Technique.execute")).toBeVisible({ timeout: 10000, }); diff --git a/frontend/e2e/flows.spec.ts b/frontend/e2e/flows.spec.ts index 104f539538..ec8a478e13 100644 --- a/frontend/e2e/flows.spec.ts +++ b/frontend/e2e/flows.spec.ts @@ -210,7 +210,7 @@ async function activateTarget( await expect(row).toBeVisible({ timeout: 10_000 }); await page.getByRole("combobox", { name: "Default objective target", exact: true }) .selectOption(targetRegistryName); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page.getByTestId("new-attack-btn")).toBeVisible({ timeout: 5_000 }); } diff --git a/frontend/e2e/routing.spec.ts b/frontend/e2e/routing.spec.ts index c308fde70f..6573c8e97d 100644 --- a/frontend/e2e/routing.spec.ts +++ b/frontend/e2e/routing.spec.ts @@ -300,7 +300,7 @@ test.describe("URL-driven routing", () => { await expect(page).toHaveURL(/\/history\/attacks$/); await expect(page.getByTestId("attacks-table")).toBeVisible(); - await page.getByTitle("Chat").click(); + await page.getByTitle("Technique.execute").click(); await expect(page).toHaveURL(/\/chat$/); await expect(page.getByRole("switch", { name: "Markdown" })).toBeChecked(); diff --git a/frontend/e2e/scenario-history.spec.ts b/frontend/e2e/scenario-history.spec.ts index 22d53393db..315b8a0738 100644 --- a/frontend/e2e/scenario-history.spec.ts +++ b/frontend/e2e/scenario-history.spec.ts @@ -543,7 +543,7 @@ test.describe("Scenario catalog, history, and live run routing", () => { expect(await primaryButtons.evaluateAll((buttons) => buttons.map((button) => button.getAttribute("aria-label")))).toEqual([ "Home", - "Chat", + "Technique.execute", "History", "Scanner", "Datasets",